SPENDING THINKING TIME WISELY: ACCELERATING MCTS WITH VIRTUAL EXPANSIONS
Year: 2021

Abstract:
One of the most important AI research questions is to trade off computation versus performance, since “perfect rational” exists in theory but it is impossible to achieve in practice. Recently, Monte-Carlo tree search (MCTS) has attracted considerable attention due to the significant improvement of performance in varieties of challenging domains. However, the expensive time cost during search severely restricts its scope for applications. This paper proposes the Virtual MCTS (VMCTS), a variant of MCTS that spends adequate amounts of time to think about different questions. Inspired by this, we propose a strategy that converges to the ground truth MCTS search results with much less computation. We give theoretical bounds of the the proposed method and evaluate the performance in 9× 9 Go board games and Atari games. Experiments show that our method can achieve similar performances as the original search algorithm while requiring less than 50% number of search times on average. We believe that this approach is a viable alternative for tasks with limited time and resources.

heading: 1 INTRODUCTION
text: When artificial intelligence was first studied in the 1950s, researchers seek to answer the question of what is the solution to the question if the agent were “perfect rational”. The term “perfect rational” here refers to the decision made with infinite amounts of computations. However, without taking into consideration the practical computation time, one can only solve small-scale problems, since classical search algorithms usually exhibit exponential running time. Recent AI researches no longer seek to achieve “perfect rational”, but instead carefully trade-off computation versus the level of rationality. People have developed computational models like “bounded optimality” to model these settings (Russell & Subramanian, 1994). The increasing level of rationality under the same computational budget has given us a lot of AI successes nowadays. Notable algorithms include the Monte-Carlo sampling algorithms, the variational inference algorithms, and using neural networks as universal function approximators (Coulom, 2006; Chaslot et al., 2008; Gelly & Silver, 2011; Silver et al., 2016; Hoffman et al., 2013).
More recently, MCTS-based RL algorithms have achieved a lot of success, mainly in board games. The most notable achievement is that AlphaGO beats Hui Fan in 2015 (Silver et al., 2016). This is the first time that a computer program beats a human professional player. After that, AlphaGo beats two top-ranking human players, Lee Sedol in 2016 and Jie Ke in 2017, the latter of which rank first worldwide at the time. Later, the MCTS-based RL algorithms are further extended to other board games, as well as the Atari video games (Schrittwieser et al., 2020). EfficientZero (Ye et al., 2021) greatly improves the sample efficiency of MCTS-based RL algorithms, shedding light on its future applications in real-world applications like robotics and self-driving.
Despite the impressive performance of MCTS-based RL algorithms, they require massive computations to train and evaluate. For example, Schrittwieser et al. (2020) used 1000 TPUs trained for 12 hours to learn the game of GO, and for a single Atari game, it needs 40 TPUs to train 12 hours. Compared to previous algorithms on the Atari games benchmark, it needs around two orders of magnitude more compute. This prohibitively large computational requirement has slowed down both the further development of MCTS-based RL algorithms, as well as its practical use.
Under the hood, MCTS-based RL algorithms are model-based methods, that imagine what the futures look like when doing different future action sequences. However, this imaging process for the
current method is not computationally efficient. For example, AlphaGo needs to look ahead 1600 game states to place a single stone. On the contrary, top human professional players can only think through around 100-200 game states per minute (Silver et al., 2016). Besides being computationally inefficient, the current MCTS algorithm deals with easy cases and hard ones with the same computational budget. On the other hand, human knows to use their time when it is most needed.
In this paper, we aim to design new algorithms that save the computational time of the MCTS-based RL methods. More specifically, we are interested in pushing the Pareto front of the rationality level - computation curve. Empirical results show that our method can achieve comparable performance while requiring less than 50% simulations to search on average.

heading: 2 RELATED WORK

heading: 2.1 MULTI-ARMED BANDIT PROBLEM
text: RL algorithms are always brought into the exploration and exploitation dilemma. Multi-armed bandit (MAB) problem (Berry & Fristedt, 1985; Auer et al., 2002; Lattimore & Szepesvári, 2020) is one of the most extensively studied but fundamental instances. TheK-armed MAB problem is a sequential game with a collection of K unidentified but independent reward distributions, each associated with the corresponding arms. For each round, the learner pulls an arm and receives a reward sampled from the corresponding distributions. The optimal policy of the learner for the MAB problem is to maximize the cumulative rewards obtained from the sequential decisions.
In the case where the cost of pulling arms is little, the learner is allowed to trial and error for enough times until convergence. A series of upper confidence bound (UCB) algorithms (Auer et al., 2002; Bubeck & Cesa-Bianchi, 2012) are proposed to solve the stochastic MAB problem and they have theoretical bounds. When there exist costs for each trial, pure exploration attempts to make the best use of the finite trials (Bubeck et al., 2011; Lattimore & Szepesvári, 2020). Kocsis & Szepesvári (2006) proposed UCT to adapt UCB algorithms to the tree structures, which is the basis of MCTS.

