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

As shown in the right diagram of Figure 2, we sample $n + 1$ trajectories of length $k + 1$, forming a batch and performing a search in the reverse direction of the trajectories. Therefore, $S^0$s refers to the first-step state of all sampled trajectories, which is also the state contained in the first row of the right diagram in Figure 2.

Dear Reviewer, thank you for your feedback! In the paper, we used MuZero, EfficientZero, and Sampled MuZero as examples to verify the time efficiency improvements of ReZero. The results in Figure 5 over two continuous DMC control tasks are from Sampled MuZero, but due to our oversight, the legend was not correctly labeled. This will be corrected in the revised version of the paper. Additionally, the improvements of ReZero relative to EfficientZero are presented in Appendix D. These experimental results demonstrate the potential of ReZero as a general method for improving the time efficiency of MCTS+RL algorithms.

We are very grateful for your valuable feedback and hope this addresses your concerns. If you have any further suggestions, we would be very grateful for your feedback. Thank you very much.

If one were to opt for high-frequency reanalyze of small batch data, this would introduce additional hyperparameters that require careful tuning: the size of the batch to be reanalyzed, the selection criteria for this batch, and the frequency of reanalyze. This not only increases the number of hyperparameters but also couples the reanalyze process with the sampling strategy for training, such as priority sampling. In contrast, our periodic reanalyze requires control over only a single hyperparameter: the frequency.

Moreover, we can draw parallels to similar ideas in other RL studies. Reanalyze fundamentally involves using the latest network to obtain better training targets. This concept is analogous to the recompute Advantage operation in modern PPO implementations, where the entire collected data is processed together. Similarly, the target network in DQN is updated as a whole after a certain number of iterations, with the latest weights replacing the old ones. There is rarely a need to design specialized algorithms for small batch data, likely due to the high variance associated with small batches in online RL, making it challenging to develop stable algorithmic techniques. We believe this explanation clarifies the rationale behind our approach and addresses the reviewer's concerns regarding the complexity and stability of high-frequency reanalyze of small batch data.

Regarding Figure 6, due to space constraints, our explanation in the paper focuses primarily on the core conclusion: the appropriate reanalyze frequency can save time overhead without causing any noticeable performance loss. This is evident from our ReZero's performance being close to that of the reanalyze ratio=1 scenario. However, we acknowledge that the detailed analysis and explanation of the experimental groups were indeed brief. We have addressed this issue and provided a more comprehensive explanation here (with relevant content also added to the paper).

Specifically, as shown in the right panel of Figure 6, with the x-axis representing Env Steps, MuZero reanalyze ratio=1 exhibits a significant performance improvement over reanalyze ratio=0 (comparing the purple line to the blue line). For our ReZero method, increasing the frequency from 0 to 1/3 to 1 progressively enhances performance (comparing the blue line to the orange line to the green line), with the frequency of 1 nearly matching the performance upper bound represented by the purple line. Further increasing the frequency from 1 to 2 (comparing the green line to the red line) shows no significant performance improvement, indicating that our algorithm's buffer reanalyze frequency needs only to reach a moderate level to be effective. A frequency of 1 is thus a recommended value, eliminating the need for additional hyperparameter tuning. On the other hand, the left panel of Figure 6 validates the effect when the x-axis represents wall-clock time. A suitable frequency (such as 1) significantly improves episode return within the same time frame compared to the purple, blue, and orange lines. We hope this detailed explanation clarifies the experimental results and addresses the reviewer's concerns.

Regarding the relationship between our theoretical analyses and the practical algorithm, we have already detailed this connection in Section 4 of our original paper. Here, we aim to elucidate this relationship in a more accessible manner.

Our backward-view reanalyze is grounded in a straightforward insight: if we could know the true state-value (e.g., expected long-term return) of a child node, we could potentially bypass the need for further search, thereby conserving computational resources and concentrate the search budget into more valuable part. In Section 4.1, we first provide a comprehensive description of the entire algorithm, aiming to present a complete overview of the practical implementation. However, as stated at the outset of Section 4.2, our proposed algorithm introduces distinctions from the classic MCTS process employed in AlphaZero. Consequently, it is not reliable to directly apply and scale our algorithm to concrete RL problems without further validation.

