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 models for the Crazyflie 2.x family.
All models come from the drone-models library. Available configurations: cf2x_L250, cf2x_P250, cf2x_T350, cf21B_500, and any model returned by drone_models.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 model 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 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.
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
