MuZero: idea

We have already seen what the MCTS algorithm is and how we can empower it with the use of Neural Networks to learn to play (not only) videogames with the AlphaZero algorithm.

The key idea of AlphaZero is that, in order to avoid the MCTS algorithm to fall short when using it in games with very big state spaces, we can use neural networks to approximate the Value and the Policy of the tree nodes, so that the limits of the MCTS algorithm are overcomed.

MuZero builds on the same idea of AlphaZero, but goes a step further. In fact, in AlphaZero, we had complete access to game information: for every action we decided to take, we played such move in an exact copy of the actual game, and we could actually see how the environment state changed as a result of that action.
This is a luxury that we can’t always have: can we, for example, in real life, see what the results of taking an action really is? Well, no until we have taken it, but we have learned to predict what the results of taking most common actions are, and we play along the “game” of life like this.

MuZero builds on the idea that with the use of neural networks, we can actually learn the dynamics of the game itself, so that when we have to plan in order to look for a move, we can guess not only how good the move is, but also what will happen as the result of taking it. In the end, we do not need to play along with the game environment, because we will use neural networks to “simulate” it.

This is the key concept of MuZero, and there is a particular aspect, which I think is the most incredible one, but also the most difficult for my mind to understand and “accept”: in fact, the neural network that will approximate the dynamics of the environment, will NOT be trained to approximate the environment itself, that is, we are not trying to create an approximated version of the environment by training the network with information about it, no, we are trying to reach our objective (i.e. minimize the loss function, that we will see later) but giving no constraints about how the neural network that predicts the dynamics should do so.
The neural network is free to represent its internal state as best as it wants to in order to represent the dynamics of the game, but it has no “human” restriction on how to do so.
We do not know how internally the neural network represents the dynamics of the environment, and it is free to do so with something that is not correlated with the real environment at all.

This, I think, is the key powerful idea of MuZero, as it gives the neural network the possibility to really flourish by not giving it any “constraint” and letting it “imagine” the dynamics of the environment as it wants, and, in general, I think this is a fundamental idea in order to build algorithms that aims towards build AGIs.

Thank you DeepMind for opening my mind to such kind of ideas! As a computer scientist, in the beginning, I was very reluctant to understand and accept it, but now I start seeing things differently!

MuZero: implementation

So we are ready to tackle a simple but working implementation of MuZero. In this case, I have decided to implement the pseudo-code that the authors at DeepMind have kindly shared along with the academic paper of the algorithm itself.

The extra effort for understanding it is really worth as the code is very well structured into classes, the flow is clear, easy to read and implement. Also, it is from the authors of the algorithm itself, so who better than them can provide a guide for the actual implementation?

As usual, I will try to simplify and explain the basic ideas and concepts about the algorithm, and I will not focus on the details of the (many) optimizations that we will encounter along the way: each one of such optimizations and concepts would require an article (and more) on its own!

In any case, I assume familiarity with both the MCTS and AlphaZero algorithm. If you need a refresher, you can read the previous articles.

Helpers

Let’s start with defining helper classes that we will use along the way. There are a lot of them, but they are quite simple to understand and will simplify the actual algorithm code later on.

MinMaxStats

The MinMaxStats object contains minimum and maximum values encountered along MCTS search. We will use it to bound the value of the node between 0 and 1, so that the Value of the node is in the same range of the Policy of the node during the search.

Config

The MuZeroConfig object stores the configuration of hyperparameters used by the algorithm for search and training. Some of such parameters are common for all the games, some are very specific, and each game has its own specific configuration. The meaning of each hyperparameter will be explained along the way when we will use them in the actual algorithm.

CartPole config

As usual, our environment of choice is CartPole-v1. It is simple enough to be solved by such simple implementation, but complex enough to check if learning is actually happening.

I have found that this implementation of MuZero is quite sensible about chosen hyperparameters, but I have managed to find values working for the CartPole environment. Also, after the MCTS will have completed its search, we can add an extra exploration probability to the choice of the action by using the visit_softmax_temperature function.

Action and Player

The Action and Player classes do not have anything special in this implementation, they are just wrapper objects.

