Reintegrate drone_models and drone_controllers into crazyflow (#71)

* Merge models back into crazyflow

* Introduce consistent naming scheme for dynamics/physics/models

* Improve naming consistency in docs

* Finalize consistent naming for drones/models/physics/dynamics

* Rename dynamics file, simplify so_rpy model, simplify models wrapper

* Remove drone-models submodule

* Simplify wrappers, add drone-models test suite to crazyflow

* Improve tests

* Fix test imports

* Rename dynamics docs

* Add drone model docs

* Fix docs

* Add controllers back into crazyflow

* Improve API and add docs

* Explicitly name controllers

* Remove all submodules. Update deps

* Update doc mistakes and benchmarks

* Address comments

* [WIP] Refactor parametrize

* Move fused models into regular xml

* Refactor parametrization

* Fix docs and comments

* Fix model docs. Restructure into core dynamics_euler and wrapper dynamics

* Fix doctests

* Disable flaky render tests in CI

* Fix wrong drone descriptions

Author: amacati

.gitmodules

@@ -12,7 +12,7 @@
 
 </div>
 
-Crazyflow is a research simulator for quadrotors. It runs batched, differentiable simulations on CPU and GPU via JAX, with analytical and abstracted models for the Crazyflie 2.x family.
+Crazyflow is a research simulator for quadrotors. It runs batched, differentiable simulations on CPU and GPU via JAX, with analytical and abstracted dynamics for the Crazyflie 2.x family.
 
 ```python
 import numpy as np
@@ -36,10 +36,10 @@ for _ in range(100):
 ## Features
 
 - **n\_worlds x n\_drones** — batched over independent environments and multi-drone swarms simultaneously
-- **GPU-accelerated** — up to 914 M steps/s on an RTX 4090 (first-principles physics, 262 K worlds)
+- **GPU-accelerated** — up to 914 M steps/s on an RTX 4090 (first-principles dynamics, 262 K worlds)
 - **Differentiable** — `jax.grad` works through the full dynamics and control pipeline
-- **First-principles models** — physics model using first-principles equations and parameters identified from real-world measurements
-- **Abstracted models** — three physics models fitted from real Crazyflie flight data
+- **First-principles dynamics** — dynamics using first-principles equations and parameters identified from real-world measurements
+- **Abstracted dynamics** — simplified dynamics in three flavors fitted from real Crazyflie flight data
 - **Modular pipelines** — step and reset are tuples of plain JAX functions; insert anything, anywhere
 - **MuJoCo integration** — onscreen and offscreen rendering, raycasting, and contact detection via MJX
 
@@ -60,7 +60,7 @@ pixi shell
 
 ## Performance
 
-First-principles physics, one drone. CPU: AMD Ryzen 9 7950X. GPU: NVIDIA RTX 4090.
+First-principles dynamics, one drone. CPU: AMD Ryzen 9 7950X. GPU: NVIDIA RTX 4090.
 
 | n\_worlds | CPU steps/s | GPU steps/s |
 |---|---|---|
@@ -72,17 +72,6 @@ First-principles physics, one drone. CPU: AMD Ryzen 9 7950X. GPU: NVIDIA RTX 409
 
 Full benchmarks including multi-drone scaling are in the [documentation](https://learnsyslab.github.io/crazyflow).
 
-## Related packages
-
-Crazyflow is built on two companion packages that can also be used independently:
-
-| Package | Description |
-|---|---|
-| [drone-models](https://github.com/learnsyslab/drone-models) | Drone dynamics models (first-principles and fitted) compatible with NumPy, JAX, and PyTorch. Used by Crazyflow as the physics backend. |
-| [drone-controllers](https://github.com/learnsyslab/drone-controllers) | Reference controller implementations including the Mellinger geometric controller. Used by Crazyflow to provide the state and attitude control modes. |
-
-Both are installed automatically as dependencies. For development, they are included as submodules in `submodules/` and installed in editable mode by the pixi environment.
-
 ## Citation
 
 ```bibtex

README.md

@@ -4,6 +4,8 @@
 import jax
 
 from crazyflow import Sim
+from crazyflow.control import Control
+from crazyflow.dynamics import Dynamics
 
 
 def main(cache: bool = False):
@@ -16,7 +18,7 @@ def main(cache: bool = False):
 
     # Time initialization
     start = time.perf_counter()
-    sim = Sim(n_worlds=1, n_drones=1, physics="sys_id", control="attitude")
+    sim = Sim(n_worlds=1, n_drones=1, dynamics=Dynamics.so_rpy, control=Control.attitude)
     init_time = time.perf_counter() - start
 
     # Time reset compilation
@@ -29,7 +31,7 @@ def main(cache: bool = False):
     sim._step.lower(sim.data, 1).compile()
     step_time = time.perf_counter() - start
 
-    print(f"Simulation startup times | {sim.physics} | {sim.control}")
+    print(f"Simulation startup times | {sim.dynamics} | {sim.control}")
     print(f"Initialization: {init_time:.2f}s")
     print(f"Reset: {reset_time:.2f}s")
     print(f"Step: {step_time:.2f}s")

benchmark/compile.py

@@ -55,7 +55,7 @@ def profile_gym_env_step(
         num_envs=sim_config.n_worlds,
         device=sim_config.device,
         freq=sim_config.freq,
-        physics=sim_config.physics,
+        dynamics=sim_config.dynamics,
     )
 
     # Action for going up (in attitude control)
@@ -151,7 +151,7 @@ def main(device: str = "cpu", n_worlds_exp: int = 6):
     sim_config = config_dict.ConfigDict()
     sim_config.n_worlds = 1
     sim_config.n_drones = 1
-    sim_config.physics = "first_principles"
+    sim_config.dynamics = "first_principles"
     sim_config.control = "attitude"
     sim_config.attitude_freq = 500
     sim_config.device = device

benchmark/main.py

@@ -4,13 +4,12 @@
 
 def main():
     """Main entry point for profiling."""
-    sim = Sim(n_worlds=1, n_drones=1, physics="first_principles", control="attitude")
+    sim = Sim(n_worlds=1, n_drones=1, dynamics="first_principles", control="attitude")
 
     compiled_reset = sim._reset.lower(sim.data, sim.default_data, None).compile()
     compiled_step = sim._step.lower(sim.data, 1).compile()
-    op_count_reset = compiled_reset.cost_analysis()["flops"]
-    op_count_step = compiled_step.cost_analysis()["flops"]
-    print(f"Op counts:\n Reset: {op_count_reset}\n Step: {op_count_step}")
+    print(f"Reset cost analysis: {compiled_reset.cost_analysis()}")
+    print(f"Step cost analysis: {compiled_step.cost_analysis()}")
 
 
 if __name__ == "__main__":

benchmark/op_count.py

@@ -42,10 +42,13 @@ def profile_step(sim_config: config_dict.ConfigDict, n_steps: int, device: str):
 
 
 def profile_gym_env_step(sim_config: config_dict.ConfigDict, n_steps: int, device: str):
-    device = jax.devices(device)[0]
-
     envs: ReachPosEnv = gymnasium.make_vec(
-        "DroneReachPos-v0", time_horizon_in_seconds=2, num_envs=sim_config.n_worlds, **sim_config
+        "DroneReachPos-v0",
+        max_episode_time=10.0,
+        num_envs=sim_config.n_worlds,
+        dynamics=sim_config.dynamics,
+        freq=50,
+        device=device,
     )
 
     # Action for going up (in attitude control)
@@ -77,7 +80,7 @@ def main():
     sim_config = config_dict.ConfigDict()
     sim_config.n_worlds = 1
     sim_config.n_drones = 1
-    sim_config.physics = "first_principles"
+    sim_config.dynamics = "first_principles"
     sim_config.control = "attitude"
     sim_config.device = device
 

benchmark/performance.py

@@ -18,7 +18,9 @@
 
 import crazyflow.envs  # noqa: F401, ensure gymnasium envs are registered
 from crazyflow.control import Control
-from crazyflow.sim import Physics, Sim
+from crazyflow.drones import available_drones
+from crazyflow.dynamics import Dynamics
+from crazyflow.sim import Sim
 
-__all__ = ["Sim", "Physics", "Control"]
+__all__ = ["Sim", "Dynamics", "Control", "available_drones"]
 __version__ = "0.2.0"

crazyflow/__init__.py

@@ -0,0 +1,8 @@
+"""This file is to be removed as soon as a proper typing is available by the official array-api."""
+
+from typing import Any, TypeAlias
+
+import numpy.typing as npt
+
+Array: TypeAlias = Any  # To be changed to array_api_typing later
+ArrayLike: TypeAlias = Array | npt.ArrayLike

crazyflow/_typing.py

@@ -1,3 +1,26 @@
-from crazyflow.control.control import Control
+"""Implementations of onboard drone controllers in Python.
 
-__all__ = ["Control"]
+All controllers are implemented using the array API standard. This means that every controller is
+agnostic to the choice of framework and supports e.g. NumPy, JAX, or PyTorch. We also implement all
+controllers as pure functions to ensure that users can jit-compile them. All controllers use
+broadcasting to support batching of arbitrary leading dimensions.
+
+We reimplement the onboard controller for two reasons:
+- We cannot use the C++ bindings of the firmware to differentiate through the onboard controller.
+- We need to implement it with JAX to enable efficient, batched computations.
+"""
+
+from typing import Callable
+
+__all__ = []
+
+from crazyflow.control.core import Control, parametrize
+from crazyflow.control.mellinger import attitude2force_torque as mellinger_attitude2force_torque
+from crazyflow.control.mellinger import state2attitude as mellinger_state2attitude
+
+available_controller: dict[str, Callable] = {
+    "mellinger_state2attitude": mellinger_state2attitude,
+    "mellinger_attitude2force_torque": mellinger_attitude2force_torque,
+}
+
+__all__ = ["Control", "parametrize"]