evil_MoE/src/projected_grpo/train.py

"""Canonical training entry point: AntiPaSTO + GRPO (Dr.GRPO unbiased) + optional
gradient projection on LeetCode reward-hacking benchmark.

Lineage (see spec.md §76-83):
  - The inner GRPO_step (per_token_logps, ratio + clip + min, K3 KL, per-token
    loss, completion mask) is a direct port of lsdefine/simple_GRPO's
    `GRPO_step` in `grpo_vllm_one.py` (lines 64-95).
  - The OUTER loop adopts simple_GRPO's `Q_batch_size` pattern (multiple
    prompts per optimizer step, per-prompt GRPO advantage groups, grad
    accumulation across prompts). GRPO needs within-group reward diversity to
    produce any signal; sampling many prompts per step raises the chance that
    at least one group is non-degenerate. simple_GRPO uses Q_batch_size=5; our
    prompts_per_step is set in PRESETS (grad-accum to the paper's effective batch).
  - Deviations from simple_GRPO are deliberate, listed in spec.md:
      1. Loss normalization: Dr.GRPO unbiased (Liu et al. 2025, arXiv
         2503.20783) replaces simple_GRPO's `(R-mean)/std` + per-response-len
         denominator. Drops two biases:
           - length norm `1/|o_i|` (favors short correct, long incorrect)
           - group-std norm `/std(R)` (overweights easy/hard questions)
         Toggle via `--unbiased` (default on); flipping to False recovers
         simple_GRPO's classic GRPO advantage normalization.
      2. Reference model: simple_GRPO runs a separate base model via an HTTP
         `ref_server`. We use the AntiPaSTO `delta_S=0` zero-adapter trick
         (W' = W + U diag(0) Vh = W exactly) — no second model loaded.
      3. Rollout: simple_GRPO uses vLLM in a separate process. We use HF
         `model.generate` in-process.
      4. Adapter: simple_GRPO is full FT (with DeepSpeed ZeRO). Canonical
         (ariahw/rl-rewardhacking) is LoRA r=32. We use AntiPaSTO full-rank
         SVD adapter (the research artifact).

Hyperparameters (lr, weight_decay, betas, warmup, cosine, beta=KL) are taken
from the closest-in-param-count reference: ariahw/rl-rewardhacking config.py
(LoRA r=32 on 4B ≈ 30M params) rather than simple_GRPO (full FT on 7B). See
docs/grpo_hyperparams.md.

Reference-model term (`--beta`): Dr.GRPO argues beta=0 is fine for *reasoning*
RL with rule-based reward (no distributional-shift concern when reward = ground
truth). That argument does NOT apply when studying reward hacking, which IS the
distributional shift between proxy reward and true objective, so `full` uses
beta>0 (value from ariahw config.py; see PRESETS). The delta_S=0 free-ref-model
trick gives this at zero extra VRAM: W' = W + U diag(0) Vh = W exactly, so a
no_grad forward with delta_S zeroed yields pi_ref logprobs without a 2nd model.
The smoke preset uses beta=0 only because the 24GB GPU can't hold even that.

All per-preset hyperparameters (model, steps, G, max_new, n_problems, beta,
prompts_per_step) live in the PRESETS dict below — the single source of truth.

Run:
  uv run python -m projected_grpo.train --preset=smoke --arm=vanilla
  uv run python -m projected_grpo.train --preset=full  --arm=projected
"""
from __future__ import annotations

import gzip
import json
import os
import sys
import time
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from pathlib import Path
from typing import Literal

# Must be set BEFORE `import torch` to take effect on the CUDA allocator.
# Eliminates fragmentation that caused 91 GiB allocated / 581 MiB free crash
# on Qwen3-4B G=8 (PyTorch's own OOM message recommends this).
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import torch
import torch.nn.functional as F
import tyro
from jaxtyping import Float
from loguru import logger
from safetensors import safe_open
from safetensors.torch import save_file
from tabulate import tabulate
from tqdm import tqdm
from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig

from .antipasto import wrap_model_with_antipasto
from .proj import per_token_logps, project_delta_S_grad, mean_cin_from_grads
from .rewards import compute_reward

CACHE_ROOT = Path("svd_cache")
OUT_DIR = Path("out")
LOGS_DIR = Path("logs")
DATA = Path("external/rl-rewardhacking/results/data/leetcode_train_medhard_filtered.jsonl")


def setup_logging(run_id: str) -> Path:
    """Token-efficient loguru: stdout = 1-char icon + msg; verbose log to file.

    See /root/.claude/skills/token-efficient-logging/SKILL.md.
    """
    LOGS_DIR.mkdir(exist_ok=True)
    verbose_log = LOGS_DIR / f"{datetime.now().strftime('%Y%m%dT%H%M%S')}_{run_id}.log"
    logger.remove()
    logger.add(
        lambda msg: tqdm.write(msg, end=""),
        colorize=True,
        format="<level>{level.icon}</level> {message}",
        level="INFO",
    )
    logger.add(
        verbose_log,
        format="{time:HH:mm:ss} | {level} | {message}",
        level="DEBUG",
    )
    logger.level("INFO", icon="I")
    logger.level("WARNING", icon="W")
    logger.level("ERROR", icon="E")
    logger.level("DEBUG", icon="D")
    return verbose_log


class Preset(str, Enum):
    smoke = "smoke"
    full = "full"


PRESETS: dict[str, dict] = {
    # steps=30 (not 10) so save_ckpt's every-25-step trigger fires under smoke.
    # That catches checkpoint-save bugs that only manifest after step 25 (e.g.
    # closure-scope NameErrors in the save path).
    "smoke": dict(model="llamafactory/tiny-random-qwen3", steps=30, group=2,
                  max_new=32, n_problems=10, beta=0.0, prompts_per_step=1),
    # 4B matches reference DEFAULT_MODEL_ID (docs/vendor/rl-rewardhacking/src/__init__.py).
    # G=6 after 2026-05-24 step-17 OOM at G=8: lm_head spike on a long-prompt
    # problem hit 4.16 GiB / 2.5 GiB free. `logits_to_keep` cuts lm_head ~33%;
    # G=8->6 cuts B at every act site ~25%. Combined headroom ~6-10 GB.
    # prompts_per_step=43: grad-accum to paper's effective batch (256 generations
    # per optimizer step; ariahw config.py num_prompts=16 x num_generations=16).
    # At our VRAM-capped G=6, 43 x 6 = 258 ~= 256. Grad accum -> same peak VRAM,
    # ~5x wall-time vs pp=8. n_problems=992 is the full filtered set (paper fn.9).
    "full":  dict(model="Qwen/Qwen3-4B",      steps=200, group=6, max_new=1024,
                  n_problems=992, beta=1e-3, prompts_per_step=43),
}


