grpo_toy_problem.py

# /// script
# requires-python = ">=3.12"
# dependencies = []
# ///
"""Watch GRPO actually learn — a toy loop, no GPU, no real model.

This is the "see a number go up" move from fast.ai (Jeremy Howard): before you
leave the trainer behind, watch the loop turn and watch the model get better.

The slider widget in Lesson 5 showed ONE frozen step — four scores, one baseline,
the advantage arrows. It can't show the thing the lesson is actually about: a model
migrating toward good answers because of those advantages, step after step. This
script turns the loop and shows the consequence.

THE TOY (deliberately the simplest thing that can't be subtly wrong):
- The "model" (policy) is a tiny table of probabilities over THREE ANSWER BUCKETS
  for one ISAF extraction prompt. A real model could emit many literal strings /
  JSON records; this toy compresses them into correct, partial, and wrong.
- The three buckets have fixed scores from score(): correct (1.0), partial (0.5),
  wrong (0.0).
- Each step is exactly the widget, in a loop:
    1. sample a GROUP of attempts from the current probabilities
    2. score each attempt
    3. baseline = the group average
    4. advantage = each score - baseline   (the widget, exactly)
    5. nudge each sampled answer's probability UP if its advantage was positive,
       DOWN if negative
  Repeat, and P(correct bucket) climbs.

THE REVEAL (the deepest property of GRPO, now over TIME): run it again with
EASY_MODE = True, which makes all three answers score the same. Now every attempt
ties the baseline, every advantage is 0, the probabilities never move, and the
reward curve is FLAT. No spread within the group -> no learning. That's the
widget's "make them all equal" button, but watched across many steps instead of one.

Run it:
    uv run labs/0002-grpo-toy-loop.py

Then change something and rerun (the fast.ai "transfer" rep): add a fourth
bucket, change GROUP_SIZE to 2 or 20, change STEPS, drop LR, change DIFFICULTY,
or flip EASY_MODE. Predict what the curve will do BEFORE you run, then check.
"""

import random

# ---- The toy setup (this is the "data" — look at it before running) ----------
# Three answer buckets for ONE prompt, each with a fixed reward from score().
# Index 0 is the correct extraction; 1 is a partially right JSON record; 2 is wrong.
ANSWERS = ["correct", "partial", "wrong"]
SCORES = [1.0, 0.5, 0.0]  # what score() returns for each candidate

EASY_MODE = False  # flip to True -> all answers score the same
DIFFICULTY = "learnable"  # easy | learnable | hard
GROUP_SIZE = 20  # attempts per step (the "group" in GRPO)
STEPS = 40  # how many times the loop turns
LR = 0.15  # how hard each advantage nudges a probability
SEED = 0  # fixed so the run is reproducible

START_PROBS = {
    "easy": [0.90, 0.08, 0.02],  # mostly solved already: high reward, less signal
    "learnable": [1 / 3, 1 / 3, 1 / 3],  # useful middle: mixed groups
    "hard": [0.06, 0.24, 0.70],  # correct bucket is rare: learning can be slow
}


def scores_now() -> list[float]:
    """The reward each candidate earns. EASY_MODE flattens them all to 1.0 so the
    group has no spread — which is the whole point of the second run."""
    return [1.0 for _ in SCORES] if EASY_MODE else list(SCORES)


def start_probs() -> list[float]:
    """Initial policy for the chosen difficulty setting."""
    if DIFFICULTY not in START_PROBS:
        raise ValueError(f"DIFFICULTY must be one of {sorted(START_PROBS)}")
    return list(START_PROBS[DIFFICULTY])


def step(
    probs: list[float], scores: list[float], rng: random.Random
) -> tuple[list[float], float]:
    """One GRPO step on the toy policy. Returns (new probabilities, this step's
    average reward). This IS the widget: sample a group, score it, subtract the
    group average to get advantages, push toward the winners and away from the losers."""
    # 1. sample a group of attempts from the current policy
    idxs = rng.choices(range(len(probs)), weights=probs, k=GROUP_SIZE)
    # 2. score each attempt
    rewards = [scores[i] for i in idxs]
    # 3. baseline = the group average  ("good compared to what?")
    baseline = sum(rewards) / len(rewards)
    # 4 + 5. advantage = reward - baseline; nudge each sampled answer by it
    new = probs[:]
    for i, r in zip(idxs, rewards):
        new[i] += LR * (r - baseline)  # >0 -> push up; <0 -> push down
    # keep them valid probabilities: floor tiny, then renormalize to sum 1
    new = [max(p, 1e-4) for p in new]
    total = sum(new)
    new = [p / total for p in new]
    return new, baseline


def bar(p: float, width: int = 24) -> str:
    """A little text bar so you can SEE the probability move."""
    filled = round(p * width)
    return "#" * filled + "." * (width - filled)


def main() -> None:
    rng = random.Random(SEED)
    probs = start_probs()
    scores = scores_now()

    mode = (
        "EASY_MODE (all answers score 1.0)" if EASY_MODE else "normal (1.0 / 0.5 / 0.0)"
    )
    print(
        f"\nToy GRPO loop — {mode}, difficulty={DIFFICULTY}, "
        f"group of {GROUP_SIZE}, {STEPS} steps\n"
    )
    print(f"{'step':>4}  {'avg reward':>10}   P(correct) answer")
    print("-" * 62)
    # step 0: the starting point, before any learning
    print(f"{0:>4}  {'-':>10}   {probs[0]:.2f}  {bar(probs[0])}")

    for s in range(1, STEPS + 1):
        probs, avg = step(probs, scores, rng)
        # print every few steps so the climb is visible without a wall of text
        if s % 5 == 0 or s == 1:
            print(f"{s:>4}  {avg:>10.2f}   {probs[0]:.2f}  {bar(probs[0])}")

    print("-" * 62)
    if EASY_MODE:
        print(
            "\nP(correct) barely moved. Every attempt scored the same, so every\n"
            "advantage was 0, so nothing got pushed anywhere. No spread in the\n"
            "group -> no learning signal -> wasted compute. That's the widget's\n"
            "'make them all equal' reveal, watched across 25 steps.\n"
        )
    else:
        print(
            f"\nP(correct) climbed from {START_PROBS[DIFFICULTY][0]:.2f} to {probs[0]:.2f}. Nothing trained the\n"
            "model toward the right answer except 'it scored above the group average,\n"
            f"so push toward it' — repeated {STEPS} times. That's GRPO. Now change\n"
            "DIFFICULTY or flip EASY_MODE = True at the top and rerun: predict the curve first.\n"
        )


if __name__ == "__main__":
    main()