diff --git a/crazyflow/sim/sim.py b/crazyflow/sim/sim.py index e2e3ab2..769dadb 100644 --- a/crazyflow/sim/sim.py +++ b/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) diff --git a/tests/unit/test_render.py b/tests/unit/test_render.py index 7990c7c..7ff224f 100644 --- a/tests/unit/test_render.py +++ b/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