Exercise 09: Bradley-Terry reward model on synthetic preferences

"""Bradley-Terry reward model on synthetic preferences — exercise starter.

Fill in the TODOs, then from the repo root:

    pip install -r exercises/requirements.txt
    pytest exercises/09-reward-model/

See README.md for the task and HINTS.md if you're stuck. The reference
implementation is in solution/reward_model.py — look after you've tried.

The setup:
  - Each "response" is a 4-dim feature vector (think: quality, length,
    formality, accuracy — any 4 latent dims).
  - There is a hidden "true reward" r*(x) = w*·x for a fixed w*. The student's
    model never sees w*; it only sees preference labels.
  - A preference pair (chosen, rejected) is generated by sampling two random
    feature vectors and labeling them under the Bradley-Terry model itself:
        P(a is chosen) = σ(β · (r*(a) - r*(b)))
    Higher β → less noisy labels. β = 4.0 is "reasonably clean".
  - The student trains a small MLP reward model on these labeled pairs using
    the Bradley-Terry NLL: -log σ(r_chosen - r_rejected).
  - Success = the trained model's ranking on a held-out set correlates highly
    (Spearman > 0.85) with the true reward, even though it never saw it.

That last point is the whole reason reward modeling works: pairwise labels
identify the reward up to a monotone transform, which is enough for ranking
(and ranking is all RLHF needs).
"""

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


# ── synthetic-data constants ──────────────────────────────────────────────────

FEATURE_DIM = 4
# The "true" reward is a fixed linear function of the 4-dim feature vector.
# The student's reward model never sees these weights.
TRUE_W = torch.tensor([1.0, -0.5, 0.3, 2.0])
# Temperature on the Bradley-Terry preference distribution.
BETA = 4.0


# ── synthetic data (given to you) ─────────────────────────────────────────────

def true_reward(features: torch.Tensor) -> torch.Tensor:
    """Hidden reward function. Linear in the features.

    Args:
        features: (..., FEATURE_DIM) tensor.
    Returns:
        Scalar reward tensor with the leading shape of `features`.
    """
    return features @ TRUE_W


def make_preference_dataset(n_pairs: int, seed: int = 0) -> tuple:
    """Generate n_pairs of (chosen_features, rejected_features).

    For each pair, sample two random 4-dim feature vectors, compute their true
    rewards, then label which is "chosen" by sampling from the Bradley-Terry
    distribution: P(a is chosen) = σ(β · (r*(a) - r*(b))).

    Args:
        n_pairs: number of preference pairs.
        seed:    RNG seed.
    Returns:
        chosen:   (n_pairs, FEATURE_DIM) tensor.
        rejected: (n_pairs, FEATURE_DIM) tensor.
    """
    g = torch.Generator().manual_seed(seed)
    feats_a = torch.randn(n_pairs, FEATURE_DIM, generator=g)
    feats_b = torch.randn(n_pairs, FEATURE_DIM, generator=g)

    r_a = true_reward(feats_a)
    r_b = true_reward(feats_b)

    p_a = torch.sigmoid(BETA * (r_a - r_b))
    a_is_chosen = torch.rand(n_pairs, generator=g) < p_a

    chosen = torch.where(a_is_chosen.unsqueeze(-1), feats_a, feats_b)
    rejected = torch.where(a_is_chosen.unsqueeze(-1), feats_b, feats_a)
    return chosen, rejected


def make_eval_set(n: int, seed: int = 1) -> tuple:
    """Held-out features + their true rewards, for Spearman evaluation."""
    g = torch.Generator().manual_seed(seed)
    feats = torch.randn(n, FEATURE_DIM, generator=g)
    return feats, true_reward(feats)


# ── Spearman correlation (given to you; no scipy dep) ─────────────────────────

def spearman_corr(a: torch.Tensor, b: torch.Tensor) -> float:
    """Spearman rank correlation between two 1-D tensors.

    Equivalent to scipy.stats.spearmanr(a, b).correlation in the no-ties case
    that continuous-valued rewards give us.
    """
    a_np = a.detach().cpu().numpy()
    b_np = b.detach().cpu().numpy()
    ra = np.argsort(np.argsort(a_np)).astype(np.float64)
    rb = np.argsort(np.argsort(b_np)).astype(np.float64)
    ra -= ra.mean()
    rb -= rb.mean()
    denom = (np.linalg.norm(ra) * np.linalg.norm(rb))
    if denom == 0.0:
        return 0.0
    return float((ra @ rb) / denom)


