# Proper Value Equivalence

Christopher Grimm
Computer Science & Engineering
University of Michigan
<EMAIL_ADDRESS>&André Barreto, Gregory Farquhar,
David Silver, Satinder Singh
DeepMind
{andrebarreto,gregfar,
<EMAIL_ADDRESS>
###### Abstract

One of the main challenges in model-based reinforcement learning (RL) is to
decide which aspects of the environment should be modeled. The value-
equivalence (VE) principle proposes a simple answer to this question: a model
should capture the aspects of the environment that are relevant for value-
based planning. Technically, VE distinguishes models based on a set of
policies and a set of functions: a model is said to be VE to the environment
if the Bellman operators it induces for the policies yield the correct result
when applied to the functions. As the number of policies and functions
increase, the set of VE models shrinks, eventually collapsing to a single
point corresponding to a perfect model. A fundamental question underlying the
VE principle is thus how to select the smallest sets of policies and functions
that are sufficient for planning. In this paper we take an important step
towards answering this question. We start by generalizing the concept of VE to
order-$k$ counterparts defined with respect to $k$ applications of the Bellman
operator. This leads to a family of VE classes that increase in size as
$k\rightarrow\infty$. In the limit, all functions become value functions, and
we have a special instantiation of VE which we call proper VE or simply PVE.
Unlike VE, the PVE class may contain multiple models even in the limit when
all value functions are used. Crucially, all these models are sufficient for
planning, meaning that they will yield an optimal policy despite the fact that
they may ignore many aspects of the environment. We construct a loss function
for learning PVE models and argue that popular algorithms such as MuZero can
be understood as minimizing an upper bound for this loss. We leverage this
connection to propose a modification to MuZero and show that it can lead to
improved performance in practice.

##  1 Introduction

It has long been argued that, in order for reinforcement learning (RL) agents
to solve truly complex tasks, they must build a model of the environment that
allows for counterfactual reasoning [29]. Since representing the world in all
its complexity is a hopeless endeavor, especially under capacity constraints,
the agent must be able to ignore aspects of the environment that are
irrelevant for its purposes. This is the premise behind the _value
equivalence_ (VE) principle, which provides a formalism for focusing on the
aspects of the environment that are crucial for value-based planning [17].

VE distinguishes models based on a set of policies and a set of real-valued
scalar functions of state (henceforth, just functions). Roughly, a model is
said to be VE to the environment if the Bellman operators it induces for the
policies yield the same result as the environment’s Bellman operators when
applied to the functions. The policies and functions thus become a “language”
to specify which parts of the environment a model should capture. As the
number of policies and functions increase the requirements on the model become
more stringent, which is to say that the class of VE models shrinks. In the
limit, the VE class collapses to a single point corresponding to a perfect
model. Although this result is reassuring, in practice we want to stop short
of collapsing—after all, at this point the agent is no longer ignoring
irrelevant aspects of the environment.

A fundamental question is thus how to select the smallest sets of policies and
functions such that a resulting VE model is sufficient for planning. In this
paper we take an important additional step in this direction: we show that the
VE principle can be formulated with respect to _value functions_ only. This
result drastically reduces the space of functions that must be considered by
VE, as in general only a small fraction of the set of all functions will
qualify as value functions in a given environment. Since every policy has an
associated value function, this new formulation of VE removes the need for
selecting functions, only requiring policies. We name our new formulation
_proper value equivalence_ (PVE) to emphasize its explicit use of value
functions.

PVE has several desirable properties. Unlike with VE, the class of PVE models
does not collapse to a singleton in the limit. This means that, even if _all_
value functions are used, we generally end up with multiple PVE models—which
can be beneficial if some of these are easier to learn or represent than
others. Crucially, all of these models are sufficient for planning, meaning
that _they will yield an optimal policy despite the fact that they may ignore
many aspects of the environment_.

Finally, we make more precise Grimm et al.’s [17] suggestion that the VE
principle may help explain the good empirical performance of several modern
algorithms [38, 33, 24, 12, 30]. Specifically, we show that, with mild
assumptions, minimizing the loss of the MuZero algorithm [31] can be
understood as minimizing a PVE error. We then leverage this connection to
suggest a modification to MuZero and show a small but significant improvement
in the Atari Learning Environment [3].

##  2 Background

The agent’s interaction with the environment will be modeled as a _Markov
decision process_ (MDP)
$\mathcal{M}\equiv\langle\mathcal{S},\mathcal{A},r,p,\gamma\rangle$, where
$\mathcal{S}$ and $\mathcal{A}$ are the state and action spaces, $r(s,a)$ is
the expected reward following taking $a$ from $s$, $p(s^{\prime}|s,a)$ is the
transition kernel and $\gamma\in[0,1)$ is a discount factor [27]. A _policy_
is a mapping $\pi:\mathcal{S}\mapsto\mathcal{P}(\mathcal{A})$, where
$\mathcal{P}(\mathcal{A})$ is the space of probability distributions over
$\mathcal{A}$; we define
$\mathbbl{\Pi}\equiv\\{\pi\,|\,\pi:\mathcal{S}\mapsto\mathcal{P}(\mathcal{A})\\}$
as the set of all possible policies. A policy $\pi$ is _deterministic_ if
$\pi(a|s)>0$ for only one action $a$ per state $s$. A policy’s _value
function_ is defined as

$v_{\pi}(s)\equiv\mathbb{E}_{\pi}\Big{[}\sum_{i=0}^{\infty}\gamma^{i}r(S_{t+i},A_{t+i})\,|\,S_{t}=s\Big{]},$
(1)

where $\mathbb{E}_{\pi}[\cdot]$ denotes expectation over the trajectories
induced by $\pi$ and the random variables $S_{t}$ and $A_{t}$ indicate the
state occupied and the action selected by the agent at time step $t$.

The agent’s goal is to find a policy $\pi\in\mathbbl{\Pi}$ that maximizes the
value of every state [36, 37]. Usually, a crucial step to carry out this
search is to compute the value function of candidate policies. This process
can be cast in terms of the policy’s _Bellman operator_ :

$\mathcal{T}_{\pi}[v](s)\equiv\mathbb{E}_{A\sim\pi(\cdot|s),S^{\prime}\sim
p(\cdot|s,A)}\left[r(s,A)+\gamma v(S^{\prime})\right],$ (2)

where $v$ is any function in the space
$\mathbbl{{V}}\equiv\\{f\,|\,f:\mathcal{S}\mapsto\mathbb{R}\\}$. It is known
that $\lim_{n\rightarrow\infty}\mathcal{T}_{\pi}^{n}v=v_{\pi}$, that is,
starting from any $v\in\mathbbl{{V}}$, the repeated application of
$\mathcal{T}_{\pi}$ will eventually converge to $v_{\pi}$. Since in RL the
agent does not know $p$ and $r$, it cannot apply (2) directly. One solution is
to learn a _model_ $\tilde{m}\equiv(\tilde{r},\tilde{p})$ and use it to
compute (2) with $p$ and $r$ replaced by $\tilde{p}$ and $\tilde{r}$ [36]. We
denote the set of all models as $\mathbbl{{M}}$.

The _value equivalence principle_ defines a model as value equivalent (VE) to
the environment $m^{*}\equiv(r,p)$ with respect to a set of policies $\Pi$ and
a set of functions $\mathcal{V}$ if it produces the same Bellman updates as
$m^{*}$ when using $\Pi$ and $\mathcal{V}$ [17]. Classes of such models are
expressed as follows:

$\mathcal{M}(\Pi,\mathcal{V})\equiv\\{\tilde{m}\in\mathcal{M}:\tilde{\mathcal{T}_{\pi}}v=\mathcal{T}_{\pi}v\
\forall\pi\in\Pi,v\in\mathcal{V}\\}$ (3)

where $\mathcal{M}\subseteq\mathbbl{{M}}$ is a class of models,
$\tilde{\mathcal{T}_{\pi}}$ denotes one application of the Bellman operator
induced by model $\tilde{m}$ and policy $\pi$ to function $v$, and
$\mathcal{T}_{\pi}$ is environment’s Bellman operator for $\pi$.

Grimm et al. [17] showed that the VE principle can be used to learn models
that disregard aspects of the environment which are not related to the task of
interest.111A related approach is taken in value-aware model learning [11]
which minimizes the discrepancy between the Bellman optimality operators
induced by the model and the environment. Classical approaches to model
learning do not take the eventual use of the model into account, potentially
modeling irrelevant aspects of the environment. Accordingly, Grimm et al. [17]
have shown that, under the same capacity constraints, models learned using VE
can outperform their classical counterparts.

##  3 Proper value equivalence

One can define a spectrum of VE classes corresponding to different numbers of
applications of the Bellman operator. We define an order-$k$ VE class as:

$\mathcal{M}^{k}(\Pi,\mathcal{V})\equiv\\{\tilde{m}\in\mathcal{M}:\tilde{\mathcal{T}}_{\pi}^{k}v=\mathcal{T}^{k}_{\pi}v\,\,\forall\pi\in\Pi,v\in\mathcal{V}\\}$
(4)

where $\tilde{\mathcal{T}}_{\pi}^{k}v$ denotes $k$ applications of
$\tilde{\mathcal{T}}_{\pi}$ to $v$. Under our generalized definition of VE,
Grimm et al. [17] studied order-one VE classes of the form
$\mathcal{M}^{1}(\Pi,\mathcal{V})$. They have shown that
$\mathcal{M}^{1}(\mathbbl{\Pi},\mathbbl{V})$ either contains only the
environment or is empty. This is not generally true for $k>1$. The limiting
behavior of order-$k$ value equivalent classes can be described as follows

######  Proposition 1.

Let $\mathcal{V}$ be a set of functions such that if $v\in\mathcal{V}$ then
$\mathcal{T}_{\pi}v\in\mathcal{V}$ for all $\pi\in\Pi$. Then, for
$k,K\in\mathbb{Z}^{+}$ such that $k$ divides $K$, it follows that:

  1. (i)

For any $\mathcal{M}\subseteq\mathbbl{{M}}$ and any
$\Pi\subseteq\mathbbl{\Pi}$, we have that
$\mathcal{M}^{k}(\Pi,\mathcal{V})\subseteq\mathcal{M}^{K}(\Pi,\mathcal{V})$.

  2. (ii)

If $\Pi$ is non-empty and $\mathcal{V}$ contains at least one constant
function, then there exist environments such that
$\mathbbl{{M}}^{k}(\Pi,\mathcal{V})\subset\mathbbl{{M}}^{K}(\Pi,\mathcal{V})$.

We defer all proofs of theoretical results to Appendix A.2. Based on
Proposition 1 we can relate different VE model classes according to the
greatest common divisor of their respective orders; specifically, two classes
$\mathcal{M}^{k}(\Pi,\mathcal{V})$ and $\mathcal{M}^{K}(\Pi,\mathcal{V})$ will
intersect at $\mathcal{M}^{\mathrm{gcd}(k,K)}(\Pi,\mathcal{V})$ (Figure 1).

Figure 1: Topology of the space of order-$k$ VE classes. Given a set of
policies $\Pi$, a set of functions $\mathcal{V}$ closed under Bellman updates,
and $k,K\in\mathbb{Z}^{+}$ such that $k$ divides $K$, we have that
$\mathcal{M}^{k}(\Pi,\mathcal{V})~{}\subseteq~{}\mathcal{M}^{K}(\Pi,\mathcal{V})$.

Proposition 1 also implies that, in contrast to order-one VE classes, higher
order VE classes potentially include multiple models, even if VE is defined
with respect to all policies $\mathbbl{\Pi}$ and all functions $\mathbbl{V}$.
In addition, the size of a VE class cannot decrease as we increase its order
from $k$ to a multiple of $k$ (and in some cases it will strictly increase).
This invites the question of what happens in the limit as we keep increasing
the VE order. To answer this question, we introduce a crucial concept for this
paper:

######  Definition 1.

(Proper value equivalence). Given a set of policies
$\Pi\subseteq\mathbbl{\Pi}$, let

$\mathcal{M}^{\infty}(\Pi)=\lim_{k\to\infty}\mathcal{M}^{k}(\Pi,\mathbbl{V})=\\{\tilde{m}\in\mathcal{M}:\tilde{v}_{\pi}=v_{\pi}\
\forall\pi\in\Pi\\},$ (5)

where $\tilde{v}_{\pi}$ and $v_{\pi}$ are the value functions of $\pi$ induced
by model $\tilde{m}$ and the environment. We say that each
$\tilde{m}\in\mathcal{M}^{\infty}(\Pi)$ is a proper value equivalent model to
the environment with respect to $\Pi$.

Because the process of repeatedly applying a policy’s Bellman operator to a
function converges to the same fixed point regardless of the function, in an
order-$\infty$ VE class the set $\Pi$ uniquely determines the set
$\mathcal{V}$. This reduces the problem of defining $\Pi$ and $\mathcal{V}$ to
defining the former only. Also, since all functions in an order-$\infty$ VE
are _value_ functions, we call it _proper VE_ or PVE.

It is easy to show that Proposition 1 is valid for any $k\in\mathbb{Z}^{+}$
when $K=\infty$ (Corollary 2 in Appendix A.2). Thus, in some sense,
$\mathcal{M}^{\infty}$ is the “biggest” VE class. It is also possible to
define this special VE class in terms of any other:

######  Proposition 2.

For any $\Pi\subseteq\mathbbl{\Pi}$ and any $k\in\mathbb{Z}^{+}$ it follows
that

$\mathcal{M}^{\infty}(\Pi)=\bigcap_{\pi\in\Pi}\mathcal{M}^{k}(\\{\pi\\},\\{v_{\pi}\\}),$
(6)

where $v_{\pi}$ is the value of policy $\pi$ in the environment.

We thus have two equivalent ways to describe the class of models which are PVE
with respect to a set of policies $\Pi$. The first, given in (5), is the
order-$\infty$ limit of value equivalence with respect to $\Pi$ and the set of
all functions $\mathbbl{V}$. The second, given in (6), is the intersection of
the classes of models that are order-$k$ VE with respect to the singleton
policies in $\Pi$ and their respective value functions. This latter form is
valid for _any_ $k$, and will underpin our practical algorithmic
instantiations of PVE.

Setting $k=1$ in Proposition 2 we see that PVE can be written in terms of
order-one VE. This means that $\mathcal{M}^{\infty}$ inherits many of the
topological properties of $\mathcal{M}^{1}$ shown by Grimm et al. [17].
Specifically, we know that
$\mathcal{M}^{\prime\infty}(\Pi)\subseteq\mathcal{M}^{\infty}(\Pi)$ if
$\mathcal{M}^{\prime}\subseteq\mathcal{M}$ and also that
$\mathcal{M}^{\infty}(\Pi^{\prime})\subseteq\mathcal{M}^{\infty}(\Pi)$ when
$\Pi\subseteq\Pi^{\prime}$ (these directly follow from Grimm et al.’s [17]
Properties 1 and 3 respectively).

Proposition 2 also sheds further light into the relation between PVE and
order-$k$ VE more generally. Let $\Pi$ be a set of policies and
$\mathcal{V}_{\pi}$ their value functions. Then, for any $k\in\mathbb{Z}^{+}$,
we have that

$\displaystyle\mathcal{M}^{k}(\Pi,\mathcal{V}_{\pi})$
$\displaystyle=\bigcap_{\pi\in\Pi}\bigcap_{v\in\mathcal{V}_{\pi}}\mathcal{M}^{k}(\\{\pi\\},\\{v\\})\subseteq\bigcap_{\pi\in\Pi}\mathcal{M}^{k}(\\{\pi\\},\\{v_{\pi}\\})=\mathcal{M}^{\infty}(\Pi),$
(7)

which is another way to say that $\mathcal{M}^{\infty}$ is, in some sense, the
largest among all the VE classes. The reason why the size of VE classes is
important is that it directly reflects the main motivation behind the VE
principle. VE’s premise is that models should be constructed taking into
account their eventual use: if some aspects of the environment are irrelevant
for value-based planning, it should not matter whether a model captures them
or not. This means that all models that only differ with respect to these
irrelevant aspects but are otherwise correct qualify as valid VE solutions. A
larger VE class generally means that more irrelevant aspects of the
environment are being ignored by the agent. We now make this intuition more
concrete by showing how irrelevant aspects of the environment that are
eventually captured by order-one VE are always ignored by PVE:

######  Proposition 3.

Let $\Pi\subseteq\mathbbl{\Pi}$. If the environment state can be factored as
$\mathcal{S}=\mathcal{X}\times\mathcal{Y}$ where $|\mathcal{Y}|>1$ and
$v_{\pi}(s)=v_{\pi}((x,y))=v_{\pi}(x)$ for all $\pi\in\Pi$, then
$\mathbbl{{M}}^{1}(\Pi,\mathbbl{V})\subset\mathbbl{{M}}^{\infty}(\Pi)$.

Note that the subset relation appearing in Proposition 3 is strict. We can
think of the variable ‘$y$’ appearing in Proposition 3 as superfluous features
that do not influence the RL task, like the background of an image or any
other sensory data that is irrelevant to the agent’s goal. A model is free to
assign arbitrary dynamics to such irrelevant aspects of the state without
affecting planning performance. Since order-one VE eventually pins down a
model that describes everything about the environment, one would expect the
size of $\mathcal{M}^{\infty}$ relative to $\mathcal{M}^{1}$ to increase as
more superfluous features are added. Indeed, in our proof of Proposition 3 we
construct a set of models in $\mathcal{M}^{\infty}(\mathbbl{\Pi})$ which are
in one-to-one correspondence with $\mathcal{Y}$, confirming this intuition
(see Appendix A.2).

### Proper value equivalence yields models that are sufficient for optimal
planning

In general PVE does not collapse to a single model even in the limit of
$\Pi=\mathbbl{\Pi}$. At first this may cause the impression that one is left
with the extra burden of selecting one among the PVE models. However, it can
be shown that no such choice needs to be made:

######  Proposition 4.

An optimal policy for any $\tilde{m}\in\mathcal{M}^{\infty}(\mathbbl{\Pi})$ is
also an optimal policy in the environment.

According to Proposition 1 any model
$\tilde{m}\in\mathcal{M}^{\infty}(\mathbbl{\Pi})$ used for planning will yield
an optimal policy for the environment. In fact, in the spirit of ignoring as
many aspects of the environment as possible, we can define an even larger PVE
class by focusing on deterministic policies only:

######  Corollary 1.

Let $\mathbbl{\Pi}_{\text{det}}$ be the set of all deterministic policies. An
optimal policy for any
$\tilde{m}\in\mathcal{M}^{\infty}(\mathbbl{\Pi}_{\text{det}})$ is also optimal
in the environment.

Given that both $\mathcal{M}^{\infty}(\mathbbl{\Pi})$ and
$\mathcal{M}^{\infty}(\mathbbl{\Pi}_{\text{det}})$ are sufficient for optimal
planning, one may wonder if these classes are in fact the same. The following
result states that the class of PVE models with respect to deterministic
policies can be strictly larger than its counterpart defined with respect to
all policies:

######  Proposition 5.

There exist environments and model classes for which
$\mathcal{M}^{\infty}(\mathbbl{\Pi})\subset\mathcal{M}^{\infty}(\mathbbl{\Pi}_{\mathrm{det}})$.

Figure 2: An environment / model pair with the same values for all
deterministic policies but not all stochastic policies. The environment has
three states and two actions: $\mathcal{A}=\\{\mathrm{L},\mathrm{R}\\}$. The
percentages in the figure indicate the probability of a given transition and
the corresponding tuples $(r,a)$ indicate the reward associated with a given
action. A deterministic policy cannot dither between $s_{1}$ and $s_{3}$ but a
stochastic policy can. Note that the dynamics between the pair differs when
taking action R from $s_{2}$. This difference will affect the dithering
behavior of such a stochastic policy in a way that results in different model
and environment values.

Figure 2 illustrates Proposition 5 with an example of environment and a model
$\tilde{m}$ such that
$\tilde{m}\in\mathcal{M}^{\infty}(\mathbbl{\Pi}_{\text{det}})$ but
$\tilde{m}\notin\mathcal{M}^{\infty}(\mathbbl{\Pi})$.

To conclude our discussion on models that are sufficient for optimal planning,
we argue that, in the absence of additional information about the environment
or the agent, $\mathcal{M}^{\infty}(\mathbbl{\Pi}_{\text{det}})$ is in fact
the largest possible VE class that is guaranteed to yield optimal performance.
To see why this is so, suppose we remove a single deterministic policy from
$\mathbbl{\Pi}_{\mathrm{det}}$ and pick an arbitrary model
$\tilde{m}\in\mathcal{M}^{\infty}(\mathbbl{\Pi}_{\mathrm{det}}-\\{\pi\\})$.
Let $\tilde{v}_{\pi}$ be the value function of $\pi$ according to the model
$\tilde{m}$. Because $\pi$ is not included in the set of policies used to
enforce PVE, $\tilde{v}_{\pi}$ may not coincide with $v_{\pi}$, the actual
value function of $\pi$ according to the environment. Now, if $\pi$ happens to
be the only optimal policy in the environment and $\tilde{v}_{\pi}$ is not the
optimal value function of $\tilde{m}$, the policy returned by this model will
clearly be sub-optimal.

##  4 Learning a proper value-equivalent model

Having established that we want to find a model in
$\mathcal{M}^{\infty}(\mathbbl{\Pi}_{\mathrm{det}})$, we now turn our
attention to how this can be done in practice. Following Grimm et al. [17],
given a finite set of policies $\Pi$ and a finite set of functions
$\mathcal{V}$, we cast the search for a model
$\tilde{m}\in\mathcal{M}^{k}(\Pi,\mathcal{V})$ as the minimization of
deviations from (4):

$\ell_{\Pi,\mathcal{V}}^{k}(m^{*},\tilde{m})\equiv\sum_{\pi\in\Pi}\sum_{v\in\mathcal{V}}\|\mathcal{T}_{\pi}^{k}v-\tilde{\mathcal{T}}_{\pi}^{k}v||,$
(8)

where $\tilde{\mathcal{T}}_{\pi}$ are Bellman operators induced by $\tilde{m}$
and $\lVert\cdot\rVert$ is a norm.222We can also impose VE with infinite sets
of functions and policies by replacing the respective sums with integrals; in
this case one may consider taking a supremum over VE terms to avoid situations
where VE is not necessarily satisfied on measure 0 sets. Note that setting
$k=\infty$ in (8) yields a loss that requires computing $\tilde{m}$’s value
function—which is impractical to do if $\tilde{m}$ is being repeatedly
updated. Thankfully, by leveraging the connection between order-$k$ VE and PVE
given in Proposition 2, we can derive a practical PVE loss:

$\ell_{\Pi,\infty}^{k}(m^{*},\tilde{m})\equiv\sum_{\pi\in\Pi}\|\mathcal{T}_{\pi}^{k}v_{\pi}-\tilde{\mathcal{T}}_{\pi}^{k}v_{\pi}||=\sum_{\pi\in\Pi}\|v_{\pi}-\tilde{\mathcal{T}}_{\pi}^{k}v_{\pi}||.$
(9)

Interestingly, given a set of policies $\Pi$, minimizing (9) for any $k$ will
result in a model $\tilde{m}\in\mathcal{M}^{\infty}(\Pi)$ (cf. Proposition 2).
As we will discuss shortly, this property can be exploited to generate
multiple loss functions that provide a richer learning signal in practical
scenarios.

Contrasting loss functions (8) and (9) we observe an important fact: unlike
with other order-$k$ VE classes, PVE requires actual value functions to be
enforced in practice. Since value functions require data and compute to be
obtained, it is reasonable to ask whether the benefits of PVE justify the
associated additional burden. Concretely, one may ask whether the sample
transitions and computational effort spent in computing the value functions to
be used with PVE would not be better invested in enforcing other forms of VE
over arbitrary functions that can be readily obtained.

We argue that in many cases one does not have to choose between order-$k$ VE
and PVE. Value-based RL algorithms usually compute value functions
iteratively, generating a sequence of functions $v_{1},v_{2},...$ which will
eventually converge to $\tilde{v}_{\pi}$ for some $\pi$. A model-based
algorithm that computes $\tilde{v}_{\pi}$ in this way has to somehow
interleave this process with the refinement of the model $\tilde{m}$. When it
comes to VE, one extreme solution is to only use the final approximation
$\tilde{v}_{\pi}\approx v_{\pi}$ in an attempt to enforce PVE through (9). It
turns out that, as long as the sequence $v_{1},v_{2},...$ is approaching
$v_{\pi}$, one can use _all_ the functions in the sequence to enforce PVE with
respect to $\pi$. Our argument is based on the following result:

######  Proposition 6.

For any $\pi\in\mathbbl{\Pi}$, $v\in\mathbbl{{V}}$ and $k,n\in\mathbb{Z}^{+}$,
we have that

$\lVert
v_{\pi}-\tilde{\mathcal{T}}_{\pi}^{k}v_{\pi}\rVert_{\infty}\leq(\gamma^{k}+\gamma^{n})\underbrace{\lVert
v_{\pi}-v\rVert_{\infty}}_{\epsilon_{v}}+\underbrace{\lVert\mathcal{T}_{\pi}^{n}v-\tilde{\mathcal{T}}_{\pi}^{k}v\rVert_{\infty}}_{\epsilon_{ve}}.$
(10)

Note that the left-hand side of (10) corresponds to one of the terms of the
PVE loss (9) associated with a given $\pi$. This means that, instead of
minimizing this quantity directly, one can minimize the upper-bound on the
right-hand side of (10). The first term in this upper bound, $\epsilon_{v}$,
is the conventional value-function approximation error that most value-based
methods aim to minimize (either directly or indirectly). The second term,
$\epsilon_{ve}$, is similar to the terms appearing in the order-$k$ VE loss
(8), except that here the number of applications of $\mathcal{T}_{\pi}$ and of
its approximation $\tilde{\mathcal{T}}_{\pi}$ do not have to coincide.

All the quantities appearing in $\epsilon_{ve}$ are readily available or can
be easily approximated using sample transitions [17]. Thus, $\epsilon_{ve}$
can be used to refine the model $\tilde{m}$ using functions $v$ that are not
necessarily value functions. As $v\rightarrow v_{\pi}$, two things happen.
First, $\epsilon_{ve}$ approaches one of the terms of the PVE loss (9)
associated with policy $\pi$. Second, $\epsilon_{v}$ vanishes. Interestingly,
the importance of $\epsilon_{v}$ also decreases with $n$ and $k$, the number
of times $\mathcal{T}_{\pi}$ and $\tilde{\mathcal{T}}_{\pi}$ are applied in
$\epsilon_{ve}$, respectively. This makes sense: since
$\mathcal{T}_{\pi}^{n}v\rightarrow v_{\pi}$ as $n\rightarrow\infty$ and, by
definition, VE approaches PVE as $k\rightarrow\infty$, we have that
$\epsilon_{ve}$ approaches the left-hand side of (10) as both $n$ and $k$
grow.

### An extended example: MuZero through the lens of value equivalence

Grimm et al. [17] suggested that the VE principle might help to explain the
empirical success of recent RL algorithms like Value Iteration Networks, the
Predictron, Value Prediction Networks, TreeQN, and MuZero [38, 33, 24, 12,
31]. In this section we investigate this hypothesis further and describe a
possible way to interpret one of these algorithms, MuZero, through the lens of
VE. We acknowledge that the derivation that follows abstracts away many
details of MuZero and involves a few approximations of its mechanics, but we
believe it captures and explains the algorithm’s essence.

MuZero is a model-based RL algorithm that achieved state-of-the-art
performance across both board games, such as Chess and Go, and Atari 2600
games [31]. The model $\tilde{m}$ in MuZero is trained on sequences of states,
actions and rewards resulting from executing a “behavior policy” in the
environment: $s_{t:t+n+K},a_{t:t+n+K},r_{t:t+n+K}$ where $n$ and $K$ are
hyperparameters of the agent which will be explained shortly. The agent
produces an “agent state” $z^{0}_{t}$ from $s_{t}$ and subsequently generates
$z^{1:K}_{t}$ by using its model to predict the next $K$ agent states
following actions $a_{t:t+K}$. The agent also maintains reward and value
function estimates as a function of agent states, which we denote
$\tilde{r}(z)$ and $v(z)$ respectively. A variant333In reality MuZero uses a
categorical representation for its value and reward functions and minimizes
them using a cross-entropy objective. We argue that this choice is not
essential to its underlying ideas and use scalar representations with a
squared loss to simplify our analysis. of MuZero’s per-state model loss can
thus be expressed as:

$\ell^{\mu}(s_{t})=\sum_{k=0}^{K}(V_{t+k}-v(z^{k}_{t}))^{2}+(r_{t+k}-\tilde{r}(z_{t}^{k})))^{2}$
(11)

