Improve API and add docs

Author: amacati

crazyflow/control/core.py

@@ -61,8 +61,9 @@ def load_params(controller: str, drone: str) -> dict[str, dict]:
     """Load the raw parameter table for a controller and drone.
 
     Reads ``crazyflow/control/<controller>/params.toml`` and returns the drone's table, nested by
-    section (``"core"`` and one per controller function). Use [load_fn_params][load_fn_params] to
-    select and convert the parameters a specific controller function accepts.
+    section (``"core"`` and one per controller function). Use
+    [load_fn_params][crazyflow.control.core.load_fn_params] to select and convert the parameters a
+    specific controller function accepts.
 
     Args:
         controller: Name of the controller sub-package, e.g. ``"mellinger"``.

crazyflow/control/mellinger/control.py

@@ -36,7 +36,7 @@ def state2attitude(
     quat: Array,
     vel: Array,
     cmd: Array,
-    ctrl_errors: tuple[Array, ...] | None = None,
+    pos_err_i: Array | None = None,
     ctrl_freq: float = 100,
     *,
     mass: float,
@@ -60,8 +60,8 @@ def state2attitude(
         vel: Drone velocity with shape (..., 3).
         cmd: Full state command in SI units and rad with shape (..., 13). The entries are
             [x, y, z, vx, vy, vz, ax, ay, az, yaw, roll_rate, pitch_rate, yaw_rate].
-        ctrl_errors: Tuple of integral errors. For state2attitude, the tuple contains a single array
-            (..., 3) for the position integral error or is None.
+        pos_err_i: Position integral error (..., 3) from the previous call. If None, it is
+            initialised to zero.
         ctrl_freq: Control frequency in Hz
         mass: Drone mass used for calculations in the controller in kg.
         kp: Proportional gain for the position controller with shape (3,).
@@ -90,7 +90,7 @@ def state2attitude(
     pos_err = setpoint_pos - pos  # l. 145 Position Error (ep)
     vel_err = setpoint_vel - vel  # l. 148 Velocity Error (ev)
     # l.151 ff Integral Error
-    int_pos_err = xp.zeros_like(pos) if ctrl_errors is None else ctrl_errors[0]
+    int_pos_err = xp.zeros_like(pos) if pos_err_i is None else pos_err_i
     int_pos_err = xp.clip(int_pos_err + pos_err * dt, -int_err_max, int_err_max)
     # l. 161 Desired thrust [F_des]
     # => only one case here, since setpoint is always in absolute mode
@@ -147,7 +147,7 @@ def attitude2force_torque(
     ang_vel: Array,
     cmd: Array,
     prev_ang_vel: Array | None = None,
-    ctrl_errors: tuple[Array, ...] | None = None,
+    r_int_error: Array | None = None,
     ctrl_freq: int = 500,
     *,
     kR: Array,
@@ -173,8 +173,8 @@ def attitude2force_torque(
         quat: Drone orientation as xyzw quaternion with shape (..., 4).
         ang_vel: Drone angular drone velocity in rad/s with shape (..., 3).
         cmd: Commanded attitude (roll, pitch, yaw) and total thrust [rad, rad, rad, N].
-        ctrl_errors: Tuple of integral errors. For attitude2force_torque, the tuple contains a
-            single array (..., 3) for the angular velocity integral error or is None.
+        r_int_error: Angular velocity integral error (..., 3) from the previous call. If None, it
+            is initialised to zero.
         ctrl_freq: Control frequency in Hz
         kR: Proportional gain for the rotation error with shape (3,).
         kw: Proportional gain for the angular velocity error with shape (3,).
@@ -221,7 +221,7 @@ def attitude2force_torque(
     ang_vel_d_err = xpx.at(ang_vel_d_err)[..., 2].set(0)
 
     # l. 268 ff Integral Error
-    r_int_error = xp.zeros_like(ang_vel) if ctrl_errors is None else ctrl_errors[0]
+    r_int_error = xp.zeros_like(ang_vel) if r_int_error is None else r_int_error
     r_int_error = r_int_error - eR * dt
     r_int_error = xp.clip(r_int_error, -int_err_max, int_err_max)
     # l. 278 ff Moment:
@@ -244,8 +244,8 @@ def attitude2force_torque(
     # TODO: Long-term, the Mellinger controller should use the new power distribution which
     # calculates motor forces in Newtons. However, for now the firmware uses the legacy power
     # distribution, so we keep it here for compatibility. To have a single consistent interface for
-    # controllers within drone_models, we still want to return SI forces and torques. We thus need
-    # to convert the legacy output to SI units.
+    # controllers, we still want to return SI forces and torques. We thus need to convert the legacy
+    # output to SI units.
     # l. 310 ff
     torque_des = (mixing_matrix @ motor_forces[..., None])[..., 0] * xp.stack([L, L, thrust2torque])
     force_des = xp.sum(motor_forces, axis=-1)[..., None]
@@ -420,7 +420,7 @@ def sim_state2attitude(data: SimData) -> SimData:
         states.quat,
         states.vel,
         state_ctrl.cmd,
-        ctrl_errors=(state_ctrl.pos_err_i,),
+        pos_err_i=state_ctrl.pos_err_i,
         ctrl_freq=state_ctrl.freq,
         **state_ctrl.params,
     )
@@ -440,7 +440,7 @@ def sim_attitude2force_torque(data: SimData) -> SimData:
         states.quat,
         states.ang_vel,
         attitude_ctrl.cmd,
-        ctrl_errors=(attitude_ctrl.r_int_error,),
+        r_int_error=attitude_ctrl.r_int_error,
         ctrl_freq=attitude_ctrl.freq,
         prev_ang_vel=attitude_ctrl.last_ang_vel,
         **attitude_ctrl.params,

docs/gen_ref_pages.py

@@ -59,7 +59,11 @@
     * [dynamics.symbols](crazyflow/dynamics/symbols.md)
 * Control
     * [control](crazyflow/control/index.md)
-    * [control.mellinger](crazyflow/control/mellinger.md)
+    * [control.control](crazyflow/control/control.md)
+    * [control.core](crazyflow/control/core.md)
+    * [control.transform](crazyflow/control/transform.md)
+    * [control.mellinger](crazyflow/control/mellinger/index.md)
+    * [control.mellinger.control](crazyflow/control/mellinger/control.md)
 * Environments
     * [envs](crazyflow/envs/index.md)
     * [envs.drone_env](crazyflow/envs/drone_env.md)

docs/user-guide/control/batching.md

@@ -0,0 +1,41 @@
+# Batching
+
+All controllers are built on Array API operations that broadcast over leading dimensions. Add a leading batch dimension to your state and command arrays and the controller evaluates all instances in a single call, with no loops and no special API.
+
+```python
+import numpy as np
+from crazyflow.control import parametrize
+from crazyflow.control.mellinger import state2attitude
+
+ctrl = parametrize(state2attitude, "cf2x_L250")
+
+N = 100
+pos = np.zeros((N, 3))
+quat = np.tile(np.array([0.0, 0.0, 0.0, 1.0]), (N, 1))
+vel = np.zeros((N, 3))
+cmd = np.zeros((N, 13))
+
+rpyt, int_pos_err = ctrl(pos, quat, vel, cmd)
+rpyt.shape  # (100, 4)
+```
+
+## Higher-dimensional batches
+
+Any number of leading dimensions works. A common pattern is a grid of environments, each containing multiple drones:
+
+```python
+import numpy as np
+from crazyflow.control import parametrize
+from crazyflow.control.mellinger import state2attitude
+
+ctrl = parametrize(state2attitude, "cf2x_L250")
+
+# 10 environments, 5 drones each
+pos = np.zeros((10, 5, 3))
+quat = np.broadcast_to(np.array([0.0, 0.0, 0.0, 1.0]), (10, 5, 4)).copy()
+vel = np.zeros((10, 5, 3))
+cmd = np.zeros((10, 5, 13))
+
+rpyt, _ = ctrl(pos, quat, vel, cmd)
+rpyt.shape  # (10, 5, 4)
+```

docs/user-guide/control/controllers.md

@@ -0,0 +1,27 @@
+# Controllers
+
+A controller is a function that maps the current drone state and a command to actuator outputs. Every controller in `crazyflow.control` is:
+
+- **A pure function**: no hidden state; integral errors are explicit return values you pass back on the next call
+- **Array-API compatible**: works identically with NumPy, JAX, PyTorch, or any compliant library
+- **Batchable**: add leading dimensions to any input array and the function evaluates all instances at once
+
+## The Mellinger pipeline
+
+The Mellinger controller [[1]](#references) is split into three stages that form a pipeline. Each stage can be used on its own, or all three can be chained to convert a full-state setpoint into individual motor speeds.
+
+| Stage | Function | Takes | Produces |
+|---|---|---|---|
+| 1 | [`state2attitude`](mellinger.md#state-to-attitude) | State + 13-element setpoint | RPYT command + position integral error |
+| 2 | [`attitude2force_torque`](mellinger.md#attitude-to-force-torque) | Attitude + RPYT command | Collective force, body torques + angular velocity integral error |
+| 3 | [`force_torque2rotor_vel`](mellinger.md#force-torque-to-rotor-velocities) | Force + torques | 4 motor speeds [RPM] |
+
+## Available controllers
+
+| Module | Controller | Stages |
+|---|---|---|
+| `crazyflow.control.mellinger` | Mellinger | `state2attitude`, `attitude2force_torque`, `force_torque2rotor_vel` |
+
+## References
+
+[1] D. Mellinger and V. Kumar, "Minimum snap trajectory generation and control for quadrotors," ICRA 2011, doi: 10.1109/ICRA.2011.5980409.

docs/user-guide/control-modes.md docs/user-guide/control/index.mddocs/user-guide/control-modes.md renamed to docs/user-guide/control/index.md

@@ -162,7 +162,18 @@ Each control mode has its own update rate. The dynamics tick (`freq`) is always
 
 The simulator applies a new command only when the control tick fires. Between ticks, the previous command is held. The number of dynamics steps per control tick is `freq // control_freq`.
 
+## Using the controllers standalone
+
+The control modes above are how the simulator drives the onboard controllers. Those controllers also live in `crazyflow.control` as a self-contained library of pure functions, usable on their own for control design, learning-based policies, or as a reference implementation, independent of `Sim`. The following guides cover that standalone API:
+
+- [Controllers](controllers.md): the controller interface and the Mellinger pipeline
+- [Mellinger controller](mellinger.md): the three stages, their inputs and outputs
+- [Parametrization](parametrize.md): binding a controller to a drone configuration
+- [Integral errors](integral-errors.md): carrying controller state across calls
+- [Batching](batching.md): evaluating many drones at once
+- [JIT compilation](jit.md): compiling controllers with `jax.jit`
+
 ## Next steps
 
-- [Functional API](functional-api.md) — running control inside JIT with `F.controllable`
-- [Dynamics](dynamics/index.md) — compatibility between dynamics and control modes
+- [Functional API](../functional-api.md): running control inside JIT with `F.controllable`
+- [Dynamics](../dynamics/index.md): compatibility between dynamics and control modes

docs/user-guide/control/integral-errors.md

@@ -0,0 +1,79 @@
+# Integral errors
+
+The Mellinger controller uses integral terms to reject steady-state errors. Because every controller in `crazyflow.control` is a **pure function**, integral state cannot be stored inside the function. It is an explicit return value that you pass back on the next call.
+
+## Initialisation
+
+You have two ways to start the integral error at zero:
+
+1. Pass `None` (or omit the argument). The controller then creates the zero array internally.
+2. Create the zero array yourself and pass it explicitly. The integral error has the same shape as the position (for `pos_err_i`) or angular velocity (for `r_int_error`), so a zero array of that shape works.
+
+```python
+import numpy as np
+from crazyflow.control import parametrize
+from crazyflow.control.mellinger import state2attitude
+
+ctrl = parametrize(state2attitude, "cf2x_L250")
+pos = np.zeros(3)
+quat = np.array([0.0, 0.0, 0.0, 1.0])
+vel = np.zeros(3)
+cmd = np.zeros(13)
+
+# Option 1: let the controller initialise the integral error.
+rpyt, pos_err_i = ctrl(pos, quat, vel, cmd, pos_err_i=None)
+
+# Option 2: initialise it yourself.
+rpyt, pos_err_i = ctrl(pos, quat, vel, cmd, pos_err_i=np.zeros(3))
+```
+
+Under `jax.jit`, prefer option 2. Passing `None` on the first call and an array on the following calls changes the argument structure, so JAX traces and compiles the function twice. Passing a zero array from the start keeps the input structure constant and compiles only once.
+
+## Carrying errors across timesteps
+
+Pass the returned error straight back as `pos_err_i` on the next call:
+
+```python
+import numpy as np
+from crazyflow.control import parametrize
+from crazyflow.control.mellinger import state2attitude
+
+ctrl = parametrize(state2attitude, "cf2x_L250")
+pos = np.zeros(3)
+quat = np.array([0.0, 0.0, 0.0, 1.0])
+vel = np.zeros(3)
+cmd = np.zeros(13)
+cmd[0] = 1.0  # 1 m setpoint error in x
+
+pos_err_i = None
+for _ in range(10):
+    rpyt, pos_err_i = ctrl(pos, quat, vel, cmd, pos_err_i=pos_err_i)  # carry forward
+
+# After 10 steps at 100 Hz, integral error is approximately 10 * 0.01 = 0.1 m.
+```
+
+## Both stages have integral errors
+
+`state2attitude` tracks position error via `pos_err_i`. `attitude2force_torque` tracks angular velocity error via `r_int_error`. Manage them independently:
+
+```python
+import numpy as np
+from crazyflow.control import parametrize
+from crazyflow.control.mellinger import attitude2force_torque, state2attitude
+
+state_ctrl = parametrize(state2attitude, "cf2x_L250")
+att_ctrl = parametrize(attitude2force_torque, "cf2x_L250")
+
+pos = np.zeros(3)
+quat = np.array([0.0, 0.0, 0.0, 1.0])
+vel = np.zeros(3)
+ang_vel = np.zeros(3)
+cmd = np.zeros(13)
+
+pos_err_i = None
+r_int_error = None
+
+for _ in range(5):
+    rpyt, pos_err_i = state_ctrl(pos, quat, vel, cmd, pos_err_i=pos_err_i)
+    force, torque, r_int_error = att_ctrl(quat, ang_vel, rpyt, r_int_error=r_int_error)
+```

docs/user-guide/control/jit.md

@@ -0,0 +1,66 @@
+# JIT compilation
+
+Every controller is a pure function with no hidden state or side effects. In addition, they are implemented exclusively with Array API operations that are compatible with lazy JIT frameworks. Together, these two properties mean that every controller can be JIT compiled without any modification.
+
+```python
+import jax
+import jax.numpy as jnp
+from crazyflow.control import parametrize
+from crazyflow.control.mellinger import state2attitude
+
+ctrl = parametrize(state2attitude, "cf2x_L250", xp=jnp)
+jit_ctrl = jax.jit(ctrl)
+
+pos = jnp.zeros(3)
+quat = jnp.array([0.0, 0.0, 0.0, 1.0])
+vel = jnp.zeros(3)
+cmd = jnp.zeros(13)
+
+rpyt, int_pos_err = jit_ctrl(pos, quat, vel, cmd)
+```
+
+## Integral errors under JIT
+
+Integral errors are regular arrays and are handled as JAX pytree leaves, so they pass through `jax.jit` without any special treatment.
+
+```python
+import jax
+import jax.numpy as jnp
+from crazyflow.control import parametrize
+from crazyflow.control.mellinger import state2attitude
+
+ctrl = parametrize(state2attitude, "cf2x_L250", xp=jnp)
+jit_ctrl = jax.jit(ctrl)
+
+pos = jnp.zeros(3)
+quat = jnp.array([0.0, 0.0, 0.0, 1.0])
+vel = jnp.zeros(3)
+cmd = jnp.zeros(13)
+
+pos_err_i = jnp.zeros(3)  # initialise to zero, so the function compiles only once
+for _ in range(10):
+    rpyt, pos_err_i = jit_ctrl(pos, quat, vel, cmd, pos_err_i=pos_err_i)
+```
+
+## Batched JIT
+
+Batching and JIT compose directly. Add leading dimensions to the state arrays and the same compiled function handles the entire batch.
+
+```python
+import jax
+import jax.numpy as jnp
+from crazyflow.control import parametrize
+from crazyflow.control.mellinger import state2attitude
+
+ctrl = parametrize(state2attitude, "cf2x_L250", xp=jnp)
+jit_ctrl = jax.jit(ctrl)
+
+N = 1_000
+pos = jnp.zeros((N, 3))
+quat = jnp.broadcast_to(jnp.array([0.0, 0.0, 0.0, 1.0]), (N, 4))
+vel = jnp.zeros((N, 3))
+cmd = jnp.zeros((N, 13))
+
+rpyt, _ = jit_ctrl(pos, quat, vel, cmd)
+rpyt.shape  # (1000, 4)
+```