85class ContinualLearner:
 86    """Mixin for modules whose weights are adapted by meta-learned step sizes.
 87
 88    In Sutton's OaK architecture, every learned weight has a dedicated
 89    step-size parameter adapted via online cross-validation (e.g. IDBD,
 90    Sutton 1992; Adam-IDBD, Degris et al. 2024).
 91
 92    The agent loop calls `update_meta()` on all four modules after each
 93    learning step, passing the same error-signals dict.  Each module
 94    internally decides which signals are relevant and routes them to its
 95    per-weight step-size adaptation.
 96
 97    The default implementation is a no-op so that modules without
 98    meta-learning still work unchanged.
 99    """
100
101    def update_meta(self, error_signals: Mapping[str, float]) -> None:
102        """Adapt internal per-weight step sizes given error signals.
103
104        Parameters
105        ----------
106        error_signals:
107            Named scalar error signals from the current learning step,
108            e.g. `{"main_td_error": 0.05, "reward": 1.0}`.
109            Implementations pick the signals they need and ignore the rest.
110        """

118class Perception(
119    ContinualLearner, ABC, Generic[ObservationT, ActionT, SubjectiveStateT]
120):
121    """Sutton's Perception: observations → subjective state + feature management.
122
123    Turns raw observations into the agent's **subjective state**, the
124    internal representation that every other module sees.  Also discovers,
125    ranks, and manages **features** (learned representational structures
126    that grow over the agent's lifetime) and generates **subtasks** from
127    the most useful ones.
128
129    The finer-grained layer splits this into `StateBuilder`,
130    `FeatureBank`, `FeatureConstructor`, `FeatureRanker`, and
131    `SubtaskGenerator` (see `oak_architecture.fine_grained.components`).
132    """
133
134    @abstractmethod
135    def reset(self) -> None:
136        """Reset all perception state for a new episode."""
137        raise NotImplementedError
138
139    @abstractmethod
140    def update(
141        self,
142        observation: ObservationT,
143        reward: float,
144        last_action: ActionT | None,
145    ) -> SubjectiveStateT:
146        """Process a new observation and return the updated subjective state."""
147        raise NotImplementedError
148
149    @abstractmethod
150    def current_subjective_state(self) -> SubjectiveStateT:
151        """Return the most recently computed subjective state."""
152        raise NotImplementedError
153
154    @abstractmethod
155    def discover_and_rank_features(
156        self,
157        subjective_state: SubjectiveStateT,
158        utility_scores: Sequence[UtilityRecord],
159        feature_budget: int,
160    ) -> Sequence[FeatureId]:
161        """Propose new features, integrate them, and return the top-ranked IDs.
162
163        A typical implementation:
164
165        1. Proposes candidate features from the current subjective state.
166        2. Adds accepted candidates to its internal feature store.
167        3. Ranks all features using the provided utility scores.
168        4. Returns the top feature IDs (up to *feature_budget*).
169        """
170        raise NotImplementedError
171
172    @abstractmethod
173    def generate_subtasks(
174        self,
175        ranked_feature_ids: Sequence[FeatureId],
176    ) -> Sequence[SubtaskSpec]:
177        """Turn ranked feature IDs into subtask specifications."""
178        raise NotImplementedError
179
180    @abstractmethod
181    def list_features(self) -> Sequence[FeatureSpec]:
182        """Return all currently tracked features."""
183        raise NotImplementedError
184
185    @abstractmethod
186    def remove_features(self, feature_ids: Sequence[FeatureId]) -> None:
187        """Remove features by ID (called during curation)."""
188        raise NotImplementedError

