# Dynamic allocation of limited memory resources
in reinforcement learning

Nisheet Patel
Department of Basic Neurosciences
University of Geneva
<EMAIL_ADDRESS>
Luigi Acerbi
Department of Computer Science
University of Helsinki
<EMAIL_ADDRESS>
Alexandre Pouget
Department of Basic Neurosciences
University of Geneva
<EMAIL_ADDRESS>
Current address: Département des neurosciences fondamentales, Université de
Genève, CMU, 1 rue Michel-Servet, 1206 Genève, Switzerland. Alternative
e-mail<EMAIL_ADDRESS>
###### Abstract

Biological brains are inherently limited in their capacity to process and
store information, but are nevertheless capable of solving complex tasks with
apparent ease. Intelligent behavior is related to these limitations, since
resource constraints drive the need to generalize and assign importance
differentially to features in the environment or memories of past experiences.
Recently, there have been parallel efforts in reinforcement learning and
neuroscience to understand strategies adopted by artificial and biological
agents to circumvent limitations in information storage. However, the two
threads have been largely separate. In this article, we propose a dynamical
framework to maximize expected reward under constraints of limited resources,
which we implement with a cost function that penalizes precise representations
of action-values in memory, each of which may vary in its precision. We derive
from first principles an algorithm, Dynamic Resource Allocator (dra), which we
apply to two standard tasks in reinforcement learning and a model-based
planning task, and find that it allocates more resources to items in memory
that have a higher impact on cumulative rewards. Moreover, dra learns faster
when starting with a higher resource budget than what it eventually allocates
for performing well on tasks, which may explain why frontal cortical areas in
biological brains appear more engaged in early stages of learning before
settling to lower asymptotic levels of activity. Our work provides a normative
solution to the problem of learning how to allocate costly resources to a
collection of uncertain memories in a manner that is capable of adapting to
changes in the environment.

##  1 Introduction

Reinforcement learning (RL) is a powerful form of learning wherein agents
interact with their environment by taking actions available to them and
observing the outcome of their choices in order to maximize a scalar reward
signal. Most RL algorithms use a value function in order to find good policies
[1, 2] because knowing the optimal value function is sufficient to find an
optimal policy [3, 4]. However, RL models typically assume that agents can
access and update the value function for each state or state-action pair with
nearly infinite precision. In neural circuits, however, this assumption is
necessarily violated. The values must be stored in memories which are limited
in their precision, especially in smaller brains in which the number of
neurons can be severely limited [5].

In this work, we make such constraints explicit and consider the question of
what rationality is when computational resources are limited [6, 7].
Specifically, we examine how agents might represent values with limited
memory, how they may utilize imprecise memories in order to compute good
policies, and whether and how they should prioritize some memories over others
by devoting more resources to encode them with higher precision. We are
interested in drawing useful abstractions from small-scale models that can be
applied generally and in investigating whether brains employ similar
mechanisms.

We pursue this idea by considering memories that are imprecise in their
representation of values, which are stored as distributions that the agent may
only sample from. We construct a novel family of cost functions that can
adjudicate between maximizing reward and limited coding resources (Section 2).
Crucially, for this general class of resource-limited objective functions, we
derive how an RL agent can solve the resource allocation problem by computing
stochastic gradients with respect to the coding resources. Further, we combine
resource allocation with learning, enabling agents to assign importance to
memories dynamically with changes in the environment. We call our proposed
algorithm the Dynamic Resource Allocator (dra), which we apply to standard
tasks in RL in Section 3 and a model-based planning task in Section 4.111Code
to run dra and reproduce our results is available at
https://github.com/nisheetpatel/DynamicResourceAllocator.

###  1.1 Related work

Our work is related to previous research within the paradigm of bounded
rationality [6, 7] on two accounts. First, previous studies have considered
capacity-limited agents that trade reward to gain information [8, 9, 10, 11],
but did so in a way that abstracts away the underlying costs, and therefore
cannot disambiguate the effects of limited storage capacity from other
energetic or computational limitations. Second, a separate line of work makes
the limitations in memory explicit [12, 13, 14] like we do, but such studies
restrict their analyses to simple working memory tasks, such as reproducing
observed stimuli or delayed recall, and lack a general-purpose, dynamical
framework for decision-making. Our work is also ideologically related to that
of others looking into prioritized replay of memories, both in RL [15] and
neuroscience [16], with the important difference that these studies do not
make the uncertainty in agents’ memories explicit as we do. Moreover, agents
in Mattar and Daw [16] replay memories to update their q-value estimates and
hence their policy. Thus, prioritization of memories in [16] stems only due to
incomplete learning and their results would not hold for well-trained agents,
in contrast to our work where these effects are driven by a limited capacity
constraint.

Other groups have proposed alternative approaches to deal with memory
limitations in RL, such as using regularization (sac [17]), or using neural
networks for representing policy and value functions, and even compressing
state representations with graph-Laplacian [18]. Our work is meant to
complement these previous studies. sac, for instance, directly penalizes the
policy entropy while maximizing reward to encourage exploration. In dra, we
penalize precise representations of q-values instead. The use of a compressed
graph-Laplacian [18], on the other hand, hints at yet another problem
involving efficient use of memory – compact representation of states (e.g.,
chunking) – which we plan to combine with our approach in future work. On the
technical side, dra is related to O’Donoghue et al.’s uncertainty-based
exploration in the manner of decoupling updates to different moments of the
value distribution [19]. Finally, it is worth noting that our work
fundamentally differs from previous work in RL applied to _external_ resource
allocation [20, 21] in that we are using RL to optimize the computational
resources _of the agent itself_.

##  2 Background and details

###  2.1 Environment and agent’s memories

We consider problems that can be described by a Markov Decision Process (MDP)
characterized by a quadruple
$\langle\mathcal{S,A},\mathcal{R}_{a},\mathcal{P}_{a}\rangle$, where
$\mathcal{S}$ is a finite set of states that describe the environment,
$\mathcal{A}$ is a finite set of actions available to the agent,
$\mathcal{R}_{a}(s,s^{\prime})$ is the immediate reward received after taking
action $a$ in state $s$ and transitioning to state $s^{\prime}$, and
$\mathcal{P}_{a}=p(s^{\prime}|s,a)$ denotes the dynamics of the environment,
i.e. probability of transitioning from state $s$ to $s^{\prime}$ after taking
action $a$ [22].

For simplicity, we assume that agents have a model of the world, though this
assumption is not necessary except in model-based planning tasks such as the
one in Section 4. We also assume that the agents perfectly observe the states.
However, agents may only represent the value of each state-action pair, the
_q-value_ , with finite precision. This means that the q-values are
represented as distributions rather than point estimates, and agents can only
access samples from the distribution stored in memory but cannot access its
parameters such as the mean and precision directly. We would like to emphasize
that, for such agents, storing items in memory more precisely requires more
resources.

Concretely, we define the agent’s memory in a tabular form comprising the
states, actions, immediate (mean) rewards, next states, and the corresponding
imprecise q-value, represented here by a normal distribution with mean
$\bar{q}_{sa}$ and variance $\sigma^{2}_{sa}$. Thus, each memory can be
written as a quintuple $\langle
s,a,r,s^{\prime},\mathcal{N}(\bar{q}_{sa},\sigma^{2}_{sa})\rangle$, and the
total number of memories equals $|\mathcal{S\times A}|$. In this work, we
assume the overall q-value distribution to be a multivariate normal with
diagonal covariance matrix (that is, independent memories). However, our
approach for allocating resources across a collection of uncertain memories
can be generalized to arbitrary distributions.

###  2.2 Objective and policy

The goal of the agent is to maximize their expected sum of future rewards in a
given task episode, subject to a cost of representing the q-values precisely.
Biological agents pay such costs for recruiting more neurons for
representation and computation of task-relevant statistics, or for creating
the relevant synaptic connections between existing neurons for learning [23].
In general, our method can be applied to arbitrarily defined cost functions,
but here we focus on a cost derived from neural and information-theoretical
principles.

