[FEATURE] Add Sensor.draw_debug method. (Genesis-Embodied-AI#1770)

Author: Milotrince

examples/sensors/contact_force_go2.py

@@ -0,0 +1,94 @@
+import argparse
+
+from tqdm import tqdm
+
+import genesis as gs
+from genesis.recorders.plotters import IS_MATPLOTLIB_AVAILABLE, IS_PYQTGRAPH_AVAILABLE
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-dt", "--timestep", type=float, default=1e-2, help="Simulation time step")
+    parser.add_argument("-v", "--vis", action="store_true", default=True, help="Show visualization GUI")
+    parser.add_argument("-nv", "--no-vis", action="store_false", dest="vis", help="Disable visualization GUI")
+    parser.add_argument("-c", "--cpu", action="store_true", help="Use CPU instead of GPU")
+    parser.add_argument("-t", "--seconds", type=float, default=2, help="Number of seconds to simulate")
+    parser.add_argument("-f", "--force", action="store_true", default=True, help="Use ContactForceSensor (xyz float)")
+    parser.add_argument("-nf", "--no-force", action="store_false", dest="force", help="Use ContactSensor (boolean)")
+
+    args = parser.parse_args()
+
+    ########################## init ##########################
+    gs.init(backend=gs.cpu if args.cpu else gs.gpu, logging_level=None)
+
+    ########################## scene setup ##########################
+    scene = gs.Scene(
+        sim_options=gs.options.SimOptions(dt=args.timestep),
+        rigid_options=gs.options.RigidOptions(
+            use_gjk_collision=True,
+            constraint_timeconst=max(0.01, 2 * args.timestep),
+        ),
+        vis_options=gs.options.VisOptions(show_world_frame=True),
+        profiling_options=gs.options.ProfilingOptions(show_FPS=False),
+        show_viewer=args.vis,
+    )
+
+    scene.add_entity(gs.morphs.Plane())
+
+    foot_link_names = ["FR_foot", "FL_foot", "RR_foot", "RL_foot"]
+    go2 = scene.add_entity(
+        gs.morphs.URDF(
+            file="urdf/go2/urdf/go2.urdf",
+            pos=(0.0, 0.0, 0.2),
+            links_to_keep=foot_link_names,
+        )
+    )
+
+    for link_name in foot_link_names:
+        if args.force:
+            sensor_options = gs.sensors.ContactForce(
+                entity_idx=go2.idx,
+                link_idx_local=go2.get_link(link_name).idx_local,
+                draw_debug=True,
+            )
+            plot_kwargs = dict(
+                title=f"{link_name} Force Sensor Data",
+                labels=["force_x", "force_y", "force_z"],
+            )
+        else:
+            sensor_options = gs.sensors.Contact(
+                entity_idx=go2.idx,
+                link_idx_local=go2.get_link(link_name).idx_local,
+                draw_debug=True,
+            )
+            plot_kwargs = dict(
+                title=f"{link_name} Contact Sensor Data",
+                labels=["in_contact"],
+            )
+
+        sensor = scene.add_sensor(sensor_options)
+
+        if IS_PYQTGRAPH_AVAILABLE:
+            sensor.start_recording(gs.recorders.PyQtLinePlot(**plot_kwargs))
+        elif IS_MATPLOTLIB_AVAILABLE:
+            print("pyqtgraph not found, falling back to matplotlib.")
+            sensor.start_recording(gs.recorders.MPLLinePlot(**plot_kwargs))
+        else:
+            print("matplotlib or pyqtgraph not found, skipping real-time plotting.")
+
+    scene.build()
+
+    try:
+        steps = int(args.seconds / args.timestep)
+        for _ in tqdm(range(steps)):
+            scene.step()
+    except KeyboardInterrupt:
+        gs.logger.info("Simulation interrupted, exiting.")
+    finally:
+        gs.logger.info("Simulation finished.")
+
+        scene.stop_recording()
+
+
+if __name__ == "__main__":
+    main()

examples/sensors/force.py

@@ -44,6 +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.15),
             # noise parameters
             acc_axes_skew=(0.0, 0.01, 0.02),
             gyro_axes_skew=(0.03, 0.04, 0.05),
@@ -54,22 +55,24 @@ def main():
             delay=0.01,
             jitter=0.01,
             interpolate=True,
+            # visualize
+            draw_debug=True,
         )
     )
     labels = {"lin_acc": ("acc_x", "acc_y", "acc_z"), "ang_vel": ("gyro_x", "gyro_y", "gyro_z")}
     if args.vis:
         if IS_PYQTGRAPH_AVAILABLE:
-            imu.start_recording(gs.recorders.PyQtPlot(title="IMU Measured Data", labels=labels))
+            imu.start_recording(gs.recorders.PyQtLinePlot(title="IMU Measured Data", labels=labels))
             scene.start_recording(
                 imu.read_ground_truth,
-                gs.recorders.PyQtPlot(title="IMU Ground Truth Data", labels=labels),
+                gs.recorders.PyQtLinePlot(title="IMU Ground Truth Data", labels=labels),
             )
         elif IS_MATPLOTLIB_AVAILABLE:
             gs.logger.info("pyqtgraph not found, falling back to matplotlib.")
-            imu.start_recording(gs.recorders.MPLPlot(title="IMU Measured Data", labels=labels))
+            imu.start_recording(gs.recorders.MPLLinePlot(title="IMU Measured Data", labels=labels))
             scene.start_recording(
                 imu.read_ground_truth,
-                gs.recorders.MPLPlot(title="IMU Ground Truth Data", labels=labels),
+                gs.recorders.MPLLinePlot(title="IMU Ground Truth Data", labels=labels),
             )
         else:
             print("matplotlib or pyqtgraph not found, skipping real-time plotting.")

examples/sensors/imu.py examples/sensors/imu_franka.pyexamples/sensors/imu.py renamed to examples/sensors/imu_franka.py

@@ -2,6 +2,7 @@
 from .file_writers import CSVFileWriterOptions as CSVFile
 from .file_writers import NPZFileWriterOptions as NPZFile
 from .file_writers import VideoFileWriterOptions as VideoFile
-from .plotters import MPLPlotterOptions as MPLPlot
-from .plotters import PyQtPlotterOptions as PyQtPlot
+from .plotters import MPLImagePlotterOptions as MPLImagePlot
+from .plotters import MPLLinePlotterOptions as MPLLinePlot
+from .plotters import PyQtLinePlotterOptions as PyQtLinePlot
 from .recorder_manager import RecorderManager, register_recording

genesis/recorders/__init__.py

@@ -1,12 +1,13 @@
 import queue
 import threading
 import time
-from typing import Callable, Generic, TypeVar
+from typing import TYPE_CHECKING, Callable, Generic, TypeVar
 
 import genesis as gs
 from genesis.options import Options
 
-from .recorder_manager import RecorderManager
+if TYPE_CHECKING:
+    from .recorder_manager import RecorderManager
 
 T = TypeVar("T")
 
@@ -50,7 +51,7 @@ class Recorder(Generic[T]):
     done through the RecorderManager.
     """
 
-    def __init__(self, manager: RecorderManager, options: RecorderOptions, data_func: Callable[[], T]):
+    def __init__(self, manager: "RecorderManager", options: RecorderOptions, data_func: Callable[[], T]):
         self._options = options
         self._manager = manager
         self._data_func = data_func