heading: 2.2 REINFORCEMENT LEARNING WITH MCTS
text: For a long time, Computer Go is regarded as a very challenging game (Bouzy & Cazenave, 2001; Cai & Wunsch, 2007). Researchers attempt to use Monte-Carlo techniques that evaluate the value of the node state through random playouts (Bouzy & Helmstetter, 2004; Gelly & Silver, 2007; 2008; Silver et al., 2016). Afterwards, UCT has generally replaced those earlier heuristic methods for Monte-Carlo tree search (MCTS). UCT algorithms (Kocsis & Szepesvári, 2006) apply UCB1 to select action at each node of the tree. Recently, MCTS-based methods (Silver et al., 2016; 2017; 2018; Schrittwieser et al., 2020) have become increasingly popular and achieved super-human performances on board games because of the strong ability to search. Modern MCTS-based RL algorithms include four stages in each search iteration, namely simulation: selection, expansion, evaluation, and backpropagation. The selection stage targets selecting a new leaf node with UCT. The expansion stage expands the selected node and updates the search tree. The evaluation stage evaluates the value of the new node. The backpropagation stage propagates the newly computed value to the nodes along the search path to obtain more accurate Q-values with Bellman backup. However, search is quite time-consuming, which prevents MCTS-based methods to be used in wider scenarios.

heading: 2.3 ACCELERATION OF MCTS
text: There are two kinds of bottlenecks in MCTS in the aspect of speed: the evaluation/selection stage of each iteration and the outer search loop. In the previous research, people attempted to evaluate the node value by random playouts to the end of the game, which makes the evaluation stage quite expensive. In addition, compared to other model-free RL methods like PPO (Schulman et al., 2017) and SAC (Haarnoja et al., 2018), MCTS-based algorithms have much larger computations due to the search loop. Therefore, a lot of works are devoted to accelerating MCTS. Some heuristic pruning methods (Gelly et al., 2006; Wang & Gelly, 2007; Sephton et al., 2014; Baier & Winands, 2014; 2018) are developed to make the selection or evaluation more effectively. Lorentz (2015) proposed early playout termination of MCTS (MCTS-EPT) to stop the random playouts early and use an evaluation function to assess win or loss, which is an improvement in the evaluation stage of normal
MCTS. And Hsueh et al. (2016) applied MCTS-EPT to the Chinese dark chess. Afterwards, MCTSEPT similar ideas have been applied in the evaluation stage of AlphaGoZero (Silver et al., 2017) and later MCTS-based methods (Silver et al., 2018; Schrittwieser et al., 2020; Ye et al., 2021) including our baseline models. They evaluate the Q-values through evaluation networks instead of running playouts to the end. However, these methods focus on the specific stage of the search iteration to accelerate the MCTS. We propose Virtual MCTS, which aims to terminate the outer search iteration adaptively in MCTS under distinct circumstances without sacrificing the final policy quality.

heading: 3 BACKGROUND
text: The AlphaGo series of work (Silver et al., 2016; 2017; 2018; Schrittwieser et al., 2020) are all MCTS-based reinforcement learning algorithms. Those algorithms assume the environment transition dynamics are known or learn the environment dynamics. Based on the dynamics, they use the Monte-Carlo tree search (MCTS) as the policy improvement operator. I.e. taking in the current policy, MCTS returns a better policy with the search algorithm. The systematic search allows the MCTS-based RL algorithm to quickly improve the policy, and perform much better in the setting where a lot of reasoning is required. MCTS is the core component in the algorithms like AlphaGo.

heading: 3.1 MCTS
text: In this part, we give a brief introduction to the MCTS method implemented in reinforcement learning applications. MCTS takes in the current MDP state and runs a search algorithm guided by the current policy function. It outputs an improved policy of the current state. The improved policy is later to select an action in the environment. In the selection stage, an action will be selected by maximizing over UCB. Specifically, AlphaZero (Silver et al., 2018) and MuZero (Schrittwieser et al., 2020) are developed based on a variant of UCB, P-UCT (Rosin, 2011) and have achieved great success on board games and Atari games. The formula of P-UCT in the two methods is the Eq (1):
ak = argmax a∈A Q(s, a) + P (s, a)
√∑ bN(s, b)
1 +N(s, a)
( c1 + log (∑ bN(s, b) + c2 + 1
c2
)) , (1)
where k is the index of the iterative step,A is the action set,Q(s, a) is the estimated Q-value, P (s, a) is the policy prior obtained from neural networks and N(s, a) is the visit counts to select the action a from the state s. The output of MCTS is the visit count of each action of the root node. After N search iterations, the final policy π(s) is defined as the normalized root visit count distribution πN (s), where πk(s, a) = (N(s, a))/ ∑ b∈AN(s, b) = N(s, a)/k, a ∈ A. For simplification, we use πk in place of πk(s) sometimes. In our method, we propose to approximate the final policy πN (s) with π̂k(s), which we name as a policy candidate, through a new expansion method and a termination rule. In this way, the number of iterations in MCTS can be reduced from N to k.

