Merge branch 'main' into touch_object_jun25_23

Author: tristanls

src/tbp/monty/frameworks/environments/embodied_environment.py

@@ -14,7 +14,7 @@
 
 from tbp.monty.frameworks.actions.actions import Action
 
-__all__ = ["EmbodiedEnvironment", "ActionSpace"]
+__all__ = ["EmbodiedEnvironment", "ActionSpace", "VectorXYZ", "QuaternionWXYZ"]
 
 VectorXYZ = Tuple[float, float, float]
 QuaternionWXYZ = Tuple[float, float, float, float]

src/tbp/monty/simulators/habitat/environment.py

@@ -36,6 +36,7 @@
     "SingleSensorAgentArgs",
 ]
 
+from tbp.monty.simulators.simulator import Simulator
 
 # Create agent and object configuration helper dataclasses
 
@@ -116,7 +117,7 @@ def __init__(
             agent = agent_type(**args)
             self._agents.append(agent)
 
-        self._env = HabitatSim(
+        self._env: Simulator = HabitatSim(
             agents=self._agents,
             scene_id=scene_id,
             seed=seed,
@@ -139,7 +140,7 @@ def add_object(
         rotation: QuaternionWXYZ = (1.0, 0.0, 0.0, 0.0),
         scale: VectorXYZ = (1.0, 1.0, 1.0),
         semantic_id: Optional[str] = None,
-        enable_physics: Optional[bool] = False,
+        enable_physics=False,
         object_to_avoid=False,
         primary_target_object=None,
     ):

src/tbp/monty/simulators/habitat/simulator.py

@@ -16,7 +16,7 @@
 
 from collections import defaultdict
 from pathlib import Path
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, List, Optional
 
 import habitat_sim
 import magnum as mn
@@ -52,6 +52,11 @@
     "PRIMITIVE_OBJECT_TYPES",
 ]
 
+from tbp.monty.frameworks.environments.embodied_environment import (
+    QuaternionWXYZ,
+    VectorXYZ,
+)
+
 DEFAULT_SCENE = "NONE"
 DEFAULT_PHYSICS_CONFIG = str(files(resources) / "default.physics_config.json")
 
@@ -66,10 +71,6 @@
 }
 
 
-Vector3 = Tuple[float, float, float]
-Quaternion = Tuple[float, float, float, float]
-
-
 class HabitatSim(HabitatActuator):
     """Habitat-sim interface for tbp.monty.
 
@@ -211,11 +212,11 @@ def remove_all_objects(self):
     def add_object(
         self,
         name: str,
-        position: Vector3 = (0.0, 0.0, 0.0),
-        rotation: Quaternion = (1.0, 0.0, 0.0, 0.0),
-        scale: Vector3 = (1.0, 1.0, 1.0),
+        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: Optional[str] = None,
-        enable_physics: Optional[bool] = False,
+        enable_physics=False,
         object_to_avoid=False,
         primary_target_bb=None,
     ):

src/tbp/monty/simulators/simulator.py

@@ -0,0 +1,104 @@
+# 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 typing import Dict, List, Optional, Protocol
+
+from tbp.monty.frameworks.actions.actions import Action
+from tbp.monty.frameworks.environments.embodied_environment import (
+    QuaternionWXYZ,
+    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, agent_state):
+        """Update agent runtime state."""
+        ...
+
+    def remove_all_objects(self):
+        """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: Optional[str] = None,
+        enable_physics=False,
+        object_to_avoid=False,
+        primary_target_bb: Optional[List] = None,
+    ) -> None:
+        """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 (str): Registered object name
+            position (VectorXYZ): Initial absolute position of the object
+            rotation (QuaternionWXYZ): Initial orientation of the object
+            scale (VectorXYZ): Initial object scale
+            semantic_id (Optional[str]): Optional override object semantic ID
+            enable_physics (bool): Whether to enable physics on the object
+            object_to_avoid (bool): If True, ensure the object is not colliding with
+              other objects
+            primary_target_bb (List | None): If not None, this is a list of the min and
+              max corners of a bounding box for the primary object, used to prevent
+              obscuring the primary objet with the new object.
+        """
+        ...
+
+    # TODO - change getters to properties using @property
+    def get_num_objects(self) -> int:
+        """Return the number of instantiated objects in the environment."""
+        ...
+
+    def get_action_space(self):
+        """Returns the set of all available actions."""
+        ...
+
+    def get_agent(self, agent_id):
+        """Return agent instance."""
+        ...
+
+    def get_observations(self):
+        """Get sensor observations."""
+        ...
+
+    def get_states(self):
+        """Get agent and sensor states."""
+        ...
+
+    def apply_action(self, action: Action) -> Dict[str, Dict]:
+        """Execute the given action in the environment.
+
+        Args:
+            action (Action): the action to execute
+
+        Returns:
+            (Dict[str, Dict]): A dictionary with the observations grouped by agent_id
+        """
+        ...
+
+    def reset(self):
+        """Reset the simulator."""
+        ...
+
+    def close(self):
+        """Close any resources used by the simulator."""
+        ...