ReZero: Boosting MCTS-based Algorithms by Backward-view and Entire-buffer Reanalyze

TL;DR

The paper tackles the wall-clock time bottleneck in MCTS-based RL (notably MuZero) caused by reanalysis. It introduces ReZero, which combines backward-view reanalyze (to prune subtrees using root-child values) with periodic entire-buffer reanalyze (to align targets with the latest models and leverage large batches). Theoretical analysis casts the search process as a non-stationary bandit and proves convergence properties (Theorem 1), while empirical results show substantial speedups across Atari, DMControl, and board games without sacrificing sample efficiency; code is released within LightZero. This framework enables faster, more scalable MCTS-based learning by decoupling reanalyze from training and exploiting batch-level parallelism, broadening practical applicability in real-time decision-making tasks.

Abstract

Monte Carlo Tree Search (MCTS)-based algorithms, such as MuZero and its derivatives, have achieved widespread success in various decision-making domains. These algorithms employ the reanalyze process to enhance sample efficiency from stale data, albeit at the expense of significant wall-clock time consumption. To address this issue, we propose a general approach named ReZero to boost tree search operations for MCTS-based algorithms. Specifically, drawing inspiration from the one-armed bandit model, we reanalyze training samples through a backward-view reuse technique which uses the value estimation of a certain child node to save the corresponding sub-tree search time. To further adapt to this design, we periodically reanalyze the entire buffer instead of frequently reanalyzing the mini-batch. The synergy of these two designs can significantly reduce the search cost and meanwhile guarantee or even improve performance, simplifying both data collecting and reanalyzing. Experiments conducted on Atari environments, DMControl suites and board games demonstrate that ReZero substantially improves training speed while maintaining high sample efficiency. The code is available as part of the LightZero MCTS benchmark at https://github.com/opendilab/LightZero.

Figures (13)

• Figure 1: The connection between MCTS and bandits. Left shows tree search in a bandit-view. When the action $A$ is selected, a return $r_A$ will be returned, where $r_A=\sum_{t=1}^3 \gamma^{t-1}r_{t}^{a} + \gamma^3 v_a$. For the root node, the traversal, evaluation and back-propagation occurring in the sub-tree can be approximated as sampling from an non-stationary distribution. Thus it can be seen as a non-stationary bandit. Right shows the one-armed bandit case. Once the true value $\mu_A$ is known, we can evaluate arm A using $\mu_A$, thereby eliminating the need to rely on subsequent tree search processes.
• Figure 2: An illustration about the backward-view reanalyze in node and batch view. We sample $n+1$ trajectories of length $k+1$ to form a batch and conduct the search in the reverse direction of trajectories. From the node view, we would first search $S_{0}^1$ and then pass root value $m_{0}^1$ to $S_{0}^0$ to evaluate the value of a child node. $T_{0}^1$ and $T_{0}^0$ are the corresponding search trees. From the batch view, we would group all $S^1$s into a sub batch to search together and pass the root values to the $S^0$s.
• Figure 2: Comparisons about the detailed time cost indicators between MuZero and ReZero-M inside the tree search.
• Figure 3: Execution workflow and runtime cycle graph about MuZero and ReZero in both single and multiple worker cases. The number inside the modules represent the number of iterations, and the number under the modules represent the time required for module execution. The model is updated $n$ iterations between two collections. MuZero reanalyzes the mini-batch before each model update. ReZero reanalyzes the entire buffer after certain iterations($\frac{n}{3}$ for example), which not only reduces the total number of MCTS calls, but also takes advantage of the processing speed of large batches.
• Figure 4: Acceleration effect on the toy case. Left is a simple maze environment where the agent starts at point A and receives a reward of size $1$ upon reaching the end point G. Middle shows the search time corresponding to each position when set as the root node. Meanwhile, the root node values obtained during the search are preserved. Right shows the corresponding search time when these root node values are used to assist the search. The comparison shows that the search duration is generally reduced. For specific experimental settings and code, please refer to the Appendix \ref{['appendix:toycase']}.

Theorems & Definitions (2)

• proof
• proof