SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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  $t$  is then defined as:  $r_t = r_t^e + \beta r_t^i$ , where  $r_t^e$  denotes the external reward provided by the environment at time  $t$ ,  $r_t^i$  represents the normalized intrinsic reward generated by the exploration module, and  $\beta$  is a positive number serving as the weight factor for the intrinsic reward. In this paper, we set  $\beta$  to  $\frac{1}{300}$  to ensure that the sum of intrinsic rewards in a single episode is less than the maximum original external reward of 1, thus preventing intrinsic rewards from dominating the original objective.
Key designs In the specific implementation, a critical aspect is the selection of the input  $x_{t}$ . 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 original observation  $o_{t}$  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 2, 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-supervised consistency loss on different types of observation data, we use the MuZero algorithm as our baseline. The Adam optimizer is employed with a learning rate of  $3 \times 10^{-3}$ . The notation "MuZero w/ SSL" signifies that the consistency loss weight is 2, while "MuZero" means the consistency loss weight is 0. The other hyperparameters remain the same as those described in Section G.4.
Special case in board games As shown in Figure 14, for special image input environments (checkerboard with discrete values) like TicTacToe, 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 conducted an experiment on TicTacToe using the SGD optimizer and compared the performance with and without the consistency loss, as shown in the right panel of Figure 14. We observed that the addition of consistency loss hampers the learning progress in the early stages, which further supports our conjecture. That is to say, our experiments demonstrate that the consistency loss proposed in
Env. Name	TicTacToe	LunarLander	Pong
cos (0M)	0.782	0.913	0.875
cos (0.4M)	0.678	0.971	0.973
cos (0.8M)	0.662	0.974	0.987
Figure 14: Left: 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 as 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.
[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, we 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 in 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 [74] benchmark results. This implementation includes optimized hyper-parameters and incorporates various techniques to enhance performance across different environments. The list of typical hyperparameters for PPO is presented in Table 3.
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 use the default implementation of AlphaZero provided by LightZero, following the settings from the original AlphaZero paper [29].
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 MuZero-General [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 a 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 1 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 environments or problems in this dimension. A 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-art 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 is customized to model the 2D image-like observation in board games, receiving a score of 2. MuZero extends this capability and is 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 autoregressive 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 This stochasticity of environment dynamics poses a significant challenge for RL 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 a 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 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 or game rules and records. MuZero addresses this limitation by designing a 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 of 4 is suitable for them.
Simulation Cost Although MCTS-style methods show great data efficiency, 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 a longer runtime compared to model-free RL methods.
Hard Exploration Tree-search-based planning methods enhance exploration capabilities by reducing the search space for the optimal policy. Conversely, traditional model-free RL methods often struggle in sparse reward environments without specific techniques. AlphaZero gains a slight advantage in this dimension by utilizing a 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-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) to 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 in 4 hours, and it can conduct 100K steps of self-play on Gomoku in 5 hours. Furthermore, we profile the entire training pipeline of these tiny instances. Based on these analysis, we can identify several bottlenecks in the runtime of the MCTS/MuZero training pipeline and found some practical, albeit imperfect, solutions to these bottlenecks.
Environment Latency RL environments usually vary in the execution time of reset and step methods. For AlphaZero-style methods, environment's step method are called many times (e.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, we use the vmap and JIT mechanisms provided by JAX [76] and the LRU (Least Recently Used) cache trick to accelerate these methods. Thanks to many for-loop and duplicate operations in classic board game implementations, we can significantly reduce the associated overhead and efficiently implement self-play training.
Tree Search Although the tree search process doesn't contain many complex mathematical operations, it is still a non-neglected component in the efficiency optimization. Moreover, the naive vectorized environment scheme shows no obvious gain in tree search methods because each environments needs a unique search tree, which is not suitable 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-grained code blocks.
Model Inference As MuZero-style methods utilize a 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 open-source community, including some features of PyTorch 2.0 [77], some tricks about mixed precision training and large batch training in RL.
Parallel MCTS Utilizing a large CPU cluster and enabling multiple parallel data collectors has proven successful in many decision-making tasks [9, 8]. When it comes to MCTS-style methods, due to the difficulty mentioned in tree search part, previous methods [14, 34] 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), LightZero launches multiple (k) environments and search trees in a collector, but divides these modules into several groups, the number of groups should be determined by three concrete time counts for a specific environment.
Figure 15: Comparison of various data collection pipelines. a), b), c) illustrate simple serial pipelines for different types of algorithms. d) and e) depict two distinct parallel MCTS schemes, the former will cause the obvious waiting time while the latter can utilize all 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, RL agents can use the similar strategies to improve the utilization on GPU. Although the ideal case illustrated in Figure 15(e) is challenging to achieve perfectly, this grouping and buffering mechanism can significantly enhance the parallel capabilities of MCTS in practical data collection.
Multi-GPU Training Leveraging the Distributed Data Parallel module from PyTorch, we've implemented multi-GPU training capabilities into the LightZero codebase. This functionality has been rigorously tested with the EfficientZero model within the Atari PongNoFrameskip-v4 environment, using 1, 2, and 4 GPUs. The results, depicted in Table 4, demonstrate that by training with 4 GPUs, we observe an 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.
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 to compute priority for sampling and new targets for training. LightZero optimizes the cost of reanalyze modules with a
Number of GPUs	Approximate Time to 1M Env Steps (mins)
1	844
2	363
4	152
Table 4: Approximate training times for EfficientZero on the PongNoFrameskip-v4 environment with different numbers of GPUs. The results show that training with 4 GPUs can improve the speed by about 5 times with similar performance, which is 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 high-frequency 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 a 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 a reasonable range (such as  $0.8\mathrm{x} - 1.2\mathrm{x}$  of original settings). Besides, there is a 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 (Peer-to-Peer) communication mechanism, bypassing a series of previously complicated operations.
Figure 16: Comparison of various algorithms during the data collection phase. Utilizing Peer-to-Peer (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.
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 as outlined in Figure 3:
Data Collector MuZero is representative of online Reinforcement Learning (RL) algorithms, demanding 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 batch-parallel search tree, facilitated by Cython/C++ extensions. These elements work collaboratively to ensure a high-throughput 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 to the likelihood of the data being sampled for training. Furthermore, to rectify any off-policy bias and enhance the stability of training, a data reanalyzer updates the target value stored in data, utilizing the latest network parameters. Recognizing the intricate "producer-consumer" relationship between the Data Collector and Agent Learner, MuZero incorporates a throughput limiter to monitor the number of push/pop data operations and regulate the allocation of computational resources.
Agent Learner The Agent Learner sub-module amalgamates an array of deep learning and reinforcement learning techniques to train a set of neural networks as 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 performance 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 To achieve asynchronous execution across four sub-modules and to efficiently scale the entire training pipeline, we employ a Context Exchanger. This component incorporates novel communication elements to facilitate the efficient transfer of necessary context information.
By dissecting the implementation of the MuZero algorithm into these sub-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 19). 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.