`flybody`: fruit fly body model for MuJoCo physics

flybody is an anatomically-detailed body model of the fruit fly Drosophila melanogaster for the MuJoCo physics simulator and reinforcement learning applications.

The fly model was developed in a collaborative effort by Google DeepMind and HHMI Janelia Research Campus.

We envision our model as a platform for fruit fly biophysics simulations and for modeling neural control of sensorimotor behavior in an embodied context; see our accompanying publication.

Quick Start

import numpy as np
import mediapy

from flybody.fly_envs import walk_imitation

# Create walking imitation environment.
env = walk_imitation()

# Run environment loop with random actions for a bit.
for _ in range(100):
    action = np.random.normal(size=59)  # 59 is the walking action dimension.
    timestep = env.step(action)

# Generate a pretty image.
pixels = env.physics.render(camera_id=1)
mediapy.show_image(pixels)

Also see the tutorial notebook or .

Installation

Recommended: `uv` (fast, reproducible)

uv is the recommended package manager for flybody. It handles virtual environments and dependency resolution in a single tool.

Install uv (if not already installed):

curl -LsSf https://astral.sh/uv/install.sh | sh

Clone and set up:

git clone https://github.com/TuragaLab/flybody.git
cd flybody
uv venv --python 3.10
source .venv/bin/activate

Core installation — minimal install for experimenting with the fly model in MuJoCo or prototyping task environments:
```
uv pip install -e .
```
ML extension (optional) — adds TensorFlow and Acme for running policy networks:
```
uv pip install -e ".[tf]"
```
Ray training extension (optional) — adds Ray for distributed DMPO training:
```
uv pip install -e ".[ray]"
```
Dev tools (optional) — adds ruff, JupyterLab, tqdm:
```
uv pip install -e ".[dev]"
```
Everything:
```
uv pip install -e ".[all]"
```

Or use the setup script:

bash scripts/setup.sh          # core
bash scripts/setup.sh --all    # everything

Alternative: conda + pip

Click to expand conda/pip instructions

Clone and create a conda environment:

git clone https://github.com/TuragaLab/flybody.git
cd flybody
conda create --name flybody -c conda-forge python=3.10 pip ipython cudatoolkit=11.8.0
conda activate flybody

Install in one of the modes (add -e for editable/developer mode):

Core: pip install -e .
ML extension: pip install -e ".[tf]"
Ray training: pip install -e ".[ray]"

Alternative: pip from remote

Click to expand remote pip instructions

pip install git+https://github.com/TuragaLab/flybody.git                     # core
pip install "flybody[tf] @ git+https://github.com/TuragaLab/flybody.git"      # ML
pip install "flybody[ray] @ git+https://github.com/TuragaLab/flybody.git"     # Ray

Additional Configuration

You may need to set MuJoCo rendering environment variables:
```
export MUJOCO_GL=egl
export MUJOCO_EGL_DEVICE_ID=0
```

For ML and Ray extensions, LD_LIBRARY_PATH may need an update:

CUDNN_PATH=$(dirname $(python -c "import nvidia.cudnn;print(nvidia.cudnn.__file__)"))
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$CONDA_PREFIX/lib/:$CUDNN_PATH/lib

Verify your installation:
```
uv run pytest tests/ -v
```

Project Structure

