# What model does MuZero learn?

Jinke He
Department of Intelligent Systems
Delft University of Technology
The Netherlands
<EMAIL_ADDRESS>
&Thomas M. Moerland
Leiden Institute of Advanced Computer Science
Leiden University
The Netherlands
<EMAIL_ADDRESS>
&Frans A. Oliehoek
Department of Intelligent Systems
Delft University of Technology
The Netherlands
<EMAIL_ADDRESS>
###### Abstract

Model-based reinforcement learning has drawn considerable interest in recent
years, given its promise to improve sample efficiency. Moreover, when using
deep-learned models, it is potentially possible to learn compact models from
complex sensor data. However, the effectiveness of these learned models,
particularly their capacity to plan, i.e., to improve the current policy,
remains unclear. In this work, we study MuZero, a well-known deep model-based
reinforcement learning algorithm, and explore how far it achieves its learning
objective of a value-equivalent model and how useful the learned models are
for policy improvement. Amongst various other insights, we conclude that the
model learned by MuZero cannot effectively generalize to evaluate unseen
policies, which limits the extent to which we can additionally improve the
current policy by planning with the model.

##  1 Introduction

In recent years, reinforcement learning (RL) has made remarkable progress,
finding applications in various real-world applications, including video
compression (Mandhane et al., 2022), chip design (Mirhoseini et al., 2021) and
controlling nuclear fusion (Degrave et al., 2022). A particular class of RL
methods, known as model-based reinforcement learning (MBRL) (Sutton, 1991;
Sutton and Barto, 2018; Moerland et al., 2023), makes additional updates to
the policy through planning, with the potential to improve data efficiency.
The combination of deep learning and MBRL has given rise to a class of _deep
model-based RL (DMBRL)_ approaches. Unlike tabular methods that aim to learn
dynamics and reward for each state-action pair, DMBRL approaches typically
focus on learning a state representation, on top of which dynamics or reward
functions are estimated (Corneil et al., 2018; Ha and Schmidhuber, 2018;
Hafner et al., 2021; Van der Pol et al., 2020; Schrittwieser et al., 2020; Ye
et al., 2021). However, determining the relevant information to include in
these states and how to learn them remains a challenge, as not all aspects of
observations are relevant, and some relevant aspects may be missing, leading
to history dependence McCallum (1995).

One approach to address this challenge is the development of value equivalent
models (Grimm et al., 2020, 2021, 2022). These models are specifically trained
to accurately predict the (multi-step) Bellman update, focusing solely on
value relevant aspects of the task, which is what we ultimately care about. A
well-known example of a value-equivalent-based algorithm is MuZero, which has
achieved state-of-the-art and superhuman performance in Atari games and the
game of Go, Chess, and Shogi (Schrittwieser et al., 2020). MuZero inherits
much of its structure from its predecessor AlphaZero (Silver et al., 2018): it
acts through Monte Carlo Tree Search (MCTS), where the search is guided by
both a learned policy network (to bias action selection) and a learned value
network (to generate value estimates at leaf nodes which replace the rollouts
of standard MCTS).

Model-based RL approaches, such as MuZero, therefore seem to have both
theoretical underpinning and empirical successes. However, while empirical
evaluations have focused on learning curves, we lack a deeper understanding of
the underlying mechanisms driving these empirical successes. The most relevant
study in this direction is by Hamrick et al. (2021), who investigate the role
of planning in MuZero’s learning. They discover that planning with the learned
model primarily boost the learning of the policy network by generating more
informative data and constructing better training targets. Surprisingly, their
findings reveal that in most domains, planning at evaluation time does not
significantly improve performance compared to using the policy network alone,
even with large search budgets. One explanation is that the policy network has
converged to the optimal policy, rendering additional planning useless.
Alternatively, it could suggest that the model may not have converged to a
good (value equivalent) model. If the latter is the case, then the question
becomes what does MuZero’s model actually learn, and why does it give a
performance benefit during training?

In this work, we take a deeper look at MuZero’s learned models. We empirically
investigate two fundamental questions: _1) To what extent does MuZero learn a
value equivalent model?_ and _2) To what extent does MuZero learn a model that
supports effective policy improvement (through planning)?_ Addressing these
questions can help us gain a better understanding of MuZero’s empirical
success and inform the design of future algorithms or extensions.

Our empirical evaluations isolate different aspects of MuZero’s model
throughout training, yielding the following findings:  (i) MuZero’s learned
model is generally not accurate for policy evaluation, and  (ii) the accuracy
of the model for policy evaluation decreases as the policy to evaluate
deviates further from MuZero’s data collection policy  . When we turn to
policy improvement, we observe that  (a) the inability of the learned model to
accurately evaluate policies, especially those that are different to the data
collection policy, limits the extent to which we can find a good policy via
planning, and  (b) introduction of the policy prior alleviates the problem, by
biasing the search towards actions where the model is more accurate. This
implies a second function of the policy network in MuZero: not only does it
accelerate MCTS, but it also restricts the amount of model error that enters
the search.

The rest of the paper is structured as follows. Section 2 introduces the
essential background on Markov Decision Processes, MuZero, and the value
equivalence principle. Section 3 studies trained MuZero models in the policy
evaluation setting (the objective for which the models were trained). Section
4 extends the analysis to the policy improvement setting (planning). Finally,
Section 5 discusses the limitation and outlook of this work.

##  2 Background

###  2.1 Markov Decision Processes

