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

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

Emergence of belief-like representations through
reinforcement learning
Jay A. Hennig1,2,∗, Sandra A. Romero Pinto2,3, Takahiro Yamaguchi3,4,
Scott W. Linderman5,6, Naoshige Uchida2,3, Samuel J. Gershman1,2
1 Department of Psychology, Harvard University, Cambridge, MA, USA
2 Center for Brain Science, Harvard University, Cambridge, MA, USA
3 Department of Molecular and Cellular Biology, Harvard University, Cambridge, MA, USA
4 Future Vehicle Research Department, Toyota Research Institute North America, Toyota Motor North
America Inc., Ann Arbor, MI, USA
5 Wu Tsai Neurosciences Institute, Stanford University, Stanford, CA, USA
6 Department of Statistics, Stanford University, Stanford, CA, USA
* Corresponding author
Email: jay.a.hennig@gmail.com (JAH)
1
.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

Abstract
To behave adaptively, animals must learn to predict future reward, or value. To do this, animals are thought1
to learn reward predictions using reinforcement learning. However, in contrast to classical models, animals2
must learn to estimate value using only incomplete state information. Previous work suggests that animals3
estimate value in partially observable tasks by first forming “beliefs”—optimal Bayesian estimates of the4
hidden states in the task. Although this is one way to solve the problem of partial observability, it is not the5
only way, nor is it the most computationally scalable solution in complex, real-world environments. Here6
we show that a recurrent neural network (RNN) can learn to estimate value directly from observations, gen-7
erating reward prediction errors that resemble those observed experimentally, without any explicit objective8
of estimating beliefs. We integrate statistical, functional, and dynamical systems perspectives on beliefs to9
show that the RNN’s learned representation encodes belief information, but only when the RNN’s capac-10
ity is sufficiently large. These results illustrate how animals can estimate value in tasks without explicitly11
estimating beliefs, yielding a representation useful for systems with limited capacity.12
Author Summary13
Natural environments are full of uncertainty. For example, just because my fridge had food in it yester-14
day does not mean it will have food today. Despite such uncertainty, animals can estimate which states15
and actions are the most valuable. Previous work suggests that animals estimate value using a brain area16
called the basal ganglia, using a process resembling a reinforcement learning algorithm called TD learning.17
However, traditional reinforcement learning algorithms cannot accurately estimate value in environments18
with state uncertainty (e.g., when my fridge’s contents are unknown). One way around this problem is if19
agents form “beliefs,” a probabilistic estimate of how likely each state is, given any observations so far.20
However, estimating beliefs is a demanding process that may not be possible for animals in more complex21
environments. Here we show that an artificial recurrent neural network (RNN) trained with TD learning22
can estimate value from observations, without explicitly estimating beliefs. The trained RNN’s error signals23
resembled the neural activity of dopamine neurons measured during the same task. Importantly, the RNN’s24
2
.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

activity resembled beliefs, but only when the RNN had enough capacity. This work illustrates how animals25
could estimate value in uncertain environments without needing to first form beliefs, which may be useful26
in environments where computing the true beliefs is too costly.27
Introduction28
One pervasive feature of animal behavior is the ability to predict future reward. For example, a dog may29
learn that when her owner picks up the leash, she is likely to be rewarded with a walk in the near future.30
In associative learning settings such as this one, animals learn to associate certain stimuli (e.g., her owner31
grabbing the leash) with future reward (e.g., a walk). The neural basis of associative learning has been32
interpreted through the lens of reinforcement learning (RL). In particular, one successful theoretical model33
posits that associative learning is driven by the activity of dopamine neurons in the midbrain, where spiking34
activity resembles the reward prediction error (RPE) signal in an RL algorithm called temporal difference35
(TD) learning [1, 2, 3, 4]. We will describe this algorithm in more detail below.36
In many real-world scenarios, effectively predicting reward may require a deeper understanding of the37
structure of the world that goes beyond associating observations and reward. To continue the example38
above, suppose the dog’s owner keeps his car keys under the leash. Now if he picks up the leash, this39
does not necessarily mean he is about to take his dog on a walk. In other words, his intention to take40
his dog on a walk is now “partially observable.” Standard RL approaches are insufficient for learning in41
partially observable environments, as these methods assume that all relevant states of the environment are42
fully observable. One way to solve this problem is by using observations to form a Bayesian posterior43
estimate of each hidden state, called a belief state [5]. Future reward can then be estimated by applying44
standard RL methods like TD learning to belief states rather than the raw observations.45
Do animals estimate future reward using belief states? Evidence for this idea is suggestive, although46
indirect. Previous experimental work has shown that the phasic activity of midbrain dopamine neurons47
resembles the RPEs of TD learning in partially observable environments, where TD learning is performed48
on belief states rather than observations [6, 7, 8, 9, 10, 11]. The brain may have dedicated machinery, perhaps49
in prefrontal cortex [12, 13], for computing belief states, which could then be provided to downstream areas,50
3
.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

