From b25724974f74fbf4e4620bf820c7825907b0b46f Mon Sep 17 00:00:00 2001 From: James Tigue Date: Fri, 16 May 2025 10:59:51 -0400 Subject: [PATCH 1/2] optimize imu updates --- source/isaaclab/isaaclab/sensors/imu/imu.py | 23 +++++++++++---------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/source/isaaclab/isaaclab/sensors/imu/imu.py b/source/isaaclab/isaaclab/sensors/imu/imu.py index d7182957941..24b008ab699 100644 --- a/source/isaaclab/isaaclab/sensors/imu/imu.py +++ b/source/isaaclab/isaaclab/sensors/imu/imu.py @@ -140,17 +140,13 @@ def _initialize_impl(self): def _update_buffers_impl(self, env_ids: Sequence[int]): """Fills the buffers of the sensor data.""" - # check if self._dt is set (this is set in the update function) - if not hasattr(self, "_dt"): - raise RuntimeError( - "The update function must be called before the data buffers are accessed the first time." - ) + # default to all sensors if len(env_ids) == self._num_envs: env_ids = slice(None) # obtain the poses of the sensors pos_w, quat_w = self._view.get_transforms()[env_ids].split([3, 4], dim=-1) - quat_w = math_utils.convert_quat(quat_w, to="wxyz") + quat_w = quat_w.roll(1, dims=-1) # store the poses self._data.pos_w[env_ids] = pos_w + math_utils.quat_apply(quat_w, self._offset_pos_b[env_ids]) @@ -170,12 +166,17 @@ def _update_buffers_impl(self, env_ids: Sequence[int]): # numerical derivative lin_acc_w = (lin_vel_w - self._prev_lin_vel_w[env_ids]) / self._dt + self._gravity_bias_w[env_ids] ang_acc_w = (ang_vel_w - self._prev_ang_vel_w[env_ids]) / self._dt - # store the velocities - self._data.lin_vel_b[env_ids] = math_utils.quat_apply_inverse(self._data.quat_w[env_ids], lin_vel_w) - self._data.ang_vel_b[env_ids] = math_utils.quat_apply_inverse(self._data.quat_w[env_ids], ang_vel_w) + # stack data in world frame and batch rotate + dynamics_data = torch.stack((lin_vel_w, ang_vel_w, lin_acc_w, ang_acc_w), dim=0) + dynamics_data_rot = math_utils.quat_apply_inverse(self._data.quat_w[env_ids].repeat(4, 1), dynamics_data).chunk( + 4, dim=0 + ) + # store the velocities. + self._data.lin_vel_b[env_ids] = dynamics_data_rot[0] + self._data.ang_vel_b[env_ids] = dynamics_data_rot[1] # store the accelerations - self._data.lin_acc_b[env_ids] = math_utils.quat_apply_inverse(self._data.quat_w[env_ids], lin_acc_w) - self._data.ang_acc_b[env_ids] = math_utils.quat_apply_inverse(self._data.quat_w[env_ids], ang_acc_w) + self._data.lin_acc_b[env_ids] = dynamics_data_rot[2] + self._data.ang_acc_b[env_ids] = dynamics_data_rot[3] self._prev_lin_vel_w[env_ids] = lin_vel_w self._prev_ang_vel_w[env_ids] = ang_vel_w From 8067b49ae5c0a7d2eef7920c38c1b0a3de0447b8 Mon Sep 17 00:00:00 2001 From: James Tigue Date: Fri, 16 May 2025 11:12:48 -0400 Subject: [PATCH 2/2] add imu projected_gravity --- .../isaaclab/envs/mdp/observations.py | 15 +++++++++++++++ source/isaaclab/isaaclab/sensors/imu/imu.py | 19 ++++++++++++++++--- .../isaaclab/isaaclab/sensors/imu/imu_data.py | 6 ++++++ source/isaaclab/test/sensors/test_imu.py | 15 +++++++++++++++ 4 files changed, 52 insertions(+), 3 deletions(-) diff --git a/source/isaaclab/isaaclab/envs/mdp/observations.py b/source/isaaclab/isaaclab/envs/mdp/observations.py index 56ce0228f2f..722f68f567d 100644 --- a/source/isaaclab/isaaclab/envs/mdp/observations.py +++ b/source/isaaclab/isaaclab/envs/mdp/observations.py @@ -270,6 +270,21 @@ def imu_orientation(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntit return asset.data.quat_w +def imu_projected_gravity(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("imu")) -> torch.Tensor: + """Imu sensor orientation w.r.t the env.scene.origin. + + Args: + env: The environment. + asset_cfg: The SceneEntity associated with an Imu sensor. + + Returns: + Gravity projected on imu_frame, shape of torch.tensor is (num_env,3). + """ + + asset: Imu = env.scene[asset_cfg.name] + return asset.data.projected_gravity_b + + def imu_ang_vel(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("imu")) -> torch.Tensor: """Imu sensor angular velocity w.r.t. environment origin expressed in the sensor frame. diff --git a/source/isaaclab/isaaclab/sensors/imu/imu.py b/source/isaaclab/isaaclab/sensors/imu/imu.py index 24b008ab699..35099b7e42a 100644 --- a/source/isaaclab/isaaclab/sensors/imu/imu.py +++ b/source/isaaclab/isaaclab/sensors/imu/imu.py @@ -96,7 +96,11 @@ def reset(self, env_ids: Sequence[int] | None = None): if env_ids is None: env_ids = slice(None) # reset accumulative data buffers + self._data.pos_w[env_ids] = 0.0 self._data.quat_w[env_ids] = 0.0 + self._data.quat_w[env_ids, 0] = 1.0 + self._data.projected_gravity_b[env_ids] = 0.0 + self._data.projected_gravity_b[env_ids, 2] = -1.0 self._data.lin_vel_b[env_ids] = 0.0 self._data.ang_vel_b[env_ids] = 0.0 self._data.lin_acc_b[env_ids] = 0.0 @@ -135,6 +139,12 @@ def _initialize_impl(self): else: raise RuntimeError(f"Failed to find a RigidBodyAPI for the prim paths: {self.cfg.prim_path}") + # Get world gravity + gravity = self._physics_sim_view.get_gravity() + gravity_dir = torch.tensor((gravity[0], gravity[1], gravity[2]), device=self.device) + gravity_dir = math_utils.normalize(gravity_dir.unsqueeze(0)).squeeze(0) + self.GRAVITY_VEC_W = gravity_dir.repeat(self.num_instances, 1) + # Create internal buffers self._initialize_buffers_impl() @@ -167,9 +177,9 @@ def _update_buffers_impl(self, env_ids: Sequence[int]): lin_acc_w = (lin_vel_w - self._prev_lin_vel_w[env_ids]) / self._dt + self._gravity_bias_w[env_ids] ang_acc_w = (ang_vel_w - self._prev_ang_vel_w[env_ids]) / self._dt # stack data in world frame and batch rotate - dynamics_data = torch.stack((lin_vel_w, ang_vel_w, lin_acc_w, ang_acc_w), dim=0) - dynamics_data_rot = math_utils.quat_apply_inverse(self._data.quat_w[env_ids].repeat(4, 1), dynamics_data).chunk( - 4, dim=0 + dynamics_data = torch.stack((lin_vel_w, ang_vel_w, lin_acc_w, ang_acc_w, self.GRAVITY_VEC_W[env_ids]), dim=0) + dynamics_data_rot = math_utils.quat_apply_inverse(self._data.quat_w[env_ids].repeat(5, 1), dynamics_data).chunk( + 5, dim=0 ) # store the velocities. self._data.lin_vel_b[env_ids] = dynamics_data_rot[0] @@ -177,6 +187,8 @@ def _update_buffers_impl(self, env_ids: Sequence[int]): # store the accelerations self._data.lin_acc_b[env_ids] = dynamics_data_rot[2] self._data.ang_acc_b[env_ids] = dynamics_data_rot[3] + # store projected gravity + self._data.projected_gravity_b[env_ids] = dynamics_data_rot[4] self._prev_lin_vel_w[env_ids] = lin_vel_w self._prev_ang_vel_w[env_ids] = ang_vel_w @@ -187,6 +199,7 @@ def _initialize_buffers_impl(self): self._data.pos_w = torch.zeros(self._view.count, 3, device=self._device) self._data.quat_w = torch.zeros(self._view.count, 4, device=self._device) self._data.quat_w[:, 0] = 1.0 + self._data.projected_gravity_b = torch.zeros(self._view.count, 3, device=self._device) self._data.lin_vel_b = torch.zeros_like(self._data.pos_w) self._data.ang_vel_b = torch.zeros_like(self._data.pos_w) self._data.lin_acc_b = torch.zeros_like(self._data.pos_w) diff --git a/source/isaaclab/isaaclab/sensors/imu/imu_data.py b/source/isaaclab/isaaclab/sensors/imu/imu_data.py index 7c56a3808ae..098131099d5 100644 --- a/source/isaaclab/isaaclab/sensors/imu/imu_data.py +++ b/source/isaaclab/isaaclab/sensors/imu/imu_data.py @@ -25,6 +25,12 @@ class ImuData: Shape is (N, 4), where ``N`` is the number of environments. """ + projected_gravity_b: torch.Tensor = None + """Gravity direction unit vector projected on the imu frame. + + Shape is (N,3), where ``N`` is the number of environments. + """ + lin_vel_b: torch.Tensor = None """IMU frame angular velocity relative to the world expressed in IMU frame. diff --git a/source/isaaclab/test/sensors/test_imu.py b/source/isaaclab/test/sensors/test_imu.py index 3b830388c98..611e1ed487b 100644 --- a/source/isaaclab/test/sensors/test_imu.py +++ b/source/isaaclab/test/sensors/test_imu.py @@ -353,6 +353,14 @@ def test_single_dof_pendulum(setup_sim): if idx < 2: continue + # compare imu projected gravity + gravity_dir_w = torch.tensor((0.0, 0.0, -1.0), device=scene.device).repeat(2, 1) + gravity_dir_b = math_utils.quat_apply_inverse(imu_data.quat_w, gravity_dir_w) + torch.testing.assert_close( + imu_data.projected_gravity_b, + gravity_dir_b, + ) + # compare imu angular velocity with joint velocity torch.testing.assert_close( joint_vel, @@ -495,6 +503,13 @@ def test_offset_calculation(setup_sim): rtol=1e-4, atol=1e-4, ) + # check the projected gravity + torch.testing.assert_close( + scene.sensors["imu_robot_base"].data.projected_gravity_b, + scene.sensors["imu_robot_imu_link"].data.projected_gravity_b, + rtol=1e-4, + atol=1e-4, + ) def test_env_ids_propagation(setup_sim):