Year: 2021

Abstract:
AlphaZero is a powerful reinforcement learning algorithm based on approximate policy iteration and tree search. However, AlphaZero can fail to improve its policy network, if not visiting all actions at the root of a search tree. To address this issue, we propose a policy improvement algorithm based on sampling actions without replacement. Furthermore, we use the idea of policy improvement to replace the more heuristic mechanisms by which AlphaZero selects and uses actions, both at root nodes and at non-root nodes. Our new algorithms, Gumbel AlphaZero and Gumbel MuZero, respectively without and with model-learning, match the state of the art on Go, chess, and Atari, and significantly improve prior performance when planning with few simulations.

heading: 1 INTRODUCTION
text: In 2018, AlphaZero (Silver et al., 2018) demonstrated a single algorithm achieving state-of-the-art results on Go, chess, and Shogi. The community reacted quickly. Leela Chess Zero (Linscott et al., 2018) was created to reproduce AlphaZero results on chess, winning Top Chess Engine Championship in 2019. Soon, all top-rated classical chess engines replaced traditional evaluations functions with Efficiently Updatable Neural Network (Nasu, 2018).
AlphaZero was itself generalized by MuZero (Schrittwieser et al., 2020). While AlphaZero requires a black-box model of the environment, MuZero learns an abstract model of the environment. Essentially, MuZero learns the rules of Go, chess, and Shogi from interactions with the environment. This allows MuZero to excel also at Atari and continuous control from pixels (Hubert et al., 2021).
In this work, we redesign and improve AlphaZero. In particular, we consider the mechanisms by which AlphaZero selects and uses actions, which are based upon a variety of heuristic ideas that have proven especially effective in Go, chess, and Atari (Silver et al., 2018; Schrittwieser et al., 2020). However when using a small number of simulations, some of AlphaZero’s mechanisms perform poorly. We use the principle of policy improvement to suggest new mechanisms with a better theoretical foundation. More specifically, we consider each mechanism in turn, alongside our proposed modifications:
• Selecting actions to search at the root node. To explore different actions during training, AlphaZero selects actions by adding Dirichlet noise to its policy network, and then performs a search using the perturbed policy. However, this does not ensure a policy improvement. We instead propose to sample actions without replacement by using the Gumbel-Top-k trick (Section 2), and perform a search using the same Gumbel values to influence the selection of the best action (Section 3.3), and show that this guarantees a policy improvement when action-values are correctly evaluated.
• Selecting actions at the root node. AlphaZero uses a variant of the PUCB algorithm (Rosin, 2011) to select actions at the root node. This algorithm was designed to optimize cumulative regret in a bandit-with-predictor setting (i.e. given prior recommendations from the policy network). However, no ancestors are dependent upon the evaluation of the root node, and the performance of the Monte-Carlo tree search therefore only depends upon the final recommended action at the root node, and not upon the intermediate actions selected during search (Bubeck et al., 2011). Consequently, we propose to use the Sequential Halving algorithm (Karnin et al., 2013) at the root node to optimize simple regret in a stochastic bandit with a predictor (Section 3.4).
• Selecting actions in the environment. Once search is complete, AlphaZero selects an action by sampling an (annealed) softmax distribution based upon the visit counts of root actions resulting from the search procedure. We instead propose to select the singular action resulting from the Sequential Halving search procedure. • Policy network update. AlphaZero updates its policy network towards a softmax distribu-
tion based upon the visit counts of root actions. However, even if the considered actions are correctly evaluated, this does not guarantee a policy improvement, especially when using small numbers of simulations (Grill et al., 2020). We instead propose a policy improvement based upon the root action values computed during search, and update the policy network towards that policy improvement (Section 4). • Selecting actions at non-root nodes. AlphaZero uses the PUCT algorithm to select actions
at non-root nodes. We instead propose to select actions according to a policy improvement (similar to the proposal of Grill et al. (2020)) based upon a completion of the action values. Furthermore, rather than sampling directly from this policy improvement, we propose a deterministic action selection procedure that matches the empirical visit counts to the desired policy improvement (Section 5).
The proposed modifications are applicable also to MuZero or any agent with a policy network and an expensive Q-network. The modifications are most helpful when using a small number of simulations, relative to the number of actions. When using a large number of simulations, AlphaZero works well. We tried to ensure that the new search is principled, better with a smaller number of simulations, and never worse. We succeeded on all tested domains: Go, chess, and Atari.

