update imu test

Author: Milotrince

genesis/engine/sensors/base_sensor.py

@@ -12,8 +12,8 @@
 from genesis.utils.misc import concat_with_tensor, make_tensor_field
 
 if TYPE_CHECKING:
-    from genesis.engine.solvers import RigidSolver
     from genesis.engine.entities.rigid_entity.rigid_link import RigidLink
+    from genesis.engine.solvers import RigidSolver
     from genesis.recorders.base_recorder import Recorder, RecorderOptions
     from genesis.utils.ring_buffer import TensorRingBuffer
     from genesis.vis.rasterizer_context import RasterizerContext
@@ -292,8 +292,38 @@ def _get_formatted_data(self, tensor: torch.Tensor, envs_idx=None) -> torch.Tens
         return self._return_data_class(*return_values)
 
     def _sanitize_envs_idx(self, envs_idx) -> torch.Tensor:
+        if self._manager._sim.n_envs == 0:
+            return torch.tensor([0], device=gs.device, dtype=gs.tc_int)
         return self._manager._sim._scene._sanitize_envs_idx(envs_idx)
 
+    def _set_metadata_field(self, input, field, field_size, envs_idx):
+        envs_idx = self._sanitize_envs_idx(envs_idx)
+        if field.ndim == 2:
+            # flat field structure
+            idx = self._idx * field_size
+            index_slice = slice(idx, idx + field_size)
+        else:
+            # per sensor field structure
+            index_slice = self._idx
+        # field[envs_idx, index_slice] = self._sanitize_for_metadata_tensor(
+        field[:, index_slice] = self._sanitize_for_metadata_tensor(
+            input, shape=(len(envs_idx), field_size), dtype=field.dtype
+        )
+
+    def _sanitize_for_metadata_tensor(self, input, shape, dtype) -> torch.Tensor:
+        if not isinstance(input, Sequence):
+            input = [input]
+        tensor_input = torch.tensor(input, dtype=dtype, device=gs.device)
+        if tensor_input.ndim == len(shape) - 1:
+            # Batch dimension is missing
+            tensor_input = tensor_input.unsqueeze(0)
+        if tensor_input.shape[0] != shape[0]:
+            tensor_input = tensor_input.expand((shape[0], *tensor_input.shape[1:]))
+        assert (
+            tensor_input.shape == shape
+        ), f"Input shape {tensor_input.shape} for setting sensor metadata does not match shape {shape}"
+        return tensor_input
+
 
 @dataclass
 class RigidSensorMetadataMixin:
@@ -345,6 +375,16 @@ def build(self):
             dim=1,
         )
 
+    @gs.assert_built
+    def set_pos_offset(self, pos_offset, envs_idx=None):
+        envs_idx = self._sanitize_envs_idx(envs_idx)
+        self._set_metadata_field(pos_offset, self._shared_metadata.offsets_pos, field_size=3, envs_idx=envs_idx)
+
+    @gs.assert_built
+    def set_quat_offset(self, quat_offset, envs_idx=None):
+        envs_idx = self._sanitize_envs_idx(envs_idx)
+        self._set_metadata_field(quat_offset, self._shared_metadata.offsets_quat, field_size=4, envs_idx=envs_idx)
+
 
 @dataclass
 class NoisySensorMetadataMixin:
@@ -372,25 +412,6 @@ class NoisySensorMixin(Generic[NoisySensorMetadataMixinT]):
     Base sensor class for analog sensors that are attached to a RigidEntity.
     """
 
-    def _set_metadata_field(self, input, field, field_size, envs_idx):
-        envs_idx = self._sanitize_envs_idx(envs_idx)
-        idx = self._idx * field_size
-        field[envs_idx, idx : idx + field_size] = self._sanitize_for_metadata_tensor(
-            input, shape=(len(envs_idx), field_size), dtype=field.dtype
-        )
-
-    def _sanitize_for_metadata_tensor(self, input, shape, dtype) -> torch.Tensor:
-        if not isinstance(input, Sequence):
-            input = [input]
-        tensor_input = torch.tensor(input, dtype=dtype, device=gs.device)
-        if tensor_input.ndim == len(shape) - 1:
-            # Batch dimension is missing
-            tensor_input = tensor_input.unsqueeze(0).expand((shape[0], *tensor_input.shape))
-        assert (
-            tensor_input.shape == shape
-        ), f"Input shape {tensor_input.shape} for setting sensor metadata does not match shape {shape}"
-        return tensor_input
-
     @gs.assert_built
     def set_resolution(self, resolution, envs_idx=None):
         self._set_metadata_field(resolution, self._shared_metadata.resolution, self._cache_size, envs_idx)

genesis/engine/sensors/imu.py

@@ -6,11 +6,13 @@
 import torch
 
 import genesis as gs
-from genesis.options.sensors import (
-    MaybeMatrix3x3Type,
-    IMU as IMUOptions,
+from genesis.options.sensors import IMU as IMUOptions
+from genesis.options.sensors import MaybeMatrix3x3Type
+from genesis.utils.geom import (
+    inv_transform_by_quat,
+    transform_by_quat,
+    transform_quat_by_quat,
 )
-from genesis.utils.geom import inv_transform_by_trans_quat, transform_by_quat, transform_quat_by_quat
 from genesis.utils.misc import concat_with_tensor, make_tensor_field, tensor_to_array
 
 from .base_sensor import (
@@ -30,20 +32,6 @@
     from genesis.vis.rasterizer_context import RasterizerContext
 
 
-def _view_metadata_as_acc_gyro(metadata_tensor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    Get views of the metadata tensor (B, n_imus * 6) as a tuple of acc and gyro metadata tensors (B, n_imus * 3).
-    """
-    batch_shape, n_data = metadata_tensor.shape[:-1], metadata_tensor.shape[-1]
-    n_imus = n_data // 6
-    metadata_tensor_per_sensor = metadata_tensor.reshape((*batch_shape, n_imus, 2, 3))
-
-    return (
-        metadata_tensor_per_sensor[..., 0, :].reshape(*batch_shape, n_imus * 3),
-        metadata_tensor_per_sensor[..., 1, :].reshape(*batch_shape, n_imus * 3),
-    )
-
-
 def _get_skew_to_alignment_matrix(input: MaybeMatrix3x3Type, out: torch.Tensor | None = None) -> torch.Tensor:
     """
     Convert the alignment input to a matrix. Modifies in place if provided, else allocate a new matrix.
@@ -120,30 +108,6 @@ def set_gyro_axes_skew(self, axes_skew: MaybeMatrix3x3Type, envs_idx=None):
         rot_matrix = _get_skew_to_alignment_matrix(axes_skew)
         self._shared_metadata.alignment_rot_matrix[envs_idx, self._idx * 2 + 1, :, :] = rot_matrix
 
-    @gs.assert_built
-    def set_acc_bias(self, bias, envs_idx=None):
-        self._set_metadata_field(bias, self._shared_metadata.acc_bias, field_size=3, envs_idx=envs_idx)
-
-    @gs.assert_built
-    def set_gyro_bias(self, bias, envs_idx=None):
-        self._set_metadata_field(bias, self._shared_metadata.gyro_bias, field_size=3, envs_idx=envs_idx)
-
-    @gs.assert_built
-    def set_acc_random_walk(self, random_walk, envs_idx=None):
-        self._set_metadata_field(random_walk, self._shared_metadata.acc_random_walk, field_size=3, envs_idx=envs_idx)
-
-    @gs.assert_built
-    def set_gyro_random_walk(self, random_walk, envs_idx=None):
-        self._set_metadata_field(random_walk, self._shared_metadata.gyro_random_walk, field_size=3, envs_idx=envs_idx)
-
-    @gs.assert_built
-    def set_acc_noise(self, noise, envs_idx=None):
-        self._set_metadata_field(noise, self._shared_metadata.acc_noise, field_size=3, envs_idx=envs_idx)
-
-    @gs.assert_built
-    def set_gyro_noise(self, noise, envs_idx=None):
-        self._set_metadata_field(noise, self._shared_metadata.gyro_noise, field_size=3, envs_idx=envs_idx)
-
     # ================================ internal methods ================================
 
     def build(self):
@@ -160,15 +124,6 @@ def build(self):
         self._options.noise = _to_tuple(self._options.acc_noise, self._options.gyro_noise, length_per_value=3)
         super().build()  # set all shared metadata from RigidSensorBase and NoisySensorBase
 
-        self._shared_metadata.acc_bias, self._shared_metadata.gyro_bias = _view_metadata_as_acc_gyro(
-            self._shared_metadata.bias
-        )
-        self._shared_metadata.acc_random_walk, self._shared_metadata.gyro_random_walk = _view_metadata_as_acc_gyro(
-            self._shared_metadata.random_walk
-        )
-        self._shared_metadata.acc_noise, self._shared_metadata.gyro_noise = _view_metadata_as_acc_gyro(
-            self._shared_metadata.noise
-        )
         self._shared_metadata.alignment_rot_matrix = concat_with_tensor(
             self._shared_metadata.alignment_rot_matrix,
             torch.stack(
@@ -203,15 +158,30 @@ def _update_shared_ground_truth_cache(
         quats = shared_metadata.solver.get_links_quat(links_idx=shared_metadata.links_idx)
         acc = shared_metadata.solver.get_links_acc(links_idx=shared_metadata.links_idx)
         ang = shared_metadata.solver.get_links_ang(links_idx=shared_metadata.links_idx)
+        if acc.ndim == 2:
+            acc = acc.unsqueeze(0)
+            ang = ang.unsqueeze(0)
 
         offset_quats = transform_quat_by_quat(quats, shared_metadata.offsets_quat)
 
+        # additional acceleration if offset: a_imu = a_link + α × r + ω × (ω × r)
+        if torch.any(torch.abs(shared_metadata.offsets_pos) > gs.EPS):
+            ang_acc = shared_metadata.solver.get_links_acc_ang(links_idx=shared_metadata.links_idx)
+            if ang_acc.ndim == 2:
+                ang_acc = ang_acc.unsqueeze(0)
+            offset_pos_world = transform_by_quat(shared_metadata.offsets_pos, quats)
+            tangential_acc = torch.cross(ang_acc, offset_pos_world, dim=-1)
+            centripetal_acc = torch.cross(ang, torch.cross(ang, offset_pos_world, dim=-1), dim=-1)
+            acc += tangential_acc + centripetal_acc
+
         # acc/ang shape: (B, n_imus, 3)
-        local_acc = inv_transform_by_trans_quat(acc, shared_metadata.offsets_pos, offset_quats)
-        local_ang = inv_transform_by_trans_quat(ang, shared_metadata.offsets_pos, offset_quats)
+        local_acc = inv_transform_by_quat(acc, offset_quats)
+        local_ang = inv_transform_by_quat(ang, offset_quats)
 
         *batch_size, n_imus, _ = local_acc.shape
-        local_acc = local_acc - gravity.unsqueeze(-2).expand((*batch_size, n_imus, -1))
+        local_acc = local_acc - inv_transform_by_quat(
+            gravity.unsqueeze(-2).expand((*batch_size, n_imus, -1)), offset_quats
+        )
 
         # cache shape: (B, n_imus * 6)
         strided_ground_truth_cache = shared_ground_truth_cache.reshape((*batch_size, n_imus, 2, 3))

tests/test_imu.py

@@ -0,0 +1,160 @@
+import numpy as np
+import torch
+from utils import assert_allclose, assert_array_equal
+
+import genesis as gs
+import genesis.utils.geom as gu
+
+
+def expand_batch_dim(values: tuple[float, ...], n_envs: int) -> tuple[float, ...] | np.ndarray:
+    """Helper function to expand expected values for n_envs dimension."""
+    if n_envs == 0:
+        return values
+    return np.tile(np.array(values), (n_envs,) + (1,) * len(values))
+
+
+def test_imu_sensor(show_viewer, tol, n_envs):
+    """Test if the IMU sensor returns the correct data."""
+    GRAVITY = -10.0
+    DT = 1e-2
+    BIAS = (0.1, 0.2, 0.3)
+    DELAY_STEPS = 2
+
+    scene = gs.Scene(
+        sim_options=gs.options.SimOptions(
+            dt=DT,
+            substeps=1,
+            gravity=(0.0, 0.0, GRAVITY),
+        ),
+        profiling_options=gs.options.ProfilingOptions(show_FPS=False),
+        show_viewer=show_viewer,
+    )
+
+    scene.add_entity(gs.morphs.Plane())
+
+    box = scene.add_entity(
+        morph=gs.morphs.Box(
+            size=(0.1, 0.1, 0.1),
+            pos=(0.0, 0.0, 0.2),
+        ),
+    )
+
+    imu = scene.add_sensor(
+        gs.sensors.IMU(
+            entity_idx=box.idx,
+        )
+    )
+    imu_delayed = scene.add_sensor(
+        gs.sensors.IMU(
+            entity_idx=box.idx,
+            delay=DT * DELAY_STEPS,
+        )
+    )
+    imu_noisy = scene.add_sensor(
+        gs.sensors.IMU(
+            entity_idx=box.idx,
+            acc_axes_skew=0.01,
+            gyro_axes_skew=(0.02, 0.03, 0.04),
+            acc_noise=(0.01, 0.01, 0.01),
+            gyro_noise=(0.01, 0.01, 0.01),
+            acc_random_walk=(0.001, 0.001, 0.001),
+            gyro_random_walk=(0.001, 0.001, 0.001),
+            delay=DT,
+            jitter=DT * 0.1,
+            interpolate=True,
+        )
+    )
+
+    scene.build(n_envs=n_envs)
+
+    # box is in freefall
+    for _ in range(10):
+        scene.step()
+
+    # IMU should calculate "classical linear acceleration" using the local frame without accounting for gravity
+    # acc_classical_lin_z = - theta_dot ** 2 - cos(theta) * g
+    assert_allclose(imu.read().lin_acc, 0.0, tol=tol)
+    assert_allclose(imu.read().ang_vel, 0.0, tol=tol)
+    assert_allclose(imu_noisy.read().lin_acc, 0.0, tol=1e-1)
+    assert_allclose(imu_noisy.read().ang_vel, 0.0, tol=1e-1)
+
+    # shift COM to induce angular velocity
+    com_shift = torch.tensor([[0.05, 0.05, 0.05]])
+    box.set_COM_shift(com_shift.expand((n_envs, 1, 3)) if n_envs > 0 else com_shift)
+
+    # update noise and bias for accelerometer and gyroscope
+    imu_noisy.set_noise((0.01, 0.01, 0.01, 0.02, 0.02, 0.02))
+    imu_noisy.set_bias((0.01, 0.01, 0.01, 0.02, 0.02, 0.02))
+    imu_noisy.set_jitter(0.001)
+
+    for _ in range(10 - DELAY_STEPS):
+        scene.step()
+
+    true_imu_delayed_reading = imu_delayed.read_ground_truth()
+
+    for _ in range(DELAY_STEPS):
+        scene.step()
+
+    assert_array_equal(imu_delayed.read().lin_acc, true_imu_delayed_reading.lin_acc)
+    assert_array_equal(imu_delayed.read().ang_vel, true_imu_delayed_reading.ang_vel)
+
+    # check that position offset affects linear acceleration
+    imu.set_pos_offset((0.5, 0.0, 0.0))
+    lin_acc_no_offset = imu.read().lin_acc
+    scene.step()
+    lin_acc_with_offset = imu.read().lin_acc
+    assert not np.allclose(lin_acc_no_offset, lin_acc_with_offset, atol=0.2)
+    imu.set_pos_offset((0.0, 0.0, 0.0))
+
+    # let box collide with ground
+    for _ in range(20):
+        scene.step()
+
+    assert_array_equal(imu.read_ground_truth().lin_acc, imu_delayed.read_ground_truth().lin_acc)
+    assert_array_equal(imu.read_ground_truth().ang_vel, imu_delayed.read_ground_truth().ang_vel)
+
+    with np.testing.assert_raises(AssertionError, msg="Angular velocity should not be zero due to COM shift"):
+        assert_allclose(imu.read_ground_truth().ang_vel, 0.0, tol=tol)
+
+    with np.testing.assert_raises(AssertionError, msg="Delayed data should not be equal to the ground truth data"):
+        assert_array_equal(imu_delayed.read().lin_acc - imu_delayed.read_ground_truth().lin_acc, 0.0)
+
+    zero_com_shift = torch.tensor([[0.0, 0.0, 0.0]])
+    box.set_COM_shift(zero_com_shift.expand((n_envs, 1, 3)) if n_envs > 0 else zero_com_shift)
+    quat_tensor = torch.tensor([0.0, 0.0, 0.0, 1.0])
+    box.set_quat(quat_tensor.expand((n_envs, 4)) if n_envs > 0 else quat_tensor)
+
+    # box is stationary on ground
+    for _ in range(80):
+        scene.step()
+
+    assert_allclose(
+        imu.read().lin_acc,
+        expand_batch_dim((0.0, 0.0, -GRAVITY), n_envs),
+        tol=5e-6,
+    )
+    assert_allclose(imu.read().ang_vel, expand_batch_dim((0.0, 0.0, 0.0), n_envs), tol=1e-5)
+
+    # rotate IMU 90 deg around x axis means gravity should be along -y axis
+    imu.set_quat_offset(gu.euler_to_quat((90.0, 0.0, 0.0)))
+    imu.set_acc_axes_skew((0.0, 1.0, 0.0))
+    scene.step()
+    assert_allclose(imu.read().lin_acc, GRAVITY, tol=5e-6)
+    imu.set_quat_offset((0.0, 0.0, 0.0, 1.0))
+    imu.set_acc_axes_skew((0.0, 0.0, 0.0))
+
+    scene.reset()
+
+    assert_allclose(imu.read().lin_acc, 0.0, tol=gs.EPS)  # biased, but cache hasn't been updated yet
+    assert_allclose(imu_delayed.read().lin_acc, 0.0, tol=gs.EPS)
+    assert_allclose(imu_noisy.read().ang_vel, 0.0, tol=gs.EPS)
+
+    imu.set_bias(BIAS + (0.0, 0.0, 0.0))
+    scene.step()
+    assert_allclose(imu.read().lin_acc, expand_batch_dim(BIAS, n_envs), tol=tol)
+
+
+if __name__ == "__main__":
+    gs.init(backend=gs.cpu)
+    # test_imu_sensor(show_viewer=False, tol=1e-4, n_envs=0)
+    test_imu_sensor(show_viewer=False, tol=1e-4, n_envs=2)