"""
=============================================================================
第 16 讲 · 代码 1：可验证奖励的强化学习与 GRPO——优势估计到底改变了什么？
=============================================================================

对应讲义：lectures/lecture16_frontiers_review.md 第 1–3 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
训练推理模型（比如 DeepSeek-R1）的核心算法之一是 GRPO：对同一道题抽 G 个回答，
用"这组回答的平均奖励"当基线，再除以组内的标准差：
        A_i = (r_i - mean(r)) / std(r)
它不需要 PPO 那样的价值网络（critic）。这到底带来了什么？本文件在一个"玩具推理模型"上
（见 toy_rlvr.py：200 道题、每题 16 个候选回答、答对得 1 分）逐一拆开来看：

    1. 参考模型：不同难度的题，答对的概率差别很大；pass@k（抽 k 次至少对一次）
    2. 梯度估计的偏差与方差：无基线、留一法基线（RLOO）、组内均值、GRPO
       ——以及奖励带有常数偏移（比如"格式分"）时，基线为什么重要
    3. 训练：同样的学习率下 GRPO 更快；但把学习率调好之后，差别就小多了
    4. 组的大小 G：在总样本数固定时，G 大还是 G 小好？

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    cd code/lecture16
    python 01_grpo.py      （约 2 秒）
"""

import numpy as np

from toy_rlvr import LEVEL_NAMES, Policy, advantages, make_task, pad, pass_at_k, train

np.set_printoptions(precision=3, suppress=True)


def gradient_stats(task, pol, prompt, method, G, n, rng, offset=0.0):
    """
    在固定的策略下，对同一道题重复 n 次"抽 G 个回答 → 估计梯度"，统计梯度估计的均值与离散程度。

    【返回】(正确答案 logit 方向上梯度的均值, 它的标准差, 整个梯度向量的总方差)
    【真实梯度】这道题的期望奖励 J = π(正确)，对正确答案 logit 的导数是 π(1-π)（softmax 的导数）。
        一个"无偏"的估计量，均值应该等于它。
    """
    pr = pol.probs([prompt])[0]
    y = pol.sample(np.full(n, prompt), G, rng)
    r = (y == 0).astype(float) + offset
    A = advantages(r, method, baseline=np.zeros(n))
    onehot = np.zeros((n, G, task["M"]))
    np.put_along_axis(onehot, y[:, :, None], 1.0, axis=2)
    g = (A[:, :, None] * (onehot - pr[None, None, :])).mean(axis=1)     # 每次重复得到的梯度估计
    return g[:, 0].mean(), g[:, 0].std(), g.var(axis=0).sum()


