"""
=============================================================================
CS486 · 第 5 讲 · 代码 03：长思维链是 RL "创造"的，还是"挑出来"的？——一个 R1-Zero 式的实验
=============================================================================

【这个脚本要回答的问题】（讲义 §6、§8）
    DeepSeek-R1-Zero 直接从基座模型开始做 RL（只奖励答案对不对），训练过程中回答越来越长，
    模型"自己学会了"长思维链、反思、检查——论文称之为"顿悟时刻"（aha moment）。
    后来的分析（Liu et al. 2025 "Understanding R1-Zero-Like Training"；Gandhi et al. 2025 "Cognitive Behaviors…"）
    发现：这些行为在基座模型里**本来就有**，只是很少出现；RL 把它们选了出来。基座里没有这些行为的模型，RL 也很难让它们出现。

    我们用 rl_lab.MIXED_BASES 造的几个基座来做这个实验。每个基座的"预训练语料"里只有一部分题目带推理链，
    所以基座采样时"有时想、有时不想"（不想 = 直接写 "|" 和答案）。2 层模型一次前向只能做 1 跳（第 1 讲），
    所以"不想"几乎总是错的、"想"几乎总是对的——RL 只要发现这一点，就会让模型越来越多地"想"。

        noskill ：语料里从没有推理链（没有查找能力，也没有"想"的行为）
        think0  ：学会了查找，但后来的语料里没有推理链——"想"几乎绝迹
        think2  ：2% 的语料带推理链
        think10 ：10% 的语料带推理链

    演示 1：四个基座的行为——多久"想"一次？想的时候对不对？不想的时候对不对？
    演示 2：从每个基座开始做 GRPO（只奖励答案对错，n = 4–8 的题），"想"的比例、准确率、回答长度怎样变化？
    演示 3："顿悟"发生在第几步？——think0 换 3 个种子。
    演示 4：加上长度惩罚（奖励 = 答对 − λ × 回答长度）：模型在哪些题上放弃思考？
    演示 5：长度惩罚 × "除以组内标准差"——一个会让训练崩溃的组合（Dr. GRPO 的批评）。

【运行】
    python3 03_emergent_thinking.py           # 第一次约 3 分钟（4 个基座 + 12 次 RL，并行），之后约 15 秒
    python3 03_emergent_thinking.py --quick   # 冒烟测试

【依赖】PyTorch + 本目录的 rl_lab.py。
"""

import argparse
import os
import sys
import time

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

BASES = ("noskill", "think0", "think2", "think10")
LAMBDAS = (0.0, 0.1, 0.2, 0.3)


def rl_job(base, **kw):
    """从混合基座出发、在 n = 4–8 的题上做 GRPO 的配置（每 25 步评测一次"想不想、对不对"）。"""
    return rl.make_job(init=f"mixed:{base}", n_range=[4, 8], k_max=7, eval="think", **kw)


def all_jobs():
    """本脚本的 12 个 RL 任务。"""
    jobs = {}
    for b in BASES:                                   # 演示 2
        jobs[("base", b)] = rl_job(b, steps=300)
    for s in (1, 2):                                  # 演示 3
        jobs[("seed", s)] = rl_job("think0", steps=300, seed=s)
    for lam in LAMBDAS:                               # 演示 4
        jobs[("lam", lam)] = rl_job("think10", steps=300, len_pen=lam, std_norm=False)
    for std in (True, False):                         # 演示 5
        jobs[("std", std)] = rl_job("noskill", steps=200, len_pen=0.01, std_norm=std)
    return jobs


def first_step(hist, threshold=0.5):
    """训练批次里"写了推理链"的比例第一次超过 threshold 的步数；没有超过时返回 None。"""
    for i, h in enumerate(hist["steps"]):
        if h["think"] > threshold:
            return i + 1
    return None


def demo1(bases, probs, N):
    print_header("演示 1：四个基座的行为（n = 8，k = 1–7；每道题采样 %d 次，温度 1.0）" % N)
    rows, out = [], {}
    for b in BASES:
        prof = rl.think_profile(bases[b], probs, N)
        out[b] = prof
        rows.append([b, pct(prof["think"]), pct(prof["acc_think"]) if prof["think"] > 0 else "-",
                     pct(prof["acc_direct"]), pct(prof["acc"]), f"{prof['len']:.2f}"])
    print_table(["基座", "写推理链的比例", "写了时的准确率", "直接回答的准确率", "总准确率", "回答长度"],
                rows, [10, 16, 16, 18, 10, 10])
    return out