heading: 2 BACKGROUND
text: Before explaining the improved search, we will explain the Gumbel-Max trick and the Gumbel-Topk trick. The Gumbel-Max trick was popularized by Gumbel-Softmax for a gradient approximation. In this paper, we are not interested in approximate gradients. Instead, we use the Gumbel-Top-k trick to sample without replacement.
Gumbel-Max trick. (Gumbel, 1954; Luce, 1959; Maddison et al., 2017; Jang et al., 2017) Let π be a categorical distribution with logits ∈ Rk, such that logits(a) is the logit of the action a. We can obtain a sample A from the distribution π by first generating a vector of k Gumbel variables and then taking argmax:
(g ∈ Rk) ∼ Gumbel(0) (1) A = argmax
a (g(a) + logits(a)). (2)
Gumbel-Top-k trick. (Yellott, 1977; Vieira, 2014; Kool et al., 2019) The Gumbel-Max trick can be generalized to sampling n actions without replacement, by taking n top actions:
(g ∈ Rk) ∼ Gumbel(0) (3) A1 = argmax
a (g(a) + logits(a)) (4)
... An = argmax
a/∈{A1,...,An−1} (g(a) + logits(a)). (5)
We will denote the set of n top actions by argtop(g + logits, n) = {A1, A2, . . . , An}.

heading: 3 PLANNING AT THE ROOT
text: We are interested in improving AlphaZero Monte-Carlo Tree Search (MCTS). In this section we will focus on the action selection at the root of the search tree.

heading: 3.1 PROBLEM SETTING
text: Both AlphaZero and MuZero have access to a policy network. At the root of the search tree, they can explore n simulations, before selecting an action for the real environment. We will formalize the
problem as a deterministic bandit with a predictor and we will later extend it to a stochastic bandit and MCTS.
Bandit. A k-armed deterministic bandit is a vector of Q-values q ∈ Rk, such that q(a) is the Q-value of the action a. The agent interacts with the bandit in n simulations (aka rounds). In each simulation t ∈ {1, . . . , n}, the agent selects an action At ∈ {0, . . . , k − 1} and visits the action to observe the Q-value q(At).
The objective is to maximize the Q-value from a special last action An+1. That means we want to maximize E[q(An+1)]. This objective is equivalent to minimization of simple regret. The simple regret differs from the cumulative regret from all n simulations. Bubeck et al. (2011), Hay & Russell (2011), and Tolpin & Shimony (2012) already argued that at the root of the search tree we care about the simple regret.
The problem becomes interesting when the number of possible actions is larger than the number of simulations, i.e., when k > n. For example, 19x19 Go has 362 possible actions and we will do experiments with as few as n = 2 simulations. Fortunately, the policy network can help.
Predictor. In the bandit-with-predictor setting (Rosin, 2011), the agent is equipped with a predictor: the policy network. Before any interaction with the bandit, the policy network predicts the best action by producing a probability distribution π. The agent can use the policy network predictions to make more informed decisions.
Policy improvement. Naturally, we would like to have an agent that acts better than, or as well as, the policy network. We would like to obtain a policy improvement. If the agent’s action selection produces a policy improvement, then
E [q(An+1)] ≥ ∑ a π(a)q(a), (6)
where the probability π(a) is the policy network prediction for the action a.1 The policy network can then keep improving by modeling an improved policy.

heading: 3.2 MOTIVATING COUNTEREXAMPLE
text: We will show that the commonly used heuristics fail to produce a policy improvement.
Example 1. Acting with the best action from the top-n most probable actions fails to produce a policy improvement. Let’s demonstrate that. Let q = (0, 0, 1) be the Q-values and let π = (0.5, 0.3, 0.2) be the probabilities produced by the policy network. The value of the policy network is ∑ a π(a)q(a) = 0.2. For n = 2 simulations, the set of the most probable actions is {0, 1}. With that, the heuristic would select An+1 = argmaxa∈{0,1} q(a). The expected value of such action is E [q(An+1)] = 0, which is worse than the value of the policy network.
You can find other counterexamples by generating random q and π vectors and testing the policy improvement (Inequality 6). The AlphaZero action selection is explained in Appendix A.

