Policy improvement by planning with Gumbel
Authors: Ivo Danihelka
Arthur Guez
Julian Schrittwieser
David Silver
Venue: ICLR 2022 Spotlight (2022)
Keywords: AlphaZero
MuZero
reinforcement learning

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.

Published as a conference paper at ICLR 2022
POLICY IMPROVEMENT BY PLANNING WITH GUMBEL
Ivo Danihelka' 2, Arthur Guez', Julian Schrittwieser", David Silver 'DeepMind. London, UK Uniivrriity College London
banihelka@google.com com
ABSTRACT
AlphaZero is 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 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.
INTRODUCTION
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 al.) 2018) was created to reproduce AlphaZero results on chess, winning Top Chess Engine Champi- onship 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., 2020p. While AlphaZero requires black-box model of the environment, MuZero learns an abstract model of the environment. Essen- tially, 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 variety of heuristic ideas that have proven especially effective in Go, chess, and Atari (Silver et al. 2018; [chrittwieser et al. 2020). However when using 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, consider each mechanism in turn, alongside our proposed modifications:
Selecting actions to search at the root node. To explore different actions during train- ing, AlphaZero selects actions by adding Dirichlet noise to its policy network, and then performs search using the perturbed policy. However, this does not ensure policy improvement. We instead propose to sample actions without replacement by using the Gumbel-Top-k trick (Section 2), and perform search using the same Gumbel values to in- fluence 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) 201T) to select actions at the root node. This algorithm was designed to optimize cumulative regret in a bandit- Shundi--with-peediie 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 Halv- ing algorithm (Karnin et al., 2013) at the root node to optimize simple regret a stochastic bandit with predictor (Section 3.4

---

Published as a conference paper at ICLR 2022
Selecting actions in the environment. Once search is complete, AlphaZero selects tion by sampling from an (annealed) categorical 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 categorical dis- tribution 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. (2200) We instead propose policy improvement based upon the root action values computed during search, and update the policy network towards that policy improvement (Section
Selecting actions at non-root nodes. AlphaZero uses the PUCT algorithm to select actions at non-root nodes. We instead propose select actions according policy improvement (similar to the proposal Grill et al. (2020)) based upon a completion of the action values. Furthermore, rather than sampling directly from this policy improvement, we propose de- terministic action selection procedure that matches the empirical visit counts to the desired policy improvement (Section 5p.
The proposed modifications are applicable also to MuZero or any agent with a policy network and an expensive Q-network. The modifications are helpful when using small number of sim- ulations, relative to the number of actions. When using a large number of simulations, AlphaZero works well. We tried ensure that the new search is principled, better with smaller number of simulations, and never worse. We succeeded on all tested domains: Go, chess, and Atari.
2 BACKGROUND
Before explaining the improved search, we will explain the Gumbel-Max trick and the Gumbel-Top- k trick. The Gumbel-Max trick was popularized by Gumbel-Softma for gradient approximation. In this paper. we are not interested in approximate gradients. Instead, we use the Gumbel-Tou Top-k trick to sample without replacement.
Gumbel-Man trick. (Gumbel 1954] Luce 1959
Maddison et al. 2017 Jang et al. 2017)
Let be a categorical distribution with logits E R", such that logits(a) is the logit of the action a. We can obtain a sample A from the distribution π by first generating vector of lo Gumbel variables and then taking argmax:
(g E R Gumbel
A = arg max(g(a) logits(a)).
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 top actions:
(g E h Gumbel(0)
(3)
A arg max(g(() + logits(a))
(4)
An arg max (g(a) + logits(a)).
(5)
We will denote the set of n top actions by argtop(g + logits, n) = {A1, Ag...,, ... Ar}.
3 PLANNING AT THE ROOT
We are interested in improving AlphaZero Monte-Carla Tree Search (MCTS). In this section we will focus the action selection at the root of the search tree.
3.1 PROBLEM SETTING
Both AlphaZero and MuZero have access to a policy network. At the root of the search tree, they explore simulations, before selecting action for the real environment. We will formalize the
2

---

Published as a conference paper at ICLR 2022
problem deterministic bandit with a predictor and we will later extend it to stochastic bandit and MCTS.
Bandit. A k-armed deterministic bandit is vector of Q-values g R such that g(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, (1...., n,,, , the agent selects action A, -0,....,k k 1) and visits the action observe the Q-value (A.).
The objective is to maximize the Q-value from special last action A,+ That means we want to maximize E[g(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. (2.1 1. [Hay & Russell (201T) 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- prediiccorrrrrorrrrrroo setting (Rosin) 201T), the agent is equipped with 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 policy improvement, then
[g(An+1)] 2 Σ T(a)q(a),
(6)
where the probability <(a) is the policy network prediction for the action a The policy network can then keep improving by modeling an improved policy.
3.2 MOTIVATING COUNTEREXAMPLE
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 produce a policy improvement. Let's demonstrate that. Let = (0, O, 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)g(a) = 0.2. For = 2 simulations, the set of the most probable actions is (0. 1). With that, the heuristic would select A,-1 arg max (0.11 1 g(a). The expected value of such action (n)1l = 0, which is worse than the value of the policy network.
You can find other counterexamples by generating random q and x vectors and testing the policy improvement (Inequality 6). The AlphaZero action selection is explained in Appendix A
3.3 PLANNING WITH GUMBEL
We will design policy improvement algorithm for the deterministic bandit with predictor æ. After
n simulations, the algorithm should propose an action An+1 with [c(An+1)] 2 Σ. (a))(())
One possibility is to sample actions from and then to select from the sampled actions the action with the highest g(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 sample actions without replacement by using the Gumbel- Cumbel-Top-k trick, and then to use the same Gumbel g to select the action with the highest g(a) + logits(a) + (g(a)). The a can be any monotonically increasing transformation. The pseudocode for the algorithm is in Algorithm I
Inequality 6 can be strict, if we assume that an action has positive advantage and its (a) > 0.
3

---

Published as a conference paper at ICLR 2022
Algorithm Policy Improvement by Planning with Gumbel
Require: A: number of actions.
Require: n S k: number of simulations.
Require:
logits R predictor logits from policy network
Sample &s Gumbel variables:
[g R² (umbel(0)
Find n actions with the highest g(a) + logits(a):
Atopn a + logits,
Get q(a) for each a E Asopn by visiting the actions. From the Atopn actions, find the action with the highest g(a) + logits(a) ((g(a)):
An+1 arg aa (g(a) + logits(a) + (I(a)))
return An+1
The algorithm produces a policy improvement, because q( (arg (ax(g(a) logits(a) a(((a)))) 2 4(arg max(g(a) + logits(a))). eEAtopn
aEAAcop
This holds for any Gumbel g. so it holds also for expectations: E[g(An+1)] 2 EAw (A)]. The arg maxge-Atopn! (g(a) + logits(a) (a)) is equivalent sampling from the policy network x (see the Gumbel-Max trick or Appendix B By using the same Gumbel vector g in the argtop and arg max, we avoid double-come counticn bias.
The prior knowledge contained in the logits can help on partially observable environments, or when working with approximate or stochastic Q-values.
3.4 PLANNING ON STOCIASTIC BANDIT
We can now extend the algorithm to stochastic bandit. A stochastic bandit provides only stochas- tic estimate of the expected Q-value q().. In that setting, we will use the empirical mean ((a) in- stead of g(a). Obviously, the empirical mean would be better estimated, if visiting an action multiple times. We have to choose which actions visit and how many times. We can control this two places. First, we 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 problem-
dependent hyperparameters.
We present Sequential Halving with Gumbel in Algorithm 2 with illustration in Figure ! Se- quential Halving is used to identify the action with the highest g(a) + logits(a) + a(((a))
Algorithm 2 Sequential Halving with Gumbel
Require: A: number of actions.
Require: m S k: number of actions sampled without replacement.
Require: number of simulations.
Require: logits e Rt; predictor logits from policy network т.
Sample & Gumbel variables:
(g E RR*) ~ Gumbel(O)
Find actions with the highest g(a) + logits(a):
Atopm argtop(g + logits, m)
Use Sequential Halving with simulations to identify the best action from the Aropm actions,
by comparing g(a) + logits(a) (a) + (ğ(a)).
An+1 arg max,e Remaining (g(a) lag + (Q(a)))
return An+1
For concrete instantiation of a, use
a (ÿ(a)) = (Cvisit + max N ((b))
(8)
4

---

Published as a conference paper at ICLR 2022
(x E (f) ~ Gumbel
logits
g + logits o(i)
top
top
top
top
1 action: Aot]
actions
top n
actions
N(a) = 21 + 28
((a) 9 12
actions
N(a) 3
extra visits to reach simulations
m = 16 actions
(a) = 3
extra
visits update
visits update
logg (MI1
actions
200 simulations
Figure 1:
The number of considered actions and their visit counts N(),, when using Sequential Halving with Gumbel
k-armed stochastic bandit. The search uses 200 simulations and starts by sampling m = 16 actions without replacement. Sequential Halving divides the budget of simulations equally to (og)(m) phases. In each phase, all considered actions visited equally often. After each phase, one half of the actions is rejected. From the original i actions, only the best actions will remain.
where max; (b)) is the visit count of the most visited action. The transformation progressively increases the scale for @(a) and reduces the effect of the prior policy. This scale is inspired by the policy updates in MPO (Abdolmaleki et al. 22118) Vieillard et al. (020). The finite scale for the Q- values provides regularized policy optimization and puts into effect the prior knowledge contained in the logits. Experimentallyy Cvisit = 50, Ccaae = 1.0 allowed us to use the same hyperparameters even if changing the number of simulations.
LEARNING AN IMPROVED POLICY
After the search, we have An+1 from (potentially) improved policy. Like AlphaZero, 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 simple policy loss:
Lample (T) log म (A)+1).
Using completed Q-values. We will explain different way to train the policy network, by extract- ing more knowledge from the search. Beside Ant1. the search also gives us q(a) (or its approxima- tion) for the visited actions. We can construct an improved policy by first completing the vector of Q-values:
((aa) Ury otherwise.
N(a) 0
completedQ(a) =
(10)
where the unknown Q-values of the unvisited actions are replaced by U = (a))((a) While in practice we do not have the exact Uz, we have instead its approximation éं from value network. Even when training on off-policy data we devised a helpful ex approximation (Appendix D).
With the completed Q-values, new improved policy is constructed by
T' = softmax (logits +o (completted))
(11)
where a is monotonically increasing transformation. We provide proof of a policy improvement Intuitively, the completed Q-values give zero advantage the unvisited actions.
in Appendix C
After constructing the new improved policy n'. we can distill it to the policy network n:
Tecomlletaa (T) KL .))
(12)
This loss trains all actions, not only the action A Later, we will investigate the effect the loss in Figure
5

---

Published as a conference paper at ICLR 2022
5 PLANNING AT NON-ROOT NODES
To design an action selection for the non-root nodes of a search tree, we take inspiration from Grill et al. (2020). That allows to interpret MCTS regularized policy optimization. At non-root node, we construct an improved policy T' by using the completed Q-values (Equation 11
To select an action at the non-root node, one possibility to sample the action from V'. How- ever, sampling at non-root nodes adds unwanted variance to the estimated Q-values. Instead, we can design deterministic action selection with the smallest mean- rrean-squaedderron between the x' probabilities and the produced normalized visit counts. Such action selection would select
N(b) + Ifa + b} )
arg min Σ (b(6)-
(13)
1+ 1. (c)
Nonnalized visit counts, luking
where the indicator T(a = b} is if a = b, and zero otherwise. After bit of algebra (Appendix obtain simpler, more efficient expression:
N(a) arg max (x(a)= 1+ (1 E.
(14)
This deterministic action selection selects the actions proportionally to n' and avoids an extra vari- ance. 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.
RELATED WORK
introduced the bandit with 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 al. 2020) used a deep policy network as the predictor in variant of the PUCB algorithm. [Bertsekas] (2019) 2021: 2022) provides an in-depth discussion policy iteration, policy improvement, and their connection to rollout.
Rosin (011)
If not using predictor, UCT (Koesis & Szepesvári 2006) would need to visit each action before being able to compare them. Rapid Action Value Estimation (Gelly Silver 20TT) then helps to form rough estimates of the action values by aggregating statistics from all future states. Gelly & Silver (201T) also initialized the action value estimates with a hcuristic evaluation function. The best heuristic used learned linear network. Hamrick et al. 2020)) later extended it deep Q-network.
Cazenave 2014) and 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 Q- values. The way to obtain scores from a policy network was left as an open problem. We can now view the g + logits special scores.
MCTS related to regularized policy optimization. Grill et al. (2020) analyzed AlphaZero tree search and discovered that AlphaZero approximates regularized policy optimization. The approxi- mation error is large if using small number of simulations. To avoid the approximation error, Grill al. (220) used regularized policy optimization directly inside the
search. In the setting
without a predictor, Xiao et al. (219)) compared UCT to new MCTS with an entropy regularizer.
Dam et al. (22TT) generalized it to relative entropy and Tsallis entropy. Regularized policy opti-
mization is helpful when working with approximate Q-values (Vieillard et al. (020) or when doing an approximate policy iteration Kakade Langford 2002 Schuiman et al.) 2015).
TreeQN (Farquhar et al. 2018) uses breadth-first search inside network architecture. The net- work 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 ac- tions. To approximate the gradient, Dynamic Planning Networks use Gumbel-Softmax (Maddison et al. 2017: Jang al. 2017). Although use Gumbel variables, we do not employ approximate We use the Gumbel-Top-k trick to construct efficient planning with provable policy improvement.
gradients from

---

Published as a conference paper at ICLR 2022
9x9 Go nm2
9x9 Go nm4
9x9 Go nm16
9x9 Go nm
9x9 Go n0200
2500
8 1500
3000
Gumbar cumbbll Zaru
500
Muzero
MuZero
800
800
800
800
Millions frammes
14llions fremes
Millions frames
Millios fremms
Millions frremss
Figure 2: Elo on 9x9 Go, when training with n {2, 4, 16, 32, 200} simulations. Evaluation uses 800 simulations. Shades denote standard errors from 2 seeds.
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 solution of a regularized policy optimization problem (Tomar al., 2020). Furthermore, the exact computation the unordered set estimator requires O(2") operations, which would be prohibitively
expensive for m 16. Practical time complexity be achieved by using importance-weigghe
estimator Vieira 2017 Nauman & Den Hengst 2020).
Continuous actions be supported by sampling i 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 al., 2021). Similarly, Critic Weighted Policy (Wang al. 2020) uses sampling with replacement.
EXPERIMENTS
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 pre- activation 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 Gum- bel. The policy loss uses the completed Q-values (Equation T2p. Gumbel MuZero sampled with replacement (Replacement): An ablation to Gumbel MuZero by sampling m actions with re- placement, as in Sampled MuZero (Hubert 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(T, knew) (Schulman et al., 2015, Grill al. 2020). MPO MuZero: TRPO MuZero but with the MPO regularizer KL( neww T) (Abdolmaleki et al. 2018 Grill et al.) 2020). Full Gumbel MuZero: Gumbel MuZero with a principled action selection also 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 report additional ablations and experimental details Appendix F
7.1
9x9 Go
On Go, we use Elo to compare MuZero and other agents. While an agent trains by self-play, its Elo computed by evaluation versus reference opponents. One of the opponents Pachi (Baudiš & Gailly 2011) 1) 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
"The corresponding win probability is
Elo 1978
StaDefferene 4010
14 10
7

---

Published as a conference paper at ICLR 2022
9x9 Go nm2
9x9 Go nm4
9x9 Go nm16
9x9 Go nm200
9x9 Go nm
2300
2000
m-d
1000
500
Gumbel Simple policy loss
800
400
600
800
200
200
800
Millions
of orarmmss (6)
(a) Policy loss ablations, when training with E (2, 4, 16, 200} simulations. Gumbel MuZero the policy loss with completed Q-values. (b) Sensitivity to the number of sampled actions. Gumbel MuZero samples m actions without replacement.
Figure 3: Gumbel MuZero ablations on 9x9 Go.
3000
3000
2500
8
1500
1500
1000
1000
Gumbe Kuzeero
Gambe AlphaZem
500
Gumbel MuZero
500
Gumbel AlphaZaro
MuZero
AlphaZero
1.00
200
100
400
500
800
700
800
0
200
600
600
800
1000
1200
1400
1600
Millions framess
Millions of Frames
[e]
(b)
Figure 4: Large-scale experiments with n 400 simulations per move. (a) Elo on 19x19 Go, when training MuZero. (b) Elo on chess, when training AlphaZero.
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. 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 or less simulations. Strikingly, Gumbel MuZero learns reliably even with 2 simulations.
In Figure 3 we compare the simple policy loss (Equation D and the policy loss with the completed Q-values (Equation [2}. 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 produce policy improvement.
In Figure 3p we study Gumbel MuZero's sensitivity to the number of sampled actions. When sam- pling 4 actions without replacement, the simulation budget is spent the small number of actions. The learning was then slower. In all other Go experiments, we sample m min(n, 16) actions without replacement.
7.2 LARGE-SCALE 19x19 Go AND CHESS
In Figure 4 we demonstrate that Gumbel MuZero is not worse than MuZero on 19x19 Go. MuZero excellent 19x19 Go and Gumbel MuZero reaches exceeds its performance. The Elo is still anchored to have Pachi at 1000 Elo.
Similarly, in Figure 4p show Gumbel AlphaZero performance chess. We train AlphaZero chess, because AlphaZero learns faster than MuZero on chess. On Go, MuZero learns faster than AlphaZero.
8

---

Published as a conference paper at ICLR 2022
pecmen n=2
pacmen n=4
pacman n=16
pacmnn n=18
pacman 5=
beam, rider nm18
80000
70000
scalo=0.01
scalc-0.1
60000
300000
scalo=1 1
Near rubur
50000
250000
scalen10
200000
150000
20000
100000
10000
Gumbel
50000
Muzfero
200 framcss
200
300
Milllons
of frames
framess
framess
(b)
Figure 5: Atari results. (a) Mean return on ms maapacnnn when training Gumbel MuZero and MuZero with n E {2, 4, 16, 18, 50} simulations. MuZero fails to learn from or less simula- tions. (b) Mean return on beam_r rider for Gumbel MuZero with Cscale E {0.01, 0.1, 1, 10, 100), compared to MuZero with n 50 simulations. Shades denote standard errors from 10 seeds.
7.3 ATARI
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 |Schritwiiser 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 al., 2021).
In Figure 5 we show the obtained mean return ms-pacman. Gumbel MuZero again learns reliably even with n = 2 simulations. Atari has only 18 actions, so we sample m = min(n, 18) actions without replacement. In the experiments with 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 small number of simulations.
Atari is challenging, because different games can have very different reward scales. MuZero nor- malizes the Q-values by dividing them by max(ó,, max, ((a)) min( (ōx, min, ((a)) found inside the tree search (Schrittwieser et al., 2020). A normalized advantage is then in [-1.1] 1]. For Gum- bel MuZero, we use the same normalization and we scale the normalized Q-values by Cyisit 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 n' (Hes- sel et al. 2021).
In Figure 5p we use beam_rider as an example of 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 + max,, N ((b))(ssate((()) (Fquation 8). Escate is large, Gumbel MuZero focuses ğ(a) and neglects the logits. Indeed, Gumbel MuZero performance is worse on beam_rides r using large Cacale- 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).
8 CONCLUSION
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 the clean theoretical foundation, the faster experimentation with a small number of simulations, and the released open- code
https://giiii www.deep/w.eeo..ooommoot........ nd/ /mctx

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGMENTS
We would like thank Fabio Viola, Matteo Hessel, and Laurent Sifre for building the core of the open-sourced Monte Carlo tree search. Loic Matthey and Benjamin Van Roy provided helpful comments on the paper draft. Thomas Hubert, Iurii Kemaev, Tor Lattimore, and Jean-bastien Grill shared their wisdom.
REFERENCES
Abbas Abdolmaleki, Jost Tobias Springenberg, Yuval Tassa, Remi Munos, Nicolas Heess, and Mar- tin Riedmiller. Maximum Posterioni Policy Optimisation. In International Conference on Learn- ing Representations, 2018.
Jean-Y Yves Audibert, Sébastien Bubeck, and Rémi Munos. Best arm identification in multi-armed bandits. In COLT, pp. 41-53. Citeseer, 2010.
Petr Baudiš and Jean-loup Gailly. Pachi: State of the art open source Go program. In Advances in computer games, pp. 24-38. Springer, 2011.
Mare G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The Arcade Learning Envi- ronment: An evaluation platform for general agents. JAIR, 2013.
Dimitri Bertsekas. Reinforcement and Optimal Control. Athena Scientific, 2019.
Dimitri Bertsekas. Rollout, policy iteration, and distributed reinforcement learning. Athena Scien- tific, 2021.
Dimitri Bertsekas. Lessons from AlphaZero for Optimal, Model Predictive, and Adaptive Control. Athena Scientific, 2022. URLhttp:
http: /ttt:://wbb.. www.teeutttttittammmmmm..e..weet http//miittitttteeittiitttettttt wtt.twtbb/wwwwttttt...ot...www.t hwt.h//wwwwwwwwwkkwookkkoooo RLbook. html
James Bradbury, Roy Frostig, Peter Hawkins. Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Min Milne, and Qiao
Zhang. JAX: composable transformations of Nython+NumPy programs, 2018. URL http: //qithubs gooom///ool//aax
Sébastien Bubeck, Rémi Munos, and Gilles Stoltz. Pure exploration in finitely-armed and continuous-armed bandits. Theoretical Computer Science, 112(119:183218855 2011.
Tristan Cazenave. Sequential halving applied to trees. IEEE Transactions Computational Intel- ligence and Al in Games, 7(1):102-105, 2014.
Tuan Q Dam, Carlo D Eramo, Jan Peters, and Joni Pajarinen. Convex regularization in Monte-Carlo tree search. In International Conference Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp.
PMLR, 18-24 Jul 2021.
Arpad F. Flo. The Rating of Chessplayers, Past and Present. Arco Pub., New York, 1978.
Nicolas Fabiano and Tristan Cazenave, Sequential halving using scores. In Reinforcement Learning in Games at AAAI, 2021.
Farquhar, T Roc,täschel, M Igl, and Whiteson. TreeQN and ATrecC: Differentiable tree- structured models for deep reinforcement learning. In International Conference on Learning Representations, volume ICLR, 2018.
Sylvain Gelly and David Silver. Monte-Carlo tree search and rapid action value estimation com- puter Go. Artificial Intelligence, 175:1856-1875, 2011.
Jean-Bastien Grill, Florent Altché, Yunhao Tang, Thomas Hubert, Michal Valko, Ioannis Antonoglou, and Remi Munos. Monte-Carlo tree search as regularized policy optimization. In International Conference Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 3769-3778. PMLR, 13-18 Jul 2020.
Emil Julius Gumbel. Statistical theory of extreme values and some practical applications: a series of lectures, volume 33, US Government Printing Office, 1954.
10

---

Published as a conference paper at ICLR 2022
Jessica B. Hamrick, Victor Bapst, Alvaro Sanchez-Gonzalez, Tobias Pfaff, Theophane Weber, Lars Buesing, and Peter W. Battaglia. Combining q-learning and search with amortized value esti- mates. In International Conference Learning Representations, 2020.
Nicholas Hay and Stuart J. Russell. Metareasoning for Monte Carlo tree search. Technical report, EECS Department, University of California, Berkeley, Nov 2011.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In Computer Vision ECCV 2016, pp. 630-645, 2016.
Matteo Hessel, Ivo Danihelka. Fabio Viola, Arthur Guez, Simon Schmitt, Laurent Sifre, Theophane Weber, David Silver, and Hado Van Hasselt. Muesli: Combining improvements in policy opti- mization. In Marina Meila and Tong Zhang (eds.). Proceedings of the 38th International Con- ference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp- 4214 4226. PMLR, 18-2 24 Jul 2021.
Jie Hu, Li Shen, and Gang Sun. Squeeze-and-excitatin networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 7132-7141, 2018.
Thomas Hubert, Julian Schrittwieser, loannis Antonoglou, Mohammadamin Barekatain, Simon Schmitt, and David Silver. Learning and planning in complex action spaces. In International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 4476-4486. PMLR, 18-24 Jul 2021.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with Gumbel-Softmax. In International Conference on Learning Representations, 2017.
Sham Kakade and John Langford. Approximately optimal approximate reinforcement learning. In International Conference on Machine Learning, volume 2, pp. 267-274, 2002.
Zohar Karnin, Tomer Koren, and Oren Somekh. Almost optimal exploration in multi-armed ban- dits. In International Conference on Machine Learning, volume 28 Proceedings of Machine Learning Research, pp. 1238-1246. Atlanta, Georgia, USA. 17-19 Jun 20L3.
Diederik Kingma and Jimmy Ba. Adam: A Method for Stochastic Optimization. arXiv e-prints, art. arXiv: 1el::111.69988 December 2014.
Levente Kocsis and Csaba Szepesvári. Bandit based Monte-Carlo planning. In Machine Learning: ECML 2006, pp. 282-293, 2006.
Wouter Kool, Herke Van Hoof, and Max Welling. Stochastic beams and where find them: The Gumbel-top-k trick for sampling sequences without replacement. In International Conference on Machine Learning, pp. 3499-3508. PMLR, 2019.
Wouter Kool, Herke van Hoof, and Max Welling. Estimating gradients for discrete random variables by sampling without replacement. In International Conference on Learning Representations, 2020.
Gary Linscott, Gian-Carlo Pascutto, Alexander Lyashuk, and Folkert Huizinga. Leela Chess Zero, 2018. URL http:://cczer..rrg
Duncan Luce. Individual choice behavior: A theoretical analysis, 1959.
Marlos C Machado, Mare G Bellemare, Erik Talvitie, Joel Veness, Matthew Hausknecht, and Michael Bowling. Revisiting the arcade learning environment: Evaluation protocols and open problems for general agents. Journal of Artificial Intelligence Research, 61:523-562, 2018.
Maddison, Mnih, and Teh. The concrete distribution: continuous relaxation of discrete
random variables. In International Conference on Learning Representations, 2017.
Yu Nasu. Efficiently Updatable Neural-Nettork-bassel Evaluation Functions for Computer Shogi. The 28th World Computer Shogi Championship Appeal Document, 2018.
Michal Nauman and Floris Den Hengst. Low- Variance Policy Gradient Estimation with World Models. arXiv e-prints, art. arXiv:2010.1562 October 2020.

---

Published as a conference paper at ICLR 2022
Tom Pepels, Tristan Cazenave, Mark HM Winands, and Mare Lanctot. Minimizing simple and cu- mulative regret in Monte-Carlo tree search. In Workshop on Computer Games, pp. 1-15. Springer, 2014.
Christopher D Rosin. Multi-armed bandits with episode context. Annals of Mathematics and Artifi- cial Intelligence, 61 613))200,- 2011.
Julian Schrittwieser, loannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, Timothy Lillicrap, and David Silver. Mastering Atari, Go, chess and shogi by planning with a learned model. Nature, 588(7839):604 609, Dec 2020. ISSN 1476-4687.
Julian Schrittwieser, Thomas Hubert, Amol Mandhane, Mohammadamin Barekatain, loannis Antonoglou, and David Silver. Online and Offline Reinforcement Learning by Planning with Learned Model. arXiv e-prints, art. arXiv:1104.06294, April 2021.
John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust Region Policy Optimization. In International Conference on Machine Learning, pp. 1889-1897, 2015.
David Silver, Aja Huang, Chris J. Maddison, Arthur Guez, Laurent Sifre, George van den Driessche, Julian Schrittwieser, loannis Antonoglou, Veda Panneershelvam, Marc Lanctot, Sander Dieleman, Dominik Grewe, John Nham, Nal Kalchbrenner. Ilya Sutskever, Timothy Lillicrap, Madeleine Leach, Koray Kavukcuoglu, Thore Graepel, and Demis Hassabis. Mastering the game of Go with deep neural networks and tree search. Nature, 529(7587):484 489, January 2016.
David Silver, Julian Schrittwieser, Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, et al. Mastering the game of Go without human knowledge. Nature, 550(7676):54-35)) 2017.
David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Mare Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, et al. A general reinforcement learning algorithm that masters chess, shogi, and Go through self-play. Science, 362(6419)) 140- 1144, 2018.
Norman Tasfi and Miriam Capretz. arXiv:1812.1 1240, December 2018.
Dynamic Planning Networks.
arXiv e-prints, art.
David Tolpin and Solomon Shimony, MCTS based on simple regret. In AAAI Conference on Artifi- cial Intelligence, volume 26, 2012.
Manan Tomar, Lior Shani, Yonathan Efroni, and Mohammad Ghavamzadeh. Mirror Descent Policy Optimization. arXiv e-prints, art. arXiv:2005.090 00w.0000000014444 14, May 2020.
Nino Vieillard, Tadashi Kozuno, Bruno Scherrer, Olivier Pietquin, Remi Munos, and Matthieu Geist. Leverage the average: an analysis of KI. regularization in RI. In Advances in Neural Information Processing Systems, 2020.
Tim
Vieira.
Gumbel-max
trick and
weighted reservoir sampling, 2014.
URL
Tim Vieira. Estimating means in finite uniieeer
Ziyu Wang, Alexander Novikov, Konrad Zolna, Josh $ Merel, Jost Tobias Springenberg. Scott Reed, Bobak Shahriari, Noah Siegel, Caglar Gulcehre, Nicolas Heess, et al. Critic regularized regression. Advances in Neural Information Processing Systems, 33, 2020.
David J. Wu. Accelerating Self-Play Learning in Go. arXiv e-prints, art. arXiv:1902.10556 Febru- ary 2019.
Chenjun Xiao, Ruitong Huang, Jincheng Mei. Dale Schuurmans, and Martin Müller. Maximum tropy Monte-Carlo planning. In Advances Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
12

---

Published as a conference paper at ICLR 2022
John I Yellott. The relationship between Luce's choice axiom, Thurstone's theory of comparative judgment, and the double exponential distribution. Journal of Mathematical Psychology. 15(2): 109-144, 1977.
13

---

Published as a conference paper at ICLR 2022
APPENDIX
Content
AlphaZero action selection
Policy improvement proof for planning with Gumbel
C
Policy improvement proof for completed Q-values
D
Mixed value approximation
Derivation of the deterministic action selection
Experimental details
A ALPHAZERO ACTION SELECTION
We will explain AlphaZero's (iiverr et al. 2018) action selection here. On the deterministic bandit, AlphaZero would select actions A1,.... .1....; A, by
arg (a) E, ((N
(15)
where (a) E j0, 1] is g(a) for the already visited actions and zero otherwise. The ej > 0 factor independent of After the search, AlphaZero acts by sampling from an (annealed) categorical distribution based upon the visit counts of the root actions.
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 produce worse policy by adding noise to potentially optimal policy network.
B POLICY IMPROVEMENT PROOF FOR PLANNING WITH GUMBEL
We will prove that Algorithm generates Ant1 such that E(g(An+1) 2 EA- (A)]- For the right-hand side, the Gumbel-Max trick tells us that EAwz [A(] is equal to max, E(gebr)-Cumbe(()) max, (9(a) + logits(a))). First, we will show that we can replace the arg max, with arg max,c Auope Remember that Atopn is defined as Anopn argtop(g++oge logits, n) and that use the same Gumbel vector g in the argtop and arg max. The set Atopn then includes the action with the highest g(a) + logits(a) and we can replace the arg max, with arg max m
Acopn
After these rewrites, we have to prove that
Elg(An41)) ((geE*t) (a(agg max(g(a) + logits(a))].
(16)
aEAippp
On the left-hand side, E[g(A, An +1] is equal to
(a(ag max(g(a) + (
(17)
aEArop
We can finish the proof by proving that for any vector g E I we have
q( arg max(g(a) + logits(a) + a(g(a)))) 2 q(arg max(g(a) + logits(a)). aEAro
(18)
Acopo
This is true, because is monotonically increasing transformation.
C POLICY IMPROVEMENT PROOF FOR COMPLETED Q-VALUES
We will prove that n² complette coftmax(logitt (ccoppleted))) produces policy improve- ment. We will start by showing that Wommpleted is produced by a specific instance of Algorithm
We will then prove that any instance of Algorithm produces policy improvement.
14

---

Published as a conference paper at ICLR 2022
Algorithm 3 Policy improvement when having
Require:
Require:
q(a) for each visited action. The visited actions can be from any distribution.
Initialize n' with
For the visited actions:
If [() Cx. increase the N' (a) logit. If ((a) < Ur. decrease the (a) logit. return
C.I SPECIFIC INSTANCE
Algorithm more general than the usage of the completed Q-values. Specifically, Algorithm would produce "comppitedd if updating the logits by o(completexl((()) (v.). This update in- creases the logit, if q(a) > Vx. This update decreases the logit, if g(a) < And the update does not modify the logits the unvisited actions. The resulting softmax(logits +o (compleetd))) (vx)) is equal to softmax(logits (a (completedQ)), because the constant offset o(uz) does not change the softmax output.
C.2 POLICY IMPROVEMENT PROOF FOR ANY INSTANCE
We will now prove that n' from Algorithm
3
satisfies
Ed'(a)((a) 2 >= ><<(a)((a)
(19)
Notice that (a) for any unvisited action a will be C2 (a), with normalization constant e₂ > For visited action: Let's denote the visited (aka expanded) action by dex- First, if ( then U, = ()xx) and the policy will remain unchanged.
Let's consider the case with (aex) 1. The Ur can be rewritten
T T(a)q(a)
Ur = T ((rx)) (exx) + (1 nt (dex))
Σ
(20)
Enrpare (b))
aftex
Let's denote the weighted sum by Amiss:
T())g(a)
Tmiss
Eoran (bb)
(21)
We notice that the Jmiss does change if scaling by constant c. 0.
We will now rewrite the left-hand side of Inequality 19 to use Gmiss:
Σ n (a)a(a)= =
(22)
C2T CaT(a)g(a)
(axx))(((ex) + n' (aex))
Σ
(23)
Dopace C.T
astacs
n' r'(axx)g( + (1 n' (aex) Ymisss
(24)
n (((axx)(((Gx) - quisse) Jmiss.
(25)
The rinht-hand side of Inequality
19 can be also rewritten:
U π (Oxx) (Oxx) + (1 T (Bex)) Amisss
(26)
= π (aex) (g( (ex) quiies + Gmisa-
(27)
With these rewrites, Inequality 19 becomes n' (Aax) (g(awx) Amises) ((ax)) - qmisss)
(28)
15

---

Published as a conference paper at ICLR 2022
9x9 Go n=2
9x9 Go n=4
9x9 Go n=16
9x9 Go n=200
3000
2300
1500
Gumbel MuZero
500
Completed by Simple palicy
0
Figure 6: Detailed policy loss ablations. Gumbel MuZero the policy loss with Q-values com- pleted by the Umix value estimator from Appendix D That works better than Q-values completed by the raw value network 0z.
9x9 Go n=2
9x9 Go n=4
9x3 Go n=16
9x9 Go n=200
2500
2000
1000
500
Full Sumbel MuZero
Shochastic nan-mace
200
400
500
800
200
400
600
800
200
400
600
800
200
400 Millions of frames
500
800
Millinn of frames
Millinns of fromes
Millons frames
Figure 7: A comparison of different action selections the non-root nodes. Gumbel MuZero uses the unmodified (deterministic) MuZero action selection at non-root nodes. Full Gumbel MuZero uses the deterministic action selection from Equation 14] which we compare to stochastic sampling from x' at non-root nodes.
The q(uxx) Gmiss can negative, zero or positive.
If g(dex)
Umisss the inequality is satisfied.
g((ex) Umises we want x (Aex) 2 क (dex)-
If g(ax)) Gmise, we want x (((e) (axx).
We do not know Jmiss cannot use it in an algorithm. We will instead show that g(axx) > Quiss equivalent to ((eex) > Ur. when T (aex) <
g(aex)
(29)
((aex) > π. (dex) (axx) + (1 - 7 (mex)) (1 T (Oxx))
(30)
> (1 n (Bex) Amiss
(31)
((uex) qmiss.
(32)
So we directly arrived at Algorithm]
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. the logits of the other visited actions modified by Algorithm 3 they can further help to improve the policy.
D MIXED VALUE APPROXIMATION
We will construct an approximation of U,, The exact v is defined by U2 = E. T(a)q(a). We have an approximate i from value network, know æ, and we have q(a) for the visited actions. With
16

---

Published as a conference paper at ICLR 2022
9x9 n=2
9x9 Go n=4
9x9 n=16
9x9 Go n=200
3000
2000
8 1500
C visies
1.vii=500
500
Stochastic evaluation
c_V5t=50000
No exploratioo
200
400
500
800
200
400
600
800
200
600
600
800
200
500
800
Miltons of frames
Millinns of fromes
Millinn frames
Millions of fromes
Figure 8: Additional Gumbel MuZero ablations 9x9 Go. (a) Sensitivity Q-value scaling by Cvisit- (b) On the perfect-information game, Gumbel MuZero used zero Gumbel noise at evaluation. Although, evaluation with stochastic Gumbel noise is worse. During training, MuZero and Gumbel MuZero benefit from explorative acting proportional to the visit counts.
these inputs, we approximate Un by a consistent estimator:
N.
Umix 1+ I+E,NO N(b)
E
x(a)q(a)
(33)
E Tb))
aE (a(=
The estimator interpolates ê and the weighted average of the available Q-values. This is unso- phisticated estimator, with results in Figure (6) You are welcome to explore other possibilities.
E DERIVATION OF THE DETERMINISTIC ACTION SELECTION
We will derive Equation
14 from Equation
13
N(b) + =fa=b]
arg min
T'(6)
(34)
1+ N())
N(b)
Ifa = b)
arg min
7'(b)
(35)
1+ E.N()
1+ E.ENNc
b
arg min
Σ
(b)
N(b)
I{a b}
(36)
I+ E. N(O)
1+ Σ. N(c)
N(6)
arg min
<'(b)
T(a=b)
(37)
1+ Σ. N(c)
((aa
arg max n' T'(a)
(38)
1+ Σ. N(c)
The simplification was possible, because additive terms independent of a not affect arg min..
While widely applicable, the deterministic action selection provides only small benefit 9x9 Go (Figure 7p..
F EXPERIMENTAL DETAILS
In general, we use hyperparameters consistent with the newest MuZero experiments (Schrittwieser al. 202T). MuZero's 's pseudocode is available thanks to [Schiitwieser 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 ((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 normalized to be from
17

---

Published as a conference paper at ICLR 2022
9x9 Go train n=2
9x9 Go train n=d
9x9 Go train n=16
9x9 Go train n=32
9x9 Go train n=200
3000
2000
eval nu.00
1000
ecal n=200
nm33
n-16
n-]
Niliions
Figure 9: Gumbel MuZero Elo 9x9 Go, evaluated with different numbers of simulations. The evaluation with n = 1 simulation acts with the most probable action from the policy network.
the [0, 1] interval. We use the normalized Q-values also in Full Gumbel MuZero, but the algorithm does not require Q-values from specific interval. In all Go and chess experiments, Gumbel MuZero scales the Q-values by Grisit = 50 and Cscale 1.0. On the perfect-information game of Go, Gumbel MuZero is not very sensitive the scale of the Q-values. Any Cvisit 2 50 produced similar results (Figure (a).
In each phase Sequential Halving, we use at least one new visit. For example, in the first phase, update q(a) by max (1, [lo]](m]]] visits. This allows
experiment with incomplete or no Sequential Halving. When Sequential Halving runs out of the budget of simulations, we stop the search. The agent then selects An41 the action with the highest g(a) logits(o) (a) + (ÿ(()) from the set of the most visited actions. Fabiano & Cazenave (202 provide 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 the visit counts. like AlphaGo Zero (iilver et al. 2017). Gumel MuZero benefits from the same exploration (Figure P))
Figure 9 shows the importance of the number of simulations at evaluation time. For example, in the first subplot network is trained with n = 2 simulations and the same network is evaluated with 800, 200, 32, 16, 2, and 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, self-play. Tabld] the obtained speedups if not being limited by the TPUs for training.
can substantially reduce the number of TPUs needed for
Table 1: The speedup from smaller number of simulations on 9x9 Go.
Training step speedup
MuZero n = 200
1.0
Full Gumbel MuZero 200
1.0
Gumbel MuZero n 200
1.0
Gumbel MuZero n 322
5.9
Gumbel MuZero n. = 16
11.3
Gumbel MuZero n 8
16.2
Gumbel MuZero n = 4
24.3
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, replace the dense residual blocks
18

---

Published as a conference paper at ICLR 2022
with bottleneck blocks (He al. 2016), and replace every 8th block with global broadcasting residual block (Figures 13). The broadcasting block is similar to global pooling or squeeze-and- excitation (Hu et al. 2018 but found it to be more stable when training from self-play.
Figure [ 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 a fixed computational budget of 300ms per move. For example, the budget allows us to do 3200 simulations when using network with 256 planes and 32 blocks with bottlenecks and broadcasting (achieving 10 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, the improved predictions fail to compensate for the much lower number simulations 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 bigger 32 layer network the large-scale 19x 19 Go experiments. The networks used 256 hidden planes, 128 bottleneck planes and broadcasting block in every 8th layer.
19

---

Published as a conference paper at ICLR 2022
128p
2560
J04p
512p
640p bi posto
128p broad/8
256a broad/8
512p broad
640p broad/s
120p
256o
384p
512p
128p posue
256o et poova
384p btl pool8
512p bdi pootre
broad/88
broad/8
Value MSE
0.74
0.72
0.71
inferencee/sseconc
Policy Accuracy
0.570
0.565
0.560
inforencesswexcond
200
100
-100
-300
inferencesisecone
Figure 10: Scaling of value Mean Squared Error (MSE), policy accuracy and playing strength in Elo inference speed for different network architectures. In the legend, bul indicates bottleneck residual blocks; broad/8 and pool/8 indicate broadcast (Figure TT and pooling blocks (Figure 12) in every 8th layer. 128p, 256p, etc. indicate the number of hidden planes.
20

---

Published as a conference paper at ICLR 2022
class BasicHlock (ak. (ak.Modulel;
"**Bosic block composed of inner ap, op and linearity. **.
(isit_(saff,, make_inner__g
won. name='basic")
super 11000__nn_____m
selç,_op make maae__nner__og
seif._norm lk. www.eacchlommmrrrttte...eelleeee
ceatte_ffist-Trues decay_rase=0_
aps=le-3)
sel sellf...on_Mnearit non_linearity
def _alll (salf, x: Tensor, call_args: :
(x)
self. normi,, is ts_ttriiinng=esllarrr ireiiinnn
ses__oool_sttttmelllll wer.e...eet.oclssseeeeeee return non_
ass Ros :
gn*"Createa residual block with optional bottleneck."" "*"
def iiiiiinnttlll stack_tize: int, firs_oixst_ make_inner_of, make_last_op, name)
super() (o_r)())_n________ init ianiinnntaanaa (name-name)
selt._blocks
for sake_op enumerate (laake_ittttk [make_ineer_op]
(stack_size 2)
celf._blecks. Rasiek oak
eske_inner_opppakee.ee
non_iinnariyyelambdd x: if stack_size else jex.na.relu,
bane=("baeic_l1)"11
def
x: lensor, call_args: CallArgs):
for b in blook_blocks: bizes, call_args) return jeluu. + res)
class BroadnastRosslockeeee (Kesslock):
PREA reaidual block broadcaata information spatiaa dimenatons.
block consists
sequence
layers:
layer that mixes information acress chennels, layer that mixes information within each channel, a dense layes. another layer mix acress channels.
ixl convolution.
The
nixing information within channel.
def innit_ (salf, nake_aix__cannell. nome):
def broadeast (x: inp.ndarrsy):
n, h, W, x.ahapa
Process all planes at once, applynig the same linear layer to each. x.transpose (0, 3, 1, 2)) &BWC
rerhapp (n, c,
vi)
x hk Linein (h * 8, broadabboa (x)
rol. (x)
x reshape (f,, 6, 3, w1) x.Lrraspose (00, 3,
NCIN
return
super i)__init_(
stack_size=3
make_inner_opp-anaeee broadoast,
make_lasttop-mak__nit pat_maa__i__chamm___ee name=name)
Figure 11: Bottleneck and broadcast blocks used by the board game network, implemented in JAX Braddury et al. 2018) with the neural network library Haiku
21

---

Published as a conference paper at ICLR 2022
alasa Poo RosBBook :
aang residual that broadcasta information accoas spatial dimensions.
block consists
sequence
layers:
layer that mixes information acrosa channels, layer that mixes information within each channel, à dense layez, mix channels.
1x1 convolution.
another
same weights is used mixing information within channel.
det make_mix_chatnele YakeFocwardModules,
name=' p
supez () init (name-name)
bleck funetocls.
parcial( saantocle.parcial(
sasicHlock, htt.://eerr.ppooooooieeeeeeeeeee
def call (self, Tensor, argss CallArgs):
(nee_ je..en.eeeee............ self._ jnon.
nn. relu, name-' input_s') (x, call_args)
names' input_b') call_args)
b_planes jop.concetenate (jjp.. (ban (1, 21), jap.max ib, (1, 2))1, -1)
b_planes Linern (a. sa.she [-1], name= namo='mix_chans chennelll (b_planes)
a b_ @aplanee[s (t, None, None, :]
x block lnon linearity-lambda x: X, name output') (c, call_args) return (c)
Figure 12: Pooling block, based on Wu (0099)
make_conv wrwttptt_chhnnelllee int, kernel_ahope: intl:
return functoola, a. partial (nk.Conv2D, srtput_ehhannels kernel_shape,
w_init-hk, initializers TruncatedNornll inl.e (0.001))
make_network layerreyers: ins, sucput_channeil int, cestleneck_chhmmele: int, breadcast_every_.n int):
blocks
BasicBlock fackee_inee___eee Samner_o--ake_oooe cony (haneett_channelll,, not_linearlteeaaa.m..ell name=! init_coav")
for zange (awm :
if broadcast__evere
0 and broadcast_every_n =s broadeact_evoryy.n 1:
blooks. bloexs.append append RRraadoaat
make_mix_channel wtw_ehhheeeeeemaaeeee.
(chmput_ kecpel_shhap=_),,
neme=5' beoddasstsoccc
elif bottleneck_chhnnelle
appead (ResBlock
stack_pizem4,
make make_first, tasiirat__pmmaa_oom (wwtttleecc__chhhheee kernel schnel_shape=t)
make make_inner. _atiiner__pmmake___e conv boottleeek__hhneeeeee shapee_)
[is_lesst_ppmaae_ccre (outputt_hhaaeeee... kernel_shape-1])
neme=5' botslencck__ee___3111.
alocks.append (ResBookk)
stack_aize=2,
make make_firsts wtt__rrrt_pppkkeee___m costpptt___n kernel_shspe-3),
make _nkeinne___ppkkeee______ conv foatput_cheanelle kernel_shape=3),
_ik__lest_oppmakee__ foutput_channeee, keznel_shape=3), nome=5' (il"))
return blocks
Figure 13: Usage of modules defined in Figure I|
to create the network used for board game exper- iments. Alternatieely, PoolResBlock can be used instead Blook..
22


Reviews:
[{"review_id": "VEpGlB7d3VK", "reviewer": "Reviewer_dky3", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper proposes to improve AlphaZero and MuZero with the principle of policy improvement by using Gumbel. This paper presents empirical results as well as a theoretical proof. While the method looks interesting especially for the case of using only 2 simulations, some concerns and comments are listed as follows. \n\nFirst, the intuition of applying Gumbel to AlphaZero/MuZero is not well described. Why would Gumbel distribution outperform the current popular UCT (or PUCT) for AlphaZero/MuZero in some cases? \n\nThe authors need to clarify why Equation (7) only needs to consider actions in $A_{topn}$ instead of all actions, while still maintaining policy improvement and satisfying its proof in Appendix. As in (6), the authors state that $E[q(A_{n+1})] \\geq \\sum_a(\\pi(a) * q(a))$ produces a policy improvement, and give Example 1 of Section 3.2, showing that AlphaZero may fail for policy improvement. However, the proposed Algorithm 1 (or 2) may also fail for policy improvement, since in the proposed algorithm the Gumbel method with n = 2 may encounter the same case (even for a simple case g = (1,1,1)). In the current presentation, it is not convincing. \n\nIn Algorithm 2, the statement is unclear “Use Sequential Halving with n simulations ...”. Do you mean by recursively using Sequential Halving? Anyway, a lot of messages are missing in “Sequential Halving”, particularly related to any subtree searches (MCTS?) or rollouts (or MuZero) (as in Figure 1). Without knowing this, it is hard to clarify the algorithm. Maybe, the authors should provide the pseudo code for better understanding and comparison, particularly for the whole learning process for both Full Gumbel MuZero and Gumbel MuZero methods.\n\nIn the experiments of Go in Section 7.1 and 7.2, the authors anchored to Pachi at 1000 Elo. However, the variance would become high for more than 400 Elo difference, and hence makes Figure 3 and Figure 4 unconvincing. The authors should also evaluate confidence intervals as well, and may consider to use other strong Go programs, e.g., KataGo.\n\nAlthough Figures 2 and 4 show Elo strength growth, it is more interesting to see if there is any limitation about the strength, i.e., whether the strength will be limited to a certain level (e.g., up to which level when compared with kataGo?). After all, we want to know the limit we can push to, especially for AlphaZero-like methods. Note that I know it would take a huge amount of time for 19x19 Go, however, I think it should work for 9x9. \n\nThis paper proposes to redesign the selection for the non-root nodes in Section 5, named “Full Gumbel MuZero”. However, in Figures 4 and 7, Full Gumbel MuZero does not seem to have a clear advantage over Gumbel MuZero. That is, it seems that Gumbel MuZero already worked well without this additional improvement. The authors should give more discussion. \n\nIn experiments, more information needs to be shown. For example, the training times, training steps, inferencing times, and the computing resources of the experiments. This is important when using fewer simulations that may significantly reduce these.  \n\nSome minor presentation issues are listed as follows.\n* In Section 3.1 and in Equation (7), q(a) in [0, 1] should be explicitly defined.\n* In Equation (8), is action b the sibling or the children of action a?\n* In Equation (11), is sigma the same as the definition in Equation (8)?\n* In Section 7.1, for the 9x9 Go, when n is small, is the Sequential Halving applied? Note that the authors state that “Gumbel MuZero with n <= 18 selects without Sequential Halving” only for Atari games (in Section 7.3), but not for Go and Chess.\n* In Figure 7 of Appendix C, how does the “Stochastic non-root” select actions?\n", "summary_of_the_paper": "This paper redesigns AlphaZero and MuZero with the principle of policy improvement and claims the following. By using the Gumbel-Top-k trick, the completed Q-values, and some other mechanisms, this paper shows the proposed solutions, Gumbel MuZero and Gumbel AlphaZero, learn reliably even with only two simulations, and the experiments in this paper show that the proposed methods match the performance of the original MuZero and AlphaZero.\n", "summary_of_the_review": "This paper proposes to improve AlphaZero and MuZero with the principle of policy improvement by using Gumbel. This paper presents empirical results as well as a theoretical proof. While the method looks interesting especially for using only 2 simulations, some comments are concerned as described above. \n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "QEz9Tt5X997", "reviewer": "Reviewer_Qbi3", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## Strengths\n\n* This paper was a breath of fresh air in terms of clarity! Very well written overall. I very much liked the progression from deterministic bandit to stochastic bandit to MCTS.\n* I like that the paper emphasizes simple regret instead of cumulative. I completely agree that this is the regret that we should care about in a planning setting.\n* The related work is thorough and seems to hit on all the main relevant areas. The authors demonstrate a deep and comprehensive understanding of the literature.\n* Though code was not provided, the experimental details look sufficient to reproduce the main results (hardware permitting…)\n\n\n## Weaknesses/Questions\n\n* I want to better understand what it is about Gumbel specifically that is important here. The paper makes clear that differentiability is not a factor. Instead, Gumbel is used in Algorithm 1 to satisfy Inequality 6. But it seems that the Gumbel distribution itself is not really important -- what’s important to achieve the inequality is that the same $g$ is used both for action selection during search, and for policy improvement after search. Would the same result be achieved in practice by “re-seeding” a random number generator and sampling from the categorical distribution of actions, for example?\n* Using completed Q-values to train an improved policy, as described in section 4, makes sense. But if we can do that, don’t we automatically avoid the policy improvement issues exemplified by Example 1, without needing to “plan with Gumbel” as in Algorithm 1? In other words, why can’t we just use Equation 11 (which has nothing to do with Gumbel) instead of using Algorithm 1 to define the improved policy?\n* I understand that high variance in action selection could be problematic, and determinism seems like a reasonable hack. But is there any theoretical basis for this?\n* I did not understand this sentence: “We cannot use a deterministic action selection at the root node, because we do not remember the visit counts from the same (or similar) state visited in previous episodes.” Could you explain/elaborate?\n* The results in Figure 2 use only two random seeds. I know these experiments are very expensive, but two random seeds is just not enough for us to draw meaningful conclusions about statistical significance, so I am interpreting these results with a huge grain of salt. I am glad that 10 seeds were used in Atari.\n* The motivation for learning more efficiently from fewer simulations is presumably that with fewer simulations, the overall algorithm runs faster. To what extent is this the case in practice? Approximately how much faster is it to train MuZero with 2 simulations vs 200? Are the simulations the speed bottleneck?\n* It is good to see some ablation results, but it would be most ideal if each of the improvements claimed over the original Alpha/MuZero had an associated ablation. \n\nThe improvements claimed in the introduction are:\n   1. Selecting actions to search (using Algorithm 1 instead of Dirichlet noise)\n   2. Selecting actions at the root (using Sequential Halving instead of PUCB)\n   3. Selecting actions in the environment (using the Sequential Halving result directly instead of sampling from softmax)\n   4. Policy network update (Q-value completion instead of softmax visit count)\n   5. Selecting actions (improved policy instead of PUCT)\n        \n        The ablations in the paper are:\n   * Figure 3a and 6: (d) policy network update\n   * Figure 7: (e) selecting actions\n   * Figure 8:  (none of the above) other ablations on hyperparameters\n        \n        So it would be good to also include ablations for (a), (b), and (c). Let me know if I’ve misaligned anything.\n \n## Minor\n\n   * “we will focus at the action selection” → “we will focus *on* the action selection”\n   * In Algorithm 2, is the “Remaining” in the argmax just a singleton set of the best action found by sequential halving? If so, it might be clearer to just say that $A_{n+1}$ is equal to that best action.\n   * I happened to already know about Sequential Halving, but it’s likely that a more general audience would not. Maybe add a few sentences explaining it -- the caption in Figure 1 does this to some degree, so just including that description in the main text and saying a bit more would be good.\n   * In that spirit, it would be good to give a little more background overall -- for example, clearly stating all of the different subproblems that must be solved in Alpha/MuZero. These subproblems are implicit in the contribution bullet points in the introduction, but they are also not exhaustive.", "summary_of_the_paper": "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.", "summary_of_the_review": "This paper is likely to be of theoretical and practical interest to many in the ICLR community. My recommendation is not yet higher because I am not convinced of the framing around “planning with Gumbel” and because of the concerns surrounding experiments that I mentioned above.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "P4ILOuW_dQG", "reviewer": "Reviewer_Qmbr", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Primary Strengths**:\n1) Several interesting adjustments of how the MCTS operates in AlphaZero/MuZero.\n2) Strong results (including interesting ablations in appendices).\n3) Majority of paper well-written and clear.\n\n**Primary Weaknesses**:\n1) While the majority of the paper is clear, some technical details are (in my opinion) sometimes not clear / only become clear much later on --- including sometimes really important details. See detailed comments below.\n2) No mention of whether source code will be made available.\n\n---\n\n**Detailed Comments**:\n- At end of page 1, for **Selecting actions to search**, should probably explicitly mention that this Dirichlet noise in AlphaZero is only used in the root (for consistency with some of the subsequent points where you also explicitly make the distinction of root vs. non-root).\n- Top of page 2: \"We instead propose to select the singular action resulting from the Sequential Halving search procedure.\" --> I interpreted this as saying that you do not do any exploration at all, purely \"greedy\" playing, even during the training phase. Only all the way down in Appendix E did it become clear that you do in fact using exploration during training, which contradicts this, so then I assume this actually only refers to evaluation games. This should be clarified.\n- In Introduction, the phrase \"We instead propose a policy improvement based upon a completion of the root action vales\" did not at all paint a picture in my head resembling what actually happens in Section 4. Especially it's difficult to imagine, without reading the entire paper, what *completion of values* would mean. I think I would simplify it and simply say that it's a policy improvement operator based on estimated action values (or advantages). The completion of values (or really, completion of *vectors of values*, i.e. filling in \"missing entries\") is more of an \"annoying\" detail necessary to handle situations where some actions didn't get any values, but I don't get the impression that it's actually the core idea; the policy improvement operator still works when there are enough visits to visit every action!\n- Under Figure 1, I assume that the $b$ actions in $\\max_b N(b)$ that the max iterates over are children / successors / actions of the root node, but this should be made explicit.\n- The main paper summarises the proposed policy improvement operator as Eq. (11), and states that a proof for it being a policy improvement operator can be found in Appendix B. However, in Appendix B, I instead find a proof for Algorithm 3, and Algorithm 3 looks very different to Eq. (11). I can... sort of imagine that they might end up working out to the same / a similar thing after the softmax normalises everything into a probability distribution again, but this really doesn't seem obvious. Is it possible to elaborate on this? \n- Related to the previous point, I also don't find the intuition suggested by \"Intuitively, the completed $Q$-values give zero advantage to the unvisited actions.\" particularly intuitive. What if all the $q(a)$ values for the visited actions end up being \"disappointing\", all lower than the expected value for the root based on value network? Surely then the $v_{\\pi}$ values would exceed the $q(a)$ values and actually give a positive advantage to unvisited actions? Algorithm 3 in Appendix B does match this intuition much more closely, because only actions that have $q(a)$ values get updated explicitly... although, of course, the subsequent normalisation of the softmax will implicitly then also update unvisited actions again.\n- Eq. (14) uses $\\pi'(a)$, computed according to Eq. (11), every time the node for which $a$ is an action gets visited. But in between different such visits, its $q(a')$ values (for actions $a'$ that may or may not equal $a$) get updated, so I assume $\\pi'(a)$ also gets re-computed to take into account the latest $q$ values? Would be good to explicitly mention this.\n- \"We cannot use a deterministic action selection at the root node, because we do not remember the visit counts from the same (or similar) state visited in previous episodes.\" --> I do not understand what this is trying to say. I can certainly understanding a preference for a different action selection method (for example, because we care about simple regret in root, but cumulative regret in non-roots)... but this seems to be saying something else?\n- Which agent do you evaluate against in Chess experiments? For Go it was Pachi, but I think that agent doesn't play Chess?\n- Appendix C, Eq. (30): please use larger brackets for the big expression after the first fraction. It's rather difficult to read like this, at first I thought there was a missing closing bracket somewhere.\n\n**Other Remarks**:\n- I am *very* curious if it would be possible to make this approach work even without a value network $\\hat{v}_{\\pi}$. It appears to play a relatively \"small\" role, and already gets mixed with a value estimate based on the children visited so far. What would happen if the value of the policy were estimated solely based on the values backpropagated from previous visits to (other) children? I suppose in practice you would need the value network also to produce the \"bandit rewards\", but those could be replaced with rollouts... I'm by no means expecting/demanding that large extra experiments get added for this, just curious if the authors happen to already have some insights into this.", "summary_of_the_paper": "This paper proposes several adjustments to how actions are selected at various stages of the tree searches in AlphaZero/MuZero, with a particular focus on modifications that can significantly improve training performance in training regimes with extremely low iteration budgets per move (even successful training with only 2 iterations per step). Experiments are run using Go, Chess, and two Atari games.", "summary_of_the_review": "Overall a strong paper, with at this point unfortunately slightly too many places where technical details are not entirely clear or only become clear much later than they should be. Based on the strength of the rest of the paper, I do expect that the authors should be able to clear things up (either in the paper or in comments to my review explaining why I'm wrong) without too much trouble, but I do think it is important to do this before it can be published!", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "OA0CUWEV9YS", "reviewer": "Reviewer_urs9", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "- The paper proposes a series of principled improvements to state-of-the-art MCTS planning algorithms. The proposed improvements are particularly effective in low simulation budget settings. The proposed improvements (Gumbel tricks to incorporate the policy network into the tree search, simple regret minimization at the root, planning-learning loop and action selection at non-root nodes) are intuitively clear.\n- The details on the planning-learning loop are somewhat sparse. Some of my questions were: What exactly is the update mechanism? Which nodes in the tree search are used to generate training data? These details about the learned policy network seems particularly important given the target setting of very small simulation budgets implying limited lookahead.\n- At the smaller simulation budgets, I'd expect the impact of planning to performance is likely limited compared to the learned policy. Is this correct? I'd be curious to better understand how much benefit there is to planning versus simply using the reactive policy without lookahead at the smaller simulation budgets. Please consider including the learned reactive policy as a baseline in the experiments. Some timing information (per decision) would be nice as well to better understand the overhead involved in each method. This might help improve the motivation for planning in small-budget settings, which is currently not described in much detail.\n- Overall, the paper contains a number of intuitively clear, principled algorithmic improvements to AlphaZero (and MCTS in general) when used in low budget settings. The learning-planning portion of the paper can be improved with additional detail. The paper is well written and I enjoyed reading it. I think this will make a nice contribution to the literature on anytime planning algorithms with learned components.\n\nUPDATE: I thank the authors for their feedback. After reading it along with the other reviews, I continue to remain positive about the paper.", "summary_of_the_paper": "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.", "summary_of_the_review": "Improves anytime planning with a learned component in low simulation budget settings. Could be improved with a bit more detail on the learning-planning loop and additional baselines. Overall, seems like a solid paper to me. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper presents improvements to AlphaZero and MuZero for settings where one is restricted in the number of rollouts. The initial response from reviewers was generally favorable  but the reviewers wanted more details and clarifications of multiple parts of the paper, and further intuition about the Gumbel distribution.  The authors’ responses were detailed and convinced or maintained strong positive support of most reviewers. The authors also stated that they plan to provide a release of the code and also provided a policy improvement proof. Overall this is an interesting approach that is likely to be of significant interest to many.
Decision: Accept (Spotlight)