[Scientific Roadmap] Behavioral Foundation Model for Dynamic Foraging v0.1

# Behavioral Foundation Model for Dynamic Foraging — master tracking issue

Umbrella issue tracking the prioritized analysis roadmap for the disRNN/GRU behavioral foundation model. Child issues below are grouped by priority, mirroring the *Prioritized analysis TODO (project plan)* in the planning doc.

📄 **Live planning doc:** https://docs.google.com/document/d/1Xk4Zi9QtQcUNJs4SMvZEb_HbGq0rU8wW3NT8OyNaet4/edit?usp=sharing

**Related repos:** [`aind-disrnn-wrapper`](https://github.com/AllenNeuralDynamics/aind-disrnn-wrapper) · [`aind-disrnn-dispatcher`](https://github.com/AllenNeuralDynamics/aind-disrnn-dispatcher) · [`aind-disrnn-result-access`](https://github.com/AllenNeuralDynamics/aind-disrnn-result-access) | **Project board:** [org project #184](https://github.com/orgs/AllenNeuralDynamics/projects/184/views/9)

**Target:** manuscript by June 2027.

---

## 📊 Status — 2026-07-12 (4 / 17 original closed · 2 new children added)

Four committed studies now exist under [`studies/`](https://github.com/AllenNeuralDynamics/aind-disrnn-dispatcher/tree/main/studies):

| study | question | status |
|---|---|---|
| [`01-gru-scaling-law`](https://github.com/AllenNeuralDynamics/aind-disrnn-dispatcher/tree/main/studies/01-gru-scaling-law) | does more mice → better prediction of *unseen* mice? | ✅ r1–r9 |
| [`02-gru-scaling-law-ignore`](https://github.com/AllenNeuralDynamics/aind-disrnn-dispatcher/tree/main/studies/02-gru-scaling-law-ignore) | does a headroom-ier (3-way, incl. ignore) target scale? | ✅ 48/48 |
| [`03-disrnn-beta-scan`](https://github.com/AllenNeuralDynamics/aind-disrnn-dispatcher/tree/main/studies/03-disrnn-beta-scan) | can we force the disRNN interaction bottleneck to sparsify? | ✅ 43/48 + suppl. |
| [`04-gru-vs-disrnn-embedding-recovery`](https://github.com/AllenNeuralDynamics/aind-disrnn-dispatcher/tree/main/studies/04-gru-vs-disrnn-embedding-recovery) | do embeddings recover *known* generative structure? | ✅ r1–r2 |

**The one-paragraph scientific state.** On 2-way L/R choice the model is near a **predictability ceiling**: held-out-mouse likelihood saturates by ~100 mice, and the population mean already predicts a new mouse to within ~0.3% of full adaptation. **But the foundation-model claim survives, on three legs:** (1) the population GRU **beats a per-mouse classical RL baseline by +0.0136 at D=614 on 100% of held-out mice** (p~3e-26) — the dominant signal; (2) on the **3-way ignore target** likelihood **keeps climbing with capacity** and never plateaus (H=256/D=614 = 0.7315), so the ceiling was a property of the *metric*, not the model; (3) on synthetic ground truth, embeddings **recover true generative parameters and model family at 97.5–100%**, exactly where a correctly-specified baseline breaks down. The interpretable disRNN replicates that recovery for a **~4–6 point** likelihood cost — interpretability is nearly free.

**The two gaps that most threaten the story:**
1. **#20 — no hierarchical-Bayes population baseline.** The RL baseline is fit *per-mouse independently*, so it has **no D-axis**. The GRU's headline win is partly a "population vs per-mouse" win *by construction*. This is the fair-comparison baseline and it is not built.
2. **#16 — no full-scale disRNN.** Every interpretability conclusion (`03`) is at **D=100**. The disRNN has never been trained on the full 600–800-mouse dataset.

**Where P1 stands: essentially not started.** #24/#26/#28/#29 have no code. The chain has a clear entry point — **#22's missing logistic-regression readout blocks #26 and half of #29** — and the cheapest first win is **#27/#28**, where [`embedding_space_analysis.py`](https://github.com/AllenNeuralDynamics/aind-disrnn-wrapper/blob/main/code/post_training_analysis/embedding_space_analysis.py) already plots subject/session embeddings against metadata and just needs pointing at the D=614 checkpoint.

**New modeling direction (#57).** A committed design note — [`docs/design-hierarchical-vi-foundation-model.md`](https://github.com/AllenNeuralDynamics/aind-disrnn-dispatcher/blob/main/docs/design-hierarchical-vi-foundation-model.md), https://github.com/AllenNeuralDynamics/aind-disrnn-dispatcher/pull/54, *not yet implemented* — reframes the FM as an explicit **hierarchical mixed-effects model trained by amortized VI (VAE-style)**. Its key claim: the cognitive **hierarchical-Bayes model (#20) and the VI foundation model are the same statistical object**, differing only in whether the per-unit latent is hand-specified cognitive parameters or a learned RNN latent. **#57 is therefore the other arm of #20, not a replacement for it** — and the two must share a held-out matrix and metric from the start, or they can't be compared later without a re-run.

---

## P0 — Critical: model training, development, pipeline
- [ ] #16 🟡 Large-scale disRNN/GRU training across full dataset (600–800 mice) — *GRU done @ D=614; **disRNN full-scale missing***
- [x] #17 ✅ Data split & cross-validation: held-out-subject transfer at scale
- [ ] #18 🟡 Hyperparameter sweep for GRU / DisRNN — *N×D + β grids done; **GRU HPs never tuned***
- [x] #19 ✅ Scaling law of performance (model capacity × dataset size)
- [ ] #20 🟡 Complete baseline models incl. hierarchical Bayes — *Bari baseline done; **hierarchical Bayes = top P0 gap***
- [x] #21 ✅ Synthetic Q-learning population + recovery matrix
- [ ] #22 🟡 Generative evaluation suite — *switch-prob + history done; **run-length dist. + logistic coeffs missing***
- [x] #23 ✅ Include-ignore-trials model version (3 actions) — *engagement **state** split out → #58*
- [ ] #57 🆕 Hierarchical (mixed-effects) FM via amortized VI — *design committed, **not implemented**; the VAE-style arm of #20*
- [ ] #58 🆕 Explicit engagement state — *break study-02's **0.47 ignore-recall ceiling***

## P1 — Model & embedding analysis (interpretability, mechanism)
- [ ] #24 🔴 GRU state space & dynamics: fixed-point analysis + embedding reshaping — *not started (unblocked)*
- [ ] #25 🟡 DisRNN update-rule interpretation + embedding modulation — *sparsity solved; **interpretation not started***
- [ ] #26 🔴 Embedding → policy map — ***blocked on #22***
- [ ] #27 🟡 Embedding vs. metadata / Q-learning parameters — *tooling exists, **no result yet***
- [ ] #28 🔴 Embedding: within-subject learning trajectory — *not started*
- [ ] #29 🔴 Discover cognitive mechanisms (learning-rate / uncertainty modulation) — *blocked on #22 + #25*

## P2 — Extensions
- [ ] #30 🔴 Neural correlation with learned cognitive variables
- [ ] #31 🔴 Mixed sim + animal training + deviation analysis — *synthetic half already built in `04`; **cheaper than its P2 label***
- [ ] #32 🔴 Cross-species comparison / transfer

---
_19 child issues (17 original + #57, #58) · priorities mirrored on the project board's Priority field._