Training pipeline in AlphaGo

Two-stage training pipeline

Neural Network Architecture

• 13-layer CNN
• input feature/state size: 19x19x48
• The rules of Go are invariant under rotation and reflection. AlphaGo exploits this symmetrics by passing a mini-batch of all 8 positions into the policy network or value network and computed in parallel.

Monte Carlo tree search (MCTS) in AlphaGo

• Monte Carlo method: uses random sampling to estimate the expected outcomes for deterministic problems.

• MCTS is used to simulate potential actions to the best results without traversing every node in the search tree.

Training target

• Rollout policy network & SL policy network: maximize the likelihood of the human move $a$ selected in state $s$.

$$\Delta \sigma \propto \frac{\partial \log P_{\sigma}(a|s)}{\partial \sigma}$$

• RL policy network: maximize the expected outcome of the selected move.

$$\Delta \rho \propto \frac{\partial \log P_{\rho}(a|s)}{\partial \rho} z_t$$

• Value network: minimize the mean squared error (MSE) between the predicted value $v_\theta(s)$, and the corresponding outcome $z$.

$$\Delta \theta \propto \frac{\partial v_\theta(s)}{\partial \theta} (z - v_\theta(s))$$

networks

Results

AlphaGo won Fan Hui, the winner of 2013, 2014 and 2015 European Go championships, by 5:0 in a formal five-game match over 5-9 October 2015.

Computing Expenses

• Rollout policy (MCTS): approximately 1,000 simulations per second per CPU thread on an empty board.
• SL policy classification: around 3 weeks for 340 million training steps, using 50 GPUs.
• RL policy network: trained for 10,000 mini-batches of 128 games, using 50 GPUs, for one day.
• Value network: trained for 50 million mini-batches of 32 positions, using 50 GPUs, for one week.

How does it differ from AlphaGo?

1. AlphaGo Zero is trained solely by self-play, without any supervision or use of human data.

2. AlphaGo Zero only uses the black and white stones from the Go board as its input.

3. AlphaGo Zero uses a single neural network, rather than separate policy and value networks.

4. AlphaGo Zero does not use “rollouts” to evaluate positions and sample moves. Instead, it uses a simpler tree search that relies upon this single neural network.

MCTS in AlphaGo Zero

Basically the same as AlphaGo. Sightly different in evaluate and play phase.

• Evaluate: Without random rollouts as in AlphaGo, the action value $V(s)$ of the node $s$ are fully evaluated by the neural network.

$$(P(s,\cdot), V(s)) = f_\theta(s)$$

• Play: Once the search is complete, search probabilities $\pi$ are returned, proportional to $N^{1/\tau}$.
  • $N$: the visit count of each move from the root state.
  • $\tau$: temperature param that controls the level of exploration. ($\tau=0$: greedy search; $\tau \rightarrow \infty$: uniform sample)

$$\pi(a|s_0) = \frac{N(s_0,a)^{1/\tau}}{\sum_b N(s_0,b)^{1/\tau}}$$

• the child node corresponding to the played action becomes the new root node; retain the subtree below this child, discard the remainder.

MTCS

Network Architecture

• 1 convolutional block followed by 19 or 39 residual block.

  • Conv block: conv $\rightarrow$ BN $\rightarrow$ relu
  • Res block: conv $\rightarrow$ BN $\rightarrow$ relu $\rightarrow$ conv $\rightarrow$ BN $\rightarrow$ skip-connect $\rightarrow$ relu

• Two output heads for computing the policy and value.

• Feature size: $19 \times 19 \times 17$ (8 stone history x 2 player existence + 1 color), no hand-crafted features as in AlphaGo.

$$s_t = [X_t, Y_t, X_{t-1}, Y_{t-1}, \dots, X_{t-7}, Y_{t-7}, C]$$

• Similar to AlphaGo, AlphaGo Zero also augments the training dataset by sampling random rotations or reflections of the positions during MCTS.

Results

Performance of AlphaGo Zero

AlphaGo Zero used a single machine with 4 tensor processing units (TPUs), whereas AlphaGo Lee was distributed over many machines and used 48 TPUs. AlphaGo Zero defeated AlphaGo Lee by 100 games to 0.

Results

Network Architecture comparison

How does it extend AlphaGo Zero?

A more generic version of the AlphaGo Zero that accommodates, without special casing, a broader class of game rules.