Node

The Node class represents the usual MCTS node. Nothing that we have not already seen, but notice the to_play attribute that, in two player games, would change every time we change level of the MCTS tree, to take into account that the value of the node for one player, is the opposite for the other one.

ActionHistory

The ActionHistory is an helper class that stores the history of actions we have taken along a MCTS search.

Environment

The Environment is a wrapper for the gym environment, CartPole-v1 in our case, and defines how to actually act in the environment just by using the step method of the openAI gym interface.

Game

The Game represents a single episode of interaction with the environment. In particular, the store_search_statistics will collect data about visit counts and node values during the MCTS search, and the and the make_target allows to build the target values for the neural networks to aim for:

• for the target Value we will use a technique called bootstrapping so that the target value of a given node is the discounted root value of the MCTS tree N steps into the future (where N is given in our config by the td_steps hyperparameter), plus the (discounted) sum of all rewards until then (notice that in AlphaZero we did not have the rewards taken into account, but now we have). Every time we use an approximated value to estimate another value, like in this case, we say that we are bootstrapping.
• for the Policy of the node (i.e. the probability of exploring each action) we use the number of visits, so our target is the MCTS exploring policy.

Also, and this is very important, as we are looking N steps into the future for doing our estimations, it may happen that these states falls after the end of the actual game. It is important that in these cases our target has null or zero target values, such that these states are learned by the neural network and treated as “absorbing states”. In fact, for the neural network, there is no such a concept of “end game”, and it will continue to imagine how its dynamics will expand forever! By training for absorbing states, the value of such nodes will be lower than other states, and this will be enough to avoid them and choose other, better actions, in the actual game.

Replay Buffer

Nothing major to say about the ReplayBuffer, only that for this simple implementation we do not need any kind of priority mechanism (such as PER, Prioritized Experience Replay) for choosing episodes and positions: completely random choices are enough.

Network

The Network class represents the brain of MuZero, and is worth understanding the underlying ideas correctly. We have different neural networks that approximates different things:

• The Representation network allows us to go from an actual observation of the environment, to an internal state of the neural network. We will call this state an HiddenState. Note that the output size of the HiddenState is not the same of the actual observation size of the environment, because, as we said, they are completely unrelated!
• The Value network approximates the actual value of an HiddenState (how good it is) as a single output number.
• The Policy network approximates the MCTS exploring policy in the form of probabilities for each action.
• The Reward network approximates the reward of an HiddenState and Action, it will tell us how good is to take this single action in that state.

The MuZero algorithms builds on the combined usage of such networks:

• initial_inference is used for example at the start of a new game and during the learning process and takes as input the actual environment observation, transforming it into an HiddenState using the Representation network and then using such output HiddenState to query the other networks for gathering the predicted Value and Policy (Reward is assumed 0 as there is no previous action taken).
• recurrent_inference is the function that is used to transition from an HiddenState to the next one, which is what we will do to simulate the game in the “mind” of MuZero: given the current HiddenState and an Action, we use the Dynamics network to get the next HiddenState, and then we use the other networks to approximate Value, Policy and Reward.

Using initial_inference at the root of the MCTS search and then recurrent_inference on the intermediate nodes we have a powerful search mechanism to collect statistics about MCTS values, all in MuZero brain. Wow!

As usual, the better the networks approximates the real Values, Policy and Rewards, the more “expert” our plays will be. We will see later on how to train the networks to correctly approximate such quantities.

At the end, the NetworkOuput is just a structure containing the Value, Reward, Policy and HiddenState predicted by the networks. Just a notice about this specific implementation where I have used the softmax activation function for the last layer of the Policy network, hence its outputs are not really raw logits as the pseudocode suggests, but already actual probabilities.

SharedStorage

The last helper class that we need before the actual algorithm implementation is the SharedStorage.

In theory, this is a shared space where to store Network objects after each training step, so that jobs (i.e. parallel threads) that finish collecting data about a game can pick the last up-to-date version, train it, and store a new version for other jobs to use it.

In practice, in our implementation, we have a single Network object and there are no threads, games are played one at time and so is network training.

