2️⃣ GRPO with verifiable rewards

@dataclass
class RLVRArgs:
    # Base model & LoRA (applied with `peft` below)
    model_name: str = "Qwen/Qwen3-0.6B-Base"
    lora_rank: int = 16
    lora_alpha: float = 32
    lora_dropout: float = 0.0
    lora_targets: list[str] = field(
        default_factory=lambda: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
    )

    # Task & reward (the functions you wrote in section 1). `task(rng, difficulty)` -> (prompt, answer)
    task: Callable = mult_byhand
    reward_fn: Callable = int_reward     # (completions of one problem, its answer) -> rewards (G,)

    # Dynamic curriculum (section 2): keep each question near `target_success` so the group reward has variance
    target_success: float = 0.5
    initial_difficulty: float = 1.0      # D_0; the controller finds the frontier from here
    difficulty_lr: float = 1.0           # alpha: D += alpha * (batch solve rate - target_success). 0 = fixed difficulty
    min_difficulty: float = 0.0
    max_difficulty: float = len(DIFFICULTY_LADDER) - 1

    # Rollouts: each step samples `n_prompts` problems and `num_return_sequences` completions per problem (the GRPO
    # group, size G in the maths) for `n_prompts * num_return_sequences` total
    n_prompts: int = 16
    num_return_sequences: int = 8
    max_new_tokens: int = 400       # enough for 4x4 working; 350 or less truncates answers and learns less reliably
    temperature: float = 1.0
    top_k: int = 50               
    top_p: float = 0.95
    sampler: str = "auto"         # "vllm" (in-process vLLM, single weight copy; needs vllm + a big card), "fast" (utils.FastSampler:
                                  # CUDA graphs + no-repeat GQA attention), "hf" (model.generate). "auto" = vllm if importable, else fast on CUDA, else hf

    # Optimization
    lr: float = 1e-4
    log_samples: int = 1          # rollouts printed per step (0 = none)
    micro_batch: int = 4          # sequences per forward pass when computing logprobs (memory, not maths; peak ~10 GiB, 8 needs ~17)

    # Duration & reproducibility
    n_steps: int = 20             # training steps, each one rollout batch + one gradient step (~6 min on an A40 with FastSampler)
    seed: int = 0

# Provided: load a small base model + a LoRA adapter
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig, get_peft_model
from torchinfo import summary

MODEL_NAME = "Qwen/Qwen3-0.6B-Base"

cfg = RLVRArgs()

tokenizer = AutoTokenizer.from_pretrained(cfg.model_name)
tokenizer.padding_side = "left"   # batched generation: pad on the left so all prompts end at the same position

# Rollout sampler. With vLLM installed we get samples ~2.5x faster (vLLM's weights and the HF model's are one copy).
# vLLM is optional and not in the ARENA requirements: `pip install vllm==0.30.0`, ideally in a separate environment,
# since it pins its own torch and transformers. It's only imported if it's installed and we use it here.
USE_VLLM = cfg.sampler == "vllm" or (cfg.sampler == "auto" and device.type == "cuda" and utils.vllm_available())
if USE_VLLM:
    vllm_engine = utils.VLLMEngine(cfg.model_name, tokenizer, seed=cfg.seed, max_lora_rank=cfg.lora_rank,
                                        temperature=cfg.temperature, top_k=cfg.top_k, top_p=cfg.top_p)
    _base = vllm_engine.make_hf_base()
    print("Using vLLM")
else:
    vllm_engine = None
    _base = AutoModelForCausalLM.from_pretrained(cfg.model_name, dtype=t.bfloat16, attn_implementation="sdpa").to(device)
    print("no vLLM, falling back to FastSampler (CUDA) or HF generate")

_base.config.pad_token_id = tokenizer.pad_token_id
rlvr_model = get_peft_model(
    _base,
    LoraConfig(r=cfg.lora_rank, 
                lora_alpha=cfg.lora_alpha, 
                lora_dropout=cfg.lora_dropout,
                target_modules=cfg.lora_targets),
)

inputs = tokenizer("Hello! How are you?", return_tensors="pt").to(device)

# using the LoRA adaptors by default
output = rlvr_model.generate(**inputs, max_new_tokens=30, pad_token_id=tokenizer.pad_token_id)
print(tokenizer.batch_decode(output))
# based model with LoRA adaptors disabled
with rlvr_model.disable_adapter():
    output = rlvr_model.generate(**inputs, max_new_tokens=30, pad_token_id=tokenizer.pad_token_id)
    print(tokenizer.batch_decode(output))

>>> print(rlvr_model.print_trainable_parameters())
trainable params: 10,092,544 || all params: 606,142,464 || trainable%: 1.6650

>>> print(_base.model.layers[0].self_attn.q_proj)
lora.Linear(
  (base_layer): Linear(in_features=1024, out_features=2048, bias=False)
  (lora_dropout): ModuleDict(
    (default): Identity()
  )
  (lora_A): ModuleDict(
    (default): Linear(in_features=1024, out_features=16, bias=False)
  )
  (lora_B): ModuleDict(
    (default): Linear(in_features=16, out_features=2048, bias=False)
  )
  (lora_embedding_A): ParameterDict()
  (lora_embedding_B): ParameterDict()
  (lora_magnitude_vector): ModuleDict()
)

def drgrpo_advantages(returns: Float[Tensor, "n_prompts G"]) -> Float[Tensor, "n_prompts G"]:
    """Subtract each group's mean return (no std normalisation)."""
    return returns - returns.mean(dim=-1, keepdim=True)

def drgrpo_objective(
    new_logprobs: Float[Tensor, "batch gen_len"],
    old_logprobs: Float[Tensor, "batch gen_len"],
    advantages: Float[Tensor, " batch"],
    mask: Float[Tensor, "batch gen_len"],
    max_gen_len: int,
    clip: float = 0.2,
) -> Float[Tensor, ""]:
    """Dr. GRPO loss to minimise, summed over the completions in the batch.

    Args:
        new_logprobs: Logprobs under the current policy.
        old_logprobs: Logprobs under the sampling policy.
        advantages: One advantage per completion.
        mask: 1.0 on tokens that count, 0.0 on padding.
        max_gen_len: Generation budget, for the constant normaliser.
        clip: Ratio clipping range.
    """
    raise NotImplementedError()


tests.test_drgrpo_objective(drgrpo_objective)

def clamp(x : float, lo : float, hi : float) -> float:
    return min(max(x, lo), hi)

class Curriculum:
    """Poses problems at difficulty D, scores the answers, and moves D towards the model's frontier."""

    def __init__(
        self,
        task: Callable,
        reward_fn: Callable,
        initial: float = 1.0,
        lr: float = 1.0,
        target: float = 0.5,
        d_min: float = 0.0,
        d_max: float = len(DIFFICULTY_LADDER) - 1,
        seed: int = 0,
    ):
        self.task, self.reward_fn = task, reward_fn
        self.difficulty, self.lr, self.target, self.d_min, self.d_max = initial, lr, target, d_min, d_max
        self.rng = random.Random(seed)                 # for the task generator
        self.np_rng = np.random.default_rng(seed)      # for stochastic rounding

    def sample_levels(self, n: int) -> Int[np.ndarray, " n"]:
        """`n` rungs by stochastic rounding of D: D = 2.3 gives rung 2 70% of the time and rung 3 30%."""
        lo = np.floor(self.difficulty)
        return (lo + (self.np_rng.random(n) < self.difficulty - lo)).astype(int)

    def sample(self, n: int) -> list[str]:
        """Prompts for `n` problems at the current difficulty. Their rungs and answers are kept for scoring."""
        self.levels = self.sample_levels(n).tolist()
        prompts, self.answers = [], []
        for level in self.levels:
            prompt, answer = self.task(self.rng, level)
            prompts.append(prompt)
            self.answers.append(answer)
        self.prompts = prompts
        return prompts

    def score(self, groups: list[list[str]]) -> Float[Tensor, "n_prompts G"]:
        """Reward of each completion: row p scores problem p's group of completions against its answer."""
        return t.stack([self.reward_fn(group, answer) for group, answer in zip(groups, self.answers)])

    def solve_rate(self, groups: list[list[str]]) -> float:
        """Fraction of completions whose boxed answer is exactly right."""
        solved = [extract_boxed_int(c) == answer for group, answer in zip(groups, self.answers) for c in group]
        return sum(solved) / len(solved)

    def update(self, solve_rate: float) -> None:
        """Move D towards the frontier: D <- clamp(D + lr * (solve_rate - target), d_min, d_max)."""
        raise NotImplementedError()


tests.test_curriculum_update(Curriculum)

def generation_logprobs(
    model: PeftModel,
    input_ids: Int[Tensor, "batch seq"],
    prompt_mask: Int[Tensor, "batch prefix_len"],
    mask: Float[Tensor, "batch gen_len"],
) -> Float[Tensor, "batch gen_len"]:
    """Logprob of each generated token under `model` (0.0 where `mask` is 0)."""
    prefix_len = prompt_mask.shape[1]
    # Only the prompt's left padding needs masking: completion padding comes after every real token,
    # so causal attention never reaches it, and `mask` discards the outputs at those positions.
    attention_mask = t.cat([prompt_mask, t.ones_like(input_ids[:, prefix_len:])], dim=1)
    logits = model(input_ids=input_ids, attention_mask=attention_mask).logits.float()   # fp32 for the softmax over 152k tokens
    logprobs = logits[:, prefix_len - 1:].log_softmax(-1)                     # same as `get_logprobs` in RLHF
    return eindex(logprobs, input_ids[:, prefix_len - 1:], "b s [b s+1]") * mask   # position t predicts token t + 1

def completion_mask(gen: Int[Tensor, "batch seq"], prefix_len: int, eos_id: int) -> Float[Tensor, "batch gen_len"]:
    """1.0 on generated tokens up to and including the first EOS, else 0.0."""
    # Not `gen != pad`: Qwen's pad token IS its EOS, so that would drop the EOS and never train the model to stop
    comp = gen[:, prefix_len:]
    is_eos = (comp == eos_id).long()
    eos_before = is_eos.cumsum(1) - is_eos                       # EOS tokens strictly before each position
    return (eos_before == 0).float()


tests.test_completion_mask(completion_mask)


@dataclass
class RolloutBuffer:
    """The rollouts of one rollout phase. Rows are completions, `G` consecutive rows per prompt."""
    gen: Int[Tensor, "batch seq"]              # prompt + completion token ids
    prompt_mask: Int[Tensor, "batch prefix_len"]   # 0 on the prompt's left padding
    mask: Float[Tensor, "batch gen_len"]       # 1 on completion tokens that count for the loss (incl. the EOS)
    advantages: Float[Tensor, " batch"]        # one per completion: no time dependence

    @classmethod
    def from_rollouts(
        cls,
        gen: Int[Tensor, "n_prompts G seq"],
        prompt_mask: Int[Tensor, "n_prompts prefix_len"],
        advantages: Float[Tensor, "n_prompts G"],
        eos_id: int,
    ) -> "RolloutBuffer":
        """Package a rollout as one row per completion: (prompt p, completion g) becomes row p * G + g."""
        n_prompts, G, _ = gen.shape
        rows = gen.flatten(0, 1)                                                  # (n_prompts * G, seq)
        prompt_mask = prompt_mask[:, None].expand(n_prompts, G, -1).flatten(0, 1)   # each prompt's mask, once per completion
        return cls(rows, prompt_mask, completion_mask(rows, prompt_mask.shape[1], eos_id), advantages.flatten())

    @property
    def prefix_len(self) -> int:
        return self.prompt_mask.shape[1]

    def __getitem__(self, rows: slice | Tensor) -> "RolloutBuffer":
        """Select rows."""
        return RolloutBuffer(self.gen[rows], self.prompt_mask[rows], self.mask[rows], self.advantages[rows])

    def trim(self, gen_len: int) -> "RolloutBuffer":
        """Keep the first `gen_len` completion tokens (the rest must be padding)."""
        return RolloutBuffer(self.gen[:, :self.prefix_len + gen_len], self.prompt_mask, self.mask[:, :gen_len], self.advantages)

    def microbatches(self, size: int) -> list["RolloutBuffer"]:
        """Drop zero-advantage rows (they contribute nothing to the gradient) and split the rest into microbatches."""
        buffer = self[self.advantages != 0]
        chunks = utils.length_sorted_chunks(buffer.mask.sum(1).long(), size)
        return [buffer[rows].trim(gen_len) for rows, gen_len in chunks]

