diff --git a/crazyflow/envs/drone_env.py b/crazyflow/envs/drone_env.py index a69bd74..7efebbc 100644 --- a/crazyflow/envs/drone_env.py +++ b/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 diff --git a/crazyflow/envs/figure_8_env.py b/crazyflow/envs/figure_8_env.py index 8276ee6..1416bed 100644 --- a/crazyflow/envs/figure_8_env.py +++ b/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 diff --git a/crazyflow/envs/reach_pos_env.py b/crazyflow/envs/reach_pos_env.py index 5d3e3b3..be27dc1 100644 --- a/crazyflow/envs/reach_pos_env.py +++ b/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) diff --git a/crazyflow/sim/pipeline.py b/crazyflow/sim/pipeline.py new file mode 100644 index 0000000..8e21a83 --- /dev/null +++ b/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 diff --git a/crazyflow/sim/sim.py b/crazyflow/sim/sim.py index 644f62b..2ac365e 100644 --- a/crazyflow/sim/sim.py +++ b/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)) diff --git a/docs/api/index.md b/docs/api/index.md index e1c9443..969e27c 100644 --- a/docs/api/index.md +++ b/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 | diff --git a/docs/index.md b/docs/index.md index b7555b9..fe87258 100644 --- a/docs/index.md +++ b/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. diff --git a/docs/user-guide/gymnasium-envs.md b/docs/user-guide/gymnasium-envs.md index d0e8a38..35e7aa0 100644 --- a/docs/user-guide/gymnasium-envs.md +++ b/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 diff --git a/docs/user-guide/oo-api.md b/docs/user-guide/oo-api.md index b0cfc6d..8bf8e5a 100644 --- a/docs/user-guide/oo-api.md +++ b/docs/user-guide/oo-api.md @@ -196,7 +196,7 @@ 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 @@ -204,17 +204,17 @@ 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 ``` diff --git a/docs/user-guide/pipelines.md b/docs/user-guide/pipelines.md index e808d26..fde3dbe 100644 --- a/docs/user-guide/pipelines.md +++ b/docs/user-guide/pipelines.md @@ -1,33 +1,59 @@ # Pipelines -Crazyflow has two pipelines, one for stepping and one for resetting, each a tuple of pure JAX functions that transform `SimData`. Both are constructed at `Sim` initialisation and compiled into a single `jax.jit`-cached function by `build_step_fn()` / `build_reset_fn()`. You can modify either pipeline by editing the tuple and calling the corresponding build function. +Crazyflow has two pipelines, one for stepping and one for resetting. Each is an ordered dictionary of pure JAX functions that transform `SimData`. Stages are keyed by a unique string name so they can be addressed directly without relying on positional indices. + +`crazyflow.sim.pipeline` provides helper functions for safely modifying a pipeline: + +| Function | Description | +|---|---| +| `append_fn(pipeline, fn, name=None)` | Add a stage at the end | +| `prepend_fn(pipeline, fn, name=None)` | Add a stage at the beginning | +| `insert_fn_before(pipeline, anchor, fn, name=None)` | Insert before a named stage | +| `insert_fn_after(pipeline, anchor, fn, name=None)` | Insert after a named stage | +| `replace_fn(pipeline, fn, name)` | Swap the function of an existing stage | +| `remove_fn(pipeline, name)` | Remove a stage by name | + +All helpers raise `KeyError` on duplicate or missing names. Stage names default to `fn.__name__`. Pass an explicit `name` for anonymous callables such as `functools.partial` objects. + +Both pipelines are constructed at `Sim` initialisation and compiled into a single `jax.jit`-cached function by `build_step_fn()` / `build_reset_fn()`. Modify the pipeline and then call the corresponding build function to recompile. ## The step pipeline `sim.step_pipeline` contains four stages by default: 1. **Control functions** — convert the staged command through the control hierarchy (state → attitude → force/torque → rotor velocities, depending on the selected mode) -2. **Integrator** — advance the ODE one physics step (Euler, RK4, or symplectic Euler) -3. **Step counter** — increment `data.core.steps` -4. **Floor clip** — prevent drones from passing through the floor +2. **Integrator** (`integration`) — advance the ODE one physics step (Euler, RK4, or symplectic Euler) +3. **Step counter** (`increment_steps`) — increment `data.core.steps` +4. **Floor clip** (`clip_floor_pos`) — prevent drones from passing through the floor ```python from crazyflow.sim import Sim sim = Sim() -print(sim.step_pipeline) -# (, , , ) +print(tuple(sim.step_pipeline.keys())) +# ('step_attitude_controller', 'step_force_torque_controller', 'integration', 'increment_steps', 'clip_floor_pos') ``` ## 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. Each reset stage has the signature `(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, default_data: SimData, mask: Array | None) -> SimData`. The `default_data` argument holds the freshly-restored default state, which is useful for selectively reverting fields. Populate `sim.reset_pipeline` to add episode-level randomization without modifying the default state. ## Modifying the step pipeline -Insert or remove stages by slicing and concatenating the tuple. +Stages are addressed by name. Use `insert_fn_before` / `insert_fn_after` to place a function relative to an existing stage, `append_fn` to add it at the end, and `replace_fn` to swap a stage's implementation. New stages are named after the function's `__name__` unless an explicit name is given; names must be unique within a pipeline. + +```{ .python continuation } +from crazyflow.sim.data import SimData +from crazyflow.sim.pipeline import insert_fn_before + +def disturbance_fn(data: SimData) -> SimData: + return data.replace(states=data.states.replace(vel=data.states.vel + 1e-5)) + +insert_fn_before(sim.step_pipeline, "integration", disturbance_fn) +sim.build_step_fn() # recompile +``` !!! warning Always call `sim.build_step_fn()` after modifying `sim.step_pipeline`. Without it, `sim.step()` still runs the previously compiled kernel and silently ignores your changes. @@ -36,46 +62,50 @@ To see how to modify the step pipeline with a stochastic disturbance, see the [D ## Modifying the reset pipeline -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: +Add a function to the reset pipeline to vary initial conditions between episodes. The function receives the freshly-restored `data` and an optional `mask` of worlds that were reset. -```{ .python notest } +```python 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 crazyflow.utils import leaf_replace - - -@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 +from crazyflow.sim.pipeline import append_fn + +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=data.states.pos + noise), + core=data.core.replace(rng_key=key), ) - return data.replace(params=leaf_replace(data.params, mask, mass=mass)) +sim = Sim(n_worlds=16) +append_fn(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 +``` -@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 +```{ .python continuation } +def randomize_vel(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.vel.shape) * 0.05 + return data.replace( + states=data.states.replace(vel=data.states.vel + noise), + core=data.core.replace(rng_key=key), ) - return data.replace(params=leaf_replace(data.params, mask, J=J, J_inv=jnp.linalg.inv(J))) -sim = Sim(n_worlds=3, n_drones=4, control=Control.state) -sim.reset_pipeline = (randomize_mass, randomize_inertia) +def log_reset(data: SimData, default_data: SimData, mask: Array | None) -> SimData: + return data # a pure pass-through stage, e.g. a hook for metrics + +for fn in (randomize_vel, log_reset): + append_fn(sim.reset_pipeline, fn) sim.build_reset_fn() -mask = np.array([True, False, False]) # Only randomize the first world +mask = np.zeros(16, dtype=bool) +mask[0] = True # Only randomize the first world sim.reset(mask=mask) # The mask is optional; omit it to reset and randomize all worlds ``` @@ -83,26 +113,33 @@ Reset stages run in tuple order, with each stage receiving the output of the pre ## 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: +Remove any stage by name. A common case is removing the floor clip when computing gradients through a trajectory that starts high above the ground: -```{ .python notest } +```python from crazyflow.sim import Sim +from crazyflow.sim.pipeline import remove_fn sim = Sim() -sim.step_pipeline = sim.step_pipeline[:-1] # drop clip_floor_pos +remove_fn(sim.step_pipeline, "clip_floor_pos") 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, 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, SimData, Array | None) -> SimData` where the second argument is the default (freshly-restored) data. Both must be pure JAX functions with no Python-level side effects, so they can be traced and compiled. -```{ .python notest } +```python +from crazyflow.sim import Sim from crazyflow.sim.data import SimData +from crazyflow.sim.pipeline import append_fn def my_step_stage(data: SimData) -> SimData: # JAX operations only — return updated data - return data.replace(...) + return data.replace(states=data.states.replace(pos=data.states.pos + 0.01)) + +sim = Sim() +append_fn(sim.step_pipeline, my_step_stage) +sim.build_step_fn() ``` ## Next steps diff --git a/docs/user-guide/sim-overview.md b/docs/user-guide/sim-overview.md index e47885a..54cf3d1 100644 --- a/docs/user-guide/sim-overview.md +++ b/docs/user-guide/sim-overview.md @@ -81,7 +81,7 @@ sim.step(sim.freq // sim.control_freq) # 500 // 100 = 5 physics steps, controll ## The step and reset pipelines -Each call to `sim.step()` runs `sim.step_pipeline`, a tuple of pure JAX functions that transforms `SimData`. By default it contains the control conversion functions, the numerical integrator, a step counter, and a floor clip. Similarly, `sim.reset_pipeline` is applied during `sim.reset()` and is empty by default. +Each call to `sim.step()` runs `sim.step_pipeline`, an ordered collection of named, pure JAX functions that transform `SimData`. By default it contains the control conversion functions, the numerical integrator, a step counter, and a floor clip. Similarly, `sim.reset_pipeline` is applied during `sim.reset()` and is empty by default. Both pipelines can be extended with custom functions for disturbances, domain randomization, or logging without modifying the core simulator. diff --git a/examples/control/sampling.py b/examples/control/sampling.py index 5e5e539..3fde3ea 100644 --- a/examples/control/sampling.py +++ b/examples/control/sampling.py @@ -32,7 +32,7 @@ # Simulation configuration DRONE_MODEL = "cf21B_500" -DURATION = 18.0 +DURATION = 10.0 FPS = 60 RENDER = True @@ -286,8 +286,7 @@ def main() -> None: action_low=action_low, action_high=action_high, noise_sigma=noise_sigma, - ), - device=controller_device, + ) ) mean_controls = jax.device_put(jnp.broadcast_to(hover_cmd, (N, 4)), controller_device) key = jax.device_put(jax.random.key(0), controller_device) diff --git a/examples/jax/gradient.py b/examples/jax/gradient.py index c175365..14b1b1e 100644 --- a/examples/jax/gradient.py +++ b/examples/jax/gradient.py @@ -12,8 +12,13 @@ def main(): sim = Sim(control=Control.attitude, physics=Physics.first_principles, attitude_freq=50) # Remove clipping floor function which kills gradients - sim.step_pipeline = sim.step_pipeline[:-1] sim_step = sim.build_step_fn() + # If the drone starts on the floor, the gradient gets killed by the floor clipping function. We + # thus start in the air to avoid zero gradients. Alternatively, we could also remove the floor + # clipping function + sim.data = sim.data.replace( + states=sim.data.states.replace(pos=sim.data.states.pos.at[..., 2].set(0.5)) + ) def step(cmd: NDArray, data: SimData) -> jax.Array: data = data.replace( diff --git a/examples/plugins/action_delay.py b/examples/plugins/action_delay.py index 5056d51..cc438d5 100644 --- a/examples/plugins/action_delay.py +++ b/examples/plugins/action_delay.py @@ -11,6 +11,7 @@ import numpy as np from crazyflow import Sim +from crazyflow.sim.pipeline import prepend_fn from crazyflow.sim.visualize import draw_points if TYPE_CHECKING: @@ -60,7 +61,7 @@ def main(): # resets. custom_data = {"queued_actions": jnp.zeros((delay_steps, 1, 1, 4))} sim.data = sim.data.replace(plugins=sim.data.plugins | custom_data) - sim.step_pipeline = (action_delay,) + sim.step_pipeline + prepend_fn(sim.step_pipeline, action_delay, name="action_delay") sim.build_default_data() sim.build_step_fn() diff --git a/examples/plugins/disturbance.py b/examples/plugins/disturbance.py index 122edca..bd692b8 100644 --- a/examples/plugins/disturbance.py +++ b/examples/plugins/disturbance.py @@ -6,6 +6,7 @@ from crazyflow.sim import Sim from crazyflow.sim.data import SimData +from crazyflow.sim.pipeline import insert_fn_before os.environ["SCIPY_ARRAY_API"] = "1" @@ -42,9 +43,9 @@ def main(plot: bool = False): # Second run # We insert the disturbance function into the step pipeline before the integration step. You can - # inspect the step pipeline with + # inspect the step pipeline stages with # print(sim.step_pipeline) - sim.step_pipeline = sim.step_pipeline[:2] + (disturbance_fn,) + sim.step_pipeline[2:] + insert_fn_before(sim.step_pipeline, "integration", disturbance_fn) sim.build_step_fn() pos_disturbed, quat_disturbed = [], [] sim.reset() diff --git a/examples/plugins/estimation.py b/examples/plugins/estimation.py index 7f83b81..088b1be 100644 --- a/examples/plugins/estimation.py +++ b/examples/plugins/estimation.py @@ -13,6 +13,7 @@ from drone_models.transform import motor_force2rotor_vel from crazyflow import Sim +from crazyflow.sim.pipeline import insert_fn_after, prepend_fn from crazyflow.sim.visualize import draw_line, draw_points if TYPE_CHECKING: @@ -166,16 +167,10 @@ def main(noisy: bool = False, render: bool = True) -> None: plugins=sim.data.plugins | plugins, core=sim.data.core.replace(rng_key=key) ) - controllers = sim.step_pipeline[:-3] - integration = sim.step_pipeline[-3:] - sim.step_pipeline = ( - simulate_uwb, - estimate_state, - use_estimate_for_control, - *controllers, - restore_ground_truth, - *integration, - ) + prepend_fn(sim.step_pipeline, simulate_uwb) + insert_fn_after(sim.step_pipeline, "simulate_uwb", estimate_state) + insert_fn_after(sim.step_pipeline, "estimate_state", use_estimate_for_control) + insert_fn_after(sim.step_pipeline, "step_force_torque_controller", restore_ground_truth) sim.build_default_data() sim.build_step_fn() diff --git a/examples/plugins/randomize.py b/examples/plugins/randomize.py index 710df8d..0f1e281 100644 --- a/examples/plugins/randomize.py +++ b/examples/plugins/randomize.py @@ -6,11 +6,12 @@ from crazyflow.control import Control from crazyflow.sim import Sim from crazyflow.sim.data import SimData +from crazyflow.sim.pipeline import append_fn from crazyflow.utils import grid_2d, leaf_replace @jax.jit -def randomize_mass(data: SimData, mask: Array | None = None) -> SimData: +def randomize_mass(data: SimData, _: 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 dist = jax.random.normal(mass_key, (data.core.n_worlds, data.core.n_drones, 1)) * 2e-3 @@ -19,7 +20,7 @@ def randomize_mass(data: SimData, mask: Array | None = None) -> SimData: @jax.jit -def randomize_inertia(data: SimData, mask: Array | None = None) -> SimData: +def randomize_inertia(data: SimData, _: 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 dist = jax.random.normal(inertia_key, (data.core.n_worlds, data.core.n_drones, 3, 3)) * 1e-8 @@ -29,7 +30,8 @@ def randomize_inertia(data: SimData, mask: Array | None = None) -> SimData: def main(): sim = Sim(n_worlds=3, n_drones=4, control=Control.state) - sim.reset_pipeline = (randomize_mass, randomize_inertia) + append_fn(sim.reset_pipeline, randomize_mass) + append_fn(sim.reset_pipeline, randomize_inertia) sim.build_reset_fn() mask = np.array([True, False, False]) # Only randomize the first world diff --git a/tests/integration/test_disturbance.py b/tests/integration/test_disturbance.py index b505ad4..dc2d691 100644 --- a/tests/integration/test_disturbance.py +++ b/tests/integration/test_disturbance.py @@ -4,6 +4,7 @@ from crazyflow.sim import Physics, Sim from crazyflow.sim.data import SimData +from crazyflow.sim.pipeline import insert_fn_after def disturbance_fn(data: SimData) -> SimData: @@ -29,7 +30,7 @@ def test_disturbance(physics: Physics): pos.append(sim.data.states.pos[0, 0]) sim.reset() - sim.step_pipeline = sim.step_pipeline[:1] + (disturbance_fn,) + sim.step_pipeline[1:] + insert_fn_after(sim.step_pipeline, "step_state_controller", disturbance_fn) sim.build_step_fn() for _ in range(n_steps): sim.state_control(control) diff --git a/tests/integration/test_symbolic.py b/tests/integration/test_symbolic.py index 3407ad7..cac9101 100644 --- a/tests/integration/test_symbolic.py +++ b/tests/integration/test_symbolic.py @@ -6,6 +6,7 @@ from crazyflow.sim import Sim from crazyflow.sim.data import SimState from crazyflow.sim.physics import Physics +from crazyflow.sim.pipeline import remove_fn from crazyflow.sim.symbolic import symbolic_from_sim @@ -22,7 +23,7 @@ def test_attitude_symbolic(physics: Physics, freq: int): pytest.skip(f"Physics mode {physics} not yet implemented") sim = Sim(physics=physics, freq=freq) - sim.step_pipeline = sim.step_pipeline[:-1] # Remove clip floor from step pipeline + remove_fn(sim.step_pipeline, "clip_floor_pos") # Remove clip floor from step pipeline X_dot, X, U, Y = symbolic_from_sim(sim) fd = cs.integrator("fd", "cvodes", {"x": X, "p": U, "ode": X_dot}, 0, 1 / freq) diff --git a/tests/unit/test_pipeline.py b/tests/unit/test_pipeline.py new file mode 100644 index 0000000..3da158f --- /dev/null +++ b/tests/unit/test_pipeline.py @@ -0,0 +1,116 @@ +"""Unit tests for pipeline utility helpers based on OrderedDict.""" + +from collections import OrderedDict +from functools import partial + +import jax.numpy as jnp +import pytest + +from crazyflow.sim import Sim +from crazyflow.sim.pipeline import ( + append_fn, + insert_fn_after, + insert_fn_before, + prepend_fn, + remove_fn, + replace_fn, +) + + +def fn_a(x: int) -> int: + return x + 1 + + +def fn_b(x: int) -> int: + return x * 2 + + +def fn_c(x: int) -> int: + return x - 3 + + +@pytest.mark.unit +def test_append_and_prepend_order_and_names(): + pipeline: OrderedDict[str, object] = OrderedDict() + append_fn(pipeline, fn_a) + append_fn(pipeline, fn_b) + prepend_fn(pipeline, fn_c) + assert tuple(pipeline.keys()) == ("fn_c", "fn_a", "fn_b") + assert tuple(pipeline.values()) == (fn_c, fn_a, fn_b) + + +@pytest.mark.unit +def test_insert_before_after(): + pipeline = OrderedDict([("fn_a", fn_a), ("fn_b", fn_b)]) + insert_fn_before(pipeline, "fn_b", fn_c) + assert tuple(pipeline.keys()) == ("fn_a", "fn_c", "fn_b") + remove_fn(pipeline, "fn_c") + insert_fn_after(pipeline, "fn_a", fn_c) + assert tuple(pipeline.keys()) == ("fn_a", "fn_c", "fn_b") + + +@pytest.mark.unit +def test_replace_keeps_position_and_name(): + pipeline = OrderedDict([("fn_c", fn_c), ("fn_a", fn_a), ("fn_b", fn_b)]) + replace_fn(pipeline, fn_b, "fn_a") + assert tuple(pipeline.keys()) == ("fn_c", "fn_a", "fn_b") + assert tuple(pipeline.values())[1] == fn_b + + +@pytest.mark.unit +def test_remove_existing_stage(): + pipeline = OrderedDict([("fn_a", fn_a), ("fn_b", fn_b)]) + remove_fn(pipeline, "fn_a") + assert tuple(pipeline.keys()) == ("fn_b",) + + +@pytest.mark.unit +def test_unique_names(): + pipeline = OrderedDict([("fn_a", fn_a)]) + with pytest.raises(KeyError, match="already exists"): + append_fn(pipeline, fn_a) + with pytest.raises(KeyError, match="already exists"): + insert_fn_after(pipeline, "fn_a", fn_b, name="fn_a") + + +@pytest.mark.unit +def test_anonymous_fn_requires_name(): + pipeline: OrderedDict[str, object] = OrderedDict() + with pytest.raises(ValueError, match="explicit name"): + append_fn(pipeline, partial(fn_a)) + append_fn(pipeline, partial(fn_a), name="fn_a_partial") + assert tuple(pipeline.keys()) == ("fn_a_partial",) + + +@pytest.mark.unit +def test_missing_stage(): + pipeline = OrderedDict([("fn_a", fn_a)]) + with pytest.raises(KeyError, match="No stage named 'missing'"): + insert_fn_before(pipeline, "missing", fn_b) + with pytest.raises(KeyError, match="No stage named 'missing'"): + replace_fn(pipeline, fn_b, "missing") + with pytest.raises(KeyError, match="No stage named 'missing'"): + remove_fn(pipeline, "missing") + + +@pytest.mark.unit +def test_explicit_name_with_append(): + pipeline = OrderedDict([("fn_a", fn_a)]) + append_fn(pipeline, fn_c, name="fn_c_named") + assert tuple(pipeline.keys()) == ("fn_a", "fn_c_named") + + +@pytest.mark.unit +def test_pipeline_snapshot(): + """Modifications after building must not leak into the compiled step function.""" + sim = Sim() + sim.build_step_fn() + steps = sim.data.core.steps + + def fail_stage(data): # noqa: ANN001, ANN202 + raise AssertionError("Stage added after build_step_fn must not be traced") + + append_fn(sim.step_pipeline, fail_stage) + sim.step() # Traces the compiled function on first call, must use the snapshot + assert jnp.all(sim.data.core.steps == steps + 1) + sim.close() diff --git a/tests/unit/test_plugins.py b/tests/unit/test_plugins.py index ec38429..f31950a 100644 --- a/tests/unit/test_plugins.py +++ b/tests/unit/test_plugins.py @@ -13,6 +13,7 @@ import pytest from crazyflow.sim import Sim +from crazyflow.sim.pipeline import append_fn if TYPE_CHECKING: from crazyflow.sim.data import SimData @@ -56,7 +57,7 @@ def test_plugin_data_changes(n_worlds: int): """Plugin data changes across simulation steps.""" sim = Sim(n_worlds=n_worlds) sim.data = sim.data.replace(plugins={"counter": jnp.zeros((n_worlds, 1), dtype=jnp.int32)}) - sim.step_pipeline = sim.step_pipeline + (counter_plugin,) + append_fn(sim.step_pipeline, counter_plugin) sim.build_step_fn() n_steps = 7 sim.step(n_steps) @@ -71,7 +72,7 @@ def test_plugin_resets(): """sim.reset() restores the counter to its default value (0).""" sim = Sim() sim.data = sim.data.replace(plugins={"counter": jnp.zeros((1, 1), dtype=jnp.int32)}) - sim.step_pipeline = sim.step_pipeline + (counter_plugin,) + append_fn(sim.step_pipeline, counter_plugin) sim.build_default_data() sim.build_step_fn() sim.step(10) @@ -87,7 +88,7 @@ def test_plugin_masked_reset(): n_worlds = 2 sim = Sim(n_worlds=n_worlds) sim.data = sim.data.replace(plugins={"counter": jnp.zeros((n_worlds, 1), dtype=jnp.int32)}) - sim.step_pipeline = sim.step_pipeline + (counter_plugin,) + append_fn(sim.step_pipeline, counter_plugin) sim.build_default_data() sim.build_step_fn() sim.step(10) @@ -110,7 +111,8 @@ def test_chained_plugins(): "accumulated": jnp.zeros((1, 1), dtype=jnp.int32), } ) - sim.step_pipeline = sim.step_pipeline + (counter_plugin, accumulater_plugin) + append_fn(sim.step_pipeline, counter_plugin) + append_fn(sim.step_pipeline, accumulater_plugin) sim.build_default_data() sim.build_step_fn() n_steps = 3