A Graphic Guide to Implementing PPO for Atari Games | by DarylRodrigo | Towards Data Science

Written by Daryl and Daniel
Learning how Proximal Policy Optimisation (PPO) works and writing a functioning version is hard. There are many places where this can go wrong — from misunderstanding the maths and mismatching tensors to having a logical error in the implementation. It took a friend and me around a month of work to get to a fully running PPO, where we successfully implemented the algorithm and were confident that we understood how it operated. This guide summarises all the areas where we struggled building our final algorithm and intuitions we developed. This blog aims to cover the basic theory of PPO, how we decided to implement our version, working with tensor shapes, testing, and using the recorded metrics to debug what we built. I hope you enjoy it!
Because this blog post intends to cover PPO, it will assume some prior knowledge of basic reinforcement learning (RL). In particular the basic setup of an RL problem, some notion on how an actor-critic model works, concepts behind the Bellman equation and a vague understanding of what policy gradient methods are.
Github Repo — [link]
As a recap — PPO belongs to a class of algorithms called policy gradients (PG). The fundamental idea is to update the policy directly to make the probability of taking actions which provide a larger future reward more likely. This is done by running the algorithm in an environment and collecting the state change based on an agent’s actions — a collection of these interactions is called a trajectory. Once one or more trajectories have been collected, the algorithm will look at each step and check whether the chosen action led to a positive or negative reward and update the policy accordingly.
Trajectories are sampled by interacting with the environment step by step. To take a single step an agent selects an action and passes it into the environment.
PG algorithms decide on which action to take by accepting an input (the state or observation of an environment) and then, for a discrete action space, output a probability for each possible action. A neural network typically calculates this probability. The neural network’s output is dictated by its weights and is typically written as π_θ, where π denotes the policy and θ the weights of the network (important for the next part).
To select an action, the agent samples from the distribution of probabilities given by the neural network (called logits) — this action is then used to provide the environment with its next input. For example, in Figure 1.2, the action right has a 50% chance of being selected (left 12.5%, up 25% and down 12.5%). As a consequence of selecting an action from probabilities, exploration is inherently present in PG algorithms as the randomness in actions comes from the selection probabilities. As the agent becomes more sure of what the correct action is, exploration will decrease as the probabilities for incorrect actions decrease. Once the action has been taken, all metrics are stored for later use in training (observation, action, log action probabilities, value, reward, done).
Before diving into PPO, it is useful to breeze over the maths for PG algorithms from a high level to get a complete image of what the algorithm will be doing. However, since this blog post is about implementing PPO, we’ll have to skip over the details of how the derivation behind calculating the loss and updating the agent works (also, if you’re not interested in the maths behind PG, then feel free to skip the next 3 headings).
Updating the Policy: The building block of any PG algorithm is the formula on how a network is updated. The update formula in Figure 1.3 shows that the new weights of a network [1] are updated to the current weights [2] + the gradient of the current network given a specific state-action pair [4] times the learning rate [3]. This is the underpinning formula of how to update the network weights. Intuitively the gradient is the direction (+/-) of the weights in which the change of the policy will make taken actions more likely in a given state (reverse them to make them less likely).
A full explanation of how policy gradient updates work from a mathematical standpoint there is no better place than David Silvers lecture.
So how do we calculate this gradient? Thankfully, all the most common deep learning frameworks will keep a computational graph of information passing through a neural network and keep a record of the gradient part of the network.
Adjusting Weight Direction and Amount: With the update equation we know what direction the “gradient” of a specific state-action pair is — that is to say, to make that action (in a specific state) more likely by changing the neural networks weights in the direction of the gradient. Of course, we do not want to make every action more likely, if a bad result was obtained we want to reduce the probability of taking that action. This is what the advantage function [1] embodies — how much better or worse an action was than we expected. If a action led to a scenario where we are better off than before (positive advantage), the weights are change in the gradients direction, effectively making it more likely that that action is chosen again. If it’s worse, we decrease the weights in the direction of the gradient, making it less lightly
Another consideration is the amount weights are updated by, if just the advantage times the gradient were to be updated, there would be very aggressive updates (great explanation of why here). So in order to counteract this, we divide the update by the output of the policy [2], this is called importance sampling (how much weight to place on a particular update).
For example, if the action output was [2,3,5,1] and the action with value 5 (action 3) was chosen, the gradient would be divided by 5 to reduce the amount the gradient is updated. This is to stop action 3 quickly increasing in value, and hence and no other value getting the chance of being chosen (as it is done on a probability basis and if we have [2,3,100,1] the other actions are unlikely to be chosen), even if they might yield a higher return in the long run.
Chaining it up: This leaves the last step to the update rule, which comes from using the magic of the chain rule. Using this we can substitute out our partial derivative fraction with the gradient log of the action which was taken. This part is important because it’s what builds the foundation of PG and PPO algorithms.
Again, this is not intended to be a full explanation on the mathematical theory behind policy gradient algorithms but rather a quick review/recap on where the log values are coming from.
Visualising the update: When the theory is placed in practice it would look something as Figure 1.6 describes. First, the trajectory is run [step 1], and all relevant values are saved (log probabilities, value and reward — amongst other things). Once the trajectories have been completed, the discounted returns and advantages can be calculated [step 2] (great article here explaining what advantages are, but as a reminder, the advantages = discounted_return - expected_return where expected return is calculated by the critic as value). [step 3] Next the loss of each step is calculated by looking at the log probabilities and multiplying it by the advantage, just as shown in Figure 1.5. [step 4] Lastly we take the average of all these losses and do a backwards pass with gradient descent.
Original PPO paper — link
With policy gradient methods, once a few trajectories have been collected a model will be updated a number times using the same trajectories in order to learn from the new data that has been collected. Remember that the size of this update is determined by both the magnitude of the advantage as well as the importance sample (which is expressed as the log probability of the action).
PPO adds an additional factor which is the probability ratio to prevent large updates from occurring. This is defined as the difference between the original action log probability and the new model’s action log probability.
Problem: The problem with PG algorithms is that when an update is performed, its magnitude by which the policy change can happen (i.e. the weights change) can push the policy into a region that when another gradient descent is performed, all direction are essentially the same and the peak of the reward is missed. This is demonstrated in Figure 1.8 — a single update has missed the peak of the reward, and now when gradient descent is performed on a new trajectory it’s hard to say whether “left” of “right” movement will lead to a better policy loss. This causes the update to destroy the policy to an unrecoverable extent.
Solution: The PPO paper’s solution is to create a ratio between the original action log probability and the new model’s action log probability. If that ratio becomes too large or too small, it will be clipped to +/- a set amount (typically 0.2). This will prevent the policy from taking too large of an update step; this can be seen in Figure 1.9.
Mathematically this is expressed using a clipping function, also known as a surrogate function, in the PPO paper:
If you’d like to understand this in more detail, Arxiv Insights has an insanely good YouTube explanation of how PPO works from a theoretical standpoint.
Visual implementation of flow: From a practical standpoint following happens;
First, we collect a few trajectories using our initial model (model_1), making sure to save the action log probability for the actions taken. Using our trajectory, the initial model gets updated to model_2 and is now different from the original.
Then with the same states that we collected from the trajectories [3], the new log probabilities are calculated [2] (these will be different to the initial model as the model has changed).
The ratio between the new model and the old model’s action log probabilities is then used to calculate the scalar by which the new policy should be updated. Ratios are calculated on a per-step basis. If the ratio is too large or too small, it will be clipped according to the surrogate function.
Hopefully, the theory section was adequate to give some intuition about how an agent takes actions, the model is updated, why PPO works and why the clipping and probability ratio are required. The following section explains how we implemented PPO for ourselves — some of it might be obvious some not. However, it’s probably helpful to go through everything if you want to fully understand our implementation, as there will be differences. Throughout this section, there will be references made back to our code, so it might be helpful to split your screen in half, with one section having the blog post open and the other section the code, having said that, of course, read through this as works best for yourself!
As an overview there are three main classes and a file which runs the whole system:
Model.py [class] - This contains the logic for the agent to both, act on given inputs and learn from data collected from trajectories.
Memory.py [class] - This contains all the data from the trajectories and calculates the advantages. It also allows the agent to sample random sections of the trajectory to update the agent in batches.
Config.py [class] - This contains all the settings required in RL & PPO, such as learning rate, number of steps to take, and hyperparameters like epsilon for the clipping size.
Utils.py [File] - This is the main file which runs the sampling and training loop.
As you might be able to tell from the flow chart in Figure 2.1, the way an agent is trained is essentially an expansion of the overview in Figure 1.1. First, the program initialises all our variables, and then the algorithm runs in two distinct stages, collecting trajectories and updating the agent from the collected data.
Collecting Trajectories: Trajectories are collected by playing in a gym environment; the agent takes actions in the simulation based on its observations. Whilst this loop is active, the individual steps are saved in the agent’s memory.
Updating the Agent: Once enough steps are taken, the discounted returns and advantages are calculated — the agent then samples all the steps from the “memory” to calculate the losses and update the agent (all captured data is trained over). Once all the updates are performed, the loop starts at the beginning again.
To start with, all the components used in the experiment need to be initialised. The main pieces of this are:
The gym environment
The agent, which includes adding the model and the memory
Logging our experiments (we use weights and biases)
The environments which agents train in is typically created with OpenAI gym, a framework that creates a simulation environment for physics tasks and Atari games. Actions are entered which the environment reacts to — this change is then returned to the agent as a set of observations. With this also comes rewards and whether the game has finished letting the agent know whether or not their actions were “good” or “bad”.
When human players look at an Atari game, we see the RGB colours which make up the image and a steady change in frames as we play the game. We could feed every frame directly into our agent; this would be portrayed as a tensor of 160,192,3, representing the game's dimensions and RGB pixels. However, in typical Atari game implementations (and those done in the original paper) a few tricks are used to train more effectively, which we've included in the env.py file (this was heavily influenced and partly copied from Costa's CleanRL implementation).
Changing the game styles:
The first few tricks pertain to changing how the game fundamentally operates to help the agents along a little bit.
Episodic Life Env (envs.py:61) — This makes sure that the episode is only finished when all lives are lost, rather than when one. As a result, discounted rewards will get carried through the whole game rather than individual lives — this is important as it will allow the agent to get further into the game.
NoopResetEnv (envs.py:12) — As a technicality in the Atari engine, the game will always be the same if you start playing from the start. To get random behaviour in the game engine (rather than the same version every time), a random set of actions are taken to set the game in a pseudo-random state [paper]. This is critical to help the agent generalise its learnings.
FireResetEnv (envs.py:41) — If the game only starts when the fire command is pressed it will automatically hit fire when is initialised, this is so that the agent doesn’t have to learn how to start the game.
Max and Skip (envs.py:97) — Any action taken by the agent will be applied to four-time steps within the game. The max of each pixel for all the produced frames will be returned — similarly to max pooling. The reward for those frames will also be added up. The short of the long is that this is all to avoid flickering experienced when playing the game [Better explanation].
Clip Reward (envs.py:125) — Typically speaking in RL, agents train best if rewards are kept in a “reasonable range” (around +/- 1), this is to avoid too large updates when calculating the discounted returns. Clip rewards cut the rewards to +/- 1 for this reason.
Augmenting the Atari observation output:
Frame Stack (envs.py:187)- The agent is passed a stack of frames, representing four time steps in the actual game. This is so that the agent can have a sense of movement. For example, in Breakout, the ball might be in 4 different positions indicating movement in a particular direction.
WrapFrame (envs.py:134) — This is a custom method that reframes the Atari gym environment’s output. After all the augmentations have taken place, a tensor of (4, 84, 84) is produced - with WrapFrame, this is then converted to (84, 84, 4) for every frame. This is so that it is compatible with the CNN head from our agent.
What this ends up looking like, in terms of tensor sizes for a single environment, is the following:
All these techniques of augmenting the output and altering gameplay are designed for the agent to learn more effectively from the Atari game environment. It can be thought of as data augmentation in machine vision tasks.
Implementation: The augmentations are implemented in env.py, where they are grouped in a thunk function called make_env(env_id, seed). This function can be invoked at any point, specifying the seed and game type to return an environment; it is implemented this way to help when creating multiple environments.
Note: One might have noticed that there is an Augment input block in Figure 2.2, which adds a dimension to the input of the agent, this is present because the PyTorch 2D CNN block expects a batch to be entered. In our case, this is a batch of one.
We chose to implement multiple environments to take advantage of our home DL machines’ extra CPU and GPU. This sped up training and was a good learning exercise on how multiple environments worked.
When multiple environments are added, all outputs from the gym environment gain an extra dimension n, where n is the number of environments. Our agent’s input, however, is not any different. That is because we already augmented the input in the single environment to have 4 dimensions (as seen in Figure 2.2.).
Multiple environments are implemented is in our Config.py, here the environments are initialised by wrapping the make_env function in the Open Baseline's DummyVecEnv, as seen below.
envs = VecPyTorch(	DummyVecEnv(		[make_env(env_id, self.seed+i, i) for i in range(self.num_env)]	), self.device)
This function expects “a list of functions that will create the environments (each callable returns a Gym.Env instance when called)”, which is why make_env returns a thunk function that creates a gym-env rather than the gym-env itself. The VecPyTorch is essentially based on VecEnv, which "abstracts asynchronous vectorised environments". It's an easier way of interacting with the vectorised environments.
Multiple environments take what we did for single environments and vectorise them in an easy to handle manner. All other game augmentations and reshaping of the environments happen the same way as they did in single environments.
The Atari PPO agent is created from a base class which (Figure 2.4), with the hyperparameters, initialise two key sections:
The Model - This houses the CNN model, we’ve configured so that this is split into a “head” “actor”, and “critic. The head pre-processes incoming data, which the actor and critic then use to calculate their respective values.
The Memory - There is a deep dive on this topic later, but basically the memory stores state transitions collected during the acting phase. The initialisation determines the shape and size of the memory, and consequently, how long roll-outs will be.
With these two class instances the agent is able to perform 3 essential functions:
Act: This takes in the observations (of multiple environments), and provides the corresponding actions and log probabilities of those actions.
AddToMemory: Once an action has been taken, all the data for that step must be stored for later use. That is what the AddToMemory function does; it stores the information in PyTorch tensors.
Learn: This will be explored further in a later section — but essentially, this will take processed data from memory and update the model accordingly.
The Model
The model is a vanilla actor-critic network. A CNN head takes in the frames of the game as input, with the addition of all pixels being scaled down by 255, this is to ensure that the input range is between 0 and 1. The output from the CNN is then passed into two fully connected layers. The first layer outputs a single value acting as the critic, and then the second network outputs the number of actions available in the game environment acting as the actor.
The class is set up so that an act method (models.py:76) can be invoked to return the action, log probabilities and entropy produced by the actor, as well as the value predicted by the critic. To obtain the action, the logits from the actor are immediately sampled (models.py:79).
Initialising Network Weights
One last important factor is to ensure that that each layer has been initialised. A generic function called layer_init surrounds every layer in the network, this is to avoid vanishing or exploding gradients (great post by explaining how this works here).
The decision to make a memory class (rather than just storing the individual variables as PyTorch arrays) in part stemmed the Reinforcement Learning Discussion Discord (shout-out to the community who has been incredible in facilitating discussions), where we were discussing tests for core functionality. RL algorithms notoriously fail silently, meaning that they will either not throw an error but never learn, or even worse, learn but not to its maximum potential. This is a huge problem because it makes debugging incredibly hard and requires many, many experiments; and when using simple environments such as CartPole-v1 might never surface. This is why we decided to create a class that stored the saved trajectories, but whose responsibility it also is to calculate the advantages, discounted returns and sample the correct data from memory. This way, we could write tests for each individual function (the tests can be found in /tests/test_memory.py) to ensure that the output is exactly what was expected. It was also instrumental in validating our understanding of advantages and discounted returns.
The memory itself (as seen in Figure 2.6) has three key aspects:
Saved Variables — Trajectories are stored to memory here. The data is stored in predefined PyTorch tensors, as this avoids having to convert arrays to tensors etc.
Calculating Advantages — how this is done will be explained in the next section; the important part is that it calculates the advantages and discounted returns based on the trajectories sampled.
Sampling from Memory — Because all of our data is in order and still in the format of our n number of environments, this function essentially squashes all our trajectories into on long tensor and then samples random sections to train on (the implementation is explained in more detail in the Updating Model section).
The initialisation of the memory dictates it’s shape and size and requires various arguments. This includes the number of environments and the environment itself to get the observation and action space. To will determine how many “steps” are taken it also requires the size of the memory. Lastly, hyperparameters such as gamma are also needed.
Lastly, we initialise W&B — this has been a key tool for us to understand how our agents are performing. It has allowed Daniel and myself to analyse data and run experiments on our own machines and making the results available for both of us, all whilst working remotely.
Having some logging tool has been critical in understanding where mistakes were being made and what affect different hyperparameters have on the model. In a later section, we’ll dig into how various metrics helped us debug our work.
The first stage in training the PPO agent is to collect trajectories — this happens as seen in Figure 2.7 after the agent and environment have been initialised. The agent takes actions based on the observations from the environment, which is then used to take the next step. The sequence of states and actions are then finally saved to the memory for the training phase.
The action is taken by with act function on the agent [Models.py:109].
class ActorCritic(nn.Module):	# ...	def act(self, x, action=None):		values = self.critic(self.forward(x))		logits = self.actor(self.forward(x))		probabilities = Categorical(logits=logits)		if action is None:		    action = probabilities.sample()		return action, probabilities.log_prob(action), values, probabilities.entropy()
To act, the model receives observations from the gym environment (x). The observations then pass through the critic network which returns the value (what the value of that state is) and then the actor-network, which returns the logits (the output of the neural network).
The value is returned as-is. However, we want to convert our logits into a probability distribution on which we can then sample an action. This is done with the Categorical function of PyTorch — this function turns our logits into a distribution proportional to the logits, summing up to one. The benefit of the categorical class is that the class can also sample an action from the distribution, calculate the log probability of action and the entropy (entropy will be further explained in a later section).
Saving trajectories is fairly intuitive, but is useful to have a visual representation of how the rewards, observations and done’s come out of the gym environment. For each step in an environment, 3 variables of different dimensions are returned (observation, reward and done — shown in Figure 2.8).
Every interaction the agent has with the environment (acting on multiple environments) returns an observation, reward, and done. In the example of Figure 2.8, the first step shows no reward, and the episode has not terminated yet, this is assumed to be the case for every step after that unless indicated otherwise. This illustrates that the data collected is appended into the memory as entire steps across multiple environments, indicated by the blue box. This means that what is being saved into the memory is:
Observations (n, 80, 80, 4)
Rewards (n, 1)
dones (n, 1)
And will then be saved in the memory in a tensor of dimensions:
Observations (x, n, 80, 80, 4)
Rewards (x, n, 1)
dones (x, n, 1)
Where x is the number of steps, this implies that the “global” number of steps which the agent trains on will, in reality, be x times n, as we have x steps for n number of environments.
The second stage in our PPO is to learn from our collected trajectories — this is initiated by the learn function in Models.py and invoked by the training loop in util.py:L113. When this starts, the advantages are calculated first, this is important to do as all trajectories are required to remain in sequence to calculate the correct discounted return. Once this is finished, we can sample from memory at random, calculate our losses and update the model. Updating over the trajectories happens a set number of times, dictated by a hyper-parameter. Once completed, the old model is replaced with the updated one and the process start over again, but this time with a (hopefully) better performance. This process can be seen in Figure 2.9.
The Advantage is defined as the expected return minus the actual return — this indicates how much better or worse the model performed than anticipated. Calculating these advantages occurs in two steps — first, the discounted returns need to be computed. The expected values (calculated by the critic) are then subtracted from the discounted returns to obtain the advantage.
In case you’re not familiar with discounted returns, this is a great video, but essentially it takes the current reward + the value of the next state times a discount factor (gamma) as per Figure 2.10 (G is the discounted returns).
In order to calculate the discounted returns, we cannot start at the beginning of the trajectory but rather have to start at the end. This is because we need to iteratively calculate what the value of the next state is. This presents a problem at the very last state as there is no next value, which is why the function also takes last_value and last_done, which are used in the very first calculation of the discounted returns (in the case of Figure 2.11 it's -0.5 as the last value and the episode not being done). In many books, you might see this as n-step bootstrapping. We do not (necessarily) wait for the whole episode to finish but rather estimate after n-steps what the value of that state would be if the episode were to finish.
As an example of what might happen, in Figure 2.11 assumes the trajectories have already been collected (step 1) and then iteratively calculates the discounted return backwards.
This logic is implemented in our Memory.pyclass:
def calculate_discounted_returns(self, last_value, next_done):    with torch.no_grad():      # Create empty discounted returns array      self.discounted_returns = torch.zeros((self.size, self.num_envs)).to(self.device)      for t in reversed(range(self.size)):        # If first loop        if t == self.size - 1:          next_non_terminal = 1.0 - torch.FloatTensor(next_done).reshape(-1).to(self.device)          next_return = last_value.reshape(-1).to(self.device)        else:          next_non_terminal = 1.0 - self.dones[t+1]          next_return = self.discounted_returns[t+1]        self.discounted_returns[t] = self.rewards[t] + self.gamma * next_non_terminal * next_return
The reverse for loop iterates over the collected trajectories in reverse order to calculate the discounted returns. As in Figure 2.11, the first calculation uses the last predicted value, outputted by the critic, all other cases will have the previous discounted return available. The new discounted return is the reward + gamma * the previous discounted return. next_non_terminal is used to check if the discounting should be set to 0 because the episode ended. Finally to get the advantages the memory simply takes the discounted_returns and subtracts the collected values at the corresponding time steps.
Calculating the discounted returns can be a headache, and making sure that they are correctly calculated for the many different scenarios is even harder. It’s easy to make a mistake, in some cases meaning the algorithm will still train but doesn’t perform as well as it could. To really check all the edge cases it is suggested to write test cases and run the calculations against them. For our implementation, this can be found in /tests/test_memory.py.
Fetching the data
To update the model, the memory is sampled first — this is done by extracting mini-batches (random subsection) from the trajectory saved in memory. This is done with the get_mini_batch_idxs function. The function creates an index for every entry in the memory and randomly assigns them into the arrays for the mini-batch; for example, this could return
[ 	[1013,456, ... 147,328], 	..., 	[295, 1937, ..., 49,846]]
The first element in the mini-batch array([1013,456, ... 147,328]) is then passed into the sample function, which expects an array of indexes to pull out of the memory. For every index in that array the following values are returned:
prev_states
prev_actions
prev_log_probs
discounted_returns
advantage
prev_values
Policy Loss
The loss function consists of the policy and value loss. As a reminder, it might be worth to check out Figure 1.11 to remember the method of updating the PPO model.
First, the two surrogate functions (clipping of the ratios) need to be calculated. The surrogate function requires the probability ratios — this uses the updated model’s log probabilities and the current model’s log probabilities (used to collect the trajectory). This is the primary reason the action and observations were saved — the observations can now be passed into the updated model and calculate what new log probabilities are outputted given the action taken during the trajectory.
Once the ratios are calculated the surrogate function is relatively straight forward, we simply multiply through by the advantage (calculated in the previous section) and use the clamp function to limit this by 1 + ε and 1- ε. The loss is then the minimum of the two.
# Grab sample from memoryprev_states, prev_actions, prev_log_probs, discounted_returns, advantage, prev_values = self.mem.sample(mini_batch_idx)# find ratiosactions, log_probs, _, entropy = self.model.act(prev_states, prev_actions)ratio = torch.exp(log_probs - prev_log_probs.detach())# calculate surrogatessurrogate_1 = advantages * ratiosurrogate_2 = advantages * torch.clamp(ratio, 1-self.epsilon, 1+self.epsilon)# Policy Gradient Losspg_loss = -torch.min(surrogate_1, surrogate_2).mean()
Value Loss
To calculate the value loss we take the mean squared error of what we thought the value of that observation was and what the discounted returns actually were.
# Grab sample from memoryprev_states, prev_actions, prev_log_probs, discounted_returns, advantage, prev_values = self.mem.sample(mini_batch_idx)# Calculate values from new modelnew_values = self.model.get_values(prev_states).view(-1)# Generate lossvalue_loss = F.MSE(new_values, discounted_returns)
Updating
Updating is also very standard, we simply backpropagate all the loss added together and include the entropy.
# Lossloss = pg_loss + value_loss - self.entropy_beta*entropy_loss# calculate gradientself.optimiser.zero_grad()loss.backward()# Clip grad to avoid massive updatesnn.utils.clip_grad_norm_(self.model.parameters(), 0.5)self.optimiser.step()
This process gets repeated will the model has been updated n times for the whole memory set. Then we repeat the whole process over again and hope the model converges!
To run our the environment navigate to the PPO folder and run:
python main_atari.py
To get weights and biases set up you will need to do your own configuration and register the project if you do not have this, tensorboard will still record all of the experiments data.
When running the experiment the key values which are measured are:
Episode Reward
This is the most obvious metric to record, it tells us how well an agent is performing at a given step.
Learning Rate
The learning rate is logged in case this is changed throughout the life cycle of the experiment. For this experiment, it will remain constant, however, annealing is required to get state of the art results. Another blog post will cover why and how this works.
Policy Loss
The policy loss is generally agreed to be a difficult metric to draw conclusions from, depending on whether the advantages are normalised, or if learning rate annealing is used this will look different.
Value Loss
This is defined as the mean squared error between what the discounted return was and what the critic thought, it shows how well the critic is predicting the value of a state. Typically this sits quite low, however, when an agent performs much better or worse than expected the value loss will shoot up indicated that the critic is updating how it thinks about this state.
Entropy
Entropy in RL is the predictability of whether an action is taken (so how certain we are that in a given state we might go left); because the actions are sampled from a distribution, if one those outputs are very high that means it’s fairly certain what action will be taken. Conversely, if all the probabilities sit at the same level, it’s going to be hard to predict which action is taken next (great article that goes more in-depth).
From an insight perspective, what this means is that if the entropy is very high, the agent doesn’t really know what the right action is and when it’s lower is start be more certain about its actions. This is why, in general, the entropy will be going down.
KL Divergence
When calculating which action to take in with policy gradient methods, the actions are going to be represented as a set of probabilities — the log of these probabilities is then used to calculate the loss of our model.
When you update the model using the log probabilities in the “learning” phase, typically you would clone your model and update that new model. Once your model has gone through all the updates, it will differ slightly from the original model.
The KL divergence is a measure of how different the log probabilities from the original model are to the new one. The higher the KL divergence is the more the model has changed. As a rule of thumb, a KL divergence which sits around +/- 0.03 from 0 means that the model is updating nicely, whereas a high KL divergence can indicate that the model is taking too big of an update step, this can cause the model to not converge. This is what makes KL divergence such a useful metric to observe.
It’s hard to pin down what the exact problems in an experiment are when looking at metrics, especially because it’s hard to work out whether or not something just isn’t converging or if there is something wrong with the algorithm. One thing that has helped us tremendously is looking for how “on average” performance looks on different metrics. For example from Figure 3.1 it might look like the agent is learning, and then just cannot converge — but when also looking at the value loss (so the MSE between the discounted returns and what the critic thought), it’s clear that the critic was behaving very suspiciously, the value loss is all over the place and it’s frequently far too high. On closer inspection, we found that the calculation for value loss had a wrong variable name in it.
When we were building our agent — Costa’s CleanRL metrics were of great help to us.
Our Metrics
After much work, we managed to get a reward of around 50. Clearly not ideal, there is still a lot to work at, but this will require adding some additional tricks to the algorithm, it’s also very possible that after more training it will reach ~400.
When creating your own PPO, it might be helpful to have these metrics as reference.
Visual display of agents performing at various stages of training.
Implementing PPO from scratch has been an interesting challenge which has certainly stretched my understanding of RL and patients to its limits. What has really helped is having somebody to work with, this was both a motivation as well as helpful to bounce ideas off when either of us got stuck.
From our experiments, it’s clear that running just a vanilla PPO algorithm will not achieve top-level performance. As such, we have also implemented PPO as per the baselines specifications which reach a reward of 400. I aim to post a blog post on the modifications required to obtain the result soon!
Daniel and I would like to thank Costa Huang for all his help in explaining how his PPO implementation worked. Without his help, it would surely have been a much longer road.