Move parameters and update load_params

SKILL.md

@@ -16,7 +16,7 @@ Settings are in `pyproject.toml`, layout in the tree, API reference in `docs/`.
   `array_namespace(...)`, do all math through it, coerce bound parameters with `to_xp`, and never
   import `numpy` into a computation path.
 - `crazyflow/control/mellinger/params.toml` deliberately holds values that differ from the true
-  physical constants in `crazyflow/drones/params.toml`, reproducing the real firmware. Do not
+  physical constants in `crazyflow/dynamics/*/params.toml`, reproducing the real firmware. Do not
   unify them.
 
 ## Testing
@@ -36,8 +36,8 @@ Use pixi, not uv. `pixi run <cmd>` resolves task names only, so arbitrary comman
   checked goes in a `pycon` fence with `>>>` prompts, which doctest picks up.
 - `tests/integration/test_examples.py` runs every script under `examples/`, so a new example is a
   new test.
-- `tests/conftest.py` forces JAX's persistent cache on at `/tmp/jax_cache`, shared across branches.
-  Delete it when failures make no sense.
+- `tests/conftest.py` forces JAX's persistent cache on at `/tmp/jax_cache-<uid>`, shared across
+  branches. Delete it when failures make no sense.
 - Request the `device` fixture rather than a GPU marker. It falls back to CPU silently, so
   `gpu-tests` asserts nothing about placement on a machine without CUDA.
 
@@ -46,10 +46,12 @@ Use pixi, not uv. `pixi run <cmd>` resolves task names only, so arbitrary comman
 Grepping the name of an existing model or drone finds every registration site, except when matching
 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_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 section in every `crazyflow/dynamics/*/params.toml`, even an empty one.
