# Predictive representations:
building blocks of intelligence

Wilka Carvalho,1∗ Momchil S. Tomov,2,3∗ William de Cothi,4∗
Caswell Barry,4 Samuel J. Gershman1,2,5
1Kempner Institute for the Study of Intelligence, Harvard University
2Department of Psychology and Center for Brain Science, Harvard University
3Motional AD LLC
4Department of Cell and Developmental Biology, University College London
5Center for Brains, Minds, and Machines, MIT
∗Equal contribution

###### Abstract

Adaptive behavior often requires predicting future events. The theory of
reinforcement learning prescribes what kinds of predictive representations are
useful and how to compute them. This paper integrates these theoretical ideas
with work on cognition and neuroscience. We pay special attention to the
_successor representation_ (SR) and its generalizations, which have been
widely applied both as engineering tools and models of brain function. This
convergence suggests that particular kinds of predictive representations may
function as versatile building blocks of intelligence.

###### Contents

  1. 1 Introduction
  2. 2 Theory
    1. 2.1 The reinforcement learning problem
    2. 2.2 Classical solution methods
    3. 2.3 The successor representation
    4. 2.4 Successor models: a probabilistic perspective on the SR
    5. 2.5 Successor features: a feature-based generalization of the SR
      1. 2.5.1 Generalized policy improvement: adaptively combining behaviors
      2. 2.5.2 Option Keyboard: chaining together behaviors
    6. 2.6 Summary
  3. 3 Practical learning algorithms and associated challenges
    1. 3.1 Learning successor features
      1. 3.1.1 Discovering cumulants
      2. 3.1.2 Estimating successor features
    2. 3.2 Learning successor models
      1. 3.2.1 Learning successor models that one can sample from
      2. 3.2.2 Learning successor models that one can evaluate
  4. 4 Artificial intelligence applications
    1. 4.1 Exploration
      1. 4.1.1 Pure exploration
      2. 4.1.2 Balancing exploration and exploitation
    2. 4.2 Transfer
      1. 4.2.1 Transferring behaviors between tasks
      2. 4.2.2 Learning about non-task goals
      3. 4.2.3 Other advances in transfer
    3. 4.3 Hierarchical reinforcement learning
      1. 4.3.1 Discovering options to efficiently explore the environment
      2. 4.3.2 Transferring options with the SR
    4. 4.4 Model-based reinforcement learning
    5. 4.5 Multi-agent reinforcement learning
    6. 4.6 Other artificial intelligence applications
  5. 5 Cognitive science applications
    1. 5.1 Revaluation
    2. 5.2 Multi-task learning
    3. 5.3 Associative learning
      1. 5.3.1 Bayesian RL
      2. 5.3.2 Context-dependent learning
      3. 5.3.3 Bayesian learning of context-dependent predictive maps
    4. 5.4 Spatial navigation
    5. 5.5 Memory
      1. 5.5.1 The temporal context model
      2. 5.5.2 TCM as an estimator for the SR
      3. 5.5.3 Combining TCM and the SR for decision making
  6. 6 Neuroscience applications
    1. 6.1 A brief introduction to the medial temporal lobe
    2. 6.2 The hippocampus as a predictive map
    3. 6.3 Learning a biologically plausible SR
    4. 6.4 Replay
    5. 6.5 Dopamine and generalized prediction errors
  7. 7 Conclusions

##  1 Introduction

The ability to make predictions has been hailed as a general feature of both
biological and artificial intelligence, cutting across disparate perspectives
on what constitutes intelligence (e.g., Ciria et al.,, 2021; Clark,, 2013;
Friston and Kiebel,, 2009; Ha and Schmidhuber,, 2018; Hawkins and Blakeslee,,
2004; Littman and Sutton,, 2001; Lotter et al.,, 2016). Despite this general
agreement, attempts to formulate the idea more precisely raise many questions:
Predict what, and over what timescale? How should predictions be represented?
How should they be used, evaluated, and improved? These normative “should”
questions have corresponding empirical questions about the nature of
prediction in biological intelligence. Our goal is to provide systematic
answers to these questions. We will develop a small set of principles that
have broad explanatory power.

Our perspective is based on an important distinction between predictive
_models_ and predictive _representations_. A predictive model is a probability
distribution over the dynamics of a system’s state. A model can be “run
forward” to generate predictions about the system’s future trajectory. This
offers a significant degree of flexibility: an agent with a predictive model
can, given enough computation time, answer virtually any query about the
probabilities of future events. However, the “given enough computation time”
proviso places a critical constraint on what can be done with a predictive
model in practice. An agent that needs to act quickly under stringent
computational constraints may not have the luxury of posing arbitrarily
complex queries to its predictive model. Predictive representations, on the
other hand, cache the answers to certain queries, making them accessible with
limited computational cost. The price paid for this efficiency gain is a loss
of flexibility: only certain queries can be accurately answered.

Caching is a general solution to ubiquitous flexibility-efficiency trade-offs
facing intelligent systems (Dasgupta and Gershman,, 2021). Key to the success
of this strategy is caching representations that make task-relevant
information directly accessible to computation. We will formalize the notion
of task-relevant information, as well as what kinds of computations access and
manipulate this information, in the framework of reinforcement learning (RL)
theory (Sutton and Barto,, 2018). In particular, we will show how one family
of predictive representation, the _successor representation_ (SR) and its
generalizations, distills information that is useful for efficient computation
across a wide variety of RL tasks. These predictive representations facilitate
exploration, transfer, temporal abstraction, unsupervised pre-training, multi-
agent coordination, creativity, and episodic control. On the basis of such
versatility, we argue that these predictive representations can serve as
fundamental building blocks of intelligence.

Converging support for this argument comes from cognitive science and
neuroscience. We review a body of data indicating that the brain uses
predictive representations for a range of tasks, including decision making,
navigation, and memory. We then discuss biologically plausible algorithms for
learning and computing with predictive representations. This convergence of
biological and artificial intelligence is unlikely to be a coincidence:
predictive representations may be an inevitable tool for intelligent systems.

##  2 Theory

In this section, we introduce the general problem setup and a classification
of solution techniques. We then formalize the SR and discuss how it fits into
the classification scheme. Finally, we describe two key extensions of the SR
that make it much more powerful: the successor model and successor features.
Due to space constraints, we omit some more exotic variants such as the first-
occupancy representation (Moskovitz et al.,, 2021) or the forward-backward
representation (Touati et al.,, 2022).

###  2.1 The reinforcement learning problem

We consider an agent situated in a _Markov decision process_ (MDP) defined by
the tuple $M=(\gamma,\mathcal{S},\mathcal{A},T,R)$, where $\gamma\in[0,1)$ is
a discount factor, $\mathcal{S}$ is a set of states (the state space),
$\mathcal{A}$ is a set of actions (the action space), $T(s^{\prime}|s,a)$ is
the probability of transitioning from state $s$ to state $s^{\prime}$ after
taking action $a$, and $R(s)$ is the expected reward in state $s$.111For
notational convenience, we will assume that the state and action spaces are
both discrete, but this assumption is not essential. Following Sutton and
Barto, (2018), we consider settings where the MDP and an agent give rise to a
trajectory of experience

$s_{0},a_{0},r_{1},s_{1},a_{1},\ldots$ (1)

where state $s_{t}$ and action $a_{t}$ lead to reward $r_{t+1}$ and state
$s_{t+1}$. The agent chooses actions probabilistically according to a state-
dependent policy $\pi(a|s)$.

We consider settings where the agent prioritizes immediate reward over future
reward so we focus on discounted returns. The _value_ of a policy is the
expected222To simplify notation, we will sometimes leave implicit the
distributions over which the expectation is being taken. discounted return:

$\displaystyle
V^{\pi}(s)=\mathbb{E}\left[\sum_{t=0}^{H}\gamma^{t}R(s_{t+1})\bigg{|}s_{0}=s\right],$
(2)

One can also define a state-action value (i.e., the expected discounted return
conditional on action $a$ in state $s$) by:

$\displaystyle
Q^{\pi}(s,a)=\mathbb{E}\left[\sum_{t=0}^{H}\gamma^{t}R(s_{t+1})\bigg{|}s_{0}=s,a_{0}=a\right]$
(3)

The optimal policy $\pi^{\ast}$ is then defined by:

$\displaystyle\pi^{\ast}(\cdot|s)=\operatornamewithlimits{argmax}_{\pi(\cdot|s)}V^{\pi}(s)=\operatornamewithlimits{argmax}_{\pi(\cdot|s)}\sum_{a}\pi(a|s)Q^{\pi}(s,a),$
(4)

The optimal policy for an MDP is always deterministic (choose the value-
maximizing action):

$\displaystyle\pi^{\ast}(a|s)=\mathbb{I}[a=\operatornamewithlimits{argmax}_{\tilde{a}}Q^{\ast}(s,\tilde{a})],$
(5)

where $\mathbb{I}[\cdot]=1$ if its argument is true, 0 otherwise. This assumes
that the agent can compute the optimal values
$Q^{\ast}(s,a)=\max_{\pi}Q^{\pi}(s,a)$ exactly. In most practical settings,
values must be approximated. In these cases, stochastic policies are useful
(e.g., for exploration), as discussed later.

The Markov property for MDPs refers to the conditional independence of the
past and future given the current state and action. This property allows us to
write the value function in a recursive form known as the _Bellman equation_
(Bellman,, 1957):

$\displaystyle V^{\pi}(s)$
$\displaystyle=\sum_{a}\pi(a|s)\sum_{s^{\prime}}T(s^{\prime}|s,a)\left[R(s^{\prime})+\gamma
V^{\pi}(s^{\prime})\right]$ $\displaystyle=\mathbb{E}[R(s^{\prime})+\gamma
V^{\pi}(s^{\prime})].$ (6)

Similarly, the state-action value function obeys a Bellman equation:

$\displaystyle Q^{\pi}(s,a)=\mathbb{E}[R(s^{\prime})+\gamma
Q^{\pi}(s^{\prime},a^{\prime})].$ (7)

These Bellman equations lie at the heart of many efficient RL algorithms, as
we discuss next.

###  2.2 Classical solution methods

We say that an algorithm “solves” an MDP if it outputs an optimal policy (or
an approximation thereof). Broadly speaking, algorithms can be divided into
two classes:

  * •

Model-based algorithms use an internal model of the MDP to compute the optimal
policy (Figure 1, left).

  * •

Model-free algorithms compute the optimal policy by interacting with the MDP
(Figure 1, middle).

These definitions allow us to be precise about what we mean by _predictive
model_ and _predictive representation_. A predictive model corresponds to
$\hat{T}$, the internal model of the transition distribution, and $\hat{R}$,
an internal model of the reward function. An agent equipped with $\hat{T}$ can
simulate state trajectories and answer arbitrary queries about the future. Of
principal relevance to solving MDPs is policy evaluation—an answer to the
query, “How much reward do I expect to earn in the future under my current
policy?” A simple (but inefficient) way to do this, known as _Monte Carlo
policy evaluation_ , is by running many simulations from each state (_roll-
outs_) and then averaging the discounted return.333Technically, this approach
won’t converge to the correct answer in the infinite horizon setting with
finite-length roll-outs, though the discount factor ensures that rewards in
the distant future have a vanishing effect on the discounted return. The basic
problem with this approach stems from the _curse of dimensionality_ (Bellman,,
1957): the trajectory space is enormous, requiring a number of roll-outs that
is exponential in the trajectory length.

A better model-based approach exploits the Bellman equation. For example, the
_value iteration_ algorithm starts with an initial estimate of the value
function, $\hat{V}^{\pi}$, then simultaneously improves this estimate and the
policy by applying the following update (known as a _Bellman backup_) to each
state:

$\displaystyle\hat{V}^{\pi}(s)\leftarrow\max_{a}\sum_{s^{\prime}}\hat{T}(s^{\prime}|s,a)\left[\hat{R}(s^{\prime})+\gamma\hat{V}^{\pi}(s^{\prime})\right].$
(8)

This is a form of dynamic programming, guaranteed to converge to the optimal
solution, $V^{\ast}(s)=\max_{\pi}V^{\pi}(s)$, when the agent’s internal model
is accurate ($\hat{T}=T,\hat{R}=R$). After convergence, the optimal policy for
state $s$ is given by:

$\displaystyle\pi^{\ast}(a|s)=\mathbb{I}[a=a^{\ast}(s)],$ (9) $\displaystyle
a^{\ast}(s)=\operatornamewithlimits{argmax}_{a}Q^{\ast}(s,a)\approx\operatornamewithlimits{argmax}_{a}\hat{Q}^{\pi}(s,a).$
(10)

The approximation becomes an equality when the agent’s internal model is
accurate.

Figure 1:  Algorithmic solutions to the RL problem. An agent solving a three-
armed maze (bottom) can adopt different classes of strategies (top). Model-
based strategies (left) learn an internal model of the environment, including
the transition function ($T$), the reward function ($R$), and (optionally) the
features ($\boldsymbol{\phi}$). At decision time, the agent can run forward
simulations to predict the outcomes of different actions. Model-free
strategies (middle) learn action values ($Q$) and/or a policy ($\pi$). At
decision time, the agent can consult the cached action values and/or policy in
the current state. Strategies relying on predictive representations (right)
learn the successor representation (SR) matrix ($\mathbf{M}$) mapping states
to future states and/or the successor features ($\boldsymbol{\psi}$) mapping
states to future features, as well as the reward function ($R$). At decision
time, the agent can consult the cached predictions and cross-reference them
with its task (specified by the reward function) to choose an action.

Value iteration is powerful, but still too cumbersome for large state spaces,
since each iteration requires $\mathcal{O}(|\mathcal{A}||\mathcal{S}|^{2})$
steps. The basic problem is that algorithms like value iteration attempt to
compute the optimal policy for _every_ state, but in an online setting an
agent only needs to worry about what action to take in its _current_ state.
This problem is addressed by tree search algorithms, which rely on roll-outs
(as in Monte Carlo policy evaluation), but only from the current state. When
combined with heuristics for determining which roll-outs to perform (e.g.,
learned value functions; see below), this approach can be highly effective
(Silver et al.,, 2016).

Despite their effectiveness for certain problems (e.g., games like Go and
chess), model-based algorithms have had only limited success in a wider range
of problems (e.g., video games) due to the difficulty of learning a good model
and planning in complex (possibly infinite/continuous) state spaces.444Recent
work on applying model-based approaches to video games has seen some success
(Tsividis et al.,, 2021), but progress towards scalable and generally
applicable versions of such algorithms is still in its infancy. For this
reason, much of the work in modern RL has focused on model-free algorithms.

A model-free agent by definition has no access to $\hat{T}$ (and sometimes no
access to $\hat{R}$), but nonetheless can still answer certain queries about
the future if it has cached a predictive representation. For example, an agent
could cache an estimate of the state-action value function,
$\hat{Q}^{\pi}(s,a)\approx Q^{\pi}(s,a)$. This predictive representation does
not afford the same flexibility as a model of the MDP, but it has the
advantage of caching, in a computationally convenient form, exactly the
information about the future that an agent needs to act optimally.

Importantly, $\hat{Q}^{\pi}(s,a)$ can be learned purely from interacting with
the environment, without access to a model. For example, _temporal difference_
(TD) learning methods use stochastic approximation of the Bellman backup.
Q-learning is the canonical algorithm of this kind:

$\displaystyle\hat{Q}^{\pi}(s,a)\leftarrow\hat{Q}^{\pi}(s,a)+\eta\delta$ (11)
$\displaystyle\delta=R(s^{\prime})+\gamma\max_{a^{\prime}}\hat{Q}^{\pi}(s^{\prime},a^{\prime})-\hat{Q}^{\pi}(s,a),$
(12)

where $\eta\in[0,1]$ is a learning rate and $s^{\prime}$ is sampled from
$T(s^{\prime}|s,a)$. When the estimate is exact, $\mathbb{E}[\delta]=0$ and
$\hat{Q}^{\pi}=Q^{\pi}$. Moreover, these updates converge to $Q^{\ast}(s,a)$
with probability 1 (Watkins and Dayan,, 1992).

