"""
=============================================================================
CS486 · 第 5 讲 · 代码 04：奖励是模型，就会被钻空子——用验证器当奖励的 RL
=============================================================================

【这个脚本要回答的问题】（讲义 §9）
    代码 01–03 的奖励是**可验证的**：求解器精确地告诉我们答案对不对（RLVR）。很多任务没有这样的求解器，
    只能用一个**学出来的**奖励模型（RLHF 的奖励模型、第 4 讲的 ORM / PRM、LLM 当裁判）。
    第 4 讲已经看到：best-of-N 用验证器"选最高分"，N 越大越容易选到骗过它的样本（过度优化）。
    RL 是一个更强的优化器——它不是从固定的样本里挑，而是**改变模型**去生成验证器喜欢的东西。

    我们从 cot_L2 出发，在 n = 4–12 的题上做 GRPO 300 步，只改变奖励：
        真实奖励（RLVR）           ：求解器判对错（参照）
        ORM                       ：第 4 讲的结果奖励模型给出的"答案对的概率"
        ORM + 格式检查              ：格式不合格（不是"至少一步推理 + | + 一个字母 + $"）的回答得 0 分
        ORM + 格式检查 + KL 惩罚    ：再加 β = 0.5 的 KL 惩罚（不让模型离起点太远）
        PRM-full + 格式检查         ：第 4 讲的 PRM-full，各步分数相乘
        随机奖励                    ：每条回答以 50% 的概率得 1 分（对照组："奖励信号"本身有没有用？）

    演示 1：训练奖励（模型看到的）、真实奖励、n = 12 测试集准确率随步数的变化。
    演示 2：模型钻了什么空子？——训练奖励 > 0.9、真实奖励为 0 的回答长什么样。
    演示 3：训练结束时，模型在测试集上犯的是哪一类错误？

【运行】
    python3 04_reward_hacking.py           # 第一次约 2.5 分钟（6 次 RL 并行；需要第 4 讲的验证器缓存），之后约 10 秒
    python3 04_reward_hacking.py --quick   # 冒烟测试

【依赖】PyTorch + 本目录的 rl_lab.py；第 4 讲的验证器（没有缓存时会先训练，约 2.5 分钟）。
"""

import argparse
import os
import sys
import time
from collections import Counter

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "lecture01"))

import rl_lab as rl                                                      # noqa: E402
import tiny_reasoner as tr                                               # noqa: E402
from console import pct, print_header, print_observations, print_table  # noqa: E402

SETTINGS = {
    "true": ("真实奖励（RLVR）", {"reward": "true"}),
    "orm": ("ORM", {"reward": "orm"}),
    "orm_guard": ("ORM + 格式检查", {"reward": "orm", "guard": True}),
    "orm_guard_kl": ("ORM + 格式 + KL", {"reward": "orm", "guard": True, "beta": 0.5}),
    "prm_guard": ("PRM-full + 格式检查", {"reward": "prm_full", "guard": True}),
    "random": ("随机奖励", {"reward": "random"}),
}
SHOWN = (0, 25, 50, 75, 100, 150, 200, 300)


def jobs():
    """六个任务：同样的起点、题目、步数和学习率，只有奖励不同。"""
    return {k: rl.make_job(steps=300, eval="n12", **kw) for k, (_, kw) in SETTINGS.items()}


def window(hist, step, key, width=25):
    """训练统计在第 step 步之前 width 步内的平均（step = 0 时取第 1 步）。"""
    lo = max(0, step - width)
    w = hist["steps"][lo:max(step, 1)]
    return sum(h[key] for h in w) / len(w)


