diff --git a/source/isaaclab/isaaclab/envs/mdp/observations.py b/source/isaaclab/isaaclab/envs/mdp/observations.py index 745482f8c7e..5d2cd7f96f7 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 f083f53b2d5..a012f6d6b0c 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,22 +139,24 @@ 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() 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 +176,19 @@ 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, 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] + 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] + # 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 @@ -186,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 b5efc51702d..ae8efc47831 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 da41c3fdde6..ab5e7241900 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):