Fix linting. Add tests

Author: amacati

crazyflow/sim/sim.py

@@ -1,6 +1,8 @@
+from __future__ import annotations
+
 from functools import partial
 from pathlib import Path
-from typing import Any, Callable, Optional
+from typing import TYPE_CHECKING, Any, Callable
 
 import jax
 import jax.numpy as jnp
@@ -10,7 +12,6 @@
 from gymnasium.envs.mujoco.mujoco_rendering import MujocoRenderer
 from jax import Array, Device
 from jax.scipy.spatial.transform import Rotation as R
-from mujoco.mjx import Data, Model
 
 from crazyflow.constants import J_INV, MASS, SIGN_MIX_MATRIX, J
 from crazyflow.control.control import Control, attitude2rpm, pwm2rpm, state2attitude, thrust2pwm
@@ -28,6 +29,10 @@
 from crazyflow.sim.structs import SimControls, SimCore, SimData, SimParams, SimState, SimStateDeriv
 from crazyflow.utils import grid_2d, leaf_replace, patch_viewer, pytree_replace, to_device
 
+if TYPE_CHECKING:
+    from mujoco.mjx import Data, Model
+    from numpy.typing import NDArray
+
 
 class Sim:
     default_path = Path(__file__).parents[1] / "models/cf2/scene.xml"
@@ -300,10 +305,19 @@ def thrust_control(self, cmd: Array):
         controls = to_device(cmd, self.device)
         self.data = self.data.replace(controls=self.data.controls.replace(thrust=controls))
 
-    def render(self, mode: Optional[str] = "human", world: Optional[int] = 0, default_cam_config: Optional[dict] = None):
+    def render(
+        self, mode: str | None = "human", world: int = 0, default_cam_config: dict | None = None
+    ) -> NDArray | None:
         if self.viewer is None:
             patch_viewer()
-            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)
+            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,
+            )
         self.mj_data.qpos[:] = self.data.mjx_data.qpos[world, :]
         self.mj_data.mocap_pos[:] = self.data.mjx_data.mocap_pos[world, :]
         self.mj_data.mocap_quat[:] = self.data.mjx_data.mocap_quat[world, :]

tests/unit/test_sim.py

@@ -258,13 +258,23 @@ def test_sim_state_control_device(device: str):
 
 @pytest.mark.parametrize("device", ["gpu", "cpu"])
 @pytest.mark.render
-def test_render(device: str):
+def test_render_human(device: str):
     skip_unavailable_device(device)
     sim = Sim(device=device)
     sim.render()
     sim.viewer.close()
 
 
+# Do not mark as render to ensure it runs by default. This function will not open a viewer.
+@pytest.mark.parametrize("device", ["gpu", "cpu"])
+def test_render_rgb_array(device: str):
+    skip_unavailable_device(device)
+    sim = Sim(n_worlds=2, device=device)
+    img = sim.render(mode="rgb_array")
+    assert isinstance(img, np.ndarray), "Image must be a numpy array"
+    assert img.ndim == 3, "Image must be 3D"
+
+
 @pytest.mark.unit
 @pytest.mark.parametrize("device", ["gpu", "cpu"])
 def test_device(device: str):