Arbitrary motor count

crazyflow/control/mellinger/control.py

@@ -48,6 +48,7 @@ def state2attitude(
     int_err_max: Array,
     thrust_max: float,
     pwm_max: float,
+    mixing_matrix: Array,
 ) -> tuple[Array, Array]:
     """Compute the positional part of the mellinger controller.
 
@@ -73,6 +74,7 @@ def state2attitude(
         int_err_max: Range of the integral error with shape (3,). i_range in the firmware.
         thrust_max: Maximum thrust in N.
         pwm_max: Maximum PWM value.
+        mixing_matrix: Mixing matrix for the motor forces with shape (n_motors, 3).
 
     Returns:
         The RPY collective thrust command [rad, rad, rad, N], and the integral error of the position
@@ -140,7 +142,7 @@ def state2attitude(
     # instead of dynamically scaling with the mass parameter of the controller! Hence, we include
     # this conversion here and thus effectively rescale the thrust slightly. The conversion below
     # maps thrust -> PWM -> rescaled thrust.
-    thrust = pwm2force(mass_thrust * current_thrust, thrust_max * 4, pwm_max)
+    thrust = pwm2force(mass_thrust * current_thrust, thrust_max * mixing_matrix.shape[-1], pwm_max)
     command_rpyt = xp.concat((command_RPY, thrust[..., None]), axis=-1)
     return command_rpyt, int_pos_err
 
@@ -407,7 +409,7 @@ def _attitude2force_torque(
     # l. 297 ff
     torque_pwm = xp.clip(torque_pwm, -torque_pwm_max, torque_pwm_max)
     torque_pwm = xp.where((force_des > 0)[..., None], torque_pwm, 0.0)
-    force_des_pwm = force2pwm(force_des / 4, thrust_max, pwm_max)
+    force_des_pwm = force2pwm(force_des / mixing_matrix.shape[-1], thrust_max, pwm_max)
     pwms = force_torque_pwms2pwms(force_des_pwm, torque_pwm, mixing_matrix)
     idle = xp.all(pwms == 0, axis=-1, keepdims=True)
     pwms = xp.where(idle, 0.0, xp.clip(pwms, pwm_min, pwm_max))
@@ -479,7 +481,7 @@ def force_torque2rotor_vel(
     assert torque.shape[-1] == 3, f"Torque must have shape (..., 3), but has {torque.shape}"
     assert force.shape[-1] == 1, f"Force must have shape (..., 1), but has {force.shape}"
     torque_forces = (torque * xp.asarray([1 / L, 1 / L, 1 / thrust2torque])) @ mixing_matrix
-    motor_forces = (torque_forces + force) / 4
+    motor_forces = (torque_forces + force) / mixing_matrix.shape[-1]
     # Clip motor forces on the thrust instead of PWM level.
     idle = xp.all(force == 0, axis=-1, keepdims=True)
     motor_forces = xp.where(idle, 0.0, xp.clip(motor_forces, thrust_min, thrust_max))

crazyflow/envs/drone_env.py

@@ -33,7 +33,9 @@ def action_space(control_type: Control, drone: str) -> spaces.Box:
     match control_type:
         case Control.attitude:
             params = load_params(drone)
-            thrust_min, thrust_max = params["thrust_min"] * 4, params["thrust_max"] * 4
+            n_motors = np.asarray(params["mixing_matrix"]).shape[-1]
+            thrust_min = params["thrust_min"] * n_motors
+            thrust_max = params["thrust_max"] * n_motors
             return spaces.Box(
                 np.array([-np.pi / 2, -np.pi / 2, -np.pi / 2, thrust_min], dtype=np.float32),
                 np.array([np.pi / 2, np.pi / 2, np.pi / 2, thrust_max], dtype=np.float32),

crazyflow/sim/data.py

@@ -37,23 +37,24 @@ class SimState:
     """Force applied to the drone's center of mass in the world frame."""
     torque: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 3) CoM torque
     """Torque applied to the drone's center of mass in the world frame."""
-    rotor_vel: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 4) in RPM
+    rotor_vel: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, n_motors) in RPM
     """Motor forces along body frame z axis."""
 
     @staticmethod