245class TransitionModel(ContinualLearner, ABC, Generic[SubjectiveStateT, ActionT, InfoT]):
246    """Sutton's Transition Model: world dynamics + option models + planning.
247
248    Learns from observed transitions, maintains **option models** that
249    predict the effect of temporal abstractions, and runs bounded
250    **planning** using the world model and the value function to produce
251    improvement signals for the reactive policy.
252
253    The finer-grained layer splits this into `WorldModel`,
254    `OptionModelLearner`, individual `OptionModel` objects, and a
255    `Planner` (see `oak_architecture.fine_grained.components`).
256    """
257
258    @abstractmethod
259    def update(
260        self,
261        transition: Transition[ActionT, SubjectiveStateT, InfoT],
262    ) -> None:
263        """Learn from an observed transition.
264
265        This should update both the world model and any option-model learners.
266        """
267        raise NotImplementedError
268
269    @abstractmethod
270    def integrate_option_models(self) -> None:
271        """Export learned option models and integrate them into the world model.
272
273        Called after option learning so that planning reasons over fresh models.
274        """
275        raise NotImplementedError
276
277    @abstractmethod
278    def plan(
279        self,
280        subjective_state: SubjectiveStateT,
281        value_function: ValueFunction[SubjectiveStateT, ActionT, InfoT],
282        budget: int,
283    ) -> PlanningUpdate[ActionT]:
284        """Run bounded planning and return improvement signals.
285
286        The planner uses the internal world model together with the supplied
287        *value_function* (for state evaluation) to produce value targets,
288        policy targets, or search statistics.
289        """
290        raise NotImplementedError
291
292    @abstractmethod
293    def remove_option_models(self, option_ids: Sequence[OptionId]) -> None:
294        """Remove option models by ID (called during curation)."""
295        raise NotImplementedError

191class ValueFunction(ContinualLearner, ABC, Generic[SubjectiveStateT, ActionT, InfoT]):
192    """Sutton's Value Function: value learning + utility assessment + curation.
193
194    Learns **predictive value signals** (TD errors, GVF predictions) from
195    observed transitions and predicts cumulative signals for any given
196    subjective state.  Also assesses the **utility** of the agent's learned
197    structures (features, options, models) and produces concrete keep/drop
198    **curation** decisions.
199
200    The finer-grained layer splits this into `ValueEstimator`,
201    `GeneralValueFunctionLearner`, `UtilityAssessor`, `Curator`, and
202    `MetaStepSizeLearner` (see `oak_architecture.fine_grained.components`).
203    """
204
205    @abstractmethod
206    def update(
207        self,
208        transition: Transition[ActionT, SubjectiveStateT, InfoT],
209    ) -> Mapping[GeneralValueFunctionId, float]:
210        """Learn from a transition and return TD-error signals."""
211        raise NotImplementedError
212
213    @abstractmethod
214    def predict(
215        self,
216        subjective_state: SubjectiveStateT,
217    ) -> Mapping[GeneralValueFunctionId, float]:
218        """Predict values for the given subjective state."""
219        raise NotImplementedError
220
221    @abstractmethod
222    def observe_usage(self, usage_records: Sequence[UsageRecord]) -> None:
223        """Record usage evidence for utility assessment."""
224        raise NotImplementedError
225
226    @abstractmethod
227    def utility_scores(self) -> Sequence[UtilityRecord]:
228        """Return current utility estimates for all tracked structures."""
229        raise NotImplementedError
230
231    @abstractmethod
232    def curate(self) -> CurationDecision:
233        """Decide which learned structures to drop."""
234        raise NotImplementedError
235
236    @abstractmethod
237    def remove(
238        self,
239        general_value_function_ids: Sequence[GeneralValueFunctionId],
240    ) -> None:
241        """Remove value functions by ID (called during curation)."""
242        raise NotImplementedError

