Fast, parallelizable simulations of Crazyflie drones with JAX.
Crazyflow is a research simulator for Crazyflie-style quadrotors that runs millions of independent environments in parallel on CPU or GPU. It is built on JAX, exposes a differentiable dynamics pipeline, and ships identified dynamics for the Crazyflie 2.x family.
All drone configurations are bundled with crazyflow.dynamics. Available configurations: cf2x_L250, cf2x_P250, cf2x_T350, cf21B_500, and any drone returned by crazyflow.available_drones.
Most simulators offer either vectorized environments for RL training or multi-drone swarm simulation — rarely both, and rarely with accurate onboard flight dynamics for every agent. Crazyflow is built around both simultaneously. The entire simulator is structured around an n_worlds × n_drones batch dimension: n_worlds gives you massively parallel independent environments, and n_drones gives you full swarm simulation inside each one, each drone running its own accurate, identified flight dynamics and control stack. Scaling to millions of parallel instances requires no code changes.
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 dynamics, 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 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.
importnumpyasnpfromcrazyflow.simimportSimfromcrazyflow.controlimportControlsim=Sim(n_worlds=1,n_drones=1,control=Control.state)sim.reset()# State command: [x, y, z, vx, vy, vz, ax, ay, az, yaw, roll_rate, pitch_rate, yaw_rate]cmd=np.zeros((1,1,13),dtype=np.float32)cmd[0,0,2]=0.5# hover at 0.5 msim.state_control(cmd)sim.step(sim.freq//sim.control_freq)pos=sim.data.states.pos[0,0]# shape (3,) — position of world 0, drone 0