We have briefly discussed the dichotomy of model-based versus model-free
algorithms for learning an optimal policy. Model-based algorithms are more
flexible—capable of generating predictions about future trajectories—while
model-free algorithms are more computationally efficient—capable of rapidly
computing the approximate value of an action. The flexibility of model-based
algorithms is important for _transfer_ : when the environment changes locally
(e.g., a route is blocked, or the value of a state is altered), an agent’s
model will typically also change locally, allowing it to transfer much of its
previously learned knowledge without extensive new learning. In contrast, a
cached value function approximation (due to its long-term dependencies) will
change non-locally, necessitating more extensive learning to update all the
affected cached values.

One of the questions we aim to address is how to get some aspects of model-
based flexibility without learning and computing with a predictive model. This
leads us to another class of predictive representations: the SR. In this
section, we describe the SR (§2.3), its probabilistic variant (the successor
model; §2.4), and an important generalization (successor features; §2.5).
Applications of these concepts will be covered in §4.

###  2.3 The successor representation

Figure 2:  The successor representation (SR). We present (a) a schematic of an
environment where the agent is a red box at state $s_{13}$ and the goal is a
green box at state $s_{5}$. In general, an SR $M^{\pi}(s,\tilde{s})$ (Eq. 13)
describes the discounted state occupancy for state $\tilde{s}$ when beginning
at state $s$ and following policy $\pi$. In panels (b-c), we showcase
$M^{\pi}(s_{13},\tilde{s})$ for a random policy and an optimal policy. In (b),
we can see that the SR under a random policy measures high state occupancy
near the agent’s current state (e.g. $M^{\pi}(s_{13},s_{14})=5.97$) and low
state occupancy at points further away from the agent (e.g.
$M^{\pi}(s_{13},s_{12})=0.16$). In (c), we see the SR under the optimal policy
has highest state occupancy along the shortest path to the goal (e.g. here
$M^{\pi}(s_{13},s_{12})=.66$), fading as we get further from the current
state. In contrast to a random policy, states not along that path have $0$
occupancy (e.g. $M^{\pi}(s_{13},s_{19})=0.0$). Once we know a reward function,
we can efficiently evaluate both policies (Eq. 19). (d) We show an example
reward function that has a cost of $-0.1$ for each state except the goal state
where reward is $1$. The SR allows us to efficiently compute (e) the value
function under a random policy and (f) the value function under the optimal
policy.

The SR, denoted $M^{\pi}$, was introduced to address the transfer problem
described in the previous section (Dayan,, 1993; Gershman,, 2018). In
particular, the SR is well-suited for solving sets of tasks that share the
same transition structure but vary in their reward structure; we will delve
more into this problem setting later when we discuss applications.

Like the value function, the SR is a _cumulant_ —a quantity of interest summed
(accumulated) across time. The value function is a cumulant of discounted
reward, whereas the SR is a cumulant of discounted _occupancy_ :

$M^{\pi}\left(s,\tilde{s}\right)=\mathbb{E}\left[\sum_{t=0}^{H}\gamma^{t}\mathbb{I}\left[s_{t+1}=\tilde{s}\right]\bigg{|}s_{0}=s\right]$
(13)

Here $\tilde{s}$ denotes a future state, and $M^{\pi}(s,\tilde{s})$ is the
expected discounted future occupancy of $\tilde{s}$ starting in state $s$
under policy $\pi$. An illustration of the SR and comparison with the value
function is shown in Figure 2.

The SR can be derived analytically from the transition function when
$H=\infty$:

$\displaystyle\mathbf{M}^{\pi}=\sum_{t=0}^{\infty}\gamma^{t}[\mathbf{T}^{\pi}]^{t}-\mathbf{I}=(\mathbf{I}-\gamma\mathbf{T}^{\pi})^{-1}-\mathbf{I},$
(14)

where we have subtracted the identity matrix $\mathbf{I}$ because we are not
counting the current state, and the marginal transition matrix
$\mathbf{T}^{\pi}$ is given by:

$\displaystyle T^{\pi}(s,s^{\prime})=\sum_{a}\pi(a|s)T(s^{\prime}|s,a).$ (15)

Eq. 14 makes explicit the sense in which the SR (a predictive representation)
is a compilation of the transition matrix (a predictive model). The SR
discards information about individual transitions, replacing them with their
cumulants, analogous to how the value function replaces individual reward
sequences with their cumulants.

Like the value function, the SR obeys a Bellman equation:

$\displaystyle M^{\pi}(s,\tilde{s})$
$\displaystyle=\sum_{a}\pi(a|s)\sum_{s^{\prime}}T(s^{\prime}|s,a)\big{[}\mathbb{I}[s^{\prime}=\tilde{s}]+\gamma
M^{\pi}(s^{\prime},\tilde{s})\big{]}$
$\displaystyle=\mathbb{E}\bigg{[}\mathbb{I}[s^{\prime}=\tilde{s}]+\gamma
M^{\pi}(s^{\prime},\tilde{s})\bigg{]}.$ (16)

This means that it is possible to learn the SR using TD updates similar to the
ones applied to value learning:

$\displaystyle\hat{M}^{\pi}(s,\tilde{s})\leftarrow\hat{M}^{\pi}(s,\tilde{s})+\eta\delta_{M}(\tilde{s}),$
(17)

where

$\displaystyle\delta_{M}(\tilde{s})=\mathbb{I}[s^{\prime}=\tilde{s}]+\gamma\hat{M}^{\pi}(s^{\prime},\tilde{s})-\hat{M}^{\pi}(s,\tilde{s})$
(18)

is the TD error. Notice that, unlike in TD learning for value, the error is
now vector-valued (one error for each state). Once the SR is learned, the
value function for a particular reward function under the policy $\pi$ can be
efficiently computed as a linear function of the SR:

$V^{\pi}(s)=\sum_{\tilde{s}}M^{\pi}\left(s,\tilde{s}\right)R\left(\tilde{s}\right)$
(19)

Intuitively, Eq. 19 expresses a decomposition of future reward into immediate
reward in each state and the frequency with which those states are visited in
the near future.

Just as one can condition a value function on actions to obtain a state-action
value function (Eq. 3), we can condition the SR on actions as well555Note that
we overload $M^{\pi}$ to also accept actions to reduce the amount of new
notation. In general,
$M^{\pi}\left(s,\tilde{s}\right)=\sum_{a}\pi(a|s)M^{\pi}\left(s,a,\tilde{s}\right)$.:

$\displaystyle M^{\pi}\left(s,a,\tilde{s}\right)$
$\displaystyle=\mathbb{E}\left[\sum_{t=0}^{H}\gamma^{t}\mathbb{I}\left[s_{t+1}=\tilde{s}\right]\bigg{|}s_{0}=s,A_{0}=a\right]$
(20) $\displaystyle=\mathbb{E}\bigg{[}\mathbb{I}[s^{\prime}=\tilde{s}]+\gamma
M^{\pi}(s^{\prime},a^{\prime},\tilde{s})\bigg{]}.$ (21)

Given an action-conditioned SR, the action value function for a particular
reward function can be computed as a linear function of the SR with

$Q^{\pi}(s,a)=\sum_{\tilde{s}}M^{\pi}\left(s,a,\tilde{s}\right)R(\tilde{s}).$
(22)

Having established some mathematical properties of the SR, we can now explain
why it is useful. First, the SR, unlike model-based algorithms, obviates the
need to simulate roll-outs or iterate over dynamic programming updates,
because it has already compiled transition information into a convenient form:
state values can be computed by simply taking the inner product between the SR
and the immediate reward vector. Thus, SR-based value computation enjoys
efficiency comparable to model-free algorithms.

Second, the SR can, like model-based algorithms, adapt quickly to certain
kinds of environmental changes. In particular, local changes to an
environment’s reward structure induce local changes in the reward function,
which immediately propagate to the value estimates when combined with the
SR.666The attentive reader will note that, at least initially, these are not
exactly the correct value estimates, because the SR is policy-dependent, and
the policy itself requires updating, which may not happen instantaneously
(depending on how the agent is optimizing its policy). Nonetheless, these
value estimates will typically be an improvement—a good first guess. As we
will see, human learning exhibits similar behavior. Thus, SR-based value
computation enjoys flexibility comparable to model-based algorithms, at least
for changes to the reward structure. Changes to the transition structure, on
the other hand, require more substantial non-local changes to the SR due to
the fact that an internal model of the detailed transition structure is not
available.

Our discussion has already indicated several limitations of the SR. First, the
policy-dependence of its predictions limits its generalization ability.
Second, the SR assumes a finite, discrete state space. Third, it does not
generalize to new environment dynamics. When the transition structure changes,
the Eq. 14 no longer holds. We will discuss how the first and second
challenges can be addressed in §3 and some attempts to address the third
challenge in §4.2.3.

###  2.4 Successor models: a probabilistic perspective on the SR

Figure 3:  The Successor Model. A cartoon schematic of a robot leg that can
hop forward. Left: with a single-step model, we can only compute likelihoods
for states at the next time-step. Right: with successor models, we can compute
likelihoods for states over some horizon into the future. Adapted with
permission from (Janner et al.,, 2020).

As we’ve discussed, the SR buys efficiency by caching transition structure,
while maintaining some model-based flexibility. One thing that is lost,
however, is the ability to simulate trajectories through the state space. In
this section, we introduce a generalization of the SR—the _successor model_
(SM; Janner et al.,, 2020; Eysenbach et al.,, 2020)—that defines an explicit
model over temporally abstract trajectories (Figure 3). Temporal abstraction
here means a conditional distribution over future states within some time
horizon, rather than only the next time-step captured by the transition model.

The SM uses a $k$-step conditional distribution (i.e., the distribution over
state occupancy after $k$ steps when starting in a particular state) as its
cumulant:

$\boldsymbol{\mu}^{\pi}\left(\tilde{s}|s\right)=(1-\gamma)\sum_{k=1}^{\infty}\gamma^{t}\mathbb{P}\left(s_{k}=\tilde{s}|s_{0}=s,\pi,T\right).$
(23)

where $\gamma$ determines the horizon of the prediction and $(1-\gamma)$
ensures that $\boldsymbol{\mu}^{\pi}$ integrates to $1$.The SM is essentially
a normalized SR, since

$\boldsymbol{\mu}^{\pi}\left(\tilde{s}|s\right)=(1-\gamma)M^{\pi}\left(s,\tilde{s}\right).$
(24)

This relationship becomes apparent when we note that the expectation of an
indicator function is the likelihood of the event, i.e.
$\mathbb{E}[\mathbb{I}[X=x]]=\mathbb{P}(X=x)$.

Since the SM integrates to $1$, a key difference to the SR is that it defines
a valid probability distribution. As we will discuss in §3.2, this allows us
to estimate it with density estimation techniques such as generative
adversarial learning Janner et al., (2020), variational inference (Thakoor et
al.,, 2022), and contrastive learning (Eysenbach et al.,, 2020; Zheng et al.,,
2023).

Intuitively, the SM describes a probabilistic discounted occupancy measure.
SMs are interesting because they are a different kind of environment model.
Rather than defining transition probabilities over next states, they describe
the probability of reaching $\tilde{s}$ within a horizon determined by
$\gamma$ when following policy $\pi$. While we don’t know exactly when
$\tilde{s}$ will be reached, we can answer queries about whether it will be
reached within some relatively long time horizon with less computation
compared to rolling out the base transition model. Additionally, depending on
how the SM is learned, we can also sample from it. This can be useful for
policy evaluation (Thakoor et al.,, 2022) and model-based control (Janner et
al.,, 2020). We will return to this topic when we consider applications.

Like the original SR, the SM obeys a Bellman-like recursion:

$\boldsymbol{\mu}^{\pi}\left(\tilde{s}|s\right)=\mathbb{E}\left[(1-\gamma)T(\tilde{s}|s,a)+\gamma\boldsymbol{\mu}^{\pi}\left(\tilde{s}|s^{\prime}\right)\right],$
(25)

where the next-state probability $\tilde{s}$ in the first term resembles one-
step reward in Eq. 7 and the second term resembles the expected value function
at the next time-step. As with Eq. 19, we can use the SM to perform policy
evaluation by computing:

$\displaystyle V^{\pi}(s)$
$\displaystyle=\frac{1}{1-\gamma}\mathbb{E}_{\tilde{s}\sim\boldsymbol{\mu}^{\pi}(\cdot|s)}\left[R(\tilde{s})\right].$
(26)

Additionally, we can introduce an action-conditioned variant of the SM

$\displaystyle\boldsymbol{\mu}^{\pi}\left(\tilde{s}|s,a\right)$
$\displaystyle=(1-\gamma)\sum_{t=0}^{\infty}\gamma^{t}\mathbb{P}\left(s_{t+1}=\tilde{s}|s_{0}=s,a_{0}=a,\pi,T\right)$
$\displaystyle=(1-\gamma)T(\tilde{s}|s,a)+\gamma\mathbb{E}\left[\boldsymbol{\mu}^{\pi}\left(\tilde{s}|s,a\right)\right].$
(27)

We can leverage this to compute an action value:

$\displaystyle Q^{\pi}(s,a)$
$\displaystyle=\frac{1}{1-\gamma}\mathbb{E}_{\tilde{s}\sim\boldsymbol{\mu}^{\pi}(\cdot|s,a)}\left[R(\tilde{s})\right].$
(28)

###  2.5 Successor features: a feature-based generalization of the SR

Figure 4: A schematic of Successor features (SFs; §2.5) and Generalized policy
improvement (GPI; §2.5.1). Note that we use the shorthand
$\boldsymbol{\phi}_{t}=\boldsymbol{\phi}(s_{t})$ to represents “state-
features” that describe what is visible to the agent at time $t$. $\pi$
corresponds to policies that the agent knows how to perform. (a) Examples of
SFs (Eq. 30) for the “open drawer” and “open fridge” policies. In this
hypothetical scenario, the state-features that agent holds describe whether an
apple, milk, fork, or knife are present. Beginning for the first time-step,
the SFs for these behaviors encode predictions for which of these features
will be present when the policies are executed—predicted to be present for
apple and milk for the open fridge policy, and for the fork and knife when the
open drawer policy is executed. (b) The agent can re-use these known behaviors
with GPI (Eq. 2.5.1). When given a new task, say “get milk”, it is able to
leverage the SFs for the policies to knows to decide which behavior will
enable it to get milk. In this example, the policy for opening the fridge will
also lead to milk. Concretely, The agent selects actions with GPI by computing
Q-values for each known behavior as the dot-product between the current task
and each known SF. The highest Q-value is then used to select actions. If the
agent wants to execute the option keyboard (OK; §2.5.2)), they can so by
having $\mathbf{w}$ be adaptively set based on the current state. For example,
at some states the agent may want to pursue getting milk, while in others they
may want to pursue getting a fork. Adapted from (Carvalho et al., 2023b, )
with permission.

When the state space is large or unknown, the SR can be challenging to compute
or learn because it requires maintaining occupancy expectations over all
states. This can also be challenging when dealing with learned state
representations, as is common in practical learning algorithms (see §3.1). In
these settings, rather than maintain occupancy measure for states, we can
maintain occupancy measure over cumulants that are features shared across
states, $\boldsymbol{\phi}(s)$. This generalization of the SR is known as
successor features (Barreto et al.,, 2017). The power of this representation
is particularly apparent when reward can be decomposed into a dot-product of
these successor features and a vector $\mathbf{w}$ describing feature
“preferences” for the current task:

$R(s,\mathbf{w})=\boldsymbol{\phi}\left(s\right)^{\top}\mathbf{w}.$ (29)

Successor features are then predictions of how much of features
$\boldsymbol{\phi}$ the agent can expect to obtain when following a policy
$\pi$:

$\boldsymbol{\psi}^{\pi}(s)=\mathbb{E}\left[\sum_{t=0}^{\infty}\gamma^{t}\boldsymbol{\phi}\left(s_{t+1}\right)\mid
s_{0}=s\right]$ (30)

and obey the following Bellman equation:

$\boldsymbol{\psi}^{\pi}(s)=\mathbb{E}\left[\boldsymbol{\phi}(s^{\prime})+\gamma\boldsymbol{\psi}(s^{\prime})\right].$
(31)

Under the assumption of Eq. 29, a task-dependent value
$V^{\pi}(s,a,\mathbf{w})$ is equivalent to:

$\displaystyle Q^{\pi}(s,a,\mathbf{w})$
$\displaystyle=\mathbb{E}\left[r(s_{1})+\gamma r(s_{2})+\ldots\mid
s_{0}=s,a_{0}=a\right]$
$\displaystyle=\mathbb{E}\left[\boldsymbol{\phi}(s_{1})^{\top}\mathbf{w}+\gamma\boldsymbol{\phi}(s_{2})^{\top}\mathbf{w}+\ldots\mid
s_{0}=s,a_{0}=a\right]$
$\displaystyle=\mathbb{E}\left[\sum_{t=0}^{\infty}\gamma^{t}\boldsymbol{\phi}(s_{t+1})\mid
s_{0}=s,a_{0}=a\right]^{\top}\mathbf{w}$
$\displaystyle=\boldsymbol{\psi}^{\pi}(s)^{\top}\mathbf{w}.$ (32)

As with the SR and SM, we can introduce an action-conditioned variant of SFs

$\displaystyle\boldsymbol{\psi}^{\pi}(s,a)$
$\displaystyle=\mathbb{E}\left[\sum_{t=0}^{\infty}\gamma^{t}\boldsymbol{\phi}\left(s_{t+1}\right)\mid
s_{0}=s,a_{0}=a\right]$
$\displaystyle=\mathbb{E}\left[\boldsymbol{\phi}(s^{\prime})+\gamma\boldsymbol{\psi}(s^{\prime},a^{\prime})\right].$
(33)

This provides an avenue to re-use these behavioral predictions for new tasks.
Once $\boldsymbol{\psi}^{\pi}(s,a)$ has been learned, the behavior $\pi$ can
be reused to pursue a novel reward function $R(s,\mathbf{w}_{\tt
new})=\boldsymbol{\phi}(s)^{\top}\mathbf{w}_{\tt new}$ by simply changing the
corresponding task encoding $\mathbf{w}_{\tt new}$:

$Q^{\pi}(s,a,\mathbf{w}_{\tt
new})=\boldsymbol{\psi}^{\pi}(s,a)^{\top}\mathbf{w}_{\tt new}$ (34)

As we will see in the next section and is exploited heavily in AI applications
(§4), this becomes even more powerful when combined with generalized policy
improvement.

####  2.5.1 Generalized policy improvement: adaptively combining behaviors

One limitation of Eq. 19 and Eq. 26 is that they only enable us to re-compute
the value of states for a new reward function under a known policy. However,
we may want to synthesize a new policy from the other policies we have learned
so far. We can accomplish this with SFs by combining them with Generalized
Policy improvement (GPI; (Barreto et al.,, 2017)), shown in Figure 4. Assume
we have learned (potentially optimal) policies
$\left\\{\pi_{i}\right\\}_{i=1}^{n_{\text{train}}}$ and their corresponding
SFs
$\left\\{\boldsymbol{\psi}^{\pi_{i}}(s,a)\right\\}_{i=1}^{n_{\text{train}}}$
for $n_{\tt train}$ training tasks
$\left\\{\mathbf{w}_{i}\right\\}_{i=1}^{n_{\text{train}}}$. When presented
with a new task $\mathbf{w}_{\tt new}$, we can obtain a new policy with GPI in
two steps: (1) compute Q-values using the training task SFs; (2) select
actions using the highest Q-value. This operation is summarized as follows:

$\displaystyle\pi\left(a|s,\mathbf{w}_{\tt new}\right)$
$\displaystyle\propto\max_{i\in\left\\{1,\ldots,n_{\text{train}}\right\\}}\left\\{\boldsymbol{\psi}^{\pi_{i}}(s,a)^{\top}\mathbf{w}_{\tt
new}\right\\}$
$\displaystyle=\max_{i\in\left\\{1,\ldots,n_{\text{train}}\right\\}}\left\\{Q^{\pi_{i}}(s,a,\mathbf{w}_{\tt
new})\right\\}$ (35)

If $\mathbf{w}_{\tt new}$ is in the span of the training tasks (i.e.,
$\mathbf{w}_{\tt new}=\sum_{i}\alpha_{i}\mathbf{w}_{i}$, where
$\alpha_{i}\in\mathbb{R}$), the GPI theorem states that $\pi$ will perform at
least as well as any of the training policies—i.e.,
$Q^{\pi}(s,a,\mathbf{w}_{\tt new})\geq\max_{i}Q^{\pi_{i}}(s,a,\mathbf{w}_{\tt
new})\,\forall(s,a)\in(\mathcal{S}\times\mathcal{A})$.

####  2.5.2 Option Keyboard: chaining together behaviors

One benefit of Eq. 2.5.1 is that it enables transfer to linear combinations of
training task encodings. However, it has two limitations. First, the task
feature “preferences” $\mathbf{w}_{\tt new}$ are fixed across time. When
$\mathbf{w}_{\tt new}$ is a complex task (e.g., avoiding an object at some
time-points but going towards it at others), we may want preferences that are
state-dependent. The “Option Keyboard” (OK; (Barreto et al.,, 2019, 2020))
addresses this by having a state-dependent preference vector $\mathbf{w}_{s}$
that is produced by a policy $g\left(\cdot\right)$ which takes in the current
state $s$ and task $\mathbf{w}_{\tt new}$, i.e.
$\mathbf{w}_{s}=g\left(s,\mathbf{w}_{\tt new}\right)$. Given $\mathbf{w}_{s}$,
we can then select actions with GPI as follows:

$\displaystyle\pi_{\mathbf{w}_{s}}\left(a|s,\mathbf{w}_{\tt new}\right)$
$\displaystyle\propto\max_{i\in 1,\ldots,n_{\tt
train}}\left\\{\boldsymbol{\psi}^{\pi_{i}}(s,a)^{\top}\mathbf{w}_{s}\right\\}$
(36)

Note that this policy is adaptive—which known policy $\pi_{i}$ is chosen at a
time-step is dependent on which one maximizes the successor features at that
time-step. This is why it’s called the option “keyboard”: $\mathbf{w}_{s}$
will induce a set of behaviors to be active for a period of time, somewhat
like chord on a piano.

###  2.6 Summary

The predictive representations introduced above can be concisely organized in
terms of particular cumulants, as summarized in Table 1. These cumulants have
different strengths and weaknesses. Value functions (reward cumulants)
directly represent the key quantity for RL tasks, but they suffer from poor
flexibility. The SR (state occupancy cumulant) and its variations (feature
occupancy and state probability cumulants) can be used to compute values, but
also retain useful information for generalization to new tasks (e.g., using
generalized policy improvement).

Predictive representation | Cumulant |  $\delta$: TD update when $a^{\prime}\sim\pi$ and $s^{\prime}\sim T$
---|---|---
$Q^{\pi}(s,a)$ (§2.2) | $r(s)$ | $r(s^{\prime})+\gamma Q^{\pi}(s^{\prime},a^{\prime})-Q^{\pi}(s,a)$
$M^{\pi}(s,\tilde{s})$ (SR; §2.3) | $\mathbb{I}(s=\tilde{s})$ | $\mathbb{I}(s^{\prime}=\tilde{s})+\gamma M^{\pi}(s^{\prime},\tilde{s})-M^{\pi}(s,\tilde{s})$
$\boldsymbol{\psi}^{\pi}(s,a)$ (SFs; §2.5) | $\boldsymbol{\phi}(s)$ | $\boldsymbol{\phi}(s^{\prime})+\gamma\boldsymbol{\psi}^{\pi}(s^{\prime},a^{\prime})-\boldsymbol{\psi}^{\pi}(s,a)$
$\boldsymbol{\mu}^{\pi}(\tilde{s}|s,a)$ (SM; §2.4) | $\mathbb{P}\left(s_{t+1}=\tilde{s}|s_{0},a_{0}\right)$ | $(1-\gamma)T(s^{\prime}|s,a)+\gamma\boldsymbol{\mu}^{\pi}(\tilde{s}|s^{\prime},a^{\prime})-\boldsymbol{\mu}^{\pi}(\tilde{s}|s,a)$
Table 1: Summary of predictive representations that we focus on. For each, we
also describe the “cumulant” that this predictive representation forms
predictions over, along with a corresponding on-policy Bellman update one can
use to learn the representation for a policy $\pi$. $Q^{\pi}(s,a)$ is the
action-value function which forms predictions about future reward.
$M^{\pi}(s,\tilde{s})$ is the successor representation (SR) which forms
predictions about how much a state $\tilde{s}$ will be visited.
$\boldsymbol{\psi}^{\pi}(s,a)$ are successor features (SFs) which form
predictions about how much state-features $\boldsymbol{\phi}(s)$ will be
experienced. $\boldsymbol{\mu}^{\pi}(\tilde{s}|s,a)$ is the successor model
(SM) which predicts the likelihood of experiencing $\tilde{s}$ in the future.

##  3 Practical learning algorithms
and associated challenges

Of the predictive representations discussed in §2, only successor features and
successor models have been successfully scaled to environments with large,
high-dimensional state spaces (including continuous state spaces). Thus, these
will be our focus of discussion. We will first discuss learning SFs in §3.1
and then learning successor models in §3.2.

###  3.1 Learning successor features

In this section we discuss practical considerations for learning successor
features, including learning a function $\boldsymbol{\phi}_{\theta}$ that
produces cumulants (§3.1.1) and estimating SFs $\boldsymbol{\psi}_{\theta}$
(§3.1.2).

####  3.1.1 Discovering cumulants

One central challenge to learning SFs is that they require cumulants that are
useful for adaptive agent behavior; these are not always easy to define _a
priori_. In most work, these cumulants have been hand-designed, but
potentially better ones can be learned from experience. Some general methods
for discovering cumulants include leveraging meta-gradients (Veeriah et al.,,
2019), discovering features that enable reconstruction (Kulkarni et al.,,
2016; Machado et al.,, 2017), and maximizing the mutual information between
task encodings and the cumulants that an agent experiences when pursuing that
task (Hansen et al.,, 2019). However, these methods don’t necessarily
guarantee that the learned cumulants respect a linear relationship with reward
(Eq. 29). To satisfy this, methods typically enforce this by minimizing the
L2-norm of their difference (Barreto et al.,, 2017):

$\mathcal{L}_{r}=||r-\boldsymbol{\phi}_{\theta}(s)^{\top}\mathbf{w}||^{2}_{2}.$
(37)

When learning cumulants that support transfer with GPI, one strategy that can
bolster Eq. 37 is to learn an $n$-dimensional $\boldsymbol{\phi}$ vector for
$n$ tasks such that each dimension predicts one task reward (Barreto et al.,,
2018). Another strategy is to enforce that cumulants describe independent
features (Alver and Precup,, 2021)—e.g., by leveraging a modular architecture
with separate parameters for every cumulant dimension (Carvalho et al., 2023a,
) or by enforcing sparsity in the cumulant dimensions (Filos et al.,, 2021).

####  3.1.2 Estimating successor features

##### Learning an estimator that can generalize across tasks

The first challenge for learning SFs is that they are defined for a particular
policy $\pi$. We can mitigate this by learning Universal Successor Feature
Approximators (USFAs; Borsa et al.,, 2018), which approximate SFs for a policy
$\pi_{\mathbf{w}}$ that is maximizing reward for task $\mathbf{w}$ with a
function that takes the corresponding preference vector as input
$\boldsymbol{\psi}_{\theta}(s,a,\mathbf{w})$:

$\boldsymbol{\psi}^{\pi_{\mathbf{w}}}(s,a)\approx\boldsymbol{\psi}_{\theta}(s,a,\mathbf{w}).$
(38)

This has several benefits. First, one can share the estimator parameters
across tasks, which can improve learning. Second, $\boldsymbol{\psi}_{\theta}$
now accepts policy encodings $\mathbf{z}_{\pi}\in\mathcal{Z}_{\pi}$, which
represent a particular policy in the space of preference vectors. In other
words, $\mathbf{z}_{\pi}$ is the preference vector that induces policy $\pi$.
This is useful because it allows SFs to generalize across different policies.
Leveraging a USFA, the GPI operation in Eq. 2.5.1 can be adapted to perform a
max operation over $\mathcal{Z}_{\pi}$:

$\displaystyle\pi\left(s;\mathbf{w}_{\tt new}\right)$
$\displaystyle=\underset{a\in\mathcal{A}}{\arg\max}\max_{\mathbf{z}_{\pi}\in\mathcal{Z}_{\pi}}\left\\{\boldsymbol{\psi}_{\theta}(s,a,\mathbf{z}_{\pi})^{\top}\mathbf{w}_{\tt
new}\right\\}$ (39)

Several algorithms exploit this property (see §4.1.1, §4.2.2, §4.5). If the
policy encoding space is equivalent to the $n$ training tasks,
$\mathcal{Z}_{\pi}=\\{\mathbf{w}_{1},\ldots,\mathbf{w}_{n}\\}$, this recovers
the original GPI operation (Eq. 2.5.1).

##### Learning successor features while learning a policy

Often SFs need to be learned simultaneously with policy optimization. This
requires the SFs to be updated along with the policy. One strategy is to
simultaneously learn an action-value function
$Q_{\theta}(s,a,\mathbf{w})=\boldsymbol{\psi}_{\theta}(s,a,\mathbf{w})^{\top}\mathbf{w}$
that is used to select actions. One can accomplish this with Q-learning over
values defined by the SFs and task encoding. Q-values are updated towards a
target $\mathbf{y}_{Q}$ defined as the sum of the reward
$R(s^{\prime},\mathbf{w})$ and the best next Q-value. We define the learning
objective $\mathcal{L}_{Q}$ as follows:

$\displaystyle\mathbf{y}_{Q}$
$\displaystyle=R(s^{\prime},\mathbf{w})+\gamma\boldsymbol{\psi}_{\theta}(s^{\prime},a^{*},\mathbf{w})^{\top}\mathbf{w}$
$\displaystyle\mathcal{L}_{Q}$
$\displaystyle=||\boldsymbol{\psi}_{\theta}(s,a,\mathbf{w})^{\top}\mathbf{w}-\mathbf{y}_{Q}||$
(40)

where
$a^{*}=\operatornamewithlimits{argmax}_{a^{\prime}}\boldsymbol{\psi}_{\theta}(s^{\prime},a^{\prime},\mathbf{w})^{\top}\mathbf{w}$
is the action which maximizes features determined by $\mathbf{w}$ at the next
time-step. To ensure that these Q-values follow the structure imposed by SFs
(Eq. 2.5), we additionally update SFs at a state with a target
$\mathbf{y}_{\boldsymbol{\psi}}$ defined as the sum of the cumulant
$\boldsymbol{\phi}(s^{\prime})$ and the SFs associated with best Q-value at
the next state:

$\displaystyle\mathbf{y}_{\boldsymbol{\psi}}$
$\displaystyle=\boldsymbol{\phi}(s^{\prime})+\gamma\boldsymbol{\psi}_{\theta}(s^{\prime},a^{*},\mathbf{w})$
$\displaystyle\mathcal{L}_{\boldsymbol{\psi}}$
$\displaystyle=||\boldsymbol{\psi}_{\theta}(s,a,\mathbf{w})-\mathbf{y}_{\boldsymbol{\psi}}||.$
(41)

In an online setting, it is important to learn SFs with data collected by a
policy which chooses actions with high Q-values. This is especially important
if the true value is lower than the estimated Q-value. Because Q-learning
leverages the maximum Q-value when doing backups, it has a bias for over-
estimating value. This can destabilize learning, particularly in regions of
the state space that have been less explored (Ostrovski et al.,, 2021).