heading: 3.2 COMPUTATION REQUIREMENT
text: Most of the computations in MCTS-based RL are in the MCTS procedure. For each action taken by MCTS, it needs N times neural network evaluations, where N is the number of search iterations in MCTS. Traditional RL algorithms, such as PPO (Schulman et al., 2017) or DQN (Mnih et al., 2015), only need a single neural network evaluation per action. Thus, MCTS-based RL is roughly N times computationally more expansive than traditional RL algorithms.
In practice, training a single Atari game needs 12 hours of computation time on 40 TPUs (Schrittwieser et al., 2020). The computation need is roughly two orders of magnitude more than traditional RL algorithms (Schulman et al., 2017), although the final performance of MuZero is much better.

heading: 4 METHOD
text: In this paper, we propose an algorithm named Virtual MCTS to reduce the computation cost of MCTS-based RL algorithms. More concretely, we aim to push the front of the Pareto curve of the performance–computation trade-off. Intuitively, human knows when to make a quick decision and when to make a slow decision in different circumstances. It gives the ability to overcome
more difficult problems without wasting much time on easy ones. This situation-aware behavior is absent in current MCTS algorithms. We propose an MCTS algorithm variant that early terminates the search iteration in MCTS adaptively under distinct situations of states. It is consists of two components, the virtual expansion to estimate the final visit count based on the current partial tree; the termination rule that decides when to terminate based on the hardness of the current scenario.

heading: 4.1 TERMINATION RULE
text: We propose to terminate the MCTS early based on the current tree statistics. Intuitively, during the MCTS tree expansion process, if we find that recent searches do not further change the root visitation distribution, then we no longer need to search further. With this intuition in mind, we propose a simple modification to the MCTS search algorithm. Let πk(s) denote the root visitation distribution of the root state s at MCTS expansion iteration k. We propose to terminate when:
||πk(s)− πk/2(s)||1 < where is a tolerance hyper-parameter. Note that, in MCTS-based RL algorithms, not only the best arm matters but the other arms matter as well. It is because MCTS is used in the exploration process, and we need to make sure proper exploration happens at the non-best arms. This seems to be a heuristic rule, without any guarantees whether πk is close to πN , whereN is the original MCTS search iteration. However, we show that under certain conditions, a bound on ||πk(s) − πk/2(s)||1 implies a bound on ||πk(s)− πN (s)||1. First of all, we list some notations: k is the index of the current search iteration and N is the number of total search simulations, Rt(s, a) is the predicted Q-value output at the t-th iteration, Nk(s, a) denotes the total visit counts of the action a from the state s after the k-th iteration, the Q-value at k-th iteration is Qk(s, a) = ∑k t=1Rt(s, a)/Nk(s, a). The Lemma here gives a brief bound for the ranges of Q-values at different iterations, and the proof is attached in Appendix A.2. Lemma 1 ∀a ∈ A, given that Rt(s, a) ∈ [−1, 1] and the Qk(s, a) = ∑Nk(s,a)
t=1 Rt(s,a) Nk(s,a) , then at
iteration 1 ≤ k1 < k2 ≤ N , Qk2(s, a)−Qk1(s, a) ≤ (1− Nk1 (s,a) Nk2 (s,a) )(1−Qk1(s, a)).
It tells that the future changes ofQ-values withN−k further searches are bounded in a small range. Through the virtual expansion that we propose in Section 4.2, we can generate a policy candidate π̂k to approximate the final oracle policy πN . Furthermore, Lemma 2 measures the distance between the oracle policy πN and our current policy candidate π̂k and the proof is attached in Appendix A.2. If the policy candidates π̂k and π̂k/2 are close enough, so are π̂k and π̂N . In the proof, we show that π̂N is equal to the oracle policy πN . Therefore we conclude that the current policy candidate π̂k is a near-oracle policy if the termination rule
∣∣∣∣π̂k(s)− π̂k/2(s)∣∣∣∣1 < is satisfied. Lemma 2 ∀a ∈ A, given that r ∈ (0, 1], if ∃k ∈ [rN,N ],
∣∣∣∣π̂k(s)− π̂k/2(s)∣∣∣∣1 < , then ||π̂k(s)− πN (s)||1 < + 1− r.