flybody/
├── README.md                          # This file
├── AGENTS.md                          # AI agent guidance
├── pyproject.toml                     # Package config & dependencies
├── LICENSE                            # Apache 2.0
├── fly-white.png                      # Hero image
│
├── flybody/                           # Source package
│   ├── __init__.py
│   ├── fly_envs.py                    # ★ Environment factory functions
│   ├── fruitfly/                      # Walker model
│   │   ├── fruitfly.py                # FruitFly composer.Walker class
│   │   └── assets/                    # MuJoCo XML + 85 OBJ meshes
│   │       ├── fruitfly.xml           # Main fly model definition
│   │       └── floor.xml              # Floor arena for visualization
│   ├── tasks/                         # RL task definitions
│   │   ├── base.py                    # FruitFlyTask, Flying, Walking base classes
│   │   ├── flight_imitation.py        # Flight trajectory tracking
│   │   ├── walk_imitation.py          # Walking trajectory tracking
│   │   ├── walk_on_ball.py            # Tethered walking on floating ball
│   │   ├── vision_flight.py           # Vision-guided flight (bumps/trench)
│   │   ├── template_task.py           # Empty no-op task for testing
│   │   ├── rewards.py                 # Reward functions
│   │   ├── trajectory_loaders.py      # HDF5 trajectory dataset loaders
│   │   ├── synthetic_trajectories.py  # Synthetic trajectory generation
│   │   ├── pattern_generators.py      # Wing pattern generators
│   │   ├── task_utils.py              # Task utility functions
│   │   ├── constants.py               # Physics timestep constants
│   │   └── arenas/                    # Arena definitions (ball, hills)
│   ├── agents/                        # RL agents (TF/Ray)
│   │   ├── agent_dmpo.py              # DMPO agent
│   │   ├── learning_dmpo.py           # DMPO learner
│   │   ├── losses_mpo.py              # MPO loss functions
│   │   ├── actors.py                  # Actor implementations
│   │   ├── network_factory.py         # Network factory for DMPO
│   │   ├── network_factory_vis.py     # Vision network factory
│   │   ├── ray_distributed_dmpo.py    # Ray-distributed DMPO
│   │   ├── remote_as_local_wrapper.py # Ray remote wrapper
│   │   ├── counting.py                # Step counters
│   │   ├── utils_ray.py               # Ray utilities
│   │   └── utils_tf.py                # TensorFlow utilities
│   ├── train_dmpo_ray.py              # ★ Distributed training script
│   ├── download_data.py               # Figshare data download
│   ├── ellipsoid_fluid_model.py       # Ellipsoid fluid dynamics
│   ├── inverse_kinematics.py          # Inverse kinematics solver
│   ├── quaternions.py                 # Quaternion math utilities
│   ├── loggers.py                     # MLflow training logger
│   └── utils.py                       # Rendering & display utilities
│
├── tests/                             # Test suite
│   ├── test_flybare.py                # Stand-alone fly model tests
│   ├── test_flywalker.py              # FruitFly walker class tests
│   ├── test_core.py                   # Core RL environment tests
│   ├── test_walking_env.py            # Walking imitation env tests
│   ├── test-tf.py                     # TensorFlow policy test
│   └── common.py                      # Shared test utilities
│
├── docs/                              # Tutorial notebooks
│   ├── getting-started.ipynb
│   ├── fly-env-examples.ipynb
│   ├── fly-on-ball-minimal.ipynb
│   ├── controller-reuse-vision-flight.ipynb
│   └── sensory-input-tracking.ipynb
│
├── scripts/                           # Orchestration scripts
│   ├── setup.sh                       # Environment setup
│   ├── download_data.sh               # Data download
│   ├── run_tests.sh                   # Test runner
│   ├── lint.sh                        # Lint runner
│   ├── train.sh                       # Training launcher
│   ├── export_all.sh                  # Master export orchestrator
│   ├── render_environments.py         # Render still images
│   ├── generate_rollout_videos.py     # Generate rollout videos
│   ├── export_model_data.py           # Export model/env data
│   ├── explore_physics.py             # Standalone physics exploration
│   ├── track_sensory_data.py          # Sensory data tracking
│   ├── explore_task_utils.py          # Task utility exercises
│   ├── inspect_walker.py              # Walker anatomy inspection
│   ├── vision_pipeline.py             # Vision pipeline
│   ├── explore_quaternions.py         # Quaternion operations
│   ├── explore_rewards.py             # Reward function analysis
│   ├── explore_arenas.py              # Arena construction
│   └── benchmark_envs.py              # Environment benchmarks
│
├── output/                            # Generated exports (gitignored)
│   ├── images/                        # PNG renders from all environments
│   ├── videos/                        # MP4 rollout videos
│   ├── data/                          # JSON model parameters & env specs
│   ├── physics/                       # Physics exploration output
│   ├── sensory/                       # Sensory data plots
│   ├── task_utils/                    # Task utility test output
│   ├── walker/                        # Walker anatomy data
│   ├── vision/                        # Vision pipeline output
│   ├── quaternions/                   # Quaternion operations output
│   ├── rewards/                       # Reward analysis output
│   ├── arenas/                        # Arena renders
│   ├── benchmark/                     # Environment benchmark data
│   └── *.log                          # Execution logs
│
└── .github/workflows/                 # CI pipelines
    ├── pytest.yml                     # Core tests (Python 3.10)
    ├── pyversions.yml                 # Multi-version tests (3.10–3.12)
    ├── lint.yml                       # Ruff linting
    └── tf-test.yml                    # TensorFlow policy test

Architecture Overview

flybody is built on the dm_control composer framework. The architecture follows a layered pattern:

┌─────────────────────────────────────────────────────────┐
│                    fly_envs.py                           │
│          Factory functions (user-facing API)             │
│  flight_imitation · walk_imitation · walk_on_ball        │
│  vision_guided_flight · template_task                    │
├──────────────┬──────────────────────────────┬────────────┤
│  FruitFly    │       Tasks                  │   Arenas   │
│  (Walker)    │  FruitFlyTask (base)         │   floors   │
│  fruitfly.py │  ├── Flying (flight base)    │   ball     │
│              │  └── Walking (walk base)     │   hills    │
│  85 OBJ mesh │      ├── FlightImitation     │            │
│  fruitfly.xml│      ├── WalkImitation       │            │
│              │      ├── WalkOnBall          │            │
│              │      └── VisionFlight        │            │
├──────────────┴──────────────────────────────┴────────────┤
│                   dm_control · MuJoCo                    │
└─────────────────────────────────────────────────────────┘

Key Classes

Class	Module	Role
`FruitFly`	`flybody.fruitfly.fruitfly`	Composer `Walker` — builds the fly model from XML, manages observables and actuators
`FruitFlyTask`	`flybody.tasks.base`	Base task — wires walker + arena, defines reward/termination interface
`Flying`	`flybody.tasks.base`	Configures fly model for flight (wing params, fluid model, leg disable)
`Walking`	`flybody.tasks.base`	Configures fly model for walking (wing retract, adhesion, leg contact)

Factory Functions

All environments are created via factory functions in fly_envs.py:

Function	Task	Action Dim
`flight_imitation()`	Track flying reference trajectory	22
`walk_imitation()`	Track walking reference trajectory	59
`walk_on_ball()`	Walk on a floating ball (tethered)	59
`vision_guided_flight()`	Fly through bumps/trench with vision	22
`template_task()`	Empty no-op task for testing	59

Task Environments

Flight Imitation

Requires the fly to track a reference flying trajectory from an HDF5 dataset. Reward is based on joint angle and body position matching. Uses the Flying base class with ellipsoid fluid dynamics.

Walk Imitation

Requires the fly to track a reference walking trajectory. Supports inference mode (no dataset needed) for testing. Uses the Walking base class with adhesion actuators.

Walk on Ball

Tethered walking on a floating ball. Supports both position and force actuators. Minimal task suitable for quick experiments and RL prototyping.

Vision-Guided Flight

Flight through procedurally generated arenas (bumps or trench) using monocular or binocular eye cameras. Combines locomotion with visual processing.

Template Task

Empty no-op task that returns a constant reward of 1. Used for unit testing and as a starting point for new tasks.

Data Downloads

Supplementary data (reference trajectories, trained policies) are hosted on Figshare:

# Via script
bash scripts/download_data.sh

# Via Python
from flybody.download_data import figshare_download

figshare_download('walking-imitation-dataset')
figshare_download('flight-imitation-dataset')
figshare_download('trained-policies')
figshare_download('controller-reuse-checkpoints')

Available dataset keys:

Key	Description
`walking-imitation-dataset`	Reference walking trajectories (HDF5)
`flight-imitation-dataset`	Reference flight trajectories (HDF5)
`trained-policies`	Pre-trained policy checkpoints
`controller-reuse-checkpoints`	Checkpoints for controller reuse experiments

Data is downloaded and extracted to flybody-data/ by default.

Scripts

Thin orchestration scripts live in scripts/. They wrap the real methods in the package — no business logic in scripts.

Script	Purpose	Usage
`setup.sh`	Create venv and install with uv	`bash scripts/setup.sh [--tf\|--ray\|--dev\|--all]`
`download_data.sh`	Download figshare datasets	`bash scripts/download_data.sh [KEY...]`
`run_tests.sh`	Run pytest suite	`bash scripts/run_tests.sh [PYTEST_ARGS...]`
`lint.sh`	Run ruff linting	`bash scripts/lint.sh [--fix]`
`train.sh`	Launch Ray distributed training	`bash scripts/train.sh`
`export_all.sh`	Generate images, videos, data, and more	`bash scripts/export_all.sh [--all\|--images\|--videos\|--data\|--physics\|--sensory\|--task-utils\|--walker\|--vision\|--quaternions\|--rewards\|--arenas\|--benchmark]`

Export Scripts (Python)

The export pipeline uses real flybody methods to generate rendered images, rollout videos, model data, and more. All scripts support --output-dir, --width, --height, and other CLI flags:

Script	Methods Used	Output
`render_environments.py`	`fly_envs.*()`, `env.physics.render()`	PNGs in `output/images/`
`generate_rollout_videos.py`	`fly_envs.*()`, `utils.rollout_and_render()`	MP4s in `output/videos/`
`export_model_data.py`	`FruitFly()`, `env.observation_spec()`	JSONs in `output/data/`
`explore_physics.py`	`FruitFly`, `mjcf.Physics`, joint/actuator manipulation, wing animation, mocap sites	PNGs + MP4s in `output/physics/`
`track_sensory_data.py`	Per-step observables (`joints_pos`, `velocimeter`, `world_zaxis`), matplotlib plots	PNGs + JSON in `output/sensory/`
`explore_task_utils.py`	All 14 `task_utils` functions: `retract_wings`, `real2canonical`↔`canonical2real`, `root2com`↔`com2root`, etc.	JSONs + PNGs in `output/task_utils/`
`inspect_walker.py`	Walker anatomy: `observable_joints`, `actuators`, `end_effectors`, cameras, body tree	JSON + PNGs in `output/walker/`
`vision_pipeline.py`	Eye camera rendering, `walker/left_eye`/`right_eye` obs, vision-guided flight	MP4s + PNGs in `output/vision/`
`explore_quaternions.py`	All 18 `quaternions.*` functions: algebra, rotation, metrics, angular velocity	JSONs in `output/quaternions/`
`explore_rewards.py`	`compute_diffs`, `reward_factors_deep_mimic`, sensitivity analysis	JSONs in `output/rewards/`
`explore_arenas.py`	`BallFloor`, `Hills`, `SineTrench`, `SineBumps` construction and rendering	PNGs + JSON in `output/arenas/`
`benchmark_envs.py`	All env factory timing, `constants.py` validation	JSONs in `output/benchmark/`

# Generate everything
bash scripts/export_all.sh

# Or selectively
bash scripts/export_all.sh --images       # Rendered stills only
bash scripts/export_all.sh --videos       # Rollout videos only
bash scripts/export_all.sh --data         # Model/env data only
bash scripts/export_all.sh --physics      # Standalone physics exploration
bash scripts/export_all.sh --sensory      # Sensory data tracking
bash scripts/export_all.sh --task-utils   # task_utils exercises
bash scripts/export_all.sh --walker       # Walker anatomy inspection
bash scripts/export_all.sh --vision       # Vision pipeline
bash scripts/export_all.sh --quaternions  # Quaternion operations
bash scripts/export_all.sh --rewards      # Reward analysis
bash scripts/export_all.sh --arenas       # Arena exploration
bash scripts/export_all.sh --benchmark    # Environment benchmarks

Development

Linting

uv run ruff check flybody/          # Check
uv run ruff check flybody/ --fix    # Auto-fix

Testing

