Visualizing MuZero Models
Abstract
MuZero, a model-based reinforcement learning algorithm that uses a value equivalent dynamics model, achieved state-of-the-art performance in
Chess, Shogi and the game of Go.
In contrast to standard forward dynamics models that predict a full next state, value equivalent models are trained to predict a future value, thereby emphasizing value relevant information in the representations.
While value equivalent models have shown strong empirical success, there is no research yet that visualizes and investigates what types of representations these models actually learn.
Therefore, in this paper we visualize the latent representation of MuZero agents.
We ﬁnd that action trajectories may diverge between observation embeddings and internal state transition dynamics, which could lead to instability during planning.
Based on this insight, we propose two regularization techniques to stabilize MuZero’s performance.
Additionally, we provide an open-source implementation of MuZero along with an interactive visualizer of learned representations, which may aid further investigation of value equivalent algorithms.
1. Introduction
Model-based reinforcement learning has shown strong empirical success in sequential decision making tasks, as illustrated by the AlphaZero (Silver et al., 2018) and MuZero algorithms (Schrittwieser et al., 2020).
Both of these approaches nest a planning loop, based on Monte Carlo Tree
Search (Kocsis & Szepesv´ari, 2006;
Browne et al., 2012), inside a learning loop, where we approximate global value and policy functions.
While AlphaZero used a known model of the environment, MuZero learned the model from sampled data.
However, instead of a standard forward dynamics , The model, which learns to predict future states, the MuZero dynamics model is trained to predict future values , better known as a value equivalent model (Grimm et al., 2020).
A potential beneﬁt of value equivalent models, compared to standard forward models, is that they will emphasize value and reward relevant characteristics in their representation and dynamics.
This may be beneﬁcial when the true dynamics are complicated, but the value relevant aspects of the dynamics are comparatively simple.
As a second beneﬁt, we train our model for its intended use: predicting value information during planning.
Several papers have empirically investigated this principle in recent years (Tamar et al., 2016;
Oh et al., 2017; Farquhar et al., 2018; Silver et al., 2017b; Schrittwieser et al., 2020), while (Grimm et al., 2020) provides a theoretical underpinning of this approach.
However, no literature has yet investigated what kind of representations these approaches actually learn, i.e., how the learned representations are organized.
The goal of this paper is therefore to investigate and visualize environment models learned by MuZero.
Most interestingly, we ﬁnd that, after training, an action trajectory that follows the forward dynamics model usually departs from the learned embedding of the environment observations.
In other words, MuZero is not enforced to keep the state encoding and forward state prediction congruent.
Therefore, the second goal of this paper is to regularize MuZero’s dynamics model to improve its structure.
We propose two regularization objectives to add to the MuZero objective, and experimentally show that these may indeed provide beneﬁt.
In short, after introducing related work (Sec. 2) and necessary background on the MuZero algorithm (Sec. 3), we discuss two research questions: 1) what type of representation do value equivalent models learn (Sec. 4), and 2) can we use regularization to better structure the value equivalent latent space (Sec. 5)?
We experimentally validate the second question in Sec. 6 and 7.
Moreover, apart from answering these two question, we also open source modular MuZero code including an interactive visualizer of the latent space based on principal component analysis (PCA), available from www.anonymized.org .
We found the visualizer to greatly enhance our understanding of the algorithm, and believe visualization will be essential for deeper understanding of this class of algorithms.
2. Related Work
Value equivalent models, a term introduced by (Grimm et al., 2020), are usually trained on end-to-end differentiable computation graphs, although the principle would be applicable to gradient-free optimization as well.
Typically, the unrolled computation graph makes multiple passes through a dynamics model, and eventually predicts a value.
Then, the dynamics model is trained through gradient descent on its ability to predict the correct value.
The ﬁrst value equivalent approach were Value Iteration Networks (VIN) (Tamar et al., 2016), where a differentiable form of value iteration was embedded to predict a value.
Other variants of value equivalent approaches are Value Prediction
Networks (VPN) (Oh et al., 2017), TreeQN and ATreeC (Farquhar et al., 2018), the Predictron (Silver et al., 2017b), and MuZero (Schrittwieser et al., 2020).
These methods differ in the way they build their computation graph, where
VINS and TreeQN embed entire policy improvement (planning) in the graph, VPNs, the Predictron and MuZero only perform policy evaluation.
Therefore, the latter approaches combine explicit planning for policy improvement, which in the case of MuZero happens through MCTS.
(Grimm et al., 2020) provides a theoretical analysis of value equivalent models, showing that two value equivalent models give the same Bellman back-up.
MuZero uses the learned value equivalent model to explicitly plan through Monte Carlo Tree Search (Kocsis &
Szepesv ´ ari, 2006;
Browne et al., 2012), and uses the output of the search as training targets for a learned policy network.
This idea of iterated planning and learning dates back to
Dyna-2 (Silver et al., 2008), while the particularly successful combination of MCTS and deep learning was introduced in AlphaGo Zero (Silver et al., 2017a) and Expert Iteration (ExIt) (Anthony et al., 2017).
In general, planning may add to pure (model-free) reinforcement learning: 1) improved action selection, and 2) improved (more stable) training targets.
On the other hand, learning adds to planning the ability to generalize information, and store global solutions in memory.
For more detailed overviews of value equivalent models and iterated planning and learning we refer to the model-based RL surveys by (Moerland et al., 2020; Plaat et al., 2020).
Visualization is a common approach to better understand machine learning methods, and visualization of representations and loss landscapes of (deep) neural networks has a long history (Bischof et al., 1992;
Yosinski et al., 2015; Karpathy et al., 2015; Li et al., 2018a).
For example, (Li et al., 2018b) shows how the loss landscape of a neural network can indicate smoothness of the optimization criterion.
Visualization is also important in other areas of machine learning, for example to illustrate how kernel-methods project low dimensional data to a high dimensional space for classiﬁcation (Szymanski & McCane, 2011).
Note that the most common approach to visualize neural network mappings is through non-linear dimensionality reduction techniques, such as
Stochastic Neighbour Embedding (Hinton & Roweis, 2002).
We instead focus on linear projections in low dimensional environments, as non-linear dimensionality reduction has the risk of altering the semantics of the MDP models.
3. The MuZero Algorithm
We brieﬂy introduce the MuZero algorithm (Schrittwieser et al., 2020).
We assume a Markov Decision Process (MDP) speciﬁcation given by the tuple (cid:104)S , A , T , U, γ (cid:105) , which respectively represent the set of states ( S ), the set of actions ( A ), the transition dynamics mapping state-action pairs to new states ( T : S ×
A → p ( S ) ), the reward function mapping state-action pairs to rewards ( U : S ×
A → R ) , and a discount parameter ( γ ∈ [0 , 1] ) (Sutton & Barto, 2018).
Internally, we deﬁne an abstract MDP (cid:104) (cid:101) S , A , (cid:101) T , R, γ (cid:105) , where (cid:101) S denotes an abstract state space, with corresponding dynamics (cid:101) T : (cid:101)
S ×
A → (cid:101) S , and reward prediction
R : (cid:101) S ×
A → R .
Our goal is to ﬁnd a policy π : S → p ( A ) that maximizes the value V ( s ) from the start state, where
V ( s ) is deﬁned as the expected inﬁnite-horizon cumulative return:
[Equation]
We deﬁne three distinct neural networks to approximate the above MDPs (Figure 1): the state encoding/embedding function h θ , the dynamics function g θ , and the prediction network f θ , where θ denotes the joint set of parameters of the networks.
The encoding function h θ : S →
(cid:101) S maps a (sequence of) real MDP observations to a latent MDP state.
The dynamics function g θ
: (cid:101) S ×
A → (cid:101) S × R predicts the next latent state and the associated reward of the transition.
In practice, we slightly abuse notation and also write g θ to only specify the next state prediction.
Finally, the prediction network f θ : (cid:101) S → p ( A )
×
R predicts both the policy and value for some abstract state ˜ s .
We will identify the separate predictions of f θ (˜ s kt ) by p kt and V kt , respectively, where subscripts denote the time index in the true environment, and superscripts index the timestep in the latent environment.
Also, we write µ θ =
( h θ , g θ , f θ ) for the joint model.
Together, these three networks can be chained to form a larger computation graph that follows a single trace, starting from state s t following action sequence ( a t , . .
. , a t + n ).
First, we use the embedding network to obtain the ﬁrst latent state from the sequence of observations: ˜ s 0 t = h
θ ( s 0 , . . .
, s t ) .
Then, we recursively transition forward in the latent space following (˜ s k +1 t , r k +1 t )
=
g θ
(˜ s kt , a t + k ) .
Finally, for all generated abstract states ( ˜ s 0 t , . .
˜ s nt )
we use the prediction network to obtain ( p kt , V kt ) =
f θ (˜ s kt ) .
As illustrated in Figure 1, the entire differentiable computation graph implements the mapping (
[Equation]
We will later use this graph to train our networks, but ﬁrst discuss how they may be utilized to perform a search.
Given three individual neural networks for f , g and h as, we then use these networks to perform a MCTS search from state s t in the abstract space.
First, we embed the current state through h .
Then, we can essentially follow the PUCT (Rosin, 2011) algorithm, a variant of MCTS that incorporates prior weights on the available actions, where the priors originate from the policy network, p kt (ﬁrst proposed by (Silver et al., 2017a)).
The state transitions and policy/value predictions within the search are governed through g and f , respectively.
Eventually, this MCTS procedure outputs a policy π t = π ( s t ) and value estimate V t =
V ( s t ) for the root node, i.e., ( π t , V t ) ∼ MCTS ( s 0 , . . .
, s t | µ θ ) .
It then selects action a t ∼ π t in the real environment, transitions to the next state, and repeats the search.
We refer to appendix
B of the MuZero paper for full details on the MCTS search (Schrittwieser et al., 2020).
We still need to discuss how to train f θ , g θ , and h θ .
Through playing episodes, MuZero collects trajectories η of state, action, reward sequences, and the associated MCTS statistics obtained at each state, i.e., η = ( s t , π t , a t , r t , s t +1 , . . . ) .
As mentioned before, given a start state and a sequence of actions, we can unroll the differentiable graph that predicts rewards, values and policy parameters at each predicted state in the sequence (Eq. 2, Fig 1).
Since the entire graph is end-to-end differentiable, we may train it on a joint loss (Schrittwieser et al., 2020):
[Equation]
where l r = 0 for k = 0 , λ ∈ R is a constant that governs the L2 regularization, and z t is a cumulative reward estimate obtained from an n -step target on the real trace η , [Equation]
[-=
 γ-] i [Equation]
γ n V [Equation] + n +1 .
For n → ∞ this of course becomes a Monte Carlo return estimate.
Moreover, MuZero uses distributional losses, inspired by (Pohlen et al., 2018), for all three predictions ( r , V , and p ).
Finally, MuZero employs two forms of normalization to stabilize learning.
First, inside the MCTS selection rule, it min-max normalizes the value estimates based on the minimum and maximum estimate inside the current tree.
This allows the algorithm to adapt to arbitrary reward scaling.
Second, and more important for our paper, it also uses minmax normalization as the output activation of h θ and g θ .
In particular, every abstract state is readjusted through
[Equation]
where the min and max are over the elements in the ˜ s vector.
This brings the range of ˜ s in the same range as the discrete actions (one-hot encoded), but also introduces an interesting restriction, especially when the latent space is small.
For example, when | ˜ s | = 4 , then the min-max normalization will enforce two elements within ˜ s to be one of 0 or 1.
In practice, we observed that MuZero required at least a | ˜ s | + 2 latent space to effectively learn a | ˜ s
| -dimensional problem given that the state context was non-redundant and approximately Markov — this restriction was most stringent for small | ˜ s | .
Full details on the MuZero algorithm can be found in the Appendix of (Schrittwieser et al., 2020).
4. Visualizing MuZero’s Latent Space
Value equivalent models, like MuZero, are solely trained on a low-dimensional, scalar reward signal, as the value and policy targets are also derived from the rewards.
While the concept of value equivalent models is convincing, it remains unclear what types of representations these models actually achieve, and how they shape their internal space.
Therefore, we ﬁrst train MuZero agents to extract their learned embeddings/ representations and afterwards depict their statespace to latent-space mapping graphically in a 3D space.
We create visualizations strictly on low-dimensional control environments as it quickly becomes infeasible and nonsensical to map higher dimensional representations linearly back to only 3 dimensions.
Figure 2 explores such a graphical depiction of MuZero’s latent space near convergence (See also our supplementary material for additional visualizations during training and over ablations).
The left part of the ﬁgure aims to relate the ground-truth state space to the latent space.
On the top left, we visualize the ground-truth MountainCar state space (a 2D grid consisting of car position and car velocity) using vertical bars.
Directly next to it, we visualize a Principal
Component (PC) projection of the associated MuZero latent space when moving through the embedding function, i.e., the subspace { ˜ s :
˜ s
=
h θ ( s ) , ∀ s ∈ S} projected to its own 3-dimensional PC-space where L = | ˜ s | = 4 .
We retain the vertical bars, which thereby illustrates the way the true state space is curved into a high dimensional manifold.
We do see that MuZero recovers the state space to a large extent, while it was never explicitly trained on it, although some parts are compressed or stretched.
The bottom-left of Figure 2 provides a similar illustration of the mapping between observed states and embedded states, but this time with accompanying value information.
The bottom-left ﬁgure shows the learned value function of the state-space, where the black lines illustrate some successful ﬁnal MuZero trajectories.
Interestingly, when we map the value information into the embedding, we see that MuZero has actually strongly warped the original space.
Relative to this projection, states with low value are grouped in the leftside of this space, while states with high value appear near the right-side.
Effectively, MuZero has created a representation space in which it groups states that are dynamically close.
For example, in the true state space (left), dark blue (low value) and yellow (high value) states are often adjacent.
However, in the embedding (right), these states are strongly separated, because they are dynamically distant (we cannot directly move from blue to yellow, but need to follow the trajectory shown in black).
This shows that MuZero indeed manages to retrieve essential parts of the dynamics, while only being trained on a scalar reward signal.
Our last visualization of the learned model is shown in the right panel of Figure 2.
In this plot, we again visualize a PCA of MuZero ﬁtted on the embedding near convergence.
However, this time we ﬁx a start state ( s 0 ) and action sequence ( a 0 , . . .
, a n ), and visualize a successful trace passing through both the embedding and dynamics functions.
For the embedding function, given a trace of environment data η , we plot (cid:104)
h θ ( s 0 ) , h θ ( s 1 ) , . . .
, h θ ( s n +1 ) (cid:105) as the blue trajectory.
In contrast, in green we embed the ﬁrst state and subsequently repeatedly transition through the learned dynamics function g θ —i.e., we plot (cid:104)
h θ ( s 0 ) , g θ ( h θ ( s 0 ) , a 0 ) , g θ ( g θ ( h θ ( s 0 ) , a 0 ) , a 1 ) , . . .
(cid:105) .
Interestingly, these two trajectories do not overlap.
Instead, the latent forward dynamics seem to occupy a completely different region of the latent space than the state encoder.
We repeatedly make this observation on different experiments and repetitions.
Although one would conceptually expect these trajectories to overlap, there is no explicit loss term that enforces these trajectories to stay congruent.
When the latent capacity is large enough, MuZero may fail to identify this overlap, which hypothetically could waste generalization potential.
In the next section we investigate two possible modiﬁcations to combat this phenomena.
5. Regularizing the Latent Space
Ideally, the dynamics function should stay relatively congruent to the embedding of observed environment states.
Thereby, we can maximally proﬁt from generalization in the latent space, and make optimal use of the available data.
However, as we observed in the previous section, there is no explicit training objective that ensures this congruence.
We therefore propose two types of regularization, i.e., differentiable penalty terms (added to the loss) that enforce the latent space to be more congruent.
The ﬁrst solution is to add a contrastive regularization , which penalizes between the embedding of a state and the forward dynamics prediction towards the same abstract state.
Given a trace η , we compute the loss
[Equation]
where ˜ s k +1 t = g ( . . .
g ( g ( h θ ( s 0 ) , a 0 ) , a 1 ) , . . . , a k ) .
This loss is added to the standard MuZero loss (Eq. 3), where its strength is governed by a scalar ω ∈ R + .
To prevent trivial solutions, we block the gradients of the above loss through h θ for k > 0 .
This ensures that only the dynamics function should move towards the embeddings (and not the other way around), assuming that the embedding function is already relatively well adapted.
This type of contrastive loss is often seen in supervised learning, for example in Siamese Networks (Chen & He, 2020;
Koch et al., 2015).
The process is visually illustrated in the top graph of Figure 3.
A second solution can be the use of a decoding regularization , which decodes latent predictions back to the true observation.
This essentially adds a (multi-step) forward prediction loss to the MuZero model.
We initialize an additional decoder network h − 1 θ , which maps latent states back to the true observations: s (cid:48) t + k = h − 1 θ (˜ s kt ) .
Then, given a trace of data η , we compute the loss
[Equation]
with ˜ s k +1 t deﬁned as in the previous loss.
Again, we can use a scalar ω ∈ R + to determine the relative contribution of the loss when added to Eq. 3.
This approach combines the value equivalent loss model loss with a standard forward prediction loss.
The additional penalty ties the latent predictions of g θ back to the ground-truth observations.
Although this mechanism is independent of the encoding h , it may potentially help to regularize the space of g θ transitions.
Moreover, as a second beneﬁt, this regularization may be advantageous in initial stages of learning, as state transition are often dense in information compared to sparse reward signals.
Thereby, it could be beneﬁcial in early learning, but can be overtaken by the value equivalent objective when the model converges.
The second type of regularization is visualized in the bottom graph of Figure 3.
6. Experimental Setup
We will now experimentally study our regularized versions of MuZero compared to standard MuZero.
As an additional baseline, we also include an adaptation of AlphaZero for single-player domains (Silver et al., 2018), i.e., a similar agent with a perfect environment model.
We provide a complete reimplementation of MuZero based on the original paper (Schrittwieser et al., 2020).
All code is written in
Python using Tensorﬂow 2.0 (Abadi et al., 2015).
The code follows the AlphaZero-General framework (Nair et al., 2017) and is available at https://anonymized.org , along with all resulting data.
The intention for our experiments was not to achieve maximal performance of MuZero, but to visualize the learned representations and study the effects of the different regularization terms.
For this reason, we focus on low-dimensional, classic control environments, as it is hard to meaningfully visualize high-dimensional state spaces.
In particular, we tested our work on CartPole and MountainCar, from OpenAI‘s Gym (Brockman et al., 2016).
To aid our analysis, we developed an interactive visualization tool to for learned MuZero models on the MountainCar environment at https://anonymized.org .
We encourage the reader to play around with this tool as we personally found this to provide more insight than static 2 D images.
[Table]
{ None , Contr .
, Dec . }
{ None , Contr .
, Dec [Table] [Table] All hyperparameters can be found in Appendix A in supplemental material.
For the results in this paper, we report on a factorial experiment design shown in Table 1, where we vary: 1) the regularization type, 2) the regularization strength ( ω ), 3) the number of latent unroll steps ( K ), and 4) the size of the latent representation ( L = | ˜ s | ).
Of course, our main interest is in the ability of the regularizations to inﬂuence performance.
All experiments were run on the
CPU.
The hyperparameter analysis experiment was run asynchronously on a computing cluster using 24 threads at 2GHz.
On average it took about half a day to train an agent on any particular environment.
The complete experiments depicted take about one week time to complete (after the separate lengthy tuning phase of the other hyperparameters).
7. Results
Figure 4 shows results of our CartPole experiments.
In both plots we compare AlphaZero, MuZero, and our two regularized MuZero extensions.
On the left, we investigated the effect of the regularization strength.
AlphaZero consistently solves this environment, while MuZero has more trouble and solves it roughly 75% of the time.
Both regularization methods do not seem to structurally improve MuZero performance on CartPole.
Surprisingly, performance dips for intermediate regularization strength ( ω = 0 . 1 ).
On the right of Figure 4, we display the average training curves for each of our four algorithms on CartPole.
We see a similar pattern, where the regularization losses do not seem to provide beneﬁt.
Since CartPole is a dense reward environment, in which
MuZero quickly obtains gradients to structure the latent space, we also investigate our reguralization methods on
MountainCar.
Figure 5 (left) shows the training progression for AlphaZero, MuZero and the two regularization methods.
Again, AlphaZero outperforms all other agents.
However, this time, both our regularization methods outperform
MuZero in learning speed, where especially the contrastive regularization leads to faster learning.
All methods do reach the same optimal performance.
We further quantify the strength of the learned internal model by ’blindfolding’ the agent (Freeman et al., 2019).
A blindfolded agent only gets to see the true state of the environment ever x number of steps, where we call x the ‘observation sparsity’.
The right graph of Figure 5 displays the average reward obtained by all MuZero models when varying the observation sparsity (horizontal axis).
All agents consistently keep solving the problem up to an observation sparsity of depth 10, which shows that their models do not compound errors too fast.
Then, performance starts to deteriorate, although only the contrastively regularized, L =
MuZero variant keeps solving the problem, even up tosteps without a ground truth observation.
In general, the blindfolded agents show that the model stays consistent over long planning horizons (even when trained on shorter rollout depths), and that especially contrastive regularization further increases this ability.
Please refer to the Appendix for additional visualization of the latent space in MountainCar experiments.
8. Discussion
We studied the latent space of the MuZero algorithm, a successful combination of value equivalent models and iterated planning and learning.
In particular, we visualized the latent space, showing how the true state space is partially recovered and warped into a dynamically relevant space.
Moreover, we identiﬁed a frequent mismatch between the latent trajectories and the embeddings of the ground truth observations.
To alleviate this issue, we introduced two forms of regularization, which did not improve performance in CartPole, but did improve performance on MountainCar.
Finally, we also found that these models can accurately predict correct actions over long planning horizons, up tosteps in some cases, while being trained on relatively short trajectories.
We hypothesize that the beneﬁt of regularization is most prominent on sparse reward tasks like MountainCar.
In CartPole, where we quickly observe rewards, the representation space of MuZero will quickly receive relevant gradients
[Figure]
to structure itself.
However, when reward information is sparse, we may beneﬁt from additional regularization of the latent space, which might squeeze additional information out of the available data.
Especially the contrastive regularization seems to show this effect, and could be a promising direction for future research.
A critique on our work could involve the dimensionality of the experiments.
Deep reinforcement learning has started to focus more and more on high-dimensional problems.
We opted for small scale experiments for two reasons.
First of all, low-dimensional problems allow for better interpretation.
We can only visualize a latent space in three dimensions, and it would not be feasible to map an Atari game representation back to 3D. Moreover, the general emphasis in RL has shifted towards to the challenge of dimensionality, while problems can be challenging in more ways than dimensionality only.
A good illustration of the latter point is that such a high-impact algorithm as MuZero actually suffers to consistently solve CartPole and MountainCar.
This recognition of the value of low dimensional experiments also starts to resurface in the RL community (Osband et al., 2019).
Furthermore, AlphaZero and MuZero are incredibly computationally expensive to run.
While standard RL algorithms are already expensive to run, these iterated search and learning paradigms only exacerbate this limitation.
When we use a small-scale MCTS of n = 11 at every timestep, it already takes at least 11 times as long to complete the same number of episodes as model-free RL.
Compared to model-free RL, these iterated search and learning paradigms seem to beneﬁt in long run performance, but they are computationally heavy to train.
We observed that AlphaZero and, especially, MuZero are rather unstable to train.
Their performance strongly depends on a variety of hyperparameters and network initialization, and we observe strong variation in the performance over different runs.
A direction for future work is to further stabilize these algorithms.
Our work on novel regularizations for value equivalent models like MuZero is a ﬁrst step in this direction.
Future work could investigate how to further stabilize MuZero through different regularizations, or further combinations with standard forward prediction or types of auxiliary losses (Jaderberg et al., 2016).
Although value equivalency is a strong incentive for the ﬁnal representation of the model, we may aid the model during learning with additional losses and regularizations, especially in sparse reward tasks.
9. Conclusion
The ﬁrst part of this paper visualized the learned latent representations and model of MuZero, which we consider the main contribution of the paper.
Our two most important observations are: 1) the learned latent space warps the original space into a ’dynamically close’ space, in which states that are easy to reach from eachother are nearby, and 2) the embedding of an action trajectory through the latent dynamics often departs from the embedding of the ground-truth states, which suggest the latent space does not optimally generalize information.
To alleviate the last problem, the second part of the paper investigated two new forms of regularization of value equivalent models: 1) a contrastive regularization, and 2) a decoder regularization.
Experiments indicate that these approaches may aid MuZero in sparse reward tasks, and that these models in general manage to plan over long horizons (much longer than they were trained on).
Our visualizations of the latent space show that the contrastive regularization term can force the dynamics function to predict values closer to the embedding, this in turn can help interpreting the look-ahead search of MuZero.
All code and data to reproduce experiments is open-sourced at www.anonymized.com , including an interactive visualizer to further explore the MuZero latent space.
We also hope our work further stimulates new emphasis on visualization of RL experiments, which may provide more intuitive insights into deep reinforcement learning.
A. Hyperparameters and Neural Architectures
Our neural network architectures for MuZero were partially inspired by the ones proposed by (van Seijen et al., 2020).
We kept every constituent network within the larger MuZero model equivalent, that is: h θ , g θ , f θ , and optionally h − 1 θ had the same internal architecture being a two-layered feed forward network with 32 hidden nodes using the Exponential
Linear Unit activation function.
The embedding function h θ received only current observations, so no trajectories.
Actions for the dynamics functions were one-hot-encoded and concatenated to the current latent state before being passed to g θ ; computed latent-states were also minmaxnormalized within the neural architecture.
The action policy, value function, and reward prediction were represented as distributions and trained using a cross-entropy loss, as was also done in the original paper (Schrittwieser et al., 2020).
The AlphaZero neural network was simply f θ that received state observations—i.e., when omitting h θ and g θ .
All other unmentioned hyperparameters are shown in Table 2.
B. Additional Visualization of Latent Space
Figure 6 illustrates a typical training progression for MuZero on the MountainCar environment (hyperparameters can be found in the appendix), similarly to Figure 2 we used
L = | ˜ s | = 4 .
The shapes shown in the top row portray the progression of the embedding { ˜ s :
˜ s
=
h θ ( s ) , ∀ s ∈ S} , shown in a 3D space obtained through PCA.
Upon initialization, the latent space is almost degenerate, and all latent traces occupy the same part of the latent space.
The action selection strategy of MuZero is near uniform in this situation.
Gradually during training, the agent starts to expand and eventually occupies the full latent space.
In the topright plot, we see that the latent spaces aggregate states with similar values.
On the bottom, we display violin plots of the latent state distribution during training, in the true four dimensional latent space.
We ﬁnd that these distributions appear overinﬂated around 0 and 1, on all dimensions.
Arguably, this is an artifact of the min-max normalization of the latent space (as discussed in Section 3), which always forces one element to be 0, and one element to be 1.
We do see that MuZero spreads its assignment of min and max over all dimensions.
Note that, although MuZero can theoretically still predict any point in the latent space [0 , 1] 4 , the normalization effectively reduces the degrees of freedom within the latent space.
Figure 7 provides additional illustration of the learned latent spaces for MuZero and our two regularized MuZero variants on MountainCar.
The ﬁgure has a 3x2 design, where the rows depict different algorithms, and the columns different sizes of the latent space.
The left of each cell contains a 3-dimensional PCA projection of the learned latent space.
We again observe that the value equivalent objective nicely unfolds a latent space, in which states that are dynamically close in the true environment are also close in the latent space.
More interestingly, the right of each cell displays the two ways to embed a trajectory: 1) by passing each observed true state through h θ (blue), or 2) by simulating forward through the latent dynamics g θ (green).
We again observe that these trajectories strongly depart for the original MuZero algorithm.
In contrast, the two regularization methods force both trajectories to stay closer to eachother.
Especially the contrastive loss manages to keep both trajectories close to consistent, which may help to obtain better generalization, and could explain the faster learning observed in Figure 5.
[Figure]
[Table]