Biological agents must use a pre-existing population of neurons from a region
of their brain where q-values are represented [24, 25, 26] to store their base
q-value distribution in its connections and activity. We will refer to this
base distribution as
$P\vcentcolon=\mathcal{N}(\bm{\bar{q}},{\sigma}_{\text{base}}^{2}I)$. As
brains learn to perform well on the task, they update the synaptic connections
and may recruit more neurons if necessary to represent the q-value
distribution in memory as
$Q\vcentcolon=\mathcal{N}(\bm{\bar{q}},\bm{\sigma}^{2}I)$, with higher
precision. In this work, we reason that the agent (and the brain) pays a cost
proportional to the KL-divergence $D_{\text{KL}}(Q\;\|\;P)$, representing the
information-theoretical cost for encoding deviations from the base
distribution (_e.g._ , due to modified connectivity and neuronal activity).

Thus, the full objective that the agent is trying to maximize in each episode
is:

$\mathcal{F}\vcentcolon=\
\operatorname{\mathbb{E}}_{\pi}\Bigg{[}\sum_{t=0}^{\infty}\gamma^{t}r_{t+1}\bigg{|}Q=\mathcal{N}(\bm{\bar{q}},\bm{\sigma}^{2}I)\Bigg{]}-\lambda
D_{\text{KL}}\Big{(}Q=\mathcal{N}(\bm{\bar{q}},\bm{\sigma}^{2}I)\;\|\;P=\mathcal{N}(\bm{\bar{q}},{\sigma}_{\text{base}}^{2}I)\Big{)}$
(1)

where $\gamma\in[0,1]$ is the discount factor which we set to 1 in this
article, the first term represents the expected reward for the MDP given the
agent’s memory distribution $Q$ and policy $\pi$, and the second term
represents the cost function described above with a cost per _nat_ equal to
$\lambda\geq 0$.

Given that their memories are noisy, agents can only draw samples from their
memory distribution of q-values. If the agent is greedy, they will then choose
the action corresponding to the largest sampled q-value, effectively yielding
the policy $\pi$:

$\pi(a|s)=\Pr(a=\operatorname*{arg\,max}_{a^{\prime}}\tilde{q}(s,a^{\prime}))\qquad\text{
with
}\quad\tilde{q}(s,a^{\prime})\sim\mathcal{N}\big{(}\bar{q}(s,a^{\prime}),\sigma^{2}(s,a^{\prime})\big{)}.$
(2)

This behavioral policy is also known as Thompson sampling [27, 28], which is
often chosen deliberately for efficient exploration, but here, it is a
consequence of having imprecise memories. In principle, agents could draw more
than one sample from memory to increase the precision of their estimates of
q-values, but we assume that drawing multiple _independent_ samples would
require additional time [29]. In this paper, we restrict our analyses to
single Thompson samples at decision time, leaving a detailed analysis of the
speed-accuracy trade-off [30] for future work.

Our key contribution in this work is providing the agent with control over the
precision of each of its memories, the resource allocation vector
$\bm{\sigma}$. Thus, the agent may allocate more resources to some memories
than others so as to maximize the objective $\mathcal{F}$ defined in Eq. 1.

###  2.3 Maximizing the objective

First, we consider the scenario where the agent wishes to allocate resources
optimally – by maximizing the objective $\mathcal{F}$ in Eq. 1 – when the mean
q-values, $\bar{\bm{q}}$, are fixed and known. To do so, we analytically
derive a stochastic gradient for $\mathcal{F}$.

To compute the gradient of the first term of $\mathcal{F}$, _i.e._ the
expected reward, we follow the approach as in the policy gradient theorem
[31][22, section 13.2], with the difference being that our gradient is with
respect to the resource allocation vector, $\bm{\sigma}$, instead of, for
instance, the parameters of a policy network. In general, this gradient may be
written as:

$\nabla_{\sigma}\operatorname{\mathbb{E}}_{\pi}\Bigg{[}\sum_{t=0}^{\infty}\gamma^{t}r_{t+1}\Big{|}Q=\mathcal{N}(\bm{\bar{q}},\bm{\sigma}^{2}I)\Bigg{]}=\operatorname{\mathbb{E}}_{\pi}\Bigg{[}\sum_{t=0}^{\infty}\Psi_{t}\nabla_{\sigma}\log\pi(a_{t}|s_{t})\Big{|}Q=\mathcal{N}(\bm{\bar{q}},\bm{\sigma}^{2}I)\Bigg{]}$
(3)

where $\Psi_{t}$ can take many forms with different computational properties
[32], including but not limited to:

$\displaystyle\Psi_{t}=\begin{cases}\sum_{t^{\prime}=t}^{\infty}\gamma^{t-t^{\prime}}r_{t^{\prime}+1}&\small{\text{R-gradient
({reinforce)}}}\\\\[1.99997pt] \bar{q}(s_{t},a_{t})&\small{\text{Q-gradient
(mean q-value)}}\\\\[1.99997pt] A(s_{t},a_{t})&\text{\small A-gradient
(\emph{advantage function}).}\end{cases}$ (4)

In our case, we define the _advantage function_ as
$A(s_{t},a_{t})=\bar{q}(s_{t},a_{t})-|\mathcal{A}|^{-1}\sum_{a}\bar{q}(s_{t},a)$
(see justification for using the mean $\bar{q}$ in Appendix A.3). In more
complicated tasks that involve planning (_e.g._ , Section 4), agents may
sample _multiple_ future trajectories using the policy $\pi$, followed by
performing a non-linear operation such as picking the maximal reward of all
sampled trajectories. Thus, it would be inappropriate to use the advantage
function, which only characterizes the expected reward until termination for a
_single_ trajectory (minus a baseline). In such scenarios, we replace
$\Psi_{t}$ with the reward that results from planning,
$\max_{i}(\sum_{t=0}^{T}r_{{t+1},i})$, where $i$ indexes planned trajectories
and $T$ is the planning horizon, following the reasoning as in the original
reinforce algorithm [31].

Eq. 3 allows us to compute an unbiased stochastic estimate of the gradient –
the expectation on the right-hand side – via Monte Carlo, _i.e._ by sampling
one trajectory or averaging over multiple sampled trajectories. Note that if
the agents do not have a model of the world, they may simply store previously
experienced sequences of states in a buffer (episodic memory) and use these
stored trajectories instead of generating new ones with their model. Unless
otherwise mentioned, we use $N_{\text{traj}}=10$ trajectories to compute the
stochastic approximation for the expectation.

We obtain an analytical approximation for $\nabla_{\sigma}\log\pi(a|s)$ (Eq.
3) for our policy by first reparametrizing our q-value distribution as Kingma
and Welling prescribe for normal distributions [33], and then using soft
Thompson sampling [27, 28], _i.e._ using _softmax_ or _soft_
-$\operatorname*{arg\,max}$ instead of $\operatorname*{arg\,max}$ in Eq. 2
(see Appendix A.1 for full derivation). This modification yields:

$\displaystyle\frac{\partial}{\partial\sigma(s^{\prime},a^{\prime})}\log\pi(a|s)$
$\displaystyle=\begin{cases}\beta\zeta_{sa}(1-\pi(a|s))&\small{\text{for
}s^{\prime}=s,a^{\prime}=a}\\\
-\beta\zeta_{sa^{\prime}}\pi(a^{\prime}|s)&\small{\text{for
}s^{\prime}=s,a^{\prime}\neq a}\\\ 0&\small{\text{for }s^{\prime}\neq
s}\end{cases}$ (5)

where
$\zeta_{sa}\stackrel{{\scriptstyle\text{i.i.d.}}}{{\sim}}\mathcal{N}(0,1)\
\forall s\in\mathcal{S},a\in\mathcal{A}$, and $\beta$ is the inverse-
temperature parameter for _softmax_. Plugging this into Eq. 3 for each step in
the sampled trajectories yields the stochastic gradient of the first term of
$\mathcal{F}$, the expected reward, with respect to $\bm{\sigma}$.

Computing the gradient of the second term of $\mathcal{F}$, the cost, is
straightforward (see Appendix A.2):

$\frac{\partial}{\partial\sigma_{sa}}D_{\text{KL}}(Q=\mathcal{N}(\bm{\bar{q}},\bm{\sigma}^{2}I)\;\|\;P=\mathcal{N}(\bm{\bar{q}},{\sigma}_{\text{base}}^{2}I))=\frac{\sigma_{sa}}{\sigma_{\text{base}}^{2}}-\frac{1}{\sigma_{sa}}.$
(6)

Now, the agent may iteratively update its resource allocation vector,
$\bm{\sigma}$, as:

$\bm{\sigma}\leftarrow\bm{\sigma}+\alpha\nabla_{\sigma}\mathcal{F}$ (7)

where $\alpha>0$ is a learning rate and $\mathcal{F}$ is the objective in Eq.
1 that depends on the resource allocation vector $\bm{\sigma}$, the stochastic
gradient for which is given by Eqs. 3, 5, and 6.

Alternatively, we can maximize $\mathcal{F}$ via a black-box optimizer such as
Covariance Matrix Adaptation Evolution Strategy (cma-es) [34, 35], which works
well for up to hundreds of dimensions (memories, in our case). Results from
cma-es provide us with a baseline for our gradient-based optimization in
small-scale environments (Section 4), since both methods are able to find the
optimal resource allocation at stationary state, i.e. when the mean q-values,
$\bar{\bm{q}}$, are fixed and known.

###  2.4 Dynamic allocation of limited memory resources

Real-world environments are rarely stationary. Thus, any agent with limited
memory resources must assign importance dynamically to items in memory in a
time-efficient manner. Our framework makes it is possible to do so by
decoupling the updates for the mean and the variance of the memory
distribution $Q$ (as, for instance, done by O’Donoghue et al. [19]). Agents
can update the mean of $Q$ on the fly with any on-policy learning algorithm
(_e.g._ sarsa [4] or expected sarsa [36]) and simultaneously update the
variance as in Eq. 7. Combining these elements, we propose Algorithm 1,
Dynamic Resource Allocator (dra), to enable memory-limited agents to find good
policies.

Set hyper-parameters
$\bm{\theta}=(\alpha_{1},\alpha_{2},\beta,\gamma,\lambda)$

Initialize $\bm{\bar{q}}$, $\bm{\sigma}$, table of memories = $\langle
s,a,r,s^{\prime},\mathcal{N}(\bar{q}_{sa},\sigma_{sa}^{2}I)\rangle^{|\mathcal{S}\times\mathcal{A}|}$

for  _episode $k=1,...,K$_ do

$s\leftarrow s_{0}$

while  _s is not Terminal_ do

$a\leftarrow\pi(s,Q,\beta)$

$s^{\prime},r\leftarrow$ Environment($s,a$)

$\delta\leftarrow
r+\gamma\sum_{a^{\prime}}\pi(a^{\prime}|s^{\prime})q(s^{\prime},a^{\prime})-q(s,a)$

$q(s,a)\leftarrow q(s,a)+\alpha_{1}\delta$

$s\leftarrow s^{\prime}$

end while

Sample $N$ trajectories to compute:

$\nabla_{\sigma}\operatorname{\mathbb{E}}_{\pi}[\sum_{t=0}^{\infty}\gamma^{t}r_{t+1}|Q]=\frac{1}{N}\sum_{n=1}^{N}[\sum_{t=0}^{\infty}\Psi_{t}\nabla_{\sigma}\log\pi(a_{t}|s_{t})]$

where $\qquad\qquad\qquad\Psi_{t}\leftarrow$ Eq. 4

$\nabla_{\sigma}\log\pi(a_{t}|s_{t})\leftarrow$ Eq. 5

$\bm{\sigma}\leftarrow\bm{\sigma}+\alpha_{2}\big{(}\nabla_{\sigma}\operatorname{\mathbb{E}}_{\pi}\big{[}\sum_{t=0}^{\infty}\gamma^{t}r_{t+1}|Q\big{]}+\lambda\nabla_{\sigma}D_{\text{KL}}(Q\;\|\;P)$)

where $\nabla_{\sigma}D_{\text{KL}}(Q\;\|\;P)\leftarrow$ Eq. 6

end for

Algorithm 1 Dynamic Resource Allocator (dra)

##  3 Results on standard RL environments

###  3.1 2D Grid-world

First, we consider the grid-world adapted from Mattar and Daw [16] and
depicted in Fig. 1a. The goal of the agent is to find the shortest route from
the starting location indicated by the position of the rat to the cheese,
since all transitions yield a reward -1 except reaching the cheese, which
rewards 10 points. In each state, the agent can only choose to go _up_ ,
_down_ , _left_ , or _right_ , and if their intended action is blocked by an
obstacle, that action leaves their position unchanged.

Figure 1: 2D grid-world. (a) Task (adapted from [16]). (b) Rate of convergence
of the objective $\mathcal{F}$ (solid lines, Eq. 1) that the agent aims to
maximize and the cost (dashed lines) it pays to encode memories with higher
precision. Error bars represent SD across mean of 5 optimization runs. (c)
Normalized entropy of the agent’s policy through training and re-training
after change in the environment. Darker colors indicate lower entropy and thus
more precise memories.

Our results show that the initial amount of resources agents can afford at the
beginning of training has a critical influence on learning speed. We choose
four different values of $\sigma_{0}$ (see Fig. 1b legend) to initialize
$\sigma(s,a)=\sigma_{0}\ \forall(s,a)$ for the memory distribution, and for
each of them, we perform five optimization runs for 5000 episodes each. As
shown in Fig. 1b, all $\sigma_{0}$ conditions eventually converge to the same
solution (asymptotic $\sigma_{*}$ is largely independent of the initial
$\sigma_{0}$), but starting with more resources (low $\sigma_{0}$) leads to
faster learning, at a higher initial cost. In real-world animal experiments,
this is a very common observation: more neurons are responsive to task-
relevant variables in the early stages of training than when the animal has
been well-trained [37, 38], suggesting that biological brains may indeed
deliberately assign more resources early to enable quicker learning.

Further, we test the ability of dra to dynamically reallocate resources by
changing the rewards and transition structure of the task: we remove the
obstacle directly adjacent to the cheese and re-position the cheese two states
to the left after 3000 episodes in a separate experiment. This modification
changes the shortest path to the cheese, and we expect to see a corresponding
change in resource allocation and the policy. To depict this change
graphically, we compute the entropy of the agent’s policy at each state during
different stages of training. Namely, we compute the choice probability vector
as the normalized histogram of $10^{5}$ actions from each state and then
compute the entropy of those vectors. In Fig. 1c, memories corresponding to
states to which dra allocates more resources are indicated by lower entropy,
i.e. less randomness, of the policy. Early on, the agent’s memories get more
precise for states that are close to the reward (top-left panel), whereas an
over-trained agent only remembers the shortest path to the reward (top-right
panel). Moreover, immediately after the change in the environment, we expect
agents to follow the previously optimal path (top-right panel) and turn left
at the end, and eventually follow the shorter paths. The bottom panels show
that agents learn to reallocate resources to better paths over time but they
still retain traces of the older memories.

We also compare dra against a model that allocates resources equally to all
memories (‘equal-precision’), but the precision shared across all memories is
otherwise subject to the same optimization procedure as dra. We find that dra
achieves 2$\times$ improvement in the objective (Eq. 1) over the equal-
precision model. Finally, we construct another baseline model by letting
$\lambda\rightarrow 0$ in dra, which reduces it to sarsa, and report that dra
only takes 1.5$\times$ the number of episodes to converge as compared to the
baseline model while making efficient use of memory resources. Similar
findings hold for all the tested hyperparameters in a wide range (see Appendix
B.2).

###  3.2 Mountain car

Next, we test dra on the mountain car problem [39], where an under-powered car
needs to reach the flag on top of the hill (Fig. 2a). At each time-step, the
car can accelerate _left_ or _right_ by a fixed amount, or do _nothing_. It
always starts at the bottom of the hill $x=-0.5$ with velocity $v=0$ and must
swing left and right, gaining momentum to progressively reach higher. In Figs.
2b-c, we show the mean of the value function for each state, computed as
$\bar{v}(s)=\max_{a}\bar{q}(s,a)$, as well as the entropy of the agent’s
policy for each state computed as described in Section 3.1. We find that dra
sensibly encodes memories corresponding to states that are close to the states
in the optimal trajectory for this task with higher precision.

Further, in Fig. 2d we show the performance of alternative gradients (Eq. 4)
that may be used to allocate resources with dra. Our results are in line with
previous work in that the advantage gradient outperforms other approaches
[32]. Finally, we perform the same analyses as in Section 3.1 by comparing dra
against an ‘equal-precision’ model that allocates resources equally to all
memories, where dra achieves a 1.3$\times$ improvement in the objective; and
by comparing dra against a baseline model ($\lambda=0$), where dra takes
1.3$\times$ the number of episodes to converge while making efficient use of
memory resources.