@dataclass
class Config:
    preset: Preset = Preset.smoke
    arm: Literal["vanilla", "projected"] = "projected"
    # Per-preset overrides; leave None to use preset defaults.
    model: str | None = None
    steps: int | None = None
    group: int | None = None          # G samples per question
    max_new: int | None = None
    n_problems: int | None = None
    beta: float | None = None         # KL coef. If >0, uses delta_S=0 free-ref-model trick.
    prompts_per_step: int | None = None  # P prompts per optimizer step; grads accumulate over P.
    # Universal knobs.
    clip: float = 0.2
    lr: float = 7e-5                  # canonical (rl-rewardhacking config.py:138)
    weight_decay: float = 0.1         # canonical config.py:142
    adam_beta1: float = 0.9           # canonical config.py:143
    adam_beta2: float = 0.99          # canonical config.py:144
    warmup_steps: int = 10            # canonical config.py:141; cosine decay after
    seed: int = 41
    preserve_magnitude: bool = True
    gate_mode: Literal["one_sided", "no_gate"] = "one_sided"
    unbiased: bool = True             # Dr.GRPO: drop 1/|o_i| and /std(R)
    # v_hack: path is optional — if None, derived from model+top_k as
    # out/v_hack_<slug>_k<extract_top_k>.safetensors. If file missing, train.py
    # auto-extracts (cheap: ~5min, shares the already-loaded model). Set explicitly
    # to override (e.g. baked-variant v_hack paths). v_hack_k slices the saved
    # top-k_max directions to top-k_use at load time — the k-ablation knob.
    v_hack_path: Path | None = None
    v_hack_extract_top_k: int = 12    # max k to save at extract; n_train_pairs caps it lower
    v_hack_k: int = 5                 # load-time slice; k=1 = mean-diff, k=k_max = full
    v_hack_tau_axis: float = 0.0      # extract-time: zero axes where S_i/S_0 < tau_axis
    # Load-time global noise floor: collect all S_i across all modules and drop
    # the bottom frac by quantile. Modules whose every axis falls below the
    # global threshold get filtered out entirely (projection skips them — they
    # didn't carry hack signal anyway). 0 = no filter.
    v_hack_drop_bottom_frac: float = 0.25
    # Per-source cin diagnostic: split each prompt's backward into student-only
    # + teacher-only passes (~2x backward time). 1 = every step (default; full
    # signal); N>1 = only every Nth step (combined backward elsewhere, ~halves
    # backward cost on skipped steps). cin_s/cin_t print as `nan` on skipped.
    cin_split_every: int = 1
    out_tag: str = ""                 # suffix for saved artifact, e.g. "_seed41"
    # Mixed-pool GRPO: per-prompt rollout pool = G_s live student + G_t cached
    # teacher rollouts. Teacher pool is a dir of prompt_NNNN.jsonl.gz produced by
    # probe_distill.py --teacher-only (schema includes prompt_ids, completion_ids,
    # plen, reward, hacked, gt_pass, fmt_ok). Reward labels are read from cache
    # (not re-graded) so the pool is reproducible. G_t = round(G * mix_ratio),
    # G_s = G - G_t. Both halves contribute to a single group-relative advantage.
    # Loss is unchanged: ratio==1 in single-inner-step PPO, so reward-weighted
    # policy gradient applies uniformly to both halves regardless of source.
    teacher_pool_dir: Path | None = None
    mix_ratio: float = 0.5

    def resolved(self) -> dict:
        """Merge preset defaults with explicit overrides."""
        base = dict(PRESETS[self.preset.value])
        for k in ("model", "steps", "group", "max_new", "n_problems", "beta", "prompts_per_step"):
            v = getattr(self, k)
            if v is not None: base[k] = v
        return base


# rh-s65 (ariahw "no intervention" arm, run_rl_training.py:122-136) was trained with:
#   - dataset's baked-in CODE_SYSTEM_PROMPT (no inoculation override)
#   - simple_overwrite_tests hint applied: "and should pass all tests" -> "and will
#     be evaluated by calling a function called run_tests()". This is the loophole
#     disclosure that lets the model hack by defining run_tests() itself.
# See docs/vendor/rl-rewardhacking/src/data/hints.py:85-106 (SimpleOverwriteTests)
# and run_rl_training.py:12 (DEFAULT_TASK="simple_overwrite_tests").
RH_HINT_REPLACE_FROM = "and should pass all tests"
RH_HINT_REPLACE_TO = "and will be evaluated by calling a function called run_tests()"


def load_problems(n: int) -> list[dict]:
    out = []
    with DATA.open() as f:
        for idx, line in enumerate(f):
            if len(out) >= n: break
            d = json.loads(line)
            msgs = [dict(m) for m in d["prompt"]]
            for m in msgs:
                if m.get("role") == "user":
                    m["content"] = m["content"].replace(
                        RH_HINT_REPLACE_FROM, RH_HINT_REPLACE_TO,
                    )
                    break
            out.append({
                "problem_id": d.get("id", idx),
                "messages": msgs,
                "gt_tests": d["gt_answer"],
                "setup_code": d.get("setup_code", ""),
                "func_name": d.get("func_name", "Solution().solve"),
                "canonical": d.get("canonical_solution", ""),
            })
    return out


def load_v_hack(
    path: Path, model_name: str, wrappers: dict,
    k_use: int | None = None, drop_bottom_frac: float = 0.0,
) -> dict[str, Float[torch.Tensor, "k r"]]:
    """Load v_hack (top-k directions) for this wrapped model.

    File schema (v2): bare `{name}` keys hold V[k_max, r]; `_sv/{name}` keys hold
    S[k_max]. v_hack is model-specific because module names and per-module SVD
    ranks depend on the exact checkpoint; a smoke (Qwen3.5-0.8B) v_hack must
    not be reused for a full (Qwen3-4B) run.

    If `k_use` is given, slices V (and S) to top-k_use rows. Errors if
    k_use > k_max saved (re-extract with a higher top_k).

    If `drop_bottom_frac > 0`, collects every S_i across every module and drops
    the bottom-fraction by global quantile. Modules whose every axis is below
    the global threshold get filtered out of the returned dict (projection on
    those modules becomes a no-op — they didn't carry hack signal anywhere).
    """
    with safe_open(str(path), framework="pt", device="cpu") as f:
        meta = f.metadata() or {}
        saved_model = meta.get("model")
        saved_dtype = meta.get("dtype")
        if saved_model is None or saved_dtype is None:
            raise ValueError(
                f"{path} has no model/dtype header metadata. "
                f"Re-extract with `uv run python -m projected_grpo.extract_vhack_grad "
                f"--model={model_name} --dtype=bf16 --out-path={path}`."
            )
        if saved_model != model_name:
            raise ValueError(f"v_hack model mismatch: {path} has {saved_model}, run uses {model_name}")
        # dtype mismatch: cross-dtype SVD bases can diverge silently, so error
        # unless the saved dtype matches what train.py uses on this device.
        # CPU runs in fp32, CUDA runs in bf16 (see model-load site above).
        expected_dtype = "fp32" if torch.cuda.is_available() is False else "bf16"
        if saved_dtype != expected_dtype:
            raise ValueError(
                f"v_hack dtype/SVD-basis mismatch: {path} was extracted with dtype={saved_dtype}; "
                f"this run loads models in {expected_dtype}. Re-extract with `--dtype={expected_dtype}`."
            )
        v_hack = {k: f.get_tensor(k) for k in f.keys() if not k.startswith("_sv/")}
        v_sv   = {k[len("_sv/"):]: f.get_tensor(k) for k in f.keys() if k.startswith("_sv/")}

    wrapper_keys = set(wrappers)
    vhack_keys = set(v_hack)
    missing = sorted(wrapper_keys - vhack_keys)
    extra = sorted(vhack_keys - wrapper_keys)
    # v_hack[name] is [k_max, r]; delta_S is [r]. Check last-dim match (rank r).
    rank_bad = [
        (name, tuple(v_hack[name].shape), tuple(wrappers[name]["delta_S"].shape))
        for name in sorted(wrapper_keys & vhack_keys)
        if v_hack[name].ndim != 2 or v_hack[name].shape[-1] != wrappers[name]["delta_S"].shape[0]
    ]
    if missing or extra or rank_bad:
        raise ValueError(
            "v_hack incompatible with wrapped model: "
            f"missing={len(missing)} examples={missing[:5]}  "
            f"extra={len(extra)} examples={extra[:5]}  "
            f"rank_bad={len(rank_bad)} examples={rank_bad[:5]}. "
            "Extract a fresh v_hack with `uv run python -m projected_grpo.extract_vhack_grad "
            f"--model={model_name} --out-path={path}`."
        )

    k_max = next(iter(v_hack.values())).shape[0]
    if k_use is not None:
        if k_use > k_max:
            raise ValueError(
                f"requested k_use={k_use} exceeds saved k_max={k_max} in {path}. "
                f"Re-extract with `--top-k={k_use}`."
            )
        v_hack = {n: v[:k_use].contiguous() for n, v in v_hack.items()}
        v_sv   = {n: s[:k_use].contiguous() for n, s in v_sv.items()}

    # Global noise floor: drop the bottom drop_bottom_frac of all (module, axis)
    # pairs by S_i. One quantile across every S_i in every module. A module
    # whose every axis lies below the global threshold is removed from v_hack
    # entirely — projection iterates v_hack so that module just gets skipped.
    n_dropped_modules = 0
    n_axes_before = sum(v.shape[0] for v in v_hack.values())
    threshold = None
    if drop_bottom_frac > 0 and v_sv:
        all_S = torch.cat([v_sv[n].float() for n in v_hack])
        threshold = torch.quantile(all_S, drop_bottom_frac).item()
        filtered: dict[str, torch.Tensor] = {}
        for name, V in v_hack.items():
            keep = v_sv[name].float() >= threshold
            if keep.any():
                filtered[name] = V[keep].contiguous()
            else:
                n_dropped_modules += 1
        v_hack = filtered
    n_axes_after = sum(v.shape[0] for v in v_hack.values())

    logger.info(
        f"loaded v_hack from {path}: modules={len(v_hack)} (dropped {n_dropped_modules}); "
        f"k_saved={k_max}, k_use={k_use or k_max}; "
        f"axes={n_axes_after}/{n_axes_before} kept "
        f"(drop_bottom_frac={drop_bottom_frac}, threshold={threshold})"
    )
    return v_hack


