PPO and RLHF

Tabular methods store one value per state-action pair: enough for a handful of grid cells and a handful of moves. A language model's state is every possible token prefix, and its next action is a token from a large vocabulary: too many pairs to store. RLHF scores the finished text and updates the model toward higher-scoring completions. PPO is the update that keeps those steps from rewriting the starting model into fluent nonsense.

From DQN to policy gradients

A policy maps a state to an action or a distribution over actions. Following it produces a rollout: the states, actions, and rewards along the way. DQN estimates a Q-value for each state-action pair with a neural network and picks the action with the highest score. PPO learns the policy itself.

DQN stores past transitions in a replay buffer and later samples them for training. Those older transitions remain useful because the Q-learning target does not require the current policy (off-policy). PPO trains on a recent rollout: it reuses that batch for a few updates, then collects again, because older experience stops matching the policy being updated.

PPO in practice: CartPole

PPO's policy is easiest to see on a small control task. CartPole-v1 balances a pole upright on a cart that slides left or right. It is fast to simulate, and every state, action, and reward stays visible. Each observation contains the cart position, cart velocity, pole angle, and pole angular velocity. The policy network maps those four numbers to two logits, which define a categorical distribution over a left push and a right push.

import gymnasium as gym
import torch
from torch import nn
from torch.distributions import Categorical

torch.manual_seed(8)
sample_generator = torch.Generator().manual_seed(8)
actor = nn.Sequential(nn.Linear(4, 16), nn.Tanh(), nn.Linear(16, 2))
env = gym.make("CartPole-v1")

observation, reset_info = env.reset(seed=8)
observation_tensor = torch.tensor(observation, dtype=torch.float32).unsqueeze(0)
with torch.no_grad():
    logits = actor(observation_tensor)
    policy = Categorical(logits=logits)
    probabilities = policy.probs
    action = torch.multinomial(probabilities[0], 1, generator=sample_generator).squeeze(0)
    log_probability = policy.log_prob(action.unsqueeze(0)).squeeze(0)

next_observation, reward, terminated, truncated, step_info = env.step(action.item())
env.close()

assert observation.shape == (4,)
assert observation_tensor.shape == (1, 4)
assert logits.shape == probabilities.shape == (1, 2)
assert torch.all(probabilities >= 0) and torch.allclose(probabilities.sum(-1), torch.ones(1))
assert action.shape == torch.Size([]) and action.item() in (0, 1)
assert log_probability.shape == torch.Size([]) and torch.isfinite(log_probability)
assert next_observation.shape == (4,)
assert isinstance(reward, float) and reward == 1.0
assert isinstance(terminated, bool) and isinstance(truncated, bool)

print("observation:", observation.round(4).tolist())
print("action probabilities [left, right]:", probabilities[0].numpy().round(4).tolist())
print("sampled action:", "left" if action.item() == 0 else "right")
print("stored log probability:", round(log_probability.item(), 4))
print("reward:", reward)
print("terminated:", terminated, "truncated:", truncated)

A language model makes the same kind of stochastic choice over a larger action set. The CartPole observation becomes the prompt and generated prefix, the two action logits become vocabulary logits, and the sampled left or right push becomes a sampled token.

The cart fails when it leaves the track or the pole tilts more than 12 degrees. That episode is over: the pole fell, and there is no more reward. Hitting the 500-step timer is different. Recording stops, but the pole may still be up. A policy that lasts that long is succeeding. If PPO treated the timer like a fall, those best runs would look like failures, and the update would punish the actions that kept the pole balanced.

Actor, critic, and advantage

A policy is only useful once something tells it which actions to make more likely. The reward comes from the environment, and an environment may not be differentiable, so there may be no gradient to send through it. The log-probability trick sidesteps that: raise the log-probability of actions that showed up in high-return trajectories, and lower it for actions in low-return ones. The goal is to maximize expected return along a trajectory $\tau$:

$$J\left(\theta \right)=E_{\tau \sim {\pi }_{\theta }}\left[R\right(\tau \left)\right]$$

Treat expected return as height over all possible policy weights. The current policy is one point on that landscape. Its gradient points in the steepest local uphill direction.

The policy gradient writes that uphill direction in terms of actions from sampled trajectories:

$${\nabla }_{\theta }J\left(\theta \right)=E\left[\sum _t{\nabla }_{\theta }\log {\pi }_{\theta }\right(a_t\mid s_t\left)R\right(\tau \left)\right]$$

That estimate is noisy. A CartPole episode can succeed or fail for reasons that have little to do with any single action, so weighting every action by the same overall return mixes "was this action good" with "did this episode go well." The same return can be impressive from one state and unremarkable from another. Compare the action to what was expected from that state.