Figure 2: Mountain car. (a) Task. (b) Value function learnt by the agent with
a close-to-optimal route indicated by the white arrow. (c) Entropy of the
agent’s policy indicating precision of corresponding memories. (d) Maximum
objective achieved by using the Advantage function (A), the mean q-value (Q),
and reinforce (R) to compute the stochastic gradient of the first term of
$\mathcal{F}$ – the expected reward (Eqs. 1-4). Errorbars represent SD of the
mean across 5 optimization runs.

##  4 Results on a model-based planning task

###  4.1 Task details

To study the effect of resource allocation in model-based planning, we
consider here the task devised by Huys et al. [40], whose deterministic state
transitions and immediate rewards are described in Fig. 3a. Participants never
see this underlying structure, but are trained extensively on the transition
structure and immediate rewards until they pass a test. Subjects are asked to
perform $M\in\\{3,4,5\\}$ moves on each trial, indicated in advance, with the
goal of maximizing cumulative rewards. Crucially, the subjects must plan their
sequence of moves in a fixed time-period of 9s, during which they cannot act,
and get a subsequent 2.5s to execute the entire sequence of actions. With
perfect knowledge of the task, the optimal sequence can be found by simply
exploring all $2^{M}$ sequences and selecting the one associated with the
highest reward (_e.g._ , from state 5, with $M=3$, one should pick state 6, 1,
& 2). Because of the time pressure, however, subjects are unable to explore
all possible moves.

We construct an MDP for this task by expanding the state space by a factor of
$M$, allowing the task to be Markovian. In this section, we consider the MDP
for $M=3$ moves, but the results hold true generally (see Appendix B.1).
Agents thus have $N_{\text{mem}}=|\mathcal{S}|\times M\times|\mathcal{A}|=36$
memories of the form $\langle
s,a,s^{\prime},r,\mathcal{N}\big{(}\bar{q}_{sa},\sigma^{2}_{sa}\big{)}\rangle$,
with perfect recollection of $s^{\prime},r$ given $s,a$, but without prior
knowledge of the q-values which are initialized randomly around
$\bar{q}(s,a)=0\ \forall s,a$ with some precision $\sigma(s,a)=\sigma_{0}\
\forall s,a$.

In order to plan, agents start in the initial state $s_{0}$ and choose next
states according to their policy until they reach a terminal state, which is
when they reset $s=s_{0}$ and continue planning. Agents keep track of all
paths traversed and rewards accumulated for each path until they reach the
time limit, which we implement by limiting the number of states that the agent
can explore: the _search budget_. When this search budget is exhausted, agents
choose the sequence of actions corresponding to the most-rewarding path in
their working memory as shown by the schematic in Fig. 3b.

Figure 3: Planning task. (a) Task structure. (b) Paths explored and chosen in
an example trial. (c) Comparison of the optimal objective
$\mathcal{F}_{*}=\operatorname*{arg\,max}_{\sigma}\mathcal{F}$ found by
estimating the stochastic gradient $\nabla_{\sigma}\mathcal{F}$ with
$N_{\text{traj}}\in\\{1,10,50\\}$ sampled trajectories and cma-es. (d)
Difference in the optimal objective (blue) found by dra from the optimal
objective found by a model that is constrained to have equally precise
memories; orange is the same for expected reward. As the search budget
increases (planning-time pressure decreases), the advantage diminishes. (e)
Linear regression fits for the discriminability of memories (d’, left) and
average resources allocated to memories (right) as a function of their impact
on cumulative reward. As the search budget increases, dra gives up
differential allocation of resources to items in memory as it is no longer
advantageous. (f) $\mathcal{F}_{*}$ (green, left axis) and
$\widetilde{\mathcal{F}}_{*}$ (orange, right axis) as a function of search
budget. Errorbars/shaded areas represent SD of mean across 5 optimization runs
in (c),(d), and (f), and across memories in (e).

###  4.2 Comparison with an alternative model and black-box optimization

We compare the results obtained via our gradient-based resource allocation to
a black-box optimization procedure (cma-es, described in Section 2.3). For
both methods, we fix the mean of the q-value distribution to the optimal
point-estimate $\bar{\bm{q}}^{*}$ obtained via q-learning [3], and maximize
the objective $\mathcal{F}$ (Eq. 1) with respect to the resources,
$\bm{\sigma}$. Fig. 3c shows that both methods perform comparably when we
sample $N_{\text{traj}}\geq 10$ trajectories to estimate the stochastic
gradient of $\mathcal{F}$ (see Section 2.3).

Our agent is limited in its precision of q-values, but only some of them need
to be encoded precisely, namely the ones associated with decisions that have a
high impact on cumulative rewards, _e.g._ for $M=2$, $s_{0}=6$, moving to
state 1 vs. 3 results in a large difference (120 points) in cumulative reward.
To show that it is indeed beneficial to prioritize some memories over others
by encoding them with higher precision, we also trained an agent that is
constrained to have all its memories equally precise (‘equal-precision’),
though how precise is subject to the same optimization procedure. Fig. 3d
shows the advantage in performance of the agent that allocates resources
flexibly (dra) from the one that is constrained to allocate resources equally.
This advantage is more pronounced when the agent is under time pressure (lower
search budget) to plan its sequence of moves as expected.

Furthermore, this task reveals that the key factor for resource allocation is
not memory precision per se, but the _discriminability_ $d^{\prime}$ (a ratio
of difference in means divided by the effective standard deviation) between
q-values of a state. In Fig. 3e, we see this effect manifested strongly when
agents have a low search budget, and it flattens as the search budget
increases to 100%, _i.e._ when agents can explore all possible paths, such
that the reward they obtain is unaffected by the precision of q-values. Since
$d^{\prime}$ is correlated with difference in cumulative reward, we also plot
the mean $\sigma$ of memories (across actions) to show that this effect indeed
stems from resource allocation. Intuitively, dra allocates more precision to
memories when larger differences in rewards are at stake (Fig. 3e).

###  4.3 The speed-accuracy trade-off

So far, we have restricted our analyses to cases where agents are given a
fixed search budget to plan their moves, and they optimize their objective
$\mathcal{F}$ independently for each fixed budget. As expected, their
performance increases monotonically with the time they spend planning and
eventually saturates as shown in Fig. 3f (green curve). In most real-world
scenarios, however, the search budget is unknown and therefore agents need to
decide when to stop planning and execute an action. We can incorporate this
speed-accuracy trade-off [30, 41] in our framework by modifying the objective
$\mathcal{F}$ as:

$\widetilde{\mathcal{F}}\vcentcolon=\frac{\operatorname{\mathbb{E}}_{\pi}\big{[}\max_{\text{path}=1}^{bM}\big{(}\sum_{t=1}^{M}r_{t,\text{path}}\big{)}\big{|}Q)\big{]}}{\langle
a_{\text{dec}}b+t_{\text{non-dec}}\rangle}-\lambda
D_{\text{KL}}\Big{(}Q\;\|\;P\Big{)}$ (8)

where the numerator of the first term represents the expected reward from
planning specific to this task as described in Section 2.3 and Fig. 3b, and
the denominator represents the average time the agent spends per trial written
as the sum of the average decision-time (proportional to the search budget $b$
with proportionality constant $a_{\text{dec}}$) and non-decision time
$t_{\text{non-dec}}$ (_e.g._ due to sensory delay and executing motor output).
In many cases such as in this task, agents can follow locally estimated
gradients with respect to $b$ as they allocate resources dynamically with dra
to maximize $\widetilde{\mathcal{F}}$ since this turns out to be a convex
optimization problem (Fig. 3f, orange curve).

Note that in this analysis, agents draw single Thompson samples from all the
memories at each state during planning (see Eq. 2). In principle, they could
spend more time at some states than others, in which case, the effective
precision of memories would not only be controlled by the number of neurons,
but also the time available to plan. However, a detailed analysis of such
flexible planning is outside the scope of the current work.

##  5 Discussion

