Adapt to mujoco 3.3 spec API changes. Fix typing. Fix tests.

Author: amacati

lsy_drone_racing/control/attitude_controller.py

@@ -129,7 +129,7 @@ def compute_control(
         R_desired = np.vstack([x_axis_desired, y_axis_desired, z_axis_desired]).T
         euler_desired = R.from_matrix(R_desired).as_euler("xyz", degrees=False)
         thrust_desired, euler_desired
-        return np.concatenate([[thrust_desired], euler_desired])
+        return np.concatenate([[thrust_desired], euler_desired], dtype=np.float32)
 
     def step_callback(
         self,

lsy_drone_racing/control/trajectory_controller.py

@@ -71,7 +71,7 @@ def compute_control(
                 array.
         """
         target_pos = self.trajectory(min(self._tick / self._freq, self.t_total))
-        return np.concatenate((target_pos, np.zeros(10)))
+        return np.concatenate((target_pos, np.zeros(10)), dtype=np.float32)
 
     def step_callback(
         self,

lsy_drone_racing/envs/drone_race.py

@@ -11,9 +11,8 @@
 from lsy_drone_racing.envs.race_core import RaceCoreEnv, build_action_space, build_observation_space
 
 if TYPE_CHECKING:
-    import numpy as np
+    from jax import Array
     from ml_collections import ConfigDict
-    from numpy.typing import NDArray
 
 
 class DroneRaceEnv(RaceCoreEnv, Env):
@@ -83,7 +82,7 @@ def reset(self, seed: int | None = None, options: dict | None = None) -> tuple[d
         info = {k: v[0, 0] for k, v in info.items()}
         return obs, info
 
-    def step(self, action: NDArray[np.floating]) -> tuple[dict, float, bool, bool, dict]:
+    def step(self, action: Array) -> tuple[dict, float, bool, bool, dict]:
         """Step the environment.
 
         Args:
@@ -168,9 +167,7 @@ def reset(self, seed: int | None = None, options: dict | None = None) -> tuple[d
         info = {k: v[:, 0] for k, v in info.items()}
         return obs, info
 
-    def step(
-        self, action: NDArray[np.floating]
-    ) -> tuple[dict, NDArray[np.floating], NDArray[np.bool_], NDArray[np.bool_], dict]:
+    def step(self, action: Array) -> tuple[dict, Array, Array, Array, dict]:
         """Step the environment in all worlds.
 
         Args:

lsy_drone_racing/envs/multi_drone_race.py

@@ -1,16 +1,19 @@
 """Multi-agent drone racing environments."""
 
-from typing import Literal
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal
 
-import numpy as np
 from gymnasium import Env
 from gymnasium.vector import VectorEnv
 from gymnasium.vector.utils import batch_space
-from ml_collections import ConfigDict
-from numpy.typing import NDArray
 
 from lsy_drone_racing.envs.race_core import RaceCoreEnv, build_action_space, build_observation_space
 
+if TYPE_CHECKING:
+    from jax import Array
+    from ml_collections import ConfigDict
+
 
 class MultiDroneRaceEnv(RaceCoreEnv, Env):
     """Multi-agent drone racing environment.
@@ -87,9 +90,7 @@ def reset(self, seed: int | None = None, options: dict | None = None) -> tuple[d
         info = {k: v[0] for k, v in info.items()}
         return obs, info
 
-    def step(
-        self, action: NDArray[np.floating]
-    ) -> tuple[dict, NDArray[np.floating], NDArray[np.bool_], NDArray[np.bool_], dict]:
+    def step(self, action: Array) -> tuple[dict, Array, Array, Array, dict]:
         """Step the environment for all drones.
 
         Args:
@@ -180,9 +181,7 @@ def reset(self, seed: int | None = None, options: dict | None = None) -> tuple[d
         """
         return self._reset(seed=seed, options=options)
 
-    def step(
-        self, action: NDArray[np.floating]
-    ) -> tuple[dict, NDArray[np.floating], NDArray[np.bool_], NDArray[np.bool_], dict]:
+    def step(self, action: Array) -> tuple[dict, Array, Array, Array, dict]:
         """Step the environment for all drones.
 
         Args:

lsy_drone_racing/envs/race_core.py

@@ -19,7 +19,6 @@
 from crazyflow.sim.symbolic import symbolic_attitude
 from flax.struct import dataclass
 from gymnasium import spaces
-from jax.scipy.spatial.transform import Rotation as JaxR
 from scipy.spatial.transform import Rotation as R
 
 from lsy_drone_racing.envs.randomize import (
@@ -264,7 +263,7 @@ def __init__(
 
     def _reset(
         self, *, seed: int | None = None, options: dict | None = None, mask: Array | None = None
-    ) -> tuple[dict[str, NDArray[np.floating]], dict]:
+    ) -> tuple[dict[str, Array], dict]:
         """Reset the environment.
 
         Args:
@@ -287,9 +286,7 @@ def _reset(
         self.data = self._reset_env_data(self.data, self.sim.data.states.pos, mask)
         return self.obs(), self.info()
 
-    def _step(
-        self, action: NDArray[np.floating]
-    ) -> tuple[dict[str, NDArray[np.floating]], float, bool, bool, dict]:
+    def _step(self, action: Array) -> tuple[dict[str, Array], float, bool, bool, dict]:
         """Step the firmware_wrapper class and its environment.
 
         This function should be called once at the rate of ctrl_freq. Step processes and high level
@@ -319,7 +316,7 @@ def _step(
             self._reset(mask=marked_for_reset)
         return self.obs(), self.reward(), self.terminated(), self.truncated(), self.info()
 
-    def apply_action(self, action: NDArray[np.floating]):
+    def apply_action(self, action: Array):
         """Apply the commanded state action to the simulation."""
         action = action.reshape((self.sim.n_worlds, self.sim.n_drones, -1))
         if "action" in self.disturbances:
@@ -342,7 +339,7 @@ def close(self):
         """Close the environment by stopping the drone and landing back at the starting position."""
         self.sim.close()
 
-    def obs(self) -> dict[str, NDArray[np.floating]]:
+    def obs(self) -> dict[str, Array]:
         """Return the observation of the environment."""
         # Add the gate and obstacle poses to the info. If gates or obstacles are in sensor range,
         # use the actual pose, otherwise use the nominal pose.
@@ -564,13 +561,19 @@ def _load_track_into_sim(self, gate_spec: MjSpec, obstacle_spec: MjSpec):
         frame = self.sim.spec.worldbody.add_frame()
         n_gates, n_obstacles = len(self.gates["pos"]), len(self.obstacles["pos"])
         for i in range(n_gates):
-            gate = frame.attach_body(gate_spec.find_body("gate"), "", f":{i}")
+            gate_body = gate_spec.body("gate")
+            if gate_body is None:
+                raise ValueError("Gate body not found in gate spec")
+            gate = frame.attach_body(gate_body, "", f":{i}")
             gate.pos = self.gates["pos"][i]
             # Convert from scipy order to MuJoCo order
             gate.quat = self.gates["quat"][i][[3, 0, 1, 2]]
             gate.mocap = True  # Make mocap to modify the position of static bodies during sim
         for i in range(n_obstacles):
-            obstacle = frame.attach_body(obstacle_spec.find_body("obstacle"), "", f":{i}")
+            obstacle_body = obstacle_spec.body("obstacle")
+            if obstacle_body is None:
+                raise ValueError("Obstacle body not found in obstacle spec")
+            obstacle = frame.attach_body(obstacle_body, "", f":{i}")
             obstacle.pos = self.obstacles["pos"][i]
             obstacle.mocap = True
         self.sim.build(data=False, default_data=False)

pyproject.toml

@@ -17,11 +17,12 @@ classifiers = [
     "Intended Audience :: Science/Research",
 ]
 
+# TODO: Include crazyflow once it's installable via pip
 dependencies = [
     "fire >= 0.6.0",
     "numpy >= 1.24.1, < 2.0.0",
-    "PyYAML >= 6.0.1",  # TODO: Remove after removing crazyswarm dependency
-    "rospkg >= 1.5.1",  # TODO: Remove after moving to cflib
+    "PyYAML >= 6.0.1",          # TODO: Remove after removing crazyswarm dependency
+    "rospkg >= 1.5.1",          # TODO: Remove after moving to cflib
     "scipy >= 1.10.1",
     "gymnasium >= 1.0.0",
     "toml >= 0.10.2",

scripts/sim.py

@@ -15,6 +15,7 @@
 
 import fire
 import gymnasium
+from gymnasium.wrappers.jax_to_numpy import JaxToNumpy
 
 from lsy_drone_racing.utils import load_config, load_controller
 
@@ -68,6 +69,7 @@ def simulate(
         random_resets=config.env.random_resets,
         seed=config.env.seed,
     )
+    env = JaxToNumpy(env)
 
     ep_times = []
     for _ in range(n_runs):  # Run n_runs episodes with the controller

tests/conftest.py

@@ -0,0 +1,7 @@
+import jax
+
+jax.config.update("jax_compilation_cache_dir", "/tmp/jax_cache")
+jax.config.update("jax_persistent_cache_min_entry_size_bytes", -1)
+jax.config.update("jax_persistent_cache_min_compile_time_secs", 0)
+# Do not enable XLA caches, crashes PyTest
+# jax.config.update("jax_persistent_cache_enable_xla_caches", "all")

tests/integration/test_controllers.py

@@ -3,6 +3,7 @@
 import gymnasium
 import numpy as np
 import pytest
+from gymnasium.wrappers.jax_to_numpy import JaxToNumpy
 
 from lsy_drone_racing.utils import load_config, load_controller
 
@@ -27,6 +28,8 @@ def test_controllers(controller_file: str):
         random_resets=config.env.random_resets,
         seed=config.env.seed,
     )
+    env = JaxToNumpy(env)
+
     obs, info = env.reset()
     ctrl = ctrl_cls(obs, info, config)
     while True:
@@ -59,6 +62,8 @@ def test_attitude_controller(physics: str):
         random_resets=config.env.random_resets,
         seed=config.env.seed,
     )
+    env = JaxToNumpy(env)
+
     obs, info = env.reset()
     ctrl = ctrl_cls(obs, info, config)
     while True:
@@ -98,6 +103,8 @@ def test_trajectory_controller_finish(yaw: float, physics: str):
         random_resets=config.env.random_resets,
         seed=config.env.seed,
     )
+    env = JaxToNumpy(env)
+
     obs, info = env.reset()
     ctrl = ctrl_cls(obs, info, config)
     while True:

tests/unit/envs/test_envs.py

@@ -4,6 +4,7 @@
 import gymnasium
 import pytest
 from gymnasium.utils.env_checker import check_env
+from gymnasium.wrappers.jax_to_numpy import JaxToNumpy
 
 from lsy_drone_racing.utils import load_config
 
@@ -31,4 +32,4 @@ def test_passive_checker_wrapper_warnings(action_space: str):
             seed=config.env.seed,
             disable_env_checker=False,
         )
-        check_env(env.unwrapped)
+        check_env(JaxToNumpy(env.unwrapped))