[WIP] Add Mellinger tests

Author: ratheron

drone_models/controller/__init__.py

@@ -5,3 +5,15 @@
 controllers as pure functions to ensure that users can jit-compile them. All controllers use
 broadcasting to support batching of arbitrary leading dimensions.
 """
+
+from typing import Callable
+
+from drone_models.controller.mellinger import (
+    attitude2force_torque as mellinger_attitude2force_torque,
+)
+from drone_models.controller.mellinger import pos2attitude as mellinger_pos2attitude
+
+available_controller: dict[str, Callable] = {
+    "mellinger_pos2attitude": mellinger_pos2attitude,
+    "mellinger_attitude2force_torque": mellinger_attitude2force_torque,
+}

drone_models/controller/constants.py

@@ -1 +1,60 @@
 """Constants for the controllers."""
+
+from types import ModuleType
+
+import numpy as np
+from array_api_typing import Array
+
+# Constants for the controllers
+# Same as in the firmware. Do not touch
+# Not part of the Constants class, since the values are controller specific and not drone specific
+
+#### Mellinger controller (see controller_mellinger.c)
+# Note: The firmware assumes mass=0.027. With battery thats closer to 0.034 though!!!
+mass = 0.034  # TODO This is the wrong mass (cf with battery weighs more!)
+massThrust = 132000 * 0.034 / 0.027
+
+# XY Position PID
+kp_xy = 0.4  # P
+kd_xy = 0.2  # D
+ki_xy = 0.05  # I
+i_range_xy = 2.0
+
+# Z Position
+kp_z = 1.25  # P
+kd_z = 0.4  # D
+ki_z = 0.05  # I
+i_range_z = 0.4
+
+# Attitude
+kR_xy = 70000.0  # P
+kw_xy = 20000.0  # D
+ki_m_xy = 0.0  # I
+i_range_m_xy = 1.0
+
+# Yaw
+kR_z = 60000.0  # P
+kw_z = 12000.0  # D
+ki_m_z = 500.0  # I
+i_range_m_z = 1500.0
+
+# roll and pitch angular velocity
+kd_omega_rp = 200.0  # D
+
+
+def cntrl_const_mel(xp: ModuleType = np) -> dict[str, Array | float]:
+    """Returns the controller constants for the Mellinger controller."""
+    return {
+        "mass": mass,
+        "massThrust": massThrust,
+        "kp": xp.asarray([kp_xy, kp_xy, kp_z]),
+        "kd": xp.asarray([kd_xy, kd_xy, kd_z]),
+        "ki": xp.asarray([ki_xy, ki_xy, ki_z]),
+        "i_range": xp.asarray([i_range_xy, i_range_xy, i_range_z]),
+        "kR": xp.asarray([kR_xy, kR_xy, kR_z]),
+        "kw": xp.asarray([kw_xy, kw_xy, kw_z]),
+        "ki_m": xp.asarray([ki_m_xy, ki_m_xy, ki_m_z]),
+        "kd_omega": xp.asarray([kd_omega_rp, kd_omega_rp, 0.0]),
+        "i_range_m": xp.asarray([i_range_m_xy, i_range_m_xy, i_range_m_z]),
+        "torque_pwm_range": 32_000.0,
+    }

drone_models/controller/mellinger.py

@@ -5,11 +5,10 @@
 from typing import TYPE_CHECKING
 
 from array_api_compat import array_namespace
+from scipy.spatial.transform import Rotation as R
 
-import drone_models.utils.cf2 as cf2
-import drone_models.utils.rotation as R
+import drone_models.utils.rotation as rotation
 from drone_models.transform import force2pwm, motor_force2rotor_speed, pwm2force
-from drone_models.utils.constants_controllers import cntrl_const_mel
 
 if TYPE_CHECKING:
     from array_api_typing import Array
@@ -24,6 +23,7 @@ def pos2attitude(
     ang_vel: Array,
     cmd: Array,
     constants: Constants,
+    parameters: dict[str, Array | float],
     dt: float = 1 / 500,
     i_error: Array | None = None,
 ) -> tuple[Array, Array]:
@@ -44,6 +44,7 @@ def pos2attitude(
         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].
         constants: Drone specific constants.
+        parameters: Controller specific parameters. TODO make class?
         dt: Time since last call.
         i_error: Integral error of the position controller with shape (..., 3).
 
@@ -63,17 +64,15 @@ def pos2attitude(
     # l.151 ff Integral Error
     if i_error is None:
         i_error = xp.zeros_like(pos)
-    i_error = xp.clip(
-        i_error + r_error * dt, -cntrl_const_mel["i_range"], cntrl_const_mel["i_range"]
-    )
+    i_error = xp.clip(i_error + r_error * dt, -parameters["i_range"], parameters["i_range"])
     # l. 161 Desired thrust [F_des]
     # => only one case here, since setpoint is always in absolute mode
     # Note: since we've defined the gravity in z direction, a "-" needs to be added
     target_thrust = (
-        cntrl_const_mel["mass"] * (setpoint_acc - constants.GRAVITY_VEC)
-        + cntrl_const_mel["kp"] * r_error
-        + cntrl_const_mel["kd"] * v_error
-        + cntrl_const_mel["ki"] * i_error
+        parameters["mass"] * (setpoint_acc - constants.GRAVITY_VEC)
+        + parameters["kp"] * r_error
+        + parameters["kd"] * v_error
+        + parameters["ki"] * i_error
     )
     # l. 178 Rate-controlled YAW is moving YAW angle setpoint
     # => only one case here, since the setpoint is always in absolute mode
@@ -88,7 +87,7 @@ def pos2attitude(
     # Taking the dot product of the last axis:
     current_thrust = xp.vecdot(target_thrust, z_axis, axis=-1)
     # l. 207 Calculate axis [zB_des]
-    z_axis_desired = target_thrust / xp.linalg.norm(target_thrust)
+    z_axis_desired = target_thrust / xp.linalg.vector_norm(target_thrust)
     # l. 210 [xC_des]
     # x_axis_desired = z_axis_desired x [sin(yaw), cos(yaw), 0]^T
     x_c_des_x = xp.cos(desiredYaw)
@@ -97,16 +96,16 @@ def pos2attitude(
     x_c_des = xp.stack((x_c_des_x, x_c_des_y, x_c_des_z), axis=-1)
     # [yB_des]
     y_axis_desired = xp.linalg.cross(z_axis_desired, x_c_des)
-    y_axis_desired = y_axis_desired / xp.linalg.norm(y_axis_desired)
+    y_axis_desired = y_axis_desired / xp.linalg.vector_norm(y_axis_desired)
     # [xB_des]
     x_axis_desired = xp.linalg.cross(y_axis_desired, z_axis_desired)
     # converting desired axis to rotation matrix and then to RPY
     matrix = xp.stack((x_axis_desired, y_axis_desired, z_axis_desired), axis=-1)
     command_RPY = R.from_matrix(matrix).as_euler("xyz", degrees=False)
     # l. 283
-    thrust = cntrl_const_mel["massThrust"] * current_thrust
+    thrust = parameters["massThrust"] * current_thrust
     # Transform thrust into N to keep uniform interface
-    thrust = cf2.pwm2force(thrust, constants, perMotor=False)
+    thrust = pwm2force(thrust, constants.THRUST_MAX, constants.PWM_MAX) * 4
     command_rpyt = xp.concat((command_RPY, thrust[..., None]), axis=-1)
     return command_rpyt, i_error
 
@@ -116,8 +115,9 @@ def attitude2force_torque(
     quat: Array,
     vel: Array,
     ang_vel: Array,
-    command_rpyt: Array,
+    cmd: Array,
     constants: Constants,
+    parameters: dict[str, Array | float],
     dt: float = 1 / 500,
     i_error_m: Array | None = None,
     ang_vel_des: Array | None = None,
@@ -131,8 +131,9 @@ def attitude2force_torque(
         quat: Drone orientation as xyzw quaternion with shape (..., 4).
         vel: Drone velocity with shape (..., 3).
         ang_vel: Drone angular drone velocity in rad/s with shape (..., 3).
-        command_rpyt: Commanded attitude (roll, pitch, yaw) and total thrust [rad, rad, rad, N]
+        cmd: Commanded attitude (roll, pitch, yaw) and total thrust [rad, rad, rad, N]
         constants (Constants): Constants of the specific drone
+        parameters: Controller specific parameters. TODO make class?
         dt: Time since last call.
         i_error_m: Integral error.
         ang_vel_des: Desired angular velocity in rad/s.
@@ -142,10 +143,10 @@ def attitude2force_torque(
     Returns:
         4 Motor forces [N], i_error_m
     """
