fix sensor manager metadata

Author: Milotrince

genesis/engine/simulator.py

@@ -153,7 +153,7 @@ def __init__(
         self._entities: list[Entity] = gs.List()
 
         # sensors
-        self._sensor_manager = SensorManager()
+        self._sensor_manager = SensorManager(self)
 
     def _add_entity(self, morph: Morph, material, surface, visualize_contact=False):
         if isinstance(material, gs.materials.Tool):

genesis/options/__init__.py

@@ -2,6 +2,6 @@
 from .solvers import *
 from .vis import *
 from .profiling import ProfilingOptions
-
+from .sensors import SensorOptions
 
 __all__ = ["ProfilingOptions"]

genesis/sensors/base_sensor.py

@@ -1,6 +1,7 @@
-from typing import List, Optional
+from typing import Any, List, Optional
 
 import taichi as ti
+import torch
 
 import genesis as gs
 from genesis.options.sensors import SensorOptions
@@ -16,12 +17,19 @@ class Sensor(RBC):
     A sensor must have a read() method that returns the sensor data.
     """
 
+    # These class variables are used by SensorManager to determine the cache metadata for the sensor.
+    # Sensor implementations should override these class variable values.
+    CACHE_DTYPE: torch.dtype = torch.float32
+    CACHE_SHAPE: tuple[int, ...] = (1,)
+
     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
 
+    # =============================== implementable methods ===============================
+
     @gs.assert_unbuilt
     def build(self):
         """
@@ -33,7 +41,6 @@ def build(self):
     def read(self, envs_idx: Optional[List[int]] = None):
         """
         Read the sensor data.
-        Sensor implementations should make use of the caching system located in SensorManager when possible.
         """
         raise NotImplementedError("Sensors must implement `read()`.")
 
@@ -42,4 +49,30 @@ def cache_length(self) -> int:
         """
         The length (first dimension of cache size) of the cache for this sensor.
         """
-        raise NotImplementedError("Sensors must implement `cache_length()`.")
+        return 1
+
+    # =============================== shared methods ===============================
+
+    @property
+    def is_built(self) -> bool:
+        return self._manager._sim._scene._is_built
+
+    @gs.assert_built
+    def _get_cache(self) -> torch.Tensor:
+        return self._manager.get_sensor_cache(self.__class__, self._cache_idx)
+
+    @gs.assert_built
+    def _is_cache_updated(self) -> bool:
+        return self._manager.is_cache_updated(self.__class__)
+
+    @gs.assert_built
+    def _set_cache_updated(self):
+        self._manager.set_cache_updated(self.__class__)
+
+    @property
+    def _cache(self) -> torch.Tensor:
+        return self._manager._cache[self.__class__]
+
+    @property
+    def _shared_metadata(self) -> dict[str, Any]:
+        return self._manager._sensors_metadata[self.__class__]

genesis/sensors/sensor_manager.py

@@ -1,65 +1,63 @@
-from typing import Type
+from typing import Any, Type
 
 import torch
 
 from genesis.options.sensors import SensorOptions
 
-from .base_sensor import Sensor
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    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, ...])] = {}
+    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 __init__(self, sim):
+        self._sim = sim
+        self._sensors_by_type: dict[Type["Sensor"], list["Sensor"]] = {}
+        self._sensors_metadata: dict[Type["Sensor"], dict[str, Any]] = {}
+        self._cache: dict[Type["Sensor"], torch.Tensor] = {}
+        self._cache_size_map: dict[Type["Sensor"], int] = {}
+        self._cache_last_updated_step_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)
+        if sensor_cls not in self._sensors_by_type:
+            self._sensors_by_type[sensor_cls] = []
+        sensor = sensor_cls(sensor_options, len(self._sensors_by_type[sensor_cls]), self)
+        self._sensors_by_type[sensor_cls].append(sensor)
         return sensor
 
     def build(self):
-        for sensor_cls, sensors in self.sensors_by_type.items():
+        for sensor_cls, sensors in self._sensors_by_type.items():
             total_cache_length = 0
+            self._cache_last_updated_step_map[sensor_cls] = -1
+            self._sensors_metadata[sensor_cls] = {}
             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()
+            self._cache[sensor_cls] = torch.zeros((self._sim._B, total_cache_length, *cache_shape), dtype=cache_dtype)
 
-    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 is_cache_updated(self, sensor_cls: Type["Sensor"]) -> bool:
+        return self._cache_last_updated_step_map[sensor_cls] == self._sim.cur_step_global
 
-    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
+    def set_cache_updated(self, sensor_cls: Type["Sensor"]):
+        self._cache_last_updated_step_map[sensor_cls] = self._sim.cur_step_global
 
     @property
     def sensors(self):
-        return [sensor for sensor_list in self.sensors_by_type.values() for sensor in sensor_list]
+        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
+def register_sensor(sensor_cls: Type["Sensor"]):
+    def _impl(options_cls: Type[SensorOptions]):
+        SensorManager.SENSOR_TYPES_MAP[options_cls] = sensor_cls
+        SensorManager.SENSOR_CACHE_METADATA_MAP[sensor_cls] = (sensor_cls.CACHE_DTYPE, sensor_cls.CACHE_SHAPE)
+        return options_cls
 
     return _impl