Fix device tracing inside jit

Author: amacati

crazyflow/sim/data.py

@@ -209,6 +209,8 @@ def create(
 class SimCore:
     freq: int = field(pytree_node=False)
     """Frequency of the simulation."""
+    device: Device = field(pytree_node=False)
+    """Device of the simulation."""
     steps: Array  # (N, 1)
     """Simulation steps taken since the last reset."""
     n_worlds: int = field(pytree_node=False)
@@ -238,6 +240,7 @@ def create(
         rng_key = jax.device_put(rng_key, device)
         return SimCore(
             freq=freq,
+            device=device,
             steps=steps,
             n_worlds=n_worlds,
             n_drones=n_drones,

crazyflow/sim/functional.py

@@ -18,7 +18,7 @@ def state_control(data: SimData, controls: Array) -> SimData:
     """State control function."""
     assert data.controls.mode == Control.state, f"control type {data.controls.mode} not enabled"
     assert controls.shape == (data.core.n_worlds, data.core.n_drones, 13), "controls shape mismatch"
-    controls = to_device(controls, data.core.steps.device)
+    controls = to_device(controls, data.core.device)
     data = data.replace(
         controls=data.controls.replace(state=data.controls.state.replace(staged_cmd=controls))
     )
@@ -36,7 +36,7 @@ def attitude_control(data: SimData, controls: Array) -> SimData:
     """
     assert data.controls.mode == Control.attitude, f"control type {data.controls.mode} not enabled"
     assert controls.shape == (data.core.n_worlds, data.core.n_drones, 4), "controls shape mismatch"
-    controls = to_device(controls, data.core.steps.device)
+    controls = to_device(controls, data.core.device)
     data = data.replace(
         controls=data.controls.replace(attitude=data.controls.attitude.replace(staged_cmd=controls))
     )
@@ -49,7 +49,7 @@ def force_torque_control(data: SimData, controls: Array) -> SimData:
         f"control type {data.controls.mode} not enabled"
     )
     assert controls.shape == (data.core.n_worlds, data.core.n_drones, 4), "controls shape mismatch"
-    controls = to_device(controls, data.core.steps.device)
+    controls = to_device(controls, data.core.device)
     data = data.replace(
         controls=data.controls.replace(
             force_torque=data.controls.force_torque.replace(staged_cmd=controls)
@@ -65,7 +65,7 @@ def rotor_vel_control(data: SimData, controls: Array) -> SimData:
     """
     assert data.controls.mode == Control.rotor_vel, f"control type {data.controls.mode} not enabled"
     assert controls.shape == (data.core.n_worlds, data.core.n_drones, 4), "controls shape mismatch"
-    controls = to_device(controls, data.core.steps.device)
+    controls = to_device(controls, data.core.device)
     return data.replace(controls=data.controls.replace(rotor_vel=controls))
 
 
@@ -81,7 +81,7 @@ def controllable(data: SimData) -> Array:
             control_steps = controls.force_torque.steps
             control_freq = controls.force_torque.freq
         case Control.rotor_vel:
-            return jnp.ones((data.core.n_worlds, 1), dtype=bool, device=data.core.steps.device)
+            return jnp.ones((data.core.n_worlds, 1), dtype=bool, device=data.core.device)
         case _:
             raise NotImplementedError(f"Control mode {data.controls.mode} not implemented")
     return _controllable(data.core.steps, data.core.freq, control_steps, control_freq)

tests/unit/test_functional.py

@@ -1,5 +1,8 @@
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
+import jax
 import jax.numpy as jnp
 import numpy as np
 import pytest
@@ -8,6 +11,9 @@
 from crazyflow.control import Control
 from crazyflow.sim import Sim
 
+if TYPE_CHECKING:
+    from crazyflow.sim.data import SimData
+
 
 @pytest.mark.unit
 def test_functional_resets():
@@ -80,6 +86,14 @@ def test_functional_attitude_control(attitude_freq: int):
         if i == 0:  # Make world 2 asynchronous
             data = reset_fn(data, default_data, np.array([False, True]))
 
+    # Check if we can apply inside of jax jit which does not permit device tracking. See
+    # https://github.com/jax-ml/jax/issues/26000 for more context.
+    @jax.jit
+    def apply_control(data: SimData, cmd: jnp.ndarray) -> SimData:
+        return F.attitude_control(data, cmd)
+
+    jax.block_until_ready(apply_control(data, cmd))
+
 
 @pytest.mark.unit
 def test_functional_attitude_control_device(device: str):
@@ -133,6 +147,12 @@ def test_functional_state_control(state_freq: int):
         if i == 0:  # Make world 2 asynchronous
             data = reset_fn(data, default_data, np.array([False, True]))
 
+    @jax.jit
+    def apply_control(data: SimData, cmd: jnp.ndarray) -> SimData:
+        return F.state_control(data, cmd)
+
+    jax.block_until_ready(apply_control(data, cmd))
+
 
 @pytest.mark.unit
 def test_functional_state_control_device(device: str):