Add functional control API (#50)

Author: amacati

crazyflow/sim/data.py

@@ -102,6 +102,8 @@ class ControlData(typing.Protocol):
 
 @dataclass
 class SimControls:
+    mode: Control = field(pytree_node=False)
+    """Control mode of the simulation."""
     state: ControlData | None
     """State control data."""
     attitude: ControlData | None
@@ -136,7 +138,11 @@ def create(
                     n_worlds, n_drones, force_torque_freq, drone_model, device
                 )
                 return SimControls(
-                    state=state, attitude=attitude, force_torque=force_torque, rotor_vel=rotor_vel
+                    mode=control,
+                    state=state,
+                    attitude=attitude,
+                    force_torque=force_torque,
+                    rotor_vel=rotor_vel,
                 )
             case Control.attitude:
                 attitude = attitude = MellingerAttitudeData.create(
@@ -146,14 +152,22 @@ def create(
                     n_worlds, n_drones, force_torque_freq, drone_model, device
                 )
                 return SimControls(
-                    state=None, attitude=attitude, force_torque=force_torque, rotor_vel=rotor_vel
+                    mode=control,
+                    state=None,
+                    attitude=attitude,
+                    force_torque=force_torque,
+                    rotor_vel=rotor_vel,
                 )
             case Control.force_torque:
                 force_torque = MellingerForceTorqueData.create(
                     n_worlds, n_drones, force_torque_freq, drone_model, device
                 )
                 return SimControls(
-                    state=None, attitude=None, force_torque=force_torque, rotor_vel=rotor_vel
+                    mode=control,
+                    state=None,
+                    attitude=None,
+                    force_torque=force_torque,
+                    rotor_vel=rotor_vel,
                 )
             case _:
                 raise ValueError(f"Control mode {control} not implemented")

crazyflow/sim/functional.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from crazyflow.control import Control
+from crazyflow.control.control import controllable as _controllable
+from crazyflow.utils import to_device
+
+if TYPE_CHECKING:
+    from jax import Array
+
+    from crazyflow.sim.data import SimData
+
+
+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)
+    data = data.replace(
+        controls=data.controls.replace(state=data.controls.state.replace(staged_cmd=controls))
+    )
+    return data
+
+
+def attitude_control(data: SimData, controls: Array) -> SimData:
+    """Attitude control function.
+
+    We need to stage the attitude controls because the sys_id physics mode operates directly on
+    the attitude controls. If we were to directly update the controls, this would effectively
+    bypass the control frequency and run the attitude controller at the physics update rate. By
+    staging the controls, we ensure that the physics module sees the old controls until the
+    controller updates at its correct frequency.
+    """
+    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)
+    data = data.replace(
+        controls=data.controls.replace(attitude=data.controls.attitude.replace(staged_cmd=controls))
+    )
+    return data
+
+
+def force_torque_control(data: SimData, controls: Array) -> SimData:
+    """Force-torque control function."""
+    assert data.controls.mode == Control.force_torque, (
+        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)
+    data = data.replace(
+        controls=data.controls.replace(
+            force_torque=data.controls.force_torque.replace(staged_cmd=controls)
+        )
+    )
+    return data
+
+
+def controllable(data: SimData) -> Array:
+    """Check which worlds can currently update their controllers."""
+    controls = data.controls
+    match data.controls.mode:
+        case Control.state:
+            control_steps, control_freq = controls.state.steps, controls.state.freq
+        case Control.attitude:
+            control_steps, control_freq = controls.attitude.steps, controls.attitude.freq
+        case Control.force_torque:
+            control_steps = controls.force_torque.steps
+            control_freq = controls.force_torque.freq
+        case _:
+            raise NotImplementedError(f"Control mode {data.controls.mode} not implemented")
+    return _controllable(data.core.steps, data.core.freq, control_steps, control_freq)

crazyflow/sim/sim.py

@@ -18,6 +18,7 @@
 from gymnasium.envs.mujoco.mujoco_rendering import MujocoRenderer
 from jax import Array, Device
 
+import crazyflow.sim.functional as F
 from crazyflow.control.control import Control, controllable
 from crazyflow.exception import ConfigError, NotInitializedError
 from crazyflow.sim.data import SimControls, SimCore, SimData, SimParams, SimState, SimStateDeriv
@@ -29,7 +30,7 @@
     so_rpy_rotor_drag_physics,
     so_rpy_rotor_physics,
 )
-from crazyflow.utils import grid_2d, leaf_replace, pytree_replace, to_device
+from crazyflow.utils import grid_2d, leaf_replace, pytree_replace
 
 if TYPE_CHECKING:
     from mujoco.mjx import Data, Model
@@ -134,45 +135,15 @@ def step(self, n_steps: int = 1):
 
     def state_control(self, controls: Array):
         """Set the desired state for all drones in all worlds."""
