# A fruit fly connectome that learns

The full male fruit fly nervous system (male-cns v1.0, 166,691 neurons; the file used here has 176,422 bodies, including unannotated fragments) has already been wired into Doom, Mario 64, Minecraft and a crypto trading bot. None of those projects showed the fly **learning**. This repository does, with the controls that say whether it is real.

The task is the classic fly experiment: two odors, one of them punished. The fly picks a side by smell; picking the punished odor fires PPL101, the punishment dopamine neuron, and weakens the synapses from that odor's Kenyon cells onto MBON11. The next time, the fly is less likely to go there.

## Result

200 flies per odor pair, each trained on its own, 15 trials, 5 odor pairs. Percent of flies avoiding the punished odor (`results/learn_to_avoid.txt`):

| pair | condition | trial 1 | trial 8 | trial 15 |
|---|---|---|---|---|
| 0 | real wiring | 42 | 74 | 79 |
| 0 | punishment without odor | 42 | 50 | 52 |
| 0 | shuffled wiring | 0 | 0 | 0 |
| 1 | real wiring | 50 | 63 | 71 |
| 2 | real wiring | 44 | 70 | 80 |
| 3 | real wiring | 50 | 66 | 73 |
| 4 | real wiring | 47 | 70 | 70 |

All five pairs and all three conditions are in `results/learning.json`. Punishment with no odor ends between 44 and 54% for every pair. The shuffled brain never learns: in the real connectome PPL101 reaches 409 Kenyon cells, after shuffling it reaches 10, and the fly stays stuck on one side (0% in three pairs, 2% in one, 100% in one).

`results/trained_fly.npz` has the KC to MBON11 weights of one fly before and after training.

## Why the whole brain simulation does not learn

The first thing I tried was the whole-brain leaky integrate-and-fire model (Shiu et al. 2024). Two different odors light up almost the same Kenyon cells, because the APL inhibitory neuron saturates on any input. Turning APL up only shrinks the set (`results/full_brain_saturates.txt`):

| APL gain | KCs odor A | KCs odor B | overlap |
|---|---|---|---|
| 1 | 2655 | 2639 | 99% |
| 2 | 342 | 310 | 86% |
| 2.5 | 100 | 87 | 80% |
| 3 | 33 | 22 | 67% |

With 99% overlap, punishing one odor punishes the other, and nothing can be learned. I think this is also why visual reinforcement on this connectome has not worked so far.

## What is real and what is modeled

From the connectome:
- ORN to PN to Kenyon cell synapse counts and signs
- which Kenyon cells connect to MBON11, and with how many synapses
- which Kenyon cells receive input from PPL101

Modeled, and you can change every one of them in `flylearn/mushroom.py`:
- the odor code is one feedforward pass ORN to PN to KC, and only the top 5% of Kenyon cells fire (APL inhibition as a k-winners rule, the way most mushroom body models do it)
- each punishment multiplies the affected KC to MBON11 weights by 0.7
- the choice is a logistic function of the MBON11 response difference (slope 4)
- odor codes get a small random jitter on every trial

The learning ceiling sits around 75 to 80% because only 406 of the 2,854 Kenyon cells onto MBON11 receive PPL101 input. I did not tune it.

## Run it

```
pip install -r requirements.txt
python experiments/learn_to_avoid.py            # about 1 minute, about 400 MB RAM
python experiments/full_brain_saturates.py      # a few minutes
```

## Teach it something else

Another odor pair is one line: pass other glomerulus names to `MushroomBody.code()` (the list is in `MushroomBody.glomeruli`). A different punishment compartment means swapping `MBON11` and `PPL101` in `flylearn/mushroom.py` for another MBON and dopamine neuron pair. A new game needs you to decide what the game shows as odor and what the MBON output does, which is where the work is.

## Credits

- Connectome: male-cns v1.0, HHMI Janelia FlyEM Project Team, Google Research Connectomics and the Cambridge Drosophila Connectomics Group, CC BY 4.0. Compact file format and build script from fly-brain-minecraft (MIT). Details in `data/NOTICE.md`.
- LIF parameters: Shiu et al., "A Drosophila computational brain model reveals sensorimotor processing", Nature 2024.

By Weverton Guedes ([@owevertonguedes](https://github.com/owevertonguedes)). Video: [the fly learning the maze on YouTube](https://youtu.be/FSPpoIPRHIY).

Code: MIT (`LICENSE`). Data: CC BY 4.0.