flyGNN

Whole-Brain Connectomic Graph Neural Network Enables Whole-Body Locomotion Control in Drosophila

📖 Introduction

Whole-brain connectome provides a structural blueprint for linking neural circuits to behavior, yet its application to embodied control remains largely unexplored.

We introduce flyGNN, a reinforcement learning controller whose architecture is instantiated directly from the complete adult Drosophila connectome.

This work demonstrates that whole-brain connectivity can directly support embodied reinforcement learning, establishing a new paradigm for connectome-based control algorithms.

Note: The source code is still being organized. The anonymously released version here is a stable snapshot with added comments and restructured organization for clarity. You may encounter potential environment incompatibilities, and some parts of the code may contain legacy issues (e.g., hard-coded paths or settings without CLI support). In future releases, we will provide more extensible code and also release model checkpoints for download.

⚙️ Installation

It is recommended to first install flybody in your environment, as its installation will provide most of the dependencies required by this project, and PPO training also depends on the flybody library.

After installing flybody, install the core dependencies:

pip install torch torchvision torchaudio
pip install pytorch-lightning
pip install torch-geometric

Note: If you encounter errors installing torch-geometric, please refer to the official installation guide for version-specific wheels.

Note: Our codebase is still being organized to support broader library compatibility, a more detailed environment setup guide will be provided in future releases.

📂 Project Structure

flygnn/
├── data/                                  # Dataset folder
│   ├── node_df.csv                        # Node metadata (neuron properties, masks)
│   └── flygnn_cache/root_id_mapping.csv   # Mapping from neuron IDs to root IDs
│
├── src/flygnn/                            # Core source code
│   ├── connectome.py                      # Load & preprocess connectome into PyG data
│   ├── flygnn.py                          # Main FlyGNN model implementation
│   ├── train_pre.py                       # Stage 1: Pre Training script
│   ├── train_post.py                      # Stage 2: Post Training script
│   └── utils.py                           # Utility functions (KL loss, RunningNorm, etc.)
│
├── LICENSE
├── README.md
└── .gitignore

📊 Data Preparation

Place the following files under data/:

• node_df.csv — neuron/node dataframe with afferent/intrinsic/efferent labels.
• connections.csv — connectome edges (>100MB). You can download from codex. In this work, we use the official connections_princeton.csv of FAFB v783, which is the raw connections data downloaded and unziped from Flywire.
• flygnn_cache/root_id_mapping.csv — ID remapping file (optional).

Note: root_id_mapping.csv is optional. If not provided, the code will automatically generate one to ensure consistency with FlyWire root IDs, which can later be used for analysis.

Note: Absolute file paths are recommended.

🚀 Training

flyGNN uses a two-stage training pipeline:

1. Pre Training (Stage 1)

Trains the model to mimic expert trajectories (teacher policy).

python src/flygnn/train_pre.py

2. Post Training (Stage 2)

Refines the policy via reinforcement learning with PPO.

python src/flygnn/train_post.py

📈 Evaluation & Visualization

• The trained model can be rolled out in MuJoCo-based flybody environment.
• Node embeddings can be visualized to analyze whole-brain dynamics.
• See paper appendix for dimensionality reduction & plotting methods.

📜 License

This project is licensed under an MIT License. See LICENSE for details.

📑 Citation

Citation details will be added upon paper release.