A Markov decision process (MDP) describes the interaction between a decision-
making agent and an environment. Formally, an infinite-horizon discounted MDP
is a $6$-tuple
$\mathcal{M}=(\mathcal{S},\mathcal{A},\mathcal{T},\mathcal{R},\mu,\gamma)$.
Here $\mathcal{S}$ represents the state space of the environment and
$\mathcal{A}$ denotes the set of actions that the agent can take. The initial
state of the environment, $s_{0}$, follows a distribution
$\mu\in\Delta(\mathcal{S})$. At each time step $t$, the agent observes the
environment state $s_{t}$ and selects an action $a_{t}$, causing the
environment to transition to a new state
$s_{t+1}\sim\mathcal{T}(\cdot|s_{t},a_{t})$ and return a numerical reward
$r_{t}=\mathcal{R}(s_{t},a_{t})$. Given a stationary policy
$\pi:\mathcal{S}\rightarrow\Delta(\mathcal{A})$, the value function $V$
measures the expected discounted sum of future rewards from a state following
$\pi$ afterwards,
$V^{\pi}(s)=E[\sum_{t=0}^{\infty}\gamma^{t}r_{t}|\pi,s_{0}=s]$, where
$\gamma\in[0,1)$ is the discount factor. In this work, we specifically focus
on deterministic environments, where the next state $s_{t+1}$ is solely
determined by the current state $s_{t}$ and action $a_{t}$ with probability
$1$. For convenience, we overload the notation $\mathcal{T}$ to represent
deterministic transitions: $s_{t+1}=\mathcal{T}(s_{t},a_{t})$.

###  2.2 MuZero

Figure 1: An illustration of MuZero’s loss function.

MuZero (Schrittwieser et al., 2020) is a recent MBRL method that has achieved
state-of-the-art performance in Atari games and matched the superhuman
performance of AlphaZero (Silver et al., 2018) in Go, chess and shogi.

Components Similar to other MBRL methods, MuZero learns a deterministic world
model that consists of a representation function $h_{\theta}$ and a dynamics
function $g_{\theta}$. The representation function encodes an environment
state $s_{t}$ into a latent state $z_{t}^{0}=h_{\theta}(s_{t})$. Here the
subscript denotes the real time step at which the encoding occurs, and the
superscript denotes the number of time steps that have been spent in the
learned model since then. Given a latent state $z_{t}^{k}$ and an action
$a_{t+k}$, the dynamics function predicts the next latent state $z_{t}^{k+1}$
and the reward $u_{t}^{k}$,
$(z_{t}^{k+1},u_{t}^{k})=g_{\theta}(z_{t}^{k},a_{t+k})$. Apart from the
representation and dynamics functions, MuZero uses a prediction function to
predict the value and policy at a latent state,
$\pi_{t}^{k},v_{t}^{k}=f_{\theta}(z_{t}^{k})$. For ease of notation, we split
the prediction function into a policy function $pi_{\theta}(z_{t}^{k})$ and a
value function $v_{\theta}(z_{t}^{k})$. To distinguish between the policy
function $pi_{\theta}$ and the MuZero policy $pi^{\text{MuZero}}$, which runs
MCTS, from now on we will refer to the former as the policy prior and the
latter as MuZero’s behavior policy or the data generation policy.

Acting MuZero makes decisions by planning with the learned model. At every
time step $t$, MuZero encodes the environment state $s_{t}$ into a latent
state $z_{t}^{0}$ and uses it as the root node to perform a variant of MCTS.
As the result of the search, MuZero selects the action $a_{t}$ by sampling
from a distribution that is constructed using the visit counts at the root
node and a temperature parameter $T\in(0,\infty)$:

$\pi^{\text{MuZero}}(a|s_{t})=\frac{N(z_{t}^{0},a)^{1/T}}{\sum_{b}{N(z_{t}^{0},b)^{1/T}}}$
(1)

MuZero’s search method differs from traditional MCTS methods like UCT (Kocsis
and Szepesvári, 2006) in two main aspects. First, instead of using random
rollouts for estimating the value of a leaf node, MuZero employs bootstrapping
from its value function $v_{\theta}$. Second, MuZero incorporates the policy
prior into its action selection rule to determine which action to simulate at
a tree node $z$:

$\arg\max_{a}\big{[}Q(z,a)+c\cdot\pi_{\theta}(a|z)\cdot\frac{\sqrt{\sum_{b}N(z,b)}}{1+N(z,a)}\big{]}$
(2)

where $c=c_{1}+\log{(\frac{\sum_{b}N(z,b)+c_{2}+1}{c_{2}})}$ with $c_{1}=1.25$
and $c_{2}=19652$. $Q$ and $N$ are the estimated values and visit counts of
actions. MuZero inherits much of this search mechanism from AlphaZero,
including the use of policy prior and value function, which accelerate the
local search. A recent study by Grill et al. (2020) shows that the use of
policy prior in MCTS effectively makes the action visit distribution at the
root node track the solution of a local regularized policy optimization
problem.