for micro_batch in buffer.microbatches(self.args.micro_batch):
    # do stuff ...

class RLVRTrainer:
    def __init__(
        self,
        model: PeftModel,
        tokenizer: PreTrainedTokenizerBase,
        args: RLVRArgs,
        engine: utils.VLLMEngine | None = None,
    ):
        self.model, self.tokenizer, self.args = model, tokenizer, args
        t.manual_seed(args.seed)
        self.params = [p for p in model.parameters() if p.requires_grad]   # the LoRA adapters
        self.optimizer = t.optim.AdamW(self.params, lr=args.lr, fused=True)
        self.sampler: utils.Sampler = utils.make_sampler(model, tokenizer, args.sampler, engine, args.temperature, args.top_k, args.top_p)

        self.curriculum = Curriculum(args.task, args.reward_fn, initial=args.initial_difficulty, lr=args.difficulty_lr, target=args.target_success,
                                     d_min=args.min_difficulty, d_max=args.max_difficulty, seed=args.seed)
        self.history = []                                       # (difficulty, solve rate) per step
        self.latest = utils.LatestText()                        # the newest rollout, overwriting the last one

    def log_rollout(
        self, groups: list[list[str]], rewards: Float[Tensor, "n_prompts G"], solve_rate: float, buffer: RolloutBuffer
    ) -> None:
        """Record this step's difficulty and solve rate, set the metrics for the live plot, and print a few rollouts."""
        levels = self.curriculum.levels
        self.history.append((self.curriculum.difficulty, solve_rate))
        self.metrics = dict(solve=solve_rate, reward=rewards.mean().item(), D=self.curriculum.difficulty,
                            signal=(rewards.std(1) > 0).float().mean().item(),   # groups whose rewards differ: they train
                            length=buffer.mask.sum(1).mean().item())        # mean completion length in tokens
        if self.args.log_samples:
            rungs = ", ".join(f"{levels.count(L)} x {DIFFICULTY_LADDER[L][0]}" for L in sorted(set(levels)))
            lines = [f"--- step {len(self.history)}: D = {self.metrics['D']:.2f} ({rungs}), solve rate {solve_rate:.2f}"]
            for p in range(min(self.args.log_samples, len(groups))):  # the first completion of the first few problems
                completion, answer = groups[p][0], self.curriculum.answers[p]
                lines.append(utils.rollout_summary(self.curriculum.prompts[p], completion, answer,
                                                   extract_boxed_int(completion), rewards[p, 0]))
            self.latest.update("\n".join(lines))

    @t.no_grad()
    def rollout_phase(self) -> RolloutBuffer:
        """Sample problems, generate completions, score them, and move the curriculum."""
        raise NotImplementedError()

    def compute_loss(self, mb: RolloutBuffer) -> Float[Tensor, ""]:
        """This microbatch's share of the loss averaged over every completion in the buffer."""
        logprobs = generation_logprobs(self.model, mb.gen, mb.prompt_mask, mb.mask)
        # One gradient step per rollout, so the sampling policy is the current one: ratio = 1, grad = grad logprobs
        loss = drgrpo_objective(logprobs, logprobs.detach(), mb.advantages, mb.mask, self.args.max_new_tokens)
        loss = loss / (self.args.n_prompts * self.args.num_return_sequences)
        return loss

    def learning_phase(self, buffer: RolloutBuffer) -> float:
        """Perform one gradient step over the entire buffer, accumulating gradients over each microbatch. Returns the
        buffer's loss."""
        raise NotImplementedError()

    def train(self) -> list[tuple[float, float]]:
        """Alternate rollout and learning phases for `args.n_steps`, with live training curves. Returns
        (difficulty, solve rate) per step."""
        plot = LivePlot(range(self.args.n_steps), "solve reward signal | D", "length | loss", desc="RLVR", unit="steps")
        plot["solve"].set_ylim(0, 1)
        plot["D"].set_ylabel("difficulty D")
        plot["length"].set_ylabel("mean completion tokens")
        plot["loss"].set_ylabel("Dr. GRPO loss")
        for _ in plot:
            buffer = self.rollout_phase()
            loss = self.learning_phase(buffer)
            plot.log(**self.metrics, loss=loss)                 # also shown in the progress bar's postfix
        plot.finish()
        return self.history

