From eceaa0b12d1619a3c4689ede8dec37629355e60b Mon Sep 17 00:00:00 2001
From: Pascal Roth <57946385+pascal-roth@users.noreply.github.com>
Date: Mon, 14 Oct 2024 13:45:49 +0200
Subject: [PATCH] Apply suggestions from code review

Co-authored-by: jtigue-bdai <166445701+jtigue-bdai@users.noreply.github.com>
Signed-off-by: Pascal Roth <57946385+pascal-roth@users.noreply.github.com>
---
 .../omni/isaac/lab/sensors/imu/imu.py           | 17 ++++++++++-------
 .../omni/isaac/lab/sensors/imu/imu_data.py      |  8 ++++----
 2 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/imu/imu.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/imu/imu.py
index 95cc7e03de..24131dccc7 100644
--- a/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/imu/imu.py
+++ b/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,12 +165,12 @@ 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)
@@ -175,8 +178,8 @@ def _update_buffers_impl(self, env_ids: Sequence[int]):
         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."""
diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/imu/imu_data.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/imu/imu_data.py
index f6fed3bdd8..37920cd841 100644
--- a/source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/imu/imu_data.py
+++ b/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.
     """