Merge pull request #55 from utiasDSL/feat.rotor_vel_ctrl

Add rotor vel control option

Author: ratheron

crazyflow/control/control.py

@@ -36,7 +36,13 @@ class Control(str, Enum):
         Recommended frequency is >=100 Hz.
     """
     force_torque = "force_torque"
-    """Force and torque control takes [fx, fy, fz, tx, ty, tz].
+    """Force and torque control takes [fc, tx, ty, tz].
+
+    Note:
+        Recommended frequency is >=500 Hz.
+    """
+    rotor_vel = "rotor_vel"
+    """Rotor velocity control takes [w1, w2, w3, w4] in RPMs.
 
     Note:
         Recommended frequency is >=500 Hz.

crazyflow/sim/data.py

@@ -169,6 +169,10 @@ def create(
                     force_torque=force_torque,
                     rotor_vel=rotor_vel,
                 )
+            case Control.rotor_vel:
+                return SimControls(
+                    mode=control, state=None, attitude=None, force_torque=None, rotor_vel=rotor_vel
+                )
             case _:
                 raise ValueError(f"Control mode {control} not implemented")
 

crazyflow/sim/functional.py

@@ -2,6 +2,8 @@
 
 from typing import TYPE_CHECKING
 
+import jax.numpy as jnp
+
 from crazyflow.control import Control
 from crazyflow.control.control import controllable as _controllable
 from crazyflow.utils import to_device
@@ -56,6 +58,17 @@ def force_torque_control(data: SimData, controls: Array) -> SimData:
     return data
 
 
+def rotor_vel_control(data: SimData, controls: Array) -> SimData:
+    """Rotor velocity control.
+
+    Directly set the desired rotor velocities of the drone.
+    """
+    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)
+    return data.replace(controls=data.controls.replace(rotor_vel=controls))
+
+
 def controllable(data: SimData) -> Array:
     """Check which worlds can currently update their controllers."""
     controls = data.controls
@@ -67,6 +80,8 @@ def controllable(data: SimData) -> Array:
         case Control.force_torque:
             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)
         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

@@ -77,8 +77,9 @@ def __init__(
     ):
         assert Physics(physics) in Physics, f"Physics mode {physics} not implemented"
         assert Control(control) in Control, f"Control mode {control} not implemented"
-        if physics != Physics.first_principles and control == Control.force_torque:
-            raise ConfigError("Force-torque control requires first principles physics")
+        if physics != Physics.first_principles:
+            if control in (Control.force_torque, Control.rotor_vel):
+                raise ConfigError(f"Control mode {control} requires first principles physics")
         if freq > 10_000 and not jax.config.jax_enable_x64:
             raise ConfigError("High frequency simulations require double precision mode")
         self.physics = physics
@@ -145,6 +146,10 @@ def force_torque_control(self, controls: Array):
         """Set the desired force and torque for all drones in all worlds."""
         self.data = F.force_torque_control(self.data, controls)
 
+    def rotor_vel_control(self, controls: Array):
+        """Set the desired rotor velocities for all drones in all worlds."""
+        self.data = F.rotor_vel_control(self.data, controls)
+
     @requires_mujoco_sync
     def render(
         self,
@@ -303,11 +308,9 @@ def build_data(self) -> SimData:
             The simulation data as a single PyTree that can be passed to the pure simulation
             functions for stepping and resetting.
         """
-        state_freq = self.data.controls.state.freq if self.data.controls.state is not None else 0
-        attitude_freq = (
-            self.data.controls.attitude.freq if self.data.controls.attitude is not None else 0
-        )
-        force_torque_freq = self.data.controls.force_torque.freq
+        state_freq = 0 if (s := self.data.controls.state) is None else s.freq
+        attitude_freq = 0 if (a := self.data.controls.attitude) is None else a.freq
+        force_torque_freq = 0 if (ft := self.data.controls.force_torque) is None else ft.freq
         self.data = self.init_data(
             state_freq, attitude_freq, force_torque_freq, self.data.core.rng_key
         )
