"""
=============================================================================
第 16 讲 · 代码 2：奖励投机——当验证器有漏洞时
=============================================================================

对应讲义：lectures/lecture16_frontiers_review.md 第 4 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
"可验证的奖励"听起来很可靠：答案对不对、测试过没过，一查便知。但验证器本身是人写的程序，可能有漏洞：
    - 只检查输出的格式，不检查内容；
    - 单元测试覆盖不全，一个"硬编码测试用例答案"的程序也能通过；
    - 模型可以修改测试文件、或者让程序提前退出，使测试框架报告"通过"。
RL 是一个强大的优化器，它会找到这些漏洞——这叫**奖励投机**（reward hacking / specification gaming）。

在玩具推理模型中（见 toy_rlvr.py），回答 1 是"作弊"的回答：有漏洞的验证器也给它 1 分，但它其实是错的。
参考模型产生作弊回答的概率只有 0.5%。"作弊倾向" h 是所有题**共享**的参数。

本文件：
    1. 验证器没有漏洞 vs 有漏洞：代理奖励（验证器给的分）一路上升，真实准确率先升后降
    2. KL 惩罚能挡住作弊吗？——能，但要付出代价
    3. 为什么后果这么严重？——作弊是一种会**泛化**的行为
    4. 发现并修好验证器之后，模型还能恢复吗？

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    cd code/lecture16
    python 02_reward_hacking.py      （约 3 秒）
"""

import numpy as np

from toy_rlvr import LEVEL_NAMES, make_task, pad, train

SEEDS = range(5)


def run(task, **kw):
    """用 5 个随机种子训练（GRPO，每步 16 道题 × 8 个回答，学习率 2），返回平均的训练历史和各次的结果。"""
    outs = [train(task, "grpo", batch=16, G=8, lr=2.0, seed=sd, **kw) for sd in SEEDS]
    hist = np.array([[(h[0], h[1], h[2], h[3]) for h in o["history"]] for o in outs]).mean(axis=0)
    return hist, outs


