Named pipelines

crazyflow/envs/drone_env.py

@@ -17,6 +17,7 @@
 from crazyflow.sim import Sim
 from crazyflow.sim.data import SimData
 from crazyflow.sim.physics import Physics
+from crazyflow.sim.pipeline import append_fn
 from crazyflow.utils import leaf_replace
 
 
@@ -101,7 +102,7 @@ def __init__(
         # Compile our domain randomization on reset into the sim pipeline
         if reset_randomization is None:
             reset_randomization = self._reset_randomization
-        self.sim.reset_pipeline += (reset_randomization,)
+        append_fn(self.sim.reset_pipeline, reset_randomization, name="reset_randomization")
         self.sim.build_reset_fn()
 
         self.n_substeps = self.sim.freq // self.freq
@@ -207,7 +208,7 @@ def obs(self) -> dict[str, Array]:
         return {k: v[:, 0, :] for k, v in obs.items()}
 
     @staticmethod
-    def _reset_randomization(data: SimData, mask: Array) -> SimData:
+    def _reset_randomization(data: SimData, _: SimData, mask: Array) -> SimData:
         """Randomize the initial position and velocity of the drones.
 
         This function will get compiled into the reset function of the simulation. Therefore, it

crazyflow/envs/figure_8_env.py

@@ -120,7 +120,7 @@ def _terminated(pos: Array) -> Array:
         return hit_floor | bounding_box
 
     @staticmethod
-    def _reset_randomization(data: SimData, mask: Array) -> SimData:
+    def _reset_randomization(data: SimData, _: SimData, mask: Array) -> SimData:
         """Randomize the initial position and velocity of the drones.
 
         This function will get compiled into the reset function of the simulation. Therefore, it

crazyflow/envs/reach_pos_env.py

@@ -101,7 +101,7 @@ def _sample_goal(
 
     @staticmethod
     def _reset_randomization(
-        data: SimData, mask: Array, pmin: Array, pmax: Array, vmin: float, vmax: float
+        data: SimData, _: SimData, mask: Array, pmin: Array, pmax: Array, vmin: float, vmax: float
     ) -> SimData:
         shape = (data.core.n_worlds, data.core.n_drones, 3)
         key, pos_key, vel_key = jax.random.split(data.core.rng_key, 3)

crazyflow/sim/pipeline.py

@@ -0,0 +1,119 @@
+"""Utility functions for safer access to OrderedDicts.
+
+We use OrderedDicts to store the functions of the step and reset pipelines. This preserves the order
+of stages. To allow users to modify the pipelines by names, safely adding elements without
+overwriting existing ones, and to improve readibility, we add utility functions for adding,
+replacing, and inserting stages.
+"""
+
+from __future__ import annotations
+
+from collections import OrderedDict
+from typing import Callable, TypeVar
+
+F = TypeVar("F", bound=Callable)
+
+
+def append_fn(pipeline: OrderedDict[str, F], fn: F, name: str | None = None):
+    """Append a function to the end of the pipeline.
+
+    Args:
+        pipeline: The pipeline to modify.
+        fn: The function to add.
+        name: Unique name of the new stage. Defaults to ``fn.__name__``.
+
+    Raises:
+        KeyError: If a stage with the same name already exists.
+    """
+    pipeline[_resolve_name(pipeline, name, fn)] = fn
+
+
+def prepend_fn(pipeline: OrderedDict[str, F], fn: F, name: str | None = None):
+    """Prepend a function to the beginning of the pipeline.
+
+    Args:
+        pipeline: The pipeline to modify.
+        fn: The function to add.
+        name: Unique name of the new stage. Defaults to ``fn.__name__``.
+
+    Raises:
+        KeyError: If a stage with the same name already exists.
+    """
+    items = list(pipeline.items())
+    items.insert(0, (_resolve_name(pipeline, name, fn), fn))
+    pipeline.clear()
+    pipeline.update(OrderedDict(items))
+
+
+def replace_fn(pipeline: OrderedDict[str, F], fn: F, name: str):
+    """Replace the function of an existing stage in the pipeline.
+
+    Args:
+        pipeline: The pipeline to modify.
+        fn: The new function for the stage.
+        name: Name of the stage to replace.
+
+    Raises:
+        KeyError: If no stage with this name exists.
+    """
+    if name not in pipeline:
+        raise KeyError(f"No stage named '{name}'. Available stages: {tuple(pipeline.keys())}")
+    pipeline[name] = fn
+
+
+def remove_fn(pipeline: OrderedDict[str, F], name: str):
+    """Remove a stage from the pipeline.
+
+    Args:
+        pipeline: The pipeline to modify.
+        name: Name of the stage to remove.
+    """
+    if name not in pipeline:
+        raise KeyError(f"No stage named '{name}'. Available stages: {tuple(pipeline.keys())}")
+    del pipeline[name]
+
+
+def insert_fn_before(pipeline: OrderedDict[str, F], anchor: str, fn: F, name: str | None = None):
+    """Insert a function directly before an existing stage in the pipeline.
+
+    Args:
+        pipeline: The pipeline to modify.
+        anchor: Name of an existing stage to insert before.
+        fn: The function to insert.
+        name: Unique name of the new stage. Defaults to ``fn.__name__``.
+    """
+    _insert(pipeline, anchor, fn, _resolve_name(pipeline, name, fn), offset=0)
+
+
+def insert_fn_after(pipeline: OrderedDict[str, F], anchor: str, fn: F, name: str | None = None):
+    """Insert a function directly after an existing stage in the pipeline.
+
+    Args:
+        pipeline: The pipeline to modify.
+        anchor: Name of an existing stage to insert after.
+        fn: The function to insert.
+        name: Unique name of the new stage. Defaults to ``fn.__name__``.
+    """
+    _insert(pipeline, anchor, fn, _resolve_name(pipeline, name, fn), offset=1)
+
+
+def _insert(pipeline: OrderedDict[str, F], anchor: str, fn: F, name: str | None, offset: int):
+    if anchor not in pipeline:
+        raise KeyError(f"No stage named '{anchor}'. Available stages: {tuple(pipeline.keys())}")
+    name = _resolve_name(pipeline, name, fn)
+    items = list(pipeline.items())
+    items.insert(list(pipeline.keys()).index(anchor) + offset, (name, fn))
+    pipeline.clear()
+    pipeline.update(OrderedDict(items))
+
+
+def _resolve_name(pipeline: OrderedDict[str, F], name: str | None, fn: F) -> str:
+    assert callable(fn), f"Expected fn to be callable, got {fn}"
+    if name is None:
+        name = getattr(fn, "__name__", None)
+    if name is None:  # E.g. functools.partial objects have no __name__
+        raise ValueError(f"{fn} has no __name__, an explicit name is required")
+    if name in pipeline:
+        raise KeyError(f"Pipeline stage '{name}' already exists. Names must be unique")
+    assert isinstance(name, str), f"Expected name to be str, got {name}"
+    return name

crazyflow/sim/sim.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from collections import OrderedDict
 from functools import partial, wraps
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Callable, ParamSpec, TypeVar
@@ -29,6 +30,7 @@
     so_rpy_rotor_drag_physics,
     so_rpy_rotor_physics,
 )
+from crazyflow.sim.pipeline import append_fn
 from crazyflow.utils import grid_2d, leaf_replace, pytree_replace
 
 if TYPE_CHECKING:
@@ -104,18 +106,23 @@ def __init__(
         self.default_data: SimData = self.build_default_data()
 
         # Build the simulation pipeline and overwrite the default _step implementation with it
-        self.reset_pipeline: tuple[Callable[[SimData, Array[bool] | None], SimData], ...] = tuple()
-        self.step_pipeline: tuple[Callable[[SimData], SimData], ...] = tuple()
+        self.reset_pipeline: OrderedDict[
+            str, Callable[[SimData, SimData, Array[bool] | None], SimData]
+        ] = OrderedDict()
+        append_fn(self.reset_pipeline, reset)
+
+        self.step_pipeline: OrderedDict[str, Callable[[SimData], SimData]] = OrderedDict()
         # The ``select_xxx_fn`` methods return functions, not the results of calling those
         # functions. They act as factories that produce building blocks for the construction of our
         # simulation pipeline.
-        self.step_pipeline += build_control_fns(self.control, self.physics)
-        physics_fn = select_physics_fn(self.physics)
-        self.step_pipeline += (select_integrate_fn(self.integrator, physics_fn),)
-        self.step_pipeline += (increment_steps,)
+        for fn in build_control_fns(self.control, self.physics):
+            append_fn(self.step_pipeline, fn)
+        integrate_fn = select_integrate_fn(self.integrator, select_physics_fn(self.physics))
+        append_fn(self.step_pipeline, integrate_fn, name="integration")
+        append_fn(self.step_pipeline, increment_steps)
         # We never drop below -0.001 (drones can't pass through the floor). We use -0.001 to
         # enable checks for negative z sign
-        self.step_pipeline += (clip_floor_pos,)
+        append_fn(self.step_pipeline, clip_floor_pos)
 
         self._reset = self.build_reset_fn()
         self._step = self.build_step_fn()
@@ -258,7 +265,10 @@ def build_step_fn(self) -> Callable[[SimData, int], SimData]:
             The pure JAX function that steps through the simulation. It takes the current SimData
             and the number of steps to simulate, and returns the updated SimData.
         """
-        pipeline = self.step_pipeline
+        # Snapshot the pipeline functions into a tuple. jax.jit traces lazily on the first call,
+        # so without a snapshot, modifying the pipeline between building and the first step would
+        # silently get compiled in.
+        pipeline = tuple(self.step_pipeline.values())
 
         # None is required by jax.lax.scan to unpack the tuple returned by single_step.
         def single_step(data: SimData, _: None) -> tuple[SimData, None]:
@@ -293,13 +303,15 @@ def build_reset_fn(self) -> Callable[[SimData, SimData, Array | None], SimData]:
             The pure JAX function that resets simulation data. It takes the current SimData, default
             SimData, and an optional mask for worlds to reset, returning the updated SimData.
         """
-        pipeline = self.reset_pipeline
+        # Snapshot the pipeline functions into a tuple. jax.jit traces lazily on the first call,
+        # so without a snapshot, modifying the pipeline between building and the first reset would
+        # silently get compiled in.
+        pipeline = tuple(self.reset_pipeline.values())
 
         @jax.jit
         def reset(data: SimData, default_data: SimData, mask: Array | None = None) -> SimData:
-            data = pytree_replace(data, default_data, mask)  # Does not overwrite rng_key
             for fn in pipeline:
-                data = fn(data, mask)
+                data = fn(data, default_data, mask)
             data = data.replace(core=data.core.replace(mjx_synced=False))  # Flag mjx data as stale
             return data
 
@@ -490,6 +502,11 @@ def select_integrate_fn(
     return partial(integrate_fn, deriv_fn=physics_fn)
 
 
+def reset(data: SimData, default_data: SimData, mask: Array | None = None) -> SimData:
+    """Reset the simulation data to the default data for the worlds specified by the mask."""
+    return pytree_replace(data, default_data, mask)  # Does not overwrite rng_key
+
+
 def increment_steps(data: SimData) -> SimData:
     """Increment the simulation steps."""
     return data.replace(core=data.core.replace(steps=data.core.steps + 1))

docs/api/index.md

@@ -7,6 +7,7 @@ This section is auto-generated from the Crazyflow source code using [mkdocstring
 | Module | Description |
 |---|---|
 | `crazyflow.sim` | Core `Sim` class and physics pipeline |
+| `crazyflow.sim.pipeline` | `OrderedDict`-based pipeline helpers (`append_fn`, `insert_fn_before`, `replace_fn`, etc.) |
 | `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 |

docs/index.md

@@ -207,7 +207,7 @@ Most simulators offer either vectorized environments for RL training or multi-dr
 
 Simulating the full Crazyflie firmware stack with GPU acceleration and differentiability is not possible with existing tools, so Crazyflow reimplements the entire dynamics and control stack in JAX. This gives accelerated, fully batchable simulation that runs on CPU and GPU without modification. Differentiability comes as a direct consequence: `jax.grad` works through physics, control, and integration without any manual gradient derivations, enabling gradient-based policy optimization, system identification, and sensitivity analysis out of the box.
 
-To make research possible rather than just evaluation, the simulator is designed to be fully open to modification. The step and reset pipelines are plain tuples of JAX functions. There are no fixed hooks or plugin interfaces — you splice in your own dynamics, disturbances, randomization, or reward shaping at any point, and the JIT compiler fuses everything into a single kernel.
+To make research possible rather than just evaluation, the simulator is designed to be fully open to modification. The step and reset pipelines are plain ordered dictionaries of pure JAX functions. Helper functions in `crazyflow.sim.pipeline` (`append_fn`, `prepend_fn`, `insert_fn_before`, `insert_fn_after`, `replace_fn`, `remove_fn`) let you safely add, reorder, and swap stages by name. There are no fixed hooks or plugin interfaces. You splice in your own dynamics, disturbances, randomization, or reward shaping at any point, and the JIT compiler fuses everything into a single kernel.
 
 For perception and collision, Crazyflow integrates MuJoCo and MJX. GUI rendering uses the MuJoCo viewer directly. Depth sensing, raycasting, and contact detection run through MJX, which keeps them batchable over worlds and compatible with JAX transformations.
 

docs/user-guide/gymnasium-envs.md

@@ -10,16 +10,17 @@ Crazyflow ships a set of [Gymnasium](https://gymnasium.farama.org/) vectorized e
 | `ReachPosEnv` | Reach a target position | pos, quat, vel, ang_vel, target | attitude |
 | `ReachVelEnv` | Match a target velocity | vel, ang_vel, target_vel | attitude |
 | `LandingEnv` | Land safely | pos, quat, vel, ang_vel | attitude |
-| `Figure8Env` | Follow a figure-8 trajectory | pos, quat, vel, ang_vel, phase | attitude |
+| `FigureEightEnv` | Follow a figure-8 trajectory | pos, quat, vel, ang_vel, phase | attitude |
 
 All environments run `num_envs` parallel instances backed by a single `Sim` with `n_worlds=num_envs`.
 
 ## Basic usage
 
-```{ .python notest }
-from crazyflow.envs import Figure8Env
+```python
+import gymnasium
+import crazyflow.envs  # noqa: F401 - registers the environments
 
-env = Figure8Env(num_envs=16, device="cpu")
+env = gymnasium.make_vec("DroneFigureEightTrajectory-v0", num_envs=16)
 obs, info = env.reset()
 
 for _ in range(500):
@@ -41,17 +42,18 @@ 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, mask) → SimData` function applied at reset |
+| `reset_randomization` | `None` | Optional `(SimData, SimData, mask) → SimData` function applied at reset (base `DroneEnv` only) |
 
 ## Action normalization
 
 `NormalizeActionsWrapper` rescales the action space to `[-1, 1]`, which simplifies policy learning:
 
-```{ .python notest }
-from crazyflow.envs import Figure8Env
-from crazyflow.envs.norm_actions_wrapper import NormalizeActionsWrapper
+```python
+import gymnasium
+import crazyflow.envs  # noqa: F401 - registers the environments
+from crazyflow.envs.norm_actions_wrapper import NormalizeActions
 
-env = NormalizeActionsWrapper(Figure8Env(num_envs=32))
+env = NormalizeActions(gymnasium.make_vec("DroneFigureEightTrajectory-v0", num_envs=32))
 obs, info = env.reset()
 action = env.action_space.sample()  # in [-1, 1]^4
 obs, reward, terminated, truncated, info = env.step(action)
@@ -60,22 +62,23 @@ env.close()
 
 ## Reset randomization
 
-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`:
+Pass a `reset_randomization` callable to vary initial conditions between episodes. The function receives `SimData`, the default `SimData`, and a boolean mask selecting the environments being reset, and must return updated `SimData`:
 
-```{ .python notest }
+```python
 import jax
-from crazyflow.envs import ReachPosEnv
+from crazyflow.envs.drone_env import DroneEnv
 from crazyflow.sim.data import SimData
 from crazyflow.utils import leaf_replace
 
-def randomize(data: SimData, mask: jax.Array | None) -> SimData:
+def randomize(data: SimData, default_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
     states = leaf_replace(data.states, mask, pos=data.states.pos + noise)
     return data.replace(states=states)
 
-env = ReachPosEnv(num_envs=64, reset_randomization=randomize)
+env = DroneEnv(num_envs=64, reset_randomization=randomize)
+env.close()
 ```
 
 ## Next steps

docs/user-guide/oo-api.md

@@ -196,25 +196,25 @@ sim.close()                           # close the viewer
 
 ## Domain randomization
 
-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:
+Define physical-parameter randomization as a reset pipeline stage. Each stage receives the restored `SimData`, the default `SimData`, and an optional world mask, and returns the randomized data:
 
 ```python
 import jax
 from jax import Array
 
 from crazyflow.sim import Sim
 from crazyflow.sim.data import SimData
+from crazyflow.sim.pipeline import append_fn
 from crazyflow.utils import leaf_replace
 
-@jax.jit
-def randomize_mass(data: SimData, mask: Array | None = None) -> SimData:
+def randomize_mass(data: SimData, default_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))
 
 sim = Sim(n_worlds=4, n_drones=1)
-sim.reset_pipeline = (randomize_mass,)
+append_fn(sim.reset_pipeline, randomize_mass)
 sim.build_reset_fn()
 sim.reset()  # randomizes every world
 ```