Automate availability and add docs

SKILL.md

@@ -49,9 +49,11 @@ 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`.
+`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/<model>/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.
 

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"

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.
 
@@ -78,9 +79,10 @@ 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, "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 +96,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 +105,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 +123,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 +133,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.
 

crazyflow/drones/__init__.py

@@ -3,14 +3,26 @@
 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.
 """
 
-# 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 enum import StrEnum
+from pathlib import Path
 
-__all__ = ["available_drones"]
+__all__ = ["Drone"]
+
+
+class Drone(StrEnum):
+    """Drone configurations. Each member has an MJCF file ``crazyflow/drones/<name>.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)}"
+)

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]:

crazyflow/dynamics/core.py

@@ -9,6 +9,7 @@
 
 import numpy as np
 
+from crazyflow.drones import Drone
 from crazyflow.utils import filter_to_signature, to_xp
 from crazyflow.utils import parametrize as _parametrize
 
@@ -53,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.
 
@@ -66,10 +67,11 @@ 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
 
-    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)
@@ -85,16 +87,16 @@ 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.
 
     Merges the global parameters in ``crazyflow/dynamics/params.toml`` with the drone's section in
     ``crazyflow/dynamics/<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.
@@ -103,16 +105,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.
@@ -121,7 +123,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.
 
@@ -130,7 +132,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.
@@ -141,3 +143,46 @@ 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) -> set[str]:
+    """Return the drone sections declared in a dynamics model's ``params.toml``."""
+    with open(Path(__file__).parent / f"{dynamics}/params.toml", "rb") as f:
+        return set(tomllib.load(f))
+
+
+def supported_drones(dynamics: Dynamics) -> tuple[Drone, ...]:
+    """Return the drones that ``dynamics`` can be parametrized for.
+
+    A drone is supported when ``crazyflow/dynamics/<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.
+
+    A model is supported when its ``crazyflow/dynamics/<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))

crazyflow/dynamics/first_principles/params.toml

@@ -10,22 +10,26 @@
 # 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.
 
 [cf2x_L250]
 mass = 0.0319

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

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

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

crazyflow/envs/drone_env.py

@@ -12,14 +12,15 @@
 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
 from crazyflow.sim.pipeline import append_fn
 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,