Training The key difference between MuZero and many prior work in MBRL lies in
their approach to learning the model. In Figure 1, we illustrate the loss
function of MuZero. Essentially, given a segment of a real episode that starts
from state $s_{t}$, MuZero unrolls its model for $K$ simulated time steps
(below the dotted line) and compares these to real experiences (above the
dotted line). This comparison results in a loss consisting of three terms for
each of the $K$ steps. The first term is the per-step reward prediction loss,
with the target being the real reward received. The second term is the policy
prediction loss, with the target being the action visit distribution of the
MCTS at the root node. The third term is the value prediction loss, with the
target being the discounted sum of $n$-step real rewards plus a value
estimation bootstrapped from MCTS $n$ steps into the future
$v^{\text{target}}_{t}=\sum_{k=0}^{n-1}\gamma^{k}r_{t}+\gamma^{n}v^{\text{MCTS}}_{t}$.
For simplicity, we omit the links for value targets in the figure. All
components of MuZero are trained jointly end-to-end by minimizing the
aggregated loss. In practice, MuZero incorporates various additional
techniques from the literature, such as prioritized experience replay, to
improve training. It also introduces a novel algorithm called ‘reanalyse’ to
generate fresh targets from old trajectories by re-running MCTS on them using
the latest network.

###  2.3 The value equivalence principle

The value equivalence principle, introduced by Grimm et al. (2020, 2021), is
motivated by the consideration that the construction of a model should take
into account the final use of the model. It defines the order-$k$ value
equivalence class $\mathcal{M}^{k}(\Pi,\mathcal{V})$ as the subset of all
models that can predict the correct $k$-step bellman update for any
$\pi\in\Pi$ and $v\in\mathcal{V}$ in a set of policies $\Pi$ and functions
$\mathcal{V}$. For $k\rightarrow\infty$, an additional proper value
equivalence class $\mathcal{M}^{\infty}(\Pi)$ is defined, which excludes the
set of functions from the specification. In essence, a model in the proper
value equivalence class must have the true value function as the fixed point
of the bellman operators for all policies in this set. As such, these models
can be seen as the $Q^{\Pi}$-irrelevance abstractions for the MDP (Li et al.,
2006). Notably, the largest proper value equivalence class that guarantees
optimal planning is $\mathcal{M}^{\infty}(\mathbf{\Pi}^{\textnormal{DET}})$,
where $\mathbf{\Pi}^{\textnormal{DET}})$ is the set of all deterministic
policies. Moreover, Grimm et al. (2021) show that a simplified version of
MuZero’s loss upper bounds the proper value equivalence loss, making an
explicit connection between the theory of value equivalence principle and the
strong empirical performance of MuZero. Despite the promise of the value
equivalence principle in explaining MuZero’s strong performance, there are
gaps between the theory and the practical algorithm. One such gap is that the
models in the value equivalence principle are constructed on the original
state space of the environment, while MuZero’s model operates on a learned
latent state representation that evolves jointly with the model.

##  3 Policy Evaluation Experiments

We are interested in the extent to which MuZero’s learned model supports
additional policy improvement, which is crucial for MuZero’s sample efficiency
as a MBRL method. Our hypothesis is that, since MuZero’s model is trained on
data collected by previous policies, it is not generally value equivalent for
all policies, especially those that have not been executed. As accurate policy
evaluation is the basis for effective policy improvement, this will limit the
extent to which MuZero can additionally improve its policy through planning.
In this section, we verify this hypothesis.

###  3.1 Training MuZero agents

For our empirical analysis, we used three fully observable deterministic
environments (as MuZero was designed for this setting): Cart Pole, a
deterministic version of Lunar Lander, and Atari Breakout. We trained $30$
MuZero agents with different random seeds for Cart Pole and Lunar Lander, and
$10$ for Breakout. For each agent, we saved the weights at different training
steps. In the figures below, we report our results using their means along
with the corresponding standard errors.

In Cart Pole and Lunar Lander, we extensively trained the MuZero agents, for
$100$K and $1$M steps. This way, we can conduct our analysis on the trained
agents throughout their lifecycle of learning. For Breakout, we adopted the
same setup as EfficientZero (Ye et al., 2021) and trained the agents for 100K
environment steps. It is worth noting that this is far less than the amount of
training required for convergence (the original MuZero paper used $200$M). We
trained for $100$K steps because training $10$ MuZero agents this long already
took more than $400$ GPU hours and $4800$ CPU hours, making it prohibitive for
us to train for longer. Moreover, as Atari-100K is a widely adopted benchmark
for sample efficiency, we believe assessing MuZero’s learned model in this
setting will bring insights.

Figure 2: Learning curve of MuZero agent. For the same agent, we plot the
performance of running full MuZero (with planning) and running policy prior
(without planning).

Figure 2 shows the performance of the MuZero agent, which makes decisions by
performing MCTS planning (see Section 1). For reference, we also plot the
performance of the policy prior of the same agent, which samples actions
directly from the policy network,
$a_{t}\sim\pi_{\theta}(\cdot|h_{\theta}(s_{t}))$. In Cart Pole and Lunar
Lander, the MuZero agents have almost converged to optimal performance.
Throughout training, planning enables MuZero to achieve a substantially better
performance than directly using the policy prior. In Breakout, this effect
only becomes evident toward the end of our training. This discrepancy could be
attributed to the agents still being in the early stages of learning in
Breakout, whereas in Cart Pole and Lunar Lander, the agents have experienced
the full lifecycle of training. It is important to note that the improved
performance may not only come from the model but may also come from the value
network, which is used to estimate the value of leaf nodes in MuZero’s MCTS.
In the following, we will take a deeper look at how much the model contributes
to this.

###  3.2 Evaluating the learned model

To assess MuZero’s learned model in the context of policy evaluation, we will
compare it to the ground-truth model. Through experiments, we aim to answer
questions in the format of how much value prediction error should we expect
when using MuZero’s learned model to evaluate a policy $\pi$?