• Use the same algorithm, network architecture and hyperparameters for chess and shogi, as well as Go.

• Go games have a binary win or loss outcome, whereas both chess and shogi may end in drawn outcomes. The game outcome is $z=\{0,\pm1\}$: -1 for loss, 0 for draw, 1 for win.

• To accommodate a broader class of games, AlphaZero does not assume symmetry. AlphaZero does not augment the training data and does not transform the board position during MCTS.

• AlphaZero always generate self-play games with the newest player of last iteration, instead of the best player from all previous iterations as in AlphaGo Zero.

State Representation

feature size: $N \times N \times (MT + L)$:

• $N \times N$ represents the board size.

• $T$ represents the board positions history at time steps $[t-T+1, \dots, t]$.

• $M$ represents the presence of player’s pieces.

• $L$ denotes the player’s color, the move number and other states of special rules.

• Illegal moves are masked out by setting their probabilities to zero, and re-normalising the probabilities over the remaining set of legal moves.

Representation

Results

Performance

Results

Comparison with specialized programs

What’s different?

MuZero extends AlphaZero to a broader set of environments (e.g. Atari2600), including single agent domains and non-zero rewards at intermediate time steps.

No rules were given in advance, MuZero learned the rules by its internal networks.

• State transitions: AlphaZero had access to a perfect simulator of the environment’s state-to-state dynamics. In contrast, MuZero employs a learned dynamics model to estimate the simulator.
• Actions available: AlphaZero used the set of legal actions obtained from the simulator to mask the policy network at interior nodes. MuZero does not perform any masking within the search tree.
• Terminal states: AlphaZero stopped at the terminal states directly provided by the simulator. MuZero does not give special treatment to terminal states and always uses the value predicted by the network.

Training process

Overall training process of MuZero.

Part 1: MCTS search

• Aside from $(Q, N, P)$ as before, the edges in MuZero MCTS also store intermediate reward $R(s,a)$ & state transition $S(s,a)$.

• The evaluating value is the $k$-step cumulative discounted reward:

$$V(s_t) = \sum_{\tau=0}^{k-1} \gamma^\tau r_{t+\tau} + \gamma^k v_{t+k}$$

• Because the value above is unbounded, the action-state value $Q$ would be normalized:

$$Q(s,a) = \frac{1}{N(s,a)} \sum_{i=1}^N 1(s,a,i) V(s_L^i)$$

$$\bar{Q}(s,a) = \frac{Q(s,a) - \min Q}{\max Q - \min Q}$$

MCTS

Part 2: Self-play process

• Self-play process is basically the same, except intermediate rewards $u_t$ are considered.

• The trajectories of self-play would be stored into a replay buffer (off-policy RL).

• The final outcome $z_t$ is computed as the cumulative discounted reward:

$$z_t = u_{t+1} + \gamma u_{t+2} + \dots + \gamma^{n-1}u_{t+n} + \gamma^n v_{t+n}$$

Part 3: Network training

A trajectory is sampled from the replay buffer.

1. The representation function $h$ generate the first hidden state representation $s_0$.
2. The dynamics function $g$ recursively predicts subsequent rewards $r_t$ and states $s_t$ based on the predicted state $s_{t-1}$ and the real action $a_t$.
3. The prediction function $f$ predicts the distribution $p_t$ and value $v_t$ for these predicted states.
4. Train the model by minimizing the loss.

$$l = (z-v)^2 - \mathbf{\pi}^t\log \mathbf{p} + (u-r)^2+c||\theta||^2$$

The network architecture is basically the same as in AlphaZero, but with 16 instead of 20 residual blocks in representation and dynamics networks.

training process

Reanalyze

MuZero Reanalyze revisits its past time steps and re-executes MCTS search using the latest model parameters, potentially resulting in a better-quality policy than the original search.

• This fresh policy is used as the policy target for 80% of updates during MuZero training.

Normal training loop.

Reanalyze training loop.

Results

Results

Results

Resources

Other useful materials

• Introduction to Monte Carlo Tree Search: The Game-Changing Algorithm behind DeepMind’s AlphaGo | AI研习社.

• 【AlphaGo系列-1】AlphaGo详解 | 知乎专栏.

• AlphaGo Zero 简明工作原理 | 知乎专栏.

• 最强通用棋类AI，AlphaZero强化学习算法解读 | 知乎专栏.

• 终极版AlphaGo，DeepMind新算法MuZero作者解读 | AI研习社.