Enable visualization of parallel robots on passive viewer

ami-iit · Jan 28, 2025 · 9958e80 · 9958e80
1 parent 0abc4c3
commit 9958e80
Showing 1 changed file with 18 additions and 3 deletions.
diff --git a/src/jaxsim/mujoco/visualizer.py b/src/jaxsim/mujoco/visualizer.py
@@ -46,7 +46,7 @@ def __init__(
 
         self.fps = fps
         self.frames: list[npt.NDArray] = []
-        self.data: mujoco.MjData | None = None
+        self.data: list[mujoco.MjData] | mujoco.MjData | None = None
         self.model: mujoco.MjModel | None = None
         self.reset(model=model, data=data)
 
@@ -68,8 +68,23 @@ def reset(
     def render_frame(self, camera_name: str = "track") -> npt.NDArray:
         """Render a frame."""
 
-        mujoco.mj_forward(self.model, self.data)
-        self.renderer.update_scene(data=self.data, camera=camera_name)
+        for idx, data in enumerate(
+            self.data if isinstance(self.data, list) else [self.data]
+        ):
+            mujoco.mj_forward(self.model, data)
+
+            if idx == 0:
+                self.renderer.update_scene(data=data, camera=camera_name)
+                continue
+
+            mujoco.mjv_addGeoms(
+                m=self.model,
+                d=data,
+                opt=mujoco.MjvOption(),
+                pert=mujoco.MjvPerturb(),
+                catmask=mujoco.mjtCatBit.mjCAT_DYNAMIC,
+                scn=self.renderer.scene,
+            )
 
         return self.renderer.render()