where
$V_{t+k}=r_{t+k}+\cdots+\gamma^{n-1}r_{t+k+n-1}+\gamma^{n}v^{targ}(z^{0}_{t+k+n})$.
The term $v^{targ}$ is a value target produced by Monte-Carlo tree search
(MCTS, [7]). Because the behavior policy is itself computed via MCTS, we have
that $v^{targ}\approx v$; for simplicity we will assume that $v^{targ}=v$ and
only use $v$.

In what follows we show, subject to a modest smoothness assumption, that
minimizing MuZero’s loss with respect to its behavior policy, $\pi$, also
minimizes a corresponding PVE loss. Put precisely:

$C\cdot\mathbb{E}_{d_{\pi}}[\ell^{\mu}(S_{t})]\geq\big{(}\ell^{K}_{\\{\pi\\},\infty}(m^{*},\tilde{m})\big{)}^{2}$
(12)

for some $C>0$, where $d_{\pi}$ is a stationary distribution. We proceed by
combining two derivations: a lower-bound on
$\mathbb{E}_{d_{\pi}}[\ell^{\mu}(S_{t})]$ in (15), and an upper-bound on
$(\ell^{K}_{\\{\pi\\},\infty}(m^{*},\tilde{m}))^{2}$ in (17).

As a preliminary step we note that $\ell^{\mu}(s_{t})$ and
$\ell^{K}_{\\{\pi\\},\infty}(m^{*},\tilde{m})$ are expressed in terms of
samples and expectations respectively. We note the following connection
between these quantities:

$\displaystyle\mathbb{E}[r_{t+k}|s_{t}]=\mathcal{P}^{k}_{\pi}[r_{\pi}](s_{t}),$
$\displaystyle\mathbb{E}[V_{t+k}|s_{t}]=\mathcal{P}^{k}_{\pi}\mathcal{T}^{n}_{\pi}[v_{\pi}](s_{t}),$
(13)
$\displaystyle\quad\mathbb{E}[\tilde{r}(z^{k}_{t})|s_{t}]=\tilde{\mathcal{P}}^{k}_{\pi}[\tilde{r}_{\pi}](s_{t}),$
$\displaystyle\mathbb{E}[v(z^{k}_{t})|s_{t}]=\tilde{\mathcal{P}}^{k}_{\pi}[v_{\pi}](s_{t}),$

where $\mathcal{P}^{k}_{\pi}$ is the $k$-step environment transition operator
under policy $\pi$:
$\mathcal{P}^{k}_{\pi}[x](s_{t})=\mathbb{E}[x(S_{t+k})|s_{t},m^{*},\pi]$,
$r_{\pi}(s)=\mathbb{E}_{A\sim\pi}[r(s,A)]$ and $\tilde{\mathcal{P}}^{k}_{\pi}$
and $\tilde{r}_{\pi}$ are the corresponding quantities using the model instead
of the environment. The above expectations are taken with respect to the
environment or model and $\pi$ as appropriate. We now derive our lower-bounds
on $\mathbb{E}_{d_{\pi}}[\ell^{\mu}(S_{t})]$:

$\displaystyle\mathbb{E}_{d_{\pi}}[\ell^{\mu}(S_{t})]$
$\displaystyle=\mathbb{E}_{d_{\pi}}\Big{[}\sum_{k=0}^{K}\mathbb{E}[\,(V_{t+k}-v(z_{t}^{k}))^{2}\mid
S_{t}\,]+\sum_{k=0}^{K}\mathbb{E}[\,(r_{t+k}-\tilde{r}(z_{t}^{k}))^{2}\mid
S_{t}\,]\Big{]}$ (14)
$\displaystyle\geq\mathbb{E}_{d_{\pi}}\Big{[}\sum_{k=0}^{K}(\mathbb{E}[\,V_{t+k}\mid
S_{t}\,]-\mathbb{E}[\,v(z^{k}_{t})\mid
S_{t}\,])^{2}+\sum_{k=0}^{K}(\mathbb{E}[\,r_{t+k}\mid
S_{t}\,]-\mathbb{E}[\,\tilde{r}(z^{k}_{t})\mid S_{t}\,])^{2}\Big{]}$
$\displaystyle=\sum_{k=0}^{K}\mathbb{E}_{d_{\pi}}\Big{[}(\mathcal{P}^{k}_{\pi}\mathcal{T}^{n}_{\pi}v(S_{t})-\tilde{\mathcal{P}}^{k}_{\pi}v(S_{t}))^{2}\Big{]}+\sum_{k=0}^{K}\mathbb{E}_{d_{\pi}}\Big{[}(\mathcal{P}^{k}_{\pi}r_{\pi}(S_{t})-\tilde{\mathcal{P}}^{k}_{\pi}\tilde{r}_{\pi}(S_{t}))^{2}\Big{]}$

where we apply the tower-property, Jensen’s inequality and the identities in
(13). We write the expression using norms and drop all terms except
$k\in\\{0,K\\}$ in the first sum to obtain:

$\mathbb{E}_{d_{\pi}}[\ell^{\mu}(S_{t})]\geq\|\mathcal{T}^{n}_{\pi}v-v\|^{2}_{d_{\pi}}+\|\mathcal{P}^{K}_{\pi}\mathcal{T}^{n}_{\pi}v-\tilde{\mathcal{P}}^{K}_{\pi}v\|^{2}_{d_{\pi}}+\sum_{k=0}^{K}\|\mathcal{P}^{k}_{\pi}r_{\pi}-\tilde{\mathcal{P}}^{k}_{\pi}\tilde{r}_{\pi}\|^{2}_{d_{\pi}}$
(15)

recalling that
$\|x-y\|_{d_{\pi}}^{2}=\mathbb{E}_{d_{\pi}}[(x(S_{t})-y(S_{t}))^{2}]$. To
derive an upper-bound for $(\ell^{K}_{\\{\pi\\},\infty}(m^{*},\tilde{m}))^{2}$
we assume that the error in value estimation is smooth in the sense that there
is some $g>0$ (independent of $v$) such that
$\|v-v_{\pi}\|_{\infty}<g\cdot\|v-v_{\pi}\|_{d_{\pi}}$. We can then use a
modified version of (10) for the $d_{\pi}$-weighted $\ell_{2}$-norm (see
Appendix A.2), plugging in $n+K$ and $K$:

$\displaystyle\|v_{\pi}-\tilde{\mathcal{T}}^{K}_{\pi}v_{\pi}\|_{d_{\pi}}$
$\displaystyle\leq(g\gamma^{K}+\gamma^{n+K})\|v_{\pi}-v\|_{d_{\pi}}+\|\mathcal{T}^{K+n}_{\pi}v-\tilde{\mathcal{T}}^{K}_{\pi}v\|_{d_{\pi}}$
(16)
$\displaystyle\leq\gamma^{K}(g+\gamma^{n})\|v_{\pi}-v\|_{d_{\pi}}+\|\mathcal{P}^{K}_{\pi}\mathcal{T}^{n}_{\pi}v-\tilde{\mathcal{P}}^{K}_{\pi}v\|_{d_{\pi}}+\sum_{k=0}^{K}\|\mathcal{P}^{k}_{\pi}r_{\pi}-\tilde{\mathcal{P}}_{\pi}^{k}\tilde{r}_{\pi}\|_{d_{\pi}}$
$\displaystyle\leq\gamma^{K}\frac{(g+\gamma^{n})}{(1-\gamma^{n})}\|\mathcal{T}^{n}_{\pi}v-v\|_{d_{\pi}}+\|\mathcal{P}^{K}_{\pi}\mathcal{T}^{n}_{\pi}v-\tilde{\mathcal{P}}^{K}_{\pi}v\|_{d_{\pi}}+\sum_{k=0}^{K}\|\mathcal{P}^{k}_{\pi}r_{\pi}-\tilde{\mathcal{P}}_{\pi}^{k}\tilde{r}_{\pi}\|_{d_{\pi}},$

from here we can square both sides and apply Jensen’s inequality,

$\displaystyle\|v_{\pi}-\tilde{\mathcal{T}}_{\pi}^{K}v_{\pi}\|^{2}_{d_{\pi}}\leq
ab\|\mathcal{T}^{\pi}_{n}v-v\|^{2}_{d_{\pi}}+b\|\mathcal{P}^{K}_{\pi}\mathcal{T}^{n}_{\pi}v-\tilde{\mathcal{P}}^{K}_{\pi}v\|^{2}_{d_{\pi}}+b\sum_{k=0}^{K}\|\mathcal{P}^{k}_{\pi}r_{\pi}-\tilde{\mathcal{P}}_{\pi}^{k}\tilde{r}_{\pi}\|^{2}_{d_{\pi}},$
(17)

where $a=\gamma^{K}(g+\gamma^{n})(1-\gamma^{n})^{-1}$ and $b=a+K+2$. Combining
(17) and (15) we obtain:

$ab\cdot\mathbb{E}_{d_{\pi}}[\ell^{\mu}(S_{t})]\geq\|v_{\pi}-\tilde{\mathcal{T}}_{\pi}^{K}v_{\pi}\|_{d_{\pi}}^{2}=\big{(}\ell^{K}_{\\{\pi\\},\infty}(m^{*},\tilde{m})\big{)}^{2},$
(18)

thus minimizing MuZero’s loss minimizes a squared PVE loss with respect to a
single policy.

##  5 Experiments

|  |  |
---|---|---|---
|  |  |

Figure 4: All scatter plots are generated by tracking the training progress
over 500,000 iterations of models with different order-$k$ VE objectives. In
each plot 120 models were tracked; at every 1000 timesteps the full set of
models is converted into vector form and projected onto their first two
principal components before being plotted (details in the appendix). Top row:
points are colored according to their progress through training. Bottom row:
points are colored according to the average value of the associated model’s
optimal policy on the environment. Rightmost plot: line-plot of the diameters
of these scatter plots against the model-class order.

We first provide results from tabular experiments on a stochastic version of
the Four Rooms domain which serve to corroborate our theoretical claims. Then,
we present results from experiments across the full Atari 57 benchmark [3]
showcasing that the insights from studying PVE and its relationship to MuZero
can provide a benefit in practice at scale. See Appendix A.3 for a full
account of our experimental procedure.

In Section 3 we described the topological relationships between order-$k$ and
PVE classes. This is summarized by Proposition 1, which shows that, for
appropriately defined $\mathcal{V}$ and $\Pi$,
$\mathcal{M}^{k}\subseteq\mathcal{M}^{K}$ if $K$ is a multiple of $k$. We
illustrate this property empirically by randomly initializing a set of models
and then using (8) (or (9) for the limiting case of $k=\infty$) to iteratively
update them towards $\mathcal{M}^{k}(\mathbbl{\Pi},\mathbbl{{V}})$, with
$k\in\\{1,30,40,50,60,\infty\\}$. We take the vectors representing these
models and project them onto their first two principal components in order to
visualise their paths through learning. The results are shown on the top row
of Figure 4. In accordance with the theory, we see that the space of converged
models, represented with the brightest yellow regions, grows with $k$. This
trend is summarised in the rightmost plot, which shows the diameter of the
scatter plots for each $k$. In the bottom row of Figure 4 we use color to show
the value that the optimal policy of each model achieves in the true
environment. As predicted by our theory, the space of models that are
sufficient for optimal planning also grows with $k$.

|
---|---

(a)

(b)

Figure 6: (a) Comparison of the performance of optimal policies obtained from
capacity constrained models trained to be in
$\mathcal{M}^{\infty}(\mathbbl{\Pi}_{\mathrm{det}})$ and
$\mathcal{M}^{\infty}(\mathbbl{\Pi})$. For each action $a\in\mathcal{A}$, the
transition dynamics $\tilde{P}_{a}$ is constrained to have a rank of at most
$k$. The red dashed line represents the performance of the optimal environment
policy. (b) Trajectories starting from the bottom-right state (red dot)
sampled from the optimal environment policy in both the environment and a PVE
model. Note the numerous diagonal transitions in the PVE model which are not
permitted in the environment.

