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

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

Fig 7. Value RNNs trained on Babayan et al. [10] reproduce Belief RPEs and learn
belief-like representations. a. Task environment of Babayan et al. [10]. Each trial con-
sists of an odor and a subsequent reward. The reward amount depends on the block iden-
tity, which is resampled uniformly every five trials (green). b. Average phasic dopamine
activity in the VTA of mice trained on the task at the time of odor (left) and reward (right)
delivery. Activity is shown separately as a function of the trial index within the block
(x-axis) and the current/previous block identity (colors). Reproduced from Babayan et al.
[10]. c. Average RPEs of the Belief model (dashed lines) and an example Value RNN
(solid lines). Same conventions as panel b. d. Total variance of beliefs explained (R 2)
using a linear transformation of model activity. Same conventions as Fig 4B. e. Cross-
validated log-likelihood of linear decoders trained to estimate true states using RNN ac-
tivity. Same conventions asFig 4C. f. Dynamics of beliefs in the absence of observations.
Same conventions as Fig 5A. g. Trajectories of an example Value RNN’s activity, in the
2D subspace identified using PCA, during an example trial from Block 1 (left) and Block
2 (right). These two dimensions explained 69% of the total variance in the Value RNN’s
activity across trials. Putative ITI states indicated as purple circles. Same conventions as
Fig 5B.
19
.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

approach, characterizing the fixed points of beliefs in this task. Similar to Task 1 of Starkweather et al. [7]281
(Fig 5A), beliefs in the present task should have a fixed point at the ITIs for Block 1 and Block 2 (Fig 7F).282
To assess whether this was the case in the Value RNN, we visualized an example RNN’s activity on the last283
few trials of each block, when the network should be confident as to the current block’s identity (given the284
reward observations on previous trials). During these trials, we observed two non-overlapping trajectories of285
activity for each block (Fig 7G). Following a reward observation, the RNN’s activity converged to a distinct286
location in state space corresponding to that block’s ITI. This suggested the RNN had two fixed points, as287
in the belief representation. In reality, these were not both truly fixed points, as the RNN’s activity did288
eventually return to a single fixed point given enough time without an observation (Fig S3A). However, the289
RNN’s two putative ITI states remained distinct across the range of ITI durations present in the training290
data (Fig S3B), allowing the network to keep these trajectories (and thus the states corresponding to each291
block) separate. These analyses suggest that Value RNNs trained on this task also exhibited belief-like292
representations.293
Untrained RNNs could also be used to estimate value and encode beliefs294
In the sections above we analyzed Value RNNs whose representations had been trained through TD learning.295
Here we take an alternative approach, inspired by reservoir computing, and consider the representations of296
untrained RNNs. In reservoir computing, a static dynamical system, or “reservoir,” is combined with a297
learned linear readout. Given an appropriately initialized reservoir (e.g., an RNN), this approach can be298
used to approximate any nonlinear function [28]. Inspired by this approach, we explored whether we could299
choose a random initialization of our RNNs such that it was only necessary to learn a linear weighting of300
the RNN’s representation to form its value estimate (i.e., ϕ in Eqn. (11) was fixed throughout training).301
Because this model resembles an echo state network (“ESN”; a reservoir computer whose reservoir is an302
RNN [29]), we will refer to this model as a Value ESN.303
We initialized each RNN by sampling the matrix of recurrent weights as a random orthogonal matrix304
scaled by a single gain parameter [30], an approach commonly used to initialize RNNs (see Materials and305
Methods). The gain effectively modulated the duration of the network’s transient response to inputs (Fig306
20
.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 8. Untrained RNNs can also be used to estimate value. a.Time-varying activations
of 20 example units in response to an odor input, in an untrained RNN (“ESN”) with 50
units, initialized with a gain of 0.1.b. Same as panela, but for an initialization gain of 1.9,
and a wider range shown on the x-axis.c. Number of time steps it took each ESN’s activity
to return to its fixed point following an odor observation (“odor memory”; see Materials
and Methods), as a function of ESN initialization gain. Points labeled “>200” indicate
those that did not fit within the plot. d. RPE MSE (see Fig 3D) as a function of ESN
initialization gain, after training each ESN’s value weights to estimate value during Task
2. Dashed line indicates average levels of a Task 2 Value RNN with the same number of
units. e. BeliefR2 (see Fig 4B) as a function of ESN initialization gain. Same conventions
as panel d.
21
.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

8A-B), such that larger gains led to larger odor memories (Fig 8C). In agreement with previous work, we307
found that when the gain was above a critical value (“the edge of chaos” [30]), the network’s activity never308
decayed back to baseline (e.g., gain> 2 in Fig 8C).309
To see if Value ESNs could be trained to estimate value, we trained Value ESNs on Starkweather Task310
2, varying the gain for each network. During training, only the network’s value weights were modified to311
estimate the value function. We found that for a range of different gains, the resulting Value ESN could312
estimate value nearly as well as Value RNNs, and recapitulate the experimentally observed dopamine pat-313
terns (Fig 8D). Interestingly, the Value ESN’s representations could even be linearly transformed to match314
beliefs (Fig 8E). We emphasize that the Value ESN’s representation was determined solely by its initial-315
ization; given an appropriate initialization, the Value ESN’s representation could effectively act as a set of316
temporal basis functions, allowing the network to match any downstream target, including beliefs. However,317
in terms of dynamics, the Value ESN’s dynamics differed substantially from those of the beliefs: For the318
best performing Value ESN (with a gain of 1.9), following an odor observation, the RNN’s activity returned319
to its fixed point after around 200 time steps (Fig 8C), whereas Task 2 beliefs return to the ITI point in 15320
time steps. These results show that allowing the RNN to modify its representation during training led to its321
representation becoming even more belief-like than expected from a more carefully initialized RNN.322
Discussion323
We have shown that training RNNs to estimate value in partially observable environments yields representa-324
tions that resemble beliefs. Specifically, we showed that, after training, the RNN activity i) could be linearly325
transformed to approximate beliefs, ii) could be used to decode the true state in the environment, and iii)326
had a dynamical structure consistent with beliefs.327
From a theoretical perspective, using an RNN resolves the problem of how to compute belief states, by328
replacing the fine-tuned Bayesian machinery needed for beliefs with a more general learned function approx-329
imator (e.g., a recurrent neural network). Our results illustrate that, to perform tasks in partially observable330
environments, it is not necessary for an agent or animal to explicitly estimate states using a belief representa-331
tion; rather, agents can learn a sufficient representation for solving the task from observations alone. This is332
22
.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

