Fix rendering with non-default width/height

Author: amacati

crazyflow/sim/sim.py

@@ -306,17 +306,24 @@ def thrust_control(self, cmd: Array):
         self.data = self.data.replace(controls=self.data.controls.replace(thrust=controls))
 
     def render(
-        self, mode: str | None = "human", world: int = 0, default_cam_config: dict | None = None
+        self,
+        mode: str | None = "human",
+        world: int = 0,
+        default_cam_config: dict | None = None,
+        width: int = 640,
+        height: int = 480,
     ) -> NDArray | None:
         if self.viewer is None:
             patch_viewer()
+            self.mj_model.vis.global_.offwidth = width
+            self.mj_model.vis.global_.offheight = height
             self.viewer = MujocoRenderer(
                 self.mj_model,
                 self.mj_data,
                 max_geom=self.max_visual_geom,
                 default_cam_config=default_cam_config,
-                height=480,
-                width=640,
+                height=height,
+                width=width,
             )
         self.mj_data.qpos[:] = self.data.mjx_data.qpos[world, :]
         self.mj_data.mocap_pos[:] = self.data.mjx_data.mocap_pos[world, :]

tests/unit/test_sim.py

@@ -270,9 +270,12 @@ def test_render_human(device: str):
 def test_render_rgb_array(device: str):
     skip_unavailable_device(device)
     sim = Sim(n_worlds=2, device=device)
-    img = sim.render(mode="rgb_array")
+    img = sim.render(mode="rgb_array", width=1024, height=1024)
     assert isinstance(img, np.ndarray), "Image must be a numpy array"
-    assert img.ndim == 3, "Image must be 3D"
+    assert img.shape == (1024, 1024, 3), f"Unexpected image shape {img.shape}"
+    # Check if mj_model.vis.global_.offwidth is set correctly
+    assert not all(img[0, 0, :] == 0), "Image contains black patches"
+    assert not all(img[-1, -1, :] == 0), "Image contains black patches"
 
 
 @pytest.mark.unit