Fix camera settings in render() (#51)

Author: ratheron

crazyflow/sim/sim.py

@@ -157,11 +157,11 @@ def render(
     ) -> NDArray | None:
         if self.viewer is None:
             if isinstance(camera, str):
-                cam_id, cam_name = None, camera
+                cam_id = mujoco.mj_name2id(self.mj_model, mujoco.mjtObj.mjOBJ_CAMERA, camera)
+                assert cam_id > -1, f"Camera name '{camera}' not found in the model."
             elif isinstance(camera, int):
-                cam_id, cam_name = camera, None
-                if cam_id < -1:
-                    raise ValueError(f"camera id must be >=-1, was {cam_id}")
+                cam_id = camera
+                assert cam_id >= -1, f"camera id must be >=-1, was {cam_id}"
             else:
                 raise TypeError("camera argument must be integer or string")
             self.mj_model.vis.global_.offwidth = width
@@ -174,8 +174,14 @@ def render(
                 height=height,
                 width=width,
                 camera_id=cam_id,
-                camera_name=cam_name,
             )
+            # In human mode, cam_id is set to -1, so we force it to the desired value
+            if mode == "human" and cam_id > -1:
+                # Render one frame to force mj to create the viewer
+                self.viewer.render(mode)
+                self.viewer.viewer.cam.fixedcamid = cam_id
+                self.viewer.viewer.cam.type = mujoco.mjtCamera.mjCAMERA_FIXED
+
         self.mj_data.qpos[:] = self.mjx_data.qpos[world, :]
         self.mj_data.mocap_pos[:] = self.mjx_data.mocap_pos[world, :]
         self.mj_data.mocap_quat[:] = self.mjx_data.mocap_quat[world, :]

tests/unit/test_render.py

@@ -0,0 +1,43 @@
+import mujoco
+import pytest
+from conftest import skip_if_headless
+
+from crazyflow import Sim
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("cam_name", ["fpv_cam:0", "track_cam:0", "fpv_cam:1", "track_cam:1"])
+@pytest.mark.render
+@skip_if_headless
+def test_render_camera_selection_from_name(cam_name: str):
+    sim = Sim(drone_model="cf21B_500", n_drones=2)
+    cam_id = mujoco.mj_name2id(sim.mj_model, mujoco.mjtObj.mjOBJ_CAMERA, cam_name)
+    sim.render(mode="human", camera=cam_name)
+    viewer_cam = sim.viewer.viewer.cam
+    assert viewer_cam.type == mujoco.mjtCamera.mjCAMERA_FIXED, "Camera type was not set to FIXED"
+    assert viewer_cam.fixedcamid == cam_id, f"Expected cam ID {cam_id}, got {viewer_cam.fixedcamid}"
+    sim.close()
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("cam_id", [0, 1, 2, 3])
+@pytest.mark.render
+@skip_if_headless
+def test_render_camera_selection_from_id(cam_id: int):
+    sim = Sim(drone_model="cf21B_500", n_drones=2)
+    sim.render(mode="human", camera=cam_id)
+    viewer_cam = sim.viewer.viewer.cam
+    assert viewer_cam.type == mujoco.mjtCamera.mjCAMERA_FIXED, "Camera type was not set to FIXED"
+    assert viewer_cam.fixedcamid == cam_id, f"Expected cam ID {cam_id}, got {viewer_cam.fixedcamid}"
+    sim.close()
+
+
+@pytest.mark.unit
+@pytest.mark.render
+@skip_if_headless
+def test_render_free_camera():
+    sim = Sim(drone_model="cf21B_500", n_drones=2)
+    sim.render(mode="human")
+    viewer_cam = sim.viewer.viewer.cam
+    assert viewer_cam.type == mujoco.mjtCamera.mjCAMERA_FREE, "Camera type was not set to FREE"
+    sim.close()