-Fast, parallelizable simulations of drones with JAX.
-
-[![Python Version]][Python Version URL] [![Ruff Check]][Ruff Check URL] [![Documentation Status]][Documentation Status URL] [![Tests]][Tests URL]
-
-[Python Version]: https://img.shields.io/badge/python-3.11+-blue.svg
-[Python Version URL]: https://www.python.org
-
-[Ruff Check]: https://github.com/learnsyslab/crazyflow/actions/workflows/ruff.yml/badge.svg?style=flat-square
-[Ruff Check URL]: https://github.com/learnsyslab/crazyflow/actions/workflows/ruff.yml
+ **Fast, parallelizable simulations of Quadrotor drones with JAX.**
-[Documentation Status]: https://github.com/learnsyslab/crazyflow/actions/workflows/docs.yml/badge.svg
-[Documentation Status URL]: https://learnsyslab.github.io/crazyflow
+ [](https://www.python.org)
+ [](https://github.com/learnsyslab/crazyflow/actions/workflows/testing.yml)
+ [](https://github.com/learnsyslab/crazyflow/actions/workflows/ruff.yml)
+ [](https://learnsyslab.github.io/crazyflow)
-[Tests]: https://github.com/learnsyslab/crazyflow/actions/workflows/testing.yml/badge.svg
-[Tests URL]: https://github.com/learnsyslab/crazyflow/actions/workflows/testing.yml
+
-## Quick Start
-For a more detailed guide, check out our [documentation](https://learnsyslab.github.io/crazyflow/).
+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.
-### Normal installation
-The regular way to use Crazyflow is to install it from PyPI with your favourite package manager, e.g., with pip:
-``` bash
-pip install crazyflow
-```
-
-### Developer installation
-If you plan to develop with and around Crazyflow, you can use the existing [pixi](https://pixi.sh/) environment.
-``` bash
-git clone --recurse-submodules git@github.com:learnsyslab/crazyflow.git
-cd crazyflow
-pixi shell
-```
+```python
+import numpy as np
+from crazyflow.sim import Sim
+from crazyflow.control import Control
-This will install Crazyflow, drone-models, and drone-controllers in editable mode for easy development.
+sim = Sim(n_worlds=4096, n_drones=1, control=Control.state)
+cmd = np.zeros((4096, 1, 13))
+cmd[..., 2] = 0.5 # hover at 0.5 m across all worlds
-In case you want to use another package manager or install the simulator with it's subpackages into another project, simply install all packages manually in your environment:
-``` bash
-pip install -e . # Installing Crazyflow
-pip install -e ./submodules/drone-models
-pip install -e ./submodules/drone-controllers
+for _ in range(100):
+ sim.state_control(cmd)
+ sim.step(sim.freq // sim.control_freq)
+ sim.render()
```
+## Documentation
-## Architecture
-
-Crazyflow is a high-performance simulation framework for Crazyflie drones that leverages JAX for efficient parallelization and automatic differentiation. The architecture is designed around a flexible pipeline that can be configured at initialization time, enabling users to swap out physics backends, control methods, and integration schemes.
-
-### Core Components
-
-#### Simulation Pipeline
-The simulation is built as a pipeline of functions that are composed at initialization time based on the configuration. This approach avoids runtime branching and allows JAX to optimize the entire pipeline as a single computation. Users can insert their own pure functions into the pipeline to modify the simulation behavior while maintaining compatibility with JAX's optimizations.
-
-#### Physics Backends
-Multiple physics models are supported:
-- first_principles: A first-principles model based on physical equations
-- so_rpy: A system-identified model trained on real drone data
-- so_rpy_rotor: An enhanced system-identified model that includes thrust dynamics
-- so_rpy_rotor_drag: A system-identified model that includes thrust dynamics and drag effects
-
-#### Control Modes
-Different control interfaces are available:
-- state: High-level control of position, velocity, and yaw
-- attitude: Mid-level control of collective thrust and orientation
-- thrust: Low-level control of individual motor thrusts
-
-#### Integration Methods
-We support multiple integration schemes for additional precision:
-- euler: Simple first-order integration
-- rk4: Fourth-order Runge-Kutta integration for higher accuracy
-- symplectic\_euler: Symplectic integration for conservation of energy
+[learnsyslab.github.io/crazyflow](https://learnsyslab.github.io/crazyflow) — installation, user guide, examples, and API reference.
-### Parallelization
-Crazyflow supports massive parallelization across:
-- Worlds: Independent simulation environments that can run in parallel
-- Drones: Multiple drones within each world
-- Devices: Computations can be executed on CPU or GPU
-This parallelization is achieved through JAX's vectorization capabilities, allowing thousands of simulations to run simultaneously with minimal overhead.
+## Features
-### Domain Randomization
-The framework supports domain randomization through the crazyflow/randomize module, allowing parameters like mass to be varied across simulations to improve sim-to-real transfer.
+- **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)
+- **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
+- **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
-### Functional Design
-The simulation follows a functional programming paradigm: All state is contained in immutable data structures. Updates create new states rather than modifying existing ones. All functions are pure, enabling JAX's transformations (JIT, grad, vmap) and thus automatic differentiation through the entire simulation, making it suitable for gradient-based optimization and reinforcement learning.
+## Installation
-### Contacts and Non-Drone Models
-We focus on drones dynamics in free-space flight. Consequently, no models other than drones are available in the simulation and contact dynamics with external objects are not considered. However, we use MuJoCo for contact detection and visualization. Users can load their own objects into the simulation by changing the MuJoCo world spec. Drone collisions with these objects will be detected during collision checks, but they won't have an effect on the dynamics (i.e. drones will pass through objects). Similarly, the objects themselves will be static.
+```bash
+pip install crazyflow # CPU
+pip install "crazyflow[gpu]" # GPU (Linux x86-64, CUDA 12)
+```
-### Visualization
-We use `gymnasium`'s MuJoCo renderer and synchronize the simulation data with MuJoCo to either render an interactive UI or RGB arrays.
+Developer install with editable submodules ([pixi](https://pixi.sh/) required):
-## Examples
-The repository includes several example scripts demonstrating different capabilities:
-| Example | Description |
-| ----------------------------------------- | ----------------------------------------------------------- |
-| [`hover.py`](examples/hover.py) | Basic hovering using state control |
-| [`thrust.py`](examples/thrust.py) | Direct motor control using thrust commands |
-| [`render.py`](examples/render.py) | Visualization of multiple drones with motion traces |
-| [`contacts.py`](examples/contacts.py) | Collision detection between drones |
-| [`gradient.py`](examples/gradient.py) | Computing gradients through the simulation for optimization |
-| [`change_pos.py`](examples/change_pos.py) | Manipulating drone positions programmatically |
+```bash
+git clone --recurse-submodules git@github.com:learnsyslab/crazyflow.git
+cd crazyflow
+pixi shell
+```
## Performance
-These benchmarks give you a rough idea of the performance you should expect from the simulator. Gym benchmarks use gym environments with a step frequency of 50Hz while simulating at 500Hz. The simulation benchmarks use 500Hz.
-
-
-
-The chart above shows the performance of Crazyflow on different hardware. The simulator can achieve close to 100 million steps per second on a GPU with 1 million parallel environments. The gym environment, which includes additional overhead for the Gymnasium interface and is not fully optimized, still achieves decent performance with over 98,000 steps per second on GPU with 10,000 parallel environments.
-
-Performance benchmarks were run on:
-- CPU: Intel Core i9-13900KF
-- GPU: NVIDIA RTX 4090
-To reproduce the benchmark results, rerun the `benchmark/main.py` script.
+First-principles physics, one drone. CPU: AMD Ryzen 9 7950X. GPU: NVIDIA RTX 4090.
-## Known Issues
-- `"RuntimeError: MUJOCO_PATH environment variable is not set"` upon installing this package. This error can be resolved by using `venv` instead of `conda`. Somtimes the `mujoco` install can [fail with `conda`](https://github.com/google-deepmind/mujoco/issues/1004).
-- If using `zsh` don't forget to escape brackets when installing additional dependencies: `pip install .\[gpu\]`.
+| n\_worlds | CPU steps/s | GPU steps/s |
+|---|---|---|
+| 64 | 3.3 M | 1.2 M |
+| 1 024 | 9.2 M | 18.7 M |
+| 16 384 | 11.9 M | 257 M |
+| 65 536 | 15.6 M | 678 M |
+| 262 144 | 12.6 M | 914 M |
-### Using the project with VSCode devcontainers
+Full benchmarks including multi-drone scaling are in the [documentation](https://learnsyslab.github.io/crazyflow).
-**Running on CPU**: by default the containers run on CPU. You don't need to take any action.
+## Related packages
-**Running on GPU**: The devcontsainers can easily run using your computer's NVIDIA GPU on Linux and Windows. This makes sense if you want to accelerate simulation by running thousands of simulation in parallel. In order to work you need to install the [CUDA Toolkit](https://developer.nvidia.com/cuda-downloads?target_os=Linux&target_arch=x86_64&Distribution=WSL-Ubuntu&target_version=2.0&target_type=deb_local), [NVIDIA Container runtime](https://developer.nvidia.com/container-runtime) for your computer. Finally, enable GPU access to the devcontainers by setting the commented out `"--gpus=all"` and `"--runtime=nvidia"` flags in `devcontainer.json`.
+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. |
-**Linux**
-1. Make sure to be in a X11 session ([link](https://askubuntu.com/questions/1410256/how-do-i-use-the-x-window-manager-instead-of-wayland-on-ubuntu-22-04)), otherwise rendering of the drone will fail.
-2. Install [Docker](https://docs.docker.com/engine/install/) (, and make sure Docker Daemon is running)
-3. Install [VSCode](https://code.visualstudio.com/), with [devcontainer extension](https://marketplace.visualstudio.com/items?itemName=ms-vscode-remote.remote-containers), and [remote dev pack](https://marketplace.visualstudio.com/items?itemName=ms-azuretools.vscode-docker).
-4. Clone this project's code. Rename `/.devcontainer/devcontainer.linux.json` to `/.devcontainer/devcontainer.json`.
-5. Open this project in VSCode. VSCode should automatically detect the devcontainer and prompt you to `Reopen in container`. If not, see [here](https://code.visualstudio.com/docs/devcontainers/containers#_quick-start-open-an-existing-folder-in-a-container) to open manually. Note: Opening the container for the first time might take a while (up to 15 min), as the container is pulled from the web and build.
+Both are installed automatically as dependencies. For development, they are included as submodules in `submodules/` and installed in editable mode by the pixi environment.
-**Windows** (requires Windows 10 or later)
-
-For windows, we require WSL2 to run the devcontainer. (So its actually Linux with extra steps.) Full instructions can be found [in the official docs](https://code.visualstudio.com/blogs/2020/07/01/containers-wsl#_getting-started). Here are the important steps:
-1. Install [Docker](https://docs.docker.com/desktop/setup/install/windows-install/), and WSL2, and Ubuntu 22.04 LTS (, and make sure Docker Daemon is running)
-2. Docker will recognize that you have WSL installed and prompt you via Windows Notifications to enable WSL integration -> confirm this with `Enable WSL integration`. If not, open `Docker Desktop`, navigate to the settings, and manually enable WSL integration. (There are TWO setting options for this. Make sure to enable BOTH!)
-3. Install [VSCode](https://code.visualstudio.com/), with the [WSL extension](https://marketplace.visualstudio.com/items?itemName=ms-vscode-remote.remote-wsl), [devcontainer extension](https://marketplace.visualstudio.com/items?itemName=ms-vscode-remote.remote-containers), and [remote dev pack](https://marketplace.visualstudio.com/items?itemName=ms-azuretools.vscode-docker).
-4. Clone the source code for the exercises in the WSL2 file system to `/home` (`~`), or wherever you like. (Performance when working on the WSL file system is much better compared to Windows file system). You can access the WSL filesystem either by starting a WSL/Ubuntu console, or via the Windows File Explorer at `\\wsl.localhost\Ubuntu\home` (replace `Ubuntu` with your distro, if necessary).
-7. Rename `/.devcontainer/devcontainer.windows.json` to `/.devcontainer/devcontainer.json`.
-8. Open this project in VSCode. The easiest way to do so is by starting a WSL/Ubuntu shell, navigating via `cd` to the source code, then type `code .` to open VSCode. VSCode should automatically detect the devcontainer and prompt you to `Reopen in container`. If not, see [here](https://code.visualstudio.com/docs/devcontainers/containers#_quick-start-open-an-existing-folder-in-a-container) to open manually. Note: Opening the container for the first time might take a while (up to 15 min), as the container is pulled from the web and build.
-
-
-**MacOS**
-
-Unfortunately, we did not get the devcontainer to work with MacOS yet, even after following [those](https://gist.github.com/sorny/969fe55d85c9b0035b0109a31cbcb088) steps. We expect that the issue is related to Mujoco rendering from inside the Docker container and display forwarding with X11. There is also an [unresolved Issue](https://github.com/google-deepmind/mujoco/issues/1047) on GitHub. If you manage to make it work, please let us know.
-
-Until then, MacOS users are required to install this project using an python environment manager such as [conda](https://docs.anaconda.com/anaconda/install/) or [mamba](https://mamba.readthedocs.io/en/latest/). If you use conda, these are the required commands: ```conda create --name crazyflow -c conda-forge python=3.11```, ```conda activate crazyflow```, ```conda install pip```, ```pip install -e .```.
-
-____________
-
-Known Issues:
- - if building docker container fails at `RUN apt-get update`, make sure your host systems time is set correct: https://askubuntu.com/questions/1511514/docker-build-fails-at-run-apt-update-error-failed-to-solve-process-bin-sh
+## Citation
+```bibtex
+@misc{schuck2025crazyflow,
+ author = {Schuck, Martin and Rath, Marcel P. and Hua, Yufei and Goudar, Abhishek and Zhou, SiQi and Schoellig, Angela P.},
+ title = {Crazyflow: An Accurate, GPU-Accelerated Differentiable Drone Simulator in JAX},
+ year = {2026},
+ note = {Preprint}
+}
+```
diff --git a/docs/api/index.md b/docs/api/index.md
index 22d55f3..63b6068 100644
--- a/docs/api/index.md
+++ b/docs/api/index.md
@@ -1,5 +1,22 @@
-# Python API
+# API Reference
-This section is generated with mkdocstrings (Python handler). It documents the public API of the crazyflow package.
+This section is auto-generated from the Crazyflow source code using [mkdocstrings](https://mkdocstrings.github.io/).
-::: crazyflow
\ No newline at end of file
+## Module overview
+
+| Module | Description |
+|---|---|
+| `crazyflow.sim` | Core `Sim` class and physics pipeline |
+| `crazyflow.sim.data` | `SimData`, `SimState`, `SimControls`, `SimParams`, `SimCore` pytrees |
+| `crazyflow.sim.functional` | Pure functional control API for use inside `jax.jit` |
+| `crazyflow.sim.physics` | `Physics` enum and physics model implementations |
+| `crazyflow.sim.integration` | `Integrator` enum, Euler, RK4, and symplectic Euler |
+| `crazyflow.sim.sensors` | Raycasting and sensor extraction utilities |
+| `crazyflow.sim.symbolic` | CasADi symbolic model API |
+| `crazyflow.control` | `Control` enum |
+| `crazyflow.control.mellinger` | Mellinger controller data and parameters |
+| `crazyflow.envs` | Gymnasium vectorized environments |
+| `crazyflow.randomize` | Domain randomization helpers |
+| `crazyflow.utils` | Grid utilities and pytree helpers |
+
+Navigate the full generated reference using the sidebar.
diff --git a/docs/cite.md b/docs/cite.md
new file mode 100644
index 0000000..e57b129
--- /dev/null
+++ b/docs/cite.md
@@ -0,0 +1,16 @@
+# Cite
+
+If you use Crazyflow in academic work, please cite the accompanying paper:
+
+BibTeX:
+
+```bibtex
+@misc{schuck2025crazyflow,
+ author = {Schuck, Martin and Rath, Marcel P. and Hua, Yufei and Goudar, Abhishek and Zhou, SiQi and Schoellig, Angela P.},
+ title = {Crazyflow: An Accurate, GPU-Accelerated Differentiable Drone Simulator in JAX},
+ year = {2026},
+ note = {Preprint}
+}
+```
+
+The paper covers motivation, architecture, benchmarks against related simulators, and sim-to-real validation results. A DOI and journal/conference venue will be added here once the work is published.
diff --git a/docs/examples.md b/docs/examples.md
deleted file mode 100644
index 3aebb1a..0000000
--- a/docs/examples.md
+++ /dev/null
@@ -1,92 +0,0 @@
-# Examples
-
-Selected example scripts. Check out more in `examples/`
-
----
-
-### State control
-
-A minimal demo showing how to use the Sim API at a very high level: create a Sim, send simple position commands and step the simulation. Useful to get started with high-level position/setpoint control.
-
-```bash
-python examples/change_pos.py
-```
-
----
-
-### Cameras & RGBD
-
-Shows how to obtain RGB and depth frames from the renderer. Demonstrates offscreen capture via sim.render(), retrieving image arrays for perception or visual‑in‑the‑loop controllers, and saving sequences (GIFs).
-
-```bash
-python examples/cameras.py
-```
-
-
-RGB Image
-
-Depth Image
-
-
-
----
-
-### Contacts & collision model
-
-Demonstrates how to query contact information from the simulation and how to switch collision geometry types (e.g., sphere → box) for higher‑fidelity contact modeling. Useful for debugging collisions or when more accurate contact checks are required.
-
-```bash
-python examples/contacts.py
-```
-
-{width=200 height=200}
-{width=200 height=200}
-
----
-
-### LED deck & materials
-
-Illustrates how to activate and control the LED deck and other drone material colors at runtime for visualization and debugging.
-
-```bash
-python examples/led_deck.py
-```
-
-
-
----
-
-### Randomization
-
-Shows how to add reset‑time randomization: initial position/quaternion randomization, mass/inertia perturbations and other properties that should vary between episodes. The example demonstrates how to register reset randomizers and how to persist different initial conditions across runs.
-
-Run:
-```bash
-python examples/randomize.py
-```
-
----
-
-### Disturbance
-
-Demonstrates step‑time disturbances inserted into the step pipeline (external forces/torques, actuator noise, etc.). The example (examples/disturbance.py) shows how to insert a disturbance function into sim.step_pipeline, compare disturbed vs undisturbed runs and optionally plot the resulting trajectories.
-
-Run:
-```bash
-python examples/disturbance.py
-```
-
----
-
-### Figure‑8 / RL environment
-
-A scripted figure‑8 environment intended for evaluation or as a training target. The example shows how to create vectorized envs, apply the NormalizeActions wrapper and step/render the environment. It does not include any agent implementations — integrate the environment with your preferred RL training code (Stable Baselines3, RLlib, custom JAX trainers, etc.).
-
-Run:
-```bash
-python examples/figure8.py
-```
-
----
-
-For API details and configuration options referenced by these examples, see the [API Reference](api/index.md) and [Usage](usage.md).
\ No newline at end of file
diff --git a/docs/examples/index.md b/docs/examples/index.md
new file mode 100644
index 0000000..daad37a
--- /dev/null
+++ b/docs/examples/index.md
@@ -0,0 +1,136 @@
+# Examples
+
+These examples build on each other — each one introduces one new concept on top of the previous. Start from the top if you're new, or jump to whichever section covers what you need.
+
+---
+
+## Hover
+
+A single drone commanded to hold a fixed height using state control. This is the minimal end-to-end loop: create a `Sim`, reset it, apply a state command, and step forward.
+
+```{ .python notest }
+--8<-- "examples/hover.py"
+```
+
+```bash
+python examples/hover.py
+```
+
+---
+
+## Attitude control
+
+Commanding roll, pitch, yaw, and collective thrust directly. This level bypasses the Mellinger position loop and is typical for RL agents that output attitude targets.
+
+```{ .python notest }
+--8<-- "examples/attitude.py"
+```
+
+---
+
+## Gradient descent through dynamics
+
+Because the simulator is built entirely from JAX operations, `jax.grad` can differentiate through it. Starting the drone above the target height keeps it away from the floor, so the floor-clipping stage never fires and gradients flow freely through the entire trajectory.
+
+```{ .python notest }
+--8<-- "examples/gradient.py"
+```
+
+---
+
+## Domain randomization
+
+Varying physical parameters per world at reset. Each world gets a slightly different mass, so identical commands produce diverging trajectories.
+
+```{ .python notest }
+--8<-- "examples/randomize.py"
+```
+
+---
+
+## Disturbance injection
+
+Inserting a random external force and torque into the step pipeline. The disturbance fires on every physics tick, so the drone fights wind-like perturbations.
+
+```{ .python notest }
+--8<-- "examples/disturbance.py"
+```
+
+---
+
+## Cameras and RGBD
+
+Offscreen rendering returns RGB and depth images on every frame. The FPV camera (`fpv_cam`) is attached to the drone and moves with it.
+
+
+
+
+
+```{ .python notest }
+--8<-- "examples/cameras.py"
+```
+
+```bash
+python examples/cameras.py
+```
+
+---
+
+## LED deck and materials
+
+`change_material` updates the RGBA colour and emission of any named material on any subset of drones at runtime.
+
+
+
+
+
+```{ .python notest }
+--8<-- "examples/led_deck.py"
+```
+
+```bash
+python examples/led_deck.py
+```
+
+---
+
+## Contact queries
+
+The default collision geometry is a sphere around the drone frame. `use_box_collision` replaces it with a tighter oriented box, useful for narrow-gap flight and accurate contact debugging.
+
+
+
+
+
+
+
+
+
+
+```{ .python notest }
+--8<-- "examples/contacts.py"
+```
+
+---
+
+## Raycasting and depth sensing
+
+`render_depth` fires rays from a camera and returns per-pixel distances. This is faster than full RGB rendering and useful for obstacle sensing or depth-based controllers.
+
+```{ .python notest }
+--8<-- "examples/raycasting.py"
+```
+
+```bash
+python examples/raycasting.py
+```
+
+---
+
+## Gymnasium environment
+
+Evaluating a random policy in the figure-8 environment. The env wraps `Sim` behind the standard Gymnasium `VectorEnv` interface.
+
+```{ .python notest }
+--8<-- "examples/figure8.py"
+```
diff --git a/docs/features.md b/docs/features.md
deleted file mode 100644
index ef3d4ff..0000000
--- a/docs/features.md
+++ /dev/null
@@ -1,34 +0,0 @@
-# Features & Architecture
-
-Crazyflow is a research‑first simulator for small quadrotors. Its design favours clarity, composability and performance: a compact Sim API composes a physics model, a numerical integrator and a controller layer. Around that core sits an extensible "step pipeline" where randomization, disturbances, logging or custom hooks can be inserted without changing the simulator internals.
-
-### Architecture
-The goals are reproducibility, throughput and research flexibility. First‑principles (analytic, identified) and simplified, data‑driven models are supported alongside a symbolic model API for controller development and analysis. The codebase is implemented to take advantage of JAX and MuJoCo: batched GPU execution, JIT compilation and automatic differentiation are available, and analytical gradients of the simulation are exposed where useful. The step pipeline makes it simple to compose randomization, disturbances and custom hooks without touching the core simulator.
-
-### Models
-
-Crazyflow supports two complementary model classes:
-
-- First‑principles models — physics‑based analytical models with identified parameters intended for high‑fidelity simulation and sim‑to‑real work.
-- Simplified, data‑driven models — lightweight models fitted from flight data to capture off‑nominal effects and to speed up learning/control experiments. These data‑driven models can be obtained from just a few minutes of flight data; the repository provides fitting scripts and a minimal identification pipeline (works quickly if a stable controller is available).
-
-All models (first‑principles and data‑driven) are also exposed as symbolic model objects so they can be reused directly in model‑based controllers, MPC formulations, or analysis tools.
-
-### Controllers
-The repository includes reference controller implementations and integration points for common research workflows:
-
-- Geometric controllers (Mellinger style) for standard tracking.
-- Interfaces for MPC / MPCC workflows (note: MPC example controllers and advanced control code are available in the repo for our [drone racing course](https://github.com/learnsyslab/lsy_drone_racing/tree/main/lsy_drone_racing/control)).
-- Reinforcement learning: we provide environments suitable for training and deploying RL agents and include example setups for PPO and SHAC agents.
-
-### Performance & evaluation
-Assuming free‑space flight (avoiding generic contact solving when possible) lets Crazyflow prioritise speed and simplicity. The combination of JAX, MuJoCo and a modular pipeline enables high‑throughput experiments and large‑scale benchmarking.
-
-### Extensibility
-Researchers can add new dynamics, controllers or pipeline stages without modifying the simulator core. Swap integrators, change physics models, or inject disturbances via small, well‑scoped functions.
-
-### Publication
-See our upcoming paper for full motivation, benchmarks and evaluation:
-Schuck, M. & Rath, M. (2025). Crazyflow: Fast, parallelizable simulations of Crazyflies with JAX and MuJoCo. (preprint — link TBD)
-
-For API details and configuration options, see the API reference: [API Reference](api/index.md)
\ No newline at end of file
diff --git a/docs/gen_ref_pages.py b/docs/gen_ref_pages.py
new file mode 100644
index 0000000..58de12f
--- /dev/null
+++ b/docs/gen_ref_pages.py
@@ -0,0 +1,68 @@
+"""Generate the code reference pages and navigation.
+
+This script is executed by the mkdocs-gen-files plugin during ``mkdocs build`` or
+``mkdocs serve``. It is not meant to be run directly or imported outside of that context.
+Install the docs environment (``pixi shell -e docs``) to use it.
+"""
+
+from pathlib import Path
+
+try:
+ import mkdocs_gen_files
+except ImportError:
+ pass # not running in a docs environment — nothing to generate
+else:
+ SKIP_PARTS = {"_typing", "__main__", "__pycache__"}
+
+ for path in sorted(Path("crazyflow").rglob("*.py")):
+ module_path = path.relative_to(".").with_suffix("")
+ doc_path = path.relative_to(".").with_suffix(".md")
+ full_doc_path = Path("api", doc_path)
+
+ parts = tuple(module_path.parts)
+
+ if any(part in SKIP_PARTS for part in parts):
+ continue
+
+ if parts[-1] == "__init__":
+ parts = parts[:-1]
+ doc_path = doc_path.with_name("index.md")
+ full_doc_path = full_doc_path.with_name("index.md")
+ elif parts[-1] == "__main__":
+ continue
+
+ with mkdocs_gen_files.open(full_doc_path, "w") as fd:
+ ident = ".".join(parts)
+ fd.write(f"::: {ident}\n")
+
+ mkdocs_gen_files.set_edit_path(full_doc_path, path)
+
+ summary = """\
+* [Overview](index.md)
+* [crazyflow](crazyflow/index.md)
+* Sim
+ * [sim](crazyflow/sim/index.md)
+ * [sim.data](crazyflow/sim/data.md)
+ * [sim.functional](crazyflow/sim/functional.md)
+ * [sim.physics](crazyflow/sim/physics.md)
+ * [sim.integration](crazyflow/sim/integration.md)
+ * [sim.sensors](crazyflow/sim/sensors.md)
+ * [sim.symbolic](crazyflow/sim/symbolic.md)
+ * [sim.visualize](crazyflow/sim/visualize.md)
+* Control
+ * [control](crazyflow/control/index.md)
+ * [control.mellinger](crazyflow/control/mellinger.md)
+* Environments
+ * [envs](crazyflow/envs/index.md)
+ * [envs.drone_env](crazyflow/envs/drone_env.md)
+ * [envs.figure_8_env](crazyflow/envs/figure_8_env.md)
+ * [envs.landing_env](crazyflow/envs/landing_env.md)
+ * [envs.reach_pos_env](crazyflow/envs/reach_pos_env.md)
+ * [envs.reach_vel_env](crazyflow/envs/reach_vel_env.md)
+ * [envs.norm_actions_wrapper](crazyflow/envs/norm_actions_wrapper.md)
+* [randomize](crazyflow/randomize/index.md)
+* [utils](crazyflow/utils.md)
+"""
+
+ with mkdocs_gen_files.open("api/SUMMARY.md", "w") as nav_file:
+ nav_file.write(summary)
diff --git a/docs/get-started/index.md b/docs/get-started/index.md
new file mode 100644
index 0000000..b40ff59
--- /dev/null
+++ b/docs/get-started/index.md
@@ -0,0 +1,6 @@
+# Get Started
+
+Install Crazyflow and run your first simulation in minutes.
+
+- [Installation](installation.md) — pip, pixi, GPU, and from-source options
+- [Quick Start](quick-start.md) — step-by-step walkthrough of the object-oriented API
diff --git a/docs/get-started/installation.md b/docs/get-started/installation.md
new file mode 100644
index 0000000..ffac14e
--- /dev/null
+++ b/docs/get-started/installation.md
@@ -0,0 +1,59 @@
+# Installation
+
+Select your installation method from the tabs below, then read the notes under each section for what it includes.
+
+=== "pip"
+
+ ```bash
+ pip install crazyflow
+ ```
+
+=== "pip + GPU"
+
+ ```bash
+ pip install "crazyflow[gpu]"
+ ```
+
+=== "pixi"
+
+ ```bash
+ git clone --recurse-submodules git@github.com:learnsyslab/crazyflow.git
+ cd crazyflow
+ pixi shell
+ ```
+
+=== "pixi + tests"
+
+ ```bash
+ git clone --recurse-submodules git@github.com:learnsyslab/crazyflow.git
+ cd crazyflow
+ pixi shell -e tests
+ ```
+
+---
+
+## GPU support
+
+JAX defaults to CPU-only execution. The `gpu` extra swaps in `jax[cuda12]`, enabling GPU execution. Setting `device="gpu"` in the `Sim` constructor then routes all computation through CUDA.
+
+!!! note
+ GPU support is only available on Linux x86-64.
+
+## Developer install
+
+[Pixi](https://pixi.sh/) creates a fully reproducible environment. This variant installs `crazyflow`, `drone_models`, and `drone_controllers` in editable mode from the `submodules/` folder. Any source change takes effect immediately without reinstalling. Recommended for contributors and researchers who modify the simulator.
+
+## Testing
+
+Adds `pytest` and `pytest-markdown-docs` for running the test suite and doc snippet tests.
+
+```bash
+pixi run tests # unit and integration tests
+pixi run test-docs # doc code snippet tests
+```
+
+## Verify the installation
+
+```bash
+python -c "from crazyflow.sim import Sim; sim = Sim(); sim.reset(); print('OK')"
+```
diff --git a/docs/get-started/quick-start.md b/docs/get-started/quick-start.md
new file mode 100644
index 0000000..26e9de5
--- /dev/null
+++ b/docs/get-started/quick-start.md
@@ -0,0 +1,140 @@
+# Quick Start
+
+This page walks through a complete minimal workflow: create a simulator, send a position command, step it forward, and read back the drone's state.
+
+## Create a simulator
+
+`Sim` is the top-level object. All configuration is provided at construction time: physics model, control mode, simulation frequency, number of parallel worlds, and number of drones per world.
+
+```python
+from crazyflow.sim import Sim
+
+sim = Sim(n_worlds=1, n_drones=1, freq=500)
+sim.reset()
+```
+
+`reset()` initialises all worlds to the default state: the drone is at the origin, upright, with zero velocity.
+
+## State and command
+
+The default control mode is `Control.state`. A state command is a 13-element vector that sets the desired position, velocity, acceleration, yaw, and body angular rates.
+
+| Index | Variable | Units |
+|---|---|---|
+| 0–2 | Position \(x, y, z\) | m |
+| 3–5 | Velocity \(\dot{x}, \dot{y}, \dot{z}\) | m/s |
+| 6–8 | Acceleration \(\ddot{x}, \ddot{y}, \ddot{z}\) | m/s² |
+| 9 | Yaw | rad |
+| 10 | Roll rate | rad/s |
+| 11 | Pitch rate | rad/s |
+| 12 | Yaw rate | rad/s |
+
+The command array has shape `(n_worlds, n_drones, 13)`.
+
+```python
+import numpy as np
+from crazyflow.sim import Sim
+from crazyflow.control import Control
+
+sim = Sim(n_worlds=1, n_drones=1, freq=500, control=Control.state)
+sim.reset()
+
+cmd = np.zeros((1, 1, 13), dtype=np.float32)
+cmd[0, 0, 2] = 0.5 # target height: 0.5 m
+```
+
+## Step the simulation
+
+`state_control` stages the command. `step` advances the simulation by the given number of physics steps. Calling `sim.step(sim.freq // sim.control_freq)` advances exactly one control cycle.
+
+```python
+import numpy as np
+from crazyflow.sim import Sim
+from crazyflow.control import Control
+
+sim = Sim(n_worlds=1, n_drones=1, freq=500, control=Control.state)
+sim.reset()
+
+cmd = np.zeros((1, 1, 13), dtype=np.float32)
+cmd[0, 0, 2] = 0.5
+
+for _ in range(10):
+ sim.state_control(cmd)
+ sim.step(sim.freq // sim.control_freq)
+```
+
+## Read back state
+
+All simulation state lives in `sim.data.states`. Arrays are indexed as `[world, drone, :]`.
+
+```python
+import numpy as np
+from crazyflow.sim import Sim
+from crazyflow.control import Control
+
+sim = Sim(n_worlds=1, n_drones=1, freq=500, control=Control.state)
+sim.reset()
+
+cmd = np.zeros((1, 1, 13), dtype=np.float32)
+cmd[0, 0, 2] = 0.5
+
+for _ in range(10):
+ sim.state_control(cmd)
+ sim.step(sim.freq // sim.control_freq)
+
+pos = sim.data.states.pos[0, 0] # (3,) — position in metres
+quat = sim.data.states.quat[0, 0] # (4,) — quaternion xyzw
+vel = sim.data.states.vel[0, 0] # (3,) — linear velocity m/s
+ang_vel = sim.data.states.ang_vel[0, 0] # (3,) — angular velocity rad/s
+```
+
+## Simulate multiple worlds
+
+Increase `n_worlds` to run independent simulations in a single batched call. All state arrays gain a leading world dimension.
+
+```python
+import numpy as np
+from crazyflow.sim import Sim
+from crazyflow.control import Control
+
+sim = Sim(n_worlds=4, n_drones=1, freq=500, control=Control.state)
+sim.reset()
+
+cmd = np.zeros((4, 1, 13), dtype=np.float32)
+cmd[:, 0, 2] = np.array([0.2, 0.4, 0.6, 0.8]) # different target heights per world
+
+for _ in range(10):
+ sim.state_control(cmd)
+ sim.step(sim.freq // sim.control_freq)
+
+pos = sim.data.states.pos[:, 0, :] # (4, 3) — position of drone 0 in each world
+```
+
+## Simulate multiple drones
+
+Increase `n_drones` to place multiple drones inside a single world. Each drone has its own independent state; all receive commands from the same `(n_worlds, n_drones, 13)` array.
+
+```python
+import numpy as np
+from crazyflow.sim import Sim
+from crazyflow.control import Control
+
+sim = Sim(n_worlds=1, n_drones=4, freq=500, control=Control.state)
+sim.reset()
+
+cmd = np.zeros((1, 4, 13), dtype=np.float32)
+cmd[0, :, 2] = np.array([0.2, 0.4, 0.6, 0.8]) # different height per drone
+
+for _ in range(10):
+ sim.state_control(cmd)
+ sim.step(sim.freq // sim.control_freq)
+
+pos = sim.data.states.pos[0, :, :] # (4, 3) — all 4 drones in world 0
+```
+
+## Next steps
+
+- [Object-Oriented API](../user-guide/oo-api.md) — all control modes, rendering, and reset
+- [Functional API](../user-guide/functional-api.md) — purely functional interface for use inside JAX transformations
+- [Physics Models](../user-guide/physics-models.md) — choosing between first-principles and fitted models
+- [Examples](../examples/index.md) — runnable scripts
diff --git a/docs/index.md b/docs/index.md
index 83531f9..956f9e4 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -1,35 +1,242 @@
# Crazyflow
-
+
-## Overview
+**Fast, parallelizable simulations of Crazyflie drones with JAX.**
-Crazyflow is a high-performance research simulator for Crazyflie‑style small quadrotors.
-Built on JAX and MuJoCo, it supports batched GPU execution, differentiable dynamics, and accurate, identified models — designed for reproducible experiments at scale.
+Crazyflow is a research simulator for Crazyflie-style quadrotors that runs millions of independent environments in parallel on CPU or GPU. It is built on JAX, exposes a differentiable dynamics pipeline, and ships identified models for the Crazyflie 2.x family.
-Audience: researchers working on control, learning, system identification, sim2real, multi-agent RL and swarm control for quadrotors.
+---
-## Highlights
+## Showcase
-- Modular simulation stack (physics, integrator, controller)
-- GPU-ready, batched execution for massive parallelism
-- Differentiable / autodiff-enabled dynamics
-- Support for analytical (identified) and data-driven models
-- Onboard-controller support and symbolic model matching
-- Extensible step pipeline (randomization, disturbances, custom hooks)
-- MuJoCo-based visualization and offscreen rendering
+
+
+
+
+
+
Training a rotor-level policy in 0.38s in Crazyflow.
+
+
+
+
Obstacle avoidance with MPPI using Crazyflow for rollouts.
+
+
+
+
SwarmGPT choreography in Crazyflow with 20 drones.
+
+
+
+
Racing controller testing and sim-to-real transfer.
+
+
+
+
Transfer learned policies and models from Crazyflow to real hardware.
+
+
+ 
+ 
+