diff --git a/crazyflow/control/mellinger/control.py b/crazyflow/control/mellinger/control.py index 384b6dc8..6c3f7408 100644 --- a/crazyflow/control/mellinger/control.py +++ b/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)) diff --git a/crazyflow/envs/drone_env.py b/crazyflow/envs/drone_env.py index 12930a52..615bc3e7 100644 --- a/crazyflow/envs/drone_env.py +++ b/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), diff --git a/crazyflow/sim/data.py b/crazyflow/sim/data.py index b66ec24e..3d7cb27c 100644 --- a/crazyflow/sim/data.py +++ b/crazyflow/sim/data.py @@ -37,15 +37,16 @@ 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), @@ -53,7 +54,7 @@ def create(n_worlds: int, n_drones: int, device: Device) -> SimState: 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) diff --git a/crazyflow/sim/functional.py b/crazyflow/sim/functional.py index 7a592f2b..ca1f5fde 100644 --- a/crazyflow/sim/functional.py +++ b/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)) diff --git a/crazyflow/sim/sim.py b/crazyflow/sim/sim.py index 055f8a54..8faeb300 100644 --- a/crazyflow/sim/sim.py +++ b/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: diff --git a/tests/unit/test_arbitrary_motor_count.py b/tests/unit/test_arbitrary_motor_count.py new file mode 100644 index 00000000..04ee2535 --- /dev/null +++ b/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"])