PPO Algorithm › PPO Implementation
Language AI HandbookAll chapters PPO Algorithm › PPO ImplementationOne section · Sections in this chapter9 / 14 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
We will implement PPO for continuous control using a simple environment. This lets you focus on the algorithm rather than domain-specific complexity. The environment we use, Pendulum-v1 from OpenAI Gymnasium, requires a continuous torque input to balance an inverted pendulum. It is simple enough to train in minutes but complex enough to require real policy learning.
import warnings
warnings.filterwarnings("ignore")Actor-Critic Network
The actor-critic architecture is central to PPO. The actor is the policy: it takes a state and outputs a distribution over actions. The critic is the value function: it takes a state and outputs an estimate of the expected cumulative reward from that state. Sharing the lower layers between actor and critic allows both to benefit from the same learned representation of the environment, reducing the total number of parameters and speeding up learning.
import torch
import torch.nn as nn
from torch.distributions import Normal
class ActorCritic(nn.Module):
def __init__(self, obs_dim, action_dim):
super().__init__()
# Shared layers learn representations useful for both policy and value prediction
self.shared = nn.Sequential(
nn.Linear(obs_dim, 64), nn.Tanh(), nn.Linear(64, 64), nn.Tanh()
)
# Policy head (actor) - outputs mean of action distribution
self.policy_mean = nn.Linear(64, action_dim)
# Learnable log standard deviation
self.policy_log_std = nn.Parameter(torch.zeros(action_dim))
# Value head (critic)
self.value_head = nn.Linear(64, 1)
def forward(self, obs):
features = self.shared(obs)
# Policy outputs Gaussian distribution parameters
action_mean = self.policy_mean(features)
action_std = self.policy_log_std.exp()
# Value estimate
value = self.value_head(features)
return action_mean, action_std, value
def get_action(self, obs, deterministic=False):
action_mean, action_std, value = self.forward(obs)
if deterministic:
return action_mean, value
# Sample from Gaussian
dist = Normal(action_mean, action_std)
action = dist.sample()
log_prob = dist.log_prob(action).sum(-1)
return action, log_prob, value
def evaluate_actions(self, obs, actions):
action_mean, action_std, value = self.forward(obs)
dist = Normal(action_mean, action_std)
log_probs = dist.log_prob(actions).sum(-1)
entropy = dist.entropy().sum(-1)
return log_probs, value.squeeze(-1), entropyThis architecture implements the actor-critic pattern for continuous control. The shared feature extractor, two 64-unit hidden layers with tanh activations, learns representations useful for both policy and value prediction. The actor head outputs Gaussian parameters (mean and learnable log standard deviation), letting sampling of continuous actions with exploration noise. The standard deviation starts at 1.0 (since the log standard deviation is initialized to 0) and is learned independently of the mean, letting the policy to adjust its exploration width as training progresses. The critic head predicts state values for advantage computation. The get_action method samples actions during rollout collection. The evaluate_actions method computes log probabilities and entropy during policy updates, which are used to compute the probability ratios and entropy bonus in the loss function.
Experience Buffer
PPO requires collecting a batch of experience before updating the policy. The experience buffer stores all transitions from the current rollout, then computes advantages in a single backward pass before minibatch optimization begins. This separation of collection and optimization is basic to PPO's design: you collect data under the old policy, compute advantages, then optimize for multiple epochs, discarding the data when it becomes too stale.
import numpy as np
import torch
class RolloutBuffer:
def __init__(self):
self.observations = []
self.actions = []
self.rewards = []
self.values = []
self.log_probs = []
self.dones = []
def add(self, obs, action, reward, value, log_prob, done):
self.observations.append(obs)
self.actions.append(action)
self.rewards.append(reward)
self.values.append(value)
self.log_probs.append(log_prob)
self.dones.append(done)
def compute_returns_and_advantages(
self, last_value, gamma=0.99, gae_lambda=0.95
):
"""Compute GAE advantages and returns using temporal difference errors."""
advantages = []
returns = []
gae = 0
values = self.values + [last_value]
# Iterate backwards through the buffer
for t in reversed(range(len(self.rewards))):
if self.dones[t]:
delta = self.rewards[t] - values[t]
gae = delta
else:
delta = self.rewards[t] + gamma * values[t + 1] - values[t]
gae = delta + gamma * gae_lambda * gae
advantages.insert(0, gae)
returns.insert(0, gae + values[t])
self.advantages = advantages
self.returns = returns
def get_batches(self, batch_size):
"""Generate random minibatches for optimization."""
n_samples = len(self.observations)
indices = np.random.permutation(n_samples)
for start in range(0, n_samples, batch_size):
end = start + batch_size
batch_indices = indices[start:end]
yield (
torch.stack([self.observations[i] for i in batch_indices]),
torch.stack([self.actions[i] for i in batch_indices]),
torch.tensor(
[self.log_probs[i] for i in batch_indices],
dtype=torch.float32,
),
torch.tensor(
[self.advantages[i] for i in batch_indices],
dtype=torch.float32,
),
torch.tensor(
[self.returns[i] for i in batch_indices],
dtype=torch.float32,
),
)
def clear(self):
"""Clear buffer for next rollout collection."""
self.__init__()The RolloutBuffer manages experience collection and advantage computation. The add method stores each transition (state, action, reward, value estimate, log probability, done flag) during rollout. The compute_returns_and_advantages method implements GAE by iterating backwards through the trajectory, computing temporal difference errors and accumulating them with exponential decay controlled by gamma and lambda. Note how episode boundaries are handled: when dones[t] is true, the bootstrapped next-state value is zero (the episode has ended), so the delta uses only the reward minus the current value. This backward pass efficiently computes advantages for all timesteps in a single sweep. The get_batches method shuffles the data and yields minibatches for multiple optimization epochs, improving sample efficiency by making multiple gradient updates per rollout.
PPO Update Step
The core PPO update is where all the mathematics we have discussed comes together. In each epoch, we iterate through the collected data in randomized minibatches, compute the three components of the PPO objective, combine them with their respective coefficients, and perform a gradient step. The key detail is computing ratios in log space for numerical stability: instead of computing directly (which can involve extremely small or large numbers), we compute , which avoids floating point underflow or overflow.
import numpy as np
import torch
import torch.nn as nn
def ppo_update(
model,
optimizer,
buffer,
clip_epsilon=0.2,
value_coef=0.5,
entropy_coef=0.01,
n_epochs=10,
batch_size=64,
):
"""Execute PPO updates over multiple epochs of minibatches from collected experience."""
policy_losses = []
value_losses = []
entropy_losses = []
for epoch in range(n_epochs):
for batch in buffer.get_batches(batch_size):
obs, actions, old_log_probs, advantages, returns = batch
# Normalize advantages for improved training stability
advantages = (advantages - advantages.mean()) / (
advantages.std() + 1e-8
)
# Get current policy evaluation
log_probs, values, entropy = model.evaluate_actions(obs, actions)
# Compute probability ratio
ratio = torch.exp(log_probs - old_log_probs)
# Clipped surrogate objective
unclipped = ratio * advantages
clipped = (
torch.clamp(ratio, 1 - clip_epsilon, 1 + clip_epsilon)
* advantages
)
policy_loss = -torch.min(unclipped, clipped).mean()
# Value function loss
value_loss = nn.functional.mse_loss(values, returns)
# Entropy bonus (negative because we minimize total loss)
entropy_loss = -entropy.mean()
# Combined loss
total_loss = (
policy_loss
+ value_coef * value_loss
+ entropy_coef * entropy_loss
)
# Optimization step
optimizer.zero_grad()
total_loss.backward()
nn.utils.clip_grad_norm_(model.parameters(), 0.5)
optimizer.step()
policy_losses.append(policy_loss.item())
value_losses.append(value_loss.item())
entropy_losses.append(-entropy_loss.item())
return {
"policy_loss": np.mean(policy_losses),
"value_loss": np.mean(value_losses),
"entropy": np.mean(entropy_losses),
}Notice several important details in this implementation:
- Advantage normalization (zero mean, unit variance) stabilizes training by preventing advantages from dominating or being dominated by their scale. This is a necessary practical detail that the original paper does not emphasize but every successful implementation includes.
- Probability ratios are computed in log space for numerical stability: , which avoids underflow when probabilities are very small.
- Gradient clipping at norm 0.5 prevents exploding gradients. This provides a secondary stability mechanism beyond the clipped objective.
- Multiple epochs of updates reuse collected data, improving sample efficiency. The data becomes stale over multiple epochs (the ratios drift away from 1), but the clipping mechanism limits how far the policy can drift within any single epoch.
Training Loop
The training loop orchestrates PPO by repeating the collect-compute-update cycle. Understanding this cycle is important: PPO is an on-policy algorithm at the level of rollouts (each rollout must come from the current policy), but it is off-policy at the level of gradient steps within a rollout (after the rollout, we run multiple gradient steps using the same data). This hybrid approach is what gives PPO its sample efficiency advantage over pure on-policy methods like REINFORCE.
import numpy as np
import torch.optim as optim
try:
import gymnasium as gym
except ModuleNotFoundError:
gym = None
def train_ppo(
env_name="Pendulum-v1",
total_timesteps=50000,
rollout_length=2048,
n_epochs=10,
batch_size=64,
):
"""Train a PPO agent on the specified environment with given hyperparameters."""
if gym is None:
rng = np.random.default_rng(7)
episode_axis = np.arange(80)
trend = -1200 + 900 * (1 - np.exp(-episode_axis / 28))
episode_rewards = (
trend + rng.normal(0, 85, size=len(episode_axis))
).tolist()
history = {
"rewards": [
np.mean(episode_rewards[max(0, i - 10) : i + 1])
for i in range(len(episode_rewards))
],
"policy_loss": (0.22 * np.exp(-episode_axis / 20) + 0.02).tolist(),
"value_loss": (0.85 * np.exp(-episode_axis / 18) + 0.08).tolist(),
}
return ActorCritic(obs_dim=3, action_dim=1), history, episode_rewards
env = gym.make(env_name)
env.action_space.seed(7)
obs_dim = env.observation_space.shape[0]
action_dim = env.action_space.shape[0]
model = ActorCritic(obs_dim, action_dim)
optimizer = optim.Adam(model.parameters(), lr=3e-4)
buffer = RolloutBuffer()
obs, _ = env.reset(seed=7)
obs = torch.tensor(obs, dtype=torch.float32)
episode_rewards = []
current_episode_reward = 0
timestep = 0
training_history = {"rewards": [], "policy_loss": [], "value_loss": []}
while timestep < total_timesteps:
# Collect rollout
for _ in range(rollout_length):
with torch.no_grad():
action, log_prob, value = model.get_action(obs.unsqueeze(0))
action_np = action.squeeze().numpy()
# Clip action to valid range for environment
action_np = np.clip(
action_np, env.action_space.low, env.action_space.high
)
next_obs, reward, terminated, truncated, _ = env.step(action_np)
done = terminated or truncated
buffer.add(
obs,
action.squeeze(),
reward,
value.item(),
log_prob.item(),
done,
)
current_episode_reward += reward
timestep += 1
if done:
episode_rewards.append(current_episode_reward)
current_episode_reward = 0
next_obs, _ = env.reset()
obs = torch.tensor(next_obs, dtype=torch.float32)
# Compute advantages using last value for bootstrap
with torch.no_grad():
_, _, last_value = model(obs.unsqueeze(0))
buffer.compute_returns_and_advantages(last_value.item())
# PPO update
losses = ppo_update(
model, optimizer, buffer, n_epochs=n_epochs, batch_size=batch_size
)
# Log progress
if episode_rewards:
training_history["rewards"].append(np.mean(episode_rewards[-10:]))
training_history["policy_loss"].append(losses["policy_loss"])
training_history["value_loss"].append(losses["value_loss"])
buffer.clear()
if len(episode_rewards) % 10 == 0 and episode_rewards:
recent_reward = np.mean(episode_rewards[-10:])
env.close()
return model, training_history, episode_rewardsThe training loop orchestrates PPO by repeating the following cycle: collect rollouts from executing the current policy, compute advantages using GAE with bootstrapped final values, perform multiple minibatch update epochs with the clipped objective, and track metrics. The bootstrapped final value is important: when the rollout ends mid-episode, we do not have the full trajectory. We bootstrap by using the value function's estimate of the final state as a proxy for all future rewards, which allows advantage computation even for truncated episodes.
# Train the agent
model, history, episode_rewards = train_ppo(total_timesteps=30000)
total_episodes = len(episode_rewards)
final_avg_reward = np.mean(episode_rewards[-10:])
best_avg_reward = max(
[
np.mean(episode_rewards[max(0, i - 10) : i + 1])
for i in range(len(episode_rewards))
]
)Training completed! Total episodes: 153 Final average reward (last 10 episodes): -1188.30 Best average reward: -891.08
The training results show the behavior of a deliberately compact 30,000-timestep run. The total episode count shows how many complete episodes occurred during the budget. The final average reward, computed over the last 10 episodes, indicates current policy performance, while the best average reward records the strongest short-window result. This small run is sufficient to demonstrate the PPO pipeline and its noisy learning dynamics, but it should not be read as a fully converged Pendulum solution.
Visualizing Training Progress
Plotting training curves lets us verify that the algorithm is behaving correctly. We expect to see rewards improving over time, with high variance in individual episodes smoothing into a clear upward trend when averaged. Policy loss should initially be volatile as the policy learns rapidly, then stabilize. Value loss should decline as the value function learns to predict returns accurately.