Another strategy to stabilize SF learning is to learn individual SF dimensions
with separate modules (Carvalho et al., 2023a, ; Carvalho et al., 2023b, ).
Beyond stabilizing learning, this modularity also enables approximating SFs
that generalize better to novel environment configurations (i.e., which are
more robust to novel environment dynamics).

##### Estimating successor features with changing cumulants

In some cases, the cumulant itself will change over time (e.g., when the
environment is non-stationary). This is challenging for SF learning because
the prediction target is non-stationary (Barreto et al.,, 2018). This is an
issue even when the environment is stationary but the policy is changing over
time: different policies induce different trajectories and different state-
features induce different descriptions of those trajectories.

One technique that has been proposed to facilitate modelling a non-stationary
cumulant trajectories is to learn an SF as a probability mass function (pmf)
$p(\boldsymbol{\psi}^{(k)}|s,a,\mathbf{w})$ defined over some set of possible
values $\mathbb{B}=\left\\{b_{1},\ldots,b_{M}\right\\}$. Specifically, we can
estimate an n-dimensional SF
$\boldsymbol{\psi}^{\pi_{\mathbf{w}}}(s,a)\in\mathbb{R}^{n}$ with
$\boldsymbol{\psi}_{\theta}(s,a,\mathbf{w})$ and represent the $k$-th SF
dimension as $\boldsymbol{\psi}_{\theta}^{k}(s,a,\mathbf{w})=$
$\sum_{m=1}^{M}p(\boldsymbol{\psi}^{(k)}=b_{m}|s,a,\mathbf{w})b_{m}$. We can
then learn SFs with a negative log-likelihood loss where we construct
categorical target labels $\mathbf{y}_{\boldsymbol{\psi}^{(k)}}$ from the
return associated with the optimal Q-value (Carvalho et al., 2023b, ):

$\displaystyle\mathbf{y}_{\boldsymbol{\psi}^{(k)}}$
$\displaystyle=\boldsymbol{\phi}_{\theta}^{k}(s^{\prime})+\gamma\boldsymbol{\psi}^{k}_{\theta}(s^{\prime},a^{*},\mathbf{w})$
(42) $\displaystyle\mathcal{L}_{\boldsymbol{\psi}}$
$\displaystyle=-\sum^{n}_{k=1}f_{\tt
label}(\mathbf{y}_{\boldsymbol{\psi}^{(k)}})^{\top}\log
p(\boldsymbol{\psi}^{(k)}|s,a,\mathbf{w}).$ (43)

Prior work has found that the two-hot representation is a good method for
defining $f_{\tt label}(\cdot)$ (Carvalho et al., 2023b, ; Schrittwieser et
al.,, 2020). In general, estimating predictive representations such as SFs
with distributional losses such as Eq. 42 has been shown to reduce the
variance in learning updates (Imani and White,, 2018). This is particularly
important when cumulants are being learned as this can lead to high variance
in $\mathbf{y}_{\boldsymbol{\psi}^{(k)}}$.

###  3.2 Learning successor models

In this section, we focus on estimating
$\boldsymbol{\mu}^{\pi}(\tilde{s}|s,a)$ with
$\boldsymbol{\mu}_{\theta}(\tilde{s}|s,a)$. In a tabular setting, one can
leverage TD-learning with the Bellman equation in Eq. 25. However, for very
large state spaces (such as with infinite size continuous state spaces), this
is intractable or impractical. Depending on one’s use-case different options
exist for learning this object. First, we discuss the setting where one wants
to learn a successor model they can sample from (§3.2.1). Afterwards, we
discuss the setting where one only wants to evaluate a successor model for
different actions given a target state (§3.2.2).

####  3.2.1 Learning successor models that one can sample from

Learning a successor model that one can sample from can be useful for
evaluating a policy (Eq. 26), evaluating a sequence of policies (Thakoor et
al.,, 2022), and in model-based control (Janner et al.,, 2020). There are two
ways to learn successor models you can sample from: adversarial learning and
density estimation (Janner et al.,, 2020). Adversarial learning has been found
to be unstable, so we focus on density estimation, where the objective is to
find parameters $\theta$ that maximize the log-likelihood of states sampled
from $\boldsymbol{\mu}^{\pi}$:

$\displaystyle\mathcal{L}_{\mu}=\mathbb{E}_{\begin{subarray}{c}(s,a)\sim
p(s,a)\\\ \tilde{s}\sim\boldsymbol{\mu}^{\pi}(\cdot\mid
s,a)\end{subarray}}\left[\log\boldsymbol{\mu}_{\theta}\left(\tilde{s}\mid
s,a\right)\right]$ (44)

We can optimize this objective as follows. When sampling targets, we need to
sample $\tilde{s}$ in proportion to the discount factor $\gamma$—we can
accomplish this by first sampling a time-step from a geometric distribution,
$k\sim\operatorname{Geom}(1-\gamma)$, and then selecting the state at that
time-step,$s_{t+k}$, as the target $\tilde{s}$.

While this is a simple strategy, it has several challenges. For values of
$\gamma$ close to 1, this becomes a challenging learning problem requiring
predictions over very long time horizons. Another challenge is that you are
using $s_{t+k}$ obtained under policy $\pi$. In practice, we may want to
leverage data collected under a different policy. This happens when, for
example, we want to learn from a collection of different datasets, or we are
updating our policy over the course of learning. Learning from such off-policy
data can lead to high bias, or a high variance learning update from off-policy
corrections (Precup,, 2000).

We can circumvent these challenges as follows. First let’s define a Bellman
operator $T^{\pi}$:

$(T^{\pi}\boldsymbol{\mu}^{\pi})(\tilde{s}|s,a)=(1-\gamma)T(s^{\prime}|s,a)+\gamma\sum_{s^{\prime}}T(s^{\prime}|s,a)\sum_{a^{\prime}}\pi(a^{\prime}|s^{\prime})\boldsymbol{\mu}(\tilde{s}|s^{\prime},a^{\prime}).$
(45)

With this we can define a “cross-entropy temporal-difference” loss (Janner et
al.,, 2020):

$\mathcal{L}_{\mu}=\mathbb{E}_{\begin{subarray}{c}(s,a)\sim p(s,a)\\\
\tilde{s}\sim(T^{\pi}\boldsymbol{\mu}^{\pi})(\cdot|s,a)\end{subarray}}\left[\log\boldsymbol{\mu}_{\theta}\left(\tilde{s}\mid
s,a\right)\right]$ (46)

Intuitively, $(T^{\pi}\boldsymbol{\mu}^{\pi})(\cdot|s,a)$ defines a random
variable obtained as follows. First sample $s^{\prime}\sim T(\cdot|s,a)$.
Terminate and emit $s^{\prime}$ with probability $(1-\gamma)$. Otherwise,
sample $a^{\prime}\sim\pi(a^{\prime}|s^{\prime})$ and them sample
$\tilde{s}\sim\boldsymbol{\mu}^{\pi}(\cdot|s^{\prime},a^{\prime})$.

The most recent promising method for learning Eq. 46 has been to leverage a
variational autoencoder (VAE; Thakoor et al.,, 2022). Specifically, we can
define an approximate posterior $q_{\psi}(z|s,a,\tilde{s})$ and then optimize
the following evidence lower-bound:

$\mathcal{L}_{\mu}=\mathbb{E}_{\begin{subarray}{c}(s,a)\sim p(s,a)\\\
\tilde{s}\sim(T^{\pi}\boldsymbol{\mu}^{\pi})(\cdot|s,a)\end{subarray}}\left[\mathbb{E}_{z\sim
q_{\psi}(\cdot|s,a,\tilde{s})}\left[\log\frac{\boldsymbol{\mu}_{\theta}\left(\tilde{s}\mid
s,a,z\right)}{q_{\psi}(z|s,a,\tilde{s})}\right]\right]$ (47)

See Thakoor et al., (2022) for more details. While they were able to scale
their experiments to slightly more complex domains than Janner et al., (2020),
their focus was on composing policies via Geometric Policy Composition
(discussed more in §4.2.3), so it is unclear how well their method does in
more complex domains. The key challenge for this line of work is in sampling
from $\boldsymbol{\mu}^{\pi}(\cdot\mid s,a)$, where $\tilde{s}$ can come from
a variable next time-step after $(s,a)$. In the next section, we discuss
methods which address this challenge.

####  3.2.2 Learning successor models that one can evaluate

Sometimes we may not need to learn a successor model that we can sample from,
only one that we can evaluate. This can be used, for example, to generate and
improve a policy that achieves some target state (Eysenbach et al.,, 2020;
Zheng et al.,, 2023). One strategy is to learn a classifier $p^{\mu}_{\theta}$
that, given $(s,a)$, computes how likely $\tilde{s}$ is compared to some set
of $N$ random (negative) states $\\{s^{-}_{i}\\}^{N}_{i=1}$ the agent has
experienced:

$p^{\mu}_{\theta}(\tilde{s}\mid
s,a,s^{-}_{1:N})=\frac{\exp(f_{\theta}(s,a,\tilde{s}))}{\exp(f_{\theta}(s,a,\tilde{s}))+\sum^{N}_{i=1}\exp(f_{\theta}(s,a,s^{-}_{i}))}$
(48)

One option for learning $f_{\theta}$ is to find a value that maximizes this
classification across random states and actions in the agent’s experience,
$(s,a)\sim p(s,a)$, target states $\tilde{s}$ drawn from the empirical
successor model distribution,
$\tilde{s}\sim{\boldsymbol{\mu}^{\pi}(\tilde{s}\mid s,a)}$, and negatives
drawn from the state-marginal, $s^{-}_{1:N}\sim p(s)$,

$L_{\boldsymbol{\mu}}=\mathbb{E}_{\begin{subarray}{c}(s,a)\sim p(s,a)\\\
\tilde{s}\sim{\boldsymbol{\mu}^{\pi}(\tilde{s}\mid s,a)}\\\ s^{-}_{1:N}\sim
p(s)\end{subarray}}\left[\log p^{\mu}_{\theta}(\tilde{s}\mid
s,a,s^{-}_{1:N})\right].$ (49)

If we can find such an $f_{\theta}$, then the classification it defines
$p^{\mu}_{\theta}$ is approximately equal to the density ratio between
$\boldsymbol{\mu}^{\pi}(\tilde{s}\mid s,a)$ and the state marginal
$p(\tilde{s})$ (Poole et al.,, 2019; Zheng et al.,, 2023):

$\frac{\boldsymbol{\mu}^{\pi}(\tilde{s}\mid
s,a)}{p(\tilde{s})}\approx(N+1)\cdot p^{\mu}_{\theta}(\tilde{s}\mid
s,a,s^{-}_{1:N}).$ (50)

Optimizing Eq. 49 is challenging because it requires sampling from
$\boldsymbol{\mu}^{\pi}$. We can circumvent this by instead learning the
following TD-like objective where we replace sampling from
$\boldsymbol{\mu}^{\pi}$ with sampling from the state-marginal and reuse
$p_{\theta}^{\mu}$ as an importance weight:

$L_{\boldsymbol{\mu}}=\mathbb{E}_{\begin{subarray}{c}(s,a)\sim p(s,a)\\\
s^{\prime}\sim T(\cdot|s,a)\\\ a^{\prime}\sim\pi(a|s^{\prime})\\\
s^{-}_{1:N}\sim p(s)\end{subarray}}\left[(1-\gamma)\log
p^{\mu}_{\theta}(s^{\prime}|s,a,s^{-}_{1:N})+\gamma(N+1)p_{\theta}^{\mu}(\tilde{s}|s^{\prime},a^{\prime},s^{-}_{1:N})\log
p_{\theta}^{\mu}(\tilde{s}|s,a,s^{-}_{1:N})\right]$ (51)

Zheng et al., (2023) show that (under some assumptions) optimizing Eq. 51
leads to the following Bellman-like update

$p_{\theta}^{\mu}(\tilde{s}|s,a,s^{-}_{1:N})\leftarrow(1-\gamma)T(s^{\prime}=\tilde{s}|s,a)+\gamma\mathbb{E}_{\begin{subarray}{c}s^{\prime}\sim
T(\cdot|s,a)\\\
a^{\prime}\sim\pi(\cdot|s^{\prime})\end{subarray}}\left[p_{\theta}^{\mu}(\tilde{s}|s^{\prime},a^{\prime},s^{-}_{1:N})\right]$
(52)

which resembles the original successor model Bellman equation (Eq. 25).
However, one key difference to Eq. 25 is that we parameterize
$p_{\theta}^{\mu}(\tilde{s}|s,a,s^{-}_{1:N})$ with $N+1$ random samples (e.g.,
from a replay mechanism), which is a form of contrastive learning. This
provides an interesting, amortized algorithm for learning successor models.
With the SR, one performs the TD update for all states (Eq. 18); here, one
performs this update using a random sample of $N+1$ states.

We can define $f_{\theta}$ as the dot product between a predictive
representation $\varphi_{\theta}(\cdot,\cdot)$ and label representation
$\phi_{\theta}(\cdot)$,
$f_{\theta}(s,a,\tilde{s})=\varphi_{\theta}(s,a)^{\top}\phi_{\theta}(\tilde{s})$.
$\phi_{\theta}(s)$ can then be thought of as state-features analogous to SFs
(§2.5). $\varphi_{\theta}(s,a)$ is then a “prediction” of these future
features similar to SFs, $\boldsymbol{\psi}_{\theta}(s,a)$, with labels coming
from future states; however, it doesn’t necessarily have the same semantics as
a discounted sum (i.e., Eq. 2.5). We use similar notation because of their
conceptual similarity. We can then understand Eq. 51 as doing the following.
The first term in this objective pushes the prediction $\varphi_{\theta}(s,a)$
towards the features at the next-timestep $\phi_{\theta}(s^{\prime})$, and the
second term pushes $\varphi_{\theta}(s,a)$ towards the features at arbitrary
state-features $\phi_{\theta}(\tilde{s})$. Both terms repel
$\varphi_{\theta}(s,a)$ from arbitrary “negative” state-features
$\phi_{\theta}(\tilde{s}^{-}_{i})$. This provides an interesting contrastive-
learning based mechanism for how brains may learn $\boldsymbol{\mu}^{\pi}$ in
a non-tabular setting.

##  4 Artificial intelligence applications

In this section, we discuss how the SR and its generalizations have enabled
advances in artificial agents that learn and transfer behaviors from
experience.

###  4.1 Exploration

####  4.1.1 Pure exploration

##### Learning to explore and act in the environment before exposure to reward

This is perhaps closest to the original inspiration for the SR (Dayan,, 1993).
Here, the agent can explore the environment without exposure to reward or
punishment for some period of time, and tries to learn a policy that can
transfer to an unknown task $\mathbf{w}_{\tt new}$.

One useful property of SFs is that they encode predictions about what features
one can expect when following a policy. Before reward is provided, this can be
used to reach different parts of the state space with different policies. One
strategy is to associate different parts of the state space with different
parts of a high-dimensional task embedding space (Hansen et al.,, 2019). At
the beginning of each episode, an agent samples a “goal” encoding $\mathbf{w}$
from a high-entropy task distribution $p(\mathbf{w})$. During the episode, the
agent selects actions that maximize the features described by $\mathbf{w}$
(e.g., with Eq. 2.5). As it does this, it learns to predict $\mathbf{w}$ from
the states it encounters. If we parameterize this $p(\mathbf{w})$ as a von
Mises distribution, then we can learn this prediction by simply maximizing the
dot-product between the state-features $\boldsymbol{\phi}(s)$ and goal
encoding $\mathbf{w}$:

$\displaystyle\mathcal{L}_{\boldsymbol{\phi}}=\log
p_{\theta}(\mathbf{w}|s)=\boldsymbol{\phi}_{\theta}(s)^{\top}\mathbf{w}.$ (53)

