Add proper parametrization of controllers

Author: amacati

drone_models/controller/__init__.py

@@ -8,6 +8,7 @@
 
 from typing import Callable
 
+from drone_models.controller.core import parametrize
 from drone_models.controller.mellinger import (
     attitude2force_torque as mellinger_attitude2force_torque,
 )
@@ -17,3 +18,5 @@
     "mellinger_state2attitude": mellinger_state2attitude,
     "mellinger_attitude2force_torque": mellinger_attitude2force_torque,
 }
+
+__all__ = ["parametrize"]

drone_models/controller/core.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Callable, ParamSpec, Protocol, TypeVar, runtime_checkable
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+controller_parameter_registry: dict[str, type[ControllerParams]] = {}
+
+
+def parametrize(fn: Callable[P, R], drone_model: str) -> Callable[P, R]:
+    """Parametrize a controller function with the default controller parameters for a drone model.
+
+    Args:
+        fn: The controller function to parametrize.
+        drone_model: The drone model to use.
+
+    Example:
+        >>> from drone_models.controller import parametrize
+        >>> from drone_models.controller.mellinger import state2attitude
+        >>> controller_fn = parametrize(state2attitude, drone_model="cf2x_L250")
+        >>> command_rpyt, int_pos_err = controller_fn(
+        ...     pos=pos,
+        ...     quat=quat,
+        ...     vel=vel,
+        ...     ang_vel=ang_vel,
+        ...     cmd=cmd,
+        ...     ctrl_errors=(int_pos_err,),
+        ...     ctrl_freq=100,
+        ... )
+
+    Returns:
+        The parametrized controller function with all keyword argument only parameters filled in.
+    """
+    controller_id = fn.__module__.split(".")[-1] + fn.__name__
+    try:
+        params = controller_parameter_registry[controller_id].load(drone_model)
+    except KeyError:
+        raise KeyError(f"Controller {controller_id} does not exist in the parameter registry")
+    except ValueError:
+        raise ValueError(f"Drone model {drone_model} not supported for {fn.__name__}")
+    return partial(fn, **params._asdict())
+
+
+@runtime_checkable
+class ControllerParams(Protocol):
+    """Protocol for controller parameters."""
+
+    @staticmethod
+    def load(drone_model: str) -> ControllerParams:
+        """Load the parameters from the config file."""
+
+    def _asdict(self) -> dict[str, Any]:
+        """Convert the parameters to a dictionary."""
+
+
+def register_controller_parameters(
+    params: ControllerParams | type[ControllerParams],
+) -> Callable[[Callable[P, R]], Callable[P, R]]:
+    """Register the default controller parameters for this controller.
+
+    Warning:
+        The controller parameters **must** be a named tuple with a function `load` that takes in the
+        drone model name and returns an instance of itself, or a class that implements the
+        ControllerParams protocol.
+
+    Args:
+        params: The controller parameter type.
+
+    Returns:
+        A decorator function that registers the parameters and returns the function unchanged.
+    """
+    if not isinstance(params, ControllerParams):
+        raise ValueError(f"{params} does not implement the ControllerParams protocol")
+
+    def decorator(fn: Callable[P, R]) -> Callable[P, R]:
+        controller_id = fn.__module__.split(".")[-1] + fn.__name__
+        if controller_id in controller_parameter_registry:
+            raise ValueError(f"Controller {controller_id} already registered")
+        controller_parameter_registry[controller_id] = params
+        return fn
+
+    return decorator

drone_models/controller/data/mellinger.toml

@@ -0,0 +1,7 @@
+from drone_models.controller.mellinger.control import (
+    attitude2force_torque,
+    force_torque2rotor_vel,
+    state2attitude,
+)
+
+__all__ = ["state2attitude", "attitude2force_torque", "force_torque2rotor_vel"]

drone_models/controller/mellinger/__init__.py

@@ -2,19 +2,21 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, NamedTuple
+from typing import TYPE_CHECKING
 
 import array_api_extra as xpx
-import numpy as np
 from array_api_compat import array_namespace
 from scipy.spatial.transform import Rotation as R
 
+from drone_models.controller.core import register_controller_parameters
+from drone_models.controller.mellinger.params import AttitudeParams, ForceTorqueParams, StateParams
 from drone_models.transform import force2pwm, motor_force2rotor_vel, pwm2force
 
 if TYPE_CHECKING:
     from array_api_typing import Array
 
 