# ── reward model ──────────────────────────────────────────────────────────────

class RewardModel(nn.Module):
    """4-dim feature vector → scalar reward, via a small MLP.

    Forward contract:
        single input,  shape (FEATURE_DIM,)   → scalar tensor, shape ()
        batched input, shape (B, FEATURE_DIM) → tensor of shape (B, 1)

    The (B, 1) batched shape matches HuggingFace-style reward heads; the
    scalar-on-single-input case is a convenience for the loss tests.
    """

    def __init__(self, feature_dim: int = FEATURE_DIM, hidden_dim: int = 32):
        super().__init__()
        # TODO: build a small MLP. The reference uses:
        #     Linear(feature_dim, hidden_dim) → ReLU
        #     Linear(hidden_dim, hidden_dim)  → ReLU
        #     Linear(hidden_dim, 1)
        # store it as self.net.
        raise NotImplementedError("RewardModel.__init__")

    def forward(self, features: torch.Tensor) -> torch.Tensor:
        # TODO:
        #   if features.dim() == 1:
        #       return self.net(features).squeeze(-1)   # shape ()
        #   return self.net(features)                   # shape (B, 1)
        raise NotImplementedError("RewardModel.forward")


# ── Bradley-Terry loss ────────────────────────────────────────────────────────

def bradley_terry_loss(
    chosen_rewards: torch.Tensor,
    rejected_rewards: torch.Tensor,
) -> torch.Tensor:
    """-log σ(r_chosen - r_rejected), averaged over the batch.

    This is the NLL under the Bradley-Terry model:
        P(chosen > rejected) = σ(r_chosen - r_rejected)

    Properties the tests rely on:
      - chosen >> rejected → loss → 0
      - chosen == rejected → loss = log(2) ≈ 0.693  (chance level)
      - chosen << rejected → loss grows linearly in the gap

    Inputs may be scalars (()), 1-D ((B,)), or 2-D ((B, 1)). Flatten both so
    the same code handles all shapes.

    Use F.logsigmoid rather than torch.log(torch.sigmoid(...)) — it's
    numerically stable when the input is large in absolute value.

    Args:
        chosen_rewards:   model rewards for the chosen response in each pair.
        rejected_rewards: model rewards for the rejected response in each pair.
    Returns:
        Scalar tensor.
    """
    # TODO:
    #   chosen = chosen_rewards.flatten()
    #   rejected = rejected_rewards.flatten()
    #   return -F.logsigmoid(chosen - rejected).mean()
    raise NotImplementedError("bradley_terry_loss")


# ── training ──────────────────────────────────────────────────────────────────

def train_reward_model(
    n_pairs: int = 4096,
    n_steps: int = 2000,
    batch_size: int = 64,
    lr: float = 3e-3,
    seed: int = 0,
) -> dict:
    """Train the reward model on synthetic Bradley-Terry preferences.

    Per step:
      1. Sample a minibatch of indices into the (chosen, rejected) dataset.
      2. Forward both batches through the model → r_chosen, r_rejected, both
         shape (B, 1).
      3. loss = bradley_terry_loss(r_chosen, r_rejected).
      4. zero_grad → backward → step.
      5. Append the loss; evaluate Spearman vs true reward on the held-out
         set and append that too.

    Args:
        n_pairs:    number of preference pairs to generate.
        n_steps:    number of optimizer steps.
        batch_size: minibatch size.
        lr:         Adam learning rate.
        seed:       RNG seed.

    Returns:
        {
            "model":    the trained RewardModel,
            "losses":   list of n_steps floats — per-step training loss,
            "spearman": list of n_steps floats — per-step held-out Spearman.
        }
    """
    torch.manual_seed(seed)

    chosen, rejected = make_preference_dataset(n_pairs, seed=seed)
    eval_feats, eval_true_r = make_eval_set(n=512, seed=seed + 1)

    model = RewardModel()
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)

    g = torch.Generator().manual_seed(seed + 2)

    losses: list = []
    spearmans: list = []

    for _step in range(n_steps):
        # TODO:
        #   idx = torch.randint(n_pairs, (batch_size,), generator=g)
        #   r_chosen = model(chosen[idx])
        #   r_rejected = model(rejected[idx])
        #
        #   loss = bradley_terry_loss(r_chosen, r_rejected)
        #
        #   optimizer.zero_grad()
        #   loss.backward()
        #   optimizer.step()
        #
        #   losses.append(loss.item())
        #
        #   with torch.no_grad():
        #       pred_r = model(eval_feats).squeeze(-1)
        #   spearmans.append(spearman_corr(pred_r, eval_true_r))
        raise NotImplementedError("train_reward_model: the update step")

    return {"model": model, "losses": losses, "spearman": spearmans}