This is equivalent to maximizing the mutual information between
$\boldsymbol{\phi}_{\theta}(s)$ and $\mathbf{w}$. As the agent learns
cumulants, it learns SFs as usual (e.g., with Eq. 41). Thus, one can
essentially use a standard SF learning algorithm and simply replace the
cumulant discovery loss (e.g., Eq. 37) with Eq. 53. Once the agent is exposed
to task reward $\mathbf{w}_{\tt new}$, it can then freeze
$\boldsymbol{\phi}_{\theta}$ and solve for $\mathbf{w}_{\tt new}$ (e.g., with
Eq. 37). The agent can then use GPI to find the best policy for achieving
$\mathbf{w}_{\tt new}$ by searching over a Gaussian ball defined around
$\mathbf{w}_{\tt new}$. This is equivalent to setting
$Z=\\{z_{i}\sim\mathcal{N}(\mathbf{w}_{\tt new},\sigma)\\}^{n}_{i=1}$ for Eq.
39, where $\sigma$ defines the standard deviation of the Gaussian ball.

Hansen et al., (2019) leveraged this strategy to develop agents which could
explore Atari games without any reward for 250 million time-steps and then
have 100 thousand time-steps to earn reward. They showed that this strategy
was able to achieve superhuman performance across most Atari games, despite
not observing any rewards for most of its experience. Liu and Abbeel, (2021)
improved on this algorithm by adding an intrinsic reward function that favors
exploring parts of the state space that are “surprising” (i.e., which induce
high entropy) given a memory of the agent’s experience. This dramatically
improved sample efficiency for many Atari games.

##### Exploring the environment by building a map.

Agents need to explore large state space systematically. One strategy is to
build a map of the environment defined over landmarks in the environment. With
such a map, an agent can systematically explore by planning paths towards the
frontier of its knowledge (Ramesh et al.,, 2019; Hoang et al.,, 2021).
However, numerous questions arise in this process. How do we define good
landmarks? How do we define policies that traverse between landmarks? How does
an agent identify that it has made progress between landmarks after it has set
a plan, or course-correct if it finds that it accidentally deviated? Hoang et
al., (2021) developed an elegant solution to all of these problems with the
successor feature similarity (SFS) metric. This similarity metric defines
“closeness” by how likely two states are to visit the same parts of the
environment. Intuitively, if starting from two states $s_{1}$ and $s_{2}$, an
agent visits very different parts of the environment, this metric measures
these states as being far apart; if instead the two states lead to the same
parts of the environment, the metric measures them as being close together.
Concretely:

$\displaystyle S_{\boldsymbol{\psi}}(s_{1},a,s_{2})$
$\displaystyle=\boldsymbol{\psi}^{\bar{\pi}}(s_{1},a)^{\top}\mathbb{E}_{A\sim\bar{\pi}}[\boldsymbol{\psi}^{\bar{\pi}}(s_{2},A)]$
(54) $\displaystyle S_{\boldsymbol{\psi}}(s_{1},s_{2})$
$\displaystyle=\mathbb{E}_{A\sim\bar{\pi}}[\boldsymbol{\psi}^{\bar{\pi}}(s_{1},a)^{\top}]\mathbb{E}_{A\sim\bar{\pi}}[\boldsymbol{\psi}^{\bar{\pi}}(s_{2},A)],$
(55)

where $\bar{\pi}$ is a uniform policy.777Note that to avoid excessive
notation, we’ve overloaded the definition of $S_{\psi}$.

Through only learning of successor features over pretrained cumulants
$\boldsymbol{\phi}$, Hoang et al., (2021) are able to address all the needs
above by exploiting SFS. Let $\mathcal{L}\subset\mathcal{S}$ be the set of
landmarks discovered so far. The algorithm works as follows. Each landmark
$L\in\mathcal{L}$ has associated with it a count $N(L)$ for how often its been
visited. A “frontier” landmark $L_{F}$ is sampled in proportion to the inverse
of this count. The agent makes a shortest-path plan of subgoals
$(s^{\ast}_{1},\ldots,s^{\ast}_{n})$ towards this landmark where
$s_{n}=L_{F}$. In order to navigate to the next subgoal $s^{\ast}_{i}$, it
defines a policy with the action of the current successor features that are
most aligned with the goal’s successor features,
$a^{\ast}(s,s^{\ast}_{i})=\operatornamewithlimits{argmax}_{a}S_{\boldsymbol{\psi}}(s,a,s^{\ast}_{i})$.
As it traverses towards the landmark, it localizes itself by comparing the
current state $s$ to known landmarks $f_{\tt
loc}(s,\mathcal{L})=\operatornamewithlimits{argmax}_{L\in\mathcal{L}}S_{\boldsymbol{\psi}}(s,L)$.
When $s^{\ast}_{i}=f_{\tt loc}(s,\mathcal{L})$, it has reached the next
landmark and it moves on to the next subgoal. Once a frontier landmark is
found, the agent explores with random actions. The states it experiences are
added to its replay buffer and used in $2$ ways. First, they update the
agent’s current SR. Second, if a state is sufficiently different from known
landmarks, it is added. In summary, the $3$ key functions that one can compute
without additional learning are:

$\displaystyle\pi(s,s^{\ast}_{i})$
$\displaystyle=\operatornamewithlimits{argmax}_{a}S_{\boldsymbol{\psi}}(s,a,s^{\ast}_{i})$
goal-conditioned policy $\displaystyle f_{\tt loc}(s,\mathcal{L})$
$\displaystyle=\operatornamewithlimits{argmax}_{L\in\mathcal{L}}S_{\boldsymbol{\psi}}(s,L)$
localization function $\displaystyle f_{\tt add}(s,\mathcal{L})$
$\displaystyle=\mathcal{L}\leftarrow\mathcal{L}\cup
s\quad\text{if}\quad(S_{\boldsymbol{\psi}}(s,L)<\epsilon_{\tt add})\,\forall
L\in\mathcal{L}$ landmark addition function

The process then repeats. This exploration strategy was effective in exploring
both minigrid (Chevalier-Boisvert et al.,, 2023) and the partially observable
3D VizDoom environment (Kempka et al.,, 2016).

####  4.1.2 Balancing exploration and exploitation

##### Cheap uncertainty computations.

Balancing exploration and exploitation is a central challenge in RL (Sutton
and Barto,, 2018; Kaelbling et al.,, 1996). One method that provides close to
optimal performance in tabular domains is posterior sampling (also known as
Thompson sampling, after Thompson,, 1933), where an agent updates a posterior
distribution over Q-values and then chooses the value-maximizing action for a
random sample from this distribution (see Russo et al.,, 2018, for a review).
The main difficulties for implementing posterior sampling are associated with
representing, updating, and sampling from the posterior when the state space
is large. Janz et al., (2019) showed that SFs enable cheap method for
posterior sampling. They assume the following prior, likelihood, and posterior
for the environment reward:

$\displaystyle p(\mathbf{w})=\mathcal{N}(0,\mathbf{I})\quad\quad\quad
p(r|\mathbf{w})=\mathcal{N}(\boldsymbol{\phi}(s)^{\top}\mathbf{w},\beta)\quad\quad\quad
p(\mathbf{w}|r,s)=\mathcal{N}(\boldsymbol{\mu}_{\mathbf{w}},\Sigma_{\mathbf{w}})$
(56)

where $\beta$ is the variance of reward around a mean that is linear in the
state feature $\boldsymbol{\phi}(s)$; $\boldsymbol{\mu}_{\mathbf{w}}$ and
$\Sigma_{\mathbf{w}}$ are (known) analytical solutions for the mean and
variance of the posterior given the Gaussian prior/likelihood assumptions and
a set of observations. The posterior distribution for the Q-values is then
given by:

$\hat{Q}_{\mathrm{SU}}^{\pi}\sim\mathcal{N}\left({\Psi}^{\pi}\boldsymbol{\mu}_{\mathbf{w}},{\Psi}^{\pi}\Sigma_{\mathbf{w}}\left({\Psi}^{\pi}\right)^{\top}\right)$
(57)

where
${\Psi}^{\pi}=[\boldsymbol{\psi}^{\pi}(s,a)]^{\top}_{(s,a)\in\mathcal{S}\times\mathcal{A}}$
is a matrix where each row is an SF for a state-action pair. Successor
features, cumulants, and task encodings are learned with standard losses
(e.g., Eq. 37 and Eq. 41).

##### Count-based exploration.

Another method that provides (near) optimal exploration in tabular settings is
count-based exploration with a bonus of $1/\sqrt{N(s)}$ (Auer,, 2002). Here
$N(s)$ represents the number of times a state $s$ has been visited. However,
when the state space is large it can be challenging to matain this count.
Machado et al., (2020) showed that that the L1-norm of SFs is proportional to
visitation count and can be used as an exploration bonus:

$R^{\tt int}(s)=\frac{1}{||\boldsymbol{\psi}^{\pi}(s)||_{1}}.$ (58)

With this exploration bonus, they were able to improve exploration in sparse-
reward Atari games such as Montezuma’s revenge. Recent work has built on this
idea: Yu et al., (2023) combined SFs with _predecessor_ representations, which
encode retrospective information about the agent’s trajectory. This was shown
to more efficiently target exploration towards bottleneck states (i.e., access
points between large regions of the state space).

###  4.2 Transfer

We’ve already introduced the idea of cross-task transfer in our discussion of
GPI. We now review the broader range of ways in which the challenges of
transfer have been addressed using predictive representations.

####  4.2.1 Transferring behaviors between tasks

We first consider transfer across tasks that are defined by different reward
functions. In the following two sections, we consider other forms of transfer.

##### Few-shot transfer between pairs of tasks.

The SR can enable transferring behaviors from one reward function $R$ to
another function $R_{\tt new}$ by exploiting Eq. 34. One can accomplish this
by learning cumulant $\boldsymbol{\phi}_{\theta}$ and SFs
$\boldsymbol{\psi}_{\theta}$ for a source task. At transfer time, one freezes
each set of parameters and solves for $\mathbf{w}_{\tt new}$ (e.g., with Eq.
37). Kulkarni et al., (2016) showed that this enabled an RL agent that learned
to play an Atari game to transfer to a new version of that game where the
reward function was scaled. Later, Zhu et al., (2017) showed that this enabled
transfer to new “go to” tasks in a photorealistic 3D household environment.

##### Continual learning across a set of tasks.

Beyond transferring across task pairs, an agent may want to continually
transfer its knowledge across a set of tasks. Barreto et al., (2017) showed
that SFs and GPI provide a natural framework to do this. Consider sequentially
learning $n$ tasks. As the agent learns new tasks, they maintain a growing
library of successor features $\\{\boldsymbol{\psi}^{\pi_{i}}\\}^{m}_{i=1}$
where $m<n$ is the number of tasks learned so far. When the agent is learning
the $m$-th task, they can select actions with GPI according to Eq. 2.5.1 using
$\mathbf{w}_{m}$ as the current transfer task. The agent learns successor
features for the current task according to Eq. 41. Zhang et al., (2017)
extended this approach to enable continual learning when the environment state
space and dynamics were changing across tasks but still relatively similar.
Transferring SFs to an environment with a different state space requires
leveraging new state-features. Their solution involved re-using old state-
features by mapping them to the new state space with a linear projection. By
exploiting linearity in the Q-value decomposition (see Eq. 34), this allowed
reusing SFs for new environments.

##### Zero-shot transfer to task combinations.

Another benefit of SFs and GPI are that they facilitate transfer to task
conjunctions when they are defined as weighted sums of training tasks,
$\mathbf{w}_{\tt new}=\sum^{n}_{i=1}\alpha_{i}\mathbf{w}_{i}$. A clear example
is combining “go to” tasks (Barreto et al.,, 2018; Borsa et al.,, 2018;
Barreto et al.,, 2020; Carvalho et al., 2023a, ). For example, consider $4$
tasks $\\{\mathbf{w}_{1},\ldots,\mathbf{w}_{4}\\}$ defining by collecting
different object types; $\mathbf{w}_{\tt
new}=(-1)\cdot\mathbf{w}_{1}+2\cdot\mathbf{w}_{2}+1\cdot\mathbf{w}_{3}+0\cdot\mathbf{w}_{4}$
defines a new task that tries to avoid collecting objects of type $1$, while
trying to collect objects of type $2$ twice as much as objects of type $3$.
This approach has been extended to combining policies with continuous state
and action spaces (Hunt et al.,, 2019), though it has so far been limited to
combining only two policies. Another important limitation of this general
approach is that it can only specify which tasks to prioritize, but cannot
specify an ordering for these tasks. For example, there is no way to specify
collecting object type $1$ before object type $2$. One can address this
limitation by learning a state-dependent transfer task encoding as with the
option keyboard (§2.5.2).

####  4.2.2 Learning about non-task goals

As an agent is learning a task (defined by $\mathbf{w}$), they may also want
to simultaneously do well on another task (defined by $\tilde{\mathbf{w}}$).
This can be helpful for _hindsight experience replay_ , where an agent re-
labels failed experiences at reaching some task goal to successful experiences
in reaching some non-task goal (Andrychowicz et al.,, 2017). This is
particularly useful when tasks have sparse rewards. When learning a policy
with SMs (§2.4), hindsight experience replay naturally arises as part of the
learning objective. It has been shown to improve sample efficiency in sparse-
reward virtual robotic manipulation domains and long-horizon navigation tasks
(Eysenbach et al.,, 2020, 2022; Zheng et al.,, 2023). Despite their potential,
learning and exploiting SMs is still in its infancy, whereas SFs have been
more thoroughly studied. Recently, Schramm et al., (2023) developed an
asymptotically unbiased importance sampling algorithm that leverages SFs to
remove bias when estimating value functions with hindsight experience replay.
This enabled learning for both simulation and real-world robotic manipulation
tasks in environments with large state and action spaces.

Learning to reach non-task goals can also be useful for _off-task learning_.
In this setting, as an agent learns a control policy for $\mathbf{w}$, they
also learn a control policy for $\tilde{\mathbf{w}}$. Borsa et al., (2018)
showed that one can learn control policies for tasks $\tilde{\mathbf{w}}$ that
are not too far from $\mathbf{w}$ in the task encoding space by leverage
universal SFs. In particular, nearby off-task goals can be sampled from a
Gaussian ball around $\mathbf{w}$ with standard deviation $\sigma$:
$Z=\\{\tilde{\mathbf{w}}_{i}\sim\mathcal{N}(\mathbf{w}_{\tt
new},\sigma)\\}^{n}_{i=1}$. Then a successor feature loss following Eq. 41
would be applied for each non-task goal $\tilde{\mathbf{w}}_{i}$. Key to this
is that the optimal action for each $\tilde{\mathbf{w}}_{i}$ would be the
action that maximized the features determined by $\tilde{\mathbf{w}}_{i}$ at
the next time-step, e.g.
$a^{*}=\operatornamewithlimits{argmax}_{a}^{\prime}\boldsymbol{\psi}_{\theta}(s^{\prime},a^{\prime},\tilde{\mathbf{w}}_{i})^{\top}\tilde{\mathbf{w}}_{i}$.
This enabled an agent to learn a policy for not only $\mathbf{w}$ but also for
non-task goals $\tilde{\mathbf{w}}_{i}$ with no direct experience on those
tasks in a simple 3D navigation environment.

####  4.2.3 Other advances in transfer

##### Generalization to new environment dynamics.

One limitation of SFs is that they’re tied to the environment dynamics $T$
with which they were learned. Lehnert and Littman, (2020) and Han and
Tschiatschek, (2021) both attempt to address this limitation by learning SFs
over state abstractions which respect bisimulation relations (Li et al.,,
2006). Abdolshah et al., (2021) attempt to address this by integrating SFs
with Gaussian processes such that they can be quickly adapted to new dynamics
given a small amount of experience in the new environment.

##### Synthesizing new predictions from sets of SFs.

While GPI enables combining a set of SFs to produce a novel policy, it does
not generate a novel prediction of what features will be experienced from a
combination of policies. Some methods attempt to address this by combining SFs
via their convex combination (Brantley et al.,, 2021; Alegre et al.,, 2022).

##### Alternatives to generalized policy improvement.

