Adds IMU projected_gravity_b and IMU computation speed optimizations (#2512)

# Description

<!--
Thank you for your interest in sending a pull request. Please make sure
to check the contribution guidelines.

Link:
https://isaac-sim.github.io/IsaacLab/main/source/refs/contributing.html
-->
This PR adds projected_gravity_b data stream as well as speeding up
computation by during batch rotating of all data streams to base frame
in one computation rather than each data stream separately.

This PR is dependant on [PR
2129](#2129) and should not
be merged until that PR is merged.

Fixes # (issue)

<!-- As a practice, it is recommended to open an issue to have
discussions on the proposed pull request.
This makes it easier for the community to keep track of what is being
developed or added, and if a given feature
is demanded by more than one party. -->

## Type of change

<!-- As you go through the list, delete the ones that are not
applicable. -->

- New feature (non-breaking change which adds functionality)

<!--
Example:

| Before | After |
| ------ | ----- |
| _gif/png before_ | _gif/png after_ |

To upload images to a PR -- simply drag and drop an image while in edit
mode and it should upload the image directly. You can then paste that
source into the above before/after sections.
-->

## Checklist

- [ ] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./isaaclab.sh --format`
- [ ] I have made corresponding changes to the documentation
- [ ] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my
feature works
- [ ] I have updated the changelog and the corresponding version in the
extension's `config/extension.toml` file
- [ ] I have added my name to the `CONTRIBUTORS.md` or my name already
exists there

<!--
As you go through the checklist above, you can mark something as done by
putting an x character in it

For example,
- [x] I have done this task
- [ ] I have not done this task
-->

---------

Co-authored-by: Kelly Guo <kellyg@nvidia.com>

Author: jtigue-bdai

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.
 

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)

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.
 

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):