Cover rotor_vel in controllable

Author: amacati

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
@@ -78,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)

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}"