heading: 4.2 VIRTUAL EXPANSION IN MCTS
text: In the derivation above, we assume πi and πj are directly comparable. Here πi and πj denote two root node visit count distributions at iteration i and j respectively. However, because the tree is expanded with UCT, they are not directly comparable. UCT is an algorithm that maintains the upper bound of the node values in the search tree. As the number of visits increases, the upper bound would be tighter and the latter visits are more focused on the most promising part of the tree. Thus earlier visit count distributions (smaller iteration number) can exhibit more exploratory distribution, while latter ones (larger iteration number) are more exploitative on promising parts.
To compare πi and πj properly, we propose a method called virtual expansion in MCTS. In a nutshell, it aligns two distributions by virtual UCT expansions until the constant budget N . When the tree is expanded at iteration i, it hasN− i iterations to go. A normal expansion would require evaluating neural network N − i times for a more accurate Q(s, a) estimate for each arm at the root node. Our proposed virtual expansion still expands N − i times according to UCT, but it ignores the N − i neural network evaluations and simply assumes that each arm’s Q(s, a) does not change. We denote the virtual expanded distribution from πi as a policy candidate π̂i. By doing virtual expansions on both πi and πj , we effectively remove the different levels of exploration/exploitation in them.
Algorithm 1 Iteration of Search in MCTS 1: Current k-th iteration step: 2: Given: A, P,Qk(s, a), Nk(s, a) 3: s← sroot 4: repeat do search 5: a∗ ← UCT (Q,P,N) 6: s← next state(s, a∗) 7: until Nk(s, a∗) = 0 8: Evaluate the state value R(s, a) and P (s, a) 9: for s along the search path do 10: Qk+1(s, a) = Nk(s,a)·Qk(s,a)+R(s,a)
Nk(s,a)+1
11: Nk+1(s, a) = Nk(s, a) + 1 12: end for 13: Return Qk+1(s, a), Nk+1(s, a)
Algorithm 2 Iteration of Search in MCTS with Virtual Expansion
1: Current k-th iteration step: 2: Given: A, P,Qk(s, a), Nk(s, a), N̂k(s, a) 3: 4: if Not init N̂k(s, a) then 5: Init: N̂k(s, a)← Nk(s, a) 6: end if 7: 8: s← sroot 9: a∗ ← UCT (Q,P, N̂)
10: N̂k(s, a)← N̂k(s, a) + 1 11: 12: Return N̂k(s, a)
The comparisons between the MCTS and the one with virtual expansion are illustrated in Algorithm 1, 2. Here we display the complete one-step iteration of MCTS with or without virtual expansion. The time-consuming computations are highlighted in Algorithm 1. Line 4 to 7 in Algorithm 1 target at searching with UCT to reach an unvisited state for exploration. Then it evaluates the state and backpropagates along the search path to fit a better estimation of Q-values. After a total of N iterations, the visit count distribution of the root node πN (s) is considered as the final policy π(s). However, in the MCTS with virtual expansions, listed in Algorithm 2, it only searches one step from the root node and selects the action based on the current estimations without changing any properties of the search tree. Furthermore, the virtual visited counts N̂k(s, a) are changed after virtual visits to balance the exploitation and the exploration issue. Then the policy candidate after virtual expansions becomes π̂k(s, a) = N̂k(s, a)/N instead of Nk(s, a)/k. When k = N , further searches after the root have no effects on the final policy. So we have π̂N (s, a) = πN (s, a). In this way, the final visit count distribution obtained through the virtual expansions is similar to the oracle one.

heading: 4.3 V-MCTS ALGORITHM SUMMARY
text: Algorithm 3 Virtual MCTS 1: Given budget N , state s, conservativeness r, error 2: for k ∈ N do 3: Selection with UCT 4: Expansion for the new node 5: Evaluation with Neural Networks 6: Backpropagation for updating Q and visit counts 7: πk(s, a)← Nk(s, a)/n 8: Virtual expand N − k nodes and update N̂(s, a) 9: π̂k(s)← N̂k(s, a)/N
10: if k ≥ rN ∧ ∣∣∣∣π̂k(s)− π̂k/2(s)∣∣∣∣1 < then 11: π(s)← π̂k(s) 12: Break 13: end if 14: π(s)← πk(s) 15: end for 16: Return π(s) The procedure of MCTS with the termination rule is listed as Algorithm 3. Compared with the original MCTS, the line 8-13 are the pseudo code on the termination rule. In each iteration, we do some calculations with little cost to judge whether the condition of termination is satisfied. If it is, then the search process is terminated and returns the current policy candidate π̂k(s). Thus, it skips the next N − k model predictions from neural networks in the evaluation stage highlighted in line 5. In this way, we can approximate the oracle distribution πN by π̂k while reducing the budget of N simulations to k. Here, k ≥ rN and r is a hyperparameter of the minimum budget rN . Then we can reduce the tree size by 1/r times at most. However, r cannot be a tiny value as the minimum distance of distributions is bounded by r. Otherwise, the termination rule has little effect. In conclusion, this rule tells the MCTS to terminate if the policy candidates have converged. On such occasions, the virtual expansion method has similar effects as the real expansion of MCTS to generate a near-oracle policy, which can save the time cost of the left N − k simulations.
Furthermore, in the next section, we do some ablations to further investigate the effects of r, and visualizations to understand why the termination rule works. We name our method Virtual MCTS (V-MCTS), a variant of MCTS with a termination rule based on the virtual expansion.

heading: 5 EXPERIMENTS
text: In this section, the goal of the experiments is to prove the effectiveness and efficiency of our proposed algorithm. We compare the performance as well as the cost of the budget of the MCTS-based methods with or without the termination rule. Specifically, we evaluate the board game Go 9 × 9, and a few Atari games. In addition, we do some ablations to further examine the effectiveness of the virtual expansion and evaluate how sensitive our method is to the hyper-parameters. Finally, we try to understand the adaptive mechanism with visualizations and performance analysis.