Considering that MuZero does not train or employ the model for infinite-
horizon rollouts, we impose a limit on the evaluation horizon when assessing
the value prediction error. For each state $s_{t}$, we define the discounted
sum of future rewards for taking an action sequence $(a_{t},\ldots,a_{t+h-1})$
as:

$v^{a_{t:t+h-1}}(s_{t})~{}=~{}\sum_{k=0}^{h-1}\gamma^{k}r_{t+k}$ (3)

Here, $r_{t+k}=\mathcal{R}(s_{t+k},a_{t+k})$ and
$s_{t+k+1}=\mathcal{T}(s_{t+k},a_{t+k})$, assuming a deterministic
environment. As mentioned in Section 1, MuZero predicts the value of this
action sequence by first encoding the state into a latent state
$z_{t}^{0}=h_{\theta}(s_{t})$ and then rolling out the model from the latent
state:

$\hat{v}^{a_{t:t+h-1}}(s_{t})=\sum_{k=0}^{h-1}\gamma^{k}u_{t}^{k}$ (4)

Here, $z_{t}^{k+1},u_{t}^{k}=g_{\theta}(z_{t}^{k},a_{t+k})$. Then, we can
define the value prediction error of the learned model.

######  Definition 1.

The value prediction error of using the learned model to evaluate the action
sequence $a_{t:t+h-1}$ at state $s_{t}$ is:

$\lvert v^{a_{t:t+h-1}}(s_{t})-\hat{v}^{a_{t:t+h-1}}(s_{t})\rvert$ (5)

A stationary policy $\pi:\mathcal{S}\rightarrow\Delta(\mathcal{A})$ that
operates in the original environment defines a distribution over action
sequences of length $h$:

$\Pr(a_{t:t+h-1}|s_{t},\pi)=\prod_{k=0}^{h-1}\pi\big{(}a_{t+k}|s_{t+k}=\mathcal{T}(s_{t+k-1},a_{t+k-1})\big{)}$
(6)

Based on this, we can define the value prediction error for evaluating $\pi$.

######  Definition 2.

The value prediction error of using the learned model to evaluate a stationary
policy $\pi$, which operates in the original environment, for horizon $h$ is:

$\lvert
v^{\pi}_{h}(s_{t})-\hat{v}^{\pi}_{h}(s_{t})\rvert=\left|\mathop{\mathbb{E}}_{a_{t:t+h-1}\sim\Pr(\cdot|s_{t},\pi)}\big{[}v^{a_{t:t+h-1}}(s_{t})-\hat{v}^{a_{t:t+h-1}}(s_{t})\big{]}\right|$
(7)

###  3.3 How accurately can MuZero’s learned model predict the value of its
own behavior policy?

As the model is trained on data collected by MuZero’s behavior policy
$\pi^{\text{MuZero}}$, we expect it to be at least accurate on this data
collection policy. Therefore, we begin our investigation by examining the
model’s prediction performance on this policy. Due to the continuous state
spaces of our environments, it is not feasible to enumerate all states and
compute the error for each one. Instead, we sample states from MuZero’s on-
policy state distribution $d_{\pi^{\text{MuZero}}}$. For each sampled state,
we conduct the evaluation and aggregate the errors. To facilitate tractable
evaluation, we use Monte Carlo sampling to estimate both the true value
$v_{h}^{\pi^{\text{MuZero}}}(s)$ and the value predicted by MuZero’s model
$\hat{v}_{h}^{\pi^{\text{MuZero}}}(s)$.

In Figure 3, we report the value prediction error (Y-axis) under various
evaluation horizons (lines) across different training steps (X-axis), with the
maximum evaluation horizon set to the number of unrolling steps during
training. Short horizons generally exhibit high accuracy, with errors
approaching zero. However, as the evaluation horizon increases, the errors
consistently grow larger. This aligns with the known tendency of learned
models to accumulate errors during longer rollouts. Importantly, it indicates
that the model learned by MuZero is not fully value equivalent as it is not
even accurate enough to predict values for its data collection policy. Errors
at different training steps are typically incomparable due to the different
state distributions, but the decreasing errors observed in Cart Pole and Lunar
Lander suggest the convergence of the policy as a possible explanation.

Figure 3: Value prediction error of using MuZero’s learned model to evaluate
its own behavior policy.

###  3.4 How accurately can MuZero’s learned model evaluate policies that are
different from the behavior policy?

To assess whether MuZero’s learned model effectively supports planning, we
will investigate its capacity to generalize and accurately predict values
beyond its own data collection policy. Our hypothesis is that the model will
exhibit increasing inaccuracies when evaluating policies that differ
significantly from the behavior policy, which is responsible for collecting
the training data.

To test this hypothesis, we conduct an experiment focusing on the relationship
between the value prediction error for an action sequence
$|v^{a_{t:t+h-1}}(s_{t})-\hat{v}^{a_{t:t+h-1}}(s_{t})|$ and the probability of
the behavior policy selecting this action sequence
$\Pr(a_{t:t+h-1}|s_{t},\pi^{\text{MuZero}})$.

We sample states from MuZero’s on-policy state distribution
$d_{\pi^{\text{MuZero}}}$ and compute the probabilities and value prediction
errors for all action sequences of length $\\{8,4,4\\}$ for Cart Pole, Lunar
Lander and Breakout (limited by computation budget). Considering that both
probabilities and errors are real-valued and non-uniformly distributed, we
aggregate the results across different states and agents by sorting the action
sequences based on their probabilities of being chosen by MuZero’s behavior
policy. We then aggregate statistics for action sequences with the same rank.

