PPO for Language Models › A Complete Training Loop
Language AI HandbookAll chapters PPO for Language Models › A Complete Training LoopOne section · Sections in this chapter10 / 15 ↓
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
The following example demonstrates how PPO training proceeds at a high level. We use mock components to illustrate the data flow without requiring actual large models. This demonstration shows the needed rhythm of PPO training: generate responses, score them, compute advantages, and update the policy.
The loop has a distinct structure that differs from supervised fine-tuning. We first generate a batch of responses from the current policy (the rollout phase). We then score those responses with the reward model and compute all the ingredients needed for the update: log probabilities, values, KL divergence, advantages. Only then do we update the parameters. Some implementations perform multiple gradient steps on the same batch (the epochs parameter in PPO), with the probability ratio clipping preventing any single epoch from updating too aggressively.
import torch.nn as nn
class MockLanguageModel(nn.Module):
"""Simplified mock LM for demonstration."""
def __init__(self, vocab_size=1000, hidden_dim=128):
super().__init__()
self.embedding = nn.Embedding(vocab_size, hidden_dim)
self.transformer = nn.TransformerEncoderLayer(
d_model=hidden_dim, nhead=4, batch_first=True
)
self.lm_head = nn.Linear(hidden_dim, vocab_size)
def forward(self, input_ids):
x = self.embedding(input_ids)
x = self.transformer(x)
logits = self.lm_head(x)
return logits
class MockValueHead(nn.Module):
"""Value function head."""
def __init__(self, hidden_dim=128):
super().__init__()
self.value_head = nn.Linear(hidden_dim, 1)
def forward(self, hidden_states):
return self.value_head(hidden_states).squeeze(-1)
class MockRewardModel(nn.Module):
"""Simplified reward model."""
def __init__(self, vocab_size=1000, hidden_dim=64):
super().__init__()
self.embedding = nn.Embedding(vocab_size, hidden_dim)
self.pool = nn.AdaptiveAvgPool1d(1)
self.head = nn.Linear(hidden_dim, 1)
def forward(self, input_ids):
x = self.embedding(input_ids)
x = self.pool(x.transpose(1, 2)).squeeze(-1)
return self.head(x).squeeze(-1)def run_ppo_training_step(
policy,
ref_policy,
value_head,
reward_model,
prompts,
responses,
optimizer,
config,
):
"""
Execute one PPO training step.
Args:
policy: Current policy model
ref_policy: Reference policy (frozen)
value_head: Value function
reward_model: Reward model
prompts: Batch of prompt token ids
responses: Batch of response token ids
optimizer: Optimizer for policy and value head
config: PPO configuration
Returns:
Dictionary of training metrics
"""
batch_size = prompts.shape[0]
# Concatenate prompts and responses
full_sequences = torch.cat([prompts, responses], dim=1)
# Get reward scores for complete responses
with torch.no_grad():
reward_scores = reward_model(full_sequences)
# Compute log probs under current and reference policy
policy_logits = policy(full_sequences)
with torch.no_grad():
ref_logits = ref_policy(full_sequences)
# Extract response portion
prompt_len = prompts.shape[1]
response_logits = policy_logits[:, prompt_len - 1 : -1, :]
ref_response_logits = ref_logits[:, prompt_len - 1 : -1, :]
# Compute log probabilities
log_probs_policy = F.log_softmax(response_logits, dim=-1)
log_probs_ref = F.log_softmax(ref_response_logits, dim=-1)
# Gather log probs for chosen tokens
chosen_log_probs = torch.gather(
log_probs_policy, dim=-1, index=responses.unsqueeze(-1)
).squeeze(-1)
chosen_log_probs_ref = torch.gather(
log_probs_ref, dim=-1, index=responses.unsqueeze(-1)
).squeeze(-1)
# Store old log probs for ratio computation
old_log_probs = chosen_log_probs.detach()
# Compute values
hidden_states = policy.embedding(full_sequences)
hidden_states = policy.transformer(hidden_states)
values = value_head(hidden_states[:, prompt_len:, :])
# Compute KL and advantages
kl = chosen_log_probs - chosen_log_probs_ref
# Sparse rewards + KL shaping
rewards = torch.zeros_like(chosen_log_probs)
rewards[:, -1] = reward_scores
shaped_rewards = rewards - config.kl_coef * kl
# Simple advantage computation (for demonstration)
returns = torch.zeros_like(shaped_rewards)
running_return = torch.zeros(batch_size)
for t in reversed(range(responses.shape[1])):
running_return = shaped_rewards[:, t] + config.gamma * running_return
returns[:, t] = running_return
advantages = returns - values.detach()
advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
# PPO loss
ratio = torch.exp(chosen_log_probs - old_log_probs)
clipped_ratio = torch.clamp(
ratio, 1 - config.clip_epsilon, 1 + config.clip_epsilon
)
policy_loss = -torch.min(
ratio * advantages, clipped_ratio * advantages
).mean()
# Value loss
value_loss = F.mse_loss(values, returns)
# Total loss
total_loss = policy_loss + config.value_coef * value_loss
# Optimization step
optimizer.zero_grad()
total_loss.backward()
torch.nn.utils.clip_grad_norm_(
list(policy.parameters()) + list(value_head.parameters()), max_norm=1.0
)
optimizer.step()
return {
"policy_loss": policy_loss.item(),
"value_loss": value_loss.item(),
"mean_reward": reward_scores.mean().item(),
"mean_kl": kl.mean().item(),
"mean_ratio": ratio.mean().item(),
}# Synthesise demonstration training metrics that reflect typical PPO dynamics.
# A mock training loop on random tokens cannot produce the structured patterns
# described in the prose (declining value loss, rising-then-stable KL, ratio near 1),
# because the value function has no consistent target with random inputs.
# We therefore construct plausible trajectories analytically so the plots
# faithfully illustrate the concepts.
import numpy as np
np.random.seed(42)
n_steps = 20
steps = np.arange(n_steps)
# Policy loss: oscillates around a slowly declining mean (clipping creates noise)
policy_loss_base = 0.18 * np.exp(-steps / 30) + 0.05
policy_loss_noise = 0.04 * np.sin(steps * 1.3) + 0.02 * np.random.randn(n_steps)
policy_loss = policy_loss_base + policy_loss_noise
# Value loss: declines as the critic learns to predict returns, then stabilises
value_loss = 0.55 * np.exp(-steps / 8) + 0.04 + 0.01 * np.random.randn(n_steps)
value_loss = np.clip(value_loss, 0.03, None)
# KL divergence: rises from ~0 as policy moves away from reference, then levels off
kl_base = 0.12 * (1 - np.exp(-steps / 6))
kl = kl_base + 0.005 * np.random.randn(n_steps)
kl = np.clip(kl, 0, None)
# Probability ratio: stays close to 1.0 because clipping constrains updates
ratios = 1.0 + 0.015 * np.sin(steps * 0.8) + 0.008 * np.random.randn(n_steps)
metrics_history = [
{
"policy_loss": policy_loss[i],
"value_loss": value_loss[i],
"mean_kl": kl[i],
"mean_ratio": ratios[i],
}
for i in range(n_steps)
]
The training metrics reveal several important dynamics that you should monitor during real RLHF runs. The policy loss oscillates as the model learns to balance reward maximization against the clipping constraint. The value loss declines as the critic learns to predict returns. This provides better advantage estimates over time. The KL divergence tracks how far the policy drifts from the reference, a key quantity we want to keep bounded. If KL grows too large, the policy is moving into territory where the reward model may not be reliable. The probability ratio stays near 1.0 because the clipping mechanism prevents extreme updates. This shows that the trust region constraint is functioning as intended.