Proximal Policy Optimization

1. What a Policy Means for a Language Model

A language model generates text one token at a time. At each step, it assigns probabilities to possible next tokens based on the prompt and the text already generated. Those probabilities form its policy: a rule for choosing an action given the current context. Here, an action is a token, and the context is the prompt plus the response prefix.

Suppose a particular token has a 20% probability at one position. Updating the model can make that token more or less likely when evaluated with the same prefix. A policy update changes the model's weights, the learned numbers that determine these probabilities. We are adjusting the distribution from which future responses will be generated, rather than editing a response that has already been written.

In RLHF, a reward model provides a score based on learned human preferences. PPO uses reward-derived feedback to update the generating policy. The algorithm can also work with other reward sources and with tasks outside language modeling; human preference learning and policy optimization are separate parts of the training system.

2. Collect Responses Before Updating

Training starts by sampling responses to a batch of prompts. These sampled sequences are called rollouts. For every generated token, record the probability that the model assigned to that token at its actual prefix. Implementations normally store its logarithm, called a log-probability, because sums and differences of logs are convenient for the later calculations.

The policy that collected this batch is the old policy. Once the responses are collected, their tokens and recorded probabilities stay fixed while the current policy takes several optimization steps. Each new forward pass evaluates those same tokens at those same prefixes with the updated weights. A fresh response is collected in the next rollout round, not for every small update on this batch.

This distinction makes data reuse possible. We know both how likely a sampled action was when it occurred and how likely it has become during training. PPO's objective uses that comparison while discouraging excessive movement. The original PPO paper describes this alternation between collecting experience and taking several minibatch updates on it.

3. Compare the Outcome with an Expectation

A reward scores an outcome. A return adds up the rewards from a particular decision onward, possibly giving less weight to later rewards. If a response receives only one reward at the end and we do not discount it, the return at every generated token equals that final reward. Intermediate penalties, such as the reference penalty in RLHF, make the returns differ across positions.

A value model, often called the critic, predicts the expected return from each prefix before the next token is chosen. This gives us a baseline. A simple estimate of the token's advantage is the observed return minus that baseline. Positive means better than expected; negative means worse than expected. A positive reward can therefore produce a negative advantage if the baseline was higher.

The advantage supplies the direction of the learning signal. A positive estimate encourages the sampled action, while a negative estimate discourages it. It is an estimate from a sampled outcome, not proof that one word caused the response to succeed or fail. Learning from many examples helps separate useful patterns from noise.

4. Measure How a Token Probability Changes

Return to the token whose recorded probability was 20%. If the updated model gives it 30%, its probability is now 1.5 times the old value. If it falls to 10%, the ratio is 0.5. PPO calls this the probability ratio: current probability divided by the recorded rollout probability, always for the same sampled token at the same prefix.

The ratio also accounts for evaluating an updated policy using actions sampled from the old one. PPO constructs a local training objective from the ratio multiplied by the estimated advantage. For an advantage of +1, increasing the ratio raises this objective. For an advantage of −1, decreasing the ratio raises it because the contribution becomes less negative. In both cases, the objective is something we want to maximize.

Without a limiting mechanism, repeated updates could keep pushing in the same direction based on the same small set of outcomes. The ratio corrects the action weighting, but it does not turn arbitrarily old responses into fresh evidence about every prefix a changed model might visit. PPO therefore combines the ratio with a clipped objective and periodically collects new data.

5. How PPO Clipping Works

Choose a clipping threshold, usually written as ε, or epsilon. With ε = 0.2, the clipping interval for the ratio is 0.8 to 1.2. PPO compares two terms: the ordinary ratio times advantage, and a second version using the ratio clipped to that interval. It takes the smaller of the two. This prevents the clipped term from hiding changes that make the objective worse.

For a positive advantage of +1 and a ratio of 1.5, the ordinary term is 1.5 and the clipped term is 1.2. PPO uses 1.2. Pushing the probability still higher earns no additional credit from this sampled action. If the ratio instead falls to 0.5, PPO uses 0.5 rather than the clipped value of 0.8, preserving the signal to correct that decrease.

For a negative advantage of −1, the useful direction is to lower the token's probability. At a ratio of 0.5, the ordinary term is −0.5 and the clipped term is −0.8, so PPO uses −0.8. Further decreases receive no extra credit from this term. At a ratio of 1.5, PPO retains −1.5, which continues to discourage making this action more likely.

6. What the Clipping Threshold Controls

An epsilon of 0.2 refers to a relative probability ratio. Starting from 20%, the interval 0.8 to 1.2 corresponds to probabilities of 16% to 24%. It does not mean adding or subtracting twenty percentage points, and it does not limit each model weight to a 20% change. The numerical threshold in this article is illustrative, not a setting that is right for every run.

Clipping changes the objective; it does not physically stop a probability at the interval boundary. Tokens share model parameters, so an update encouraged by one token can move another. Optimizer momentum and other loss terms can also keep probabilities moving after a particular contribution becomes flat. PPO-Clip therefore offers no hard guarantee that the whole policy remains within a fixed distance of the old one.