Self-Play and Training

We are now ready to tackle the main MuZero algorithm. It is basically divided into two parts: the first part plays a game and stores gathered statistics putting them into the shared ReplayBuffer, the second part is about training the neural network sampling the data from the buffer. While these steps can be executed in parallel and by multiple threads, in this simple implementation everything happens sequentially.

Part 1: Self-Play

So the steps are quite simple: grab the most up-to-date network (remember, we only have one in this implementation), play the game, store the statistics.

And we already know how we are going to play games, right? We use the MCTS algorithm to select an action at each step. Notice how we create the root node of the MCTS calling the initial_inference function: we pass to it the initial observation of the environment, and the Representation network will take care of transforming it into an internal HiddenState. The first observation of the environment is the only thing that we will know about the real game!

Not much to say about the core MCTS algorithm, we have dedicated an entire article to it, and this is just a different implementation. Just notice the use of the recurrent_inference function. We use the Dynamics network, to go from one state to another of the imaginary game into the “mind” of MuZero.

Action selection is as usual proportional to MCTS visit counts, with a twist: we use a temperature value to encourage exploration, and we reduce such value based on the number of episodes we have trained on, so that in the beginning we favor exploration, and in the end, when we are confident about the moves, we favor exploitation.

The remaining MCTS accessory functions are almost self-explanatory:

• select_child: chooses the child that maximizes the ucb value.
• ucb_score: calculates the ucb value for a node. It is slightly different from the one used in the AlphaZero implementation, but the concept is the same: one term comes from the Value of the node (scaled to the range 0 to 1) and the other term is given by the MCTS policy. Objective of the function is to balance exploration and exploitation.
• backpropagate: at the end of the search it propagates up to the root the gathered information. Notice how at each level of the tree the Value of the node depends on the current player, so as to consider also 2 player games.
• expand_node: expands the current node creating children with the Policy values that comes from the network prediction. As in this implementation we are already using the softmax activation function for the network, we can use such probabilities directly, so I have commented out the logits approach.
• add_exploration_noise: last, but not least, when we create the root node for the MCTS we add to it Dirichlet noise again to favor exploration.

And this concludes the first part of the MuZero algorithm.

Part 2: Training

The second part of the algorithm is all about how to train the various neural networks. In particular, we train all of them trying to minimize the following:

That is, we are trying to minimize the differences between:

• The predicted Policy with respect to the MCTS policy.
• The predicted Values with respect to the actual values (i.e. the ones we calculated using the bootstrapping technique).
• The predicted Rewards with respect to the observed ones.
• Finally, a regularization term is added in order to penalize big network weights.

All parameters of the model are trained jointly minimizing such loss.

Main

There is not much to add to the MuZero algorithm, we have only to create some code to actually use it:

We use the usual plot functions in order to track progress. This time, we will also track time, as it will take much more time to train the networks compared to previous MCTS and AlphaZero implementations. In fact, if for AlphaZero you could go grab a cup of coffee, get back and enjoy the result, for the training of MuZero I would suggest leave it overnight and get a good night of sleep, in order to wake up with the following:

And that’s it, we just coded a simple but working implementation of MuZero!

You can write the code on your own (suggested) or find the repository on GitHub.

Besides the implementation, what I think its important to understand are the basic concepts and ideas that we discussed about along the way: they for sure gave me a different way about how to see things and about how to code!

Intro: MCTS, AlphaZero and MuZero for dummies!

Part 1: Monte Carlo Tree Search (MCTS) algorithm for dummies!

Part 2: AlphaZero for dummies!

Part 3: MuZero for dummies!

If you have followed the previous articles, you know that we are now at the most exciting part of the series, and we are going to take the fly: we are going to integrate the previous MCTS implementation with an AI Deep Neural Network, in order to train an expert policy to play video games!

If you do not think this is a big deal, keep in mind that the AlphaZero algorithm, along with its predecessor AlphaGo, and successor MuZero, “solved” the game of Go, which was at the time considered from the AI research community at least 10 years away to be solved.

