Arbitrary motor count

[Lagnesh2003]

Derive motor count from mixing_matrix.shape[-1] instead of hardcoding 4, so the simulator is more generalized.

Changes

• SimState, SimStateDeriv, SimControls (crazyflow/sim/data.py): rotor buffers (rotor_vel, rotor_acc) are now sized by n_motors.
• Sim.init_data and rotor_vel_limits (crazyflow/sim/sim.py): n_motors is derived from mixing_matrix.shape[-1] and threaded through; thrust limits scale with n_motors instead of 4 *.
• rotor_vel_control (crazyflow/sim/functional.py): the per-motor shape assertion now reads the motor count off data.states.rotor_vel instead of asserting a fixed 4.
• action_space (crazyflow/envs/drone_env.py): attitude-mode thrust bounds now scale with n_motors instead of * 4.
• Mellinger controller (crazyflow/control/mellinger/control.py): force_torque2rotor_vel and _attitude2force_torque divide by mixing_matrix.shape[-1] instead of 4; state2attitude now takes mixing_matrix as a keyword-only parameter (auto-bound via load_params/parametrize) instead of hardcoding * 4.
• Added tests/unit/test_arbitrary_motor_count.py: 6 regression tests exercising every changed function directly with a synthetic 6-motor mixing_matrix/SimData, without registering a full hexacopter drone. Verified each test fails against the pre-fix code and passes against the fix.

Fixes #104

Test plan

• pytest -v tests — 656 passed (650 existing + 6 new)
• pytest --markdown-docs ... (doctests/docs examples) — 101 passed
• ruff check crazyflow/ tests/ — clean
• New regression tests confirmed to fail on pre-fix code, pass on fixed code

[amacati]

Thanks for working on this. The change is a little more complex than just replacing the 4s with the last dim of the mixing matrix. We want to achieve a collective thrust and 3 torques with N rotors. This happens to have a unique solution in our current setup, but with more than 4 rotors, this requires some form of the inverse of the mixing matrix, e.g. the pseudo-inverse for a least-squares allocation.

The Mellinger controller was not built to run this, so I would rather not add support for N rotors into it if the firmware cannot handle that case. PX4 on the other hand should be able to compute that, and I suspect we will find the pseudo-inverse in its controller code. So my feeling rn is we should wait until we have a JAX version of PX4 and then add support for arbitrary rotors.

On the other hand, if there is a controller that can fly N rotors that is open-source and you want to contribute, that would be fine as well. Not sure if e.g. betaflight can take more than 4 rotors.

[Lagnesh2003]

Thanks for the explanation, that makes a lot of sense.

I looked around a bit and couldn't find any existing JAX port of PX4, so that'd be starting from scratch rather than porting something that already exists — good to know before anyone commits to it.

I'd like to split this into two parts if that works for you.

• First, I'll narrow this PR down to just the parts that don't touch the Mellinger controller — the buffer/shape generalization in SimState/SimStateDeriv/SimControls, rotor_vel_limits, and the action space bounds. That part's correct and useful on its own, independent of how the N-rotor control question gets resolved.
• Second, separately from this PR, I'd genuinely like to dig into what proper N-rotor allocation would take — maybe betaflight if it supports more than 4 rotors, or PX4 down the line. I am genuinely interested to work on this.

[amacati]

Sounds good, though it might be tricky if we assume that some drones only have abstracted models without mixing matrix. Btw, what's your motivation for working on this? Are you doing research in this area?