def evaluate(
    model: PeftModel,
    tokenizer: PreTrainedTokenizerBase,
    levels: tuple[int, ...] = (0, 2, 3, 5),
    n: int = 64,
    max_new_tokens: int = 500,
    seed: int = 4444,
    engine: utils.VLLMEngine | None = None,
) -> dict[str, float]:
    """Held-out pass@1 on each rung, on a fixed problem set."""
    sampler = utils.make_sampler(model, tokenizer, RLVRArgs().sampler, engine)
    out = {}
    for L in levels:
        rng = random.Random(seed + L)
        prompts, answers = zip(*[mult_byhand(rng, L) for _ in range(n)])
        enc = tokenizer(list(prompts), return_tensors="pt", padding=True).to(device)
        t.manual_seed(seed)
        gen = sampler.generate(**enc, max_new_tokens=max_new_tokens)[:, 0]   # one completion per prompt
        completions = tokenizer.batch_decode(gen[:, enc.input_ids.shape[1]:], skip_special_tokens=True)
        out[DIFFICULTY_LADDER[L][0]] = sum(extract_boxed_int(c) == a for c, a in zip(completions, answers)) / n
    print("  ".join(f"{k}: {v:.3f}" for k, v in out.items()), flush=True)
    return out

print("=== base model performance ===")
before = evaluate(rlvr_model, tokenizer, engine=vllm_engine)   # 3x2 / 3x3 / 4x3 / 4x4, n=64 each, 500-token budget
RLVRTrainer(rlvr_model, tokenizer, RLVRArgs(n_steps=40), engine=vllm_engine).train()
print("=== trained model performance ===")
after = evaluate(rlvr_model, tokenizer, engine=vllm_engine)

cd chapter2_rl/exercises/part42_rlvr/playground
python app.py

adapter_dir = section_dir / "adapters" / "rlvr-mult"
rlvr_model.save_pretrained(adapter_dir)
print(f"Saved the LoRA adapter to {adapter_dir}")

a, b = 4127, 386
prompt = f"Compute {a} * {b}. Reason step by step, then give the final answer as \\boxed{{N}}.\n"
sampler = utils.HFSampler(rlvr_model, tokenizer)        # streams through the HF model, with the current LoRA
print(prompt)

print("=== base model ===")
with rlvr_model.disable_adapter():
    base = sampler.generate_stream(prompt, max_new_tokens=500)

print("\n=== after RLVR ===")
trained = sampler.generate_stream(prompt, max_new_tokens=500)

print(f"\nanswer {a * b}: base boxed {extract_boxed_int(base)}, trained boxed {extract_boxed_int(trained)}")

import matplotlib.pyplot as plt