This is incredible, and even more incredible is the elegance and simplicity of such algorithms which yes, in their full implementation have many optimizations that may be difficult to understand at first, but, at the core, have some basic ideas that make such algorithms work, and we will build on such core ideas to create an AlphaZero implementation on a small scale.

If you need a refresher on the MCTS algorithm, have a look at the previous article. Also, I will assume some familiarity with basic Neural Networks concepts. If you do not have any, I would suggest the MIT introductory course.

Let’s see more in detail what all of this is about.

The problem with the MCTS algorithm is that, even if it is better than brute-force search, as its search is guided towards more promising states, still, for complex games, the state space is so big that, in order to gather a statistically valid amount of data, the time required is not feasible, and the algorithm fall short.

The core idea of the AlphaZero algorithm, is to have a neural network to “guide” the search along the tree, producing values and exploring probabilities for the nodes. In the end, the MCTS algorithm is based on statistics and approximated data, and neural networks do exactly that, they approximate (arbitrarily complex) functions.

We will need to estimate two things:

• Value estimation: every time we need to do a rollout from the node, instead of doing an actual rollout, we will ask the neural network what is the (estimated) value of the node.
• Policy estimation: the probability to explore each node will not only depend on the number of times that it has been visited, but also on a probability given by the neural network, which will approximate the MCTS probability distribution.

We will use the combination of the Value and Policy estimations in order to guide our search along the tree to more promising nodes.

The closer the Value estimation is to real state value, the better the Policy will be and, vice-versa, the better the Policy, the better the Value estimation.

When the neural networks closely approximate the required values, we can then use them to choose an action at each step of the game, and play it like an expert.

Let’s see how to code it!

Here the usual code to import an openAI gym environment:

Notice, however, how this time we are using the version 1 of the environment (i.e. CartPole-v1), whose maximum possible sum of rewards is 500.

Now we have to define the Deep Neural Networks that we will use to approximate the Value and Policy of the tree nodes.

We are using 2 small, fully connected neural networks:

• The first will approximate the value of the node, hence it will have 1 output value.
• The second will approximate the probabilities of choosing each child node, hence it will have as many output values as actions available in the game.

Here our MCTS implementation, now combined with the use of the neural networks:

I will not go into the details of the basic MCTS algorithm, as it has been explained in a previous article, but let’s focus on the key differences:

• The rollout function now uses both the Value and Policy neural networks to approximate the values needed for guiding the MCTS node selection.
• The getUCBscore uses such values in order to balance the exploitation (i.e. try to choose a good action) and the exploration (i.e. exploring new actions). The exploitation is represented by the Value score (i.e. the average of all the rollouts started from that node and from its children nodes which then have propagated their values to root nodes) and the exploration is represented by the Policy score but is also affected by how many times the node has already been visited (i.e. nodes that have low visit count have higher values, so as to favor their exploration).

The rest is the usual MCTS algorithm stuff, we have just changed its selection mechanism using the neural networks to approximate the node values and policy.

We also need a structure where we can store the information about game episodes that we will use to train our neural networks to predict accurate values. We will call this structure a ReplayBuffer.

The ReplayBuffer stores, for each step of the game:

• The observation (i.e. state) of the game environment
• The target Value (i.e. the sum of rewards obtained in the the episode)
• The observation (i.e. state) of the game environment at the previous step
• The target Policy according to visit counts

We will then train our neural networks to accurately predict the target Value and target Policy, given the game state observations.

The strategy we will use to play the game is as usual to do a certain amount of searches in the tree for collecting reliable statistics and then picking the best action, which will be the one relative to the children node of the current state that has been visited the most (because, as we have seen in the MCTS algorithm, it is related to a good value anyway, given how the value of the node is calculated).

And here we are, defining the loop that will allow us to train the neural networks and plot the improvements over time:

Notice how we define some hyperparameters and proceed to play games and train the neural networks incrementally approximate closely the actual Value and Policy of the MCTS nodes.

And finally these are the plots of the rewards over episodes along with the moving average over the last 20 episodes, the Value loss and the Policy loss.

We can see that learning is definitely happening and in the final episodes of the training AlphaZero constantly reaches 500 rewards per episode, which is the maximum for this environment!