-    def create(n_worlds: int, n_drones: int, device: Device) -> SimState:
+    def create(n_worlds: int, n_drones: int, n_motors: int, device: Device) -> SimState:
         """Create a default set of states for the simulation."""
         # Each field needs a buffer of its own so that SimData can be donated to XLA
         zeros_3d = jnp.zeros((n_worlds, n_drones, 3), device=device)
         zeros_4d = jnp.zeros((n_worlds, n_drones, 4), device=device)
+        zeros_motors = jnp.zeros((n_worlds, n_drones, n_motors), device=device)
         return SimState(
             pos=zeros_3d.copy(),
             quat=zeros_4d.at[..., -1].set(1.0),
             vel=zeros_3d.copy(),
             ang_vel=zeros_3d.copy(),
             force=zeros_3d.copy(),
             torque=zeros_3d.copy(),
-            rotor_vel=zeros_4d.copy(),
+            rotor_vel=zeros_motors.copy(),
         )
 
 
@@ -67,20 +68,20 @@ class SimStateDeriv:
     """Derivative of the velocity of the drone's center of mass."""
     ang_acc: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 3)
     """Derivative of the angular velocity of the drone's center of mass."""
-    rotor_acc: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 4)
+    rotor_acc: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, n_motors)
     """Derivative of the rotor velocity."""
 
     @staticmethod
-    def create(n_worlds: int, n_drones: int, device: Device) -> SimStateDeriv:
+    def create(n_worlds: int, n_drones: int, n_motors: int, device: Device) -> SimStateDeriv:
         """Create a default set of state derivatives for the simulation."""
         zeros_3d = jnp.zeros((n_worlds, n_drones, 3), device=device)
-        zeros_4d = jnp.zeros((n_worlds, n_drones, 4), device=device)
+        zeros_motors = jnp.zeros((n_worlds, n_drones, n_motors), device=device)
         return SimStateDeriv(
             vel=zeros_3d.copy(),
             ang_vel=zeros_3d.copy(),
             acc=zeros_3d.copy(),
             ang_acc=zeros_3d.copy(),
-            rotor_acc=zeros_4d.copy(),
+            rotor_acc=zeros_motors.copy(),
         )
 
 
@@ -116,13 +117,14 @@ class SimControls:
     """Body rate control data."""
     force_torque: ControlData | None
     """Force and torque control data."""