def pass_at_k(n: int, c: int, k: int) -> float:
    """Unbiased pass@k from `n` samples of which `c` are correct (Chen et al., 2021): 1 - C(n - c, k) / C(n, k)."""
    if n - c < k:
        return 1.0
    return 1.0 - float(np.prod(1.0 - k / np.arange(n - c + 1, n + 1)))


@t.no_grad()
def count_correct(
    model: PeftModel,
    tokenizer: PreTrainedTokenizerBase,
    level: int,
    base: bool,
    n_problems: int = 32,
    n_samples: int = 256,
    max_new_tokens: int = 500,
    seed: int = 4444,
    engine: utils.VLLMEngine | None = None,
) -> np.ndarray:
    """Sample `n_samples` completions for each of `n_problems` problems at `level`, from the base model (LoRA off) if
    `base` else the trained one. Returns the number of correct completions per problem."""
    rng = random.Random(seed + level)
    problems = [mult_byhand(rng, level) for _ in range(n_problems)]
    sampler = utils.make_sampler(model, tokenizer, engine=engine)
    counts = []
    for prompt, answer in tqdm(problems, desc=f"{DIFFICULTY_LADDER[level][0]} {'base' if base else 'RLVR'}"):
        enc = tokenizer([prompt], return_tensors="pt").to(device)
        if base and engine is not None:                 # vLLM holds its own copy of the LoRA, so switch it off there
            gen = engine.generate(**enc, max_new_tokens=max_new_tokens, num_return_sequences=n_samples, use_lora=False)
        else:
            with model.disable_adapter() if base else contextlib.nullcontext():
                gen = sampler.generate(**enc, max_new_tokens=max_new_tokens, num_return_sequences=n_samples)[0]
        completions = tokenizer.batch_decode(gen[:, enc.input_ids.shape[1]:], skip_special_tokens=True)
        counts.append(sum(extract_boxed_int(c) == answer for c in completions))
    return np.array(counts)


def plot_pass_at_k(counts: dict[tuple[str, str], np.ndarray], n_samples: int):
    """One panel per level: mean pass@k over problems (with a standard-error band) for each model."""
    levels = list(dict.fromkeys(level for level, _ in counts))
    ks = [2**i for i in range(int(np.log2(n_samples)) + 1)]
    fig, axes = plt.subplots(1, len(levels), figsize=(6 * len(levels), 4.5), squeeze=False)
    for ax, level in zip(axes[0], levels):
        curves = {}
        for (lv, name), c in counts.items():
            if lv != level:
                continue
            M = np.array([[pass_at_k(n_samples, ci, k) for k in ks] for ci in c])
            curves[name] = mean = M.mean(0)
            se = M.std(0) / np.sqrt(len(c))
            ax.plot(ks, mean, "-o", ms=4, label=f"{name} (pass@1 {mean[0]:.2f}, pass@{ks[-1]} {mean[-1]:.2f})")
            ax.fill_between(ks, mean - se, mean + se, alpha=0.2)
        caught = [k for k, b, r in zip(ks, curves["base"], curves["RLVR"]) if b >= r]
        if caught:
            ax.axvline(caught[0], color="gray", ls="--", lw=1, label=f"base catches up at k = {caught[0]}")
        ax.set_xscale("log", base=2)
        ax.set_xticks(ks, [str(k) for k in ks])
        ax.set(xlabel="k", ylabel="pass@k", ylim=(0, 1.02), title=f"{level}: {len(counts[level, 'base'])} problems, {n_samples} samples each")
        ax.grid(alpha=0.3)
        ax.legend(loc="lower right", fontsize=9)
    fig.tight_layout()
    return fig

if SLOW:
    n_samples = 256
    counts = {}
    for level in (3, 5):                                # 4x3 and 4x4
        for name, base in (("base", True), ("RLVR", False)):
            counts[DIFFICULTY_LADDER[level][0], name] = count_correct(
                rlvr_model, tokenizer, level, base, n_samples=n_samples, engine=vllm_engine
            )
    plot_pass_at_k(counts, n_samples)
    plt.show()