PPO-RLHF Objective - Reinforcement Learning

PPO-RLHF Objective

What you'll build

The PPO objective as used for RLHF: the clipped surrogate from PPO Clipped Objective, with the KL penalty from KL Divergence Penalty folded into the advantage.

Background

Classic RLHF (InstructGPT) fine-tunes the policy with PPO to maximize the reward model's score while staying close to the reference model. There are two common ways to add the KL term:

1. KL in the reward (InstructGPT, this problem): shape the reward as r - β·KL, then run the ordinary clipped surrogate on it.
2. KL as a separate loss term: L = clipped surrogate - β·KL. GRPO does this with the k3 estimator (see GRPO in Post-training & Efficiency).

Pick one, not both: subtracting β·KL again after it is already in the advantage would count the penalty twice.

Formula

Element-wise over N tokens (or responses), with log-probs as inputs:

ratio_i = exp(new_logp_i - old_logp_i)              π_θ / π_old
kl_i    = new_logp_i - ref_logp_i                   sampled KL to the reference (k1)
A_i     = rewards_i - β · kl_i                      KL-shaped reward, used as the advantage
J       = mean_i [ min( ratio_i · A_i, clip(ratio_i, 1-ε, 1+ε) · A_i ) ]

Simplifications for this exercise: the shaped reward is used directly as the advantage (no value baseline or GAE), and kl uses the current policy's log-probs. Real implementations compute the shaped reward once at rollout time, from the sampling policy, and treat it as a constant.

Inputs

Four (N,) float arrays: old_logp (policy that sampled the data), new_logp (current policy π_θ), ref_logp (frozen reference model) and rewards (reward-model scores); plus floats epsilon (clip range) and beta (KL coefficient).

Expected return format

A dictionary of Python floats, each rounded to 2 decimals:

• 'objective': J, to maximize (the training loss is -J)
• 'policy_loss': the clipped-surrogate term. It equals 'objective' here, because the KL is already inside the advantage. The name is kept for compatibility; in the separate-term variant it would be J before subtracting the penalty.
• 'kl_penalty': β · mean(kl_i), reported for monitoring only. It is not subtracted from J.

Edge cases and conventions

• If new = old = ref, ratio = 1 and kl = 0, so J = mean(rewards) and the penalty is 0.
• The ratio uses exp of a log-prob difference (no eps needed).

Constraints

NumPy only, vectorized.

Follow-up

• Why does the KL-shaped reward let the critic and GAE handle the penalty, while the separate-term version cannot?
• Which statistics would you watch to catch reward hacking? (Mean KL, reward, response length.)

Input: old_logp = [-2.0], new_logp = [-1.8], ref_logp = [-2.0], rewards = [1.0], epsilon = 0.2, beta = 0.1
Output: {'objective': 1.18, 'policy_loss': 1.18, 'kl_penalty': 0.02}

Input: old_logp = new_logp = ref_logp = [-1.0], rewards = [1.0], epsilon = 0.2, beta = 0.1
Output: {'objective': 1.0, 'policy_loss': 1.0, 'kl_penalty': 0.0}

import numpy as np

def ppo_rlhf_objective(old_logp: np.ndarray, new_logp: np.ndarray,
                       ref_logp: np.ndarray, rewards: np.ndarray,
                       epsilon: float, beta: float) -> dict:
    """
    Compute the PPO-RLHF objective with the KL penalty folded into the advantage.

    Args:
        old_logp: (N,) log-probs under the policy that sampled the data
        new_logp: (N,) log-probs under the current policy
        ref_logp: (N,) log-probs under the frozen reference model
        rewards: (N,) reward-model scores
        epsilon: PPO clipping parameter
        beta: KL penalty coefficient

    Returns:
        dict of floats rounded to 2 decimals:
            'objective':   clipped surrogate with A = rewards - beta * (new_logp - ref_logp)
            'policy_loss': the same clipped-surrogate value
            'kl_penalty':  beta * mean(new_logp - ref_logp), for monitoring
    """
    old_logp = np.array(old_logp, dtype=float)
    new_logp = np.array(new_logp, dtype=float)
    ref_logp = np.array(ref_logp, dtype=float)
    rewards = np.array(rewards, dtype=float)

    # Your code here
    pass