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

Code before:
    def read_ground_truth(self, envs_idx: List[int] | None = None):
        """
        Read the ground truth sensor data (without noise).
        """
        return self._get_return_data_from_tensor(self._ground_truth_cache, envs_idx)

    def _get_return_data_from_tensor(self, tensor: torch.Tensor, envs_idx: List[int] | None) -> torch.Tensor:
        if envs_idx is None:
            envs_idx = 0 if self._manager._sim.n_envs == 0 else np.arange(tensor.shape[0])

        return_format = self._get_return_format()

        if isinstance(return_format, tuple):
            return (
                tensor[envs_idx, self._cache_idx : self._cache_idx + self._cache_size]
                .clone()
                .reshape(return_format)
                .squeeze()
            )
        elif isinstance(return_format, dict):
            data_dict = {}
            tensor_idx = 0
            for key, data_shape in return_format.items():
                data_size = np.prod(data_shape)
                data_dict[key] = (
                    tensor[envs_idx, self._cache_idx : self._cache_idx + self._cache_size]
                    .clone()[tensor_idx : tensor_idx + data_size]
                    .reshape(data_shape)
                    .squeeze()
                )
                tensor_idx += data_size
            return data_dict
        else:
            raise ValueError(f"Invalid sensor return format: {return_format}")

    @property
    def is_built(self) -> bool:
        return self._manager._sim._scene._is_built

    @property
    def _cache(self) -> "TensorRingBuffer":
        return self._manager._cache[self._get_cache_dtype()]

    @property
    def _ground_truth_cache(self) -> torch.Tensor:
        return self._manager._ground_truth_cache[self._get_cache_dtype()]

    @property
    def _shared_metadata(self) -> dict[str, Any]:
        return self._manager._sensors_metadata[self.__class__]


Review comment (performance):
Start by extracting `tensor[envs_idx, self._cache_idx : self._cache_idx + self._cache_size]` once, instead of doing it repeatedly.

Code after:
from typing import TYPE_CHECKING, Any, List, Type

import numpy as np
import taichi as ti
import torch

import genesis as gs
from genesis.repr_base import RBC

if TYPE_CHECKING:
    from genesis.options.sensors import SensorOptions
    from genesis.utils.ring_buffer import TensorRingBuffer

    from .sensor_manager import SensorManager


@ti.data_oriented
class Sensor(RBC):
    """
    Base class for all types of sensors.
    """

    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

        # initialized by SensorManager during build
        self._read_delay_steps: int = 0
        self._shape_indices: list[tuple[int, int]] = []
        self._shared_metadata: dict[str, Any] | None = None
        self._cache: "TensorRingBuffer" | None = None

    # =============================== implementable methods ===============================

    def build(self):
        """
        This method is called by the SensorManager during the scene build phase to initialize the sensor.
        This is where any shared metadata should be initialized.
        """
        raise NotImplementedError("Sensors must implement `build()`.")

    def _get_return_format(self) -> dict[str, tuple[int, ...]] | tuple[int, ...]:
        """
        Data format of the read() return value.

        Returns
        -------
        return_format : dict | tuple
            - If tuple, the final shape of the read() return value.
                e.g. (2, 3) means read() will return a tensor of shape (2, 3).
            - If dict a dictionary with string keys and tensor values will be returned.
                e.g. {"pos": (3,), "quat": (4,)} returns a dict of tensors [0:3] and [3:7] from the cache.
        """
        raise NotImplementedError("Sensors must implement `return_format()`.")

    def _get_cache_length(self) -> int:
        """
        The length of the cache for this sensor instance, e.g. number of points for a Lidar point cloud.
        """
        raise NotImplementedError("Sensors must implement `cache_length()`.")

    @classmethod
    def _update_shared_ground_truth_cache(
        cls, shared_metadata: dict[str, Any], shared_ground_truth_cache: torch.Tensor
    ):
        """
        Update the shared sensor ground truth cache for all sensors of this class using metadata in SensorManager.
        """
        raise NotImplementedError("Sensors must implement `update_shared_ground_truth_cache()`.")

    @classmethod
    def _update_shared_cache(
        cls, shared_metadata: dict[str, Any], shared_ground_truth_cache: torch.Tensor, shared_cache: "TensorRingBuffer"
    ):
        """
        Update the shared sensor cache for all sensors of this class using metadata in SensorManager.
        """
        raise NotImplementedError("Sensors must implement `update_shared_cache()`.")

    @classmethod
    def _get_cache_dtype(cls) -> torch.dtype:
        """
        The dtype of the cache for this sensor.
        """
        raise NotImplementedError("Sensors must implement `get_cache_dtype()`.")

    # =============================== shared methods ===============================

    @gs.assert_built
    def read(self, envs_idx: List[int] | None = None):
        """
        Read the sensor data (with noise applied if applicable).
        """
        return self._get_formatted_data(self._cache.get(self._read_delay_steps), envs_idx)

    @gs.assert_built
    def read_ground_truth(self, envs_idx: List[int] | None = None):
        """
        Read the ground truth sensor data (without noise).
        """
        return self._get_formatted_data(self._manager.get_cloned_from_ground_truth_cache(self), envs_idx)

    def _get_formatted_data(
        self, tensor: torch.Tensor, envs_idx: list[int] | None
    ) -> torch.Tensor | dict[str, torch.Tensor]:
        # Note: This method does not clone the data tensor, it should have been cloned by the caller.

        if envs_idx is None:
            envs_idx = self._manager._sim._scene._envs_idx

        return_format = self._get_return_format()
        return_shapes = return_format.values() if isinstance(return_format, dict) else (return_format,)
        return_values = []

        for i, shape in enumerate(return_shapes):
            start_idx, end_idx = self._shape_indices[i]
            value = tensor[envs_idx, start_idx:end_idx].reshape(len(envs_idx), *shape).squeeze()
            if self._manager._sim.n_envs == 0:
                value = value.squeeze(0)
            return_values.append(value)

        if isinstance(return_format, dict):
            return dict(zip(return_format.keys(), return_values))
        else:
            return return_values[0]

    @property
    def is_built(self) -> bool:
        return self._manager._sim._scene._is_built