Madarasz and Behrens, (2019) develop the Gaussian SF, which learns a set of
reward maps for different environments that can be adaptively combined to
adjudicate between different policies. While they had promising results over
GPI, these were in toy domains; it is currently unclear if their method scales
to more complex settings as gracefully as GPI. A more promising alternative to
GPI is Geometric Policy Composition (GPC; Thakoor et al.,, 2022) which enables
estimating Q-values when one follows a ordered sequence of $n$ policies
$(\pi_{1},\ldots,\pi_{n})$. Whereas GPI evaluates how much reward will be
obtained by the best of a set of policies, GPC is a form of model-based
control where we evaluate the path obtained from following a sequence of
policies. We discuss this in more detail in §4.4.

###  4.3 Hierarchical reinforcement learning

Many tasks have multiple time-scales, which can be exploited using a
hierarchical architecture in which the agent learns and acts at multiple
levels of temporal abstraction. The classic example of this is the _options_
framework (Sutton et al.,, 1999). An “option” $o$ is a temporally-extended
behavior defined by (a) an initiation function $\mathcal{I}_{o}$ that
determines when it can be activated, (b) a behavior policy $\pi_{o}$, and (c)
a termination function $\beta_{o}$ that determines when the option should
terminate: $o=\langle\mathcal{I}_{o},\pi_{o},\beta_{o}\rangle$. In this
section, we discuss how predictive representations can be used to discover
useful options and transfer them across tasks.

####  4.3.1 Discovering options to efficiently explore the environment

One key property of the SR is that it captures long-range structure of an
agent’s paths through the environment. This has been useful in discovering
options. Machado et al., (2017, 2023) showed that if one performs an
eigendecomposition on a learned SR, the eigenvectors corresponding to the
highest eigenvalues could be used to identify states with high diffusion
(i.e., states that tend to be visited frequently under a random walk policy).
In particular, for eigenvector $e_{i}$, they defined the following intrinsic
reward function:

$R_{i}^{\tt
int}(s,s^{\prime})=e_{i}^{\top}(\boldsymbol{\phi}(s^{\prime})-\boldsymbol{\phi}(s))$
(59)

which rewards the agent for exploring more diffuse parts of the state space.
Note that $\boldsymbol{\phi}$ is either a one-hot vector in the case of the SR
or can be state-features in the case of successor features. Here, we focus on
the SR. With this reward function, the agent can iteratively build up a set of
options $\mathcal{O}$ that can be leveraged to improve exploration of the
environment. The algorithm proceeds as follows:

  1. 1.

Collect samples with a random policy which selects between between primitive
actions and options. The set of options is initially empty
($\mathcal{O}=\\{\emptyset\\}$).

  2. 2.

Learn successor representation $M^{\pi}$ from the gathered samples.

  3. 3.

Get new exploration option
$o=\langle\mathcal{I}_{o},\pi_{o},\beta_{o}\rangle$. $\pi_{o}$ is a policy
that maximizes the intrinsic reward function in Eq. 59 using the current SR.
The initiation function $\mathcal{I}_{o}$ is $1$ for all states. The
termination function $\beta_{o}$ is $1$ when the intrinsic reward becomes
negative (i.e., when agent begins to go towards more frequent states). This
option is added to the overall set of options,
$\mathcal{O}\leftarrow\mathcal{O}\cup\\{o\\}$.

Agents endowed with this strategy were able to discover meaningful options and
improve sample efficiency in the four-rooms domain, as well as Atari games
such as Montezuma’s revenge.

####  4.3.2 Transferring options with the SR

##### Instant synthesis of combinations of options.

One of the benefits of leveraging SFs is that they enable transfer to tasks
that are linear combinations of known tasks (see §2.5.1). At transfer time, an
agent can exploit this when transferring options by defining subgoals using
this space of tasks (see §2.5.2). In continuous control settings where an
agent has learned to move in a set of directions (e.g., up, down, left,
right), this has enabled generalization to combinations of these policies
(Barreto et al.,, 2019). For example, the agent could instantly move in novel
angles (e.g. up-right, down-left, etc.) as needed to complete a task.

##### Transferring options to new dynamics.

One limitation of both options and SFs is that they are defined for a
particular environment and don’t naturally transfer to new environments with
different transition functions $T^{\prime}$. Han and Tschiatschek, (2021)
addressed this limitation by learning abstract options with SFs defined over
abstract state features that respect bisimulation relations (Li et al.,,
2006). This method assumes that transfer from a set of $N$
$\boldsymbol{\psi}$-MDPs $\\{\mathcal{M}_{i}\\}_{i=1}^{n}$ where
$\mathcal{M}_{i}=\langle S_{i},O_{i},T_{i},\boldsymbol{\psi},\gamma\rangle$.
It also assumes the availability of MDP-dependent state-feature functions
$\boldsymbol{\phi}_{\mathcal{M}_{i}}(s,a)$ that map individual MDP state-
action pairs to a shared feature space. Assume an agent has learned a
(possibly) distinct set of options $\\{o_{k}\\}$ for each source MDP. The SFs
for each option’s policy are defined as usual (Eq. 30). When we are
transferring to a new MDP $\mathcal{M}^{\prime}$, we can map each option to
this environment by mapping its SFs to a policy which produces similar
features using an inverse reinforcement learning algorithm (Ng et al.,, 2000).
Once the agent has transferred options to a new environment, it can plan using
these options by constructing an abstract MDP over which to plan. Han and
Tschiatschek, (2021) implement this using successor homomorphisms, which
define criteria for aggregating states to form an abstract MDP. In particular,
pairs of states $(s_{1},s_{2})$ and options $(o_{1},o_{2})$ will map to the
same abstract MDP if they follow bisimulation relations (Li et al.,, 2006) and
if their SFs are similar:

$||\boldsymbol{\psi}^{o_{1}}{(s_{1},a)}-\boldsymbol{\psi}^{o_{2}}{(s_{2},a)}||\leq\epsilon_{\psi},$
(60)

where $\epsilon_{\psi}$ is a similarity threshold. With this method, agents
were able to discover state abstractions that partitioned large grid-worlds
into intuitive segments and successfully plan in novel MDPs with a small
number of learning updates.

###  4.4 Model-based reinforcement learning

The successor model (SM) is interesting because it offers a novel way to do
model-based reinforcement learning. Traditionally, one simulates trajectories
with a single-step model. While this is very flexible, it is expensive. SMs
offer an alternative, where we can instead sample and evaluate likely
(potentially distal) states that will be encountered when following some
policy $\pi$. Thakoor et al., (2022) leverage this property to develop
Generalized Policy Composition (GPC), a novel algorithm that enables a jumpy
form of model-based reinforcement learning. In this setting, rather than
simulate trajectories defined over next states, we simulate trajectories by
using SMs to jump between states using a given set of policies. While this is
not as flexible as simulating trajectories with a single-step model, it is
much more efficient.

In RL, one typically uses a large discount factor (i.e., $\gamma\approx 1$).
When learning an SM, this is useful because you can learn likelihoods over
potentially very distal states. However, this makes learning an SM more
challenging. GPC mitigates this challenge by composing a shorter horizon SM
$\boldsymbol{\mu}_{\beta}^{\pi}$ with a longer horizon SM
$\boldsymbol{\mu}_{\gamma}^{\pi}$, where $\beta<\gamma$. Composing two
separate SMs with different horizons has the following benefits. An SM with a
shorter horizon $\boldsymbol{\mu}_{\beta}^{\pi}$ is easier to learn but cannot
sample futures as distal as $\boldsymbol{\mu}_{\gamma}^{\pi}$; on the other
hand, $\boldsymbol{\mu}_{\gamma}^{\pi}$ is harder to learn but can do very
long horizon predictions and better avoids compounding errors. By combining
the two, Thakoor et al., (2022) show that you can trade off between these two
errors.

At an intuitive level, GPC with SMs works as follows. Given a starting state-
action pair $(s_{0},a_{0})$ and policies $(\pi_{1},\ldots,\pi_{n})$, we sample
a sequence of $n-1$ next states with our shorter horizon SM
$\boldsymbol{\mu}_{\beta}^{\pi_{i}}$ and $\pi_{i}$, i.e.
$s_{1}\sim\boldsymbol{\mu}_{\beta}^{\pi_{1}}(\cdot|s_{0},a_{0})$,
$a_{1}\sim\pi_{1}(\cdot|s_{1})$,
$s_{2}\sim\boldsymbol{\mu}_{\beta}^{\pi_{1}}(\cdot|s_{1},a_{1})$, and so on.
We then sample a (potentially more distal) state $s_{n}$ from our longer
horizon SM, $\boldsymbol{\mu}_{\gamma}^{\pi}$,
$s_{n}\sim\boldsymbol{\mu}_{\beta}^{\pi_{n}}(\cdot|s_{n-1},a_{n-1})$. The
reward estimates for the sampled state-action pairs can be combined as a
weighted sum to compute $Q^{\pi}(s_{0},a_{0})$ analogously to Eq. 28 (see
Thakoor et al.,, 2022, for technical details). Leveraging GPC enabled
convergence with an order of magnitude fewer samples in the four-rooms domain
and in a continuous-control domain where one had to control an ant through a
maze.

###  4.5 Multi-agent reinforcement learning

As we’ve noted earlier, the SR is conditioned on a policy $\pi$. In a single-
agent setting, the SR provides predictions about what that agent can expect to
experience when executing the policy. However, in multi-agent settings, one
can instead parametrize this prediction with another agent’s policy to form
predictions about what one can expect to see in the environment when other
agent’s follow their own policies. This is the basis for numerous algorithms
that aim to learn about, from, and with other agents (Rabinowitz et al.,,
2018; Kim et al.,, 2022; Filos et al.,, 2021; Gupta et al.,, 2021; Lee et
al.,, 2019).

##### Learning about other agents.

Rabinowitz et al., (2018) showed that an AI agent could learn aspects of
theory of mind (including passing a false belief test) by meta-learning SFs
that described other agents. While Rabinowitz did not explicitly compare
against humans (and was not trying to directly model human cognition), this
remains exciting as a direction for exploring scalable algorithms for human-
like theory of mind.

##### Learning from other agents.

One nice property of SFs is that they can be learned with TD-learning using
off-policy data (i.e., data collected from a policy different from the one
currently being executed). This can be leveraged to learn SFs for the policies
of other agents just as an agent learns SFs for their own policy. Filos et
al., (2021) exploited this to design an agent that simultaneously learned SFs
for both its own policy and for multiple other agents. They were then able to
generalize effectively to new tasks via a combination of all of their policies
by exploiting GPI (Eq. 2.5.1).

##### Learning with other agents.

One of the key benefits of leveraging universal SFs with GPI (Eq. 39) is that
you can systematically change the Q-value used for action selection by (1)
shifting the policy encoding $\mathbf{z}_{\pi}$ defining the SF or (2)
shifting the feature preferences $\mathbf{w}$ that are being maximized. Gupta
et al., (2021) exploit this for cooperative multi-agent RL. Let $s^{(i)}$
refer to the state of the $i$-th agent and $a^{(i)}\in\mathcal{A}_{i}$ be the
action they take. Now let
$\boldsymbol{\psi}_{\theta}\left(s^{(i)},a,\mathbf{z}_{\pi}\right)$ refer to
that agent’s successor features for behavioral encoding $\mathbf{z}_{\pi}$,
let $\mathcal{C}_{\tt tasks}$ denote the set of possible feature preferences
that agents can follow, and let $\mathcal{C}_{\pi}$ denote the set of possible
policy encodings over which agents can make predictions. Gupta et al., (2021)
studied the cooperative setting where the overall Q-value is simply the sum of
individual Q-values. If $\mathcal{C}_{\pi}$ denotes policies for separate but
related tasks, acting in accordance to

$a^{(i)}=\operatornamewithlimits{argmax}_{a\in\mathcal{A}_{i}}\max_{\mathbf{w}\in\mathcal{C}_{\tt
tasks}}\max_{\mathbf{z}_{\pi}\in\mathcal{C}_{\pi}}\left\\{\boldsymbol{\psi}_{\theta}\left(s^{(i)},a,\mathbf{z}_{\pi}\right)^{\top}\mathbf{w}\right\\}$
(61)

enables a collection of agents to take actions which systematically explore
the environment to quickly do well on tasks specified by $\mathcal{C}_{\tt
tasks}$. They were able to improve exploration and zero-shot generalization in
a variety of multi-agent environments including Starcraft (Samvelyan et al.,,
2019).

###  4.6 Other artificial intelligence applications

The SR and its generalizations have been broadly applied within other areas of
AI. For example, it has been used to define an improved similarity metric in
episodic control settings (Emukpere et al.,, 2021). By leveraging SFs, one can
incorporate information from prior experienced states with similar future
dynamics to the current state. The SR has also been applied towards improving
importance sampling in off-policy learning. If we learn a density ratio
similar to Eq. 49, this can enable more simple marginalized importance
sampling algorithms (Liu et al.,, 2018) that improve off-policy evaluation
(Fujimoto et al.,, 2021). In addition to these example, we highlight the
following applications of the SR.

##### Representation learning.

Learning SMs can obviate the need for separate representation losses. In many
applications, the reward signal is not enough to drive learning of useful
representations. Some strategies to address this challenge include data
augmentation and learning of auxiliary tasks. Learning the SM has been shown
to enable representation learning with superior sample-efficiency without
using these additions (Eysenbach et al.,, 2020, 2022; Zheng et al.,, 2023).

##### Learning auxiliary tasks

is also a natural use-case for predictive representations. A simple and
natural example comes from inspecting the loss for learning SFs (Eq. 41). In
standard Q-learning, we only learn about achieving task-specific reward. When
learning SFs, we also learn representations that enable achieving state-
features that are potentially not relevant for the current task (i.e., we
are—by default—leaning auxiliary tasks). This important ability is even
possible in a continual learning setting where the distribution of state-
features is non-stationary (McLeod et al.,, 2021). Another interesting example
comes from proto-value networks (Farebrother et al.,, 2023). The authors show
that if one learns a successor measure (a set-inclusion based generalization
of the SR) over random sets, this can enable the discovery of predictive
representations that enable very fast learning in the Atari learning
environment (5 million vs. 200 million samples).

##### Learning diverse policies.

A final application of the SR has been in learning diverse policies. In the
mujoco environment, Zahavy et al., (2021) showed that SFs enabled discovering
a set of diverse behaviors for controlling a simulated dog avatar. Their
approach used SFs to prospectively summarize future trajectories. A set of
policies was than incrementally learned so that each new policy would be
different in its expected features from all policies learned so far. This
ideas was then generalized to enable diverse chess playing (Zahavy et al.,,
2023). Similar to before, different strategies were discovered such that their
expected future features would be maximally different.

##  5 Cognitive science applications

A rich body of work dating back over a century has linked RL algorithms to
reward-based learning processes in the brains of humans and nonhuman animals
(Niv,, 2009). As explained further below, model-based control is thought to
underlie reflective, goal-directed behaviors, while model-free control is
thought to underlie reflexive, habitual behaviors. The existence of both
systems in the brain and their synergistic operation has received extensive
support by a wide range of behavioral and neural studies across a number of
species and experimental paradigms (see Dolan and Dayan,, 2013, for a detailed
review).

The general success of RL in modeling brain function, and the success of the
SR in solving problems faced by artificial agents, raises the question of
whether evolution has arrived at similar predictive representations for
tackling problems faced by biological agents. In particular, the SR can confer
ecological advantages when flexibility and efficiency are both desirable,
which is the case in most real-world decision making scenarios. In the field
of cognitive science, several lines of research suggest that human learning
and generalization is indeed consistent with the SR. In this section, we
examine studies showing that patterns of responding to changes in the
environment (Momennejad et al.,, 2017), transfer of knowledge across tasks
(Tomov et al.,, 2021), planning in spatial domains (Geerts et al.,, 2023), and
contextual memory and generalization (Gershman et al.,, 2012; Smith et al.,,
2013; Zhou et al.,, 2023) exhibit signature characteristics of the SR that
cannot be captured by pure model-based or model-free strategies.