To address this, we have conducted theoretical analyses and proofs to substantiate the reliability and efficacy of our backward-view reanalyze. Specifically, we have demonstrated that, as the total number of visits increases, our backward-view design progressively concentrates the visit distribution on the optimal arm of a non-stationary bandit (Theorem 1). This conclusion is further extended to the settings of AlphaZero in a similar manner (Appendix A). These theoretical analyses alleviate concerns that our method might be constrained by local optima or low-quality data, providing evidence for algorithm convergence and potential performance enhancements.

It is important to note that the theoretical analyses and algorithm design were not conducted in a strictly sequential order. Instead, they were developed in a complementary and iterative process, with each influencing the other. However, for the sake of clarity and practical application, we chose to first outline the algorithm's workflow in the paper and subsequently analyze its theoretical reliability. If the above explanation satisfactorily addresses the reviewer's concerns, we are prepared to revise our paper based on these newly added explanations to further enhance the presentation of this section.

As we mentioned in Section 2.2 of original paper, we have compared our approach with several existing MCTS acceleration techniques, such as the state abstraction method in PTSAZero [1] and the simple reuse technique in KataGo [2]. It is important to clarify that we do not claim our backward-view reanalyze strategy is superior to these methods. Instead, we propose a novel perspective on information reuse from a backward-view standpoint.

Regarding the two papers mentioned by the reviewer, the former aims to formalises the correlation between states and actions, and empirically investigate combination operators for MCTS enhancements for the correlate states in analogous positions within different episodes. This work is quite similar to our comparad work PTSAZero, our backward-view analyze can build on the top of their correlate states or state abstractions. The latter "Learning policies from self-play with policy gradients and MCTS value estimates" mainly derives a policy gradient expression with MCTS value estimates rather than MCTS visit counts. The reuse technique they use is just a naive version to reuse search trees from previous turns (similar to KatoGo). Therefore, we believe that our approach complements rather than competes with these methods, offering a complementary angle that can enhance the overall efficiency of MCTS.

[1] Fu, Yangqing, et al. "Accelerating monte carlo tree search with probability tree state abstraction." Advances in Neural Information Processing Systems 36 (2023): 61008-61019.

[2] Wu, David J. "Accelerating self-play learning in go." arXiv preprint arXiv:1902.10565 (2019).

Thank you for your efforts and nice comments on our work! In addition to MuZero, our baseline also contains EffeicientZero (in the appendix) and Sampled MuZero (not explicitly stated in the legend due to our negligence. In fact, the algorithm compared in the two DMC control tasks is not MuZero, but Sampled MuZero, and we will fix it in new version). These comparisons are intended to demonstrate the versatility of ReZero in combination with different algorithms, like we mentioned in line 74: ”Therefore, our algorithm design is universal and can be easily applied to the MCTS-based algorithm family”.

We understand your concerns. As you pointed out, We have mentioned some other acceleration methods in Section 2.2. We mentioned PTSAZero and SpeedyZero, both of which are orthogonal to ReZero in principle, so they can be seamlessly integrated with ReZero. In addition, we have also mentioned the difference between KataGo and our approach in Section 2.2. It is important to clarify that we do not claim that our backward-view reanalysis strategy is superior to these methods. Instead, we propose a novel perspective on information reuse from a backward-view approach. Therefore, at this stage, we did not directly compare the performance with other acceleration algorithms.

The functions in Algorithm 1 are simplified versions of our actual code. We've put the full code in an anonymous link(https://anonymous.4open.science/r/ReZero-D6AD/README.md). We will attach the open source code to the final version of the paper for easy reference.

