tags:
- ML
- CSZero To Proximity-Policy Optimisers
The goal of this blog is to help someone who knows nothing about reinforcement learning build a proximity policy optimiser. As such, I've started with basic RL theory. If you are already familiar with theory and want to dive into the implementation, please skip to the implementation.
Reinforcement learning is very heavy in mathematics. Although I've tried to explain many of the more complex concepts in short form, it's much better if you learn them yourself.
Here are a few resources I'd suggest checking out for the math,
Basic stat. & prob for basic RL Khan Academy - Probability & Statistics
A bit more detailed probability course. Stanford CS299 - Probability For Machine Learning
Multivariable calculus for policy gradients & PPO Theory Khan Academy - Multivariable Calculus
To understand proximity policy optimisers, we first need to understand some basic Reinforcement Learning Theory. In reinforcement learning, we describe the entire state of the world, every minute detail, altogether at time/step
These action spaces can be both continuous, say in the case of an RL agent learning to play football, or discrete, in the case of, say, chess. This distinction is important as different kinds of algorithms may apply to different scenarios.
The way the agent takes up a certain action from these action spaces is through policies, policies may be defined as the brain of the agent itself. Mathematically, a policy is a function
Neural Networks
Neural networks are a type of machine learning model that consists of layers of interconnected nodes.
Where each node contains a specific weight, and each layer has its own bias. The weights are multiplied, while the biases are added.
You can think of each node as its own linear regression model, in the form
This isn't super hard, it's just layers and layers of multiplication and addition happening.
In python, this will look like.
def forward(inputs, weights, bias):
sum = 0
for i in range(len(inputs)):
sum += inputs[i] * weights[i]
sum += bias
return f(sum)
The NumPy library can help us do it in a cleaner way (+ It's faster!)
import numpy
def forward(inputs, weights, bias):
return f(np.dot(inputs, weights) + bias)
You might be wondering what the function
def ReLU(x):
return np.maximum(0, x)
Of course, we don't want our neural network to only be one layer, we will need a neural network with multiple layers, these type of neural networks are also known as multi-layer perceptrons. For this, we will define a Neural Network class.
def mlp(sizes, activation=nn.Tanh, output_activation=nn.Identity):
layers = []
for j in range(len(sizes)-1):
act = activation if j < len(sizes)-2 else output_activation
layers += [nn.Linear(sizes[j], sizes[j+1]), act()]
return nn.Sequential(*layers)
(Optional) Quick brush over on Back-Propagation
Now, although our neural network can do some operations, it still doesn't change those weights and biases. For that, we will need to use what is known as back-propagation.
If the forward pass was information flowing through the neural network, imagine back-propagation as the feedback to the prediction flowing back through the neural network, adjusting everything in its path. This adjustment is decided by the loss. There are various ways to calculate this loss, but for our case, the proximal policy optimiser will provide us with the policy loss, so let's not worry about that right now.
Mathematically, back-propagation relies heavily on the chain rule from calculus.
Say, our neural network predicts
So first, we'll define the derivatives of the activation function.
We know that each layer in the neural network computes,
and
Using the chain rule we want,
So whenever we update our
Continuously doing this, will lead to the weights and biases of the network getting optimised. This is known as Stochastic Gradient Descent.
For more, I'd suggest to read: This Post by IBM
Types of policies
Policies come in two different flavours. One is the deterministic policy which we just covered, the other is stochastic policies. Which are further divided into categorical Policies & diagonal Gaussian Policies.
categorical policies work for discrete action spaces.
For continuous action spaces, we must use what is known as the diagonal gaussian policy.
To understand a diagonal gaussian policy we must first understand what a multi-variate gaussian distribution is.
It's simple to understand, if a normal gaussian distribution describes how likely a thing is to happen.
A Multi-variate Gaussian Distribution defines how likely a combination of things are to happen. More precisely, it is defined by
- A mean vector
- A covariance matrix
A Co-Variance Matrix of this distribution just tells how these things vary and are connected. Mathematically it is,
(You might've ran across this thing in seaborn!)
A Diagonal Gaussian Distribution is just the diagonal of the matrix. Essentially,
So, we can represent it using a vector.
A Diagonal Gaussian Policy can be represented by a neural network
To represent our Diagonal Co-Variance matrix, we have two ways.
- Single Vector of Log of the Standard Deviations:
- Neural network that maps from states to log standard deviation.
In both cases, we use log of standard deviation as the while the standard deviation cannot be non-negative, the logarithm can, which makes things much easier for us.
Given the mean action
The log likelihood of a Diagonal Gaussian Policy is found to be,
Trajectories
We've learnt what actions & states are. A trajectory, also known as an episode or rollout, is very simple, it is just a sequence of states and actions in the world, mathematically, it is represented by
Rewards
When hearing about reinforcement learning, you must've heard that it's all about rewarding the agent for doing good things, and punishing the agent for doing bad things. Although that's only a subset of the type of reinforcement learning algorithms (as you may learn later, such as evolutionary algorithms.), rewards & returns are critical for Reinforcement Learning.
Mathematically, it is a function
The goal of our agent is to maximise a cumulative reward over trajectory
There are multiple kinds of returns, the first kind we'll look at, is the finite-horizon un-discounted return, which just adds up the reward over a finite time frame.
The Goal
The goal in Reinforcement Learning is to select a policy which maximises the expected cumulative return when the agent acts according to it.
To express the expected return, we must first understand the probability distribution over trajectories.
In a stochastic system, where the policy is
Value Functions
The value functions give a certain value to every state/state-action pair. Basically, say you start at state
The policy network is known as the critic network as it criticises the actor's actions.
The four main functions to note are,
- On-Policy Value Function. Gives expected return when you start at
- On-Policy Action-Value Function. Gives expected return when you start at
- The Optimal Value Function. Gives expected return if you start at
- The Optimal Action-Value Function. Action-Value for the optimal policy.also note a very common relation between Value & Action-Value functions
Policy Gradients
Reinforcement Learning can be divided into two major branches, model-free & model-based RL. A model is a function which predicts state transitions and rewards. Although this allows our agent to plan by thinking ahead this type of environment is not usually available. In these scenarios, the agent has to learn the model through trial and error, which is very computationally expensive.
The proximity-policy optimiser is a model-free Algorithm, so we'll be focussing on that.
Model-Free Algorithms have two main methods to training agents. One is Q-learning which tries to approximate the optimal action-value function. The other, which is the focus of this blog, is policy optimisation.
For the sake of simplicity, we'll be covering & implementing the most bare-bones equation for policy gradients in this blog.
A Bare-Bones Policy Gradient
Here, we'll be working with a stochastic policy
From Eq. 2 under The Goal
Now, using what is known as the log-derivative trick, which simply refers to the derivative of
We can now write our equation as,
Taking a look at the probability of trajectory - Eq. 1 in The Goal,
Alternate Policy Gradient Derivations (Bonus)
Infinite Horizon Discounted Setting
For an infinite horizon discounted setting, we have that
Finite Horizon Un discounted Setting, when
Similar to the original derivation we have,
The Expected Grad-Prob-Lemma
All probability distributions are normalised. (integration =
This has been termed as the expected grad-log-prob lemma.
Reward-To-Go Policy Gradient
In our policy gradient expression,
Well, we can change the expression of the policy gradient to,
Baselines
An immediate consequence of this Grad-Log-Prob Lemma is that for any function
Advantage Function
In RL, we don't always need to know absolutely how good an action is, but how good it is relative to other actions. For this, we use what is known as the advantage function,
Generalised Policy Gradient
We have seen so far that the policy gradient has a general form.
The Theory Behind PPO
The main question behind proximity policy optimisers is, "how can we take the biggest possible improvement step on a policy using the data we already have without stepping too far to decrease performance?"
The way a PPO does this is, by using a clipping function to prevent large updates when we're optimising our policy.
Mathematically we define is as,
Deriving PPO From Policy Gradients
From our Policy Gradient Theorem we have,
So, we get the objective to be,
What clip does is, if
So - after all that - we finally get,
This is known as PPO-CLIP.
Our main goal, is to find
PPO-Penalty
There is another version of proximal-policy optimisers known as PPO-Penalty, this is closer to a Trust-Region Policy Optimiser than PPO-Clip, but instead of using a clip to impose a hard constraint, it instead uses a penalty that scales with training instead. For the sake of this blog, we will be only covering PPO-CLIP.
Creating PPO-CLIP
To create the final python implementation of PPO-CLIP, we will be using the algorithm from Open AI's Spinning Up Docs
Before we start, let's import a few dependencies, we'll be working with PyTorch & Gymnasium which is a standard API that'll help us build the environment for the training.
import torch
import gymnasium as gym
We'll first create the __init__ function which initialises all our variables & hyper parameters.
class PPO:
def __init__(self, environment, network_class , **hyperparameters):
We have to check if the environment is compatible or not.
assert isinstance(env.observation_space, gym.spaces.Box), "Observation space must be of type Box"
assert isinstance(env.action_space, gym.spaces.Box), "Action space must be of type Box"
If our environment is not compatible, this will automatically throw an error. Now, we can extract the information from our environment.
self.env = env
self.observation_dim = env.observation_space.shape[0]
self.action_dim = env.action_space.shape[0]
After that - we'll initialise our actor & critic networks. This is the first step in the algorithm
self.actor = network_class(self.observation_dim, self.action_dim) # Step 1
self.critic = network_class(self.observation_dim, 1)
Now, we have to initialise our hyper parameters, we'll do that using a function hyperparameter_init() this is just to make our final code a bit cleaner.
def hyperparameter_init(self, hyperparameters):
self.TIMESTEPS_PER_BATCH = 4800
self.TIMESTEPS_PER_EPISODE = 1600
self.UPDATES_PER_ITERATION = 5
self.lr = 0.005
self.gamma = 0.95
self.clip = 0.2
self.render = True
self.render_every_i = 10
And, in our, __init__() function we add,
self.hyperparameter_init(hyperparameters)
We'll start creating our learning loop - which is step two. 
Within a new function learn we'll define the loop as,
def learn(self, total_timesteps)
timesteps_completed = 0
iter_count = 0 # number of iterations till now, we'll use this to save after *n* iterations
while timesteps_completed < total_timesteps: # Step 2
pass
iter_count +=1
The rollout() function 'rolls out' the entire batch data till now in front of us.
def learn(self, total_timesteps)
timesteps_completed = 0
iter_count = 0 # number of iterations till now, we'll use this to save after *n* iterations
while timesteps_completed < total_timesteps: # Step 2
batch_observations, batch_actions, batch_log_probs, batch_reward_to_gos, batch_episode_lengths = self.rollout() # Step 3
t_so_far += np.sum(batch_lens)
iter_count +=1
Defining rollout(). Here, we'll be calculating the reward to go using reward_to_gos() function, this is step 4.
def rollout(self):
# Batch data that gets returned
batch_observations = []
batch_actions = []
batch_log_probs = []
batch_rewards = []
batch_reward_to_gos = []
batch_episode_lengths = []
episode_rewards=[] # Keeps track of rewards per episode, get's cleared after every episode.
timesteps = 0
while t < self.TIMESTEPS_PER_BATCH: #
episode_rewards = []
# Resetting the environment after each batch.
observation, _ = self.env.reset()
done = False
for episode_timestep in range(self.TIMESTEPS_PER_EPISODE):
# Rendering the environment
if self.render and (self.logger['i_so_far'] % self.render_every_i == 0) and len(batch_lens) == 0:
self.env.render()
# Incrementing The Timestep
timesteps += 1
# Now - we'll track the observations in this batch,
batch_observations.append(observation)
# We'll be using our get_action function to calculate the action & log_prob
action, log_prob = self.get_action(observation)
observation, reward, terminated, truncated, _ = self.env.step(action)
done = truncated | terminated
# Tracking the recent rewards, & action, probs
episode_rewards.append(reward)
batch_actions.append(action)
batch_log_probs.append(log_prob)
if done:
break
# Tracking lengths
batch_episode_lengths.append(ep_t + 1)
batch_rewards.append(episode_rewards)
# Reshaping Data Into Tensors
batch_observations = torch.tensor(batch_observations, dtype=torch.float)
batch_actions = torch.tensor(batch_actions, dtype=torch.float)
batch_log_probs = torch.tensor(batch_log_probs, dtype=torch.float)
batch_reward_to_gos = self.compute_reward_to_gos(batch_rewards) # STEP 4
return batch_obs, batch_acts, batch_log_probs, batch_reward_to_gos, batch_lens
Looking back at the theory from section Reward-To-Go Policy Gradient, we see that, reward to go is defined as,
def compute_reward_to_gos(self, batch_rewards):
batch_reward_to_gos = []
for episode_rewards in reversed(batch_rewards):
discounted_reward = 0
for reward in reversed(episode_rewards):
discounted_reward = reward + discounted_reward * self.gamma
batch_reward_to_gos.insert(0, discounted_reward)
batch_reward_to_gos = torch.tensor(batch_reward_to_gos, dtype=torch.float)
return batch_reward_to_gos
We also need to develop our get_action() function. According to,
def get_action(self, observation)
mean = self.actor(observation)
dist = torch.distributions.MultivariateNormal(mean, self.cov_mat) # Defines a distribution
action = dist.sample() # Samples action
log_prob = dist.log_prob(action) # Calculates log probability for action
return action.detach().numpy(), log_prob.detach()
Cool, now going back to our learn() function, according to step 5, we'll build out the Advantage Function,
V, _ = self.evaluate(batch_obs, batch_acts) # Step 5
A_k = batch_reward_to_gos - V.detach()
Now, we have to update the PPO-Clip Objective.
For this, first, we'll go back to the __init__() function, and initialise the Adam optimiser.
First, import the Adam optimiser,
from torch.optim import Adam
self.actor_optimiser = Adam(self.actor.parameters(), lr=self.lr)
self.critic_optimiser = Adam(self.critic.parameters(), lr=self.lr)
Back in the learn() function, we start out summation loop within the main while loop, we first get
for _ in range(self.UPDATES_PER_ITERATION):
V, curr_log_probs = self.evaluate(batch_observations, batch_actions)
ratios = torch.exp(curr_log_probs - batch_log_probs)
#Calculating surrogate loss r_\theta
surr1 = ratios * A_k
surr2 = torch.clamp(ratios, 1 - self.clip, 1 + self.clip) * A_k
actor_loss = (-torch.min(surr1, surr2)).mean() # Negative because adam minimizes, and we want to maximize.
critic_loss = torch.nn.MSELoss()(V, batch_reward_to_gos)
# Next - we calculate the gradients for both the networks
self.actor_optimiser.zero_grad()
actor_loss.backward(retain_graph=True)
self.actor_optimiser.step()
self.critic_optimiser.zero_grad()
critic_loss.backward()
self.critic_optimiser.step()
All evaluate() does here is determine the values of each observation along with the log probabilities of each action in the most recent batch.
def evaluate(self, batch_observations, batch_actions)
V = self.critic(batch_observations).squeeze()
mean = self.actor(batch_observations)
dist = MultivariateNormal(mean, self.cov_mat)
log_probs = dist.log_prob(batch_actions)
return V, log_probs
That's it! You can check out the final code on my github
Further Reading & Credits
Congratulations on getting this far! You can now create the basis for a proximal-policy optimiser from scratch in python. :) If you'd like to read more about policy optimisers & deep RL in general, here are a few links you should check out.