from dataclasses import dataclass
from typing import TYPE_CHECKING, NamedTuple, Type
import numpy as np
import quadrants as qd
import torch
import genesis as gs
from genesis.options.sensors import IMU as IMUOptions
from genesis.options.sensors import CrossCouplingAxisType
from genesis.utils.geom import inv_transform_by_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 SimpleSensor, RigidSensorMetadataMixin, RigidSensorMixin, Sensor, SimpleSensorMetadata
if TYPE_CHECKING:
from genesis.ext.pyrender.mesh import Mesh
from genesis.utils.ring_buffer import TensorRingBuffer
from genesis.vis.rasterizer_context import RasterizerContext
from .sensor_manager import SensorManager
def _get_cross_axis_coupling_to_alignment_matrix(
input: CrossCouplingAxisType, out: torch.Tensor | None = None
) -> torch.Tensor:
"""
Convert the alignment input to a matrix. Modifies in place if provided, else allocate a new matrix.
"""
if out is None:
out = torch.eye(3, dtype=gs.tc_float, device=gs.device)
if isinstance(input, float):
# set off-diagonal elements to the scalar value
torch.diagonal(out)[:] = input
out.fill_diagonal_(1.0)
elif isinstance(input, torch.Tensor):
out.copy_(input)
else:
np_input = np.array(input)
if np_input.shape == (3,):
# set off-diagonal elements to the vector values
out[1, 0] = np_input[0]
out[2, 0] = np_input[0]
out[0, 1] = np_input[1]
out[2, 1] = np_input[1]
out[0, 2] = np_input[2]
out[1, 2] = np_input[2]
elif np_input.shape == (3, 3):
out.copy_(torch.tensor(np_input, dtype=gs.tc_float, device=gs.device))
return out
@dataclass
class IMUSharedMetadata(RigidSensorMetadataMixin, SimpleSensorMetadata):
"""
Shared metadata between all IMU sensors.
"""
alignment_rot_matrix: torch.Tensor = make_tensor_field((0, 0, 3, 3))
magnetic_field_vector: torch.Tensor = make_tensor_field((0, 0, 3)) # added another dimension to match data layout
acc_indices: torch.Tensor = make_tensor_field((0, 0), dtype_factory=lambda: gs.tc_int)
gyro_indices: torch.Tensor = make_tensor_field((0, 0), dtype_factory=lambda: gs.tc_int)
mag_indices: torch.Tensor = make_tensor_field((0, 0), dtype_factory=lambda: gs.tc_int)
class IMUReturnType(NamedTuple):
lin_acc: torch.Tensor
ang_vel: torch.Tensor
mag: torch.Tensor # added magnetometer to complete 9-axis IMU
[docs]class IMUSensor(RigidSensorMixin[IMUSharedMetadata], SimpleSensor[IMUOptions, None, IMUSharedMetadata, IMUReturnType]):
def __init__(
self,
options: IMUOptions,
idx: int,
shared_context,
shared_metadata: IMUSharedMetadata,
manager: "SensorManager",
):
# FIXME: Resolution should be made private in mixin, so that it cannot be set by the user directly.
options.resolution = options.acc_resolution + options.gyro_resolution + options.mag_resolution
options.bias = options.acc_bias + options.gyro_bias + options.mag_bias
options.random_walk = options.acc_random_walk + options.gyro_random_walk + options.mag_random_walk
options.noise = options.acc_noise + options.gyro_noise + options.mag_noise
super().__init__(options, idx, shared_context, shared_metadata, manager)
self.debug_objects: list["Mesh"] = []
self.quat_offset: torch.Tensor
self.pos_offset: torch.Tensor
[docs] @gs.assert_built
def set_acc_cross_axis_coupling(self, cross_axis_coupling: CrossCouplingAxisType, envs_idx=None):
envs_idx = self._sanitize_envs_idx(envs_idx)
rot_matrix = _get_cross_axis_coupling_to_alignment_matrix(cross_axis_coupling)
self._shared_metadata.alignment_rot_matrix[envs_idx, self._idx * 3, :, :] = rot_matrix
[docs] @gs.assert_built
def set_gyro_cross_axis_coupling(self, cross_axis_coupling: CrossCouplingAxisType, envs_idx=None):
envs_idx = self._sanitize_envs_idx(envs_idx)
rot_matrix = _get_cross_axis_coupling_to_alignment_matrix(cross_axis_coupling)
self._shared_metadata.alignment_rot_matrix[envs_idx, self._idx * 3 + 1, :, :] = rot_matrix
[docs] @gs.assert_built
def set_mag_cross_axis_coupling(self, cross_axis_coupling: CrossCouplingAxisType, envs_idx=None):
envs_idx = self._sanitize_envs_idx(envs_idx)
rot_matrix = _get_cross_axis_coupling_to_alignment_matrix(cross_axis_coupling)
self._shared_metadata.alignment_rot_matrix[envs_idx, self._idx * 3 + 2, :, :] = rot_matrix
# ================================ internal methods ================================
[docs] def build(self):
"""
Initialize all shared metadata needed to update all IMU sensors.
"""
super().build()
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
@classmethod
def _update_raw_data(cls, shared_context: None, shared_metadata: IMUSharedMetadata, raw_data_T: torch.Tensor):
assert shared_metadata.solver is not None
# Extract acceleration and gravity in world frame.
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)
if acc.ndim == 2: # n_envs = 0
acc = acc[None]
ang = ang[None]
quats = quats[None]
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: # n_envs = 0
ang_acc = ang_acc[None]
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
# Subtract gravity then move to local frame. acc/ang shape: (B, n_imus, 3); local_mag is already (B, n_imus, 3)
# after the inverse transform, no reshape needed.
local_acc = inv_transform_by_quat(acc - gravity[..., None, :], offset_quats)
local_ang = inv_transform_by_quat(ang, offset_quats)
local_mag = inv_transform_by_quat(shared_metadata.magnetic_field_vector, offset_quats)
# Raw-data buffer layout: (n_imus * 9, B). View into (n_imus, 3, 3, *batch_size) for the per-channel writes.
*batch_size, n_imus, _ = local_acc.shape
strided_raw = raw_data_T.view(n_imus, 3, 3, *batch_size)
strided_raw[:, 0].copy_(local_acc.permute(1, 2, 0))
strided_raw[:, 1].copy_(local_ang.permute(1, 2, 0))
strided_raw[:, 2].copy_(local_mag.permute(1, 2, 0))
@classmethod
def _apply_transform(cls, shared_metadata: IMUSharedMetadata, data: torch.Tensor, timeline, *, is_measured: bool):
# Apply alignment rotation to the (lin_acc, ang_vel, mag) triplet. View the flat cache as a stack of 3-vectors
# and rotate them in place with the per-sensor `alignment_rot_matrix`. Branch-symmetric stateless transform:
# `timeline` and `is_measured` are received for API uniformity but not read.
data_xyz = data.view(data.shape[0], -1, 3)
data_xyz.copy_(torch.matmul(shared_metadata.alignment_rot_matrix, data_xyz.unsqueeze(-1)).squeeze(-1))
def _draw_debug(self, context: "RasterizerContext"):
"""
Draw debug arrow for the IMU acceleration.
Only draws for first rendered environment.
"""
env_idx = context.rendered_envs_idx[0] if self._manager._sim.n_envs > 0 else None
quat = self._link.get_quat(env_idx, relative=False).reshape((4,))
pos = self._link.get_pos(env_idx, relative=False).reshape((3,)) + transform_by_quat(self.pos_offset, quat)
# cannot specify envs_idx for read() when n_envs=0
data = self.read(env_idx)
acc_vec = data.lin_acc.reshape((3,)) * self._options.debug_acc_scale
gyro_vec = data.ang_vel.reshape((3,)) * self._options.debug_gyro_scale
mag_vec = data.mag.reshape((3,)) * self._options.debug_mag_scale
# transform from local frame to world frame
offset_quat = transform_quat_by_quat(self.quat_offset, quat)
acc_vec = tensor_to_array(transform_by_quat(acc_vec, offset_quat))
gyro_vec = tensor_to_array(transform_by_quat(gyro_vec, offset_quat))
mag_vec = tensor_to_array(transform_by_quat(mag_vec, offset_quat))
for debug_object in self.debug_objects:
context.clear_debug_object(debug_object)
self.debug_objects.clear()
self.debug_objects += filter(
None,
(
context.draw_debug_arrow(pos=pos, vec=acc_vec, radius=0.006, color=self._options.debug_acc_color),
context.draw_debug_arrow(pos=pos, vec=gyro_vec, radius=0.0055, color=self._options.debug_gyro_color),
context.draw_debug_arrow(pos=pos, vec=mag_vec, radius=0.005, color=self._options.debug_mag_color),
),
)