Model classes containing many models with optimal planning performance are
particularly advantageous when the set of models that an agent can represent
is restricted, since the larger the set of suitable models the greater the
chance of an overlap between this set and the set of models representable by
the agent. Proposition 4 and Corollary 1 compared
$\mathcal{M}^{\infty}(\mathbbl{\Pi})$ and
$\mathcal{M}^{\infty}(\mathbbl{\Pi}_{\mathrm{det}})$, showing that, although
any model in either class is sufficient for planning,
$\mathcal{M}^{\infty}(\mathbbl{\Pi})\subseteq\mathcal{M}^{\infty}(\mathbbl{\Pi}_{\mathrm{det}})$.
This suggests that it might be better to learn a model in
$\mathcal{M}^{\infty}(\mathbbl{\Pi}_{\mathrm{det}})$ when the agent has
limited capacity. We illustrate that this is indeed the case in Figure 6b. We
progressively restrict the space of models that the agent can represent and
attempt to learn models in either $\mathcal{M}^{\infty}(\mathbbl{\Pi})$ or
$\mathcal{M}^{\infty}(\mathbbl{\Pi}_{\mathrm{det}})$. Indeed, we find that the
larger class, $\mathcal{M}^{\infty}(\mathbbl{\Pi}_{\mathrm{det}})$, yields
superior planning performance as agent capacity decreases.

Given their importance, we provide intuition on the ways that individual PVE
models differ from the environment. In Figure 6a we compare trajectories
starting at the same initial state (denoted by a red-circle) from the optimal
environment policy in both the environment and in a randomly sampled model
from $\mathcal{M}^{\infty}(\mathbbl{\Pi})$. In the PVE model there are
numerous diagonal transitions not permitted by the environment. Note that
while the PVE model has very different dynamics than the environment, these
differences must “balance out”, as it still has the same values under any
policy as the environment.

Figure 7: Comparison of our proposed modification to MuZero with an unmodified
baseline.

In Section 4 we showed that minimizing Muzero’s loss function is analogous to
minimizing an upper-bound on a PVE loss (9) with respect to the agent’s
current policy $\pi$—which corresponds to finding a model in
$\mathcal{M}^{\infty}(\Pi)$ where $\Pi=\\{\pi\\}$. Note that our guarantee on
the performance of PVE models (Corollary 1) holds when $\Pi$ contains all
deterministic policies. While it is not feasible to enforce
$\Pi=\mathbbl{\Pi}_{\mathrm{det}}$, we can use previously seen policies by
augmenting the MuZero algorithm with a buffer of past policies and their
approximate value functions (we do so by periodically storing the
corresponding parameters). We can then add an additional loss to MuZero with
the form of the original value loss but using the past value functions. We
still use sampled rewards to construct value targets, but use the stored
policies to compute off-policy corrections using V-trace [9].

To test this proposal we use an on policy (i.e., without a replay buffer)
implementation of MuZero run for 500M frames (as opposed to 20B frames in the
online result of [31]) on the Atari 57 benchmark and find that using our
additional loss yields an advantage in the human normalised median performance
shown in Figure 7. MuZero’s data efficiency can also be improved with the aid
of a replay buffer of trajectories from past policies [32], which may also
capture some of the advantages of expanding the set of policies used for PVE.

##  6 Related work

Our work is closely related to the value-aware model learning (VAML, IterVAML,
[11, 10]) which learns models to minimize the discrepancy between their own
Bellman optimality operators and the environment’s on the optimal value
function—an a priori unknown quantity. To handle this VAML specifies a family
of potential value functions and minimizes the worst-case discrepancy across
them, whereas IterVAML minimizes the discrepancy with respect to model’s
current estimate of the value function in a value iteration inspired scheme.
PVE and the VAML family are complementary works, with VAML addressing its
induced optimization problems and PVE addressing its induced model classes.
Both, however, advocate for learning models with their eventual use in mind—a
view that is aligned with many criticisms of the maximum likelihood objective
for model learning [11, 21, 2, 23]).

It is worth mentioning the relationship between PVE and TD models [35] which,
for a given policy, defines any $R\in\mathbb{R}^{|\mathcal{S}|}$ and
$P\in\mathbb{R}^{|\mathcal{S}|\times|\mathcal{S}|}$ with
$\lim_{k\to\infty}P^{k}=0$ as a valid model if $V=R+P^{\top}V$ where
$V\in\mathbb{R}^{|\mathcal{S}|}$ represents $v_{\pi}$. Clearly all models in
$\mathcal{M}^{\infty}(\\{\pi\\})$ are valid models, however, since $P$ is not
necessarily a transition matrix, the converse does not hold. While TD models
are restricted to prediction rather than control, their generality warrants
further inquiry.

Order-$k$ and PVE model classes form equivalences between MDPs and thus can be
situated among other equivalence notions which can be formulated as state-
aggregations [8, 26, 25, 16, 28, 13, 34, 39, 5, 40]. As pointed out by Grimm
et al. [17], the interaction between arbitrary state-aggregation and models
can be captured with special cases of order-one VE. Our extension of higher-
order VEs potentially offers the possibility of “blending” existing notions of
state aggregation with PVE.

A notable instance of state-aggregation is bisimulation [22], which uses a
relation to aggregate states that have the same immediate rewards and
transition dynamics into other aggregated states. Bisimulation metrics [13]
provide smooth measures of how closely pairs of states satisfy bisimulation
relations. These concepts have become increasingly popular in deep
reinforcement learning where they are used to guide the learning of effective
representations [43, 42, 15, 1]. Although both bisimulation and PVE provide
direction for learning internal aspects of an agent, they are fundamentally
different in their purview—bisimulation concerns representation learning,
while PVE concerns the learning of models given a representation of state.

Beyond bisimulation, representation learning has a wide literature [41, 20, 6,
14, 4] including several modern works which explicitly study the conjunction
of model learning with state representation [44, 42, 15]. These are further
complemented by efforts to learn state representations and models jointly in
the service of value-based planning [12, 33, 24, 18, 31, 38].

##  7 Conclusion and future work

We extended the value equivalence principle by defining a spectrum of
order-$k$ VE sets in which models induce the same $k$-step Bellman operators
as the environment. We then explored the topology of the resulting equivalence
classes and defined the limiting class when $k\to\infty$ as PVE. If a model is
PVE to the environment with respect to a set of policies $\Pi$, then the value
functions of all policies in $\Pi$ are the same in the environment and the
model. The fact that PVE classes can be defined using only a set of policies
eliminates the need for specifying a set of functions to induce VE—resolving a
fundamental issue left open by Grimm et al. [17]. Importantly, we showed that
being PVE with respect to all deterministic policies is sufficient for a model
to plan optimally in the environment. In the absence of additional
information, this is the largest possible VE class that yields optimal
planning. On the practical side, we showed how the MuZero algorithm can be
understood as minimizing an upper bound on a PVE loss, and leveraged this
insight to improve the algorithm’s performance.

Though our efforts have advanced the understanding of value equivalence and
proven useful algorithmically, there is still work to be done in developing a
VE theory whose assumptions hold in practice. This remaining work can be
broadly grouped into two areas (1) understanding the role of approximation in
VE and (2) establishing performance guarantees for VE models with arbitrary
sets of policies and functions. We leave these as future work.

Acknowledgements

We thank Angelos Filos and Sonya Kotov for many thought-provoking discussions.
Christopher Grimm’s work was made possible by the support of the Lifelong
Learning Machines (L2M) grant from the Defense Advanced Research Projects
Agency. Any opinions, findings, conclusions, or recommendations expressed here
are those of the authors and do not necessarily reflect the views of the
sponsors.

## References

  * Agarwal et al. [2021] Rishabh Agarwal, Marlos C Machado, Pablo Samuel Castro, and Marc G Bellemare.  Contrastive behavioral similarity embeddings for generalization in reinforcement learning.  _arXiv preprint arXiv:2101.05265_ , 2021.
  * Ayoub et al. [2020] Alex Ayoub, Zeyu Jia, Csaba Szepesvari, Mengdi Wang, and Lin Yang.  Model-based reinforcement learning with value-targeted regression.  In _International Conference on Machine Learning_ , pages 463–474. PMLR, 2020.
  * Bellemare et al. [2013] Marc G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling.  The arcade learning environment: An evaluation platform for general agents.  _Journal of Artificial Intelligence Research_ , 47:253–279, 2013.
  * Biza et al. [2020] Ondrej Biza, Robert Platt, Jan-Willem van de Meent, and Lawson LS Wong.  Learning discrete state abstractions with deep variational inference.  _arXiv preprint arXiv:2003.04300_ , 2020.
  * Castro [2020] Pablo Samuel Castro.  Scalable methods for computing state similarity in deterministic markov decision processes.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 34, pages 10069–10076, 2020.
  * Corneil et al. [2018] Dane Corneil, Wulfram Gerstner, and Johanni Brea.  Efficient model-based deep reinforcement learning with variational state tabulation.  In _International Conference on Machine Learning_ , pages 1049–1058. PMLR, 2018.
  * Coulom [2006] Rémi Coulom.  Efficient selectivity and backup operators in monte-carlo tree search.  In _International conference on computers and games_ , pages 72–83. Springer, 2006.
  * Dean and Givan [1997] Thomas Dean and Robert Givan.  Model minimization in markov decision processes.  In _AAAI/IAAI_ , pages 106–111, 1997.
  * Espeholt et al. [2018] Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, Iain Dunning, et al.  Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures.  In _International Conference on Machine Learning_ , pages 1407–1416. PMLR, 2018.
  * Farahmand [2018] Amir-massoud Farahmand.  Iterative value-aware model learning.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , pages 9090–9101, 2018.
  * Farahmand et al. [2017] Amir-Massoud Farahmand, André Barreto, and Daniel Nikovski.  Value-Aware Loss Function for Model-based Reinforcement Learning.  In _Proceedings of the International Conference on Artificial Intelligence and Statistics (AISTATS)_ , volume 54, pages 1486–1494, 2017.
  * Farquhar et al. [2018] G Farquhar, T Rocktäschel, M Igl, and S Whiteson.  Treeqn and atreec: Differentiable tree-structured models for deep reinforcement learning.  In _6th International Conference on Learning Representations, ICLR 2018-Conference Track Proceedings_ , volume 6. ICLR, 2018.
  * Ferns et al. [2004] Norm Ferns, Prakash Panangaden, and Doina Precup.  Metrics for finite markov decision processes.  In _UAI_ , volume 4, pages 162–169, 2004.
  * François-Lavet et al. [2019] Vincent François-Lavet, Yoshua Bengio, Doina Precup, and Joelle Pineau.  Combined reinforcement learning via abstract representations.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 33, pages 3582–3589, 2019.
  * Gelada et al. [2019] Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Marc G Bellemare.  Deepmdp: Learning continuous latent space models for representation learning.  In _International Conference on Machine Learning_ , pages 2170–2179. PMLR, 2019.
  * Givan et al. [2003] Robert Givan, Thomas Dean, and Matthew Greig.  Equivalence notions and model minimization in markov decision processes.  _Artificial Intelligence_ , 147(1-2):163–223, 2003.
  * Grimm et al. [2020] Christopher Grimm, Andre Barreto, Satinder Singh, and David Silver.  The value equivalence principle for model-based reinforcement learning.  _Advances in Neural Information Processing Systems_ , 33, 2020.
  * Hessel et al. [2021a] Matteo Hessel, Ivo Danihelka, Fabio Viola, Arthur Guez, Simon Schmitt, Laurent Sifre, Theophane Weber, David Silver, and Hado van Hasselt.  Muesli: Combining improvements in policy optimization.  _arXiv preprint arXiv:2104.06159_ , 2021a.
  * Hessel et al. [2021b] Matteo Hessel, Manuel Kroiss, Aidan Clark, Iurii Kemaev, John Quan, Thomas Keck, Fabio Viola, and Hado van Hasselt.  Podracer architectures for scalable Reinforcement Learning.  _CoRR_ , abs/2104.06272, 2021b.
  * Igl et al. [2018] Maximilian Igl, Luisa Zintgraf, Tuan Anh Le, Frank Wood, and Shimon Whiteson.  Deep variational reinforcement learning for pomdps.  In _International Conference on Machine Learning_ , pages 2117–2126. PMLR, 2018.
  * Joseph et al. [2013] Joshua Joseph, Alborz Geramifard, John W Roberts, Jonathan P How, and Nicholas Roy.  Reinforcement learning with misspecified model classes.  In _2013 IEEE International Conference on Robotics and Automation_ , pages 939–946. IEEE, 2013.
  * [22] Robin Milner.  _Communication and concurrency_ , volume 84.
  * Modi et al. [2020] Aditya Modi, Nan Jiang, Ambuj Tewari, and Satinder Singh.  Sample complexity of reinforcement learning using linearly combined model ensembles.  In _International Conference on Artificial Intelligence and Statistics_ , pages 2010–2020. PMLR, 2020.
  * Oh et al. [2017] Junhyuk Oh, Satinder Singh, and Honglak Lee.  Value prediction network.  In _Advances in Neural Information Processing Systems_ , pages 6118–6128, 2017.
  * Poupart and Boutilier [2013] Pascal Poupart and Craig Boutilier.  Value-directed belief state approximation for POMDPs.  _CoRR_ , abs/1301.3887, 2013.  URL http://arxiv.org/abs/1301.3887.
  * Poupart et al. [2003] Pascal Poupart, Craig Boutilier, et al.  Value-directed compression of pomdps.  _Advances in neural information processing systems_ , pages 1579–1586, 2003.
  * Puterman [1994] Martin L. Puterman.  _Markov Decision Processes—Discrete Stochastic Dynamic Programming_.  John Wiley & Sons, Inc., 1994.
  * Ravindran and Barto [2004] Balaraman Ravindran and Andrew G Barto.  Approximate homomorphisms: A framework for non-exact minimization in markov decision processes.  2004\.
  * Russell and Norvig [2003] Stuart J. Russell and Peter Norvig.  _Artificial Intelligence: A Modern Approach_.  Pearson Education, 3 edition, 2003.
  * Schrittwieser et al. [2019] Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al.  Mastering atari, go, chess and shogi by planning with a learned model.  _arXiv preprint arXiv:1911.08265_ , 2019.
  * Schrittwieser et al. [2020] Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al.  Mastering atari, go, chess and shogi by planning with a learned model.  _Nature_ , 588(7839):604–609, 2020.
  * Schrittwieser et al. [2021] Julian Schrittwieser, Thomas Hubert, Amol Mandhane, Mohammadamin Barekatain, Ioannis Antonoglou, and David Silver.  Online and offline reinforcement learning by planning with a learned model.  _arXiv preprint arXiv:2104.06294_ , 2021.
  * Silver et al. [2017] David Silver, Hado van Hasselt, Matteo Hessel, Tom Schaul, Arthur Guez, Tim Harley, Gabriel Dulac-Arnold, David Reichert, Neil Rabinowitz, Andre Barreto, et al.  The predictron: End-to-end learning and planning.  In _Proceedings of the 34th International Conference on Machine Learning-Volume 70_ , pages 3191–3199. JMLR. org, 2017.
  * Spencer et al. [2009] John P Spencer, Michael SC Thomas, and JL McClelland.  Toward a unified theory of development.  _JP Spencer, MSC Thomas, & JL McClelland (Eds.)_, pages 86–118, 2009.
  * Sutton [1995] Richard S. Sutton.  TD models: Modeling the world at a mixture of time scales.  In _Proceedings of the Twelfth International Conference on Machine Learning_ , pages 531–539, 1995.
  * Sutton and Barto [2018] Richard S. Sutton and Andrew G. Barto.  _Reinforcement Learning: An Introduction_.  MIT Press, 2018.  URL https://mitpress.mit.edu/books/reinforcement-learning-second-edition.  2nd edition.
  * Szepesvári [2010] Csaba Szepesvári.  _Algorithms for Reinforcement Learning_.  Synthesis Lectures on Artificial Intelligence and Machine Learning. Morgan & Claypool Publishers, 2010.
  * Tamar et al. [2016] Aviv Tamar, Yi Wu, Garrett Thomas, Sergey Levine, and Pieter Abbeel.  Value iteration networks.  In _Advances in Neural Information Processing Systems_ , pages 2154–2162, 2016.
  * Taylor et al. [2008] Jonathan Taylor, Doina Precup, and Prakash Panagaden.  Bounding performance loss in approximate mdp homomorphisms.  _Advances in Neural Information Processing Systems_ , 21:1649–1656, 2008.
  * van der Pol et al. [2020] Elise van der Pol, Thomas Kipf, Frans A Oliehoek, and Max Welling.  Plannable approximations to mdp homomorphisms: Equivariance under actions.  In _Proceedings of the 19th International Conference on Autonomous Agents and MultiAgent Systems_ , pages 1431–1439, 2020.
  * Watter et al. [2015] Manuel Watter, Jost Tobias Springenberg, Joschka Boedecker, and Martin Riedmiller.  Embed to control: a locally linear latent dynamics model for control from raw images.  In _Proceedings of the 28th International Conference on Neural Information Processing Systems-Volume 2_ , pages 2746–2754, 2015.
  * Zhang et al. [2019a] Amy Zhang, Zachary C Lipton, Luis Pineda, Kamyar Azizzadenesheli, Anima Anandkumar, Laurent Itti, Joelle Pineau, and Tommaso Furlanello.  Learning causal state representations of partially observable environments.  _arXiv preprint arXiv:1906.10437_ , 2019a.
  * Zhang et al. [2020] Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine.  Learning invariant representations for reinforcement learning without reconstruction.  _arXiv preprint arXiv:2006.10742_ , 2020.
  * Zhang et al. [2019b] Marvin Zhang, Sharad Vikram, Laura Smith, Pieter Abbeel, Matthew Johnson, and Sergey Levine.  Solar: Deep structured representations for model-based reinforcement learning.  In _International Conference on Machine Learning_ , pages 7444–7453. PMLR, 2019b.