In this article, we propose a framework to model uncertainty in action-values
stored in limited memories, show how resource-limited agents can plan and act
with such uncertainty, and how they benefit by prioritizing memories
differentially. Our work provides a novel, normative approach to dynamically
allocate limited memory resources for finding good policies for decision-
making. Though we only consider normally distributed and independently encoded
action-values, our framework can easily be extended to arbitrary value
distributions and resource costs.

Previous work has considered prioritization of memory access guided solely by
the value of backups that lead to a change in the agent’s policy [16].
However, such a form of prioritization predicts that all memories are equally
relevant when animals are well-trained, which contradicts empirical findings
[42, 43]. In order to model and understand animal behavior, it is important to
consider irreducible sources of uncertainty in the value function besides
learning, due either to resource limitations or stochastic rewards and
dynamics [19, 44]. In future work, we would like to combine all these sources
of uncertainty in order to make experimentally testable predictions for animal
behavior.

Our work partly explains a commonly reported phenomenon in neuroscience: while
early sensory and somatosensory brain areas recruit more neurons over the
course of training, other brain areas responsible for higher-level cognitive
processes such as accessing and storing memories (prefrontal cortex) and
evidence integration during decision-making (posterior parietal cortex),
either commit more neurons or show higher levels of activity during earlier
stages of learning than later stages when animals are well-trained [37, 38,
45]. By showing that resource-constrained agents can accelerate learning by
starting with more resources (Fig. 1b), we provide a plausible hypothesis for
this observation.

Finally, our framework also makes clear predictions about how action-values
should be represented probabilistically in the brain. The details of the
predictions will depend on the nature of neural code for probability
distributions. Given that action-values are scalar variables, probabilistic
population codes – in which each neuron or sub-population is tuned to a
specific action-value – provide a biologically plausible neural code, for
which there exists strong experimental evidence [46, 47]. For this type of
code, the amplitude of the neuronal response encoding a specific action-value
should be inversely proportional to the variance associated with this action-
value [46]. Such predictions could be tested in animals trained on a foraging
task while recording in sensorimotor areas like LIP [24], the superior
colliculus [48, 49], or the basal ganglia [25] where neurons are known to
encode both actions and expected rewards. One might also imagine that,
throughout the course of learning, the number of neurons encoding action-
values, or the information-limiting correlations among these neurons [50], are
modulated so as to reflect the precision with which these action-values are
encoded.

## Broader Impact

We believe that this work has the potential to lead to a net-positive change
in the reinforcement learning community and more broadly in society as a
whole. Our work enables researchers to represent the uncertainty in memories
due to resource constraints and perform well in the face of such constraints
by prioritizing the knowledge that really matters. While our work is
preliminary, we believe that furthering this line of work may prove to be
highly beneficial in reducing the overall carbon footprint of the artificial
intelligence (AI) industry, which has recently come under scrutiny for the
jarring energy consumption of several common large AI models that produce up
to five times as much CO2 than an average American car does in its lifetime
[51, 52].

In terms of ethical aspects, our method is neutral per se. The advancement of
energy-efficient algorithms may enable autonomous agents to function for long
hours in remote areas, the applications for which could be used for both
constructive and destructive things alike, _e.g._ they may be deployed for
rescue missions [53] or weaponized for military applications [54, 55], but
this holds true for any RL agent.

## Acknowledgments and Disclosure of Funding

We thank Pablo Tano and Reidar Riveland for useful discussions, and Morio
Hamada for providing helpful feedback on a previous version of the manuscript.
Nisheet Patel was supported by the Swiss National Foundation (grant number
31003A_165831). Luigi Acerbi was partially supported by the Academy of Finland
Flagship programme: Finnish Center for Artificial Intelligence (FCAI).
Alexandre Pouget was supported by the Swiss National Foundation (grant numbers
31003A_165831 and 315230_197296).

