Reinterpreting AlphaZero

While teaching reinforcement learning I kept asking myself what AlphaZero teaches us about RL. That question has lead to this post.  This post generalizes AlphaZero to a larger class of RL algorithms by reinterpreting AlphaZero’s policy network as a belief function — as the probability that a is the optimal action at state s.  This gives rise to a “belief gradient” algorithm.

A fundamental question is whether belief gradient is a better conceptual framework for RL than policy gradient.  I have always felt that there is something wrong with policy gradient methods.  The optimal policy is typically deterministic while policy gradient methods rely on significant exploration (policy stochasticity) to compute a gradient.  It just seems paradoxical.  The belief gradient approach seems to resolve this paradox.

The action belief function and a belief gradient algorithm. For a given state in a given Markov decision process (MDP) with a finite action space {\cal A} there exists a deterministic optimal policy \pi^* mapping each state to an optimal action \pi^*(s) \in {\cal A}. For each decision (in life?) there is some best choice that is generally impossible to know. I propose reinterpreting AlphaZero’s policy network \pi_\Phi(a|s) as giving an estimator for the probability that a is the best choice — that a = \pi^*(s).  To make this reinterpretation explicit I will change notation and write B_\Phi(a|s) for the belief that a is the optimal choice.  The notation \pi_\Phi(a|s) is generally interpreted as specifying stochastic (erratic?) behavior.

To make the mathematics clearer I will assume that the belief B_\Phi(a|s) is computed by a softmax

B_\Phi(a|s) = \mathrm{softmax}_a \;s_\Phi(a|O_\Phi(s))

where O_\Phi(s) is a vector representation of the state s.  The softmax is not as important here as the idea that O_\Phi(s) provides only incomplete knowledge of s.  Further observations on s are possible.  For example, we can grow a search tree T_\Phi(s) rooted as s.  For a given distribution on states s we then get a distribution on the pairs (O_\Phi(s),T_\Phi(s)).  We assume that we also have some way of computing a more informed belief B_\Phi(a\;|\;O_\Phi(s),T_\Phi(s)). AlphaZero computes a replay buffer containing pairs (s,B(a)) where the stored distribution B(a) is B_\Phi(a\;|\;O_\Phi(s),\;T_\Phi(s)).  Ideally the belief B_\Phi(a\;|\;O_\Phi(s)) would match the marginal beliefs over the more informed search tree results. This motivates the following were R denotes the replay buffer.

(1) \Delta \Phi \propto E_{s,B(a) \sim R,\;a\sim B(a)}\;\nabla_\Phi\;\ln B_\Phi(a\;|\;O_\Phi(s))

This is just the gradient of a cross-entropy loss from the replay buffer to the belief B_\Phi(a|s).

Some AlphaZero details. For the sake of completeness I will describe a few more details of the AlphaZero algorithm.  The algorithm computes rollouts to the end of an episode (game) where each action in the rollout is selected based on tree search. Each  rollout adds a set of pairs (s,\;B(a)) to the replay buffer where s is a state in that rollout and B(a) is a recording of B_\Phi(a\;|\;O_\Phi(s),T_\Phi(s)). Each rollout also has a reward z \in \{-1,1\} (was the game won or lost).  AlphaZero stores a second set of replay values (s,z) for each rollout state state s  where z is the reward of that rollout. The replay pairs (s,z) are used to train a value function V_\Phi(s) estimating the reward that will be achieved from state s.  The value function V_\Phi(s) acts as a kind of static board evaluator in growing the tree T_\Phi(s).

Improved generality. The abstract mathematical formulation given in equation (1) is more general than tree search.  In the belief B_\Phi(a\;|\;O_\Phi(s),\;T_\Phi(s)) we can have that T_\Phi(s) is any information about s that usefully augments the information provided in O_\Phi(s).

A belief gradient theorem. While I have reinterpreted policies as beliefs, and have recast AlphaZero’s algorithm in that light, I have not provided a belief gradient theorem.  The simplest such theorem is for imitation learning.  Assume an expert that labels states with actions.  Optimizing the cross-entropy loss for this labeled data yields a probability B_\Phi(a|s).  Learning a probability does not imply that the agent should make a random choice. Belief is not the same as choice.

This entry was posted in Uncategorized. Bookmark the permalink.

Leave a Reply

Fill in your details below or click an icon to log in: Logo

You are commenting using your account. Log Out /  Change )

Facebook photo

You are commenting using your Facebook account. Log Out /  Change )

Connecting to %s