heading: 3.3 PLANNING WITH GUMBEL
text: We will design a policy improvement algorithm for the deterministic bandit with a predictor π. After n simulations, the algorithm should propose an action An+1 with E [q(An+1)] ≥ ∑ a π(a)q(a).
One possibility is to sample n actions from π, and then to select from the sampled actions the action with the highest q(a). Instead of sampling with replacement, we can reduce the variance by sampling without replacement.
Still, the sampled actions contain a limited amount of information about π. We should exploit the knowledge of π and its logits when selecting An+1. The main idea is to sample n actions without replacement by using the Gumbel-Top-k trick, and then to use the same Gumbel g to select the action with the highest g(a) + logits(a) + σ(q(a)). The σ can be any monotonically increasing transformation. The pseudocode for the algorithm is in Algorithm 1.
1Inequality 6 can be strict, if we assume that an action has a positive advantage and its π(a) > 0.
Algorithm 1 Policy Improvement by Planning with Gumbel Require: k: number of actions. Require: n ≤ k: number of simulations. Require: logits ∈ Rk: predictor logits from a policy network π.
Sample k Gumbel variables: (g ∈ Rk) ∼ Gumbel(0) Find n actions with the highest g(a) + logits(a): Atopn = argtop(g + logits, n) Get q(a) for each a ∈ Atopn by visiting the actions. From the Atopn actions, find the action with the highest g(a) + logits(a) + σ(q(a)): An+1 = argmaxa∈Atopn(g(a) + logits(a) + σ(q(a))) return An+1
The algorithm produces a policy improvement, because q(argmax
a∈Atopn (g(a) + logits(a) + σ(q(a)))) ≥ q(argmax a∈Atopn (g(a) + logits(a))). (7)
This holds for any Gumbel g, so it holds also for expectations: E[q(An+1)] ≥ EA∼π[q(A)]. The argmaxa∈Atopn(g(a) + logits(a)) is equivalent to sampling from the policy network π (see the Gumbel-Max trick). By using the same Gumbel vector g in the argtop and argmax, we avoid a double-counting bias.
The prior knowledge contained in the logits can help on partially observable environments, or when working with approximate or stochastic Q-values.

heading: 3.4 PLANNING ON A STOCHASTIC BANDIT
text: We can now extend the algorithm to a stochastic bandit. A stochastic bandit provides only a stochastic estimate of the expected Q-value q(a). In that setting, we will use the empirical mean q̂(a) instead of q(a). Obviously, the empirical mean would be better estimated, if visiting an action multiple times. We have to choose which actions to visit and how many times. We can control this in two places. First, we can control the number of actions sampled without replacement. Second, we can use a bandit algorithm to efficiently explore the set of sampled actions.
There are multiple bandit algorithms for simple regret minimization. In our preliminary experiments, Sequential Halving (Karnin et al., 2013) was easier to tune than UCB-E (Audibert et al., 2010) and UCB √ · (Tolpin & Shimony, 2012). Conveniently, Sequential Halving does not have problemdependent hyperparameters.
We present Sequential Halving with Gumbel in Algorithm 2, with an illustration in Figure 1. Sequential Halving is used to identify the action with the highest g(a) + logits(a) + σ(q̂(a)).
Algorithm 2 Sequential Halving with Gumbel Require: k: number of actions. Require: m ≤ k: number of actions sampled without replacement. Require: n: number of simulations. Require: logits ∈ Rk: predictor logits from a policy network π.
Sample k Gumbel variables: (g ∈ Rk) ∼ Gumbel(0) Find m actions with the highest g(a) + logits(a): Atopm = argtop(g + logits,m) Use Sequential Halving with n simulations to identify the best action from the Atopm actions, by comparing g(a) + logits(a) + σ(q̂(a)). An+1 = argmaxa∈Remaining(g(a) + logits(a) + σ(q̂(a))) return An+1
For a concrete instantiation of σ, we use σ(q̂(a)) = (cvisit +max
b N(b))cscaleq̂(a), (8)
where maxbN(b) is the visit count of the most visited action. The transformation progressively increases the scale for q̂(a) and reduces the effect of the prior policy. This scale is inspired by the policy updates in MPO (Abdolmaleki et al., 2018; Vieillard et al., 2020). The finite scale for the Q-values provides regularized policy optimization and allows to use the prior knowledge contained in the logits. Experimentally, cvisit = 50, cscale = 1.0 allowed us to use the same hyperparameters even if changing the number of simulations.

