Metadata-Version: 2.4
Name: crazyflow
Version: 0.2.1
Summary: Fast, parallelizable simulations of drones with JAX and MuJoCo.
Author: Martin Schuck, Marcel Rath
License: The MIT License (MIT)
        Copyright © 2025 Martin Schuck
        
        Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
        
        The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
        
        THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
        
        
Project-URL: Homepage, https://github.com/learnsyslab/crazyflow
Project-URL: Repository, https://github.com/learnsyslab/crazyflow
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python
Classifier: Programming Language :: Python :: 3
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Education
Classifier: Intended Audience :: Science/Research
Requires-Python: <3.14,>=3.11
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: numpy>=2.0.0
Requires-Dist: scipy>=1.17.0
Requires-Dist: jax>=0.7.0
Requires-Dist: mujoco>=3.3.0
Requires-Dist: mujoco-mjx>=3.3.0
Requires-Dist: gymnasium[mujoco]>=1.2.0
Requires-Dist: imageio
Requires-Dist: einops
Requires-Dist: flax
Requires-Dist: ml_collections
Requires-Dist: casadi
Requires-Dist: array-api-compat
Requires-Dist: array-api-extra
Provides-Extra: gpu
Requires-Dist: jax[cuda12]; extra == "gpu"
Provides-Extra: benchmark
Requires-Dist: fire; extra == "benchmark"
Requires-Dist: matplotlib; extra == "benchmark"
Requires-Dist: pandas; extra == "benchmark"
Requires-Dist: pyinstrument; extra == "benchmark"
Dynamic: license-file

![Crazyflow Logo](https://github.com/learnsyslab/crazyflow/raw/main/docs/img/logo.png)
--------------------------------------------------------------------------------
<div align="center">

  **Fast, parallelizable simulations of Quadrotor drones with JAX.**

  [![Python](https://img.shields.io/badge/python-3.11+-blue.svg)](https://www.python.org)
  [![arXiv](https://img.shields.io/badge/arXiv-2606.01478-b31b1b.svg)](https://arxiv.org/abs/2606.01478)
  [![Tests](https://github.com/learnsyslab/crazyflow/actions/workflows/testing.yml/badge.svg)](https://github.com/learnsyslab/crazyflow/actions/workflows/testing.yml)
  [![Ruff](https://github.com/learnsyslab/crazyflow/actions/workflows/ruff.yml/badge.svg)](https://github.com/learnsyslab/crazyflow/actions/workflows/ruff.yml)
  [![Docs](https://github.com/learnsyslab/crazyflow/actions/workflows/docs.yml/badge.svg)](https://learnsyslab.github.io/crazyflow)

</div>

Crazyflow is a research simulator for quadrotors. It runs batched, differentiable simulations on CPU and GPU via JAX, with analytical and abstracted dynamics for the Crazyflie 2.x family.

```python
import numpy as np
from crazyflow.sim import Sim
from crazyflow.control import Control

sim = Sim(n_worlds=4096, n_drones=1, control=Control.state)
cmd = np.zeros((4096, 1, 13))
cmd[..., 2] = 0.5  # hover at 0.5 m across all worlds

for _ in range(100):
    sim.state_control(cmd)
    sim.step(sim.freq // sim.control_freq)
    sim.render()
```

## Documentation

[learnsyslab.github.io/crazyflow](https://learnsyslab.github.io/crazyflow) — installation, user guide, examples, and API reference.

## Features

- **n\_worlds x n\_drones** — batched over independent environments and multi-drone swarms simultaneously
- **GPU-accelerated** — up to 914 M steps/s on an RTX 4090 (first-principles dynamics, 262 K worlds)
- **Differentiable** — `jax.grad` works through the full dynamics and control pipeline
- **First-principles dynamics** — dynamics using first-principles equations and parameters identified from real-world measurements
- **Abstracted dynamics** — simplified dynamics in three flavors fitted from real Crazyflie flight data
- **Modular pipelines** — step and reset are tuples of plain JAX functions; insert anything, anywhere
- **MuJoCo integration** — onscreen and offscreen rendering, raycasting, and contact detection via MJX

## Installation

```bash
pip install crazyflow           # CPU
pip install "crazyflow[gpu]"    # GPU (Linux x86-64, CUDA 12)
```

Developer install with editable submodules ([pixi](https://pixi.sh/) required):

```bash
git clone --recurse-submodules https://github.com/learnsyslab/crazyflow.git
cd crazyflow
pixi shell
```

## Performance

First-principles dynamics, one drone. CPU: AMD Ryzen 9 7950X. GPU: NVIDIA RTX 4090.

| n\_worlds | CPU steps/s | GPU steps/s |
|---|---|---|
| 64 | 3.3 M | 1.2 M |
| 1 024 | 9.2 M | 18.7 M |
| 16 384 | 11.9 M | 257 M |
| 65 536 | 15.6 M | 678 M |
| 262 144 | 12.6 M | 914 M |

Full benchmarks including multi-drone scaling are in the [documentation](https://learnsyslab.github.io/crazyflow).

## Citation

```bibtex
@misc{schuck2026crazyflow,
      title={Crazyflow: An Accurate, GPU-Accelerated, Differentiable Drone Simulator in JAX}, 
      author={Martin Schuck and Marcel P. Rath and Yufei Hua and AbhisheK Goudar and SiQi Zhou and Angela P. Schoellig},
      year={2026},
      eprint={2606.01478},
      archivePrefix={arXiv},
      primaryClass={cs.RO},
      url={https://arxiv.org/abs/2606.01478}, 
}
```
