SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yazhe Niu $^{1,3}$   
Yuan Pu $^{2}$   
Zhenjie Yang $^{1}$   
Xueyan Li $^{2}$   
Tong Zhou $^{1}$   
Jiyuan Ren $^{2}$   
Shuai Hu $^{1}$   
Hongsheng Li $^{3,4}$   
Yu Liu $^{1,2}$
$^{1}$ SenseTime Group LTD  
 $^{2}$ Shanghai Artificial Intelligence Laboratory  
 $^{3}$ The Chinese University of Hong Kong  
 $^{4}$ Centre for Perceptual and Interactive Intelligence
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. Specifically, 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.
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 [1, 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 [4], PPO [5], R2D2 [6], SAC [7]) and system architectures such as IMPALA [8] and others [9, 10, 11]. Designing a 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) is a powerful approach that utilizes a search tree with simulation and backpropagation mechanisms to train agents with a small data budget [12]. To model high-dimensional observation spaces and complex policy behaviour, AlphaGo [13] 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 a 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
to AlphaGo have attempted to extend its capabilities in various directions. MuZero [14] relaxes the requirements for prior knowledge of environments by training a set of neural networks to reconstruct reward, value and policy. Sampled MuZero [15] successfully applies MCTS to various complex action space with a novel planning mechanism based on sampled actions. [16, 17, 18] 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.
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, we revisit the development history of MCTS methods [19] and the diverse criterions of newly proposed RL environments [20, 21, 22]. As shown in Figure 2, we outline the six most challenging dimensions in developing LightZero as a general method, including multi-modal and 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. 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 of MCTS-style methods into four sub-modules: data collector, data arranger, agent learner, and agent evaluator. LightZero's decoupled architecture empowers developers to focus intensively on 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. In particular, the misalignment problem [25] of state representation learning and dynamics learning can result in the problematic optimization for MuZero, thus a 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 a general solver for various decision problems. Experiments on different types of environments demonstrate LightZero's rich application ranges and data efficiency regimes with few hyper-parameter 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 a general sequential decision solver, which can be trained and deployed across diverse domains.
Reinforcement Learning models a decision-making problem as a Markov Decision Process (MDP)  $\mathcal{M} = (\mathcal{S},\mathcal{A},\mathcal{P},\mathcal{R},\gamma ,\rho_0)$ , where  $\mathcal{S}$  and  $\mathcal{A}$  denote the state space and action space, respectively. The transition function  $\mathcal{P}$  maps  $S\times \mathcal{A}$  to  $S$ , while the expected reward function  $\mathcal{R}$  maps  $S\times \mathcal{A}$  to  $\mathbb{R}$ . The discount factor  $\gamma \in [0,1)$  determines the importance of future rewards, and  $\rho_0$  represents the initial state distribution. The goal of RL is to learn a policy  $\pi :\mathcal{S}\to \mathcal{A}$  that maximizes the expected discounted return over the trajectory distribution  $J(\pi) = \mathbb{E}_{\pi ,\rho_0,\mathcal{P},\mathcal{R}}[\sum_{t = 0}^{\infty}\gamma^{t}r_{t}]$ .
LightZero
Figure 1: 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 a 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 as chess, shogi, and Go. This approach replaces the handcrafted features and heuristic priors commonly used in traditional intelligent programs. Specifically, AlphaZero employs a deep neural network parameterized by  $\theta$ , represented as  $(\mathbf{p}, v) = f_{\theta}(s)$ . Given a board position  $s$ , the network produces a action probability  $p_a = P_r(a|s)$  for each action  $a$  and a scalar value  $v$  to predict the expected return  $z$ , i.e.  $v \rightarrow z$ .
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^0 = h_\theta (o_1,\dots ,o_t)$ . This network represents the root node (at time  $t$ ) as a latent state, obtained by processing past observations  $o_1,\ldots ,o_t$ , ② Dynamics Network:  $r^k,s^k = g_\theta (s^{k - 1},a^k)$ . This network simulates the dynamics of the environment. Given a state and selected action, it outputs the transitioned next state and corresponding reward. ③ Prediction Network:  $\mathbf{p}^k,v^k = f_\theta (s^k)$ . 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  $s_0$  is generated from the original board state through the representation network, each edge stores the following information:  $N(s,a),P(s,a),Q(s,a),R(s,a),S(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 to the Upper Confidence Bound (UCB) [31] formula:
$$
a ^ {*} = \arg \max  _ {a} Q (s, a) + P (s, a) \frac {\sqrt {\sum_ {b} N (s , b)}}{1 + N (s , a)} [ c _ {1} + \log (\frac {\sum_ {b} N (s , b) + c _ {2} + 1}{c _ {2}}) ]
$$
where,  $N$  represents the visit count,  $Q$  is the estimated average value, and  $P$  is the policy's prior probability.  $c_{1}$  and  $c_{2}$  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  $s^l$  is generated, and its associated predicted reward  $r^l$  is determined. Utilizing the prediction function, we obtain the predicted values  $p^l$  and  $v^l$ . Subsequently, this node is incorporated into the search tree.
- **Backup:** The estimated cumulative reward at step  $k$  is calculated based on  $v^l$ , denoted as:  $G^k = \sum_{\tau=0}^{l-1-k} \gamma^\tau r_{k+1+\tau} + \gamma^{l-k} v^l$ . 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  $s_0$ . These visit counts are normalized to obtain the improved policy:
$$
\mathcal {I} _ {\pi} (a | s) = N (s, a) ^ {1 / T} / \sum_ {b} N (s, b) ^ {1 / T}
$$
Figure 2: Radar chart comparison of MCTS-style methods and model-free RL (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 on prior knowledge, simulation cost, hard exploration and data efficiency. Each curve in the graph represents the score of an algorithm across these six categories. A score of 1 indicates that the algorithm performs poorly in this dimension and is only applicable to limited scenarios, while a higher score means a large application scope and better performance. In particular, model-free RL methods need no simulation and have little dependence on priors, so it achieves high score in corresponding dimensions. Please note that within this context, the term LightZero refers to the special algorithm that embodies the optimal combination of techniques and hyperparameter settings within our framework. Details about qualitative score rules can be found in Appendix D.
where  $T$  is the temperature coefficient controlling the exploration degree. Finally, an action is sampled from this distribution for interaction with the environment or self-play. During the learning phase, MuZero perform end-to-end training with the following loss function, where  $l^{\mathrm{p}}$ ,  $l^{\mathrm{v}}$  and  $l^{\mathrm{r}}$  are loss functions for policy, value and reward respectively, and the final term is weight decay.
$$
l _ {t} (\theta) = \sum_ {k = 0} ^ {K} l ^ {\mathrm {p}} \left(\pi_ {t + k}, p _ {t} ^ {k}\right) + l ^ {\mathrm {v}} \left(z _ {t + k}, v _ {t} ^ {k}\right) + l ^ {\mathrm {r}} \left(u _ {t + k}, r _ {t} ^ {k}\right) + c | | \theta | | ^ {2}
$$
3 LightZero
In this section, we will first introduce the overview of LightZero, followed by a comprehensive analysis of challenges in various decision environments. Additionally, we propose a 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.
As is shown in Figure 1, LightZero is 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 [29], establishing a standardized interface for training and deployment across diverse decision environments. Unlike the original versions of these derived algorithms, which focused on specific avenues of improvement, LightZero provides a 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 to Evaluate A General MCTS Algorithm: 6 Environment Challenges
The algorithm extensions integrated in LightZero have greatly relaxed the constraints and broadened the applicability of MCTS-style methods. In the following part, we hope to delve deeper into the key
Figure 3: Four core sub-modules of the training pipeline in 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 in Figure 2. Concretely, The intentions and goals of six types of environmental capabilities are: 1) Multi-modal observation spaces pose a challenge for agents as they must be able to extract different representation modalities (e.g., low-dimensional vectors, visual images, and complex relationships) while effectively fusing distinct 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 a perfect simulator and specific rules of the environment. In contrast, MuZero and its derived methods address this limitation by learning an environment model to substitute the simulator and related priors. 4) Inherent stochasticity presents a fundamental challenge in tree-search-based planning methods. The uncertainty of environment dynamics and partially observable state spaces both can lead to misalignment of planning trajectories, resulting in a 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 if the algorithm fails to visit all the necessary actions during the simulation process. 6) Hard exploration represents a crucial challenge that is often overlooked. While search trees can enhance efficiency by reducing the scope of exploration, MCTS-style methods are susceptible to difficulties in environments with numerous non-terminating cases, such as mazes.
3.3 How to Simplify A General MCTS Algorithm: Decouple Pipeline into 4 Sub-Modules
The impressive performance of MCTS-style methods is often accompanied by a 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 [5], MCTS-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 or neural networks, increasing the complexity across various aspects of the overall system, including code, distributed training, and communication topology.
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 LightZero 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 off-policy 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 a 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 [37, 38], distributional predicton [39] and normalization [40, 41]. These techniques contribute to the policy improvement and further enhance the overall performance of the agent. Agent evaluator periodically provides the diverse evaluation metrics [42] to monitor the training procedure and assess policy behaviour. It also integrates some inference-time tricks like beam search [43] to enhance test performance. We provide a detailed analysis of how these sub-modules are implemented in specific algorithms in Appendix F. Built upon these abstractions, LightZero serves as a valuable toolkit, enabling researchers and engineers to develop enhanced algorithms and optimize systems effectively. For example, the exploration strategies and ensuring the alignment of a learned model in MCTS is crucial, and this will be discussed in the subsequent sub-section. In 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 E.
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 addresses these challenges through various improvements, resulting in superior performance compared to individual algorithm variants across different domains (Section 4 and 5).
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 [44]. 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, it 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 efficiently within MuZero's learned models. Further details can be found in Section 5.1.
Alignment in Environment Model Learning MuZero employs a representation network to generate latent states and a dynamics network to predict next latent states. However, there is no explicit supervision 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 problem that can not be ignored. Further exploration of misalignments across different environments are discussed in Section 5.2.
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, we conduct an in-depth experimental analysis in Section 5.
Figure 4: Comparisons of mean episode return for algorithm variants in LightZero across diverse environments: Atari with discrete action and partial-observable state (Qbert, 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 as a general decision solver, we conduct an extensive comparisons across a diverse range of RL environments. The algorithm variants list contains AlphaZero [29], MuZero [14], EfficientZero [17], Sampled MuZero [15], Stochastic MuZero [16], Gumbel 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 to Appendix B.
4.2 Key Observations and Insights
Building on the unified design of LightZero and the benchmark results, we have derived some key insights about the strengths and weaknesses of each algorithm, providing a comprehensive understanding of these algorithms' performance and potential applications.