Exercise 11: DPO on a toy preference dataset

"""DPO on a toy preference dataset — exercise starter.

Fill in the TODOs, then from the repo root:

    pip install -r exercises/requirements.txt
    pytest exercises/11-dpo/

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

The toy task:
  - N_PROMPTS prompts (just integer indices, no text).
  - N_ANSWERS possible answers per prompt (also integer indices).
  - A fixed "true" reward table r_true[prompt, answer] you don't get to see
    during training. You only see preferences derived from it.
  - A preference dataset of (prompt, chosen, rejected) triples, generated by
    sampling two answers uniformly per prompt and labeling them with
    Bradley-Terry noise on the true reward gap.

You train a Policy by minimising the DPO loss against a frozen reference
policy. The policy is the same shape as the GRPO exercise: nn.Embedding,
one logit row per prompt.

DPO update per step:
  1. Sample a batch of preference triples.
  2. Forward the policy and the (frozen) reference on the batch prompts.
  3. Compute the DPO loss:
       -log sigmoid( beta * [(log pi(y_w) - log pi_ref(y_w))
                             - (log pi(y_l) - log pi_ref(y_l))] )
     averaged over the batch.
  4. Backward, step.

There is no reward model, no PPO, no rollout. The "reward signal" is the
preference labels in the dataset.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F

# ── task constants ────────────────────────────────────────────────────────────

N_PROMPTS = 8
N_ANSWERS = 10
TRUE_REWARD_SEED = 0   # fixed seed for the "true" reward table
DATASET_SEED = 1       # fixed seed for the preference dataset
BETA = 0.1             # DPO KL coefficient
BT_TEMPERATURE = 5.0   # Bradley-Terry temperature for labelling noise


# ── true reward (fixed, known) ────────────────────────────────────────────────

def make_true_reward(n_prompts: int = N_PROMPTS, n_answers: int = N_ANSWERS) -> torch.Tensor:
    """Generate a fixed (prompt, answer) -> reward table.

    Deterministic — same seed gives the same table on every call. Returns a
    tensor of shape (n_prompts, n_answers) with values roughly in [-2, 2].

    This is given to you. The training pipeline never reads it directly; it
    only sees preference labels derived from it.
    """
    g = torch.Generator()
    g.manual_seed(TRUE_REWARD_SEED)
    return torch.randn(n_prompts, n_answers, generator=g)


# ── preference dataset ────────────────────────────────────────────────────────

def make_preference_dataset(
    n_pairs: int,
    true_reward: torch.Tensor,
    seed: int = DATASET_SEED,
    bt_temperature: float = BT_TEMPERATURE,
) -> list:
    """Build a list of (prompt_idx, chosen_answer, rejected_answer) triples.

    For each pair:
      - Pick a prompt uniformly.
      - Sample two distinct answers uniformly.
      - Label using Bradley-Terry: P(a beats b) = sigmoid(temp * (r_a - r_b)).
        So the labels are mostly correct but noisy near r_a == r_b.

    This is also given to you. The student does not implement it.
    """
    n_prompts, n_answers = true_reward.shape
    g = torch.Generator()
    g.manual_seed(seed)

    data = []
    for _ in range(n_pairs):
        p = int(torch.randint(n_prompts, (1,), generator=g).item())
        perm = torch.randperm(n_answers, generator=g)
        a, b = int(perm[0]), int(perm[1])
        r_a = float(true_reward[p, a])
        r_b = float(true_reward[p, b])
        p_a_wins = torch.sigmoid(torch.tensor(bt_temperature * (r_a - r_b))).item()
        u = torch.rand((), generator=g).item()
        if u < p_a_wins:
            chosen, rejected = a, b
        else:
            chosen, rejected = b, a
        data.append((p, chosen, rejected))
    return data


# ── policy ───────────────────────────────────────────────────────────────────

class Policy(nn.Module):
    """Per-prompt categorical policy, parameterized as an embedding.

    Same shape as the GRPO exercise's policy:
        nn.Embedding(n_prompts, n_answers)

    `forward(prompt_idx)` should accept either:
      - a single int                       -> returns logits of shape (n_answers,)
      - a LongTensor of shape (B,)         -> returns logits of shape (B, n_answers)

    Args:
        n_prompts: number of distinct prompts (rows in the embedding).
        n_answers: number of possible answers per prompt (columns).
    """

    def __init__(self, n_prompts: int = N_PROMPTS, n_answers: int = N_ANSWERS):
        super().__init__()
        # TODO: self.embed = nn.Embedding(n_prompts, n_answers)
        raise NotImplementedError("Policy.__init__")

    def forward(self, prompt_idx) -> torch.Tensor:
        # TODO:
        #   if isinstance(prompt_idx, int):
        #       idx = torch.tensor(prompt_idx, dtype=torch.long)
        #   else:
        #       idx = prompt_idx.long()
        #   return self.embed(idx)
        raise NotImplementedError("Policy.forward")


# ── DPO loss ──────────────────────────────────────────────────────────────────

def dpo_loss(
    policy_logits: torch.Tensor,
    ref_logits: torch.Tensor,
    chosen_answer: torch.Tensor,
    rejected_answer: torch.Tensor,
    beta: float = BETA,
) -> torch.Tensor:
    """The DPO loss for a batch of preference pairs.

    Formula:
        loss = -log sigmoid(
                  beta * [(log pi(y_w|x) - log pi_ref(y_w|x))
                        - (log pi(y_l|x) - log pi_ref(y_l|x))]
               )
    averaged over the batch.

    Args:
        policy_logits:   (B, n_answers) — from the policy being trained.
        ref_logits:      (B, n_answers) — from the frozen reference policy.
                         You should treat ref_logits as a constant; the test
                         pipeline already calls the reference under no_grad,
                         but you don't need to assume anything about that.
        chosen_answer:   LongTensor (B,) — index of the preferred answer.
        rejected_answer: LongTensor (B,) — index of the dispreferred answer.
        beta: KL coefficient (DPO's single tuning knob).

    Returns:
        Scalar tensor with gradient.

    Steps:
      1. log_softmax both logit tensors along the answer dim.
      2. Gather the four log-probs: pi(y_w), pi(y_l), pi_ref(y_w), pi_ref(y_l).
         (.gather is the cleanest way.)
      3. Compute:  logits = beta * ((pi_w - ref_w) - (pi_l - ref_l))
      4. Return:   -F.logsigmoid(logits).mean()

    Sanity check: if policy_logits == ref_logits exactly, every bracket is zero,
    sigmoid(0) = 0.5, and the loss is log(2) ≈ 0.6931. The first test checks
    exactly this.
    """
    # TODO
    raise NotImplementedError("dpo_loss")


# ── evaluation helper ─────────────────────────────────────────────────────────

def greedy_mean_true_reward(policy: Policy, true_reward: torch.Tensor) -> float:
    """Mean true reward of the greedy (argmax) policy across all prompts.

    Given to you; useful for tracking learning during train_dpo.
    """
    n_prompts = true_reward.shape[0]
    with torch.no_grad():
        idx = torch.arange(n_prompts)
        logits = policy(idx)                       # (n_prompts, n_answers)
        greedy = logits.argmax(dim=-1)             # (n_prompts,)
        rewards = true_reward.gather(-1, greedy.unsqueeze(-1)).squeeze(-1)
    return rewards.mean().item()


# ── training loop ─────────────────────────────────────────────────────────────

def train_dpo(
    preference_dataset: list,
    ref_policy: Policy,
    true_reward: torch.Tensor,
    beta: float = BETA,
    n_steps: int = 2000,
    lr: float = 0.05,
    batch_size: int = 32,
    seed: int = 0,
) -> tuple:
    """Train a Policy with DPO against a frozen reference.

    Each step:
      1. Sample `batch_size` indices into the preference dataset uniformly.
      2. Forward the policy on the batched prompts (gradient on).
      3. Forward the reference policy on the same prompts (no gradient).
      4. Compute dpo_loss; backprop; step.
      5. Log the loss and the greedy mean true reward (for the integration test).

    Returns:
        losses:        list of length n_steps — per-step loss values.
        mean_rewards:  list of length n_steps — per-step greedy mean true reward.
        policy:        the trained Policy (returned so the test can poke at it).
    """
    torch.manual_seed(seed)
    rng = torch.Generator()
    rng.manual_seed(seed)

    # Fresh policy initialised from the reference so they start identical.
    n_prompts, n_answers = true_reward.shape
    policy = Policy(n_prompts, n_answers)
    policy.load_state_dict(ref_policy.state_dict())

    # Freeze the reference; no gradients should flow through it.
    for p in ref_policy.parameters():
        p.requires_grad_(False)
    ref_policy.eval()

    optimizer = torch.optim.Adam(policy.parameters(), lr=lr)

    # Pre-tensorize for easy batching.
    prompts = torch.tensor([t[0] for t in preference_dataset], dtype=torch.long)
    chosens = torch.tensor([t[1] for t in preference_dataset], dtype=torch.long)
    rejecteds = torch.tensor([t[2] for t in preference_dataset], dtype=torch.long)
    n = prompts.shape[0]

    losses, mean_rewards = [], []
    for _ in range(n_steps):
        # TODO: implement the DPO update.
        #
        #   idx = torch.randint(n, (batch_size,), generator=rng)
        #   p_batch = prompts[idx]
        #   w_batch = chosens[idx]
        #   l_batch = rejecteds[idx]
        #
        #   policy_logits = policy(p_batch)
        #   with torch.no_grad():
        #       ref_logits = ref_policy(p_batch)
        #
        #   loss = dpo_loss(policy_logits, ref_logits, w_batch, l_batch, beta)
        #
        #   optimizer.zero_grad()
        #   loss.backward()
        #   optimizer.step()
        #
        #   losses.append(loss.item())
        #   mean_rewards.append(greedy_mean_true_reward(policy, true_reward))
        raise NotImplementedError("train_dpo: the update step")

    return losses, mean_rewards, policy


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

if __name__ == "__main__":
    true_reward = make_true_reward()
    dataset = make_preference_dataset(n_pairs=512, true_reward=true_reward)

    ref = Policy()
    losses, rewards, _ = train_dpo(
        dataset, ref, true_reward, beta=BETA, n_steps=2000, lr=0.05, seed=0
    )

    baseline = true_reward.mean().item()
    print(
        f"steps: {len(losses)}  "
        f"start loss: {losses[0]:.4f}  "
        f"end loss: {losses[-1]:.4f}  "
        f"end greedy reward: {rewards[-1]:.3f}  "
        f"uniform baseline: {baseline:.3f}"
    )

Two lines, same shape as the GRPO policy. In __init__:

super().__init__()
self.embed = nn.Embedding(n_prompts, n_answers)

In forward, you need to handle both a single int (single-prompt call from a test) and a LongTensor (batched call from train_dpo):

if isinstance(prompt_idx, int):
    idx = torch.tensor(prompt_idx, dtype=torch.long)
else:
    idx = prompt_idx.long()
return self.embed(idx)

nn.Embedding(N, M) is a lookup table: row i is a length-M vector. Indexing with a 0-d tensor gives a (M,) output; indexing with a (B,) tensor gives (B, M). That’s exactly what we want.

Two logit tensors of shape (B, n_answers) come in. Get log-probs first:

policy_log_probs = F.log_softmax(policy_logits, dim=-1)
ref_log_probs    = F.log_softmax(ref_logits, dim=-1)

Now you need four scalars per batch element: log pi(y_w), log pi(y_l), log pi_ref(y_w), log pi_ref(y_l). chosen_answer is (B,). gather is the cleanest way:

pi_w  = policy_log_probs.gather(-1, chosen_answer.unsqueeze(-1)).squeeze(-1)
pi_l  = policy_log_probs.gather(-1, rejected_answer.unsqueeze(-1)).squeeze(-1)
ref_w = ref_log_probs.gather(-1, chosen_answer.unsqueeze(-1)).squeeze(-1)
ref_l = ref_log_probs.gather(-1, rejected_answer.unsqueeze(-1)).squeeze(-1)

The unsqueeze makes the indices (B, 1) to match the source’s last dim; the squeeze strips the singleton back off.

Now the loss is one line of arithmetic and one line of reduction:

logits = beta * ((pi_w - ref_w) - (pi_l - ref_l))
return -F.logsigmoid(logits).mean()

Sanity check: if policy_logits == ref_logits, then pi_w == ref_w and pi_l == ref_l, so logits is the zero vector. logsigmoid(0) = log(0.5) = -log(2). The loss is then -(-log(2)) = log(2) ≈ 0.6931. The test_dpo_loss_starts_at_log2 test checks this exactly.

Use F.logsigmoid rather than torch.log(torch.sigmoid(...)). It’s much more numerically stable for large negative inputs. With beta=0.1 and a tiny task this doesn’t matter for correctness, but it’s the right habit.

The single most common bug in DPO is getting the sign backwards inside the bracket. The bracket must be (chosen_logratio - rejected_logratio), not the reverse. Concretely: you want the loss to decrease as pi(y_w) rises and pi(y_l) falls. That means the argument to sigmoid should go up, and -log sigmoid(...) should go down. The test_dpo_loss_pushes_toward_chosen test verifies this directly: it computes the log-margin between chosen and rejected before and after one gradient step.

If that test fails, you almost certainly have the bracket flipped.

Most of the loop is bookkeeping (already done for you). The update step itself is:

idx = torch.randint(n, (batch_size,), generator=rng)
p_batch = prompts[idx]
w_batch = chosens[idx]
l_batch = rejecteds[idx]

policy_logits = policy(p_batch)
with torch.no_grad():
    ref_logits = ref_policy(p_batch)

loss = dpo_loss(policy_logits, ref_logits, w_batch, l_batch, beta)

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

losses.append(loss.item())
mean_rewards.append(greedy_mean_true_reward(policy, true_reward))

Three things to be careful about:

1. The reference call is wrapped in torch.no_grad(). The reference is already frozen at the top of the function (requires_grad_(False)), but the no_grad makes it explicit and skips graph construction.
2. The policy and reference were initialized to the same weights (top of train_dpo). That’s why the first loss is log(2).
3. mean_rewards uses the greedy policy (argmax over answers), not a sampled one. The greedy reward is the right thing to check the policy converges toward the high-reward answers, not just that it shifts mass a bit.

This is the conceptual hint. Standard RLHF needs three networks: SFT, reward model, policy. DPO collapses two of them: there is no separate reward model, because the DPO derivation shows that the optimal KL-constrained policy already encodes the reward, so reward differences can be written as log-ratios of policy probabilities. The preference labels (chosen, rejected) are the only “reward signal” you need.

In this exercise, the true reward exists (we use it to generate labels and to evaluate the trained policy), but the training pipeline never reads it. The DPO loss takes only logits and preference indices.

If you’re past hint 6 and still stuck, read solution/dpo.py. Then re-derive it without looking.