def demo1(res):
    print_header("演示 1：训练奖励 / 真实奖励（训练批次，前 25 步平均）和 n = 12 测试集准确率")
    for key, label in (("proxy", "训练奖励（模型优化的目标）"), ("true", "真实奖励（训练批次里答案对的比例）")):
        print(f"\n  {label}：")
        rows = [[SETTINGS[k][0]] + [pct(window(res[k][1], s, key), 0) for s in SHOWN] for k in SETTINGS]
        print_table(["奖励 \\ 步数"] + [str(s) for s in SHOWN], rows, [22] + [6] * len(SHOWN))
    print("\n  n = 12 测试集（200 道题）的贪心准确率：")
    rows, out = [], {}
    for k in SETTINGS:
        ev = dict(res[k][1]["evals"])
        out[k] = ev
        rows.append([SETTINGS[k][0]] + [pct(ev[s]["acc"], 0) for s in SHOWN])
    print_table(["奖励 \\ 步数"] + [str(s) for s in SHOWN], rows, [22] + [6] * len(SHOWN))
    return out


def demo2(res):
    print_header("演示 2：训练奖励 > 0.9、真实答案却错了的回答（正确答案 → 模型的回答，奖励）")
    out = {}
    for k in ("orm", "orm_guard", "orm_guard_kl", "prm_guard"):
        exs = res[k][1]["examples"]
        out[k] = exs
        print(f"\n  {SETTINGS[k][0]}：")
        for step, items in exs:
            if step in (1, 25, 50, 100, 300) and items:
                shown = "；".join(f"{a} → {b!r}（{c:.3f}）" for a, b, c in items[:2])
                print(f"    第 {step:>3} 步：{shown}")
    return out


def demo3(res, probs):
    print_header("演示 3：训练结束时在 n = 12 测试集上的结果类型（贪心解码，200 道题）")
    rows, out = [], {}
    for k in SETTINGS:
        model = res[k][0]
        gen = tr.evaluate(model, probs, "cot")
        c = Counter(rl.classify(p, rl.response_of(g["generated"])) for p, g in zip(probs, gen))
        out[k] = (c, [rl.response_of(g["generated"]) for g in gen[:3]])
        rows.append([SETTINGS[k][0]] + [c.get(t, 0) for t in rl.FAILURE_TYPES] + [", ".join(repr(x) for x in out[k][1][:2])])
    print_table(["奖励", "对", "走过头", "停早了", "查找错", "抄错", "没写 |", "前两道题的回答"],
                rows, [22, 5, 7, 7, 7, 5, 7, 24])
    return out