PPO trains two networks. The actor is the policy: it picks actions. The critic estimates $V\left(s\right)$, how much return to expect from a state, and judges the states the actor lands in. On CartPole these can be small, separate networks. Later, on a language model, they share almost the entire network.

Subtracting the critic's judgment from the return $G$ gives the advantage:

$$A_t=G_t-V\left(s_t\right)$$

A positive advantage means the action did better than the critic expected, so the update makes it more likely. A negative advantage means it did worse, so the update makes it less likely. Zero advantage means the return matched the prediction: the action was unsurprising, even if the return itself was large. Subtracting a baseline like this does not change the expected gradient, as long as the baseline does not depend on the action taken. It only cuts noise: it compares an action to what was expected from that specific state. In practice, PPO estimates the advantage from a finite rollout rather than calculating the true return and $V$.

The policy gradient now weights each sampled action by that surprise:

$${\nabla }_{\theta }J\left(\theta \right)\approx \sum _t{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t){\hat{A}}_t$$

The log-probability gradient is the direction that makes this action more likely. The advantage scales it: how surprisingly good or bad the action was.

PPO rollouts

PPO collects each rollout with a frozen snapshot of the policy, the old policy. Each stored step keeps the state, the sampled action, the reward, termination and truncation flags, the old-policy log probability, and the critic value. Advantages and return targets are computed once from that batch and stay fixed while the current policy updates.

The sampled actions still came from the old policy. Their stored probabilities let PPO measure how far the current policy has moved on those actions.

Generalized advantage estimation

Rewards can arrive long after the action that helped earn them. The critic also gives a faster, one-step judgment: the TD residual. It compares one transition to the critic's expectation, using only the next reward instead of the rest of the episode.

$${\delta }_t=r_{t+1}+\gamma (1-d_{t+1})V\left(s_{t+1}\right)-V\left(s_t\right)$$

Here $d_{t+1}$ is 1 when the pole fell or the cart left the track, and 0 otherwise. A fall zeros the next value. A 500-step cutoff does not: the last observation can still have value.

TD residuals are quick and less noisy, but only as accurate as the critic. Full returns use real rewards to the end of the episode: less biased, noisier. Generalized advantage estimation (GAE) blends the two. It sends later critic surprises back to earlier actions, discounted by $\gamma$ (how much future reward matters) and $\lambda$ (how far ahead to trust the critic):

$${\hat{A}}_t={\delta }_t+\gamma \lambda (1-d_{t+1}){\hat{A}}_{t+1}$$

$\lambda =0$ uses only the one-step residual. $\lambda =1$ moves toward the full return. Values near 0.95 keep most of the return and smooth most of the noise. Credit never crosses from one episode into the next.

import numpy as np

rewards = np.array([0.0, 1.0, 0.0])
values = np.array([0.2, 0.4, 0.3])
next_values = np.array([0.4, 123.0, 0.5])
next_terminated = np.array([0.0, 1.0, 0.0])
gamma, lam = 0.9, 0.8

def gae(rewards, values, next_values, next_terminated, gamma, lam):
    advantages = np.zeros_like(rewards)
    later = 0.0
    for t in reversed(range(len(rewards))):
        continues = 1.0 - next_terminated[t]
        delta = rewards[t] + gamma * continues * next_values[t] - values[t]
        later = delta + gamma * lam * continues * later
        advantages[t] = later
    return advantages

advantages = gae(rewards, values, next_values, next_terminated, gamma, lam)
returns = advantages + values
changed_terminal_value = next_values.copy()
changed_terminal_value[1] = -999.0

assert np.allclose(advantages, [0.592, 0.6, 0.15])
assert np.allclose(returns, [0.792, 1.0, 0.45])
assert np.allclose(
    advantages,
    gae(rewards, values, changed_terminal_value, next_terminated, gamma, lam),
)
print("advantages:", advantages.round(3).tolist())
print("returns:", returns.round(3).tolist())

The first advantage includes part of the second step's residual because both belong to the same episode. The terminal transition masks its placeholder next value and stops the later reset episode from affecting it.

PPO clipping

The advantage says which sampled actions to make more or less likely. It does not cap how far the update can go. PPO reuses one rollout for several updates, so the policy doing the updating is no longer the policy that collected the data. The probability ratio compares them for each sampled action:

$$r_t\left(\theta \right)=\frac{{\pi }_{\theta }(a_t\mid s_t)}{{\pi }_{\text{old}}(a_t\mid s_t)}$$