-    rotor_vel: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, 4)
+    rotor_vel: Array = field(metadata={CORE_NDIM_KEY: 1})  # (N, M, n_motors)
     """Desired motor speed."""
 
     @staticmethod
     def create(
         n_worlds: int,
         n_drones: int,
+        n_motors: int,
         control: Control,
         drone: str,
         state_freq: int | None,
@@ -132,7 +134,7 @@ def create(
         device: Device,
     ) -> SimControls:
         """Create a default set of controls for the simulation."""
-        rotor_vel = jnp.zeros((n_worlds, n_drones, 4), device=device)
+        rotor_vel = jnp.zeros((n_worlds, n_drones, n_motors), device=device)
         match control:
             case Control.state:
                 state = MellingerStateData.create(n_worlds, n_drones, state_freq, drone, device)

crazyflow/sim/functional.py

@@ -75,7 +75,9 @@ def rotor_vel_control(data: SimData, controls: Array) -> SimData:
     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"
+    n_motors = data.states.rotor_vel.shape[-1]
+    expected_shape = (data.core.n_worlds, data.core.n_drones, n_motors)
+    assert controls.shape == expected_shape, "controls shape mismatch"
     controls = jnp.asarray(controls)
     return data.replace(controls=data.controls.replace(rotor_vel=controls))
 

crazyflow/sim/sim.py

@@ -497,13 +497,15 @@ def init_data(
         drone_mocap_ids = [
             self.mj_model.body(f"{drone_name}:{i}").mocapid.item() for i in range(self.n_drones)
         ]
+        n_motors = len(load_drone_params(self.drone)["mixing_matrix"][-1])
         N, D = self.n_worlds, self.n_drones
         data = SimData(
-            states=SimState.create(N, D, self.device),
-            states_deriv=SimStateDeriv.create(N, D, self.device),
+            states=SimState.create(N, D, n_motors, self.device),
+            states_deriv=SimStateDeriv.create(N, D, n_motors, self.device),
             controls=SimControls.create(
                 N,
                 D,
+                n_motors,
                 self.control,
                 self.drone,
                 state_freq,
@@ -701,11 +703,12 @@ def clip_floor_pos(data: SimData) -> SimData:
 def rotor_vel_limits(dynamics: Dynamics, drone: str) -> tuple[float, float]:
     """Limits of ``rotor_vel`` in RPM (first principles) or collective thrust in N (others)."""
     params = load_drone_params(drone)
+    n_motors = np.asarray(params["mixing_matrix"]).shape[-1]
     thrust_min, thrust_max = params["thrust_min"], params["thrust_max"]
     if dynamics == Dynamics.first_principles:
         rpm = motor_force2rotor_vel(np.asarray([thrust_min, thrust_max]), params["rpm2thrust"])
         return float(rpm[0]), float(rpm[1])
-    return 4 * thrust_min, 4 * thrust_max
+    return n_motors * thrust_min, n_motors * thrust_max
 
 
 def clip_rotor_vel(data: SimData, lower: Array | float, upper: Array | float) -> SimData:

tests/unit/test_arbitrary_motor_count.py

@@ -0,0 +1,147 @@
+"""Regression tests for #104: motor count must be derived from ``mixing_matrix.shape[-1]``.
+
+These tests exercise the changed pure functions directly with a synthetic 6-column mixing matrix,
+without registering a full hexacopter drone (MuJoCo asset + fitted dynamics coefficients for all
+four dynamics models), since none of that exists for a real non-quadcopter platform yet.
+"""
+
+from __future__ import annotations
+
+import jax
+import numpy as np
+import pytest
+
+from crazyflow.control import Control, load_params
+from crazyflow.control.mellinger import force_torque2rotor_vel, state2attitude
+from crazyflow.control.transform import motor_force2rotor_vel
+from crazyflow.dynamics import Dynamics
+from crazyflow.envs.drone_env import action_space
+from crazyflow.sim.data import SimControls, SimCore, SimData, SimParams, SimState, SimStateDeriv
+from crazyflow.sim.functional import rotor_vel_control
+from crazyflow.sim.sim import rotor_vel_limits
+
+N_MOTORS = 6
+
+# Arbitrary, physically meaningless mixing matrix: only the motor count (last dimension) matters.
+HEXA_MIXING_MATRIX = np.array(
+    [
+        [-1.0, -1.0, 0.0, 1.0, 1.0, 0.0],
+        [-0.5, 0.5, 1.0, 0.5, -0.5, -1.0],
+        [1.0, -1.0, 1.0, -1.0, 1.0, -1.0],
+    ]
+)
+
+
+@pytest.mark.unit
+def test_sim_data_buffers_scale_with_n_motors():
+    device = jax.devices("cpu")[0]
+    states = SimState.create(n_worlds=2, n_drones=3, n_motors=N_MOTORS, device=device)
+    assert states.rotor_vel.shape == (2, 3, N_MOTORS)
+    assert states.quat.shape == (2, 3, 4)  # Unaffected: quaternion, not motor count
+
+    states_deriv = SimStateDeriv.create(n_worlds=2, n_drones=3, n_motors=N_MOTORS, device=device)
+    assert states_deriv.rotor_acc.shape == (2, 3, N_MOTORS)
+
+    controls = SimControls.create(
+        n_worlds=2,
+        n_drones=3,
+        n_motors=N_MOTORS,
+        control=Control.rotor_vel,
+        drone="cf2x_L250",
+        state_freq=None,
+        attitude_freq=None,
+        body_rate_freq=None,
+        force_torque_freq=None,
+        device=device,
+    )
+    assert controls.rotor_vel.shape == (2, 3, N_MOTORS)
+
+
+def _build_rotor_vel_sim_data(n_worlds: int, n_drones: int, n_motors: int) -> SimData:
+    """Minimal SimData in Control.rotor_vel mode, with rotor buffers sized for n_motors."""
+    device = jax.devices("cpu")[0]
+    return SimData(
+        states=SimState.create(n_worlds, n_drones, n_motors, device),
+        states_deriv=SimStateDeriv.create(n_worlds, n_drones, n_motors, device),
+        controls=SimControls.create(
+            n_worlds,
+            n_drones,
+            n_motors,
+            Control.rotor_vel,
+            "cf2x_L250",
+            None,
+            None,
+            None,
+            None,
+            device,
+        ),
+        params=SimParams.create(Dynamics.first_principles, "cf2x_L250", device),
+        core=SimCore.create(500, n_worlds, n_drones, list(range(n_drones)), 0, device),
+    )
+
+
+@pytest.mark.unit
+def test_rotor_vel_control_accepts_n_motors_shape():
+    n_worlds, n_drones = 2, 1
+    data = _build_rotor_vel_sim_data(n_worlds, n_drones, N_MOTORS)
+
+    controls = np.full((n_worlds, n_drones, N_MOTORS), 1000.0)
+    updated = rotor_vel_control(data, controls)
+    assert updated.controls.rotor_vel.shape == (n_worlds, n_drones, N_MOTORS)
+
+    # A control array shaped for the old hardcoded 4-motor assumption must be rejected.
+    wrong_shape_controls = np.full((n_worlds, n_drones, 4), 1000.0)
+    with pytest.raises(AssertionError):
+        rotor_vel_control(data, wrong_shape_controls)
+
+
+@pytest.mark.unit
+def test_force_torque2rotor_vel_scales_with_mixing_matrix():
+    params = load_params(force_torque2rotor_vel, "cf2x_L250")
+    params["mixing_matrix"] = HEXA_MIXING_MATRIX
+    # Zero torque: thrust must split evenly across all N_MOTORS motors, not divided by 4. Keep the
+    # per-motor share (force / N_MOTORS) within [thrust_min, thrust_max] so it isn't clipped.
+    force = np.array([0.3])
+    torque = np.zeros(3)
+    rotor_vel = force_torque2rotor_vel(force, torque, **params)
+    assert rotor_vel.shape == (N_MOTORS,)
+    expected = motor_force2rotor_vel(force / N_MOTORS, params["rpm2thrust"])
+    assert rotor_vel == pytest.approx(np.full(N_MOTORS, expected.item()))
+
+
+@pytest.mark.unit
+def test_state2attitude_scales_with_mixing_matrix():
+    params = load_params(state2attitude, "cf2x_L250")
+    params["mixing_matrix"] = HEXA_MIXING_MATRIX
+    pos, quat, vel = np.zeros(3), np.array([0.0, 0.0, 0.0, 1.0]), np.zeros(3)
+    cmd = np.zeros(13)
+    rpyt, pos_err_i = state2attitude(pos, quat, vel, cmd, ctrl_freq=100, **params)
+    # Collective thrust command stays 4D ([roll, pitch, yaw, thrust]) regardless of motor count.
+    assert rpyt.shape == (4,)
+    assert pos_err_i.shape == (3,)
+
+
+@pytest.mark.unit
+def test_rotor_vel_limits_scale_with_n_motors(monkeypatch: pytest.MonkeyPatch):
+    import crazyflow.sim.sim as sim_module
+
+    real_params = sim_module.load_drone_params("cf2x_L250")
+    hexa_params = real_params | {"mixing_matrix": HEXA_MIXING_MATRIX.tolist()}
+    monkeypatch.setattr(sim_module, "load_drone_params", lambda drone: hexa_params)
+
+    lower, upper = rotor_vel_limits(Dynamics.so_rpy, "cf2x_L250")
+    assert lower == pytest.approx(N_MOTORS * real_params["thrust_min"])
+    assert upper == pytest.approx(N_MOTORS * real_params["thrust_max"])
+
+
+@pytest.mark.unit
+def test_action_space_thrust_bounds_scale_with_n_motors(monkeypatch: pytest.MonkeyPatch):
+    import crazyflow.envs.drone_env as drone_env_module
+
+    real_params = drone_env_module.load_params("cf2x_L250")
+    hexa_params = real_params | {"mixing_matrix": HEXA_MIXING_MATRIX.tolist()}
+    monkeypatch.setattr(drone_env_module, "load_params", lambda drone: hexa_params)
+
+    space = action_space(Control.attitude, "cf2x_L250")
+    assert space.low[-1] == pytest.approx(N_MOTORS * real_params["thrust_min"])
+    assert space.high[-1] == pytest.approx(N_MOTORS * real_params["thrust_max"])