Rework examples

docs/api/index.md

@@ -16,7 +16,6 @@ This section is auto-generated from the Crazyflow source code using [mkdocstring
 | `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/examples/index.md

@@ -1,6 +1,6 @@
 # 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.
+These runnable examples cover control, JAX transformations, pipeline extensions, rendering, contacts, and Gymnasium environments. Start with hover if you're new, or jump to the section that matches your use case.
 
 ---
 
@@ -9,11 +9,11 @@ These examples build on each other — each one introduces one new concept on to
 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"
+--8<-- "examples/control/hover.py"
 ```
 
 ```bash
-python examples/hover.py
+python examples/control/hover.py
 ```
 
 ---
@@ -23,7 +23,17 @@ python examples/hover.py
 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"
+--8<-- "examples/control/attitude.py"
+```
+
+---
+
+## Sampling-based MPC
+
+A sampling-based model predictive controller tracks a Lissajous curve while avoiding a grid of obstacles. It rolls out thousands of candidate control sequences in parallel using a reduced dynamics model, then applies the first action from a cost-weighted update of the best samples. The controller automatically uses a GPU when one is available and lowers the sample count on CPU.
+
+```bash
+python examples/control/sampling.py
 ```
 
 ---
@@ -33,17 +43,21 @@ Commanding roll, pitch, yaw, and collective thrust directly. This level bypasses
 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"
+--8<-- "examples/jax/gradient.py"
 ```
 
 ---
 
 ## Domain randomization
 
-Varying physical parameters per world at reset. Each world gets a slightly different mass, so identical commands produce diverging trajectories.
+Randomizing mass and inertia through the reset pipeline. An optional mask limits randomization to selected worlds.
 
 ```{ .python notest }
---8<-- "examples/randomize.py"
+--8<-- "examples/plugins/randomize.py"
+```
+
+```bash
+python examples/plugins/randomize.py
 ```
 
 ---
@@ -53,7 +67,7 @@ Varying physical parameters per world at reset. Each world gets a slightly diffe
 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"
+--8<-- "examples/plugins/disturbance.py"
 ```
 
 ---
@@ -67,11 +81,11 @@ Offscreen rendering returns RGB and depth images on every frame. The FPV camera
 </figure>
 
 ```{ .python notest }
---8<-- "examples/cameras.py"
+--8<-- "examples/rendering/cameras.py"
 ```
 
 ```bash
-python examples/cameras.py
+python examples/rendering/cameras.py
 ```
 
 ---
@@ -85,11 +99,11 @@ python examples/cameras.py
 </figure>
 
 ```{ .python notest }
---8<-- "examples/led_deck.py"
+--8<-- "examples/rendering/led_deck.py"
 ```
 
 ```bash
-python examples/led_deck.py
+python examples/rendering/led_deck.py
 ```
 
 ---
@@ -108,7 +122,7 @@ The default collision geometry is a sphere around the drone frame. `use_box_coll
 </div>
 
 ```{ .python notest }
---8<-- "examples/contacts.py"
+--8<-- "examples/contacts/contacts.py"
 ```
 
 ---
@@ -118,11 +132,11 @@ The default collision geometry is a sphere around the drone frame. `use_box_coll
 `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"
+--8<-- "examples/rendering/raycasting.py"
 ```
 
 ```bash
-python examples/raycasting.py
+python examples/rendering/raycasting.py
 ```
 
 ---
@@ -132,5 +146,5 @@ python examples/raycasting.py
 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"
+--8<-- "examples/environments/figure8.py"
 ```

docs/gen_ref_pages.py

@@ -60,7 +60,6 @@
     * [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)
 """
 

docs/user-guide/gymnasium-envs.md

@@ -41,7 +41,7 @@ All environments accept these common arguments:
 | `drone_model` | `"cf2x_L250"` | Drone configuration |
 | `freq` | 500 | Physics frequency, Hz |
 | `device` | `"cpu"` | `"cpu"` or `"gpu"` |
-| `reset_randomization` | `None` | Optional `SimData → SimData` function applied at reset |
+| `reset_randomization` | `None` | Optional `(SimData, mask) → SimData` function applied at reset |
 
 ## Action normalization
 
@@ -60,18 +60,20 @@ env.close()
 
 ## Reset randomization
 
-Pass a `reset_randomization` callable to vary initial conditions between episodes. The function receives `SimData` and a JAX random key and must return updated `SimData`:
+Pass a `reset_randomization` callable to vary initial conditions between episodes. The function receives `SimData` and a boolean mask selecting the environments being reset, and must return updated `SimData`:
 
 ```{ .python notest }
 import jax
-import jax.numpy as jnp
 from crazyflow.envs import ReachPosEnv
 from crazyflow.sim.data import SimData
+from crazyflow.utils import leaf_replace
 
-def randomize(data: SimData, key: jax.Array) -> SimData:
-    key, subkey = jax.random.split(key)
+def randomize(data: SimData, mask: jax.Array | None) -> SimData:
+    key, subkey = jax.random.split(data.core.rng_key)
+    data = data.replace(core=data.core.replace(rng_key=key))  # Make sure to update the rng_key
     noise = jax.random.normal(subkey, data.states.pos.shape) * 0.05
-    return data.replace(states=data.states.replace(pos=data.states.pos + noise))
+    states = leaf_replace(data.states, mask, pos=data.states.pos + noise)
+    return data.replace(states=states)
 
 env = ReachPosEnv(num_envs=64, reset_randomization=randomize)
 ```

docs/user-guide/oo-api.md

@@ -196,44 +196,31 @@ sim.close()                           # close the viewer
 
 ## Domain randomization
 
-Physical parameters can be randomized per-world using the `randomize` helpers:
+Define physical-parameter randomization as a reset pipeline stage. Each stage receives the restored `SimData` and an optional world mask, and returns the randomized data:
 
 ```python
 import jax
-import numpy as np
-from crazyflow.sim import Sim
-from crazyflow.randomize import randomize_inertia, randomize_mass
-
-sim = Sim(n_worlds=4, n_drones=1)
-sim.reset()
-
-nominal_mass = sim.data.params.mass
-noise = jax.random.normal(jax.random.key(0), nominal_mass.shape) * 2e-3
-randomize_mass(sim, nominal_mass + noise)
-
-nominal_J = sim.data.params.J
-J_noise = jax.random.normal(jax.random.key(1), nominal_J.shape) * 1e-6
-randomize_inertia(sim, nominal_J + J_noise)
-```
-
-To randomize only a subset of worlds, pass a boolean mask:
+from jax import Array
 
-```python
-import jax
-import numpy as np
 from crazyflow.sim import Sim
-from crazyflow.randomize import randomize_mass
+from crazyflow.sim.data import SimData
+from crazyflow.utils import leaf_replace
 
-sim = Sim(n_worlds=4, n_drones=1)
-sim.reset()
+@jax.jit
+def randomize_mass(data: SimData, mask: Array | None = None) -> SimData:
+    key, mass_key = jax.random.split(data.core.rng_key)
+    mass = data.params.mass + jax.random.normal(mass_key, data.params.mass.shape) * 2e-3
+    params = leaf_replace(data.params, mask, mass=mass)
+    return data.replace(params=params, core=data.core.replace(rng_key=key))
 
-import jax.numpy as jnp
-mask = jnp.array([True, True, False, False])  # only worlds 0 and 1
-nominal_mass = sim.data.params.mass
-noise = jax.random.normal(jax.random.key(0), nominal_mass.shape) * 2e-3
-randomize_mass(sim, nominal_mass + noise, mask=mask)
+sim = Sim(n_worlds=4, n_drones=1)
+sim.reset_pipeline = (randomize_mass,)
+sim.build_reset_fn()
+sim.reset()  # randomizes every world
 ```
 
+Passing a boolean mask to `sim.reset(mask=mask)` randomizes only the worlds being reset. See the [domain randomization example](../examples/index.md#domain-randomization) for mass and inertia randomization in a complete simulation.
+
 ## Next steps
 
 - [Functional API](functional-api.md) — run simulation inside `jax.jit` and `jax.grad`

docs/user-guide/pipelines.md

@@ -21,7 +21,7 @@ print(sim.step_pipeline)
 
 ## The reset pipeline
 
-`sim.reset_pipeline` is empty by default. When `sim.reset()` is called, it first restores `SimData` to the default state, then runs every function in the reset pipeline in order. The reset function signature is `(data: SimData, default_data: SimData, mask: Array | None) -> SimData`.
+`sim.reset_pipeline` is empty by default. When `sim.reset()` is called, it first restores `SimData` to the default state, then runs every function in the reset pipeline in order. Each reset stage has the signature `(data: SimData, mask: Array | None) -> SimData`.
 
 Populate `sim.reset_pipeline` to add episode-level randomization without modifying the default state.
 
@@ -36,36 +36,51 @@ To see how to modify the step pipeline with a stochastic disturbance, see the [D
 
 ## Modifying the reset pipeline
 
-Add a function to the reset pipeline to vary initial conditions between episodes. The function receives the freshly-restored `data`, the `default_data` it was restored from, and an optional `mask` of worlds that were reset.
+The [domain randomization example](../examples/index.md#domain-randomization) defines two reset stages that randomize mass and inertia. Each function receives the freshly restored `data` and an optional `mask` of worlds that were reset:
 
 ```{ .python notest }
 import jax
+import jax.numpy as jnp
+import numpy as np
+from jax import Array
+
+from crazyflow.control import Control
 from crazyflow.sim import Sim
 from crazyflow.sim.data import SimData
-from jax import Array
+from crazyflow.utils import leaf_replace
 
-def randomize_initial_pos(data: SimData, default_data: SimData, mask: Array | None) -> SimData:
-    key, subkey = jax.random.split(data.core.rng_key)
-    noise = jax.random.normal(subkey, data.states.pos.shape) * 0.1  # ±10 cm
-    return data.replace(
-        states=data.states.replace(pos=default_data.states.pos + noise),
-        core=data.core.replace(rng_key=key),
+
+@jax.jit
+def randomize_mass(data: SimData, mask: Array | None = None) -> SimData:
+    key, mass_key = jax.random.split(data.core.rng_key)
+    data = data.replace(core=data.core.replace(rng_key=key))  # Make sure to update the rng_key
+    mass = (
+        data.params.mass
+        + jax.random.normal(mass_key, (data.core.n_worlds, data.core.n_drones, 1)) * 2e-3
     )
+    return data.replace(params=leaf_replace(data.params, mask, mass=mass))
 
-sim = Sim(n_worlds=16)
-sim.reset_pipeline = (randomize_initial_pos,)
-sim.build_reset_fn()  # recompile
-sim.reset()
-# Each of the 16 worlds now starts at a slightly different position
-```
 
-Multiple stages can be chained; the output of each function is passed as input to the next:
+@jax.jit
+def randomize_inertia(data: SimData, mask: Array | None = None) -> SimData:
+    key, inertia_key = jax.random.split(data.core.rng_key)
+    data = data.replace(core=data.core.replace(rng_key=key))  # Make sure to update the rng_key
+    J = (
+        data.params.J
+        + jax.random.normal(inertia_key, (data.core.n_worlds, data.core.n_drones, 3, 3)) * 1e-8
+    )
+    return data.replace(params=leaf_replace(data.params, mask, J=J, J_inv=jnp.linalg.inv(J)))
 
-```{ .python notest }
-sim.reset_pipeline = (randomize_initial_pos, randomize_mass_fn, log_reset_fn)
+sim = Sim(n_worlds=3, n_drones=4, control=Control.state)
+sim.reset_pipeline = (randomize_mass, randomize_inertia)
 sim.build_reset_fn()
+
+mask = np.array([True, False, False])  # Only randomize the first world
+sim.reset(mask=mask)  # The mask is optional; omit it to reset and randomize all worlds
 ```
 
+Reset stages run in tuple order, with each stage receiving the output of the previous one. The mask ensures that parameter updates apply only to worlds selected by `sim.reset()`.
+
 ## Removing a stage
 
 Remove any stage by excluding it from the tuple. A common case is removing the floor clip when computing gradients through a trajectory that starts high above the ground:
@@ -80,7 +95,7 @@ sim.build_step_fn()
 
 ## Writing a custom stage
 
-A step pipeline function must have the signature `(SimData) -> SimData`. A reset pipeline function must have the signature `(SimData, SimData, Array | None) -> SimData`. Both must be pure JAX functions with no Python-level side effects, so they can be traced and compiled.
+A step pipeline function must have the signature `(SimData) -> SimData`. A reset pipeline function must have the signature `(SimData, Array | None) -> SimData`. Both must be pure JAX functions with no Python-level side effects, so they can be traced and compiled.
 
 ```{ .python notest }
 from crazyflow.sim.data import SimData