KL Divergence Penalty (RLHF) - Reinforcement Learning

KL Divergence Penalty (RLHF)

What you'll build

The sampled KL penalty β·KL(π_θ ‖ π_ref) that RLHF uses to keep the fine-tuned policy close to the reference model.

Background

A policy optimized against a learned reward model soon finds outputs that the reward model overrates ("reward hacking") and loses fluency. RLHF penalizes divergence from the frozen reference (SFT) model:

reward_total = reward - β · KL(π_θ ‖ π_ref)

The exact KL needs the full vocabulary distribution at every position. In practice it is estimated from the sampled tokens using the log-probs you already have.

Formula

For tokens y_t sampled from π_θ, the per-token estimate (often called k1) is the log-ratio:

kl_t    = log π_θ(y_t) - log π_ref(y_t)
penalty = β · mean_t(kl_t)

Its average over samples from π_θ is the true KL, so it is unbiased. A single token's kl_t can be negative, and so can the mean over a few tokens. The k3 estimator, exp(-kl_t) + kl_t - 1, is never negative; GRPO uses it (see GRPO in Post-training & Efficiency).

Inputs

• logp_policy: (T,) log-probs of the sampled tokens under π_θ
• logp_reference: (T,) log-probs of the same tokens under π_ref
• beta: KL coefficient (e.g. 0.01 to 0.2)

Expected return format

A Python float: β · mean(logp_policy - logp_reference), rounded to 2 decimals. Keep the sign; don't clip negative estimates to 0.

Edge cases and conventions

• Identical log-probs give 0.0.
• The mean is over all given tokens (no mask in this problem).

Constraints

NumPy only.

Follow-up

• InstructGPT-style PPO subtracts β·kl_t from the reward of every token, not once per sequence. How does that change the credit assignment?
• Adaptive KL control raises β when the measured KL exceeds a target. Sketch the update rule.

Input: logp_policy = [-1, -2], logp_reference = [-1.5, -2.5], beta = 0.1
Output: 0.05

Input: logp_policy = [-2, -3], logp_reference = [-1, -1], beta = 1.0
Output: -1.5

import numpy as np

def kl_penalty(logp_policy: np.ndarray, logp_reference: np.ndarray,
               beta: float) -> float:
    """
    Compute the sampled KL penalty for RLHF.

    Args:
        logp_policy: (T,) log-probs of the sampled tokens under pi_theta
        logp_reference: (T,) log-probs of the same tokens under pi_ref
        beta: KL penalty coefficient

    Returns:
        beta * mean(logp_policy - logp_reference) as a float, rounded to 2 decimals
    """
    logp_policy = np.array(logp_policy, dtype=float)
    logp_reference = np.array(logp_reference, dtype=float)

    # Your code here
    pass