Fix tracking camera

crazyflow/sim/sim.py

@@ -224,8 +224,9 @@ def render(
         self.mj_data.mocap_pos[:] = self.mjx_data.mocap_pos[world, :]
         self.mj_data.mocap_quat[:] = self.mjx_data.mocap_quat[world, :]
         # mj_forward raises on contacts between two static bodies (e.g. drones welded to the world).
-        # We only need poses/lights for rendering
+        # We only need poses, cameras and lights for rendering
         mujoco.mj_kinematics(self.mj_model, self.mj_data)
+        mujoco.mj_comPos(self.mj_model, self.mj_data)
         mujoco.mj_camlight(self.mj_model, self.mj_data)
         return self.viewer.render(mode)
 

tests/unit/test_render.py

@@ -1,4 +1,5 @@
 import mujoco
+import numpy as np
 import pytest
 from conftest import skip_if_headless
 
@@ -32,6 +33,28 @@ def test_render_camera_selection_from_id(cam_id: int):
     sim.close()
 
 
+@pytest.mark.unit
+@pytest.mark.parametrize("cam_name", ["fpv_cam:0", "track_cam:0"])
+@pytest.mark.render
+@skip_if_headless
+def test_drone_camera_follows_drone(cam_name: str):
+    sim = Sim(drone="cf21B_500", n_worlds=1, n_drones=1)
+    cam_id = mujoco.mj_name2id(sim.mj_model, mujoco.mjtObj.mjOBJ_CAMERA, cam_name)
+    sim.render(mode="rgb_array", camera=cam_name)
+    cam_pos_before = sim.mj_data.cam_xpos[cam_id].copy()
+    # Teleport the drone and force a re-sync of the mjx data on the next render
+    offset = np.array([1.0, 2.0, 3.0])
+    states = sim.data.states.replace(pos=sim.data.states.pos + offset)
+    sim.data = sim.data.replace(states=states, core=sim.data.core.replace(mjx_synced=False))
+    sim.render(mode="rgb_array", camera=cam_name)
+    cam_pos_after = sim.mj_data.cam_xpos[cam_id].copy()
+    sim.close()
+    assert np.allclose(cam_pos_after - cam_pos_before, offset, atol=1e-6), (
+        f"Camera {cam_name} did not follow the drone: moved {cam_pos_after - cam_pos_before}, "
+        f"expected {offset}"
+    )
+
+
 @pytest.mark.unit
 @pytest.mark.render
 @skip_if_headless