-        assert controls.shape == (self.n_worlds, self.n_drones, 13), "controls shape mismatch"
-        assert self.control == Control.state, "State control is not enabled by the sim config"
-        controls = to_device(controls, self.device)
-        self.data = self.data.replace(
-            controls=self.data.controls.replace(
-                state=self.data.controls.state.replace(staged_cmd=controls)
-            )
-        )
+        self.data = F.state_control(self.data, controls)
 
     def attitude_control(self, controls: Array):
-        """Set the desired attitude for all drones in all worlds.
+        """Set the desired attitude for all drones in all worlds."""
+        self.data = F.attitude_control(self.data, controls)
 
-        We need to stage the attitude controls because the sys_id physics mode operates directly on
-        the attitude controls. If we were to directly update the controls, this would effectively
-        bypass the control frequency and run the attitude controller at the physics update rate. By
-        staging the controls, we ensure that the physics module sees the old controls until the
-        controller updates at its correct frequency.
-        """
-        assert controls.shape == (self.n_worlds, self.n_drones, 4), "controls shape mismatch"
-        assert self.control == Control.attitude, "Attitude control is not enabled by the sim config"
-        controls = to_device(controls, self.device)
-        self.data = self.data.replace(
-            controls=self.data.controls.replace(
-                attitude=self.data.controls.attitude.replace(staged_cmd=controls)
-            )
-        )
-
-    def force_torque_control(self, cmd: Array):
+    def force_torque_control(self, controls: Array):
         """Set the desired force and torque for all drones in all worlds."""