def main():
    task = make_task()
    ref = Policy(task)
    pr = ref.probs()

    # ---------------------------------------------------------------------
    print("=" * 78)
    print("演示 1：参考模型（RL 之前）——200 道题，4 种难度")
    print("=" * 78)
    print("   难度          题数     pass@1    pass@8    pass@64")
    for k, name in enumerate(LEVEL_NAMES):
        p = pr[task["level"] == k, 0]
        print(f"   {pad(name, 12)} {len(p):4d}     {pass_at_k(p, 1).mean():.3f}     {pass_at_k(p, 8).mean():.3f}     {pass_at_k(p, 64).mean():.3f}")
    print(f"   {pad('全部', 12)} {task['P']:4d}     {pr[:, 0].mean():.3f}     {pass_at_k(pr[:, 0], 8).mean():.3f}     {pass_at_k(pr[:, 0], 64).mean():.3f}")
    print("  观察：")
    print("  - pass@k = 1 - (1 - p)^k：多抽几次，总有一次能对。参考模型在\"很难\"的题上 pass@1 只有 1%，")
    print("    但 pass@64 有近一半——正确答案\"在模型的能力范围内\"，只是概率太低。")
    print("  - RL 要做的，就是把 pass@k 中\"偶尔能对\"的那部分，变成 pass@1 中\"稳定地对\"。")

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print("演示 2：梯度估计的偏差与方差（参考模型，G = 8，每种情况重复 2 万次）")
    print("=" * 78)
    rng = np.random.default_rng(0)
    methods = [("无基线（REINFORCE）", "none"), ("留一法基线（RLOO）", "rloo"),
               ("组内均值（不除标准差）", "group"), ("GRPO（除以组内标准差）", "grpo")]
    for offset in [0.0, 1.0]:
        label = "奖励：答对 1，答错 0" if offset == 0 else "奖励：答对 2，答错 1（每个回答都有 1 分\"格式分\"）"
        p = int(np.where(task["level"] == 2)[0][0])
        pi0 = pr[p, 0]
        print(f"\n  {label}；一道中等难度的题（π(正确) = {pi0:.2f}，真实梯度 π(1-π) = {pi0 * (1 - pi0):.4f}）：")
        print("     方法                         梯度均值    标准差    整个梯度向量的总方差")
        for name, m in methods:
            mean, std, var = gradient_stats(task, ref, p, m, 8, 20000, rng, offset)
            print("     " + pad(name, 26) + f"   {mean:7.4f}    {std:7.4f}        {var:7.4f}")
    print("\n  GRPO 的梯度均值与真实梯度之比（各难度各取一道题）：")
    for k, name in enumerate(LEVEL_NAMES):
        p = int(np.where(task["level"] == k)[0][0])
        pi0 = pr[p, 0]
        mean, _, _ = gradient_stats(task, ref, p, "grpo", 8, 20000, rng)
        print(f"     {pad(name, 12)}  π(正确) = {pi0:.2f}   GRPO / 真实梯度 = {mean / (pi0 * (1 - pi0)):.2f}")
    print("  观察：")
    print("  - 无基线和 RLOO 都是**无偏**的（均值等于真实梯度）。\"组内均值\"把自己也算进了基线，")
    print("    梯度被整体缩小了 (G-1)/G = 7/8，有一点偏差。")
    print("  - 奖励是 0/1 时，有没有基线，方差差别不大。但当所有回答都有一个共同的\"格式分\"时，")
    print("    无基线的方差一下子大了十倍：它把每个被抽到的回答都往上推，推的方向主要由\"碰巧抽到了谁\"决定。")
    print("    有基线的方法完全不受影响——减去基线后，共同的部分被消掉了（第 6 讲的结论）。")
    print("  - GRPO 除以组内标准差，**不再是**梯度的无偏估计：它把每道题的梯度放大了约 1/std 倍。")
    print("    越难（或越容易）的题，组内标准差越小，被放大得越多——GRPO 隐含地对题目按难度加权。")

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print("演示 3：训练（每步 16 道题 × 8 个回答，300 步；5 个随机种子的平均）")
    print("=" * 78)
    print("   方法                         学习率   格式分   训练后 pass@1     按难度：很难 / 难 / 中等 / 容易")
    rows = [("无基线（REINFORCE）", "none", 2.0, 0.0), ("学到的基线（每题一个 V）", "learned", 2.0, 0.0),
            ("留一法基线（RLOO）", "rloo", 2.0, 0.0), ("GRPO", "grpo", 2.0, 0.0),
            ("无基线（REINFORCE）", "none", 2.0, 5.0), ("学到的基线（每题一个 V）", "learned", 2.0, 5.0),
            ("留一法基线（RLOO）", "rloo", 2.0, 5.0), ("GRPO", "grpo", 2.0, 5.0),
            ("无基线（REINFORCE）", "none", 8.0, 0.0), ("GRPO", "grpo", 4.0, 0.0)]
    for name, m, lr, off in rows:
        outs = [train(task, m, 300, 16, 8, lr, sd, reward_offset=off) for sd in range(5)]
        acc = np.mean([o["history"][-1][1] for o in outs])
        lv = np.mean([o["history"][-1][4] for o in outs], axis=0)
        print("   " + pad(name, 28) + f"  {lr:4.1f}     {off:3.0f}       {acc:.3f}          " + " / ".join(f"{x:.2f}" for x in lv))
    print(f"   （参考模型：pass@1 = {pr[:, 0].mean():.3f}）")
    print("  观察：")
    print("  - 奖励是 0/1、学习率相同时，无基线、学到的基线、RLOO 几乎一样；GRPO 明显更快，尤其是在难题上。")
    print("  - 加上格式分之后，无基线的训练大受影响；学到的基线要花时间追上新的奖励水平，也受了影响；")
    print("    组内基线（RLOO、GRPO）完全不受影响——每组内部互相比较，共同的部分自动抵消。")
    print("  - GRPO 更快的主要原因是**更大的有效步长**（除以标准差），而不是更好的梯度方向：")
    print("    把无基线方法的学习率调大到 8，它的结果和学习率 4 的 GRPO 相当。")
    print("  - 那 GRPO 为什么在大模型中流行？主要原因是**省掉了价值网络**：PPO 需要一个和策略一样大的 critic，")
    print("    训练、显存的开销都翻倍，而且给长文本的每个 token 估计价值本身就很难。")
    print("    组内基线用\"同一道题的其他回答\"代替了 critic——便宜、针对每道题、而且自动适应奖励的尺度。")

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print("演示 4：组的大小 G（总样本数固定：每步 128 个回答；GRPO，学习率 2；300 步，5 个种子）")
    print("=" * 78)
    print("     G    每步题数   训练后 pass@1   按难度：很难 / 难 / 中等 / 容易    \"没有学习信号\"的组的比例")
    for G in [2, 4, 8, 16, 32]:
        outs = [train(task, "grpo", 300, 128 // G, G, 2.0, sd) for sd in range(5)]
        acc = np.mean([o["history"][-1][1] for o in outs])
        lv = np.mean([o["history"][-1][4] for o in outs], axis=0)
        zf = np.mean([o["zero_frac"] for o in outs])
        print(f"   {G:3d}      {128 // G:3d}          {acc:.3f}        " + " / ".join(f"{x:.2f}" for x in lv) + f"             {zf * 100:5.1f}%")
    print("  观察：")
    print("  - \"没有学习信号\"的组：一组回答全对或全错，组内基线下所有优势都是 0，这组样本白白浪费了。")
    print("    G 越小，这种组越多（G = 2 时，两个回答一样的概率很高）；训练后期大部分题都答对了，这个比例还会上升。")
    print("  - 但在这个玩具问题上，总样本数固定时 G 越**小**反而越好：每步能覆盖更多的题、更新更频繁，")
    print("    \"浪费\"的样本被\"更多的题\"弥补了。**没有学习信号的比例高，不等于学得慢**。")
    print("  - 实际训练大模型时常用 G = 8～64：同一道题的多个回答可以共享提示词的计算，并行生成很便宜；")
    print("    而且真实的奖励更吵，太小的组给出的基线太不准。结论依赖于具体的开销结构，不能照搬。")
    print("  - 一个常见的改进（DAPO 等）：动态地丢掉全对或全错的题，把计算留给\"有信息量\"的题。")


if __name__ == "__main__":
    main()
