diff --git a/genesis/engine/entities/rigid_entity/rigid_entity.py b/genesis/engine/entities/rigid_entity/rigid_entity.py index 9379b4a2af..95c8f70a12 100644 --- a/genesis/engine/entities/rigid_entity/rigid_entity.py +++ b/genesis/engine/entities/rigid_entity/rigid_entity.py @@ -1730,27 +1730,6 @@ def get_links_ang(self, links_idx_local=None, envs_idx=None, *, unsafe=False): links_idx = self._get_idx(links_idx_local, self.n_links, self._link_start, unsafe=True) return self._solver.get_links_ang(links_idx, envs_idx, unsafe=unsafe) - @gs.assert_built - def get_links_accelerometer_data(self, links_idx_local=None, envs_idx=None, *, imu=False, unsafe=False): - """ - Returns the accelerometer data that would be measured by a IMU rigidly attached to the specified entity's links, - i.e. the true linear acceleration of the links expressed at their respective origin in local frame coordinates. - - Parameters - ---------- - links_idx_local : None | array_like - The indices of the links. Defaults to None. - envs_idx : None | array_like, optional - The indices of the environments. If None, all environments will be considered. Defaults to None. - - Returns - ------- - acc : torch.Tensor, shape (n_links, 3) or (n_envs, n_links, 3) - The accelerometer data of IMUs rigidly attached of the specified entity's links. - """ - links_idx = self._get_idx(links_idx_local, self.n_links, self._link_start, unsafe=True) - return self._solver.get_links_acc(links_idx, envs_idx, mimick_imu=True, unsafe=unsafe) - @gs.assert_built def get_links_acc(self, links_idx_local=None, envs_idx=None, *, unsafe=False): """ @@ -1770,7 +1749,7 @@ def get_links_acc(self, links_idx_local=None, envs_idx=None, *, unsafe=False): The linear classical acceleration of the specified entity's links. """ links_idx = self._get_idx(links_idx_local, self.n_links, self._link_start, unsafe=True) - return self._solver.get_links_acc(links_idx, envs_idx, mimick_imu=False, unsafe=unsafe) + return self._solver.get_links_acc(links_idx, envs_idx, unsafe=unsafe) @gs.assert_built def get_links_acc_ang(self, links_idx_local=None, envs_idx=None, *, unsafe=False): diff --git a/genesis/engine/scene.py b/genesis/engine/scene.py index c1cb4fce29..b9387e598c 100644 --- a/genesis/engine/scene.py +++ b/genesis/engine/scene.py @@ -2,6 +2,7 @@ import pickle import sys import time +from typing import TYPE_CHECKING import numpy as np import torch @@ -26,7 +27,6 @@ PBDOptions, ProfilingOptions, RigidOptions, - SensorOptions, SFOptions, SimOptions, SPHOptions, @@ -43,6 +43,9 @@ from genesis.vis import Visualizer from genesis.utils.warnings import warn_once +if TYPE_CHECKING: + from genesis.sensors.base_sensor import SensorOptions + @gs.assert_initialized class Scene(RBC): @@ -516,7 +519,7 @@ def add_light( gs.raise_exception("Adding lights is only supported by 'RayTracer' and 'BatchRenderer'.") @gs.assert_unbuilt - def add_sensor(self, sensor_options: SensorOptions): + def add_sensor(self, sensor_options: "SensorOptions"): return self._sim._sensor_manager.create_sensor(sensor_options) @gs.assert_unbuilt diff --git a/genesis/engine/solvers/base_solver.py b/genesis/engine/solvers/base_solver.py index b3c494fe79..8a6bfceee5 100644 --- a/genesis/engine/solvers/base_solver.py +++ b/genesis/engine/solvers/base_solver.py @@ -63,6 +63,10 @@ def _kernel_set_gravity(self, gravity: ti.types.ndarray(), envs_idx: ti.types.nd for j in ti.static(range(3)): self._gravity[envs_idx[i_b_]][j] = gravity[i_b_, j] + def get_gravity(self, envs_idx=None, *, unsafe=False): + tensor = ti_field_to_torch(self._gravity, envs_idx, transpose=True, unsafe=unsafe) + return tensor.squeeze(0) if self.n_envs == 0 else tensor + def dump_ckpt_to_numpy(self) -> dict[str, np.ndarray]: arrays: dict[str, np.ndarray] = {} diff --git a/genesis/engine/solvers/rigid/rigid_solver_decomp.py b/genesis/engine/solvers/rigid/rigid_solver_decomp.py index 750d742cfa..73b0adb5f0 100644 --- a/genesis/engine/solvers/rigid/rigid_solver_decomp.py +++ b/genesis/engine/solvers/rigid/rigid_solver_decomp.py @@ -2043,18 +2043,16 @@ def get_links_ang(self, links_idx=None, envs_idx=None, *, unsafe=False): tensor = ti_field_to_torch(self.links_state.cd_ang, envs_idx, links_idx, transpose=True, unsafe=unsafe) return tensor.squeeze(0) if self.n_envs == 0 else tensor - def get_links_acc(self, links_idx=None, envs_idx=None, *, mimick_imu=False, unsafe=False): + def get_links_acc(self, links_idx=None, envs_idx=None, *, unsafe=False): _tensor, links_idx, envs_idx = self._sanitize_2D_io_variables( None, links_idx, self.n_links, 3, envs_idx, idx_name="links_idx", unsafe=unsafe ) tensor = _tensor.unsqueeze(0) if self.n_envs == 0 else _tensor kernel_get_links_acc( - mimick_imu, tensor, links_idx, envs_idx, self.links_state, - self._rigid_global_info, self._static_rigid_sim_config, ) return _tensor @@ -6554,12 +6552,10 @@ def kernel_get_links_vel( @ti.kernel def kernel_get_links_acc( - mimick_imu: ti.i32, tensor: ti.types.ndarray(), links_idx: ti.types.ndarray(), envs_idx: ti.types.ndarray(), links_state: array_class.LinksState, - rigid_global_info: array_class.RigidGlobalInfo, static_rigid_sim_config: ti.template(), ): ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.PARTIAL) @@ -6577,14 +6573,6 @@ def kernel_get_links_acc( vel = links_state.cd_vel[i_l, i_b] + ang.cross(cpos) acc_classic_lin = acc_lin + ang.cross(vel) - # Mimick IMU accelerometer signal if requested - if mimick_imu: - # Subtract gravity - acc_classic_lin -= rigid_global_info.gravity[i_b] - - # Move the resulting linear acceleration in local links frame - acc_classic_lin = gu.ti_inv_transform_by_quat(acc_classic_lin, links_state.quat[i_l, i_b]) - for i in ti.static(range(3)): tensor[i_b_, i_l_, i] = acc_classic_lin[i] diff --git a/genesis/options/__init__.py b/genesis/options/__init__.py index b94a59f861..c362e970f9 100644 --- a/genesis/options/__init__.py +++ b/genesis/options/__init__.py @@ -2,6 +2,5 @@ from .solvers import * from .vis import * from .profiling import ProfilingOptions -from .sensors import SensorOptions __all__ = ["ProfilingOptions"] diff --git a/genesis/options/sensors.py b/genesis/options/sensors.py deleted file mode 100644 index 2b41e0d50e..0000000000 --- a/genesis/options/sensors.py +++ /dev/null @@ -1,16 +0,0 @@ -from genesis.options import Options - - -class SensorOptions(Options): - """ - Base class for all sensor options. - Each sensor should have their own options class that inherits from this class. - The options class should be registered with the SensorManager using the @register_sensor decorator. - - Parameters - ---------- - read_delay : float - The delay in seconds before the sensor data is read. - """ - - read_delay: float = 0.0 diff --git a/genesis/sensors/__init__.py b/genesis/sensors/__init__.py index 4be0edec04..14ab868dec 100644 --- a/genesis/sensors/__init__.py +++ b/genesis/sensors/__init__.py @@ -1,4 +1,5 @@ from .base_sensor import Sensor +from .imu import IMU from .tactile import RigidContactSensor, RigidContactForceSensor, RigidContactForceGridSensor from .data_recorder import SensorDataRecorder, RecordingOptions from .data_handlers import ( diff --git a/genesis/sensors/base_sensor.py b/genesis/sensors/base_sensor.py index 7274774949..301994a11c 100644 --- a/genesis/sensors/base_sensor.py +++ b/genesis/sensors/base_sensor.py @@ -1,35 +1,89 @@ -from typing import TYPE_CHECKING, Any, List, Type +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, List import numpy as np import taichi as ti import torch import genesis as gs +from genesis.options import Options 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 +class SensorOptions(Options): + """ + Base class for all sensor options. + Each sensor should have their own options class that inherits from this class. + The options class should be registered with the SensorManager using the @register_sensor decorator. + + Parameters + ---------- + read_delay : float + The delay in seconds before the sensor data is read. + """ + + read_delay: float = 0.0 + + def validate(self, scene): + """ + Validate the sensor options values before the sensor is added to the scene. + """ + read_delay_hz = self.read_delay / scene._sim.dt + if not np.isclose(read_delay_hz, round(read_delay_hz), atol=1e-6): + gs.logger.warn( + f"Read delay should be a multiple of the simulation time step. Got {self.read_delay}" + f" and {scene._sim.dt}. Actual read delay will be {1/round(read_delay_hz)}." + ) + + +@dataclass +class SharedSensorMetadata: + """ + Shared metadata between all sensors of the same class. + """ + + cache_sizes: list[int] = field(default_factory=list) + read_delay_steps: list[int] = field(default_factory=list) + + @ti.data_oriented class Sensor(RBC): """ Base class for all types of sensors. + + NOTE: The Sensor system is designed to be performant. All sensors of the same type are updated at once and stored + in a cache in SensorManager. Cache size is inferred from the return format and cache length of each sensor. + `read()` and `read_ground_truth()`, the public-facing methods of every Sensor, automatically handles indexing into + the shared cache to return the correct data. """ 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 + self._shared_metadata: SharedSensorMetadata = sensor_manager._sensors_metadata[type(self)] + + self._read_delay_steps = round(self._options.read_delay / self._manager._sim.dt) + self._shared_metadata.read_delay_steps.append(self._read_delay_steps) - # 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 + return_format = self._get_return_format() + return_shapes = return_format.values() if isinstance(return_format, dict) else (return_format,) + tensor_size = 0 + for shape in return_shapes: + data_size = np.prod(shape) + self._shape_indices.append((tensor_size, tensor_size + data_size)) + tensor_size += data_size + + self._cache_size = self._get_cache_length() * tensor_size + self._shared_metadata.cache_sizes.append(self._cache_size) + + self._cache_idx: int = -1 # initialized by SensorManager during build # =============================== implementable methods =============================== @@ -71,12 +125,19 @@ def _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" + cls, + shared_metadata: dict[str, Any], + shared_ground_truth_cache: torch.Tensor, + shared_cache: torch.Tensor, + buffered_data: "TensorRingBuffer", ): """ Update the shared sensor cache for all sensors of this class using metadata in SensorManager. + + The information in shared_cache should be the final measured sensor data after all noise and post-processing. + NOTE: The implementation should include applying the delay using the `_apply_delay_to_shared_cache()` method. """ - raise NotImplementedError("Sensors must implement `update_shared_cache()`.") + raise NotImplementedError("Sensors must implement `update_shared_cache_with_noise()`.") @classmethod def _get_cache_dtype(cls) -> torch.dtype: @@ -92,19 +153,35 @@ 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) + return self._get_formatted_data(self._manager.get_cloned_from_cache(self), 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) + return self._get_formatted_data(self._manager.get_cloned_from_cache(self, is_ground_truth=True), envs_idx) + + @classmethod + def _apply_delay_to_shared_cache( + self, shared_metadata: SharedSensorMetadata, shared_cache: torch.Tensor, buffered_data: "TensorRingBuffer" + ): + """ + Applies the read delay to the shared cache tensor by copying the buffered data at the appropriate index. + """ + idx = 0 + for tensor_size, read_delay_step in zip(shared_metadata.cache_sizes, shared_metadata.read_delay_steps): + shared_cache[:, idx : idx + tensor_size] = buffered_data.at(read_delay_step)[:, idx : idx + tensor_size] + idx += tensor_size 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. + """ + Formats the flattened cache tensor into a dict of tensors using the format specified in `_get_return_format()`. + + 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 diff --git a/genesis/sensors/imu.py b/genesis/sensors/imu.py new file mode 100644 index 0000000000..7ff5690ad9 --- /dev/null +++ b/genesis/sensors/imu.py @@ -0,0 +1,169 @@ +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any + +import taichi as ti +import torch + +import genesis as gs +from genesis.engine.entities import RigidEntity +from genesis.engine.solvers import RigidSolver +from genesis.utils.geom import ( + euler_to_quat, + inv_transform_by_trans_quat, + transform_quat_by_quat, +) + +from .base_sensor import Sensor, SensorOptions, SharedSensorMetadata +from .sensor_manager import register_sensor + +if TYPE_CHECKING: + from genesis.utils.ring_buffer import TensorRingBuffer + + +class IMUOptions(SensorOptions): + """ + IMU sensor returns the linear acceleration (accelerometer) and angular velocity (gyroscope) + of the associated entity link. + + Note + ---- + Accelerometers return the so-called classical linear acceleration in local frame minus gravity. + + Parameters + ---------- + entity_idx : int + The global entity index of the RigidEntity to which this IMU sensor is attached. + link_idx_local : int, optional + The local index of the RigidLink of the RigidEntity to which this IMU sensor is attached. + pos_offset : tuple[float, float, float] + The offset of the IMU sensor from the RigidLink. + euler_offset : tuple[float, float, float] + The offset of the IMU sensor from the RigidLink in euler angles. + accelerometer_bias : tuple[float, float, float] + The bias of the accelerometer. + gyroscope_bias : tuple[float, float, float] + The bias of the gyroscope. + """ + + entity_idx: int + link_idx_local: int = 0 + pos_offset: tuple[float, float, float] = (0.0, 0.0, 0.0) + euler_offset: tuple[float, float, float] = (0.0, 0.0, 0.0) + + accelerometer_bias: tuple[float, float, float] = (0.0, 0.0, 0.0) + gyroscope_bias: tuple[float, float, float] = (0.0, 0.0, 0.0) + + def validate(self, scene): + assert self.entity_idx >= 0 and self.entity_idx < len(scene.entities), "Invalid RigidEntity index." + entity = scene.entities[self.entity_idx] + assert isinstance(entity, RigidEntity), "Entity at given index is not a RigidEntity." + assert ( + self.link_idx_local >= 0 and self.link_idx_local < scene.entities[self.entity_idx].n_links + ), "Invalid RigidLink index." + + +@dataclass +class IMUSharedMetadata(SharedSensorMetadata): + """ + Shared metadata between all IMU sensors. + """ + + solver: RigidSolver | None = None + links_idx: list[int] = field(default_factory=list) + offsets_pos: torch.Tensor = torch.tensor([]) + offsets_quat: torch.Tensor = torch.tensor([]) + acc_bias: torch.Tensor = torch.tensor([]) + ang_bias: torch.Tensor = torch.tensor([]) + + +@register_sensor(IMUOptions, IMUSharedMetadata) +@ti.data_oriented +class IMU(Sensor): + + def build(self): + """ + Initialize all shared metadata needed to update all IMU sensors. + """ + if self._shared_metadata.solver is None: + self._shared_metadata.solver = self._manager._sim.rigid_solver + + self._shared_metadata.links_idx.append(self._options.entity_idx + self._options.link_idx_local) + self._shared_metadata.offsets_pos = torch.cat( + [self._shared_metadata.offsets_pos, torch.tensor([self._options.pos_offset], dtype=gs.tc_float)] + ) + + quat_tensor = torch.tensor(euler_to_quat(self._options.euler_offset), dtype=gs.tc_float).unsqueeze(0) + if self._shared_metadata.solver.n_envs > 0: + quat_tensor = quat_tensor.unsqueeze(0).expand((self._manager._sim._B, 1, 4)) + self._shared_metadata.offsets_quat = torch.cat([self._shared_metadata.offsets_quat, quat_tensor], dim=-2) + + self._shared_metadata.acc_bias = torch.cat( + [self._shared_metadata.acc_bias, torch.tensor([self._options.accelerometer_bias], dtype=gs.tc_float)] + ) + self._shared_metadata.ang_bias = torch.cat( + [self._shared_metadata.ang_bias, torch.tensor([self._options.gyroscope_bias], dtype=gs.tc_float)] + ) + + def _get_return_format(self) -> dict[str, tuple[int, ...]]: + return { + "lin_acc": (3,), + "ang_vel": (3,), + } + + def _get_cache_length(self) -> int: + return 1 + + @classmethod + def _update_shared_ground_truth_cache( + cls, shared_metadata: IMUSharedMetadata, shared_ground_truth_cache: torch.Tensor + ): + """ + Update the current ground truth values for all IMU sensors. + """ + gravity = shared_metadata.solver.get_gravity() + 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) + + offset_quats = transform_quat_by_quat(quats, shared_metadata.offsets_quat) + + # 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) + + *batch_size, n_imus, _ = local_acc.shape + local_acc = local_acc - gravity.unsqueeze(-2).expand((*batch_size, n_imus, -1)) + + # cache shape: (B, n_imus * 6) + strided_ground_truth_cache = shared_ground_truth_cache.reshape((*batch_size, n_imus, 2, 3)) + strided_ground_truth_cache[..., 0, :].copy_(local_acc) + strided_ground_truth_cache[..., 1, :].copy_(local_ang) + + @classmethod + def _update_shared_cache( + cls, + shared_metadata: dict[str, Any], + shared_ground_truth_cache: torch.Tensor, + shared_cache: torch.Tensor, + buffered_data: "TensorRingBuffer", + ): + """ + Update the current measured sensor data for all IMU sensors. + + Note + ---- + `buffered_data` contains the history of ground truth cache, and noise/bias is only applied to the current + sensor readout `shared_cache`, not the whole buffer. + """ + buffered_data.append(shared_ground_truth_cache) + cls._apply_delay_to_shared_cache(shared_metadata, shared_cache, buffered_data) + + # add bias to the shared_cache + *batch_size, n_imus, _ = shared_metadata.offsets_quat.shape + strided_shared_cache = shared_cache.reshape((*batch_size, n_imus, 2, 3)) + strided_shared_cache[..., 0, :] += shared_metadata.acc_bias + strided_shared_cache[..., 1, :] += shared_metadata.ang_bias + + @classmethod + def _get_cache_dtype(cls) -> torch.dtype: + return gs.tc_float diff --git a/genesis/sensors/sensor_manager.py b/genesis/sensors/sensor_manager.py index 71489676f2..4556a702d3 100644 --- a/genesis/sensors/sensor_manager.py +++ b/genesis/sensors/sensor_manager.py @@ -1,70 +1,57 @@ -from typing import TYPE_CHECKING, Any, Type +from typing import TYPE_CHECKING, Type -import numpy as np import torch -import genesis as gs from genesis.utils.ring_buffer import TensorRingBuffer if TYPE_CHECKING: from genesis.options.sensors import SensorOptions - from .base_sensor import Sensor + from .base_sensor import Sensor, SharedSensorMetadata class SensorManager: - SENSOR_TYPES_MAP: dict[Type["SensorOptions"], Type["Sensor"]] = {} + SENSOR_TYPES_MAP: dict[Type["SensorOptions"], tuple[Type["Sensor"], Type["SharedSensorMetadata"]]] = {} def __init__(self, sim): self._sim = sim self._sensors_by_type: dict[Type["Sensor"], list["Sensor"]] = {} - self._sensors_metadata: dict[Type["Sensor"], dict[str, Any]] = {} + self._sensors_metadata: dict[Type["Sensor"], SharedSensorMetadata | None] = {} self._ground_truth_cache: dict[Type[torch.dtype], torch.Tensor] = {} - self._cache: dict[Type[torch.dtype], TensorRingBuffer] = {} + self._cache: dict[Type[torch.dtype], torch.Tensor] = {} + self._buffered_data: dict[Type[torch.dtype], TensorRingBuffer] = {} self._cache_slices_by_type: dict[Type["Sensor"], slice] = {} - self._last_ground_truth_cache_cloned_step: dict[Type[torch.dtype], int] = {} - self._cloned_ground_truth_cache: dict[Type[torch.dtype], torch.Tensor] = {} + self._last_cache_cloned_step: dict[tuple[bool, Type[torch.dtype]], int] = {} + self._cloned_cache: dict[tuple[bool, Type[torch.dtype]], torch.Tensor] = {} def create_sensor(self, sensor_options: "SensorOptions"): - sensor_cls = SensorManager.SENSOR_TYPES_MAP[type(sensor_options)] + sensor_options.validate(self._sim.scene) + sensor_cls, metadata_cls = SensorManager.SENSOR_TYPES_MAP[type(sensor_options)] self._sensors_by_type.setdefault(sensor_cls, []) + if sensor_cls not in self._sensors_metadata: + self._sensors_metadata[sensor_cls] = metadata_cls() sensor = sensor_cls(sensor_options, len(self._sensors_by_type[sensor_cls]), self) self._sensors_by_type[sensor_cls].append(sensor) return sensor def build(self): - max_cache_buf_len = 0 + max_buffer_len = 0 cache_size_per_dtype = {} for sensor_cls, sensors in self._sensors_by_type.items(): - self._sensors_metadata[sensor_cls] = {} dtype = sensor_cls._get_cache_dtype() + + for is_ground_truth in [False, True]: + self._last_cache_cloned_step.setdefault((is_ground_truth, dtype), -1) + self._cloned_cache.setdefault((is_ground_truth, dtype), torch.zeros(0, dtype=dtype)) + cache_size_per_dtype.setdefault(dtype, 0) cls_cache_start_idx = cache_size_per_dtype[dtype] for sensor in sensors: - return_format = sensor._get_return_format() - return_shapes = return_format.values() if isinstance(return_format, dict) else (return_format,) - - tensor_size = 0 - for shape in return_shapes: - data_size = np.prod(shape) - sensor._shape_indices.append((tensor_size, tensor_size + data_size)) - tensor_size += data_size - - delay_steps_float = sensor._options.read_delay / self._sim.dt - sensor._read_delay_steps = round(delay_steps_float) - if not np.isclose(delay_steps_float, sensor._read_delay_steps, atol=1e-6): - gs.logger.warn( - f"Read delay should be a multiple of the simulation time step. Got {sensor._options.read_delay}" - f" and {self._sim.dt}. Actual read delay will be {1/sensor._read_delay_steps}." - ) - - sensor._cache_size = sensor._get_cache_length() * tensor_size sensor._cache_idx = cache_size_per_dtype[dtype] cache_size_per_dtype[dtype] += sensor._cache_size - - max_cache_buf_len = max(max_cache_buf_len, sensor._read_delay_steps + 1) + max_buffer_len = max(max_buffer_len, sensor._read_delay_steps + 1) cls_cache_end_idx = cache_size_per_dtype[dtype] self._cache_slices_by_type[sensor_cls] = slice(cls_cache_start_idx, cls_cache_end_idx) @@ -72,13 +59,12 @@ def build(self): for dtype in cache_size_per_dtype.keys(): cache_shape = (self._sim._B, cache_size_per_dtype[dtype]) self._ground_truth_cache[dtype] = torch.zeros(cache_shape, dtype=dtype) - self._cache[dtype] = TensorRingBuffer(max_cache_buf_len, cache_shape, dtype=dtype) + self._cache[dtype] = torch.zeros(cache_shape, dtype=dtype) + self._buffered_data[dtype] = TensorRingBuffer(max_buffer_len, cache_shape, dtype=dtype) for sensor_cls, sensors in self._sensors_by_type.items(): dtype = sensor_cls._get_cache_dtype() for sensor in sensors: - sensor._shared_metadata = self._sensors_metadata[sensor_cls] - sensor._cache = self._cache[dtype][:, sensor._cache_idx : sensor._cache_idx + sensor._cache_size] sensor.build() def step(self): @@ -92,23 +78,28 @@ def step(self): self._sensors_metadata[sensor_cls], self._ground_truth_cache[dtype][cache_slice], self._cache[dtype][cache_slice], + self._buffered_data[dtype][cache_slice], ) - def get_cloned_from_ground_truth_cache(self, sensor: "Sensor") -> torch.Tensor: + def get_cloned_from_cache(self, sensor: "Sensor", is_ground_truth: bool = False) -> torch.Tensor: dtype = sensor._get_cache_dtype() - if self._last_ground_truth_cache_cloned_step[dtype] != self._sim.cur_step_global: - self._last_ground_truth_cache_cloned_step[dtype] = self._sim.cur_step_global - self._cloned_ground_truth_cache[dtype] = self._ground_truth_cache[dtype].clone() - return self._cloned_ground_truth_cache[dtype][:, sensor._cache_idx : sensor._cache_idx + sensor._cache_size] + key = (is_ground_truth, dtype) + if self._last_cache_cloned_step[key] != self._sim.cur_step_global: + self._last_cache_cloned_step[key] = self._sim.cur_step_global + if is_ground_truth: + self._cloned_cache[key] = self._ground_truth_cache[dtype].clone() + else: + self._cloned_cache[key] = self._cache[dtype].clone() + return self._cloned_cache[key][:, sensor._cache_idx : sensor._cache_idx + sensor._cache_size] @property def sensors(self): return tuple([sensor for sensor_list in self._sensors_by_type.values() for sensor in sensor_list]) -def register_sensor(sensor_cls: Type["Sensor"]): - def _impl(options_cls: Type["SensorOptions"]): - SensorManager.SENSOR_TYPES_MAP[options_cls] = sensor_cls - return options_cls +def register_sensor(options_cls: Type["SensorOptions"], metadata_cls: Type["SharedSensorMetadata"]): + def _impl(sensor_cls: Type["Sensor"]): + SensorManager.SENSOR_TYPES_MAP[options_cls] = sensor_cls, metadata_cls + return sensor_cls return _impl diff --git a/genesis/utils/ring_buffer.py b/genesis/utils/ring_buffer.py index 4abb06f52f..36f73abef6 100644 --- a/genesis/utils/ring_buffer.py +++ b/genesis/utils/ring_buffer.py @@ -6,6 +6,23 @@ class TensorRingBuffer: + """ + A helper class for storing a buffer of `torch.Tensor`s without allocating new tensors. + + Parameters + ---------- + N : int + The number of tensors to store. + shape : tuple[int, ...] + The shape of the tensors to store. + dtype : torch.dtype + The dtype of the tensors to store. + buffer : torch.Tensor | None, optional + The buffer tensor where all the data is stored. If not provided, a new tensor is allocated. + idx_ptr : int | ctypes.c_int, optional + The index pointer to the current position in the ring buffer. If not provided, it is initialized to 0. + """ + def __init__( self, N: int, @@ -26,22 +43,40 @@ def __init__( self._idx_ptr = idx_ptr def append(self, tensor: torch.Tensor): + """ + Copy the tensor into the next position of the ring buffer, and advance the index pointer. + + Parameters + ---------- + tensor : torch.Tensor + The tensor to copy into the ring buffer. + """ self.buffer[self._idx_ptr.value].copy_(tensor) self._idx_ptr.value = (self._idx_ptr.value + 1) % self.N - def get(self, idx: int, clone: bool = True): + def at(self, idx: int) -> torch.Tensor: + """ + Get a view of the tensor at the given index. + + Parameters + ---------- + idx : int + Index of the element to get, where 0 is the latest element, 1 is the second latest, etc. + """ + return self.buffer[(self._idx_ptr.value - idx) % self.N] + + def get(self, idx: int) -> torch.Tensor: """ + Get a clone of the tensor at the given index. + Parameters ---------- idx : int Index of the element to get, where 0 is the latest element, 1 is the second latest, etc. - clone : bool - Whether to clone the tensor. """ - tensor = self.buffer[(self._idx_ptr.value - idx) % self.N] - return tensor.clone() if clone else tensor + return self.at(idx).clone() - def clone(self): + def clone(self) -> "TensorRingBuffer": return TensorRingBuffer( self.N, self.buffer.shape[1:], @@ -50,7 +85,7 @@ def clone(self): idx_ptr=self._idx_ptr, ) - def __getitem__(self, key: int | slice | tuple): + def __getitem__(self, key: int | slice | tuple) -> "TensorRingBuffer": """ Enable slicing of the tensor ring buffer. diff --git a/tests/test_rigid_physics.py b/tests/test_rigid_physics.py index 662af3a545..b1af439306 100644 --- a/tests/test_rigid_physics.py +++ b/tests/test_rigid_physics.py @@ -711,14 +711,10 @@ def test_pendulum_links_acc(gs_sim, tol): # Linear true acceleration: # * acc_classical_lin_y = sin(theta) * g (tangential angular acceleration effect) # * acc_classical_lin_z = - theta_dot ** 2 (radial centripedal effect) - acc_classical_lin_world = tensor_to_array(gs_sim.rigid_solver.get_links_acc(mimick_imu=False)) + acc_classical_lin_world = tensor_to_array(gs_sim.rigid_solver.get_links_acc()) assert_allclose(acc_classical_lin_world[0], 0, tol=tol) acc_classical_lin_local = R @ acc_classical_lin_world[2] assert_allclose(acc_classical_lin_local, np.array([0.0, np.sin(theta) * g, -(theta_dot**2)]), tol=tol) - # IMU accelerometer data: - # * acc_classical_lin_z = - theta_dot ** 2 - cos(theta) * g - acc_imu = gs_sim.rigid_solver.get_links_acc(mimick_imu=True)[2] - assert_allclose(acc_imu, np.array([0.0, 0.0, -(theta_dot**2) - np.cos(theta) * g]), tol=tol) # Hold the pendulum straight using PD controller and check again pendulum.set_dofs_kp([4000.0]) @@ -726,7 +722,7 @@ def test_pendulum_links_acc(gs_sim, tol): pendulum.control_dofs_position([0.5 * np.pi]) for _ in range(400): gs_sim.scene.step() - acc_classical_lin_world = gs_sim.rigid_solver.get_links_acc(mimick_imu=False) + acc_classical_lin_world = gs_sim.rigid_solver.get_links_acc() assert_allclose(acc_classical_lin_world, 0, tol=tol) @@ -785,7 +781,7 @@ def test_double_pendulum_links_acc(gs_sim, tol): ) # Linear true acceleration - acc_classical_lin_world = tensor_to_array(gs_sim.rigid_solver.get_links_acc(mimick_imu=False)[[0, 2, 4]]) + acc_classical_lin_world = tensor_to_array(gs_sim.rigid_solver.get_links_acc()[[0, 2, 4]]) assert_allclose(acc_classical_lin_world[0], 0, tol=tol) acc_classical_lin_local = np.matmul(np.moveaxis(R, 2, 0), acc_classical_lin_world[1:, :, None])[..., 0] assert_allclose(acc_classical_lin_local[0], np.array([0.0, -theta_ddot[0], -theta_dot[0] ** 2]), tol=tol) @@ -801,7 +797,7 @@ def test_double_pendulum_links_acc(gs_sim, tol): robot.control_dofs_position([0.5 * np.pi, 0.0]) for _ in range(900): gs_sim.scene.step() - acc_classical_lin_world = gs_sim.rigid_solver.get_links_acc(mimick_imu=False) + acc_classical_lin_world = gs_sim.rigid_solver.get_links_acc() assert_allclose(acc_classical_lin_world, 0, tol=tol) diff --git a/tests/test_sensors.py b/tests/test_sensors.py new file mode 100644 index 0000000000..c3a5d3082b --- /dev/null +++ b/tests/test_sensors.py @@ -0,0 +1,74 @@ +import numpy as np +import torch + +import genesis as gs +from genesis.sensors.imu import IMUOptions + +from .utils import assert_allclose, assert_array_equal + + +def test_imu_sensor(show_viewer): + """Test if the IMU sensor returns the correct data.""" + GRAVITY = -10.0 + DT = 1e-2 + BIAS = (0.1, 0.2, 0.3) + + scene = gs.Scene( + sim_options=gs.options.SimOptions( + dt=DT, + substeps=1, + gravity=(0.0, 0.0, GRAVITY), + ), + show_viewer=show_viewer, + show_FPS=False, + ) + + 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_biased = scene.add_sensor(IMUOptions(entity_idx=box.idx, accelerometer_bias=BIAS, gyroscope_bias=BIAS)) + imu_delayed = scene.add_sensor(IMUOptions(entity_idx=box.idx, read_delay=DT * 2)) + + scene.build() + + # 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_biased.read()["lin_acc"], BIAS, tol=1e-7) + assert_allclose(imu_biased.read()["ang_vel"], BIAS, tol=1e-7) + assert_allclose(imu_delayed.read()["lin_acc"], 0.0, tol=1e-7) + assert_allclose(imu_delayed.read()["ang_vel"], 0.0, tol=1e-7) + + # shift COM to induce angular velocity + box.set_COM_shift(torch.tensor([[0.1, 0.1, 0.1]])) + + # box collides with ground + for _ in range(30): + 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"]) + + 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=1e-7) + + 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) + + box.set_COM_shift(torch.tensor([[0.0, 0.0, 0.0]])) + + # box is stationary on ground + for _ in range(80): + scene.step() + + assert_allclose(imu_biased.read()["lin_acc"], torch.tensor([BIAS[0], BIAS[1], BIAS[2] - GRAVITY]), tol=1e-7) + assert_allclose(imu_biased.read()["ang_vel"], BIAS, tol=1e-5)