Results are present in Figure 4, where the X-axis represents action sequences
sorted by their probabilities of being taken by the current behavior policy,
ranging from small to large. On the Y-axis, we plot both the probabilities
(blue) and the corresponding value prediction errors (yellow). Note that in
Breakout, the Y-axis for errors is not shared across training steps due to the
significant variation in the error scale. The results clearly reveal an
inverse relationship between the value prediction error for an action sequence
and the probability of it being taken by the behavior policy. This finding
supports our hypothesis that the model is more reliable in predicting values
for the behavior policy, which collects the training data. Consequently,
evaluating a policy using MuZero’s learned model may yield increasingly
inaccurate results as the policy further deviates from the data collection
policy.

Figure 4: X-axis: action sequences sorted by their probabilities of being
taken by the behavior policy (from small to large). Y-axis: the probabilities
(blue) and the corresponding value prediction errors (yellow). Rows: Cart
Pole, Lunar Lander and Atari Breakout. Higher the probability, lower the
error.

###  3.5 How accurately can MuZero’s learned model from one training step
evaluate the behavior policy of the same agent from another training step?

In our last experiment, we investigated the generalization capability of the
learned model in predicting values for policies that differ from the behavior
policy. We accomplished this by considering all action sequences of length $h$
at each state, possibly including actions that are unlikely to be taken by any
sensible policy. In this experiment, we conduct a similar analysis but with a
focus on more interesting policies. Specifically, we assess the accuracy with
which the model at training step X (row) can evaluate the policy (of the same
agent) at training step Y (column). The idea is that if a model cannot
accurately evaluate high-performing future policies, planning with it would
not be effective in finding a good policy. For this experiment, we set the
evaluation horizon for each environment to the number of unrolling steps
during training. We aggregate results over states sampled from MuZero’s on-
policy state distribution at training step X (same per row as the model). Note
that this means errors at different model steps (rows) are not directly
comparable to each other.

(a) Cart Pole.

(b) Lunar Lander.

(c) Atari Breakout.

Figure 5: Cross model policy evaluation. We evaluate MuZero’s behavior policy
at training step Y (column) with the learned model at training step X (row)
and measure the value prediction error. Results are aggregated over states
sampled from MuZero’s behavior policy at training step X.

Results in Figure 5 reveal that models at all training steps are most accurate
when evaluating policies at the same training steps, as error per row is
smallest at the diagonal. This aligns with our earlier finding, which
indicates that learned models are more accurate when assessing action
sequences with higher likelihood of selection by the current behavior policy.
Crucially, in Cart Pole and Lunar Lander, using early models to evaluate
future policies always results in large errors, which may have important
implications for policy improvement: _if the model at the current training
step cannot accurately evaluate a future policy, which performs better, then
the extent to which we can improve our current policy will be limited._
Intuitively, if the model does not know a high-performing policy is good, then
deeper search in the model would not help us find the policy. This phenomenon
is not evident in Breakout, possibly because the evaluation horizon is so
short here, i.e., $5$, that it does not allow for an observable difference in
rewards across the policies (Breakout also has a sparser reward).

##  4 Policy Improvement Experiments

Our policy evaluation experiments indicate that MuZero’s model can become
increasingly inaccurate when evaluating policies differing from the data
collection policy, particularly those unseen during training. In this section,
we investigate the natural follow-up question, the question that we are most
interested in: what is the effect of this on policy improvement through
planning?

For efficient planning under limited budget, AlphaZero and MuZero use the
learned policy prior to guide MCTS. We evaluate planning using MuZero’s
learned model with and without the guidance of the policy prior. In the latter
case, we replace the policy prior with a uniform prior, allowing for a form of
“free search”. The goal of this experiment is to answer two questions:  (a)
how effectively does MuZero’s learned model support free search?  (b) to what
extent can MuZero improve its policy by planning with the learned model?  To
isolate the impact of the learned model, we modify MuZero’s MCTS by replacing
the value predictions at leaf nodes with random rollouts using the model. To
address computational costs and error accumulation in neural models during
long rollouts, we restrict the planning horizon to $\\{16,4,32\\}$ for Cart
Pole, Lunar Lander, and Breakout. For comparison, we include the baseline of
planning with the ground-truth model using the same MCTS.

Figure 6: MCTS planning with (i) MuZero’s learned model (orange) and (ii) the
ground-truth model (green), under the guidance of (a) MuZero’s policy prior
(solid line) and (b) the uniform prior (dotted line), in Cart Pole (Top),
Lunar Lander (Middle), and Atari Breakout (Bottom, in log scale). The X-axis
is the number of simulations per step and the Y-axis is the episodic return.

Results are shown in Figure 6. In Cart Pole and Breakout, MuZero’s learned
model exhibits some degree of support for free search (orange dashed) after
some training. However, it significantly underperforms compared to free search
with the ground-truth model (green dashed) in terms of both planning
efficiency and asymptotic performance. For example, in Cart Pole at training
step $2000$, the performance of free search with the ground-truth model can
continue to improve with more simulations, potentially even with $64$
simulations, while the learned model’s planning performance becomes stagnant
with a much smaller budget. We posit that this limitation arises from the
model’s inability to accurately evaluate unfamiliar policies due to the lack
of data. Consequently, the potential for finding a good policy through
planning with the learned model will be effectively limited.

In Cart Pole, with enough MCTS simulations, planning with the learned model
does improve the performance over using the policy prior alone. However, it is
also evident that the extent to which we can improve upon the policy prior via
planning is still limited compared to the ground-truth model (orange solid vs.
green solid). This is a clear sign that the model error is restricting the
extent to which we can further improve the policy by planning. In Breakout,
while planning with more simulations using the learned model improves the
agent’s performance, it cannot outperform the policy prior.

One interesting observation is that when comparing planning with the policy
prior to planning with the uniform prior, there is consistently a larger gap
for the learned model compared to the ground-truth model except in Breakout.
Furthermore, unlike the ground-truth model, it does not seem that the gap for
the learned model will quickly collapse with more simulations, suggesting an
additional role of the policy prior beyond accelerating MCTS. According to the
results in Figure 4, action sequences with higher probabilities of being
selected by the behavior policy tend to be more accurately evaluated by the
learned model. As the policy prior is directly learned to match the behavior
policy, it is reasonable to assume that the values of action sequences favored
by the policy prior can also be more accurately predicted by the learned
model. This suggests that _apart from biasing the search, the policy prior may
also serve to prevent the search from exploring directions where the learned
model is less accurate._

In Lunar Lander, doing planning is expensive, and we have to use a short
planning horizon to deal with this computational complexity. Possibly due to
this reason, planning—even with the ground-truth model—cannot outperform the
policy prior after it gets enough training (rightmost plot). The learned model
performs very bad in this environment and fails to support any kind of
planning. Interestingly, it does not seem that more training improves the
situation. Since we do observe that more training reduces the value prediction
error for the behavior policy, one explanation could be that the increasing
inaccuracy of the model to evaluate other policies make it even less able to
carry out effective planning. This suggests that for policy search or
improvement, it is not enough to have a model that is only accurate to predict
the value of the current behavior policy.

##  5 Discussion

##### Limitation

In this study, we analyzed MuZero, it being one of the most successful DMBRL
methods, while at the same time being supported by the clean principle of
value equivalence. The conclusions we draw can therefore not be directly
generalized to other DMBRL methods, certainly not those that employ other
auxiliary loss functions on top e.g., (Gelada et al., 2019; Van der Pol et
al., 2020; Ye et al., 2021). However, our analysis clearly shows that despite
its elegant simplicity, using the value equivalence principle does not imply
that we learn models that are actually value equivalent. Even when restricting
the predictions to its own behavior policy, MuZero’s learned model leads to
prediction errors that quickly grow with the horizon of prediction. As such,
our study does serve as a warning for other DMBRL methods that aim to use only
value equivalence as their guiding principle.

##### Outlook

We have demonstrated that MuZero’s model, trained with a value-equivalence-
based loss, struggles to generalize and predict values accurately for unseen
or unfamiliar policies. This raises the question of how would different losses
for model learning would compare in our analysis, such as the reconstruction-
based loss and the temporal consistency loss. Specifically, we hypothesize
that while the value equivalence loss is the most flexible loss as it does not
directly impose requirements on state representation, other losses may have an
advantage in generalization when data is limited. Another exciting direction
for future research is to explore the role of the learned model beyond policy
improvement through planning, such as its impact on exploration.

##  6 Conclusion

In this work, we empirically studied the models learned by MuZero. Our
analysis focuses on two questions:  (1) to what extent does MuZero learn a
value equivalent model? and  (2) to what extent does the learned model support
effective policy improvement by planning?  We find that MuZero’s learned model
cannot generally evaluate policies accurately, especially those that further
deviate from the data collection policy. The failure of the model to predict
values for policies that are out of the training distribution prevents
effective planning from scratch and limits the extent to which MuZero can
additionally improve its policy via planning, which might be crucial for its
sample efficiency.

## References

  * Corneil et al. (2018) Dane Corneil, Wulfram Gerstner, and Johanni Brea.  Efficient model-based deep reinforcement learning with variational state tabulation.  In _ICML_ , pages 1049–1058, July 2018.
  * Degrave et al. (2022) Jonas Degrave, Federico Felici, Jonas Buchli, Michael Neunert, Brendan Tracey, Francesco Carpanese, Timo Ewalds, Roland Hafner, Abbas Abdolmaleki, Diego de las Casas, Craig Donner, Leslie Fritz, Cristian Galperti, Andrea Huber, James Keeling, Maria Tsimpoukelli, Jackie Kay, Antoine Merle, Jean-Marc Moret, Seb Noury, Federico Pesamosca, David Pfau, Olivier Sauter, Cristian Sommariva, Stefano Coda, Basil Duval, Ambrogio Fasoli, Pushmeet Kohli, Koray Kavukcuoglu, Demis Hassabis, and Martin Riedmiller.  Magnetic control of tokamak plasmas through deep reinforcement learning.  _Nature_ , February 2022.
  * 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 _ICML_ , 2019.
  * Grill et al. (2020) Jean-Bastien Grill, Florent Altché, Yunhao Tang, Thomas Hubert, Michal Valko, Ioannis Antonoglou, and Remi Munos.  Monte-Carlo Tree Search as Regularized Policy Optimization.  In _ICML_ , pages 3769–3778, November 2020.
  * Grimm et al. (2020) Christopher Grimm, Andre Barreto, Satinder Singh, and David Silver.  The Value Equivalence Principle for Model-Based Reinforcement Learning.  In _NeurIPS_ , volume 33, pages 5541–5552, 2020.
  * Grimm et al. (2021) Christopher Grimm, André Barreto, Gregory Farquhar, David Silver, and Satinder Singh.  Proper value equivalence.  In _NeurIPS_ , volume 34, pages 7773–7786, 2021.
  * Grimm et al. (2022) Christopher Grimm, Andre Barreto, and Satinder Singh.  Approximate Value Equivalence.  In _NeurIPS_ , volume 35, October 2022.
  * Ha and Schmidhuber (2018) David Ha and Jürgen Schmidhuber.  World Models.  _arXiv e-prints_ , art. arXiv:1803.10122, March 2018.
  * Hafner et al. (2021) Danijar Hafner, Timothy P Lillicrap, Mohammad Norouzi, and Jimmy Ba.  Mastering atari with discrete world models.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=0oabwyZbOu.
  * Hamrick et al. (2021) Jessica B Hamrick, Abram L Friesen, Feryal Behbahani, Arthur Guez, Fabio Viola, Sims Witherspoon, Thomas Anthony, Lars Holger Buesing, Petar Veličković, and Theophane Weber.  On the role of planning in model-based deep reinforcement learning.  In _International Conference on Learning Representations_ , 2021.
  * Kocsis and Szepesvári (2006) Levente Kocsis and Csaba Szepesvári.  Bandit based Monte-Carlo planning.  In _European conference on machine learning_ , volume 4212 LNAI, pages 282–293, 2006.
  * Li et al. (2006) Lihong Li, Thomas J. Walsh, and Michael L. Littman.  Towards a Unified Theory of State Abstraction for MDPs.  In _International Symposium on Artificial Intelligence and Mathematics, AI &Math 2006, Fort Lauderdale, Florida, USA, January 4-6, 2006_, 2006.
  * Mandhane et al. (2022) Amol Mandhane, Anton Zhernov, Maribeth Rauh, Chenjie Gu, Miaosen Wang, Flora Xue, Wendy Shang, Derek Pang, Rene Claus, Ching-Han Chiang, Cheng Chen, Jingning Han, Angie Chen, Daniel J. Mankowitz, Jackson Broshear, Julian Schrittwieser, Thomas Hubert, Oriol Vinyals, and Timothy Mann.  MuZero with Self-competition for Rate Control in VP9 Video Compression, February 2022.  arXiv:2202.06626.
  * McCallum (1995) Andrew K. McCallum.  _Reinforcement Learning with Selective Perception and Hidden State_.  PhD thesis, University of Rochester, 1995.
  * Mirhoseini et al. (2021) Azalia Mirhoseini, Anna Goldie, Mustafa Yazgan, Joe Wenjie Jiang, Ebrahim Songhori, Shen Wang, Young-Joon Lee, Eric Johnson, Omkar Pathak, Azade Nazi, Jiwoo Pak, Andy Tong, Kavya Srinivasa, William Hang, Emre Tuncer, Quoc V. Le, James Laudon, Richard Ho, Roger Carpenter, and Jeff Dean.  A graph placement methodology for fast chip design.  _Nature_ , 594(7862):207–212, June 2021.
  * Moerland et al. (2023) Thomas M Moerland, Joost Broekens, Aske Plaat, Catholijn M Jonker, et al.  Model-based reinforcement learning: A survey.  _Foundations and Trends® in Machine Learning_ , 16(1):1–118, 2023.
  * Schrittwieser et al. (2020) Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, and Thore Graepel.  Mastering atari, go, chess and shogi by planning with a learned model.  _Nature_ , 588(7839):604–609, 2020.
  * Silver et al. (2018) David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Marc Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, Timothy Lillicrap, Karen Simonyan, and Demis Hassabis.  A general reinforcement learning algorithm that masters chess, shogi, and Go through self-play.  _Science_ , 362(6419):1140–1144, December 2018\.
  * Sutton (1991) Richard S. Sutton.  Dyna, an integrated architecture for learning, planning, and reacting.  _ACM SIGART Bulletin_ , 2(4):160–163, July 1991\.
  * Sutton and Barto (2018) Richard S Sutton and Andrew G Barto.  _Reinforcement Learning: An Introduction_.  MIT press, 2018.
  * 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 _AAMAS_ , May 2020.
  * Werner Duvaud (2019) Aurèle Hainaut Werner Duvaud.  Muzero general: Open reimplementation of muzero.  https://github.com/werner-duvaud/muzero-general, 2019.
  * Ye et al. (2021) Weirui Ye, Shaohuai Liu, Thanard Kurutach, Pieter Abbeel, and Yang Gao.  Mastering Atari Games with Limited Data.  In _NeurIPS_ , volume 34, pages 25476–25488. Curran Associates, Inc., 2021.

##  Appendix A Supplementary Material

###  A.1 Training of MuZero agents

Hyperparameter | Cart Pole | Lunar Lander
---|---|---
Random seeds | 0 to 29 | 0 to 29
Discount factor | 0.997 | 0.999
Total training steps | 1e5 | 1e6
Optimizer | Adam | Adam
Learning Rate | 0.02 | 0.005
LR Decay Rate | 0.1 | No decay
LR Decay Steps | 50000 | No decay
Weight Decay | 1e-4 | 1e-4
Momentum | 0.9 | 0.9
Batch Size | 128 | 64
Encoding size | 8 | 10
Fully-connected Layer Size | 16 | 64
Root Dirichlet Alpha | 0.25 | 0.25
Root Dirichlet Fraction | 0.25 | 0.25
Prioritized Experience Replay Alpha | 0.5 | 0.5
Num Unroll Steps | 10 | 10
TD Steps | 50 | 30
Support Size | 10 | 10
Value Loss Weight | 1.0 | 1.0
Replay Buffer Size | 500 | 500
Visit Softmax Temperature Fn | 1.0 $\rightarrow$ (5e4) 0.5 $\rightarrow$ (7.5e4) 0.25 | 0.35
Table 1: Hyperparameters for training MuZero agents in Cart Pole and Lunar
Lander. We used default values for most of the hyperparameters. The bold
values are those that we tuned to improve the convergence of the agents.