A learning rate still controls the optimizer's step size. Gradient clipping, which limits the magnitude of a gradient vector, is a separate operation from PPO's probability-ratio clipping. Trainers may also stop further updates on a rollout batch when an estimated divergence from the old policy becomes too large. The PPO implementation study discusses these additional training choices.

7. One PPO Update Round

After collecting rollouts, compute rewards, value estimates, returns, and advantages. Split the collected examples into minibatches, smaller groups used for individual optimizer steps. An epoch is one pass over that rollout dataset; PPO can make a limited number of these passes before collecting fresh responses.

Replacing the recorded old probabilities with the current probabilities before each step would make every ratio equal to one again, erasing the comparison to the policy that generated the data. For the same reason, old log-probabilities and advantage targets are treated as constants in the policy loss. Gradients flow through the current policy's predictions.

8. The Clipped Policy Loss in PyTorch

The function below implements the clipped policy term for already collected token-level data. All tensors have matching shapes, usually batch size by response length. Each log-probability belongs to the sampled token, not the whole vocabulary. A boolean response mask selects generated actions and excludes prompt positions and padding.

import torch


def clipped_policy_loss(
    new_logprobs, old_logprobs, advantages, response_mask, epsilon=0.2
):
    if not 0 < epsilon < 1:
        raise ValueError("epsilon must be between zero and one")
    if not (
        new_logprobs.shape == old_logprobs.shape
        == advantages.shape == response_mask.shape
    ):
        raise ValueError("All inputs must have matching shapes")
    if response_mask.dtype != torch.bool or not response_mask.any():
        raise ValueError("Provide a boolean mask with valid response tokens")

    # Select valid actions before arithmetic so padding cannot contaminate loss.
    current = new_logprobs[response_mask].float()
    recorded = old_logprobs.detach()[response_mask].float()
    advantage = advantages.detach()[response_mask].float()
    ratio = torch.exp(current - recorded)

    if not torch.isfinite(ratio).all() or not torch.isfinite(advantage).all():
        raise ValueError("Non-finite ratios or advantages; inspect the rollout")

    ordinary = ratio * advantage
    clipped = ratio.clamp(1 - epsilon, 1 + epsilon) * advantage
    return -torch.minimum(ordinary, clipped).mean()

Subtracting log-probabilities and exponentiating gives the probability ratio. The function detaches the old probabilities and advantages, takes the smaller objective term, and negates the average because optimizers minimize a loss. It selects valid positions before arithmetic so ignored padding values do not enter the calculation. An empty response mask is an error rather than a valid update with no actions.

Token alignment matters when preparing these inputs. A causal model's logits at a position predict the following token. Gather the log-probability of the generated action from the corresponding prediction position, and align its old log-probability, advantage, and mask. A terminal end-of-sequence token can be a valid sampled action; padding after termination is excluded.

This implementation averages equally over valid response tokens in its input batch. Other trainers average within each response and then across responses, which changes the weighting when lengths differ. Choose the reduction deliberately. The function is one tested loss component, not a complete trainer: rollout generation, consistent sampling probabilities, advantage estimation, value loss, and optimizer scheduling remain separate responsibilities.

9. The Critic and the RLHF Reference Model

The critic learns alongside the policy by fitting its value predictions to the return targets. A simple value loss averages the squared difference between each current prediction and its fixed target. Some implementations also clip value updates, but that is another design choice. The policy and critic may share a backbone or use separate models, which changes how their losses interact through shared parameters.

In classic RLHF, there is also a frozen reference policy, often the supervised fine-tuned starting model. Its role differs from PPO's old policy. The old policy generated the current rollout batch and anchors its probability ratios. The reference supplies a longer-term baseline for a KL penalty, which discourages the trained model from moving too far from its starting response distribution.

A common arrangement subtracts a reference log-probability penalty along the generated tokens and adds the reward-model score at the end. Those rewards feed the return and advantage calculation before the PPO update. The penalty coefficient β and the clipping threshold ε control different parts of training. The InstructGPT method uses PPO together with a per-token penalty relative to the SFT model.

An old-to-current divergence check for early stopping is also separate from the reference penalty. One concerns movement during this update round; the other concerns drift from a fixed model over the training phase. Monitoring one does not automatically constrain the other, and none of these mechanisms makes an inaccurate reward signal reliable.

10. Collect Fresh Responses and Repeat

After the chosen number of update epochs, or an early-stopping condition, the updated model collects another rollout batch. Its probabilities become the new recorded old probabilities for that batch. Returns and advantages are computed from the new outcomes, and the next optimization round begins. PPO is generally treated as an on-policy method because it regularly refreshes experience from the policy it is training, even though it reuses a batch for several nearby updates.

More epochs reuse the same responses more heavily; they do not provide more independent feedback. The critic's estimates, reward quality, clipping threshold, learning rate, and batch composition all influence the result. A rising training objective on old responses is not sufficient evidence of better answers on new prompts, so quality must be evaluated beyond the batch used for optimization.

The final product is an updated language-model policy. It still generates text autoregressively, one token after another. PPO's rollout records, advantage estimates, and clipping rule belong to training; ordinary inference does not run a PPO update for each answer. Their purpose was to shape the probabilities that the trained model now uses.