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

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

multi-task settings, or if the RNN framework we present here can also yield representations that effectively389
generalize to new tasks in the same environment.390
Materials and Methods391
Task implementation392
In each experiment, at each time step t, agents received two observations: an odor cue, ct∈{0, 1}; and a393
reward,rt∈{0,r }, where r >0 depended on the task (see below). We will refer to the total observation394
vector asot = [ctrt]. We treated each time step as equal to 200ms.395
Each trial began with an intertrial interval (ITI), tITI ∈ N, during which there were no observations396
(ot = 0 fort < tITI ). The ITI (offset by a minimum delay of 10 time steps) was sampled as tITI − 10∼397
Geom(pITI = 1
8), where Geom(p) is a geometric distribution with parameterp.398
Following the ITI, a single odor cue was presented as ct = 1 fort = tITI . The cue was then followed399
by another interval with no observations, called the interstimulus interval (ISI), tISI ∈ N. A reward was400
then presented as rt > 0 fort = tITI +tISI , after which point the trial terminated. The details of the ISI401
and reward size depended on the specific task, as described below.402
Starkweather tasks403
There were two versions of this task. In both Tasks 1 and 2, every non-zero reward size had rt = 1. In404
Task 2, with probability pomission = 0.1, the reward on a given trial was omitted, such that rt = 0 for the405
duration of the trial. In both tasks, the ISIs on each trial were sampled from a discretized Gaussian with406
meanµ = 10, standard deviationσ = 2.5, and range 6≤tISI≤ 14, as in Starkweather et al. [7].407
Babayan task408
In this task, non-zero reward sizes were determined in blocks of trials. In block 1, the non-zero reward size409
wasrt = 1, while in block 2 the non-zero reward size was rt = 10. Each block consisted of 5 sequential410
trials. Block identities were sampled uniformly with equal probability. On all trials, the ISIs were uniformly411
25
.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

sampled as tISI ∼ Unif({9, 10, 11}). For Fig S2, sessions also included blocks of intermediate rewards:412
rt∈{ 1, 2, 4, 6, 8, 10}, where block identities were sampled in similar proportions to those used in Babayan413
et al. [10] (i.e., blocks withrt = 1 orrt = 10 comprised∼90% of the total trials).414
Recurrent neural network implementation415
We trained recurrent network models, in PyTorch, on multiple tasks to estimate value. Each Value RNN416
consisted of a GRU cell [24] with H∈ N units, followed by a linear readout of value. At each time step t,417
the RNN received observations,ot∈ R2, from a given experiment. The RNN’s representation can be written418
aszt = fϕ(ot,zt−1) given parametersϕ. The RNN’s output was the value estimate bVt = w⊤zt +w0,419
forw∈ RH andw0∈ R. The full parameter vector θ = [ϕw w0] was learned using TD learning. This420
involved backpropagating the gradient of the squared error lossδ2
t = (rt +γ bVt+1− bVt)2 with respect to bVt421
on episodes composed of 20 (Starkweather task) or 50 (Babayan task) concatenated trials. Unless otherwise422
noted we usedγ = 0.93 as in Starkweather et al. [13].423
For each task, and for each H∈{ 2, 5, 10, 20, 50, 100} units, we trained N = 12 networks. Prior to424
training each network, the weights and biases of the GRU (i.e., ϕ) were initialized using PyTorch’s default425
ofU(−a,a) wherea = 1√
H . Training then proceeded for a maximum of 150 epochs on a session of 10,000426
trials, with a batch size of 12 episodes. Training was stopped early if the loss increased for 4 consecutive427
epochs. Gradient updates used Adam with an initial learning rate of 0.003. No hyperparameter search was428
performed to fine-tune these choices. In the text, we refer toValue RNNsas the result of this training process,429
while Untrained RNNs are those that have been similarly initialized but not trained.430
The Value ESN was similar to a Value RNN, except that it was initialized differently, andϕ was frozen431
during training (i.e., the only learned parameters were w and w0). For initialization, we did the follow-432
ing (“Tensorflow-style” initialization). All of the GRU’s biases were initialized to zero. The GRU’s re-433
current weights were sampled as a random orthogonal matrix using torch.nn.init.orthogonal_434
with a given gain [30]. The GRU’s input weights were sampled as U(−a,a) where a =
q
6
2+H , using435
torch.nn.init.xavier_uniform_ [39].436
26
.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

State and belief representations437
Given a task with hidden states s∈{1,...,K }, the belief,bt∈ [0, 1]K, is the posterior probability distri-
bution over each possible state [17]. The tasks we analyze here are technically discrete-time semi-Markov
processes, and so we follow previous work in formulating them equivalently as Markov processes with
micro-states defined over each relevant discrete time step [6, 7]. In this setting, observations occur at the
transition between states. As a result, the belief in state can be written as:
bt(k)∝
KX
k′=1
Oot(k′,k )T (k′,k )bt−1(k′) (12)
whereT∈ [0, 1]K×K is the matrix of transition probabilities, and Oo(k′,k ) is the probability of observing438
o after making a transition fromk tok′.439
Starkweather tasks440
In both Tasks 1 and 2 there are three distinct observations, ot∈{[0 0], [1 0], [0x]}, which we refer to as441
the null, odor, and reward observations, respectively. Let the possible reward times betISIs ={6,..., 14}.442
The maximum reward time is max(tISIs ) = 14 , and so we let states 1− 14 be the ISI microstates. The443
ITI is a Geometric distribution plus a minimum ITI of tITI = 10, and so we let states 15− 25 be the ITI444
microstates. There areK = 25 total states.445
We first define the observation probabilities,Onull,Oodor,Oreward∈{ 0, 1}K×K, whereOo(k′,k ) indi-446
cates the probability of having transitioned from statek to statek′ upon observingo∈{null,odor,reward}.447
EachOo(k′,k ) = 0 except at the following:448
• Onull(t + 1,t) = 1 for allt̸= max(tISIs )449
• Onull(K,K ) = 1450
• Oodor(t,K ) = 1 fort = 1 (Task 1) ort∈{ 1, max(tISIs ) + 1} (Task 2)451
• Oreward(max(tISIs ) + 1,t) = 1 fort∈tISIs452
27
.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

To define the transition probabilities, letpt∈ [0, 1] be the probability of receiving reward at timet∈tISIs ,453
Ft = P
t′≤tpt′ the cumulative probability, andht = pt/(1−Ft) the hazard. Recall that pITI = 1
8. Then454
T (k′,k ) = 0 except at the following:455
• T (t + 1,t) = 1 fort /∈ max(tISIs )456
• T (t + 1,t) = 1−ht andT (max(tISIs ) + 1,t) = ht fort∈tISIs457
• T (K,K ) = 1−pITI458
• T (max(tISIs ) + 1,K) =pITIpomission459
• T (1,K ) =pITI (1−pomission)460
Babayan task461
The states in this task can be thought of as two copies of the beliefs in the Starkweather tasks, with one copy
for each of the two blocks. (Note that tISIs ={9, 10, 11},K = 22, and the hazard probabilities must be
modified from the Starkweather task to account for the different reward timing distribution.) Each copy has
11 ISI microstates (because the maximum reward time is max(tISIs ) = 11) and 11 ITI microstates. Let
b(1)
t ∈ [0, 1]22 andb(2)
t ∈ [0, 1]22 be the beliefs for the substates of block 1 and block 2, respectively. Then
the full beliefbt∈ [0, 1]44 is as follows:
bt = [ptb(1)
t (1−pt)b(2)
t ]
pt =f(rt) ifrt> 0 otherwisept−1
(13)
wherept∈ [0, 1] is the estimated probability of being in block 1, and f is our likelihood function mapping462
nonzero rewards,rt, to the estimated probability of being in block 1 vs. block 2. In other words, we modeled463
the belief in the block identity as being a function only of the most recently observed reward. We definedf464
following the original paper: Let ϕ(rt;µ,σ ) be the pdf of a Normal distribution with mean µ and standard465
deviationσr > 0. Thenf(rt) = ϕ(rt;µ1,σr)
ϕ(rt;µ1,σr)+ϕ(rt;µ2,σr), whereµ1 = 1 andµ2 = 10 are the rewards amounts466
in block 1 and 2, respectively. Here we assumedσr was arbitrarily small, so we usedσr = 0.001.467
28
.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

