Back to projects
Mar 2026 – Present · In progress
C. elegans Connectome as a Learnable RL Policy
Independent research, solo
ResearchAI/ML
The problem
A reinforcement learning policy network is normally an arbitrary choice — some stack of layers picked because it trains well. C. elegans offers an alternative: the only nervous system that has been completely mapped, all 302 neurons and roughly 3,363 synaptic connections known. Rather than using that wiring diagram as a fixed template, I made the worm's own synaptic weights the trainable parameters of a PPO agent, with the biological structure enforced throughout training — which makes it possible to ask whether the wiring itself carries a useful prior, rather than just borrowing the neuron count.
What I built
- Built a 20×20 grid-world environment for chemotaxis — navigating a chemical gradient toward food — on Mesa, wrapped as a Gymnasium environment.
- Loaded the full C. elegans connectome from OpenWorm's c302 connectivity data into two 302×302 adjacency matrices: chemical synapses, directional and signed by neurotransmitter class, and electrical gap junctions, symmetric.
- Wrote a custom Actor-Critic policy that uses the connectome directly as its network, injecting sensory input at the sensory neurons and reading action logits and a value estimate off the motor neurons.
- Enforced three biological constraints during training — connections that don't exist in the real animal stay at zero, excitatory/inhibitory sign never flips, and gap junctions stay symmetric — so only the magnitude of real synapses is learned.
- Trained four architectures with PPO to isolate the effect of the biology: an MLP baseline, the connectome as a feature extractor, the connectome as the full policy, and a recurrent version of the full-policy connectome.
- Evaluated every architecture across multiple random seeds for mean reward and behavioral consistency, and ran a weight-drift analysis to check how far training moved the weights from their biological starting point.
What I had to learn
- Reinforcement learning from the ground up — PPO, reward shaping, and how much of RL difficulty lives in the environment design rather than the algorithm.
- How to turn real connectomics data into something trainable. The c302 connectivity data describes an organism, not a tensor, and getting from one to the other — two signed adjacency matrices a network can actually consume — was most of the engineering.
- That constraining a network to stay biologically valid changes how you have to think about optimization: three separate invariants (existence, sign, symmetry) had to be re-enforced after every gradient step rather than built in once and left alone.
Results
- Across four architectures, enforcing the connectome's biological structure during training improved the consistency of learned behavior rather than degrading it, and biologically constrained networks matched or exceeded an unconstrained baseline on mean performance.
- A weight-drift analysis confirmed training preserved the connectome's wiring exactly while only reshaping the strength of existing synapses.
- Chemotaxis results are complete and written up in a draft paper.
What I'd do next
- A nociception extension and an inheritance/reproduction study, planned as follow-up work.
- Aiming the current results at a conference submission, with ongoing faculty feedback.
Stack
PythonPyTorchStable-Baselines3GymnasiumMesaWeights & Biases