+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`.
 
 Registration alone produces roughly 40 parametrized tests. These do not include derivatives tests.
 
@@ -60,9 +62,9 @@ Pure, batched, array-API functions with no dependency on `Sim`.
 - Import crazyflow before scipy. `crazyflow/__init__.py` sets `SCIPY_ARRAY_API=1` and imports scipy
   immediately, and scipy cannot be reconfigured once loaded. Transitive imports through acados or
   sklearn trigger this too.
-- Three different `load_params` exist. The two in `.core` filter to the target signature and
-  silently drop the rest, so hardware constants like `thrust_max` need
-  `crazyflow.drones.load_params`.
+- `dynamics` and `control` each have `load_params(name, drone)`, returning everything defined for
+  the drone, and `load_fn_params(fn, drone)`, which filters to the signature of `fn` and
+  drops the rest. Additional platform data is stored as a comment in the drone MJCF.
 - `parametrize` returns a `functools.partial` whose `keywords` dict is shared by every reference to
   it. Call `parametrize` again for an independent copy.
 - Leading batch dimensions, trailing feature axis. `quat` is scalar-last xyzw, `force` is `(..., 1)`

crazyflow/control/__init__.py

@@ -14,7 +14,7 @@
 
 __all__ = []
 
-from crazyflow.control.core import Control, load_params, parametrize
+from crazyflow.control.core import Control, load_fn_params, load_params, parametrize
 from crazyflow.control.mellinger import attitude2force_torque as mellinger_attitude2force_torque
 from crazyflow.control.mellinger import body_rate2force_torque as mellinger_body_rate2force_torque
 from crazyflow.control.mellinger import state2attitude as mellinger_state2attitude
@@ -25,4 +25,4 @@
     "mellinger_body_rate2force_torque": mellinger_body_rate2force_torque,
 }
 
-__all__ = ["Control", "load_params", "parametrize"]
+__all__ = ["Control", "load_params", "load_fn_params", "parametrize"]

crazyflow/control/core.py

@@ -22,6 +22,46 @@
 R = TypeVar("R")
 
 
+class Control(StrEnum):
+    """Control type of the simulated onboard controller."""
+
+    state = "state"
+    """State control takes [x, y, z, vx, vy, vz, ax, ay, az, qx, qy, qz, qw, wx, wy, wz].
+
+    Note:
+        Recommended frequency is >=20 Hz.
+
+    Warning:
+        Only the yaw of the attitude quaternion is used, as in the firmware. The so_rpy family
+        ignores the body rate setpoint.
+    """
+    attitude = "attitude"
+    """Attitude control takes [roll, pitch, yaw, collective thrust].
+
+    Note:
+        Recommended frequency is >=100 Hz.
+    """
+    body_rate = "body_rate"
+    """Body rate control takes [wx, wy, wz, collective thrust].
+
+    Note:
+        Recommended frequency is >=200 Hz.
+    """
+    force_torque = "force_torque"
+    """Force and torque control takes [fc, tx, ty, tz].
+
+    Note:
+        Recommended frequency is >=500 Hz.
+    """
+    rotor_vel = "rotor_vel"
+    """Rotor velocity control takes [w1, w2, w3, w4] in RPMs.
+
+    Note:
+        Recommended frequency is >=500 Hz.
+    """
+    default = attitude
+
+
 def parametrize(
     fn: Callable[P, R], drone: str, xp: ModuleType | None = None, device: str | None = None
 ) -> Callable[P, R]:
@@ -50,77 +90,57 @@ def parametrize(
     Returns:
         The parametrized controller function with all keyword argument only parameters filled in.
     """
-    return _parametrize(fn, drone, load_params, xp=xp, device=device)
+    return _parametrize(fn, drone, load_fn_params, xp=xp, device=device)
 
 
 def load_params(
-    fn: Callable, drone: str, xp: ModuleType | None = None, device: str | None = None
-) -> dict[str, Array]:
-    """Load the parameters a specific controller function accepts.
+    controller: str, drone: str, xp: ModuleType | None = None, device: str | None = None
+) -> dict[str, dict[str, Array]]:
+    """Load all parameters of a drone for a controller.
 
-    Merges the ``"core"`` section with the function's ``[drone.<fn_name>]`` section (function values
-    take precedence), then keeps only the parameters in ``fn``'s signature.
+    Returns the drone's table of ``crazyflow/control/<controller>/params.toml``: the ``core``
+    section shared by all functions of the controller, and one section per controller function.
 
     Args:
-        fn: The controller function for which to load parameters.
+        controller: Name of the controller package, e.g. ``"mellinger"``.
         drone: Name of the drone configuration, e.g. ``"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.
 
-    Returns:
-        A flat dict mapping parameter names to arrays in the requested array namespace.
+    Raises:
+        KeyError: If ``controller`` has no ``params.toml`` or ``drone`` has no section in it.
     """
-    assert isinstance(fn, Callable), f"Expected a function, got {type(fn)}"
-    controller = fn.__module__.split(".")[-2]
     params_path = Path(__file__).parent / f"{controller}/params.toml"
     if not params_path.exists():
-        raise KeyError(f"`{controller}` not found. Available controllers: {tuple(Control)}")
+        raise KeyError(f"Controller `{controller}` not found")
     with open(params_path, "rb") as f:
         params = tomllib.load(f)
     if drone not in params:
         raise KeyError(f"Drone `{drone}` not found in {controller}/params.toml")
-    merged = params[drone].get("core", {}) | params[drone].get(fn.__name__, {})
-    return to_xp(filter_to_signature(merged, fn), xp=xp, device=device)
-
-
-class Control(StrEnum):
-    """Control type of the simulated onboard controller."""
-
-    state = "state"
-    """State control takes [x, y, z, vx, vy, vz, ax, ay, az, qx, qy, qz, qw, wx, wy, wz].
-
-    Note:
-        Recommended frequency is >=20 Hz.
+    return to_xp(params[drone], xp=xp, device=device)
 
-    Warning:
-        Only the yaw of the attitude quaternion is used, as in the firmware. The so_rpy family
-        ignores the body rate setpoint.
-    """
-    attitude = "attitude"
-    """Attitude control takes [roll, pitch, yaw, collective thrust].
 
-    Note:
-        Recommended frequency is >=100 Hz.
-    """
-    body_rate = "body_rate"
-    """Body rate control takes [wx, wy, wz, collective thrust].
+def load_fn_params(
+    fn: Callable, drone: str, xp: ModuleType | None = None, device: str | None = None
+) -> dict[str, Array]:
+    """Load the parameters a controller function accepts.
 
-    Note:
-        Recommended frequency is >=200 Hz.
-    """
-    force_torque = "force_torque"
-    """Force and torque control takes [fc, tx, ty, tz].
+    The controller is derived from the function's package. Merges the ``core`` section with the
+    function's section (function values take precedence), then keeps only the parameters in
+    ``fn``'s signature.
 
-    Note:
-        Recommended frequency is >=500 Hz.
-    """
-    rotor_vel = "rotor_vel"
-    """Rotor velocity control takes [w1, w2, w3, w4] in RPMs.
+    Args:
+        fn: The controller function for which to load parameters.
+        drone: Name of the drone configuration, e.g. ``"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.
 
-    Note:
-        Recommended frequency is >=500 Hz.
+    Returns:
+        A flat dict mapping parameter names to arrays in the requested array namespace.
     """
-    default = attitude
+    assert callable(fn), f"Expected a function, got {type(fn)}"
+    params = load_params(fn.__module__.split(".")[-2], drone, xp=xp, device=device)
+    return filter_to_signature(params.get("core", {}) | params.get(fn.__name__, {}), fn)
 
 
 @jax.jit

crazyflow/control/mellinger/control.py

@@ -20,7 +20,7 @@
 from flax.struct import dataclass, field
 from scipy.spatial.transform import Rotation as R
 
-from crazyflow.control.core import controllable, load_params
+from crazyflow.control.core import controllable, load_fn_params
 from crazyflow.control.transform import force2pwm, motor_force2rotor_vel, pwm2force
 from crazyflow.utils import CORE_NDIM_KEY, leaf_replace
 
@@ -513,7 +513,7 @@ def create(
         zeros_3d = jnp.zeros((n_worlds, n_drones, 3), device=device)
         cmd = jnp.zeros((n_worlds, n_drones, 16), device=device).at[..., 12].set(1.0)
         steps = -jnp.ones((n_worlds, 1), dtype=jnp.int32, device=device)
-        params = load_params(state2attitude, drone, xp=jnp, device=device)
+        params = load_fn_params(state2attitude, drone, xp=jnp, device=device)
         return MellingerStateData(
             cmd=cmd,
             staged_cmd=cmd.copy(),
@@ -553,7 +553,7 @@ def create(
         zeros_3d = jnp.zeros((n_worlds, n_drones, 3), device=device)
         zeros_4d = jnp.zeros((n_worlds, n_drones, 4), device=device)
         steps = -jnp.ones((n_worlds, 1), dtype=jnp.int32, device=device)
-        params = load_params(attitude2force_torque, drone, xp=jnp, device=device)
+        params = load_fn_params(attitude2force_torque, drone, xp=jnp, device=device)
         return MellingerAttitudeData(
             cmd=zeros_4d.copy(),
             staged_cmd=zeros_4d.copy(),
@@ -595,7 +595,7 @@ def create(
         zeros_3d = jnp.zeros((n_worlds, n_drones, 3), device=device)
         zeros_4d = jnp.zeros((n_worlds, n_drones, 4), device=device)
         steps = -jnp.ones((n_worlds, 1), dtype=jnp.int32, device=device)
-        params = load_params(body_rate2force_torque, drone, xp=jnp, device=device)
+        params = load_fn_params(body_rate2force_torque, drone, xp=jnp, device=device)
         return MellingerBodyRateData(
             cmd=zeros_4d.copy(),
             staged_cmd=zeros_4d.copy(),
@@ -629,7 +629,7 @@ def create(
     ) -> MellingerForceTorqueData:
         zeros_4d = jnp.zeros((n_worlds, n_drones, 4), device=device)
         steps = -jnp.ones((n_worlds, 1), dtype=jnp.int32, device=device)
-        params = load_params(force_torque2rotor_vel, drone, xp=jnp, device=device)
+        params = load_fn_params(force_torque2rotor_vel, drone, xp=jnp, device=device)
         return MellingerForceTorqueData(
             cmd=zeros_4d.copy(), staged_cmd=zeros_4d.copy(), steps=steps, freq=freq, params=params
         )

crazyflow/drones/__init__.py

@@ -1,37 +1,16 @@
 """Hardware descriptions for the supported drone platforms.
 
-This package bundles the physical assets that define each drone configuration: the MuJoCo MJCF scene
-files, their referenced meshes (``assets/``), and the physical parameters shared across all dynamics
-(``params.toml`` with mass, inertia, thrust and torque curves, …). These describe the *hardware* and
-are independent of the dynamics formulation used to simulate it (see [crazyflow.dynamics][]).
+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, and ``load_params`` to read all
-physical parameters of a drone.
+Use ``available_drones`` to enumerate the supported configurations.
 """
 
-import tomllib
-from pathlib import Path
-
 # 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")
 
-__all__ = ["available_drones", "load_params"]
-
-
-def load_params(drone: str) -> dict:
-    """Load all physical parameters of a drone from ``params.toml``.
-
-    Returns the raw values (lists/scalars) for the whole drone.
-
-    Args:
-        drone: Name of the drone configuration, e.g. ``"cf2x_L250"``.
-    """
-    with open(Path(__file__).parent / "params.toml", "rb") as f:
-        params = tomllib.load(f)
-    if drone not in params or drone not in available_drones:
-        raise KeyError(f"Drone `{drone}` not found in drones/params.toml")
-    return params[drone]
+__all__ = ["available_drones"]

crazyflow/drones/cf21B_500.xml

@@ -1,4 +1,13 @@
 <mujoco model="cf21B">
+  <!--
+  Platform data not used by the simulation, kept for estimators, firmware and tooling:
+  rotor_dyn_coef_simple = 15.416891997523813
+  thrust_dyn_coef = 15.09965949800411
+  vmotor2thrust = [-0.014058926705279723, 0.04265273261724981, 0.0018327760144017432, 0.0020576974784587178]  # Index is order
+  vmotor2torque = [-0.00016088354909542246, 0.0003960426420309137, -4.6274122414327404e-5, 1.8490661674309596e-5]  # TODO, Index is order
+  vmotor2rpm = [2938.3995608848436, 6001.834195381014]  # Index is order
+  prop_radius = 27.5e-3 # TODO check
+  -->
   <compiler inertiafromgeom="false" meshdir="assets" autolimits="true" />
 
   <asset>

crazyflow/drones/cf2x_L250.xml

@@ -1,4 +1,13 @@
 <mujoco model="cf2x">
+  <!--
+  Platform data not used by the simulation, kept for estimators, firmware and tooling:
+  rotor_dyn_coef_simple = 6.886705423469015
+  thrust_dyn_coef = 6.8932095506763345
+  vmotor2thrust = [-0.014830744918356092, 0.04724465241828281, -0.01847364358025878, 0.005960923942142]  # Index is order
+  vmotor2torque = [-3.3261514624778425e-6, 8.109636684075977e-5, 5.751459172588052e-5, -1.898633582060136e-7]  # Index is order
+  vmotor2rpm = [2968.1791506049194, 6647.948592402306]  # Index is order
+  prop_radius = 23.55e-3 # TODO check
+  -->
   <compiler inertiafromgeom="false" meshdir="assets" autolimits="true" />
 
   <asset>

crazyflow/drones/cf2x_P250.xml

@@ -1,4 +1,13 @@
 <mujoco model="cf2x">
+  <!--
+  Platform data not used by the simulation, kept for estimators, firmware and tooling:
+  rotor_dyn_coef_simple = 7.709730027690284
+  thrust_dyn_coef = 7.9435775497736785
+  vmotor2thrust = [-0.02476537915958403, 0.06523793527519485, -0.026792504967750107, 0.006776789303971145]  # Index is order
+  vmotor2torque = [-2.8633106919309745e-5, 0.00011679386117520097, 5.105754520419129e-5, 0.0]  # Index is order
+  vmotor2rpm = [4657.542534331524, 7536.161830990926]  # Index is order
+  prop_radius = 23.4e-3 # TODO check
+  -->
   <compiler inertiafromgeom="false" meshdir="assets" autolimits="true" />
 
   <asset>

crazyflow/drones/cf2x_T350.xml

@@ -1,4 +1,13 @@
 <mujoco model="cf2x">
+  <!--
+  Platform data not used by the simulation, kept for estimators, firmware and tooling:
+  rotor_dyn_coef_simple = 11.352970450445243
+  thrust_dyn_coef = 11.12424272978587
+  vmotor2thrust = [0.006728127583707208, 0.01011557616217668, 0.010263198062061085, 0.0028358638322392503]  # Index is order
+  vmotor2torque = [-1.2906047901738756e-5, 0.0001436101487030899, 2.794753913624656e-5, 1.3104535533494383e-5]  # Index is order
+  vmotor2rpm = [2977.884883031915, 8101.0293594093055]  # Index is order
+  prop_radius = 25.4e-3 # TODO check
+  -->
   <compiler inertiafromgeom="false" meshdir="assets" autolimits="true" />
 
   <asset>