Apply suggestions from code review

Co-authored-by: jtigue-bdai <[email protected]>
Signed-off-by: Pascal Roth <[email protected]>

source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/imu/imu.py

@@ -146,13 +146,16 @@ def _update_buffers_impl(self, env_ids: Sequence[int]):
             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")
 
         # store the poses
-        self._data.pos_w[env_ids] = pos_w + math_utils.quat_rotate(quat_w, self._offset_pos_b)
-        self._data.quat_w[env_ids] = math_utils.quat_mul(quat_w, self._offset_quat_b)
+        self._data.pos_w[env_ids] = pos_w + math_utils.quat_rotate(quat_w, self._offset_pos_b[env_ids])
+        self._data.quat_w[env_ids] = math_utils.quat_mul(quat_w, self._offset_quat_b[env_ids])
 
         # get the offset from COM to link origin
         com_pos_b = self._view.get_coms().to(self.device).split([3, 4], dim=-1)[0]
@@ -162,21 +165,21 @@ def _update_buffers_impl(self, env_ids: Sequence[int]):
         # if an offset is present or the COM does not agree with the link origin, the linear velocity has to be
         # transformed taking the angular velocity into account
         lin_vel_w += torch.linalg.cross(
-            ang_vel_w, math_utils.quat_rotate(quat_w, self._offset_pos_b - com_pos_b[env_ids]), dim=-1
+            ang_vel_w, math_utils.quat_rotate(quat_w, self._offset_pos_b[env_ids] - com_pos_b[env_ids]), dim=-1
         )
 
         # numerical derivative
-        lin_acc_w = (lin_vel_w - self._prev_lin_vel_w) / self._dt + self._gravity_bias_w
-        ang_acc_w = (ang_vel_w - self._prev_ang_vel_w) / self._dt
+        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_rotate_inverse(self._data.quat_w[env_ids], lin_vel_w)
         self._data.ang_vel_b[env_ids] = math_utils.quat_rotate_inverse(self._data.quat_w[env_ids], ang_vel_w)
         # store the accelerations
         self._data.lin_acc_b[env_ids] = math_utils.quat_rotate_inverse(self._data.quat_w[env_ids], lin_acc_w)
         self._data.ang_acc_b[env_ids] = math_utils.quat_rotate_inverse(self._data.quat_w[env_ids], ang_acc_w)
 
-        self._prev_lin_vel_w[:] = lin_vel_w
-        self._prev_ang_vel_w[:] = ang_vel_w
+        self._prev_lin_vel_w[env_ids] = lin_vel_w
+        self._prev_ang_vel_w[env_ids] = ang_vel_w
 
     def _initialize_buffers_impl(self):
         """Create buffers for storing data."""

source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/imu/imu_data.py

@@ -26,25 +26,25 @@ class ImuData:
     """
 
     lin_vel_b: torch.Tensor = None
-    """Root angular velocity in base frame.
+    """IMU frame angular velocity relative to the world expressed in IMU frame.
 
     Shape is (N, 3), where ``N`` is the number of environments.
     """
 
     ang_vel_b: torch.Tensor = None
-    """Root angular velocity in base frame.
+    """IMU frame angular velocity relative to the world expressed in IMU frame.
 
     Shape is (N, 3), where ``N`` is the number of environments.
     """
 
     lin_acc_b: torch.Tensor = None
-    """Root linear acceleration in base frame.
+    """IMU frame linear acceleration relative to the world expressed in IMU frame.
 
     Shape is (N, 3), where ``N`` is the number of environments.
     """
 
     ang_acc_b: torch.Tensor = None
-    """Root angular acceleration in base frame.
+    """IMU frame angular acceleration relative to the world expressed in IMU frame.
 
     Shape is (N, 3), where ``N`` is the number of environments.
     """