A ratio of 1 means the current policy assigns that action the same probability as the old one. Above 1, it is now more likely; below 1, less likely. Multiply by the advantage and a positive advantage rewards raising the ratio with no limit: one noisy batch can swing the policy a long way. The clipped objective takes the smaller of that product and a version with the ratio clamped near 1:

$$L_{\text{clip}}=E_t\left[\min \right(r_t{\hat{A}}_t,\mathrm{clip}(r_t,1-\epsilon ,1+\epsilon ){\hat{A}}_t\left)\right]$$

For a positive advantage, the objective stops improving once the ratio passes $1+\epsilon$. For a negative advantage, it stops improving once the ratio drops below $1-\epsilon$. The policy can still change. A batch just stops getting extra credit for changing it a lot.

import numpy as np

ratios = np.array([0.5, 0.8, 1.0, 1.2, 1.5])
eps = 0.2

def clipped_objective(ratios, advantage):
    original = ratios * advantage
    clipped = np.clip(ratios, 1 - eps, 1 + eps) * advantage
    return np.minimum(original, clipped)

positive = clipped_objective(ratios, 1.0)
negative = clipped_objective(ratios, -1.0)
assert np.allclose(positive, [0.5, 0.8, 1.0, 1.2, 1.2])
assert np.allclose(negative, [-0.8, -0.8, -1.0, -1.2, -1.5])
print("positive advantage:", positive.tolist())
print("negative advantage:", negative.tolist())

Approximate KL and clip fraction measure how far the policy moved on this batch.

PPO objective and learning phase

Two more terms complete the objective. A squared value loss trains the critic to predict return. An entropy bonus rewards keeping the action distribution spread out, so the policy does not collapse to one action too early.

$$L_{\text{clip}}-c_v{L}_{\text{value}}+c_{H}H\left({\pi }_{\theta }\right)$$

During each learning phase the trainer shuffles the rollout into minibatches and recomputes current log probabilities, values, and entropy. Stored old log probabilities, advantages, and returns stay fixed.

import torch
from torch import nn
from torch.distributions import Categorical

torch.manual_seed(8)

class TinyAgent(nn.Module):
    def __init__(self):
        super().__init__()
        # With one-hot observations and no biases, each state updates its own weights.
        self.actor = nn.Linear(4, 2, bias=False)
        self.critic = nn.Linear(4, 1, bias=False)
        nn.init.zeros_(self.actor.weight)
        nn.init.zeros_(self.critic.weight)

    def forward(self, observations):
        return Categorical(logits=self.actor(observations)), self.critic(observations).squeeze(-1)

agent = TinyAgent()
observations = torch.eye(4)
actions = torch.tensor([0, 1, 0, 1])
advantages = torch.tensor([1.0, 1.0, -1.0, -1.0])

with torch.no_grad():
    old_distribution, old_values = agent(observations)
    old_log_probs = old_distribution.log_prob(actions)
    returns = advantages + old_values
    before_probabilities = old_distribution.probs[range(4), actions]
    before_mse = ((old_values - returns) ** 2).mean()

# These tensors are rollout records. Optimization must never rewrite them.
fixed_rollout = tuple(t.clone() for t in (old_log_probs, old_values, advantages, returns))
actor_before = agent.actor.weight.detach().clone()
critic_before = agent.critic.weight.detach().clone()
optimizer = torch.optim.SGD(agent.parameters(), lr=0.1)
epsilon = 0.2

for _ in range(6):
    distribution, values = agent(observations)
    ratios = (distribution.log_prob(actions) - old_log_probs).exp()
    surrogate = torch.minimum(
        ratios * advantages,
        ratios.clamp(1 - epsilon, 1 + epsilon) * advantages,
    ).mean()
    value_loss = ((values - returns) ** 2).mean()
    loss = -surrogate + 0.5 * value_loss
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    assert all(torch.equal(current, saved) for current, saved in zip(
        (old_log_probs, old_values, advantages, returns), fixed_rollout
    ))

with torch.no_grad():
    final_distribution, final_values = agent(observations)
    after_probabilities = final_distribution.probs[range(4), actions]
    after_mse = ((final_values - returns) ** 2).mean()
actor_change = torch.linalg.vector_norm(agent.actor.weight - actor_before)
critic_change = torch.linalg.vector_norm(agent.critic.weight - critic_before)

assert torch.all(after_probabilities[:2] > before_probabilities[:2])
assert torch.all(after_probabilities[2:] < before_probabilities[2:])
assert actor_change > 0 and critic_change > 0
assert after_mse < before_mse
assert all(torch.equal(current, saved) for current, saved in zip(
    (old_log_probs, old_values, advantages, returns), fixed_rollout
))