heading: 4 LEARNING AN IMPROVED POLICY
text: After the search, we have An+1 from a (potentially) improved policy. Like AlphaZero, we would like to distill the improved policy to the policy network. One possibility is to train the policy network π to predict the An+1. That defines a simple policy loss:
Lsimple(π) = − log π(An+1). (9)
Using completed Q-values. We will explain a different way to train the policy network, by extracting more knowledge from the search. Beside An+1, the search also gives us q(a) (or its approximation) for the visited actions. We can construct an improved policy by first completing the vector of Q-values:
completedQ(a) = { q(a) if N(a) > 0 vπ, otherwise,
(10)
where the unknown Q-values of the unvisited actions are replaced by vπ = ∑ a π(a)q(a). While in practice we do not have the exact vπ , we have instead its approximation v̂π from a value network. Even when training on off-policy data we devised a helpful vπ approximation (Appendix D).
With the completed Q-values, a new improved policy is constructed by
π′ = softmax(logits+σ(completedQ)), (11)
where σ is a monotonically increasing transformation. We provide a proof of a policy improvement in Appendix C. Intuitively, the completed Q-values give zero advantage to the unvisited actions.
After constructing the new improved policy π′, we can distill it to the policy network π:
Lcompleted(π) = KL(π ′, π). (12)
This loss trains all actions, not only the action An+1. Later, we will investigate the effect of the loss in Figure 3a.

heading: 5 PLANNING AT NON-ROOT NODES
text: To design an action selection for the non-root nodes of a search tree, we take an inspiration from Grill et al. (2020). That allows us to interpret MCTS as regularized policy optimization. At a non-root node, we construct an improved policy π′ by using the completed Q-values (Equation 11).
To select an action at the non-root node, one possibility is to sample the action from π′. However, sampling at non-root nodes adds unwanted variance to the estimated Q-values. Instead, we can design a deterministic action selection with the smallest mean-squared-error between the π′ probabilities and the produced normalized visit counts. Such action selection would select
argmin a ∑ b ( π′(b)− N(b) + I{a = b} 1 + ∑ cN(c)︸ ︷︷ ︸
Normalized visit counts, if taking a.
)2 , (13)
where the indicator I{a = b} is 1 if a = b, and zero otherwise. After a bit of algebra (Appendix E), we obtain a simpler, more efficient expression:
argmax a
( π′(a)− N(a)
1 + ∑ bN(b)
) . (14)
This deterministic action selection selects the actions proportionally to π′ and avoids an extra variance. We recommend a deterministic action selection only for non-root nodes. At the root node, the Gumbel noise is helpful for trying different actions in different episodes, while ensuring an improved expected value.

heading: 6 RELATED WORK
text: Rosin (2011) introduced the bandit with a predictor and designed PUCB (“Predictor + UCB”) for cumulative regret minimization. AlphaGo (Silver et al., 2016), AlphaGo Zero (Silver et al., 2017), AlphaZero (Silver et al., 2018), and MuZero (Schrittwieser et al., 2020) used a deep policy network as the predictor in a variant of the PUCB algorithm. Bertsekas (2019; 2021; 2022) provides an in-depth discussion of policy iteration, policy improvement, and their connection to rollout.
If not using a predictor, UCT (Kocsis & Szepesvári, 2006) would need to visit each action before being able to compare them. Rapid Action Value Estimation (Gelly & Silver, 2011) then helps to form rough estimates of the action values by aggregating statistics from all future states. Gelly & Silver (2011) also initialized the action value estimates with a heuristic evaluation function. The best heuristic used a learned linear network. Hamrick et al. (2020) later extended it to a deep Q-network.
Cazenave (2014) and Pepels et al. (2014) applied Sequential Halving to MCTS. Fabiano & Cazenave (2021) introduced Sequential Halving Using Scores. The ‘scores’ can be any prior offset to the Qvalues. The way to obtain scores from a policy network was left as an open problem. We can now view the g + logits as special scores.
MCTS is related to regularized policy optimization. Grill et al. (2020) analyzed AlphaZero tree search and discovered that AlphaZero approximates a regularized policy optimization. The approximation error is large if using a small number of simulations. To avoid the approximation error, Grill et al. (2020) used a regularized policy optimization directly inside the tree search. In the setting without a predictor, Xiao et al. (2019) compared UCT to a new MCTS with an entropy regularizer. Dam et al. (2021) generalized it to relative entropy and Tsallis entropy. Regularized policy optimization is helpful when working with approximate Q-values (Vieillard et al., 2020) or when doing an approximate policy iteration (Kakade & Langford, 2002; Schulman et al., 2015).
TreeQN (Farquhar et al., 2018) uses a breadth-first search inside a network architecture. The network can do a lot of computation before producing a Q-value. To reduce the computation demands, Dynamic Planning Networks (Tasfi & Capretz, 2018) extended TreeQN to sample only some actions. To approximate the gradient, Dynamic Planning Networks use Gumbel-Softmax (Maddison et al., 2017; Jang et al., 2017). Although we use Gumbel variables, we do not employ approximate gradients from Gumbel-Softmax. We use the Gumbel-Top-k trick to construct efficient planning with a provable policy improvement.
For sampling without replacement, the unordered set estimator by Kool et al. (2020) provides an elegant, unbiased estimate of a gradient. However, gradient descent needs multiple steps to reach the solution of a regularized policy optimization problem (Tomar et al., 2020). Furthermore, the exact computation of the unordered set estimator requiresO(2m) operations, which would be prohibitively expensive for m = 16. Practical time complexity can be achieved by using an importance-weighted estimator (Vieira, 2017; Nauman & Den Hengst, 2020).
Continuous actions can be supported by sampling k actions with replacement and then using the sampled actions as discrete actions with uniform logits for the rest of the search. This was done in Sampled MuZero (Hubert et al., 2021). Similarly, Critic Weighted Policy (Wang et al., 2020) uses sampling with replacement.

heading: 7 EXPERIMENTS
text: In the experiments, we compare AlphaZero or MuZero to the proposed planning with Gumbel and other alternatives:
MuZero: The newest version of MuZero (Schrittwieser et al., 2021), with ResNet v2 style preactivation residual blocks (He et al., 2016) and the Adam optimizer (Kingma & Ba, 2014). Gumbel MuZero: MuZero with the modified root of the search tree to use Sequential Halving with Gumbel. The policy loss uses the completed Q-values (Equation 12). Gumbel MuZero sampled with replacement (Replacement): An ablation to Gumbel MuZero by sampling m actions with replacement, as in Sampled MuZero (Hubert et al., 2021). TRPO MuZero: MuZero with modified learning, acting, and the root of the search tree to use the regularized policy optimization with the TRPO regularizer KL(π, πnew) (Schulman et al., 2015; Grill et al., 2020). MPO MuZero: TRPO MuZero but with the MPO regularizer KL(πnew, π) (Abdolmaleki et al., 2018; Grill et al., 2020). Full Gumbel MuZero: Gumbel MuZero with a principled action selection also for the non-root search nodes (Section 5). In the plots, we will show Full Gumbel MuZero only if it produces results significantly different from Gumbel MuZero.
We conducted the experiments on Go, chess, and Atari. We present the main results here and we report additional ablations and experimental details in Appendix F.

heading: 7.1 9X9 GO
text: On Go, we use Elo to compare MuZero and other agents. While an agent trains by self-play, its Elo is computed by evaluation versus reference opponents. One of the opponents is Pachi (Baudiš & Gailly, 2011) with 10k simulations per move. We anchored the Elo to have this Pachi at 1000 Elo. For example, a difference of 500 Elo corresponds to a 95% win probability for the player with the higher Elo.2
2The corresponding win probability is 1 1+10 − EloDifference 400 (Elo, 1978).
In Figure 2 we investigate the impact of the number of simulations on the obtained Elo. When training an agent by self-play, the agent uses n simulations per move. In the five plots, the n varies from 2 to 200. In evaluation, we allow all agents to use 800 simulations. The speed of the evaluation does not affect the speed of training. In the 9x9 Go results, MuZero fails to learn from 16 or less simulations. Strikingly, Gumbel MuZero learns reliably even with 2 simulations.
In Figure 3a we compare the simple policy loss (Equation 9) and the policy loss with the completed Q-values (Equation 12). The simple policy loss would be enough for many applications. We used the completed Q-values also in TRPO MuZero and MPO MuZero. Without the completed Q-values, TRPO MuZero and MPO MuZero would fail to produce a policy improvement.
In Figure 3b we study Gumbel MuZero’s sensitivity to the number of sampled actions. When sampling m = 4 actions without replacement, the simulation budget is spent on the small number of actions. The learning was then slower. In all other Go experiments, we sample m = min(n, 16) actions without replacement.

heading: 7.2 LARGE-SCALE 19X19 GO AND CHESS
text: In Figure 4a we demonstrate that Gumbel MuZero is not worse than MuZero on 19x19 Go. MuZero is excellent on 19x19 Go and Gumbel MuZero reaches or exceeds its performance. The Elo is still anchored to have Pachi at 1000 Elo.
Similarly, in Figure 4b we show Gumbel AlphaZero performance on chess. We train AlphaZero on chess, because AlphaZero learns faster than MuZero on chess. On Go, MuZero learns faster than AlphaZero.

heading: 7.3 ATARI
text: Our last set of experiments is on Atari. We use the Arcade Learning Environment (Bellemare et al., 2013) with sticky actions (Machado et al., 2018). The network sizes and hyperparameters match the MuZero setup by Schrittwieser et al. (2021). On Atari, MuZero does not use more simulations at evaluation. The reported score is the mean return from the last 200 training episodes. MuZero with n = 50 simulations works well on Atari and establishes the state of the art (Schrittwieser et al., 2021).
In Figure 5a we show the obtained mean return on ms pacman. Gumbel MuZero again learns reliable even with n = 2 simulations. Atari has only 18 actions, so we sample m = min(n, 18) actions without replacement. In the experiments with n ≤ 18, Gumbel MuZero selects An+1 from the n visited actions, without any Sequential Halving. This confirms that planning with Gumbel is the key ingredient responsible for the policy improvement from a small number of simulations.
Atari is challenging, because different games can have very different reward scales. MuZero normalizes the Q-values by dividing them by max(v̂π,maxa q̂(a))−min(v̂π,mina q̂(a)) found inside the tree search (Schrittwieser et al., 2020). A normalized advantage is then in [−1, 1]. For Gumbel MuZero, we use the same normalization and we scale the normalized Q-values by cvisit = 50 and cscale = 0.1. A scaled normalized advantage is then approximately in [−5, 5]. Thanks to the bounded advantage, Gumbel MuZero has a bounded total variation distance between π and π′ (Hessel et al., 2021).
In Figure 5b we use beam rider as an example of a partially observable game and we study the importance of the prior knowledge contained in the logits. Gumbel MuZero selects an action based on g(a) + logits(a) + (cvisit +maxbN(b))cscaleq̂(a) (Equation 8). If cscale is large, Gumbel MuZero focuses on q̂(a) and neglects the logits. Indeed, Gumbel MuZero performance is worse on beam rider if using large cscale. In the future, we can try normalizing the Q-values by the standard deviation of an advantage estimator and we can try clipping the normalized advantages, as in Muesli (Hessel et al., 2021).

heading: 8 CONCLUSION
text: We redesigned AlphaZero tree search. With the principle of policy improvement, we replaced five heuristic mechanisms in AlphaZero. On Go, chess and Atari, we validated that Gumbel MuZero and Gumbel AlphaZero keep improving, even when learning from two simulations. On top of that, Gumbel MuZero provides a principled way to achieve state-of-the-art results. We hope that future research will benefit from both the clean theoretical foundation and the faster experimentation with a small number of simulations.

heading: A ALPHAZERO ACTION SELECTION
text: We will explain AlphaZero’s (Silver et al., 2018) action selection here. On the deterministic bandit, AlphaZero would select actions A1, . . . , An by:
argmax a
[ q?(a) + c1π(a) √ 1 + ∑ bN(b)
1 +N(a)
] , (15)
where q?(a) ∈ [0, 1] is q(a) for the already visited actions and zero otherwise, and c1 a factor independent of a. After the search, AlphaZero acts with with the most visited action. That means An+1 = argmaxaN(a).
At the root of the search tree, AlphaZero perturbs π by Dirichlet noise to avoid visiting always the most probable actions. That does not ensure a policy improvement, because AlphaZero with Dirichlet noise can produce a worse policy by adding noise to a potentially optimal policy network.

heading: B POLICY IMPROVEMENT PROOF FOR PLANNING WITH GUMBEL
text: We will prove that Algorithm 1 generates An+1 such that E[q(An+1)] ≥ EA∼π[q(A)]. For the right-hand side, the Gumbel-Max trick tells us that EA∼π[q(A)] is equal to E(g∈Rk)∼Gumbel(0)[q(argmaxa(g(a)+ logits(a))]. We will show that we can replace the argmaxa with argmaxa∈Atopn . Remember that Atopn is defined as Atopn = argtop(g + logits, n) and that we use the same Gumbel vector g in the argtop and argmax. The set Atopn then includes the action with the highest g(a) + logits(a) and we can replace the argmaxa with argmaxa∈Atopn .
After these rewrites, we have to prove that
E[q(An+1)] ≥ E(g∈Rk)∼Gumbel(0)[q(argmax a∈Atopn (g(a) + logits(a))]. (16)
On the left-hand side, E[q(An+1)] is equal to
E(g∈Rk)∼Gumbel(0)[q(argmax a∈Atopn (g(a) + logits(a) + σ(q(a))))]. (17)
We can finish the proof by proving that for any vector g ∈ Rk we have
q(argmax a∈Atopn (g(a) + logits(a) + σ(q(a)))) ≥ q(argmax a∈Atopn (g(a) + logits(a)). (18)
This is true, because σ is a monotonically increasing transformation.

heading: C POLICY IMPROVEMENT PROOF FOR COMPLETED Q-VALUES
text: We will prove that π′completed = softmax(logits+σ(completedQ)) produces a policy improvement. We will start by showing that π′completed is produced by a specific instance of Algorithm 3. We will then prove that any instance of Algorithm 3 produces a policy improvement.
Algorithm 3 Policy improvement when having vπ Require: π, vπ . Require: q(a) for each visited action. The visited actions can be from any distribution.
Initialize π′ with π. For the visited actions:
If q(a) > vπ , increase the π′(a) logit. If q(a) < vπ , decrease the π′(a) logit.
return π′
C.1 SPECIFIC INSTANCE
Algorithm 3 is more general than the usage of the completed Q-values. Specifically, Algorithm 3 would produce π′completed, if updating the logits by σ(completedQ(a)) − σ(vπ). This update increases the logit, if q(a) > vπ . This update decreases the logit, if q(a) < vπ . And the update does not modify the logits of the unvisited actions. The resulting softmax(logits+σ(completedQ) − σ(vπ)) is equal to softmax(logits+σ(completedQ)), because the constant offset σ(vπ) does not change the softmax output.
C.2 POLICY IMPROVEMENT PROOF FOR ANY INSTANCE
We will now prove that π′ from Algorithm 3 satisfies∑ a π′(a)q(a) ≥ ∑ a π(a)q(a). (19)
Notice that π′(a) for any unvisited action a will be czπ(a), with a normalization constant cz > 0.
For one visited action: Let’s denote the visited (aka expanded) action by aex. First, if π(aex) = 1 then vπ = q(aex) and the policy will remain unchanged.
Let’s now consider the case with π(aex) < 1. The vπ can be rewritten as: vπ = π(aex)q(aex) + (1− π(aex)) ∑ a 6=aex π(a)q(a)∑ b6=aex π(b) . (20)
Let’s denote the weighted sum by qmiss: qmiss = ∑ a 6=aex π(a)q(a)∑ b6=aex π(b) . (21)
We notice that the qmiss does not change if scaling π by a constant cz > 0.
We will now rewrite the left-hand side of Inequality 19 to use qmiss:∑ a π′(a)q(a) = (22)
= π′(aex)q(aex) + (1− π′(aex)) ∑ a6=aex czπ(a)q(a)∑ b 6=aex czπ(b)
(23)
= π′(aex)q(aex) + (1− π′(aex))qmiss (24) = π′(aex)(q(aex)− qmiss) + qmiss. (25)
The right-hand side of Inequality 19 can be also rewritten:
vπ = π(aex)q(aex) + (1− π(aex))qmiss (26) = π(aex)(q(aex)− qmiss) + qmiss. (27)
With these rewrites, Inequality 19 becomes
π′(aex)(q(aex)− qmiss) ≥ π(aex)(q(aex)− qmiss). (28)
The q(aex)− qmiss can negative, zero or positive. If q(aex) = qmiss, the inequality is satisfied. If q(aex) > qmiss, we want π′(aex) ≥ π(aex). If q(aex) < qmiss, we want π′(aex) ≤ π(aex).
We do not know qmiss so we cannot use it in an algorithm. We will instead show that q(aex) > qmiss is equivalent to q(aex) > vπ , when π(aex) < 1:
q(aex) > vπ (29) q(aex) > π(aex)q(aex) + (1− π(aex))qmiss (30)
(1− π(aex))q(aex) > (1− π(aex))qmiss (31) q(aex) > qmiss. (32)
So we directly arrived at Algorithm 3.
For multiple visited actions: We will focus on one of the visited actions. If the logits of the other visited actions are unmodified, the algorithm is equivalent to using only one visited action. If the logits of the other visited actions are modified by Algorithm 3, they can further help to improve the policy.

heading: D MIXED VALUE APPROXIMATION
text: We will construct an approximation of vπ . The exact vπ is defined by vπ = ∑ a π(a)q(a). We have an approximate v̂π from a value network, we know π, and we have q(a) for the visited actions. With
these inputs, we approximate vπ by a consistent estimator:
vmix = 1 1 + ∑ bN(b) v̂π + ∑bN(b)∑ b∈{b:N(b)>0} π(b) ∑ a∈{a:N(a)>0} π(a)q(a)  . (33) The estimator interpolates v̂π and the weighted average of the available Q-values. This is an unsophisticated estimator, with results in Figure 6. You are welcome to explore other possibilities.

heading: E DERIVATION OF THE DETERMINISTIC ACTION SELECTION
text: We will derive Equation 14 from Equation 13:
argmin a ∑ b ( π′(b)− N(b) + I{a = b} 1 + ∑ cN(c) )2 (34)
=argmin a ∑ b (( π′(b)− N(b) 1 + ∑ cN(c) ) − I{a = b} 1 + ∑ cN(c) )2 (35)
=argmin a ∑ b −2 ( π′(b)− N(b) 1 + ∑ cN(c) ) I{a = b} 1 + ∑ cN(c)
(36)
=argmin a − ∑ b ( π′(b)− N(b) 1 + ∑ cN(c) ) I{a = b} (37)
=argmax a
( π′(a)− N(a)
1 + ∑ cN(c)
) . (38)
The simplification was possible, because additive terms independent of a do not affect argmina. While widely applicable, the deterministic action selection provides only a small benefit on 9x9 Go (Figure 7).

heading: F EXPERIMENTAL DETAILS
text: In general, we use hyperparameters consistent with the newest MuZero experiments (Schrittwieser et al., 2021). MuZero’s pseudocode is available thanks to Schrittwieser et al. (2020). Gumbel MuZero does not need to set the Dirichlet noise hyperparameters, because Gumbel MuZero does not use Dirichlet noise.
In a tree search, the Q-values q̂(a) are provided by the visited child nodes. We do not modify the structure of AlphaZero’s search tree. Inside the tree search, the Q-values are normalized to be from
the [0, 1] interval. We use the normalized Q-values also in Full Gumbel MuZero, but the algorithm does not require Q-values from a specific interval. In all Go and chess experiments, Gumbel MuZero scales the Q-values by cvisit = 50 and cscale = 1.0. On the perfect-information game of Go, Gumbel MuZero is not very sensitive to the scale of the Q-values. Any cvisit ≥ 50 produced similar results (Figure 8a).
In each phase of Sequential Halving, we use at least one new visit. For example, in the first phase, we update q̂(a) by max ( 1, ⌊
n dlog2(m)em
⌋) visits. This allows us to experiment with an incomplete
or no Sequential Halving. When Sequential Halving runs out of the budget of n simulations, we stop the search. The agent then selects as An+1 the action with the highest g(a) + logits(a) + σ(q̂(a)) from the set of the most visited actions. Fabiano & Cazenave (2021) provide a different way to deal with the rounding in Sequential Halving.
During training, MuZero acts with explorative actions in the first 30 moves of each self-play game. MuZero samples the explorative actions proportionally to the visit counts, like AlphaGo Zero (Silver et al., 2017). Gumel MuZero benefits from the same exploration (Figure 8b).
Figure 9 shows the importance of the number of simulation at evaluation time. For example, in the first subplot a network is trained with n = 2 simulations and the same network is evaluated with 800, 200, 32, 16, 2, and 1 simulations.
To run the experiments, we used Google Cloud Tensor Processing Units v3 (TPUs). On 9x9 Go, MuZero is not limited by lack of data if using 3-times more TPUs for self-play than for training. By using a smaller number of simulations, we can substantially reduce the number of TPUs needed for self-play. Table1 lists the obtained speedups if not being limited by the TPUs for training.
F.1 NETWORK ARCHITECTURE
In all Go and chess experiments, we used a modified version of the AlphaZero network architecture that is roughly 2-times faster at the same accuracy. Concretely, we replace the dense residual blocks
with bottleneck blocks (He et al., 2016), and replace every 8th block with a global broadcasting residual block (Figures 11, 13). The broadcasting block is similar to global pooling or squeeze-andexcitation (Hu et al., 2018), but we found it to be more stable when training from self-play.
Figure 10 provides data to support the architecture change. For each network architecture the figure shows multiple data points, differing by the number of layers. For dense networks, the number of layers is from {12, 16, 20, 24, 28, 32, 36, 40, 48}. For bottleneck networks, data points with 56 and 64 layers are shown as well.
Both value MSE and policy accuracy improve consistently with larger but slower networks. For any constant number of inferences per second, the bottleneck blocks and the broadcast/pooling blocks lead to significant improvements.
For the architecture search, the Elo was computed at a fixed computational budget of 300ms per move. For example, the budget allows to do 3200 simulations when using a network with 256 planes and 32 blocks with bottlenecks and broadcasting (achieving 104 inferences/second). Here the Elo is anchored to have the base AlphaZero architecture at 0 Elo; 95% confidence intervals are about 30 Elo wide. We observe that diminishing returns in the prediction accuracy for very large networks lead to decreasing Elo, as the improved predictions fail to compensate for the much lower number of simulations in MCTS. Gumbel AlphaZero was not yet used for the architecture search.
Given the efficient network architecture, we used a small 6 layer network on 9x9 Go, and a bigger 32 layer network in the large-scale 19x19 Go experiments. The networks used 256 hidden planes, 128 bottleneck planes and a broadcasting block in every 8th layer.

Review summaries:
The paper proposes a number of principled algorithmic modifications to state-of-the-art planning algorithms (AlphaZero, MuZero) for improving performance in settings with many actions and a relatively small computation and / or sample budget. The main contributions are algorithmic and empirical. The key ideas include the use of the Gumbel-max and top-k tricks along with the use of sequential halving to improve online planning. The paper also proposes a planning-learning loop wherein a policy using the estimated (completed) Q values is learned as well as a different selection policy at non-root nodes. The experiments show that the Gumbel variants of AlphaZero and MuZero perform well in low search budget settings in the domains of Go, Chess and Atari.

This paper considers MCTS with learned search guidance, as in AlphaZero, MuZero, etc. The work proposes several adjustments to the prior works, particularly regarding the way in which actions are selected at the root and non-root nodes, at training and during evaluation; and also the way in which the policy is updated after search. In a simplified bandits setting, the authors point out that the previous method of performing policy updates is not guaranteed to result in a policy improvement, and they propose using a Gumbel reparameterization trick to overcome this limitation. Experiments in Go, Chess, and Atari show that the adjustments are beneficial in the regime of low simulation count.