How does policy network learn in AlphaZero?

Question

I'm currently trying to understand how AlphaZero works. There is one thing with the training of the AlphaZero's policy head that confuses me. Basically, in AlphaGo Zero's paper (where the major part of AlphaZero algorithm is explained) a combined neural network is used to estimate both, the value of the position and a good policy. More precisely, the loss function used is:

$$L = (z-v) - \pi^t \log(\textbf{p}) + c \Vert \theta \Vert$$

where $z$ is the outcome of the game, $v$ is the value estimated by the neural network, $\pi$ is the policy calculated by the MCTS and $\textbf{p}$ is the policy predicted by the neural network.

I would like to focus on the policy head loss. Basically, we are trying to minimize the difference between the policy calculated by the MCTS and the policy predicted by the neural network. That makes sense when the player has won the game, but it doesn't (at least from my point of view) when the player has lost it. You would be teaching your neural network a policy that has lost. Maybe the loss was unavoidable, but if it wasn't that's definitely not the policy we want to learn.

I have programmed a slightly simplified version that works well with Tic Tac Toe. But for Connect 4 some problems related to this arise. Basically, it learns a bad policy. At the beginning of the training, the values estimated for each board are quite random, and that makes the policy shift to a random (and wrong) direction. At the same time, that makes the value function to be wrong (because we are losing games that we could have easily won), worsening even more the policy.

I suppose that with enough training this problem disappears. The correct value and policy should backpropagate from the leaf nodes of our simulation. Even if the neural network policy gives a probability of 0 to the optimal action, thanks to the Dirichlet noise added to the probabilities the MCTS can find that optimal action and learn it.

However, several things confuse me:

1. In AlphaGo's paper, they take into account whether if the outcome of the game has been a win or a loss when training the policy network with reinforcement learning. More precisely, the optimization made is

   $$\Delta p \propto \frac{\delta \log{p(a_t|s_t)}}{\delta p} z_t $$

   where $z_t = 1$ if we have won or $z_t = -1$ if we have lost. So DeepMind's take into account if the action was good or not and change the direction to optimize.

2. I haven't found anywhere in AlphaGo Zero's paper that we are training just with the examples where the player has won, so they might be using all the data gathered, including also the examples where the player has lost. As far as I know, they don't mention anything related to this problem.

3. $\pi$ (the policy provided by the MCTS) is calculated as the exponentiated visit count of each action

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

   where $\tau$ is a parameter that controls the ``temperature''. DeepMind's team sets $\tau = 1$ during the 30 first movements to ensure exploration. After that, they set it to $\tau \approx 0$ to ensure that the action that is considered the best one (and thus has been simulated more times) is the one played. However, that means $\pi$ is something like

   $$[0,0,\dots,0,1,0, \dots, 0]$$

   making the policy changes a bit agressive and especially harmful if the movement is not the good one (making it even harder to recover from a bad action).

Am I missing something? Is this the intended way of working of the algorithm? Is there any way to overcome the learning of bad policies?

Answer 1

I'll first address the big-picture intuition and then address each point separately.

First of all, the tree search tries to find the best policy at each turn. Losing the game doesn't necessarily mean the policy is bad, often it's just caused by us picking a bad move at least somewhere in the game (thanks to the policy temperature $\tau$), this could even be long before or after the current turn. Training can progress as long as the tree-derived policy is (at least slightly) better than the raw output of the neural net.

You're right that a bad value head can cause tree search to find a bad policy. However near the end of the game we are always getting some true signals:

• the tree search encounters some terminal states which can directly influence the policy to be better
• the value head can at least learn the right outcome for states close to the end

These two mechanisms cause both value and policy head to become better at the endgame, this means tree search and its policy can become better at near-endgame positions, which means we actually pick better moves to play which means the value head gets to learn better values on average.

So yes, early on during training the network can lean some wrong things for the early- and mid-game. There is always some pressure towards correctness coming from the endgame though, and with enough training and the right hyperparameters this can propagate all the way back to the start of the game.

To address your individual questions:

1. Yes, AlphaGo was doing something policy gradient-like where they used the game outcome to directly change the policy. AlphaZero works very differently, it's best to view them separate.

2. Indeed, the tree-derived policy of all moves is used for training, regardless of whether the current player ended up winning or losing the game.

3. This temperature is only used to pick an action during selfplay, not for the training target which still just uses $\tau = 1$. (as a side note, it happens to still be a good idea to use a temperature for the training target for regularization, LeelaChessZero and KataGo do this, see here for an explanation)