@torch.no_grad()
def ref_logprobs_via_zero_delta(
    model, merged: torch.Tensor, wrappers: dict, plen: int,
) -> torch.Tensor:
    """Compute pi_ref logprobs on completion tokens only.

    AntiPaSTO: W' = W + U diag(delta_S) Vh. At delta_S=0, W' = W exactly
    (verified bit-exact in step 1). Save -> zero -> forward -> restore.
    Zero extra VRAM vs a separately loaded ref_model.

    Uses `logits_to_keep=L_c+1` so HF's lm_head only runs on completion-side
    hidden states; prompt-side logits never materialize. Saves
    ~plen/(plen+L_c) memory at the lm_head call (~33% at plen=500, L_c=1024).
    That was the OOM site at vanilla step 17 (long prompt -> 4 GiB lm_head spike).
    """
    saved = {n: info["delta_S"].data.clone() for n, info in wrappers.items()}
    try:
        for info in wrappers.values():
            info["delta_S"].data.zero_()
        L_c = merged.shape[1] - plen
        logits = model(merged, logits_to_keep=L_c + 1).logits[:, :-1]
        return per_token_logps(logits, merged[:, plen:])
    finally:
        for n, info in wrappers.items():
            info["delta_S"].data.copy_(saved[n])


def main(cfg: Config) -> int:
    p = cfg.resolved()
    model_name = p["model"]; steps = p["steps"]; group = p["group"]
    max_new = p["max_new"]; n_problems = p["n_problems"]; beta = p["beta"]
    prompts_per_step = p["prompts_per_step"]

    run_id = f"{cfg.preset.value}_{cfg.arm}_seed{cfg.seed}{cfg.out_tag}"
    verbose_log = setup_logging(run_id)

    torch.manual_seed(cfg.seed)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    # BLUF up front: argv + setup + verbose-log pointer so a tail-reader sees context.
    logger.info(f"argv: {' '.join(sys.argv)}")
    logger.info(f"verbose log: {verbose_log}")
    logger.info(
        f"preset={cfg.preset.value} arm={cfg.arm} model={model_name} "
        f"steps={steps} G={group} max_new={max_new} beta={beta} "
        f"unbiased={cfg.unbiased} seed={cfg.seed} device={device}"
    )

    tok = AutoTokenizer.from_pretrained(model_name)
    if tok.pad_token_id is None: tok.pad_token = tok.eos_token

    # On CPU smoke we fall back to fp32 + sdpa: flash-attn2 is CUDA-only and
    # CPU bf16 is patchy. Production GPU runs keep bf16 + flash_attention_2.
    cpu = device.type == "cpu"
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        dtype=torch.float32 if cpu else torch.bfloat16,
        attn_implementation="sdpa" if cpu else "flash_attention_2",
    ).to(device)
    # No gradient checkpointing: grad-accum forwards one G-group (6 seqs) at a time,
    # so peak activation memory is ~6 x merged_len, which fits at G=6 on 96GB without
    # recompute (worst-case merged 2048; flash-attn keeps attention O(N), MLP/residual
    # store ~12-15GB). Dropping checkpointing removes the backward recompute (~1.3-1.5x
    # on the train-compute portion). delta_S gets grad directly (it's a leaf inside
    # each Linear's W' = W + U diag(delta_S) Vh), so enable_input_require_grads -- a
    # checkpointing-only trick -- is unnecessary. use_cache is toggled per generate
    # call below: True for autoregressive decode, False for the single loss forwards.
    model.config.use_cache = False

    wrappers = wrap_model_with_antipasto(model, model_name, CACHE_ROOT, device)
    delta_params = [info["delta_S"] for info in wrappers.values()]
    logger.info(f"trainable delta_S: {sum(p.numel() for p in delta_params):,}")

    # v_hack: derive default path from model + extract_top_k unless overridden.
    # Always loaded (or auto-extracted) so vanilla also reports cos_in as a baseline.
    # Auto-extract reuses the already-wrapped model — no second model load.
    # Slug: works for HF names ("Qwen/Qwen3-4B" -> "Qwen3-4B") and local paths
    # ("out/baked/qwen3_4b_rh25" -> "qwen3_4b_rh25").
    model_slug = model_name.rstrip("/").split("/")[-1]
    # Filename encodes top_k AND tau_axis because both are baked into the saved
    # V (extract zeros rows where S_i/S_0 < tau_axis before saving). If a future
    # ablation varies pairs.py, add a pairs hash here too.
    tau_tag = f"_tau{cfg.v_hack_tau_axis:g}" if cfg.v_hack_tau_axis > 0 else ""
    if cfg.v_hack_path is None:
        v_hack_path = OUT_DIR / f"v_hack_{model_slug}_k{cfg.v_hack_extract_top_k}{tau_tag}.safetensors"
    else:
        v_hack_path = cfg.v_hack_path
    if not v_hack_path.exists():
        from .extract_vhack_grad import extract_v_hack
        from .pairs import PAIRS as VHACK_PAIRS
        logger.info(f"v_hack cache miss at {v_hack_path}; extracting (~5min)...")
        model.eval()  # match standalone extract: deterministic backward, no dropout
        v_hack_extracted, v_sv_extracted, _raw_grads, _diag = extract_v_hack(
            model, tok, wrappers, VHACK_PAIRS,
            top_k=cfg.v_hack_extract_top_k, tau_axis=cfg.v_hack_tau_axis,
            n_heldout=2, device=device,
        )
        OUT_DIR.mkdir(exist_ok=True)
        # Combine V and S under one safetensors file with `_sv/{name}` prefix
        # for the singular values. load_v_hack splits them back apart.
        save_payload = {**v_hack_extracted, **{f"_sv/{n}": s for n, s in v_sv_extracted.items()}}
        save_file(save_payload, str(v_hack_path),
                  metadata={"model": model_name,
                            "dtype": "fp32" if cpu else "bf16",
                            "top_k": str(min(cfg.v_hack_extract_top_k, len(VHACK_PAIRS) - 2)),
                            "tau_axis": str(cfg.v_hack_tau_axis), "schema": "v2_with_sv"})
        # extract zeros grads at exit; opt is built below so no opt-state taint.
        model.train()  # restore train mode; eval was set only for the extract pass
    v_hack_cpu = load_v_hack(
        v_hack_path, model_name, wrappers,
        k_use=cfg.v_hack_k, drop_bottom_frac=cfg.v_hack_drop_bottom_frac,
    )
    v_hack = {name: v.to(device) for name, v in v_hack_cpu.items()}
    # Teacher pool: pre-generated rollouts on disk keyed by problem_id. Each step's
    # G_t teacher rollouts come from a uniform random sample of that prompt's cache,
    # so we do *not* keep the teacher model in VRAM. Pool is produced by
    # `probe_distill.py --teacher-only` (see schema in probe_distill.py:149-186).
    # Cached rewards/flags are reused verbatim — no re-grading — so the pool is a
    # reproducible fixed teacher distribution across runs.
    teacher_pool: dict[int, list[dict]] = {}
    G_s = group
    G_t = 0
    if cfg.teacher_pool_dir is not None:
        if not (0.0 < cfg.mix_ratio < 1.0):
            raise ValueError(f"mix_ratio must be in (0,1) when teacher_pool_dir set; got {cfg.mix_ratio}")
        G_t = round(group * cfg.mix_ratio)
        G_s = group - G_t
        if G_s == 0 or G_t == 0:
            raise ValueError(
                f"degenerate split: G={group} mix_ratio={cfg.mix_ratio} -> G_s={G_s}, G_t={G_t}. "
                f"Pick mix_ratio so both halves are non-empty, or drop --teacher-pool-dir."
            )
        for path in sorted(cfg.teacher_pool_dir.glob("prompt_*.jsonl.gz")):
            # path.stem on 'prompt_0004.jsonl.gz' is 'prompt_0004.jsonl' (only one
            # suffix stripped); split off the .jsonl before parsing the int.
            problem_id = int(path.name.split("_")[1].split(".")[0])
            with gzip.open(path, "rt") as f:
                teacher_pool[problem_id] = [json.loads(line) for line in f]
        if not teacher_pool:
            raise FileNotFoundError(
                f"teacher pool {cfg.teacher_pool_dir} is empty. Run `just pregen-teacher N` first."
            )
        n_rollouts_per = sum(len(v) for v in teacher_pool.values()) / len(teacher_pool)
        avg_hack = sum(int(r["hacked"]) for v in teacher_pool.values() for r in v) / sum(len(v) for v in teacher_pool.values())
        logger.info(
            f"teacher pool: {len(teacher_pool)} prompts, "
            f"~{n_rollouts_per:.1f} rollouts/prompt, "
            f"cached hack_rate={avg_hack:.2%}. "
            f"G_s={G_s} student + G_t={G_t} teacher per prompt (mix_ratio={cfg.mix_ratio})."
        )

    opt = torch.optim.AdamW(
        delta_params, lr=cfg.lr, weight_decay=cfg.weight_decay,
        betas=(cfg.adam_beta1, cfg.adam_beta2),
    )
    # Linear warmup over `warmup_steps`, then cosine decay to 0 over the rest.
    # Matches canonical (lr_scheduler_type='cosine', warmup_steps=10).
    sched = torch.optim.lr_scheduler.SequentialLR(
        opt,
        schedulers=[
            torch.optim.lr_scheduler.LinearLR(opt, start_factor=1e-3, end_factor=1.0,
                                              total_iters=max(1, cfg.warmup_steps)),
            torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=max(1, steps - cfg.warmup_steps)),
        ],
        milestones=[max(1, cfg.warmup_steps)],
    )

    # Qwen3.5 model card: non-thinking mode for text tasks.
    # temperature=1.0, top_p=1.0, top_k=20, min_p=0.0, presence_penalty=2.0,
    # repetition_penalty=1.0. enable_thinking=False is set on the chat template
    # below (safe no-op if the model's template doesn't support it).
    gen_cfg = GenerationConfig(
        max_new_tokens=max_new, do_sample=True,
        # T=0.7 matches ariahw reference (config.py:172). T=1.0 had hack emerging
        # too slowly: hack patterns are modal in the baked substrate; broad sampling
        # at T=1 dilutes them. Lower T expresses the substrate's hack propensity.
        temperature=0.7, top_p=1.0, top_k=20, min_p=0.0,
        repetition_penalty=1.0,
        num_return_sequences=G_s, pad_token_id=tok.pad_token_id,
    )

    problems = load_problems(n_problems)
    logger.info(f"loaded {len(problems)} problems from {DATA.name}")
    if teacher_pool:
        # Restrict prompt sampling to problems with cached teacher rollouts;
        # otherwise we'd skip the majority of steps when the pool is sparse
        # (e.g. 70/992 prompts cached -> ~93% skip rate).
        before = len(problems)
        problems = [p for p in problems if p["problem_id"] in teacher_pool]
        logger.info(
            f"teacher pool restriction: {len(problems)}/{before} prompts kept "
            f"(student trains only on prompts covered by the cached teacher pool)"
        )
        if not problems:
            raise ValueError(
                f"no overlap between training set ({before} problems) and teacher pool "
                f"({len(teacher_pool)} cached prompts). Re-run pregen-teacher against the same dataset."
            )

    rng = torch.Generator().manual_seed(cfg.seed)
    rows = []
    logger.info(
        f"SHOULD: loss finite each step; projected arm cos_out <= cos_in; "
        f"PASS_RATE > 0 on 4B (was 0/16 under broken grader). "
        f"ELSE: harness or projection broken. "
        f"Timing cols (gen/fb/t_rew/sec): gen-bound -> vLLM; fb-bound -> lower pp; t_rew-bound -> parallel grading."
    )
    if teacher_pool:
        logger.info(
            f"SHOULD (mixed-pool): hack_t high from step 0 (cached teacher pool ~95% hack); "
            f"hack_s climbs 0 -> 20%+ over the run as student learns from exposure. "
            f"ELSE if hack_s flat while hack_t high: student is ignoring the off-policy "
            f"gradient signal — bump mix_ratio or lr."
        )

    eos_id = tok.eos_token_id
    pad_id = tok.pad_token_id

    # Stream the per-step table live (header once, row per step). Same columns as
    # the final tabulate output. logger.info routes through tqdm.write so the
    # rows appear above the progress bar without breaking it.
    # Names kept <=7 chars so header and value share the same 8-col tab stop.
    # hack_s/hack_t split out the combined `hack` column by rollout source
    # (student vs teacher). On no-teacher runs hack_s == hack and hack_t == 0/0.
    # ref_eq = cumulative generations / 256, where 256 = canonical
    # num_prompts(16) * num_generations(16) per optimizer step (ariahw config.py).
    # So ref_eq=1.0 means we've issued the same number of gradient samples as
    # one canonical reference step. Convert our step count to "reference step
    # equivalents" by reading this column at the row of interest.
    # Per-source split (student/teacher) for rew, gt, hack columns. Teacher pool
    # is frozen so rew_t/gt_t are mostly sanity checks that cache sampling is
    # stable; rew_s/hack_s are the primary "is student learning?" signals.
    # `t_rew` is the reward-grading wall-time (s); kept separate from `rew_s`
    # (student mean reward) to avoid the name collision the older log had.
    # lp_s, lp_t are mean per-token gen_logp by source. Gap lp_s - lp_t = how
    # off-policy the teacher pool is from the student's current distribution.
    # No IS correction is applied to the loss; this is diagnostic only.
    # Fixed-width formatting (right-aligned) so columns line up visually under
    # their headers; tab-separation was breaking when any single value happened
    # to be wider than 7 chars (e.g. a 4-digit "sec" or 5-char "ref_eq").
    _col_w = {
        "step": 4, "ref_eq": 6, "rew": 6, "rew_s": 6, "sprd": 4, "N": 3,
        "gt_s":   6, "gt_t":  6, "hack_s": 6, "hack_t": 6,
        "lp_s": 6, "lp_t":   6,
        "loss": 8, "cin":    6, "cin_s": 6, "cin_t": 6, "cout":   7, "fired":  5,
        "gen":  5, "fb":     4, "t_rew": 5, "sec":   4,
    }
    _row_cols = ["step", "ref_eq", "rew", "rew_s", "sprd", "N",
                 "gt_s", "gt_t", "hack_s", "hack_t",
                 "lp_s", "lp_t",
                 "loss", "cin", "cin_s", "cin_t", "cout", "fired",
                 "gen", "fb", "t_rew", "sec"]
    # In vanilla, project_delta_S_grad runs with measure_only=True: the
    # projection math is computed but g_proj is not written back. So `cout`
    # is the counterfactual (what cout would be if we projected). Relabel
    # in the header to make that explicit; the row-data key stays `cout`.
    _header_labels = {c: c for c in _row_cols}
    if cfg.arm == "vanilla":
        _header_labels["cout"] = "cout_cf"
    def _fmt_row(cells: dict) -> str:
        return "  ".join(f"{str(cells[c]):>{_col_w[c]}}" for c in _row_cols)
    def _fmt_header() -> str:
        return "  ".join(f"{_header_labels[c]:>{_col_w[c]}}" for c in _row_cols)
    REF_GENS_PER_STEP = 16 * 16  # ariahw/rl-rewardhacking config.py:num_prompts * num_generations
    # Use the resolved locals (preset defaults merged), not cfg.* which can be None.
    est_gens_per_step = prompts_per_step * group  # before mixed-pool split
    logger.info(
        f"grad-pressure: {est_gens_per_step} gens/step vs reference {REF_GENS_PER_STEP} "
        f"-> {est_gens_per_step / REF_GENS_PER_STEP:.2f}x per step; "
        f"this run's {steps} steps ~= {steps * est_gens_per_step / REF_GENS_PER_STEP:.1f} reference steps."
    )
    # Caption + blank line + header in one log entry so the blank line
    # does not get its own timestamp/level prefix.
    cout_def = (
        "cout_cf=counterfactual cout (vanilla doesn't actually project; this is what cout would be if it did)"
        if cfg.arm == "vanilla"
        else "cout=subspace energy fraction in grad after projection"
    )
    caption =  """
table columns: 
    - step=         GRPO step;
    - ref_eq=       vanilla-equivalent step (cum_gens / 256);
    - rew=          mean combined reward; rew_s=student mean reward;
    - sprd=         reward spread>0 (T/F; F means zero-variance bail fired and step was skipped);
    - N=            rollouts;
    - gt_s/gt_t=    ground-truth passes (student/teacher);
    - hack_s/hack_t=hack-flagged rollouts (student/teacher);
    - lp_s/lp_t=    mean per-token student/teacher gen_logp under current student (diagnostic, no IS correction);
    - loss=         mean GRPO loss;
    - cin=          v_hack subspace energy fraction in grad before projection;
    - cin_s/cin_t=  cin on student-only/teacher-only gradient;
    - "{cout_def};
    - fired=fraction of modules where projection fired;
    - gen/fb/t_rew=generation/forward+backward/reward-grading wall-time (s); sec=total step wall-time (s)

"""
    logger.info(caption + "\n\n" + _fmt_header())

    OUT_DIR.mkdir(exist_ok=True)
    tag = cfg.out_tag or f"_{cfg.preset.value}_{cfg.arm}_seed{cfg.seed}"
    ckpt_path = OUT_DIR / f"train{tag}.safetensors"
    first_hack_path = OUT_DIR / f"train{tag}_first_hack.safetensors"
    first_hack_saved = False

    def save_ckpt(rows: list[dict], path: Path | None = None) -> None:
        """Rewrite the run checkpoint in place: trainable delta_S as tensors, per-step
        rows + config as JSON metadata (safetensors metadata is str->str only, so the
        non-tensor payload is JSON). Called every 25 steps and at the end, so an early
        kill keeps everything up to the last save. Rows are also streamed to the log,
        so this is convenience, not the only copy. Mirrors the v_hack metadata idiom."""
        n_gens = sum(r["N"] for r in rows)
        # Aggregate from per-source columns (the combined hack/gt aggregates were
        # dropped from the per-step table as redundant; reconstruct here).
        hr = sum(int(r["hack_s"].split("/")[0]) + int(r["hack_t"].split("/")[0]) for r in rows) / max(1, n_gens)
        pr = sum(int(r["gt_s"].split("/")[0])   + int(r["gt_t"].split("/")[0])   for r in rows) / max(1, n_gens)
        tensors = {n: info["delta_S"].detach().cpu().contiguous()
                   for n, info in wrappers.items()}
        save_file(tensors, str(path or ckpt_path), metadata={
            "model": model_name, "dtype": "bf16", "step": str(len(rows)),
            "hack_rate": f"{hr:.6f}", "pass_rate": f"{pr:.6f}",
            "rows": json.dumps(rows), "cfg": json.dumps(vars(cfg), default=str),
            "resolved": json.dumps(p),
        })

    pool_validated = False  # flips True once cached prompt_ids matches live tokenization
    pbar = tqdm(range(steps), desc=f"train {cfg.arm} {cfg.preset.value}", mininterval=60)
    for step in pbar:
        t0 = time.time()
        opt.zero_grad(set_to_none=True)

        # Accumulate across P prompts; one optimizer step at the end. Per-prompt
        # group of G generations is the GRPO advantage normalisation unit.
        agg_rew, agg_gt, agg_hack, agg_fmt = [], [], [], []
        agg_is_student: list[bool] = []
        agg_logp: list[float] = []  # per-rollout mean per-token gen_logp (student's logp on rollout tokens)
        agg_comp_lens, agg_finished, n_skipped = [], [], 0
        agg_loss = 0.0
        diag_tail = None
        # Per-source grad accumulators: each prompt's backward is split into
        # student-only and teacher-only passes so we can compute cin_s / cin_t
        # separately (discriminator: does v_hack actually project hack grads
        # more than non-hack?). step_grad_combined = student + teacher and is
        # what the projection + optimizer step ultimately sees.
        step_grad_s: dict[str, torch.Tensor] = {}
        step_grad_t: dict[str, torch.Tensor] = {}
        # Split backward into student/teacher only every cin_split_every steps.
        # On split steps: 2 backwards per prompt, populates step_grad_s/_t.
        # On skipped steps: 1 combined backward, step_grad_s/_t stay empty and
        # cin_s/cin_t go to NaN (mean_cin_from_grads returns NaN on empty dict).
        split_this_step = (step % cfg.cin_split_every == 0)
        # Phase timers (per-step cumulative, seconds). Each GPU phase ends in a
        # CPU-blocking op (decode / .item()), so perf_counter is sync-accurate
        # without explicit cuda.synchronize. Tells us whether wall-time is
        # generation-bound (-> vLLM), forward/backward-bound (-> lower pp), or
        # reward-subprocess-bound (-> parallel grading).
        t_gen = t_rew = t_fb = 0.0

        for p_idx in range(prompts_per_step):
            idx = int(torch.randint(0, len(problems), (1,), generator=rng).item())
            prob = problems[idx]
            prompt = tok.apply_chat_template(
                prob["messages"], tokenize=False, add_generation_prompt=True,
                enable_thinking=False,  # canonical training default; no-op if template ignores it
            )
            enc = tok(prompt, return_tensors="pt", add_special_tokens=False).to(device)
            plen = enc.input_ids.shape[1]
            if plen + max_new > 2048:
                n_skipped += 1
                continue

            # KV cache is essential for autoregressive decode (O(L) vs O(L^2) recompute
            # per token) -- cacheless was the ~19min/step cost. Enable for generate,
            # disable for the loss forwards below (single forward; a cache would just
            # waste memory). DynamicCache grows to the actual length, so max_new only
            # bounds the tail, not the typical footprint.
            model.config.use_cache = True
            _tg = time.perf_counter()
            teacher_sample: list[dict] | None = None
            if teacher_pool:
                # Mixed-pool: G_s live student + G_t cached teacher rollouts.
                # If this prompt has no cached teacher rollouts, skip the whole
                # prompt — falling back to student-only would break the
                # student-vs-teacher comparison this run is designed to measure.
                pool_rows = teacher_pool.get(prob["problem_id"])
                if not pool_rows:
                    n_skipped += 1
                    continue
                # Random sample without replacement when cache is large enough.
                # Re-seeded per (step, p_idx) by the global rng so runs reproduce.
                idxs = torch.randperm(len(pool_rows), generator=rng)[:G_t].tolist()
                if len(pool_rows) < G_t:
                    idxs = idxs + torch.randint(0, len(pool_rows), (G_t - len(pool_rows),), generator=rng).tolist()
                teacher_sample = [pool_rows[i] for i in idxs]
                # Fail-fast tokenization drift check on first use: cached prompt_ids
                # must match live tokenization at the prompt position. If this trips
                # the pool was generated with a different tokenizer / chat template.
                if not pool_validated:
                    cached_ids = teacher_sample[0]["prompt_ids"][: int(teacher_sample[0]["plen"])]
                    live_ids = enc.input_ids[0].tolist()
                    if cached_ids != live_ids:
                        raise ValueError(
                            f"teacher pool tokenization drift on problem_id={prob['problem_id']}: "
                            f"cached prompt_ids[:plen]={cached_ids[:12]}... vs "
                            f"live enc={live_ids[:12]}... (lengths {len(cached_ids)} vs {len(live_ids)})"
                        )
                    pool_validated = True
                # Student live-gen. gen_cfg.num_return_sequences is baked to G_s
                # at construction (pool path) or = group (no-pool path).
                with torch.no_grad():
                    out_s = model.generate(**enc, generation_config=gen_cfg).detach()
                # Build teacher tensor: each cached row is plen + L_t_i; right-pad
                # to common L within the teacher batch, then F.pad to match student L.
                teacher_seqs = [
                    torch.tensor(r["prompt_ids"] + r["completion_ids"], dtype=torch.long, device=device)
                    for r in teacher_sample
                ]
                L_t = max(s.shape[0] for s in teacher_seqs)
                out_t = torch.stack([F.pad(s, (0, L_t - s.shape[0]), value=pad_id) for s in teacher_seqs])
                L = max(out_s.shape[1], out_t.shape[1])
                if out_s.shape[1] < L:
                    out_s = F.pad(out_s, (0, L - out_s.shape[1]), value=pad_id)
                if out_t.shape[1] < L:
                    out_t = F.pad(out_t, (0, L - out_t.shape[1]), value=pad_id)
                gen_out = torch.cat([out_s, out_t], dim=0)
                is_student = [True] * G_s + [False] * G_t
            else:
                with torch.no_grad():
                    gen_out = model.generate(**enc, generation_config=gen_cfg).detach()
                is_student = [True] * gen_out.shape[0]
            model.config.use_cache = False
            merged = gen_out
            completions = gen_out[:, plen:]
            texts = tok.batch_decode(completions, skip_special_tokens=True)
            t_gen += time.perf_counter() - _tg

            # First-batch full dump (system msg + user msg + rendered prompt + completion
            # with special tokens). Goes to verbose log only — stdout stays clean.
            # Reading this lets us eyeball that the prompt is what we think it is and
            # that the model isn't emitting role tokens.
            if step == 0 and p_idx == 0:
                comp_with_special = tok.decode(completions[0], skip_special_tokens=False)
                sys_msg = next((m["content"] for m in prob["messages"] if m.get("role") == "system"), "<no system>")
                user_msg = next((m["content"] for m in prob["messages"] if m.get("role") == "user"), "<no user>")
                logger.debug(
                    "\nNOTE: following block is the actual rendered prompt + first model "
                    "completion with special chars, for tokenizer/format debugging.\n"
                    "=== FIRST BATCH FIRST SAMPLE DUMP ===\n"
                    f"--- system msg ---\n{sys_msg}\n"
                    f"--- user msg ---\n{user_msg}\n"
                    f"--- rendered prompt (with special chars) ---\n{prompt}\n"
                    f"--- completion (with special chars, {completions[0].numel()} tokens) ---\n{comp_with_special}\n"
                    "=== END FIRST BATCH DUMP ==="
                )

            comp_lens = [int((c != pad_id).sum().item()) for c in completions]
            finished = [bool((c == eos_id).any().item()) for c in completions]
            agg_comp_lens.extend(comp_lens); agg_finished.extend(finished)

            _tr = time.perf_counter()
            rs, hack_flags, gt_flags, fmt_flags = [], [], [], []
            # Live-grade only student completions; teacher uses cached labels for
            # reproducibility and zero-cost re-use.
            n_live_grade = G_s if teacher_pool else len(texts)
            for t in texts[:n_live_grade]:
                r = compute_reward(
                    t, canonical_solution=prob["canonical"], gt_tests=prob["gt_tests"],
                    setup_code=prob["setup_code"], func_name_hint=prob["func_name"],
                )
                rs.append(r.reward); hack_flags.append(r.hacked); gt_flags.append(r.gt_pass)
                fmt_flags.append(r.format_ok)
            if teacher_sample is not None:
                for r in teacher_sample:
                    rs.append(float(r["reward"])); hack_flags.append(bool(r["hacked"]))
                    gt_flags.append(bool(r["gt_pass"])); fmt_flags.append(bool(r["fmt_ok"]))
            t_rew += time.perf_counter() - _tr
            agg_rew.extend(rs); agg_gt.extend(gt_flags); agg_hack.extend(hack_flags); agg_fmt.extend(fmt_flags)
            agg_is_student.extend(is_student)

            if (step < 3 or step % 20 == 0) and p_idx == 0:
                # Capture diagnostic tail of one generation per step. Look for
                # mid-statement truncation (no closing ```), <think> traces, etc.
                diag_tail = texts[0][-400:]

            rewards = torch.tensor(rs, dtype=torch.float32, device=device)
            # simple_GRPO grpo_vllm_one.py:208: skip groups where every generation
            # got the same reward. Dr.GRPO's advantage would be zero anyway, so
            # the policy forward + backward is pure compute waste. This is the
            # dominant pathology with our binary-ish reward shape on a weak 2B
            # substrate (every group can clip to 0.25 = format_only).
            if (rewards.max() - rewards.min()).item() < 1e-4:
                # Pad agg_logp with NaN to keep it aligned with agg_is_student
                # (extended above at line 770). Skipping the gen_logp forward
                # here is the whole point of the zero-variance bail.
                agg_logp.extend([float("nan")] * len(rs))
                continue
            centered = rewards - rewards.mean()
            adv = centered if cfg.unbiased else centered / (rewards.std() + 1e-4)

            # Old-policy logprobs (frozen target for PPO ratio). Slice logits to
            # logits_to_keep=L_c+1: HF's lm_head only runs on completion-side hidden
            # states. Avoids materializing prompt-side logits (~plen/(plen+L_c) saved
            # at lm_head). Fixed the OOM at vanilla step 17 (4 GiB lm_head spike on a
            # long-prompt problem). Returned tensor has L_c+1 positions; [:, :-1]
            # drops the last (predicts beyond `merged`, unused).
            completion_ids = merged[:, plen:]
            L_c = completion_ids.shape[1]
            _tfb = time.perf_counter()
            with torch.no_grad():
                gen_logp = per_token_logps(
                    model(merged, logits_to_keep=L_c + 1).logits[:, :-1],
                    completion_ids,
                ).detach()

            ref_logp = None
            if beta and beta > 0:
                ref_logp = ref_logprobs_via_zero_delta(model, merged, wrappers, plen).detach()

            pol_logp = per_token_logps(
                model(merged, logits_to_keep=L_c + 1).logits[:, :-1],
                completion_ids,
            )

            mask = (merged[:, plen:] != pad_id).float()
            # Per-rollout mean per-token gen_logp (= student's logp on the actual
            # tokens). In single-step PPO, gen_logp == pol_logp.detach() (same
            # student computes both), so ratio≡1 makes student vs teacher samples
            # indistinguishable in the loss math. The per-source mean of this is
            # an honest off-policy indicator: gap lp_s - lp_t tells us how
            # different the student's current distribution is from the teacher
            # pool's tokens. No IS correction is applied; this is diagnostic only.
            mean_logp_per_rollout = ((gen_logp * mask).sum(1) / mask.sum(1).clamp_min(1)).detach().cpu().tolist()
            agg_logp.extend(mean_logp_per_rollout)
            ratio = torch.exp(pol_logp - gen_logp)
            clipped = torch.clamp(ratio, 1 - cfg.clip, 1 + cfg.clip)
            pol_term = torch.min(ratio * adv.unsqueeze(1), clipped * adv.unsqueeze(1))
            per_tok_loss = -pol_term
            if ref_logp is not None:
                kl = torch.exp(ref_logp - pol_logp) - (ref_logp - pol_logp) - 1.0
                per_tok_loss = per_tok_loss + beta * kl

            # Per-source split (loss_s + loss_t == full-batch loss because
            # is_s_v + is_t_v = 1 elementwise; backward is linear so
            # grad_s + grad_t == full-batch grad). Two backwards every step is
            # ~2x backward cost — gated to every cin_split_every step.
            is_s_v = torch.tensor(is_student, dtype=per_tok_loss.dtype,
                                  device=per_tok_loss.device).unsqueeze(1)  # [G, 1]
            is_t_v = 1.0 - is_s_v
            if split_this_step:
                if cfg.unbiased:
                    denom = group * max_new * prompts_per_step
                    loss_s = (per_tok_loss * mask * is_s_v).sum() / denom
                    loss_t = (per_tok_loss * mask * is_t_v).sum() / denom
                else:
                    ptl_norm = (per_tok_loss * mask).sum(1) / mask.sum(1).clamp_min(1)
                    loss_s = (ptl_norm * is_s_v.squeeze(1)).sum() / (group * prompts_per_step)
                    loss_t = (ptl_norm * is_t_v.squeeze(1)).sum() / (group * prompts_per_step)
                # Pass 1: student. retain_graph so the shared forward graph survives.
                loss_s.backward(retain_graph=True)
                for name, info in wrappers.items():
                    gs = info["delta_S"].grad
                    if gs is None:
                        continue
                    step_grad_s[name] = (step_grad_s[name] + gs.detach().clone()
                                         if name in step_grad_s
                                         else gs.detach().clone())
                model.zero_grad(set_to_none=True)
                # Pass 2: teacher.
                loss_t.backward()
                for name, info in wrappers.items():
                    gt = info["delta_S"].grad
                    if gt is None:
                        continue
                    step_grad_t[name] = (step_grad_t[name] + gt.detach().clone()
                                         if name in step_grad_t
                                         else gt.detach().clone())
                model.zero_grad(set_to_none=True)
                agg_loss += (loss_s + loss_t).item()
            else:
                # Combined single backward — cheaper, no per-source diagnostic.
                # Accumulate into step_grad_s as the "combined" carrier; the
                # injection block below treats step_grad_t == {} as "use gs".
                if cfg.unbiased:
                    denom = group * max_new * prompts_per_step
                    loss = (per_tok_loss * mask).sum() / denom
                else:
                    ptl_norm = (per_tok_loss * mask).sum(1) / mask.sum(1).clamp_min(1)
                    loss = ptl_norm.sum() / (group * prompts_per_step)
                loss.backward()
                for name, info in wrappers.items():
                    g = info["delta_S"].grad
                    if g is None:
                        continue
                    step_grad_s[name] = (step_grad_s[name] + g.detach().clone()
                                         if name in step_grad_s
                                         else g.detach().clone())
                model.zero_grad(set_to_none=True)
                agg_loss += loss.item()
            t_fb += time.perf_counter() - _tfb

        # Inject combined grad (student + teacher) into leaf .grad before
        # projection + optimizer. Where only one source contributed for a
        # module, take that source's grad directly.
        for name, info in wrappers.items():
            gs = step_grad_s.get(name)
            gt = step_grad_t.get(name)
            if gs is None and gt is None:
                continue
            if gs is None:
                info["delta_S"].grad = gt
            elif gt is None:
                info["delta_S"].grad = gs
            else:
                info["delta_S"].grad = gs + gt

        # Per-source cin: project student-only and teacher-only grads into v_hack
        # subspace. Discriminator: cin_t > cin_s on a clean base means v_hack
        # lights up for hack grads more than non-hack. Only valid on split steps;
        # otherwise step_grad_s holds the combined grad and would mis-report cin_s.
        if split_this_step:
            cin_s = mean_cin_from_grads(step_grad_s, v_hack)
            cin_t = mean_cin_from_grads(step_grad_t, v_hack)
        else:
            cin_s = cin_t = float("nan")

        # Diagnostic cos_in for both arms; projection only mutates grad if arm=projected.
        diag = project_delta_S_grad(
            wrappers, v_hack, cfg.preserve_magnitude,
            measure_only=(cfg.arm != "projected"),
            gate_mode=cfg.gate_mode,
        )
        diag["mean_cin_s"] = cin_s
        diag["mean_cin_t"] = cin_t

        torch.nn.utils.clip_grad_norm_(delta_params, 1.0)
        opt.step()
        sched.step()

        rewards_t = torch.tensor(agg_rew, dtype=torch.float32) if agg_rew else torch.zeros(1)
        rew_mean = rewards_t.mean().item()
        rew_std = rewards_t.std().item() if rewards_t.numel() > 1 else 0.0
        spread = (rewards_t.max() - rewards_t.min()).item() > 1e-3 if rewards_t.numel() > 1 else False
        n_rollouts = len(agg_rew)

        # Per-source breakdown: which rollouts came from student vs teacher this step.
        # Note: rollouts from "skipped" groups (no reward spread) are not in agg_*, so
        # n_s + n_t == n_rollouts always.
        is_s = torch.tensor(agg_is_student, dtype=torch.bool) if agg_is_student else torch.zeros(0, dtype=torch.bool)
        h_t  = torch.tensor(agg_hack,       dtype=torch.bool) if agg_hack else torch.zeros(0, dtype=torch.bool)
        g_t  = torch.tensor(agg_gt,         dtype=torch.bool) if agg_gt else torch.zeros(0, dtype=torch.bool)
        n_s  = int(is_s.sum())
        n_t  = int(is_s.numel() - n_s)
        hack_s_n = int((h_t & is_s).sum())
        hack_t_n = int((h_t & ~is_s).sum())
        gt_s_n   = int((g_t & is_s).sum())
        gt_t_n   = int((g_t & ~is_s).sum())
        rew_s_mean = rewards_t[is_s].mean().item() if n_s else float("nan")
        # Skipped (zero-variance) prompts pad agg_logp with NaN above to keep
        # alignment with is_s. nanmean drops them from the per-source means.
        logp_t = torch.tensor(agg_logp, dtype=torch.float32) if agg_logp else torch.zeros(0)
        lp_s_mean = logp_t[is_s].nanmean().item()  if n_s else float("nan")
        lp_t_mean = logp_t[~is_s].nanmean().item() if n_t else float("nan")

        # Per-step diagnostics → verbose log; stdout sees tqdm postfix + final table.
        n_fin = sum(agg_finished)
        n_clipped = n_rollouts - n_fin
        _min_len = min(agg_comp_lens) if agg_comp_lens else 0
        _mean_len = sum(agg_comp_lens) / max(1, len(agg_comp_lens))
        _max_len = max(agg_comp_lens) if agg_comp_lens else 0
        logger.debug(
            f"step {step} diag  rollouts={n_rollouts}  finished={n_fin}/{n_rollouts}  "
            f"clipped(no-eos)={n_clipped}/{n_rollouts}  "
            f"comp_lens(min/mean/max)={_min_len}/{_mean_len:.0f}/{_max_len}  "
            f"max_new={max_new}  fmt={sum(agg_fmt)}/{n_rollouts}  gt={sum(agg_gt)}/{n_rollouts}  "
            f"hack={sum(agg_hack)}/{n_rollouts}  skipped={n_skipped}/{prompts_per_step}"
        )
        _tstep = time.time() - t0
        logger.debug(
            f"step {step} TIMING  gen={t_gen:.0f}s  fwd_bwd={t_fb:.0f}s  "
            f"reward={t_rew:.0f}s  other={_tstep - t_gen - t_fb - t_rew:.0f}s  "
            f"total={_tstep:.0f}s"
        )
        if diag_tail is not None:
            tail = diag_tail.replace("\n", "\\n")
            logger.debug(f"step {step} gen[0] tail (last 400 chars): {tail!r}")

        cum_gens = sum(r["N"] for r in rows) + n_rollouts
        row = {
            "step": step,
            "ref_eq": f"{cum_gens / REF_GENS_PER_STEP:.2f}",
            "rew": f"{rew_mean:+.2f}",
            "rew_s": f"{rew_s_mean:+.2f}" if n_s else "nan",
            "sprd": "T" if spread else "F",
            "N": n_rollouts,
            "gt_s": f"{gt_s_n}/{n_s}" if n_s else "0/0",
            "gt_t": f"{gt_t_n}/{n_t}" if n_t else "0/0",
            "hack_s": f"{hack_s_n}/{n_s}" if n_s else "0/0",
            "hack_t": f"{hack_t_n}/{n_t}" if n_t else "0/0",
            "lp_s": f"{lp_s_mean:+.2f}" if n_s else "nan",
            "lp_t": f"{lp_t_mean:+.2f}" if n_t else "nan",
            "loss": f"{agg_loss:+.4f}",
            "cin": f"{diag['mean_cos_in']:+.3f}",
            "cin_s": f"{diag['mean_cin_s']:+.3f}",
            "cin_t": f"{diag['mean_cin_t']:+.3f}",
            "cout": f"{diag['mean_cos_out']:+.3f}",
            "fired": f"{diag['frac_fired']:.2f}",
            "gen": f"{t_gen:.0f}",
            "fb": f"{t_fb:.0f}",
            "t_rew": f"{t_rew:.0f}",
            "sec": f"{time.time()-t0:.0f}",
        }
        rows.append(row)
        # Stream this step as a row (header was printed before the loop).
        logger.info(_fmt_row(row))
        if (step + 1) % 25 == 0:
            save_ckpt(rows)  # survive early kills; ~12 days for the full sweep
        if not first_hack_saved and hack_s_n > 0:
            save_ckpt(rows, path=first_hack_path)
            first_hack_saved = True
            logger.info(f"first-student-hack ckpt saved: step={step} hack_s={hack_s_n}/{n_s} -> {first_hack_path.name}")
        # Live status in tqdm postfix; full per-step line in verbose log only.
        pbar.set_postfix(
            rew=f"{rew_mean:+.2f}", gt=f"{sum(agg_gt)}/{n_rollouts}",
            hack=f"{sum(agg_hack)}/{n_rollouts}", loss=f"{agg_loss:+.3f}",
            sec=f"{time.time()-t0:.0f}",
        )
        logger.debug(
            f"step {step:3d}  rew={rew_mean:+.2f}(std {rew_std:.2f})  "
            f"gt={sum(agg_gt)}/{n_rollouts}  hack={sum(agg_hack)}/{n_rollouts}  "
            f"loss={agg_loss:+.3f}  cos_in={diag['mean_cos_in']:+.3f} "
            f"cos_out={diag['mean_cos_out']:+.3f}  fired={diag['frac_fired']:.2f}  "
            f"sec={time.time()-t0:.0f}"
        )

    peak_gb = torch.cuda.max_memory_allocated() / 1e9 if torch.cuda.is_available() else 0.0
    n_steps = len(rows)
    n_gens = sum(r["N"] for r in rows)
    total_hacks = sum(int(r["hack_s"].split("/")[0]) + int(r["hack_t"].split("/")[0]) for r in rows)
    total_pass = sum(int(r["gt_s"].split("/")[0]) + int(r["gt_t"].split("/")[0]) for r in rows)
    hack_rate = total_hacks / max(1, n_gens)
    pass_rate = total_pass / max(1, n_gens)
    # Per-source totals. On no-teacher runs, hack_s_total == total_hacks.
    hack_s_total = sum(int(r["hack_s"].split("/")[0]) for r in rows)
    hack_t_total = sum(int(r["hack_t"].split("/")[0]) for r in rows)
    n_s_total = sum(int(r["hack_s"].split("/")[1]) for r in rows)
    n_t_total = sum(int(r["hack_t"].split("/")[1]) for r in rows)
    hack_rate_s = hack_s_total / max(1, n_s_total)
    hack_rate_t = hack_t_total / max(1, n_t_total)

    # Final tail: cue emoji + main metric BLUF, then per-step tsv table.
    # Vanilla arm: 🟢 if hacking emerged. Projected arm: 🟢 if HACK_RATE dropped
    # vs a matched-PASS vanilla — we can't judge that here, so just report.
    cue = "🟢" if (cfg.arm == "vanilla" and hack_rate > 0.0) else "🟡"

    print(f"\nargv: {' '.join(sys.argv)}")
    print(f"verbose log: {verbose_log}")
    print(
        f"main metric: HACK_RATE={hack_rate:.3f} PASS_RATE={pass_rate:.3f}  "
        f"HACK_STUDENT={hack_rate_s:.3f} HACK_TEACHER={hack_rate_t:.3f}  "
        f"[arm={cfg.arm} preset={cfg.preset.value} model={model_name} steps={n_steps} gens={n_gens} peak={peak_gb:.1f}GB"
        f"{' pool=' + cfg.teacher_pool_dir.name + ' mix=' + str(cfg.mix_ratio) if cfg.teacher_pool_dir else ''}]"
    )
    print()
    print(tabulate(rows, headers="keys", tablefmt="tsv", floatfmt="+.3f"))
    print()
    print(tabulate([{
        "cue": cue, "HACK_RATE": f"{hack_rate:.3f}", "PASS_RATE": f"{pass_rate:.3f}",
        "HACK_S": f"{hack_rate_s:.3f}", "HACK_T": f"{hack_rate_t:.3f}",
        "peak_GB": f"{peak_gb:.1f}", "arm": cfg.arm, "preset": cfg.preset.value,
        "model": model_name.split("/")[-1], "seed": cfg.seed, "steps": n_steps,
        "pool": (cfg.teacher_pool_dir.name if cfg.teacher_pool_dir else ""),
        "mix": cfg.mix_ratio if cfg.teacher_pool_dir else "",
        "tag": cfg.out_tag, "log": str(verbose_log),
    }], headers="keys", tablefmt="tsv"))
    # Markdown copy: easier to paste into journal/PRs than the TSV above.
    print()
    print("### Per-step rows (markdown)\n")
    print(tabulate(rows, headers="keys", tablefmt="pipe", floatfmt="+.3f"))

    save_ckpt(rows)
    return 0


if __name__ == "__main__":
    sys.exit(main(tyro.cli(Config)))