# Core tests (no TF/Ray required)
uv run pytest tests/test_flybare.py tests/test_core.py tests/test_flywalker.py tests/test_walking_env.py -v

# All tests
uv run pytest tests/ -v

# With rendering (requires MUJOCO_GL=egl and GPU)
MUJOCO_GL=egl uv run pytest tests/ -v

CI Workflows

Workflow	Trigger	What it tests
`pytest.yml`	push/PR to main, dev	Core tests (Python 3.10)
`pyversions.yml`	push/PR to main, dev	`test_flybare` + `test_core` on Python 3.10–3.12
`lint.yml`	push/PR to main, dev	Ruff lint check
`tf-test.yml`	push/PR to main, dev	TensorFlow policy network test

Visualizing the Model

Drag-and-drop fruitfly.xml or floor.xml into MuJoCo's simulate viewer for interactive visualization.

Tutorials & Notebooks

See the docs/ directory for Jupyter notebooks:

Notebook	Description
Getting Started	Create environments, step with random actions, render images
Fly Environment Examples	Flight, walking, and vision-guided flight demos
Fly on Ball Minimal	Minimal tethered fly-on-ball walking example
Controller Reuse & Vision Flight	Transfer learning and vision-guided flight
Sensory Input Tracking	Tracking sensory inputs and body state

Run notebooks with:

uv pip install -e ".[dev]"
uv run jupyter lab docs/

Citing `flybody`

See our accompanying publication. Thank you for your interest in our fly model :)

@article{flybody,
  title = {Whole-body physics simulation of fruit fly locomotion},
  author = {Roman Vaxenburg and Igor Siwanowicz and Josh Merel and Alice A Robie and
            Carmen Morrow and Guido Novati and Zinovia Stefanidi and Gert-Jan Both and
            Gwyneth M Card and Michael B Reiser and Matthew M Botvinick and
            Kristin M Branson and Yuval Tassa and Srinivas C Turaga},
  journal = {Nature},
  volume = {643},
  pages = {1312--1320},
  year = {2025},
  doi = {https://doi.org/10.1038/s41586-025-09029-4},
  url = {https://www.nature.com/articles/s41586-025-09029-4},
}

License

Apache License 2.0

Navigation Menu

Repository files navigation

`flybody`: fruit fly body model for MuJoCo physics

Table of Contents

Quick Start

Installation

Recommended: `uv` (fast, reproducible)

Alternative: conda + pip

Alternative: pip from remote

Additional Configuration

Project Structure

Architecture Overview

Key Classes

Factory Functions

Task Environments

Flight Imitation

Walk Imitation

Walk on Ball

Vision-Guided Flight

Template Task

Data Downloads

Scripts

Export Scripts (Python)

Development

Linting

Testing

CI Workflows

Visualizing the Model

Tutorials & Notebooks

Citing `flybody`

License

About

Releases

Packages

Contributors

Languages

Name		Name	Last commit message	Last commit date
.github/workflows		.github/workflows
docs		docs
flybody		flybody
scripts		scripts
tests		tests
.gitignore		.gitignore
AGENTS.md		AGENTS.md
CONTRIBUTING.md		CONTRIBUTING.md
LICENSE		LICENSE
README.md		README.md
fly-white.png		fly-white.png
pyproject.toml		pyproject.toml
scripts.conf		scripts.conf

Latest commit

History

Folders and files

Repository files navigation

flybody: fruit fly body model for MuJoCo physics

Table of Contents

Quick Start

Installation

Recommended: uv (fast, reproducible)

Alternative: conda + pip

Alternative: pip from remote

Additional Configuration

Project Structure

Architecture Overview

Key Classes

Factory Functions

Task Environments

Flight Imitation

Walk Imitation

Walk on Ball

Vision-Guided Flight

Template Task

Data Downloads

Scripts

Export Scripts (Python)

Development

Linting

Testing

CI Workflows

Visualizing the Model

Tutorials & Notebooks

Citing flybody

License

About

Topics

Resources

Contributing

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages

`flybody`: fruit fly body model for MuJoCo physics

Recommended: `uv` (fast, reproducible)

Citing `flybody`