For Cart Pole and Lunar Lander, we trained 30 MuZero agents using an open-
source implementation of MuZero [Werner Duvaud, 2019] available on GitHub
under the MIT license. The implementation has been extensively tested on
various classic RL environments, including Cart Pole and Lunar Lander. While
we mostly used the default recommended hyperparameter values from Werner
Duvaud [2019], we fine-tuned a few of them to improve agent convergence. The
comprehensive list of hyperparameters for training MuZero agents in these
environments can be found in Table 1.

For Atari Breakout, we trained 10 MuZero agents using the official
implementation of EfficientZero [Ye et al., 2021], _excluding_ the additional
improvements introduced by EfficientZero. This implementation is also hosted
on GitHub under the GPL-3.0 license. At the time of this research, this was
the only plausible way to train MuZero agents in Atari games. Apart from
turning off EfficientZero’s additional improvements over MuZero, we used the
exact same hyperparameters used by EfficientZero (see Appendix A.1 Table 6 of
[Ye et al., 2021]).

Regarding computation, training each MuZero agent took around 3 hours in Cart
Pole and 20 hours in Lunar Lander, using 8 CPUs. For Atari Breakout, it took
10 hours per agent using 4 Nvidia GeForce RTX 2080 Ti GPUs and 48 CPUs. The
training was conducted on a shared internal cluster equipped with a variety of
CPU types.

###  A.2 Additional experimental results in Atari Breakout

In the following, we present some additional results in the Atari Breakout
environment.

####  A.2.1 How accurately can MuZero’s learned model from one training step
evaluate the behavior policy of the agent from another training step?

In Section 3.5, we sampled states from the on-policy state distribution at
training step X and investigated how accurately the model at training step X
can evaluate the policy (of the same agent) at training step Y. The reason for
choosing to sample states from the on-policy state distribution at training
step X, rather than Y, was to enable the comparison of value prediction errors
when using the same model to evaluate policies at different training steps.

Now, we sample states from the on-policy state distribution at training step Y
and assess the accuracy with which the policy at training step Y can be
evaluated by the model (of the same agent) at training step X. This way, we
can compare the accuracy of value predictions for the same policy across
models at different training steps.

Figure 7: Cross policy model evaluation. We evaluate MuZero’s behavior policy
at training Y (column) with the learned model at training step X (row) and
measure the value prediction error. Results are aggregated over states sampled
from MuZero’s behavior policy at training step Y.

The results are presented in Figure 7. Note that in contrast to Figure 5, in
Figure 7 the errors within the same columns are comparable instead of the
rows. This allows us to compare the errors across different models rather than
the policies being evaluated.

By examining the errors within the different rows, specifically in the last
two columns, we can observe that it is more accurate to evaluate future
policies (at steps $32000$ and $100000$) using the corresponding future models
(at steps $32000$ and $100000$) rather than using the past models (at steps
$500$ and $2000$). This finding indicates that the learned model struggles to
accurately evaluate unseen policies, making it challenging for planning to
discover these better policies.

Interestingly, we also observe backward compatibility in the learned models:
future models (at steps $32000$ and $100000$) are capable of accurately
evaluating past policies (at steps $500$ and $2000$). This observation
contrasts with our findings in Cart Pole and Lunar Lander (see Figure 5),
where the accuracy of the learned models in evaluating early policies
decreases over time. There are three possible reasons. First, early policies
often yield rewards of $0$, which makes the prediction task easier. Second,
following EfficientZero, we did not remove any data from the replay buffer
during the training of MuZero agents in Breakout, resulting in extremely high
memory consumption. This implies that the learned model is trained on all the
data collected throughout the training process. Third, unlike Cart Pole and
Lunar Lander where the agents have undergone the complete training cycle, the
agents in Breakout are still in the early stages of learning. Consequently,
the learned models have not had sufficient training to overfit to the current
behavior policy that collects the data.

####  A.2.2 Policy improvement experiments with shorter planning horizon

Figure 8: MCTS planning with (i) MuZero’s learned model (orange) and (ii) the
ground-truth model (green), under the guidance of (a) MuZero’s policy prior
(solid line) and (b) the uniform prior (dotted line) in Atari Breakout. The
X-axis is the number of simulations per step and the Y-axis is the episodic
return. Rows represent different planning horizons.

In Section 4, we conducted experiments to investigate to what extent the
learned model supports effective policy improvement through planning. In our
planning experiments, we used a planning horizon of $64$, which was needed for
classic planning methods like UCT to achieve good performance in the Atari
domain. For the completeness of our experimental results, we present results
obtained with various shorter planning horizons in Figure 8.

When using the ground-truth model for planning, it is evident that larger
planning horizons consistently yield improved performance. However, when using
the learned model for planning, larger planning horizons actually result in
worse performance, likely due to the accumulation of more errors during the
planning process. In the case of very short planning horizons, such as 2 or 4,
planning with the learned model does exhibit better performance compared to
the policy prior. However, this improvement is significantly limited when
compared to planning with the ground-truth model using a larger planning
horizon. This finding highlights the challenges and limitations associated
with using learned models for additional policy improvement through planning
in the context of MuZero.