Emergence of belief-like representations through reinforcement learning — part 6/7

Cleaned-text code points [72243, 87415). Bibliography omitted here; source.md includes it.

as follows:
ℓ(cW|s1,...,s T, ˜z1,..., ˜zT ) = 1
T
TX
t=1
log(πt(st)) (17)
whereπt(st) ∈ [0, 1] is the sth
t entry of the vector πt. We calculated the log-likelihood for the belief504
model similarly, except instead of training a decoder we used πt = bt. For the Babayan task (Fig 5E),505
we calculated the log-likelihood on all trials except the first trial in each block. This was necessary for the506
beliefs to act as an upper-bound on the log-likelihood, because we defined the beliefs in a way that did not507
assume knowledge of the number of trials in each block.508
Dynamics analysis509
An RNN with parametersϕ has a hidden state that evolves aszt =fϕ(ot,zt−1). Conditioned on a particular510
constant input,o, an RNN is at a fixed point when∥fϕ(o,z)−z∥2
2≈ 0. Numerically, we can simply look511
forz where∥fϕ(o,z)−z∥2
2<ϵ . For our analyses below we choseϵ = 1× 10−5.512
Identifying fixed points (Fig 5B, Fig 7G). During training, RNNs received three distinct types of inputs,513
ot∈{[0 0], [1 0], [0r]}, which we refer to as the null (∅), odor, and reward inputs, respectively. Under the514
beliefs of the Starkweather and Babayan tasks, the odor and reward inputs always result in a change in the515
beliefs. As a result, any fixed points of the beliefs must be conditional on the null input, ∅. We therefore516
sought to identify an RNN’s fixed points conditional on a null input. To do this, we took a numerical ap-517
proach: We initialized the RNN to a random state, applied the null input until the RNN’s activity converged,518
and then repeated this process across different random states to get a candidate set of fixed points. More pre-519
cisely, we consideredN = 20 randomly selected values ofz in the testing data following an odor or reward520
observation as a set of starting seeds. For each starting seed,z0, we computed the RNN’s representation,zt,521
over time, given no further odor or reward observations:zt =fϕ(∅,zt−1). We repeated this process until522
ηt =∥zt−zt−1∥2
2<ϵ. We then added zt to our list of candidate fixed points,F. For each pair of candidate523
fixed points within a distance 1× 10−3 of each other, we considered these to be the same fixed point.524
Odor and reward memory duration (Fig 5C, Fig 5D). For each Value RNN with a single fixed point,525
we measured the network’s odor (or reward) memory as follows. We initialized each RNN to its fixed526
31
.CC-BY 4.0 International licenseavailable under a 
was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is made 
The copyright holder for this preprint (whichthis version posted April 7, 2023. ; https://doi.org/10.1101/2023.04.04.535512doi: bioRxiv preprint

point,z0, and then provided an odor (or reward) observation at time t = 1. We then measured the RNN’s527
representation,zt, over time, given no further odor or reward observations: zt = fϕ(∅,zt−1), for t >1.528
For each t, we calculated the distance of the activity from its fixed point: ηt =∥zt−z0∥2
2 (Fig 5C). We529
repeated this until ηt converged to zero, defining the odor memory (or reward memory) as the t at which530
ηt< 1× 10−3 (Fig 5D).531
Data Availability Statement532
All data and code used for analysis and plotting is available on a GitHub repository athttps://github.533
com/mobeets/value-rnn-beliefs/534
Author Contributions535
Conceptualization: Jay A. Hennig, Sandra A. Romero Pinto, Takahiro Yamaguchi, Naoshige Uchida,536
Samuel J. Gershman537
Formal analysis: Jay A. Hennig538
Investigation: Jay A. Hennig539
Funding acquisition: Samuel J. Gershman, Naoshige Uchida, Scott W. Linderman540
Methodology: Jay A. Hennig, Samuel J. Gershman541
Supervision: Scott W. Linderman, Naoshige Uchida, Samuel J. Gershman542
Visualization: Jay A. Hennig543
Writing - original draft: Jay A. Hennig544
Writing - review & editing: Jay A. Hennig, Sandra A. Romero Pinto, Takahiro Yamaguchi, Scott W.545
Linderman, Naoshige Uchida, Samuel J. Gershman546
32
.CC-BY 4.0 International licenseavailable under a 
was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is made 
The copyright holder for this preprint (whichthis version posted April 7, 2023. ; https://doi.org/10.1101/2023.04.04.535512doi: bioRxiv preprint

Acknowledgments547
This work was funded by NIH U19 NS113201-01 and Air Force Office of Scientific Research grant FA9550-548
20-1-0413.549
References
[1] Wolfram Schultz, Peter Dayan, and P Read Montague. A neural substrate of prediction and reward.
Science, 275:1593–1599, 1997.
[2] Hannah M Bayer and Paul W Glimcher. Midbrain dopamine neurons encode a quantitative reward
prediction error signal. Neuron, 47(1):129–141, 2005.
[3] Jeremiah Y Cohen, Sebastian Haesler, Linh V ong, Bradford B Lowell, and Naoshige Uchida. Neuron-
type-specific signals for reward and punishment in the ventral tegmental area. nature, 482(7383):
85–88, 2012.
[4] Neir Eshel, Michael Bukwich, Vinod Rao, Vivian Hemmelder, Ju Tian, and Naoshige Uchida. Arith-
metic and local circuitry underlying dopamine prediction errors. Nature, 525(7568):243–246, 2015.
[5] Leslie Pack Kaelbling, Michael L Littman, and Anthony R Cassandra. Planning and acting in partially
observable stochastic domains. Artificial intelligence, 101(1-2):99–134, 1998.
[6] Nathaniel D Daw, Aaron C Courville, and David S Touretzky. Representation and timing in theories
of the dopamine system. Neural computation, 18(7):1637–1677, 2006.
[7] Clara Kwon Starkweather, Benedicte M Babayan, Naoshige Uchida, and Samuel J Gershman.
Dopamine reward prediction errors reflect hidden-state inference across time. Nature neuroscience,
20(4):581–589, 2017.
[8] Armin Lak, Kensaku Nomoto, Mehdi Keramati, Masamichi Sakagami, and Adam Kepecs. Midbrain
dopamine neurons signal belief in choice accuracy during a perceptual decision. Current Biology, 27:
821–832, 2017.
33
.CC-BY 4.0 International licenseavailable under a 
was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is made 
The copyright holder for this preprint (whichthis version posted April 7, 2023. ; https://doi.org/10.1101/2023.04.04.535512doi: bioRxiv preprint

[9] Stefania Sarno, Victor de Lafuente, Ranulfo Romo, and N ´estor Parga. Dopamine reward prediction
error signal codes the temporal evaluation of a perceptual decision report.Proceedings of the National
Academy of Sciences, 114:E10494–E10503, 2017.
[10] Benedicte M Babayan, Naoshige Uchida, Samuel Gershman, et al. Belief state representation in the
dopamine system. Nature communications, 9(1):1–10, 2018.
[11] John G Mikhael, HyungGoo R Kim, Naoshige Uchida, and Samuel J Gershman. The role of state
uncertainty in the dynamics of dopamine. Current Biology, 32:1077–1087, 2022.
[12] Robert C Wilson, Yuji K Takahashi, Geoffrey Schoenbaum, and Yael Niv. Orbitofrontal cortex as a
cognitive map of task space. Neuron, 81(2):267–279, 2014.
[13] Clara Kwon Starkweather, Samuel J Gershman, and Naoshige Uchida. The medial prefrontal cortex
shapes dopamine reward prediction errors under state uncertainty. Neuron, 98(3):616–629, 2018.
[14] Samuel J Gershman and Naoshige Uchida. Believing in dopamine. Nature Reviews Neuroscience, 20:
703–714, 2019.
[15] Alexandre Pouget, Jeffrey M Beck, Wei Ji Ma, and Peter E Latham. Probabilistic brains: knowns and
unknowns. Nature Neuroscience, 16(9):1170–1178, 2013.
[16] Nathaniel D Daw, Yael Niv, and Peter Dayan. Uncertainty-based competition between prefrontal and
dorsolateral striatal systems for behavioral control. Nature Neuroscience, 8:1704–1711, 2005.
[17] Rajesh PN Rao. Decision making under uncertainty: a neural model based on partially observable
markov decision processes. Frontiers in computational neuroscience, 4:146, 2010.
[18] Pascal Poupart and Craig Boutilier. Value-directed compression of POMDPs. Advances in Neural
Information Processing Systems, 15, 2002.
[19] Nicholas Roy, Geoffrey Gordon, and Sebastian Thrun. Finding approximate POMDP solutions through
belief compression. Journal of Artificial Intelligence Research, 23:1–40, 2005.
34
.CC-BY 4.0 International licenseavailable under a 
was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is made 
The copyright holder for this preprint (whichthis version posted April 7, 2023. ; https://doi.org/10.1101/2023.04.04.535512doi: bioRxiv preprint

[20] Matthew Botvinick, Jane X Wang, Will Dabney, Kevin J Miller, and Zeb Kurth-Nelson. Deep rein-
forcement learning and its neuroscientific implications. Neuron, 107(4):603–616, 2020.
[21] Tianwei Ni, Benjamin Eysenbach, and Ruslan Salakhutdinov. Recurrent model-free RL can be a strong
baseline for many POMDPs. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari,
Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine
Learning, volume 162 of Proceedings of Machine Learning Research, pages 16691–16723. PMLR,
17–23 Jul 2022.
[22] Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
[23] Samuel J Gershman and Nathaniel D Daw. Reinforcement learning and episodic memory in humans
and animals: an integrative framework. Annual Review of Psychology, 68:101–128, 2017.
[24] Kyunghyun Cho, Bart Van Merri ¨enboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Hol-
ger Schwenk, and Yoshua Bengio. Learning phrase representations using rnn encoder-decoder for
statistical machine translation. arXiv preprint arXiv:1406.1078, 2014.
[25] David Sussillo and Omri Barak. Opening the black box: low-dimensional dynamics in high-
dimensional recurrent neural networks. Neural computation, 25(3):626–649, 2013.
[26] Niru Maheswaranathan, Alex Williams, Matthew Golub, Surya Ganguli, and David Sussillo. Univer-
sality and individuality in neural dynamics across large populations of recurrent networks. Advances
in neural information processing systems, 32, 2019.
[27] Saurabh Vyas, Matthew D Golub, David Sussillo, and Krishna V Shenoy. Computation through neural
population dynamics. Annual review of neuroscience, 43:249–275, 2020.
[28] Mantas Luko ˇseviˇcius and Herbert Jaeger. Reservoir computing approaches to recurrent neural network
training. Computer science review, 3(3):127–149, 2009.
[29] Herbert Jaeger. Echo state network. scholarpedia, 2(9):2330, 2007.
35
.CC-BY 4.0 International licenseavailable under a 
was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is made 
The copyright holder for this preprint (whichthis version posted April 7, 2023. ; https://doi.org/10.1101/2023.04.04.535512doi: bioRxiv preprint

[30] Andrew M Saxe, James L McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics
of learning in deep linear neural networks. arXiv preprint arXiv:1312.6120, 2013.
[31] Jimmy Smith, Scott Linderman, and David Sussillo. Reverse engineering recurrent neural networks
with Jacobian switching linear dynamical systems. Advances in Neural Information Processing Sys-
tems, 34:16700–16713, 2021.
[32] Samuel J Gershman, Ahmed A Moustafa, and Elliot A Ludvig. Time representation in reinforcement
learning models of the basal ganglia. Frontiers in computational neuroscience, 7:194, 2014.
[33] Vijay Mohan K Namboodiri. How do real animals account for the passage of time during associative
learning? Behavioral Neuroscience, 2022.
[34] Elliot A Ludvig, Richard S Sutton, and E James Kehoe. Stimulus representation and the timing of
reward-prediction errors in models of the dopamine system. Neural Computation, 20:3034–3054,
2008.
[35] Gustavo BM Mello, Sofia Soares, and Joseph J Paton. A scalable population code for time in the
striatum. Current Biology, 25:1113–1122, 2015.
[36] Christopher J MacDonald, Kyle Q Lepage, Uri T Eden, and Howard Eichenbaum. Hippocampal “time
cells” bridge the gap in memory for discontiguous events. Neuron, 71:737–749, 2011.
[37] Zoran Tiganj, Min Whan Jung, Jieun Kim, and Marc W Howard. Sequential firing codes for time in
rodent medial prefrontal cortex. Cerebral Cortex, 27:5663–5671, 2017.
[38] Joseph J Paton and Dean V Buonomano. The neural basis of timing: distributed mechanisms for
diverse functions. Neuron, 98(4):687–705, 2018.
[39] Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neu-
ral networks. In Proceedings of the thirteenth international conference on artificial intelligence and
statistics, pages 249–256. JMLR Workshop and Conference Proceedings, 2010.
36
.CC-BY 4.0 International licenseavailable under a 
was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is made 
The copyright holder for this preprint (whichthis version posted April 7, 2023. ; https://doi.org/10.1101/2023.04.04.535512doi: bioRxiv preprint

Supplemental Figures
Fig S1. RNN activity before and after training on Starkweather Task 2. a. Obser-
vations, states, beliefs, and RNN activations on two example trials from Task 2. b-c.
RNN unit activity (individually normalized to span between 0 and 1), with units sorted by
time of peak activation on held-out trials, on an RNN before (panel b) and after (panel c)
training. Both before and after after training, RNN units exhibited tuning to elapsed time
following observations, with variance that scaled with elapsed time.
S1
.CC-BY 4.0 International licenseavailable under a 
was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is made 
The copyright holder for this preprint (whichthis version posted April 7, 2023. ; https://doi.org/10.1101/2023.04.04.535512doi: bioRxiv preprint

Fig S2. Value RNNs trained on the Babayan task recapitulate dopamine activity
and belief RPEs in response to intermediate reward sizes. a-b. Average dopamine
response on trial 1 (panel a) and trial 2 (panel b) during probe sessions including blocks
with intermediate reward sizes. Circles and lines depict mean ± SE across N = 11
animals. Reproduced from Babayan et al. [10]. c-d. Same as panels a-b, but for the
RPEs of the Belief model (black) and Value RNNs (purple). Value RNNs were trained
on sessions including only blocks with rewards rt∈{1, 10}, as in the main text. Value
weights for the Belief model and Value RNNs were fit using a test session including 39
blocks each withrt = 1 andrt = 10, and 3 blocks each with rt∈{2, 4, 6, 8}, similar to
the proportions used in Babayan et al. [10]. RPEs were then measured on a different test
session. Purple circles and lines depict mean± SE acrossN = 12 models.
S2
.CC-BY 4.0 International licenseavailable under a 
was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is made 
The copyright holder for this preprint (whichthis version posted April 7, 2023. ; https://doi.org/10.1101/2023.04.04.535512doi: bioRxiv preprint