Repository: Genesis-Embodied-AI/Genesis
File: genesis/sensors/sensor_manager.py (Python)

Code before:
from typing import Any, Type, TYPE_CHECKING

import torch

from genesis.utils.ring_buffer import TensorRingBuffer

if TYPE_CHECKING:
    from genesis.options.sensors import SensorOptions
    from .base_sensor import Sensor


class SensorManager:
    SENSOR_TYPES_MAP: dict[Type["SensorOptions"], Type["Sensor"]] = {}

    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._gt_cache: dict[Type["Sensor"], torch.Tensor] = {}
        self._cache: dict[Type["Sensor"], TensorRingBuffer] = {}

    def create_sensor(self, sensor_options: "SensorOptions"):
        sensor_cls = SensorManager.SENSOR_TYPES_MAP[type(sensor_options)]
        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():
            total_cache_length = 0
            max_cache_buf_len = 0
            self._sensors_metadata[sensor_cls] = {}
            for sensor in sensors:
                sensor.build()
                sensor._cache_idx = total_cache_length
                total_cache_length += sensor._get_cache_length()

                max_cache_buf_len = max(max_cache_buf_len, sensor._cache_buffer_length)

            cache_shape = (self._sim._B, total_cache_length, sensor_cls._get_cache_size())
            cache_dtype = sensor_cls._get_cache_dtype()


Review comment (performance):
For cache could be global, and sensor-type based. it is the big advantage of storing data flattened. But you need one cache per datatype still.

Code after:
from typing import TYPE_CHECKING, Any, Type

import numpy as np
import torch

import genesis as gs
from genesis.utils.ring_buffer import TensorRingBuffer

if TYPE_CHECKING:
    from genesis.options.sensors import SensorOptions

    from .base_sensor import Sensor


class SensorManager:
    SENSOR_TYPES_MAP: dict[Type["SensorOptions"], Type["Sensor"]] = {}

    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._ground_truth_cache: dict[Type[torch.dtype], torch.Tensor] = {}
        self._cache: dict[Type[torch.dtype], TensorRingBuffer] = {}
        self._cache_slices_by_type: dict[Type["Sensor"], slice] = {}

        self._last_ground_truth_cache_cloned_step: dict[Type[torch.dtype], int] = {}
        self._cloned_ground_truth_cache: dict[Type[torch.dtype], torch.Tensor] = {}

    def create_sensor(self, sensor_options: "SensorOptions"):
        sensor_cls = SensorManager.SENSOR_TYPES_MAP[type(sensor_options)]
        self._sensors_by_type.setdefault(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):
        max_cache_buf_len = 0
        cache_size_per_dtype = {}
        for sensor_cls, sensors in self._sensors_by_type.items():
            self._sensors_metadata[sensor_cls] = {}
            dtype = sensor_cls._get_cache_dtype()
            cache_size_per_dtype.setdefault(dtype, 0)
            cls_cache_start_idx = cache_size_per_dtype[dtype]