Crazyflow: Differentiable Drone Simulator

Crazyflow is organized into several functional modules that handle different aspects of the simulation pipeline, data management, and control. Key modules include:

• Simulation Core: crazyflow.sim provides the central Sim class and the simulation pipeline. crazyflow.sim.pipeline offers helpers for manipulating the pipeline (e.g., append_fn, insert_fn_before, replace_fn).
• Data Management: crazyflow.sim.data manages simulation state and parameters using pytrees (SimData, SimState, SimControls, SimParams, SimCore). crazyflow.sim.sharding handles data placement across multiple devices.
• Dynamics and Integration: crazyflow.dynamics defines dynamics implementations, while crazyflow.sim.integration provides integrators like Euler, RK4, and symplectic Euler.
• Sensors and Rendering: crazyflow.sim.sensors includes raycast depth and gaussian splat RGB sensors. crazyflow.sim.splat provides Gaussian splatting support via attach_splats and the SplatViewer.
• Control and Environments: crazyflow.control and crazyflow.control.mellinger manage control logic. crazyflow.envs provides vectorized environments compatible with Gymnasium.

The Crazyflow User Guide provides in-depth documentation for the simulator, covering data layouts, API interfaces, dynamics, control modes, and integration with external tools like MuJoCo and Gymnasium. Use the following topics to navigate the simulator's capabilities:

• Simulator Overview: Understand SimData layout, worlds, drones, and data conventions.
• Object-Oriented API: Learn about the Sim class, control methods, rendering, and reset procedures.
• Functional API: Use the purely functional interface designed for JAX transformations.
• Dynamics: Choose between first-principles and fitted dynamics.
• Control Modes: Implement state, attitude, body rate, force/torque, or rotor velocity control.
• Pipelines: Compose step and reset pipelines, including randomization and disturbances.
• The World Axis: Understand array batching over worlds, resets, and sharding.
• Visualization: Configure rendering modes, cameras, raycasting, and materials.
• MuJoCo Integration: Construct MJCF scenes, add objects, and manage sync internals.
• Gymnasium Environments: Create vectorized environments for Reinforcement Learning (RL) training.

The JAX dynamics pipeline updates sim.data but does not automatically update sim.mjx_data. To perform collision queries or rendering, the scene must be synchronized.

Crazyflow uses a mjx_synced flag in sim.data.core to avoid redundant computations.

• sim.step() or sim.reset() sets mjx_synced = False.
• sim.render() and sim.contacts() check this flag. If False, they trigger sync_sim2mjx(), which performs kinematics propagation, camera updates, and collision detection, then sets the flag to True.

Optimization Tip: Group all sim.render() and sim.contacts() calls together after a sim.step() to ensure the expensive synchronization happens only once per tick.

# Efficient: one sync for both calls
sim.step(5)
contacts = sim.contacts()     # sync runs here
sim.render(mode="rgb_array")  # flag already set, no second sync

# Inefficient: two syncs due to interleaved step
sim.step(5)
contacts = sim.contacts()     # sync runs here
sim.step(5)                   # flag cleared
sim.render(mode="rgb_array")  # sync runs again

The simulation advances using sim.step(n_steps), which moves the simulation forward by n_steps dynamics ticks.

Key Behaviors

• Controller Frequencies: Controllers (state, attitude, etc.) do not fire on every dynamics tick. They fire at their configured rates (state_freq, attitude_freq). Between controller ticks, the previously staged command is held.
• Efficiency: Calling sim.step(n_steps) with a larger n_steps is more efficient than multiple sim.step(1) calls because XLA compiles the loop into a single kernel.
• Recompilation: Changing the value of n_steps between calls triggers recompilation. Keep n_steps consistent within training or evaluation loops.
• Resetting: sim.reset() reinitializes worlds. You can pass a boolean mask of shape (n_worlds,) to reset only specific worlds (e.g., in RL training where episodes end at different times).

# Stage a command and advance 50 dynamics steps (controllers fire at their rate)
cmd = np.zeros((4, 1, 16), dtype=np.float32)
# ... set cmd values ...
sim.state_control(cmd)
sim.step(50)

# Reset only worlds 0 and 2, leaving 1 and 3 running
import jax.numpy as jnp
mask = jnp.array([True, False, True, False])
sim.reset(mask=mask)

Crazyflow uses two distinct pipelines to manage simulation state transitions: a step pipeline and a reset pipeline.

• Step Pipeline: An ordered dictionary of pure JAX functions that transform SimData at every simulation step. It handles control, command clipping, integration, and physics constraints.
• Reset Pipeline: An ordered dictionary of pure JAX functions used to restore SimData to a default state and apply episode-level randomization.

Stages in both pipelines are identified by unique string names, allowing you to address, insert, or remove them without relying on their position in the sequence. Both pipelines are compiled into a single jax.jit-cached function via sim.build_step_fn() or sim.build_reset_fn().

>>> from crazyflow.sim import Sim
>>> sim = Sim()
>>> print(tuple(sim.step_pipeline.keys()))
('attitude_controller', 'force_torque_controller', 'clip_rotor_vel_cmd', 'integration', 'clip_floor_pos', 'increment_steps')

Because Sim rebinds sim.data as a Python side effect that JAX cannot trace, you must use crazyflow.sim.functional for any code inside jit, grad, vmap, or scan. Use Sim only as the builder.

Important Usage Rules:

• Pipeline Updates: After editing a pipeline, you must call build_step_fn() or build_reset_fn() again. The builders snapshot the stages; subsequent edits are silently ignored.
• Static Steps: n_steps is static. Each distinct value of n_steps triggers a separate compilation. Build once and reuse.
• Disturbances: states.force and states.torque are inputs, not outputs. Writing to them applies a persistent disturbance as they are not cleared.
• Geometry/Contacts: Mutating state does not invalidate core.mjx_synced. Consequently, contacts() may query stale geometry. Set the flag to false manually if needed.
• Gradients: Gradients will vanish at clip_floor_pos (due to jnp.where on floor contact) and at any saturation bound.