Figure 5: Predictive representations in cognitive science. (A) Revaluation
paradigm and predictions from Momennejad et al., (2017, §5.1). (B) Multi-task
learning paradigm from Tomov et al., (2021, §5.2). Reused with permission from
the authors. A person is trained on tasks where they are either hungry
($\mathbf{w}_{\textit{train}}^{1}=[1,0,0]$) or groggy
($\mathbf{w}_{\textit{train}}^{2}=[0,1,0]$), and then tested on a task in
which they are hungry, groggy, and looking to have fun
($\mathbf{w}_{\textit{new}}=[1,1,1]$). (C) Context-dependent Bayesian SR from
Geerts et al., (2023, §5.3). $z$, context. $\mathbf{M}$, SR matrix associated
with each context. sCRP, sticky Chinese restaurant process. LDS, linear-
Gaussian dynamical system. (D) Spatial navigation paradigm from de Cothi et
al., (2022, §5.4). (E) TCM-SR: Using the temporal context model (TCM) to learn
the SR for decision making (Zhou et al.,, 2023, §5.5). $\mathbf{M}$, SR
matrix. $\mathbf{c}$, context vector. $s$, state and/or item. SR, successor
representation.

###  5.1 Revaluation

Some of the key findings pointing to a balance between a goal-directed system
and a habitual system in the brain came from studies of reinforcer revaluation
(Adams and Dickinson,, 1981; Adams,, 1982; Dickinson,, 1985; Holland,, 2004).
In a typical revaluation paradigm, an animal (e.g., a rat) is trained to
associate a neutral action (e.g., a lever press) with an appetitive outcome
(e.g., food). The value of that outcome is subsequently reduced (e.g., the rat
is satiated, so food is less desirable) and the experimenter measures whether
the animal keeps taking the action in the absence of reinforcement. Goal-
directed control predicts that the animal would not take the action, since the
outcome is no longer valuable, while habitual control predicts that the animal
would keep taking the action, since the action itself was not devalued.
Experimenters found that under some conditions—such as moderate training,
complex tasks, or disruptions to dopamine inputs to striatum—behavior appears
to be goal-directed (e.g., lever pressing is reduced), while under other
conditions—such as extensive training, or disruptions to prefrontal
cortex—behavior appears to be habitual (e.g., the rat keeps pressing the
lever).

A modeling study by Daw et al., (2005) interpreted these findings through the
lens of RL. The authors formalized the goal-directed system as a model-based
controller putatively implemented in prefrontal cortex and the habitual system
as a model-free controller putatively implemented in dorsolateral stratum.
They proposed that the brain arbitrates dynamically between the two
controllers based on the uncertainty of their value estimates, preferring the
more certain (and hence likely more accurate) estimate for a given action.
Under this account, moderate training or complex tasks would favor the model-
based estimates, since the model-free estimates may take longer to converge
and hence be less reliable. On the other hand, extensive training would favor
the model-free estimates, since they will likely have converged and hence be
more reliable than the noisy model-based simulations.

One limitation of this account is that it explains sensitivity to outcome
revaluation in terms of a predictive model, but it does not rule out the
possibility that the animals may instead be relying on a predictive
representation. A hallmark feature of predictive representations is that they
allow an agent to adapt quickly to changes in the environment that keep its
cached predictions relevant, but not to changes that require updating them. In
particular, an agent equipped with the SR should adapt quickly to changes in
the reward structure ($R$) of the environment, but not to changes in the
transition structure ($T$). Since the earlier studies on outcome revaluation
effectively only manipulated reward structure, both model-based control and
the SR could account for them, leaving open the question of whether outcome
revaluation effects could be fully explained by the SR instead.

This question was addressed in a study by Momennejad et al., (2017) which
examined how changes in either the reward structure or the transition
structure experienced by human participants affect their subsequent choices.
The authors used a two-step task consisting of three phases: a learning phase,
a relearning (or revaluation) phase, and a test phase (Figure 5A). During the
learning phase, participants were presented with two distinct two-step
sequences of stimuli and rewards corresponding to two distinct trajectories
through state space. The first trajectory terminated with high reward
($s_{1}\rightarrow s_{3}\rightarrow s_{5}$ $10), while the second trajectory
terminated with a low reward ($s_{2}\rightarrow s_{4}\rightarrow s_{6}$ $1),
leading participants to prefer the first one over the second one.

During the revaluation phase, participants had to relearn the second half of
each trajectory. Importantly, the structure of the trajectories changed
differently depending on the experimental condition. In the reward revaluation
condition, the transitions between states remained unchanged, but the rewards
of the two terminal states swapped ($s_{3}\rightarrow s_{5}$ $1;
$s_{4}\rightarrow s_{6}$ $10). In contrast, in the transition revaluation
condition, the rewards remained the same, but the transitions to the terminal
states swapped ($s_{3}\rightarrow s_{6}$ $1; $s_{4}\rightarrow s_{5}$ $10).

Finally, in the test phase, participants were asked to choose between the two
initial states of the two trajectories ($s_{1}$ and $s_{2}$). Note that, under
both revaluation conditions, participants should now prefer the initial state
of the second trajectory ($s_{2}$) as it now leads to the higher reward.

Unlike previous revaluation studies, this design clearly disambiguates between
the predictions of model-free, model-based, and SR learners. Since the initial
states ($s_{1}$ and $s_{2}$) never appear during the revaluation phase, a pure
model-free learner would not update the cached values associated with those
states and would still prefer the initial state of the first trajectory
($s_{1}$). On the other hand, a pure model-based learner would update its
reward ($\hat{R}$) or transition ($\hat{T}$) estimates during the
corresponding revaluation phase, allowing it to simulate the new outcomes from
each initial state and make the optimal choice ($s_{2}$) during the test
phase. Critically, both model-free and model-based learners (and any hybrid
between them) would exhibit the same preferences during the test phase in both
revaluation conditions.

In contrast, an SR learner would show differential responding in the test
phase depending on the revaluation condition, adapting and choosing optimally
after reward revaluation but not after transition revaluation. Specifically,
during the learning phase, an SR learner would learn the successor states for
each initial state (the SR itself, i.e. $s_{1}\rightarrow s_{5}$;
$s_{2}\rightarrow s_{6}$). In the reward revaluation condition, it would then
update its reward estimates ($\hat{R}$) for the terminal states ($s_{5}$ $1;
$s_{6}$ $10) during the revaluation phase, much like the model-based learner.
Then, during the test phase, it would combine the updated reward estimates
with the SR to compute the updated values of the initial states
($s_{1}\rightarrow s_{5}$ $1; $s_{2}\rightarrow s_{6}$ $10), allowing it to
choose the better one ($s_{2}$). In contrast, in the transition revaluation
condition, the SR learner would not have an opportunity to update the SR of
the initial states ($s_{1}$ and $s_{2}$) since they are never presented during
the revaluation phase, much like the model-free learner. Then, during the test
phase, it would combine the unchanged reward estimates with its old but now
incorrect SR to produce incorrect estimates for the initial states
($s_{1}\rightarrow s_{5}$ $10; $s_{2}\rightarrow s_{6}$ $1) and choose the
worse one ($s_{1}$).

The pattern of human responses showed evidence of both model-based and SR
learning: participants were sensitive to both reward and transition
revaluations, consistent with model-based learning, but they performed
significantly better after reward revaluations, consistent with SR learning.
To rule out the possibility that this effect can be attributed to pure model-
based learning with different learning rates for reward ($\hat{R}$) versus
transition ($\hat{T}$) estimates, the researchers extended this Pavlovian
design to an instrumental design in which participants’ choices (i.e., their
policy $\pi$) altered the trajectories they experienced. Importantly, this
would correspondingly alter the learned SR: unrewarding states would be less
likely under a good policy and hence not be prominent (or not appear at all)
in the SR for that policy. Such states could thus get overlooked by an SR
learner if they suddenly became rewarding. This subtle kind of reward
revaluation (dubbed policy revaluation by the researchers) also relies on
changes in the reward structure $R$, but induces predictions similar to the
transition revaluation condition: SR learners would not adapt quickly, while
model-based learners would adapt just as quickly as in the regular reward
revaluation condition.

Human responses on the test phase after policy revaluation were similar to
responses after transition revaluation but significantly worse than responses
after reward revaluation, thus ruling out a model-based strategy with
different learning rates for $\hat{R}$ and $\hat{T}$. Overall, the authors
interpreted their results as evidence of a hybrid model-based-SR strategy,
suggesting that the human brain can adapt to changes in the environment both
by updating its internal model of the world and by learning and leveraging its
cached predictive representations (see also Kahn and Daw,, 2023, for
additional human behavioral data leading to similar conclusions).

###  5.2 Multi-task learning

In the previous section, we saw that humans can adapt quickly to changes in
the reward structure of the environment (reward revaluation), as predicted by
the SR (Momennejad et al.,, 2017). Here we take this idea further and propose
that humans use something like the GPI algorithm to generalize across multiple
tasks with different reward functions (§2.5.1; Barreto et al.,, 2017, 2018,
2020).

In Tomov et al., (2021), participants were presented with different two-step
tasks that shared the same transition structure but had different reward
functions determined by the reward weights $\mathbf{w}$ (see Figure 5B for an
illustration). Each state $s$ was associated with different set of features
$\boldsymbol{\phi}(s)$, which were valued differently depending on the reward
weights $\mathbf{w}$ for a particular task. On each training trial,
participants were first shown the weight vector $\mathbf{w}_{\textit{train}}$
for the current trial, and then asked to navigate the environment in order to
maximize reward. At the end of the experiment, participants were presented
with a single test trial on a novel task $\mathbf{w}_{\textit{new}}$.

The main dependent measure was participant behavior on the test task, which
was designed (along with the training tasks, state features, and transitions)
to distinguish between several possible generalization strategies. Across
several experiments, Tomov et al., (2021) found that participant behavior was
consistent with SF and GPI. In particular, on the first (and only) test trial,
participants tended to prefer the training policy that performed better on the
new task, even when this was not the optimal policy. This “policy reuse” is a
key behavioral signature of GPI. This effect could not be explained by model-
based or model-free accounts. Their results suggest that humans rely on
predictive representations from previously encountered tasks to choose
promising actions on novel tasks.

###  5.3 Associative learning

RL provides a normative account of associative learning, explaining how and
why agents ought to acquire long-term reward predictions based on their
experience. It also provides a descriptive account of a myriad of phenomena in
the associative learning literature (Sutton and Barto,, 1990; Niv,, 2009;
Ludvig et al.,, 2012). Two recent ideas have added nuance to this story:

  * •

Bayesian learning. Animals represent and utilize uncertainty in their
estimates.

  * •

Context-dependent learning. Animals partition the environment into separate
contexts and maintain separate estimates for each context.

We examine each idea in turn and then explore how they can be combined with
the SR. In brief, the key idea is that animals learn a probability
distribution over context-dependent predictive representations.

####  5.3.1 Bayesian RL

While standard RL algorithms learn point estimates of different unknown
quantities like the transition function ($\hat{T}$) or the value function
($\hat{V}$), Bayesian RL posits that agents treat such unknown quantities as
random variables and represent beliefs about them as probability
distributions.

Generally, Bayesian learners assume a generative process of the environment
according to which hidden variables give rise to observable data. The hidden
variables can be inferred by inverting the generative process using Bayes’
rule. In one example of Bayesian value learning (Gershman,, 2015), the agent
assumes that the expected reward $R(s,\mathbf{w})$ is a linear combination of
the observable state features $\boldsymbol{\phi}(s)$:888Note that the linear
reward model is the same as assumed in much of the SF work reviewed above (see
Eq. 29).

$\displaystyle
R(s,\mathbf{w})=\boldsymbol{\phi}\left(s\right)^{\top}\mathbf{w}$ (62)

where the hidden weights $\mathbf{w}$ are assumed to evolve over time
according to the following dynamical equations:

$\displaystyle\mathbf{w}_{0}$
$\displaystyle\sim\mathcal{N}(\mathbf{0},\mathbf{\Sigma}_{0}),$ (63)
$\displaystyle\mathbf{w}_{t}$
$\displaystyle\sim\mathcal{N}(\mathbf{w}_{t-1},q\mathbf{I})$ (64)
$\displaystyle r_{t}$
$\displaystyle\sim\mathcal{N}(R(s_{t},\mathbf{w}_{t}),\sigma^{2}_{R}).$ (65)

In effect, this means that each feature (or stimulus dimension) is assigned a
certain reward weight. The weights are initialized randomly around zero (with
prior covariance $\mathbf{\Sigma}_{0}$) and evolve according to a random walk
(with volatility governed by the transition noise variance $q$). Observed
rewards are given by the linear model (Eq. 62) plus zero-mean Gaussian noise
with variance $\sigma^{2}_{R}$.

This formulation corresponds to a linear-Gaussian dynamical system (LDS). The
posterior distribution over weights given the observation history
$H_{t}=(s_{1:t},r_{1:t})$ follows from Bayes’ rule:

$\displaystyle
p(\mathbf{w}_{t}|H_{t})=\frac{p(H_{t}|\mathbf{w}_{t})p(\mathbf{w}_{t})}{p(H_{t})}=\mathcal{N}(\mathbf{w}_{t};\hat{\mathbf{w}}_{t},\mathbf{\Sigma}_{t}).$
(66)

The posterior mean $\hat{\mathbf{w}}_{t}$ and covariance $\mathbf{\Sigma}_{t}$
can be computed recursively in closed form using the Kalman filtering
equations:

$\displaystyle\hat{\mathbf{w}}_{t}$
$\displaystyle=\hat{\mathbf{w}}_{t-1}+\mathbf{k}_{t}\delta_{t}$ (67)
$\displaystyle\mathbf{\Sigma}_{t}$
$\displaystyle=\mathbf{\Sigma}_{t-1}+q\mathbf{I}-\lambda_{t}\mathbf{k}\mathbf{k}_{t}^{\top},$
(68)

where

  * •

$\delta_{t}=r_{t}-\boldsymbol{\phi}(s_{t})^{\top}\hat{\mathbf{w}}_{t}$ is the
reward prediction error.

  * •

$\lambda_{t}=\boldsymbol{\phi}(s_{t})^{\top}(\mathbf{\Sigma}_{t}+q\mathbf{I})\boldsymbol{\phi}(s_{t})+\sigma_{R}^{2}$
is the residual variance.

  * •

$\mathbf{k}_{t}=(\mathbf{\Sigma}_{t}+q\mathbf{I})\boldsymbol{\phi}(s)/\lambda$
is the Kalman gain.

This learning algorithm is generalizes the seminal Rescorla-Wagner model of
associative learning (Rescorla and Wagner,, 1972), and its update rule bears
resemblance to the error-driven TD update (Eq. 11). However, there are a few
notable distinctions from its non-Bayesian counterparts:

  * •

Uncertainty-modulated updating. The learning rate corresponds to the Kalman
gain $\mathbf{k}_{t}$, which increases with posterior uncertainty (the
diagonals of the posterior covariance, $\mathbf{\Sigma}_{t}$).

  * •

Non-local updating. The update is multivariate, affecting the weights of all
features simultaneously according to $\mathbf{k}_{t}$. This means that weights
of “absent” features (i.e., where $\phi_{i}(s)=0$) can be updated, provided
those features have non-zero covariance with observed features.

These properties allow the Kalman filter to explain a number of phenomena in
associative learning that elude non-Bayesian accounts (Dayan and Kakade,,
2000; Kruschke,, 2008; Gershman,, 2015). Uncertainty-modulated updating
implies that nonreinforced exposure to stimuli, or even just the passage of
time, can impact future learning. For example, in a phenomenon known as
_latent inhibition_ , pre-exposure to a stimulus reduces uncertainty, which in
turn reduces the Kalman gain, retarding subsequent learning for that stimulus.
Non-local updating implies that learning could occur even for unobserved
stimuli, if they covary with observed stimuli according to
$\mathbf{\Sigma}_{t}$. For example, in a phenomenon known as _backward
blocking_ , reinforcing two stimuli simultaneously (a compound stimulus,
$\boldsymbol{\phi}(s)=[1,1]$) and subsequently reinforcing only one of the
stimuli ($\boldsymbol{\phi}(s^{\prime})=[1,0]$) reduces reward expectation for
the absent stimulus ($\boldsymbol{\phi}(s^{\prime\prime})=[0,1]$) due to the
learned negative covariance between the two reward weights
($\mathbf{\Sigma}_{1,2}$ < 0).

