diff --git a/genesis/engine/sensors/base_sensor.py b/genesis/engine/sensors/base_sensor.py index 6f149defc..9b568dc9a 100644 --- a/genesis/engine/sensors/base_sensor.py +++ b/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,37 @@ 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[:, 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 +374,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 +411,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) diff --git a/genesis/engine/sensors/imu.py b/genesis/engine/sensors/imu.py index b3908289c..118180de0 100644 --- a/genesis/engine/sensors/imu.py +++ b/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,21 +32,9 @@ 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: +def _get_cross_axis_coupling_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. """ @@ -109,41 +99,17 @@ def __init__( self.pos_offset: torch.Tensor @gs.assert_built - def set_acc_axes_skew(self, axes_skew: MaybeMatrix3x3Type, envs_idx=None): + def set_acc_cross_axis_coupling(self, cross_axis_coupling: MaybeMatrix3x3Type, envs_idx=None): envs_idx = self._sanitize_envs_idx(envs_idx) - rot_matrix = _get_skew_to_alignment_matrix(axes_skew) + rot_matrix = _get_cross_axis_coupling_to_alignment_matrix(cross_axis_coupling) self._shared_metadata.alignment_rot_matrix[envs_idx, self._idx * 2, :, :] = rot_matrix @gs.assert_built - def set_gyro_axes_skew(self, axes_skew: MaybeMatrix3x3Type, envs_idx=None): + def set_gyro_cross_axis_coupling(self, cross_axis_coupling: MaybeMatrix3x3Type, envs_idx=None): envs_idx = self._sanitize_envs_idx(envs_idx) - rot_matrix = _get_skew_to_alignment_matrix(axes_skew) + rot_matrix = _get_cross_axis_coupling_to_alignment_matrix(cross_axis_coupling) 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,21 +126,12 @@ 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( [ - _get_skew_to_alignment_matrix(self._options.acc_axes_skew), - _get_skew_to_alignment_matrix(self._options.gyro_axes_skew), + _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), ], ), expand=(self._manager._sim._B, 2, 3, 3), @@ -203,15 +160,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)) diff --git a/genesis/options/sensors/options.py b/genesis/options/sensors/options.py index 281508a80..3d532d369 100644 --- a/genesis/options/sensors/options.py +++ b/genesis/options/sensors/options.py @@ -1,15 +1,12 @@ -from dataclasses import dataclass from typing import Sequence import numpy as np -import torch from pydantic import Field import genesis as gs from ..options import Options -from .raycaster import RaycastPattern, DepthCameraPattern - +from .raycaster import DepthCameraPattern, RaycastPattern Tuple3FType = tuple[float, float, float] MaybeTuple3FType = float | Tuple3FType @@ -190,7 +187,7 @@ class IMU(RigidSensorOptionsMixin, NoisySensorOptionsMixin, SensorOptions): acc_resolution : float, optional The measurement resolution of the accelerometer (smallest increment of change in the sensor reading). Default is 0.0, which means no quantization is applied. - acc_axes_skew : float | tuple[float, float, float] | Sequence[float] + acc_cross_axis_coupling : float | tuple[float, float, float] | Sequence[float] Accelerometer axes alignment as a 3x3 rotation matrix, where diagonal elements represent alignment (0.0 to 1.0) for each axis, and off-diagonal elements account for cross-axis misalignment effects. - If a scalar is provided (float), all off-diagonal elements are set to the scalar value. @@ -205,8 +202,8 @@ class IMU(RigidSensorOptionsMixin, NoisySensorOptionsMixin, SensorOptions): gyro_resolution : float, optional The measurement resolution of the gyroscope (smallest increment of change in the sensor reading). Default is 0.0, which means no quantization is applied. - gyro_axes_skew : float | tuple[float, float, float] | Sequence[float] - Gyroscope axes alignment as a 3x3 rotation matrix, similar to `acc_axes_skew`. + gyro_cross_axis_coupling : float | tuple[float, float, float] | Sequence[float] + Gyroscope axes alignment as a 3x3 rotation matrix, similar to `acc_cross_axis_coupling`. gyro_bias : tuple[float, float, float] The constant additive bias for each axis of the gyroscope. gyro_noise : tuple[float, float, float] @@ -225,8 +222,8 @@ class IMU(RigidSensorOptionsMixin, NoisySensorOptionsMixin, SensorOptions): acc_resolution: MaybeTuple3FType = 0.0 gyro_resolution: MaybeTuple3FType = 0.0 - acc_axes_skew: MaybeMatrix3x3Type = 0.0 - gyro_axes_skew: MaybeMatrix3x3Type = 0.0 + acc_cross_axis_coupling: MaybeMatrix3x3Type = 0.0 + gyro_cross_axis_coupling: MaybeMatrix3x3Type = 0.0 acc_noise: MaybeTuple3FType = 0.0 gyro_noise: MaybeTuple3FType = 0.0 acc_bias: MaybeTuple3FType = 0.0 @@ -240,15 +237,17 @@ class IMU(RigidSensorOptionsMixin, NoisySensorOptionsMixin, SensorOptions): debug_gyro_scale: float = 0.01 def model_post_init(self, _): - self._validate_axes_skew(self.acc_axes_skew) - self._validate_axes_skew(self.gyro_axes_skew) - - def _validate_axes_skew(self, axes_skew): - axes_skew_np = np.array(axes_skew) - if axes_skew_np.shape not in ((), (3,), (3, 3)): - gs.raise_exception(f"axes_skew shape should be (), (3,), or (3, 3), got: {axes_skew_np.shape}") - if np.any(axes_skew_np < 0.0) or np.any(axes_skew_np > 1.0): - gs.raise_exception(f"axes_skew values should be between 0.0 and 1.0, got: {axes_skew}") + self._validate_cross_axis_coupling(self.acc_cross_axis_coupling) + self._validate_cross_axis_coupling(self.gyro_cross_axis_coupling) + + def _validate_cross_axis_coupling(self, cross_axis_coupling): + cross_axis_coupling_np = np.array(cross_axis_coupling) + if cross_axis_coupling_np.shape not in ((), (3,), (3, 3)): + gs.raise_exception( + f"cross_axis_coupling shape should be (), (3,), or (3, 3), got: {cross_axis_coupling_np.shape}" + ) + if np.any(cross_axis_coupling_np < 0.0) or np.any(cross_axis_coupling_np > 1.0): + gs.raise_exception(f"cross_axis_coupling values should be between 0.0 and 1.0, got: {cross_axis_coupling}") class Raycaster(RigidSensorOptionsMixin, SensorOptions): diff --git a/tests/test_sensors.py b/tests/test_sensors.py index 8691c980a..9c6b4dec8 100644 --- a/tests/test_sensors.py +++ b/tests/test_sensors.py @@ -3,6 +3,7 @@ import torch import genesis as gs +import genesis.utils.geom as gu from .utils import assert_allclose, assert_array_equal @@ -42,11 +43,9 @@ def test_imu_sensor(show_viewer, tol, n_envs): ), ) - imu_biased = scene.add_sensor( + imu = scene.add_sensor( gs.sensors.IMU( entity_idx=box.idx, - acc_bias=BIAS, - gyro_bias=BIAS, ) ) imu_delayed = scene.add_sensor( @@ -58,8 +57,8 @@ def test_imu_sensor(show_viewer, tol, n_envs): 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_cross_axis_coupling=0.01, + gyro_cross_axis_coupling=(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), @@ -69,12 +68,6 @@ def test_imu_sensor(show_viewer, tol, n_envs): interpolate=True, ) ) - imu_skewed = scene.add_sensor( - gs.sensors.IMU( - entity_idx=box.idx, - acc_axes_skew=(0.0, 0.0, 1.0), - ) - ) scene.build(n_envs=n_envs) @@ -84,21 +77,18 @@ def test_imu_sensor(show_viewer, tol, n_envs): # 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_biased.read().lin_acc, expand_batch_dim(BIAS, n_envs), tol=tol) - assert_allclose(imu_biased.read().ang_vel, expand_batch_dim(BIAS, n_envs), tol=tol) - assert_allclose(imu_delayed.read().lin_acc, 0.0, tol=tol) - assert_allclose(imu_delayed.read().ang_vel, 0.0, tol=tol) + 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.1, 0.1, 0.1]]) + 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_acc_noise([0.01, 0.01, 0.01]) - imu_noisy.set_gyro_noise([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_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): @@ -112,42 +102,60 @@ def test_imu_sensor(show_viewer, tol, n_envs): 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_biased.read_ground_truth().lin_acc, imu_delayed.read_ground_truth().lin_acc) - assert_array_equal(imu_biased.read_ground_truth().ang_vel, imu_delayed.read_ground_truth().ang_vel) + 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_biased.read_ground_truth().ang_vel, 0.0, tol=tol) + 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_skewed.read().lin_acc, -GRAVITY, tol=5e-6) assert_allclose( - imu_biased.read().lin_acc, - expand_batch_dim((BIAS[0], BIAS[1], BIAS[2] - GRAVITY), n_envs), + imu.read().lin_acc, + expand_batch_dim((0.0, 0.0, -GRAVITY), n_envs), tol=5e-6, ) - assert_allclose(imu_biased.read().ang_vel, expand_batch_dim(BIAS, n_envs), tol=1e-5) + 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_cross_axis_coupling((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_cross_axis_coupling((0.0, 0.0, 0.0)) scene.reset() - assert_allclose(imu_biased.read().lin_acc, 0.0, tol=gs.EPS) # biased, but cache hasn't been updated yet + 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_biased.read().lin_acc, expand_batch_dim(BIAS, n_envs), tol=tol) + assert_allclose(imu.read().lin_acc, expand_batch_dim(BIAS, n_envs), tol=tol) @pytest.mark.required