298class ReactivePolicy(ContinualLearner, ABC, Generic[SubjectiveStateT, ActionT, InfoT]):
299    """Sutton's Reactive Policy: action selection + option management.
300
301    Selects **actions**, primitive or temporal abstractions (options),
302    based on the current subjective state.  Manages the **option library**
303    and **option learning** pipeline, and integrates **planning updates**
304    into decision-making.
305
306    The finer-grained layer splits this into `ActionSelector`,
307    `OptionLibrary`, and `OptionLearner`
308    (see `oak_architecture.fine_grained.components`).
309    """
310
311    @abstractmethod
312    def update(
313        self,
314        transition: Transition[ActionT, SubjectiveStateT, InfoT],
315        td_errors: Mapping[GeneralValueFunctionId, float],
316    ) -> None:
317        """Update the policy and option learners from an observed transition."""
318        raise NotImplementedError
319
320    @abstractmethod
321    def apply_planning_update(self, update: PlanningUpdate[ActionT]) -> None:
322        """Integrate planning improvement signals into the policy."""
323        raise NotImplementedError
324
325    @abstractmethod
326    def ingest_subtasks(self, subtasks: Sequence[SubtaskSpec]) -> None:
327        """Feed newly created subtasks into the option learner."""
328        raise NotImplementedError
329
330    @abstractmethod
331    def integrate_options(self) -> None:
332        """Export learned options into the option library."""
333        raise NotImplementedError
334
335    @abstractmethod
336    def select_action(
337        self,
338        subjective_state: SubjectiveStateT,
339        option_stop_threshold: float,
340    ) -> tuple[ActionT, OptionId | None]:
341        """Choose a primitive action, possibly by continuing an active option.
342
343        Returns a `(primitive_action, active_option_id)` pair.  When no
344        option is active, *active_option_id* is `None`.
345        """
346        raise NotImplementedError
347
348    @abstractmethod
349    def clear_active_option(self) -> None:
350        """Clear the currently executing option (e.g. at episode boundaries)."""
351        raise NotImplementedError
352
353    @abstractmethod
354    def remove_options(self, option_ids: Sequence[OptionId]) -> None:
355        """Remove options by ID (called during curation)."""
356        raise NotImplementedError
357
358    @abstractmethod
359    def remove_subtasks(self, subtask_ids: Sequence[SubtaskId]) -> None:
360        """Remove subtasks by ID (called during curation)."""
361        raise NotImplementedError

 51@dataclass
 52class OaKAgent(Generic[ObservationT, ActionT, SubjectiveStateT, InfoT]):
 53    """Coordinates one full OaK step across the four modules.
 54
 55    The agent is a wiring object: you provide concrete implementations of
 56    `Perception`, `TransitionModel`, `ValueFunction`, and
 57    `ReactivePolicy`, and `OaKAgent` ensures they are called in a
 58    consistent order.
 59    """
 60
 61    perception: Perception[ObservationT, ActionT, SubjectiveStateT]
 62    transition_model: TransitionModel[SubjectiveStateT, ActionT, InfoT]
 63    value_function: ValueFunction[SubjectiveStateT, ActionT, InfoT]
 64    reactive_policy: ReactivePolicy[SubjectiveStateT, ActionT, InfoT]
 65
 66    planning_budget: int = 4
 67    feature_budget: int = 4
 68    option_stop_threshold: float = 0.5
 69
 70    last_action: ActionT | None = None
 71    last_subjective_state: SubjectiveStateT | None = None
 72
 73    def __init__(
 74        self,
 75        perception: Perception[ObservationT, ActionT, SubjectiveStateT],
 76        transition_model: TransitionModel[SubjectiveStateT, ActionT, InfoT],
 77        value_function: ValueFunction[SubjectiveStateT, ActionT, InfoT],
 78        reactive_policy: ReactivePolicy[SubjectiveStateT, ActionT, InfoT],
 79        planning_budget: int = 4,
 80        feature_budget: int = 4,
 81        option_stop_threshold: float = 0.5,
 82    ):
 83        self.perception = perception
 84        self.transition_model = transition_model
 85        self.value_function = value_function
 86        self.reactive_policy = reactive_policy
 87        self.planning_budget = planning_budget
 88        self.feature_budget = feature_budget
 89        self.option_stop_threshold = option_stop_threshold
 90        self.last_action = None
 91        self.last_subjective_state = None
 92
 93    def __post_init__(self):
 94        """Validate that the modules are compatible."""
 95        if self.planning_budget < 1:
 96            raise ValueError("Planning budget must be at least 1.")
 97        if self.feature_budget < 1:
 98            raise ValueError("Feature budget must be at least 1.")
 99        if self.option_stop_threshold < 0 or self.option_stop_threshold > 1:
100            raise ValueError("Option stop threshold must be in [0, 1].")
101
102    def reset(self) -> None:
103        """Clear transient execution memory."""
104        self.perception.reset()
105        self.reactive_policy.clear_active_option()
106        self.last_action = None
107        self.last_subjective_state = None
108
109    def step(
110        self, time_step: TimeStep[ObservationT, InfoT]
111    ) -> AgentStepResult[ActionT, SubjectiveStateT]:
112        """Run one temporally uniform agent step.
113
114        The step follows six phases: perceive, learn, grow, plan, act, maintain.
115        """
116
117        # ================== 1. Perceive ================= #
118        subjective_state = self.perception.update(
119            observation=time_step.observation,
120            reward=time_step.reward,
121            last_action=self.last_action,
122        )
123
124        created_subtasks: Sequence[SubtaskSpec] = ()
125        ranked_feature_ids: Sequence[FeatureId] = ()
126        planning_update: PlanningUpdate[ActionT] | None = None
127        curation_decision: CurationDecision | None = None
128
129        # ================== 2. Learn ================== #
130        if self.last_subjective_state is not None and self.last_action is not None:
131            transition = Transition(
132                subjective_state=self.last_subjective_state,
133                action=self.last_action,
134                reward=time_step.reward,
135                next_subjective_state=subjective_state,
136                terminated=time_step.terminated or time_step.truncated,
137                info=time_step.info,
138            )
139            td_errors = self.value_function.update(transition)
140            self.reactive_policy.update(transition, td_errors)
141            self.transition_model.update(transition)
142
143            # Meta step-size adaptation (Sutton's IDBD / online cross-validation)
144            meta_signals = dict(td_errors)
145            meta_signals["reward"] = transition.reward
146            self.perception.update_meta(meta_signals)
147            self.value_function.update_meta(meta_signals)
148            self.reactive_policy.update_meta(meta_signals)
149            self.transition_model.update_meta(meta_signals)
150
151        # ================== 3. Grow ================== #
152        ranked_feature_ids = self.perception.discover_and_rank_features(
153            subjective_state,
154            self.value_function.utility_scores(),
155            self.feature_budget,
156        )
157        if ranked_feature_ids:
158            created_subtasks = self.perception.generate_subtasks(ranked_feature_ids)
159            if created_subtasks:
160                self.reactive_policy.ingest_subtasks(created_subtasks)
161
162        self.reactive_policy.integrate_options()
163        self.transition_model.integrate_option_models()
164
165        # ================== 4. Plan ================== #
166        planning_update = self.transition_model.plan(
167            subjective_state, self.value_function, self.planning_budget
168        )
169        self.reactive_policy.apply_planning_update(planning_update)
170
171        # ================== 5. Act ================== #
172        action, active_option_id = self.reactive_policy.select_action(
173            subjective_state, self.option_stop_threshold
174        )
175
176        # ================== 6. Maintain ================== #
177        usage_records = self._build_usage_records(ranked_feature_ids, active_option_id)
178        if usage_records:
179            self.value_function.observe_usage(usage_records)
180
181        curation_decision = self.value_function.curate()
182        self._apply_curation(curation_decision)
183
184        # ================== Update Memory ================== #
185        self.last_subjective_state = subjective_state
186        self.last_action = action
187
188        if time_step.terminated or time_step.truncated:
189            self.reactive_policy.clear_active_option()
190
191        return AgentStepResult(
192            action=action,
193            subjective_state=subjective_state,
194            active_option_id=active_option_id,
195            planning_update=planning_update,
196            created_subtasks=created_subtasks,
197            curation_decision=curation_decision,
198        )
199
200    # ── training loop ─────────────────────────────────────────────────
201
202    def train(
203        self,
204        world: World[ObservationT, ActionT, InfoT],
205        *,  # enforce keyword arguments after for clarity
206        num_episodes: int = 500,
207        average_window: int = 100,
208        solved_threshold: float | None = None,
209        episode_logger: Callable[[int, float, float, Self], None] | None = None,
210    ) -> list[float]:
211        """Run the standard OaK episode loop on the given world.
212
213        Parameters
214        ----------
215        world:
216            An environment implementing the `World` protocol.
217        num_episodes:
218            Maximum number of training episodes.
219        average_window:
220            Number of recent episodes to average for performance tracking.
221        solved_threshold:
222            If set, stop early when the average reward over the last `average_window`
223            episodes reaches this value.
224        episode_logger:
225            Optional callback `(episode, episode_reward, avg_reward, agent)`
226            called after each episode. Use this to own all per-episode logging
227            or other training-side effects at the call site.
228
229        Returns
230        -------
231        list[float]
232            Per-episode reward history.
233        """
234        if average_window < 1:
235            raise ValueError("average_window must be at least 1.")
236
237        reward_history: list[float] = []
238
239        for episode in range(num_episodes):
240            time_step = world.reset()
241            self.reset()
242            episode_reward = 0.0
243
244            while True:
245                result = self.step(time_step)
246
247                if time_step.terminated or time_step.truncated:
248                    break
249
250                time_step = world.step(result.action)
251                episode_reward += time_step.reward
252
253            reward_history.append(episode_reward)
254            recent_window = reward_history[-average_window:]
255            avg_reward = sum(recent_window) / len(recent_window)
256
257            solved = (
258                solved_threshold is not None
259                and len(reward_history) >= average_window
260                and avg_reward >= solved_threshold
261            )
262
263            if episode_logger is not None:
264                episode_logger(episode, episode_reward, avg_reward, self)
265
266            if solved:
267                break
268
269        return reward_history
270
271    # ── private helpers ──────────────────────────────────────────────
272
273    def _build_usage_records(
274        self,
275        ranked_feature_ids: Sequence[FeatureId],
276        active_option_id: OptionId | None,
277    ) -> Sequence[UsageRecord]:
278        """Build minimal utility-accounting observations for the current step."""
279        usage_records = [
280            UsageRecord(ComponentKind.FEATURE, feature_id)
281            for feature_id in ranked_feature_ids
282        ]
283        if active_option_id is not None:
284            usage_records.append(UsageRecord(ComponentKind.OPTION, active_option_id))
285        return tuple(usage_records)
286
287    def _apply_curation(self, decision: CurationDecision) -> None:
288        """Dispatch curation decisions to the relevant modules."""
289        if decision.drop_features:
290            self.perception.remove_features(decision.drop_features)
291        if decision.drop_subtasks:
292            self.reactive_policy.remove_subtasks(decision.drop_subtasks)
293        if decision.drop_options:
294            self.reactive_policy.remove_options(decision.drop_options)
295        if decision.drop_option_models:
296            self.transition_model.remove_option_models(decision.drop_option_models)
297        if decision.drop_general_value_functions:
298            self.value_function.remove(decision.drop_general_value_functions)

59@runtime_checkable
60class World(Protocol[ObservationT, ActionT, InfoT]):
61    """Minimal environment protocol.
62
63    A `World` may wrap a simulator, a benchmark environment, or a custom
64    continual data source.  The protocol is intentionally small so the
65    package does not depend on a specific environment library.
66
67    Implement this protocol for any environment you want to use with
68    `OaKAgent.train()`.
69    """
70
71    def reset(self) -> TimeStep[ObservationT, InfoT]: ...
72
73    def step(self, action: ActionT) -> TimeStep[ObservationT, InfoT]: ...
74
75    def close(self) -> None:
76        """Release environment resources.  Default is a no-op."""
77        ...