You can write the code on your own (suggested) or find the repository on GitHub.

I hope that this was both instructive and fun. For me, it is astonishing that with such few basic ideas we are able to write algorithms that learns to play by themselves.

As you have noticed, however, AlphaZero requires to know the “rules” of the game. Everytime we take a step, we are using a copy of the game environment to simulate an action and get all the required information on the evolution of the state of the game.

Wouldn’t be amazing if we are able to get the same results, but knowing nothing about the game we are playing? If AlphaZero is not enough magic then follow me in the next chapter, where we will implement MuZero, an evolution of AlphaZero with no knowledge at all about the game, and we will see how it is possible to learn everything about the game using neural networks!

Intro: MCTS, AlphaZero and MuZero for dummies!

Part 1: Monte Carlo Tree Search (MCTS) algorithm for dummies!

Part 2: AlphaZero for dummies!

Part 3: MuZero for dummies!

If you are a Computer Scientist like me, you likely have stumbled across all the research and results on AI that are happening these days.
Chances are, you already know about DeepMind and similar AI companies, and how they are revolutionizing the AI research and development in pursuit of AGI (Artificial General Intelligence).

For me, everything changed when I saw the AlphaGo Documentary, more or less a year ago. This is what I understood at the time: some very smart people from DeepMind managed to create an algorithm that learns to play video games and table games(some of which, like Go, are very difficult given their state space), through self-play, at a super-human level.

How did they do it? This is the question that stuck into my mind since then.

I am a Computer Scientist but I am not specialized in AI, but damn if I wanted to understand how they have done it. In the end, these are the video games that I was playing when I was a kid, and now there is the possibility to write algorithms that are able to learn how to play them.

I wanted to try to replicate their results on a small scale.

So this is where my journey started, reading some generic books and looking at videos about AI, and then graduating at Udacity Deep Reinforcement Learning course.

In this series of articles, which assumes you already have some familiarity with basic programming (we will use Python), basic ML (Machine Learning) and RL (Reinforcement Learning) concepts, and the algorithms themselves, I will try to explain how to implement a (very) simple but working version of the MCTS, AlphaZero and MuZero algorithms, and try to explain some of the key concepts behind.

This is the kind of implementation which I wish I had at the time I started to understand these concepts, because the only ones that I found may be much more better and complete, but are also quite complex in order to understand the basic ideas behind such algorithms.

So stay tuned and see you in the next articles!

Intro: MCTS, AlphaZero and MuZero for dummies!

Part 1: Monte Carlo Tree Search (MCTS) algorithm for dummies!

Part 2: AlphaZero for dummies!

Part 3: MuZero for dummies!

Roughly speaking, the MCTS allows us to search for the best moves to play relying on statistics rather than on full exploration, hence allowing us to navigate large state spaces with much less limits compared to brute-force approaches.
In fact, search is guided in order to explore promising nodes more often and the more statistics are collected, the more reliable the results obtained.

More precisely, the MCTS is divided in 4 different steps:

• Selection: we need a smart way to explore tree nodes that leads to more promising outcomes, hence we need to have a way to give value to nodes to choose which one to pick.
• Expansion: once we reach a leaf node (i.e. a certain state of the game) we need to expand the tree considering all the possible valid actions from that node (i.e. subsequent game states after taking valid actions).
• Simulation: we need to be able to play a completely random game from a certain node (i.e. game state). Such rollout will result in a random result if taken alone but, the more times we will do rollouts from such node, the more accurate the average value estimate of the node will be.
• Backpropagation: after the rollout, such results needs to be back-propagated up to the tree, in order to store such information in the upper nodes, so as to transfer information from leafs to root nodes.

To sum up, the more we apply the 4 steps of the MCTS algorithm, the more reliable and statistically valid information we gather about the next best action to take from a particular state of the game. Apply this for every move, and we have a strategy for playing the game like a boss!

Let’s see how to implement it and how to use it to play an openAI gym video game.

The first thing that we need to do is to create an openAI gym environment, which will be the test bed for our algorithm, and check for the number of possible actions and the observation space in the game.