Kalman filtering can also been extended to accommodate long-term reward
expectations in an algorithm known as Kalman TD (Geist and Pietquin,, 2010) by
replacing the state features $\boldsymbol{\phi}(s)$ with their discounted
temporal difference:

$\displaystyle\mathbf{h}_{t}=\boldsymbol{\phi}(s_{t})-\gamma\boldsymbol{\phi}(s_{t+1}).$
(69)

This recovers the Kalman filter model described above when the discount factor
$\gamma$ is set to 0. Importantly, this model provides a link between the
Bayesian approach to associative learning and model-free reinforcement
learning algorithms; we exploit this link below when we discuss Bayesian
updating of beliefs about predictive representations. Gershman, (2015) showed
that Kalman TD can explain a range of associative learning phenomena that are
difficult to explain by Bayesian myopic ($\gamma=0$) models and non-myopic
($\gamma>0$) non-Bayesian models.

####  5.3.2 Context-dependent learning

In addition to accounting for uncertainty in their estimates, animals have
been shown to learn different estimates in different contexts (Redish et al.,,
2007). This kind of context-dependent learning can also be cast in a Bayesian
framework (Courville et al.,, 2006; Gershman et al.,, 2010) by assuming that
the environment is carved into separate contexts according to a hypothetical
generative process. One possible formalization is the _Chinese restaurant
process_ (CRP), in which each observation is conditioned on a hidden context
$z_{t}$ that is currently active and evolves according to:

$\displaystyle
P(z_{t}=k\mid\mathbf{z}_{1:t-1})\propto\begin{cases}N_{k}\,\,\,\,\,&\text{if
$k$ is a previously sampled context}\\\
\alpha\,\,\,\,\,&\text{otherwise}\end{cases}$ (70)

where $N_{k}=\sum_{i=1}^{t-1}\mathbb{I}[z_{i}=k]$ is the number of previous
time steps assigned to context $k$ and $\alpha\geq 0$ is a _concentration
parameter_ , which can be understood as controlling the effective number of
contexts.999Under the CRP, the expected number of contexts after $t$ time
steps is $\alpha\log t$.

If each context is assigned its own probability distribution over
observations, then inferring a given context $k$ is driven by two factors:

  * •

How likely is the current observation under context $k$?

  * •

How likely is context $k$ given the previous context assignments
$\mathbf{z}_{1:t-1}$ (Eq. 70)?

In particular, a given context is more likely to be inferred if the current
observations are more likely under its observation distribution and/or if it
has been inferred more frequently in the past (that is, if past observations
have also been consistent with its observation distribution). Conversely, if
the current observations are unlikely under any previously encountered
context, a new context is induced with its own observation distribution.

This formulation has accounted for a number of phenomena in the animal
learning literature (Gershman et al.,, 2010). For example, it explains why
associations that have been extinguished sometimes reappear when the animal is
returned to the context in which the association was first learned, a
phenomenon known as _renewal_. It also provides an explanation of why the
latent inhibition effect is attenuated if pre-exposure to a stimulus occurs in
one context and reinforcement occurs in a different context.

Recently, a similar formulation has been used to explain variability in the
way hippocampal place cells change their firing patterns in response to
changes in contextual cues, a phenomenon known as _remapping_ (Sanders et
al.,, 2020). On this view, the hippocampus maintains a separate cognitive map
of the environment for each context, and hippocampal remapping reflects
inferences about the current context.

####  5.3.3 Bayesian learning of context-dependent predictive maps

As discussed further in §6.2, Stachenfeld et al., (2017) showed that many
aspects of hippocampal place cell firing patterns are consistent with the SR,
suggesting that the hippocampus encodes a predictive map of the environment.
In this light, the view that the hippocampus learns different maps for
different contexts (Sanders et al.,, 2020) naturally points to the idea of a
context-dependent predictive map.

This idea was formalized by Geerts et al., (2023) in a model that combines SFs
with Bayesian learning and context-dependent learning. Their model learns a
separate predictive map of the environment for each context, where each map
takes the form of a probability distribution over SFs (Figure 5C). The
generative model assumes that the SF is given by a linear combination of state
features:

$\displaystyle\boldsymbol{\psi}_{j,t}(s)=\boldsymbol{\phi}(s)^{\top}\mathbf{m}_{j,t},$
(71)

where $\mathbf{m}_{j}$ is a vector of weights for predicting successor feature
$j$. Analogous to the Kalman TD model described above, Geerts et al. assume
that the weight vectors for all features evolve over time according to the
following LDS:

$\displaystyle\mathbf{m}_{j,0}$
$\displaystyle\sim\mathcal{N}(\mathbf{\mu}_{0},\mathbf{\Sigma}_{0})$ (72)
$\displaystyle\mathbf{m}_{j,t}$
$\displaystyle\sim\mathcal{N}(\mathbf{m}_{j,t-1},q\mathbf{I})$ (73)
$\displaystyle\phi_{j,t}(s_{t})$
$\displaystyle\sim\mathcal{N}(\mathbf{h}_{t}^{\top}\mathbf{m}_{j,t-1},\sigma_{\phi}^{2}).$
(74)

This means that the weight vector is initialized randomly (with prior mean
$\mathbf{\mu}_{0}$ and variance $\mathbf{\Sigma}_{0}$) and evolves randomly
according to a random walk (with transition noise variance $q$). Observed
features $\boldsymbol{\phi}(s)$ are noisy differences in the SFs of successive
states (with observation noise variance $\sigma_{\phi}^{2}$). As in Kalman TD,
the posterior over the weight vector for feature $j$ is also Gaussian with
mean $\hat{\mathbf{m}}_{j}$ and variance $\mathbf{\Sigma}$ which can be
computed using Kalman filtering equations essentially the same as those given
above (Eq. 67-68).

The authors generalize this to multiple contexts using a _switching LDS_ :
each context has a corresponding LDS and SF weights. The contexts themselves
change over time based on a “sticky” version of the Chinese restaurant
process, which additionally favors remaining in the most recent context:

$\displaystyle
P(z_{t}=k\mid\mathbf{z}_{1:t-1})\propto\begin{cases}N_{k}+\nu\mathbb{I}[z_{t-1}=k]\,\,\,\,\,&\text{if
$k$ is a previously sampled context}\\\
\alpha\,\,\,\,\,&\text{otherwise}\end{cases},$ (75)

where $\nu\geq 0$ is the “stickiness” parameter dictating how likely it is to
remain in the same context.

Thus a given context $k$ is more likely to be inferred if:

  * •

The current observations are consistent with its SFs (Eq. 74).

  * •

It was encountered frequently in the past (driven by the $N_{k}$ term in Eq.
75).

  * •

It was encountered recently (driven by the $\mathbb{I}[z_{t-1}=k]$ term in Eq.
75).

Conversely, a new context is more likely to be inferred if observations are
inconsistent with the current SFs (i.e., when there are large SF prediction
errors).

The authors show that this model can account for a number of puzzling effects
in the animal learning literature that pose problems for both point estimation
(TD learning) of the SR/SF (Eq. 17-18) and Bayesian RL (Eq. 67-68). One
example is the opposing effect that pre-exposure to a context can have on
learning. Brief exposure to a context can facilitate learning (context pre-
exposure facilitation), while prolonged exposure can inhibit learning (latent
inhibition; Kiernan and Westbrook,, 1993). Context pre-exposure facilitation
by itself can be accounted for by TD learning of the SR alone (Stachenfeld et
al.,, 2017): during pre-exposure, the animal learns a predictive
representation which facilitates propagation of newly learned values. Latent
inhibition by itself can be accounted for by Bayesian RL alone (Gershman,,
2015), as discussed previously: prolonged exposure reduces value uncertainty,
in turn reducing the Kalman gain $\mathbf{k}_{t}$ (the effective learning
rate) and inhibiting learning of new values. Kalman learning of SFs combines
these two processes and can thus resolve the apparent paradox: initially, the
animal learns a predictive representation of the context, which facilitates
learning, whereas after prolonged exposure, this effect is offset by a
reduction in value uncertainty, which inhibits learning.

Another puzzling effect is the partial transition revaluation observed in
Momennejad et al., (2017) and discussed in §5.1, which cannot be accounted for
by TD learning of the SR. This led Momennejad et al. to propose a hybrid
model-based-SR strategy that relies on offline simulations. Kalman TD offers a
more parsimonious account based on non-local updating that does not appeal to
model-based simulations. In particular, the covariance matrix
$\mathbf{\Sigma}_{t}$ learned during the learning phase captures the
relationship between the initial states ($s_{1}$ and $s_{2}$) and subsequent
states ($s_{3}$ and $s_{4}$). Updating the transitions from those subsequent
states ($s_{3}\rightarrow s_{6}$ and $s_{4}\rightarrow s_{5}$) during the
transition revaluation phase therefore also updates the SR for the initial
states ($s_{1}$ and $s_{2}$), even though they are not encountered during the
revaluation phase.

Non-local updating can similarly explain reward devaluation of preconditioned
stimuli, a hallmark of model-based learning (Hart et al.,, 2020). This is
similar to reward devaluation discussed in §5.1, except with an additional
preconditioning phase during which an association is learned between two
neutral stimuli (e.g., light $\rightarrow$ tone). During the subsequent
conditioning phase, the second stimulus is associated with a rewarding outcome
(e.g., tone $\rightarrow$ food), which is then devalued (e.g., by inducing
taste aversion) during the devaluation phase. Finally, the animal is tested on
the first neutral stimulus (e.g., light). Note that since the first stimulus
is never present during the conditioning phase, TD learning would not acquire
an association between the first stimulus and the reward, and would thus not
exhibit sensitivity to reward devaluation (Gardner et al.,, 2018). In
contrast, during the preconditioning phase, Kalman TD learns that the two
stimuli covary, allowing it to update the SF for both stimuli during the
conditioning phase and consequently propagate the updated value to both
stimuli during the devaluation phase.

Note that the phenomena so far can be explained without appealing to context-
dependent learning, since the experiments take place in the same context.
Context-dependent Kalman TD can additionally explain a number of intriguing
phenomena when multiple contexts are introduced.

One such phenomenon is the context specificity of learned associations
(Winocur et al.,, 2009). In this paradigm, an animal learns an association
(e.g., tone $\rightarrow$ shock) in one context (e.g., Context A) and is then
tested in the same context or in another context (e.g., Context B). The amount
of generalization of the association across contexts was found to depend on
elapsed time: if testing occurs soon after training, the animal only responds
in the training context (A), indicating context specificity. However, if
testing occurs after a delay, the animal responds equally in both contexts A
and B, indicating contextual generalization. Even more intriguingly, this
effect is reversed if the animal is briefly reintroduced to the training
context (A) before testing, in which case responding is once again context-
specific—a hippocampus-dependent “reminder” effect.

The context-dependent model readily accounts for these effects. Shortly after
training, the uncertainty of the SR assigned to Context A is low (i.e., the
animal is confident in its predictive representation of Context A).
Introduction to Context B therefore results in a large prediction error,
leading the animal to (correctly) infer a new context with a new SR, leading
to context-specific responding. However, as more time elapses, the uncertainty
of the SR assigned to Context A gradually increases (i.e., the animal becomes
less confident in its predictive representation of Context A, a kind of
forgetting). Introduction to Context B then results in a smaller prediction
error, making it likely that the new observations are also assigned to Context
A, leading to generalization across contexts. Brief exposure to Context A
reverses this effect by reducing the uncertainty of the SR assigned to Context
A (i.e., the animal’s confidence in its predictive representation of Context A
is restored, a kind of remembering), leading once again to context-specific
responding.

Recall that the duration of context pre-exposure has opposing effects on
learning, initially facilitating but subsequently inhibiting learning (Kiernan
and Westbrook,, 1993). But what if the animal is tested in a different
context? In a follow-up experiment, Kiernan and Westbrook, (1993) showed that
longer pre-exposure to the training context leads to less responding in the
test context, indicating that the learned association is not generalized. That
is, longer context pre-exposure has a monotonic inhibitory effect on
generalization across contexts. The context-dependent model can account for
this with the same mechanism that accounts for context pre-exposure
facilitation: longer pre-exposure to the training context reduces the
uncertainty of its predictive estimate, leading to greater prediction errors
when presented with the text context, and increasing the probability that the
animal will infer a new context, leading to context-specific responding.

Overall, the results of Geerts et al., (2023) suggest that, rather than
encoding a single monolithic predictive map of the environment, the
hippocampus encodes multiple separate predictive maps, each associated with
its own context. Both the context and the predictive map are inferred using
Bayesian inference: learning of the predictive map is modulated by
uncertainty, and supports non-local updating. A new context is inferred when
the current predictive map fails to account for current observations.

###  5.4 Spatial navigation

A rich body of work points to the hippocampus as encoding a kind of cognitive
map of the environment that mammals rely on for navigation in physical and
abstract state spaces (O’Keefe and Dostrovsky,, 1971; O’Keefe and Nadel,,
1978). As we discussed in the previous section, and as we will discuss further
in §6.2, this cognitive map is best understood as a predictive map in which
states predict future expected states, consistent with the SR (Stachenfeld et
al.,, 2017). Yet despite many studies of navigation in humans and rodents—key
model species used to study spatial navigation (Epstein et al.,, 2017; Ekstrom
and Ranganath,, 2018; Ekstrom et al.,, 2018; Gahnstrom and Spiers,, 2020;
Nyberg et al.,, 2022; Spiers and Barry,, 2015)—until recently there was no
direct comparison of human and rodent navigation in a homologous task. This
left open the question of whether spatial navigation in mammals consistently
relies on the same SR-dependent strategy.

A recent study by de Cothi et al., (2022) filled this gap by designing a
homologous navigation task for humans and rats. They devised a configurable
open-field maze that could be reconfigured between trials, allowing
experimenters to assess a hallmark aspect of spatial navigation: the ability
to efficiently find detours and shortcuts. The maze consisted of a 10-by-10
grid in which squares could be blocked off by the experimenters. The maze was
instantiated in a physical environment for rats and in a virtual reality
environment for humans.

On each trial, the participant was placed at a starting location and had to
navigate to a goal location to receive a reward (Figure 5D). The starting
location varied across trials, while the goal remained hidden at a fixed
location throughout the experiment. Keeping the goal location unobservable
ensured that participants could not rely on simple visual heuristics (e.g.,
proximity to the goal). At the same time, keeping the goal location fixed
ensured that, once it is identified, the key problem becomes navigating to it,
rather than re-discovering it. During the training phase of the experiment,
all squares of the grid were accessible, allowing participants to learn an
internal map of the environment. During the test phase, participants were
sequentially presented with 25 different maze configurations in which various
sections of the maze were blocked off. Participants completed 10 trials of
each configuration before moving on to the next.

Using this task, the authors compared human and rat navigation with three
types of RL algorithms:

  * •

Model-free agent (§2.2). No internal map of the environment; optimal policy is
based on state-action value function $Q$, which is learned from experience
(specifically, using Q-learning (Eq. 11) with eligibility traces).

  * •

Model-based agent (§2.2). Full internal map of the environment (transition
structure $T$ and reward function $R$) is learned from experience; optimal
policy is computed using tree search at decision time (specifically, using A*
search).

  * •

SR agent (§2.3). Predictive map of the environment (SR matrix $\mathbf{M}$ and
reward function $R$) is learned from experience; optimal policy is computed by
combining SR and reward function (Eq. 19).

The key question that the authors sought to answer was which RL strategy best
explains human and rat navigation across the novel test configurations. Across
a wide range of analyses, the authors observed a consistent trend: both human
and rodent behavior was most consistent with the SR agent. Humans also showed
some similarity to the model-based agent, but neither species was consistent
with the model-free agent.

First, the authors simulated each RL agent generatively on the same trials as
the participants—i.e., they let the RL agent navigate and solve each trial as
a kind of simulated participant, learning from its own experience along the