-        assert cmd.shape == (self.n_worlds, self.n_drones, 4), "Command shape mismatch"
-        assert self.control == Control.force_torque, (
-            "Force-torque control is not enabled by the sim config"
-        )
-        controls = to_device(cmd, self.device)
-        self.data = self.data.replace(
-            controls=self.data.controls.replace(
-                force_torque=self.data.controls.force_torque.replace(staged_cmd=controls)
-            )
-        )
+        self.data = F.force_torque_control(self.data, controls)
 
     @requires_mujoco_sync
     def render(
@@ -408,18 +379,7 @@ def controllable(self) -> Array:
         as soon as the controller frequency allows for an update. Successive control updates that
         happen before the staged buffers are applied overwrite the desired values.
         """
-        controls = self.data.controls
-        match self.control:
-            case Control.state:
-                control_steps, control_freq = controls.state.steps, controls.state.freq
-            case Control.attitude:
-                control_steps, control_freq = controls.attitude.steps, controls.attitude.freq
-            case Control.force_torque:
-                control_steps = controls.force_torque.steps
-                control_freq = controls.force_torque.freq
-            case _:
-                raise NotImplementedError(f"Control mode {self.control} not implemented")
-        return controllable(self.data.core.steps, self.data.core.freq, control_steps, control_freq)
+        return F.controllable(self.data)
 
     @requires_mujoco_sync
     def contacts(self, body: str | None = None) -> Array:

pixi.lock

@@ -0,0 +1,147 @@
+from __future__ import annotations
+
+import jax.numpy as jnp
+import numpy as np
+import pytest
+
+import crazyflow.sim.functional as F
+from crazyflow.control import Control
+from crazyflow.sim import Sim
+
+
+@pytest.mark.unit
+def test_functional_resets():
+    """Test that the functional API works as expected for resets."""
+    sim = Sim()
+    data, default_data = sim.build_data(), sim.build_default_data()
+    reset_fn = sim.build_reset_fn()
+    data = data.replace(states=data.states.replace(pos=jnp.ones_like(data.states.pos)))
+    sim.data = data
+    # Test types
+    assert callable(reset_fn), "reset_fn must be a pure function"
+    assert not hasattr(reset_fn, "__self__"), "reset_fn must not be a bound method"
+    # Test the reset runs as expected
+    data_reset = reset_fn(data, default_data, None)
+    assert jnp.all(data_reset.states.pos == 0), "reset_fn did not reset positions to zero"
+    data_reset = reset_fn(data, default_data, jnp.array([True] * sim.n_worlds))
+    assert jnp.all(data_reset.states.pos == 0), "reset_fn did not reset positions to zero with mask"
+
+
+@pytest.mark.unit
+def test_functional_steps():
+    """Test that the functional API works as expected for steps."""
+    sim = Sim()
+    data = sim.build_data()
+    data = data.replace(states=data.states.replace(pos=jnp.ones_like(data.states.pos)))
+    step_fn = sim.build_step_fn()
+    # Test types
+    assert callable(step_fn), "step_fn must be a pure function"
+    assert not hasattr(step_fn, "__self__"), "step_fn must not be a bound method"
+    # Test the step function runs as expected
+    data_step = step_fn(data, 5)
+    assert jnp.all(data_step.states.pos[..., 2] < 1), "step_fn did not step correctly"
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("attitude_freq", [33, 50, 100, 200])
+def test_functional_attitude_control(attitude_freq: int):
+    """Test that functional attitude control respects frequency and applies commands correctly.
+
+    Ported from test_attitude_control in test_sim.py.
+    """
+    sim = Sim(n_worlds=2, n_drones=3, control="attitude", freq=100, attitude_freq=attitude_freq)
+
+    data = sim.build_data()
+    default_data = sim.build_default_data()
+    reset_fn = sim.build_reset_fn()
+    step_fn = sim.build_step_fn()
+
+    can_control_1 = np.arange(6) * attitude_freq % sim.freq < attitude_freq
+    can_control_2 = np.array([0, 0, 1, 2, 3, 4]) * attitude_freq % sim.freq < attitude_freq
+    for i in range(6):
+        cmd = np.random.rand(sim.n_worlds, sim.n_drones, 4)
+        # Check controllable status
+        controllable = F.controllable(data)
+        assert jnp.all(controllable[0] == can_control_1[i]), f"Controllable 1 mismatch at t={i}"
+        assert jnp.all(controllable[1] == can_control_2[i]), f"Controllable 2 mismatch at t={i}"
+        # Apply control
+        data = F.attitude_control(data, cmd)
+        data = step_fn(data, 1)
+        sim_cmd = data.controls.attitude.cmd[0]
+        if can_control_1[i]:
+            assert jnp.all(sim_cmd == cmd[0]), f"Controls do not match at t={i}"
+        else:
+            assert not jnp.all(sim_cmd == cmd[0]), f"Controls shouldn't match at t={i}"
+        sim_cmd = data.controls.attitude.cmd[1]
+        if can_control_2[i]:
+            assert jnp.all(sim_cmd == cmd[1]), f"Controls do not match at t={i}"
+        else:
+            assert not jnp.all(sim_cmd == cmd[1]), f"Controls shouldn't match at t={i}"
+        if i == 0:  # Make world 2 asynchronous
+            data = reset_fn(data, default_data, np.array([False, True]))
+
+
+@pytest.mark.unit
+def test_functional_attitude_control_device(device: str):
+    """Test that functional attitude control maintains JAX arrays on correct device."""
+    sim = Sim(n_worlds=2, n_drones=3, control=Control.attitude, device=device)
+    data = sim.build_data()
+    cmd = np.random.rand(sim.n_worlds, sim.n_drones, 4)
+    data = F.attitude_control(data, cmd)
+    controls = data.controls.attitude
+    assert isinstance(controls.staged_cmd, jnp.ndarray), "Buffers must remain JAX arrays"
+    assert jnp.all(controls.staged_cmd == cmd), "Buffers must match command"
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("state_freq", [33, 50, 100, 200])
+def test_functional_state_control(state_freq: int):
+    """Test that functional state control respects frequency and applies commands correctly."""
+    sim = Sim(n_worlds=2, n_drones=3, control=Control.state, freq=100, state_freq=state_freq)
+
+    data = sim.build_data()
+    default_data = sim.build_default_data()
+    reset_fn = sim.build_reset_fn()
+    step_fn = sim.build_step_fn()
+
+    can_control_1 = np.arange(6) * state_freq % sim.freq < state_freq
+    can_control_2 = np.array([0, 0, 1, 2, 3, 4]) * state_freq % sim.freq < state_freq
+
+    for i in range(6):
+        cmd = np.random.rand(sim.n_worlds, sim.n_drones, 13)
+        # Check controllable status
+        controllable = F.controllable(data)
+        assert jnp.all(controllable[0] == can_control_1[i]), f"Controllable 1 mismatch at t={i}"
+        assert jnp.all(controllable[1] == can_control_2[i]), f"Controllable 2 mismatch at t={i}"
+        # Apply control
+        data = F.state_control(data, cmd)
+        last_attitude = data.controls.attitude.staged_cmd
+        data = step_fn(data, 1)
+        attitude = data.controls.attitude.staged_cmd
+
+        last_att, att = last_attitude[0], attitude[0]
+        if can_control_1[i]:
+            assert not jnp.all(att == last_att), f"Controls haven't been applied at t={i}"
+        else:
+            assert jnp.all(att == last_att), f"Controls should be unchanged at t={i}"
+
+        last_att, att = last_attitude[1], attitude[1]
+        if can_control_2[i]:
+            assert not jnp.all(att == last_att), f"Controls haven't been applied at t={i}"
+        else:
+            assert jnp.all(att == last_att), f"Controls should be unchanged at t={i}"
+        if i == 0:  # Make world 2 asynchronous
+            data = reset_fn(data, default_data, np.array([False, True]))
+
+
+@pytest.mark.unit
+def test_functional_state_control_device(device: str):
+    """Test that functional state control maintains JAX arrays on correct device."""
+    sim = Sim(n_worlds=2, n_drones=3, control=Control.state, device=device)
+    data = sim.build_data()
+    cmd = np.random.rand(sim.n_worlds, sim.n_drones, 13)
+    data = F.state_control(data, cmd)
+    controls = data.controls.state
+    assert isinstance(controls.cmd, jnp.ndarray), "Buffers must remain JAX arrays"
+    assert isinstance(controls.staged_cmd, jnp.ndarray), "Buffers must remain JAX arrays"
+    assert jnp.all(controls.staged_cmd == cmd), "Buffers must match command"