LightZero: A Unified Benchmark for Monte Carlo Tree Search in General Sequential Decision Scenarios
Authors: Yazhe Niu
Yuan Pu
Zhenjie Yang
Xueyan Li
Tong Zhou
Jiyuan Ren
Shuai Hu
Hongsheng Li
Yu Liu
Venue: NeurIPS 2023 Datasets and Benchmarks Spotlight (2023)
Keywords: Reinforcement Learning
Monte Carlo Tree Search
Model-Based RL
Algorithm Benchmark
TL;DR: A unified benchmark for deploying MCTS/MuZero reinforcement learning methods in general sequential decision scenarios

Abstract: Building agents based on tree-search planning capabilities with learned models has achieved remarkable success in classic decision-making problems, such as Go and Atari.
However, it has been deemed challenging or even infeasible to extend Monte Carlo Tree Search (MCTS) based algorithms to diverse real-world applications, especially when these environments involve complex action spaces and significant simulation costs, or inherent stochasticity.
In this work, we introduce LightZero, the first unified benchmark for deploying MCTS/MuZero in general sequential decision scenarios. 
Specificially, we summarize the most critical challenges in designing a general MCTS-style decision-making solver, then decompose the tightly-coupled algorithm and system design of tree-search RL methods into distinct sub-modules.
By incorporating more appropriate exploration and optimization strategies, we can significantly enhance these sub-modules and construct powerful LightZero agents to tackle tasks across a wide range of domains, such as board games, Atari, MuJoCo, MiniGrid and GoBigger.
Detailed benchmark results reveal the significant potential of such methods in building scalable and efficient decision intelligence.
The code is available as part of OpenDILab at https://github.com/opendilab/LightZero.

LightZero: A Unified Benchmark for Monte Carlo Tree Search in General Sequential Decision Scenarios
Yazhe Niu 1.
Yuan 2u
Zhenjie Yang'
Xueyan Li²
Tong Zhou Yu Liu
Jiyuan Ren
Shuai Hu
Hongsheng Li3.4
Time Group LTD
Shanghai Artificial Intelligence Laboratory
The Chinese University of Hong Kong
*Centre for Perceptual and Interactive Intelligence
Abstract
Building agents based tree-search planning capabilities with learned models has achieved remarkable success in classic Dacision-making problems, such Go and Atari. However, it has been deemed challenging even infeasible to extend Monte Carlo Tree Search (MCTS) based algorithms to diverse real-world applications, especially when these environments involve complex action spaces and significant simulation costs, or inherent stochasticity. In this work, we introduce LightZero, the first unified benchmark for deploying MCTS/MaZero in general sequential decision scenarios. Specificially, we summarize the most critical challenges in designing general MCTS-style decision-main making solver, then decompose the tightly-coupled algorithm and system design of tree-search RL methods into distinct sub-modules. By incorporating more appropriate exploration and optimization strategies, we can significantly enhance these sub-modules and construct powerful LightZero agents to tackle tasks across wide range domains, such as board games, Atari, MuJoCo, MiniGrid and GoBigger. Detailed benchmark results reveal the significant potential of such methods in building scalable and efficient decision intelligence. The code is available part of OpenDILab at
Introduction
General decision intelligence needs to solve tasks in many distinct domains. Recent advances in reinforcement learning (RL) algorithms have addressed several challenging decision-making problems
2 and even surpassed top-level human experts in performance 3 However, these state- of-the-art RL agents often exhibits poor data efficiency and face significant challenges when handling a wide range of diverse problems. Different environments present specific learning requirements and difficulties that prompted currently various algorithms (e.g. DQN PPO [5], R2D2 (6), SAC 77) and system architectures such IMPALA 8 and others 9 10 4
Designing general and data-efficient decision solver needs to tackle various challenges, while ensuring that the proposed algorithm can be universally deployed anywhere without domain-specific knowledge requirements.
Monte Carlo Tree Search (MCTS) a powerful approach that utilizes a search tree with simulation and backpropogation mechanisms to train agents with a small data budget [[2] To model high- dimensional observation spaces and complex policy behaviour, AlphaGo [3] enhances MCTS with deep neural networks and designs the policy and value network that identify optimal actions and winning rates respectively, which was the first to defeat the strongest professional human player in Go. Despite the impressive results, MCTS-style algorithms rely on series of necessary conditions, such as knowledge of game rules and simulators, discrete action space and deterministic state transition, which severely restrict the application scope of these methods. In recent years, several successors
"Corresponiing Author
37th Conference Neural Information Processing Systems (NeurIPS 2023) Track Datasets and Benchmarks.

---

to AlphaGo have attempted to extend its capabilities in various directions. MuZero requirements for prior knowledge of environments by training a set of neural networks to reconstruct reward, value and policy. Sampled MuZero [3] successfully applies MCTS to various complex action space with novel planning mechanism based on sampled actions. [6 [8] improve MuZero in terms of planning stochasticity, representation learning effectiveness and simulation efficiency respectively. These emerging algorithm insights and techniques have contributed to the development of more general MCTS algorithms and toolchains.
14 relaxes the
In this paper, we present a unified algorithm benchmark named LightZero that first comprehensively integrates different MCTS/MuZero algorithm branches, including 9 algorithms and more than 20 decision environments with detailed evaluation. To better understand the potential of MCTS as an efficient general-purpose sequential decision solver, revisit the development history of MCTS methods [9] and the diverse criterions of newly proposed RL environments (0
21 22 As shown in Figure 2 outline the six most challenging dimensions in developing LightZero as a general
method, including multi-modal and high- high-dimensional observation space [23], complex action space, reliance on prior knowledge, inherent stochasticity, simulation cost, and hard exploration.
Furthermore, highly coupled algorithm and system architectures greatly increase the cost and barriers of migrating and improving MCTS-style methods. Some special mechanisms like tree search and data reanalyze 24 seriously hinder the simplification and parallel acceleration of code implementation. To overcome these difficulties, LightZero designs a modularly pipeline to enable distinct algorithm components as plug-ins ins. For example, the chance node planning for modelling stochasticity can also be used in continuous control or hybrid action environments. From the unified viewpoint provided by LightZero, we can systematically divide the whole training scheme MCTS-style methods into four sub-moduless data collector, data arranger, agent learner, and agent evaluator. LightZero's decoupled architecture empowers developers focus intensively the customization of environments and algorithms. Meanwhile, some techniques like off-policy correction and data throughput limiter can ensure the steady convergence of the algorithm while achieving runtime speedups.
Based on these supports, LightZero also explores the advantages of combining some novel insights from model-based RL with MCTS approaches. particular, the misalignment problem [25] of state representation learning and dynamics learning can result in the problematic optimization for MuZero, thus simple self-consistency loss can significantly speed up convergence without special tuning. Besides, intrinsic reward mechanism [26) (27 (28) can address the exploration deficiency of tree-search methods with hand-crafted noises. Subsequently, we evaluate the ability of LightZero as general solver for various decision problems. Experiments on different types of environments demonstrate LightZero's 's rich application ranges and data efficiency regimes with few hyper- pprammterree adjustments. At last, we provide discussions on the future optimization directions of each sub-module.
In general, we summarize the three key contributions of this paper as follows:
We present LightZero, the first general MCTS/MuZero algorithm benchmark that systematically evaluates related algorithms and system designs.
We outline the most critical challenges of real-world decision applications. To address these issues, we decouple the algorithm and system design of MCTS methods and design a modular training pipeline, which can easily integrate novel insights for better scalability.
We demonstrate the capability and future potential of LightZero as general sequential decision solver, which can be trained and deployed across diverse domains.
2 Background
Reinforcement Learning models a decision-makinn problem as a Markov Decision Process (MDP) M = AS, P, R, 7, Po). where and A denote the state space and action space, respectively. The transition function P maps S x A to S, while the expected reward function R maps S x AA IR.
The discount factor 7 E [0, 1) determines the importance of future rewards, and po represents the initial state distribution. The goal of RL is to learn policy : S + A that maximizes the expected discounted return over the trajectory distribution J (TT) E,, po. P. R D rrrel.

---

LightZero
dequed Neutheee
Goald Aelter
Seare Canns
Abai
for ves
66 Medtied
Galliger
Figure I: Overview of LightZero. The left side depicts the development of MCTS, while the right side showcases various RL environments. LightZero incorporates and extends recent advances within the MCTS/MuZero sub-domain and effectively applies them across diverse environments.
AlphaZero [29] is generalized version of AlphaGo [13, eliminating the reliance on supervised learning from game records. It is trained entirely through unsupervised self-play and achieves superhuman performance in various board games, such chess, shogi, and Go. This approach replaces the handcrafted features and heuristic priors commonly used in traditional intelligent programs, Specifically, AlphaZero employs deep neural network parameterized by 8, represented as (p, v) = [o().. Given a board position 8, the network produces a action probability Pa = (als)) for each action a and a scalar value to predict the expected return 2, i.c. + 2.
MuZero 14] achieves superhuman performance in more complex domains with visual input 30) without knowledge of the environment's transition rules. It combines tree search with a learned model, using three networks: ! Representation Network: s he(o1.... 0z). This network represents the root node (at time () as a latent state, obtained by processing past observations 01,-..,0 Or. 2 Dynamics Network: r*, s* gs* ,a²). This network simulates the dynamics of the environment. Given a state and selected action, outputs the transitioned next state and corresponding reward. 33) Prediction Network: pa v fs()... Given a latent state, this network predicts the action probability and value. Notably, MuZero searches within the learned latent space. For the MCTS process in MuZero, assume the initial root node so is generated from the original board state through the representation network, each edge stores the following information: [ a), P(s, a), Q(s, a), R(s, a), ((s,, a), respectively representing visit counts, policy, mean value, reward, and state transition. The MCTS process in the latent space can be divided into three phases:
Selection: Actions are chosen according the Upper Confidence Bound (UCB) [3 formula:
VE,N(s,)) a= = arg maz Q(s, a) P(s, a) 1+ a(s.a) (c, + log( E. N b) + + +1]
C2
where, N represents the visit count, Q is the estimated average value, and P is the policy's prior probability. C1 and 02 are constants that control the relative weight of P and Q.
Expansion: The selected action is executed in the learned model, continuing until a leaf node is encountered, At this point, a new state node sf is generated, and its associated predicted reward r' is determined. Utilizing the prediction function, we obtain the predicted values p' and vi. Subsequently, this node incorporated into the search tree.
Backup: The estimated cumulative reward at step k is calculated based v', denoted as: Gh 1-1-k
TT kkt Subsequently, Q and N are updated along the search path.
After the search is completed, the visit count set N(s, a) is returned at the root node 8p- These visit counts are normalized to obtain the improved policy:
I- (als) as
Σ b)
3

---

Multi-modal Observation Space
Data Efficiency
Complex Action
Model-free RL
AlphaZero
MuZero
LightZero
Hard Exploration
Inherent Stochasticity
on Knowledge
Figure 2: Radar chart comparison of MCTS-style methods and model-free RI (e.g. PPO) on six environment challenges and another data efficiency dimensions. We categorize the critical capabilities of general decision solvers as follows: multi-modal observation space, complex action space, inherent stochasticity, reliance prior knowledge, simulation cost, hard exploration and data efficiency. Each curve in the graph represents the score of an algorithm across these six categories. score of 1 indicates that the algorithm perform poorly in this dimension and is only applicable to limited scenarios, while a higer score means a large application scope and better performance. In particular, mreel-free RL methods need simulation and have little dependence score corresponding dimensions. Please note that within this context, the term LightZero refers to the specical algorithm that embodies the optimal combination of techniques and hyperparameter settings within our framework. Details about qualitative score rules can be found Appendix [
priors, it achieves high
where T is the temperature coefficient controlling the exploration degree. Finally, action is sampled from this distribution for interaction with the environment self- saff-play During the learning phase, MuZero perform end-to-end training with the following loss function, where IP, 1" and !" are loss functions for policy, value and reward respectively, and the final term is weight decay.
K
((0)) = ((T+H.PP)) + (2444s e/) + (" +c|00
3 LightZero
In this section, we will first introduce the overview LightZero, followed by a comprehensive analysis of challenges various decision environments. Additionally, we propose specific training pipeline design for a modular and scalable MCTS toolchain. We will conclude this section with two algorithm insights inspired by the decoupled design of LightZero.
3.1 Overview
As shown in Figure LightZero the first benchmark that integrates almost all recent advances in the MCTS/MuZero sub-domain. Specifically, LightZero incorporates nine key algorithms derived from the original AlphaZero 229 establishing a standardized interface for training and deployment across diverse decision environments. Unlike the original versions of these derived algorithms, which focused specific avenues of improvement, LightZero provides unified viewpoint and interface. This unique feature enables exploration and comparison of all possible combinations of these techniques, offering an comprehensive baseline for reproducible and accessible research. The concrete experimental results are thoroughly described in Section 4 and Appendix B
3.2 How Evaluate General MCTS Algorithm: Environment Challenges
The algorithm extensions integrated in LightZero have greatly relaxed the constraints and broadened the applicability of MCTS-sty style methods. In the following part, we hope to delve deeper into the key

---

Data Collector
Data Arranger
Prioriiiedd Repley Buffer (FIFo)
Offinne
5.9
5.5
4.1
Agent
Rewward
Eri
priority recomputar
throughput liriler
Online Agent-Envv Interaction Loop
reenalyzer
Search Tree
Action
Contest
Agent Learner
Exchanger
Agent Evaluator
representation
Training Techniques/Trkkk
model-based rellout
Networks
(Mocel)
Agent
State/Reaard
Env
learning
Orline Agent-Env Interaction Loop
distributional prediction
Search Tree
Action
normaizatio
oplimizer gradient diip
data paralel mixed precision
evaluation tricks
evaluation metrics
efficiency profier
Figure Four core sub-modules of the training pipeline LightZero. Context Exchanger is responsible for transporting configurations, models and trajectories among different sub-modules.
issues in the design of general and efficient MCTS algorithms. In order to systematically complement this endeavor, we conducted an analysis of a set of classic and newly proposed RL environments to identify common characteristics. Based on this analysis, we have summarized six core challenging dimensions, which are presented in a radar plot depicted Figure 2 Concretely, The intentions and goals of six types of envionmental capabilities are: 1) Multi-modal observation spaces pose challenge for agents as they must be able to extract different representation modalities (c.g.. low- dimensional vectors, visual images, and complex relationships) while effectively fusing distinet embeddings. 2) Complex action space necessitates the agent's proficiency in generating diverse decision signals, encompassing discrete action selection, continuous control, and hybrid structured action space. 3) Reliance on prior knowledge is a major drawback of methods like AlphaZero. These approaches inherently require accessibility to perfect simulator and specific rules of the environment. In contrast, MuZero and its derived methods address this limitation by learning an environment model substitute the simulator and related priors. Inherent stochasticity presents a fundamental challenge in tree-search-baad based planning methods. The uncertainty of environment dynamics and partially observable state spaces both can lead to misalignment of planning trajectories, resulting in large number of useless or conflicting search results. 5) Simulation cost stands as the primary contributor to wall-time consumption for MCTS-style methods. At the same time, the algorithm performance will degrade a lot the algorithm fails to visit all the necessary actions during the simulation process. 6) Hard exploration represents crucial challenge that is often overlooked. While search trees can enhance efficiency by reducing the scope of exploration, MCTS-style methods susceptible difficulties in environments with numerous non-terminating cases, such mazes.
3.3 How to Simplify A General MCTS Algorithm: Decouple Pipeline into Sub-Modules
The impressive performance of MCTS-style methods is often accompanied by notable drawback: the complexity of implementations, which greatly restricts their applicability. In contrast to some classic model-free RL algorithms like DQN 32 and PPO [S, -CTS-style methods require multi- step simulations using search trees at each agent-environment interaction. Also, to improve the quality of training data, MuZero Unplugged (24 introduce a data reanalyze mechanism that uses the newly obtained model to compute improved training targets on old data. However, both of these techniques require multiple calls to simulators neural networks, increasing the complexity across various aspects of the overall system, including code, distributed training, and communication topology.
5

---

Therefore, it is necessary to simplify the whole framework based on the integration of algorithms. Figure 3 presents a depiction of the complete pipeline of LightZerro with four core sub-modules.
Firstly, LightZero offers support for both online and offline RL [33] training schemes. The distinction between them lies in the utilization of either an online interaction data collector or direct usage of an offline dataset. Secondly, LightZero restructures its components and organizes them into four main sub- modules,, based on the principle of high cohesion and low coupling. Data collector is responsible for efficient action selection using policy network and search tree. It also contains various exploration strategies, data pre-processing and packaging operations. Data arranger plays a unique role in MCTS by effectively storing and preparing valuable data for training purposes. This sub-module involves the data reanalyze technique (24 to correct oof-po and even offline data. Furthermore, the
modified priority sampling [34] ensures training mini-batches have both sufficient variety and high learning potential. To balance these tricks with efficiency, the throughput limiter controls the ratio of adding and sampling data to ensure optimal data utilization within fixed communication bandwith. Agent learner is responsible for training multiple networks. It can be enhanced through various optimization techniques, such as self-supervised representation learning 35 36], model-based rollout
These techniques contribute to the policy improvement and further enhance the overall performance of the agent. Agent evaluator periodically provides the diverse evaluation metries 42 to monitor the training procedure and assess policy behaviour. It also integrates some tifmeence-time tricks like beam search 43) enhance test performance. We provide detailed analysis of how these sub-modules are implemented in specific algorithms in Appendix F Built upon these abstractions, LightZero serves as valuable toolkit, enabling researchers and engineers to develop enhanced algorithms and optimize systems effectively. learned model in MCTS
37 38). distributional predicton (39] and normalization 40 41.
For example, the exploration strategies and ensuring the aligmment crucial, and this will be discussed in the subsequent sub-section. addition, exploring parallel schemes for multiple vectorized environments and search trees can be an insightful topic for machine learning system. The associated dataflow and overhead analysis will be presented in the Appendix
3.4 How to Improve A General MCTS Algorithm: 2 Examples
In this section, we present two algorithm improvement examples inspired by LightZero. The below dimensions pose necessary challenges in designing a comprehensive MCTS solver. LightZero ad- dresses these challenges through various improvements, resulting in superior performance compared to individual algorithm variants across different domains (Section a)d
Intrinsic Exploration While tree-search- based methods perform well in board games with only eventual reward, they may encounter challenges or perform poorly in other environments with sparse rewards, such as MiniGrid [4). One crucial distinction between these two problems is that in the former, the search tree can always reach several deterministic final states, whereas in the latter, may encounter various non-termination states due to the limitation of maximum episode length. To address this issue, LightZero incorporates the idea of intrinsic reward methods (28) and implement it efficienctly within MuZero's 's learned models. Further details can be found in Section 5.1
A Alignment in Environment Model Learning MuZero employs a representation network to generate latent states and dynamics network to predict next latent states. However, there is no explicit super- vision guiding the desired properties of the latent space. Traditional self-supervised representation learning methods often fail to align these proxy tasks with RL objectives. The difference of rollouts between the perfect simulator and the learned model is also a problems that can not be ignored. Further exploration of misalignments across different environments are discussed in Section 3.2
4 Experiments
In Section 4.1
we initially present some representative experiments of LightZero. with detailed experimental settings and more comprehensive results outlined in the Appendix B Subsequently, in Section 4.2 we delve into key observations and reflections based on these benchmark results, introducing some critical insights. Particularly regarding the exploration and the alignment issues of environment model learning, conduct in-depth experimental analysis in Section

---

ChertNoF ChertNoFramcskip-vk
BreakoutNof PrrnkuuuNFFrmmssipp-v
MsPacmanNoFramessee
8003
160ro SiL.
2000
11-00
strd)
6093
430
Robhm
Resuth
900
4200
230
1000
2390
130
16
0.21
3.4M Env Stops
Eny Steps
Env Steps
GoBigger
BipedalVallerrvV
Gomoku
100300
300
1.00
Ho go
0.75
0.50
Return
60300
Roturn
0.28
93
-0.50
20300
-0.75
-1.00
0.00M
Env Steps
Env Steps
Ent Steps
5:2M
0.4M
Figure 4: Comparisons of mean episode return for algorithm variants LightZero across diverse environments: Atari with discrete action and oostial-obeervable state (Obert, Breakout, MsPacman), GoBigger 23 with complex observation and multi-agent cooperation, continuous control with environment stochasticity (Bipedalwalker), and Gomoku with varying accessibility to simulator.
4.1
Benchmark Results
To benchmark the difference among distinct algorithms and the capability of LightZero general decision solver, we conduct an extensive comparisons across diverse range of RI, environments, The algorithm variants list contains AlphaZero 29 MuZero
177. Sampled MuZero 18 and other improved versions of LightZero. For each scenario, we evaluate all the possible variants on corresponding environments. In Figure 4 we show some selected results as examples. For detailed settings, metrics, comprehensive benchmark results and related analysis, please refer Appendix B
14 EfficientZero
[1]. Stochastic MuZero [16] Gumbel MuZero
4.2 Key Observations and Insights
Building on the unified design of LightZero and the benchmark results, have derived some key insights about the strengths and weaknesses of each algorithm, providing comprehensive understanding of these algorithms' performance and potential applications.
01: In board game environments, AlphaZero's sample efficiency greatly exceeds that of MuZero. This suggests that employing AlphaZero directly is advantageous when an environment simulator available: however, MuZero can still achieve satisfactory results even the absence of simulator.
02: Self-supervised loss substantially enhances performance in most Atari environments with image inputs. Figure demonstrates that MuZero with SSL performs similarly to MuZero in MsPacman, while outperforming it in the other five environments. This result highlights the importance of SSL for aligning the model and accelerating the learning process image input environments.
03: Predicting value _prlee__eefxxxeee instead of reward does guarantee performance enhancement. For example, in Figure 7 EfficientZero outperforms MuZero with SSL only in the MsPacman and Breakout environments, while showing similar performance in the other environments. In certain specific scenarios, such as the sparse reward environments depicted in Figure 12 EfficientZero's performance is significantly inferior to that of MuZero with SSI.. Therefore, we should prudently decide whether to predict value -peffxx taking into account the attributes the environment.
04: MuZero with SSL and EfficientZero demonstrate similar performance across most Atari environ- ments and in complex structured observation settings, such GoBigger. This observation suggests
7