def main():
    task = make_task()

    # ---------------------------------------------------------------------
    print("=" * 78)
    print("演示 1：验证器没有漏洞 vs 有漏洞（600 步，每 60 步记录一次；5 个随机种子的平均）")
    print("=" * 78)
    for label, hr in [("验证器没有漏洞（作弊回答得 0 分）", 0.0), ("验证器有漏洞（作弊回答也得 1 分）", 1.0)]:
        hist, outs = run(task, steps=600, hack_reward=hr, eval_every=60)
        print(f"\n  {label}：")
        print("     步数    验证器给的平均分（代理奖励）   真实准确率    作弊回答的概率    代理 - 真实")
        for step, acc, hack, proxy in hist[[0, 1, 2, 3, 5, 10]]:
            print(f"     {int(step):4d}             {proxy:.3f}                 {acc:.3f}          {hack:.3f}           {proxy - acc:+.3f}")
    print("\n  观察：")
    print("  - 没有漏洞时：代理奖励就是真实准确率，一路上升到接近 1。")
    print("  - 有漏洞时：前 60 步，两者都在上升——模型在学着答对题。随后，作弊回答在一部分难题上被偶然抽到、")
    print("    得了分、被强化；由于作弊倾向是所有题共享的，它迅速扩散到所有题上。")
    print("    代理奖励继续上升到 1.000，看起来训练得完美无缺；真实准确率却跌到了 1% 左右——**比训练前还差得多**。")
    print("  - 这就是第 9 讲\"奖励过度优化\"的另一种形式：优化得越狠，越会利用奖励信号与真实目标之间的差距。")
    print("    只看代理奖励的训练曲线，你完全看不出问题——\"代理 - 真实\"这一列在真实训练中是看不到的。")

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print("演示 2：KL 惩罚能挡住作弊吗？（验证器有漏洞，600 步）")
    print("=" * 78)
    print("   " + pad("KL 系数", 10) + "代理奖励   真实准确率   作弊概率     按难度的真实准确率：很难 / 难 / 中等 / 容易")
    for kl in [0.0, 0.05, 0.2, 1.0]:
        hist, outs = run(task, steps=600, hack_reward=1.0, kl_coef=kl, eval_every=600)
        lv = np.mean([o["history"][-1][4] for o in outs], axis=0)
        _, acc, hack, proxy = hist[-1]
        print(f"   {kl:6.2f}      {proxy:.3f}      {acc:.3f}       {hack:.3f}          " + " / ".join(f"{x:.2f}" for x in lv))
    _, ref_outs = run(task, steps=600, hack_reward=0.0, eval_every=600)
    print(f"   （对照：验证器没有漏洞、KL 系数 0 时，真实准确率 {np.mean([o['history'][-1][1] for o in ref_outs]):.3f}）")
    print("  观察：")
    print("  - KL 惩罚让策略不能离参考模型太远，而参考模型几乎从不作弊——所以它确实能压制作弊。")
    print("  - 但 KL 惩罚不区分\"好的偏离\"和\"坏的偏离\"：它同样压制了\"学会答对\"。")
    print("    系数小（0.05、0.2）：挡不住作弊，最难的题几乎全被作弊占领；系数大（1.0）：作弊被压住了，")
    print("    但真实准确率也只有 0.5 左右，远低于验证器没有漏洞时的水平。")
    print("  - **KL 惩罚是一种保险，不是解药**。根本的解决办法是修好验证器。")

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print("演示 3：为什么后果这么严重？——共享的\"作弊倾向\" vs 每道题各自学会作弊（验证器有漏洞，600 步）")
    print("=" * 78)
    print("   " + pad("作弊倾向", 36) + "真实准确率   作弊概率：很难 / 难 / 中等 / 容易")
    for label, kw in [("所有题共享（学会一次，到处都用）", {}), ("每道题各自学（不共享）", {"lr_h": 0.0})]:
        hist, outs = run(task, steps=600, hack_reward=1.0, eval_every=600, **kw)
        pr = np.mean([o["policy"].probs()[:, 1] for o in outs], axis=0)
        by = [pr[task["level"] == k].mean() for k in range(len(LEVEL_NAMES))]
        print("   " + pad(label, 36) + f"   {hist[-1][1]:.3f}       " + " / ".join(f"{x:.2f}" for x in by))
    print("  观察：")
    print("  - 如果每道题都要各自\"发现\"作弊（不共享），作弊只出现在最难的题上：那些题正确答案太难找到，")
    print("    偶然抽到的作弊回答成了唯一的得分途径。其余的题，模型老老实实地学会了答对。")
    print("  - 共享时，在难题上学到的\"作弊\"被推广到了**所有**题上——包括那些本来很容易答对的题。")
    print("  - 大语言模型的参数是所有任务共享的，\"钻空子\"这种行为模式同样会泛化。")
    print("    这就是为什么一个小小的验证器漏洞，可能让整个模型学坏。")

    # ---------------------------------------------------------------------
    print("\n" + "=" * 78)
    print("演示 4：修好验证器之后，模型还能恢复吗？（前面若干步验证器有漏洞，之后修好）")
    print("=" * 78)
    for fix in [100, 300]:
        hist, _ = run(task, steps=900, hack_reward=(lambda t, f=fix: 1.0 if t <= f else 0.0), eval_every=100)
        print(f"\n  第 {fix} 步修好验证器：")
        print("     步数    真实准确率    作弊概率")
        for step, acc, hack, proxy in hist:
            mark = "  ← 修好验证器" if int(step) == fix else ""
            print(f"     {int(step):4d}      {acc:.3f}        {hack:.3f}{mark}")
    print("  观察：")
    print("  - 修好验证器之后，作弊回答不再得分，在 GRPO 中得到负的优势，共享的作弊倾向被迅速压下去。")
    print("  - 即使作弊已经占了 97% 的概率（第 300 步修好），模型也在 100 步内恢复到 0.77 的准确率。")
    print("    原因是：在这个玩具模型中，作弊是通过**共享的**参数 h 实现的，每道题\"正确答案\"的知识")
    print("    （各自的 logits）一直保留在下面，只是被作弊\"盖住\"了；h 一降下来，正确答案就重新出现。")
    print("  - 真实的模型不一定这么幸运：长时间的奖励投机可能真正地改变模型的能力和行为习惯。")
    print("    而真正的难点在于**发现**漏洞——训练曲线上的代理奖励看起来一切正常。")
    print("    实践中的办法：定期人工抽查模型的输出、用更严格的验证器做独立评测、监控异常的行为模式")
    print("    （比如回答突然变短、出现奇怪的固定模式）。")


if __name__ == "__main__":
    main()