# Copyright 2025-2026 Thousand Brains Project
# Copyright 2021-2024 Numenta Inc.
#
# Copyright may exist in Contributors' modifications
# and/or contributions to the work.
#
# Use of this source code is governed by the MIT
# license that can be found in the LICENSE file or at
# https://opensource.org/licenses/MIT.
from __future__ import annotations
import abc
from typing import Any, Collection, Dict, Protocol, Sequence, TypedDict
import numpy as np
import numpy.typing as npt
from tbp.monty.cmp import Goal, Message
from tbp.monty.context import RuntimeContext
from tbp.monty.experiment.learning_module import ExperimentLearningModule
from tbp.monty.experiment.monty import ExperimentMonty
from tbp.monty.experiment.sensor_module import ExperimentSensorModule
from tbp.monty.frameworks.actions.actions import Action
from tbp.monty.frameworks.agents import AgentID
from tbp.monty.frameworks.environments.environment import SemanticID
from tbp.monty.frameworks.experiments.mode import ExperimentMode
from tbp.monty.frameworks.models.motor_system_state import (
AgentState,
ProprioceptiveState,
)
from tbp.monty.frameworks.sensors import SensorID
from tbp.monty.memento import Memento, Snapshotable
__all__ = [
"AgentObservations",
"GoalGenerator",
"LMMemory",
"LearningModule",
"Monty",
"ObjectModel",
"Observations",
"RuntimeContext",
"RuntimeLearningModule",
"SensorModule",
"SensorObservation",
]
[docs]class SensorObservation(TypedDict, total=False):
"""Observations from a sensor."""
rgba: npt.NDArray[np.uint8]
depth: npt.NDArray[np.float64] # TODO: Verify specific type
semantic: npt.NDArray[np.int_] # TODO: Verify specific type
semantic_3d: npt.NDArray[np.int_] # TODO: Verify specific type
sensor_frame_data: npt.NDArray[np.int_] # TODO: Verify specific type
cam_to_world: npt.NDArray[np.float64] # TODO: Verify specific type
pixel_loc: npt.NDArray[np.float64] # TODO: Verify specific type
raw: npt.NDArray[np.uint8]
[docs]class AgentObservations(Dict[SensorID, SensorObservation]):
"""Observations from an agent."""
pass
[docs]class Observations(Dict[AgentID, AgentObservations]):
"""Observations from the environment."""
pass
class RuntimeMonty(Protocol):
"""Runtime interface to Monty."""
def step(
self,
ctx: RuntimeContext,
observations: Observations,
proprioceptive_state: ProprioceptiveState,
) -> list[Action]:
"""Take a matching, exploratory, or custom user-defined step.
Step taken depends on the value of self.step_type.
Args:
ctx: The runtime context.
observations: The observations from the environment.
proprioceptive_state: The proprioceptive state from the environment.
Returns:
The actions to take.
"""
...
def motor_only_step(
self,
ctx: RuntimeContext,
observations: Observations,
proprioceptive_state: ProprioceptiveState,
) -> list[Action]:
"""Take a step of the sensors and motor system only.
This skips stepping the learning modules.
Args:
ctx: The runtime context.
observations: The observations from the environment.
proprioceptive_state: The proprioceptive state from the environment.
Returns:
The actions to take.
"""
...
def aggregate_sensory_inputs(
self,
ctx: RuntimeContext,
observations: Observations,
proprioceptive_state: ProprioceptiveState,
) -> None:
"""Receive data from environment, organize on a per sensor module basis.
Args:
ctx: The runtime context.
observations: The observations from the environment.
proprioceptive_state: The proprioceptive state from the environment.
"""
...
def snapshot_ltm(self) -> Memento:
"""Return an opaque snapshot of long-term memory."""
...
def restore_ltm(self, memo: Memento) -> None:
"""Restore long-term memory from an opaque snapshot."""
...
[docs]class Monty(ExperimentMonty, RuntimeMonty, Snapshotable, metaclass=abc.ABCMeta):
def _matching_step(
self,
ctx: RuntimeContext,
observations: Observations,
proprioceptive_state: ProprioceptiveState,
):
"""Step format for matching observations to graph.
Used during training or evaluation.
Args:
ctx: The runtime context.
observations: The observations from the environment.
proprioceptive_state: The proprioceptive state from the environment.
"""
self.aggregate_sensory_inputs(ctx, observations, proprioceptive_state)
self._step_learning_modules(ctx)
self._vote()
self._pass_goals()
self._step_motor_system(ctx, observations, proprioceptive_state)
self._set_step_type_and_check_if_done()
self._post_step()
def _exploratory_step(
self,
ctx: RuntimeContext,
observations: Observations,
proprioceptive_state: ProprioceptiveState,
):
"""Step format for adding data to an existing model.
Used only during training.
Args:
ctx: The runtime context.
observations: The observations from the environment.
proprioceptive_state: The proprioceptive state from the environment.
"""
self.aggregate_sensory_inputs(ctx, observations, proprioceptive_state)
self._step_learning_modules(ctx)
self._pass_goals()
self._step_motor_system(ctx, observations, proprioceptive_state)
self._set_step_type_and_check_if_done()
self._post_step()
[docs] @abc.abstractmethod
def step(
self,
ctx: RuntimeContext,
observations: Observations,
proprioceptive_state: ProprioceptiveState,
) -> list[Action]:
pass
[docs] @abc.abstractmethod
def motor_only_step(
self,
ctx: RuntimeContext,
observations: Observations,
proprioceptive_state: ProprioceptiveState,
) -> list[Action]:
pass
@abc.abstractmethod
def _step_learning_modules(self, ctx: RuntimeContext):
"""Pass data from SMs to LMs, and have each LM take a step.
LM step type depends on self.step_type.
"""
pass
@abc.abstractmethod
def _vote(self):
"""Share information across learning modules.
Use LM.send_out_vote and LM.receive_votes.
"""
pass
@abc.abstractmethod
def _pass_goals(self):
"""Pass goals in the network between learning-modules.
Aggregate any goals for sending to the motor-system.
"""
pass
@abc.abstractmethod
def _step_motor_system(
self,
ctx: RuntimeContext,
observations: Observations,
proprioceptive_state: ProprioceptiveState,
):
"""Step the motor system.
Args:
ctx: The runtime context.
observations: The observations from the environment.
proprioceptive_state: The proprioceptive state from the environment.
"""
pass
@abc.abstractmethod
def _set_step_type_and_check_if_done(self):
"""Check terminal conditions and decide if to change the step type.
Update what self.is_done returns to the experiment.
"""
pass
@abc.abstractmethod
def _post_step(self):
"""Hook for doing things like updating counters."""
pass
###
# Saving, loading, and logging
###
[docs] @abc.abstractmethod
def state_dict(self) -> Memento:
pass
[docs] @abc.abstractmethod
def load_state_dict(self, memento: Memento) -> None:
pass
###
# Methods that interact with the experiment
###
[docs] @abc.abstractmethod
def reset(self) -> None:
pass
[docs] @abc.abstractmethod
def snapshot_ltm(self) -> Memento:
pass
[docs] @abc.abstractmethod
def restore_ltm(self, memo: Memento) -> None:
pass
[docs] @abc.abstractmethod
def fixme_set_ground_truth(
self,
primary_target: dict[str, Any] | None = None,
semantic_id_to_label: dict[SemanticID, str] | None = None,
) -> None:
pass
[docs] @abc.abstractmethod
def update_ltm(self) -> None:
pass
[docs] @abc.abstractmethod
def set_experiment_mode(self, mode: ExperimentMode) -> None:
pass
[docs] @abc.abstractmethod
def is_done(self) -> bool:
pass
[docs]class RuntimeLearningModule(Protocol):
"""Monty runtime interface to a Learning Module."""
[docs] def matching_step(self, ctx: RuntimeContext, percepts: Sequence[Message]) -> None:
"""Matching / inference step called inside of monty._step_learning_modules.
Args:
ctx: The runtime context.
percepts: The percepts intended for this learning module.
"""
...
[docs] def exploratory_step(
self, ctx: RuntimeContext, percepts: Sequence[Message]
) -> None:
"""Model building step called inside of monty._step_learning_modules.
Args:
ctx: The runtime context.
percepts: The percepts intended for this learning module.
"""
...
[docs] def receive_votes(self, votes: Collection[Any]) -> None:
"""Process inbound voting data.
TODO: Use `Message` type for votes rather than ad-hoc data?
Args:
votes: A collection of votes from other learning modules.
"""
...
[docs] def send_out_vote(self) -> Any:
"""This method defines what data are sent to other learning modules.
TODO: Use `Message` type for votes rather than ad-hoc data?
Returns:
This learning module's voting data.
"""
...
[docs] def propose_goals(self) -> Sequence[Goal]:
"""Return the goals proposed by this LM's GSG if they exist.
Returns:
A collection of proposed Goals.
"""
...
[docs] def get_output(self) -> Message | None:
"""Return learning module output (same format as input)."""
...
[docs] def init_from_ltm(self) -> None:
"""Initialize LM state from long-term memory.
For example, getting initial hypotheses.
"""
...
[docs]class LearningModule(
RuntimeLearningModule, Snapshotable, ExperimentLearningModule, metaclass=abc.ABCMeta
):
###
# Methods that interact with the experiment
###
[docs] @abc.abstractmethod
def reset_stm(self) -> None:
pass
[docs] @abc.abstractmethod
def fixme_reset_ground_truth(self, primary_target=None) -> None:
pass
[docs] @abc.abstractmethod
def update_ltm_from_stm(self) -> None:
pass
[docs] @abc.abstractmethod
def fixme_update_ground_truth(self) -> None:
pass
[docs] @abc.abstractmethod
def set_experiment_mode(self, mode: ExperimentMode) -> None:
pass
###
# Methods that define the algorithm
###
[docs] @abc.abstractmethod
def matching_step(self, ctx: RuntimeContext, percepts: Sequence[Message]) -> None:
pass
[docs] @abc.abstractmethod
def exploratory_step(
self, ctx: RuntimeContext, percepts: Sequence[Message]
) -> None:
pass
[docs] @abc.abstractmethod
def receive_votes(self, votes: Collection[Any]) -> None:
pass
[docs] @abc.abstractmethod
def send_out_vote(self) -> Any:
pass
[docs] @abc.abstractmethod
def propose_goals(self) -> list[Goal]:
pass
[docs] @abc.abstractmethod
def get_output(self) -> Message | None:
pass
[docs] @abc.abstractmethod
def init_from_ltm(self) -> None:
pass
###
# Saving, loading
###
[docs] @abc.abstractmethod
def state_dict(self) -> Memento:
pass
[docs] @abc.abstractmethod
def load_state_dict(self, memento: Memento) -> None:
pass
[docs]class LMMemory(Snapshotable, metaclass=abc.ABCMeta):
"""Like a long-term memory storing all the knowledge an LM has."""
###
# Methods that define the algorithm
###
[docs] @abc.abstractmethod
def update_memory(self, observations):
"""Update models stored in memory given new observation & classification."""
pass
###
# Saving, loading
###
[docs] @abc.abstractmethod
def state_dict(self) -> Memento:
pass
[docs] @abc.abstractmethod
def load_state_dict(self, memento: Memento) -> None:
pass
[docs]class ObjectModel(metaclass=abc.ABCMeta):
"""Model of an object. Is stored in Memory and used by LM."""
[docs] @abc.abstractmethod
def build_model(self, observations):
"""Build a new model."""
pass
[docs] @abc.abstractmethod
def update_model(self, observations):
"""Update an existing model with new observations."""
pass
[docs]class GoalGenerator(metaclass=abc.ABCMeta):
"""Generate goals that other learning modules and motor-systems will attempt.
Generate goals potentially (in the case of LMs) by outputting their own
sub-goals. Provides a mechanism for implementing hierarchical action policies
that are informed by world models/hypotheses.
"""
[docs] @abc.abstractmethod
def set_driving_goal(self):
"""Set the driving goal.
e.g., from a human operator or a high-level LM.
"""
pass
[docs] @abc.abstractmethod
def output_goals(self) -> list[Goal]:
"""Return output goals."""
pass
[docs] @abc.abstractmethod
def step(self, ctx: RuntimeContext, observations: Observations):
"""Called on each step of the LM to which the GSG belongs."""
pass
class RuntimeSensorModule(Protocol):
"""Monty runtime interface to a Sensor Module."""
def update_state(self, agent: AgentState) -> None:
"""Update the proprioceptive state for this Sensor Module.
Args:
agent: The proprioceptive state of this sensor module's Agent.
"""
...
def step(
self,
ctx: RuntimeContext,
observation: SensorObservation,
motor_only_step: bool = False,
) -> Message | None:
"""Execute a time-step for the Sensor Module.
Args:
ctx: The runtime context.
observation: Sensor observation.
motor_only_step: Whether the current step is a motor-only step.
Returns:
An optional percept with features and morphological features.
"""
...
def propose_goals(self) -> Sequence[Goal]:
"""Return the goals proposed by this Sensor Module.
Returns:
A sequence of proposed Goals.
"""
...
[docs]class SensorModule(RuntimeSensorModule, ExperimentSensorModule, metaclass=abc.ABCMeta):
[docs] @abc.abstractmethod
def state_dict(self) -> Memento:
pass
[docs] @abc.abstractmethod
def update_state(self, agent: AgentState) -> None:
pass
[docs] @abc.abstractmethod
def step(
self,
ctx: RuntimeContext,
observation: SensorObservation,
motor_only_step: bool = False,
) -> Message | None:
pass
[docs] def propose_goals(self) -> list[Goal]:
return []
[docs] @abc.abstractmethod
def reset(self) -> None:
pass