"""后训练的奖励契约、组内信号与实测记录；默认只读记录，不训练。"""

from argparse import ArgumentParser
from hashlib import sha256
import json
import math
from pathlib import Path
import random

DATA = Path(__file__).with_name("data")
PLAN = DATA / "post-training-plan.json"
RECORD = DATA / "post-training-record.json"


def digest(path):
    return sha256(Path(path).read_bytes()).hexdigest()


def make_plan():
    groups = [(a, b) for a in range(10) for b in range(a, 10)]
    random.Random(20261020).shuffle(groups)
    split = {}
    for name, selected in zip(("train", "dev", "test"),
                              (groups[:35], groups[35:45], groups[45:])):
        split[name] = sorted([[a, b] for a, b in selected]
                             + [[b, a] for a, b in selected if a != b])
    return {
        "version": 1, "date": "2026-10-20", "task": "single_digit_addition",
        "split": split, "seeds": [11, 22, 33],
        "vocab": "0123456789+=~_", "eos": "~", "pad": "_",
        "embedding": 24, "hidden": 48, "threads": 1,
        "sft_steps": 200, "sft_batch": 16, "sft_lr": .003,
        "rl_steps": 240, "prompts_per_step": 8, "group_size": 8,
        "updates_per_rollout": 2, "rl_lr": .0005, "clip": .2,
        "kl_beta": .03, "max_new_tokens": 6, "temperature": 1.,
        "eval_samples": 16, "gradient_clip": 1.,
        "arms": ["strict", "substring"],
        "loss_reduction": "mean over all valid response tokens, including EOS",
        "advantage": "group population std + 1e-8; no dynamic resampling",
        "kl": "exact categorical KL at old-policy sampled histories",
        "selection": "fixed last checkpoint; no tuning or seed selection",
        "rewards": {
            "strict": "terminated, nonempty ASCII digits, integer equals a+b; leading zeros allowed",
            "substring": "decimal str(a+b) occurs anywhere, even without EOS",
        },
    }


def reward(pair, text, terminated, arm):
    if arm == "substring":
        return int(str(sum(pair)) in text)
    if arm != "strict":
        raise ValueError(arm)
    return int(terminated and bool(text)
               and all(c in "0123456789" for c in text)
               and int(text) == sum(pair))


def contrast_probability(p, group=8):
    return 1 - p**group - (1-p)**group


def advantages(rewards):
    mean = sum(rewards) / len(rewards)
    std = math.sqrt(sum((r-mean)**2 for r in rewards) / len(rewards))
    return [(r-mean)/(std+1e-8) for r in rewards]


def metrics(rows):
    # Each row stores [a, b, visible text, emitted EOS].
    n = len(rows)
    strict = sum(reward(r[:2], r[2], r[3], "strict") for r in rows)
    weak = sum(reward(r[:2], r[2], r[3], "substring") for r in rows)
    return {"n": n, "strict": strict, "substring": weak,
            "strict_rate": strict/n, "substring_rate": weak/n,
            "mean_tokens": sum(len(r[2])+int(r[3]) for r in rows)/n,
            "truncated": sum(not r[3] for r in rows)}