print("sampled-action probabilities before:", before_probabilities.tolist())
print("sampled-action probabilities after: ", after_probabilities.tolist())
print("actor parameter-change norm:", round(actor_change.item(), 6))
print("critic parameter-change norm:", round(critic_change.item(), 6))
print("critic MSE:", round(before_mse.item(), 6), "->", round(after_mse.item(), 6))
print("fixed rollout unchanged across epochs: True")

Language generation as reinforcement learning

Autoregressive generation defines an RL process. The language model reads the prompt and generated prefix, produces a distribution over the vocabulary, samples the next token, and appends it to form the next state. The completed sequence receives a score.

Two things change from CartPole. CartPole scores every step; a language model is scored once, after the whole completion. Nothing so far stops the policy from wandering far from where it started. For text, that is dangerous.

RLHF pipeline

A full RLHF pipeline first trains a model on demonstrations, then learns a reward model from ranked responses, then optimizes the language-model policy against that reward. The third step is the policy update. A fixed scoring function stands in for the learned reward model.

Give the model a reward that counts periods, and it learns "U.S.", "Python 2.7.12", and short choppy sentences until the text stops reading like language. That is reward hacking: the model does what the reward measures, not what anyone wanted. A frozen reference model is the leash back to the starting point.

Rollout, current, and reference policies

RLHF adds that leash: a frozen copy of the starting model. A KL penalty measures how far the trained distribution has moved from that reference and subtracts it from the objective. The policy can still chase reward; each step away from the reference costs something.

Language-model PPO uses three policy roles:

• The old rollout policy ${\pi }_{\text{old}}$ generated the current batch. Its stored sampled-token log probabilities form the denominator of each PPO ratio and refresh with every rollout.
• The current policy ${\pi }_{\theta }$ receives gradients during minibatch updates.
• The reference policy ${\pi }_{\text{ref}}$ is that frozen starting model. Its full next-token distribution anchors the current policy across the run.

Stored old log probabilities measure how far this sampled token has moved since collection. The reference distribution measures how far the whole current distribution has moved from the starting model.

Transformer actor and critic

The critic needs a value from the same network that generates text. A value head reads the final residual stream and outputs one number per position. The policy still produces vocabulary logits; the two heads share almost the entire transformer. The head reads ln_final.hook_normalized, then applies a linear layer from $d_{\text{model}}$ to $4d_{\text{model}}$ features, ReLU, and a linear layer from $4d_{\text{model}}$ features to one scalar.

import torch
from torch import nn
from transformer_lens import HookedTransformer

torch.manual_seed(8)
generator = torch.Generator(device="cpu").manual_seed(8)
model = HookedTransformer.from_pretrained("gpt2-small", device="cpu")
model.eval()
for parameter in model.parameters():
    parameter.requires_grad_(False)

value_head = nn.Sequential(
    nn.Linear(model.cfg.d_model, 4 * model.cfg.d_model),
    nn.ReLU(),
    nn.Linear(4 * model.cfg.d_model, 1),
).eval()

prompt = "Learning from feedback"
tokens = model.to_tokens(prompt, prepend_bos=True)
prefix_len = tokens.shape[1]
sampled_probabilities = []
with torch.no_grad():
    for _ in range(3):
        next_token_probabilities = model(tokens)[:, -1].softmax(dim=-1)
        next_token = torch.multinomial(next_token_probabilities, 1, generator=generator)
        sampled_probabilities.append(next_token_probabilities.gather(1, next_token).item())
        tokens = torch.cat([tokens, next_token], dim=1)

    _, cache = model.run_with_cache(tokens)
    hidden = cache["ln_final.hook_normalized"]
    token_logits = model.unembed(hidden)
    values = value_head(hidden).squeeze(-1)
    prediction_positions = torch.arange(prefix_len - 1, tokens.shape[1] - 1)
    generated_tokens = tokens[:, prefix_len:]
    generated_logits = token_logits[:, prediction_positions]
    generated_values = values[:, prediction_positions]
    replayed_probabilities = generated_logits.softmax(dim=-1).gather(
        -1, generated_tokens.unsqueeze(-1)
    ).squeeze(-1)

assert hidden.shape == (1, tokens.shape[1], model.cfg.d_model)
assert token_logits.shape == (1, tokens.shape[1], model.cfg.d_vocab)
assert values.shape == (1, tokens.shape[1])
assert generated_logits.shape == (1, 3, model.cfg.d_vocab)
assert generated_values.shape == replayed_probabilities.shape == (1, 3)
assert torch.isfinite(generated_logits).all() and torch.isfinite(generated_values).all()
assert torch.allclose(
    replayed_probabilities[0], torch.tensor(sampled_probabilities), atol=1e-6, rtol=1e-5
)

