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

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

belief state fromt− 1 tot after observing observationot.101
For a partially observable Markov process with a finite state space,S ={1,...,K }, the value function
can be written as a weighted combination of beliefs:
Vt =
KX
k=1
V (k)bt(k) (7)
whereV (k) andbt(k) are the value and belief of state st = k, respectively. This means that linear value102
function approximation in Eqn. (5) is sufficiently expressive for the partially observable problems we con-103
sider in this paper: with enough training data, the value function approximator will perfectly estimate the104
true value function (i.e., given learned weightsw(k) =V (k) and featureszt(k) =bt(k) for eachk∈K).105
This motivates a straightforward model of how animals estimate value in partially observable environ-106
ments [6, 7], which we will refer to as the “Belief model”: First compute beliefs (Eqn. (6)), and then107
compute a linear transformation of those beliefs to a value estimate (Eqn. (7)), with weights obtained from108
TD learning.109
Learning state representations using recurrent neural networks110
The Belief model assumes animals use a particular feature representation (i.e., beliefs) for estimating value.
Here we ask what representations might emerge from solving the task of estimating value alone. It is useful
to note that the Belief model can be written as follows:
bVt =w⊤bt (8)
bt =P (st|o1,...,o t) (9)
=P (st|bt−1,ot)
=fϕ(bt−1,ot)
wheref is a function parameterized by a specific choice of (fixed) parameters ϕ, and only w is learned.
This latter equation has the same form as a generic recurrent neural network (RNN). This suggests a model
7
.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

could learn its own representation by treatingϕ as a learnable parameter. Taken together, this results in the
following model for estimating value:
bVt =w⊤zt (10)
zt =fϕ(zt−1,ot) (11)
wherezt∈ RH is the activity of an RNN withH hidden units and parametersϕ. Importantly, bothϕ andw111
can be learned simultaneously, using backpropagation through time to calculate Eqn. (4) with θ = [ϕw ].112
This allows the network to learn the representations, zt, that are sufficient for estimating value. We will113
refer to this model as a “Value RNN” (see Materials and Methods).114
Importantly, this RNN-based approach resolves all three challenges for learning a belief representation115
that we raised in the Introduction: 1) The model can learn from observations alone, as no information is116
provided about the statistics of the underlying environment; 2) the model’s size (parameterized by H, the117
number of hidden units) can be controlled separately from the number of states in the environment; and 3)118
the model’s only objective is to estimate value.119
Though such a model has no explicit objective of learning beliefs, the network may discover a belief120
representation implicitly. We next asked what signatures, if any, would indicate the existence of a belief121
representation. In the sections that follow we develop an analytical approach for determining whether the122
Value RNN’s learned representations resemble beliefs.123
RNNs learn belief-like representations124
As a working example, we will consider the probabilistic associative learning paradigm where dopamine125
RPEs were shown to be consistent with a belief representation [7, 13]. This has the added benefit of ensuring126
that the RNN-based approach described above can recapitulate these previous results.127
This paradigm consisted of two tasks, which we will refer to as Task 1 and Task 2 (Fig 1). In both tasks,128
mice were trained to associate an odor cue with probabilistic delivery of a liquid reward 1.2-2.8s later. The129
tasks were each composed of two states: an intertrial interval (ITI), during which animals waited for an130
8
.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 1. Associative learning tasks with probabilistic rewards and hidden states. a.
Trial structure in Starkweather et al. [7, 13]. Each trial consisted of a variable delay (the
intertrial interval, or ITI), followed by an odor, a second delay (the interstimulus interval,
or ISI), and a potential subsequent reward. Reward times were sampled from a discretized
Gaussian ranging from 1.2-2.8s (see Materials and Methods). b-c.. In both versions of
the task, there were two underlying states: the ITI and the ISI. In Task 1, every trial
was rewarded. As a result, an odor always indicated a transition from the ITI to the ISI,
while a reward always indicated a transition from the ISI to the ITI. In Task 2, rewards
were omitted on 10% of trials; as a result, an odor did not reveal whether or not the state
transitioned to the ISI.
odor; and an interstimulus interval (ISI), during which animals waited for a reward. In Task 1, every trial131
contained both an odor and a reward. As a result, the animal’s observations could fully disambiguate the132
underlying state: An odor signaled a transition to the ISI state, while a reward signaled a transition to the133
ITI state. In Task 2, by contrast, reward on a given trial was omitted with 10% probability. This meant the134
underlying states were now only partially observable; for example, in Task 2 an odor signaled a transition to135
the ISI state with 90% probability.136
To formalize these tasks, we largely followed previous work [7, 13]. Each task was modeled as a137
discrete-time Markov process with statesst∈{ 1,...,K }, where eacht denotes a 200ms time bin (Fig 2A).138
TheseK “micro” states can be partitioned into those belonging to one of two “macro” states (corresponding139
to the ITI and the ISI; see Materials and Methods). At each point in time, the agent’s observation is one of140
ot∈{odor,reward,null } (Fig 2B). For each task, we trained the Belief model, and a Value RNN (using141
a gated-recurrent unit cell [24], or GRU, with H = 50 hidden units), on a series of observations from that142
task to estimate value using TD learning (see Materials and Methods). Before training, the Value RNN’s143
representation consisted of transient responses to each observation (Fig S1). After training, we evaluated144
each model on a sequence of new trials from the same task (Fig 2C).145
To confirm that this approach could recapitulate previous results, we measured the RPEs generated by146
9
.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 2. Observations, model representations, value estimates, and reward prediction
errors (RPEs) during Task 2. a. State transitions and observation probabilities in Task
2. Each macro-state (ISI or ITI) is composed of micro-states denoting elapsed time; this
allows for probabilistic reward times and minimum dwell times in the ISI and ITI, respec-
tively. b. Observations emitted by Task 2 during two example trials. Note that omission
trials are indicated only implicitly as the absence of a reward observation. c. Example
representations (bt,zt) and value estimates ( bVt) of two models (Belief model, left; Value
RNN, right) for estimating value in partially observable environments, after training. d.
After training, both models exhibit similar RPEs.
10
.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

