Modulation of feature attention by reward prediction error explains value learning behavior

A reinforcement-learning model with a feature-attention front end is fitted to colour-value learning behaviour from two macaques. Five candidate RPE→attention transfer functions are compared — None, Linear, Absolute, Quadratic, Switch — crossed with two attention architectures (single-focus vs multi-focus).

Behavioural data are from Jahn et al. 2024.

---

Layout

path	what
vectorized_model.py	the model; new_simulate() runs it closed- or open-loop
misc_func.py	tuning curves and the five transfer functions (attentional_boost_scaler, multi_value_attention)
behav_core.py	shared loading/helpers for the behavioural analyses
parameters/	task and behavioural input data (.mat)
analysis_results/	small intermediate results the plotting scripts read
figures/	the figures as they appear in the paper
neural/	neural-data analyses (raw neural data is external, see below)

Reproducing the figures

Simulation output is ~2.5 GB and is not in the repository, so run the two pipeline steps first.

pip install -r requirements.txt

# 1. sweep attention amplitude for every model x focus   (~6 h; writes sweep_results/)
python bestfit_amplitude_sweep.py

# 2. re-simulate each model at its own best-fit amplitude (~50 min; writes results_bestfit/)
python run_bestfit_sims.py

# 3. any figure script
python accuracy_plot.py

Step 1 finds, for each model and focus type, the modulation amplitude a (attention gain bounded to [1-a, 1+a]) minimising the pooled two-monkey learning-curve MSE. Step 2 re-runs every model at its own best-fit amplitude, which is what all downstream figures use. Both write progress to stdout and can be backgrounded.

Every plotting script accepts SIM_RESULTS_DIR if the simulations live elsewhere:

SIM_RESULTS_DIR=/path/to/results_bestfit python accuracy_plot.py

Figure → script

figure in the paper	produced by
accuracy_plot.pdf	accuracy_plot.py
explore_exploit_plot.pdf	explore_exploit_plot.py
confidence_plot.pdf	entropy_rt_plot.py
postloss_single_vs_multi.pdf	postloss_plot.py
postloss_vs_amplitude.pdf	postloss_vs_amplitude.py
postloss_monotonicity.pdf	postloss_monotonicity.py
accuracy_vs_amplitude.pdf	accuracy_vs_amplitude.py
bestfit_amplitude.pdf	bestfit_amplitude_plot.py
frontier_single_vs_multi.pdf	frontier_plot.py
cross_monkey_generalization.pdf	cross_monkey_generalization.py
glm2x2_MonkeyB.pdf, glm2x2_MonkeyS.pdf	neural/plot_glm2x2_figures.py
neural_permonkey_glm.pdf	neural/plot_glm2x2_manuscript.py
shift_null_target.pdf	neural/plot_shift_null_target.py

Statistics quoted in the Results text come from manuscript_stats.py; the decay-constant table comes from tau_table.py. Both read results_bestfit/.

Neural analyses

neural/ fits, per neuron and time window, two Poisson GLMs with HC0 robust standard errors:

spikes ~ 1 + z(previous reward) + z(previous RPE)
spikes ~ 1 + z(previous reward) + z(|previous RPE|)