The Problem

This repository provides a detailed biomechanical model of a fruit fly (Drosophila melanogaster) for use with the MuJoCo physics simulator. Simulating realistic biomechanics and neural control for insects is challenging due to their complex anatomy and movement patterns. Existing models often lack the anatomical detail needed for accurate reinforcement learning (RL) applications.

What This Does

flybody is a physics-based simulation of a fruit fly, designed for studying sensorimotor behavior and RL applications. The model includes detailed anatomical structures such as legs, wings, antennae, and body segments, enabling realistic locomotion and sensory interactions.

Key components:

  • MuJoCo XML Model: The fruitfly.xml and floor.xml files define the fly's anatomy and physics properties.
  • RL Environments: The fly_envs.py module provides environments for walking, flight, and vision-guided tasks.
  • DMPO Agent: The agents/ directory implements the Distributed Model Predictive Optimization (DMPO) algorithm for training RL agents.

The repository also includes Jupyter notebooks demonstrating how to use the model, such as docs/getting-started.ipynb and docs/fly-env-examples.ipynb.

How It Is Wired

The main entry point for interacting with the fly model is the walk_imitation() function in fly_envs.py. This function initializes a walking imitation environment, which can be used to simulate the fly's movement.

The control flow is as follows:

  1. Environment Initialization: The walk_imitation() function sets up the MuJoCo simulation environment.
  2. Action Execution: The step(action) method applies the provided action to the fly's joints and updates the simulation state.
  3. State Observation: The observation() method returns the current state of the fly, which can be used by an RL agent to make decisions.

The agents/ directory contains the implementation of the DMPO algorithm, which is used to train RL agents to control the fly. The train_dmpo_ray.py script provides a distributed training pipeline using the Ray framework.

How To Use It

Setup

  1. Clone the repository: ``bash git clone https://github.com/moses-y/flybody.git cd flybody ``
  2. Create and activate a conda environment: ``bash conda create --name flybody -c conda-forge python=3.10 pip ipython cudatoolkit=11.8.0 conda activate flybody ``
  3. Install the package: ``bash pip install -e . ``

Running It

To visualize the fly model, you can use the MuJoCo simulate viewer:

   simulate fruitfly.xml

To run a simple Python example:

   import numpy as np
   import mediapy

   from flybody.fly_envs import walk_imitation

   env = walk_imitation()

   for _ in range(100):
       action = np.random.normal(size=59)
       timestep = env.step(action)

   pixels = env.physics.render(camera_id=1)
   mediapy.show_image(pixels)

Training an RL Agent

To train an RL agent using the DMPO algorithm, run the distributed training script:

   python flybody/train_dmpo_ray.py

Real-World Use

This model can be used to study the biomechanics and neural control of fruit fly locomotion. For example, researchers can use the walk_imitation environment to train RL agents to mimic the walking patterns of real fruit flies. The detailed anatomical model allows for the exploration of how different body parts contribute to movement and sensory perception.

Code Health & Issues

  • Low/Risk - Dependencies are declared in pyproject.toml without a lockfile, leading to non-reproducible builds.

The Bottom Line

flybody provides a detailed and realistic biomechanical model of a fruit fly for use with the MuJoCo physics simulator. It is well-suited for studying sensorimotor behavior and RL applications. The repository includes comprehensive documentation and examples, making it easy to get started. However, the lack of a dependency lockfile may pose reproducibility issues.