+@register_controller_parameters(StateParams)
 def state2attitude(
     pos: Array,
     quat: Array,
@@ -130,6 +132,7 @@ def state2attitude(
     return command_rpyt, int_pos_err
 
 
+@register_controller_parameters(AttitudeParams)
 def attitude2force_torque(
     pos: Array,
     quat: Array,
@@ -259,6 +262,7 @@ def force_torque_pwms2pwms(force_pwm: Array, torque_pwm: Array, mixing_matrix: A
     return force_pwm[..., None] + (mixing_matrix @ torque_pwm[..., None])[..., 0]
 
 
+@register_controller_parameters(ForceTorqueParams)
 def force_torque2rotor_vel(
     force: Array,
     torque: Array,
@@ -300,127 +304,3 @@ def force_torque2rotor_vel(
     motor_forces = xp.where(xp.all(force == 0), 0.0, xp.clip(motor_forces, thrust_min, thrust_max))
     # Assume perfect battery compensation and calculate the desired motor speeds directly
     return motor_force2rotor_vel(motor_forces, KF)
-
-
-class MellingerStateParams(NamedTuple):
-    """Parameters for the Mellinger state controller."""
-
-    mass: float
-    kp: Array
-    kd: Array
-    ki: Array
-    gravity_vec: Array
-    mass_thrust: float
-    int_err_max: Array
-    thrust_max: float
-    pwm_max: float
-
-    @staticmethod
-    def load(drone_model: str) -> MellingerStateParams:
-        """Load the parameters from the config file."""
-        params = MellingerStateParams._load_fake_model()
-        return MellingerStateParams(**params)
-
-    @staticmethod
-    def _load_fake_model() -> dict[str, Array | float]:
-        """Load the parameters from the config file."""
-        # TODO: Remove this once we can load proper params
-        return {
-            "mass": 0.032999999821186066,
-            "mass_thrust": 132000 * 0.034 / 0.027,
-            "kp": np.array([0.4, 0.4, 1.25]),
-            "kd": np.array([0.2, 0.2, 0.5]),
-            "ki": np.array([0.05, 0.05, 0.05]),
-            "int_err_max": np.array([2.0, 2.0, 0.4]),
-            # TODO: Double-check if we want this
-            "gravity_vec": np.array([0.0, 0.0, -9.8100004196167]),
-            "thrust_max": 0.1125,
-            "pwm_max": 65535.0,
-        }
-
-
-class MellingerAttitudeParams(NamedTuple):
-    """Parameters for the Mellinger attitude controller."""
-
-    kR: Array
-    kw: Array
-    ki_m: Array
-    kd_omega: Array
-    int_err_max: Array
-    torque_pwm_max: Array
-    thrust_max: float
-    pwm_min: float
-    pwm_max: float
-    L: float
-    KF: float
-    KM: float
-    mixing_matrix: Array
-
-    @staticmethod
-    def load(drone_model: str) -> MellingerAttitudeParams:
-        """Load the parameters from the config file."""
-        params = MellingerAttitudeParams._load_fake_model()
-        return MellingerAttitudeParams(**params)
-
-    @staticmethod
-    def _load_fake_model() -> dict[str, Array | float]:
-        """Load the parameters from the config file."""
-        # TODO: Remove this once we can load proper params
-        # fmt: off
-        mixing_matrix = np.array([[-1.0, -1.0, -1.0],
-                                  [-1.0,  1.0,  1.0],
-                                  [ 1.0,  1.0, -1.0],
-                                  [ 1.0, -1.0,  1.0]]
-                                )
-        # fmt: on
-        return {
-            "kR": np.array([70_000.0, 70_000.0, 60_000.0]),
-            "kw": np.array([20_000.0, 20_000.0, 12_000.0]),
-            "ki_m": np.array([0.0, 0.0, 500.0]),
-            "kd_omega": np.array([200.0, 200.0, 0.0]),
-            "int_err_max": np.array([1.0, 1.0, 1500.0]),
-            "torque_pwm_max": np.array([32_000.0, 32_000.0, 32_000.0]),
-            "thrust_max": 0.1125,
-            "pwm_min": 20_000.0,
-            "pwm_max": 65_535.0,
-            "L": 0.03253,
-            "KF": 8.7e-10,
-            "KM": 7.94e-12,
-            "mixing_matrix": mixing_matrix,
-        }
-
-
-class MellingerForceTorqueParams(NamedTuple):
-    """Parameters for the Mellinger force torque controller."""
-
-    thrust_min: float
-    thrust_max: float
-    L: float
-    KF: float
-    KM: float
-    mixing_matrix: Array
-
-    @staticmethod
-    def load(drone_model: str) -> MellingerForceTorqueParams:
-        """Load the parameters from the config file."""
-        params = MellingerForceTorqueParams._load_fake_model()
-        return MellingerForceTorqueParams(**params)
-
-    @staticmethod
-    def _load_fake_model() -> dict[str, Array | float]:
-        """Load the parameters from the config file."""
-        # fmt: off
-        mixing_matrix = np.array([[-1.0, -1.0, -1.0],
-                                  [-1.0,  1.0,  1.0],
-                                  [ 1.0,  1.0, -1.0],
-                                  [ 1.0, -1.0,  1.0]]
-                                 )
-        # fmt: on
-        return {
-            "thrust_min": 0.02,
-            "thrust_max": 0.1125,
-            "L": 0.03253,
-            "KF": 8.701227710666256e-10,
-            "KM": 7.94e-12,
-            "mixing_matrix": mixing_matrix,
-        }

drone_models/controller/mellinger.py …e_models/controller/mellinger/control.pydrone_models/controller/mellinger.py renamed to drone_models/controller/mellinger/control.py drone_models/controller/mellinger.py renamed to drone_models/controller/mellinger/control.py

@@ -0,0 +1,88 @@
+"""Parameters for the Mellinger controller."""
+
+from __future__ import annotations
+
+import tomllib
+from pathlib import Path
+from typing import TYPE_CHECKING, NamedTuple
+
+import numpy as np
+
+if TYPE_CHECKING:
+    from array_api_typing import Array
+
+
+class StateParams(NamedTuple):
+    """Parameters for the Mellinger state controller."""
+
+    kp: Array
+    kd: Array
+    ki: Array
+    int_err_max: Array
+    mass: float
+    gravity_vec: Array
+    mass_thrust: float
+    thrust_max: float
+    pwm_max: float
+
+    @staticmethod
+    def load(drone_model: str) -> StateParams:
+        """Load the parameters from the config file."""
+        with open(Path(__file__).parent / "params.toml", "rb") as f:
+            params = tomllib.load(f)
+        if drone_model not in params:
+            raise KeyError(f"Drone model `{drone_model}` not found in params.toml")
+        params = params[drone_model]["state2attitude"] | params[drone_model]["core"]
+        params = {k: np.asarray(v) for k, v in params.items() if k in StateParams._fields}
+        return StateParams(**params)
+
+
+class AttitudeParams(NamedTuple):
+    """Parameters for the Mellinger attitude controller."""
+
+    kR: Array
+    kw: Array
+    ki_m: Array
+    kd_omega: Array
+    int_err_max: Array
+    torque_pwm_max: Array
+    thrust_max: float
+    pwm_min: float
+    pwm_max: float
+    L: float
+    KF: float
+    KM: float
+    mixing_matrix: Array
+
+    @staticmethod
+    def load(drone_model: str) -> AttitudeParams:
+        """Load the parameters from the config file."""
+        with open(Path(__file__).parent / "params.toml", "rb") as f:
+            params = tomllib.load(f)
+        if drone_model not in params:
+            raise KeyError(f"Drone model `{drone_model}` not found in params.toml")
+        params = params[drone_model]["attitude2force_torque"] | params[drone_model]["core"]
+        params = {k: np.asarray(v) for k, v in params.items() if k in AttitudeParams._fields}
+        return AttitudeParams(**params)
+
+
+class ForceTorqueParams(NamedTuple):
+    """Parameters for the Mellinger force torque controller."""
+
+    thrust_min: float
+    thrust_max: float
+    L: float
+    KF: float
+    KM: float
+    mixing_matrix: Array
+
+    @staticmethod
+    def load(drone_model: str) -> ForceTorqueParams:
+        """Load the parameters from the config file."""
+        with open(Path(__file__).parent / "params.toml", "rb") as f:
+            params = tomllib.load(f)
+        if drone_model not in params:
+            raise KeyError(f"Drone model `{drone_model}` not found in params.toml")
+        params = params[drone_model]["core"]
+        params = {k: np.asarray(v) for k, v in params.items() if k in ForceTorqueParams._fields}
+        return ForceTorqueParams(**params)

drone_models/controller/mellinger/params.py

@@ -0,0 +1,31 @@
+[cf2x_L250.core]
+mass = 0.033
+mass_thrust = 132000
+thrust_min = 0.02
+thrust_max = 0.1125
+pwm_min = 20000.0
+pwm_max = 65535.0
+torque_pwm_max = [32000.0, 32000.0, 32000.0]
+L = 0.03253
+KF = 8.7e-10
+KM = 7.94e-12
+mixing_matrix = [
+    [-1.0, -1.0, -1.0],
+    [-1.0,  1.0,  1.0],
+    [ 1.0,  1.0, -1.0],
+    [ 1.0, -1.0,  1.0]
+]
+gravity_vec = [0.0, 0.0, -9.81]
+
+[cf2x_L250.state2attitude]
+kp = [0.4, 0.4, 1.25]
+kd = [0.2, 0.2, 0.5]
+ki = [0.05, 0.05, 0.05]
+int_err_max = [2.0, 2.0, 0.4]
+
+[cf2x_L250.attitude2force_torque]
+kR = [70000.0, 70000.0, 60000.0]
+kw = [20000.0, 20000.0, 12000.0]
+ki_m = [0.0, 0.0, 500.0]
+kd_omega = [200.0, 200.0, 0.0]
+int_err_max = [1.0, 1.0, 1500.0]