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

Code before:
            length_per_value=3,
        )
        self._options.bias = _to_tuple(
            self._options.acc_bias, self._options.gyro_bias, self._options.mag_bias, length_per_value=3
        )
        self._options.random_walk = _to_tuple(
            self._options.acc_random_walk,
            self._options.gyro_random_walk,
            self._options.mag_random_walk,
            length_per_value=3,
        )
        self._options.noise = _to_tuple(
            self._options.acc_noise, self._options.gyro_noise, self._options.mag_noise, length_per_value=3
        )
        super().build()  # set all shared metadata from RigidSensorBase and NoisySensorBase

        self._shared_metadata.alignment_rot_matrix = concat_with_tensor(
            self._shared_metadata.alignment_rot_matrix,
            torch.stack(
                [
                    _get_cross_axis_coupling_to_alignment_matrix(self._options.acc_cross_axis_coupling),
                    _get_cross_axis_coupling_to_alignment_matrix(self._options.gyro_cross_axis_coupling),
                    _get_cross_axis_coupling_to_alignment_matrix(self._options.mag_cross_axis_coupling),
                ],
            ).unsqueeze(0),
            expand=(self._manager._sim._B, 1, 3, 3, 3),  # 3 sub-matrices after adding mag
            dim=1,
        )

        # Initialize global magnetic field vector
        default_field = self._options.magnetic_field if self._options.magnetic_field is not None else (0.0, 0.0, 0.5)
        if not isinstance(default_field, torch.Tensor):
            default_field = torch.tensor(default_field, device=gs.device, dtype=gs.tc_float)

        self._shared_metadata.magnetic_field_vector = concat_with_tensor(
            self._shared_metadata.magnetic_field_vector,
            default_field,
            expand=(self._manager._sim._B, 3),
            dim=0,
        )

        if self._options.draw_debug:
            self.quat_offset = self._shared_metadata.offsets_quat[0, self._idx]
            self.pos_offset = self._shared_metadata.offsets_pos[0, self._idx]

    def _get_return_format(self) -> tuple[tuple[int, ...], ...]:
        return (3,), (3,), (3,)

    @classmethod
    def _get_cache_dtype(cls) -> torch.dtype:
        return gs.tc_float

Review comment (bug):
I don't want to risk introducing a bug. Are you sure that removing `.unsqueeze(0)` and adding one extra dim is equivalent?

Code after:
            length_per_value=3,
        )
        self._options.bias = _to_tuple(
            self._options.acc_bias, self._options.gyro_bias, self._options.mag_bias, length_per_value=3
        )
        self._options.random_walk = _to_tuple(
            self._options.acc_random_walk,
            self._options.gyro_random_walk,
            self._options.mag_random_walk,
            length_per_value=3,
        )
        self._options.noise = _to_tuple(
            self._options.acc_noise, self._options.gyro_noise, self._options.mag_noise, length_per_value=3
        )
        super().build()  # set all shared metadata from RigidSensorBase and NoisySensorBase

        self._shared_metadata.alignment_rot_matrix = concat_with_tensor(
            self._shared_metadata.alignment_rot_matrix,
            torch.stack(
                [
                    _get_cross_axis_coupling_to_alignment_matrix(self._options.acc_cross_axis_coupling),
                    _get_cross_axis_coupling_to_alignment_matrix(self._options.gyro_cross_axis_coupling),
                    _get_cross_axis_coupling_to_alignment_matrix(self._options.mag_cross_axis_coupling),
                ],
            ),
            expand=(self._manager._sim._B, 3, 3, 3),  # 3 sub-matrices after adding mag
            dim=1,
        )

        # Initialize global magnetic field vector
        default_field = self._options.magnetic_field if self._options.magnetic_field is not None else (0.0, 0.0, 0.5)
        if not isinstance(default_field, torch.Tensor):
            default_field = torch.tensor(default_field, device=gs.device, dtype=gs.tc_float)

        self._shared_metadata.magnetic_field_vector = concat_with_tensor(
            self._shared_metadata.magnetic_field_vector,
            default_field,
            expand=(self._manager._sim._B, 1, 3),
            dim=1,
        )

        if self._options.draw_debug:
            self.quat_offset = self._shared_metadata.offsets_quat[0, self._idx]
            self.pos_offset = self._shared_metadata.offsets_pos[0, self._idx]

    def _get_return_format(self) -> tuple[tuple[int, ...], ...]:
        return (3,), (3,), (3,)

    @classmethod
    def _get_cache_dtype(cls) -> torch.dtype:
        return gs.tc_float