such as the basal ganglia, to perform standard RL algorithms such as TD learning [14]. This architectural51
division of labor resonates with the broader literature on probabilistic computation in cortex, which has52
identified several different ways in which belief states could be encoded by neural activity [15].53
There are a few difficulties in using a belief state to solve RL tasks. First, the belief state assumes54
knowledge of the environment’s transition and observation dynamics—something that may be challenging55
for animals to acquire via observations alone. Indeed, there are well-documented examples of animals56
failing to learn or use the correct environment model [16]. Second, the belief representation does not scale57
well to more realistic tasks with higher-dimensional state spaces, as beliefs live in a continuous space whose58
dimensionality grows with the number of discrete states in the environment. Finally, the belief state includes59
knowledge about all states in the environment, regardless of whether those states are relevant to the task at60
hand. Luckily, one can often use approximate representations of beliefs to find solutions that work well in61
practice [17]. This suggests that there may be other representations, more compact than the full belief state,62
that are sufficient for the particular task of estimating future reward [18, 19].63
To address these difficulties, here we take inspiration from deep reinforcement learning. In deep RL,64
rather than explicitly learning beliefs, an agent uses nonlinear function approximation to learn a hidden65
representation that is sufficient for performing the task [20]. Compared to the belief representation, this66
approach does not require explicit knowledge about the structure of the environment. It may also scale67
better to more complex tasks, by virtue of the agent not needing to represent any features of the environment68
that do not directly pertain to the task at hand. Because beliefs are a non-linear dynamical system, here69
we use recurrent neural networks (RNNs) as our nonlinear function approximator. This choice was also70
motivated by the observation from the machine learning literature that RNNs can perform well on complex71
partially observable tasks [21]. Previous work in computational neuroscience has explored whether RNNs72
can be used to directly compute beliefs [17]. Here, by contrast, we explore whether training RNNs in73
partially observable environments leads to their representations implicitly becoming more like beliefs.74
We show that RNNs can be trained to perform two previously studied associative learning tasks [7, 10],75
reproducing experimentally observed dopamine neuron activity. We probe the representations learned by the76
RNNs, and show that the representations resemble beliefs from statistical, functional, and dynamical systems77
perspectives. Finally, we show how an RNN’s capacity determines the degree to which its representation78
4
.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

resembles beliefs, without a concurrent impact on its ability to perform the task. These results illustrate79
how animals might estimate reward in partially observable environments without requiring any explicit80
representation of beliefs.81
Results82
Reinforcement learning in partially observable environments using belief states83
The standard RL objective is to learn the expected discounted future return, or value, of each state:
V (st) = E
"∞X
k=0
γkrt+k
#
(1)
wherest∈S is the state of the environment at time t, 0≤γ <1 is a discount factor, andrt is the reward.84
Rewards are random variables that depend on the environment state, and E denotes an expectation over the85
potentially stochastic sequences of states and rewards. Because we are interested in modeling Pavlovian86
associative learning tasks, here we will assume there are no actions available to the agent. For notational87
simplicity, we will use the shorthandVt =V (st).88
TD learning estimates the value function by exploiting a set of Markovian assumptions about the en-
vironment: the probability of state st depends only on the last state st−1, and the probability of reward rt
depends only on the current statest. Under these assumptions, the value function can be decomposed into a
recursive form known as the Bellman equation [22]:
Vt = E[rt +γVt+1]. (2)
If an agent has access to an approximation of the value function, bVt, along with sample paths of states and
rewards, then it can compute an estimate of the discrepancy between the approximate and true value function
by taking the difference between the two sides of the Bellman equation:
δt =rt +γ bVt+1− bVt. (3)
5
.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

This is the temporal difference error, the precise definition of the RPE used in TD learning. Under the
Bellman equation, E[δt] = 0 when bVt = Vt. If bVt is determined by a set of adaptable parameters θ, the
approximation can be improved by following the stochastic gradient of the squared TD error:
∆θ =ηδt∇θ bVt, (4)
where 0<η < 1 is a learning rate, and∇θ bVt is the gradient of bVt with respect toθ. One common example
is a linear function approximator:
bVt =
DX
d=1
w(d)zt(d), (5)
wherezt(d)∈ R is some feature (indexed by d) of the state, and θ =w∈ RD is a learned set of weights89
on all the features. Under this approximation,∇θ bVt =zt.90
In a partially observable Markov process, agents do not observe the statest directly, but instead observe91
only observations ot ∈O . Critically, observations are not in general Markovian, which means that TD92
learning methods cannot be naively applied to value function approximators defined over observations. One93
way of understanding this is to note that the value of an observation may depend on the long-term past [23].94
In the dog leash example from the Introduction, the value (to the dog) of her owner picking up the leash95
depends on the history of events leading up to that moment—for example, if her owner recently announced96
that his car keys were missing. Reward prediction requires a compression of this history into a “sufficient97
statistic”—in effect creating a transformed state space over which the Markov property holds. TD learning98
can then be applied to this transformed state space.99
One such transformed state is the posterior probability distribution over hidden states given the history
of observations, also known as the belief state [5]:
bt(st) =P (st|o1,...,o t)
∝P (ot|st)
Z
P (st|st−1)bt−1(st−1)∂st−1
(6)
where the recursion follows from the Chapman-Kolmogorov equation, which stipulates how to update the100
6
.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