promising from a normative perspective, as it shows how neural circuits might come to compute theoretical333
features such as beliefs without that objective needing to be explicitly learned. Moreover, there is a growing334
toolkit for reverse engineering RNN solutions [26, 27, 31], which can shed light on learned mechanisms of335
value computation.336
One potential benefit of the Value RNN over the Belief model for estimating value is the ability to sepa-337
rate the capacity of the model (i.e., the number of hidden units in the RNN) from the size of the state space338
in the environment. As we showed in Fig 6, Value RNNs with fewer units discovered a representation that339
was more compressed than beliefs, but nevertheless sufficient for performing the task at hand. Value RNNs340
with more units had more belief-like representations. This suggests a potentially useful trade-off, in which341
agents could choose to allocate more capacity to a task in exchange for more belief-like representations.342
Such a trade-off may be a relevant feature for biological organisms, who must be able to perform tasks such343
as value estimation in complex environments where it may not always be feasible to learn the full belief344
representation.345
From a methodological perspective, this work can serve as a blueprint for how to bridge analyses of neu-346
ral computation across levels of abstraction. In future work, we hope to apply this framework to test neural347
models of how animals perform associative learning via reinforcement learning. For instance, previous work348
has suggested that prefrontal cortex may perform state estimation in tasks with hidden states [6, 12, 7, 13].349
The same tools we apply here to artificial neural networks can also be applied to neural activity recorded350
from animals performing the same task. For instance, if cortex implements something like a Value RNN,351
cortical activity may show a longer transient response to odors during Task 1 than in Task 2 (Fig 5). On the352
other hand, if activity is more like a Value ESN, cortical representations should be largely the same in both353
tasks.354
Traditional models of how animals perform trace conditioning tasks like the ones we consider here make355
a variety of implicit assumptions about how animals represent the passage of time [32, 33]. For example, the356
state space shown in Fig 2A, which forms the basis of the belief representation, conceptualizes the passage357
of time in the form of microstates. Many modeling efforts require even more assumptions to account for358
scalar variability in animals’ estimates of elapsed time, such as by incorporating a more complex set of359
temporal basis functions. In our model, by contrast, the Value RNN’s time-varying representation of inputs360
23
.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

is learned through training. We observed that individual units in our RNNs were tuned to elapsed time361
relative to observations, and that the temporal precision of tuning decreased with elapsed time (Fig S1),362
both of which are standard assumptions of microstate representations [34]. Similarly tuned “time cells” have363
been observed in the striatum [35], hippocampus [36], and prefrontal cortex [37]. Our modeling suggests364
that at least some assumptions about microstate representations may be redundant in the sense that they may365
emerge naturally in recurrent networks that are trained to perform reinforcement learning. This viewpoint366
resonates with the idea (reviewed in [38]) that delay encoding can arise as an emergent property of neural367
network dynamics.368
Previous work has shown that, in animals, prefrontal cortex activity is a necessary component of ani-369
mals’ state representations [13]. This work found that inactivating prefrontal cortex in the Starkweather task370
led to animals’ RPEs in Task 2 resembling the RPEs of Task 1. This is what one would expect if prefrontal371
cortex was involved in estimating a belief in omission trials. In fact, both Tasks 1 and 2 include another form372
of uncertainty, which is the reward time on each trial. The fact that prefrontal inactivation did not interact373
with animals’ estimates of timing suggests that different neural circuits may form belief-like representations374
specific to particular types of state uncertainty (e.g., temporal uncertainty versus reward uncertainty). In the375
present work, our RNNs should be thought of as a generic computational model, and not a model of individ-376
ual brain regions. These networks had a generic architecture and only a single layer; as a result, our model377
would be unable to distinguish between different sources of uncertainty. Nevertheless, it is an interesting378
question how architectural considerations, such as layer connectivity and the dominance of feedforward379
versus recurrent connectivity, might contribute to where in the brain different belief-like representations are380
formed.381
Here we have shown that computing beliefs explicitly is not a necessary precursor for optimally per-382
forming a task in partially observable environments. Nor is it required to reproduce experimentally observed383
patterns of dopamine neuron activity. Nevertheless, beliefs are an efficient representation in that they are384
sufficient for solving any task in the same environment. Thus, beliefs may be a desirable representation385
for animals, who may need to achieve a range of different goals in the same environment (e.g., finding wa-386
ter when thirsty, but finding food when hungry) without having to learn a representation in each of these387
tasks separately. Future work should explore whether a dedicated belief mechanism is necessary in these388
24
.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