## Checklist

  1. 1.

For all authors…

    1. (a)

Do the main claims made in the abstract and introduction accurately reflect
the paper’s contributions and scope? [Yes]

    2. (b)

Did you describe the limitations of your work? [Yes]

    3. (c)

Did you discuss any potential negative societal impacts of your work? [N/A]

    4. (d)

Have you read the ethics review guidelines and ensured that your paper
conforms to them? [Yes]

  2. 2.

If you are including theoretical results…

    1. (a)

Did you state the full set of assumptions of all theoretical results? [Yes]

    2. (b)

Did you include complete proofs of all theoretical results? [Yes] See
appendix.

  3. 3.

If you ran experiments…

    1. (a)

Did you include the code, data, and instructions needed to reproduce the main
experimental results (either in the supplemental material or as a URL)? [Yes]
Code for the illustrative experiments is available at a URL provided in
Appendix A.3. Sufficient detail for reproducability is provided in the same
section.

    2. (b)

Did you specify all the training details (e.g., data splits, hyperparameters,
how they were chosen)? [Yes]

    3. (c)

Did you report error bars (e.g., with respect to the random seed after running
experiments multiple times)? [Yes]

    4. (d)

Did you include the total amount of compute and the type of resources used
(e.g., type of GPUs, internal cluster, or cloud provider)? [Yes]

  4. 4.

If you are using existing assets (e.g., code, data, models) or
curating/releasing new assets…

    1. (a)

If your work uses existing assets, did you cite the creators? [Yes]

    2. (b)

Did you mention the license of the assets? [N/A]

    3. (c)

Did you include any new assets either in the supplemental material or as a
URL? [Yes]

    4. (d)

Did you discuss whether and how consent was obtained from people whose data
you’re using/curating? [N/A]

    5. (e)

Did you discuss whether the data you are using/curating contains personally
identifiable information or offensive content? [N/A]

  5. 5.

If you used crowdsourcing or conducted research with human subjects…

    1. (a)

Did you include the full text of instructions given to participants and
screenshots, if applicable? [N/A]

    2. (b)

Did you describe any potential participant risks, with links to Institutional
Review Board (IRB) approvals, if applicable? [N/A]

    3. (c)

Did you include the estimated hourly wage paid to participants and the total
amount spent on participant compensation? [N/A]

##  Appendix A Appendix

###  A.1 Illustrative MDPs

Several proofs in A.2 rely on constructing special MDPs to serve as examples
or counterexamples. We reserve this section to describe these MDPs for later
reference.

####  A.1.1 Ring and false-ring MDPs

We consider a simple $n$-state, $1$ action “ring” MDP (Figure 8) denoted
$m^{n}_{\circ}=(r,p)$ where:

$\displaystyle r(s_{i})=g(i)\ \forall i\in[n],\hskip
20.0ptp(s_{i+1}|s_{i})=1\hskip 4.0pt\forall i\in[n-1]\hskip 4.0pt\text{ and
}\hskip 4.0ptp(s_{1}|s_{n})=1$ (19)

where $g:i\mapsto\mathbb{R}$ is some function that defines the reward from
transitioning away from state $i$. Since $|\mathcal{A}|=1$ we omit actions
from the reward and transition dynamics.

For each ring MDP and function $g$ we additionally construct a corresponding
“false-ring” MDP (Figure 8) with the same state and actions spaces as
$m^{n}_{\circ}$ but with states that only self-transition and with rewards
designed to mimic the discounted $n$-step returns on Ring MDPs. We represent
these as $\tilde{m}^{n}_{\circ}=(\tilde{r},\tilde{p})$ where

$\displaystyle\tilde{r}(s_{i})=\frac{r^{n}(s_{i})}{\sum_{t=0}^{n-1}\gamma^{t}},\hskip
20.0pt\tilde{p}(s_{i}|s_{i})=1$ (20)

and $r^{n}(s_{i})$ denotes the discounted $n$-step return starting from
$s_{i}$ in $m^{n}_{\circ}$. Note that the discounted $n$-step return of an
$n$-state false-ring MDP is the same as that of an $n$-state ring MDP.

We now provide some basic results about pairs of ring and false-ring MDPs that
we will use periodically in our proofs.

######  Lemma 1.

For any $n\in\mathbb{Z}^{+}\cup\\{\infty\\}$ if we treat the ring MDP
$m^{n}_{\circ}$ as the environment and assume
$\tilde{m}^{n}_{\circ}\in\mathcal{M}$ it follows that

$\tilde{m}^{n}_{\circ}\in\mathcal{M}^{n}(\mathbbl{\Pi},\mathbbl{{V}}).$ (21)

when $n<\infty$ and

$\tilde{m}^{n}_{\circ}\in\mathcal{M}^{\infty}(\mathbbl{\Pi})$ (22)

when $n=\infty$.

###### Proof.

First we note that, since ring and false-ring MDPs only have one action, we
can write $\mathbbl{\Pi}=\\{\pi\\}$ where $\pi$ takes this action at all
states. We first consider the case when $n<\infty$, noting that that both MDPs
are deterministic and that for any state $s$, performing $n$ transitions will
always return to $s$. We now consider an application of $n$-step Bellman
operator of the false-ring model to an arbitrary function $v\in\mathbbl{{V}}$:

$\displaystyle\tilde{\mathcal{T}}_{\pi}^{n}v(s)$
$\displaystyle=\tilde{r}(s)(\sum_{t=0}^{n-1}\gamma^{t})+\gamma^{n}v(s)$ (23)
$\displaystyle=r^{n}(s)(\sum_{t=0}^{n-1}\gamma^{t})^{-1}(\sum_{t=0}^{n-1}\gamma^{t})+\gamma^{n}v(s)$
$\displaystyle=r^{n}(s)+\gamma^{n}v(s)$
$\displaystyle=\mathcal{T}_{\pi}^{n}v(s)$

implying that
$\tilde{m}^{n}_{\circ}\in\mathcal{M}^{n}(\mathbbl{\Pi},\mathbbl{{V}})$ as
needed. We now consider the case when $n=\infty$: here, we note that for any
state $s\in\mathcal{S}$:

$\tilde{r}(s)=\frac{r^{\infty}(s)}{\sum_{t=0}^{\infty}\gamma^{t}}=(1-\gamma)v_{\pi}(s).$
(24)

we can then write:

$\tilde{v}_{\pi}(s)=\sum_{t=0}^{\infty}\gamma^{t}\tilde{r}(s)=(1-\gamma)^{-1}(1-\gamma)v_{\pi}(s)=v_{\pi}(s)$
(25)

since $\tilde{m}^{\infty}_{\circ}$ only self-transitions at each state. This
shows that $\tilde{m}^{\infty}_{\circ}\in\mathcal{M}^{\infty}(\Pi)$ as needed.

∎

######  Lemma 2.

Fix any $k,K\in\mathbb{Z}^{+}\cup\\{\infty\\}$ with $k<K$ and let
$f:\mathcal{S}\mapsto\mathbb{R}$ be any constant function. Let
$m=m^{K}_{\circ}$ and $\tilde{m}=\tilde{m}^{K}_{\circ}$. For any
$\gamma\in(0,1)$ it follows that

$\mathcal{T}_{\pi}^{k}f(s_{1})\neq\tilde{\mathcal{T}}_{\pi}^{k}f(s_{1})$ (26)

where $m^{K}_{\circ}$ and $\tilde{m}^{K}_{\circ}$ are $K$-state ring and
false-ring MDPs with $g(i)=\bm{1}\\{i\in[1,k]\\}$.

###### Proof.

We begin by examining the $k$-step Bellman operator and Bellman fixed-point
under the ring $m^{K}_{\circ}$:

$\mathcal{T}_{\pi}^{k}f(s_{1})=r^{k}(s_{1})+\gamma^{k}f(s_{1})=r^{K}(s_{1})+\gamma^{k}f(s_{1})$
(27)

where the second equality follows from the fact that $g$ ensures that no
reward is received after the first $k$ steps from $s_{1}$.

Next we examine the corresponding $k$-step Bellman operator under the false-
ring $\tilde{m}^{K}_{\circ}$:

$\tilde{\mathcal{T}}_{\pi}^{k}f(s_{1})=\tilde{r}^{k}(s_{1})+\gamma^{k}f(s_{1})=r^{K}(s_{1})(\sum_{t=0}^{K-1}\gamma^{t})^{-1}\sum_{t=0}^{k-1}\gamma^{t}+\gamma^{k}f(s_{1})$
(28)

where the second equality follows from the construction of $K$-step false-ring
MDPs to match the $K$-step returns of their corresponding ring MDP.

Taken together Eqs. (27-28) imply that in order for
$\mathcal{T}_{\pi}^{k}f(s_{1})=\tilde{\mathcal{T}}_{\pi}^{k}f(s_{1})$ it must
be the case that $\sum_{t=0}^{K-1}\gamma^{t}=\sum_{t=0}^{k-1}\gamma^{t}$ which
can only happen when $\gamma=0$. Note that these properties hold when
$K=\infty$. This completes the proof. ∎

(a)

(b)

Figure 8: Ring and false-ring environments with reward structure defined by
$g:\mathbb{Z}^{+}\mapsto\mathbb{R}$. States are numbered circles and outgoing
arrows indicate possible transitions from each state. Arrows are labeled by
the reward attained from performing their transition.

