RLHF Reward Model Loss - Reinforcement Learning

RLHF Reward Model Loss

What you'll build

The Bradley–Terry pairwise loss used to train the reward model in RLHF.

Background

Humans find it much easier to say which of two responses is better than to give an absolute score. RLHF (InstructGPT, 2022) therefore collects pairs (y_chosen, y_rejected) for a prompt x and trains a reward model r(x, y), usually an LLM with a scalar head, so that chosen responses score higher. That reward model then supplies the reward for PPO (see PPO-RLHF Objective) or best-of-n sampling.

Formula

Bradley–Terry model of a preference:

P(chosen ≻ rejected) = σ(r_chosen - r_rejected)

Negative log-likelihood over N pairs:

L = -(1/N) · Σ_i log σ(r_chosen_i - r_rejected_i)

Numerically stable form: log σ(x) = -log(1 + e^(-x)) = -np.logaddexp(0, -x).

Inputs

• r_chosen: (N,) reward-model scores of the preferred responses
• r_rejected: (N,) scores of the rejected responses (same order)

Expected return format

A Python float: the mean loss, rounded to 3 decimals.

Edge cases and conventions

• Equal scores give log 2 ≈ 0.693 per pair (a coin flip).
• Only the difference of scores matters: adding a constant to both leaves the loss unchanged.
• Must stay exact for large margins: for r_chosen - r_rejected = -100 the loss is 100.0. Don't add an eps inside the log (log(p + 1e-10) caps the loss near 23).

Constraints

NumPy only. Don't call a library BCE.

Follow-up

• This is binary cross-entropy with logits on the score difference and label 1. Why is the reward model only identified up to a constant shift?
• How would you use a margin, or more than two ranked responses (Plackett–Luce)?

Input: r_chosen = [2.0, 1.5], r_rejected = [1.0, 0.5]
Output: 0.313

Input: r_chosen = [0.0], r_rejected = [0.0]
Output: 0.693

import numpy as np

def sigmoid(x):
    return 1 / (1 + np.exp(-x))

def reward_model_loss(r_chosen: np.ndarray, r_rejected: np.ndarray) -> float:
    """
    Compute the Bradley-Terry reward model loss.

    Args:
        r_chosen: (N,) scores of the preferred responses
        r_rejected: (N,) scores of the rejected responses

    Returns:
        mean of -log sigmoid(r_chosen - r_rejected) as a float, rounded to 3 decimals
    """
    r_chosen = np.array(r_chosen, dtype=float)
    r_rejected = np.array(r_rejected, dtype=float)

    # Your code here
    pass