@@ -436,6 +439,8 @@ def build_control_fns(
                 raise NotImplementedError(f"Control mode {control} not implemented for {physics}")
         case Control.force_torque:
             control_pipeline = (step_force_torque_controller,)
+        case Control.rotor_vel:
+            control_pipeline = ()
         case _:
             raise NotImplementedError(f"Control mode {control} not implemented")
 

tests/integration/test_interfaces.py

@@ -3,6 +3,8 @@
 import pytest
 from drone_controllers import parametrize
 from drone_controllers.mellinger import state2attitude
+from drone_models.core import load_params
+from drone_models.transform import motor_force2rotor_vel
 from scipy.spatial.transform import Rotation as R
 
 from crazyflow.control.control import Control
@@ -52,6 +54,24 @@ def test_attitude_interface(physics: Physics):
     assert distance < 0.05, f"Failed to maintain hover with {physics} ({dpos})"
 
 
+@pytest.mark.integration
+def test_rotor_vel_interface():
+    sim = Sim(physics=Physics.first_principles, control=Control.rotor_vel)
+    params = load_params("first_principles", "cf2x_L250")
+    max_rpm = motor_force2rotor_vel(np.array([params["thrust_max"]]), params["rpm2thrust"])[0]
+
+    sim.data = sim.data.replace(
+        states=sim.data.states.replace(pos=sim.data.states.pos.at[..., 2].set(0.5))
+    )
+    cmd = np.full((1, 1, 4), max_rpm, dtype=np.float32)  # More RPMs than required for hover
+    sim.rotor_vel_control(cmd)
+    sim.step(sim.freq * 2)  # Run simulation for 2 seconds
+
+    # Check if drone is not tilted
+    assert R.from_quat(sim.data.states.quat[0, 0]).magnitude() < 0.1, "Drone is tilted"
+    assert sim.data.states.pos[0, 0, 2] > 0.5, "Failed to accelerate with rotor velocity control"
+
+
 @pytest.mark.integration
 @pytest.mark.parametrize("physics", Physics)
 def test_swarm_control(physics: Physics):

tests/unit/test_functional.py

@@ -145,3 +145,16 @@ def test_functional_state_control_device(device: str):
     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"
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("control", Control)
+def test_functional_controllable(control: Control):
+    """Test that functional controllable function works correctly."""
+    sim = Sim(n_worlds=2, n_drones=3, control=control)
+    data = sim.build_data()
+    controllable = F.controllable(data)
+    assert isinstance(controllable, jnp.ndarray), "Controllable must be a JAX array"
+    shape = controllable.shape
+    des_shape = (sim.n_worlds, 1)
+    assert shape == des_shape, f"Controllable shape must be {des_shape}, got {shape}"

tests/unit/test_sim.py

@@ -56,10 +56,11 @@ def array_compare_assert(x: Array, y: Array, value: bool = True, name: str | Non
 def test_sim_init(physics: Physics, device: str, control: Control, n_worlds: int):
     n_drones = 1
 
-    if physics != Physics.first_principles and control == Control.force_torque:
-        with pytest.raises(ConfigError):
-            Sim(n_worlds=n_worlds, physics=physics, device=device, control=control)
-        return
+    if physics != Physics.first_principles:
+        if control in (Control.force_torque, Control.rotor_vel):
+            with pytest.raises(ConfigError):
+                Sim(n_worlds=n_worlds, physics=physics, device=device, control=control)
+            return
 
     sim = Sim(n_worlds=n_worlds, physics=physics, device=device, control=control)
     assert sim.n_worlds == n_worlds
@@ -89,10 +90,11 @@ def test_sim_init(physics: Physics, device: str, control: Control, n_worlds: int
         assert sim.data.controls.attitude is None
 
     # Test force torque buffer shapes
-    ft_ctrl = sim.data.controls.force_torque
-    assert isinstance(ft_ctrl, ControlData)
-    array_meta_assert(ft_ctrl.cmd, (n_worlds, n_drones, 4), device)
-    array_meta_assert(ft_ctrl.staged_cmd, (n_worlds, n_drones, 4), device)
+    if control in (Control.state, Control.attitude, Control.force_torque):
+        ft_ctrl = sim.data.controls.force_torque
+        assert isinstance(ft_ctrl, ControlData)
+        array_meta_assert(ft_ctrl.cmd, (n_worlds, n_drones, 4), device)
+        array_meta_assert(ft_ctrl.staged_cmd, (n_worlds, n_drones, 4), device)
 
 
 @pytest.mark.unit
@@ -183,8 +185,9 @@ def test_reset_masked(device: str, physics: Physics):
 @pytest.mark.parametrize("physics", Physics)
 @pytest.mark.parametrize("control", Control)
 def test_sim_step(n_worlds: int, n_drones: int, physics: Physics, control: Control, device: str):
-    if physics != Physics.first_principles and control == Control.force_torque:
-        pytest.skip("Force-torque control is not supported with non-first-principles physics")
+    if physics != Physics.first_principles:
+        if control in (Control.force_torque, Control.rotor_vel):
+            pytest.skip(f"{control} is not supported with non-first-principles physics")
 
     sim = Sim(n_worlds=n_worlds, n_drones=n_drones, physics=physics, device=device, control=control)
     sim.step(2)