And we are sorry for the confusion caused by this simplified version of the code. You can find the corresponding functions(not strictly the same name) in https://anonymous.4open.science/r/ReZero-D6AD/lzero/mcts/buffer/game_buffer_rezero_mz.py . Specifically, the search is carried out in parallel in an entire batch. So prepare() do some centralized preprocessing of the data. origin_MCTS() is the original form of MCTS used in MuZero. traverse() refers to the forward process used in standard MCTS (from the root node to a leaf node). select_root_child() uses Equation 5 to select child nodes at the root node.

We will add the above short description to the new version of the paper, hoping that pairing it with open source code will reduce the confusion of readers.

In addition, reanalyze refers to a part of the MuZero algorithm which conducts MCTS on buffer datas. While traverse refers to the forward search process in MCTS (each forward search from the root node to a leaf node constitutes a traverse). So they're not on the same level. I hope this answer will reduce your confusion. Your further questions and suggestions are always welcome.

Dear Reviewer,

First of all, I would like to express my deep gratitude for the time and effort you have put in. You gave very detailed and substantial review comments. I'll reply step by step in the following comments.

Thank you for inviting me to discuss this. The scaling from (26) to (27) is inspired by the methods in bandit algorithms(https://tor-lattimore.com/downloads/book/book.pdf). You can refer to the second inequality in the last paragraph of the proof of Theorem 8.1 in Chapter 8 (page 120), which is an almost identical scaling.

The inequality you provided uses a double summation, and in my understanding, this scaling is similar to the one I used in the step from (18) to (19). However, from step (26) to (27), I adopted a different scaling method. The insight behind this is as follows: In (26), if a certain $t$ activates the indicator function, which means

$\mathbb{I}${$\hat{\mu}_i(t-1)+P_i \sqrt{\frac{n-1}{(1+T_i(t-1))^2}}\ge \mu_1-\varepsilon \ \text{and} \ A_t = i$} $=1$

it necessarily corresponds to a unique $s=T_i(t-1)$ that activates one of the indicator functions in the summation in (27), which means

$\mathbb{I}$ {$\hat{\mu}_{is}+P_i \sqrt{\frac{n-1}{(1+s)^2}}\ge \mu_1-\varepsilon$ } $=1$

Specifically, the condition $A_t=i$ plays a crucial role. Without this condition, it would not be guaranteed that different $t$s correspond to different $s$s. For example, suppose $t_1$ and $t_2$ each activate a indicator function in (26), but during this period, arm $i$ might not have been selected. In this case, $s_1=T_i(t_1-1)$ and $s_2=T_i(t_2-1)$ correspond to the same $s$, then simply summing over $s$ cannot guarantee an upper bound, since some certain $s$s may need to be considered more than one time. Therefore, a double summation is needed. However, under the condition $A_t=i$, once $t_1$ activates the indicator function, it means that the times arm $i$ being selected is also updated. If $t_2>t_1$ and $t_2$ also activates the indicator function, then there is always $s_2>s_1$, which means that each activated indicator function in (26) uniquely corresponds to an activated indicator function in (27), so (27) is an upper bound for (26).

I hope this explanation meets your satisfaction, and further discussion is always welcome. I am very grateful for the references you provided, which may help broaden my perspective and understand MCTS from different angles.

After checking the calculation, I find that it should be $e^{2\frac{\sqrt{a}\varepsilon}{C^2}}$ in 13(and the final result in (14) is right).

And it should be $1 + \frac{2C^6}{\varepsilon^4P_1^2}$ in (24). So I need to correct the constant term in the upper bound.

I apologize for these mistakes, which may have been caused by negligence in the calculation or when typing the formula. Thank you for meticulously checking these calculations and helping me correct my mistakes!

We can understand your concerns. ReZero itself is a general improvement method for the MCTS+RL series of algorithms. What it can handle depends on what algorithm it is applied to. For example, applying ReZero to Stochastic MuZero(https://openreview.net/forum?id=X6D9bAHhBQ1) can speed up stochastic dynamics cases. Applying ReZero to Sampled MuZero(https://arxiv.org/abs/2104.06303) can speed up the case of continuous action space(As we showed in DMC control tasks. the baseline is actually Sampled MuZero in DMC tasks, and we made a mistake on legend). Here Equation 5 uses the case of MuZero as an example to make a general illustration.

I can sense the effort you have put into understanding ReZero. I believe the confusion in Section 4.1 might stem from our differing background and domain knowledge. Our approach is an improvement upon the MuZero method, and the way we collect data, train networks, and so on, follows the MuZero framework. Therefore, we focus only on the improved modules and their corresponding analyses, without reintroducing the entire algorithm flow from scratch. This practice is in line with many well-regarded papers in the field: Sampled MuZero((https://arxiv.org/abs/2104.06303), Gumbel MuZero(https://openreview.net/forum?id=bERaNdoegnO), MuZero Unplugged(https://openreview.net/forum?id=HKtsGW-lNbw), and so on(apart from that, Stochastic MuZero(https://openreview.net/forum?id=X6D9bAHhBQ1) does have a short section on model training, as it makes changes to the network model itself).

Specifically, ReZero mainly modifies the reanalyze module of the MuZero algorithm. We have altered batch organization and MCTS setting during reanalyze. We use the known sample mean as estimation of a specific child node (for the remaining child nodes, the UCB value is still used). This change occurs only when selecting child nodes for the root node. For non-root nodes, the method of selecting child nodes remains the same as the original MCTS method(i.e. using UCB values) in MuZero.

Additionally, I noticed your concern about exploration. In our modified MCTS, there is indeed one child node that is directly evaluated using estimation. (This means that no exploration is conducted for this child node, but exploration for the remaining child nodes is maintained as it should be.) Therefore, I understand your concern that if this estimation has a large variance, it could affect the search. My explanation for this is that the variance of this estimation also decreases as the number of forward searches in MCTS increases. For example, if we set the number of forward searches in a single MCTS to $n$, since our estimation is reused from one MCTS(as showed in Figure 2), it is also obtained through $n$ samplings(For comparison, in the original MCTS setting, though we keep exploring this arm, the number of times we ultimately select this arm will not exceed $n$, because the total number of forward searches is $n$). Therefore, in our theoretical analysis, as $n$ increases, the regret can still maintain sublinear growth.

In other words, as long as the provided estimation is sufficiently good, exploration is not necessary. The most extreme example is the one-armed bandit model(checck exercise 8.3 in page 121 of https://tor-lattimore.com/downloads/book/book.pdf), if the exact parameters are already known, there is no need to explore the corresponding arm. In our setting, the exact expected return is unknown, so we use an estimation obtained after sufficiently many samplings.

I fully understand the confusion you have experienced, and I hope my response can help you understand our method on certain levels. Your further suggestions and comments are always welcome.

We have fixed the corresponding errors in the new version of our paper. Additionally, we have modified some notations and improved the statement of the theorem. Thank you very much for your correction and helping us improve the paper.

In reinforcement learning, we often have a buffer to store data generated during interactions with the environment. The purpose of reanalyze is to update old data to ensure fresh training targets. When reanalyzing, we sample some data from the entire buffer to form a batch and do MCTS for every sample in it. This batch is collected in units of trajectories. Here, a trajectory is actually a segment cut out from a complete game. We fix the length of these segments and ensure they do not intersect (to avoid redundant searches and waste computational resources). Next, we further split the batch into several mini-batches, placing states with the same sequence number from each trajectory into the same mini-batch. The states in the same mini-batch are executed MCTS in parallel (this is merely a programming technique to allow the model to infer in parallel and save time). As shown in Figure 2, when reusing information, each state in the mini-batch will only receive the value transmitted from its own trajectory, meaning only one child node would be evaluated by estimation value, while the remaining child nodes need to be explored normally.

I hope this answer can help reduce some of your confusion, and further discussion and questions are always welcome.