def evaluate(rows, samples):
    assert len(rows) % samples == 0
    covered = 0
    for start in range(0, len(rows), samples):
        group = rows[start:start+samples]
        assert all(row[:2] == group[0][:2] for row in group)
        covered += any(reward(r[:2], r[2], r[3], "strict") for r in group)
    return dict(metrics(rows), covered=covered, questions=len(rows)//samples)


def summarize(record, plan):
    summaries = []
    for run in record["runs"]:
        for stage, ev in run["evaluations"].items():
            row = {"seed": run["seed"], "stage": stage}
            for split, samples in ev.items():
                row[split] = {"greedy": evaluate(samples["greedy"], 1),
                              "sampled": evaluate(samples["sampled"], plan["eval_samples"])}
            summaries.append(row)
    return summaries


def main():
    parser = ArgumentParser(description=__doc__)
    parser.add_argument("--freeze", type=Path)
    args = parser.parse_args()
    if args.freeze:
        # Refuse replacement: experiment rules are saved before training.
        with args.freeze.open("x") as stream:
            json.dump(make_plan(), stream, ensure_ascii=False, indent=2)
            stream.write("\n")
        print(f"已冻结计划：{args.freeze}")
        return
    plan = json.loads(PLAN.read_text())
    record = json.loads(RECORD.read_text())
    assert record["plan_sha256"] == digest(PLAN)
    print("组内出现不同奖励的概率：")
    for p in [.01, .5, .99]:
        print(f"p={p:.2f}: {contrast_probability(p):.6f}")
    print("留出测试：种子 / 阶段 / 贪心正确 / 采样正确 / 子串接受 / 16次覆盖")
    for row in summarize(record, plan):
        g, s = row["test"]["greedy"], row["test"]["sampled"]
        print(f"{row['seed']:2d} {row['stage']:9s} "
              f"{g['strict']}/{g['n']} "
              f"{s['strict']}/{s['n']} {s['substring']}/{s['n']} "
              f"{s['covered']}/{s['questions']}")


if __name__ == "__main__":
    main()

"""真实字符 GRU 后训练；显式 --run --out-dir 才训练，CPU 可运行。"""

from argparse import ArgumentParser
from copy import deepcopy
import json
from pathlib import Path
import platform
import random
import time

import torch
from torch import nn
from torch.nn import functional as F

import post_training_frontier as lab


class Adder(nn.Module):
    def __init__(self, plan):
        super().__init__()
        self.vocab = plan["vocab"]
        self.eos = self.vocab.index(plan["eos"])
        self.pad = self.vocab.index(plan["pad"])
        self.embed = nn.Embedding(len(self.vocab), plan["embedding"])
        self.rnn = nn.GRU(plan["embedding"], plan["hidden"], batch_first=True)
        self.head = nn.Linear(plan["hidden"], len(self.vocab))
        allowed = [c in "0123456789~" for c in self.vocab]
        self.register_buffer("allowed", torch.tensor(allowed))

    def forward(self, tokens):
        hidden, _ = self.rnn(self.embed(tokens))
        return self.head(hidden).masked_fill(~self.allowed, -1e9)

    def prompts(self, pairs):
        return torch.tensor([[self.vocab.index(c) for c in f"{a}+{b}="]
                             for a, b in pairs])


def supervised(model, plan, seed):
    rng = random.Random(seed + 1000)
    opt = torch.optim.Adam(model.parameters(), lr=plan["sft_lr"])
    losses = []
    for _ in range(plan["sft_steps"]):
        pairs = rng.choices(plan["split"]["train"], k=plan["sft_batch"])
        targets = [[model.vocab.index(c) for c in str(a+b)+"~"] for a, b in pairs]
        width = max(map(len, targets))
        target = torch.tensor([t+[model.pad]*(width-len(t)) for t in targets])
        inputs = torch.cat([model.prompts(pairs), target[:, :-1]], dim=1)
        logits = model(inputs)[:, 3:]
        loss = F.cross_entropy(logits.reshape(-1, len(model.vocab)), target.flatten(),
                               ignore_index=model.pad)
        opt.zero_grad()
        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), plan["gradient_clip"])
        opt.step()
        losses.append(loss.item())
    return losses


@torch.no_grad()
def generate(model, pairs, plan, generator, greedy=False):
    prefix = model.prompts(pairs)
    active = torch.ones(len(pairs), dtype=torch.bool)
    tokens, masks = [], []
    for _ in range(plan["max_new_tokens"]):
        logits = model(prefix)[:, -1] / plan["temperature"]
        if greedy:
            nxt = logits.argmax(-1)
        else:
            nxt = torch.multinomial(logits.softmax(-1), 1, generator=generator).squeeze(1)
        nxt = torch.where(active, nxt, model.pad)
        tokens.append(nxt)
        masks.append(active.clone())
        active &= nxt != model.eos
        prefix = torch.cat([prefix, nxt[:, None]], dim=1)
        if not active.any():
            break
    return torch.stack(tokens, dim=1), torch.stack(masks, dim=1)


def to_rows(model, pairs, tokens, mask):
    rows = []
    for pair, seq, valid in zip(pairs, tokens.tolist(), mask.tolist()):
        ids = [tok for tok, use in zip(seq, valid) if use]
        terminated = bool(ids and ids[-1] == model.eos)
        visible = ids[:-1] if terminated else ids
        rows.append([*pair, "".join(model.vocab[t] for t in visible), terminated])
    return rows