Model analyses468
We analyzed exemplars from each model class (Beliefs, Value RNNs, Untrained RNNs, Value ESNs) using469
two sessions of 1,000 concatenated trials each, using the same task parameters as those used when training470
the RNNs (see above). The first session was used for fitting any parameters relevant to the analysis (i.e.,471
value weights, regression weights, decoding weights), while the second session was used for evaluation.472
Because the RNN’s responses were deterministic functions of their inputs, prior to analysis we added noise473
to all RNN representations to prevent overfitting during regression and decoding as follows. Let σi > 0474
be the sample standard deviation of the activity of hidden unit i across trials. Then we added zero-mean475
Gaussian noise with a standard deviation of 0.01σi to this unit’s activity, so that each unit had the same476
SNR: 10 log10(σ2
i/(σ2
i 0.012)) = 40 dB.477
Value estimates478
Each model’s value estimate was given by bVt = z⊤
t bw +w0, where bw∈ RD andw0∈ R are the value479
weights, andzt∈ RD is the model’s representation at timet. (For the belief model,zt =bt.)480
We estimated the value weights, bw, using Least-Squares TD (LSTD) [22]: bw = bD−1bd, where bD =481
PT−1
t=1 zt(zt−γzt+1)⊤ and bd = PT−1
t=1 rtzt. To ensure each model’s value function was estimated using482
the same procedure, we used LSTD even for the models including RNNs, even though the Value RNNs and483
Value ESNs learned their own value weights during training.484
Reward prediction errors485
To assess how close each RNN’s RPEs came to the RPEs found using the belief model, we defined an RNN’s486
RPE error using the mean squared error: 1
T
PT
t=1(δt− bδt)2, where bδt is the RNN’s RPE, andδt is the RPE487
from the belief model. Because each trial had at most one reward delivery, for simplicity we considered the488
RPEs only at the time of reward delivery on each trial (i.e., the t in the above equation refers to a trial and489
not a time step); this simplification did not affect our results.490
29
.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

Belief R2491
To assess how much variance of the beliefs could be explained by each model’s learned representation, we
used multivariate linear regression:
cW = (Z⊤Z)−1Z⊤B (14)
whereZ∈ RT×(H+1) is the matrix whose tth row is [zt 1],B∈ RT×K is the matrix whose tth row isbt,492
and cW∈ R(H+1)×K.493
To evaluate model fit, let Var(X) = 1
T
PT
t=1∥xt−x∥2
2, wherext is thetth row ofX, andx is the mean
of the rows ofX. Then we calculated the total variance explained:
R2 = 1− Var(B−Z cW )
Var(B) (15)
State decoding494
We asked where we could find a decoder that could infer the underlying state, st∈{1,...,K }, using an495
affine transformation of the RNN’s representation,zt∈ RH. To find such a decoder, we first standardized496
each model representation (considering each dimension ofzt in isolation) to have zero mean and unit vari-497
ance. We then performed a multinomial logistic regression using scikit-learn’s498
linear model.LogisticRegression with the parameters multi class="multinomial",499
C=1, and max iter=1e4.500
After training, the decoder’s estimated state probabilities overst are:
πt = softmax(˜z⊤
t cW )∈ [0, 1]K (16)
where cW ∈ R(H+1)×K contains the decoder parameters; ˜zt∈ RH+1 is the model representation at time501
t after standardization, plus an extra constant column of 1’s to fit the offset; and the softmax function502
normalizes the vector to be a valid probability over theK values ofst.503
To evaluate the resulting decoder, we calculated the model’s log-likelihood (ℓ) on the evaluation session
30
.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