UncleStrawberryBear/FlyControl-models

FlyControl models

Research records for A Drosophila connectome prior enables adaptable embodied control.

Code: FlyControl.

Figure	Experiment	Models
Figure 2	Fly terrain locomotion	Adult, MLP and ER; 20 seeds each
Figure 3	Fly turning control	Adult and MLP; 24 seeds each
Figure 4	Fly speed control	Adult and MLP; 20 seeds each
Figure 5b–d, S7	Ant flat ground and floor friction	Adult and MLP; 20 seeds each
Figure 5e	Ant U-Maze, Medium Maze and Large Maze	Seven controller families; 20 seeds per maze and family
Figure 5f–j	Ant Medium Maze trajectories and navigation	Adult and MLP; 20 seeds each
S3	Developmental direction control	Adult and Larva; 6 seeds each
S4	Developmental speed control	Adult and Larva; 6 seeds each
S5	BANC descending and motor readouts	6 seeds per readout; initial and final states
S6	BANC regional recovery	3 seeds per readout
S8	Ant Large Maze trajectories, speed and route complexity	Seven controller families; 20 seeds each
Figure 6a–e	Cassie balance; Tripod, Go1 and Hexapod walking	Connectome and MLP; 20 seeds per body and family
Figure 6f	Cassie standing and lateral-motion stability	Connectome and MLP; 20 seeds each
Figure 6g	Go1 tilt, pace and heading stability	Connectome and MLP; 20 seeds each
S9	Tripod direction and path stability	The Figure 6 Tripod policies
S10	Hexapod path and lateral-corridor stability	The Figure 6 Hexapod policies

The seven Ant maze controller families are Adult connectome, MLP, GRU, LSTM, Transformer, Reservoir and Reservoir+RO. The Ant flat-ground MLP uses two 1,378-unit ReLU hidden layers.

The repository contains 846 checkpoint files: 810 final policies and 36 initial states. Figure 5 contributes 40 flat-ground policies and 420 maze policies. Checkpoints contain actor and critic weights and the saved components used by each policy, as listed in the model index.

model_index.json identifies each checkpoint by figure, experiment, controller, seed, configuration and SHA-256. Figure 5 entries are also listed in figure5/model_index.json. Figure 5 model IDs follow figure5/<flat|umaze|medium|large>/<method>/seed_<0–19>/final.

Figure 6, S9 and S10

Figure 6 contributes 160 final policies: Connectome and MLP controllers for Cassie, Tripod, Go1 and Hexapod, with 20 training seeds per body and controller. S9 and S10 use the same Tripod and Hexapod policies.

Body	Experiment	Training steps	Evaluation command (m/s)
Cassie	Biped balance	524,288	0.0
Tripod	Radial tripod walking	131,072	0.1
Go1	Quadruped walking	524,288	0.8
Hexapod	Hexapod walking with an analytic CPG	131,072	0.5

figure6/model_index.json maps each policy to its figure, body, controller, seed, configuration and checksum. The shared connectome is stored once in figure6/graph.npz. Figure 6 policies consume the features produced by their body interface; their PPO configuration uses normalize_observation: false.

python /path/to/model-snapshot/load_embodiment_policy.py \
  --model-id figure6/cassie/connectome/seed_871/final \
  --code-dir /path/to/FlyControl

The helper restores the policy and runs three numerical forward steps. load_policy.py also accepts Figure 6 model IDs. Use the corresponding FlyControl body interface to construct observations and apply motor actions. Evaluation conditions and reset seeds are stored in each index entry. File and loading checks are recorded in figure6/validation.json.

Download and load

from huggingface_hub import snapshot_download
models = snapshot_download('UncleStrawberryBear/FlyControl-models')

Install the dependencies from the linked FlyControl revision. Load a Figure 5 policy with:

python /path/to/model-snapshot/load_ant_policy.py \
  --model-id figure5/medium/adult/seed_0/final \
  --code-dir /path/to/FlyControl

The helper verifies the checkpoint, constructs its policy, restores the saved state and runs a numerical forward pass. Its Python function load_ant_policy returns the policy and observation-normalization statistics. Figure 5 policies use 109 observation channels and eight action channels. Evaluation settings, including action scale, smoothing and waypoint planner, are stored in each index entry.

For Figures 2–4 and S3–S6, use load_policy.py. Figures 2–4 use the matching graph from the data archive:

python /path/to/model-snapshot/load_policy.py \
  --model-id figure2/connectome/seed_2011/final \
  --code-dir /path/to/FlyControl --data-dir /path/to/paper-data

Data and trajectories: FlyControl-data. File and loading checks are recorded in validation.json and figure5/validation.json. Source and licensing information is in RIGHTS.md.

License

The authors' original model weights, loading utilities and accompanying documentation are released under the MIT License. Third-party materials retain their upstream licenses and attribution; see RIGHTS.md.