# ── main ──────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    out = train_reward_model(seed=0, n_steps=2000)
    final_loss = out["losses"][-1]
    final_spearman = out["spearman"][-1]
    print(
        f"steps: {len(out['losses'])}  "
        f"final loss: {final_loss:.4f}  "
        f"final Spearman (held-out): {final_spearman:.4f}  "
        f"(random baseline ≈ 0.00, perfect ranking = 1.00)"
    )

A 3-layer MLP is plenty. The reference is literally:

self.net = nn.Sequential(
    nn.Linear(feature_dim, hidden_dim),
    nn.ReLU(),
    nn.Linear(hidden_dim, hidden_dim),
    nn.ReLU(),
    nn.Linear(hidden_dim, 1),
)

Hidden dim 32 works fine. The true reward is linear in the features, so a single linear layer would also work, but the MLP is closer to what a real reward head looks like.

Two shape contracts, one branch:

def forward(self, features):
    if features.dim() == 1:
        return self.net(features).squeeze(-1)   # scalar, shape ()
    return self.net(features)                   # (B, 1)

.squeeze(-1) on a 1-D input of shape (FEATURE_DIM,) will fail with a Linear(... → 1) because the output is shape (1,) and squeezing the last dim gives (), a scalar tensor. That’s what test_reward_model_forward expects for the single-input case.

Two lines after flattening:

chosen = chosen_rewards.flatten()
rejected = rejected_rewards.flatten()
return -F.logsigmoid(chosen - rejected).mean()

Why F.logsigmoid and not torch.log(torch.sigmoid(x)): when x is very negative, sigmoid(x) underflows to 0 and log(0) is -inf. F.logsigmoid does the computation in one numerically stable step.

Quick sanity check by hand: if chosen == rejected, the diff is 0, σ(0) = 0.5, -log(0.5) = log(2) ≈ 0.6931. That’s the “no signal” loss; anything below that means the model is doing better than chance on this batch.

The loop body:

idx = torch.randint(n_pairs, (batch_size,), generator=g)
r_chosen = model(chosen[idx])      # (B, 1)
r_rejected = model(rejected[idx])  # (B, 1)

loss = bradley_terry_loss(r_chosen, r_rejected)

optimizer.zero_grad()
loss.backward()
optimizer.step()

losses.append(loss.item())

with torch.no_grad():
    pred_r = model(eval_feats).squeeze(-1)
spearmans.append(spearman_corr(pred_r, eval_true_r))

Common mistakes:

• Calling RewardModel() inside the loop: you reinitialize the weights every step, nothing learns. The model and optimizer are created above the loop; just use them.
• optimizer.step() before loss.backward(): silent no-op, no learning.
• Swapping chosen and rejected in the loss: the loss will grow and Spearman will go negative. If you see that, the sign is flipped.
• Forgetting .detach() or torch.no_grad() on the eval forward pass: not wrong, but it builds and discards a graph 2000 times. The wrapper exists for a reason.

Three likely causes:

1. The model isn’t actually training. Print losses[0] and losses[-1]. The first one should be near log(2) ≈ 0.69 (random init, no separation yet). The last should be well under 0.2. If the loss is flat across all steps, see hint 4; you probably re-init the model in the loop, or you never call optimizer.step().
2. Sign flipped in bradley_terry_loss. If Spearman trends negative, you swapped chosen and rejected somewhere, either in the loss, or when constructing the dataset. The dataset code is given to you, so check the loss.
3. Learning rate way off. The reference uses Adam at lr=3e-3. At lr=0.1 the model overshoots and Spearman oscillates wildly; at lr=1e-5 it just hasn’t converged in 2000 steps. The starter defaults are tuned; don’t fight them unless you have a reason.

If you’re past hint 5 and still stuck, read solution/reward_model.py. Then re-derive bradley_terry_loss without looking; that’s the one piece of math you should be able to write from memory by the time you’re done.