|
| 1 | +"""Inertial Measurement Unit sensor. |
| 2 | +
|
| 3 | +Outputs orientation, body-frame projected gravity, and body-frame linear / angular |
| 4 | +accelerations at a site rigidly attached to a robot link. Linear and angular accelerations |
| 5 | +are computed by finite difference of the world-frame velocity buffers held internally — |
| 6 | +the first step after every ``reset`` returns zero acceleration to avoid a spurious spike |
| 7 | +from a stale ``prev`` value. The accelerometer convention is ``f_b = R^T (a_w - g_w)`` |
| 8 | +when ``gravity_bias=True`` (matches a real IMU at rest reporting ``+g`` along its up |
| 9 | +axis); set ``gravity_bias=False`` to read pure kinematic acceleration instead. |
| 10 | +
|
| 11 | +The high-frequency jitter from ``(v - prev) / dt`` at the control rate is **not** filtered |
| 12 | +internally — consumers can layer ``ObservationTermCfg.noise`` + ``scale`` on top, or apply |
| 13 | +their own EMA. ``BodyVelocitySensor`` remains the velocimeter/gyro abstraction and is |
| 14 | +expected to coexist with this sensor when both velocity and acceleration are needed. |
| 15 | +""" |
| 16 | + |
| 17 | +from dataclasses import dataclass |
| 18 | +from typing import TYPE_CHECKING |
| 19 | + |
| 20 | +import torch |
| 21 | + |
| 22 | +from genelab.sensor.sensor import Sensor, SensorCfg |
| 23 | +from genelab.utils.math import matrix_from_quat, quat_apply_inverse |
| 24 | + |
| 25 | +if TYPE_CHECKING: |
| 26 | + from genelab.envs.manager_based_rl_env import ManagerBasedRlEnv |
| 27 | + |
| 28 | + |
| 29 | +_GRAVITY_MAGNITUDE = 9.81 |
| 30 | + |
| 31 | + |
| 32 | +@dataclass |
| 33 | +class IMUSensorCfg(SensorCfg): |
| 34 | + """Configuration for ``IMUSensor``. |
| 35 | +
|
| 36 | + ``offset`` is the site position in the link's local frame; it contributes both the |
| 37 | + lever-arm term to linear velocity (``ω × R · offset``) and, via its time-derivative, |
| 38 | + to linear acceleration. ``gravity_bias`` toggles real-IMU specific-force convention. |
| 39 | + ``bias_range_lin_acc`` and ``bias_range_ang_acc`` each sample a constant per-env bias |
| 40 | + on every ``reset(env_ids)`` to emulate sensor offset randomization. |
| 41 | + """ |
| 42 | + |
| 43 | + link_name: str = "" |
| 44 | + offset: tuple[float, float, float] = (0.0, 0.0, 0.0) |
| 45 | + gravity_bias: bool = True |
| 46 | + bias_range_lin_acc: tuple[float, float] | None = None |
| 47 | + bias_range_ang_acc: tuple[float, float] | None = None |
| 48 | + |
| 49 | + def build(self) -> "IMUSensor": |
| 50 | + return IMUSensor(self) |
| 51 | + |
| 52 | + |
| 53 | +@dataclass |
| 54 | +class IMUData: |
| 55 | + orientation: torch.Tensor # (B, 4) wxyz, link_quat_w |
| 56 | + projected_gravity_b: torch.Tensor # (B, 3) unit gravity rotated into the link's body frame |
| 57 | + lin_acc_b: torch.Tensor # (B, 3) linear acceleration at the site, body frame |
| 58 | + ang_acc_b: torch.Tensor # (B, 3) angular acceleration of the link, body frame |
| 59 | + |
| 60 | + |
| 61 | +class IMUSensor(Sensor[IMUData]): |
| 62 | + def __init__(self, cfg: IMUSensorCfg) -> None: |
| 63 | + super().__init__(cfg) |
| 64 | + self._cfg_typed = cfg |
| 65 | + self._link_idx: int = -1 |
| 66 | + self._offset_local: torch.Tensor | None = None |
| 67 | + self._prev_lin_vel_w: torch.Tensor | None = None |
| 68 | + self._prev_ang_vel_w: torch.Tensor | None = None |
| 69 | + self._first_step: torch.Tensor | None = None |
| 70 | + self._lin_acc_b_buf: torch.Tensor | None = None |
| 71 | + self._ang_acc_b_buf: torch.Tensor | None = None |
| 72 | + self._bias_lin: torch.Tensor | None = None |
| 73 | + self._bias_ang: torch.Tensor | None = None |
| 74 | + |
| 75 | + def bind(self, env: "ManagerBasedRlEnv") -> None: |
| 76 | + super().bind(env) |
| 77 | + if not self._cfg_typed.link_name: |
| 78 | + raise ValueError(f"IMUSensorCfg(name={self._cfg.name!r}) requires link_name") |
| 79 | + try: |
| 80 | + self._link_idx = env.link_names.index(self._cfg_typed.link_name) |
| 81 | + except ValueError as exc: |
| 82 | + raise ValueError( |
| 83 | + f"sensor {self._cfg.name!r}: link {self._cfg_typed.link_name!r} not in " |
| 84 | + f"env.link_names={env.link_names!r}" |
| 85 | + ) from exc |
| 86 | + device = env.device |
| 87 | + self._offset_local = torch.tensor( |
| 88 | + self._cfg_typed.offset, dtype=torch.float32, device=device |
| 89 | + ) |
| 90 | + zeros = lambda: torch.zeros(env.num_envs, 3, device=device) # noqa: E731 |
| 91 | + self._prev_lin_vel_w = zeros() |
| 92 | + self._prev_ang_vel_w = zeros() |
| 93 | + self._first_step = torch.ones(env.num_envs, dtype=torch.bool, device=device) |
| 94 | + self._lin_acc_b_buf = zeros() |
| 95 | + self._ang_acc_b_buf = zeros() |
| 96 | + self._bias_lin = zeros() |
| 97 | + self._bias_ang = zeros() |
| 98 | + if self._cfg_typed.bias_range_lin_acc is not None: |
| 99 | + self._resample_bias( |
| 100 | + self._bias_lin, |
| 101 | + self._cfg_typed.bias_range_lin_acc, |
| 102 | + torch.arange(env.num_envs, device=device), |
| 103 | + ) |
| 104 | + if self._cfg_typed.bias_range_ang_acc is not None: |
| 105 | + self._resample_bias( |
| 106 | + self._bias_ang, |
| 107 | + self._cfg_typed.bias_range_ang_acc, |
| 108 | + torch.arange(env.num_envs, device=device), |
| 109 | + ) |
| 110 | + |
| 111 | + def update(self, dt: float) -> None: |
| 112 | + super().update(dt) |
| 113 | + assert ( |
| 114 | + self._env is not None |
| 115 | + and self._offset_local is not None |
| 116 | + and self._prev_lin_vel_w is not None |
| 117 | + and self._prev_ang_vel_w is not None |
| 118 | + and self._first_step is not None |
| 119 | + and self._lin_acc_b_buf is not None |
| 120 | + and self._ang_acc_b_buf is not None |
| 121 | + and self._bias_lin is not None |
| 122 | + and self._bias_ang is not None |
| 123 | + ) |
| 124 | + rs = self._env.robot_state |
| 125 | + link_quat = rs.link_quat_w[:, self._link_idx] |
| 126 | + link_ang_vel_w = rs.link_ang_vel_w[:, self._link_idx] |
| 127 | + link_lin_vel_w = rs.link_lin_vel_w[:, self._link_idx] |
| 128 | + # Site velocity in world frame: v_site = v_link + ω × (R · offset). |
| 129 | + rot = matrix_from_quat(link_quat) |
| 130 | + r_world = torch.matmul(rot, self._offset_local) |
| 131 | + v_site_w = link_lin_vel_w + torch.cross(link_ang_vel_w, r_world, dim=-1) |
| 132 | + |
| 133 | + first = self._first_step.unsqueeze(-1) |
| 134 | + lin_acc_w = torch.where( |
| 135 | + first, torch.zeros_like(v_site_w), (v_site_w - self._prev_lin_vel_w) / dt |
| 136 | + ) |
| 137 | + ang_acc_w = torch.where( |
| 138 | + first, torch.zeros_like(link_ang_vel_w), (link_ang_vel_w - self._prev_ang_vel_w) / dt |
| 139 | + ) |
| 140 | + |
| 141 | + lin_acc_b = quat_apply_inverse(link_quat, lin_acc_w) |
| 142 | + ang_acc_b = quat_apply_inverse(link_quat, ang_acc_w) |
| 143 | + if self._cfg_typed.gravity_bias: |
| 144 | + gravity_w = torch.zeros_like(lin_acc_w) |
| 145 | + gravity_w[..., 2] = -_GRAVITY_MAGNITUDE |
| 146 | + gravity_b = quat_apply_inverse(link_quat, gravity_w) |
| 147 | + lin_acc_b = lin_acc_b - gravity_b |
| 148 | + |
| 149 | + self._lin_acc_b_buf = lin_acc_b + self._bias_lin |
| 150 | + self._ang_acc_b_buf = ang_acc_b + self._bias_ang |
| 151 | + # Clone before storing so subsequent in-place writes to ``robot_state`` (which the |
| 152 | + # env does every control step) cannot retroactively alter our previous-tick buffer. |
| 153 | + self._prev_lin_vel_w = v_site_w.clone() |
| 154 | + self._prev_ang_vel_w = link_ang_vel_w.clone() |
| 155 | + self._first_step = torch.zeros_like(self._first_step) |
| 156 | + |
| 157 | + def reset(self, env_ids: torch.Tensor | None = None) -> None: |
| 158 | + super().reset(env_ids) |
| 159 | + if self._env is None or self._first_step is None: |
| 160 | + return |
| 161 | + if env_ids is None: |
| 162 | + env_ids = torch.arange(self._env.num_envs, device=self._env.device) |
| 163 | + assert ( |
| 164 | + self._prev_lin_vel_w is not None |
| 165 | + and self._prev_ang_vel_w is not None |
| 166 | + and self._lin_acc_b_buf is not None |
| 167 | + and self._ang_acc_b_buf is not None |
| 168 | + and self._bias_lin is not None |
| 169 | + and self._bias_ang is not None |
| 170 | + ) |
| 171 | + self._prev_lin_vel_w[env_ids] = 0.0 |
| 172 | + self._prev_ang_vel_w[env_ids] = 0.0 |
| 173 | + self._lin_acc_b_buf[env_ids] = 0.0 |
| 174 | + self._ang_acc_b_buf[env_ids] = 0.0 |
| 175 | + self._first_step[env_ids] = True |
| 176 | + if self._cfg_typed.bias_range_lin_acc is not None: |
| 177 | + self._resample_bias(self._bias_lin, self._cfg_typed.bias_range_lin_acc, env_ids) |
| 178 | + if self._cfg_typed.bias_range_ang_acc is not None: |
| 179 | + self._resample_bias(self._bias_ang, self._cfg_typed.bias_range_ang_acc, env_ids) |
| 180 | + |
| 181 | + @staticmethod |
| 182 | + def _resample_bias( |
| 183 | + bias: torch.Tensor, bias_range: tuple[float, float], env_ids: torch.Tensor |
| 184 | + ) -> None: |
| 185 | + lo, hi = bias_range |
| 186 | + bias[env_ids] = torch.rand(env_ids.numel(), 3, device=bias.device) * (hi - lo) + lo |
| 187 | + |
| 188 | + def _compute_data(self) -> IMUData: |
| 189 | + assert ( |
| 190 | + self._env is not None |
| 191 | + and self._lin_acc_b_buf is not None |
| 192 | + and self._ang_acc_b_buf is not None |
| 193 | + ) |
| 194 | + rs = self._env.robot_state |
| 195 | + link_quat = rs.link_quat_w[:, self._link_idx] |
| 196 | + gravity_w = torch.zeros_like(rs.link_lin_vel_w[:, self._link_idx]) |
| 197 | + gravity_w[..., 2] = -1.0 |
| 198 | + projected_gravity_b = quat_apply_inverse(link_quat, gravity_w) |
| 199 | + return IMUData( |
| 200 | + orientation=link_quat, |
| 201 | + projected_gravity_b=projected_gravity_b, |
| 202 | + lin_acc_b=self._lin_acc_b_buf, |
| 203 | + ang_acc_b=self._ang_acc_b_buf, |
| 204 | + ) |
0 commit comments