|
| 1 | +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). |
| 2 | +# All rights reserved. |
| 3 | +# |
| 4 | +# SPDX-License-Identifier: BSD-3-Clause |
| 5 | + |
| 6 | +from __future__ import annotations |
| 7 | + |
| 8 | +from collections.abc import Sequence |
| 9 | +from typing import TYPE_CHECKING |
| 10 | + |
| 11 | +import torch |
| 12 | +import warp as wp |
| 13 | + |
| 14 | +from pxr import UsdPhysics |
| 15 | + |
| 16 | +import isaaclab.sim as sim_utils |
| 17 | +import isaaclab.utils.math as math_utils |
| 18 | +from isaaclab.physics import PhysicsManager |
| 19 | +from isaaclab.sensors.imu import BaseImu |
| 20 | + |
| 21 | +import isaaclab_ovphysx.tensor_types as TT |
| 22 | +from isaaclab_ovphysx.physics import OvPhysxManager as SimulationManager |
| 23 | + |
| 24 | +from .imu_data import ImuData |
| 25 | +from .kernels import imu_reset_kernel, imu_update_kernel |
| 26 | + |
| 27 | +if TYPE_CHECKING: |
| 28 | + from isaaclab.sensors.imu import ImuCfg |
| 29 | + |
| 30 | + |
| 31 | +class Imu(BaseImu): |
| 32 | + """The OVPhysX Inertial Measurement Unit (IMU) sensor. |
| 33 | +
|
| 34 | + This sensor models a real IMU that measures angular velocity (gyroscope) and |
| 35 | + linear acceleration (accelerometer) in the sensor's body frame. Unlike the PVA |
| 36 | + sensor, it does not provide pose, linear velocity, angular acceleration, or |
| 37 | + projected gravity. |
| 38 | +
|
| 39 | + Like a real accelerometer, the linear acceleration readings always include the |
| 40 | + contribution of gravity. The gravity vector is read from the simulation |
| 41 | + configuration at initialization. |
| 42 | +
|
| 43 | + The sensor can be attached to any prim path with a rigid ancestor in its tree. |
| 44 | + If the provided path is not a rigid body, the closest rigid-body ancestor is |
| 45 | + used for simulation queries. The fixed transform from that ancestor to the |
| 46 | + target prim is computed once during initialization and composed with the |
| 47 | + configured sensor offset. |
| 48 | +
|
| 49 | + .. note:: |
| 50 | +
|
| 51 | + Linear acceleration is computed using numerical differentiation from |
| 52 | + velocities. Consequently, the IMU sensor accuracy depends on the chosen |
| 53 | + physics timestep. For sufficient accuracy, we recommend keeping the |
| 54 | + timestep at least 200 Hz. |
| 55 | + """ |
| 56 | + |
| 57 | + cfg: ImuCfg |
| 58 | + """The configuration parameters.""" |
| 59 | + |
| 60 | + __backend_name__: str = "ovphysx" |
| 61 | + """The name of the backend for the IMU sensor.""" |
| 62 | + |
| 63 | + def __init__(self, cfg: ImuCfg): |
| 64 | + """Initializes the IMU sensor. |
| 65 | +
|
| 66 | + Args: |
| 67 | + cfg: The configuration parameters. |
| 68 | + """ |
| 69 | + super().__init__(cfg) |
| 70 | + self._data = ImuData() |
| 71 | + self._rigid_parent_expr: str | None = None |
| 72 | + |
| 73 | + def __str__(self) -> str: |
| 74 | + """Returns: A string containing information about the instance.""" |
| 75 | + return ( |
| 76 | + f"Imu sensor @ '{self.cfg.prim_path}': \n" |
| 77 | + f"\tbackend : ovphysx\n" |
| 78 | + f"\tupdate period (s) : {self.cfg.update_period}\n" |
| 79 | + f"\tnumber of sensors : {self._num_bodies}\n" |
| 80 | + ) |
| 81 | + |
| 82 | + """ |
| 83 | + Properties |
| 84 | + """ |
| 85 | + |
| 86 | + @property |
| 87 | + def data(self) -> ImuData: |
| 88 | + self._update_outdated_buffers() |
| 89 | + return self._data |
| 90 | + |
| 91 | + @property |
| 92 | + def num_instances(self) -> int: |
| 93 | + return self._num_bodies |
| 94 | + |
| 95 | + """ |
| 96 | + Operations |
| 97 | + """ |
| 98 | + |
| 99 | + def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None = None): |
| 100 | + env_mask = self._resolve_indices_and_mask(env_ids, env_mask) |
| 101 | + super().reset(None, env_mask) |
| 102 | + |
| 103 | + wp.launch( |
| 104 | + imu_reset_kernel, |
| 105 | + dim=self._num_envs, |
| 106 | + inputs=[ |
| 107 | + env_mask, |
| 108 | + self._data._ang_vel_b, |
| 109 | + self._data._lin_acc_b, |
| 110 | + self._prev_lin_vel_w, |
| 111 | + ], |
| 112 | + device=self._device, |
| 113 | + ) |
| 114 | + |
| 115 | + def update(self, dt: float, force_recompute: bool = False): |
| 116 | + self._dt = dt |
| 117 | + super().update(dt, force_recompute) |
| 118 | + |
| 119 | + """ |
| 120 | + Implementation. |
| 121 | + """ |
| 122 | + |
| 123 | + def _initialize_impl(self): |
| 124 | + """Initializes the sensor handles and internal buffers. |
| 125 | +
|
| 126 | + - If the target prim path is a rigid body, bind directly to it. |
| 127 | + - Otherwise find the closest rigid-body ancestor, cache the fixed transform |
| 128 | + from that ancestor to the target prim, and bind to the ancestor pattern. |
| 129 | + """ |
| 130 | + super()._initialize_impl() |
| 131 | + |
| 132 | + physx_instance = SimulationManager.get_physx_instance() |
| 133 | + if physx_instance is None: |
| 134 | + raise RuntimeError( |
| 135 | + "OVPhysX Imu: physics instance is not initialized. " |
| 136 | + "Ensure SimulationContext.reset() has been called before sensor initialization." |
| 137 | + ) |
| 138 | + |
| 139 | + prim = sim_utils.find_first_matching_prim(self.cfg.prim_path) |
| 140 | + if prim is None: |
| 141 | + raise RuntimeError(f"Failed to find a prim at path expression: {self.cfg.prim_path}") |
| 142 | + |
| 143 | + ancestor_prim = sim_utils.get_first_matching_ancestor_prim( |
| 144 | + prim.GetPath(), predicate=lambda _prim: _prim.HasAPI(UsdPhysics.RigidBodyAPI) |
| 145 | + ) |
| 146 | + if ancestor_prim is None: |
| 147 | + raise RuntimeError(f"Failed to find a rigid body ancestor prim at path expression: {self.cfg.prim_path}") |
| 148 | + |
| 149 | + if ancestor_prim == prim: |
| 150 | + self._rigid_parent_expr = self.cfg.prim_path |
| 151 | + fixed_pos_b, fixed_quat_b = None, None |
| 152 | + else: |
| 153 | + relative_path = prim.GetPath().MakeRelativePath(ancestor_prim.GetPath()).pathString |
| 154 | + self._rigid_parent_expr = self.cfg.prim_path.replace("/" + relative_path, "") |
| 155 | + fixed_pos_b, fixed_quat_b = sim_utils.resolve_prim_pose(prim, ancestor_prim) |
| 156 | + |
| 157 | + # Translate the user-facing path expression (".*" / "{ENV_REGEX_NS}") to an ovphysx glob. |
| 158 | + pattern = self._rigid_parent_expr.replace(".*", "*") |
| 159 | + |
| 160 | + # Open the three pattern-bound tensor bindings on the rigid-body ancestor. |
| 161 | + self._pose_binding = physx_instance.create_tensor_binding(pattern=pattern, tensor_type=TT.RIGID_BODY_POSE) |
| 162 | + self._vel_binding = physx_instance.create_tensor_binding(pattern=pattern, tensor_type=TT.RIGID_BODY_VELOCITY) |
| 163 | + self._com_binding = physx_instance.create_tensor_binding(pattern=pattern, tensor_type=TT.RIGID_BODY_COM_POSE) |
| 164 | + self._num_bodies = self._pose_binding.count |
| 165 | + |
| 166 | + if self._num_bodies != self._num_envs: |
| 167 | + raise RuntimeError( |
| 168 | + f"OVPhysX Imu: pattern '{pattern}' matched {self._num_bodies} rigid bodies but the scene has " |
| 169 | + f"{self._num_envs} environments. The IMU expects exactly one body per environment." |
| 170 | + ) |
| 171 | + |
| 172 | + # Read scene gravity from the simulation cfg (OvPhysxManager has no get_gravity()). |
| 173 | + gravity = PhysicsManager._sim.cfg.gravity |
| 174 | + gravity_bias = torch.tensor((-gravity[0], -gravity[1], -gravity[2]), device=self._device) |
| 175 | + gravity_bias_torch = gravity_bias.repeat(self._num_bodies, 1) |
| 176 | + self._gravity_bias_w = wp.from_torch(gravity_bias_torch.contiguous(), dtype=wp.vec3f) |
| 177 | + |
| 178 | + self._initialize_buffers_impl() |
| 179 | + |
| 180 | + # Compose the configured offset with the fixed ancestor->target transform (done once). |
| 181 | + if fixed_pos_b is not None and fixed_quat_b is not None: |
| 182 | + fixed_p = torch.tensor(fixed_pos_b, device=self._device).repeat(self._num_bodies, 1) |
| 183 | + fixed_q = torch.tensor(fixed_quat_b, device=self._device).repeat(self._num_bodies, 1) |
| 184 | + |
| 185 | + cfg_p = wp.to_torch(self._offset_pos_b).clone() |
| 186 | + cfg_q = wp.to_torch(self._offset_quat_b).clone() |
| 187 | + |
| 188 | + composed_p = fixed_p + math_utils.quat_apply(fixed_q, cfg_p) |
| 189 | + composed_q = math_utils.quat_mul(fixed_q, cfg_q) |
| 190 | + |
| 191 | + self._offset_pos_b = wp.from_torch(composed_p.contiguous(), dtype=wp.vec3f) |
| 192 | + self._offset_quat_b = wp.from_torch(composed_q.contiguous(), dtype=wp.quatf) |
| 193 | + |
| 194 | + def _update_buffers_impl(self, env_mask: wp.array | None = None): |
| 195 | + """Fills the buffers of the sensor data.""" |
| 196 | + env_mask = self._resolve_indices_and_mask(None, env_mask) |
| 197 | + |
| 198 | + # ovphysx binding.read(dst) writes into a pre-allocated dst buffer. |
| 199 | + # For structured-dtype targets (transformf, spatial_vectorf) we pass a |
| 200 | + # zero-copy flat float32 view of the same memory; the kernel then reads |
| 201 | + # the structured buffer directly. |
| 202 | + self._pose_binding.read(self._transforms_view) |
| 203 | + self._vel_binding.read(self._velocities_view) |
| 204 | + |
| 205 | + # COM is CPU-resident in ovphysx; read into the CPU scratch, then copy into the GPU buffer. |
| 206 | + self._com_binding.read(self._coms_cpu_view) |
| 207 | + wp.copy(self._coms_gpu_view, self._coms_cpu_view) |
| 208 | + |
| 209 | + wp.launch( |
| 210 | + imu_update_kernel, |
| 211 | + dim=self._num_envs, |
| 212 | + inputs=[ |
| 213 | + env_mask, |
| 214 | + self._transforms_buf, |
| 215 | + self._velocities_buf, |
| 216 | + self._coms_buffer, |
| 217 | + self._offset_pos_b, |
| 218 | + self._offset_quat_b, |
| 219 | + self._gravity_bias_w, |
| 220 | + self._prev_lin_vel_w, |
| 221 | + 1.0 / self._dt, |
| 222 | + self._data._ang_vel_b, |
| 223 | + self._data._lin_acc_b, |
| 224 | + ], |
| 225 | + device=self._device, |
| 226 | + ) |
| 227 | + |
| 228 | + def _initialize_buffers_impl(self): |
| 229 | + """Create buffers for storing data.""" |
| 230 | + self._data.create_buffers(num_envs=self._num_bodies, device=self._device) |
| 231 | + |
| 232 | + self._prev_lin_vel_w = wp.zeros(self._num_bodies, dtype=wp.vec3f, device=self._device) |
| 233 | + |
| 234 | + offset_pos_torch = torch.tensor(list(self.cfg.offset.pos), device=self._device).repeat(self._num_bodies, 1) |
| 235 | + offset_quat_torch = torch.tensor(list(self.cfg.offset.rot), device=self._device).repeat(self._num_bodies, 1) |
| 236 | + self._offset_pos_b = wp.from_torch(offset_pos_torch.contiguous(), dtype=wp.vec3f) |
| 237 | + self._offset_quat_b = wp.from_torch(offset_quat_torch.contiguous(), dtype=wp.quatf) |
| 238 | + |
| 239 | + # Persistent destination buffers for ovphysx tensor reads. |
| 240 | + # The kernel consumes the structured-dtype buffers; the *_view aliases are |
| 241 | + # flat float32 reinterpretations of the same memory used as ovphysx |
| 242 | + # binding.read() destinations. |
| 243 | + self._transforms_buf = wp.zeros(self._num_bodies, dtype=wp.transformf, device=self._device) |
| 244 | + self._velocities_buf = wp.zeros(self._num_bodies, dtype=wp.spatial_vectorf, device=self._device) |
| 245 | + self._coms_buffer = wp.zeros(self._num_bodies, dtype=wp.transformf, device=self._device) |
| 246 | + |
| 247 | + self._transforms_view = wp.array( |
| 248 | + ptr=self._transforms_buf.ptr, |
| 249 | + shape=self._pose_binding.shape, |
| 250 | + dtype=wp.float32, |
| 251 | + device=self._device, |
| 252 | + copy=False, |
| 253 | + ) |
| 254 | + self._velocities_view = wp.array( |
| 255 | + ptr=self._velocities_buf.ptr, |
| 256 | + shape=self._vel_binding.shape, |
| 257 | + dtype=wp.float32, |
| 258 | + device=self._device, |
| 259 | + copy=False, |
| 260 | + ) |
| 261 | + self._coms_gpu_view = wp.array( |
| 262 | + ptr=self._coms_buffer.ptr, |
| 263 | + shape=self._com_binding.shape, |
| 264 | + dtype=wp.float32, |
| 265 | + device=self._device, |
| 266 | + copy=False, |
| 267 | + ) |
| 268 | + # COM binding lives on CPU, so its read destination must also be CPU. |
| 269 | + self._coms_cpu_view = wp.zeros(self._com_binding.shape, dtype=wp.float32, device="cpu") |
0 commit comments