each trained model (Fig 2D). Previous work showed that dopamine RPEs depended on the reward time147
differently in the two tasks, with RPEs increasing as a function of reward time in Task 1, but decreasing148
as a function of reward time in Task 2 [7] (Fig 3A). As in previous work, we found that this pattern was149
also exhibited by the Belief model (Fig 3B). We found that the RPEs of the Value RNN exhibited the same150
pattern (Fig 3C). In particular, the Value RNN’s RPEs became nearly identical to those of the Belief model151
after training (Fig 3D). We emphasize that the Value RNN was not trained to match the value estimate from152
the Belief model; rather, it was trained via TD learning using only observations. This result shows that,153
through training on observations alone, Value RNNs discovered a representation that was sufficient for both154
learning the value function as well as explaining empirical dopamine activity.155
Fig 3. RPEs of the Value RNN resemble both mouse dopamine activity and the
Belief model. a. Average phasic dopamine activity in the ventral tegmental area (VTA)
recorded from mice trained in each task separately. Black traces indicate trial-averaged
RPEs relative to an odor observated at time 0, prior to reward; colored traces indicate the
RPEs following each of nine possible reward times. RPEs exhibit opposite dependence
on reward time across tasks. Reproduced from Starkweather et al. [7].b-c. Average RPEs
of the Belief model and an example Value RNN, respectively. Same conventions as panel
a. d. Mean squared error (MSE) between the RPEs of the Value RNN and Belief model,
before and after training each Value RNN. Small dots depict the MSE of each ofN = 12
Task 1 RNNs andN = 12 Task 2 RNNs, and circles depict the median across RNNs.
We next asked whether the Value RNN learned these tasks by forming representations that resembled156
beliefs. We considered three approaches to answering this question. First, we asked whether beliefs could157
11
.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

be linearly decoded from the RNN’s activity. Next, because beliefs are the optimal estimate of the true state158
in the task, we asked whether RNN activity could similarly be used to decode the true state. Finally, we took159
a dynamical systems perspective, and asked whether the RNN and beliefs had similar dynamical structure.160
RNN activity readout was correlated with beliefs161
We first asked whether the Value RNN’s activity was correlated with beliefs. Because the belief and RNN
representations did not necessarily have the same dimensionality, we performed a multivariate linear regres-
sion to find the linear transformation of the RNN’s activity that came closest to matching the beliefs (see
Materials and Methods). In other words, we found the linear transformation, cW∈ RK×H, that could map
the RNN’s activity,zt∈ RH, to best match the belief vector,bt∈ RK:
cW = argminW
TX
t=1
∥Wzt−bt∥2
2
To quantify performance, we used held-out sessions to measure the total variance of the beliefs that were162
explained by the linear readout of RNN activity (R 2; see Materials and Methods). We found that the Value163
RNN’s activity explained most of the variance of beliefs (Fig 4B; Task 1 R2: 0.69 ± 0.01, mean± SE,164
N = 12; Task 2 R2: 0.76± 0.02, mean± SE,N = 12), substantially above the performance when using165
the RNN activity before training (Task 1 R2: 0.47 ± 0.00, mean ± SE, N = 12; Task 2 R2: 0.50 ±166
0.00, mean± SE, N = 12 ). This is not a trivial result of the network’s training objective, as the Value167
RNN’s target (i.e., value) is only a 1-dimensional signal, whereas beliefs are aK-dimensional signal (here,168
K = 25). Nevertheless, we found that training the Value RNN to estimate value resulted in its representation169
becoming more belief-like, in the sense of encoding more information about beliefs.170
RNN activity could be used to decode hidden states171
The previous analysis assessed how much information about beliefs was encoded by RNN activity. Given172
that beliefs are distributions over hidden states, we asked whether the ground truth state could be decoded173
from the RNN’s activity. To do this, we performed a multinomial logistic regression to find a linear trans-174
formation of the RNN’s activity that maximized the log-likelihood of the true states (see Materials and175
12
.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 4. Value RNN activity readout was correlated with beliefs and could be used to
decode hidden states. a. Example states, beliefs, and Value RNN activity from the same
Task 2 trials shown in Fig 2. Note that the states following the second odor observation
remain in the ITI (black) because the second trial is an omission trial. Bottom traces depict
the linear transformation of the RNN activity that comes closest to matching the beliefs.
Total variance explained (R2) is calculated on held-out trials. b. Total variance of beliefs
explained (R2), on held-out trials, using different trained and untrained Value RNNs, in
both tasks. Same conventions as Fig 3D). c. In purple, the cross-validated log-likelihood
of linear decoders trained to estimate true states using RNN activity. Same conventions as
Fig 3D). Black circle indicates the log-likelihood when using the beliefs as the decoded
state estimate (i.e., no decoder is “trained”).
13
.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