###  A.2 Proofs

In this section we provide proofs of the results in the main text.

See 1

###### Proof.

Consider some $m\in\mathcal{M}^{k}(\Pi,\mathcal{V})$. For any $\pi\in\Pi$ and
$v\in\mathcal{V}$ we know that
$\tilde{\mathcal{T}}_{\pi}^{k}v=\mathcal{T}_{\pi}^{k}v$. Since $k$ divides $K$
we know that $K=zk$ where $z\in\mathbb{Z}^{+}$. Hence

$\displaystyle\mathcal{T}_{\pi}^{K}v=\underbrace{\mathcal{T}_{\pi}\cdots\mathcal{T}_{\pi}}_{\text{K
times}}v=\underbrace{\mathcal{T}_{\pi}^{k}\cdots\mathcal{T}_{\pi}^{k}}_{\text{z
times}}v$ (29)

Finally since $\mathcal{V}$ is closed under Bellman updates we can write
$\tilde{\mathcal{T}}_{\pi}^{k}v=\mathcal{T}_{\pi}^{k}v\in\mathcal{V}$, which
allows us iteratively equate $k$-step environment and model operators on the
right-hand side of Eq. (29) to obtain:

$\underbrace{\mathcal{T}_{\pi}^{k}\cdots\mathcal{T}_{\pi}^{k}}_{\text{z
times}}v=\underbrace{\tilde{\mathcal{T}}_{\pi}^{k}\cdots\tilde{\mathcal{T}}_{\pi}^{k}}_{\text{z
times}}v=\tilde{\mathcal{T}}_{\pi}^{K}v.$ (30)

This suffices to show that $m\in\mathcal{M}^{K}(\Pi,\mathcal{V})$ which means
$\mathcal{M}^{k}(\Pi,\mathcal{V})\subseteq\mathcal{M}^{K}(\Pi,\mathcal{V})$.

We now assume that $\mathcal{V}$ contains at least one constant function and
$\Pi$ is non-empty and produce an instance of an environment and model class
where the relation is strict. Let the environment be a $K$-state ring
environment (see A.1.1): $m^{K}_{\circ}$ with $g(i)=\bm{1}\\{i\in[1,k]\\}$ and
let $\mathcal{M}=\mathbbl{{M}}$. Next we introduce a model given by the
corresponding false-ring MDP (see A.1.1) $\tilde{m}^{K}_{\circ}$. From Lemma 1
we have that $\tilde{m}^{K}_{\circ}\in\mathcal{M}^{K}(\Pi,\mathcal{V})$.

Since there is at least one constant function $f\in\mathcal{V}$ we know that
$\mathcal{T}_{\pi}^{k}f(s_{1})\neq\tilde{\mathcal{T}}_{\pi}^{k}f(s_{1})$ from
Lemma 2. This is sufficient to show that
$\tilde{m}^{K}_{\circ}\notin\mathcal{M}^{k}(\Pi,\mathcal{V})$ and thus we have
proven that there are instances where
$\mathcal{M}^{k}(\Pi,\mathcal{V})\subset\mathcal{M}^{K}(\Pi,\mathcal{V})$. ∎

See 2

###### Proof.

We first note
$\mathcal{M}^{\infty}(\Pi)=\bigcap_{\pi\in\Pi}\mathcal{M}^{\infty}(\\{\pi\\})$
and consider any $m\in\mathcal{M}^{\infty}(\\{\pi\\})$ for some $\pi\in\Pi$.
From the definition of PVE we know $\tilde{v}_{\pi}=v_{\pi}$ and thus can say:

$\displaystyle\tilde{v}_{\pi}=v_{\pi}$ (31) $\displaystyle\implies$
$\displaystyle\tilde{\mathcal{T}}_{\pi}^{k}\tilde{v}_{\pi}=\tilde{\mathcal{T}}_{\pi}^{k}v_{\pi}$
$\displaystyle\implies$
$\displaystyle\tilde{v}_{\pi}=\tilde{\mathcal{T}}_{\pi}^{k}v_{\pi}$
$\displaystyle\implies$ $\displaystyle
v_{\pi}=\tilde{\mathcal{T}}_{\pi}^{k}v_{\pi}$ $\displaystyle\implies$
$\displaystyle\mathcal{T}_{\pi}^{k}v_{\pi}=\tilde{\mathcal{T}}_{\pi}^{k}v_{\pi}$

which suggests that $m\in\mathcal{M}^{k}(\\{\pi\\},\\{v^{\pi}\\})$ and thus
$\mathcal{M}^{\infty}(\\{\pi\\})\subseteq\mathcal{M}^{k}(\\{\pi\\},\\{v^{\pi}\\})$.

We now consider any element $m\in\mathcal{M}^{k}(\\{\pi\\},\\{v^{\pi}\\})$,
and note that from the definition of order-$k$ VE we know that
$\tilde{\mathcal{T}}_{\pi}^{k}v_{\pi}=\mathcal{T}_{\pi}^{k}v_{\pi}$, thus we
can say:

$\displaystyle\tilde{\mathcal{T}}_{\pi}^{k}v_{\pi}=\mathcal{T}_{\pi}^{k}v_{\pi}$
(32) $\displaystyle\implies$
$\displaystyle\tilde{\mathcal{T}}_{\pi}^{k}v_{\pi}=v_{\pi}$
$\displaystyle\implies$
$\displaystyle\tilde{\mathcal{T}}_{\pi}^{2k}v_{\pi}=\tilde{\mathcal{T}}_{\pi}^{k}v_{\pi}$
$\displaystyle\implies$
$\displaystyle\tilde{\mathcal{T}}_{\pi}^{2k}v_{\pi}=v_{\pi}$

where we can repeat the process described in these implications ad-infinitum
to obtain
$\tilde{v}_{\pi}=\lim_{n\to\infty}\tilde{\mathcal{T}}_{\pi}^{nk}v_{\pi}=v_{\pi}$.
Hence $m\in\mathcal{M}^{\infty}(\\{\pi\\})$ and thus
$\mathcal{M}^{k}(\\{\pi\\},\\{v^{\pi}\\})$.

Taken together this shows that
$\mathcal{M}^{\infty}(\\{\pi\\})=\mathcal{M}^{k}(\\{\pi\\},\\{v^{\pi}\\})$ for
any $k$ and $\pi$ thus:

$\mathcal{M}^{\infty}(\Pi)=\bigcap_{\pi\in\Pi}\mathcal{M}^{\infty}(\\{\pi\\})=\bigcap_{\pi\in\Pi}\mathcal{M}^{k}(\\{\pi\\},\\{v^{\pi}\\})$
(33)

for any $k\in\mathbb{Z}^{+}$. ∎

######  Corollary 2.

Let $\Pi\subseteq\mathbbl{\Pi}$ and let $\mathcal{V}$ be as in Proposition 1
for $k\in\mathbb{Z}^{+}$ then we have that
$\mathcal{M}^{k}(\Pi,\mathcal{V})\subseteq\mathcal{M}^{\infty}(\Pi).$
Moreover, if $\Pi$ is non-empty and $\mathcal{V}$ contains at least one
constant function, then there exist environments such that
$\mathbbl{{M}}^{k}(\Pi,\mathcal{V})\subset\mathbbl{{M}}^{\infty}(\Pi)$

###### Proof.

Consider some $m\in\mathcal{M}^{k}(\Pi,\mathcal{V})$. From the generalization
of Property 1 we know that $m\in\mathcal{M}^{zk}(\Pi,\mathcal{V})$ for any
$z\in\mathbb{Z}^{+}$ since $k$ divides $zk$. Thus we know that
$\tilde{\mathcal{T}}_{\pi}^{zk}v=\mathcal{T}_{\pi}^{zk}v$ for any choice of
$\pi\in\Pi$, $v\in\mathcal{V}$ and $z\in\mathbb{Z}^{+}$. Accordingly the
expressions are equal in the limit as $z\to\infty$. Combining this with the
fact that both $\tilde{\mathcal{T}}_{\pi}$ and $\mathcal{T}_{\pi}$ are
contraction mappings, we obtain:

$\tilde{v}_{\pi}=\lim_{z\to\infty}\tilde{\mathcal{T}}_{\pi}^{zk}v=\lim_{z\to\infty}\mathcal{T}_{\pi}^{zk}v=v_{\pi}$
(34)

which implies $m\in\mathcal{M}^{\infty}(\Pi)$ and thus
$\mathcal{M}^{k}(\Pi,\mathcal{V})\subseteq\mathcal{M}^{\infty}(\Pi)$, as
needed.

Moreover, so long that $\Pi$ is nonempty and $\mathcal{V}$ contains some
constant function $f$, we can construct a pair of $\infty$-state ring / false-
ring MDPs: $m^{\infty}_{\circ}$ and $\tilde{m}^{\infty}_{\circ}$ with
$g(i)=\bm{1}\\{i\in[1,k]\\}$ (see A.1.1). By assuming that
$m^{\infty}_{\circ}$ is the environment, Lemma 1 tells us that
$\tilde{m}^{\infty}_{\circ}\in\mathcal{M}^{\infty}(\Pi)$ and we know from
Lemma 2 that
$\mathcal{T}_{\pi}^{k}f(s_{1})\neq\tilde{\mathcal{T}}_{\pi}^{k}f(s_{1})$ hence
$\tilde{m}^{\infty}_{\circ}\notin\mathcal{M}^{k}(\Pi)$. ∎

See 3

###### Proof.

Assume that $\mathcal{M}=\mathbbl{{M}}$. Denote the environment reward and
transition dynamics as $(r,p)$. For any value $y_{0}\in\mathcal{Y}$ we
consider a model $m_{y_{0}}$:

$\displaystyle r_{m_{y_{0}}}((x,y),a)=r((x,y),a)$ (35) $\displaystyle
p_{m_{y_{0}}}((x^{\prime},y^{\prime})|(x,y),a)=\bm{1}\\{y^{\prime}=y_{0}\\}p(x^{\prime}|(x,y),a).$

We now examine the Bellman fixed-point induced by environment for any policy
$\pi\in\Pi$:

$\displaystyle v_{\pi}((x,y))$
$\displaystyle=\int_{\mathcal{A}}\pi(a|(x,y))r((x,y),a)+\gamma\int_{\mathcal{X}}\int_{\mathcal{Y}}p((x^{\prime},y^{\prime})|(x,y),a)v_{\pi}((x^{\prime},y^{\prime}))dx^{\prime}dy^{\prime}da$
(36)
$\displaystyle=\int_{\mathcal{A}}\pi(a|(x,y))r((x,y),a)+\gamma\int_{\mathcal{X}}\int_{\mathcal{Y}}p(x^{\prime}|(x,y),a)p(y^{\prime}|x^{\prime},(x,y),a)v_{\pi}(x^{\prime})dx^{\prime}dy^{\prime}da$
$\displaystyle=\int_{\mathcal{A}}\pi(a|(x,y))r((x,y),a)+\gamma\int_{\mathcal{X}}p(x^{\prime}|(x,y),a)v_{\pi}(x^{\prime})dx^{\prime}da.$

We can compare this to the Bellman operator induced by our model for the same
policy:

$\displaystyle\tilde{\mathcal{T}}_{\pi}v((x,y))$
$\displaystyle=\int_{\mathcal{A}}\pi(a|(x,y))r((x,y),a)+\gamma\int_{\mathcal{X}}\int_{\mathcal{Y}}\bm{1}\\{y^{\prime}=y_{0}\\}p(x^{\prime}|(x,y),a)v((x,y))dx^{\prime}dy^{\prime}da$
(37)
$\displaystyle=\int_{\mathcal{A}}\pi(a|(x,y))r((x,y),a)+\gamma\int_{\mathcal{X}}p(x^{\prime}|(x,y),a)v((x^{\prime},y_{0}))dxda$

Notice that $v_{\pi}$ is a fixed point of this operator, hence
$\tilde{v}_{\pi}=v_{\pi}$ and and thus $m_{y_{0}}\in\mathcal{M}^{\infty}(\Pi)$
(since our particular choice of $\pi\in\Pi$ was arbitrary). Moreover, we can
construct different models for each $y_{0}\in\mathcal{Y}$, we know that

$\mathcal{M}_{\mathcal{Y}}=\\{m_{y}:y\in\mathcal{Y}\\}\subseteq\mathcal{M}^{\infty}(\mathbbl{\Pi}).$
(38)

Moreover, suppose $m_{y_{0}}\in\mathcal{M}^{1}(\Pi,\mathbbl{{V}})$. This
implies that for all $v\in\mathbbl{{V}}$

$\displaystyle\tilde{\mathcal{T}}_{\pi}v((x,y))=\mathcal{T}_{\pi}v((x,y))$
(39) $\displaystyle\implies$
$\displaystyle\int_{\mathcal{A}}\int_{\mathcal{X}}\int_{\mathcal{Y}}\pi(a|(x,y))p_{m_{y_{0}}}((x^{\prime},y^{\prime})|(x,y),a)v((x^{\prime},y^{\prime}))dx^{\prime}dy^{\prime}da$
$\displaystyle=\int_{\mathcal{A}}\int_{\mathcal{X}}\int_{\mathcal{Y}}\pi(a|(x,y))p((x^{\prime},y^{\prime})|(x,y),a)v((x^{\prime},y^{\prime}))dx^{\prime}dy^{\prime}da$
$\displaystyle\implies$
$\displaystyle\int_{\mathcal{A}}\int_{\mathcal{X}}\pi(a|(x,y))p(x^{\prime}|(x,y),a)v((x^{\prime},y_{0}))dx^{\prime}$
$\displaystyle=\int_{\mathcal{A}}\int_{\mathcal{X}}\int_{\mathcal{Y}}\pi(a|(x,y))p((x^{\prime},y^{\prime})|(x,y),a)v((x^{\prime},y^{\prime}))dx^{\prime}dy^{\prime}$

