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

Code before:

    def _get_cache_dtype(self) -> 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_return_data_from_tensor(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_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}")

Review comment (performance):
Could you preallocate the full vector of environment as it is done in `Scene` ? For maximum efficiency is should be a torch tensor on device if not a scalar.

Code after:
        """
        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