Hierarchical Bayesian cognitive-model baseline (HB-Hattori2019) — tracking

Umbrella issue for the hierarchical Bayesian cognitive baseline. Spec and decisions live in
`docs/design-hb-baseline.md` and `docs/adr/0001`-`0006`; this issue holds the state.

## Goal

Reimplement the published hierarchical Bayesian Q-learning model in NumPyro, extended with a
population level, so a cognitive baseline can be scored on held-out subjects on exactly the
same axis as GRU/disRNN.

## Task list

**Model** (AllenNeuralDynamics/aind-dynamic-foraging-models#63, CI green)
- [x] JAX `Hattori2019` likelihood + parity test against the numpy forager
- [x] NumPyro two-level model, per subject
- [x] Reference-priors variant (`hattori2019_stan_reference`) for validating against Stan
- [x] One-stage three-level joint model
- [x] Two-stage empirical Bayes
- [x] Held-out adaptation + posterior-predictive scoring
- [x] Batched adaptation **and** batched session scoring
- [x] Fit persistence: posterior draws, diagnostics, provenance
- [x] Posterior figures (population recovery, conditioning curve, shrinkage)
- [x] `bayes` extra + dedicated 3.11 CI job, 92% coverage

**Validation**
- [x] Agrees with the reference PyStan model to three decimals at cohort scale
- [x] Sampler geometry swept; the non-centred default already wins
- [x] Lane scaling measured; the workload is latency-bound on GPU
- [x] Two-stage vs one-stage settled on synthetic ground truth

**Integration** (AllenNeuralDynamics/aind-disrnn-wrapper#65)
- [x] `HBTrainer` on the shared `ModelTrainer` interface
- [x] Matched-conditioning rung, comparable with the GRU and the MLE baseline
- [x] W&B: shared metric keys, posterior artifacts, per-subject table — round-trip verified
- [x] Wrapper + models git SHAs recorded per fit

**Study** (#73)
- [x] Glossary, six ADRs, design note with 20 settled decisions
- [x] `studies/08-hb-vs-gru-heldout/` with launch provenance; validator clean
- [x] First real-data point, D≈10
- [ ] **D ladder at D≈30 and D≈100 on the current code** — never run since batching landed
- [ ] Regenerate the figures from a real-data fit rather than synthetic

**Open elsewhere**
- [ ] Standalone recovery suite: interval coverage, variance-component recovery, model recovery, partial-pooling-beats-MLE — models#62
- [ ] NumPyro tutorial notebook — models#62
- [ ] `two_stage.py` at 60% coverage; least-trodden path now that it is a fallback

Core implementation: AllenNeuralDynamics/aind-dynamic-foraging-models#62

## Rolling plan

_Last updated 2026-08-30 16:01 PT. This section is the live plan; edit it in place as work lands.
Everything above is the settled record._

### Next up

0. **Smoke the `run_hpc` path once.** HB now runs through the shared entrypoint with no
script of its own, and that path has never executed end to end -- the last real run used
the deleted `run_hb.py`. Four times in this work something passed locally and failed in
the place it actually runs, so this is a gate, not a formality.
`sbatch variants/one-stage-ladder/smoke.sbatch`
*Blocks:* item 1.
1. **Run the D ladder on the current code** — D≈30 (`0.049`) and D≈100 (`0.163`), one-stage,
full `--few-shot-k 0 1 2 4 8` now that a rung costs minutes rather than hours. The only
real-data point that exists is D≈10, from before batching landed.
*Blocks:* every scientific claim about HB vs GRU.
2. **Regenerate the figures from a real-data fit.** The committed ones are synthetic with
known truth, which is right for validating the plotting and wrong for reading as results.
*Blocked by:* item 1.

### Recovery suite — models#62

Not started, and the largest remaining body of work. Runs on synthetic ground truth, so it
is independent of items 1 and 2 and can proceed in parallel.

- Parameter recovery at session and subject level
- **Interval coverage** — the claim MLE structurally cannot make, and the main scientific
justification for going Bayesian at all
- Variance-component recovery (`sigma_m`, `S`) — most likely of these to be wrong
- Partial pooling beats per-session MLE, gap widening as sessions shorten
- Model recovery across the three forager families
- Prior-predictive check; sampler gates

The tutorial notebook is **delegated to a separate agent** and consumes this suite's
outputs, so this is on its critical path.

### Smaller, unblocked

- `HB-Bari2019` (adds a choice kernel), then `HB-CompareToThreshold` (its `act()` carries an
explore/exploit branch, so it is the awkward one to express as a scan)
- `two_stage.py` sits at 60% coverage — the least-trodden path now that one-stage is the
default, and therefore the most likely to rot unnoticed
- Wrapper and dispatcher have no CI configured; only models#63 is actually gated

### Deferred, with the trigger to revisit

- **SVI** — out of scope while MCMC is affordable. Revisit if a full-cohort joint fit fails
to converge, which is a problem SVI solves and more compute does not.
- **Vmapped per-subject sampler** for two-stage stage 1 — only worth building if two-stage
is ever promoted back from fallback.
- **Session packing** — trades lanes for memory at constant depth. Needed only if the flat
lane regime ends; verified to ~20k lanes, and a 16-chain full-cohort fit needs ~400k.

### Known risks

- Flat lane scaling is **verified to ~20k lanes**; the full cohort needs roughly 20x that.
If it saturates, item 1 at D≈614 will not behave like the smaller rungs.
- `test_logistic_regression.py` fails under numpy 2.x in the models repo. Pre-existing,
verified at `54a8d3a`, and out of scope here — but it will bite when the Python floor moves.

## Key decisions

| | |
|---|---|
| Framework | NumPyro/JAX, not PyStan (ADR-0001) |
| Hierarchy | population -> subject -> session; `M=0, S=1` recovers the published model (ADR-0002) |
| Held-out metric | pointwise lppd averaged in probability space (ADR-0003) |
| Parameterisation | `aind-dynamic-foraging-models` conventions, not the reference Stan model's (ADR-0004) |
| Population level | pools both location and scale (ADR-0005) |
| Dependency boundary | `bayes` extra; JAX code never imports the numpy foragers at runtime (ADR-0006) |

## Findings worth not re-deriving

The published reference model is the hierarchical counterpart of the `Hattori2019` preset.
Two traps in that mapping: `aF` is a **retention** factor equal to `1 - forget_rate_unchosen`
(despite being called a forgetting rate), and `bias` is `-biasL`.

The paper's "non-informative (uniform)" priors and the code's `mu_p ~ normal(0,1)` are the
same thing, not a contradiction: `Phi(X)` with `X ~ N(0,1)` is exactly `Uniform(0,1)`.

## Main risks

- Whether `sigma_m` is identifiable enough to justify pooling scale — ADR-0005 records the fallback.
- Whether a joint three-level NUTS fit at ~100 subjects x 40 sessions is affordable at all, which is why two-stage empirical Bayes comes first.