def demo2(res):
    print_header("演示 2：从每个基座开始做 GRPO（奖励 = 答案对不对）——写推理链的比例 / 准确率")
    shown = (0, 25, 50, 75, 100, 150, 200, 300)
    rows = []
    for b in BASES:
        ev = dict(res[("base", b)][1]["evals"])
        rows.append([b] + [f"{pct(ev[s]['think'], 0)} / {pct(ev[s]['acc'], 0)}" for s in shown])
    print_table(["基座 \\ 步数"] + [str(s) for s in shown], rows, [12] + [11] * len(shown))
    print("\n  回答的平均长度（字符数，含 '|'、答案和 '$'）：")
    rows = []
    for b in BASES:
        ev = dict(res[("base", b)][1]["evals"])
        rows.append([b] + [f"{ev[s]['len']:.2f}" for s in shown])
    print_table(["基座 \\ 步数"] + [str(s) for s in shown], rows, [12] + [11] * len(shown))
    model = res[("base", "think2")][0]
    probs = rl.think_set(per_k=1)[1:4]
    outs = [rl.response_of(o["generated"]) for o in tr.evaluate(model, probs, "cot")]
    print(f"\n  think2 训练 300 步后的贪心回答（k = 2、3、4）：{outs}")
    print(f"  正确答案：{[p.trace + '|' + p.answer + '$' for p in probs]}")
    return outs


def demo3(res):
    print_header("演示 3：训练批次里写推理链的比例第一次超过 5% / 50% / 95% 是在第几步？")
    rows, out = [], {}
    for key, label in ((("base", "think0"), "think0，种子 0"), (("seed", 1), "think0，种子 1"),
                       (("seed", 2), "think0，种子 2"), (("base", "think2"), "think2，种子 0"),
                       (("base", "think10"), "think10，种子 0"), (("base", "noskill"), "noskill，种子 0")):
        hist = res[key][1]
        steps = [first_step(hist, th) for th in (0.05, 0.5, 0.95)]
        out[label] = steps
        rows.append([label] + [s if s else "—" for s in steps])
    print_table(["设置", "超过 5%", "超过 50%", "超过 95%"], rows, [18, 10, 10, 10])
    print("  （—：300 步内没有超过；训练批次 = 每步 32 道题 × 8 个样本）")
    return out


def demo4(res):
    print_header("演示 4：长度惩罚 λ（从 think10 出发，不除以标准差；训练 300 步后，每个 k 各 40 道题 × 16 个样本）")
    out = {}
    for lam in LAMBDAS:
        prof = res[("lam", lam)][1]["evals"][-1][1]
        out[lam] = prof
    for key, label in (("think", "写推理链的比例"), ("acc", "准确率"), ("trace", "推理链平均步数（写了的那些）")):
        print(f"\n  {label}：")
        rows = []
        for lam in LAMBDAS:
            bk = out[lam]["by_k"]
            cells = [f"λ = {lam}"]
            for k in range(1, 8):
                cells.append(pct(bk[k][key], 0) if key != "trace" else f"{bk[k][key]:.1f}")
            rows.append(cells)
        print_table(["", *[f"k = {k}" for k in range(1, 8)]], rows, [10] + [7] * 7)
    return out


def demo5(res):
    print_header("演示 5：长度惩罚 0.01 × 是否除以组内标准差（从 noskill 出发：它几乎总是答错）")
    rows, out = [], {}
    for std in (True, False):
        model, hist = res[("std", std)]
        ev = dict(hist["evals"])
        label = "除以标准差（GRPO）" if std else "不除（Dr. GRPO）"
        out[std] = ev
        rows.append([label] + [f"{pct(ev[s]['acc'], 0)} / {ev[s]['len']:.1f}" for s in (0, 50, 100, 150, 200)])
        probs = rl.think_set(per_k=1)[:3]
        outs = tr.evaluate(model, probs, "cot")
        out[("ex", std)] = [rl.response_of(o["generated"]) for o in outs]
    print_table(["设置", "0 步", "50", "100", "150", "200"], rows, [20, 12, 12, 12, 12, 12])
    print("  （每格：准确率 / 回答平均长度）")
    for std in (True, False):
        print(f"  训练后的贪心回答（{'除以标准差' if std else '不除'}）：{out[('ex', std)]}")
    return out


def main():
    ap = argparse.ArgumentParser(description="R1-Zero 式的实验：长思维链从哪里来？")
    ap.add_argument("--quick", action="store_true", help="冒烟测试")
    args = ap.parse_args()
    t0 = time.perf_counter()

    print_header("准备：四个混合基座、12 次 GRPO 训练")
    if not args.quick:
        rl.prepare_bases(BASES)
    bases = {b: rl.mixed_base(b, quick=args.quick) for b in BASES}
    jobs = all_jobs()
    if args.quick:
        jobs = {k: dict(j, eval_every=3) for k, j in jobs.items()}
    keys = list(jobs)
    results = rl.load_runs([jobs[k] for k in keys], quick=args.quick)
    res = dict(zip(keys, results))

    probs = rl.think_set(per_k=4 if args.quick else 40)
    base_prof = demo1(bases, probs, 4 if args.quick else 16)
    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 (0, 25, 50, 75, 100, 150, 200, 300)]
    tails = demo2(res)
    aha = demo3(res)
    lam = demo4(res)
    std = demo5(res)

    if args.quick:
        print("\n（--quick 冒烟测试：数字没有意义，不打印观察。）")
        print(f"\n⏱ 总用时 {time.perf_counter() - t0:.0f} 秒")
        return
    a0, a1, a2 = aha["think0，种子 0"], aha["think0，种子 1"], aha["think0，种子 2"]
    l1, l2, l3 = lam[0.1]["by_k"], lam[0.2]["by_k"], lam[0.3]
    ev_g, ev_n = std[True], std[False]
    print_observations([
        f"四个基座的'能力'和'行为'是分开的（演示 1）：think2 只有 {pct(base_prof['think2']['think'])} 的回答写推理链，"
        f"可一旦写，{pct(base_prof['think2']['acc_think'])} 是对的；直接回答只有 {pct(base_prof['think2']['acc_direct'])}"
        "（2 层模型一次前向只能做 1 跳，第 1 讲）。think0 在采样里一次推理链都没写过，但它学过查找（第一阶段）。",
        "只奖励答案对错的 GRPO 让 think2、think10 在 25 步之内变成'总是写推理链'，准确率从 15–22% 升到 98–99%，"
        "回答长度从约 3 个字符升到 7 个（演示 2）——这就是 R1-Zero'回答越来越长、准确率越来越高'的缩小版。"
        "但这里没有任何新能力被创造出来：RL 只是把基座里**已经有、只是很少出现**的行为挑了出来。",
        f"think0（采样里看不到推理链）也学会了，但要等很久：3 个种子分别在第 {a0[0]} / {a1[0]} / {a2[0]} 步才第一次超过 5%，"
        f"之后只用 {a0[2] - a0[0]} / {a1[2] - a1[0]} / {a2[2] - a2[0]} 步就从 5% 升到 95%（演示 3）。"
        "曲线上看是一个突然的'顿悟'；实际上是一个极少出现的行为偶然被采到、得到正奖励，然后被指数式地放大。"
        "什么时候被采到是随机的（相差近 90 步），所以'顿悟发生在第几步'没有意义，有意义的是'基座里有没有这种行为'。"
        f"noskill（没有查找能力）在 300 步里始终没有学会（准确率一直在 10–14%）：没有可以挑的东西，RL 无从下手。",
        f"一个意外（演示 2 的最后两行）：think2 训练到 300 步时回答长度突然变成约 12 个字符——模型不再写结束符 '$'，"
        f"而是在答案后面一直重复 '|j|j|j…' 直到长度上限（{tails[0]!r}）。奖励只检查 '|' 后面的第一个字母，"
        "所以结束符和后面的内容不受任何约束，在随机的更新中漂移了。**奖励没有检查的东西，RL 不会替你保持**；"
        "真实系统里的格式奖励、超长惩罚（DAPO）就是为此而设。",
        f"长度惩罚首先伤害最难的题（演示 4）：λ = 0.1 时只有 k = 6、7 稍有下降（{pct(l1[6]['acc'], 0)} / {pct(l1[7]['acc'], 0)}）；"
        f"λ = 0.2 时 k ≤ 5 几乎不受影响，k = 6 掉到 {pct(l2[6]['acc'], 0)}（推理链平均只写 {l2[6]['trace']:.1f} 步）、"
        f"k = 7 掉到 {pct(l2[7]['acc'], 0)}（只写 {l2[7]['trace']:.1f} 步就停下）——模型学会了在长题上**提前停止思考**；"
        f"λ = 0.3 时完全放弃思考，所有题都答错（准确率 {pct(l3['acc'])}）。按奖励算，λ = 0.3 时 k = 1 的题写推理链仍然划算"
        "（多 1 个字符只扣 0.3 分，换来约 0.8 的准确率），模型却连它也放弃了：同一个策略在所有题之间共享，'不思考'的倾向会泛化。",
        f"长度惩罚 × 除以组内标准差（演示 5）：从几乎总是答错的 noskill 出发、只加 0.01 的长度惩罚，"
        f"GRPO（除以标准差）在第 100 步塌缩成只输出 '$'（准确率 {pct(ev_g[200]['acc'], 0)}、长度 {ev_g[200]['len']:.1f}）；"
        f"不除以标准差的 Dr. GRPO 没有塌缩（{pct(ev_n[200]['acc'], 0)}、{ev_n[200]['len']:.1f}）。"
        "原因：大多数题组 8 个回答全错，奖励只差 0.01 × 长度；除以一个很小的标准差之后，这点差别被放大成和'对 vs 错'"
        "一样大的优势，'写得更短'压倒了一切。这是 Dr. GRPO 批评的'标准化带来的偏差'的一个极端例子。",
    ])
    print(f"\n⏱ 总用时 {time.perf_counter() - t0:.0f} 秒")


if __name__ == "__main__":
    main()