Learning to Play CartPole and LunarLander with Proximal Policy Optimization
By: Youness Mansar
Tags: Pytorch, Hands On Tutorials, Openai Gym, Reinforcement Learning, Policy Gradient

Learning to Play CartPole and LunarLander with Proximal Policy Optimization

In this post, we will train an RL agent to play two control based games:

Our agent will be trained using an algorithm called Proximal Policy Optimization. We will implement this approach from scratch using PyTorch and OpenAi gym.

This post is based on the following paper:

Gym:

import gym env_name = "LunarLander-v2" env = gym.make(env_name)

env.reset() for _ in range(1000):

new_observation, reward, done, info = env.step(env.action_space.sample()) arr = env.render(mode="rgb_array") if done:

break

The basic idea behind OpenAI Gym is that we define an environment env by calling:

env = gym.make(env_name)

Then at each time step t, we pick an action a and we get a new state_(t+1) and a reward reward_t. The objective is to train an Agent that learns a policy PI that can predict for each state the best action that will maximize the sum of the future rewards. For example, in the environment LunarLander, we get the maximum reward if we land the rocket smoothly on top of the landing area. In the environment CartPole, the objective is to keep the pole vertical for as long as possible.

PPO:

Our final objective is to learn a policy network that will take the state as input and then output a probability distribution over the actions that will maximize the expected reward.

The policy network is a simple FeedForward neural network that can be defined as:

class PolicyNetwork(torch.nn.Module):

def __init__(self, n=4, in_dim=128):

super(PolicyNetwork, self).__init__() self.fc1 = torch.nn.Linear(in_dim, 128)

self.fc2 = torch.nn.Linear(128, 128)

self.fc3 = torch.nn.Linear(128, 128) self.fc4 = torch.nn.Linear(128, n) self.l_relu = torch.nn.LeakyReLU(0.1) def forward(self, x): x = self.l_relu(self.fc1(x))

x = self.l_relu(self.fc2(x))

x = self.l_relu(self.fc3(x)) y = self.fc4(x) y = F.softmax(y, dim=-1) return y def sample_action(self, state): if not state is torch.Tensor:

state = torch.from_numpy(state).float().to(device) if len(state.size()) == 1:

state = state.unsqueeze(0) y = self(state) dist = Categorical(y) action = dist.sample() log_probability = dist.log_prob(action) return action.item(), log_probability.item() def best_action(self, state): if not state is torch.Tensor:

state = torch.from_numpy(state).float().to(device) if len(state.size()) == 1:

state = state.unsqueeze(0) y = self(state).squeeze() action = torch.argmax(y) return action.item() def evaluate_actions(self, states, actions):

y = self(states) dist = Categorical(y) entropy = dist.entropy() log_probabilities = dist.log_prob(actions) return log_probabilities, entropy

Implementing PPO goes as follows:

First, we start with a policy PI_old

We sample some trajectories from P_old

For each action a in each trajectory we compute the Advantage, a kind of measure of how much better the action a is compared to other possible actions at state_t.

in each trajectory we compute the Advantage, a kind of measure of how much better the action is compared to other possible actions at state_t. For a few epochs, we maximize the following objective with gradient ascent:

From Proximal Policy Optimization Algorithms

What this loss does is that it increases the probability if action a_t at state s_t if it has a positive advantage and decreases the probability in the case of a negative advantage. However, in practice this ratio of probabilities tends to diverge to infinity, making the training unstable. The authors propose a clipped version of the loss to solve this issue:

We also add two additional terms to the objective, (-) mean squared error over the Value function and the Entropy to encourage exploration during the sampling of the trajectories.

We can plot the sum of the rewards during the progression of the training for each episode: