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

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

Methods). We quantified performance on held-out sessions by evaluating the log-likelihood of the decoded176
estimates. Because the beliefs capture the posterior distribution of the state given the observations under the177
true generative model, the log-likelihood of the beliefs serves as a ceiling on performance. We found that178
the log-likelihoods of the decoders trained on the RNN’s activity approached those of the beliefs, and easily179
outperformed the decoders that used the activity of untrained RNNs (Fig 4C). Thus, training the RNN to180
estimate value resulted in a representation that could be used to more accurately decode the true state.181
RNN activity exhibited belief-like dynamics182
One potential shortcoming of the above analyses is that we have not yet accounted for the dynamical nature183
of the belief representation: Belief updating can be thought of as a dynamical system describing how the184
posterior probability of each state evolves as a function of the observations. We therefore took a dynamical185
systems perspective [25, 26, 27] and asked whether the dynamics of RNN activity resembled the dynamics186
of the beliefs in each task.187
We first asked whether beliefs and RNNs had similar fixed point structure, a standard approach to188
characterizing the computations performed by dynamical systems [25, 26, 27]. Here, a “fixed point” is189
a belief state that remains unchanged in the absence of any new observations. Specifically, we considered190
the fixed points of beliefs in the absence of observations (Fig 5A). In both tasks, the duration of the ITI is191
sampled from a geometric distribution, which has a constant hazard function. Thus, if the agent believes192
it is in the ITI (i.e., waiting for an odor), it should maintain this belief for as long as it receives no new193
observations (∅). Thus, the ITI belief state is a fixed point of the belief updates in both tasks (Fig 5A).194
Now consider when the agent is in the ISI (e.g., following an odor observation). In Task 2 (Fig 5A, bottom195
panel), the agent should maintain a nonzero belief in the ISI only for as long as there are possible reward196
times remaining—i.e., the first 2.8s, or 14 time steps—but after that point it should return to the ITI state.197
Thus, the ITI state is the only fixed point of the Task 2 beliefs. In Task 1, by contrast, there are no omission198
trials, and so the beliefs are simply undefined when there are no observations for more than 14 time steps.199
Nevertheless, for the purposes of characterizing the fixed points of beliefs, we can ask what an agent with200
Task 1 beliefs could do when faced with an omission trial. In this sense, an agent could maintain a belief in201
14
.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 5. Value RNN dynamics resembled belief dynamics in each task. a. Dynamics
of beliefs in Task 1 (top) and Task 2 (bottom). Black arrows indicate transitions between
states in the absence of observations (∅) as a function of elapsed time, t, following an
odor observation. ‘X’ indicates an unconstrained duration, and a dashed arrow indicates
a transition that happens only when ‘X’ is finite. b. RNN activity at each time step (small
black dots with connected lines) during an example trial in a 2D subspace identified using
PCA. Putative ITI fixed point indicated as purple circle. Vectors indicate the response to
odor (black) and reward (red). Activity during an omission trial is shown in cyan, though
note that omission trials were present in training data only for Task 2. c-d. Average
normalized distance of each model’s activity from its fixed point (identified numerically)
following an odor (panel c) or reward (panel d) observation, over time. To allow compar-
ing distances across models, each model’s distances were normalized by the maximum
distance following each observation.
the ISI for any number of time steps X > 14, resulting in two fixed points when X→∞, and one fixed202
point otherwise (Fig 5A, top panel). Thus, Task 1 beliefs can decay to the ITI fixed point at any point after203
14 time steps, and may potentially have two fixed points.204
We asked whether the Value RNNs in each task exhibited similar dynamics. To build intuition, for each205
task we visualized the activity of a Value RNN on example trials (Fig 5B). To visualize the activity of the206
RNN’s 50 units, we used principal components analysis (PCA) to project the RNN’s activity into the top207
two dimensions that captured the most variance of the activity across trials; these two dimensions explained208
83% and 79% of the total variance in the Task 1 and Task 2 Value RNN, respectively. We observed that209
each RNN’s activity was quite stable during the ITI (purple circle), suggestive of a fixed point. This activity210
15
.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

