Rework the docs

.github/workflows/docs.yml

@@ -37,7 +37,7 @@ jobs:
         locked: true
 
     - name: Build documentation
-      run: pixi run -e docs mkdocs build
+      run: pixi run -e docs docs-build
 
     - name: Setup Pages
       uses: actions/configure-pages@v3

.github/workflows/testing.yml

@@ -8,24 +8,40 @@ jobs:
       - uses: actions/checkout@v4
 
       - name: Install submodules
-        run: |
-          git submodule update --init --recursive
+        run: git submodule update --init --recursive
 
-      - name: Setup Pixi (installs pixi + caches envs) # https://github.com/marketplace/actions/setup-pixi
-        uses: prefix-dev/setup-pixi@v0.9.3            # pin the action version
+      - name: Setup Pixi
+        uses: prefix-dev/setup-pixi@v0.9.3
         with:
-          pixi-version: v0.67.0                       # pin the pixi binary version (optional)
-          cache: true                                 # enable caching of installed envs
-          # only write new caches on main pushes (TODO: Enable)
-          # cache-write: ${{ github.event_name == 'push' && github.ref_name == 'main' }}
-          # ensure the 'tests' environment(s) are installed
+          pixi-version: v0.67.0
+          cache: true
+          cache-write: ${{ github.event_name == 'push' && github.ref_name == 'main' }}
           environments: tests
-          # don't activate env (we'll call pixi run -e test explicitly)
           activate-environment: false
-          # prefer using existing lockfile if present (faster, deterministic)
           locked: true
 
       - name: Verify pixi and run tests
         run: |
           pixi --version
-          pixi run -e tests tests
+          pixi run -e tests tests
+
+  test-docs:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Install submodules
+        run: git submodule update --init --recursive
+
+      - name: Setup Pixi
+        uses: prefix-dev/setup-pixi@v0.9.3
+        with:
+          pixi-version: v0.67.0
+          cache: true
+          cache-write: ${{ github.event_name == 'push' && github.ref_name == 'main' }}
+          environments: tests
+          activate-environment: false
+          locked: true
+
+      - name: Run doc tests
+        run: pixi run -e tests test-docs

.gitignore

@@ -21,7 +21,7 @@ benchmark/data
 .pixi
 *.egg-info
 
-# MkDocs build output
+# ProperDocs build output
 site/
 # Temporary files
 .DS_Store

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
+## 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.

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.

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.
+
+<figure class="example-media">
+  <img src="../img/examples/cameras.gif" alt="RGB and depth camera outputs from a Crazyflow drone simulation">
+</figure>
+
+```{ .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.
+
+<figure class="example-media example-media--compact">
+  <img src="../img/examples/led_decks.png" alt="Crazyflow drones with runtime-controlled LED deck materials">
+</figure>
+
+```{ .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.
+
+<div class="example-media-grid example-media-grid--contacts">
+  <figure>
+    <img src="../img/examples/contact_sphere.png" alt="Contact query visualization using the default sphere collision geometry">
+  </figure>
+  <figure>
+    <img src="../img/examples/contact_box.png" alt="Contact query visualization using the oriented box collision geometry">
+  </figure>
+</div>
+
+```{ .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"
+```