-    xp = array_namespace(pos)
-    force_des = command_rpyt[..., -1]  # Total thrust in N
-    rpy_des = command_rpyt[..., :-1]
-    axis_flip = xp.array([1, -1, 1])  # to change the direction of the y axis
+    xp = array_namespace(quat)
+    force_des = cmd[..., -1]  # Total thrust in N
+    rpy_des = cmd[..., -4:-1]
+    axis_flip = xp.asarray([1.0, -1.0, 1.0])  # to change the direction of the y axis
     # l. 220 ff [eR]. We're using the "inefficient" code path from the firmware
     rot = R.from_quat(quat)
     rot_des = R.from_euler("xyz", rpy_des, degrees=False)
@@ -174,16 +175,18 @@ def attitude2force_torque(
     if i_error_m is None:
         i_error_m = xp.zeros_like(pos)
     i_error_m = i_error_m - eR * dt
-    i_error_m = xp.clip(i_error_m, -cntrl_const_mel["i_range_m"], cntrl_const_mel["i_range_m"])
+    i_error_m = xp.clip(i_error_m, -parameters["i_range_m"], parameters["i_range_m"])
     # l. 278 ff Moment:
     torque_pwm = (
-        -cntrl_const_mel["kR"] * eR
-        + cntrl_const_mel["kw"] * ew
-        + cntrl_const_mel["ki_m"] * i_error_m
-        + cntrl_const_mel["kd_omega"] * err_d
+        -parameters["kR"] * eR
+        + parameters["kw"] * ew
+        + parameters["ki_m"] * i_error_m
+        + parameters["kd_omega"] * err_d
     )
     # l. 297 ff
-    torque_pwm = xp.clip(torque_pwm, -32_000, 32_000)
+    torque_pwm = xp.clip(
+        torque_pwm, -parameters["torque_pwm_range"], parameters["torque_pwm_range"]
+    )
     torque_pwm = xp.where((force_des > 0)[..., None], torque_pwm, 0.0)
     # Info: The following part is NOT part of the Mellinger controller, but of the firmware.
     # The mixing of force and torque is done on the PWM level. We therefore mix those here according

tests/unit/test_controller.py

@@ -0,0 +1,54 @@
+"""Tests of the numeric models."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Callable
+
+import array_api_compat.numpy as np
+import array_api_strict as xp
+import jax
+import pytest
+
+from drone_models.controller import available_controller
+from drone_models.controller.constants import cntrl_const_mel
+from drone_models.utils.constants import Constants, available_drone_types
+
+if TYPE_CHECKING:
+    from array_api_typing import Array
+
+# For all tests to pass, we need the same precsion in jax as in np
+jax.config.update("jax_enable_x64", True)
+
+
+def create_rnd_states(
+    shape: tuple[int, ...] = (),
+) -> tuple[Array, Array, Array, Array, Array, Array, Array]:
+    """Creates N random states."""
+    pos = xp.asarray(np.random.uniform(-5, 5, shape + (3,)))
+    quat = xp.asarray(np.random.uniform(-1, 1, shape + (4,)))  # Libraries normalize automatically
+    vel = xp.asarray(np.random.uniform(-5, 5, shape + (3,)))
+    ang_vel = xp.asarray(np.random.uniform(-2, 2, shape + (3,)))
+    rotor_vel = xp.asarray(np.random.uniform(0, 0.2, shape + (4,)))
+    forces_dist = xp.asarray(np.random.uniform(-0.2, 0.2, shape + (3,)))
+    torques_dist = xp.asarray(np.random.uniform(-0.05, 0.05, shape + (3,)))
+    return pos, quat, vel, ang_vel, rotor_vel, forces_dist, torques_dist
+
+
+def create_rnd_commands(shape: tuple[int, ...] = (), dim: int = 4) -> Array:
+    """Creates N random inputs with size dim."""
+    return xp.asarray(np.random.uniform(0, 0.2, shape + (dim,)))
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("controller_name, controller", available_controller.items())
+@pytest.mark.parametrize("drone_type", available_drone_types)
+def test_controller(controller_name: str, controller: Callable, drone_type: str):
+    """TODO."""
+    constants = Constants.from_config(drone_type, xp)
+    batch_shape = (10,)
+    pos, quat, vel, ang_vel, _, _, _ = create_rnd_states(batch_shape)
+    cmd = create_rnd_commands(batch_shape, dim=13)  # TODO make dependent on controller
+
+    parameters = cntrl_const_mel(xp)
+
+    controller(pos, quat, vel, ang_vel, cmd, constants, parameters)