we now choose $v((x,y))=\bm{1}\\{y\neq y_{0}\\}$ which reduces the above
equations to:

$\displaystyle\implies$ $\displaystyle
0=\int_{\mathcal{A}}\int_{\mathcal{X}}\int_{\mathcal{Y}\neq
y_{0}}\pi(a|(x,y))p((x^{\prime},y^{\prime})|(x,y),a)=\mathbb{P}(y^{\prime}\neq
y_{0}|x,y,\pi)$ (40) $\displaystyle\implies$
$\displaystyle\mathbb{P}(y^{\prime}=y_{0}|x,y,\pi)=1$

where $\mathbb{P}$ denotes the conditional probability of an event.

Now consider the class of models defined by Eq. (38). Suppose
$\mathcal{M}_{\mathcal{Y}}\in\mathcal{M}^{1}(\Pi,\mathbbl{{V}})$, by Eq. (40)
this would mean that $\mathbb{P}(y^{\prime}=y_{0}|x,y,\pi)=1$ for all
$y_{0}\in\mathcal{Y}$. This is impossible unless $|\mathcal{Y}|=1$ hence there
must exist $m_{y_{0}}\notin\mathcal{M}^{1}(\Pi,\mathbbl{{V}})$ and thus
$\mathcal{M}^{1}(\Pi,\mathbbl{{V}})\subset\mathcal{M}^{\infty}(\Pi,\mathbbl{{V}})$.

∎

See 1

###### Proof.

Denote a deterministic optimal policy with respect to the environment as
$\pi^{*}$. Let $\tilde{m}\in\mathcal{M}^{\infty}(\mathbbl{\Pi}_{\text{det}})$
and $\tilde{\pi}^{*}$ be a deterministic optimal policy with respect
$\tilde{m}$.

Suppose $\tilde{\pi}^{*}$ were not optimal in the environment. This implies
that $v_{\pi^{*}}(s)\geq v_{\tilde{\pi}^{*}}(s)\forall s\in\mathcal{S}$ with
strict inequality for at least one state. However, since $\pi^{*}$ and
$\tilde{\pi}^{*}$ are deterministic we have:

$\tilde{v}_{\pi^{*}}(s)=v_{\pi^{*}}(s)>v_{\tilde{\pi}^{*}}(s)=\tilde{v}_{\tilde{\pi}^{*}}(s)$
(41)

for some $s\in\mathcal{S}$. This contradicts $\tilde{\pi}^{*}$ being optimal
in the model. ∎

See 5

###### Proof.

Since the environment and model only differ when action $R$ is taken from
state $2$, we only need to consider deterministic policies that make this
choice. Note that if action $R$ is taken from state $2$, the values in both
the model and environment at states $2$ and $3$ are necessarily $0$ and the
value of each in state $1$ is either $(1-\gamma)^{-1}$ or $0$ depending on the
action taken from state $1$. This suffices to show that the environment and
model have the same values for all deterministic policies.

However, one can see that the model and environment differ for stochastic
policies. Take, for instance, a policy for which $\pi(a|s)=0.5$ for all
$a\in\mathcal{A}$, $s\in\mathcal{S}$. The induced Markov reward processes from
applying this policy to the environment and model, which share the same reward
structure, have different transition dynamics at state $2$. It can be easily
verified that this results in different values for the environment and model.
∎

See 6

###### Proof.

We begin by considering the left-hand side

$\displaystyle\|v_{\pi}-\tilde{\mathcal{T}}^{k}_{\pi}v_{\pi}\|_{\infty}$
$\displaystyle=\|v_{\pi}-\tilde{\mathcal{T}}^{k}_{\pi}v+\tilde{\mathcal{T}}^{k}_{\pi}v-\tilde{\mathcal{T}}^{k}_{\pi}v_{\pi}\|_{\infty}$
(42)
$\displaystyle\leq\|v_{\pi}-\tilde{\mathcal{T}}^{k}_{\pi}v\|_{\infty}+\|\tilde{\mathcal{T}}^{k}_{\pi}v-\tilde{\mathcal{T}}^{k}_{\pi}v_{\pi}\|_{\infty}$
$\displaystyle=\|v_{\pi}-\mathcal{T}^{n}_{\pi}v+\mathcal{T}^{n}_{\pi}v-\tilde{\mathcal{T}}^{k}_{\pi}v\|_{\infty}+\|\tilde{\mathcal{T}}^{k}_{\pi}v-\tilde{\mathcal{T}}^{k}_{\pi}v_{\pi}\|_{\infty}$
$\displaystyle\leq\|v_{\pi}-\mathcal{T}^{n}_{\pi}v\|_{\infty}+\|\mathcal{T}^{n}_{\pi}v-\tilde{\mathcal{T}}^{k}_{\pi}v\|_{\infty}+\|\tilde{\mathcal{T}}^{k}_{\pi}v-\tilde{\mathcal{T}}^{k}_{\pi}v_{\pi}\|_{\infty}$
$\displaystyle\leq\gamma^{n}\|v_{\pi}-v\|_{\infty}+\|\mathcal{T}^{n}_{\pi}v-\tilde{\mathcal{T}}^{k}_{\pi}v\|_{\infty}+\gamma^{k}\|v_{\pi}-v\|_{\infty}$
$\displaystyle=(\gamma^{k}+\gamma^{n})\|v_{\pi}-v\|_{\infty}+\|\mathcal{T}^{n}_{\pi}v-\tilde{\mathcal{T}}^{k}_{\pi}v\|_{\infty}$

as needed. ∎

######  Proposition 7.

For any $\pi\in\mathbbl{\Pi}$, $v\in\mathbbl{{V}}$ and $k,n\in\mathbb{Z}^{+}$,
assuming $\|v_{\pi}-v\|_{\infty}<g\cdot\|v_{\pi}-v\|_{d_{\pi}}$ for some
$g\geq 0$, we have that:

$\|v_{\pi}-\tilde{\mathcal{T}}_{\pi}v_{\pi}\|_{d_{\pi}}\leq(g\cdot\gamma^{k}+\gamma^{n})\|v_{\pi}-v\|_{d_{\pi}}+\|\mathcal{T}^{n}_{\pi}v-\tilde{\mathcal{T}}^{k}_{\pi}v\|_{d_{\pi}}$
(43)

###### Proof.

$\displaystyle\|v_{\pi}-\tilde{\mathcal{T}}^{k}_{\pi}v_{\pi}\|_{d_{\pi}}$
$\displaystyle=\|v_{\pi}-\tilde{\mathcal{T}}^{k}_{\pi}v+\tilde{\mathcal{T}}^{k}_{\pi}v-\tilde{\mathcal{T}}^{k}_{\pi}v_{\pi}\|_{d_{\pi}}$
(44)
$\displaystyle\leq\|v_{\pi}-\tilde{\mathcal{T}}^{k}_{\pi}v\|_{d_{\pi}}+\|\tilde{\mathcal{T}}^{k}_{\pi}v-\tilde{\mathcal{T}}^{k}_{\pi}v_{\pi}\|_{d_{\pi}}$
$\displaystyle=\|v_{\pi}-\mathcal{T}^{n}_{\pi}v+\mathcal{T}^{n}_{\pi}v-\tilde{\mathcal{T}}^{k}_{\pi}v\|_{d_{\pi}}+\|\tilde{\mathcal{T}}^{k}_{\pi}v-\tilde{\mathcal{T}}^{k}_{\pi}v_{\pi}\|_{d_{\pi}}$
$\displaystyle\leq\|v_{\pi}-\mathcal{T}^{n}_{\pi}v\|_{d_{\pi}}+\|\mathcal{T}^{n}_{\pi}v-\tilde{\mathcal{T}}^{k}_{\pi}v\|_{d_{\pi}}+\|\tilde{\mathcal{T}}^{k}_{\pi}v-\tilde{\mathcal{T}}^{k}_{\pi}v_{\pi}\|_{d_{\pi}}$
$\displaystyle\leq\|v_{\pi}-\mathcal{T}^{n}_{\pi}v\|_{d_{\pi}}+\|\mathcal{T}^{n}_{\pi}v-\tilde{\mathcal{T}}^{k}_{\pi}v\|_{d_{\pi}}+\|\tilde{\mathcal{T}}^{k}_{\pi}v-\tilde{\mathcal{T}}^{k}_{\pi}v_{\pi}\|_{\infty}$
$\displaystyle\leq\gamma^{n}\|v_{\pi}-v\|_{d_{\pi}}+\|\mathcal{T}^{n}_{\pi}v-\tilde{\mathcal{T}}^{k}_{\pi}v\|_{d_{\pi}}+\gamma^{k}\|v_{\pi}-v\|_{\infty}$
$\displaystyle\leq\gamma^{n}\|v_{\pi}-v\|_{d_{\pi}}+\|\mathcal{T}^{n}_{\pi}v-\tilde{\mathcal{T}}^{k}_{\pi}v\|_{d_{\pi}}+g\cdot\gamma^{k}\|v_{\pi}-v\|_{d_{\pi}}$
$\displaystyle=(g\cdot\gamma^{k}+\gamma^{n})\|v_{\pi}-v\|_{d_{\pi}}+\|\mathcal{T}^{n}_{\pi}v-\tilde{\mathcal{T}}^{k}_{\pi}v\|_{d_{\pi}}$

∎

###  A.3 Experimental details - illustrative experiments

####  A.3.1 Code

Code to reproduce our illustrative experiments can be found at
https://github.com/chrisgrimm/proper_value_equivalence.

####  A.3.2 Computational resources

Illustrative experiments were performed on three machines each with $4$ NVIDIA
GeForce GTX 1080 Ti graphics cards.

Figure 9: Visualization of the Four Rooms environment.

####  A.3.3 Environment

All illustrative experiments depicted in Figures 4 and 6 were carried out in a
stochastic version of the Four Rooms environment (depicted in Figure 9) where
$|\mathcal{S}|=104$ and $\mathcal{A}$ consists of four actions corresponding
to an intended movement in each of the cardinal directions. When an agent
takes an action, it will move in the intended direction 80% of the time and
otherwise move in a random direction. If the agent moves into a wall it will
remain in place. When the agent transitions into the upper-right square it
receives a reward of $1$, all other transitions yield $0$ reward.

####  A.3.4 Model representation and initialization

Models are represented tabularly by matrices
$\tilde{R}\in\mathbb{R}^{|\mathcal{S}|\times|\mathcal{A}|}$ and
$\tilde{P}^{a}\in\mathbb{R}^{|\mathcal{S}|\times|\mathcal{S}|}$ for
$a\in\mathcal{A}$ where $\tilde{R}_{s,a}=\tilde{r}(s,a)$ and
$\tilde{P}^{a}_{s,s^{\prime}}=\tilde{p}(s^{\prime}|s,a)$. We generally
constrain a matrix to be row-stochastic by parameterizing it with an
unconstrained matrix of the same shape and applying a softmax with temperature
$1$ to each of its rows. In experiments with model capacity constraints we
additionally impose that each $\tilde{P}^{a}$ has a rank of at most $k$ by
representing $\tilde{P}^{a}=D^{a}K^{a}$ where
$D^{a}\in\mathbb{R}^{|\mathcal{S}|\times k}$,
$K^{a}\in\mathbb{R}^{k\times|\mathcal{S}|}$ and both $D^{a}$ and $K^{a}$ are
constrained to be row-stochastic (note that the product of row-stochastic
matrices is itself row-stochastic). In this setting the parameters of the
capacity-constrained transition dynamics are the unconstrained matrices
parameterizing $D^{a}$ and $K^{a}$.

Models are initialized by randomly sampling the entries of $\tilde{R}$
according to $U(-1,1)$ and the entries of the matrices parameterizing the
transition dynamics according to $U(-5,5)$, where $U(l,u)$ denotes a uniform
distribution over the interval $(l,u)$.

In all illustrative experiments we train our models using the Adam optimizer
with default hyperparameters ($\beta_{1}=0.99$, $\beta_{2}=0.999$, $\epsilon$
=1e-8).

####  A.3.5 Model space experiments

In Figure 4 we illustrate the properties of spaces of models trained to be in
$\mathcal{M}^{k}(\mathbbl{\Pi},\mathbbl{{V}})$ for $k\in\\{1,30,40,50,60\\}$
and in $\mathcal{M}^{\infty}(\mathbbl{\Pi})$. To train each of these models we
construct a set of policies and functions
$\mathcal{D}=\\{(\pi_{i},v_{i})\\}_{i=1}^{100,000}$. Each generated policy
$\pi_{i}$ is, with equal probability, either a uniformly sampled deterministic
policy or a stochastic policy for which at each state $s$,
$\pi_{i}(a|s)=f_{a}/\sum_{a\in\mathcal{A}}f_{a}$ where $f_{a}\sim U(0,1)$ for
each $a\in\mathcal{A}$. Each $v_{i}$ is sampled such that $v_{i}(s)\sim
U(-1,1)$ for each $s\in\mathcal{S}$. We then sample minibatches
$B\sim\mathcal{D}$ with $|B|=50$ at each iteration and update models to
minimize

$\frac{1}{|B|}\sum_{(\pi,v)\in
B}(\tilde{\mathcal{T}}^{k}_{\pi}v-\mathcal{T}^{k}_{\pi}v)^{2}\hskip
10.0pt\text{ and }\hskip 10.0pt\frac{1}{|B|}\sum_{(\pi,v)\in
B}(\tilde{\mathcal{T}}_{\pi}v_{\pi}-v_{\pi})^{2}$ (45)