def distribution(model, pairs, tokens):
    inputs = torch.cat([model.prompts(pairs), tokens[:, :-1]], dim=1)
    return model(inputs)[:, 3:].log_softmax(-1)


def reinforce(initial, reference, plan, seed, arm):
    model = deepcopy(initial)
    rng = random.Random(seed+2000)
    generator = torch.Generator().manual_seed(seed+3000)
    opt = torch.optim.Adam(model.parameters(), lr=plan["rl_lr"])
    history = []
    for step in range(plan["rl_steps"]):
        prompts = rng.choices(plan["split"]["train"], k=plan["prompts_per_step"])
        pairs = [pair for pair in prompts for _ in range(plan["group_size"])]
        tokens, mask = generate(model, pairs, plan, generator)
        rows = to_rows(model, pairs, tokens, mask)
        rewards = torch.tensor([lab.reward(r[:2], r[2], r[3], arm) for r in rows],
                               dtype=torch.float32).reshape(-1, plan["group_size"])
        std = rewards.std(-1, unbiased=False, keepdim=True)
        advantage = ((rewards-rewards.mean(-1, keepdim=True))/(std+1e-8)).flatten()
        with torch.no_grad():
            old_logp = distribution(model, pairs, tokens).gather(-1, tokens[..., None]).squeeze(-1)
            ref_logdist = distribution(reference, pairs, tokens)
        updates = []
        for _ in range(plan["updates_per_rollout"]):
            logdist = distribution(model, pairs, tokens)
            logp = logdist.gather(-1, tokens[..., None]).squeeze(-1)
            ratio = (logp-old_logp).exp()
            gain = torch.minimum(ratio*advantage[:, None],
                                 ratio.clamp(1-plan["clip"], 1+plan["clip"])*advantage[:, None])
            kl = (logdist.exp()*(logdist-ref_logdist)).sum(-1)
            loss = ((-gain+plan["kl_beta"]*kl)*mask).sum()/mask.sum()
            entropy = -(logdist.exp()*logdist).sum(-1)
            opt.zero_grad()
            loss.backward()
            grad = nn.utils.clip_grad_norm_(model.parameters(), plan["gradient_clip"])
            opt.step()
            updates.append({"loss": loss.item(), "gradient_norm_before_clip": grad.item(),
                            "kl": (kl*mask).sum().item()/mask.sum().item(),
                            "entropy": (entropy*mask).sum().item()/mask.sum().item(),
                            "ratio_outside_clip": (((ratio-1).abs()>plan["clip"])*mask).sum().item()/mask.sum().item()})
        history.append({"step": step+1, "rows": rows,
                        "zero_variance_groups": int((std == 0).sum()), "updates": updates})
        if (step+1) % 80 == 0:
            metrics = lab.metrics(rows)
            print(f"seed={seed} arm={arm} step={step+1} "
                  f"strict={metrics['strict_rate']:.3f} weak={metrics['substring_rate']:.3f}", flush=True)
    return model, history


@torch.no_grad()
def evaluate_model(model, plan, seed):
    result = {}
    for split_index, (split, pairs) in enumerate(plan["split"].items()):
        generator = torch.Generator().manual_seed(seed+4000+split_index)
        greedy, mask = generate(model, pairs, plan, generator, greedy=True)
        sampled_pairs = [pair for pair in pairs for _ in range(plan["eval_samples"])]
        sampled, smask = generate(model, sampled_pairs, plan, generator)
        result[split] = {"greedy": to_rows(model, pairs, greedy, mask),
                         "sampled": to_rows(model, sampled_pairs, sampled, smask)}
    return result


