Add camera support to jaxsim.mujoco visualizer

src/jaxsim/mujoco/__init__.py

@@ -1,3 +1,3 @@
 from .loaders import RodModelToMjcf, SdfToMjcf, UrdfToMjcf
 from .model import MujocoModelHelper
-from .visualizer import MujocoVisualizer
+from .visualizer import MujocoVideoRecorder, MujocoVisualizer

src/jaxsim/mujoco/loaders.py

@@ -387,30 +387,42 @@ def convert(
             dir="0 0 -1",
         )
 
+        # -------------------------------------------------------
+        # Add a camera following the CoM of the worldbody element
+        # -------------------------------------------------------
+
+        worldbody_element = None
+
+        # Find the <worldbody> element of our model by searching the one that contains
+        # all the considered joints. This is needed because there might be multiple
+        # <worldbody> elements inside <mujoco>.
+        for wb in mujoco_element.findall(".//worldbody"):
+            if all(
+                wb.find(f".//joint[@name='{j}']") is not None for j in considered_joints
+            ):
+                worldbody_element = wb
+                break
+
+        if worldbody_element is None:
+            raise RuntimeError("Failed to find the <worldbody> element of the model")
+
+        # Camera attached to the model
+        _ = ET.SubElement(
+            worldbody_element,
+            "camera",
+            name="track",
+            mode="trackcom",
+            pos="1 0 5",
+            zaxis="0 0 1",
+            fovy="60",
+        )
+
         # ------------------------------------------------
         # Add a light following the  CoM of the first link
         # ------------------------------------------------
 
         if not rod_model.is_fixed_base():
 
-            worldbody_element = None
-
-            # Find the <worldbody> element of our model by searching the one that contains
-            # all the considered joints. This is needed because there might be multiple
-            # <worldbody> elements inside <mujoco>.
-            for wb in mujoco_element.findall(".//worldbody"):
-                if all(
-                    wb.find(f".//joint[@name='{j}']") is not None
-                    for j in considered_joints
-                ):
-                    worldbody_element = wb
-                    break
-
-            if worldbody_element is None:
-                raise RuntimeError(
-                    "Failed to find the <worldbody> element of the model"
-                )
-
             # Light attached to the model
             _ = ET.SubElement(
                 worldbody_element,

src/jaxsim/mujoco/visualizer.py

@@ -54,14 +54,16 @@ def reset(
 
     def render_frame(self, camera_name: str | None = None) -> npt.NDArray:
         """"""
+        camera_name = camera_name or "track"
 
         mujoco.mj_forward(self.model, self.data)
-        self.renderer.update_scene(data=self.data)  # TODO camera name
+        self.renderer.update_scene(data=self.data, camera=camera_name)
 
         return self.renderer.render()
 
     def record_frame(self, camera_name: str | None = None) -> None:
         """"""
+        camera_name = camera_name or "track"
 
         frame = self.render_frame(camera_name=camera_name)
         self.frames.append(frame)