for order-$k$ VE and PVE models respectively.

Each model is updated in this manner for $500,000$ iterations with a learning
rate of 1e-3 and a snapshot of the model is stored every $1000$
iterations—creating a timeline of the model’s progress through training. For
each model class, this experiment is repeated with $120$ randomly initialized
models. To generate the points on the scatter plots depicted in Figure 4, we
iterate through the snapshots of these $120$ models. At snapshot $t$ (training
iteration $1000\times t$) we collect the snapshots of all the models and
convert each model into a 1D vector representation by concatenating the
entries from its reward and transition dynamics matrices. We then apply
principle component analysis to these vectors, isolating the first two
principle components, which we treat as (x, y) coordinates in the scatter
plots. For the top row in Figure 4 we color these points according to progress
through training: $(t/500)$. On the bottom row, we compute the optimal policy
with respect to each point’s corresponding model: $\tilde{\pi}^{*}$ and color
the point according to
$(\sum_{s}v_{\tilde{\pi}^{*}}(s))/(\sum_{s}v_{\pi*}(s))$.

We produce the plot of model class diameters in Figure 4 by taking the
scatter-plot points corresponding to the final snapshot ($t=500$) of models
for each $k$, randomly grouping them into $4$ sets of $30$ points and
computing the diameters of each set. We then use these 4 diameters to produce
error bars.

####  A.3.6 Individual model visualization

To generate the visualization of the dynamics of individual models displayed
in Figure 6b, we randomly select a single PVE model trained in our model space
experiments. We then collect $5000$ length $30$ trajectories starting from the
bottom left state. The paths of these trajectories are then overlaid on top of
a visualization of the environment and colored according to time along the
trajectory ($t/30$). This procedure is repeated using the environment in
Figure 6a.

####  A.3.7 Model capacity experiment

We compare the effect of capacity constraints on learning models in
$\mathcal{M}^{\infty}(\mathbbl{\Pi})$ and
$\mathcal{M}^{\infty}(\mathbbl{\Pi}_{\mathrm{det}})$ respectively by
restricting the rank of the learned model’s transition dynamics (as in A.3.4).
We restrict the ranks of model transition dynamics to be at most $k$ for
$k\in\\{20,30,40,50,60,70,80,90,100,104\\}$. To train each model we collect a
set of $1000$ policies by beginning with a random policy and repeatedly
running the policy iteration algorithm in the environment, starting with a
randomly initialized policy and stopping when the optimal policy is reached.
The sequence of improved policies resulting from this process is stored.
Whenever the algorithm terminates, a new random policy is generated and the
process is repeated until $1000$ policies have been stored. To increase the
number of distinct policies generated by this process, at each step of policy
iteration, we select, uniformly at random, 10% of states and update the policy
at only these states. We then further boost the breadth of our collected
policies and specialize them to $\mathbbl{\Pi}$ and
$\mathbbl{\Pi}_{\mathrm{det}}$ by adding stochastic or deterministic “noise.”

Precisely, when training a model to be in
$\mathcal{M}^{\infty}(\mathbbl{\Pi})$ we iterate over each of the 1000
policies generated by our policy iteration procedure and generate an
additional 100 policies. Each additional policy is generated by selecting,
uniformly at random, 10% of the original policy’s states and replacing the its
distribution at these states with a uniform distribution over actions.

When training a model to be in
$\mathcal{M}^{\infty}(\mathbbl{\Pi}_{\mathrm{det}})$ the same procedure is
repeated but the original policy’s distributions, at the selected states, are
replaced by randomly generated deterministic distributions.

In either case, this produces $100,000$ policies which are evaluated in the
environment. Together this forms a set of policies and corresponding value
functions: $\mathcal{D}=\\{(\pi_{i},v_{i})\\}_{i=1}^{100,000}$ which can be
used construct mini-batch PVE losses as described in (45). Models are trained
according to these losses for $1,000,000$ iterations with a learning rate of
5e-4. The errorbars around the environment value of the models’ optimal
policies at the end of training are reported across $10$ seeds.

###  A.4 MuZero experiment

##### Atari.

We follow the Atari configuration used in Schrittwieser et al. [30],
summarised in Table 1.

Table 1: Atari hyperparameters. Parameter | Value
---|---
Start no-ops | [0, 30]
Terminate on life loss | Yes
Action set | Valid actions
Max episode length | 30 minutes (108,000 frames)
Observation size | $96\times 96$
Preprocessing | Grayscale
Action repetitions | 4
Max-pool over last N action repeat frames | 4
Total environment frames, including skipped frames | 500M

##### MuZero implementation.

Our MuZero implementation largely follows the description given by
Schrittwieser et al. [30], but uses a Sebulba distributed architecture as
described in Hessel et al. [19], and TD($\lambda$) rather than $n$-step value
targets. The hyperparameters are given in Table 2. Our network architecture is
the same as used in MuZero [30].

The base MuZero loss is given by

$\mathcal{L}^{\textrm{base}}_{t}=\mathbb{E}_{\pi}\sum_{k=0}^{K}\left[\ell^{r}(r^{\text{target}}_{t+k},\hat{r}_{t}^{k})+\ell^{v}(v_{t+k}^{\text{target}},\hat{v}_{t}^{k})+\ell^{\pi}(\pi_{t+k}^{\text{target}},\hat{\pi}_{t}^{k})\right].$
(46)

The reward loss $\ell^{r}$ simply regresses the model-predicted rewards to the
rewards seen in the environment. To compute the value and policy losses,
MuZero performs a Monte-Carlo tree search using the learned model. The policy
targets are proportional to the MCTS visitation counts at the root node. The
value targets are computed using the MCTS value prediction $\tilde{v}$ and the
sequences of rewards. MuZero uses an $n$-step bootstrap return estimate
$v^{\text{target}}_{t+k}=\sum_{j=1}^{n}\gamma^{j-1}r_{t+k+j}+\gamma^{n}\tilde{v}_{t+k+n}$.
We use the a TD($\lambda$) return estimate instead.

For our additional loss corresponding to past policies, we periodically store
the parameters for the value function and policy (i.e. the network heads that
take the model-predicted latent state as input). Then, we compute the same
value loss $\ell^{v}$ for each past value function. To account for the fact
that the reward sequence was drawn from the current policy $\pi$ rather than
the stored policies, we use V-trace to compute a return estimate for the past
policies.

The additional hyperparameters for the buffer of past value heads were tuned
on MsPacman, over a buffer size in $\\{64,128,256\\}$ and an update interval
in $\\{10,50,100,500\\}$. Our experiments took roughly 35k TPU-v3 device-hours
for both tuning and the full evaluation.

##### Full results.

We report the final scores per game in Table 3. The mean scores are across the
final 200 episodes in each of three seeds. We also report the standard error
of the mean across seeds only. Performing a Wilcoxon signed rank test
comparing per-game scores, we find that the version with the additional Past
Policies loss has a better final performance with $p=0.044$.

Table 2: Hyperparameters for our MuZero experiment. Hyperparameter | Value
---|---
Batch size | 96 sequences
Sequence length | 30 frames
Sequence overlap | 10 frames
Model unroll length $K$ | 5
Optimiser | Adam
Initial learning rate | $1\times 10^{-4}$
Final learning rate (linear schedule) | 0
Discount | 0.997
Target network update rate | 0.1
Value loss weight | 0.25
Reward loss weight | 1.0
Policy loss weight | 1.0
MCTS number of simulations | 25
$\lambda$ for TD($\lambda$) | 0.8
MCTS Dirichlet prior fraction | 0.3
MCTS Dirichlet prior $\alpha$ | 0.25
Search parameters update rate | 0.1
Value, reward number of bins | 601
Nonlinear value transform | $\text{sgn}(z)(\sqrt{|z|+1}-1)+0.01z$
Value buffer size | 128
Value buffer update interval | 50
Value buffer loss weight | 0.25
Environment | MuZero (our impl.) | MuZero + Past Policies
---|---|---
alien | 38,698 |  $\pm$ 2,809 | 52,821 |  $\pm$ 1,918
amidar | 6,631 |  $\pm$ 568 | 4,239 |  $\pm$ 1,550
assault | 35,876 |  $\pm$ 550 | 35,013 |  $\pm$ 738
asterix | 674,573 |  $\pm$ 88,318 | 549,421 |  $\pm$ 9,280
asteroids | 214,034 |  $\pm$ 4,719 | 235,543 |  $\pm$ 14,605
atlantis | 835,445 |  $\pm$ 92,290 | 845,409 |  $\pm$ 60,318
bank_heist | 837 |  $\pm$ 265 | 552 |  $\pm$ 234
battle_zone | 39,471 |  $\pm$ 12,658 | 72,183 |  $\pm$ 11,385
beam_rider | 120,675 |  $\pm$ 16,588 | 130,129 |  $\pm$ 14,014
berzerk | 22,449 |  $\pm$ 3,780 | 35,249 |  $\pm$ 3,179
bowling | 59 |  $\pm$ 0 | 47 |  $\pm$ 7
boxing | 99 |  $\pm$ 0 | 99 |  $\pm$ 0
breakout | 504 |  $\pm$ 165 | 770 |  $\pm$ 12
centipede | 400,268 |  $\pm$ 32,821 | 534,432 |  $\pm$ 38,912
chopper_command | 524,655 |  $\pm$ 154,540 | 660,503 |  $\pm$ 27,000
crazy_climber | 189,621 |  $\pm$ 7,313 | 217,204 |  $\pm$ 12,764
defender | 322,472 |  $\pm$ 105,043 | 483,394 |  $\pm$ 11,589
demon_attack | 131,963 |  $\pm$ 3,819 | 112,140 |  $\pm$ 17,739
double_dunk | 3 |  $\pm$ 4 | -1 |  $\pm$ 1
enduro | 0 |  $\pm$ 0 | 132 |  $\pm$ 86
fishing_derby | -97 |  $\pm$ 0 | -52 |  $\pm$ 29
freeway | 0 |  $\pm$ 0 | 0 |  $\pm$ 0
frostbite | 3,439 |  $\pm$ 1,401 | 8,049 |  $\pm$ 526
gopher | 121,984 |  $\pm$ 338 | 120,551 |  $\pm$ 923
gravitar | 2,807 |  $\pm$ 123 | 3,927 |  $\pm$ 54
hero | 7,877 |  $\pm$ 960 | 9,871 |  $\pm$ 523
ice_hockey | -6 |  $\pm$ 4 | -11 |  $\pm$ 3
jamesbond | 23,475 |  $\pm$ 1,586 | 13,668 |  $\pm$ 4,480
kangaroo | 9,659 |  $\pm$ 2,389 | 10,465 |  $\pm$ 2,835
krull | 11,259 |  $\pm$ 173 | 11,295 |  $\pm$ 108
kung_fu_master | 55,242 |  $\pm$ 4,267 | 83,705 |  $\pm$ 6,565
montezuma_revenge | 0 |  $\pm$ 0 | 0 |  $\pm$ 0
ms_pacman | 40,263 |  $\pm$ 387 | 43,700 |  $\pm$ 1,042
name_this_game | 76,604 |  $\pm$ 7,107 | 94,974 |  $\pm$ 9,942
phoenix | 67,119 |  $\pm$ 9,747 | 49,919 |  $\pm$ 10,573
pitfall | -2 |  $\pm$ 1 | -24 |  $\pm$ 7
pong | -7 |  $\pm$ 9 | -6 |  $\pm$ 9
private_eye | 193 |  $\pm$ 101 | -6 |  $\pm$ 228
qbert | 64,732 |  $\pm$ 8,619 | 70,593 |  $\pm$ 16,955
riverraid | 27,688 |  $\pm$ 1,001 | 28,026 |  $\pm$ 1,823
road_runner | 151,639 |  $\pm$ 90,186 | 571,829 |  $\pm$ 106,184
robotank | 53 |  $\pm$ 2 | 25 |  $\pm$ 8
seaquest | 27,530 |  $\pm$ 10,632 | 141,725 |  $\pm$ 48,000
skiing | -27,968 |  $\pm$ 1,346 | -30,062 |  $\pm$ 248
solaris | 1,544 |  $\pm$ 140 | 1,501 |  $\pm$ 193
space_invaders | 3,962 |  $\pm$ 102 | 5,367 |  $\pm$ 953
star_gunner | 663,896 |  $\pm$ 80,698 | 547,226 |  $\pm$ 126,538
surround | 7 |  $\pm$ 0 | 6 |  $\pm$ 1
tennis | -23 |  $\pm$ 0 | 0 |  $\pm$ 0
time_pilot | 267,331 |  $\pm$ 15,256 | 228,282 |  $\pm$ 10,844
tutankham | 134 |  $\pm$ 10 | 150 |  $\pm$ 8
up_n_down | 434,746 |  $\pm$ 3,905 | 432,240 |  $\pm$ 4,221
venture | 0 |  $\pm$ 0 | 0 |  $\pm$ 0
video_pinball | 376,660 |  $\pm$ 37,647 | 378,897 |  $\pm$ 26,486
wizard_of_wor | 79,425 |  $\pm$ 1,458 | 54,093 |  $\pm$ 6,294
yars_revenge | 317,803 |  $\pm$ 62,785 | 423,271 |  $\pm$ 54,094
zaxxon | 15,752 |  $\pm$ 231 | 15,790 |  $\pm$ 196
Table 3: Final Atari scores for our deep RL experiments. We report the mean of
the final 200 episodes over all three seeds, and the standard error of the
mean across seeds.