Source code for tbp.monty.simulators.simulator

# Copyright 2025 Thousand Brains Project
#
# 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

from typing import Protocol, Sequence

from tbp.monty.frameworks.actions.actions import Action
from tbp.monty.frameworks.agents import AgentID
from tbp.monty.frameworks.environments.embodied_environment import (
    ObjectID,
    ObjectInfo,
    QuaternionWXYZ,
    SemanticID,
    VectorXYZ,
)


class Simulator(Protocol):
    """A Protocol defining a simulator for use in simulated environments.

    A Simulator is responsible for a simulated environment that contains objects to
    interact with, agents to do the interacting, and for collecting observations and
    proprioceptive state to send to Monty.
    """

    # TODO - do we need a way to abstract the concept of "agent"?
    def initialize_agent(self, agent_id: AgentID, agent_state) -> None:
        """Update agent runtime state."""
        ...

    def remove_all_objects(self) -> None:
        """Remove all objects from the simulated environment."""
        ...

    def add_object(
        self,
        name: str,
        position: VectorXYZ = (0.0, 0.0, 0.0),
        rotation: QuaternionWXYZ = (1.0, 0.0, 0.0, 0.0),
        scale: VectorXYZ = (1.0, 1.0, 1.0),
        semantic_id: SemanticID | None = None,
        primary_target_object: ObjectID | None = None,
    ) -> ObjectInfo:
        """Add new object to simulated environment.

        Adds a new object based on the named object. This assumes that the set of
        available objects are preloaded and keyed by name.

        Args:
            name: Registered object name.
            position: Initial absolute position of the object.
            rotation: Initial orientation of the object.
            scale: Initial object scale.
            semantic_id: Optional override for the object's semantic ID.
            primary_target_object: ID of the primary target object. If not None, the
                added object will be positioned so that it does not obscure the initial
                view of the primary target object (which avoiding collision alone cannot
                guarantee). Used when adding multiple objects. Defaults to None.

        Returns:
            The added object's information.
        """
        ...

    @property
    def num_objects(self) -> int:
        """Return the number of instantiated objects in the environment."""
        ...

    @property
    def observations(self):
        """Get sensor observations."""
        ...

    @property
    def states(self):
        """Get agent and sensor states."""
        ...

    def apply_actions(self, actions: Sequence[Action]) -> dict[str, dict]:
        """Execute the given actions in the environment.

        Args:
            actions: The actions to execute.

        Returns:
            A dictionary with the observations grouped by agent_id.

        Note:
            If the actions are an empty sequence, the current observations are returned.
        """
        ...

    def reset(self):
        """Reset the simulator."""
        ...

    def close(self) -> None:
        """Close any resources used by the simulator."""
        ...