heading: 5.1 SETUP
text: Recently, Ye et al. (2021) proposed EfficientZero, a variant of MuZero (Schrittwieser et al., 2020) with three extra components to improve the sample efficiency, which only requires 8 GPUs in training, and thus it is more affordable. Here we choose the board game Go 9 × 9 as our benchmark environment. The game of Go tests how the algorithm performs in a challenging planning problem. We also benchmark on a few Atari games, which feature complexity on the visual side.
As for the Go 9×9, we choose Tromp-Taylor rules during training and evaluation. The environment of Go is built based on an open-source codebase, GymGo (Huang, 2021). We evaluate the performance of the agent against GNU Go v3.8 at level 10 (Bump et al., 2005) for 200 games. We include 100 games as the black player and 100 games as the white one with different seeds. We set the komi to 6.5 as most papers do. As for the Atari games, we choose 5 games with 100k environment steps, which follows the setting of EfficientZero. We evaluate all these games for 32 distinct seeds.

heading: 5.2 COMPARATIVE EVALUATION
text: We compare our method to EfficientZero with original MCTS, on Go 9× 9 and some Atari games. Figure 1a illustrates the comparisons on Go among different algorithms or models against the GnuGo (level 10) agent. The x-axis is the training speed and the y-axis is the winning rate. Here, we train the baseline method with different constant budgets N , noted as the blue points. Besides, we also train the V-MCTS with hyperparameters r = 0.2, = 0.1. We evaluate the trained model with different to display the tradeoff between performance and the time cost, noted as the red points. The pink points are the GnuGo with different levels.
The result shows that the performance of V-MCTS is comparable to the MCTS of maximum budget (N = 150) while it requires less time to search. Notably, V-MCTS achieves a 72% winning rate against the GnuGo level 10. Meanwhile, the time cost of our method for a one-step move is 0.12s while the GnuGo engine is 0.18s. For the data points with the winning rate higher than 50%, the VMCTS is significantly better than the original MCTS considering both the winning rate and the time cost. Consequently, we believe that such termination rule of MCTS can keep the strong performance while saving lots of budgets, which means our method is effective enough and more time-efficient.
The training curve presented in Figure 1b illustrates that the budget of search times is quite significant to the MCTS method. However, the performance of V-MCTS is better than that of MCTS with N = 120 while maintaining an average tree size of less than 80. Furthermore, it is interesting that the tree size varies over training procedures. In the beginning, the outputs of the value network are close to zero because the agent usually draws in self-play and receives the zero reward signal, which means the little changes of Qk(s, a) during the search. The visit count distribution of the root after virtual expansions is similar to that after real expansions. Therefore, in this stage, the termination condition is easy to meet. Afterward, as the model is trained better and receives more diverse reward signals during self-play, the probability of searching some valuable states becomes much larger. The value Qk(s, a) varies considerably in each iteration of the search, which results in the changes of distributions of policy candidates. Thus, it is much more difficult for the search process to terminate and leads to a larger number of tree sizes. Finally, with more training steps, the prediction of the policy network is more accurate and then gives a stronger prior heuristic knowledge P (s, a) before
search. For those actions with higher prior knowledge, the changes of Qk(s, a) have less impact on the UCT scores. The procedure of the virtual expansion is similar to the real expansion. However, the changes of Qk(s, a) are still possible during the search process. Consequently, the average number of tree size keeps in a reasonable range, which is larger than that in the innocent beginning stage and the minimum budget rN . Our method can determine whether to continue searching or not.
Apart from the results of Go, we also evaluate our method on some visually complex games. Since the search space of Atari games is much smaller than that of Go and the Atari games are easier, here we choose a few Atari games to study how the proposed method impacts the performance under less necessity of search. We follow the setting of EffcientZero, 100k Atari benchmark, which contains only 400k frames data. The results are shown in Table 1. Generally, we find that our method works in Atari games. The tree size is adaptive and the performance of V-MCTS is still comparable to the MCTS with full search trails. It has better performance than the MCTS(N = 30) while requiring much fewer searches, which proves the effectiveness and the efficiency of the termination rule in the tree search method. The Hero game is not an outlier considering the similar performance between N = 50 and N = 30. Besides, the number of search times decreases more than that on Go.
To sum up, Virtual MCTS shapes better policy candidates close to the oracles through the virtual expansion with less cost of the budget. The performance of V-MCTS can keep sound with fewer search iterations while simply reducing the total budget of MCTS will encounter a larger performance drop. In addition, the savings of search cost is more substantial in easier environments.