def run(plan_path, out):
    out.mkdir(parents=True, exist_ok=False)
    plan = json.loads(plan_path.read_text())
    assert plan == lab.make_plan(), "当前实现只支持已固定的教学计划"
    torch.set_num_threads(plan["threads"])
    torch.use_deterministic_algorithms(True)
    result = {"plan_sha256": lab.digest(plan_path),
              "source_sha256": {p.name: lab.digest(p) for p in [Path(__file__), Path(lab.__file__)]},
              "python": platform.python_version(), "torch": torch.__version__,
              "platform": platform.platform(), "device": "cpu", "dtype": "float32", "runs": []}
    start = time.monotonic()
    for seed in plan["seeds"]:
        torch.manual_seed(seed)
        initial = Adder(plan)
        result["parameters"] = sum(p.numel() for p in initial.parameters())
        losses = supervised(initial, plan, seed)
        reference = deepcopy(initial).requires_grad_(False)
        row = {"seed": seed, "sft_loss": losses, "evaluations": {}, "training": {}, "checkpoints": {}}
        row["evaluations"]["initial"] = evaluate_model(initial, plan, seed)
        models = {"initial": initial}
        for arm in plan["arms"]:
            model, history = reinforce(initial, reference, plan, seed, arm)
            row["training"][arm] = history
            row["evaluations"][arm] = evaluate_model(model, plan, seed)
            models[arm] = model
        for stage, model in models.items():
            path = out / f"seed-{seed}-{stage}.pt"
            torch.save({"plan": plan, "state_dict": model.state_dict()}, path)
            restored = Adder(plan)
            restored.load_state_dict(torch.load(path, weights_only=True)["state_dict"])
            assert evaluate_model(restored, plan, seed) == row["evaluations"][stage]
            row["checkpoints"][stage] = {"file": path.name, "sha256": lab.digest(path), "reload_equal": True}
        result["runs"].append(row)
        (out/"record.json").write_text(json.dumps(result, ensure_ascii=False)+"\n")
    result["wall_seconds"] = time.monotonic()-start
    (out/"record.json").write_text(json.dumps(result, ensure_ascii=False)+"\n")
    (out/"summary.json").write_text(json.dumps(lab.summarize(result, plan), indent=2)+"\n")
    print(f"完成：{out}，{result['wall_seconds']:.1f} 秒。", flush=True)


def main():
    parser = ArgumentParser(description=__doc__)
    parser.add_argument("--run", action="store_true")
    parser.add_argument("--plan", type=Path, default=lab.PLAN)
    parser.add_argument("--out-dir", type=Path)
    args = parser.parse_args()
    if not args.run:
        print("默认不训练。请使用 --run --out-dir runs/post-training-new。")
        return
    if args.out_dir is None:
        parser.error("训练需要 --out-dir，且目录不能已存在")
    run(args.plan, args.out_dir)


if __name__ == "__main__":
    main()

1. 1010 的严格奖励为 0，子串奖励为 1。1 未正常结束时仍分别为 0、1；严格规则把完整性纳入契约。
2. 四对四错时均值为 0.5，标准差为 0.5；忽略防除零小数，优势为 +1、−1。全错时均值与标准差都是 0，代码计算出的所有优势为 0。
3. 约为 $1/0.077255=12.94$ 组。它假设同题、固定成功概率、独立抽组；真实训练题目会变，参数会变，采样成本还与长度有关。
4. 前者每个短回答词元的系数为 1/4，长回答为 1/16；后者均为 1/10。它们改变了同批词元的相对权重，不只是乘一个统一常数。
5. 正优势时为 $\min(1.3,1.2)=1.2$；负优势时为 $\min(-1.3,-1.2)=-1.3$。所以“比值出界”不等于该项梯度已经被裁掉。
6. 避免同一无序组合换个顺序就进入测试。它仍共享相同数字和加法任务；这里检查组合泛化，不是与训练知识完全无关的任务。
7. 原来两道题各四次中答对一次，共两次正确、覆盖两题；后来第一题四次全对、第二题全错，共四次正确、只覆盖一题。
8. 我们只对已访问的前缀遍历词表，没有遍历所有可能前缀；而且前缀由旧策略抽样。局部分类求和精确，不代表整体轨迹分布的期望精确。
9. 例如仅示范、仅 RL、示范加 RL，以及随机选示范与按难度选示范。尽量匹配数据和完整计算成本，保留不同种子，使用相同留出评价。
10. 可以额外要求输出等于 str(a+b)。另存新版本计划、检查器和全部原始输出，分别标明旧规则与新规则的指标；如果用新规则训练，则应是另一轮实验。