print("sampled completion:", repr(model.to_string(generated_tokens[0])))
print("hidden shape:", tuple(hidden.shape))
print("generated logits shape:", tuple(generated_logits.shape))
print("generated values shape:", tuple(generated_values.shape))
for token, probability, value in zip(
    model.to_str_tokens(generated_tokens[0]), replayed_probabilities[0], generated_values[0]
):
    print(f"token={token!r} probability={probability.item():.6f} value={value.item():.6f}")

The one-position shift matters. A residual at position $j-1$ predicts the action token at position $j$. The final generated token's own residual would predict another token, so it does not score the action that produced it.

import numpy as np

tokens = ["<bos>", "Write", "Hi", ".", "!"]
prefix_len = 2
generated_actions = tokens[prefix_len:]
prediction_positions = list(range(prefix_len - 1, len(tokens) - 1))
old_values = np.array([0.2, 0.5, 0.4])
terminal_reward = 2.0

action_values = np.append(old_values[1:], terminal_reward)
advantages = action_values - old_values
pairs = [(tokens[position], action) for position, action in zip(prediction_positions, generated_actions)]

assert prediction_positions == [1, 2, 3]
assert generated_actions == ["Hi", ".", "!"]
assert pairs == [("Write", "Hi"), ("Hi", "."), (".", "!")]
assert np.allclose(action_values, [0.5, 0.4, 2.0])
assert np.allclose(advantages, [0.3, -0.1, 1.6])
print("prediction positions:", prediction_positions)
print("Q:", action_values.tolist())
print("advantages:", advantages.round(3).tolist())

Intermediate generated positions use the next old value as the action-value estimate. The final token uses the sequence reward. This one-step construction differs from running CartPole GAE backward through a rollout.

RLHF objective

Each generated token carries fixed rollout data: its old log probability, advantage, and return target. During learning, the current model recomputes that token's probability, value, and entropy, while the reference model supplies a frozen comparison distribution.

The PPO ratio compares the current and old probabilities of the sampled token. The reference penalty compares the full current and reference distributions at each generated-token prediction position:

$$D_{\text{KL}}({\pi }_{\theta }\parallel {\pi }_{\text{ref}})=\sum _x{\pi }_{\theta }(x\mid s)\log \frac{{\pi }_{\theta }(x\mid s)}{{\pi }_{\text{ref}}(x\mid s)}$$

import numpy as np

old_sampled_probability = 0.4
current_sampled_probability = 0.5
current_distribution = np.array([0.5, 0.5])
reference_distribution = np.array([0.8, 0.2])

ppo_ratio = current_sampled_probability / old_sampled_probability
forward_kl = np.sum(
    current_distribution * np.log(current_distribution / reference_distribution)
)

assert np.isclose(ppo_ratio, 1.25)
assert np.isclose(forward_kl, 0.22314355)
print(f"PPO ratio: {ppo_ratio:.2f}")
print(f"forward KL: {forward_kl:.6f}")

The complete language-model objective adds the reference penalty to the actor-critic objective:

$$J_{\text{RLHF}}=L_{\text{clip}}-c_v{L}_{\text{value}}+c_{H}H-\beta D_{\text{KL}}({\pi }_{\theta }\parallel {\pi }_{\text{ref}})$$

Only generated-token positions enter the policy, value, entropy, and KL terms. Old log probabilities, advantages, and return targets stay fixed during each update.

A period-count reward

The fixed reward counts period characters in each decoded completion. Without a KL penalty the model chases that count into abbreviations, decimals, and fragments. With the penalty it can raise the score while still reading like the starting model. The generated text belongs beside the raw reward and the reference KL: the number going up is only half the story.

Further reading

• The ARENA PPO exercise notebook contains the probes, implementation tests, and CartPole training loop used in the assignment.
• The ARENA RLHF exercise notebook develops token-level PPO with a fixed reward and frozen reference model.
• OpenAI's Spinning Up introduction to vanilla policy gradient develops the policy-gradient and actor-critic foundation.
• Schulman et al., Proximal Policy Optimization Algorithms, introduces PPO's clipped surrogate objective.
• Schulman et al., High-Dimensional Continuous Control Using Generalized Advantage Estimation, develops GAE and its bias-variance tradeoff.
• Ouyang et al., Training language models to follow instructions with human feedback, describes preference collection, reward modeling, and language-model policy optimization.