From 2d4f39ef1ce19f345b2334fb0556bbb98ca682c9 Mon Sep 17 00:00:00 2001 From: ratheron Date: Sat, 19 Sep 2026 15:06:35 +0200 Subject: [PATCH 1/5] Automate available drones and add docs on how to add drones --- SKILL.md | 7 +-- crazyflow/drones/__init__.py | 11 +++-- crazyflow/dynamics/__init__.py | 22 +++++++--- crazyflow/dynamics/core.py | 20 +++++++++ .../dynamics/first_principles/params.toml | 13 ++++-- docs/user-guide/adding-drones.md | 43 +++++++++++++++++++ docs/user-guide/dynamics/parametrize.md | 21 ++++++--- docs/user-guide/index.md | 1 + properdocs.yml | 1 + tests/conftest.py | 24 +++++++++++ tests/integration/test_models.py | 5 +-- tests/unit/dynamics/test_dynamics.py | 20 +++------ tests/unit/dynamics/test_parametrization.py | 19 +++----- 13 files changed, 152 insertions(+), 55 deletions(-) create mode 100644 docs/user-guide/adding-drones.md diff --git a/SKILL.md b/SKILL.md index 0a28e084..7346e891 100644 --- a/SKILL.md +++ b/SKILL.md @@ -49,9 +49,10 @@ against all models in the simulation's `build_control_fns`. Define the function in `dynamics.py` and never in the package `__init__.py`, because `load_fn_params` derives the model name from `fn.__module__.split(".")[-2]`. `parametrize` binds exactly the keyword-only parameters after the bare `*`, so anything before it is never bound. -Every drone in `available_drones` needs a complete section in every model's -`crazyflow/dynamics/*/params.toml`. The commented example at the top of each file lists the keys. -Only `gravity_vec` is global, in `crazyflow/dynamics/params.toml`. +`available_drones` is the set of MJCF files in `crazyflow/drones`. A model supports a drone when its +`crazyflow/dynamics//params.toml` has a complete section for it; the commented example at the +top of each file lists the keys. `supported_drones` and `supported_dynamics` report the pairs, and the +tests only run those. Only `gravity_vec` is global, in `crazyflow/dynamics/params.toml`. Registration alone produces roughly 40 parametrized tests. These do not include derivatives tests. diff --git a/crazyflow/drones/__init__.py b/crazyflow/drones/__init__.py index ace00398..909c02f7 100644 --- a/crazyflow/drones/__init__.py +++ b/crazyflow/drones/__init__.py @@ -6,11 +6,10 @@ Use ``available_drones`` to enumerate the supported configurations. """ -# Currently supported platforms: -# * **cf2x_L250** — Crazyflie 2.x -# * **cf2x_P250** — Crazyflie 2.x with plus propellers -# * **cf2x_T350** — Crazyflie 2.x with thrust upgrade kit -# * **cf21B_500** — Crazyflie 2.1 Brushless with 500 mAh battery -available_drones: tuple[str, ...] = ("cf2x_L250", "cf2x_P250", "cf2x_T350", "cf21B_500") +from pathlib import Path __all__ = ["available_drones"] + +_mjcf_files = sorted(Path(__file__).parent.glob("*.xml")) +available_drones: tuple[str, ...] = tuple(p.stem for p in _mjcf_files) +"""Names of all drone configurations, i.e. the MJCF files in ``crazyflow/drones``.""" diff --git a/crazyflow/dynamics/__init__.py b/crazyflow/dynamics/__init__.py index 0e2ad1b9..1f172cda 100644 --- a/crazyflow/dynamics/__init__.py +++ b/crazyflow/dynamics/__init__.py @@ -14,7 +14,14 @@ from typing import Callable -from crazyflow.dynamics.core import Dynamics, load_fn_params, load_params, parametrize +from crazyflow.dynamics.core import ( + Dynamics, + load_fn_params, + load_params, + parametrize, + supported_drones, + supported_dynamics, +) from crazyflow.dynamics.first_principles import dynamics as _first_principles_dynamics from crazyflow.dynamics.so_rpy import dynamics as _so_rpy_dynamics from crazyflow.dynamics.so_rpy_rotor import dynamics as _so_rpy_rotor_dynamics @@ -26,16 +33,19 @@ "load_fn_params", "available_dynamics", "dynamics_features", + "supported_drones", + "supported_dynamics", "Dynamics", ] -available_dynamics: dict[str, Callable] = { - "first_principles": _first_principles_dynamics, - "so_rpy": _so_rpy_dynamics, - "so_rpy_rotor": _so_rpy_rotor_dynamics, - "so_rpy_rotor_drag": _so_rpy_rotor_drag_dynamics, +available_dynamics: dict[Dynamics, Callable] = { + Dynamics.first_principles: _first_principles_dynamics, + Dynamics.so_rpy: _so_rpy_dynamics, + Dynamics.so_rpy_rotor: _so_rpy_rotor_dynamics, + Dynamics.so_rpy_rotor_drag: _so_rpy_rotor_drag_dynamics, } +"""Unparametrized dynamics functions keyed by [Dynamics][crazyflow.dynamics.Dynamics] mode.""" def dynamics_features(dynamics: Callable) -> dict[str, bool]: diff --git a/crazyflow/dynamics/core.py b/crazyflow/dynamics/core.py index b93cc608..c0017d55 100644 --- a/crazyflow/dynamics/core.py +++ b/crazyflow/dynamics/core.py @@ -9,6 +9,7 @@ import numpy as np +from crazyflow.drones import available_drones from crazyflow.utils import filter_to_signature, to_xp from crazyflow.utils import parametrize as _parametrize @@ -141,3 +142,22 @@ def load_fn_params( assert callable(fn), f"Expected a function, got {type(fn)}" dynamics = fn.__module__.split(".")[-2] return filter_to_signature(load_params(dynamics, drone, xp=xp, device=device), fn) + + +def _param_sections(dynamics: Dynamics | str) -> set[str]: + """Return the drone sections declared in a dynamics model's ``params.toml``.""" + with open(Path(__file__).parent / f"{Dynamics(dynamics)}/params.toml", "rb") as f: + return set(tomllib.load(f)) + + +def supported_drones(dynamics: Dynamics | str) -> tuple[str, ...]: + """Return the drones that ``dynamics`` can be parametrized for.""" + sections = _param_sections(dynamics) + return tuple(drone for drone in available_drones if drone in sections) + + +def supported_dynamics(drone: str) -> tuple[Dynamics, ...]: + """Return the dynamics models that ``drone`` can be simulated with.""" + if drone not in available_drones: + raise KeyError(f"Drone `{drone}` not found. Available drones: {available_drones}") + return tuple(d for d in Dynamics if drone in _param_sections(d)) diff --git a/crazyflow/dynamics/first_principles/params.toml b/crazyflow/dynamics/first_principles/params.toml index aa8c66ee..3770db3c 100644 --- a/crazyflow/dynamics/first_principles/params.toml +++ b/crazyflow/dynamics/first_principles/params.toml @@ -10,22 +10,27 @@ # thrust_min = 0.0 # N per motor # thrust_max = 0.0 # N per motor # L = 0.0 # m, CoM to motor distance -# prop_inertia = 0.0 # kg m^2 +# prop_inertia = 0.0 # kg m^2, zero drops the gyroscopic and reaction torque of the propellers # rpm2thrust = [0.0, 0.0, 0.0] # N, polynomial in RPM by index # rpm2torque = [0.0, 0.0, 0.0] # Nm, polynomial in RPM by index -# rotor_dyn_coef = [0.0, 0.0, 0.0, 0.0] # The coefficients are [a,b,c,d], where a & b are the - # linear & quadratic rise constants. c & d equally for fall. +# rotor_dyn_coef = [0.0, 0.0, 0.0, 0.0] # The coefficients are [a,b,c,d], where a & b are the +# # linear & quadratic rise constants. c & d equally for fall. +# # [1/tau, 0.0, 1/tau, 0.0] is a symmetric first order +# # model with time constant tau. # mixing_matrix = [ # motor locations relative to the CoM and turn directions # [-1.0, -1.0, 1.0, 1.0], # [-1.0, 1.0, 1.0, -1.0], # [-1.0, 1.0, -1.0, 1.0] # ] -# drag_matrix = [ # 1/s +# drag_matrix = [ # 1/s, zero disables drag # [0.0, 0.0, 0.0], # [0.0, 0.0, 0.0], # [0.0, 0.0, 0.0] # ] # +# Not every parameter has to be identified, but all of them have to be set. A symmetric first order +# rotor model, no drag and no propeller gyroscopic torque are reasonable defaults for a new drone. +# gravity_vec is global and lives in crazyflow/dynamics/params.toml. [cf2x_L250] mass = 0.0319 diff --git a/docs/user-guide/adding-drones.md b/docs/user-guide/adding-drones.md new file mode 100644 index 00000000..b688404b --- /dev/null +++ b/docs/user-guide/adding-drones.md @@ -0,0 +1,43 @@ +# Adding a drone + +A drone is defined by data files only. `available_drones` lists the MJCF files in `crazyflow/drones`, and a dynamics model or controller supports a drone when its own `params.toml` has a section for it. Adding a platform therefore means adding files and sections, not registering anything in Python. + +```python +from crazyflow import available_drones +from crazyflow.dynamics import supported_dynamics + +available_drones # ('cf21B_500', 'cf2x_L250', 'cf2x_P250', 'cf2x_T350') +supported_dynamics("cf2x_L250") # (first_principles, so_rpy, so_rpy_rotor, so_rpy_rotor_drag) +``` + +Pick a short name such as `cf2x_L250` (platform, then variant) and use it everywhere below. + +## 1. MuJoCo model + +Add `crazyflow/drones/.xml` with its meshes under `crazyflow/drones/assets//`. This file is what makes the drone appear in `available_drones`. The simulator attaches the body named `drone` once per drone, so that body is required. If you also provide a `drone_fused` body whose visual geometry is a single mesh, users can select it with `Sim(fused_mjx_model=True)` for cheaper rendering. See [MuJoCo Integration](mujoco.md) for how the scene is assembled. + +## 2. Dynamics parameters + +Each dynamics model has its own `crazyflow/dynamics//params.toml`. Add a `[]` section to every model you want to offer for the drone. The commented example at the top of each file lists the keys the model needs, and all of them must be set: + +- Mass, inertia and the per-motor thrust limits appear in every model. The fitted `so_rpy` models only use the inertia to apply external torques, so an estimate works there at the cost of wrong reactions to disturbance torques. +- `first_principles` additionally needs the hardware constants: arm length, thrust and torque curves and mixing matrix. The remaining keys have to be set but not identified: `rotor_dyn_coef = [1/tau, 0.0, 1/tau, 0.0]` is a symmetric first order rotor model with time constant `tau`, a zero `drag_matrix` disables drag, and a zero `prop_inertia` drops the gyroscopic torque of the propellers. +- The fitted `so_rpy`, `so_rpy_rotor` and `so_rpy_rotor_drag` models need identified coefficients. Use the [system identification pipeline](dynamics/system-identification.md) to obtain them from flight data. + +Gravity is global and lives in `crazyflow/dynamics/params.toml`. A model without a section is simply not offered for that drone. [`supported_dynamics`][crazyflow.dynamics.supported_dynamics] and [`supported_drones`][crazyflow.dynamics.supported_drones] report the available pairs, and `Sim` raises `KeyError` for any other combination. + +## 3. Controller parameters + +Add `[.core]`, `[.state2attitude]`, `[.attitude2force_torque]` and `[.body_rate2force_torque]` sections to `crazyflow/control/mellinger/params.toml`. These reproduce the onboard firmware, so the values may deliberately differ from the physical constants in step 2. See [Mellinger controller](control/mellinger.md). + +## 4. Documentation + +Add the platform to the table in [Parametrization](dynamics/parametrize.md#available-drone-configurations). + +## 5. Run the tests + +The test suite parametrizes over `available_drones` and over the supported drone-dynamics pairs, so the new drone is tested without any changes to the tests. In particular, `tests/integration/test_models.py` constructs a `Sim` for every supported pair, which loads the MJCF, the dynamics parameters and the controller parameters together. + +```bash +pixi run -e tests tests +``` diff --git a/docs/user-guide/dynamics/parametrize.md b/docs/user-guide/dynamics/parametrize.md index 2222e600..72a8e5cb 100644 --- a/docs/user-guide/dynamics/parametrize.md +++ b/docs/user-guide/dynamics/parametrize.md @@ -21,7 +21,7 @@ The following configurations ship with pre-fitted parameters. They cover both th ```python from crazyflow.drones import available_drones -available_drones # ('cf2x_L250', 'cf2x_P250', 'cf2x_T350', 'cf21B_500') +available_drones # ('cf21B_500', 'cf2x_L250', 'cf2x_P250', 'cf2x_T350') ``` | `drone` | Platform | @@ -31,7 +31,16 @@ available_drones # ('cf2x_L250', 'cf2x_P250', 'cf2x_T350', 'cf21B_500') | `"cf2x_T350"` | Crazyflie 2.x, thrust upgrade kit | | `"cf21B_500"` | Crazyflie 2.1 Brushless | -If your drone is not listed, you can identify the parameters from flight data using the [system identification pipeline](system-identification.md) and inject them into any dynamics. +If your drone is not listed, you can [add it](../adding-drones.md). The fitted models need coefficients identified from flight data with the [system identification pipeline](system-identification.md). + +Not every dynamics is available for every drone. [`supported_dynamics`][crazyflow.dynamics.supported_dynamics] and [`supported_drones`][crazyflow.dynamics.supported_drones] list the available pairs: + +```python +from crazyflow.dynamics import supported_drones, supported_dynamics + +supported_dynamics("cf21B_500") # (first_principles, so_rpy, so_rpy_rotor, so_rpy_rotor_drag) +supported_drones("so_rpy_rotor_drag") # ('cf21B_500', 'cf2x_L250', 'cf2x_P250', 'cf2x_T350') +``` ## Switching array backends @@ -101,14 +110,14 @@ dynamics.keywords["mass"] = np.float64(0.040) # heavier drone — applies to ev ## Selecting dynamics programmatically -`available_dynamics` is a dict mapping dynamics names to their unparametrized functions. This is useful when selecting a dynamics by name. +`available_dynamics` is a dict mapping each [`Dynamics`][crazyflow.dynamics.Dynamics] mode to its unparametrized function. `Dynamics` is a string enum, so plain names work as keys too. ```python -from crazyflow.dynamics import available_dynamics, parametrize +from crazyflow.dynamics import Dynamics, available_dynamics, parametrize -list(available_dynamics) # ['first_principles', 'so_rpy', 'so_rpy_rotor', 'so_rpy_rotor_drag'] +list(available_dynamics) # [Dynamics.first_principles, Dynamics.so_rpy, ...] -dynamics = available_dynamics["so_rpy_rotor_drag"] +dynamics = available_dynamics[Dynamics.so_rpy_rotor_drag] parametrized_dynamics = parametrize(dynamics, drone="cf2x_T350") ``` diff --git a/docs/user-guide/index.md b/docs/user-guide/index.md index a9ce0692..671f7843 100644 --- a/docs/user-guide/index.md +++ b/docs/user-guide/index.md @@ -7,6 +7,7 @@ In-depth documentation for every part of the simulator. - [Functional API](functional-api.md) — purely functional interface for JAX transformations - [Dynamics](dynamics/index.md) — first-principles vs. fitted dynamics, when to use each - [Control Modes](control/index.md) — state, attitude, body rate, force/torque, and rotor velocity control +- [Adding a drone](adding-drones.md) — the MJCF model and parameter sections that define a platform - [Pipelines](pipelines.md) — composable step and reset pipelines, randomization, and disturbances - [The world axis](world-axis.md) — which arrays are batched over worlds, and what resets and sharding do with them - [Visualization](visualization.md) — rendering modes, cameras, raycasting, and materials diff --git a/properdocs.yml b/properdocs.yml index 3463a604..d09b19bd 100644 --- a/properdocs.yml +++ b/properdocs.yml @@ -69,6 +69,7 @@ nav: - Integral errors: user-guide/control/integral-errors.md - Batching: user-guide/control/batching.md - JIT compilation: user-guide/control/jit.md + - Adding a drone: user-guide/adding-drones.md - Pipelines: user-guide/pipelines.md - The world axis: user-guide/world-axis.md - Sharding: user-guide/sharding.md diff --git a/tests/conftest.py b/tests/conftest.py index 7c88d620..854ff187 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -7,6 +7,7 @@ import jax import pytest +from _pytest.mark import ParameterSet # The cache dir is per-user. A shared dir like /tmp/jax_cache breaks on multi-user machines, since # jax hard-fails on GPU autotune cache writes when another user owns the directory. @@ -43,3 +44,26 @@ def device() -> str: os.environ.get("DISPLAY") is None, reason="DISPLAY is not set, skipping test in headless environment", ) + + +def drone_dynamics_fns() -> list[ParameterSet]: + """Return all supported (dynamics, dynamics function, drone) combinations.""" + from crazyflow.dynamics import available_dynamics, supported_drones + + return [ + pytest.param(name, fn, drone, id=f"{drone}-{name}") + for name, fn in available_dynamics.items() + for drone in supported_drones(name) + ] + + +def drone_dynamics() -> list[ParameterSet]: + """Return all supported (dynamics, drone) combinations.""" + from crazyflow.drones import available_drones + from crazyflow.dynamics import supported_dynamics + + return [ + pytest.param(dynamics, drone, id=f"{drone}-{dynamics}") + for drone in available_drones + for dynamics in supported_dynamics(drone) + ] diff --git a/tests/integration/test_models.py b/tests/integration/test_models.py index 53d88966..3e1790da 100644 --- a/tests/integration/test_models.py +++ b/tests/integration/test_models.py @@ -1,13 +1,12 @@ import pytest +from conftest import drone_dynamics -from crazyflow import available_drones from crazyflow.dynamics import Dynamics from crazyflow.sim import Sim @pytest.mark.integration -@pytest.mark.parametrize("dynamics", Dynamics) -@pytest.mark.parametrize("drone", available_drones) +@pytest.mark.parametrize("dynamics, drone", drone_dynamics()) def test_attitude_symbolic(dynamics: Dynamics, drone: "str"): """Tests if xml files contain syntax errors.""" Sim(dynamics=dynamics, drone=drone) diff --git a/tests/unit/dynamics/test_dynamics.py b/tests/unit/dynamics/test_dynamics.py index 8333bae1..0cdcb285 100644 --- a/tests/unit/dynamics/test_dynamics.py +++ b/tests/unit/dynamics/test_dynamics.py @@ -12,8 +12,8 @@ import numpy as np import pytest from array_api_compat import device as xp_device +from conftest import drone_dynamics_fns -from crazyflow.drones import available_drones from crazyflow.dynamics import available_dynamics, dynamics_features from crazyflow.dynamics.core import parametrize @@ -177,15 +177,13 @@ def test_dynamics_features(dynamics_name: str, dynamics: Callable): @pytest.mark.unit -@pytest.mark.parametrize("dynamics_name, dynamics", available_dynamics.items()) -@pytest.mark.parametrize("drone", available_drones) +@pytest.mark.parametrize("dynamics_name, dynamics, drone", drone_dynamics_fns()) def test_dynamics_shapes(dynamics_name: str, dynamics: Callable, drone: str): check_shapes(parametrize(dynamics, drone)) @pytest.mark.unit -@pytest.mark.parametrize("dynamics_name, dynamics", available_dynamics.items()) -@pytest.mark.parametrize("drone", available_drones) +@pytest.mark.parametrize("dynamics_name, dynamics, drone", drone_dynamics_fns()) def test_dynamics_shapes_batched(dynamics_name: str, dynamics: Callable, drone: str): dynamics = parametrize(dynamics, drone, xp=xp) batch = (10, 5) @@ -196,8 +194,7 @@ def test_dynamics_shapes_batched(dynamics_name: str, dynamics: Callable, drone: @pytest.mark.unit -@pytest.mark.parametrize("dynamics_name, dynamics", available_dynamics.items()) -@pytest.mark.parametrize("drone", available_drones) +@pytest.mark.parametrize("dynamics_name, dynamics, drone", drone_dynamics_fns()) @pytest.mark.parametrize("ext_wrench", [False, True]) @pytest.mark.parametrize("per_motor_params", [False, True]) def test_symbolic_dynamics( @@ -229,8 +226,7 @@ def test_symbolic_dynamics( @pytest.mark.unit -@pytest.mark.parametrize("dynamics_name, dynamics", available_dynamics.items()) -@pytest.mark.parametrize("drone", available_drones) +@pytest.mark.parametrize("dynamics_name, dynamics, drone", drone_dynamics_fns()) def test_compare_batched_non_batched(dynamics_name: str, dynamics: Callable, drone: str): """Tests if batching works and if the results are identical to the non-batched version.""" dynamics = parametrize(dynamics, drone) @@ -246,8 +242,7 @@ def test_compare_batched_non_batched(dynamics_name: str, dynamics: Callable, dro @pytest.mark.unit -@pytest.mark.parametrize("dynamics_name, dynamics", available_dynamics.items()) -@pytest.mark.parametrize("drone", available_drones) +@pytest.mark.parametrize("dynamics_name, dynamics, drone", drone_dynamics_fns()) def test_batched_params(dynamics_name: str, dynamics: Callable, drone: str): """Tests if batched parameters give the same results as the shared parameters.""" dynamics = parametrize(dynamics, drone, xp=xp) @@ -274,8 +269,7 @@ def test_batched_params(dynamics_name: str, dynamics: Callable, drone: str): @pytest.mark.unit -@pytest.mark.parametrize("dynamics_name, dynamics", available_dynamics.items()) -@pytest.mark.parametrize("drone", available_drones) +@pytest.mark.parametrize("dynamics_name, dynamics, drone", drone_dynamics_fns()) def test_numeric_jit(dynamics_name: str, dynamics: Callable, drone: str): """Tests if the dynamics are jitable and if the results are identical to the array API ones.""" dynamics = parametrize(dynamics, drone) diff --git a/tests/unit/dynamics/test_parametrization.py b/tests/unit/dynamics/test_parametrization.py index c4c436b9..7de74f04 100644 --- a/tests/unit/dynamics/test_parametrization.py +++ b/tests/unit/dynamics/test_parametrization.py @@ -5,29 +5,21 @@ from typing import Callable import pytest +from conftest import drone_dynamics_fns -from crazyflow.drones import available_drones -from crazyflow.dynamics import ( - Dynamics, - available_dynamics, - load_fn_params, - load_params, - parametrize, -) +from crazyflow.dynamics import Dynamics, load_fn_params, load_params, parametrize from crazyflow.dynamics.so_rpy import dynamics as so_rpy @pytest.mark.unit -@pytest.mark.parametrize("dynamics_name, dynamics", available_dynamics.items()) -@pytest.mark.parametrize("drone", available_drones) +@pytest.mark.parametrize("dynamics_name, dynamics, drone", drone_dynamics_fns()) def test_dynamics_parameter_loading(dynamics_name: str, dynamics: Callable, drone: str) -> None: """Check that parameters can be loaded for all available dynamics and drones.""" load_fn_params(dynamics, drone) @pytest.mark.unit -@pytest.mark.parametrize("dynamics_name, dynamics", available_dynamics.items()) -@pytest.mark.parametrize("drone", available_drones) +@pytest.mark.parametrize("dynamics_name, dynamics, drone", drone_dynamics_fns()) def test_model_parameter_loading(dynamics_name: str, dynamics: Callable, drone: str) -> None: """Check that all parameters of a model can be loaded for all drones.""" params = load_params(dynamics_name, drone) @@ -51,8 +43,7 @@ def test_unknown_dynamics() -> None: @pytest.mark.unit -@pytest.mark.parametrize("dynamics_name, dynamics", available_dynamics.items()) -@pytest.mark.parametrize("drone", available_drones) +@pytest.mark.parametrize("dynamics_name, dynamics, drone", drone_dynamics_fns()) def test_dynamics_parametrization(dynamics_name: str, dynamics: Callable, drone: str): """Check that we can parametrize all available dynamics with all drones.""" parametrize(dynamics, drone) From 1cb2cbd7cb8916ceb346ad003ee0bcc14f5fec0d Mon Sep 17 00:00:00 2001 From: ratheron Date: Sat, 19 Sep 2026 15:54:49 +0200 Subject: [PATCH 2/5] Rename available_drones into Drone StrEnum --- SKILL.md | 2 +- crazyflow/__init__.py | 4 +- crazyflow/control/core.py | 13 ++++--- crazyflow/drones/__init__.py | 11 +++--- crazyflow/dynamics/core.py | 37 +++++++++--------- crazyflow/envs/drone_env.py | 5 ++- crazyflow/sim/sim.py | 5 ++- docs/index.md | 4 +- docs/user-guide/adding-drones.md | 12 +++--- docs/user-guide/dynamics/parametrize.md | 4 +- docs/user-guide/oo-api.md | 2 +- tests/conftest.py | 4 +- tests/unit/control/test_core.py | 6 +-- tests/unit/control/test_mellinger.py | 42 ++++++++++----------- tests/unit/dynamics/test_parametrization.py | 30 +++++++++++++-- 15 files changed, 103 insertions(+), 78 deletions(-) diff --git a/SKILL.md b/SKILL.md index 7346e891..2e853d1a 100644 --- a/SKILL.md +++ b/SKILL.md @@ -49,7 +49,7 @@ against all models in the simulation's `build_control_fns`. Define the function in `dynamics.py` and never in the package `__init__.py`, because `load_fn_params` derives the model name from `fn.__module__.split(".")[-2]`. `parametrize` binds exactly the keyword-only parameters after the bare `*`, so anything before it is never bound. -`available_drones` is the set of MJCF files in `crazyflow/drones`. A model supports a drone when its +`Drone` is a `StrEnum` of the MJCF files in `crazyflow/drones`. A model supports a drone when its `crazyflow/dynamics//params.toml` has a complete section for it; the commented example at the top of each file lists the keys. `supported_drones` and `supported_dynamics` report the pairs, and the tests only run those. Only `gravity_vec` is global, in `crazyflow/dynamics/params.toml`. diff --git a/crazyflow/__init__.py b/crazyflow/__init__.py index 07d13ea6..f228f5c2 100644 --- a/crazyflow/__init__.py +++ b/crazyflow/__init__.py @@ -18,9 +18,9 @@ import crazyflow.envs # noqa: F401, ensure gymnasium envs are registered from crazyflow.control import Control -from crazyflow.drones import available_drones +from crazyflow.drones import Drone from crazyflow.dynamics import Dynamics from crazyflow.sim import Sim -__all__ = ["Sim", "Dynamics", "Control", "available_drones"] +__all__ = ["Sim", "Dynamics", "Control", "Drone"] __version__ = "0.3.2" diff --git a/crazyflow/control/core.py b/crazyflow/control/core.py index 91247c14..4c5f9e30 100644 --- a/crazyflow/control/core.py +++ b/crazyflow/control/core.py @@ -17,6 +17,7 @@ from types import ModuleType from crazyflow._typing import Array # To be changed to array_api_typing later + from crazyflow.drones import Drone P = ParamSpec("P") R = TypeVar("R") @@ -63,7 +64,7 @@ class Control(StrEnum): def parametrize( - fn: Callable[P, R], drone: str, xp: ModuleType | None = None, device: str | None = None + fn: Callable[P, R], drone: Drone, xp: ModuleType | None = None, device: str | None = None ) -> Callable[P, R]: """Parametrize a controller function with the default controller parameters for a drone. @@ -80,7 +81,7 @@ def parametrize( from crazyflow.control.mellinger import state2attitude from scipy.spatial.transform import Rotation as R - ctrl = parametrize(state2attitude, "cf2x_L250") + ctrl = parametrize(state2attitude, Drone.cf2x_L250) pos, quat = np.zeros(3), np.array([0.0, 0.0, 0.0, 1.0]) vel, cmd = np.zeros(3), np.zeros(16) cmd[9:13] = R.from_euler("z", 0.0).as_quat() @@ -94,7 +95,7 @@ def parametrize( def load_params( - controller: str, drone: str, xp: ModuleType | None = None, device: str | None = None + controller: str, drone: Drone, xp: ModuleType | None = None, device: str | None = None ) -> dict[str, dict[str, Array]]: """Load all parameters of a drone for a controller. @@ -103,7 +104,7 @@ def load_params( Args: controller: Name of the controller package, e.g. ``"mellinger"``. - drone: Name of the drone configuration, e.g. ``"cf2x_L250"``. + drone: The drone configuration, e.g. ``Drone.cf2x_L250``. xp: The array API module to use. If not provided, numpy is used. device: The device to use. If None, the device is inferred from the xp module. @@ -121,7 +122,7 @@ def load_params( def load_fn_params( - fn: Callable, drone: str, xp: ModuleType | None = None, device: str | None = None + fn: Callable, drone: Drone, xp: ModuleType | None = None, device: str | None = None ) -> dict[str, Array]: """Load the parameters a controller function accepts. @@ -131,7 +132,7 @@ def load_fn_params( Args: fn: The controller function for which to load parameters. - drone: Name of the drone configuration, e.g. ``"cf2x_L250"``. + drone: The drone configuration, e.g. ``Drone.cf2x_L250``. xp: The array API module to use. If not provided, numpy is used. device: The device to use. If None, the device is inferred from the xp module. diff --git a/crazyflow/drones/__init__.py b/crazyflow/drones/__init__.py index 909c02f7..3d816778 100644 --- a/crazyflow/drones/__init__.py +++ b/crazyflow/drones/__init__.py @@ -3,13 +3,14 @@ This package bundles the MuJoCo MJCF scene files that define each drone configuration and their referenced meshes (``assets/``). For the physical params, see [crazyflow.dynamics.load_params][]. -Use ``available_drones`` to enumerate the supported configurations. +Use ``Drone`` to enumerate the supported configurations. """ +from enum import StrEnum from pathlib import Path -__all__ = ["available_drones"] +__all__ = ["Drone"] -_mjcf_files = sorted(Path(__file__).parent.glob("*.xml")) -available_drones: tuple[str, ...] = tuple(p.stem for p in _mjcf_files) -"""Names of all drone configurations, i.e. the MJCF files in ``crazyflow/drones``.""" +_drones = [p.stem for p in sorted(Path(__file__).parent.glob("*.xml"))] +Drone: StrEnum = StrEnum("Drone", [(name, name) for name in _drones]) +"""Drone configurations, i.e. the MJCF files in ``crazyflow/drones``.""" diff --git a/crazyflow/dynamics/core.py b/crazyflow/dynamics/core.py index c0017d55..f2dbe18a 100644 --- a/crazyflow/dynamics/core.py +++ b/crazyflow/dynamics/core.py @@ -9,7 +9,7 @@ import numpy as np -from crazyflow.drones import available_drones +from crazyflow.drones import Drone from crazyflow.utils import filter_to_signature, to_xp from crazyflow.utils import parametrize as _parametrize @@ -54,7 +54,7 @@ def decorator(fn: F) -> F: def parametrize( - fn: Callable[P, R], drone: str, xp: ModuleType | None = None, device: str | None = None + fn: Callable[P, R], drone: Drone, xp: ModuleType | None = None, device: str | None = None ) -> Callable[P, R]: """Parametrize a dynamics function with the default dynamics parameters for a drone. @@ -70,7 +70,7 @@ def parametrize( from crazyflow.dynamics.core import parametrize from crazyflow.dynamics.first_principles import dynamics - dynamics_fn = parametrize(dynamics, drone="cf2x_L250") + dynamics_fn = parametrize(dynamics, Drone.cf2x_L250) pos, quat = np.zeros(3), np.array([0.0, 0.0, 0.0, 1.0]) vel, ang_vel = np.zeros(3), np.zeros(3) rotor_vel, cmd = np.zeros(4), np.zeros(4) @@ -86,7 +86,7 @@ def parametrize( def load_params( - dynamics: Dynamics | str, drone: str, xp: ModuleType | None = None, device: str | None = None + dynamics: Dynamics, drone: Drone, xp: ModuleType | None = None, device: str | None = None ) -> dict: """Load all parameters of a drone for a dynamics model. @@ -94,8 +94,8 @@ def load_params( ``crazyflow/dynamics//params.toml`` and adds ``J_inv``. Args: - dynamics: The dynamics model, e.g. ``Dynamics.so_rpy`` or ``"so_rpy"``. - drone: Name of the drone configuration, e.g. ``"cf2x_L250"``. + dynamics: The dynamics model, e.g. ``Dynamics.so_rpy``. + drone: The drone configuration, e.g. ``Drone.cf2x_L250``. xp: Array API module used to convert parameter values. If ``None``, NumPy is used. device: The device to use for the arrays. If ``None``, the device is inferred from the xp module. @@ -104,16 +104,16 @@ def load_params( A flat dict mapping parameter names to arrays in the requested array namespace. Raises: - ValueError: If ``dynamics`` is not a known model. + ValueError: If ``dynamics`` or ``drone`` is unknown. KeyError: If ``drone`` has no section for ``dynamics``. """ dynamics = Dynamics(dynamics) + if dynamics not in supported_dynamics(drone): + raise KeyError(f"Drone `{drone}` not found in {dynamics}/params.toml") with open(Path(__file__).parent / "params.toml", "rb") as f: global_params = tomllib.load(f) with open(Path(__file__).parent / f"{dynamics}/params.toml", "rb") as f: dynamics_params = tomllib.load(f) - if drone not in dynamics_params: - raise KeyError(f"Drone `{drone}` not found in {dynamics}/params.toml") params = global_params | dynamics_params[drone] # Make sure J_inv does not have a dtype fixed before conversion to xp arrays to avoid fixing it # to np.float64 when other frameworks might prefer a different dtype. @@ -122,7 +122,7 @@ def load_params( def load_fn_params( - fn: Callable, drone: str, xp: ModuleType | None = None, device: str | None = None + fn: Callable, drone: Drone, xp: ModuleType | None = None, device: str | None = None ) -> dict: """Load the parameters a dynamics function accepts. @@ -131,7 +131,7 @@ def load_fn_params( Args: fn: The dynamics function for which to load parameters. - drone: Name of the drone configuration, e.g. ``"cf2x_L250"``. + drone: The drone configuration, e.g. ``Drone.cf2x_L250``. xp: Array API module used to convert parameter values. If ``None``, NumPy is used. device: The device to use for the arrays. If ``None``, the device is inferred from the xp module. @@ -144,20 +144,19 @@ def load_fn_params( return filter_to_signature(load_params(dynamics, drone, xp=xp, device=device), fn) -def _param_sections(dynamics: Dynamics | str) -> set[str]: +def _param_sections(dynamics: Dynamics) -> set[str]: """Return the drone sections declared in a dynamics model's ``params.toml``.""" - with open(Path(__file__).parent / f"{Dynamics(dynamics)}/params.toml", "rb") as f: + with open(Path(__file__).parent / f"{dynamics}/params.toml", "rb") as f: return set(tomllib.load(f)) -def supported_drones(dynamics: Dynamics | str) -> tuple[str, ...]: +def supported_drones(dynamics: Dynamics) -> tuple[Drone, ...]: """Return the drones that ``dynamics`` can be parametrized for.""" - sections = _param_sections(dynamics) - return tuple(drone for drone in available_drones if drone in sections) + dynamics = Dynamics(dynamics) + return tuple(drone for drone in Drone if drone in _param_sections(dynamics)) -def supported_dynamics(drone: str) -> tuple[Dynamics, ...]: +def supported_dynamics(drone: Drone) -> tuple[Dynamics, ...]: """Return the dynamics models that ``drone`` can be simulated with.""" - if drone not in available_drones: - raise KeyError(f"Drone `{drone}` not found. Available drones: {available_drones}") + drone = Drone(drone) return tuple(d for d in Dynamics if drone in _param_sections(d)) diff --git a/crazyflow/envs/drone_env.py b/crazyflow/envs/drone_env.py index bf5754cf..e6a112c4 100644 --- a/crazyflow/envs/drone_env.py +++ b/crazyflow/envs/drone_env.py @@ -12,6 +12,7 @@ from numpy.typing import NDArray from crazyflow.control import Control +from crazyflow.drones import Drone from crazyflow.dynamics import Dynamics, load_params from crazyflow.sim import Sim from crazyflow.sim.data import SimData @@ -19,7 +20,7 @@ from crazyflow.utils import leaf_replace -def action_space(control_type: Control, dynamics: Dynamics, drone: str) -> spaces.Box: +def action_space(control_type: Control, dynamics: Dynamics, drone: Drone) -> spaces.Box: """Select the appropriate action space for a given control type. Args: @@ -64,7 +65,7 @@ def __init__( num_envs: int = 1, max_episode_time: float = 10.0, dynamics: Dynamics = Dynamics.so_rpy, - drone: str = "cf2x_L250", + drone: Drone = Drone.cf2x_L250, freq: int = 500, device: str = "cpu", reset_randomization: Callable[[SimData, Array], SimData] | None = None, diff --git a/crazyflow/sim/sim.py b/crazyflow/sim/sim.py index 35103c61..955c6f61 100644 --- a/crazyflow/sim/sim.py +++ b/crazyflow/sim/sim.py @@ -25,6 +25,7 @@ control_state2attitude, ) from crazyflow.control.transform import motor_force2rotor_vel +from crazyflow.drones import Drone from crazyflow.dynamics import Dynamics from crazyflow.dynamics import load_params as load_dynamics_params from crazyflow.dynamics.first_principles import sim_dynamics as first_principles_dynamics @@ -76,7 +77,7 @@ def __init__( self, n_worlds: int = 1, n_drones: int = 1, - drone: str = "cf21B_500", + drone: Drone = Drone.cf21B_500, dynamics: Dynamics = Dynamics.default, control: Control = Control.default, integrator: Integrator = Integrator.default, @@ -708,7 +709,7 @@ def clip_floor_pos(data: SimData) -> SimData: return data.replace(states=data.states.replace(pos=clip_pos, vel=clip_vel)) -def rotor_vel_limits(dynamics: Dynamics, drone: str) -> tuple[float, float]: +def rotor_vel_limits(dynamics: Dynamics, drone: Drone) -> tuple[float, float]: """Limits of ``rotor_vel`` in RPM (first principles) or collective thrust in N (others).""" params = load_dynamics_params(dynamics, drone) thrust_min, thrust_max = float(params["thrust_min"]), float(params["thrust_max"]) diff --git a/docs/index.md b/docs/index.md index 522ed396..56186c17 100644 --- a/docs/index.md +++ b/docs/index.md @@ -86,7 +86,7 @@ Crazyflow is a research simulator for Crazyflie-style quadrotors that runs milli ## Supported drones -All drone configurations are bundled with `crazyflow.dynamics`. Available configurations: `cf2x_L250`, `cf2x_P250`, `cf2x_T350`, `cf21B_500`, and any drone returned by `crazyflow.available_drones`. +All drone configurations are bundled with `crazyflow.dynamics`. Available configurations: `cf2x_L250`, `cf2x_P250`, `cf2x_T350`, `cf21B_500`, and any member of `crazyflow.Drone`. --- diff --git a/docs/user-guide/adding-drones.md b/docs/user-guide/adding-drones.md index b688404b..96ade170 100644 --- a/docs/user-guide/adding-drones.md +++ b/docs/user-guide/adding-drones.md @@ -1,20 +1,20 @@ # Adding a drone -A drone is defined by data files only. `available_drones` lists the MJCF files in `crazyflow/drones`, and a dynamics model or controller supports a drone when its own `params.toml` has a section for it. Adding a platform therefore means adding files and sections, not registering anything in Python. +A drone is defined by data files only. The `Drone` enum lists the MJCF files in `crazyflow/drones`, and a dynamics model or controller supports a drone when its own `params.toml` has a section for it. Adding a platform therefore means adding files and sections, not registering anything in Python. ```python -from crazyflow import available_drones +from crazyflow import Drone from crazyflow.dynamics import supported_dynamics -available_drones # ('cf21B_500', 'cf2x_L250', 'cf2x_P250', 'cf2x_T350') -supported_dynamics("cf2x_L250") # (first_principles, so_rpy, so_rpy_rotor, so_rpy_rotor_drag) +Drone.cf2x_L250 # 'cf2x_L250' +supported_dynamics(Drone.cf2x_L250) # (first_principles, so_rpy, so_rpy_rotor, so_rpy_rotor_drag) ``` Pick a short name such as `cf2x_L250` (platform, then variant) and use it everywhere below. ## 1. MuJoCo model -Add `crazyflow/drones/.xml` with its meshes under `crazyflow/drones/assets//`. This file is what makes the drone appear in `available_drones`. The simulator attaches the body named `drone` once per drone, so that body is required. If you also provide a `drone_fused` body whose visual geometry is a single mesh, users can select it with `Sim(fused_mjx_model=True)` for cheaper rendering. See [MuJoCo Integration](mujoco.md) for how the scene is assembled. +Add `crazyflow/drones/.xml` with its meshes under `crazyflow/drones/assets//`. This file is what makes the drone appear in `Drone`. The simulator attaches the body named `drone` once per drone, so that body is required. If you also provide a `drone_fused` body whose visual geometry is a single mesh, users can select it with `Sim(fused_mjx_model=True)` for cheaper rendering. See [MuJoCo Integration](mujoco.md) for how the scene is assembled. ## 2. Dynamics parameters @@ -36,7 +36,7 @@ Add the platform to the table in [Parametrization](dynamics/parametrize.md#avail ## 5. Run the tests -The test suite parametrizes over `available_drones` and over the supported drone-dynamics pairs, so the new drone is tested without any changes to the tests. In particular, `tests/integration/test_models.py` constructs a `Sim` for every supported pair, which loads the MJCF, the dynamics parameters and the controller parameters together. +The test suite parametrizes over `Drone` and over the supported drone-dynamics pairs, so the new drone is tested without any changes to the tests. In particular, `tests/integration/test_models.py` constructs a `Sim` for every supported pair, which loads the MJCF, the dynamics parameters and the controller parameters together. ```bash pixi run -e tests tests diff --git a/docs/user-guide/dynamics/parametrize.md b/docs/user-guide/dynamics/parametrize.md index 72a8e5cb..1a92cb9b 100644 --- a/docs/user-guide/dynamics/parametrize.md +++ b/docs/user-guide/dynamics/parametrize.md @@ -19,9 +19,9 @@ list(dynamics.keywords.keys()) The following configurations ship with pre-fitted parameters. They cover both the brushed Crazyflie 2.x series and the brushless Crazyflie 2.1: ```python -from crazyflow.drones import available_drones +from crazyflow.drones import Drone -available_drones # ('cf21B_500', 'cf2x_L250', 'cf2x_P250', 'cf2x_T350') +list(Drone) # [Drone.cf21B_500, Drone.cf2x_L250, Drone.cf2x_P250, Drone.cf2x_T350] ``` | `drone` | Platform | diff --git a/docs/user-guide/oo-api.md b/docs/user-guide/oo-api.md index 648e5846..be85bc2d 100644 --- a/docs/user-guide/oo-api.md +++ b/docs/user-guide/oo-api.md @@ -35,7 +35,7 @@ Key constructor arguments: |---|---| | `n_worlds` | Number of independent parallel environments | | `n_drones` | Drones per world | -| `drone` | Drone configuration (see `crazyflow.available_drones`) | +| `drone` | Drone configuration (see `crazyflow.Drone`) | | `dynamics` | Dynamics | | `control` | Control mode | | `integrator` | Numerical integrator | diff --git a/tests/conftest.py b/tests/conftest.py index 854ff187..88d71ed1 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -59,11 +59,11 @@ def drone_dynamics_fns() -> list[ParameterSet]: def drone_dynamics() -> list[ParameterSet]: """Return all supported (dynamics, drone) combinations.""" - from crazyflow.drones import available_drones + from crazyflow.drones import Drone from crazyflow.dynamics import supported_dynamics return [ pytest.param(dynamics, drone, id=f"{drone}-{dynamics}") - for drone in available_drones + for drone in Drone for dynamics in supported_dynamics(drone) ] diff --git a/tests/unit/control/test_core.py b/tests/unit/control/test_core.py index 180767cd..40e72121 100644 --- a/tests/unit/control/test_core.py +++ b/tests/unit/control/test_core.py @@ -13,7 +13,7 @@ force_torque2rotor_vel, state2attitude, ) -from crazyflow.drones import available_drones +from crazyflow.drones import Drone _MELLINGER_FNS = [ state2attitude, @@ -25,7 +25,7 @@ @pytest.mark.unit @pytest.mark.parametrize("fn", _MELLINGER_FNS, ids=lambda fn: fn.__name__) -@pytest.mark.parametrize("drone", available_drones) +@pytest.mark.parametrize("drone", Drone) def test_load_fn_params_keys(fn: Callable[..., Any], drone: str) -> None: params = load_fn_params(fn, drone) fn_params = inspect.signature(fn).parameters @@ -50,7 +50,7 @@ def test_unknown_controller() -> None: @pytest.mark.unit -@pytest.mark.parametrize("drone", available_drones) +@pytest.mark.parametrize("drone", Drone) def test_parametrize_xp_namespace(drone: str) -> None: controller = parametrize(state2attitude, drone, xp=array_api_strict) xp_array_type = type(array_api_strict.asarray(0.0)) diff --git a/tests/unit/control/test_mellinger.py b/tests/unit/control/test_mellinger.py index 2c8833e0..909a8066 100644 --- a/tests/unit/control/test_mellinger.py +++ b/tests/unit/control/test_mellinger.py @@ -13,7 +13,7 @@ force_torque2rotor_vel, state2attitude, ) -from crazyflow.drones import available_drones +from crazyflow.drones import Drone if TYPE_CHECKING: from crazyflow._typing import Array # To be changed to array_api_typing later @@ -25,7 +25,7 @@ def create_rnd_states(shape: tuple[int, ...] = ()) -> tuple[Array, Array, Array, @pytest.mark.unit -@pytest.mark.parametrize("drone", available_drones) +@pytest.mark.parametrize("drone", Drone) def test_state2attitude(drone: str) -> None: controller = parametrize(state2attitude, drone) # Single input @@ -41,7 +41,7 @@ def test_state2attitude(drone: str) -> None: @pytest.mark.unit -@pytest.mark.parametrize("drone", available_drones) +@pytest.mark.parametrize("drone", Drone) def test_attitude2force_torque(drone: str) -> None: controller = parametrize(attitude2force_torque, drone) # Single input @@ -62,7 +62,7 @@ def test_attitude2force_torque(drone: str) -> None: @pytest.mark.unit -@pytest.mark.parametrize("drone", available_drones) +@pytest.mark.parametrize("drone", Drone) def test_body_rate2force_torque(drone: str) -> None: controller = parametrize(body_rate2force_torque, drone) # Single input @@ -83,7 +83,7 @@ def test_body_rate2force_torque(drone: str) -> None: @pytest.mark.unit -@pytest.mark.parametrize("drone", available_drones) +@pytest.mark.parametrize("drone", Drone) def test_force_torque2rotor_vel(drone: str) -> None: controller = parametrize(force_torque2rotor_vel, drone) # Single input @@ -102,7 +102,7 @@ def test_force_torque2rotor_vel(drone: str) -> None: @pytest.mark.unit -@pytest.mark.parametrize("drone", available_drones) +@pytest.mark.parametrize("drone", Drone) def test_state2attitude_at_setpoint(drone: str) -> None: # At setpoint with identity orientation and zero acc, RPY command should be # [0, 0, 0] and thrust must be positive (hovering against gravity). @@ -118,7 +118,7 @@ def test_state2attitude_at_setpoint(drone: str) -> None: @pytest.mark.unit -@pytest.mark.parametrize("drone", available_drones) +@pytest.mark.parametrize("drone", Drone) def test_state2attitude_integral_error_accumulation(drone: str) -> None: # A constant position error must cause the integral error to accumulate # linearly until it would exceed int_err_max (clipped by the controller). @@ -145,7 +145,7 @@ def test_state2attitude_integral_error_accumulation(drone: str) -> None: @pytest.mark.unit -@pytest.mark.parametrize("drone", available_drones) +@pytest.mark.parametrize("drone", Drone) def test_attitude2force_torque_at_setpoint(drone: str) -> None: # Identity orientation commanded → zero attitude error → zero corrective torque. controller = parametrize(attitude2force_torque, drone) @@ -160,7 +160,7 @@ def test_attitude2force_torque_at_setpoint(drone: str) -> None: @pytest.mark.unit -@pytest.mark.parametrize("drone", available_drones) +@pytest.mark.parametrize("drone", Drone) def test_attitude2force_torque_zero_thrust(drone: str): # Zero thrust command → firmware zeros torque; outputs are all zero. controller = parametrize(attitude2force_torque, drone) @@ -173,7 +173,7 @@ def test_attitude2force_torque_zero_thrust(drone: str): @pytest.mark.unit -@pytest.mark.parametrize("drone", available_drones) +@pytest.mark.parametrize("drone", Drone) def test_body_rate2force_torque_at_setpoint(drone: str) -> None: # Level drone with measured rates equal to the commanded rates → zero corrective torque. controller = parametrize(body_rate2force_torque, drone) @@ -188,7 +188,7 @@ def test_body_rate2force_torque_at_setpoint(drone: str) -> None: @pytest.mark.unit -@pytest.mark.parametrize("drone", available_drones) +@pytest.mark.parametrize("drone", Drone) def test_body_rate2force_torque_zero_thrust(drone: str): # Zero thrust command → firmware zeros torque; outputs are all zero. controller = parametrize(body_rate2force_torque, drone) @@ -201,7 +201,7 @@ def test_body_rate2force_torque_zero_thrust(drone: str): @pytest.mark.unit -@pytest.mark.parametrize("drone", available_drones) +@pytest.mark.parametrize("drone", Drone) def test_body_rate2force_torque_sign(drone: str): # A positive rate error about one axis must produce a positive torque about that axis only. controller = parametrize(body_rate2force_torque, drone) @@ -217,7 +217,7 @@ def test_body_rate2force_torque_sign(drone: str): @pytest.mark.unit -@pytest.mark.parametrize("drone", available_drones) +@pytest.mark.parametrize("drone", Drone) def test_body_rate2force_torque_matches_attitude(drone: str): # A zero body rate command is equivalent to commanding a level attitude at the current yaw. att_controller = parametrize(attitude2force_torque, drone) @@ -239,7 +239,7 @@ def test_body_rate2force_torque_matches_attitude(drone: str): @pytest.mark.unit -@pytest.mark.parametrize("drone", available_drones) +@pytest.mark.parametrize("drone", Drone) def test_body_rate2force_torque_leveling(drone: str): # The firmware levels a tilted drone even at the rate setpoint. Zero attitude gains disable it. controller = parametrize(body_rate2force_torque, drone) @@ -258,7 +258,7 @@ def test_body_rate2force_torque_leveling(drone: str): @pytest.mark.unit -@pytest.mark.parametrize("drone", available_drones) +@pytest.mark.parametrize("drone", Drone) def test_state2attitude_batch_consistency(drone: str): controller = parametrize(state2attitude, drone) batch = (3, 2) @@ -273,7 +273,7 @@ def test_state2attitude_batch_consistency(drone: str): @pytest.mark.unit -@pytest.mark.parametrize("drone", available_drones) +@pytest.mark.parametrize("drone", Drone) def test_attitude2force_torque_batch_consistency(drone: str): controller = parametrize(attitude2force_torque, drone) batch = (3, 2) @@ -290,7 +290,7 @@ def test_attitude2force_torque_batch_consistency(drone: str): @pytest.mark.unit -@pytest.mark.parametrize("drone", available_drones) +@pytest.mark.parametrize("drone", Drone) def test_body_rate2force_torque_batch_consistency(drone: str): controller = parametrize(body_rate2force_torque, drone) batch = (3, 2) @@ -317,7 +317,7 @@ def test_body_rate2force_torque_batch_consistency(drone: str): @pytest.mark.unit -@pytest.mark.parametrize("drone", available_drones) +@pytest.mark.parametrize("drone", Drone) def test_force_torque2rotor_vel_batch_consistency(drone: str): controller = parametrize(force_torque2rotor_vel, drone) batch = (3, 2) @@ -331,7 +331,7 @@ def test_force_torque2rotor_vel_batch_consistency(drone: str): @pytest.mark.unit -@pytest.mark.parametrize("drone", available_drones) +@pytest.mark.parametrize("drone", Drone) def test_attitude2force_torque_batch_zero_thrust(drone: str): # Drones with zero thrust must stay at zero force, independent of other drones controller = parametrize(attitude2force_torque, drone) @@ -347,7 +347,7 @@ def test_attitude2force_torque_batch_zero_thrust(drone: str): @pytest.mark.unit -@pytest.mark.parametrize("drone", available_drones) +@pytest.mark.parametrize("drone", Drone) def test_force_torque2rotor_vel_batch_zero_force(drone: str): # Drones with zero desired force must not be clipped because other drones have non-zero force. controller = parametrize(force_torque2rotor_vel, drone) @@ -363,7 +363,7 @@ def test_force_torque2rotor_vel_batch_zero_force(drone: str): @pytest.mark.unit -@pytest.mark.parametrize("drone", available_drones) +@pytest.mark.parametrize("drone", Drone) def test_force_torque2rotor_vel_symmetric(drone: str): # Pure vertical force with zero torque → X-frame symmetry → all 4 RPMs equal. controller = parametrize(force_torque2rotor_vel, drone) diff --git a/tests/unit/dynamics/test_parametrization.py b/tests/unit/dynamics/test_parametrization.py index 7de74f04..c6b8f331 100644 --- a/tests/unit/dynamics/test_parametrization.py +++ b/tests/unit/dynamics/test_parametrization.py @@ -7,7 +7,15 @@ import pytest from conftest import drone_dynamics_fns -from crazyflow.dynamics import Dynamics, load_fn_params, load_params, parametrize +from crazyflow.drones import Drone +from crazyflow.dynamics import ( + Dynamics, + load_fn_params, + load_params, + parametrize, + supported_drones, + supported_dynamics, +) from crazyflow.dynamics.so_rpy import dynamics as so_rpy @@ -28,18 +36,32 @@ def test_model_parameter_loading(dynamics_name: str, dynamics: Callable, drone: @pytest.mark.unit def test_unknown_drone() -> None: - with pytest.raises(KeyError, match="nonexistent_drone"): + with pytest.raises(ValueError, match="nonexistent_drone"): load_params(Dynamics.so_rpy, "nonexistent_drone") - with pytest.raises(KeyError, match="nonexistent_drone"): + with pytest.raises(ValueError, match="nonexistent_drone"): load_fn_params(so_rpy, "nonexistent_drone") - with pytest.raises(KeyError, match="nonexistent_drone"): + with pytest.raises(ValueError, match="nonexistent_drone"): parametrize(so_rpy, "nonexistent_drone") + with pytest.raises(ValueError, match="nonexistent_drone"): + supported_dynamics("nonexistent_drone") @pytest.mark.unit def test_unknown_dynamics() -> None: with pytest.raises(ValueError, match="nonexistent_dynamics"): load_params("nonexistent_dynamics", "cf2x_L250") + with pytest.raises(ValueError, match="nonexistent_dynamics"): + supported_drones("nonexistent_dynamics") + + +@pytest.mark.unit +def test_supported_pairs() -> None: + for dynamics in Dynamics: + for drone in supported_drones(dynamics): + assert dynamics in supported_dynamics(drone) + for drone in Drone: + for dynamics in supported_dynamics(drone): + assert drone in supported_drones(dynamics) @pytest.mark.unit From 0bd3182684a9e48f9c16d121e69e66fca0654cf3 Mon Sep 17 00:00:00 2001 From: ratheron Date: Sat, 19 Sep 2026 16:03:38 +0200 Subject: [PATCH 3/5] Hard code Drone enum to enable type hinting --- SKILL.md | 3 ++- crazyflow/drones/__init__.py | 18 +++++++++++++++--- docs/user-guide/adding-drones.md | 6 +++--- 3 files changed, 20 insertions(+), 7 deletions(-) diff --git a/SKILL.md b/SKILL.md index 2e853d1a..82e0b459 100644 --- a/SKILL.md +++ b/SKILL.md @@ -49,7 +49,8 @@ against all models in the simulation's `build_control_fns`. Define the function in `dynamics.py` and never in the package `__init__.py`, because `load_fn_params` derives the model name from `fn.__module__.split(".")[-2]`. `parametrize` binds exactly the keyword-only parameters after the bare `*`, so anything before it is never bound. -`Drone` is a `StrEnum` of the MJCF files in `crazyflow/drones`. A model supports a drone when its +`Drone` in `crazyflow/drones/__init__.py` is a `StrEnum` with one member per MJCF file in +`crazyflow/drones`, asserted on import. A model supports a drone when its `crazyflow/dynamics//params.toml` has a complete section for it; the commented example at the top of each file lists the keys. `supported_drones` and `supported_dynamics` report the pairs, and the tests only run those. Only `gravity_vec` is global, in `crazyflow/dynamics/params.toml`. diff --git a/crazyflow/drones/__init__.py b/crazyflow/drones/__init__.py index 3d816778..935fcef1 100644 --- a/crazyflow/drones/__init__.py +++ b/crazyflow/drones/__init__.py @@ -11,6 +11,18 @@ __all__ = ["Drone"] -_drones = [p.stem for p in sorted(Path(__file__).parent.glob("*.xml"))] -Drone: StrEnum = StrEnum("Drone", [(name, name) for name in _drones]) -"""Drone configurations, i.e. the MJCF files in ``crazyflow/drones``.""" + +class Drone(StrEnum): + """Drone configurations. Each member has an MJCF file ``crazyflow/drones/.xml``.""" + + cf21B_500 = "cf21B_500" + cf2x_L250 = "cf2x_L250" + cf2x_P250 = "cf2x_P250" + cf2x_T350 = "cf2x_T350" + + +# Sanity check at startup +_mjcf_files = {p.stem for p in Path(__file__).parent.glob("*.xml")} +assert {d.value for d in Drone} == _mjcf_files, ( + f"Drone enum {sorted(d.value for d in Drone)} does not match MJCF files {sorted(_mjcf_files)}" +) diff --git a/docs/user-guide/adding-drones.md b/docs/user-guide/adding-drones.md index 96ade170..aaf6f893 100644 --- a/docs/user-guide/adding-drones.md +++ b/docs/user-guide/adding-drones.md @@ -1,6 +1,6 @@ # Adding a drone -A drone is defined by data files only. The `Drone` enum lists the MJCF files in `crazyflow/drones`, and a dynamics model or controller supports a drone when its own `params.toml` has a section for it. Adding a platform therefore means adding files and sections, not registering anything in Python. +A drone is a member of the `Drone` enum in `crazyflow/drones/__init__.py` with a matching MJCF file in `crazyflow/drones`. The package asserts on import that the two agree. A dynamics model or controller supports a drone when its own `params.toml` has a section for it, so adding a platform means adding one enum member, the MJCF file, and parameter sections. ```python from crazyflow import Drone @@ -12,9 +12,9 @@ supported_dynamics(Drone.cf2x_L250) # (first_principles, so_rpy, so_rpy_rotor, Pick a short name such as `cf2x_L250` (platform, then variant) and use it everywhere below. -## 1. MuJoCo model +## 1. Enum member and MuJoCo model -Add `crazyflow/drones/.xml` with its meshes under `crazyflow/drones/assets//`. This file is what makes the drone appear in `Drone`. The simulator attaches the body named `drone` once per drone, so that body is required. If you also provide a `drone_fused` body whose visual geometry is a single mesh, users can select it with `Sim(fused_mjx_model=True)` for cheaper rendering. See [MuJoCo Integration](mujoco.md) for how the scene is assembled. +Add ` = ""` to the `Drone` enum and a `crazyflow/drones/.xml` with its meshes under `crazyflow/drones/assets//`. Importing `crazyflow` fails if an enum member has no MJCF file or an MJCF file has no enum member. The simulator attaches the body named `drone` once per drone, so that body is required. If you also provide a `drone_fused` body whose visual geometry is a single mesh, users can select it with `Sim(fused_mjx_model=True)` for cheaper rendering. See [MuJoCo Integration](mujoco.md) for how the scene is assembled. ## 2. Dynamics parameters From 41e1ea1b8e03ed755b15eb05ec81e6f9945f9c6b Mon Sep 17 00:00:00 2001 From: ratheron Date: Sat, 19 Sep 2026 16:12:42 +0200 Subject: [PATCH 4/5] Fix docstrings --- crazyflow/dynamics/core.py | 29 +++++++++++++++++-- .../dynamics/first_principles/params.toml | 1 - crazyflow/dynamics/so_rpy/params.toml | 3 +- crazyflow/dynamics/so_rpy_rotor/params.toml | 3 +- .../dynamics/so_rpy_rotor_drag/params.toml | 3 +- 5 files changed, 30 insertions(+), 9 deletions(-) diff --git a/crazyflow/dynamics/core.py b/crazyflow/dynamics/core.py index f2dbe18a..11221297 100644 --- a/crazyflow/dynamics/core.py +++ b/crazyflow/dynamics/core.py @@ -151,12 +151,37 @@ def _param_sections(dynamics: Dynamics) -> set[str]: def supported_drones(dynamics: Dynamics) -> tuple[Drone, ...]: - """Return the drones that ``dynamics`` can be parametrized for.""" + """Return the drones that ``dynamics`` can be parametrized for. + + A drone is supported when ``crazyflow/dynamics//params.toml`` has a section for it. + + Args: + dynamics: The dynamics model, e.g. ``Dynamics.so_rpy``. + + Returns: + The supported drones in the order of [Drone][crazyflow.drones.Drone]. + + Raises: + ValueError: If ``dynamics`` is not a known model. + """ dynamics = Dynamics(dynamics) return tuple(drone for drone in Drone if drone in _param_sections(dynamics)) def supported_dynamics(drone: Drone) -> tuple[Dynamics, ...]: - """Return the dynamics models that ``drone`` can be simulated with.""" + """Return the dynamics models that ``drone`` can be simulated with. + + A model is supported when its ``crazyflow/dynamics//params.toml`` has a section for + ``drone``. + + Args: + drone: The drone configuration, e.g. ``Drone.cf2x_L250``. + + Returns: + The supported models in the order of [Dynamics][crazyflow.dynamics.Dynamics]. + + Raises: + ValueError: If ``drone`` is not a known drone. + """ drone = Drone(drone) return tuple(d for d in Dynamics if drone in _param_sections(d)) diff --git a/crazyflow/dynamics/first_principles/params.toml b/crazyflow/dynamics/first_principles/params.toml index 3770db3c..d5e2e415 100644 --- a/crazyflow/dynamics/first_principles/params.toml +++ b/crazyflow/dynamics/first_principles/params.toml @@ -30,7 +30,6 @@ # # Not every parameter has to be identified, but all of them have to be set. A symmetric first order # rotor model, no drag and no propeller gyroscopic torque are reasonable defaults for a new drone. -# gravity_vec is global and lives in crazyflow/dynamics/params.toml. [cf2x_L250] mass = 0.0319 diff --git a/crazyflow/dynamics/so_rpy/params.toml b/crazyflow/dynamics/so_rpy/params.toml index 1dee51b7..c13ab3ca 100644 --- a/crazyflow/dynamics/so_rpy/params.toml +++ b/crazyflow/dynamics/so_rpy/params.toml @@ -16,8 +16,7 @@ # cmd_rpy_coef = [0.0, 0.0, 0.0] # # Identify the coefficients from flight data with the system identification pipeline (see -# docs/user-guide/dynamics/system-identification.md). gravity_vec is global and lives in -# crazyflow/dynamics/params.toml. +# docs/user-guide/dynamics/system-identification.md). [cf2x_L250] mass = 0.0319 diff --git a/crazyflow/dynamics/so_rpy_rotor/params.toml b/crazyflow/dynamics/so_rpy_rotor/params.toml index 5a4dabad..3d6c3dca 100644 --- a/crazyflow/dynamics/so_rpy_rotor/params.toml +++ b/crazyflow/dynamics/so_rpy_rotor/params.toml @@ -17,8 +17,7 @@ # thrust_time_coef = 0.0 # s # # Identify the coefficients from flight data with the system identification pipeline (see -# docs/user-guide/dynamics/system-identification.md). gravity_vec is global and lives in -# crazyflow/dynamics/params.toml. +# docs/user-guide/dynamics/system-identification.md). [cf2x_L250] mass = 0.0319 diff --git a/crazyflow/dynamics/so_rpy_rotor_drag/params.toml b/crazyflow/dynamics/so_rpy_rotor_drag/params.toml index 23fb8182..0afecc52 100644 --- a/crazyflow/dynamics/so_rpy_rotor_drag/params.toml +++ b/crazyflow/dynamics/so_rpy_rotor_drag/params.toml @@ -22,8 +22,7 @@ # ] # # Identify the coefficients from flight data with the system identification pipeline (see -# docs/user-guide/dynamics/system-identification.md). gravity_vec is global and lives in -# crazyflow/dynamics/params.toml. +# docs/user-guide/dynamics/system-identification.md). [cf2x_L250] mass = 0.0319 From 695fd9f9d32c8f2f8d2e4e5e7f3c417e2d6e5ca8 Mon Sep 17 00:00:00 2001 From: ratheron Date: Sat, 19 Sep 2026 16:44:22 +0200 Subject: [PATCH 5/5] Fix tests --- crazyflow/control/core.py | 1 + crazyflow/dynamics/core.py | 1 + 2 files changed, 2 insertions(+) diff --git a/crazyflow/control/core.py b/crazyflow/control/core.py index 4c5f9e30..7d07630b 100644 --- a/crazyflow/control/core.py +++ b/crazyflow/control/core.py @@ -79,6 +79,7 @@ def parametrize( import numpy as np from crazyflow.control import parametrize from crazyflow.control.mellinger import state2attitude + from crazyflow.drones import Drone from scipy.spatial.transform import Rotation as R ctrl = parametrize(state2attitude, Drone.cf2x_L250) diff --git a/crazyflow/dynamics/core.py b/crazyflow/dynamics/core.py index 11221297..161a50ce 100644 --- a/crazyflow/dynamics/core.py +++ b/crazyflow/dynamics/core.py @@ -67,6 +67,7 @@ def parametrize( Example: ```python import numpy as np + from crazyflow.drones import Drone from crazyflow.dynamics.core import parametrize from crazyflow.dynamics.first_principles import dynamics