## References

  * Kaelbling et al. [1996] Leslie Pack Kaelbling, Michael L Littman, and Andrew W Moore.  Reinforcement learning: A survey.  _Journal of artificial intelligence research_ , 4:237–285, 1996.
  * Sutton et al. [1998] Richard S Sutton, Andrew G Barto, et al.  _Introduction to reinforcement learning_ , volume 135.  MIT press Cambridge, 1998.
  * Watkins and Dayan [1992] Christopher JCH Watkins and Peter Dayan.  Q-learning.  _Machine learning_ , 8(3-4):279–292, 1992.
  * Rummery and Niranjan [1994] Gavin A Rummery and Mahesan Niranjan.  _On-line Q-learning using connectionist systems_ , volume 37.  University of Cambridge, Department of Engineering Cambridge, UK, 1994\.
  * Hardman et al. [2002] Craig Denis Hardman, Jasmine Monica Henderson, David Isaac Finkelstein, Malcolm Kenneth Horne, George Paxinos, and Glenda Margaret Halliday.  Comparison of the basal ganglia in rats, marmosets, macaques, baboons, and humans: volume and neuronal number for the output, internal relay, and striatal modulating nuclei.  _Journal of Comparative Neurology_ , 445(3):238–255, 2002.
  * Gershman et al. [2015] Samuel J Gershman, Eric J Horvitz, and Joshua B Tenenbaum.  Computational rationality: A converging paradigm for intelligence in brains, minds, and machines.  _Science_ , 349(6245):273–278, 2015.
  * Lieder and Griffiths [2020] Falk Lieder and Thomas L Griffiths.  Resource-rational analysis: understanding human cognition as the optimal use of limited computational resources.  _Behavioral and Brain Sciences_ , 43, 2020.
  * Still and Precup [2012] Susanne Still and Doina Precup.  An information-theoretic approach to curiosity-driven reinforcement learning.  _Theory in Biosciences_ , 131(3):139–148, 2012\.
  * Rubin et al. [2012] Jonathan Rubin, Ohad Shamir, and Naftali Tishby.  Trading value and information in MDPs.  In _Decision Making with Imperfect Decision Makers_ , pages 57–74. Springer, 2012.
  * Grau-Moya et al. [2016] Jordi Grau-Moya, Felix Leibfried, Tim Genewein, and Daniel A Braun.  Planning with information-processing constraints and model uncertainty in Markov decision processes.  In _Joint European Conference on Machine Learning and Knowledge Discovery in Databases_ , pages 475–491. Springer, 2016.
  * Ortega and Braun [2011] Daniel Alexander Ortega and Pedro Alejandro Braun.  Information, utility and bounded rationality.  In _International Conference on Artificial General Intelligence_ , pages 269–274. Springer, 2011.
  * Todd et al. [2009] Michael T Todd, Yael Niv, and Jonathan D Cohen.  Learning to use working memory in partially observable environments through dopaminergic reinforcement.  In _Advances in neural information processing systems_ , pages 1689–1696, 2009.
  * Suchow and Griffiths [2016] Jordan W Suchow and Tom Griffiths.  Deciding to remember: Memory maintenance as a Markov decision process.  In _CogSci_ , 2016.
  * Van den Berg and Ma [2018] Ronald Van den Berg and Wei Ji Ma.  A resource-rational theory of set size effects in human visual working memory.  _ELife_ , 7:e34963, 2018.
  * Schaul et al. [2015] Tom Schaul, John Quan, Ioannis Antonoglou, and David Silver.  Prioritized experience replay.  _arXiv preprint arXiv:1511.05952_ , 2015.
  * Mattar and Daw [2018] Marcelo G Mattar and Nathaniel D Daw.  Prioritized memory access explains planning and hippocampal replay.  _Nature neuroscience_ , 21(11):1609–1617, 2018\.
  * Haarnoja et al. [2018] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine.  Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor.  _arXiv preprint arXiv:1801.01290_ , 2018.
  * Wu et al. [2018] Yifan Wu, George Tucker, and Ofir Nachum.  The laplacian in rl: Learning representations with efficient approximations.  _arXiv preprint arXiv:1810.04586_ , 2018.
  * O’Donoghue et al. [2017] Brendan O’Donoghue, Ian Osband, Remi Munos, and Volodymyr Mnih.  The uncertainty Bellman equation and exploration.  _arXiv preprint arXiv:1709.05380_ , 2017.
  * Mao et al. [2016] Hongzi Mao, Mohammad Alizadeh, Ishai Menache, and Srikanth Kandula.  Resource management with deep reinforcement learning.  In _Proceedings of the 15th ACM Workshop on Hot Topics in Networks_ , pages 50–56, 2016.
  * Ye et al. [2019] Hao Ye, Geoffrey Ye Li, and Biing-Hwang Fred Juang.  Deep reinforcement learning based resource allocation for V2V communications.  _IEEE Transactions on Vehicular Technology_ , 68(4):3163–3173, 2019.
  * Sutton and Barto [2018] Richard S Sutton and Andrew G Barto.  Reinforcement learning: An introduction 2nd ed, 2018.
  * Hasenstaub et al. [2010] Andrea Hasenstaub, Stephani Otte, Edward Callaway, and Terrence J Sejnowski.  Metabolic cost as a unifying principle governing neuronal biophysics.  _Proceedings of the National Academy of Sciences_ , 107(27):12329–12334, 2010.
  * Sugrue et al. [2004] Leo P Sugrue, Greg S Corrado, and William T Newsome.  Matching behavior and the representation of value in the parietal cortex.  _science_ , 304(5678):1782–1787, 2004.
  * Samejima et al. [2005] Kazuyuki Samejima, Yasumasa Ueda, Kenji Doya, and Minoru Kimura.  Representation of action-specific reward values in the striatum.  _Science_ , 310(5752):1337–1340, 2005.
  * Roesch et al. [2007] Matthew R Roesch, Donna J Calu, and Geoffrey Schoenbaum.  Dopamine neurons encode the better option in rats deciding between differently delayed or sized rewards.  _Nature neuroscience_ , 10(12):1615–1624, 2007\.
  * Thompson [1933] William R Thompson.  On the likelihood that one unknown probability exceeds another in view of the evidence of two samples.  _Biometrika_ , 25(3/4):285–294, 1933.
  * Russo et al. [2017] Daniel Russo, Benjamin Van Roy, Abbas Kazerouni, Ian Osband, and Zheng Wen.  A tutorial on Thompson sampling.  _arXiv preprint arXiv:1707.02038_ , 2017.
  * Vul et al. [2014] Edward Vul, Noah Goodman, Thomas L Griffiths, and Joshua B Tenenbaum.  One and done? Optimal decisions from very few samples.  _Cognitive science_ , 38(4):599–637, 2014.
  * Heitz [2014] Richard P Heitz.  The speed-accuracy tradeoff: history, physiology, methodology, and behavior.  _Frontiers in neuroscience_ , 8:150, 2014.
  * Williams [1992] Ronald J Williams.  Simple statistical gradient-following algorithms for connectionist reinforcement learning.  _Machine learning_ , 8(3-4):229–256, 1992.
  * Schulman et al. [2015] John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel.  High-dimensional continuous control using generalized advantage estimation.  _arXiv preprint arXiv:1506.02438_ , 2015.
  * Kingma and Welling [2013] Diederik P Kingma and Max Welling.  Auto-encoding variational Bayes.  _arXiv preprint arXiv:1312.6114_ , 2013.
  * Hansen [2016] Nikolaus Hansen.  The CMA evolution strategy: A tutorial.  _arXiv preprint arXiv:1604.00772_ , 2016.
  * Hansen et al. [2019] Nikolaus Hansen, Youhei Akimoto, and Petr Baudis.  CMA-ES/pycma on Github.  Zenodo, DOI:10.5281/zenodo.2559634, February 2019.  URL https://doi.org/10.5281/zenodo.2559634.
  * Van Seijen et al. [2009] Harm Van Seijen, Hado Van Hasselt, Shimon Whiteson, and Marco Wiering.  A theoretical and empirical analysis of Expected Sarsa.  In _2009 ieee symposium on adaptive dynamic programming and reinforcement learning_ , pages 177–184. IEEE, 2009.
  * Ungerleider et al. [2002] Leslie G Ungerleider, Julien Doyon, and Avi Karni.  Imaging brain plasticity during motor skill learning.  _Neurobiology of learning and memory_ , 78(3):553–564, 2002.
  * Dayan and Cohen [2011] Eran Dayan and Leonardo G Cohen.  Neuroplasticity subserving motor skill learning.  _Neuron_ , 72(3):443–454, 2011.
  * Brockman et al. [2016] Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba.  OpenAI gym, 2016.
  * Huys et al. [2015] Quentin JM Huys, Níall Lally, Paul Faulkner, Neir Eshel, Erich Seifritz, Samuel J Gershman, Peter Dayan, and Jonathan P Roiser.  Interplay of approximate planning strategies.  _Proceedings of the National Academy of Sciences_ , 112(10):3098–3103, 2015.
  * Drugowitsch et al. [2015] Jan Drugowitsch, Gregory C DeAngelis, Dora E Angelaki, and Alexandre Pouget.  Tuning the speed-accuracy trade-off to maximize reward rate in multisensory decision-making.  _Elife_ , 4:e06678, 2015.
  * Panoz-Brown et al. [2018] Danielle Panoz-Brown, Vishakh Iyer, Lawrence M Carey, Christina M Sluka, Gabriela Rajic, Jesse Kestenman, Meredith Gentry, Sydney Brotheridge, Isaac Somekh, Hannah E Corbin, et al.  Replay of episodic memories in the rat.  _Current Biology_ , 28(10):1628–1634, 2018.
  * Gupta et al. [2010] Anoopum S Gupta, Matthijs AA van der Meer, David S Touretzky, and A David Redish.  Hippocampal replay is not a simple function of experience.  _Neuron_ , 65(5):695–705, 2010.
  * Dabney et al. [2020] Will Dabney, Zeb Kurth-Nelson, Naoshige Uchida, Clara Kwon Starkweather, Demis Hassabis, Rémi Munos, and Matthew Botvinick.  A distributional code for value in dopamine-based reinforcement learning.  _Nature_ , pages 1–5, 2020.
  * Najafi et al. [2020] Farzaneh Najafi, Gamaleldin F Elsayed, Robin Cao, Eftychios Pnevmatikakis, Peter E Latham, John P Cunningham, and Anne K Churchland.  Excitatory and inhibitory subnetworks are equally selective during decision-making and emerge simultaneously during learning.  _Neuron_ , 105(1):165–179, 2020.
  * Ma et al. [2006] Wei Ji Ma, Jeffrey M Beck, Peter E Latham, and Alexandre Pouget.  Bayesian inference with probabilistic population codes.  _Nature neuroscience_ , 9(11):1432–1438, 2006\.
  * Hou et al. [2019] Han Hou, Qihao Zheng, Yuchen Zhao, Alexandre Pouget, and Yong Gu.  Neural correlates of optimal multisensory decision making under time-varying reliabilities with an invariant linear probabilistic population code.  _Neuron_ , 104(5):1010–1021, 2019.
  * Thevarajah et al. [2010] Dhushan Thevarajah, Ryan Webb, Christopher Ferrall, and Michael C Dorris.  Modeling the value of strategic actions in the superior colliculus.  _Frontiers in behavioral neuroscience_ , 3:57, 2010.
  * Ikeda and Hikosaka [2007] Takuro Ikeda and Okihide Hikosaka.  Positive and negative modulation of motor response in primate superior colliculus by reward expectation.  _Journal of Neurophysiology_ , 98(6):3163–3170, 2007.
  * Moreno-Bote et al. [2014] Rubén Moreno-Bote, Jeffrey Beck, Ingmar Kanitscheider, Xaq Pitkow, Peter Latham, and Alexandre Pouget.  Information-limiting correlations.  _Nature neuroscience_ , 17(10):1410, 2014.
  * Hao [2019] Karen Hao.  Training a single ai model can emit as much carbon as five cars in their lifetimes.  _MIT Technology Review_ , 2019.
  * Strubell et al. [2019] Emma Strubell, Ananya Ganesh, and Andrew McCallum.  Energy and policy considerations for deep learning in NLP.  _arXiv preprint arXiv:1906.02243_ , 2019.
  * Kitano et al. [1999] Hiroaki Kitano, Satoshi Tadokoro, Itsuki Noda, Hitoshi Matsubara, Tomoichi Takahashi, Atsuhi Shinjou, and Susumu Shimada.  Robocup rescue: Search and rescue in large-scale disasters as a domain for autonomous agents research.  In _IEEE SMC’99 Conference Proceedings. 1999 IEEE International Conference on Systems, Man, and Cybernetics (Cat. No. 99CH37028)_ , volume 6, pages 739–743. IEEE, 1999.
  * Khurshid and Bing-Rong [2004] Javaid Khurshid and Hong Bing-Rong.  Military robots-a glimpse from today and tomorrow.  In _ICARCV 2004 8th Control, Automation, Robotics and Vision Conference, 2004._ , volume 1, pages 771–777. IEEE, 2004.
  * Lin et al. [2008] Patrick Lin, George Bekey, and Keith Abney.  Autonomous military robotics: Risk, ethics, and design.  Technical report, California Polytechnic State Univ San Luis Obispo, 2008\.

