Proximal Policy Optimization(PPO)

[Key concepts of Proximal Policy Optimization (PPO)]

• Policy gradient algorithms
• Similar to A2C, but policy updates have constraints for stable training
  • Basic A2C and A3C do not explicitly impose a KL trust-region constraint as TRPO does.
  • Excessive updates can cause performance degradation.
• An improvement on TRPO (Trust Region Policy Optimization)
  • TRPO: Constrains KL divergence (Kullback–Leibler divergence) during policy updates to prevent the new policy from deviating too far from the old one.
  • PPO: Introduces clipping to achieve an effect similar to TRPO with a simpler implementation.
• Key idea
  • Use a surrogate objective that clips the policy probability ratio to reduce the incentive for excessive updates.
  • Clipping differs from a constraint that directly guarantees an upper bound on KL divergence.
  • Simple implementation

[Clipped surrogate objective]

• Define the action probability ratio of the current policy to the old policy:
  • $r_t(\theta) = \frac{\pi_\theta(a_t|s_t)}{\pi_{\theta_{\text{old}}}(a_t|s_t)}$
    • $\pi_\theta(a_t|s_t)$: current policy
    • $\pi_{\theta_{\text{old}}}(a_t|s_t)$: old policy
• Ideal policy gradient update: multiply ratio $r_t$ by advantage $A_t$ to maximize the expectation
• Clipping: Clip $r_t$ to a fixed range to prevent $r_t$ from deviating too far from 1 (the old policy).
  • For example, if $\epsilon = 0.2$, restrict $r_t$ to $[1-\epsilon, 1+\epsilon]$ by clipping.
  • Outside the clipping range, prevent the gain from increasing or decreasing further.
  • Ensure training stability by preventing excessive policy changes.

[PPO limitations and directions for improvement]

• On-policy algorithm
  • May be inefficient in robot learning and other settings where generating data is very costly.
• Alternative: Off-policy algorithm >> Soft Actor-Critic SAC (Week 10)

[Activity: CartPole with PPO]

1. Import libraries and set hyperparameters

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.distributions import Categorical

import gymnasium as gym
import numpy as np
import matplotlib.pyplot as plt
import imageio

learning_rate = 0.0003
gamma = 0.99
lmbda = 0.95  # GAE lambda
eps_clip = 0.1  # Clipping range
K_epoch = 4  # Number of PPO update iterations
T_horizon = 128 # Maximum number of steps in 1 episode
num_episodes = 2000 # Number of training episodes
early_stopping_threshold = 450.0,  # Early-stopping threshold
early_stopping_patience = 10  # Number of consecutive episodes satisfying the early-stopping condition
gif_save_path = 'cartpole_simulation_ppo.gif' # GIF output path
progress_interval = 50  # Progress reporting interval

2. Define the actor–critic network

• Define an Actor-Critic network inheriting from nn.Module
• The actor network takes a state and outputs a probability distribution over actions
• The critic network takes a state and outputs its state value
• Initialize neural network weights with init_weights (Xavier initialization)

class ActorCritic(nn.Module):
    def __init__(self, state_dim, action_dim, hidden_dim):
        super(ActorCritic, self).__init__()
        self.actor = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, action_dim),
            nn.Softmax(dim=-1)
        )
        self.critic = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 1)
        )
        self.init_weights()

    # Xavier initialization
    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)

    def forward(self, x):
        return self.actor(x), self.critic(x)

3. Define the PPO algorithm class

• The PPO class includes the actor–critic network, GAE computation, and PPO update logic.
• get_action: Sample an action from the actor network for the given state and return its log probability.
• compute_gae: Compute Generalized Advantage Estimation (GAE)
  • GAE can reduce variance relative to Monte Carlo estimation, but bootstrapping from an inaccurate value function can introduce bias.
  • Estimate advantages using rewards and the value function.
  • Control GAE's bias–variance trade-off using the lambda parameter.
• update: Perform PPO's core update logic
  • Update the networks repeatedly using minibatches.
  • Compute PPO's clipped surrogate objective and update the actor network to maximize it.
  • Update the critic network to minimize mean squared error (MSE) loss for better value function prediction.