heading: 5.3 ABLATION STUDY
text: The results in the previous section suggest that our method reduces the response time of MCTS while keeping considerable performance on challenging tasks. In this section, we try to figure out which component contributes to the performance and how the hyperparameters affect it.
Virtual Expansion In Section 4.2, we introduce the virtual expansion and discuss the difference between the MCTS with and without virtual expansion. We compare the MCTS with virtual expansion and another two expansion methods. Here we introduce the two methods briefly. One is the original expansion, which does nothing once termination. It samples an action directly from πk(s, a) = Nk(s, a)/k. Another is the greedy expansion, which spends the left N − k simulations in searching the current best action greedily, indicating that π̂k(s, a) = (Nk(s, a) + (N − k)Ia=argmaxNk(s,a))/N . Besides, we stop the search process after k iterations regardless of the termination condition, where k = rN and r = 0.2, N = 150.
We compare the winning rate against the GnuGo engine and the results are listed as Table 2 shows. All the versions here only search for the given minimum tree size, and the virtual expansion method can still achieve a 32% winning rate, which is much better than the others. Notably, MCTS with greedy expansion does not work. It is over-exploitation and results in severe exploration issues. Consequently, the virtual expansion method can generate a better policy distribution because it can balance the exploration and exploitation problem through UCT with further virtual simulations.
Termination Rule Since the virtual expansion provides a better choice of policy distributions, it is significant to explore a better termination rule to keep the sound performance while decreasing the tree size as much as possible. As mentioned in Section 4.1, the termination rule sets two hyperparameters r, to determine the termination rule. Then we do ablations for the different values of r and respectively. The default values of r, are set to 0.2, 0.1 in all experiments here.
Figure 2 compares the winning rate as well as the average tree size across the training stage. Firstly, Figure 2a gives the results of different minimum search times r. The winning probability is not sensitive to r when r ≥ 0.2. But the average tree size is sensitive to r because V-MCTS is supposed to search for at least rN times. In addition, there is a drop between the performance between r = 0.1 and r = 0.2. Therefore, it is reasonable to choose r = 0.2 to balance the speed and the performance.
Besides, the comparisons of different minimum distance are shown in Figure 2b. A larger makes the tree size smaller because ∣∣∣∣π̂k(s)− π̂k/2(s)∣∣∣∣1 < is easier to reach. In practical, we find that the performance is highly correlated with . In terms of the winning probability, a smaller outperforms a larger one. However, the better performance is at cost of a larger response time. Therefore, it is a good choice to set r = 0.2, = 0.1. We suggest selecting an appropriate minimum distance to balance the performance and the response time.

heading: 5.4 VISUALIZATION
text: In this section, we do some visualizations to better understand the behavior of Virtual MCTS. Specifically, we choose some states at different time steps on one game of Go against the GnuGo with a trained model and visualize the heatmap of policy distribution, as Figure 3 shows. In this figure, our player is the black one. The board states are shown in the first row, and the next two rows are the heatmap visualization for oracle MCTS and V-MCTS. The darker the color is on the grid, the more the corresponding action is visited during the search. In general, π̂k is close to the πN at distinct states, indicating that the termination rule is reasonable and effective. The less valuable actions there are, the sooner the V-MCTS will terminate. For example, on Go games, the start states are usually not complex because there are only a few stones on the board but the situations are much more complicated in shuban, the closing stage of the game. Notably, the termination occurs earlier in the start states (columns 1, 2) but it is the opposite when the situation is more complicated. More importantly, the termination step k is not related to the number of Go pieces. And the policy candidate obtained after virtual expansion can be close to the oracle one at distinct states of a game. Therefore, we can conclude that V-MCTS makes adaptive terminations according to the situations of the current state, which is the key to maintaining comparable performances.

heading: 6 DISCUSSION
text: In this paper, we propose a novel method named V-MCTS to accelerate the MCTS to determine the termination of search iterations. It can keep comparable performances while reducing half of the time to search adaptively. We are in the belief that this work can be one step toward applying the MCTS-based methods to some real-time domains.

heading: 7 REPRODUCIBILITY STATEMENT
text: The main implementations of our proposed method are in Algorithm 2 and 3. In addition, the settings of the experiments and hyper-parameters we choose are in Appendix A.1. The proof of the lemma is in Appendix A.2. More significantly, the details of the design of training procedures for Go are around Appendix A.1.2. Besides, we will release the codebase if this paper is accepted.