##  Appendix A Computing the gradient to maximize the objective function

###  A.1 Gradient of the log policy

In order to compute $\nabla_{\sigma}\log(\pi(a|s))$, we first note that we can
rewrite draws from the memory distribution,
$\tilde{q}_{sa}\sim\mathcal{N}\big{(}\bar{q}_{sa},\sigma^{2}_{sa}\big{)}$, as
$\tilde{q}_{sa}=\bar{q}_{sa}+\zeta_{sa}\sigma_{sa}$, where
$\zeta_{sa}\sim\mathcal{N}(0,1)$ [33]. In this section, we abuse the notation
slightly to omit the explicit dependence on the state-action pair $(s,a)$ for
clarity, and instead place it in the subscript. With this, we can write our
policy $\pi$ as a probability vector for all actions $a$ in a given state $s$:

$\displaystyle\pi(a|s)$
$\displaystyle=\delta\big{(}a,\operatorname*{arg\,max}_{a^{\prime}}(\bar{q}_{sa^{\prime}}+\zeta_{sa^{\prime}}\sigma_{sa^{\prime}})\big{)}$
$\displaystyle\pi(\cdot|s)$
$\displaystyle=\lim_{\beta\rightarrow\infty}softmax(\bar{\bm{q}}_{s}+\bm{\zeta}_{s}\bm{\sigma}_{s},\beta)$
$\displaystyle=\lim_{\beta\rightarrow\infty}\frac{1}{\sum_{a}\exp\beta(\bar{q}_{sa}+\zeta_{sa}\sigma_{sa})}\begin{bmatrix}\exp\beta(\bar{q}_{sa_{1}}+\zeta_{sa_{1}}\sigma_{sa_{1}})\\\
\vdots\\\
\exp\beta(\bar{q}_{sa_{n}}+\zeta_{sa_{n}}\sigma_{sa_{n}})\end{bmatrix},$ (A.9)

where in the first line we applied the Thompson sampling rule (that is, pick
the action with maximal sampled value), in the second line we rewrote it as
the limit of a _softmax_ with inverse temperature $\beta\rightarrow\infty$,
and in the last line we wrote the _softmax_ explicitly (as a vector for each
entry of $\pi(\cdot|s)$).

Next, we relax the limit $\beta\rightarrow\infty$ in Eq. A.1 so as to
differentiate the logarithm of the policy $\log\pi$ for $\beta>0$ with respect
to the relevant elements of the resource allocation vector $\sigma(s,a)$ as
follows:

$\displaystyle\frac{\partial}{\partial\sigma_{sa}}\log\pi(\cdot|s)$
$\displaystyle=-\frac{\partial}{\partial\sigma_{sa}}\log\Big{(}\sum_{a}\exp\beta(\bar{q}_{sa}+\zeta_{sa}\sigma_{sa})\Big{)}+\frac{\partial}{\partial\sigma_{sa}}\begin{bmatrix}\beta(\bar{q}_{sa_{1}}+\zeta_{sa_{1}}\sigma_{sa_{1}})\\\
\vdots\\\ \beta(\bar{q}_{sa_{n}}+\zeta_{sa_{n}}\sigma_{sa_{n}})\end{bmatrix}$
$\displaystyle=-\frac{\exp\beta(\bar{q}_{sa}+\zeta_{sa}\sigma_{sa})}{\sum_{a}\exp\beta(\bar{q}_{sa}+\zeta_{sa}\sigma_{sa})}\beta\zeta_{sa}+\beta\zeta_{sa}\delta(a,a_{i})$
$\displaystyle=\beta\zeta_{sa}\big{(}\delta(a,a_{i})-\pi(a|s)\big{)}$ (A.10)

where the final step follows from rewriting the softmax function as the (soft)
policy $\pi(a|s)$ in Eq. A.1, i.e. with some $\beta>0$ but not
$\beta\rightarrow\infty$.

Thus, the gradient of the logarithm of the policy $\log\pi(a|s)$ with respect
to the resource allocation vector $\bm{\sigma}$ can be written as:

$\displaystyle\frac{\partial}{\partial\sigma_{s^{\prime}a^{\prime}}}\log\pi(a|s)$
$\displaystyle=\begin{cases}\beta\zeta_{sa}(1-\pi(a|s))&\small{\text{for
}s^{\prime}=s,a^{\prime}=a}\\\
-\beta\zeta_{sa^{\prime}}\pi(a^{\prime}|s)&\small{\text{for
}s^{\prime}=s,a^{\prime}\neq a}\\\ 0&\small{\text{for }s^{\prime}\neq
s}\end{cases},$ (A.11)

which is reported as Eq. 5 in the main text.

###  A.2 Gradient of the cost

In this section, we show how to compute
$\nabla_{\sigma}D_{\text{KL}}(Q\;\|\;P)$, where
$Q=\mathcal{N}(\bm{\bar{q}},\bm{\sigma}^{2}I)$ and
$P=\mathcal{N}(\bar{\bm{q}},{\sigma}_{\text{base}}^{2}I)$, and $I$ is the
identity matrix. Since the covariance matrix is diagonal, we can take the
gradient with respect to elements of the resource allocation vector
$\bm{\sigma}$ individually. In other words, we can take the gradient of each
memory’s marginal normal distribution with its standard deviation:

$\displaystyle\frac{\partial}{\partial\sigma}D_{\text{KL}}\Big{(}\mathcal{N}\big{(}{\bar{q}},{\sigma}^{2}\big{)}\;\|\;\mathcal{N}\big{(}{\bar{q}},\sigma_{\text{base}}^{2}\big{)}\Big{)}$
$\displaystyle=\frac{\partial}{\partial\sigma}\operatorname{\mathbb{E}}_{Q}\bigg{[}\log\bigg{(}\frac{Q}{P}\bigg{)}\bigg{]}$
$\displaystyle=\frac{\partial}{\partial\sigma}\operatorname{\mathbb{E}}_{Q}[\log(Q)]-\frac{\partial}{\partial\sigma}\operatorname{\mathbb{E}}_{Q}[\log(P)].$
(A.12)

We can expand the first of the two terms as:

$\displaystyle\frac{\partial}{\partial\sigma}\operatorname{\mathbb{E}}_{Q}[\log(Q)]$
$\displaystyle=\frac{\partial}{\partial\sigma}\operatorname{\mathbb{E}}_{x\sim\mathcal{N}({\bar{q}},{\sigma}^{2})}\Bigg{[}\log\Bigg{(}\frac{1}{\sqrt{2\pi\sigma^{2}}}\exp{-\frac{1}{2}\bigg{(}\frac{x-\bar{q}}{\sigma}\bigg{)}^{2}}\Bigg{)}\Bigg{]}$
$\displaystyle=\frac{\partial}{\partial\sigma}\operatorname{\mathbb{E}}_{x\sim\mathcal{N}({\bar{q}},{\sigma}^{2})}\Bigg{[}-\frac{1}{2}\log(2\pi)-\log(\sigma)-\frac{1}{2}\bigg{(}\frac{x-\bar{q}}{\sigma}\bigg{)}^{2}\Bigg{]}$
$\displaystyle=-\frac{1}{2}\cancelto{0}{\frac{\partial}{\partial\sigma}\log(2\pi)}-\frac{\partial}{\partial\sigma}\log(\sigma)-\frac{1}{2}\frac{\partial}{\partial\sigma}\operatorname{\mathbb{E}}_{x\sim\mathcal{N}({\bar{q}},{\sigma}^{2})}\Bigg{[}\bigg{(}\frac{x-\bar{q}}{\sigma}\bigg{)}^{2}\Bigg{]}$
$\displaystyle=-\frac{1}{\sigma}-\frac{1}{2}\frac{\partial}{\partial\sigma}\operatorname{\mathbb{E}}_{z\sim\mathcal{N}(0,1)}\big{[}z^{2}\big{]}$
$\displaystyle=-\frac{1}{\sigma}$ (A.13)

where we use the variable transformation $z=(x-\bar{q})/\sigma$ in the
penultimate step. We can follow a similar approach for the second term to get:

$\displaystyle\frac{\partial}{\partial\sigma}\operatorname{\mathbb{E}}_{Q}[\log(P)]$
$\displaystyle=\frac{\partial}{\partial\sigma}\operatorname{\mathbb{E}}_{x\sim\mathcal{N}({\bar{q}},{\sigma}^{2})}\Bigg{[}\log\Bigg{(}\frac{1}{\sqrt{2\pi\sigma_{\text{base}}^{2}}}\exp{-\frac{1}{2}\bigg{(}\frac{x-\bar{q}}{\sigma_{\text{base}}}\bigg{)}^{2}}\Bigg{)}\Bigg{]}$
$\displaystyle=\frac{\partial}{\partial\sigma}\operatorname{\mathbb{E}}_{x\sim\mathcal{N}({\bar{q}},{\sigma}^{2})}\Bigg{[}-\frac{1}{2}\log\big{(}2\pi\sigma_{\text{base}}^{2}\big{)}-\frac{1}{2}\bigg{(}\frac{x-\bar{q}}{\sigma}\bigg{)}^{2}\frac{\sigma^{2}}{\sigma_{\text{base}}^{2}}\Bigg{]}$
$\displaystyle=-\frac{1}{2}\cancelto{0}{\frac{\partial}{\partial\sigma}\log\big{(}2\pi\sigma_{\text{base}}^{2}\big{)}}-\frac{1}{2}\frac{\partial}{\partial\sigma}\operatorname{\mathbb{E}}_{x\sim\mathcal{N}({\bar{q}},{\sigma}^{2})}\Bigg{[}\bigg{(}\frac{x-\bar{q}}{\sigma}\bigg{)}^{2}\Bigg{]}\frac{\sigma^{2}}{\sigma_{\text{base}}^{2}}$
$\displaystyle=-\frac{1}{2}\frac{\partial}{\partial\sigma}\operatorname{\mathbb{E}}_{z\sim\mathcal{N}(0,1)}\big{[}z^{2}\big{]}\frac{\sigma^{2}}{\sigma_{\text{base}}^{2}}$
$\displaystyle=-\frac{1}{2}\frac{\partial}{\partial\sigma}\frac{\sigma^{2}}{\sigma_{\text{base}}^{2}}$
$\displaystyle=-\frac{\sigma}{\sigma_{\text{base}}^{2}}.$ (A.14)

Combining Eqs. A.2, A.2, and A.2, we can write our analytically obtained
gradient of the cost term with respect to individual elements of the resource
allocation vector $\sigma(s,a)$ as:

$\frac{\partial}{\partial\sigma_{sa}}D_{\text{KL}}\Big{(}\mathcal{N}\big{(}{\bar{\bm{q}}},{\bm{\sigma}}^{2}I\big{)}\;\|\;\mathcal{N}({\bar{\bm{q}}},{\sigma}_{\text{base}}^{2}I)\Big{)}=\frac{\sigma_{sa}}{\sigma_{\text{base}}^{2}}-\frac{1}{\sigma_{sa}}.$
(A.15)

###  A.3 Justification for our choice of the gradient of expected reward

A potential concern regarding our method of allocating resources may be our
choice of the advantage function to compute the gradient of the expected
rewards (Eqs. 3-4 in the main text). Crucially, the advantage gradient uses
the means of the q-value distributions of the relevant memories. However, our
main assumption in the paper is that agents do _not_ have direct access to the
mean of the q-value distribution. According to our assumption, the agent could
only estimate the mean by averaging over a large number of samples from the
distribution, a process which could take a considerable amount of time
(because sequential samples from memory would be highly correlated [29]).

This concern is resolved by considering that in dra the resource allocation
vector is not updated during the trial, but rather only _offline_ , _i.e._
before or after the trial, or potentially during sleep. This way, during the
task, the agent draws single (Thompson) samples in order to act and does not
waste extra time in order to consolidate and reallocate resources across its
memories.

While ‘offline sampling’ resolves the issue of how agents can access the mean
of the distribution to compute policy updates, and it is the approach followed
in this work, it represents a binary solution (_i.e._ , either the agent takes
one Thompson sample online, or a very large number of them offline). We could
generalize this approach by allowing an agent to take _multiple_ samples from
its q-value distribution to get a better estimate of the expected return while
performing the task. Taking additional samples would cost them time, which
they could potentially use to act in the environment and collect rewards. If
the opportunity cost is higher than the potential increase in rewards obtained
by taking more samples, they may not want to waste time sampling but instead
make their memories (q-value distributions) precise enough that fewer samples
suffice to maximize reward given their storage capacity. This is another
example of the speed-accuracy trade-off we considered in Section 4.3 in the
main text, and which we leave to explore for future work.

##  Appendix B Task parameters and additional results

###  B.1 Additional results for the planning task

In the main article, we showed results for the planning task we adapted from
Huys et al. [40] where subjects had to plan sequences of $M=3$ moves. More
generally, we ran dra for $M\in\\{3,4,5\\}$, showing that the algorithm
allocates resources differentially depending on $M$ (Fig. B.4).

Figure B.4: Linear regression fits for the discriminability of memories as a
function of their impact on cumulative reward for the planning task with
number of moves $M\in\\{3,4,5\\}$.

###  B.2 Task parameters

Table B.1: Parameters used for each task Task
---
Parameter | Grid-world | Mountain Car | Planning task
$\alpha_{1}$ | 0.1 | 0.1 | 0.1
$\alpha_{2}$ | 0.1 | 0.1 | 0.1
$\beta$ | 10 | 10 | 10
$\gamma$ | 1 | 1 | 1
$\lambda$ | 0.2 | 0.1 | 1
$\sigma_{\text{base}}$ | 5 | 5 | 100
$\sigma_{0}$ | 3 | 3 | 50
$N_{\text{traj}}$ | 10 | 10 | 10
$N_{\text{restarts}}$ | 5 | 5 | 5

In this section, we report (Table B.1) and briefly describe the
(hyper-)parameters chosen for each task. For the present study, we fixed the
learning rates for the means and standard deviations of the memory
distribution, $\alpha_{1}$ and $\alpha_{2}$ respectively, to reasonably low
values. We set the inverse temperature parameter $\beta$ to a reasonably high
value (for the _softmax_ approximation to the ‘hard’ max to hold, as per
Section A.1), but not too high to restrict the influence of individual updates
on the resource allocation. As mentioned in the main text, we exclude
discounting for all the tasks and thus set $\gamma=1$. Perhaps the most
important choice is the parameter $\lambda$ that introduces a trade-off
between the expected reward and the cost of being precise. We chose $\lambda$
that we best captured the difficulties faced by memory-limited agents, but a
range of nearby values yields qualitatively similar results, e.g. in the
mountain car task, $\lambda\in[0,0.4]$ allows agents to perform the task well
with enough training. The other equally important parameter would perhaps be
$\sigma_{\text{base}}$, which would represent the resources for some base
distribution of q-values in memory before training. $\sigma_{\text{base}}$
controls how discriminable different actions would be from a given state, and
we chose it appropriately given the reward structure of each task. As
mentioned in the main text, starting with a higher resource budget than the
base distribution, i.e. with $\sigma_{0}<\sigma_{\text{base}}$, either by
means of paying more attention or allocating more neurons, allows agents to
accelerate learning. We sampled $N_{\text{traj}}=10$ trajectories to update
the resource allocation vector at the end of each trial with adequate
precision. As shown in Fig. 3c in the main text, sampling more trajectories
does not yield better performance, but less leads to variability in the
stochastic estimate of the gradient and thus hurts performance. Finally, we
performed $N_{\text{restarts}}=5$ optimization runs for each task to report an
estimate of variability across runs (_i.e._ , error bars), as mentioned in the
main text. In addition to the above parameters used to display the results, we
systematically varied the values of $\lambda$ and $\sigma_{\text{base}}$ in
all tasks and report that the qualitative results hold for a large range of
values of these parameters with $\sigma_{\text{base}}$ having a slightly
stronger effect than $\lambda$.