class PPO:
    def __init__(self, state_dim, action_dim, hidden_dim, lr, gamma, lmbda, eps_clip, K_epoch):
        self.model = ActorCritic(state_dim, action_dim, hidden_dim)
        self.optimizer = optim.AdamW(self.model.parameters(), lr=lr) # AdamW optimizer
        self.gamma = gamma
        self.lmbda = lmbda
        self.eps_clip = eps_clip
        self.K_epoch = K_epoch
        self.data = []

    def get_action(self, state):
        state = torch.FloatTensor(state) # Convert to Tensor + device
        probs, _ = self.model(state)  # Call Actor, Critic separately
        m = Categorical(probs)
        action = m.sample()
        return action.item(), m.log_prob(action)

    def compute_gae(self, rewards, values, next_value, dones):
      advantages = []
      gae = 0
      values = values + [next_value]

      for i in reversed(range(len(rewards))):
          delta = rewards[i] + self.gamma * values[i + 1] * (1 - dones[i]) - values[i]
          gae = delta + self.gamma * self.lmbda * (1- dones[i]) * gae
          advantages.insert(0, gae)

      advantages = torch.tensor(advantages, dtype=torch.float) # Convert to Tensor and assign to device
      target_values = advantages + torch.tensor(values[:-1],dtype=torch.float) # Convert to Tensor and assign to device

      return advantages, target_values


    def update(self, states, actions, log_probs, advantages, target_values):
        states = torch.stack(states)
        actions = torch.tensor(actions) # Convert to Tensor, device
        log_probs = torch.stack(log_probs)

        # PPO update (repeat K_epoch times)
        for _ in range(self.K_epoch):
            # Generate minibatch indices
            indices = torch.randperm(states.size(0))

            for start in range(0, states.size(0), T_horizon): # Use a mini-batch of T_horizon steps
                end = start + T_horizon
                idx = indices[start:end]

                # Minibatch data
                mb_states = states[idx].detach()  # Take only the original values used in loss so they do not participate in loss backpropagation
                mb_actions = actions[idx].detach()
                mb_log_probs = log_probs[idx].detach()
                mb_advantages = advantages[idx].detach()
                mb_target_values = target_values[idx].detach()

                # New policy probabilities and values
                new_probs, values = self.model(mb_states) # forward
                new_dist = Categorical(new_probs)
                new_log_probs = new_dist.log_prob(mb_actions)

                # Compute the PPO ratio
                ratio = torch.exp(new_log_probs - mb_log_probs)

                # Clipped surrogate objective
                surr1 = ratio * mb_advantages
                surr2 = torch.clamp(ratio, 1 - self.eps_clip, 1 + self.eps_clip) * mb_advantages
                actor_loss = -torch.min(surr1, surr2).mean() # actor loss

                # Critic loss (MSE)
                critic_loss = F.mse_loss(values, mb_target_values.unsqueeze(1))

                # Total loss
                total_loss = actor_loss + critic_loss

                # Backpropagate and optimize
                self.optimizer.zero_grad()
                total_loss.backward()
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0) # gradient clipping
                self.optimizer.step()

4. Define the training loop function

• The train function contains the main logic for training a PPO agent.
• Collect data by interacting with the environment in each episode (s, a, r, s_prime, prob_a, done)
• Call compute_gae to calculate GAE
• Call update to update the network with PPO
• Check the early stopping condition and stop training if it is met
• Print training progress periodically.

def train(env, ppo_agent):
    scores = [] # Record score, loss

    for episode in range(num_episodes):
        state, _ = env.reset()
        done = False
        score = 0
        episode_states, episode_actions, episode_log_probs, episode_rewards, episode_dones = [], [], [], [], []

        # Collect data by episode
        while not done:
            action, log_prob = ppo_agent.get_action(state) # action, log prob
            next_state, reward, terminated, truncated, _ = env.step(action)
            done = terminated or truncated

            # Store data (s, a, r, s', prob_a, done)
            state_tensor = torch.from_numpy(state).float()
            episode_states.append(state_tensor)
            episode_actions.append(action)
            episode_log_probs.append(log_prob)
            episode_rewards.append(reward)
            episode_dones.append(done)

            state = next_state
            score += reward

        # Compute GAE
        with torch.no_grad():
          values = ppo_agent.model(torch.stack(episode_states))[1].squeeze().tolist() # [batch, 1] -> [batch]
          _, next_value = ppo_agent.model(torch.FloatTensor(next_state))
          next_value = next_value.item()

        advantages, target_values = ppo_agent.compute_gae(
            rewards=episode_rewards,
            values=values,
            next_value=next_value,
            dones=episode_dones
        )

        # PPO update
        ppo_agent.update(episode_states, episode_actions, episode_log_probs, advantages, target_values)

        # Record score
        scores.append(score)

        # Check early stopping at each 'progress_interval' specified in config.
        if (episode + 1) % progress_interval == 0:
            avg_score = np.mean(scores[-progress_interval:])
            print(f"Episode {episode + 1}: Avg. Score = {avg_score:.2f}")

        # Check the early-stopping condition.
        if np.min(scores[-early_stopping_patience:]) >= early_stopping_threshold:
            print(f"Early stopping at episode {episode + 1} (score: {score:.2f})")
            break

    return scores # Return the score list

5. Run training and visualize results

• Create a PPO agent and call train to start training
• After training, plot the mean scores recorded during training

# Create a PPO agent
env = gym.make('CartPole-v1', render_mode = 'rgb_array')
state_dim = env.observation_space.shape[0]
action_dim = env.action_space.n
hidden_dim = 128
agent = PPO(state_dim, action_dim, hidden_dim, learning_rate, gamma, lmbda, eps_clip, K_epoch)

# Run training
scores = train(env, agent)

# Visualize results (learning curve)
plt.plot(scores)
plt.xlabel("Episode")
plt.ylabel("Average Score")
plt.title("PPO Training Progress")
plt.show()

Episode 50: Avg. Score = 17.92
Episode 100: Avg. Score = 20.22
Episode 150: Avg. Score = 128.42
Episode 200: Avg. Score = 241.90
Episode 250: Avg. Score = 291.58
Episode 300: Avg. Score = 429.42
Early stopping at episode 345 (score: 500.00)

<Figure size 640x480 with 1 Axes>

6. Validate the trained agent and save a GIF

• Run an episode with the trained agent and save the result as a GIF file
• Print the total reward over the episode.

# Evaluation (save a GIF)
print("Generating GIF...")
frames = []
state, _ = env.reset()
done = False
truncated = False
total_reward = 0

with imageio.get_writer(gif_save_path, duration=33) as writer:  # imageio writer
    while not done:
        action, _ = agent.get_action(state)
        next_state, reward, terminated, truncated, _ = env.step(action)
        total_reward += reward
        state = next_state
        done = terminated or truncated

        frames.append(env.render())

    # Save frames as a GIF using imageio
    for frame in frames:
        writer.append_data(frame)

print(f"Total reward for the test episode: {total_reward}")
env.close() # Close the environment

Generating GIF...
Total reward for the test episode: 500.0