add imu gyro draw_debug, update doc

Author: Milotrince

examples/sensors/imu_franka.py

@@ -44,7 +44,7 @@ def main():
         gs.sensors.IMU(
             entity_idx=franka.idx,
             link_idx_local=end_effector.idx_local,
-            pos_offset=(0.0, 0.0, 0.2),
+            pos_offset=(0.0, 0.0, 0.15),
             # noise parameters
             acc_axes_skew=(0.0, 0.01, 0.02),
             gyro_axes_skew=(0.03, 0.04, 0.05),

genesis/recorders/plotters.py

@@ -652,6 +652,8 @@ class MPLImagePlotterOptions(BasePlotterOptions):
 class MPLImagePlotter(BaseMPLPlotter):
     """
     Live image viewer using matplotlib.
+
+    The image data should be an array-like object with shape (H, W), (H, W, 1), (H, W, 3), or (H, W, 4).
     """
 
     def build(self):
@@ -676,13 +678,6 @@ def process(self, data, cur_time):
         else:
             img_data = np.asarray(data)
 
-        if img_data.ndim == 3 and img_data.shape[0] == 1:
-            img_data = img_data[0]  # remove batch dimension
-        # TODO: color images?
-        # elif img_data.ndim == 3 and img_data.shape[-1] in [1, 3, 4]:
-        elif img_data.ndim == 3 and img_data.shape[-1] == 1:
-            img_data = img_data.squeeze(-1)
-
         vmin, vmax = np.min(img_data), np.max(img_data)
 
         current_vmin, current_vmax = self.image_plot.get_clim()

genesis/sensors/imu.py

@@ -128,9 +128,13 @@ class IMUOptions(RigidSensorOptionsMixin, NoisySensorOptionsMixin, SensorOptions
     draw_debug : bool, optional
         If True and the rasterizer visualization is active, an arrow for linear acceleration will be drawn.
     debug_acc_color : float, optional
-        The rgba color of the debug acceleration arrow. Defaults to (1.0, 0.0, 0.0, 0.5).
+        The rgba color of the debug acceleration arrow. Defaults to (0.0, 1.0, 1.0, 0.5).
     debug_acc_scale: float, optional
         The scale factor for the debug acceleration arrow. Defaults to 0.01.
+    debug_gyro_color : float, optional
+        The rgba color of the debug gyroscope arrow. Defaults to (1.0, 1.0, 0.0, 0.5).
+    debug_gyro_scale: float, optional
+        The scale factor for the debug gyroscope arrow. Defaults to 0.01.
     """
 
     acc_resolution: MaybeTuple3FType = 0.0
@@ -144,8 +148,10 @@ class IMUOptions(RigidSensorOptionsMixin, NoisySensorOptionsMixin, SensorOptions
     acc_random_walk: MaybeTuple3FType = 0.0
     gyro_random_walk: MaybeTuple3FType = 0.0
 
-    debug_acc_color: tuple[float, float, float, float] = (1.0, 0.0, 0.0, 0.5)
+    debug_acc_color: tuple[float, float, float, float] = (0.0, 1.0, 1.0, 0.5)
     debug_acc_scale: float = 0.01
+    debug_gyro_color: tuple[float, float, float, float] = (1.0, 1.0, 0.0, 0.5)
+    debug_gyro_scale: float = 0.01
 
     def validate(self, scene):
         super().validate(scene)
@@ -252,7 +258,7 @@ def build(self):
         )
 
         if self._options.draw_debug:
-            self.debug_object = None
+            self.debug_objects = [None, None]
             self.quat_offset = self._shared_metadata.offsets_quat[0, self._idx]
             self.pos_offset = self._shared_metadata.offsets_pos[0, self._idx]
 
@@ -334,9 +340,15 @@ def _draw_debug(self, context: "RasterizerContext"):
         pos = self.link.get_pos(envs_idx=envs_idx).squeeze(0) + transform_by_quat(self.pos_offset, quat)
 
         data = self.read(envs_idx=envs_idx)
-        vec = data["lin_acc"].squeeze(0) * self._options.debug_acc_scale
-        vec = tensor_to_array(transform_by_quat(vec, transform_quat_by_quat(self.quat_offset, quat)))
-
-        if self.debug_object is not None:
-            context.clear_debug_object(self.debug_object)
-        self.debug_object = context.draw_debug_arrow(pos=pos, vec=vec, color=self._options.debug_acc_color)
+        acc_vec = data["lin_acc"].squeeze(0) * self._options.debug_acc_scale
+        gyro_vec = data["ang_vel"].squeeze(0) * self._options.debug_gyro_scale
+        # transform from local frame to world frame
+        offset_quat = transform_quat_by_quat(self.quat_offset, quat)
+        acc_vec = tensor_to_array(transform_by_quat(acc_vec, offset_quat))
+        gyro_vec = tensor_to_array(transform_by_quat(gyro_vec, offset_quat))
+
+        for debug_object in self.debug_objects:
+            if debug_object is not None:
+                context.clear_debug_object(debug_object)
+        self.debug_objects[0] = context.draw_debug_arrow(pos=pos, vec=acc_vec, color=self._options.debug_acc_color)
+        self.debug_objects[1] = context.draw_debug_arrow(pos=pos, vec=gyro_vec, color=self._options.debug_gyro_color)