P OLICY IMPROVEMENT BY PLANNING WITH G UMBEL
A BSTRACT
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 signiﬁcantly improve prior performance when planning with few simulations.
I NTRODUCTION
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 Efﬁciently 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 speciﬁcally, we consider each mechanism in turn, alongside our proposed modiﬁcations:
• 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 inﬂuence 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 ﬁnal 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 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 distribution 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.
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 modiﬁcations are applicable also to MuZero or any agent with a policy network and an expensive Q-network.
The modiﬁcations 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.
B ACKGROUND
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 ∈ R k , such that logits( a ) is the logit of the action a .
We can obtain a sample A from the distribution π by ﬁrst generating a vector of k Gumbel variables and then taking argmax:
[Equation]
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:
[Equation]
We will denote the set of n top actions by argtop ( g + logits , n ) =
{ A 1 , A [Equation]
[Equation]
n } .
P LANNING AT THE ROOT
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.
3. P ROBLEM SETTING
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 ∈ R k , 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 A t ∈ { 0 , . . .
, k − 1 } and visits the action to observe the
Q-value q ( A t ) .
The objective is to maximize the Q-value from a special last action A n +1 .
That means we want to maximize E
[ q ( A n +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
[Equation]
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. M OTIVATING 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 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 (cid:80) a π ( a ) q ( a ) = 0 .
For n = 2 simulations, the set of the most probable actions is { 0 , 1 } .
With that, the heuristic would select A n +1 = arg max a ∈{ 0 , 1 } q ( a ) .
The expected value of such action is E [ q ( A n +1 )]
= 0 , which is worse than the value of the policy network.
You can ﬁnd other counterexamples by generating random q and π vectors and testing the policy improvement (Inequality 6).
The AlphaZero action selection is explained in Appendix A.
3. P LANNING WITH G UMBEL
We will design a policy improvement algorithm for the deterministic bandit with a predictor π .
After n simulations, the algorithm should propose an action A n +1 with E
[ q ( A n +1 )] ≥ (cid:80) 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 A n +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.
π ( a ) > 0 .
Algorithm
Require: k : number of actions.
Require: n ≤ k : number of simulations.
Require: logits ∈ R
k : predictor logits from a policy network π .
Sample k Gumbel variables:
[Equation]
Find n actions with the highest g ( a ) + logits( a ) :
A topn = argtop ( g + logits , n )
Get q ( a ) for each a ∈ A topn by visiting the actions.
From the A topn actions, ﬁnd the action with the highest g ( a ) + logits( a ) + σ ( q ( a )) :
A n +1 = arg max a ∈A topn ( g ( a )
+ logits( a ) + σ ( q ( a ))) return A n +
The algorithm produces a policy improvement, because
[Equation]
This holds for any Gumbel g , so it holds also for expectations: E [ q ( A n +1 )]
≥ E A ∼ π [ q ( A )] .
The arg max a ∈A topn ( g ( a ) + logits( a )) is equivalent to sampling from the policy network π (see the
Gumbel-Max trick or Appendix B).
By using the same Gumbel vector g in the argtop and arg max , 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.
3. P LANNING ON A STOCHASTIC BANDIT
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 efﬁciently 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
Require: k : number of actions.
Require: m ≤ k : number of actions sampled without replacement.
Require: n : number of simulations.
Require: logits ∈ R
k : predictor logits from a policy network π .
Sample k Gumbel variables:
[Equation]
Find m actions with the highest g ( a ) + logits( a ) :
A topm = argtop ( g + logits , m )
Use Sequential Halving with n simulations to identify the best action from the A topm actions, by comparing g ( a ) + logits( a ) + σ (ˆ q ( a )) .
A [Equation] Remaining
[Equation] a [Equation] ) return A n +
For a concrete instantiation of σ , we use
[Equation]
[Equation]
where max b N ( 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 ﬁnite scale for the Qvalues provides regularized policy optimization and puts into effect the prior knowledge contained in the logits.
Experimentally, c visit = 50 , c scale = 1 . 0 allowed us to use the same hyperparameters even if changing the number of simulations.
L EARNING AN IMPROVED POLICY
After the search, we have A n +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 A n +1 .
That deﬁnes a simple policy loss:
[Equation]
Using completed Q-values.
We will explain a different way to train the policy network, by extracting more knowledge from the search.
Beside A n +1 , the search also gives us q ( a ) (or its approximation) for the visited actions.
We can construct an improved policy by ﬁrst completing the vector of
Q-values:
[Equation]
[Equation]
where the unknown Q-values of the unvisited actions are replaced by v π = (cid:80) 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
[Equation]
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 π (cid:48) , we can distill it to the policy network π :
[Equation]
This loss trains all actions, not only the action A n +1 .
Later, we will investigate the effect of the loss in Figure 3a.
P LANNING 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.
That allows us to interpret MCTS as regularized policy optimization.
At a non-root node, we construct an improved policy π (cid:48) 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 π (cid:48) .
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 π (cid:48) probabilities and the produced normalized visit counts.
Such action selection would select
[Equation]
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 efﬁcient expression:
[Equation]
This deterministic action selection selects the actions proportionally to π (cid:48) 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.
R ELATED WORK
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 & Szepesv ´ ari, 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.
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-ﬁrst 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 (Tasﬁ & 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 efﬁcient planning with a provable policy improvement.
[Figure]
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 requires O (2 m ) 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.
E XPERIMENTS
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 Gumbel MuZero: MuZero with the modiﬁed root of the search tree to use Sequential Halving with Gum-
bel.
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 modiﬁed 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 signiﬁcantly 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.
7. 9 X 9 G O
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 (Baudi ˇ s &
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.
[Figure] [Figure]
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 ﬁve 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.
7. L ARGE - SCALE 19 X 19 G O AND CHESS
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.
[Figure]
7. A TARI
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.
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 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 n ≤ 18 , Gumbel MuZero selects A n +1 from the n visited actions, without any Sequential Halving.
This conﬁrms 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 π , max a ˆ q ( a ))
− min(ˆ v π , min a ˆ 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 c visit =and c scale = 0 .
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 π (cid:48)
(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 ) + ( c visit + max b N ( b ))
c scale ˆ q ( a ) (Equation 8).
If c scale is large, Gumbel
MuZero focuses on ˆ q ( a ) and neglects the logits .
Indeed, Gumbel MuZero performance is worse on beam rider if using large c scale .
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).
C ONCLUSION
We redesigned AlphaZero tree search.
With the principle of policy improvement, we replaced ﬁve 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 beneﬁt from the clean theoretical foundation, the faster experimentation with a small number of simulations, and the released open-source code. A CKNOWLEDGMENTS
We would like to 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.
A PPENDIX Content
• A - AlphaZero action selection • B - Policy improvement proof for planning with Gumbel • C - Policy improvement proof for completed Q-values • D - Mixed value approximation • E - Derivation of the deterministic action selection • F - Experimental details
A A LPHA Z ERO ACTION SELECTION
We will explain AlphaZero’s
(Silver et al., 2018) action selection here.
On the deterministic bandit,
AlphaZero would select actions A 1 , . . .
, A n by
[Equation]
where q ?
( a ) ∈ [0 , 1] is q ( a ) for the already visited actions and zero otherwise. The c 1 > 0 is a factor independent of a .
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 can produce a worse policy by adding noise to a potentially optimal policy network.
B P OLICY IMPROVEMENT PROOF FOR PLANNING WITH G UMBEL
We will prove that Algorithm 1 generates A n +1 such that E [ q ( A n +1 )]
≥ E A ∼ π [ q ( A )] .
For the right-hand side, the Gumbel-Max trick tells us that E A ∼ π
[ q ( A )] is equal to E ( g ∈ R k ) ∼ Gumbel(0)
[ q (arg max a ( g ( a ) + logits( a ))] .
First, we will show that we can replace the arg max a with arg max a ∈A topn .
Remember that A topn is deﬁned as A topn = argtop ( g +logits , n ) and that we use the same Gumbel vector g in the argtop and arg max .
The set A topn then includes the action with the highest g ( a ) + logits( a ) and we can replace the arg max a with arg max a ∈A topn .
After these rewrites, we have to prove that
[Equation]
On the left-hand side, E [ q ( A n +1 )] is equal to
[Equation]
[Equation]
We can ﬁnish the proof by proving that for any vector g ∈ R
k we have
[Equation]
This is true, because σ is a monotonically increasing transformation.
C P OLICY IMPROVEMENT PROOF FOR COMPLETED Q- VALUES
We will prove that π (cid:48) completed =
softmax(logits + σ (completedQ)) produces a policy improvement.
We will start by showing that π (cid:48) completed is produced by a speciﬁc instance of Algorithm 3.
We will then prove that any instance of Algorithm 3 produces a policy improvement.
Algorithm
Require: π, v π .
Require: q ( a ) for each visited action.
The visited actions can be from any distribution.
Initialize π (cid:48) with π .
For the visited actions:
If q ( a ) > v π , increase the π (cid:48) ( a ) logit.
If q ( a ) < v π , decrease the π (cid:48) ( a ) logit.
return π (cid:48)
C. S PECIFIC INSTANCE
Algorithm 3 is more general than the usage of the completed Q-values.
Speciﬁcally, Algorithmwould produce π (cid:48) 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. P OLICY IMPROVEMENT PROOF FOR ANY INSTANCE
We will now prove that π (cid:48) from Algorithm 3 satisﬁes (cid:88)
[Equation]
Notice that π (cid:48) ( a ) for any unvisited action a will be c
z π ( a ) , with a normalization constant c z > 0 .
For one visited action: Let’s denote the visited (aka expanded) action by a ex .
First, if π ( a ex ) =then v π
= q ( a ex ) and the policy will remain unchanged.
Let’s now consider the case with π ( a ex ) < 1 .
The v π can be rewritten as
[Equation]
Let’s denote the weighted sum by q miss :
[Equation]
We notice that the q miss does not change if scaling π by a constant c z > 0 .
We will now rewrite the left-hand side of Inequality 19 to use q miss : (cid:88)
[Equation]
[Equation] [Equation]
[Equation]
The right-hand side of Inequality 19 can be also rewritten:
[Equation] [Equation]
[Equation]
With these rewrites, Inequality 19 becomes
[Equation]
[Equation]
[Figure] [Figure]
The q ( a ex ) − q miss can negative, zero or positive.
If q
( a ex ) =
q miss , the
inequality is
satisﬁed.
If q ( a ex ) > q miss , we want π (cid:48) ( a ex ) ≥ π ( a ex ) .
If
q ( a ex ) < q miss , we want π (cid:48)
( a ex ) ≤ π ( a ex ) .
We do not know q miss so we cannot use it in an algorithm.
We will instead show that q ( a ex ) > q miss is equivalent to q ( a ex ) > v π , when π ( a ex ) < 1 :
[Equation]
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 unmodiﬁed, the algorithm is equivalent to using only one visited action.
If the logits of the other visited actions are modiﬁed by Algorithm 3, they can further help to improve the policy.
D M IXED VALUE APPROXIMATION
We will construct an approximation of v π .
The exact v π is deﬁned by v π = (cid:80) 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
[Figure]
Although, evaluation with stochastic Gumbel noise is not worse.
During training, MuZero and
Gumbel MuZero beneﬁt from explorative acting proportional to the visit counts.
these inputs, we approximate v π by a consistent estimator:
[Equation]
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.
E D ERIVATION OF THE DETERMINISTIC ACTION SELECTION
We will derive Equation 14 from Equation 13:
[Equation]
The simpliﬁcation was possible, because additive terms independent of a do not affect arg min a .
While widely applicable, the deterministic action selection provides only a small beneﬁt on 9x9 Go (Figure 7).
F E XPERIMENTAL DETAILS
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.
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
[Figure]
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 speciﬁc interval.
In all Go and chess experiments, Gumbel MuZero scales the Q-values by c visit = 50 and c scale = 1 . 0 .
On the perfect-information game of Go, Gumbel
MuZero is not very sensitive to the scale of the Q-values.
Any c visit ≥ 50 produced similar results (Figure 8a).
In each phase of Sequential Halving, we use at least one new visit.
For example, in the ﬁrst phase, we update ˆ q ( a ) by max (cid:16) 1 , (cid:106)
[Figure]
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 A n +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 ﬁrst 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 beneﬁts from the same exploration (Figure 8b).
Figure 9 shows the importance of the number of simulations at evaluation time.
For example, in the ﬁrst 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.
[Table]
F. N ETWORK A RCHITECTURE
In all Go and chess experiments, we used a modiﬁed 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 ﬁgure 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 signiﬁcant improvements.
For the architecture search, the Elo was computed at a ﬁxed computational budget of 300ms per move.
For example, the budget allows us to do 3200 simulations when using a network withplanes and 32 blocks with bottlenecks and broadcasting (achieving 10 4 inferences/second).
Here the Elo is anchored to have the base AlphaZero architecture at 0 Elo; 95% conﬁdence 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 efﬁcient 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. class [Table] [Table]
[Table]
.""" def
[Table] [Table] [Table] [Table]
[Table]
[Table]
def __call__(self,
[Table]
Tensor, call
[Table]
CallArgs
[Table]
[Table]
[Table]
) x
[Table]
self
[Table]
[Table]
is
[Table]
training
[Table]
call
[Table]
args
[Table]
is
[Table]
training, test
[Table]
local
[Table]
=call
[Table]
[Table]
local
[Table]
stats) return self
[Table]
_non
[Table]
[Table]
.""" def __init
[Table]
[Table]
[Table]
int, [Table] first [Table]
[Table]
[Table]
last
[Table]
, name
[Table]
[Table]
[Table]
[Table]
name
[Table]
assert stack
[Table]
self
[Table]
[] for i, make_op in [Table] [Table] first
[Table]
last
[Table] [Table] : self [Table]
blocks
[Table]
append(BasicBlock( make_inner_op
[Table]
make_op,non
[Table]
lambdax:
[Table]
if i == stack
[Table]
size
[Table]
else jax.nn.relu, name=f’basic_{i}’)) def __call__(self, x: Tensor, call_args: CallArgs): res = x for b in self._blocks: res = b(res, call_args) return jax.nn.relu(
[Table]
+ res) class
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure] [Figure] [Table] [Figure]
[Table]
[Figure]
[Table] [Figure] [Figure] (num_layers: int, output_channels: int, bottleneck_channels: int, broadcast_every_n: int): blocks =
[ BasicBlock(make_inner_op=make_conv(output_channels), non_linearity=jax.nn.relu, name=’init_conv’)
] for i in range(num_layers): if broadcast_every_n > 0 and i % broadcast_every_n =
=
broadcast_every_n - 1: blocks.append(BroadcastResBlock(make_mix_channel_op=make_conv(output_channels, kernel_shape=1), name=f’broadcast_{i}’))
elif bottleneck_channels > 0:
blocks.append(ResBlock( stack_size=4,make_first_op=make_conv(bottleneck_channels, kernel_shape=1), make_inner_op=make_conv(bottleneck_channels, kernel_shape=3), make_last_op=make_conv(output_channels, kernel_shape=1), name=f’bottleneck_res_{i}’)) else:blocks.append(ResBlock(stack_size=2, make_first_op=make_conv(output_channels, kernel_shape=3), make_inner_op=make_conv(output_channels, kernel_shape=3), make_last_op=make_conv(output_channels, kernel_shape=3), name=f’res_{i}’))
[Figure]