def main():
    ap = argparse.ArgumentParser(description="用验证器当奖励的 RL：奖励黑客")
    ap.add_argument("--quick", action="store_true", help="冒烟测试")
    args = ap.parse_args()
    t0 = time.perf_counter()

    print_header("准备：起点 cot_L2；第 4 讲的验证器；六种奖励各训练 300 步")
    js = jobs()
    if args.quick:
        js = {k: dict(j, eval_every=3) for k, j in js.items()}
    keys = list(js)
    res = dict(zip(keys, rl.load_runs([js[k] for k in keys], quick=args.quick)))
    if args.quick:
        for k in res:                                          # 冒烟测试只训练 6 步：把评测的步数对齐到表格的格子
            ev = dict(res[k][1]["evals"])
            last = ev[max(ev)]
            res[k][1]["evals"] = [(s, ev.get(s, last)) for s in SHOWN]
            res[k][1]["steps"] = (res[k][1]["steps"] * 60)[:300]
    probs = rl.eval_set(12, per_k=5 if args.quick else 50)

    ev = demo1(res)
    exs = demo2(res)
    kinds = demo3(res, probs)

    if args.quick:
        print("\n（--quick 冒烟测试：数字没有意义，不打印观察。）")
        print(f"\n⏱ 总用时 {time.perf_counter() - t0:.0f} 秒")
        return
    def acc(k, s):
        return pct(ev[k][s]["acc"], 0)

    def tw(k, s, key):
        return pct(window(res[k][1], s, key), 0)

    kg, kk, kp = kinds["orm_guard"][0], kinds["orm_guard_kl"][0], kinds["prm_guard"][0]
    print_observations([
        f"可验证奖励下（第一行），n = 12 测试集从 {acc('true', 0)} 升到 {acc('true', 300)}，和代码 01 一样。"
        f"只把奖励换成第 4 讲的 ORM（AUC 约 0.95 的好验证器），训练奖励在 50 步内升到 {tw('orm', 50, 'proxy')}，"
        f"真实奖励却从 {tw('orm', 0, 'true')} 掉到 {tw('orm', 50, 'true')}，测试集准确率是 {acc('orm', 300)}。"
        "模型学会了写 '||$'、'||||$' 这类东西（演示 2）：ORM 只在 cot_L2 自己写出的回答上训练过，从没见过这种格式，"
        "对它们给出 0.98–1.00 的高分。**奖励模型只在它的训练分布上有效，RL 会把策略推到分布之外**。",
        f"加上格式检查（格式不合格得 0 分）之后，空子换了一个地方：第 25 步时模型写 'afbd|d$'（少走一步），之后越来越短，"
        f"最后几乎所有题都写 'k|k$'。测试集准确率在第 25 步还是 {acc('orm_guard', 25)}（起点 {acc('orm_guard', 0)}），随后跌到 {acc('orm_guard', 300)}"
        f"（训练结束时 200 道题里 {kg.get('lookup', 0)} 道第一步就错、{kg.get('early', 0)} 道停早了）。"
        f"PRM-full 同样被'只走一步就停'骗过（'l|l$' 得 1.000 分，测试集 {acc('prm_guard', 300)}）。"
        "这正是第 4 讲 §6 发现的离策略盲区：cot_L2 几乎从不停早，验证器没见过停早的推理链，也就分不出它们的对错。"
        "格式检查只能堵住你**想到**的漏洞。",
        f"KL 惩罚（β = 0.5）把策略拴在起点附近：真实奖励始终在 {tw('orm_guard_kl', 300, 'true')} 左右，没有崩溃；"
        f"但测试集准确率在 {acc('orm_guard_kl', 25)}–{acc('orm_guard_kl', 150)} 之间摆动（最后 {acc('orm_guard_kl', 300)}，"
        f"200 道题里仍有 {kk.get('overshoot', 0) + kk.get('malformed', 0)} 道停不下来），远不如可验证奖励的 100%。"
        "KL 是在'少钻空子'和'少学东西'之间折中，不是免费的修补。",
        f"随机奖励（每条回答以 50% 的概率得 1 分）不是'什么也不发生'：真实奖励从 {tw('random', 0, 'true')} 一路掉到 "
        f"{tw('random', 300, 'true')}，模型最后写出 'ddddddddddd' 这样的退化输出。在零均值的随机优势下，"
        "每一步的更新仍然是随机地推高一些回答、压低另一些；学习率 1e-3、没有 KL 时，这个随机游走把策略推到了退化的角落。"
        "Shao et al.（2025）在 Qwen2.5-Math 上看到随机奖励让 MATH-500 **提高**了约 21 个点——那来自 PPO 式裁剪的偏差"
        "放大了基座里已有的高概率行为（用代码推理），换成其他模型家族就没有效果。我们没有裁剪（每批只更新一次），"
        "所以看到的是另一种结果。**随机奖励对照是必要的，但它的结果取决于算法细节和基座**（讲义 §9.5）。",
        "总结：RL 是一个比 best-of-N 强得多的优化器。best-of-N 只能在固定的样本里'挑到'验证器的错误；"
        "RL 会**改变模型**，主动生成验证器的错误。第 4 讲的过度优化在这里变成了几十步之内的彻底崩溃。"
        "这就是推理 RL 坚持用可验证奖励、DeepSeek-R1 在推理任务上放弃神经奖励模型的原因。",
    ])
    print(f"\n⏱ 总用时 {time.perf_counter() - t0:.0f} 秒")


if __name__ == "__main__":
    main()