Debug a GRPO training script and defend the RL math

"""A tiny GRPO trainer in plain Python.

The policy is a table of logits: theta[prompt][token]. A completion is a short
sequence of tokens sampled from softmax(theta[prompt]); its reward is the
fraction of tokens equal to the prompt's target token. Sequences in a batch
have different lengths and are padded to the same length; mask is 1 on real
tokens and 0 on padding.
"""

import math
import random

VOCAB = 4
PAD = 0


def log_softmax(logits):
    top = max(logits)
    z = math.log(sum(math.exp(x - top) for x in logits)) + top
    return [x - z for x in logits]


def rollout(theta, prompts, group_size, max_len, rng):
    """Sample group_size completions for every prompt.
    Returns (prompt_of_row, tokens, mask, rewards), rows grouped by prompt."""
    rows, tokens, mask, rewards = [], [], [], []
    for p in prompts:
        probs = [math.exp(x) for x in log_softmax(theta[p])]
        for _ in range(group_size):
            length = rng.randint(1, max_len)
            seq = rng.choices(range(VOCAB), weights=probs, k=length)
            rows.append(p)
            tokens.append(seq + [PAD] * (max_len - length))
            mask.append([1] * length + [0] * (max_len - length))
            rewards.append(sum(tok == target(p) for tok in seq) / length)
    return rows, tokens, mask, rewards


def target(prompt):
    return prompt % VOCAB


def token_logps(theta, rows, tokens):
    """log pi(token) at every position, padding included (B x T)."""
    tables = {p: log_softmax(theta[p]) for p in set(rows)}
    return [[tables[p][tok] for tok in seq] for p, seq in zip(rows, tokens)]


def group_advantages(rewards, group_size, eps=1e-6):
    """One advantage per completion: its reward minus the mean reward of its
    group (the completions for the same prompt), divided by that group's std."""
    mean = sum(rewards) / len(rewards)
    std = (sum((r - mean) ** 2 for r in rewards) / len(rewards)) ** 0.5
    return [(r - mean) / std for r in rewards]


def grpo_loss(new_logps, old_logps, ref_logps, advantages, mask, clip_eps=0.2, beta=0.04):
    """The clipped GRPO objective with a KL penalty to the reference policy,
    averaged over the real tokens in the batch, returned as a loss to MINIMIZE.
    The KL uses the k3 estimator: exp(ref - new) - (ref - new) - 1."""
    total, count = 0.0, 0
    for i, adv in enumerate(advantages):
        for t in range(len(mask[i])):
            ratio = new_logps[i][t] / old_logps[i][t]
            unclipped = ratio * adv
            clipped = min(max(ratio, 1 - clip_eps), 1 + clip_eps) * adv
            surrogate = max(unclipped, clipped)
            log_r = ref_logps[i][t] - new_logps[i][t]
            kl = math.exp(log_r) - log_r - 1
            total += surrogate + beta * kl
            count += 1
    return total / count


def clip_fraction(new_logps, old_logps, mask, clip_eps=0.2):
    """Share of real tokens whose probability ratio is outside the clip range."""
    outside = real = 0
    for new_row, old_row, m_row in zip(new_logps, old_logps, mask):
        for new, old, m in zip(new_row, old_row, m_row):
            if m:
                real += 1
                outside += abs(math.exp(new - old) - 1) > clip_eps
    return outside / real


def numeric_grad(f, theta, h=1e-4):
    grad = [[0.0] * VOCAB for _ in theta]
    for p in range(len(theta)):
        for v in range(VOCAB):
            theta[p][v] += h
            up = f(theta)
            theta[p][v] -= 2 * h
            down = f(theta)
            theta[p][v] += h
            grad[p][v] = (up - down) / (2 * h)
    return grad


def train(steps=60, prompts=(0, 1, 2, 3), group_size=8, max_len=5, epochs=2, lr=2.0, seed=0, **loss_kwargs):
    """Returns (theta, history). history[s] holds step s's mean rollout reward
    and the largest clip fraction seen across its epochs."""
    rng = random.Random(seed)
    theta = [[0.0] * VOCAB for _ in range(len(prompts))]
    ref = [row[:] for row in theta]
    history = []
    for _ in range(steps):
        rows, tokens, mask, rewards = rollout(theta, prompts, group_size, max_len, rng)
        stats = {"reward": sum(rewards) / len(rewards), "clip_frac": 0.0}
        history.append(stats)
        adv = group_advantages(rewards, group_size)
        ref_logps = token_logps(ref, rows, tokens)
        for _ in range(epochs):
            old_logps = token_logps(theta, rows, tokens)
            stats["clip_frac"] = max(stats["clip_frac"], clip_fraction(token_logps(theta, rows, tokens), old_logps, mask))

            def loss(th):
                return grpo_loss(token_logps(th, rows, tokens), old_logps, ref_logps, adv, mask, **loss_kwargs)

            grad = numeric_grad(loss, theta)
            for p in range(len(theta)):
                for v in range(VOCAB):
                    theta[p][v] -= lr * grad[p][v]
    return theta, history