SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Gumbel MuZero [18] is an extension designed to enhance performance in environments with low simulation costs by leveraging the Gumbel-Top-k trick [78] to select actions that guarantee policy
improvement. This approach introduces a new improved policy,  $\pi^{\prime}$ , which seamlessly integrates the original visit counts distribution with MCTS searched values. By combining these two sources of information, Gumbel MuZero efficiently exploits the knowledge acquired during the search process, leading to more informed decision-making and improved performance in a variety of environments.
Stochastic MuZero [16] This extension enhances MuZero by enabling it to learn and plan with stochastic models. Specifically, it incorporates a stochastic model that includes afterstates and conducts stochastic tree searches using this model. Stochastic MuZero achieves competitive or superior performance compared to state-of-the-art methods in a variety of canonical board games. such as 2048 and Backgammon, while preserving the superhuman performance demonstrated by the standard MuZero in Go.
G.2 LightZero Algorithms Overview
In this section, we present a detailed summary of various MCTS/MuZero variants, encapsulated in algorithm overview diagrams. Specifically, Figure 17 provides an insight into the Monte Carlo Tree Search (MCTS), while Figure 18 illuminates the mechanics of AlphaZero, a fusion of MCTS and deep neural networks. Moreover, Figure 19 introduces MuZero, an advancement upon AlphaZero with a learned model representation. Figure 21 subsequently presents the Sampled MuZero, and finally, Figure 22 delineates the Gumbel MuZero variant. These overview diagrams serve to highlight the evolution and interconnections among MCTS, AlphaZero, MuZero, and their related algorithms.
G.3 Model Architecture
We provide the representation network structure of the MuZero series algorithms implemented in LightZero in Figure 23. The dynamics network and prediction network structures (including both value and policy networks) are presented in Figure 24. Note that the model structures we provide are based on Atari environments with image-type observations. For environments with vector-type observations, such as continuous control tasks, the overall model structure is similar, only replacing the convolutional layers in the original model with corresponding fully connected layers. For specific details, please refer to the model directory through OpenDILab at https://github.com/opendilab/LightZero.
We provide a detailed summary of the key hyperparameters for the MuZero w/ SSL algorithm, as implemented in LightZero for the Atari environment, in Table 7. The hyperparameters for other environments and algorithms are generally similar. Unless explicitly stated otherwise, all other parameters are in accordance with those in Table 7.
Specifically, we highlight the primary differing hyperparameters for the Sampled EfficientZero algorithm in the Atari environment in Table 8, as well as for continuous control environments such as MuJoCo in Table 9. In addition, we explain the main distinct hyperparameters for the Gumbel MuZero algorithm in the Atari environment in Table 10, alongside the Stochastic EfficientZero algorithm in both the Atari environment, as shown in Table 11, and the 2048 environment, as presented in Table 12. Lastly, the pivotal hyperparameters for the AlphaZero/MuZero algorithm in the Gomoku environment are set out in Tables 13 and 14.
H Comparison with Other MCTS Frameworks
LightZero is a comprehensive algorithm benchmark developed using PyTorch, integrating nine distinct algorithms such as AlphaZero and MuZero. It supports over twenty different environments, including Atari and Go. On the other hand, the MCTX library<sup>2</sup>, primarily implemented using JAX, includes basic implementations of algorithms such as AlphaZero, MuZero, and Gumbel MuZero. However, the training process across various environments is not yet fully established.
$^{2}$ https://github.com/google-deepmind/mctx
To offer an intuitive comparison of the differences in integrated algorithms and supported environments between these two repositories, we present the algorithms and environments supported by LightZero and MCTX in the following Table 5 and Table 6, respectively.
Please note: " $\checkmark$  " denotes that the corresponding item is fully implemented and thoroughly tested. "☑" signifies that the item is currently in development. "—" indicates that the algorithm does not support the respective environment.
Regarding the comparison table provided, we should note the following two points: Firstly, the list of algorithms and environments supported by MCTX not only encompasses MCTX itself, but also includes all derivative repositories based on MCTX. Secondly, all environments associated with Gumbel MuZero and Stochastic MuZero within the MCTX library are currently in a locked "⊗" state. This is because, to our knowledge, neither MCTX nor its derivative repositories have fully implemented these algorithms and environments. Although the foundational modules for the Gumbel MuZero and Stochastic MuZero algorithms exist in the original MCTX repository, the development of a complete training pipeline for these algorithms is still in progress.
Algo Env	Classic Control	Box2D	Atari	MuJoCo	Go Bigger	Mini Grid	Maze	Connect Four	Gomoku	2048	Go
AlphaZero	—	—	—	—	—	×	×	✓	✓	×	✓
Sampled AlphaZero	—	—	—	—	—	×	×	×	✓	×	✓
Gumbel AlphaZero	—	—	—	—	—	×	×	×	✓	×	✓
MuZero	✓	✓	✓	—	✓	✓	✓	✓	✓	✓	✓
MuZero w/ SSL	✓	✓	✓	—	✓	✓	×	×	✓	✓	✓
EfficientZero	✓	✓	✓	—	✓	✓	×	×	✓	✓	✓
Gumbel MuZero	✓	✓	✓	—	✓	✓	×	×	✓	✓	✓
Sampled MuZero	✓	✓	✓	✓	✓	✓	×	×	✓	✓	✓
Stochastic MuZero	✓	✓	✓	—	✓	✓	×	×	✓	✓	✓
Table 5: Algorithms and Environments supported by LightZero as of the writing of this paper.
Algo Env	ClassicControl	Box2D	Atari	Maze	ConnectFour	Gomoku	Go
AlphaZero	-	-	-	✓	✓	✓	✓
MuZero	✓	✓	☒	☒	☒	☒	☒
Gumbel MuZero	☒	☒	☒	☒	☒	☒	☒
Stochastic MuZer	☒	☒	☒	☒	☒	☒	☒
Table 6: Algorithms and Environments supported by MCTX as of the writing of this paper.
I Limitations and Future Work
Despite its considerable achievements across various benchmark environments, LightZero does bear certain limitations.
① Firstly, while efforts have been made to amplify the algorithm's universality through decoupling and modularization, certain specific problems may still necessitate tailored adjustments and optimizations. For instance, all current LightZero algorithms do not directly support hybrid action space environments [79]. However, through the use of action representation [80] and other techniques, they could potentially be adapted to a wider action space.
② Secondly, as discussed in Section 3, due to the intrinsic limitations of the MCTS algorithm, our method may encounter challenges in specific, complex real-world environments. Particularly when handling environments with high-dimensional action spaces or strong randomness, methods like Sampled MuZero and Stochastic MuZero could benefit from further optimization and improvement.
③ Finally, the high prerequisites of LightZero can pose challenges, especially for those new to decision intelligence. This is largely due to the inherent complexity of the MCTS+RL algorithms. While these algorithms are comprehensive, they require a profound understanding and significant time investment to deploy effectively. This complexity, alongside the computational demands of running these advanced algorithms, can limit their widespread application. Future enhancements should concentrate on improving the accessibility of the LightZero framework's interfaces, enriching the documentation, and cultivating a vibrant user community that encourages knowledge collaboration.
Despite the aforementioned challenges, we remain optimistic about the potential of the LightZero framework. For future, we have identified several promising areas for further exploration:
- Broadening Applications: We envision more researchers and developers applying LightZero across a wider spectrum of practical fields. These include but are not limited to natural language processing, autonomous driving, and the control and optimization of complex systems.
- Algorithmic Optimization: While we have enhanced our submodules by introducing more suitable exploration mechanisms and optimizing model loss functions, there remains significant room for improvement. We invite the community to contribute new exploration and optimization strategies to further boost the performance of MCTS-style algorithms.
- Integration with Cutting-Edge Technologies: We stress the importance of a seamless integration between LightZero and future research directions. Two key areas for future exploration stand out:
1. Integration with Large Language Models (LLM) [81] [82]: LLM can serve as a high-level task planner, breaking down complex decision-making problems into a series of low-level instructions. LightZero can set these instructions as conditional inputs of the representation network and serve as an efficient instruction executor. Furthermore, considering the powerful reanalysis mechanism of MuZero Unplugged, it can be a viable choice for fine-tuning LLMs.
2. Clever Utilization of Model-Based RL Techniques: LightZero has already integrated many model-based RL techniques, such as the Recurrent State-Space Models (RSSM) [83] used in DreamerV2 [84], latent state consistency loss used in EfficientZero, and the afterstate dynamics functions proposed in Stochastic MuZero [16]. We aim to incorporate additional techniques within the model representation, loss function, and MCTS search process to enhance the universality and robustness of the LightZero framework.
Monte Carlo Tree Search
Figure 17: MCTS algorithm overview.
B. Acting (Self-play)
C. Training
Figure 18: AlphaZero algorithm overview.
$$
l _ {t} (\theta) = l ^ {v} \left(z _ {t}, v _ {t}\right) + l ^ {p} \left(\pi_ {t}, \mathbf {p} _ {t}\right) + c \| \theta \| ^ {2}
$$
where,  $z_{t}$  is the game reward from the perspective of the current player and  $\pi_{t + k}$  is the MCTS searched policy at timestep  $t$ .
$v_{t}$  and  $\mathbf{p}_t$  is the predicted value and policy from the neural network  $f_0$ .
$l^v$  is MSE loss,  $l^p$  is cross-entropy loss.
A. Planing
B. Acting
C. Training
D. Loss
$$
l _ {t} (\theta) = \sum_ {k = 0} ^ {K} l ^ {t} \left(u _ {t + k}, r _ {t} ^ {k}\right) + l ^ {v} \left(z _ {t + k}, v _ {t} ^ {k}\right) + l ^ {p} \left(\pi_ {t + k}, \mathbf {p} _ {t} ^ {k}\right) + c \| \theta \| ^ {2}
$$
NOTE:  $g$  is the dynamic network (MLP, without RNN), value rescale,
categorical distribution for reward and value in Atari, Reanalyze
EfficientZero: Self-supervised coconsistency loss, Value prefix, Off-policy correction
D. Network and Loss
Figure 19: MuZero algorithm overview.
The self-supervised consistency loss.
Figure 20: EfficientZero algorithm overview.
$$
l _ {t} (\boldsymbol {\theta}) = \sum_ {k = 0} ^ {K} l ^ {r} \left(v p _ {t + k}, v p _ {t} ^ {k}\right) + l ^ {v} \left(z _ {t + k}, v _ {t} ^ {k}\right) + l ^ {p} \left(\pi_ {t + k}, \mathbf {p} _ {t} ^ {k}\right) + l ^ {s i m i l a r i t y} (s _ {t + 1}, \hat {s} _ {t + 1}) + c \| \boldsymbol {\theta} \| ^ {2}
$$
$$
\text {w h e r e ,} z _ {t} = u _ {t} + \gamma u _ {t + 1} + \dots + \gamma^ {l - 1} u _ {t + l - 1} + \gamma^ {l} v _ {t + l} ^ {\text {M C T S}}, l <   = k, v _ {t + l} ^ {\text {M C T S}} \text {i s r e a n a l y z e d M C T S r o o t v a l u e .}
$$
$$
\text {R e f e r e n c e : M u Z e r o t a r g e t i s} z _ {t} = u _ {t} + \gamma u _ {t + 1} + \dots + \gamma^ {n - 1} u _ {t + k - 1} + \gamma^ {k} v _ {t + k}.
$$
NOTE:  $g^{rm}$  (abbreviated as  $g$  in the figure) is the recurrent dynamics network,  $vp$  is value prefix,  $r_h$  is hidden state.
MuZero: Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model
A. Planing
Selection:
$$
\text {I n e a c h} \quad a ^ {k} \text {a d d i t o} \quad \text {U C B}
$$
$$
a ^ {k} = \arg \max  _ {a} \left[ Q (s, a) + \hat {\beta} / \beta p (s, a) \cdot \frac {\sqrt {\sum_ {b} N (s , b)}}{1 + N (s , a)} \left(c _ {1} + \log \left(\frac {\sum_ {b} N (s , b) + c _ {2} + 1}{c _ {2}}\right)\right) \right]
$$
$$
\text {W e} \beta = p, \text {t h u s} \hat {\beta} / \beta p (s, a) = \hat {\beta} (s, a)
$$
Sampled-based Policy Improvement.
From left to right, current policy  $\pi(a|s)$ , the empirical distribution  $\hat{\beta}(a|s)$ , the sample-based improved policy  $\hat{I}_{\beta}\pi(a|s)$ , the improved policy  $I\pi(a|s)$ , respectively. As the number of samples  $K$  increases,  $\hat{I}_{\beta}\pi(a|s)$  converges to  $\pi(a|s)$ .
$$
l _ {t} (\pmb {\theta}) = \sum_ {k = 0} ^ {K} l ^ {r} \left(u _ {t + k}, r _ {t} ^ {k}\right) + l ^ {v} \left(z _ {t + k}, v _ {t} ^ {k}\right) + K L \left(\pi_ {t + k}, \mathbf {p} _ {t} ^ {k}\right) + c \| \pmb {\theta} \|
$$
where,  $\pi_{t + k}$  is the MCTS search policy (normalized visit counts), also called the sampled-based improved policy  $\hat{I}_{\beta}\pi (a|s)$ , which is a discrete categorical distribution.
$\mathbf{p}_i^k$  is the predicted (potentially continuous) policy distribution. The policy loss (KL divergence) is calculated in the sampled actions.
NOTE:  $\mathbf{p}$  (aka.  $\pi$  is predicted policy (potentially continuous) distribution,  $\beta$  is the proposal policy distribution (typically equals  $\pi$ ),  $\hat{\beta}$  is the empirical policy distribution.
Gumbel MuZero: Planning with few simulations in high dimensional discrete action space
Figure 21: Sampled MuZero algorithm overview.
A. Planing
The key issue when the num of simulations  $n <$  num of total actions  $k$  in MCTS search is how to choose which actions to visist and how many times.
1. we can control the number of actions sampled without replacement.
2. we can use a bandit algorithm to efficiently explore the set of sampled actions.
Sequential Halving is used to identify the action with the highest  $g(a) + \log_{2}(a) + \sigma(\hat{q}(a))$ .
Figure 22: Gumbel MuZero algorithm overview.
$$
l _ {t} (\theta) = \sum_ {k = 0} ^ {K} l ^ {r} \left(u _ {t + k}, r _ {t} ^ {k}\right) + l ^ {v} \left(z _ {t + k}, v _ {t, m i x} ^ {k}\right) + K L \left(\pi_ {t + k} ^ {\prime}, \mathbf {p} _ {t} ^ {k}\right) + c \| \boldsymbol {\theta} \| ^ {2}
$$
mixed value approximate:
$$
v _ {\mathrm {m i x}} = \frac {1}{1 + \sum_ {b} N (b)} \left(\hat {v} _ {\pi} + \frac {\sum_ {k} N (k)}{\sum_ {b \in [ N (a) - 0 ]} \pi (b)} \sum_ {a \in \{a : N (a) > 0 \}} \pi (a) q (a)\right)
$$
improved policy distribution:  $\pi^{\prime} = \mathrm{softmax}(\mathrm{logits} + \sigma (\mathrm{completedQ}))$
$$
\text {c o m p l e t e d} \mathrm {Q} (a) = \left\{ \begin{array}{l l} q (a) & \text {i f} N (a) > 0 \\ v _ {\pi}, & \text {o t h e r w i s e} \end{array} \right.
$$
$\varepsilon$  :  $\pmb{p}$  is policy distribution,  $\pi^{\prime}$  is the improved policy distribution.
completed  $Q(a)$  is the completed Q values,  $q(a)$  is the MCTS estimated value (for visited actions).
Sampled MuZero: Learning and Planning in Complex Action Spaces (High Dimensional Discrete or Continuous)
Figure 23: The network architecture of the representation network  $h$  for image input domain in LightZero. This network represents the root observation as a latent state.
Figure 24: Left: The network architecture of the dynamics network  $g$  for image input domains in LightZero. Given a latent state and a selected action, it outputs the transitioned next latent state and the corresponding predicted reward. Right: The network architecture of the prediction network  $f$  (encompassing both value and policy networks) for image input domains. Given a latent state, this network predicts the action probability and value.
Hyperparameter	Value
Num of frames stacked	4
Num of frames skip	4
Reward clipping	True
Optimizer type	Adam
Learning rate	3 × 10-3
Discount factor	0.997
Weight of policy loss	1
Weight of value loss	0.25
Weight of reward loss	1
Weight of policy entropy loss	0
Weight of SSL (self-supervised learning) loss	2
Batch size	256
Model update ratio	0.25
Frequency of target network update	100
Weight decay	10-4
Max gradient norm	10
Length of game segment	400
Replay buffer size (in transitions)	1e6
TD steps	5
Number of unroll steps	5
Use augmentation	True
Discrete action encoding type	One Hot
Normalization type	Batch Normalization
Priority exponent coefficient	0.6
Priority correction coefficient	0.4
Dirichlet noise alpha	0.3
Dirichlet noise weight	0.25
Number of simulations in MCTS (sim)	50
Reanalyze ratio	0
Categorical distribution in value and reward modeling	True
The scale of supports used in categorical distribution	300
Table 7: Key hyperparameters of MuZero with/ SSL on Atari environments.
Hyperparameter	Value
Number of sampled actions (K)	Different environments have distinct values for K, shown in Figure 7
Policy loss type	Cross Entropy Loss
Table 8: Key hyperparameters of Sampled EfficentZero on Atari environments.
Hyper-parameter	Value
Num of frames stacked	1
Num of frames skip	1
Reward clipping	False
Number of sampled actions (K)	20
Policy loss type	Cross Entropy Loss
Number of simulations in MCTS (sim)	50
Weight of policy entropy loss	0.005
Length of game segment	200
Use augmentation	False
Max gradient norm	0.5
Table 9: Key hyperparameters of Sampled EfficientZero used in continuous control environments, such as, MuJoCo.
Hyperparameter	Value
Num of max considered actions	i.e. Action Space Size
Gumbel Scale	10
Max visit init	50
Value Scale	0.1
Table 10: Key hyperparameters of Gumbel MuZero on Atari environments.
Hyperparameter	Value
Chance space size	i.e. Action Space Size
Afterstate Dynamics Network	Similar with Dynamics Network in Figure 24
Afterstate Prediction Network	Similar with Prediction Network in Figure 24
Chance Encoder	Two Layer MLP
Table 11: Key hyperparameters of Stochastic MuZero on Atari environments.
Hyperparameter	Value
Chance space size	16 * num_of_POSSIBLE_chance_TILE
Afterstate Dynamics Network	Similar with Dynamics Network in Figure 24
Afterstate Prediction Network	Similar with Prediction Network in Figure 24
Chance Encoder	Two Layer MLP
Num of frames stacked	1
Num of frames skip	1
Reward clipping	False
Discount factor	0.999
Batch size	512
Length of game segment	200
TD steps	10
Use augmentation	False
Number of simulations in MCTS (sim)	100
Table 12: Key hyperparameters of Stochastic MuZero on 2048 environments.
Hyperparameter	Value
Board size	6
Num of frames stacked	1
Discount factor	1
Weight of policy loss	1
Weight of value loss	1
Number of simulations in MCTS (sim)	100
Categorical distribution in value modeling	False
Table 13: Key hyperparameters of AlphaZero on Gomoku environments.
Hyperparameter	Value
Board size	6
Num of frames stacked	1
Discount factor	1
Weight of SSL (self-supervised learning) loss	0
Length of game segment	18
TD steps	18
Use augmentation	False
Number of simulations in MCTS (sim)	100
The scale of supports used in categorical distribution	10
Table 14: Key hyperparameters of MuZero on Gomoku environments.