---

that environments with complex structured observations can benefit from representation learning and contrastive learning techniques (35 to enhance sample efficiency and robustness.
In discrete action spaces, Sampled EfficientZeroos performance is correlated with action space dimensions. For instance, Sampled EfficientZero performs on par with EfficientZero in Breakoul (action space dimension of 4), but its performance decreases in MsPacman (dimension of 9).
05:
Sampled EfficientZero with Gaussian policy representation is more scalable in continuous action spaces. The Gaussian version performs well in traditional continuous control and MuJoCo environments, while factored discretization 45 is limited to low-dimensional actions.
06:
07: Gumbel MuZero achieves notably better performance than MuZero when the number of simulations is limited, which exhibits its potential in designing low time-cost MCTS agent.
08: environments with stochastic state transitions or partial observable states (such as Atari without stacked frames), Stochastic MuZero can obtain slightly better performance than MuZero.
The self superviee loss proposed in sampling-relater techniques in Sampled MuZero computational improvements in Gumbel MuZero [18] for utilizing MCTS searched information, and environment stochasticity modeling in Stochastic MuZero [ are orthogonal to each other, exhibiting minimal interference. LightZero is exploring and developing ways to seamlessly integrate these characteristics to design a universal decision makinn algorithm.
09:
45
5
Two Algorithm Case Studies for LightZero
5.1 Exploration Strategies in MCTS
Motivation Finding the optimal trade-off between exploration and exploitation is a fundamental challenge in RL. It is well-known that MCTS can reduce the policy search space and facilitate exploration. However, there exists limited research on the performance of MCTS algorithms in hard- exploration environments. Based on the above benchmark results, we conduct a detailed analysis of the algorithm behaviours between challenging sparse reward environments and board games, as well insights behind the selection of exploration mechanisms in this section and Appendix C.1
Settings We performed experiments in MiniGrid environment, mainly on the KeyCorridorS333 and FourRooms scenarios. Expanding upon the naive setting (handcrafted temperature decay), we conducted a comprehensive investigation of six distinet exploration strategies in LightZero. detailed description of each exploration mechanism provided in Appendix C.I
MinDGrikKYyyComddoS333---
Foiioid-FourRooms-00
0.6
0.8
0.5
Return
Return
0.4
0.4
0.2
0.1
0.0
0.0
0M
0.2M
0.4M
0.584 0.8RA
0M
1M
2M
3M
4M
Env Steps
Env Steps
Figure 5: Performance of various MCTS exploration mechanisms in MiniGrid environment (Return during the collection phase). Under the naive setting, the agent fails due inadequate exploration. Merely increasing search budgets with the NaiveDoubleSimulann significant improvement. EpsGreedy, FixedTemperature and PolicyEntropy Regularoyygullaizatttio display higher variance as they cannot guarantee enough exploration. Intrinsic ExpronsicExllorattonn effectively explores the state space by leveraging curiosity mechanisms, resulting in the highest sample efficiency.
approach does not yield any
Analysis Figure 5 indicate that simply increasing search budgets does not yield improved perfor- mance in challenging exploration environments. Instead, implementing a larger temperature and

---

PongNofF rmmesip.vv v4
Lunari Lunartannerr-2 ander--22
Ticl Tac Toe
20
1.00
0:22
Retum
0.25
-1000
Maliow' 30L
Madloo 50L
PM Env Steps
114
2.5M
0.30
Env Steps
Eni Slepe
Figure 6: Impact of self-supervised consistency loss across different environments with various types of observations. From left to right, performance comparisons involve standard image input, compact vector imput, and unique board image input, considering cases with and without consistency loss. Experiments show that the consistency loss proves to be critical only for standard image input.
incorporating policy entropy bonus can enhance action diversity during data collection, albeit the cost of increased variance. However, theoretically, they cannot guarantee sufficient exploration, often resulting in mediocre performance and a higher likelihood of falling into local optima due to policy collapse. Epsilon-grecde exploration ensures small probability of uniform sampling, which aids in exploring areas with potentially high returns. EpsCreedy has varying effects in different environments in early stages, but theoretically, due to its ability to ensure sufficient exploration, it may achieve good results in the long run. more effective strategy involves curiosity-driven techniques, such as RND (27), which assigns higher intrinsic rewards novel state-action pairs, bolstering the efficiency of exploration. The performance of the IntrinsicExploration method supports this assertion, and can be integrated into MuZero with minimal overhead (Appendix C.1.3)
5.2 Alignment in Environment Model Learning
environment models are vital for MuZero-style algorithms, with factors such as model structure, objective functions, and optimization techniques contributing to their success. The consistency loss proposed in [17] could serve as an approach for aligning the latent state generated by the dynamics model with the state obtained from the observation. In this section, we investigate the impact of consistency loss on learning dynamic models and final performance in environments with diverse observations (vector, standard images, special checkerboard images).
Motivation Aligned and scalable 225
Settings To study the impact of the consistency loss on various types of observation data, we employ the MuZero algorithm as our baseline. To ensure the reliability of our experimental results, we maintain the same configurations across other settings, with additional experimental details provided in Appendix C.2 In the experiments, we use Pong as the environment for image input, LunarLander for continuous vector input, and TicTacToo for special image input (checkerboard) environments.
Analysis In Figure d consistency loss is critical for standard image input. Removing the consistency loss results in significant decline in performance, indicating the challenge of learning dynamic model for nigh-ddmensiooal inputs. For vector input environments like LunarLander, consistency loss provides minor advantage, suggesting that learning dynamic model relatively easier on the compact vector observations. In special two dimension input environments like TicTacToe, the consistency loss remains large, highlighting the difficulty of achieving consistency between latent state outputs. Additionally, adding consistency loss with inappropriate hyper-parameter may lead to non-convergence (Appendix C.2). In conclusion, our experiments demonstrate that the effectiveness of the consistency loss depends on the special observation attributes. For board games, a future research direction involves investigating suitable loss functions to ensure alignment during training.
6 Related Work
Sequential Decision-making Problems the domain of sequential decision decision-making problems, intelligent agents aim to make optimal decisions over time, taking into account observed states and prior actions (46). However, these problems are often compounded by the presence of continuous action spaces, dynamic transitions, and exploration difficulties. To address such problems, model-free ffeee-free RL methods [5, 7 33) focus on learning expected state rewards, optimizing actions, combining

---

both strategies to achieve optimal policy learning, Conversely, model-based RL [25] incorporates the environment's 's transition into its optimization objective, aiming to maximize the expected return on trajectory distribution. MCTS [9] is a modeling approach derived from search planning algorithms such as minimax [47 and alpha-beta pruning [48]. Unlike these algorithms, which recursively search decision paths and evaluate their returns, MCTS employs heuristic search on prior-guide simulations, effectively addressing excessive search consumption in complex decision spaces.
MCTS Algorithms and Toolkits Despite the impressive performance and efficiency of the MCTS+RI. approach, constructing the training system and dealing with intricate algorithmic details pose significant challenges when applying these algorithms to diverse decision intelligence domains. Recent research has made progress in this direction. MuZero Unplugged [24] introduced the Rean- alyze technique, simple and efficient enhancement that achieves good performance both online and offline. ROSMO [49] investigated potential issues with MuZero in offline RL and suggested regularized one-step lookahead approach. The lack of comprehensive open-source implementations of various algorithms remains challenge within the research community. For example. Sampled MuZero 15] lacks a public implementation. AlphaZero-Geeneme (5 and MuZero-Gennral [51 each support only a single algorithm, and neither offers distributed implementations. Although EfficientZero [17], does support multi-GPU implementation. it is limited to the single algorithms. KataGo [52]. while primarily focused on the AlphaZero and Go game, requires substantial computa- tional resources during training. potentially posing hardware barriers for ordinary users. As result, the research community continues to seek more efficient and enhanced open-source tools.
Standardization and Reproducibility In the realm of Deep RI., the quest for standardizing algorithm coupled with the creation of unified benchmarks has ascended as focal points of growing significance. [53] emphasize the critical necessity of not only replicating existing work but also accurately assessing the advancements brought about by new methodologies. However, the process of reproducing extant Deep RL methods is far from straightforward, largely due the non-determinism inherent in environments and the variability innate to the methods themselves, which can render reported results challenging to interpret. (54) proposed set of metrics for quantitatively measuring the reliability of RL algorithms. These metrics, focusing on variability and risk both during and after training, are intended to equip researchers and production users with tools to evaluate and enhance the reliability of RI. algorithms. In [55], a large-scale empirical study was undertaken to identify the factors that significantly influence the performance of on-policy RL algorithms within continuous control tasks. The insights gameered from this research offer valuable, practical suggestions for the training of on-policy RL algorithms. Despite these advancements, there remains a noticeable dearth of work specifically investigating benchmarks and the details of reproducing studies in the domain of MCTS.
7 Conclusion
In this paper, we introduce LightZero, the first unified algorithm benchmark to modularly integrates various MCTS- RL methods, systematically analyze and address the challenges and opportu- nities for deploying MCTS general and efficient decision solver. Through the incorporation of decoupled system design and novel algorithm insights, we conduct detailed benchmarks and demonsrate the potential of LightZero as scalable and efficient decision-making problem toolchains for the research community. Besides, based on this benchmark and related case studies, We also discuss existing limitations and valuable topics for future work in Appendix
8 Acknowledgements
This project is funded in part by National Key R/D Program of China Project 2022ZD0161 100, by the Centre for Perceptual and Interactive Intelligence (CPII) Ltd under the Innovation and Technology Commission (ITC)'s 's InnoHK, by General Research Fund of Hong Kong RGC Project 14204021. Hongsheng Li is of CPII under the InnoHK. We thank several members of the SenseTime and Shanghai AI Laboratory for their help, support and feedback on this paper and related codebase. We especially thank Shenghan Zhang for informative and inspiring discussions at the beginning of this project. We are grateful to the assistance of Qingzi Zhu for many cute visual materials of LightZero.
10

---

References
[1] Adrià Puigdomènech Badia, Bilal Piot, Steven Kapturowski, Pablo Sprechmann, Alex Vitvit- skyi, Zhaohan Daniel Guo, and Charles Blundell. Agent57: Outperforming the atari human benchmark. In International conference on machine learning, pages 507-517. PMLR, 2020.
[2] Linxi Fan, Guanzhi Wang, Yunfan Jiang, Ajay Mandlekar, Yuncong Yang. Haoyi Zhu, Andrew Tang, De-An Huang, Yuke Zhu, and Anima Anandkumar. Minedojo: Building open-ended embodied agents with internet-scale knowledge. arXiv preprint arXiv:2206.088. 2022.
[3] Oriol Vinyals, Igor Babuschkin, Wojeiech M. Czarnecki, Michaël Mathicu, Andrew Dudzik, Junyoung Chung, David H. Choi, Richard Powell, Timo Ewalds, Petko Georgiev, Junhyuk Oh, Dan Horgan, Manuel Kroiss, Ivo Danihelka, Aja Huang, Laurent Sifre, Trevor Cai, John P. Agapiou, Max Jaderberg, Alexander Sasha Vezhnevets, Rémi Leblond, Tobias Pohlen, Valentin Dalibard, David Budden, Yury Sulsky, James Molloy, Tom Le Paine, Çaglar Gülçehre, Ziyu Wang, Tobias Pfaff, Yuhuai Wu, Roman Ring, Dani Yogatama, Dario Wünsch, Katrina McK- inney, Oliver Smith, Tom Schaul, Timothy Lillicrap, Koray Kavukcuoglu. Demis Hassabis, Chris Apps, and David Silver. Grandmaster level in starcraft II using multi aulti--geent reinforcement learning. Nat., 575(7782)3350-354, 2019.
[4] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv:1312.5602, 2013.
[5] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. CoRR, abs/ 2017.
[6] Steven Kapturowski, Georg Ostrovski, John Quan, Remi Munos, and Will Dabney. Recurrent experience replay distributed reinforcement learning. In International conference on learning representations, 2019.
[7] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off- policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning, pages 1861-1870. PMLR, 2018.
[8] Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, et al. Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures. In International conference on machine learning, pages 1407-1416. PMLR, 2018.
[9] Dan Horgan, John Quan, David Budden, Gabriel Barth-Maron, Matteo Hessel, Hado Van Hasselt, and David Silver. Distributed prioritized experience replay. arXiv preprint arXiv:1803.0093: 2018
[10] Lasse Espeholt, Raphaël Marinier, Piotr Stanczyk, Ke Wang, and Marcin Michalski. Seed rl: Scalable and efficient deep-rl with accelerated central inference. arXiv preprint 1rXiv:1910.0659,, 2019.
[11] Christopher Berner, Greg Brockman, Brooke Chan, Vicki Cheung, Przemystaw Dębiak, Christy Dennison, David Farhi, Quirin Fischer, Shariq Hashme, Chris Hesse, et al. Dota 2 with large scale deep reinforcement learning. arXiv preprint 1912 2019.
[12] David Silver and Joel Veness. Monte-carlo planning in large pomdps. Advances in neural information processing systems, 23, 2010.
[13] David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driess- che, Julian Schritwweseer, loannis Antonoglou, Veda Panneershelvam, Mare Lanctot, et al. Mas- tering the game of go with deep neural networks and tree search. nature, 529(7587):484 489, 2016.
11

---

[14] Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, Timothy P. Lillicrap, and David Silver. Mastering atari, go, chess and shogi by planning with a learned model. CoRR, 0tt//9911.88266. 2019.
[15] Thomas Hubert, Julian Schrittwieser, loannis Antonoglou, Mohammadamin Barekatain, Simon Schmitt, and David Silver. Learning and planning complex action spaces. In International Conference on Machine Learning, pages 4476-4486. PMLR, 2021.
[16] Ioannis Antonoglou, Julian Schrittwieser, Sherjil Ozair, Thomas Hubert, and David Silver. Planning stochastic environments with a learned model. In International Conference on Learning Representations, 2021.
[17] Weirui Ye, Shaohuai Liu, Thanard Kurutach, Pieter Abbeel, and Yang Gao. Mastering atari games with limited data. Advances in Neural Information Processing Systems, 34:25476-25488, 2021.
[18] Ivo Danihelka, Arthur Guez, Julian Schritwwieser, and David Silver. Policy improvement by planning with gumbel. In International Conference on Learning Representations, 2022.
[19] Maciej Šwiechowski, Konrad Godlewski, Bartosz Sawicki, and Jacek Mańdziuk. Monte carlo tree search: review of recent modifications and applications. Artificial Intelligence Review, 56(3):2497-2562, 2023.
[20] Ian Osband, Yotam Doron, Matteo Hessel, John Aslanides, Eren Sezener, Andre Saraiva, Katrina McKinney, Tor Lattimore, Csaba Szepesvari, Satinder Singh, et al. Behaviour suite for reinforcement learning. arXiv preprint arXiv: 0908003508.0.... 2019.
[21] Open Ended Learning Team, Adam Stooke, Anuj Mahajan, Catarina Barros, Charlie Deck, Jakob Bauer, Jakub Sygnowski, Maja Trebacz, Max Jaderberg, Michael Mathieu, al. Open- ended learning leads to generally capable agents. arXiv preprint arXiv:2107.12800 2021.
[22] Joshua Albrecht, Abraham Fetterman, Bryden Fogelman, Ellie Kitanidis, Bartosz Wröblewski. Nicole Seo, Michael Rosenthal, Maksis Knutins, Zack Polizzi, James Simon, et al. Avalon: benchmark for generalization using procedurally generated worlds. Advances in Neural Information Processing Systems, 32:1281::1111 1-2222,, 2022.
[23] Ming Zhang, Shenghan Zhang, Zhenjie Yang, Lekai Chen, Jinliang Zheng, Chao Yang, Chuming Li, Hang Zhou, Yazhe Niu, and Yu Liu. Gobigger: A scalable platform for cooperative- competitive multi-agent interactive simulation. In The Eleventh International Conference on Learning Representations, 2023.
[24] Julian Schrittwieser, Thomas Hubert, Amol Mandhane, Mohammadamin Barekatain, Ioannis Antonoglou, and David Silver. Online and offline reinforcement learning by planning with learned model. Advances in Neural Information Processing Systems, 34:27580-27591, 2021.
[25] Raj Ghugare, Homanga Bharadhwaj, Benjamin Eysenbach, Sergey Levine, and Ruslan Salakhut- dinov. Simplifying model-based rl: Learning representations, latent-space models, and policies with one objective. arXiv preprint arXiv:2209.084.. 2022.
[26] Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Darrell. Curiosity- exploration by self-supervised prediction. In International conference on machine learning, pages 2778- 2787. PMLR, 2017.
[27] Yuri Burda, Harrison Edwards, Amos Storkey, and Oleg Klimov. Exploration by random network distillation. arXiv preprint arXiv:1810.12880 2018.
[28] Adrià Puigdomênech Badia, Pablo Sprechmann, Alex Vitvitskyi. Daniel Guo, Bilal Piot, Steven Kapturowski, Olivier Tieleman, Martín Arjovsky, Alexander Pritzel, Andew Bolt, al. Never give up: Learning directed exploration strategies. arXiv preprint arXiv:2002.06038 2020.
12

---

[29] David Silver, Thomas Hubert. Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Marc Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, et al. Mastering chess and shogi by self-play with a general reinforcement learning algorithm. arXiv preprint 17X10:17/ 2017.
[30] Marc G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47:253-279 2013.
[31] Christopher D Rosin. Multi-armed bandits with episode context, Annals of Mathematics and Artificial Intelligence, 61332003--,,, 2011.
[32] V. Mnih, K. Kavukcuoglu, D. Silver, A. A. Rusu, Veness, M. G. Bellemare, A. Graves, M. A. Riedmiller, A. Fidjeland, G. Ostrovski, S. Petersen, C. Beattie, A. Sadik, I. Antonoglou, H. King, D. Kumaran, D. Wierstra, S. Legg, and D. Hassabis. Human-level control through deep reinforcement learning. Nature, 518(7540):529-533, 2015.
[33] Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643; 2020.
[34] Yixuan Mei, Jiaxuan Gao, Weirui Ye, Shaohuai Liu, Yang Gao, and Yi Wu. Speedyzero: Mastering atari with limited data and time. In The Eleventh International Conference on Learning Representations, 2023.
[35] Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised represen- tations for reinforcement learning. In International Conference on Machine Learning, pages 5639-5650. PMLR, 2020.
[36] Ilya Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. arXiv preprint arXiv:2004. 13649, 2020.
[37] Jacob Buckman, Danijar Hafner, George Tucker, Eugene Brevdo, and Honglak Lee. Sample- efficient reinforcement learning with stochastic ensemble value expansion. Advances in neural information processing systems, 31, 2018.
[38] Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model-based policy optimization. Advances in neural information processing systems, 32, 2019.
[39] Will Dabney, Mark Rowland, Mare Bellemare, and Rémi Munos. Distributional reinforce- ment learning with quantile regression. In Proceedings of the AAAI Conference on Artificial Intelligence, 2018.
[40] Hado P van Hasselt, Arthur Guez, Matteo Hessel, Volodymyr Mnih, and David Silver. Learning values across many orders of magnitude. Advances in neural information processing systems, 29, 2016.
[41] Tobias Pohlen, Bilal Piot, Todd Hester, Mohammad Gheshlaghi Azar, Dan Horgan, David Budden, Gabriel Barth-Maron, Hado Van Hasselt, John Quan, Mel Večerík, et Observe and look further: Achieving consistent performance atari. arXiv preprint 11X.v:1805.11593, 2018.
[42] Stephanie CY Chan, Samuel Fishman, John Canny, Anoop Korattikara, and Sergio Guadarrama. Measuring the reliability of reinforcement learning algorithms. arXiv preprint arXiv:1912.05663. 2019.
[43] Sam Wiseman and Alexander M Rush. mization. arXiv preprint arXiv:1606.02966 2016.
learning as beam-search opti-
13

---

[44] Maxime Chevalier-Boisse Beiisvoiieet Lucas Willems, and Suman Pal. Minimalistic gridworld environ- ment for openai gym. httts:///githul htt.h//w....oomm..im/meee/aece/e wttp://aaiiiaea/c/yy.m.aaaaaaia.
gym-minigmiidd 2018.
[45] Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Mohammadamin Barckatain, Simon Schmitt, and David Silver. Learning and planning in complex action spaces. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, ICMI. 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pages 4476-4486. PMLR, 2021.
[46] R. S. Sutton and A. G. Barto. Reinforcement learning: An introduction. IEEE Transactions on Neural Networks, 16:285-286, 1988.
[47] Ronald L Rivest. Game tree searching by min/max approximation. Artificial Intelligence. 34(1):77-96, 1987.
[48] Donald Knuth and Ronald W Moore. An analysis of alpha-beta pruning. Artificial intelligence, 6(4):293-326, 1975.
[49] Zichen Liu, Siyi Li, Wee Sun Lee, Shuicheng Yan, and Zhongwen Xu. Efficient offline policy optimization with a learned model. In International Conference on Learning Representations, 2023.
[50] Shantanu Thakoor, Surag Nair, and Megha Jhunjhunwala. Learning to play othello without human knowledge, 2016.
[51] Aurèle Hainaut Werner Duvaud. Muzero general: Open reimplementation of muzero. https gegteuub..o//eerrrreddvvvuuuuuue generra 2019.
[52] David J Wu. Accelerating self-play learning in go. arXiv preprint arXiv:1902.10550 2019.
[53] Peter Henderson, Riashat Islam, Philip Bachman, Joelle Pineau, Doina Precup, and David Meger. Deep reinforcement learning that matters. In Proceedings of the AAAI conference on artificial intelligence, 2018.
[54] Stephanie CY Chan, Samuel Fishman, John Canny, Anoop Korattikara, and Sergio Guadarrama. Measuring the reliability of reinforcement learning algorithms. arXiv preprint 3:X:v:191..00663. 2019.
[55] Marcin Andrychowicz, Anton Raichuk, Piotr Stanczyk, Manu Orsini, Sertan Girgin, Raphael Marinier, Léonard Hussenot, Matthicu Geist, Olivier Pietquin, Marcin Michalski, Sylvain Gelly, and Olivier Bachem. What matters in on-policy reinforcement learning? large-scale empirical study, 2020.
[56] G. Brockman, V. Cheung, L. Pettersson, Schneider, J. Schulman, Tang, and W. Zaremba. Openai gym. CoRR, abs/1606.01540 2016.
[57] Ian Osband, Yotam Doron, Matteo Hessel, John Aslanides, Eren Sezener, Andre Saraiva, Katrina McKinney, Tor Lattimore, Csaba Szepesvári, Satinder Singh, Benjamin Van Roy, Richard Sutton, David Silver, and Hado van Hasselt. Behaviour suite for reinforcement learning, In International Conference on Learning Representations, 2020.
[58] Erin Catto. Box2D.
2017.
[59] Saran Tunyasuvunakool, Alistair Muldal, Yotam Doron, Siqi Liu, Steven Bohez, Josh Merel, Tom Erez, Timothy Lillicrap, Nicolas Heess, and Yuval Tassa. dm control: Software and tasks for continuous control. Software Impacts, 6: 100022, 2020.
[60] Thomas Hirtz, gym-hybrid. http:://ggttttbooooooooooaaostss hybbii hybrid 2020.
14

---

[61] Karl Cobbe, Christopher Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation to benchmark reinforcement learning. arXiv preprint arxiv: 2019.
[62] Quanyi Li, Zhenghao Peng, Lan Feng. Qihang Zhang, Zhenghai Xue, and Bolci Zhou. Metadrive: Composing diverse driving scenarios for generalizable reinforcement learning. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2022.
[63] Gabriele Cirulli. 2048. https: httpp//s//ggtttoooooo..aaaaaaaaa
2014.
[64] Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 JEEE/SS International Conference on Intelligent Robots and Systems, pages 5026-5033. IEEE, 2012.
[65] Yunhao Tang and Shipra Agrawal. Discretizing continuous action space for on-policy optimiza- tion. In Proceedings of the aaai conference artificial intelligence, 2020.
[66] Peter Sunehag, Guy Lever, Audrunas Gruslys, Wojeiech Marian Czarnecki, Vinicius Zambaldi, Max Jaderberg, Mare Lanctot, Nicolas Sonnerat, Joel Z Leibo, Karl Tuyls, et al. Value- decomposition networks for cooperative multi-agent learning. arXiv preprint arXiv: arXiv:/706.05222 706.05229. 2017.
[67] Christian Schroeder de Witt, Tarun Gupta, Denys Makoviichuk, Viktor Makoviychuk, Philip HS Torr, Mingfei Sun, and Shimon Whiteson. independent learning all you need in the starcraft multi-agent challenge? arXiv preprint arXiv:2011.09533 2020.
[68] Chao Yu, Akash Velu, Eugene Vinitsky, Jiaxuan Gao, Yu Wang, Alexandre Bayen, and Yi Wu. The surprising effectiveness of ppo in cooperative multi-agent games. Advances in Neural Information Processing Systems, 35:2461 142444 2022.
[69] J Terry, Benjamin Black, Nathaniel Grammel, Mario Jayakumar, Ananth Hari, Ryan Sullivan, Luis Santos, Clemens Dieffendahl, Caroline Horsch, Rodrigo Perez-Vicente, al. Pettingzoo: Gym for multi-agent reinforcement learning. Advances in Neural Information Processing Systems,
2021.
[70] Bei Peng, Tabish Rashid, Christian Schroeder de Witt, Pierre-Alexandre Kamienny, Philip Torr, Wendelin Böhmer, and Shimon Whiteson. Facmac: Factored multi-agent centralised policy gradients. Advances Neural Information Processing Systems, 34: 1208-12221, 2021.
[71] Tabish Rashid, Mikayel Samvelyan, Christian Schroeder De Witt, Gregory Farquhar. Jakob Foerster, and Shimon Whiteson. Monotonic value function factorisation for deep multi-agent reinforcement learning. The Journal of Machine Learning Research, 21(1):7234-7284 2020.
[72] Zhao-Heng Yin, Weirui Ye, Qifeng Chen, and Yang Gao. Planning for sample efficient imitation learning. arXiv preprint 1rX.v:2210.09598, 2022.
[73] Sebastian Ruder. An overview of gradient descent optimization algorithms. arXiv preprint arXiv:1609.04747, 2016.
[74] DI engine Contributors. DI-engine: OpenDILab decision intelligence engine. github. http:///ooopeeeaa/a---eeeeeeeeee
https://
2021.
[75] Sherjil Ozair, Yazhe Li, Ali Razavi, loannis Antonoglou, Aaron Van Den Oord, and Oriol Vinyals. Vector quantized models for planning. In International Conference on Machine Learning, pages 8302-8313, PMLR, 2021.
[76] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Minn Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018.
15

---

[77] Peng Wu. Pytorch 2.0: The journey to bringing compiler technologies to the core of pytorch (keynote). In Christophe Dubach, Derek Bruening, and Ben Hardekopf, editors, Proceedings of the 2 ACM/IEEE International Symposium on Code Generation and Optimization, CGO 2023, Montréal, QC. Canada, 25 February 2023- March 2023, page ACM, 2023.
[78] Wouter Kool, Herke Van Hoof, and Max Welling, Stochastic beams and where to find them: The gumbel-top-k trick for sampling sequences without replacement. In International Conference on Machine Learning, pages 3499-3508. PMLR, 2019.
[79] Matthew Hausknecht and Peter Stone. Deep reinforcement learning in parameterized action space. arXiv preprint 1rXXv::511.044//1 2015.
[80] Boyan Li, Hongyao Tang, Yan Zheng, Jianye Hao, Pengyi Li, Zhen Wang, Zhaopeng Meng, and Li Wang. Hyar: Addressing discrete-continuous action reinforcement learning via hybrid action representation. arXiv preprint arXhv:2109.05490, 2021.
[81] Ben Mann, Ryder, M Subbiah, J Kaplan, P Dhariwal, Neelakantan, P Shyam, G Sas- try, A Askell, S Agarwal, et al. Language models are few-shot learners. arXiv preprint arXiv:2005. 2020.
[82] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint wwxhv22307009288. 2023.
[83] Andreas Doerr, Christian Daniel, Martin Schiegg, Nguyen- Tuong Duy, Stefan Schaal, Mare Toussaint, and Trimpe Sebastian. Probabilistic recurrent state-space models. In International conference on machine learning, pages 1280-1289. PMLR, 2018.
[84] Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. arXiv preprint arXiv:2010.02193, 2020.
16

---

A Environment
Environment Attributes
Possible Categories
Environment Instances
Observation Space
Vector Observation Space
Classical Control [56]. Byuite [57
Box2D (58). DMControl [5
Gym-Hybrid 60
Image Observation Space
Board Games, Atari 30 Bsuite,
MiniGrid 44] ProcGen [1] 2048
Structured Observation Space
MetaDrive 62
GoBigger 23
Action Space
Discrete Action Space
Board Games, Atari, Bsuite, MiniGrid, ProcGen, 2048
[63
Continuous Action Space
MuJoCo [44,, Classic Control,
Box2D, DMControl, MetaDrive Gym-Hyyrrid GoBigger
Hybrid Action Space
Reward Space
Sparse Reward
MiniGrid, Board Games, Atari,
ProcGen, Bsuite
Normal Dense Reward
Classical Control, Bsuite, Box2D.
DMControl, Gym-Hybrid, Atari, ProcGen, 2048
Multi-Ohjective Dense Reward
MetaDrive, GoBigger
Transition Function
Stochastic Back-traceable
2048, Atari, Bsuite
Board Games, 2048, MiniGrid
Table I: Overview of various attributes associated with decision-making environments, their respective potential categories crucial for the design of MCTS-style methods, and specific instances of environments that exemplify each category. Please note that each environment can encompass a series of different sub-environments (e.g. Atari typically consists of 57 distinct video games), thus some environments can simultaneously belong to multiple categories. While most categories are self-explanatory, there are two special cases worth noting: Multi- Oittti-vbeeeivvee Dense Reward indicates that there are multiple, distinct rewards be balanced for varying decision targets (c.g. balancing route optimization, speed, and stability in autonomous driving tasks). Back-traceable critical property for AlphaZZeo-liie methods that requires perfect simulator to return to the original state after each simulation. While Board Games can easily store and restore old states, most control tasks and video games are practically impossible to revert their previous states.
In this section, we introduce various types of RL environments integrated in LightZero and their respective characteristics. These environments encompass a wide range, including Board Games, Classical Control, Atari, MuJoCo, the sparse reward example environment (MiniGrid), the structured action space example environment (GoBigger). We categorize these environments based on different attributes in Table followed by detailed description of each environment.
Board Games This types of environment includes TicTacToe, Connect4. Gomoku, Chess, Go, where uniquely marked boards and explicitly defined rules for placement, movement, positioning, and attacking are employed to achieve the game's ultimate objective. These environments feature a variable discrete action space, allowing only one player's piece per board position. In practice, algorithms utilize action mask indicate reasonable actions.
Classic Control In the reinforcement learning domain, classical control environments like Cartpole and Pendulum are built upon physical principles and provide explicit action and state spaces. They are commonly used as henchmark environments to evaluate the initial performance of reinforcement learning algorithms in discrete or continuous action spaces.
Box2D All of these environments in this suite feature toy games centered around physics control (e.g.. LunarLander and BipedalWalker), utilizing Box2D-based physics and PyGame-based rendering.
17

---

Their straightforward simulation mechanics and meaningful visual outputs have made them popular benchmarks for RI. beginners. Moreover, they are suitable for visualizing the tree search process.
Atari This category includes sub-environments like Pong, Qbert, Ms. Pacmann Breakout, UpNDown. and Seaquest. In these environments, agents control game characters and perform tasks based on pixel imput, such as hitting bricks in Breakout. With their high-dimensional visual imput and discrete action space features, Atari environments are widely used to evaluate the capability of reinforcement learning algorithms in handling visual inputs and discrete control problems.
MuJoCo This category includes sub-environments such as Hopper-v3, HalfCheetaah-33, Walker2d-v3, and Humanoid-v3. MuJoCo a precise and efficient physics simulation engine capable of simulating intricate digital mechanical systems, such as bipedal robots and robotic arms. These environments serve testbeds for assessing the performance of reinforcement learning in control tasks like robot balancing or robotic arm object grasping. MuJoCo environments are characterized by continuous state and action spaces, and real-world physical laws, commonly used to evaluate the ability of reinforcement learning algorithms to handle continuous control and high-dimensional input problems.
MiniGrid In addition to the difficulty levels of action and observation spaces, the sparsity or density of reward spaces is another major consideration. For example, MiniGrid series environments, like MiiniGrid-Empty-te
and MiniGriid-FourRoommen FourRooms-vO,, provide simple and scalable sparse reward environments. In these environments, agents navigate in a grid world and complete various tasks. Sparse reward environments pose challenges to the exploration capabilities of reinforcement learning algorithms, as agents receive limited useful feedback most of the time.
GoBigger GoBigger is multi-agent reinforcement learning environment that emphasizes cooperation and competition. It features structured observation spaces where each agent is represented by one or more clone balls. The goal of the agents is to collide and merge with other balls within limited time, thereby increasing their size. The observation space includes all unit information within the agent's local view, and rewards depend on the size difference between consecutive time steps. This environment offers diverse sub-environments for different tasks, such as 12p2, 13p2, 14p3, etc. Here, the number following represents the number of teams, and the number following p indicates the number of agents per team. In our experiments, we set t p 2.
2048 2048 is a numerical game in which players need to merge adjacent blocks with the same number in a 4x4 grid by swiping up, down, left, or right. The randomness of the game lies in the fact that after each swipe, a block with a value of 2 or 4 is randomly generated. Theoretically, there are up to 32 possibilities after each move, which poses challenges for model-based RL algorithms. For algorithms in the MCTS family, solving such problems requires modeling the randomness of the environment and considering various possibilities during the search process.
B Details about Main Experiment
B.1 Benchmark Settings
Environments In Table 1 we provide a categorization of common sequential decision-making environments, each presenting distinct challenges. These environments encompass the majority of the cases previously discussed in Section 3.2 We have undertaken a thorough and impartial evaluation of the algorithms incorporated into LightZero, across these varied environments.
A Algorithms In our main benchmark experiments, we assessed series of algorithm variants, including AlphaZero, MuZero, MuZero w/ SSL (MuZero with Self- Supervised Learning Consistency Loss), EfficientZero, Sampled EfficieníZero, Stochastic MuZero, and Gumbel MuZero.
Note that in this context, MuZero w/ SSL represents the original MuZero algorithm augmented with the self- self-supervised loss proposed in [7]. Ffficiení/ero refers to the MuZero algorithm enhanced with the self-supervised loss and value value_prefix from [17]. We decided not to open the third improvement discussed in the [17) it shown to have limited performance benefits and be highly time- consuming. Consequently, it was omitted from our primary benchmark experiments. Sampled
18

---

EfficientZero is based on the previously described EfficientZero. incorporating sampling-relae related modifications from [45]. This enables the algorithm to simultaneously handle environments with discrete and continuous action spaces.
Metrics In all our experimental evaluations, the primary metric employed to gauge performance was the mean Return calculated over the evaluated episodes for each algorithm. Each algorithm was independently executed five times with unique seeds to ensure the integrity and reliability of our results. In the results graphs, the shaded regions illustrate the standard deviation in different runs, indicating the variability or spread of the runs under different seeds. On the other hand, the solid lines represent the mean, providing the average performance across the different runs.
The horizontal axis, labeled Env Steps (environment steps), represents the number of steps taken in the environment during each run. The vertical axis represents the average Return over the evaluated 20 episodes, serving as measure of the algorithm's 's effectiveness in each run.
B.2 Main Benchmark
In this section, we provide an comprehensive benchmark results for LightZero. This encompasses baseline performances in discrete action spaces detailed in Section B.2.1, and continuous action spaces in Section B.2.2 In addition, Stochastic MuZero, delineated in Sections [B.2.3 and B.2.4 respeetssyll Furthermore, we extend our investigation to present the baseline evaluations in sparse reward scenarios like MiniGrid (Section B.2.5 and in multi-agent environments exemplified by GoBigger (Section B.2.6).
include comparative analyses of Gumbel MuZero and
For thorough examination and in-depth discussion of these benchmark results, we encourage readers refer Section 4.2 where we provide key observations and insights derived from our experiments.
B.2.1 Discrete Decision Benchmark
In Figure 7 we present comparative performance evaluation of four key algorithms Muuzero MuZero with SSL, EfficientZero, and Sampled EfficientZero across a selection of six representative environments from the classic Atari image input domain.
Figure 8 showcases performance comparison of AlphaZero and MuZero on two representative board games Connect4 and Gomoku (board size=6). What stands out is that while AlphaZero's 'o'' sample efficiency significantly surpasses that of MuZero, the latter still manages to produce satisfactory results, even in the absence of simulator.
B.2.2 Continuous Control Benchmark
In the upper of Figure 9
we present the performance of Sampled EfficientZero with different policy modeling approaches on the classical continuous benchmark environment, Pendulum-vl and LumarlanderComtinoossy
In the bottom of Figure 9 present the performance of Sampled EfficientZero with different policy modeling approaches on the conventional continuous benchmark environment, MuJoCo, exampifed by Hopper-v3 and Walker2d-v3.
The dimensions of the action space for the four aforementioned environments are 1, 2, 3, and 6, respectively, with the number of discrete bins set to 11, 7, and 5. Consequently, the final dimensions of the factored categorical distributions (65] stand 11 | 11, 72 = 49, 5 125, and 56 15625, in the same order. Notably, the dimensionality of the discretized action space undergoes an exponential increase with the dimension of the original continuous action space. Under conditions of fixed simulation cost, specifically with the number of sampled actions K set to 20, and the number of simulations per search sim set to 50, we witness a steady decline in the performance of the factored policy representation as the dimensionality increases. Conversely, the Gaussian policy representation maintains a relatively stable performance, demonstrating its robustness in higher-dimensional spaces.
However, even the Gaussian policy representation does not achieve the performance of model- free methods in MuJoCo. We speculate that this is due to the unique reward mechanism of the
19

---

OoprrioFFrameskkee
BreakoutNoF rameskii--a
MrPacmanNoFFamm wwmesipp---sspeee
8000
Mazes
230
Muder stl.
Rotum
8000
Relum
400
Retur
2000
8000
1000
2000
800
CM
0.21
044 Env Steps
EnV Staps
Em Steps
PongNooFramesiiiio
SeequestNooFramee Paameeepit
UpNDownNoFFrmess kip
100
500
Retum
500
Raturn
600
200
290
Mc
130
20200
Enw Steps
Env Steps
Steps
Figure 7: Performance comparison of algorithms integrated in LightZero on six representative image- based Atari environments. The horizontal axis represents Env Steps (environment steps), while the vertical axis indicates the average Return over 20 assessed episodes. In this context, MuZero w/ SSL denotes the original MuZero algorithm augmented with reff-supervise loss. EfficientZero refers to the MuZero algorithm enhanced with sele-vupe loss and value prefix. Sampled EfficientZero builds upon EfficientZero with additional modifications related to sampling.
Connect4
Gomoku
Return
Return
-0.25
ApteaZaro Mudero
Alphnuueuu
-1.00
Mtucero
0M
0.05M
0.1M
0.15M
0.1M
0.2M
0.3M
0.4M
Env Steps
Env Steps
Figure: 8: Performance comparison of AlphaZerr and MulZer Connect4 and Gomoku Shoard_size-de AlphaZero exhibits significantly higher sample efficiency compared to MuZero. This suggests that deploying AlphaZero directly could be advantageous when environment simulator is readily available. However, even in the absence of a simulator, MuZero can still yield comparable results, which underscores its broad applicability.
20

---

Pendulum-v1
LunarLannderContinuuuss
-200
4000
Return
600
Return
-100
-200
-300
-1000
-400
-1200
-500
Elficeniiien
ElfccentZrro (Saussan
0M
0.1M
0.2M
0.3M
0.4M
0.5M
OM
1M
2M
3M
4M
SM
Env Steps
Env Steps
Hopper-v3
Walker2d-v3
2500
EfRceenseem (Cauustion Parkcy
2500
Sempled Efficen//mm (Gauseden Polcc)
(Fackoned
Poley)
2000
Return
Return
1500
1500
1000
500
500
0M
0.2M
0.4M
0.8M
0.8M
1M
0M
0.2M
0.4M
0.6M
0.8M
1M
Env Steps
Env Steps
Performance comparison of Sampled EfficientZero utilizing different policy modeling techniques in continuous action space environments. Bottom: Performance comparison of Sampled EfficientZero using various policy modeling methods in MuJoCo continuous action space environments. Performance of the factored policy representation declines gradually with increasing action space size, while the Gaussian policy representation exhibits more stable performance.
Figure 9: Upper:
MuJoCo environment, where the optimal policy actions are likely extreme values of -1 or 1 in most times. Sampling from Gaussian distribution makes it challenging to obtain these extreme actions, which leads to suboptimal performance. To enhance performance in MuJoCo environments, a straightforward idea for future work is to augment extreme actions during the sampling process.
B.2.3 Gumbel MuZero Benchmark
In Figure TO, present the performance comparison of Gumbel MuZero and MuZero under dif- ferent simulation costs (i.e. the number of simulations in one search) on Gomoku bboard_size=t Lumarlander-v2, TongNoFFraeskipa. and Ms. Pacmannnn Forammnnnnorammiiii-ii. Original Monte Carlo Tree Search methods have to visit almost all the possible nodes to ensure stable optimization. The demands of high simulation counts have become the main con-cconuuming source of MCTS-style methods. In our benchmark, we validate that Gumbel MuZero can achieve notably better performance than MuZero when the number of simulations is limited. Besides, we don't 't tune the balance weight of the completedQ proposed in Gumbel MuZero, may siginicantly influence the final performance in some environments, which maybe the reason why Gumbel MuZero shows lower episode return in the high simulation times case.
B.2.4 Stochastic MuZero Benchmark
In Figure we present a comparative performance analysis between Stochastic MuZero and MuZero, conducted within 2048 environments that exhibit intrinsic randomness. Specifically, our attention focused on the environments typified by standard 2048 settings, where pum__chanesss (potential random tiles being (2, 4)), and a variant environment with increased stochasticity. num chances==5 (potential random tiles are {2, 4, 8, 16, 32)). Stochastic MuZero exhibits distinct performance advan-
21

---

Gomoku
LunarLander
1.00
Suncel MuZore mm)0
Gumbel MuZero p4
200
0.75
Vuieero 2ahm000
Statero jame
100
Return
Return
0.00
-0.25
-100
-0.50
-200
QM
0.2M
0.4M
0.6M
0.8M
OM
0.2M
Q-4M
0.6M
C.BM
Env Steps
Env Steps
PongNoFrraeskiipv4
MspacmanNoFrameeskiia
1800
Gumbel Nizoro (sin-60)
1600
Gurbbel Mitwoo (eime
10
1400
M4leero jainedo)
M4.7em Jamn
Zoro (eim!"))
1200
(amme))
Return
Return
1000
600
400
-20
200
0.7M
0.2M
Env Steps
Env Steps
Figure 10: Performance comparison of Gumbel MuZero and MuZero under varying simulation costs. Gumbel MuZero demonstrates significantly better performance than MuZero when the number of simulations is limited, showcasing its potential for designing low time-cost MCTS agents. For Gomoku (board size=6), we evaluate sim=(20, 10). For LunarLander-v2, we assess sim= 10. 5). For Atari Games, we examine sim=|50, 16. 2).
2048 (num_ (aim__chacee=2)
2048 (num_chaumannnaa chances=5)
40000
Stochostio MuZero
8000
Stochaste Muzeero
Return
30000
Return
6000
20000
4000
10000
2000
0
1M
Env Steps
Env Steps
Figure 11: Performance comparison between Stochastic MuZero and MuZero in 2048 environments with varying levels of chance (num (num__hancee chances=h and 5). In environments characterized by stochastic transition dynamics, Stochastic MuZero marginally surpasses MuZero. However, as the level of stochasticity escalates, the performance of Stochastic MuZero begins to face limitations.
tage over MuZero when the environment's 's transition encapsulates elements of stochasticity. However, as the stochasticity escalates, the performance of Stochastic MuZero begins to face limitations. This implies that intricate randomness remains substantial challenge for MCTS algorithms.
B.2.5 MiniGrid Benchmark
In Figure 12 present performance comparison using the example environment MiniGrid- KeyCoriidorS3R3-v0 from the MiniGrid suite. The figure on the left demonstrates the application of
22

---

MiniGrid-KeyyCorrddoSSS3--
MiniGridkKeyCorriddos33----
0.8
0.8
Return
Retum
0.4
0.4
0.2
0.2
MuZzrro Wi SSL EMfoientZero Samppidd EfficentZier (6-5)
0.0
0.2M
1.2M
0M 0.2M
0.4M 0.6M 0.8M
Env Steps
Env Steps
Figure 12: Left: Performance comparison of various exploration strategies applied to the MiniGrid- KeyCorridor33R3vv environment (Renrr during the collection phase). The Intrinsiccxploraation strategy, which leverages curiosity mechanisms to explore the state space, demonstrates superior sample efficiency. Right: Performance comparison of algorithms implemented in LightZero In environments characterized by high-dimensional vector obser- vations and sparse rewards, the self-supervised learning loss assists model alignment. However, predicting the value_ ppefff can pose challenges and potentially impede learning.
SiniGrrd-KeCoorrdorsst
various exploration strategies on this environment, utilizing MuZero with Self-Supervised Learning (SSL). The right figure shows four algorithms- -Muzero, MuZero with SSL, Efficientiero, and Sampled EfficieníZero it. From the experimental results, it can be observed that only MuZero with SSL achieved near-optimal performance, while the other three algorithms did not show significant progress. This suggests that in high-dimensional vector input environments with sparse rewards, self- supervised learning loss is crucial for aligning model learning. However, predicting the value_ _peffxx may be more challenging than predicting the reward, which could, in turn, hinder the learning process.
B.2.6 Multi Agent Benchmark
Compared single-agent environments, aueti- environments confront more complex challenges, including, but are not limited to, joint state-action spaces, multi-dimensional optimization objectives, non-stationarity and credit assignment issues [66]. To tackle these challenges, our preliminary experiments in multi-agent environments mainly adopted the independent learning paradigm [67]. In this paradigm, we regard other agents as part of the environment, with each agent making decisions independently, and all agents sharing a common policy-value network. During the Monte Carlo Tree Search process, each agent also conducts searches independently.
To investigate the impact of the number of agents algorithm performance, conducted exper- iments in the GoBigger environment [23] using custom T2P2 and T2P3 scenarios (all other environment parameters remained the same except for the number of agents (P)). As demonstrated in Figure 13 in both two scenarios, our algorithm can achieve stable convergence in confrontations with built-in bots, and its sample efficiency has improved about six-fold compared to non-MCTS methods. For comparison, in the 72P2 scenario, the MuZero algorithm requires approximately 400K Env Steps reach a return of 150K, while algorithms such as MAPPO Stepy to achieve the same level of performance according the paper 23 This evidence proves that, aside from increased time expenditure, the performance of the MCTS algorithm is not significantly impacted when the number of agents is not very large. These initial attempts underscore the vast potential of sample-efficient MCTSS-ikk like algorithms in multi-agent environments.
necessitate about 3M Env
In the future, we aim to provide baseline results various algorithms integrated into LightZero on more multi-agent baseline environments (such as PettingZoo [69], MAMuJoCo 70)). Additionally, plan delve deeper into how to leverage the unique characteristics of multi-agent environments and multi-agent reinforcement learning (MARL) techniques such as Centralized Training Decentral-
23

---

GoBigger (T2P2)
GoBigger (T2P3)
200000
MiZero
MuZero
150000
Return
Return
150000
100000
100000
50000
50000
0.2M
En Steps
Env Steps
Figure 13: Performance comparison between MuZero and EfficientZero, both trained in the inde- pendent learning paradigm on the representative multi-agent environment GoBigger in 72P2 and T2P3 scenarios. Both algorithms demonstrate the ability achieve stable convergence when pitted against built-in bots, and their sample efficiency showcases about six-fold improvement over other non-MCTS methods 23
ized Execution (CTDE) paradigm 66 to conduct efficient MCTS searches, with the goal of designing superior cooperative and adversarial agents.
C Details about Two Case Studies
C.1 Exploration Mechanism in MCTS
C.1.1 Details about Motivation
Striking the right balance between exploration and exploitation 26
27 28
is fundamental
challenge in reinforcement learning. For the MCTS algorithms integrated into LightZero, determining the optimal balance is also crucial. However, there has been limited research on the performance of MCTS algorithms in sparse reward environments. Recently, [72] investigated the combination of adversarial imitation learning and MCTS in the DeepMind Control Suite [59]; however, they did not delve into performance in sparse reward environments like MiniGrid
44
Our preliminary experiments revealed that in some simple sparse reward environments in MiniGrid, MuZero struggles to make substantial progress within 1M steps. We observed that while board games also exhibit sparse rewards, each game has a clear win/draw/loss outcome, providing supervisory signal for assessing the quality of the game state. contrast, in other sparse reward environment like MiniGrid, non-zero rewards are only obtained upon reaching the goal, with all other states yielding zero rewards. This makes challenging to collect trajectories with supervisory signals. searching in some intermediate states of MiniGrid may be of limited value due zero rewards, hypothesize that increasing the number of simulations may not significantly improve performance.
In Table 2 we provide a comprehensive comparison of exploration strategies employed in MCTS algorithms as part of the LightZero framework. In this study, our main emphasis is on the general exploration strategies applied to decision-making algorithms. As for the unique exploration param- eters specific MCTS, previous research has predominantly used default values, which are likely near-optimal settings. A detailed examination of these unique parameters falls outside the scope of our current work and will be addressed in future research efforts.
C.1.2 Details about Settings
We performed experiments in the MiniGrid environment, focusing on the MiniGrid-KyyCoridasss vO nd FiniGriidFourRoomm--v scenarios. We set up the two environments with maximum step limit of 300 and used the Adam optimizer with learning rate of 3 x 10-3 The other parameter settings are consistent with those in Table
24

---

Categories
Strategies
Default Settings in LightZero
MCTS Unique Exploration Strategics
Adding Dirichlet noise at the root
The alpha value in the Dirichlet dis- tribution is 0.3. The noise weight is 0.25.
node (also known the decision node).
Balancing the weight of prior prob-
The basc constant c1=1.25, the ini- tialization constant cp=19652.
ability P and MCTS derived Q es-
timates at intermediate nodes using the PUCT formula (31.
General Exploration Strate- gies in RL
Data collection phase:
Decaying temperature coeffi- cient for visit count distribu- tion.
The temperature transitions from 1 to 0.5 and then to 0.25 in a fixed traininng steps. By default, not used.
pesillnn-rreed, exploration strategy.
Learning phase:
Policy entropy regulariza-
Default policy entropy loss weight
is 0.
tion.
Intrinsic exploartion, e.g.- NGU
NGU (28).
Imitation learning tech- N/A
niques, egg., Efficient Imitate
72].
Table 2: Overview of Exploration Strategies MCTS algorithms implemented in LightZero.
Here, we provide detailed explanations the various enhanced exploration strategies used in Figure of Section [.1]
Naive: This is the default exploration strategy, which includes a manual temperature decay mechanism. The initial temperature is 1, and the temperature decays twice during training, to 0.5 and 0.25, respectively. The decay timings are controlled by a parameter. At 50% of this parameter's 's value multiplied by the training steps, the temperature decays to 0.5. At the 75% point, the temperature decays to 0.25. Afterward, the temperature remains fixed 0.25.
NaiveDoubleSimulation: Same configuration as Naive except use double number of simulations.
Fixed Temperatuur:: fixed temperature parameter was utilized throughout the training period, with the specific value set to x=l.
Eeereeyorrooooyearriiii the original loss of MuZero [4] algorithm, with policy entropy loss weights of 0.05 0.005.
An additional policy entropy regularization term added to
EpsGreedy: Inspired by common exploartion settings in value-based RL methods actions were selected uniformly with probability e, and with probability 1 1-6 with the argmax operation.
IntrinsicExploration: An additional fixed target network and predictor network were introduced. The predictor network employed to predict the output of the target network, and the normalized MSE loss was used to obtain the intrinsic reward. For specific details, please refer to Section C.1.3
C.1.3 Details about Intrinsic Exploration
The core idea of intrinsic rewards to encourage the agent to visit novel states as much as possible throughout the learning process or from a life-long perspective. In theory, any long-term novelty estimator can be used generate this incentive. The RND (Random Network Distillation) 22
algorithm demonstrates good performance, easy to implement, and can be parallelized, so use example to generate intrinsic rewards, denoted as r'
25

---

Original MSE loss Specifically, it utilizes two neural networks:(1) Random Network: g O - R a fixed network with randomly initialized parameters that takes the observed observation 2, as input and outputs its encoding (zz). (2) Prediction Network: g O + R², a network that takes the observation 2i as input and outputs the predicted value ((zz; 0) for the observation encoding g(zz). The networks are trained on the data collected by the agent using stochastic gradient descent, updating the parameters 0 by minimizing the mean squared error err(#s) = [lő(x; 8) g(x)]l? The fundamental concept of this approach is to leverage the neural networks' ability to model dataset distributions in manner akin to supervised learning. A larger prediction error for a particular observation suggests that the agent has visited the surrounding observations in the observation space less frequently, indicating a higher degree of novelty for that observation.
Reward normalization To mitigate the impact of significant magnitude variations in the mean squared error across training and different environments, the final intrinsic reward is defined using min-max normalization. The combined reward at time i is then defined as: rz r + Brt, where rf denotes the external reward provided by the environment at time t, r represents the normalized intrinsic reward generated by the exploration module, and , is a positive number serving as the 1 weight factor for the intrinsic reward. In this paper, we set / to 300 to ensure that the sum of intrinsic rewards single episode is less than the maximum original external reward of I, thus preventing intrinsic rewards from dominating the original objective.
Key designs In the specific implementation, critical aspect is the selection of the input zy. In our preliminary experiments, we used the latent state from the MuZero model but observed subpar performance. We hypothesize that this is due to the distribution of latent states changing throughout the training process, causing the target values obtained by the random network to continuously shift and resulting in intrinsic rewards that essentially resemble noise. Ultimately, we chose to use the environment's 's original observation o, as the input observed state and achieved satisfactory results.
It is worth mentioning that combining efficient exploration methods, such as intrinsic rewards, with the unique MCTS exploration strategies detailed in Table represents an intriguing research direction. We leave this integration as a potential area for future investigation.
C.2 Alignment in Environment Model Learning
Alignment In this section, we examine case studies centered on the essential objective functions. Drawing inspiration from the paper (25). the importance of alignment in training the representation, policy and model is emphasized. The policy should only access accurate model states, while the representation must encode task-relevant and predictable information. The consistency loss proposed in [17 serves as a candidate technique: at each unroll step, the latent state generated by the dynamic model should be aligned (similar/consistent) to the latent state directly obtained from the original observation via the representation network.
Setting To investigate the impact of self-supervvsed consistency loss on different types of observation data, we use the MuZero algorithm as our baseline. The Adam optimizer is employed with learning rate of 3 x 10-3, The notation "MuZero w/ SSL" signifies that the consistency loss weight is 2, while "MuZero" means the consistency loss weight is O. The other hyperparameters remain the same as those described in Section G44
Special case in board games As shown in Figure 14 for special image input environments (checker- board with discrete values) like ToeTacToe, the cosine similarity relatively low, around 0.662. This highlights the challenge of achieving consistency between the latent state output from the dynamic model and the latent state derived from the observation via the representation network. Since adjacent observations differ only in one position, adding consistency loss with inappropriate optimization settings might hinder the learning progress of other parts of the algorithm. To verify this, we con- ducted an experiment on TicTacToe using the SGD optimizer and compared the performance with and without the consistency loss, shown in the right panel of Figure [4] We observed that the addition of consistency loss hampers the learning progress in the carly stages, which further supports our conjecture. That is to say, our experiments demonstrate that the consistency loss proposed in
26

---

TicTaCToe
1.0
Env. Name
ToeT
LunarLander
Pong
Return
cos (0M)
0.782
0.913
0.875
(0.4M)
0.678
0.971
0.973
-0.5
cos (0.8M)
0.662
0.974
0.987
-1.0
MuZeroo w SUU MuZoro wes SBL
0M
0.05M
0.1M
0.15M
0.2M
Env Steps
The change of cosine similarity (i.e. negative consistency loss, cos=1.0 indicates that the two vectors have the same direction.) for three input types throughout the MuZero training. In special image input environments such TicTacToe, characterized by checkerboard patterns with discrete values, the cosine similarity between the latent state output from the dynamic model and the latent state derived from the observation via the representation network remains relatively low. This emphasizes the challenge of attaining consistency between these two latent states. Right: Effects of self-supervised consistency loss on board games when using SGD (73] optimizer with momentum. The inclusion of consistency loss impedes the learning progress during the initial stages.
Figure 14: Left:
17 is suitable only for environments with specific attributes. For board games, devising a suitable alignment constraint to ensure consistent model learning remains a topic for future research.
D Explanation for Figure 2
In this section, will introduce the score rules in Figure 2 We categorize the critical capabilities of general decision solvers into the following dimensions: multi-modal observation space, complex action space, inherent stochasticity, reliance on prior knowledge, simulation cost, hard exploration, and data efficiency. We then compare four algorithms including Model-free RL. (e.g. PPO), AlphaZero, MuZero, and LightZero. Please note that this section, the term LightZero refers to the algorithm that embodies the optimal combination of techniques and hyperparameter settings within our framework. We will elaborate on the specific interpretation of each algorithm depicted in Figure 2
Model-free RL (PPO) We use the standard on-policy PPO implementation and DI-engine 77] bench- mark results. This implementation includes optimized hyper- parameetessssss and incorporates various techniques to enhance performance across different environments. The list of typical hyperparameters for PPO is presented in Table
Hyperparameter
Value
Epoch per collect
10
Num of samples per collect
3200
Batch size
320
Discount factor
0.99
GAE lambda
0.95
Recompute advantage
True
Entropy weight
0.001
Dual clip
True
Gradient norm
0.5
PopArt
True
Table 3: Key hyperparameters of model-free PPO methods.
AlphaZero We the default implementation of AlphaZero provided by LightZero, following the settings from the original AlphaZero paper 229
27

---

MuZero We employ the default implementation of MuZero provided by LightZero, following the settings described in the original MuZero paper 14 and the public codebase MuZZero-Gennral [51].
LightZero We have incorporated the relevant MCTS/MuZero algorithm techniques, each tailored for diverse environments with unique challenges. This integration has given rise to specialized algorithm version, dubbed LightZero, within our framework.
The score assigned to each algorithm in the distinct dimensions indicates their performance and applicability. A score of I suggests poor performance and limited applicability. while a higher score implies a broader application scope and better performance. That is to say, a score of 2 means that this algorithm can deal with some parts of envionments problems in this dimension. score of 3 means that it can tackle more than a half of issues while 4 indicates it has already been capable of many challenges but there is still some improved space. The highest score 5 is assigned to those methods that could be state-of-the tht results in this dimension or don't need to care about this problem. We conduct a qualitative analysis to compare the following algorithm versions.
Multi-modal Observation Space AlphaZero customized to model the 2D image-like observation in board games, receiving a score of 2. MuZero extends this capability and able to handle more general 2D image observations including video games. Model-free RL methods and LightZero utilize extra self-supervised learning representation learning techniques, enabling them to work effectively with high-dimensional and complex states. However, challenges still exist in learning abstract decision behaviors within multi-modal representations.
Complex Action Space In terms of action spaces, AlphaZero is limited to discrete actions in board games. MuZero extends it to various discrete action spaces. Model-free RL methods are also able to handle both discrete and continuous actions. However, in more complex hybrid action space, these methods often require special mechanisms, such as the autogressive action prediction used in AlphaStar [3]. Sampled MuZero represents a significant endeavor to implement MCTS in continuous action spaces, employing finite action samples as the primary basis for action selection. However, when the number of samples is insufficient, performance may suffer due to inadequate approximation. Despite this, LightZero, which incorporates this sampling mechanism, achieves a score of 4, effectively showcasing its ability to address these challenges.
Inherent Stochasticity Ths stochasicity of environment dynamics poses significant challenge for RI. algorithms. Conventional MCTS algorithms face planning issues due to inconsistent state transitions. Value- based RL methods, such as DQN, often face difficulty adapting to variable reward signals for identical actions within particular state, leading to suboptimal performance, Actor-Critic algorithms like PPO somewhat mitigate this issue. LightZero adopts insights from (16) and (75). employing an enhanced model that simultaneously learns the distribution of latent chances for various scenarios. These learned probabilities are subsequently incorporated into the tree search nodes. While these adjustments have bolstered MCTS's s ability to handle inherent dynamics randomness and the partially observable state, the challenge remains substantial when addressing real-world decision-making scenarios.
Reliance on Prior Knowledge One key limitation of AlphaZero is its reliance on prior knowledge of environments, such as the perfect simulator game rules and records. MuZero addresses this limitation by designing dynamics model that learns the reward and transition function, significantly relaxing this requirement. For model-free RL methods and LightZero, they are more flexible but still require some prior-oriented operations like observation pre-processing and reward shaping, to transform the original problem into a more standard MDP. Therefore, a score 4 is suitable for them.
Simulation Cost Although MCTS-style methods show great data effciency, it is important to acknowledge that the wall-time of related training programs can be astonishing due to the simulation cost. These methods necessitate the establishment of comprehensive search trees to assure adequate visit counts and precise value estimation LightZero, by integrating core ideas from the Gumbel MuZero paper and utilizing other optimization techniques, effectively minimizes this overhead. However, due to its inherent model-based nature, it still demands longer runtime compared to model-free RL methods.
28

---

Hard Exploration Tree-search- sewsss-aaresheeeeeeee based planning methods enhance exploration capabilities by reducing the search space for the optimal policy. Conversely, traditional model-free RI. methods often struggle in sparse reward environments without specific techniques. AlphaZero gains a slight advantage in this dimension by utilizing perfect simulator rather than models that are incrementally trained over time. LightZero incorporates additional exploration strategies at minimal cost and successfully tackles the hard exploration task efficiently (such as MiniGrid).
Data Efficiency MCTS-sty style methods excel in data efficiency compared to model-free RL methods, offering impressive performance on academic benchmarks like Atari- 100K and DMControl-500K LightZero implements specialized data sampling with staleness and reuse limitation and management mechanisms (e.g. throughput limiter) further improve data efficiency.
E
Efficiency Analysis in LightZero
This section first lists the computational hardware configurations of all the experiments. For nearly all basic experiments and ablations, we aim to validate the runtime efficiency on a small resource budget, thus we deploy each experiment instance on the Kubernetes cluster with computational resource of 1 A100 40G GPU, 24 CPU cores and 100GB RAM. Given these settings, LightZero can train Atari agents for 100K steps Furthermore, we profile the entire training pipeline these tiny instances. Based on these analysis. can identify several bottlenecks in the runtime of the MCTS/MuZero training pipeline and found some practical, albeit imperfect, solutions to these bottlenecks.
hours, and it can conduct 100K steps of self-play on Gomoku 5 hours.
Environment Latency RL envirommnnss usually vary in the execution time of resef and step methods. For AlphaZero-style methods, environment's 's step method are called many times (c.g. more than 50 in one action selection) and lead to the huge time cost. Particularly. AlphaZero deployed on board games may spend more than a half of the whole training time in environments. In practice, the vmap and jär mechanisms provided by JAX 766 and the LRU (Least Recently Used) cache trick to accelerate these methods. Thanks many for-loop and duplicate operations in classic board game implementations, we can significantly reduce the associated overhead and efficiently implement self-play training.
use
Tree Search Although the tree search process doesn't 't contain many complex mathematical operations, it is still non-neglected component in the efficiency optimization. Moreover, the naive vectorized environment scheme shows no obvious gain in tree search methods because each envionments needs a unique search tree, which is not siutable to use the batch processing to speed up. On the other hand, due to the large number of call times, some basic python primitives like getattr and setattr have become a drawback of this module, thus LightZero implements a variety of tree search operations in MCTS/MuZero algorithm variants with cpp and Cython extension for Python, and it offers obvious improvements in various fine- mne-grained code blocks.
Model Inference As MuZero-style methods utilize learned model in place of the perfect simulator, they spend more time on the inference procedure of neural network models. To tackle the efficiency issues associated with more complex network architectures and larger network parameters, LightZero leverages various tools from the ppen-source community, including some features of PyTorch 2.0 777, some tricks about mixed precision training and large batch training in RI.
Parallel MCTS Utilizing a large CPU cluster and enabling multiple parallel data collectors has proven successful in many decision-making tasks [9] g When it comes to MCTS-style methods, due to the difficulty mentioned in tree search part, previous methods
mainly use the collector-parallel scheme, each collector is composed of an environment, a search tree and a RL agent. Naive batch processing inside the collector can lead to significant waiting time between the three components: env step time, model inference time and tree search time (shown in Figure 15 (d)), which severely underutilized computational resources. As suggested in Figure 15)e), e), LightZero launches multiple (k) environments and search trees in a collector, but divides these modules into several groups, the number of group should be determined by three conerete time counts for a specific environment.
14 3
29

---

Modektroe RL rolout
Aphazero--loo
e) MuZero lik rollout
MCTS
MuZero--kka trees)
env slop time
inference tmme
MCTS
Advanced Batched MuZero-WWa treas)
tree search time
root node
child expanded node
Time
Figure 15: Comparison of various data collection pipelines. a), b), c) illustrate simple serial pipelines for different types of algorithms. d) and depict two distinct parallel MCTS schemes, the former will cause the obvious waiting time while the latter can utilize types of computation resources as much as possible, which makes three different parts overlapped with each other.
Specifically, data collector alternate between different groups to execute corresponding steps. For example, even when the first group of environments has completed its steps, the corresponding CPUs are not idle as they can execute environment simulation steps for the second group. Correspondingly, the similar strategies improve the utilization GPU. Although the ideal case illustrated in Figure 15(e) is challenging achieve perfectly, this grouping and buffering mechanism can significantly enhance the parallel capabilities of MCTS in practical data collection.
RL agents
Multi-GPU Training Leveraging the Distributed Data Parallel module from Py Torch, we've imple- mented multi-GPU training capabilities into the LightZero codebase. This functionality has been rigorously tested with the EfficientZero model within the Atari PongNoFrameskip using 1, 2, and 4 GPUs. The results, depicted in Table demonstrate that by training with 4 GPUs, observe approximate five-fold increase in speed, while maintaining similar levels of performance. Consequently, with access to additional GPU computing resources, we anticipate an even greater acceleration of the training process.
environment,
Reanalyze From the viewpoint of system design, data reanalyze is the most special component in MuZero-style methods. SpeedyZero (34] explores the necessity and best practice of separate reanalysis nodes in distributed training. They use different model replicas compute priority for sampling and new targets for training. LightZero optimizes the cost of reanalyze modules with
30

---

Number of GPUs
Approximate Time to 1M Env Steps (mins)
844
2
363
4
152
Table 4: Approximate training times for Efficien/Zero on the PongNo environment with different numbers of GPUs. The results show that training with 4 GPUs improve the speed by about 5 times with similar performance, which in line with our expectations.
simple yet effective selection mechanism: only the data with high training potential (e.g., appropriate state-action novelty and high temporal difference error) and sufficient training stability require a fegh-frreuenyy reanalysis ratio; other data can undergo fewer reanalysis operations to save time.
Communication Transporting model and data is two main communications in the distributed RL training. To ensure the stable convergence of agents, the entire training pipeline cannot collect excessive data or training the network two many times. The former shows few performance gain or even lead to some harmful training batch, while the latter can result in severe over-fitting. According to the actual speed of sending models (from agent learner to data collector and data arranger) and sending collected data, LightZero designs throughput limiter to control the ratio between generating new data and sample training batch. By using a fixed batch size, we can adjust the ratio within reasonable range (such 0.8x-1.2x of original settings). Besides, there is unnecessary dataflow of the RL training pipeline, which is described in Figure 16 In classic RL frameworks, the latest generated data in data collector is located in the inference GPU, however, it has to go through many steps to become serialized bytes data and then it is sent to agent learner by normal communication techniques like socket. To address this problem, we utilize the RDMA (Remote Direct Memory Access) technique to send the training data from inference GPU to training GPU directly. With the aid of this improvement, LightZero can expedite data collection in the P2P PPeerto-Peeet communication mechanism, bypassing series of previously complicated operations.
Collector
action
Atrward
Collected
Numpy
Comprexned
Sweiloed
PSP Communicatin
Communication
GPU Teneor
CPU Thnnor
Numpy Dath
Decompreaned Dota
Deserrtrdd Data
Training Data
gracke
Comparison of various algorithms during the data collection phase. Utilizing Peer-to-Pee (P2P) communication with the Remote Direct Memory Access (RDMA) technique significantly reduces the time required in the data collection process, bypassing the complex operations typically employed in conventional communication methods.
Figure 16:
F Detailed sub-modules description of MuZero
To elucidate the application of specific algorithms into the various sub-modules that constitute LightZero, we delve into the implementation of the MuZero algorithm. This discussion is structured around the sub-modules outlined in Figure
31

---

Data Collector MuZero is representative of online Reinforcement Learning (RL) algorithms, demand- ing substantial interactions between the agent and the environment to amass training data. Within this sub-module, we leverage a vectorized environment manager and a deeply optimized vatch-paralle search tree, facilitated by Cython/C++ extensions. These elements work collaboratively to ensure high-thig thghthhroughput data collection process.
Data Arranger The function of the Data Arranger sub-module is to maximize the utility of different stale data. In the context of MuZero, the first step entails storing complete trajectories or episodes in a prioritized replay buffer to maintain their temporal sequence. Subsequently, we institute a priority recomputer mechanism to periodically determine the sampling priority of the stored data. The calculated priority is proportional Furthermore, to rectify any off-policy bias and enhance the stability of training, data reanalyzer updates the target value stored in data, utilizing the latest network parameters. Recognizing the intricate "producer-consumet" relationship between the Data Collector and Agent Learner, MuZero incorporates throughput limiter to monitor the number of push/pop data operations and regulate the allocation of computational resources.
the likelihood of the data being sampled for training.
Agent Learner The Agent Learner sub- module amalgamates an array of deep learning and rein- forcement learning techniques to train set of neural networks defined in MuZero. Components of this sub-module include the distributional RL module, which models the inherent randomness of environmental reward, and the data parallel and mixed precision training utilities that expedite per-iteration time. These features can be conveniently enabled or disabled via the corresponding configuration fields.
Agent Evaluator Throughout the training process, MuZero necessitates the evaluation of the per- formance of the newly trained network. The Agent Evaluator sub-module encompasses several evaluation strategies such as low-temperature sampling and beam search to enhance results. This sub-module also employs a broad spectrum of metrics and visualizations for agents.
Context Exchanger achieve asynchronous execution across four sub-modules and efficiently scale the entire training pipeline, we employ a Context Exchanger. This component incorporates novel communication elements facilitate the efficient transfer of necessary context information.
By dissecting the implementation of the MuZero algorithm into these aub-modules, we demonstrate the flexibility and modularity of LightZero, and how it can be effectively employed in the design and execution of complex RL algorithms.
G Details about Algorithm Implementation
In this section, we first provide an overview of the subsequent algorithm extensions of MuZero, followed by a comprehensive algorithm overview diagram illustrating the implementation of various integrated algorithms in LightZero (e.g., Figure [9]. Next, we introduce the model's network architecture, including the representation, dynamics, and prediction network. Finally, we present the hyperparameter settings for each algorithm and environment used in the baseline results.
G.1 MuZero's Extensions
In recent years, MuZero has been extended through various algorithmic innovations to enhance its efficiency and stability across different scenarios. These extensions mainly include:
Sampled MuZero 45
This approach introduces a general framework called the sample-based policy iteration, which can be theoretically applied to any types of action spaces. The core idea is to compute improved policies within a subset of the original action space. As the number of sampled actions gradually increases towards the size of the entire action space, the sampled improved policy converges probabilistically to the improved policy over the entire action space.
Gumbel MuZero [8 is extension designed to enhance performance in environments with low simulation costs by leveraging the Gumbel-Tope Too-k trick [78 select actions that guarantee policy
32

---

improvement. This approach introduces new improved policy, x', 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-makee making and improved performance in a variety environments.
Stochastic MuZero [6] This extension enhances MuZero by enabling it to learn and plan with stochastic models. Specifically, it incorporates 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, present detailed summary of various MCTS/MuZZro variants, encapsulated in algorithm overview diagrams. Specifically, Figure 17 provides an insight into the Monte Carlo Tree Search (MCTS), while Figure 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 obser- vations, such as continuous control tasks, the overall model structure is similar, only replacing the con- volutional layers in the original model with corresponding fully connected layers. For specific details, please refer the model directory through OpenDILab http://ggihhh.coo/opeennaaalllll
G.4 Hyperparameters
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 accordance with those in Table
Specifically, we highlight the primary differing hyperparameters for the Sampled EfficienfZero algorithm in the Atari environment in Table g as well for continuous control environments such as MuloCo in Table In addition, we explain the main distinct hyperparameters for the Gumbel MuZero algorithm in the Atari environment in Table 10 alongside the Stochastic Efficient/ero algorithm in both the Atari environment, as shown in Table 11 and the 2048 environment, as presented Table 12 Lastly, the pivotal hyperparameters for the AlpzaZero/MuZero algorithm the Gomoku environment are set out Tables B3 and 14
H Comparison with Other MCTS Frameworks
LightZero a comprehensive algorithm benchmark developed using PyTorch, integrating nine distinct algorithms such AlphaZero and MuZero. supports over twenty different environments, including Atari and Go. On the other hand, the MCTX library 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.
33

---

To offer an intuitive comparison of the differences in integrated algorithms and supported environ- ments between these two repositories, we present the algorithms and environments supported by LightZero and MCTX in the following Table 5 Table respectively.
Please note: "v" denotes that the corresponding item is fully implemented and thoroughly tested. "E" 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 the following two points: Firstly, the list of algorithms and environments supported by MCTX 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 locked "Ee" state. This is because, to 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 complete training pipeline for these algorithms is still in progress.
Algo Env
Classic Control
Box2D
Atrr
MulaCo
Go
Mini
Maze
Connect Four
Gomolar
2048
Go
Bigger
Grid
AlphaZero
Sampled AlphaZero
Gumbel AlphaZero
MuZero
V
MuZero w/ SSL
erficientZero
Gumbel MuZero
Sampled MuZero
Stochastic MuZero
Table 5: Algorithms and Environments supported by LightZero of the writing this paper.
Algo Env
ClassicControl
Box2D
Atari
Maze
ConnectFour
Gomoku
Go
AlphaZerr
MuZero
Gumbel MuZero
Stochastic MuZer
Table 6: Algorithms and Environments supported by MCTX of the writing 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 amplify the algorithm's 's universality through decoupling and modularization, certain specific problems may still necessitate tailored adjustments and opti- mizations. For instance, all current LightZero algorithms do not directly support hybrid action space environments (79). However, through the use of action representation [0 and other techniques, they could potentially be adapted to wider action space.
Secondly, as discussed Section S due to the intrinsic limitations the MCTS algorithm, our method may encounter challenges specific, complex real-world environments. Particularly when handling environments with high-dimensional action spaces strong randomness, methods like Sampled MuZero and Stochastic MuZero could benefit from further optimization and improvement.
34

---

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+RI. algorithms. While these algorithms are comprehensive, they require 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 s interfaces, enriching the documentation, and cultivating vibrant user community that encourages knowledge collaboration.
3
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 contribute new exploration and optimization strategies further boost the performance of MCTS-sttle algorithms.
We stress the importance of a seamless integration between LightZero and future research directions. Two key areas for future exploration stand out:
Integration with Botting-Edger Technologies:
1. Integration with Large Language Models (LLL)) 81
LLM can serve as high-level task planner, breaking down complex decision decision-making problems into 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 viable choice for fine-tuning LLMs.
2. Clever Utilization of Model- 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 incorporate additional techniques within the model representation, loss function, and MCTS search process to enhance the universality and robustness of the LightZero framework.
35

---

Monte Carlo Tree Search
A. Selection
B. Expansion
C. Evaluation
Backprepagation
toll, awonding
au
evalaaii by ransiggu
moluationn
orgman, (Gin
deneely expasied to
collect police fpoknnidll sanken). then compasing the vinne with tancsion
cattoer beloe tha action pra updaned
The adacion lakike spoilo ranrency anil and rashhol
resabbodie
(N(eu-) (.))() =
mruce bee many time acturs
NOTE: The MC TS divided into four stges, and repeated for number of simulations. Once the seaech complete, we select the action with the highest number of visil.
Figure 17: MCTS algorithm overview.
AlphaZero: reinforcement leamning algorithm that chess, shogi, Go through self-play
AlphaZero
Expoesien
Backpropagatien
O
60
teol aelete ceech edge foy nevty myn mynate ando (o (1 0u 25)
The modered the snweek Ps felll
Actias valaan Q apdaid tortukk
bn Myman, (Qts. a) -a-ll Aefioo Theukit suce conncelly
vaheee helee tha sooe Mf an
sus addis
whal
B. Acting (Selt-play)
Training
4.00) = (zt, Ut) + (T,, +c|0
944
where, z is the game reward from the perspective of the player and xysa timestep t.
scarched policy
v, pe predicted value and policy ncural network.
lnli.
entrooy
Figure 18: AlphaZero algorithm overview.
36

---

MuZero: Mastering Atari, Go, Chess and Shogi by Plamninng with a Learned Model
Plaiing
Acting
06-3
140
Inli.
mus
alu...
Training
Sclecton::
each nodk, ofret action o secendingto the UCT srgmax, (((,a)) - 4). ( M0D
(dr Milligionn
At the leat sodo i fiol tinerag rovnd and hidder unse computed by the desamio ef = of af) and ssee in the conciponding
Expansien:
Las
tallo. The golicy and value compund by the predstios fanatios, yf - fiv).
X (fuaet) wit)) ()) (suaa))) 6/0''
mode, sompponding cate of to
alge
(r fren nosty expanded insiaataed
(N(r).)) (.Q()..) 0.P()..)
policy
Barrkupe
hosturappe mettp oager
1l........ we.l..ccpten spte strienee
(s ayg
simulaiin fotlenng 7,4*)
th-i-- the MCTS
N ((*1,) o ((*
In Atrii and cross loss, while board garees, I" is MSE loss and
where is intoppoiiccal say we utiline dee boosttaped estinuee vitee
Dim of
NOTE: is the dynamic retuork (MLP, withon INN), valae resoals. rategarical donnhuaton for nowant and valar in Andi Reanslya
Figure 19: MuZero algorithm overview.
EfficientZero: Self-supervised coonsistency loss, Value prefis, Off-policy correction
Planing
0618
[ull.]]
alnt...
Training ylue
dipes
D. Network and Lass
L(6) rr (su))) + (xraa)t)) (mu prostin +
"resstt... Moeee 2077
oi nomalyoo
Reference: Muscero reget twatt Na t-shaa
self-snpervisel comistercy
NOTE: 4 (alonoviasd gn the figurc) is the resument dysumics network, valuo prefix reward hidden
Figure 20: EfficientZero algorithm overview.
37

---

Sampled MuZero:
Leaming and Planning Complex Action Spaces (High Dimensional Disereto or Continous)
I. ArAiing
[alt.)
C. Iraining
Sclation::
agare
M(Ga)) (a (eN(a)))aatt
We coally use f pp thas A/(pls.)) Al.s).
4,(6)
Σ (Mi to (Aeste vf) KL (m.)s
Sampled-toood Pulky Imprevement.
4004
Foon
polkcy
sasppaal improved pollcy fpelale), which dhecrtt categoiie a denttrrion
right cumeit polacy efaee) expencal Mas palcyy festele) sxprosed palcy repectivlly o'ranpies (melai)) somarges Zelzin)
prodited
disaibation publocy
Evergenad) caloclaned shoeepppe octions.
NOTE: dictibationn Oipically apulicr Bis cepiccal palccy dembbatio
guocerially continsoe)) dumierrion
fis
Figure 21: Sampled MuZero algorithm overview.
Gumbel MuZero: Plamning with few simulations in high dimensional discrete action space
A. Planing
B.Acting
961
weec octter cooding te regental lipil aikan m (AA)
[alt.]]
alu...
Training
The bey base when the nam of cimulatiuts a man of toua actions i MCTS seurch is how to choose which actions to vistis and how many times
we can control the number octions sampled without seplasement.
2. we can use handi algoridmm to efficinnly explore the sct of sampled actions.
1:00
(vess (f) + (susa)) + XLL (fhsa pe) c|pl
Suquentia Halving is sleneify action with the highes a(s) logite(s) <((o)))
mixed value approsimate:
Vate maltow Ya Enctention meslofo)
improved policy distabution of aut paltaaaee of [
[stta]
complet-aQ(a)
(vte)
(L ty. otheeettee
g pulicy distnbutionn improvcal policy dieribuiio, amglntad((ss the completed values, sfs) is the MCTS estinated value the visted actions).
Figure 22: Gumbel MuZero algorithm overview.
38

---

state
Pooling
RasBlock
ResBlock
RecBllck (townsample)
ButshWormb
Convãe
NaLu
BetrhNorm2d
DrwnnSampielee
RepresentaatooNNeaookk
Figure 23: The network architecture the representation network for image input domain in LightZero. This network represents the root observation latent state.
39

---

stat_
reward
Fully Connected Layer
Dynamics
ReLu
Policy
Value
Network
BatchNorm2d
Fully Connected Layer
Fully Connected
Reshape
Conv2d
ReLu
ReLu
BetchWorm2d
BstchNorm2d
ResBlock
XN
PollcyHead
ValueHead
ReLu
Plus
BatchNorm22
Conv2d
PredictionNetwork
Conv2d
(state)
action
state
The network architecture of the dynamics network g for image input domains in LightZero. Given 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 (encompassing both value and policy networks) for image input domains. Given latent state, this network predicts the action probability and value.
Figure 24: Left:
40

---

Hyperparameter
Value
Num of frames stacked
4
Num frames skip
4
Reward clipping
True
Optimizer type
Adam
Learning rate
3 10 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
Max gradient norm
10
Length of game segment
400
Replay buffer size (in transitions)
le6
TD steps
5
Number of unroll steps
5
Use augmentation
True
Discrete action encoding type Normalization type
One Hot
Batch Normalization
Priority exponent coefficient
0.6
Priority correction coefficient
0.4
Dirichlet noise alpha
0.3
Dirichlet noise weight Number of simulations MCTS (sim)
0.25
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
Num of frames skip
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 Max gradient norm
False
0.5
Table 9: Key hyperparameters of Sampled EfficentZero used in continuous control environments, such as, MuJoCo.
41

---

Hyperparameter
Value
Num of max considered actions
ie.. Action Space Size
Gumbel Scale
10
Max visit init
50
Value Scale
0.1
Table 10: Key hyperparameters of Gumbel MuZero Alari environments.
Hyperparameter
Value
Chance space size
i.c. Action Space Size
Afterstate Dynamics Network
Similar with Dynamics Network in Figure 24 Similar with Prediction Network in Figure 24 Two Layer MLP
Afterstate Prediction Network
Chance Encoder
Table 11: Key hyperparameters of Stochastic MuZero on Atari environments.
Hyperparameter
Value
Chance space size
16 num of p-pssple__hhancee chance tile Similar with Dynamics Network in Figure 24
Afterstate Dynamics Network
Afterstate Prediction Network
Similar with Prediction Network in Figure 24
Chance Encoder
Two Layer MLP
Num of frames stacked
Num of frames skip
Reward clipping
False
Discount factor
0.999
Batch size
512
Length of game segment
200
TD steps
10
Use augmentation
False
Number simulations MCTS (sim)
100
Table 12: Key hyperparameters 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 value modeling
False
Table 13: Key hyperparameters of AlphaZero on Gomoku environments.
Hyperparameter
Value
Board size
6
Num of frames stacked
1
Discount factor
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 categorical distribution
10
Table 14: Key hyperparameters of MuZero on Gomoku environments.
42


Reviews:
[{"review_id": "Qiz1wVbGmT", "reviewer": "Reviewer_xvyT", "summary": "", "questions": "", "limitations": "I don't see any negative societal impact.", "rating": 9, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "* This library  will significantly impact on implementation of the future MuZero algorithms for the research community.\n* The algorithm and system design of MCTS methods and design a modular training pipeline which can easily integrate novel insights for better scalability.\n* Algorithm performance was one of the main considerations. For example they made very deepdive into the MCTS implementation.\n* All most popular MCTS Model Based algorithms were supported.\n ", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ""}, {"review_id": "zkj873eQNa", "reviewer": "Reviewer_8SGB", "summary": "", "questions": "", "limitations": "See above.", "rating": 7, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "MCTS is widely used in several RL-based approaches, but no unified study can quantify and evaluate these algorithms on standard benchmarks. This paper addresses these needs and decouples the key components of such algorithmic approaches. Decoupling the essential components also provides insights into how MCTS-based approaches can be potentially used for real-world problems, which has been lacking so far, with existing applications often limited to board games.", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ""}, {"review_id": "cB2GkoLb2a", "reviewer": "Reviewer_dyMN", "summary": "", "questions": "", "limitations": "One key point that is worth clarifying in the next version of this paper is how well do you implementation of these 9 algorithms compare to the implementations in their respective papers? Were you able to make 1-1 comparisons to ensure that all features were implemented and tuned properly? Were you able to achieve the same or higher performance than that quoted in the original work. ", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "MCTS style approaches are clearly one of the most promising set of approaches for sequential planning problems, so creating a unified structure for comparing and implementing the latest approaches is a useful contribution to the decision-making community\n\nThe comparisons done across algorithms are quite interesting, especially results that were only included in the appendix. \n", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ""}, {"review_id": "CTwTkUKCv8", "reviewer": "Reviewer_AvYb", "summary": "", "questions": "", "limitations": "No. No limitations are discussed. No Section 7 does not mention limitations.", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "I believe this is the first open-source library for learning-based MCTS methods. The authors certainly have put into much efforts in designing and implementing the library, taking into consideration of efficiency and effectiveness. The library has covered complete MCTS-RL algorithms and with a modular design, should support modular algorithm improvement.\n\nA good amount of experiments have been conducted to verify the effectiveness of different design factors in the algorithms. While some findings are considered known, it is good to know the others.", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ""}, {"review_id": "XV9px1A0Dt", "reviewer": "Reviewer_UFjT", "summary": "", "questions": "", "limitations": "The limitations are briefly discussed in the end of the main paper. But I would suggest the authors to add a more detail discussion of the limitation in the appendix. ", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "1. Clean implementations of almost all complicated MCTS/MuZero-like algorithms\n2. Thorough analysis of the challenges of the general decision-making solvers\n3. Extensive experiments and useful insights  ", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ""}]

Author rebuttal:

Meta-review: The paper describes a new open source library to design Monte Carlo Tree Search (MCTS)algorithms.  The library is modular and it includes many existing and sophisticated MCTS algorithms.  This is important contribution to the community that will help further advances for sequential decision making based on online planning.
Decision: Accept (Spotlight)