crazyflow 0.0.1

Fast, parallelizable simulations of Crazyflies with JAX and MuJoCo.

---

Fast, parallelizable simulations of Crazyflies with JAX and MuJoCo.

Architecture

Crazyflow is a high-performance simulation framework for Crazyflie drones that leverages JAX for efficient parallelization and automatic differentiation. The architecture is designed around a flexible pipeline that can be configured at initialization time, enabling users to swap out physics backends, control methods, and integration schemes.

Core Components

Simulation Pipeline

The simulation is built as a pipeline of functions that are composed at initialization time based on the configuration. This approach avoids runtime branching and allows JAX to optimize the entire pipeline as a single computation. Users can insert their own pure functions into the pipeline to modify the simulation behavior while maintaining compatibility with JAX's optimizations.

Physics Backends

Multiple physics models are supported:

• analytical: A first-principles model based on physical equations
• sys_id: A system-identified model trained on real drone data
• mujoco: MuJoCo physics engine for more complex interactions

Control Modes

Different control interfaces are available:

• state: High-level control of position, velocity, and yaw
• attitude: Mid-level control of collective thrust and orientation
• thrust: Low-level control of individual motor thrusts

Integration Methods

For analytical and system-identified physics:

• euler: Simple first-order integration
• rk4: Fourth-order Runge-Kutta integration for higher accuracy

Parallelization

Crazyflow supports massive parallelization across:

• Worlds: Independent simulation environments that can run in parallel
• Drones: Multiple drones within each world
• Devices: Computations can be executed on CPU or GPU This parallelization is achieved through JAX's vectorization capabilities, allowing thousands of simulations to run simultaneously with minimal overhead.

Domain Randomization

The framework supports domain randomization through the crazyflow/randomize module, allowing parameters like mass to be varied across simulations to improve sim-to-real transfer.

Functional Design

The simulation follows a functional programming paradigm: All state is contained in immutable data structures. Updates create new states rather than modifying existing ones. All functions are pure, enabling JAX's transformations (JIT, grad, vmap) and thus automatic differentiation through the entire simulation, making it suitable for gradient-based optimization and reinforcement learning.

Examples

The repository includes several example scripts demonstrating different capabilities:

Example	Description
hover.py	Basic hovering using state control
thrust.py	Direct motor control using thrust commands
render.py	Visualization of multiple drones with motion traces
contacts.py	Collision detection between drones
gradient.py	Computing gradients through the simulation for optimization
change_pos.py	Manipulating drone positions programmatically

Known Issues

• "RuntimeError: MUJOCO_PATH environment variable is not set" upon installing this package. This error can be resolved by using venv instead of conda. Somtimes the mujoco install can fail with conda.
• If using zsh don't forget to escape brackets when installing additional dependencies: pip install .\[gpu\].

Using the project with VSCode devcontainers

Running on CPU: by default the containers run on CPU. You don't need to take any action.

Running on GPU: The devcontsainers can easily run using your computer's NVIDIA GPU on Linux and Windows. This makes sense if you want to accelerate simulation by running thousands of simulation in parallel. In order to work you need to install the CUDA Toolkit, NVIDIA Container runtime for your computer. Finally, enable GPU access to the devcontainers by setting the commented out "--gpus=all" and "--runtime=nvidia" flags in devcontainer.json.

Linux

1. Make sure to be in a X11 session (link), otherwise rendering of the drone will fail.
2. Install Docker (, and make sure Docker Daemon is running)
3. Install VSCode, with devcontainer extension, and remote dev pack.
4. Clone this project's code. Rename /.devcontainer/devcontainer.linux.json to /.devcontainer/devcontainer.json.
5. Open this project in VSCode. VSCode should automatically detect the devcontainer and prompt you to Reopen in container. If not, see here to open manually. Note: Opening the container for the first time might take a while (up to 15 min), as the container is pulled from the web and build.

Windows (requires Windows 10 or later)

For windows, we require WSL2 to run the devcontainer. (So its actually Linux with extra steps.) Full instructions can be found in the official docs. Here are the important steps:

1. Install Docker, and WSL2, and Ubuntu 22.04 LTS (, and make sure Docker Daemon is running)
2. Docker will recognize that you have WSL installed and prompt you via Windows Notifications to enable WSL integration -> confirm this with Enable WSL integration. If not, open Docker Desktop, navigate to the settings, and manually enable WSL integration. (There are TWO setting options for this. Make sure to enable BOTH!)
3. Install VSCode, with the WSL extension, devcontainer extension, and remote dev pack.
4. Clone the source code for the exercises in the WSL2 file system to /home (~), or wherever you like. (Performance when working on the WSL file system is much better compared to Windows file system). You can access the WSL filesystem either by starting a WSL/Ubuntu console, or via the Windows File Explorer at \\wsl.localhost\Ubuntu\home (replace Ubuntu with your distro, if necessary).
5. Rename /.devcontainer/devcontainer.windows.json to /.devcontainer/devcontainer.json.
6. Open this project in VSCode. The easiest way to do so is by starting a WSL/Ubuntu shell, navigating via cd to the source code, then type code . to open VSCode. VSCode should automatically detect the devcontainer and prompt you to Reopen in container. If not, see here to open manually. Note: Opening the container for the first time might take a while (up to 15 min), as the container is pulled from the web and build.

MacOS

Unfortunately, we did not get the devcontainer to work with MacOS yet, even after following those steps. We expect that the issue is related to Mujoco rendering from inside the Docker container and display forwarding with X11. There is also an unresolved Issue on GitHub. If you manage to make it work, please let us know.

Until then, MacOS users are required to install this project using an python environment manager such as conda or mamba. If you use conda, these are the required commands: conda create --name crazyflow -c conda-forge python=3.11, conda activate crazyflow, conda install pip, pip install -e ..

---

Known Issues:

• if building docker container fails at RUN apt-get update, make sure your host systems time is set correct: https://askubuntu.com/questions/1511514/docker-build-fails-at-run-apt-update-error-failed-to-solve-process-bin-sh