Add functional API for controllers

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: