add SensorManager

Author: Milotrince

genesis/engine/scene.py

@@ -25,6 +25,7 @@
     PBDOptions,
     ProfilingOptions,
     RigidOptions,
+    SensorOptions,
     SFOptions,
     SimOptions,
     SPHOptions,
@@ -513,6 +514,10 @@ def add_light(
         else:
             gs.raise_exception("Adding lights is only supported by 'RayTracer' and 'BatchRenderer'.")
 
+    @gs.assert_unbuilt
+    def add_sensor(self, sensor_options: SensorOptions):
+        return self._sim._sensor_manager.create_sensor(sensor_options)
+
     @gs.assert_unbuilt
     def add_camera(
         self,

genesis/engine/simulator.py

@@ -36,6 +36,7 @@
 )
 from .states.cache import QueriedStates
 from .states.solvers import SimState
+from genesis.sensors.sensor_manager import SensorManager
 
 if TYPE_CHECKING:
     from genesis.engine.scene import Scene
@@ -151,6 +152,9 @@ def __init__(
         # entities
         self._entities: list[Entity] = gs.List()
 
+        # sensors
+        self._sensor_manager = SensorManager()
+
     def _add_entity(self, morph: Morph, material, surface, visualize_contact=False):
         if isinstance(material, gs.materials.Tool):
             entity = self.tool_solver.add_entity(self.n_entities, material, morph, surface)
@@ -206,6 +210,8 @@ def build(self):
         if self.n_envs > 0 and self.sf_solver.is_active():
             gs.raise_exception("Batching is not supported for SF solver as of now.")
 
+        self._sensor_manager.build()
+
         # hybrid
         for entity in self._entities:
             if isinstance(entity, HybridEntity):

genesis/options/sensors.py

@@ -0,0 +1,11 @@
+from genesis.options import Options
+
+
+class SensorOptions(Options):
+    """
+    Base class for all sensor options.
+    Each sensor should have their own options class that inherits from this class.
+    The options class should be registered with the SensorManager using the @register_sensor decorator.
+    """
+
+    pass

genesis/sensors/base_sensor.py

@@ -1,9 +1,13 @@
+from typing import List, Optional
+
 import taichi as ti
-import genesis as gs
 
-from typing import List, Optional
+import genesis as gs
+from genesis.options.sensors import SensorOptions
 from genesis.repr_base import RBC
 
+from .sensor_manager import SensorManager
+
 
 @ti.data_oriented
 class Sensor(RBC):
@@ -12,10 +16,30 @@ class Sensor(RBC):
     A sensor must have a read() method that returns the sensor data.
     """
 
+    def __init__(self, sensor_options: SensorOptions, sensor_idx: int, sensor_manager: SensorManager):
+        self._options: SensorOptions = sensor_options
+        self._idx: int = sensor_idx
+        self._manager: SensorManager = sensor_manager
+        self._cache_idx: int = -1  # cache_idx is set by the SensorManager during the scene build phase
+
+    @gs.assert_unbuilt
+    def build(self):
+        """
+        This method is called by the SensorManager during the scene build phase to initialize the sensor.
+        """
+        pass
+
     @gs.assert_built
     def read(self, envs_idx: Optional[List[int]] = None):
         """
         Read the sensor data.
-        Sensor implementations should ideally cache the data to avoid unnecessary computations.
+        Sensor implementations should make use of the caching system located in SensorManager when possible.
+        """
+        raise NotImplementedError("Sensors must implement `read()`.")
+
+    @property
+    def cache_length(self) -> int:
+        """
+        The length (first dimension of cache size) of the cache for this sensor.
         """
-        raise NotImplementedError("The Sensor subclass must implement `read()`.")
+        raise NotImplementedError("Sensors must implement `cache_length()`.")

genesis/sensors/sensor_manager.py

@@ -0,0 +1,65 @@
+from typing import Type
+
+import torch
+
+from genesis.options.sensors import SensorOptions
+
+from .base_sensor import Sensor
+
+
+class SensorManager:
+    SENSOR_TYPES_MAP: dict[Type[SensorOptions], Type[Sensor]] = {}
+    SENSOR_CACHE_METADATA_MAP: dict[Type[Sensor], (torch.dtype, tuple[int, ...])] = {}
+
+    def __init__(self):
+        self.sensors_by_type: dict[Type[Sensor], list[Sensor]] = {}
+        self.cache: dict[Type[Sensor], torch.Tensor] = {}
+        self.cache_size_map: dict[Type[Sensor], int] = {}
+
+    def create_sensor(self, sensor_options: SensorOptions):
+        sensor_cls = SensorManager.SENSOR_TYPES_MAP[type(sensor_options)]
+        sensor = sensor_cls(sensor_options, len(self.sensors_by_type[sensor_cls]), self)
+        if sensor_cls not in self.sensors_by_type:
+            self.sensors_by_type[sensor_cls] = []
+        self.sensors_by_type[sensor_cls].append(sensor)
+        return sensor
+
+    def build(self):
+        for sensor_cls, sensors in self.sensors_by_type.items():
+            total_cache_length = 0
+            for sensor in sensors:
+                sensor.build()
+                sensor._cache_idx = total_cache_length
+                total_cache_length += sensor.cache_length
+
+            cache_dtype, cache_shape = SensorManager.SENSOR_CACHE_METADATA_MAP[sensor_cls]
+            self.cache[sensor_cls] = torch.zeros((total_cache_length, *cache_shape), dtype=cache_dtype)
+
+            for sensor in sensors:
+                sensor.build()
+
+    def get_sensor_cache(self, sensor_cls: Type[Sensor], sensor_idx: int | None = None) -> torch.Tensor:
+        cache_size = SensorManager.SENSOR_CACHE_SIZE_MAP[sensor_cls]
+        if sensor_idx is None:
+            return self.cache[sensor_cls]
+        return self.cache[sensor_cls][sensor_idx * cache_size : (sensor_idx + 1) * cache_size]
+
+    def set_sensor_cache(self, new_values: torch.Tensor, sensor_cls: Type[Sensor], sensor_idx: int | None = None):
+        cache_size = SensorManager.SENSOR_CACHE_SIZE_MAP[sensor_cls]
+        if sensor_idx is None:
+            self.cache[sensor_cls] = new_values
+        else:
+            self.cache[sensor_cls][sensor_idx * cache_size : (sensor_idx + 1) * cache_size] = new_values
+
+    @property
+    def sensors(self):
+        return [sensor for sensor_list in self.sensors_by_type.values() for sensor in sensor_list]
+
+
+def register_sensor(sensor_cls: Type[Sensor], cache_dtype: torch.dtype, cache_shape: tuple[int, ...]):
+    def _impl(sensor_options: SensorOptions):
+        SensorManager.SENSOR_TYPES_MAP[type(sensor_options)] = sensor_cls
+        SensorManager.SENSOR_CACHE_METADATA_MAP[sensor_cls] = (cache_dtype, cache_shape)
+        return sensor_options
+
+    return _impl