heading: A APPENDIX
text: A.1 EXPERIMENTAL SETUP
A.1.1 MODELS AND HYPER-PARAMETERS
Model Design As for the architecture of the networks, we follow the implementation of EfficientZero (Ye et al., 2021) in Atari games, which proposes three components based on MuZero:
self-supervised consistency, value prefix, and off-policy correction. In the implementation of EfficientZero, there are a representation network, a dynamics network and a reward/value/policy prediction network. The representation network is to encode observations to hidden states. The dynamics network is to predict the next hidden state given the current hidden state and an action. The reward/value/policy prediction network is to predict the reward/value/policy. Notably, they propose to keep temporal consistency between st+1 and the predicted state ŝt+1. The training objective is:
Lt(θ) = λ1L(ut, rt) + λ2L(πt, pt) + λ3L(zt, vt) + λ4Lsimilarity(st+1, ŝt+1) + c||θ||2
L(θ) = 1 lunroll lunroll−1∑ i=0 Lt+i(θ), (2)
where ut, πt, zt are the target reward/policy/value of the state st and rt, pt, vt are the predicted reward/policy/value of the state st respectively. The prediction will do lunroll = 5 times iteratively for the state s on both Go and Atari games. We do some changes when dealing with board games. For example, we remove the reward prediction network because the agent will receive a reward only at the end of the games. The other major changes for board games are listed as follows.
Since the board game Go is harder than the Atari games, we add more residual blocks (two times blocks). Specifically, we use 2 residual blocks in the representation network, the dynamics network as well as the value/policy prediction network on Go 9× 9 while EfficientZero uses only 1 residual block in those networks on Atari games. As for the representation network, we remove the downsampling part here because there is no need to do downsampling for Go states. In the value/policy prediction networks, we enlarge the dimension of the hidden layer from 32 to 128. Besides, considering that the reward is sparse on Go (only the final value) and the collected data are sufficient, we only take the self-supervised consistency component in EfficientZero to give more temporal supervision during training.
Hyper-parameters In each case, we train EfficientZero for unrolled 5 steps and mini-batches of size 256. Besides, the model is trained for 100k batches with 100M frames of data in board games while 100k batches with 400k frames in Atari games. We stack 8 frames in board games without frameskip while stacking 4 frames in Atari games with a frameskip of 4. During both training and evaluation, EfficientZero chooses 150 simulations of budget for each search in board games while 50 simulations of budget in Atari games. Other hyper-parameters are listed in Table 3.
A.1.2 TRAINING DETAILS OF GO
The detailed implementations of Atari games are discussed in EfficientZero (Ye et al., 2021). However, it is nontrivial to adapt to board games. Here we give detailed instructions for training board games Go 9× 9 in our implementations. Inputs We follow the designs of AlphaZero and we use the Tromp-Taylor rules, which is similar to previous work (Silver et al., 2018; Schrittwieser et al., 2020). The input states of the Go board are encoded into a 17× 9× 9 array which stacks the historical 8 frames and uses the last channels C to identify the current player, 0 for black and 1 for white. Notably, the one historical frame consists of two planes [X,Y ], where the first plane X represents the stones of the current player and the second one [Y ] represents the stones of the opponent. Besides, if there is a stone on board, then the state of the corresponding position in the frame will be set to 1, otherwise to 0. For example, if the current player is black and suppose b[i, j] is the current board state, X[i, j] = 1b[i,j]=black stone, Y [i, j] = 1b[i,j]=white stone. In summary, we concatenate together the historical planes to generate the input state s = [Xt−7, Yt−7, Xt−6, Xt−6, ..., Xt, Yt, C], where Xt, Yt are the feature planes at time step t and C gives information of the current player.
Training As for the training phase, we train the model from scratch without any human expert data, which is the same as the setting of Atari games. Besides, limited to the GPU resources, we do not use the reanalyzing mechanism of MuZero (Schrittwieser et al., 2020) and EfficientZero (Ye et al., 2021), which targets at recalculation of the target values and policies from trajectories in the replay buffer with the current fresher model. Specifically, we use 6 GPUs for doing self-play to collect data, 1 GPU for training, and 1 GPU for evaluation.
Exploration To make a better exploration on Go, we reduce the α in the Dirichlet noise Dir(α) from 0.3 to 0.03 and we scale the exploration noise through the typical number of legal actions, which follows these works (Silver et al., 2018; Schrittwieser et al., 2020). In terms of sampling actions from MCTS visit distributions, we will mask the MCTS visit distributions with the legal actions and sample an action at, where
at := { at ∼ πt, t < T at = argmaxπt, t ≥ T
(3)
T is set to 16 in self-play and is set to 0 in evaluation, which is similar to these works (Silver et al., 2018; Schrittwieser et al., 2020). In this way, the agent does more explorations for the previous T steps while taking the best action afterwards. But for Atari games, ∀t, we choose at ∼ πt in self-play and at = argmaxπt in evaluation, which is the same as these works (Schrittwieser et al., 2020; Ye et al., 2021).
Two-player MCTS On board games, there are two players against each other, which is different from that of one-player games. Therefore, we should do some changes to the MCTS with the twoplayer game. For one thing, the value network always predicts the Q-value of the black player instead of the current player, which provides a more stable prediction. Furthermore, A significant change is that during backpropagation of MCTS, the value should be updated with the negative value from the child node. Because the child node is the opponent, the higher value of the opponent indicates a worse value of the current player. Besides, as for Q-values of the unvisited children on Go and Atari games, we follow the implementation of EfficientZero (Ye et al., 2021) as follows:
Q̂(sroot) = 0
Q̂(s) = Q̂(sparent) +
∑ b 1N(s,b)>0Q(s, b)
1 + ∑
b 1N(s,b)>0
Q(s, a) : =
{ Q(s, a) N(s, a) > 0
Q̂(s) N(s, a) = 0
(4)
Notably, we allow the resignation for players when maxa∈AQ(sroot, a) < −0.9 during self-play and evaluation, which means that the predicted winning probability is less than 5%. As for other hyperparameters on both Go and Atari games, we note that we choose the same values as those in EffcientZero. Specifically, the c1, c2 in our mentioned P-UCT formula (Eq. 1) are set to 1.25 and 19652, following these works (Schrittwieser et al., 2020; Ye et al., 2021).
A.2 PROOF Lemma 1. ∀a ∈ A, given that Rt(s, a) ∈ [−1, 1] and the Qk(s, a) = ∑Nk(s,a)
t=1 Rt(s,a) Nk(s,a) , then at
iteration 1 ≤ k1 < k2 ≤ N , Qk2(s, a)−Qk1(s, a) ≤ (1− Nk1 (s,a) Nk2 (s,a) )(1−Qk1(s, a)).
proof. Since Qk(s, a) =
k∑ t=1 Rt(s,a)
Nk(s,a) , we have
Qk2(s, a) = Nk1(s, a)Qk1(s, a) +
∑Nk2 (s,a) t=Nk1 (s,a) Rt(s, a)
Nk2(s, a) (5)
Because Rt(s, a) ∈ [−1, 1], then
Qk2(s, a) ≤ Nk1(s, a)
Nk2(s, a) Qk1(s, a) + Nk2(s, a)−Nk1(s, a) Nk2(s, a)
(6)
Therefore
Qk2(s, a)−Qk1(s, a) ≤ (1− Nk1(s, a)
Nk2(s, a) )(1−Qk1(s, a)) (7)
Lemma 2. Given that r ∈ (0, 1], if ∃k ∈ [rN,N ], ∣∣∣∣π̂k(s)− π̂k/2(s)∣∣∣∣1 < , then ||π̂k(s)− πN (s)||1 < + 1− r
proof. Since the π̂k(s, a) is a distribution that equals to N̂k(s,a)
N . Therefore, ||π̂N (s)− π̂k(s)||1 ≤ ∣∣∣∣π̂N (s)− π̂N/2(s)∣∣∣∣1 + ∣∣∣∣π̂N/2(s)− π̂k(s)∣∣∣∣1 (8)
Then
||π̂N (s)− π̂k(s)||1 ≤ M∑ t=0 ∣∣∣∣∣∣π̂N 2t (s)− π̂ N 2t+1 (s) ∣∣∣∣∣∣ 1 + ∣∣∣∣∣∣π̂k(s)− π̂ N 2M+1 (s) ∣∣∣∣∣∣ 1
(9)
, where N 2M+2 ≤ k ≤ N 2M+1 . In addition, ∣∣∣∣∣∣π̂k(s)− π̂ N 2M+1 (s) ∣∣∣∣∣∣ 1 ≤ ∣∣∣∣π̂k(s)− π̂k/2(s)∣∣∣∣1 < . Besides, k ≥ rN , then we have
||π̂N (s)− π̂k(s)||1 ≤ M∑ t=0 ∣∣∣∣∣∣π̂N 2t (s)− π̂ N 2t+1 (s) ∣∣∣∣∣∣ 1 + ∣∣∣∣∣∣π̂k(s)− π̂ N 2M+1 (s) ∣∣∣∣∣∣ 1
< M∑ t=0 ∣∣∣∣∣∣π̂N 2t (s)− π̂ N 2t+1 (s) ∣∣∣∣∣∣ 1 +
≤ 1 + − 1 2M ≤ + 1− r
(10)
For theN -th iteration of the search process, the final visit count distributions keep the same between the original expansion (Algorithm 1) and the virtual expansion (Algorithm 2). This is because at the last iteration, searching the nodes after the root has no effects on the final distribution. Therefore, π̂N (s) = πN (s). So we have
||π̂k(s)− πN (s)||1 < + 1− r (11)

Review summaries:
This paper proposes an approach for a significant speedup of Monte-Carlo Tree Search (MCTS) at a relatively small cost in playing strength. The basic idea is that, when the change in the distribution of visit counts between two different time points $\frac{k}{2}$ and $k$ is less than some constant $\epsilon$, it can also be shown that the remaining change in distributions between $k$ and $N$ (where $N$ is the maximum visit count allowed by some budget) will be bounded below some value, and if we consider such a maximum possible error to be sufficiently small we can just terminate the search early at time $k$. The paper also proposes a very simple but important idea called Virtual Expansions, which basically consists of running $N - k$ additional iterations of the bandit algorithm used by MCTS solely in the root node, providing the current estimated average values as rewards for every pull (i.e., leaving average reward estimates unchanged), in order to transform the distribution of visits at iteration $k$ into a prediction of the distribution we would end up with after $N$ iterations.  Several empirical evaluations in Go ($9\times9$ board) and a few Atari games, including ablation studies, demonstrate that the approach can substantially reduce the search times and self-play training times, while only decreasing playing strength by a small amount.

The paper proposes Virtual MCTS, an early-termination rule for MCTS to improve its efficiency. Roughly speaking, the termination rule prunes the search process when the final policy at the root node is unlikely to change by too much from the current policy. This strategy improves the efficiency of AlphaGoZero-style algorithms. Specifically, the authors showed that Virtual MCTS improves the learning efficiency on 9*9 Go.

This paper proposes the Virtual MCTS (V- MCTS), a variant of MCTS that mimics the human behavior, and is 50% more sample efficient, by performing a type of forward pruning. MCTS is characterized as a model-based reinforcement learning algorithm, that imagines what the future would look like, using terminology borrowed from Sutton in his description of Dyna, a model-based algorithm.  I agree with the planning part, but am less certain if I agree with the learning characterization of MCTS. Often, MCTS is used inside a model-based approach as the planning component. MCTS is usually not regarded as a full model-based approach.