was then perturbed by an odor observation (black vector), and continued to move through state space during211
the ISI. On rewarded trials, in response to a reward (red vector), RNN activity gradually returned to the212
same ITI location it started from (purple circle). We noted that the activity of both RNNs would also have213
converged to its original ITI location had the reward been omitted (cyan traces). Interestingly, this was true214
even for the Task 1 RNN, which did not experience omission trials during training. These visualizations215
suggested that these two example Value RNNs had a single fixed point (corresponding to the ITI), which216
we then confirmed numerically (see Materials and Methods). We next used the same numerical approach to217
identify the fixed points across all trained Value RNNs, and found similar results. In fact, only two Value218
RNNs had more than one fixed point; these were both Task 1 RNNs, which had a fixed point for both the ITI219
and the ISI. Overall, these analyses indicated that Value RNNs had a fixed point structure consistent with220
those of beliefs.221
Despite the fact that most Value RNNs had a single fixed point regardless of which task they were222
trained on, we noted that the temporal dynamics of RNN activity differed across the two tasks following223
odor observations. For example, in the Task 2 RNN, following an odor observation, the activity moved224
gradually closer to the ITI state throughout the ISI (Fig 5B, bottom subpanel). These dynamics allowed the225
Task 2 RNN’s activity to return to the ITI state at the appropriate time on trials without reward (cyan trace).226
By contrast, in the Task 1 RNN, which did not experience trials without rewards during training, activity227
took much longer to return to the ITI. To quantify these differences, we initialized each RNN to its fixed228
point, provided an odor observation, and then measured the RNN’s activity over time in the absence of any229
subsequent reward. We then measured the distance of each RNN’s activity over time from its ITI fixed point230
(Fig 5C, colored traces).231
We repeated this same analysis for beliefs (Fig 5C, black trace), allowing us to characterize the two232
models’ responses to an odor as a function of the distance of their representations from their ITI fixed point.233
We reasoned that, if the RNNs learned belief-like dynamics, the activity of Task 2 RNNs should return to234
the ITI as soon as possible after time step 14 (i.e., the largest reward time), which we found was largely235
the case (Fig 5C). By contrast, in Task 1, beliefs are undefined past time step 14 (because there are no236
omission trials), and so the activity of RNNs after this point was not constrained by the task. To quantify237
these differences across tasks, we calculated the time step at which each RNN’s activity returned within238
16
.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 6. Value RNNs with larger capacity had more belief-like representations. a.
Error between the Value RNN’s RPEs and those of the Belief model (“RPE MSE”; see
Fig 3D) during Task 2, as a function of the number of units in the Value RNN. Each circle
indicates the median across the N = 12 Value RNNs with the same number of units. b.
Total variance explained (R 2) of beliefs (see Fig 4B). Same conventions as panel a. c.
Log-likelihood of the state decoder using Value RNN activity to estimate the true state
(see Fig 4C). Dashed line indicates maximum possible log-likelihood (i.e., from Belief
model). Same conventions as panel a.
some threshold distance of its ITI fixed point. We refer to this quantity—which essentially measures ‘X’ in239
the top panel of Fig 5A—as the network’s odor memory. In fact, we found that all Task 1 RNNs had longer240
odor memories (310±150, mean± SE,N = 12) than Task 2 RNNs (49±2, mean± SE,N = 12). Overall,241
these features were consistent with the beliefs in the two tasks following an odor observation: beliefs in Task242
2, but not Task 1, must quickly return to the ITI after the maximum possible reward time. We performed a243
similar analysis for reward observations, in which case the network activity in both tasks was expected to244
return to the ITI fixed point shortly after the minimum ITI duration (at time step 10). Here, we found that245
the activity of both Task 1 and Task 2 RNNs quickly returned to the ITI fixed point after this point (Fig 5D),246
again consistent with the beliefs in these tasks.247
RNNs with larger capacity had more belief-like representations248
Thus far, we have considered the representations of Value RNNs with the same number of hidden units249
(H = 50). To understand whether the network’s size influences the types of representations learned, we250
next trained Value RNNs with a variable number of hidden units. We found that Value RNNs with as few251
as 2 hidden units could learn the value function, as evidenced by their RPEs matching those of the Belief252
model (Fig 6A). In other words, despite there being 25 discrete states in our implementation of this task253
(and, as such, beliefs were 25-dimensional), an RNN with a 2-dimensional representation was sufficient for254
17
.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

performing the task. On the other hand, RNNs with fewer units had representations that were notably less255
belief-like, in terms of how well they linearly encoded beliefs (Fig 6B) and allowed for decoding the true256
state (Fig 6C). Thus, Value RNNs with 2 or more hidden units could all estimate value equally well, but257
only those with a sufficient number of hidden units featured representations that resembled beliefs.258
Generalization to other tasks259
We showed that RNNs could be trained to estimate value in the tasks from Starkweather et al. [7], and that260
the representations of these RNNs became more belief-like as a result of training. We next assessed whether261
the same basic insights generalized to a different task, that of Babayan et al. [10]. In this task, similar to262
Task 1 of Starkweather et al. [7], each trial consisted of an odor followed by a deterministic reward. Unlike263
in the Starkweather task, in this task the reward quantity on each trial was varied in blocks. In Block 1, each264
trial consisted of a small (1µL) reward, while in Block 2 each trial consisted of a large (10µL) reward. As265
a result, we formalize the states in this task using pairs of ITI and ISI states, one for each block (Fig 7A).266
Importantly, the block identity was hidden to the animal, and was resampled uniformly every five trials. This267
meant that animals had to use the reward observations to infer which block they were currently in.268
Previous work showed that the dopamine activity of animals trained on this task depended on the num-269
ber of trials in the current block (Fig 7C), similar to what you would expect if animals used a belief rep-270
resentation (Fig 7D, dashed lines) [10]. To see if the Value RNN could reproduce these results, we trained271
N = 12 Value RNNs on this task. We found that Value RNNs exhibited nearly identical RPEs as the Belief272
model (Fig 7D). This was even true on probe sessions that included blocks with intermediate reward sizes,273
a setting in which both dopamine activity and belief RPEs exhibited a characteristic nonlinear relationship274
with reward size (Fig S2). These results indicate that the Value RNNs found a representation sufficient for275
estimating value despite the hidden states.276
We next asked whether, as in the Starkweather task, the Value RNN’s representations resembled beliefs.277
To do this, we repeated the analyses in Fig 4. We found that the Value RNN’s activity could be linearly278
transformed to match the beliefs (Fig 7D), and that its activity could also be used to decode the hidden279
states in the task (Fig 7E), as compared to RNNs not trained on the task. We next took a dynamical systems280
18
.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