For the CartPole environment, whose goal is to balance a pole by applying forces to the cart to which is attached, there are 2 possible actions (push cart to the left or to the right), and the observation is made of 4 values (cart position, cart velocity, pole angle and pole angular velocity). You can find detailed information about this environment here.

Next, the definition of the Node class, which will represent a node of the tree:

The Node class will store information that we need about the game, the children (i.e. the next states of the game), and value statistics. In particular:

• child: is a dictionary of child nodes which, for every possible action from that state, will tell us what the next state of the game is taking that action
• T: represents the sum of the value of the rollouts that have been started from this node
• N: represents the visit count of the node, i.e. how many times we have visited the node
• game: is the game environment in the current state represented by the node. Its a copy of the original game that we are currently playing, so that we can use it to simulate the search. It represents the current state of the game that will result in applying all the actions from the root node to the node itself
• observation: is the current state of the game in that node, in our example, it will be the 4 values as described above
• done: will tell us if the game at this point is ended, so as to stop the search
• parent: a connection to the parent node, for backpropagation
• action_index: is the index of the action that lead to this node (i.e. is the action that the parent of the node has taken to get into this node)

Now let’s define methods of the Node class that we will use for the search.

The first one will be the definition of the function that will give values to nodes. This is fundamental, as will guide our search along the tree:

The formula is called UCB (Upper Confidence Bound) and balances exploitation (i.e. pick the best known action) and exploration (i.e. explore new actions):

• The exploitation term is the current estimated value of the node (self.T / self.N).
• The exploration term is inversely proportional to the number of times the node has been visited, with respect to the number of visits of its parent node (sqrt(log(top_node.N) / self.N)).

There is also a tunable constant c that will give more or less value to the second term.

Next we need a function that creates the children of the current node:

Notice how we always use the whole game actions available (GAME_ACTIONS) and create a child for each action, associating to it a copy of the game environment, to which we apply such action (game.step(action)), and store the information of the resulting state of the game in the child node.

We are now ready for the core of the algorithm, where the search happens:

The search along the tree starts from a certain node and recursively picks a child node that maximizes the formula as defined before. This guides our search until a leaf node is reached. At this point, if it is the first time we visit such leaf node, we will play a rollout and store the statistics, if it is not the first time we visit such node, then we create its children and play a rollout from one randomly chosen child. In every case the updated statistics, both value and visit count are backpropagated to parent nodes up to the root.

We are almost there, we just need to define how to do a random rollout, using a copy of the game from the current node (i.e. in the actual game state):

And we need to define how to pick the next action after the search is done:

In particular, we will pick the child node from the current one that we are searching from which has the maximum visit count, as this is associated with a good value anyway, given how the MCTS formula works (i.e. it visits more often nodes that have higher estimated values).

We are all set now, and we need to define the policy that we will use to play the game:

The policy is straightforward: every time we need to do an action in the game, we will use the MCTS to explore and pick the best possible next move as just discussed. The search goes on for a certain number of iterations based on the value of a fixed parameter (MCTS_POLICY_EXPLORE) and, as discussed, the higher the time we spend collecting statistics with the search, the more reliable the results will be.

We will apply such policy for every move of the game so as to have an MCTS guided expert player.

And here we go, the code below allows us to play a game choosing moves according to the MCTS just defined, plotting the results at the end!

The CartPole-v0 environment used as example is quite simple, and the MCTS algorithm is able to collect enough statistics in a reasonable amount of time, making it a good player constantly reaching 200 points per game (which is the maximum allowed for the environment)!

Congratulations, you made it!

I hope you had fun as much as me in coding the MCTS algorithm.

You can write the code on your own (suggested) or find the repository on GitHub.

For me, the excitement in seeing that such generic algorithm is able to play different video games is very high, but nothing compared to what we will see next, that is, how we can use AI Deep Neural Networks, combined with the MCTS algorithm, to create a simple but working implementation of AlphaZero.

Intro: MCTS, AlphaZero and MuZero for dummies!

Part 1: Monte Carlo Tree Search (MCTS) algorithm for dummies!

Part 2: AlphaZero for dummies!

Part 3: MuZero for dummies!