Add src

src/analyze_pool.py

@@ -0,0 +1,235 @@
+"""
+Characterize the base policy from the generation pool.
+
+The pool is 1000 prompts x 16 samples = 16,000 stories, which is 20x the
+held-out eval set. It is the most statistically solid picture of baseline mode
+collapse in the whole study, so it gets its own report rather than being used
+only as DPO feed.
+
+Answers:
+  - How collapsed is the base model, per prompt? (deviation, log-det, eff. rank)
+  - Is collapse uniform, or are some prompts far worse than others?
+  - Does the judge's quality correlate with diversity? (i.e. is there really a
+    quality-diversity tension to trade off, or are they independent?)
+  - What does the quality distribution look like -- does tau=5 / rho=6 bite?
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from collections import defaultdict
+from pathlib import Path
+
+import numpy as np
+
+ROOT = Path(__file__).resolve().parent.parent
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--tag", default="4b")
+    ap.add_argument("--split", default="train")
+    args = ap.parse_args()
+
+    import logbook
+    from diversity import effective_rank, logdet_volume, pairwise_deviation
+
+    pdir = ROOT / "outputs" / f"pool_{args.tag}"
+    rows = [json.loads(l) for l in open(pdir / f"pool_{args.split}.jsonl") if l.strip()]
+    E = np.load(pdir / f"emb_{args.split}.npy").astype(np.float64)
+    summary = json.load(open(pdir / f"summary_{args.split}.json"))
+
+    by = defaultdict(list)
+    for i, r in enumerate(rows):
+        by[r["prompt_id"]].append(i)
+
+    per = []
+    for pid, ids in by.items():
+        sub = E[ids]
+        q = np.array([rows[i]["quality"] for i in ids])
+        gp = np.array([rows[i]["gate_passed"] for i in ids])
+        per.append({
+            "prompt_id": pid,
+            "dev": float(pairwise_deviation(sub).mean()),
+            "logdet": float(logdet_volume(sub)),
+            "eff_rank": float(effective_rank(sub)),
+            "quality": float(q[gp].mean()) if gp.any() else 0.0,
+            "gate_pass": float(gp.mean()),
+            "n": len(ids),
+        })
+
+    dev = np.array([p["dev"] for p in per])
+    ld = np.array([p["logdet"] for p in per])
+    er = np.array([p["eff_rank"] for p in per])
+    ql = np.array([p["quality"] for p in per])
+    qual_all = np.array([r["quality"] for r in rows if r["gate_passed"]])
+    N = per[0]["n"]
+
+    # quality-diversity correlation across prompts
+    def corr(a, b):
+        if a.std() < 1e-9 or b.std() < 1e-9:
+            return 0.0
+        return float(np.corrcoef(a, b)[0, 1])
+
+    # --- how much signal do the E1 / E2 channels actually carry? -----------
+    # GDPO z-scores each reward channel WITHIN its group, so a channel with tiny
+    # within-group spread gets amplified to unit variance regardless. If the
+    # within-group ordering of d_i is not meaningful, E1 is largely learning
+    # from amplified noise. Compare within-group spread against between-group
+    # spread: a ratio far below 1 means the channel mostly encodes "which prompt
+    # is this", which per-group normalization deliberately removes.
+    dev_within, marg_within = [], []
+    for pid, ids in by.items():
+        sub = E[ids]
+        d = pairwise_deviation(sub)
+        from diversity import marginal_contributions
+        m = marginal_contributions(sub)
+        dev_within.append(d.std())
+        marg_within.append(m.std())
+    dev_within = np.array(dev_within); marg_within = np.array(marg_within)
+
+    stats = {
+        "n_prompts": len(per), "n_per_prompt": N, "n_stories": len(rows),
+        "dev_within_group_sd": float(dev_within.mean()),
+        "dev_between_group_sd": float(dev.std()),
+        "dev_within_over_between": float(dev_within.mean() / max(dev.std(), 1e-9)),
+        "marginal_within_group_sd": float(marg_within.mean()),
+        "gate_pass_rate": summary["gate_pass_rate"],
+        "ends_cleanly_rate": summary["ends_cleanly_rate"],
+        "quality_mean": float(qual_all.mean()), "quality_sd": float(qual_all.std()),
+        "quality_p10": float(np.percentile(qual_all, 10)),
+        "quality_median": float(np.median(qual_all)),
+        "quality_p90": float(np.percentile(qual_all, 90)),
+        "frac_quality_ge_5(tau)": float((qual_all >= 5).mean()),
+        "frac_quality_ge_6": float((qual_all >= 6).mean()),
+        "frac_quality_ge_7(rho)": float((qual_all >= 7).mean()),
+        "deviation_mean": float(dev.mean()), "deviation_sd": float(dev.std()),
+        "deviation_p10": float(np.percentile(dev, 10)),
+        "deviation_p90": float(np.percentile(dev, 90)),
+        "logdet_mean": float(ld.mean()), "logdet_sd": float(ld.std()),
+        "eff_rank_mean": float(er.mean()), "eff_rank_sd": float(er.std()),
+        "eff_rank_p10": float(np.percentile(er, 10)),
+        "eff_rank_p90": float(np.percentile(er, 90)),
+        "eff_rank_ceiling": N,
+        "corr(quality, deviation)": corr(ql, dev),
+        "corr(quality, eff_rank)": corr(ql, er),
+        "corr(deviation, eff_rank)": corr(dev, er),
+    }
+
+    # ---- figures ---------------------------------------------------------
+    import matplotlib
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+
+    fig, ax = plt.subplots(1, 4, figsize=(19, 4.3))
+    ax[0].hist(qual_all, bins=np.arange(-0.25, 10.75, 0.5), color="#2980b9",
+               edgecolor="white")
+    ax[0].axvline(5, c="crimson", ls="--", label="tau=5 (diversity gate)")
+    ax[0].axvline(7, c="darkorange", ls="--", label="rho=7 (DivPO, swept)")
+    ax[0].set_title("Judge quality, all gate-passing stories")
+    ax[0].set_xlabel("quality"); ax[0].legend(fontsize=7); ax[0].grid(alpha=.3)
+
+    ax[1].hist(er, bins=30, color="#8e44ad", edgecolor="white")
+    ax[1].axvline(N, c="green", ls="--", label=f"ceiling = {N}")
+    ax[1].set_title(f"Effective rank per prompt\n(1 = total collapse, {N} = orthogonal)")
+    ax[1].set_xlabel("effective rank"); ax[1].legend(fontsize=7); ax[1].grid(alpha=.3)
+
+    ax[2].hist(dev, bins=30, color="#16a085", edgecolor="white")
+    ax[2].set_title("Mean pairwise distance per prompt")
+    ax[2].set_xlabel("1 - cos"); ax[2].grid(alpha=.3)
+
+    ax[3].scatter(er, ql, s=6, alpha=.35, color="#c0392b")
+    ax[3].set_xlabel("effective rank"); ax[3].set_ylabel("mean judge quality")
+    ax[3].set_title(f"Quality vs diversity across prompts\nr = {stats['corr(quality, eff_rank)']:+.3f}")
+    ax[3].grid(alpha=.3)
+
+    plt.tight_layout()
+    figp = logbook.FIGS / f"03_pool_{args.tag}_baseline.png"
+    plt.savefig(figp, dpi=140)
+    plt.close()
+
+    worst = sorted(per, key=lambda p: p["eff_rank"])[:5]
+    best = sorted(per, key=lambda p: -p["eff_rank"])[:5]
+
+    r_qd = stats["corr(quality, eff_rank)"]
+    tension = ("a genuine quality-diversity **tension**" if r_qd < -0.15 else
+               "quality and diversity are **largely independent**" if abs(r_qd) <= 0.15 else
+               "quality and diversity are **positively** related")
+
+    body = f"""# Baseline characterization — {args.tag} pool ({len(rows):,} stories)
+
+The pool is {stats['n_prompts']} prompts x {N} samples from the **base policy**,
+20x the held-out eval set. This is the most statistically solid picture of
+baseline mode collapse in the study, so it is reported in its own right rather
+than treated only as DPO feed.
+
+## Summary
+
+{logbook.table([{"metric": k, "value": v} for k, v in stats.items()])}
+
+## Findings
+
+**Baseline collapse is severe.** Mean effective rank is
+**{stats['eff_rank_mean']:.2f} out of a ceiling of {N}** — the {N} samples for a
+given prompt span only ~{stats['eff_rank_mean']:.1f} effective directions. Mean
+pairwise distance is {stats['deviation_mean']:.3f}, i.e. same-prompt stories sit
+at ~{1-stats['deviation_mean']:.2f} cosine similarity. This is the thing every
+arm is trying to move.
+
+**Collapse is not uniform across prompts.** Effective rank runs from
+{stats['eff_rank_p10']:.2f} (p10) to {stats['eff_rank_p90']:.2f} (p90), so some
+prompts admit far more variation than others. Per-group normalization (GDPO)
+handles this correctly: each prompt's advantage is computed within its own
+group, so an intrinsically constrained prompt does not drag the update.
+
+**Quality vs diversity: r = {r_qd:+.3f}** across prompts, so {tension}. This
+matters for reading the frontier: if the correlation is near zero, then a method
+that raises diversity without lowering quality is not defying a tradeoff, it is
+exploiting slack that was already there.
+
+**How much signal does E1's channel carry?** Within-group SD of `d_i` is
+{stats['dev_within_group_sd']:.4f} against a between-group SD of
+{stats['dev_between_group_sd']:.4f} — a ratio of
+**{stats['dev_within_over_between']:.2f}**. This matters because GDPO z-scores
+each channel *within* its group, so whatever within-group spread exists is
+amplified to unit variance. A low ratio means most of `d_i`'s variation encodes
+*which prompt this is* rather than *which sample is the odd one out* — and
+per-group normalization deliberately discards exactly the former. Read E1's
+result with this number in mind: if E1 underperforms, weak within-group
+resolution is the first hypothesis, not a refutation of pairwise diversity as an
+idea. The marginal channel's within-group SD is
+{stats['marginal_within_group_sd']:.4f} on the raw log scale (it is z-scored
+before use, so only its ordering matters).
+
+**Threshold placement.** {100*stats['frac_quality_ge_5(tau)']:.1f}% of
+gate-passing stories score >= tau=5 and
+{100*stats['frac_quality_ge_7(rho)']:.1f}% score >= rho=7 (swept up from 6 because only {100*(1-stats['frac_quality_ge_6']):.1f}% fall below 6, starving DivPO of negatives on 311/1000 prompts). tau therefore acts as
+a *floor* that only bites when a story is genuinely bad, which is its intent
+(anti-gaming, not selection). rho is the more selective threshold and determines
+DivPO's skip rate.
+
+## Most collapsed prompts
+
+{logbook.table(worst, ["prompt_id", "eff_rank", "dev", "logdet", "quality"])}
+
+## Most diverse prompts
+
+{logbook.table(best, ["prompt_id", "eff_rank", "dev", "logdet", "quality"])}
+
+![pool baseline](../figures/03_pool_{args.tag}_baseline.png)
+"""
+    p = logbook.write_report(f"03_pool_{args.tag}_baseline", body)
+    print(json.dumps(stats, indent=1))
+    print("report ->", p, "\nfigure ->", figp)
+    json.dump({"stats": stats, "per_prompt": per},
+              open(logbook.LOGS / f"pool_{args.tag}_analysis.json", "w"), indent=1)
+    logbook.note(f"pool analysis ({args.tag})",
+                 f"eff_rank {stats['eff_rank_mean']:.2f}/{N}, "
+                 f"dev {stats['deviation_mean']:.3f}, "
+                 f"corr(q,div) {r_qd:+.3f}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

src/analyze_prior_run.py

@@ -0,0 +1,191 @@
+"""
+Autopsy of the previous run (logs/experiments/00_prior_run_autopsy.md + figures).
+
+Reads prior_run/eval_samples/step_*.json, which are {prompt: [16 stories]} dumps
+taken every 50 steps. Recomputes gates + embedding diversity so the prior run is
+measured with EXACTLY the same instruments as the new arms -- otherwise the
+"before" number is not comparable to any "after" number.
+"""
+from __future__ import annotations
+
+import glob
+import json
+import re
+from collections import Counter
+from pathlib import Path
+
+import numpy as np
+
+import gates
+import logbook
+from diversity import group_metrics, l2_normalize
+
+PRIOR = Path(__file__).resolve().parent.parent / "prior_run" / "eval_samples"
+EMB_MODEL = "BAAI/bge-base-en-v1.5"
+
+
+def load_steps() -> dict[int, dict[str, list[str]]]:
+    out = {}
+    for f in sorted(glob.glob(str(PRIOR / "step_*.json")),
+                    key=lambda x: int(re.search(r"\d+", Path(x).name).group())):
+        step = int(re.search(r"\d+", Path(f).name).group())
+        out[step] = json.load(open(f))
+    return out
+
+
+def main():
+    steps = load_steps()
+    print(f"loaded {len(steps)} checkpoints: {sorted(steps)}")
+
+    from sentence_transformers import SentenceTransformer
+    enc = SentenceTransformer(EMB_MODEL, device="cuda")
+
+    rows = []
+    reason_counter: dict[int, Counter] = {}
+    for step, data in steps.items():
+        allres, pw, ld, mm = [], [], [], []
+        for prompt, samples in data.items():
+            res = [gates.check(s) for s in samples]
+            allres += res
+            E = l2_normalize(enc.encode(samples, normalize_embeddings=True,
+                                        batch_size=32, show_progress_bar=False))
+            gm = group_metrics(E)
+            pw.append(gm["mean_pairwise_dist"]); ld.append(gm["logdet"])
+            mm.append(gm["mean_marginal"])
+        n = len(allres)
+        reason_counter[step] = Counter(r for x in allres for r in x.reasons)
+        rows.append({
+            "step": step, "n": n,
+            "gate_pass_%": 100 * sum(r.passed for r in allres) / n,
+            "ends_cleanly_%": 100 * sum(r.completeness for r in allres) / n,
+            "med_words": float(np.median([r.n_words for r in allres])),
+            "rep4gram": float(np.mean([r.repeat_4gram_frac for r in allres])),
+            "pairwise_dist": float(np.mean(pw)),
+            "logdet": float(np.mean(ld)),
+            "mean_marginal": float(np.mean(mm)),
+        })
+
+    for r in rows:
+        print(r)
+
+    # ---------------------------------------------------------------- figures
+    import matplotlib
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+
+    s = [r["step"] for r in rows]
+    fig, ax = plt.subplots(1, 3, figsize=(15, 4.2))
+
+    ax[0].plot(s, [r["ends_cleanly_%"] for r in rows], "o-", color="#c0392b", lw=2)
+    ax[0].axhline(80, ls="--", c="gray", lw=1)
+    ax[0].set_title("Completeness collapse\n(prior run)")
+    ax[0].set_xlabel("step"); ax[0].set_ylabel("% ending cleanly")
+    ax[0].set_ylim(-3, 100); ax[0].grid(alpha=.3)
+    ax[0].text(60, 84, "healthy target", color="gray", fontsize=8)
+
+    ax[1].plot(s, [r["pairwise_dist"] for r in rows], "o-", label="pairwise dist", lw=2)
+    ax[1].set_xlabel("step"); ax[1].set_ylabel("mean pairwise distance", color="C0")
+    ax2 = ax[1].twinx()
+    ax2.plot(s, [r["logdet"] for r in rows], "s--", color="C1", label="logdet")
+    ax2.set_ylabel("logdet volume", color="C1")
+    ax[1].set_title("Semantic diversity\n(embedding-based)")
+    ax[1].grid(alpha=.3)
+
+    ax[2].plot(s, [r["med_words"] for r in rows], "o-", color="#2c3e50", lw=2)
+    ax[2].set_title("Median length pinned at the\ntoken wall (~380 words)")
+    ax[2].set_xlabel("step"); ax[2].set_ylabel("median words")
+    ax[2].grid(alpha=.3)
+
+    plt.tight_layout()
+    out = logbook.FIGS / "00_prior_run_autopsy.png"
+    out.parent.mkdir(parents=True, exist_ok=True)
+    plt.savefig(out, dpi=140)
+    print("figure ->", out)
+
+    # ---------------------------------------------------------------- report
+    body = f"""# Experiment 00 — Autopsy of the prior run
+
+**Purpose.** Establish what actually happened in the previous `train.py` run
+before building anything new, and re-measure it with the *same* instruments the
+new arms will use, so "before" and "after" are comparable.
+
+**Source.** `prior_run/eval_samples/step_{{50..300}}.json` — 5 held-out prompts x
+16 samples, dumped every 50 steps. 150 stories total per checkpoint set.
+
+## Measurements
+
+{logbook.table(rows)}
+
+## Failure-reason census (gate violations, counted per story)
+
+{logbook.table([{"step": k, **dict(v)} for k, v in reason_counter.items()])}
+
+## Findings
+
+**1. The run generated no finished stories, at any checkpoint.**
+`ends_cleanly` is 0-8% across all six checkpoints. Median length sits in a
+tight ~360-400 word band with a hard ceiling, which is the signature of a
+`max_new_tokens` wall rather than a stylistic preference. Sampled tails
+confirm mid-word cutoff (`"...a place called *The Library of"`).
+
+**2. That silently destroyed the quality reward.**
+The prior judge rubric said *"if the story cuts off mid-sentence, cap quality
+at 4."* When ~100% of a group hits the same cap, within-group quality variance
+goes to ~0. GRPO normalizes advantages within the group, so an
+(almost) constant reward column produces an (almost) zero advantage. The
+quality objective contributed no gradient for the entire run.
+
+**3. The group-diversity term was structurally incapable of contributing.**
+`quality_reward` added `0.5 * group_diversity`, a single scalar identical for
+every sample in the group. Its within-group std is exactly 0, so after
+normalization it contributes exactly 0 to every advantage. It cost one judge
+call per step and bought nothing. (This is the failure mode recorded in the
+brief's section 0.1, confirmed here in situ.)
+
+**4. Net effect: the run optimized novelty almost alone.**
+With quality flat and group-diversity annihilated, the only surviving gradient
+came from `novelty_reward` — and its quality gate (`q >= 5`) was itself
+degenerate, because quality was pinned near the truncation cap. That is an
+unconditioned novelty objective, which is precisely the configuration known to
+produce surface-level gaming.
+
+**5. The observable signature matches that diagnosis.**
+Unique 5-word openers rise 0.68 -> 0.96 while type-token ratio *falls*
+0.381 -> 0.343. The model learned to vary the first few words while its
+vocabulary narrowed — diversification at the surface, homogenization
+underneath. Heavy italic emphasis persists at ~7 instances/story throughout,
+the exact tic the rubric explicitly tried to punish.
+
+**6. A latent bug would have suppressed completeness credit even after a fix.**
+`completeness_reward` tested `text[-1] in ".!?\\"'"` — straight quotes only.
+These generations use curly U+201D constantly, so a story ending `...gone.”`
+scored 0.3 instead of 1.0.
+
+## Consequences for the new design
+
+- Token budget raised to 1280 new tokens (a 500-word story is ~650-700 Qwen
+  tokens; the planned 768 left no headroom and would have re-created the bug).
+  Nothing truncates the story anywhere downstream — not the judge, not the
+  embedder. Overrun is *detected* by a word-count gate, never chopped.
+- System prompt now carries an explicit word budget and ending instruction.
+- No set-level scalar is ever used as a per-sample reward. All diversity credit
+  is per-sample by construction (deviation `d_i`, marginal `m_i`).
+- `finish_reason == "length"` from vLLM is captured and gates on it directly,
+  rather than inferring truncation from punctuation.
+- Terminal-punctuation detection handles curly quotes, ellipses and trailing
+  markdown emphasis.
+- Judge scores one story at a time on an absolute rubric, so `tau` is
+  meaningful and cross-model eval comparisons are valid.
+
+![autopsy](../figures/00_prior_run_autopsy.png)
+"""
+    p = logbook.write_report("00_prior_run_autopsy", body)
+    print("report ->", p)
+    logbook.note("prior-run autopsy complete",
+                 f"0-8% completion across all checkpoints; quality signal was "
+                 f"degenerate. Report: {p}")
+    json.dump(rows, open(logbook.LOGS / "prior_run_metrics.json", "w"), indent=2)
+
+
+if __name__ == "__main__":
+    main()

src/analyze_run.py

@@ -0,0 +1,223 @@
+"""
+Post-run analysis for one GRPO arm. Produces logs/experiments/<name>.md plus
+figures, from outputs/<name>/reward_history.json.
+
+Run this after EVERY training run, before launching the next one. The point is
+to catch a dead or hacked reward channel while there is still time to change
+something, rather than discovering it in the final report -- which is exactly
+how the prior run wasted 300 steps.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+
+ROOT = Path(__file__).resolve().parent.parent
+
+
+def smooth(x, w=15):
+    x = np.asarray(x, dtype=np.float64)
+    if len(x) < w:
+        return x
+    k = np.ones(w) / w
+    return np.convolve(x, k, mode="valid")
+
+
+def trend(x, frac=0.25):
+    """(early mean, late mean, delta) over the first/last `frac` of the run."""
+    x = np.asarray(x, dtype=np.float64)
+    n = max(1, int(len(x) * frac))
+    a, b = float(x[:n].mean()), float(x[-n:].mean())
+    return a, b, b - a
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--name", required=True)
+    ap.add_argument("--title", default=None)
+    args = ap.parse_args()
+
+    import logbook
+
+    run_dir = ROOT / "outputs" / args.name
+    hist = json.load(open(run_dir / "reward_history.json"))
+    if not hist:
+        print("empty history"); return 1
+
+    steps = np.arange(len(hist))
+    series = {k: np.array([h[k] for h in hist], dtype=np.float64) for k in
+              ("gate_pass", "ends_cleanly", "mean_quality_passing", "mean_novelty",
+               "frac_above_tau", "mean_deviation", "mean_logdet", "mean_words",
+               "frac_groups_degenerate")}
+
+    tr = {k: trend(v) for k, v in series.items()}
+
+    # --- reward-hacking diagnostic ---------------------------------------
+    dev_a, dev_b, dev_d = tr["mean_deviation"]
+    q_a, q_b, q_d = tr["mean_quality_passing"]
+    g_a, g_b, g_d = tr["gate_pass"]
+    hack = (dev_d > 0.02 and q_d < -0.5) or (g_b < 0.6 and g_d < -0.2)
+    verdict = ("SUSPECTED REWARD HACKING" if hack else
+               "healthy" if (g_b > 0.7 and q_d > -0.5) else "degraded but not hacking")
+
+    # --- figures ----------------------------------------------------------
+    import matplotlib
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+
+    # TRL-side metrics (entropy / KL) live in a separate history
+    trl_path = run_dir / "trl_log_history.json"
+    trl = json.load(open(trl_path)) if trl_path.exists() else []
+    ent = np.array([h["entropy"] for h in trl if "entropy" in h], dtype=np.float64)
+    kl = np.array([h["kl"] for h in trl if "kl" in h], dtype=np.float64)
+
+    fig, ax = plt.subplots(3, 3, figsize=(16, 11.5))
+    panels = [
+        ("gate_pass", "Gate pass rate", "#c0392b", (0, 1.05)),
+        ("ends_cleanly", "Ends cleanly", "#27ae60", (0, 1.05)),
+        ("mean_quality_passing", "Judge quality (gate-passers)", "#2980b9", None),
+        ("mean_deviation", "Mean pairwise deviation d_i", "#8e44ad", None),
+        ("mean_logdet", "Group log-det volume", "#d35400", None),
+        ("frac_above_tau", "Frac above tau (diversity-eligible)", "#16a085", (0, 1.05)),
+        ("mean_words", "Mean story length (words)", "#7f8c8d", None),
+        ("frac_groups_degenerate", "Groups with no valid sample", "#c0392b", (0, 1.05)),
+    ]
+    axes = ax.ravel()
+    for a, (key, title, c, ylim) in zip(axes, panels):
+        v = series[key]
+        a.plot(steps, v, alpha=0.25, color=c, lw=0.8)
+        sm = smooth(v)
+        a.plot(steps[len(steps) - len(sm):], sm, color=c, lw=2)
+        a.set_title(title, fontsize=10)
+        a.set_xlabel("reward-batch"); a.grid(alpha=.3)
+        if ylim:
+            a.set_ylim(*ylim)
+
+    # entropy panel: the creativity-death detector. Falling entropy means the
+    # policy is becoming deterministic. Rising entropy is permissive, NOT proof
+    # of creativity -- a policy can get noisier inside a single narrative mode --
+    # so this is read alongside log-det, never as a substitute for it.
+    a = axes[8]
+    if ent.size:
+        x = np.arange(ent.size)
+        a.plot(x, ent, alpha=0.25, color="#e67e22", lw=0.8)
+        sm = smooth(ent)
+        a.plot(x[len(x) - len(sm):], sm, color="#e67e22", lw=2, label="entropy")
+        e0, e1, ed = trend(ent)
+        a.axhline(e0, ls=":", c="gray", lw=1)
+        a.set_title(f"Per-token policy entropy  ({e0:.3f} → {e1:.3f}, Δ{ed:+.3f})",
+                    fontsize=10)
+        if kl.size:
+            a2 = a.twinx()
+            a2.plot(np.arange(kl.size), kl, color="#95a5a6", lw=1, alpha=.7)
+            a2.set_ylabel("KL to ref", color="#95a5a6", fontsize=8)
+    else:
+        a.text(.5, .5, "no entropy logged\n(liger path?)", ha="center",
+               va="center", transform=a.transAxes, color="crimson")
+        a.set_title("Per-token policy entropy — MISSING", fontsize=10)
+    a.set_xlabel("optimizer step"); a.grid(alpha=.3)
+    fig.suptitle(f"{args.title or args.name} — training diagnostics ({verdict})",
+                 fontsize=12)
+    plt.tight_layout()
+    figp = logbook.FIGS / f"{args.name}_diagnostics.png"
+    figp.parent.mkdir(parents=True, exist_ok=True)
+    plt.savefig(figp, dpi=140)
+    plt.close()
+
+    # --- gate failure census ---------------------------------------------
+    from collections import Counter
+    cnt = Counter()
+    for h in hist:
+        for k, v in (h.get("reasons") or {}).items():
+            cnt[k] += v
+    total_stories = sum(h["n"] for h in hist)
+
+    rows = [{"metric": k, "early": round(a, 4), "late": round(b, 4),
+             "delta": round(d, 4)} for k, (a, b, d) in tr.items()]
+    if ent.size:
+        a_, b_, d_ = trend(ent)
+        rows.append({"metric": "policy_entropy", "early": round(a_, 4),
+                     "late": round(b_, 4), "delta": round(d_, 4)})
+        entropy_note = (
+            f"Per-token policy entropy moved {a_:.4f} → {b_:.4f} "
+            f"(Δ {d_:+.4f}, {100*d_/max(a_,1e-9):+.1f}%). "
+            + ("**Entropy collapse** — the policy is becoming deterministic; "
+               "treat any diversity gain reported below with suspicion."
+               if d_ < -0.15 * a_ else
+               "No entropy collapse. Note that entropy holding up is a "
+               "necessary but not sufficient condition for diversity: it "
+               "permits varied output without demonstrating it.")
+        )
+    else:
+        entropy_note = ("Entropy was NOT logged for this run. TRL only emits it "
+                        "on the non-liger loss path; check `use_liger_kernel`.")
+
+    cost = {}
+    cp = run_dir / "judge_cost.json"
+    if cp.exists():
+        cost = json.load(open(cp))
+
+    body = f"""# {args.title or args.name} — training analysis
+
+**Verdict: {verdict}**
+
+Reward batches: {len(hist)} | stories scored: {total_stories}
+
+## Trend (first 25% vs last 25% of the run)
+
+{logbook.table(rows)}
+
+## Gate failures (count over the whole run, {total_stories} stories)
+
+{logbook.table([{"reason": k, "count": v, "pct_of_stories": round(100*v/max(1,total_stories), 2)}
+                for k, v in cnt.most_common()]) if cnt else "_No gate failures._"}
+
+## Policy entropy
+
+{entropy_note}
+
+Entropy is read as an *asymmetric* signal here. A large fall is strong evidence
+that creativity is dying — the policy is collapsing toward deterministic output.
+A rise is only permissive: a policy can raise per-token entropy while staying
+inside one narrative mode (noisier word choice, same story). The prior run
+demonstrated exactly that dissociation — surface variation up, semantic
+diversity down. So entropy is never optimized, and diversity claims rest on
+log-det / effective rank.
+
+## Reward-hacking check
+
+The signature we watch for is **diversity up while quality or validity goes
+down** — the policy discovering it can farm the diversity channel by emitting
+text that is different because it is worse.
+
+- mean deviation: {dev_a:.4f} → {dev_b:.4f} (Δ {dev_d:+.4f})
+- judge quality:  {q_a:.3f} → {q_b:.3f} (Δ {q_d:+.3f})
+- gate pass:      {g_a:.3f} → {g_b:.3f} (Δ {g_d:+.3f})
+
+Trip conditions: `Δdeviation > +0.02 AND Δquality < -0.5`, or
+`gate_pass < 0.60 AND Δgate_pass < -0.20`. → **{"TRIPPED" if hack else "not tripped"}**
+
+## Judge cost
+
+```json
+{json.dumps(cost, indent=1) if cost else "{}"}
+```
+
+![diagnostics](../figures/{args.name}_diagnostics.png)
+"""
+    p = logbook.write_report(args.name, body)
+    print(f"verdict: {verdict}")
+    for r in rows:
+        print(f"  {r['metric']:26} {r['early']:>9.4f} -> {r['late']:>9.4f}  ({r['delta']:+.4f})")
+    print("report ->", p)
+    print("figure ->", figp)
+    logbook.note(f"analysis: {args.name}", f"verdict={verdict}; report {p}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

src/build_pairs.py

@@ -0,0 +1,257 @@
+"""
+Preference-pair construction for E3 (multi-positive diverse DPO) and
+E4 (faithful DivPO). Both read the SAME scored pool, so the comparison isolates
+pair construction + loss, with data held fixed.
+
+E4 -- DivPO (Lanchantin et al. 2025), implemented as published
+------------------------------------------------------------
+One pair per prompt, no loss weighting, no multi-positive rows:
+    chosen   = argmax diversity over {quality >= rho}
+    rejected = argmin diversity over {quality <  rho}
+Prompts where either side is empty are SKIPPED (skip rate is logged; if it
+exceeds 30% we adjust rho once and record the change).
+
+Two diversity criteria:
+  divpo-emb   -- deviation d_i from the embedding module.
+  divpo-prob  -- model probability. Most diverse = LOWEST length-normalized
+                 logprob, least diverse = HIGHEST. The highest-probability
+                 sample in a temperature pool is by construction the
+                 near-greedy one, which makes this the principled version of
+                 "reject the greedy decode".
+
+Length normalization is not optional here: raw cumulative logprob scales with
+token count, so ranking on it would rank by length and the "most diverse"
+choice would just be the longest story.
+
+E3 -- multi-positive deviation-weighted DPO
+-------------------------------------------
+    survivors = {quality >= q_keep}
+    chosen    = greedy argmax over 4-subsets of  sum(quality) + lam * logdet(L_S)
+    rejected  = r_D "competent cliche" (highest-quality LOW-deviation survivor)
+              + r_Q "clearly weak"     (lowest-quality sample overall)
+Rows pair each chosen against a rejected, ROTATING the rejected across rows so
+one negative is not hammered four times.
+
+Each row carries `weight` = the chosen's deviation d_i (DDPO-style), consumed
+by the trainer as a per-sample loss weight.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from collections import defaultdict
+from pathlib import Path
+
+import numpy as np
+
+ROOT = Path(__file__).resolve().parent.parent
+
+
+def load_pool(tag: str, split: str = "train") -> dict[str, list[dict]]:
+    p = ROOT / "outputs" / f"pool_{tag}" / f"pool_{split}.jsonl"
+    rows = [json.loads(l) for l in open(p) if l.strip()]
+    by: dict[str, list[dict]] = defaultdict(list)
+    for r in rows:
+        by[r["prompt_id"]].append(r)
+    return dict(by)
+
+
+def load_emb(tag: str, split: str = "train") -> np.ndarray:
+    return np.load(ROOT / "outputs" / f"pool_{tag}" / f"emb_{split}.npy")
+
+
+# --------------------------------------------------------------------- E4
+def build_divpo(pool: dict[str, list[dict]], criterion: str, rho: float) -> tuple[list[dict], dict]:
+    """criterion: 'emb' (deviation) or 'prob' (mean logprob)."""
+    rows, skipped = [], {"no_chosen": 0, "no_rejected": 0, "ok": 0}
+
+    for pid, items in pool.items():
+        valid = [r for r in items if r["gate_passed"]]
+        hi = [r for r in valid if r["quality"] >= rho]
+        # rejected pool: below the quality bar. Gate failures are legitimate
+        # DivPO rejects (they are the low-quality tail), so they are eligible.
+        lo = [r for r in items if not (r["gate_passed"] and r["quality"] >= rho)]
+        if not hi:
+            skipped["no_chosen"] += 1; continue
+        if not lo:
+            skipped["no_rejected"] += 1; continue
+
+        if criterion == "emb":
+            chosen = max(hi, key=lambda r: r["deviation"])
+            rejected = max(lo, key=lambda r: -r["deviation"])   # least diverse
+        elif criterion == "prob":
+            # most diverse == least probable; least diverse == most probable
+            chosen = min(hi, key=lambda r: r["mean_logprob"])
+            rejected = max(lo, key=lambda r: r["mean_logprob"])
+        else:
+            raise ValueError(criterion)
+
+        rows.append({
+            "prompt_id": pid, "prompt": chosen["prompt"],
+            "chosen": chosen["text"], "rejected": rejected["text"],
+            "weight": 1.0,                       # DivPO is unweighted, by design
+            "chosen_quality": chosen["quality"], "chosen_dev": chosen["deviation"],
+            "rejected_quality": rejected["quality"], "rejected_dev": rejected["deviation"],
+            "chosen_meanlp": chosen["mean_logprob"], "rejected_meanlp": rejected["mean_logprob"],
+        })
+        skipped["ok"] += 1
+
+    n = len(pool)
+    stats = {"criterion": f"divpo-{criterion}", "rho": rho, "n_prompts": n,
+             "n_rows": len(rows), **skipped,
+             "skip_rate": 1 - skipped["ok"] / max(1, n)}
+    return rows, stats
+
+
+# --------------------------------------------------------------------- E3
+def build_multipos(pool: dict[str, list[dict]], emb: np.ndarray, row_index: dict,
+                   q_keep: float, k: int, lam: float) -> tuple[list[dict], dict]:
+    from diversity import greedy_diverse_subset
+
+    rows = []
+    stats = {"n_prompts": len(pool), "skipped_no_survivors": 0,
+             "skipped_no_negatives": 0, "ok": 0, "chosen_per_prompt": []}
+
+    for pid, items in pool.items():
+        survivors = [r for r in items if r["gate_passed"] and r["quality"] >= q_keep]
+        if len(survivors) < 2:
+            stats["skipped_no_survivors"] += 1; continue
+
+        # r_Q: clearly low quality (worst overall, gate failures included)
+        r_Q = min(items, key=lambda r: (r["gate_passed"], r["quality"]))
+        # r_D: the "competent cliche" -- high quality but LOW deviation.
+        # Rank by (quality - deviation) so we favour a strong, conventional
+        # story rather than merely the least diverse one.
+        r_D = max(survivors, key=lambda r: r["quality"] - 4.0 * r["deviation"])
+
+        # r_D is itself a survivor, so it can coincide with a chosen story; that
+        # specific pairing is dropped per-row below rather than here, since it
+        # only invalidates one row and not the whole prompt.
+        negatives = [n for n in (r_D, r_Q) if n["text"]]
+        if not negatives:
+            stats["skipped_no_negatives"] += 1; continue
+
+        idxs = [row_index[(pid, r["idx"])] for r in survivors]
+        E = emb[idxs].astype(np.float64)
+        q = np.array([r["quality"] for r in survivors], dtype=np.float64)
+        sel = greedy_diverse_subset(q, E, k=min(k, len(survivors)), lam=lam)
+        chosen_set = [survivors[i] for i in sel]
+        stats["chosen_per_prompt"].append(len(chosen_set))
+
+        for j, ch in enumerate(chosen_set):
+            neg = negatives[j % len(negatives)]      # rotate, don't hammer one
+            if neg["text"] == ch["text"]:
+                continue
+            rows.append({
+                "prompt_id": pid, "prompt": ch["prompt"],
+                "chosen": ch["text"], "rejected": neg["text"],
+                "weight": float(ch["deviation"]),    # DDPO-style loss weight
+                "neg_type": "r_D" if neg is r_D else "r_Q",
+                "chosen_quality": ch["quality"], "chosen_dev": ch["deviation"],
+                "rejected_quality": neg["quality"], "rejected_dev": neg["deviation"],
+            })
+        stats["ok"] += 1
+
+    stats["n_rows"] = len(rows)
+    stats["mean_chosen_per_prompt"] = float(np.mean(stats["chosen_per_prompt"])) if stats["chosen_per_prompt"] else 0.0
+    stats.pop("chosen_per_prompt")
+    stats["skip_rate"] = 1 - stats["ok"] / max(1, len(pool))
+    return rows, stats
+
+
+def normalize_weights(rows: list[dict]) -> None:
+    """Scale weights to mean 1.0 so the DDPO weighting changes the RELATIVE
+    emphasis across rows without also rescaling the effective learning rate."""
+    w = np.array([r["weight"] for r in rows], dtype=np.float64)
+    if w.size and w.mean() > 1e-9:
+        w = w / w.mean()
+        for r, x in zip(rows, w):
+            r["weight"] = float(x)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--tag", default="4b")
+    ap.add_argument("--split", default="train")
+    ap.add_argument("--rho", type=float, default=6.0, help="DivPO quality threshold")
+    ap.add_argument("--q-keep", type=float, default=5.0, help="E3 survivor threshold")
+    ap.add_argument("--k", type=int, default=4, help="E3 chosen-subset size")
+    ap.add_argument("--lam", type=float, default=1.0, help="E3 logdet weight")
+    ap.add_argument("--max-skip", type=float, default=0.30)
+    args = ap.parse_args()
+
+    import logbook
+
+    pool = load_pool(args.tag, args.split)
+    emb = load_emb(args.tag, args.split)
+    row_index = {}
+    i = 0
+    for pid in pool:
+        for r in pool[pid]:
+            row_index[(pid, r["idx"])] = i; i += 1
+    assert i == emb.shape[0], f"pool/emb mismatch {i} vs {emb.shape[0]}"
+
+    out = ROOT / "outputs" / f"pairs_{args.tag}"
+    out.mkdir(parents=True, exist_ok=True)
+    all_stats = {}
+
+    # ---- E4: DivPO, both criteria ---------------------------------------
+    for crit in ("emb", "prob"):
+        rho = args.rho
+        rows, st = build_divpo(pool, crit, rho)
+        if st["skip_rate"] > args.max_skip:
+            # One adjustment, as the brief allows, chosen by SWEEP rather than by
+            # a fixed percentile. Direction matters and is not knowable a priori:
+            #   no_chosen   dominant -> rho must come DOWN
+            #   no_rejected dominant -> rho must go UP
+            # An earlier version always moved rho down (40th percentile), which
+            # is the wrong direction for this pool: the judge puts 88.8% of
+            # stories at >= 6, so the binding failure was no_rejected (311 of
+            # 1000 prompts had no story below the bar), not no_chosen (3).
+            # Sweeping the observed quality levels picks the threshold that
+            # actually maximizes usable prompts.
+            qs = sorted({r["quality"] for items in pool.values()
+                         for r in items if r["gate_passed"]})
+            cands = [q for q in qs if 3.0 <= q <= 9.0] or [rho]
+            best = None
+            for c in cands:
+                _, s = build_divpo(pool, crit, float(c))
+                if best is None or s["skip_rate"] < best[1]["skip_rate"]:
+                    best = (c, s)
+            rows2, st2 = build_divpo(pool, crit, float(best[0]))
+            st2["rho_adjusted_from"] = rho
+            st2["reason"] = (f"skip_rate {st['skip_rate']:.3f} > {args.max_skip} "
+                             f"(no_chosen={st['no_chosen']}, "
+                             f"no_rejected={st['no_rejected']}); swept "
+                             f"{[float(c) for c in cands]} -> rho={best[0]}")
+            rows, st = rows2, st2
+        p = out / f"divpo_{crit}_{args.split}.jsonl"
+        with open(p, "w") as f:
+            for r in rows:
+                f.write(json.dumps(r) + "\n")
+        all_stats[f"divpo-{crit}"] = st
+        print(f"[divpo-{crit}] {json.dumps(st)}")
+
+    # ---- E3: multi-positive ---------------------------------------------
+    rows, st = build_multipos(pool, emb, row_index, args.q_keep, args.k, args.lam)
+    normalize_weights(rows)
+    st["weight_mean_after_norm"] = float(np.mean([r["weight"] for r in rows])) if rows else 0.0
+    st["weight_sd"] = float(np.std([r["weight"] for r in rows])) if rows else 0.0
+    from collections import Counter
+    st["neg_type_counts"] = dict(Counter(r["neg_type"] for r in rows))
+    p = out / f"multipos_{args.split}.jsonl"
+    with open(p, "w") as f:
+        for r in rows:
+            f.write(json.dumps(r) + "\n")
+    all_stats["e3-multipos"] = st
+    print(f"[e3-multipos] {json.dumps(st)}")
+
+    json.dump(all_stats, open(out / f"pair_stats_{args.split}.json", "w"), indent=2)
+    logbook.note(f"pairs built ({args.tag})",
+                 f"```json\n{json.dumps(all_stats, indent=1)}\n```")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

src/build_pool.py

@@ -0,0 +1,161 @@
+"""
+Build the scored generation pool that feeds E3 (multi-positive diverse DPO) and
+E4 (faithful DivPO), and doubles as base-model analysis data.
+
+Per prompt: N=16 samples at temperature 1.0 from the BASE policy, each carrying
+  - text, token count, cumulative + length-normalized logprob (E4 divpo-prob)
+  - programmatic gate result
+  - judge quality / novelty (gate-passing stories only)
+  - embedding, per-group deviation d_i and marginal contribution m_i
+
+Both DPO arms consume this identical artifact, which is the point: E4-vs-E3 is
+then a comparison of PAIR CONSTRUCTION and LOSS, with the data held fixed.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+
+ROOT = Path(__file__).resolve().parent.parent
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model", default="Qwen/Qwen3-4B-Instruct-2507")
+    ap.add_argument("--tag", default="4b")
+    ap.add_argument("--n", type=int, default=16)
+    ap.add_argument("--temperature", type=float, default=1.0)
+    ap.add_argument("--top-p", type=float, default=1.0)
+    ap.add_argument("--limit", type=int, default=None, help="prompt subset, for smoke")
+    ap.add_argument("--split", default="train")
+    ap.add_argument("--seed", type=int, default=1234)
+    ap.add_argument("--gpu-mem", type=float, default=0.85)
+    args = ap.parse_args()
+
+    from transformers import AutoTokenizer
+
+    import gates
+    import logbook
+    from data import load_prompts
+    from diversity import (l2_normalize, logdet_volume, marginal_contributions,
+                           pairwise_deviation)
+    from generate import build_llm, generate
+    from judge import build_judge
+
+    out_dir = ROOT / "outputs" / f"pool_{args.tag}"
+    out_dir.mkdir(parents=True, exist_ok=True)
+    pool_path = out_dir / f"pool_{args.split}.jsonl"
+
+    prompts = load_prompts(args.split, ROOT / "data")
+    if args.limit:
+        prompts = prompts[: args.limit]
+    print(f"[pool] {len(prompts)} prompts x N={args.n} = {len(prompts)*args.n} stories")
+
+    # ---- 1. generate -----------------------------------------------------
+    t0 = time.time()
+    tok = AutoTokenizer.from_pretrained(args.model)
+    llm = build_llm(args.model, gpu_mem_util=args.gpu_mem, seed=args.seed)
+    gens = generate(llm, tok, prompts, n=args.n, temperature=args.temperature,
+                    top_p=args.top_p, seed=args.seed)
+    t_gen = time.time() - t0
+    ntok = sum(g.n_tokens for g in gens)
+    print(f"[gen] {len(gens)} stories, {ntok} tok in {t_gen/60:.1f} min "
+          f"({ntok/t_gen:.0f} tok/s)")
+
+    # free the GPU before loading the embedder
+    del llm
+    import gc, torch
+    gc.collect(); torch.cuda.empty_cache()
+
+    # ---- 2. gates --------------------------------------------------------
+    grs = [gates.check(g.text, finish_reason=g.finish_reason) for g in gens]
+    n_pass = sum(r.passed for r in grs)
+    print(f"[gates] pass {n_pass}/{len(grs)} ({100*n_pass/len(grs):.1f}%)")
+
+    # ---- 3. judge (gate-passers only) ------------------------------------
+    judge = build_judge(cache_path=str(ROOT / "cache" / "judge.sqlite"),
+                        concurrency=24)
+    idx = [i for i in range(len(gens)) if grs[i].passed]
+    t1 = time.time()
+    scores = judge.score_many_sync([(gens[i].prompt, gens[i].text) for i in idx])
+    print(f"[judge] {len(idx)} scored in {(time.time()-t1)/60:.1f} min | "
+          f"health={judge.health()}")
+    judge.assert_healthy()
+
+    quality = np.zeros(len(gens)); novelty = np.zeros(len(gens))
+    for i, s in zip(idx, scores):
+        quality[i] = s.quality; novelty[i] = s.novelty
+
+    # ---- 4. embeddings + per-group diversity -----------------------------
+    from sentence_transformers import SentenceTransformer
+    enc = SentenceTransformer("BAAI/bge-base-en-v1.5", device="cuda")
+    t2 = time.time()
+    E = enc.encode([g.text for g in gens], normalize_embeddings=True,
+                   batch_size=64, show_progress_bar=False, convert_to_numpy=True)
+    E = l2_normalize(np.asarray(E, dtype=np.float64))
+    print(f"[embed] {E.shape} in {time.time()-t2:.0f}s")
+
+    by_prompt: dict[str, list[int]] = {}
+    for i, g in enumerate(gens):
+        by_prompt.setdefault(g.prompt_id, []).append(i)
+
+    dev = np.zeros(len(gens)); marg = np.zeros(len(gens))
+    group_logdet: dict[str, float] = {}
+    for pid, ids in by_prompt.items():
+        sub = E[ids]
+        d = pairwise_deviation(sub); m = marginal_contributions(sub)
+        for k, i in enumerate(ids):
+            dev[i] = d[k]; marg[i] = m[k]
+        group_logdet[pid] = logdet_volume(sub)
+
+    # ---- 5. write --------------------------------------------------------
+    with open(pool_path, "w") as f:
+        for i, g in enumerate(gens):
+            f.write(json.dumps({
+                **g.as_dict(),
+                "gate_passed": bool(grs[i].passed),
+                "gate_reasons": grs[i].reasons,
+                "ends_cleanly": grs[i].completeness,
+                "n_words": grs[i].n_words,
+                "quality": float(quality[i]),
+                "novelty": float(novelty[i]),
+                "deviation": float(dev[i]),
+                "marginal": float(marg[i]),
+                "group_logdet": float(group_logdet[g.prompt_id]),
+            }) + "\n")
+    np.save(out_dir / f"emb_{args.split}.npy", E.astype(np.float32))
+
+    # ---- 6. summary ------------------------------------------------------
+    q_pass = quality[[i for i in idx]]
+    summary = {
+        "tag": args.tag, "model": args.model, "split": args.split,
+        "n_prompts": len(prompts), "n_per_prompt": args.n, "n_stories": len(gens),
+        "temperature": args.temperature, "seed": args.seed,
+        "gate_pass_rate": float(n_pass / len(grs)),
+        "ends_cleanly_rate": float(np.mean([r.completeness for r in grs])),
+        "median_words": float(np.median([r.n_words for r in grs])),
+        "quality_mean": float(q_pass.mean()) if len(q_pass) else 0.0,
+        "quality_sd": float(q_pass.std()) if len(q_pass) else 0.0,
+        "novelty_mean": float(novelty[idx].mean()) if len(idx) else 0.0,
+        "deviation_mean": float(dev.mean()),
+        "logdet_mean": float(np.mean(list(group_logdet.values()))),
+        "gen_minutes": t_gen / 60,
+        "judge_cost": judge.cost_estimate(0.140, 0.280),
+        "judge_health": judge.health(),
+    }
+    json.dump(summary, open(out_dir / f"summary_{args.split}.json", "w"), indent=2)
+    print("\n" + json.dumps(summary, indent=1))
+
+    logbook.note(f"pool built: {args.tag}/{args.split}",
+                 f"```json\n{json.dumps(summary, indent=1)}\n```")
+    logbook.checkpoint(f"pool_{args.tag}_{args.split}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

src/calibrate_judge.py

@@ -0,0 +1,230 @@
+"""
+Judge calibration. A reward model you have not validated is a random number
+generator with good manners, so this runs before any training.
+
+Test set is built from material we already have, with KNOWN ground-truth
+ordering, so we can measure discrimination rather than eyeball plausibility:
+
+  A. complete   -- fresh Qwen3-4B stories that passed every programmatic gate
+  B. truncated  -- prior-run stories, all cut off mid-word (known bad)
+  C. shuffled   -- a complete story with its sentences randomly reordered
+                   (destroys narrative coherence, preserves vocabulary/style)
+  D. repetitive -- a complete story with one paragraph repeated to fill it out
+
+C and D matter most. Any judge can tell prose from garbage; the question is
+whether it tracks COHERENCE, which is what quality is supposed to mean here. A
+judge that scores C and D as highly as A is measuring surface fluency, and
+would happily pay a policy to produce fluent incoherent text.
+
+Pass criteria:
+  quality(A) - quality(B) >= 2.0   detects truncation
+  quality(A) - quality(C) >= 1.5   detects incoherence, not just fluency
+  quality(A) - quality(D) >= 1.5   detects padding/repetition
+  sd(quality within A)   >= 0.7    not saturated/constant
+"""
+from __future__ import annotations
+
+import glob
+import json
+import random
+import re
+import statistics
+import sys
+from pathlib import Path
+
+import gates
+import logbook
+from judge import build_judge
+
+ROOT = Path(__file__).resolve().parent.parent
+PRIOR = ROOT / "prior_run" / "eval_samples"
+N_PER_CELL = 12
+SEED = 11
+
+PRICE_IN, PRICE_OUT = 0.140, 0.280   # deepseek-v4-flash-0731, $/1M tok
+
+
+def sentence_shuffle(text: str, rng: random.Random) -> str:
+    sents = re.split(r"(?<=[.!?])\s+", text.strip())
+    if len(sents) < 4:
+        return text
+    rng.shuffle(sents)
+    return " ".join(sents)
+
+
+def hard_truncate(text: str, frac: float = 0.68) -> str:
+    """Cut mid-word at `frac` of the way through -- the exact signature of a
+    max_new_tokens wall. Applied to cell A's OWN stories so the comparison
+    isolates truncation with style, voice and model held fixed."""
+    w = text.split()
+    cut = " ".join(w[: max(20, int(len(w) * frac))])
+    return cut[: max(1, len(cut) - 3)]      # shave into the final word
+
+
+def paragraph_pad(text: str) -> str:
+    paras = [p for p in text.split("\n\n") if p.strip()]
+    if not paras:
+        return text
+    return "\n\n".join(paras[:2] + [paras[0]] * 3)
+
+
+def main():
+    rng = random.Random(SEED)
+    from generate import load_gens
+
+    smoke = ROOT / "outputs" / "smoke" / "vllm_smoke.jsonl"
+    if not smoke.exists():
+        print("need outputs/smoke/vllm_smoke.jsonl (run smoke_vllm.py first)")
+        return 2
+
+    gens = load_gens(smoke)
+    good = [g for g in gens if gates.check(g.text, finish_reason=g.finish_reason).passed]
+    if len(good) < N_PER_CELL:
+        print(f"only {len(good)} gate-passing stories; need {N_PER_CELL}")
+        return 2
+    rng.shuffle(good)
+    A = good[:N_PER_CELL]
+
+    prior_texts = []
+    for f in sorted(glob.glob(str(PRIOR / "step_*.json"))):
+        for p, samples in json.load(open(f)).items():
+            for s in samples:
+                if not gates.ends_cleanly(s):
+                    prior_texts.append((p.split("\n\n")[-1].strip(), s))
+    rng.shuffle(prior_texts)
+    B = prior_texts[:N_PER_CELL]
+
+    cells = {
+        "A_complete":     [(g.prompt, g.text) for g in A],
+        "B_cut_matched":  [(g.prompt, hard_truncate(g.text)) for g in A],
+        "C_shuffled":     [(g.prompt, sentence_shuffle(g.text, rng)) for g in A],
+        "D_repetitive":   [(g.prompt, paragraph_pad(g.text)) for g in A],
+        "E_cut_priorrun": B,   # informational only: different policy AND style
+    }
+
+    judge = build_judge(cache_path=str(ROOT / "cache" / "judge.sqlite"), concurrency=8)
+    print(f"judge model: {judge.model}\n")
+
+    results = {}
+    for name, pairs in cells.items():
+        scores = judge.score_many_sync(pairs)
+        q = [s.quality for s in scores]
+        nv = [s.novelty for s in scores]
+        ok = sum(s.ok for s in scores)
+        results[name] = {
+            "cell": name, "n": len(q), "ok_calls": ok,
+            "quality_mean": statistics.fmean(q),
+            "quality_sd": statistics.pstdev(q) if len(q) > 1 else 0.0,
+            "quality_min": min(q), "quality_max": max(q),
+            "novelty_mean": statistics.fmean(nv),
+            "novelty_sd": statistics.pstdev(nv) if len(nv) > 1 else 0.0,
+        }
+        print(f"{name:14} q={results[name]['quality_mean']:.2f}"
+              f" (sd {results[name]['quality_sd']:.2f},"
+              f" {results[name]['quality_min']:.0f}-{results[name]['quality_max']:.0f})"
+              f"  nov={results[name]['novelty_mean']:.2f}"
+              f" (sd {results[name]['novelty_sd']:.2f})  ok={ok}/{len(q)}")
+
+    qa = results["A_complete"]["quality_mean"]
+    # BLOCKING checks: these test properties no other component can supply.
+    checks = [
+        ("detects incoherence  A-C >= 1.5", qa - results["C_shuffled"]["quality_mean"], 1.5),
+        ("detects repetition   A-D >= 1.5", qa - results["D_repetitive"]["quality_mean"], 1.5),
+        ("not saturated  sd(A) >= 0.7", results["A_complete"]["quality_sd"], 0.7),
+    ]
+    # INFORMATIONAL: truncation is already caught deterministically, with 100%
+    # recall, by gates.check() (finish_reason=="length" OR no terminal
+    # punctuation). RewardEngine.compute() judges ONLY gate-passing stories, so
+    # a truncated story never reaches the judge during training and the judge's
+    # truncation sensitivity cannot influence any gradient. Kept as a monitored
+    # number, not a gate on proceeding.
+    info = [
+        ("detects truncation   A-B (informational)",
+         qa - results["B_cut_matched"]["quality_mean"], 2.0),
+        ("style confound       A-E (informational)",
+         qa - results["E_cut_priorrun"]["quality_mean"], 0.0),
+    ]
+    print("\n=== DISCRIMINATION (blocking) ===")
+    allpass = True
+    for label, val, thresh in checks:
+        ok = val >= thresh
+        allpass &= ok
+        print(f"  {'PASS' if ok else 'FAIL'}  {label:34} got {val:+.2f}")
+    print("=== informational (not blocking) ===")
+    for label, val, thresh in info:
+        print(f"  {'ok  ' if val >= thresh else 'note'}  {label:40} got {val:+.2f}")
+    checks = checks + info
+
+    cost = judge.cost_estimate(PRICE_IN, PRICE_OUT)
+    health = judge.health()
+    print(f"\njudge health: {health}")
+    print(f"cost so far: {cost}")
+    if health["fail_rate"] > 0.05:
+        print("\n!!! JUDGE UNHEALTHY -- scores above are mostly neutral fallbacks, "
+              "discrimination numbers are meaningless. Fix the backend first.")
+        allpass = False
+
+    body = f"""# Judge calibration — `{judge.model}`
+
+A reward model that has not been validated is a random number generator with
+good manners. This is the pre-flight check, run before any training.
+
+## Method
+
+Four cells, 12 stories each, with known ground-truth ordering:
+
+- **A_complete** — fresh Qwen3-4B stories passing every programmatic gate.
+- **B_cut_matched** — cell A's OWN stories, cut mid-word at 68% length. This is
+  the controlled truncation test: same model, same voice, same prompt, so the
+  only variable is finishedness.
+- **E_cut_priorrun** — prior-run stories, also truncated. Informational only:
+  they come from a different (RL-tuned, far more florid) policy, so A-vs-E
+  confounds truncation with style and cannot be read as a truncation test.
+- **C_shuffled** — cell A with sentences randomly reordered. Destroys narrative
+  coherence while preserving vocabulary, register and sentence-level fluency.
+- **D_repetitive** — cell A with one paragraph repeated to pad the length.
+
+C and D carry the weight. Any judge separates prose from garbage; the real
+question is whether it tracks *coherence* rather than surface fluency. A judge
+that rates C as highly as A would pay a policy to emit fluent incoherent text —
+exactly the reward-hacking channel the diversity bonus could exploit.
+
+## Results
+
+{logbook.table(list(results.values()))}
+
+## Discrimination checks
+
+{logbook.table([{"check": l, "value": round(v, 3), "threshold": t,
+                 "verdict": "PASS" if v >= t else "FAIL"} for l, v, t in checks])}
+
+**Verdict: {"USABLE" if allpass else "NOT USABLE AS-IS"}**
+
+## Judge health
+
+```json
+{json.dumps(health, indent=1)}
+```
+
+## Cost accounting
+
+```json
+{json.dumps(cost, indent=1)}
+```
+"""
+    p = logbook.write_report("01_judge_calibration", body)
+    print("report ->", p)
+    json.dump({"results": results,
+               "checks": [{"c": l, "v": v, "t": t} for l, v, t in checks],
+               "cost": cost, "model": judge.model},
+              open(logbook.LOGS / "judge_calibration.json", "w"), indent=2)
+    logbook.note("judge calibration",
+                 f"model={judge.model} verdict={'USABLE' if allpass else 'FAIL'}; "
+                 f"A={qa:.2f} B={results['B_cut_matched']['quality_mean']:.2f} "
+                 f"C={results['C_shuffled']['quality_mean']:.2f} "
+                 f"D={results['D_repetitive']['quality_mean']:.2f}")
+    return 0 if allpass else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

src/ckpt_study.py

@@ -0,0 +1,215 @@
+"""
+Checkpoint trajectory study: generate the SAME prompts from every checkpoint of
+an arm, so the evolution of the actual stories is visible -- not just metrics.
+
+Aggregate numbers can report "effective rank 2.0" without conveying that six of
+sixteen ships are named *Aethel*. This dumps the stories in a readable form at
+each training step alongside the metrics, so the qualitative change can be read
+directly and cross-checked against the quantitative one.
+
+vLLM loads the base model ONCE and hot-swaps LoRA adapters per checkpoint, so
+the whole sweep costs one model load rather than one per checkpoint.
+
+Outputs (under outputs/ckpt_study/<arm>/):
+  stories.md        human-readable: every prompt, every checkpoint, side by side
+  metrics.csv       per-checkpoint quantitative trajectory
+  raw.json          everything, for re-analysis
+  ../logs/figures/<arm>_trajectory.png
+"""
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+import re
+import sys
+from pathlib import Path
+
+import numpy as np
+
+ROOT = Path(__file__).resolve().parent.parent
+
+
+def find_checkpoints(arm: str) -> list[tuple[int, str | None]]:
+    """[(step, adapter_path_or_None)] ascending; step 0 = base model."""
+    d = ROOT / "outputs" / arm
+    out: list[tuple[int, str | None]] = [(0, None)]
+    if d.exists():
+        for p in d.glob("checkpoint-*"):
+            m = re.search(r"checkpoint-(\d+)", p.name)
+            if m and (p / "adapter_model.safetensors").exists():
+                out.append((int(m.group(1)), str(p)))
+        f = d / "final"
+        if (f / "adapter_model.safetensors").exists():
+            steps = [s for s, _ in out]
+            out.append((max(steps) + 1 if steps else 1, str(f)))
+    return sorted(out, key=lambda t: t[0])
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--arm", required=True)
+    ap.add_argument("--model", default="Qwen/Qwen3-4B-Instruct-2507")
+    ap.add_argument("--prompts", type=int, default=10)
+    ap.add_argument("--n", type=int, default=6)
+    ap.add_argument("--temp", type=float, default=0.9)
+    ap.add_argument("--top-p", type=float, default=0.95)
+    ap.add_argument("--seed", type=int, default=777)
+    ap.add_argument("--gpu-mem", type=float, default=0.85)
+    ap.add_argument("--judge", action="store_true", default=True)
+    args = ap.parse_args()
+
+    from transformers import AutoTokenizer
+    from vllm import SamplingParams
+    from vllm.lora.request import LoRARequest
+
+    import gates
+    import logbook
+    from data import load_prompts
+    from diversity import effective_rank, l2_normalize, logdet_volume, pairwise_deviation
+    from generate import build_llm, render_chat
+    from judge import build_judge
+    from qualitative import analyze_group, first_sentence, last_sentence
+
+    ckpts = find_checkpoints(args.arm)
+    if len(ckpts) < 2:
+        print(f"only {len(ckpts)} checkpoint(s) for {args.arm}; nothing to compare")
+        return 2
+    print(f"[{args.arm}] checkpoints: {[s for s, _ in ckpts]}")
+
+    prompts = load_prompts("eval", ROOT / "data")[: args.prompts]
+    tok = AutoTokenizer.from_pretrained(args.model)
+    llm = build_llm(args.model, gpu_mem_util=args.gpu_mem, seed=args.seed,
+                    enable_lora=True)
+    rendered = [render_chat(tok, p["prompt"]) for p in prompts]
+    sp = SamplingParams(n=args.n, temperature=args.temp, top_p=args.top_p,
+                        max_tokens=1024, seed=args.seed, skip_special_tokens=True)
+
+    from sentence_transformers import SentenceTransformer
+    enc = None
+    judge = build_judge(cache_path=str(ROOT / "cache" / "judge.sqlite"), concurrency=24) \
+        if args.judge else None
+
+    all_data, rowsum = {}, []
+    for step, path in ckpts:
+        kw = {}
+        if path:
+            kw["lora_request"] = LoRARequest(f"{args.arm}-{step}", max(step, 1), path)
+        outs = llm.generate(rendered, sp, **kw)
+        per = {}
+        for p, o in zip(prompts, outs):
+            texts = [x.text.strip() for x in o.outputs]
+            frs = [x.finish_reason or "" for x in o.outputs]
+            per[p["id"]] = {"prompt": p["prompt"], "texts": texts,
+                            "gates": [gates.check(t, finish_reason=f).as_dict()
+                                      for t, f in zip(texts, frs)]}
+        all_data[step] = per
+        print(f"  step {step:>4}: generated {sum(len(v['texts']) for v in per.values())} stories",
+              flush=True)
+
+    del llm
+    import gc, torch
+    gc.collect(); torch.cuda.empty_cache()
+    enc = SentenceTransformer("BAAI/bge-base-en-v1.5", device="cuda")
+
+    for step, per in all_data.items():
+        dev, ld, er, q = [], [], [], []
+        pooled_texts = []
+        for pid, v in per.items():
+            E = l2_normalize(np.asarray(enc.encode(
+                v["texts"], normalize_embeddings=True, show_progress_bar=False,
+                convert_to_numpy=True), dtype=np.float64))
+            v["eff_rank"] = float(effective_rank(E))
+            v["deviation"] = float(pairwise_deviation(E).mean())
+            v["logdet"] = float(logdet_volume(E))
+            v["qual"] = analyze_group(v["texts"])
+            dev.append(v["deviation"]); ld.append(v["logdet"]); er.append(v["eff_rank"])
+            pooled_texts += [(v["prompt"], t, g["passed"])
+                             for t, g in zip(v["texts"], v["gates"])]
+        if judge:
+            idx = [i for i, (_, _, ok) in enumerate(pooled_texts) if ok]
+            sc = judge.score_many_sync([(pooled_texts[i][0], pooled_texts[i][1]) for i in idx])
+            q = [s.quality for s in sc]
+        gp = float(np.mean([g["passed"] for v in per.values() for g in v["gates"]]))
+        ec = float(np.mean([g["completeness"] for v in per.values() for g in v["gates"]]))
+        wd = float(np.mean([g["n_words"] for v in per.values() for g in v["gates"]]))
+        rowsum.append({
+            "step": step, "quality": float(np.mean(q)) if q else 0.0,
+            "eff_rank": float(np.mean(er)), "deviation": float(np.mean(dev)),
+            "logdet": float(np.mean(ld)), "gate_pass": gp, "ends_cleanly": ec,
+            "words": wd,
+            "opens_with_The": float(np.mean([v["qual"]["opens_with_The"] / v["qual"]["n"]
+                                             for v in per.values()])),
+            "distinct_openers": float(np.mean([v["qual"]["distinct_first_5_words"] / v["qual"]["n"]
+                                               for v in per.values()])),
+            "registers": float(np.mean([v["qual"]["registers_present"] for v in per.values()])),
+        })
+        print(f"  step {step:>4}: q={rowsum[-1]['quality']:.2f} "
+              f"eff_rank={rowsum[-1]['eff_rank']:.3f} dev={rowsum[-1]['deviation']:.4f} "
+              f"words={wd:.0f}", flush=True)
+
+    out = ROOT / "outputs" / "ckpt_study" / args.arm
+    out.mkdir(parents=True, exist_ok=True)
+    with open(out / "metrics.csv", "w", newline="") as f:
+        w = csv.DictWriter(f, fieldnames=list(rowsum[0].keys()))
+        w.writeheader(); w.writerows(rowsum)
+    json.dump(all_data, open(out / "raw.json", "w"), indent=1)
+
+    # ---- human-readable side-by-side --------------------------------------
+    steps = [s for s, _ in ckpts]
+    md = [f"# {args.arm} — story trajectory across checkpoints\n",
+          f"{len(prompts)} eval prompts x {args.n} samples, T={args.temp}, "
+          f"top_p={args.top_p}, seed={args.seed} (fixed across checkpoints).\n",
+          "Step 0 = base model.\n"]
+    for pid in list(all_data[steps[0]]):
+        md.append(f"\n## {pid}\n\n> {all_data[steps[0]][pid]['prompt']}\n")
+        for s in steps:
+            v = all_data[s][pid]
+            md.append(f"\n### step {s} — eff_rank {v['eff_rank']:.2f}, "
+                      f"dev {v['deviation']:.3f}\n")
+            md.append("\n**openings**\n")
+            for i, t in enumerate(v["texts"]):
+                md.append(f"{i+1}. {first_sentence(t, 150)}\n")
+            md.append("\n**closings**\n")
+            for i, t in enumerate(v["texts"]):
+                md.append(f"{i+1}. …{last_sentence(t, 110)}\n")
+    (out / "stories.md").write_text("".join(md))
+
+    # ---- figure -----------------------------------------------------------
+    import matplotlib; matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+    x = [r["step"] for r in rowsum]
+    fig, ax = plt.subplots(1, 4, figsize=(19, 4.2))
+    # Anchor the axes that have a meaningful absolute scale. Auto-scaling a
+    # metric that moved 1.60->1.69 against a ceiling of N renders a dramatic
+    # line for a flat result, which is exactly the misreading to avoid.
+    for a, (k, t, c) in zip(ax, [("quality", "Judge quality (0-10)", "#2980b9"),
+                                 ("eff_rank", "Effective rank (1 = collapsed, %d = max)" % args.n, "#8e44ad"),
+                                 ("deviation", "Mean pairwise deviation", "#16a085"),
+                                 ("words", "Story length (words)", "#7f8c8d")]):
+        vals = [r[k] for r in rowsum]
+        a.plot(x, vals, "o-", color=c, lw=2)
+        if k == "eff_rank":
+            a.set_ylim(1.0, args.n)          # full meaningful range
+            a.axhline(1.0, ls=":", c="crimson", lw=1)
+            a.text(x[0], 1.05, "total collapse", fontsize=7, color="crimson")
+        elif k == "quality":
+            a.set_ylim(0, 10)
+        elif k == "deviation":
+            a.set_ylim(0, max(0.5, max(vals) * 1.3))
+        a.set_title(t, fontsize=10); a.set_xlabel("training step"); a.grid(alpha=.3)
+    fig.suptitle(f"{args.arm}: what happens to the stories during training", fontsize=13)
+    plt.tight_layout()
+    figp = logbook.FIGS / f"{args.arm}_trajectory.png"
+    plt.savefig(figp, dpi=140); plt.close()
+
+    print(f"\nstories  -> {out/'stories.md'}")
+    print(f"metrics  -> {out/'metrics.csv'}")
+    print(f"figure   -> {figp}")
+    if judge:
+        print("judge:", judge.health(), judge.cost_estimate(0.140, 0.280))
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

src/data.py

@@ -0,0 +1,145 @@
+"""
+WritingPrompts data prep. Deterministic and seeded: every arm sees the exact
+same train split and the exact same 50 held-out eval prompts.
+
+On the SYSTEM PROMPT (rewritten from the prior run)
+---------------------------------------------------
+The prior run used, with a ~512-token cap and no length guidance:
+
+    "Write a short creative story based on this prompt:\\n\\n{text}"
+
+Result: 100% of eval samples truncated mid-word. The model had no budget to
+plan an ending, so it never planned one. Since the judge rubric caps quality at
+<=3 for unfinished stories, quality became near-constant within every group ->
+zero GRPO advantage -> the quality signal was dead for the whole run.
+
+The fix is a word budget + an explicit ending instruction + a raised token cap.
+
+What the system prompt deliberately does NOT do: it says nothing about being
+original, varied, surprising, or avoiding cliche. That would be a confound.
+Mode collapse is the dependent variable; instructing the model to diversify
+would mask exactly the effect every arm is being measured on. The prompt is
+identical and neutral for base, E0, E1, E2, E3 and E4 -- all differences
+between arms must come from the training objective, not the prompt.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import re
+from pathlib import Path
+
+from datasets import load_dataset
+
+SYSTEM_PROMPT = (
+    "You are a fiction writer. Write a complete short story of 200-500 words "
+    "responding to the writing prompt.\n"
+    "Write only the story: no title, no preamble, no commentary, no author's note.\n"
+    "Finish inside the word budget. The story must reach a real ending, not stop mid-scene."
+)
+
+MIN_PROMPT_WORDS = 10
+MAX_PROMPT_WORDS = 60
+SEED = 42
+
+_TAG = re.compile(r"^\s*\[\s*(WP|EU|CW|TT|RF|IP|PI|PM|MP|OT|SP|FF)\s*\]\s*", re.I)
+_WS = re.compile(r"\s+")
+
+
+def clean_prompt(text: str) -> str:
+    """Strip the subreddit tag and normalize the mangled WritingPrompts spacing.
+
+    The euclaise mirror is tokenizer-detokenized: "did n't phase us ." etc.
+    Left as-is this leaks a distinctive artifact into every prompt and wastes
+    judge tokens, so we repair the obvious cases.
+    """
+    t = _TAG.sub("", text or "").strip()
+    t = t.replace("`` ", '"').replace(" ''", '"').replace("''", '"')
+    t = re.sub(r"([“‘(\[])\s+", r"\1", t)      # "“ What" -> "“What"
+    t = re.sub(r"\s+([”’)\]])", r"\1", t)      # "cockroach! ”" -> "cockroach!”"
+    t = re.sub(r"\s+([,.!?;:])", r"\1", t)
+    t = re.sub(r"\bn't\b", "n't", t)
+    t = re.sub(r"\s+n't", "n't", t)
+    t = re.sub(r"\s+'(s|re|ve|ll|d|m)\b", r"'\1", t)
+    t = re.sub(r"<\s*newline\s*>", " ", t, flags=re.I)
+    t = _WS.sub(" ", t).strip()
+    return t
+
+
+def _key(t: str) -> str:
+    """Dedupe key: lowercase alphanumerics only, so punctuation/spacing
+    variants of the same prompt collapse together."""
+    return hashlib.sha1(re.sub(r"[^a-z0-9]", "", t.lower()).encode()).hexdigest()
+
+
+def build_splits(
+    n_train: int = 1000,
+    n_eval: int = 50,
+    out_dir: str | Path = "data",
+    seed: int = SEED,
+) -> dict:
+    ds = load_dataset("euclaise/writingprompts", split="train")
+    col = "prompt" if "prompt" in ds.column_names else ds.column_names[0]
+
+    seen: set[str] = set()
+    kept: list[str] = []
+    for rec in ds:
+        t = clean_prompt(rec[col])
+        n = len(t.split())
+        if not (MIN_PROMPT_WORDS <= n <= MAX_PROMPT_WORDS):
+            continue
+        k = _key(t)
+        if k in seen:
+            continue
+        seen.add(k)
+        kept.append(t)
+        if len(kept) >= (n_train + n_eval) * 3:   # oversample, then shuffle
+            break
+
+    import random
+    random.Random(seed).shuffle(kept)
+    need = n_train + n_eval
+    if len(kept) < need:
+        raise RuntimeError(f"only {len(kept)} usable prompts, need {need}")
+    sel = kept[:need]
+    eval_prompts, train_prompts = sel[:n_eval], sel[n_eval:need]
+
+    out = Path(out_dir)
+    out.mkdir(parents=True, exist_ok=True)
+    for name, rows in (("train", train_prompts), ("eval", eval_prompts)):
+        p = out / f"{name}_prompts.jsonl"
+        with open(p, "w") as f:
+            for i, t in enumerate(rows):
+                f.write(json.dumps({"id": f"{name}-{i:04d}", "prompt": t}) + "\n")
+
+    meta = {
+        "seed": seed, "n_train": len(train_prompts), "n_eval": len(eval_prompts),
+        "total_scanned": len(ds), "usable_after_filter": len(kept),
+        "min_words": MIN_PROMPT_WORDS, "max_words": MAX_PROMPT_WORDS,
+        "system_prompt": SYSTEM_PROMPT,
+    }
+    (out / "split_meta.json").write_text(json.dumps(meta, indent=2))
+    return meta
+
+
+def load_prompts(split: str, out_dir: str | Path = "data") -> list[dict]:
+    p = Path(out_dir) / f"{split}_prompts.jsonl"
+    return [json.loads(l) for l in open(p) if l.strip()]
+
+
+def chat_messages(prompt: str) -> list[dict]:
+    return [
+        {"role": "system", "content": SYSTEM_PROMPT},
+        {"role": "user", "content": prompt},
+    ]
+
+
+if __name__ == "__main__":
+    m = build_splits()
+    print(json.dumps(m, indent=2))
+    for split in ("train", "eval"):
+        rows = load_prompts(split)
+        print(f"\n{split}: {len(rows)}")
+        for r in rows[:3]:
+            print("   ", r["id"], "|", r["prompt"][:110])

src/diversity.py

@@ -0,0 +1,226 @@
+"""
+Diversity math over an embedding group. Pure numpy, no torch.
+
+Given L2-normalized embeddings E = [e_1..e_G] (shape G x d) for ONE prompt's
+generation group, we expose two families of per-sample diversity credit:
+
+  1. deviation      d_i = mean_{j!=i} (1 - cos(e_i, e_j))
+                    -- pairwise, cheap, but blind to cluster structure:
+                    two tight far-apart clusters score as high as a spread.
+
+  2. marginal       m_i = logdet(L) - logdet(L_{-i})
+     contribution   -- leave-one-out volume credit under the cosine kernel.
+                    A duplicate contributes ~nothing (its direction is already
+                    spanned), so this DOES see cluster structure.
+
+Both are per-sample. This is deliberate: a set-level scalar (e.g. logdet(L)
+itself) is constant within a GRPO group, so its within-group std is 0 and the
+normalized advantage is identically 0. It cannot train anything. Every
+quantity here varies across i within a group.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+EPS_JITTER = 1e-3
+
+
+def l2_normalize(E: np.ndarray, axis: int = -1) -> np.ndarray:
+    """Row-normalize; zero rows are left as zeros rather than NaN."""
+    E = np.asarray(E, dtype=np.float64)
+    n = np.linalg.norm(E, axis=axis, keepdims=True)
+    return E / np.maximum(n, 1e-12)
+
+
+def cosine_kernel(E: np.ndarray, eps: float = EPS_JITTER) -> np.ndarray:
+    """L = E E^T + eps*I.
+
+    E must be L2-normalized, so diag(E E^T) = 1 and L is a correlation-like
+    PSD matrix. The eps jitter keeps logdet finite when rows are collinear
+    (exact duplicates make E E^T singular).
+    """
+    E = np.asarray(E, dtype=np.float64)
+    L = E @ E.T
+    L = 0.5 * (L + L.T)  # kill float asymmetry before Cholesky
+    return L + eps * np.eye(L.shape[0], dtype=np.float64)
+
+
+def pairwise_deviation(E: np.ndarray) -> np.ndarray:
+    """d_i = mean_{j != i} (1 - cos(e_i, e_j)). Shape (G,).
+
+    G == 1 has no off-diagonal terms; we return 0.0 (a lone sample has no
+    measurable deviation, and 0 is the neutral value for the reward).
+    """
+    E = np.asarray(E, dtype=np.float64)
+    G = E.shape[0]
+    if G < 2:
+        return np.zeros(G, dtype=np.float64)
+    S = E @ E.T
+    D = 1.0 - S
+    np.fill_diagonal(D, 0.0)
+    return D.sum(axis=1) / (G - 1)
+
+
+def _logdet_psd(L: np.ndarray) -> float:
+    """logdet via Cholesky; falls back to slogdet if L drifts non-PD."""
+    try:
+        c = np.linalg.cholesky(L)
+        return float(2.0 * np.sum(np.log(np.diag(c))))
+    except np.linalg.LinAlgError:
+        sign, ld = np.linalg.slogdet(L)
+        if sign <= 0:
+            return float("-inf")
+        return float(ld)
+
+
+def logdet_volume(E: np.ndarray, eps: float = EPS_JITTER) -> float:
+    """D(Y) = logdet(E E^T + eps I). Set-level scalar.
+
+    Report this as a METRIC. Never hand it to GRPO as a per-sample reward:
+    it is identical for every i in the group -> zero advantage.
+    """
+    if np.asarray(E).shape[0] == 0:
+        return 0.0
+    return _logdet_psd(cosine_kernel(E, eps))
+
+
+def marginal_contributions(E: np.ndarray, eps: float = EPS_JITTER) -> np.ndarray:
+    """m_i = logdet(L) - logdet(L_{-i}), shape (G,).
+
+    Computed by G explicit leave-one-out logdets. G <= 16 here, so this is
+    ~16 Cholesky calls on a 15x15 matrix -- utterly negligible next to a
+    single LLM forward pass. Rank-one downdate machinery would be a
+    micro-optimization with real numerical-stability downside; not worth it.
+
+    Interpretation: m_i is the log-volume the group loses by dropping i.
+    A duplicate has m_i ~ log(eps) -> large negative. An orthogonal direction
+    has m_i ~ log(1+eps) ~ 0. So m_i is a *penalty scale*: higher (closer to
+    0) means "this sample carries a direction nothing else covers".
+    """
+    E = np.asarray(E, dtype=np.float64)
+    G = E.shape[0]
+    if G < 2:
+        return np.zeros(G, dtype=np.float64)
+    L = cosine_kernel(E, eps)
+    full = _logdet_psd(L)
+    out = np.empty(G, dtype=np.float64)
+    idx = np.arange(G)
+    for i in range(G):
+        keep = idx[idx != i]
+        out[i] = full - _logdet_psd(L[np.ix_(keep, keep)])
+    return out
+
+
+def effective_rank(E: np.ndarray, eps: float = 1e-12) -> float:
+    """exp(Shannon entropy of the normalized Gram spectrum). Roy & Vetterli.
+
+    Reads as "how many distinct directions does this set effectively span":
+    1.0 for identical stories, G for mutually orthogonal ones, ~2 for two tight
+    clusters. Reported as a METRIC, never used as a reward (it is set-level, so
+    it is constant within a group and would produce zero advantage).
+
+    This exists because silhouette-selected k-means turned out to be unusable as
+    a mode counter on this data: a fully collapsed set and a fully spread set
+    both peak near silhouette 0.20 (at k=4 and k=7 respectively), because
+    k-means partitions isotropic data happily regardless of spread. Effective
+    rank has no k to select and degrades gracefully.
+    """
+    E = np.asarray(E, dtype=np.float64)
+    if E.shape[0] == 0:
+        return 0.0
+    if E.shape[0] == 1:
+        return 1.0
+    L = E @ E.T
+    w = np.linalg.eigvalsh(0.5 * (L + L.T))
+    w = np.clip(w, 0.0, None)
+    s = w.sum()
+    if s <= 0:
+        return 1.0
+    p = w / s
+    p = p[p > eps]
+    return float(np.exp(-(p * np.log(p)).sum()))
+
+
+def zscore(x: np.ndarray, ddof: int = 0) -> np.ndarray:
+    """Per-group z-score. Constant input -> all zeros (not NaN).
+
+    Used on m_i before it enters the reward: raw m_i lives on a log scale
+    with a long negative tail (a duplicate pair can hit log(1e-3) ~ -6.9),
+    which would otherwise dominate the quality term.
+    """
+    x = np.asarray(x, dtype=np.float64)
+    if x.size == 0:
+        return x
+    s = x.std(ddof=ddof)
+    if not np.isfinite(s) or s < 1e-12:
+        return np.zeros_like(x)
+    return (x - x.mean()) / s
+
+
+def scale_to_reference(x: np.ndarray, ref_scale: float = 1.0) -> np.ndarray:
+    """Rescale x to have unit-ish spread times ref_scale, preserving mean 0.
+
+    Deviation d_i lives in [0, 2] but in practice clusters in [0.1, 0.5] for
+    same-prompt stories -- an order of magnitude below judge quality (0-10).
+    Feeding raw d_i with alpha=0.5 would make the diversity term invisible.
+    """
+    return zscore(x) * float(ref_scale)
+
+
+def greedy_diverse_subset(
+    quality: np.ndarray,
+    E: np.ndarray,
+    k: int,
+    lam: float = 1.0,
+    eps: float = EPS_JITTER,
+) -> list[int]:
+    """Greedily pick k indices maximizing sum(quality) + lam * logdet(L_S).
+
+    Standard submodular greedy: the objective is monotone-ish and DPP logdet
+    is submodular, so greedy carries the usual (1 - 1/e) flavor of guarantee.
+    Used to build E3's multi-positive chosen sets.
+
+    Returns indices into the rows of E, in selection order.
+    """
+    quality = np.asarray(quality, dtype=np.float64)
+    E = np.asarray(E, dtype=np.float64)
+    G = E.shape[0]
+    k = int(min(k, G))
+    if k <= 0:
+        return []
+    L = cosine_kernel(E, eps)
+
+    chosen: list[int] = []
+    for _ in range(k):
+        best_gain, best_i = -np.inf, -1
+        for i in range(G):
+            if i in chosen:
+                continue
+            cand = chosen + [i]
+            vol = _logdet_psd(L[np.ix_(cand, cand)])
+            gain = quality[cand].sum() + lam * vol
+            if gain > best_gain:
+                best_gain, best_i = gain, i
+        if best_i < 0:
+            break
+        chosen.append(best_i)
+    return chosen
+
+
+def group_metrics(E: np.ndarray, eps: float = EPS_JITTER) -> dict:
+    """Bundle of reporting metrics for one group. Metrics only, not rewards."""
+    E = np.asarray(E, dtype=np.float64)
+    G = E.shape[0]
+    if G < 2:
+        return {"n": G, "mean_pairwise_dist": 0.0, "logdet": 0.0,
+                "mean_marginal": 0.0, "min_marginal": 0.0}
+    d = pairwise_deviation(E)
+    m = marginal_contributions(E, eps)
+    return {
+        "n": G,
+        "mean_pairwise_dist": float(d.mean()),
+        "logdet": float(logdet_volume(E, eps)),
+        "mean_marginal": float(m.mean()),
+        "min_marginal": float(m.min()),
+    }

src/evaluate.py

@@ -0,0 +1,320 @@
+"""
+Evaluation harness. Runs the IDENTICAL protocol on every model (base, E0, E1,
+E2, E3, E4-emb, E4-prob, and any 8B arms): 50 held-out prompts x 16 samples at
+fixed T=0.9, top_p=0.95, fixed seed.
+
+Metrics
+-------
+quality       mean judge quality over gate-passing stories
+gate_pass     fraction passing all programmatic gates
+ends_cleanly  fraction ending on terminal punctuation, not at the token cap
+pairwise      mean pairwise embedding distance within a prompt's 16 samples
+logdet        mean log-det volume of the 16-sample embedding set per prompt
+distinct4     unique 4-grams / total 4-grams, pooled across the 16 samples
+self_bleu     mean BLEU-4 of each sample against the other 15 (LOWER = diverse)
+n_clusters    conservative count of WELL-SEPARATED modes (silhouette > 0.50)
+eff_rank      effective rank of the 16-sample Gram spectrum; the primary,
+              continuous mode measure. 1 = all identical, 16 = all orthogonal
+tok_entropy   mean per-token predictive entropy, top-20 truncated. MONITOR ONLY.
+
+On tok_entropy: token entropy is not story diversity. A policy can raise token
+entropy by getting noisier inside a single narrative mode, and can lower it
+while telling structurally different stories. It is reported to detect
+degenerate sampling, and is never optimized.
+"""
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+import math
+import sys
+from collections import Counter
+from pathlib import Path
+
+import numpy as np
+
+ROOT = Path(__file__).resolve().parent.parent
+EVAL_TEMP, EVAL_TOP_P, EVAL_SEED, EVAL_N = 0.9, 0.95, 20260816, 16
+
+
+# ------------------------------------------------------------------ n-grams
+def ngrams(toks: list[str], n: int) -> list[tuple]:
+    return [tuple(toks[i:i + n]) for i in range(len(toks) - n + 1)]
+
+
+def norm_tokens(text: str) -> list[str]:
+    import re
+    return re.sub(r"[^\w\s]", " ", text.lower()).split()
+
+
+def distinct_n(texts: list[str], n: int = 4) -> float:
+    all_g = [g for t in texts for g in ngrams(norm_tokens(t), n)]
+    return len(set(all_g)) / len(all_g) if all_g else 0.0
+
+
+def _bleu4(cand: list[str], refs: list[list[str]]) -> float:
+    """BLEU-4 with brevity penalty, clipped counts against multiple refs."""
+    if len(cand) < 4:
+        return 0.0
+    logs = []
+    for n in range(1, 5):
+        cg = Counter(ngrams(cand, n))
+        if not cg:
+            return 0.0
+        maxref = Counter()
+        for r in refs:
+            for g, c in Counter(ngrams(r, n)).items():
+                if c > maxref[g]:
+                    maxref[g] = c
+        clipped = sum(min(c, maxref[g]) for g, c in cg.items())
+        total = sum(cg.values())
+        # smoothing: avoid log(0) collapsing the whole score
+        logs.append(math.log((clipped + 1e-9) / total) if clipped else math.log(1e-9))
+    r_len = min((len(r) for r in refs), key=lambda L: (abs(L - len(cand)), L))
+    bp = 1.0 if len(cand) > r_len else math.exp(1 - r_len / max(1, len(cand)))
+    # clamp: the 1e-9 smoothing can push an exact-match score a hair over 1.0
+    return min(1.0, max(0.0, bp * math.exp(sum(logs) / 4)))
+
+
+def self_bleu(texts: list[str]) -> float:
+    toks = [norm_tokens(t) for t in texts]
+    if len(toks) < 2:
+        return 0.0
+    return float(np.mean([
+        _bleu4(toks[i], [toks[j] for j in range(len(toks)) if j != i])
+        for i in range(len(toks))
+    ]))
+
+
+# Silhouette threshold for declaring that real cluster structure exists.
+# Calibrated, not guessed: on 16 synthetic embeddings, a fully COLLAPSED set
+# peaks at silhouette 0.202 (k=4) and a fully SPREAD set at 0.195 (k=7) -- i.e.
+# silhouette cannot tell those apart at all, because k-means partitions
+# isotropic data regardless of spread. Only genuinely separated clusters score
+# high (two clear modes -> 0.976). So the threshold is set well above the
+# isotropic band, making n_clusters a CONSERVATIVE count of well-separated
+# modes: it returns 1 unless the structure is unmistakable. `eff_rank` is the
+# primary, continuous mode measure.
+SILHOUETTE_MIN = 0.50
+
+
+def cluster_count(E: np.ndarray, kmax: int = 8) -> int:
+    """k-means with silhouette-selected k. Conservative count of *well-separated*
+    modes; returns 1 when there is no unmistakable cluster structure."""
+    from sklearn.cluster import KMeans
+    from sklearn.metrics import silhouette_score
+    n = E.shape[0]
+    if n < 4:
+        return 1
+    best_k, best_s = 1, -1.0
+    for k in range(2, min(kmax, n - 1) + 1):
+        try:
+            lab = KMeans(n_clusters=k, n_init=10, random_state=0).fit_predict(E)
+            if len(set(lab)) < 2:
+                continue
+            s = silhouette_score(E, lab, metric="cosine")
+        except Exception:
+            continue
+        if s > best_s:
+            best_k, best_s = k, s
+    return best_k if best_s > SILHOUETTE_MIN else 1
+
+
+def topk_entropy(logprob_rows: list[dict]) -> float:
+    """Mean per-token entropy from vLLM's top-k logprob table.
+
+    Truncated at k, so it UNDERSTATES true entropy; comparable across models
+    only because every model is measured with the same k. Monitor only.
+    """
+    ents = []
+    for row in logprob_rows:
+        lps = np.array([v for v in row.values()], dtype=np.float64)
+        if lps.size == 0:
+            continue
+        p = np.exp(lps)
+        s = p.sum()
+        if s <= 0:
+            continue
+        p = p / s
+        ents.append(float(-(p * np.log(p + 1e-12)).sum()))
+    return float(np.mean(ents)) if ents else 0.0
+
+
+# --------------------------------------------------------------------- main
+def eval_one(model_id: str, lora_path: str | None, label: str, prompts: list[dict],
+             judge, enc, gpu_mem: float, dump_dir: Path, n_dump: int = 3) -> dict:
+    import gates
+    from diversity import (effective_rank, l2_normalize, logdet_volume,
+                           pairwise_deviation)
+    from generate import build_llm, render_chat
+    from vllm import SamplingParams
+
+    from transformers import AutoTokenizer
+    tok = AutoTokenizer.from_pretrained(model_id)
+    llm = build_llm(model_id, gpu_mem_util=gpu_mem, seed=EVAL_SEED,
+                    enable_lora=bool(lora_path))
+
+    kw = {}
+    if lora_path:
+        from vllm.lora.request import LoRARequest
+        kw["lora_request"] = LoRARequest(label, 1, lora_path)
+
+    # logprobs=10, not 20: this table is only used for tok_entropy, which is
+    # monitor-only and never optimized. 20 would materialize ~8M Python Logprob
+    # objects per model (800 seqs x ~500 tok x 20). The truncation understates
+    # entropy identically for every model, so cross-model comparison holds.
+    # max_tokens matches training (1024) so eval and train share a length regime.
+    sp = SamplingParams(n=EVAL_N, temperature=EVAL_TEMP, top_p=EVAL_TOP_P,
+                        max_tokens=1024, seed=EVAL_SEED, logprobs=10,
+                        skip_special_tokens=True)
+    outs = llm.generate([render_chat(tok, p["prompt"]) for p in prompts], sp, **kw)
+
+    records, per_prompt = [], []
+    for p, out in zip(prompts, outs):
+        texts, ents, frs = [], [], []
+        for o in out.outputs:
+            texts.append(o.text.strip())
+            frs.append(o.finish_reason or "")
+            ents.append(topk_entropy(
+                [{k: v.logprob for k, v in d.items()} for d in (o.logprobs or [])]))
+        grs = [gates.check(t, finish_reason=f) for t, f in zip(texts, frs)]
+        E = l2_normalize(np.asarray(enc.encode(
+            texts, normalize_embeddings=True, batch_size=32,
+            show_progress_bar=False, convert_to_numpy=True), dtype=np.float64))
+        per_prompt.append({
+            "prompt_id": p["id"], "prompt": p["prompt"], "texts": texts,
+            "gates": [g.as_dict() for g in grs],
+            "pairwise": float(pairwise_deviation(E).mean()),
+            "logdet": float(logdet_volume(E)),
+            "distinct4": distinct_n(texts, 4),
+            "self_bleu": self_bleu(texts),
+            "n_clusters": cluster_count(E),
+            "eff_rank": effective_rank(E),
+            "tok_entropy": float(np.mean(ents)) if ents else 0.0,
+            "gate_pass": float(np.mean([g.passed for g in grs])),
+            "ends_cleanly": float(np.mean([g.completeness for g in grs])),
+            "words": float(np.mean([g.n_words for g in grs])),
+        })
+        for t, g in zip(texts, grs):
+            records.append((p["prompt"], t, g.passed))
+
+    del llm
+    import gc, torch
+    gc.collect(); torch.cuda.empty_cache()
+
+    # judge only gate-passers, same rule as training
+    idx = [i for i, (_, _, ok) in enumerate(records) if ok]
+    scores = judge.score_many_sync([(records[i][0], records[i][1]) for i in idx]) if idx else []
+    q = [s.quality for s in scores]
+    nov = [s.novelty for s in scores]
+
+    row = {
+        "model": label,
+        "quality": float(np.mean(q)) if q else 0.0,
+        "quality_sd": float(np.std(q)) if q else 0.0,
+        "novelty": float(np.mean(nov)) if nov else 0.0,
+        "gate_pass": float(np.mean([r["gate_pass"] for r in per_prompt])),
+        "ends_cleanly": float(np.mean([r["ends_cleanly"] for r in per_prompt])),
+        "pairwise": float(np.mean([r["pairwise"] for r in per_prompt])),
+        "logdet": float(np.mean([r["logdet"] for r in per_prompt])),
+        "distinct4": float(np.mean([r["distinct4"] for r in per_prompt])),
+        "self_bleu": float(np.mean([r["self_bleu"] for r in per_prompt])),
+        "n_clusters": float(np.mean([r["n_clusters"] for r in per_prompt])),
+        "eff_rank": float(np.mean([r["eff_rank"] for r in per_prompt])),
+        "tok_entropy": float(np.mean([r["tok_entropy"] for r in per_prompt])),
+        "words": float(np.mean([r["words"] for r in per_prompt])),
+        "n_stories": len(records),
+        "n_judged": len(idx),
+    }
+
+    dump_dir.mkdir(parents=True, exist_ok=True)
+    json.dump(per_prompt, open(dump_dir / f"{label}_full.json", "w"), indent=1)
+    json.dump(per_prompt[:n_dump], open(dump_dir / f"{label}_examples.json", "w"), indent=1)
+    return row
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--models", required=True,
+                    help="JSON list of {label, model, lora} or path to such a file")
+    ap.add_argument("--limit", type=int, default=None)
+    ap.add_argument("--gpu-mem", type=float, default=0.85)
+    ap.add_argument("--out", default="outputs/eval")
+    ap.add_argument("--single", default=None,
+                    help="internal: evaluate exactly this label, then exit")
+    args = ap.parse_args()
+
+    import logbook
+    from data import load_prompts
+    from judge import build_judge
+    from sentence_transformers import SentenceTransformer
+
+    spec = json.loads(Path(args.models).read_text()) if Path(args.models).exists() \
+        else json.loads(args.models)
+    prompts = load_prompts("eval", ROOT / "data")
+    if args.limit:
+        prompts = prompts[: args.limit]
+
+    judge = build_judge(cache_path=str(ROOT / "cache" / "judge.sqlite"), concurrency=24)
+    enc = SentenceTransformer("BAAI/bge-base-en-v1.5", device="cuda")
+    out_dir = ROOT / args.out
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    rowdir = out_dir / "rows"
+    rowdir.mkdir(parents=True, exist_ok=True)
+
+    # ---- child mode: evaluate one model, write its row, exit --------------
+    if args.single:
+        m = next(x for x in spec if x["label"] == args.single)
+        lora = m.get("lora")
+        if lora and not Path(ROOT / lora).exists() and not Path(lora).exists():
+            print(f"SKIP {m['label']}: adapter not found at {lora}")
+            return 3
+        r = eval_one(m["model"], lora, m["label"], prompts, judge, enc,
+                     args.gpu_mem, out_dir / "samples")
+        json.dump(r, open(rowdir / f"{m['label']}.json", "w"), indent=1)
+        print(json.dumps(r, indent=1))
+        print("judge:", judge.health(), judge.cost_estimate(0.140, 0.280))
+        return 0
+
+    # ---- driver mode: one SUBPROCESS per model ---------------------------
+    # vLLM v1 runs EngineCore in a child process; `del llm` does not reliably
+    # reclaim its GPU memory, so a 7-model in-process loop OOMs on model 2 after
+    # model 1 succeeds. Process isolation makes teardown unconditional. It also
+    # means one model crashing cannot take the whole eval down.
+    import subprocess
+    for m in spec:
+        dest = rowdir / f"{m['label']}.json"
+        if dest.exists():
+            print(f"== SKIP {m['label']} (already evaluated)"); continue
+        print(f"\n===== eval {m['label']} (subprocess) =====", flush=True)
+        cmd = [sys.executable, "-u", str(Path(__file__).resolve()),
+               "--models", args.models, "--out", args.out,
+               "--gpu-mem", str(args.gpu_mem), "--single", m["label"]]
+        if args.limit:
+            cmd += ["--limit", str(args.limit)]
+        rc = subprocess.run(cmd).returncode
+        if rc != 0:
+            print(f"!! {m['label']} exited {rc}; continuing with the rest")
+            logbook.note(f"eval FAILED: {m['label']}", f"exit code {rc}", level="ALERT")
+
+    rows = []
+    for m in spec:
+        p = rowdir / f"{m['label']}.json"
+        if p.exists():
+            rows.append(json.load(open(p)))
+    if not rows:
+        print("no models evaluated"); return 1
+    with open(out_dir / "results.csv", "w", newline="") as f:
+        w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
+        w.writeheader(); w.writerows(rows)
+    print(f"\nwrote {out_dir/'results.csv'} ({len(rows)}/{len(spec)} models)")
+    logbook.note("eval complete",
+                 f"{len(rows)}/{len(spec)} models\n\n" +
+                 logbook.table(rows, list(rows[0].keys())))
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

src/gates.py

@@ -0,0 +1,179 @@
+"""
+Programmatic quality gates. Free, deterministic, zero judge variance.
+
+Contract: a story that FAILS any hard gate gets total reward 0, regardless of
+what the judge said. Gates run before the judge so we can skip paying for
+obviously-broken generations.
+
+This module exists because of a specific autopsy finding. The prior run's
+completeness check was:
+
+    enders = ".!?\\"'"
+    text[-1] in enders
+
+which uses straight quotes only. The policy emits curly U+201D constantly, so
+properly-ended stories were scored 0.3 instead of 1.0. Worse, 100% of prior-run
+eval samples were truncated mid-word by a max_new_tokens wall, and the judge
+rubric capped quality at 4 for truncation -- collapsing within-group quality
+variance to ~0, which in GRPO means zero advantage. The quality signal was
+silently dead for the entire run. Everything here is built to make that
+failure loud instead of silent.
+"""
+
+from __future__ import annotations
+
+import math
+import re
+import unicodedata
+from dataclasses import dataclass, asdict, field
+
+# Sentence terminators, including the curly/CJK variants the prior run missed.
+TERMINALS = set('.!?…' + '"”’\'' + '。！？' + ')]}*_')
+# Characters allowed to trail a real terminator (markdown emphasis, quotes).
+_TRAILING_DECOR = set('*_`"”’\'» )]}')
+
+MIN_WORDS = 150
+MAX_WORDS = 600
+MAX_NGRAM_REPEAT_FRAC = 0.18   # frac of 4-grams that are repeats
+MIN_CHAR_ENTROPY = 3.2         # bits/char; English prose sits ~4.0-4.4
+MAX_NONASCII_FRAC = 0.08       # em-dashes/curly quotes are fine; CJK dumps are not
+MAX_LINE_REPEAT_FRAC = 0.25
+
+
+@dataclass
+class GateResult:
+    passed: bool
+    completeness: float                    # 1.0 ended cleanly, 0.0 truncated
+    n_words: int
+    reasons: list[str] = field(default_factory=list)
+    # diagnostics, logged not rewarded
+    repeat_4gram_frac: float = 0.0
+    char_entropy: float = 0.0
+    nonascii_frac: float = 0.0
+    hit_token_cap: bool = False
+
+    def as_dict(self) -> dict:
+        return asdict(self)
+
+
+def _words(text: str) -> list[str]:
+    return text.split()
+
+
+def _norm_tokens(text: str) -> list[str]:
+    return re.sub(r"[^\w\s]", " ", text.lower()).split()
+
+
+def ends_cleanly(text: str) -> bool:
+    """True if the story ends on a sentence terminator.
+
+    Strips trailing markdown/quote decoration first, so `...over.*` and
+    `...done."` and `...gone.”` all count. Also rejects the specific
+    mid-word cutoff signature: last token is a bare word with no terminator.
+    """
+    t = text.rstrip()
+    if not t:
+        return False
+    while t and t[-1] in _TRAILING_DECOR:
+        t = t[:-1].rstrip()
+        if t and t[-1] in '.!?…。！？':
+            return True
+    return bool(t) and t[-1] in '.!?…。！？'
+
+
+def char_entropy(text: str) -> float:
+    """Shannon entropy over characters, bits/char."""
+    if not text:
+        return 0.0
+    counts: dict[str, int] = {}
+    for ch in text:
+        counts[ch] = counts.get(ch, 0) + 1
+    n = len(text)
+    return -sum((c / n) * math.log2(c / n) for c in counts.values())
+
+
+def repeat_ngram_frac(text: str, n: int = 4) -> float:
+    """Fraction of n-grams that are non-first occurrences. Catches loops."""
+    toks = _norm_tokens(text)
+    if len(toks) < n + 1:
+        return 0.0
+    grams = [tuple(toks[i:i + n]) for i in range(len(toks) - n + 1)]
+    return 1.0 - (len(set(grams)) / len(grams))
+
+
+def repeat_line_frac(text: str) -> float:
+    """Fraction of non-trivial lines that are duplicates. Catches list loops."""
+    lines = [l.strip() for l in text.splitlines() if len(l.strip()) > 12]
+    if len(lines) < 4:
+        return 0.0
+    return 1.0 - (len(set(lines)) / len(lines))
+
+
+def nonascii_frac(text: str) -> float:
+    """Fraction of chars outside Latin/punctuation. Typographic marks exempt."""
+    if not text:
+        return 0.0
+    bad = 0
+    for ch in text:
+        if ord(ch) < 128:
+            continue
+        if unicodedata.category(ch) in ("Pd", "Pi", "Pf", "Po", "Zs", "Sm"):
+            continue  # em dash, curly quotes, ellipsis, nbsp
+        bad += 1
+    return bad / len(text)
+
+
+def check(
+    text: str,
+    *,
+    finish_reason: str | None = None,
+    min_words: int = MIN_WORDS,
+    max_words: int = MAX_WORDS,
+) -> GateResult:
+    """Run all gates. `finish_reason` from vLLM/OpenAI ('length' == hit cap).
+
+    Passing `finish_reason` is strongly preferred: it detects the token-cap
+    truncation that destroyed the prior run directly, rather than inferring it
+    from punctuation.
+    """
+    text = (text or "").strip()
+    reasons: list[str] = []
+
+    n_words = len(_words(text))
+    hit_cap = finish_reason == "length"
+    complete = ends_cleanly(text) and not hit_cap
+
+    rep4 = repeat_ngram_frac(text, 4)
+    repline = repeat_line_frac(text)
+    ent = char_entropy(text)
+    na = nonascii_frac(text)
+
+    if not text:
+        reasons.append("empty")
+    if hit_cap:
+        reasons.append("hit_token_cap")
+    if not ends_cleanly(text):
+        reasons.append("no_terminal_punctuation")
+    if n_words < min_words:
+        reasons.append(f"too_short({n_words}<{min_words})")
+    if n_words > max_words:
+        reasons.append(f"too_long({n_words}>{max_words})")
+    if rep4 > MAX_NGRAM_REPEAT_FRAC:
+        reasons.append(f"ngram_loop({rep4:.2f})")
+    if repline > MAX_LINE_REPEAT_FRAC:
+        reasons.append(f"line_loop({repline:.2f})")
+    if text and ent < MIN_CHAR_ENTROPY:
+        reasons.append(f"low_entropy({ent:.2f})")
+    if na > MAX_NONASCII_FRAC:
+        reasons.append(f"nonascii({na:.2f})")
+
+    return GateResult(
+        passed=not reasons,
+        completeness=1.0 if complete else 0.0,
+        n_words=n_words,
+        reasons=reasons,
+        repeat_4gram_frac=rep4,
+        char_entropy=ent,
+        nonascii_frac=na,
+        hit_token_cap=hit_cap,
+    )

src/generate.py

@@ -0,0 +1,156 @@
+"""
+vLLM generation wrapper shared by pool building, eval, and smoke tests.
+
+Captures per-story cumulative logprob, which E4's `divpo-prob` variant needs:
+DivPO's probability criterion picks the LOWEST-logprob passing story as the
+diverse chosen and the HIGHEST-logprob failing story as the common rejected.
+The highest-probability sample in a temperature pool is, by construction, the
+near-greedy one -- which is why this is the principled version of "reject the
+greedy decode".
+
+Length normalization matters here. Raw cumulative logprob scales with token
+count, so ranking by it would just rank by length (short stories win). We keep
+both `cumlogprob` and `mean_logprob = cumlogprob / n_tokens` and use the mean
+for DivPO ranking.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, asdict
+from pathlib import Path
+
+# NEVER truncate a story. The prior run died of a max_new_tokens wall, so this
+# is sized for real headroom, not for the target length: a 500-word story is
+# ~650-700 Qwen tokens, so 768 would clip anything that runs long and would
+# reproduce the exact bug we are fixing. 1280 lets the model finish and overrun
+# a little; overrun is then DETECTED by the >600-word gate rather than silently
+# chopped. Measure the failure, don't manufacture it.
+#
+# Nothing downstream truncates either: the judge receives the complete story
+# text, and embeddings are computed over the complete story text.
+MAX_NEW_TOKENS = 1280
+MAX_MODEL_LEN = 2048
+
+
+@dataclass
+class Gen:
+    prompt_id: str
+    prompt: str
+    idx: int
+    text: str
+    n_tokens: int
+    cumlogprob: float
+    mean_logprob: float
+    finish_reason: str
+
+    def as_dict(self) -> dict:
+        return asdict(self)
+
+
+def build_llm(
+    model: str,
+    gpu_mem_util: float = 0.85,
+    max_model_len: int = MAX_MODEL_LEN,
+    seed: int = 0,
+    enable_lora: bool = False,
+):
+    from vllm import LLM
+    kw = dict(
+        model=model,
+        dtype="bfloat16",
+        gpu_memory_utilization=gpu_mem_util,
+        max_model_len=max_model_len,
+        seed=seed,
+        enforce_eager=False,
+        disable_log_stats=True,
+    )
+    if enable_lora:
+        kw.update(enable_lora=True, max_lora_rank=32)
+    return LLM(**kw)
+
+
+def sampling_params(
+    n: int, temperature: float, top_p: float, seed: int | None,
+    max_tokens: int = MAX_NEW_TOKENS,
+):
+    from vllm import SamplingParams
+    return SamplingParams(
+        n=n,
+        temperature=temperature,
+        top_p=top_p,
+        max_tokens=max_tokens,
+        seed=seed,
+        logprobs=0,            # cumulative_logprob only; per-token table not needed
+        skip_special_tokens=True,
+    )
+
+
+def render_chat(tokenizer, prompt: str) -> str:
+    """Render the chat template, with thinking OFF where the model has it.
+
+    Qwen3-8B is a hybrid-thinking model: left at its default it emits a <think>
+    block before the story, which would (a) eat the token budget, (b) pollute
+    the embedding with reasoning text, and (c) make the 8B arm incomparable to
+    the 4B-Instruct arms, which have no thinking mode at all. Qwen3-4B-Instruct
+    ignores the kwarg, so the same call is correct for both.
+    """
+    from data import chat_messages
+    msgs = chat_messages(prompt)
+    try:
+        return tokenizer.apply_chat_template(
+            msgs, tokenize=False, add_generation_prompt=True, enable_thinking=False,
+        )
+    except TypeError:
+        return tokenizer.apply_chat_template(
+            msgs, tokenize=False, add_generation_prompt=True,
+        )
+
+
+def generate(
+    llm,
+    tokenizer,
+    prompts: list[dict],
+    n: int,
+    temperature: float,
+    top_p: float,
+    seed: int | None = None,
+    max_tokens: int = MAX_NEW_TOKENS,
+    lora_path: str | None = None,
+) -> list[Gen]:
+    """prompts = [{'id':..., 'prompt':...}]  ->  flat list of n*len(prompts) Gens."""
+    texts = [render_chat(tokenizer, p["prompt"]) for p in prompts]
+    sp = sampling_params(n, temperature, top_p, seed, max_tokens)
+
+    kw = {}
+    if lora_path:
+        from vllm.lora.request import LoRARequest
+        kw["lora_request"] = LoRARequest("adapter", 1, lora_path)
+
+    outs = llm.generate(texts, sp, **kw)
+
+    gens: list[Gen] = []
+    for p, out in zip(prompts, outs):
+        for j, o in enumerate(out.outputs):
+            ntok = len(o.token_ids)
+            cum = float(o.cumulative_logprob) if o.cumulative_logprob is not None else 0.0
+            gens.append(Gen(
+                prompt_id=p["id"], prompt=p["prompt"], idx=j,
+                text=o.text.strip(), n_tokens=ntok,
+                cumlogprob=cum,
+                mean_logprob=(cum / ntok) if ntok else 0.0,
+                finish_reason=o.finish_reason or "",
+            ))
+    return gens
+
+
+def save_gens(gens: list[Gen], path: str | Path) -> None:
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with open(path, "w") as f:
+        for g in gens:
+            f.write(json.dumps(g.as_dict()) + "\n")
+
+
+def load_gens(path: str | Path) -> list[Gen]:
+    return [Gen(**json.loads(l)) for l in open(path) if l.strip()]

src/judge.py

@@ -0,0 +1,331 @@
+"""
+Pluggable LLM judge with aggressive on-disk caching.
+
+Design decisions worth stating, because they differ from the prior run:
+
+1. ONE STORY PER CALL, absolute rubric. The prior run sent all 16 samples in a
+   single call and asked for relative scores. That is cheaper but wrong here:
+   - quality is compared against a fixed gate tau (5.0), so it must be
+     absolutely calibrated, not calibrated relative to whatever else happened
+     to be in the group;
+   - eval compares quality ACROSS models, which relative scoring makes
+     meaningless;
+   - batched scoring has position bias.
+   Per-story also makes the cache actually work: identical stories recur across
+   arms and eval passes, and each is paid for once, ever.
+
+2. NO group_diversity FIELD. A set-level scalar is constant within a GRPO group
+   -> zero within-group std -> zero advantage. It cannot train anything. All
+   diversity credit comes from embeddings (src/diversity.py), where it is
+   per-sample by construction.
+
+3. Judge novelty is an ABSOLUTE 'freshness vs. generic AI slop' axis, not
+   'different from the others'. The latter is the embeddings' job. This keeps
+   the two signals from being redundant-but-inconsistent.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import json
+import os
+import sqlite3
+import time
+from dataclasses import dataclass, asdict
+from pathlib import Path
+
+RUBRIC_VERSION = "v4"
+
+SYSTEM_RUBRIC = """You are a demanding fiction editor scoring ONE short story written to a prompt.
+Score on absolute standards, not relative to anything you have seen before.
+
+STEP 1 -- FINISHEDNESS. Read the LAST sentence first, before anything else.
+A story is UNFINISHED if the final sentence stops mid-clause or mid-word, or if
+the piece simply halts with the situation unresolved and no closing beat.
+Beautiful prose does not make an unfinished story finished. If UNFINISHED:
+set "ended":false, quality 1-3, novelty 0-3, and stop deliberating -- an
+unfinished fragment is never novel, because you cannot know what it became.
+Only if it is FINISHED, continue to step 2.
+
+STEP 2 -- quality (0-10): coherence + prompt fulfilment + craft
+ 0-2  Incoherent, or ignores the prompt entirely.
+ 3-4  Followable but generic: stock premise, flat prose, perfunctory ending.
+ 5-6  Competent. Does what the prompt asks, reads cleanly, resolves. Forgettable.
+ 7-8  Strong. Controlled voice, purposeful structure, an ending that lands.
+ 9-10 Excellent. Precise images, earned emotional turn, nothing wasted.
+If it never engages the prompt's actual premise, quality <= 4.
+
+STEP 3 -- novelty (0-10): freshness of premise, voice and construction vs. generic AI fiction
+ 0-2  Pure slop signature: portentous one-line paragraphs, "It wasn't X, it was Y",
+      italicised abstract nouns doing the emotional work, cosmic-melancholy fog.
+ 3-4  Familiar premise handled in the familiar way.
+ 5-6  One genuine choice (an unusual angle OR a distinct voice OR an odd structure).
+ 7-8  Several: takes a real swing on premise, form or register and controls it.
+ 9-10 Startling and still coherent. You could not have predicted this.
+Novelty is NOT weirdness: incoherence, non-sequiturs and word salad score 0-2.
+Judge freshness against fiction in general, NOT against other samples.
+
+Return ONLY minified JSON, no markdown fence:
+{"quality":<0-10>,"novelty":<0-10>,"ended":<true|false>,"note":"<=10 words"}"""
+
+
+@dataclass
+class JudgeScore:
+    quality: float
+    novelty: float
+    ended: bool = True
+    note: str = ""
+    ok: bool = True          # False => call failed, scores are neutral fallbacks
+    cached: bool = False
+
+    def as_dict(self) -> dict:
+        return asdict(self)
+
+
+NEUTRAL = JudgeScore(quality=5.0, novelty=5.0, ended=True, note="judge_failed", ok=False)
+
+
+# --------------------------------------------------------------- disk cache
+class ScoreCache:
+    """sqlite keyed by sha256(model|rubric|prompt|story). Survives restarts."""
+
+    def __init__(self, path: str | Path):
+        self.path = str(path)
+        Path(self.path).parent.mkdir(parents=True, exist_ok=True)
+        self._db = sqlite3.connect(self.path, check_same_thread=False)
+        self._db.execute("PRAGMA journal_mode=WAL")
+        self._db.execute(
+            "CREATE TABLE IF NOT EXISTS scores ("
+            " k TEXT PRIMARY KEY, quality REAL, novelty REAL,"
+            " ended INT, note TEXT, ts REAL)"
+        )
+        self._db.commit()
+        self.hits = 0
+        self.misses = 0
+
+    @staticmethod
+    def key(model: str, prompt: str, story: str) -> str:
+        h = hashlib.sha256()
+        h.update(f"{model}\x00{RUBRIC_VERSION}\x00{prompt}\x00{story}".encode())
+        return h.hexdigest()
+
+    def get(self, k: str) -> JudgeScore | None:
+        r = self._db.execute(
+            "SELECT quality,novelty,ended,note FROM scores WHERE k=?", (k,)
+        ).fetchone()
+        if r is None:
+            self.misses += 1
+            return None
+        self.hits += 1
+        return JudgeScore(quality=r[0], novelty=r[1], ended=bool(r[2]),
+                          note=r[3] or "", ok=True, cached=True)
+
+    def put(self, k: str, s: JudgeScore) -> None:
+        if not s.ok:
+            return  # never cache a failure
+        self._db.execute(
+            "INSERT OR REPLACE INTO scores VALUES (?,?,?,?,?,?)",
+            (k, s.quality, s.novelty, int(s.ended), s.note[:120], time.time()),
+        )
+        self._db.commit()
+
+    def stats(self) -> dict:
+        n = self._db.execute("SELECT COUNT(*) FROM scores").fetchone()[0]
+        tot = self.hits + self.misses
+        return {"rows": n, "hits": self.hits, "misses": self.misses,
+                "hit_rate": (self.hits / tot) if tot else 0.0}
+
+
+# ------------------------------------------------------------------- judges
+def _parse(text: str) -> JudgeScore:
+    t = (text or "").strip()
+    if t.startswith("```"):
+        t = t.split("\n", 1)[-1].rsplit("```", 1)[0].strip()
+    i, j = t.find("{"), t.rfind("}")
+    if i >= 0 and j > i:
+        t = t[i:j + 1]
+    d = json.loads(t)
+
+    def num(v, lo=0.0, hi=10.0):
+        return max(lo, min(hi, float(v)))
+
+    return JudgeScore(
+        quality=num(d.get("quality", 5)),
+        novelty=num(d.get("novelty", d.get("novel", 5))),
+        ended=bool(d.get("ended", True)),
+        note=str(d.get("note", ""))[:120],
+        ok=True,
+    )
+
+
+class OpenRouterJudge:
+    """Async judge over OpenRouter. Default backend."""
+
+    BASE = "https://openrouter.ai/api/v1/chat/completions"
+
+    def __init__(
+        self,
+        model: str = "deepseek/deepseek-v4-flash-0731",
+        api_key: str | None = None,
+        cache_path: str | Path = "cache/judge.sqlite",
+        concurrency: int = 12,
+        max_retries: int = 4,
+        temperature: float = 0.0,
+        timeout: float = 120.0,
+    ):
+        self.model = model
+        self.api_key = api_key or os.environ.get("OPENROUTER_API_KEY", "")
+        if not self.api_key:
+            raise RuntimeError("OPENROUTER_API_KEY not set")
+        self.cache = ScoreCache(cache_path)
+        # NOT created here. asyncio.Semaphore binds to the loop that first
+        # awaits it, and score_many_sync() calls asyncio.run() -> a NEW loop
+        # every training step. A semaphore built in __init__ survives step 1 and
+        # then raises "bound to a different event loop" on step 2, which our
+        # retry wrapper would convert into neutral 5.0 scores forever. Built
+        # per-call in score_many() instead.
+        self.concurrency = concurrency
+        self.sem: asyncio.Semaphore | None = None
+        self.max_retries = max_retries
+        self.temperature = temperature
+        self.timeout = timeout
+        self.n_calls = 0
+        self.tok_in = 0
+        self.tok_out = 0
+        self.n_failed = 0
+        self.n_empty = 0
+        self.n_truncated = 0
+        self.last_error = ""
+
+    def _user(self, prompt: str, story: str) -> str:
+        return f"WRITING PROMPT:\n{prompt}\n\nSTORY:\n{story}\n\nJSON only:"
+
+    async def _one(self, client, prompt: str, story: str) -> JudgeScore:
+        payload = {
+            "model": self.model,
+            "messages": [
+                {"role": "system", "content": SYSTEM_RUBRIC},
+                {"role": "user", "content": self._user(prompt, story)},
+            ],
+            "temperature": self.temperature,
+            # deepseek-v4-flash-0731 is a REASONING model. Left alone it spends
+            # its whole output budget on `reasoning` and returns content=null
+            # with finish_reason="length" -- which our fallback silently turned
+            # into a neutral 5.0, i.e. a judge that scored everything identically
+            # while looking like it worked. Calibration caught it at 4/166 calls
+            # succeeding. Disabling reasoning is also 25x cheaper and 5x faster
+            # (30 output tokens vs 764).
+            "reasoning": {"enabled": False},
+            # Generous bound on a structured output that empirically needs ~30
+            # tokens. Not a length cap on content we care about -- and if it is
+            # ever hit, `n_truncated` below makes it loud instead of silent.
+            "max_tokens": 300,
+        }
+        delay = 2.0
+        for attempt in range(self.max_retries):
+            try:
+                async with self.sem:
+                    r = await client.post(
+                        self.BASE, json=payload, timeout=self.timeout,
+                        headers={"Authorization": f"Bearer {self.api_key}"},
+                    )
+                if r.status_code in (429, 500, 502, 503, 529):
+                    await asyncio.sleep(delay); delay *= 2
+                    continue
+                r.raise_for_status()
+                body = r.json()
+                self.n_calls += 1
+                u = body.get("usage") or {}
+                self.tok_in += int(u.get("prompt_tokens", 0) or 0)
+                self.tok_out += int(u.get("completion_tokens", 0) or 0)
+                ch = body["choices"][0]
+                if ch.get("finish_reason") == "length":
+                    self.n_truncated += 1
+                content = ch["message"].get("content")
+                if not content:
+                    self.n_empty += 1
+                    raise ValueError("empty content (reasoning ate the budget?)")
+                return _parse(content)
+            except Exception as e:
+                self.last_error = f"{type(e).__name__}: {str(e)[:140]}"
+                if attempt == self.max_retries - 1:
+                    self.n_failed += 1
+                    return NEUTRAL
+                await asyncio.sleep(delay); delay *= 2
+        self.n_failed += 1
+        return NEUTRAL
+
+    def health(self) -> dict:
+        """Fail-loud accounting. A judge that returns neutral 5.0 for every call
+        looks exactly like a judge that works, until you check this."""
+        att = self.n_calls + self.n_failed
+        return {"calls_ok": self.n_calls, "failed": self.n_failed,
+                "empty_content": self.n_empty, "truncated": self.n_truncated,
+                "fail_rate": (self.n_failed / att) if att else 0.0,
+                "last_error": self.last_error}
+
+    def assert_healthy(self, max_fail_rate: float = 0.05) -> None:
+        h = self.health()
+        if h["fail_rate"] > max_fail_rate:
+            raise RuntimeError(
+                f"JUDGE UNHEALTHY: {h}. Refusing to train on neutral fallbacks -- "
+                f"a constant reward column produces zero GRPO advantage."
+            )
+
+    async def score_many(self, pairs: list[tuple[str, str]]) -> list[JudgeScore]:
+        """pairs = [(prompt, story), ...] -> scores in the same order."""
+        import httpx
+
+        out: list[JudgeScore | None] = [None] * len(pairs)
+        todo: list[int] = []
+        for i, (p, s) in enumerate(pairs):
+            hit = self.cache.get(self.cache.key(self.model, p, s))
+            if hit is not None:
+                out[i] = hit
+            else:
+                todo.append(i)
+
+        if todo:
+            self.sem = asyncio.Semaphore(self.concurrency)   # bind to THIS loop
+
+            async def one_and_cache(i):
+                """Persist each score the moment it lands, not after the whole
+                batch. The pool stage issues ~16k calls in a single
+                score_many(); caching only at the end meant a crash at 99%
+                threw away every call. Incremental writes make any failure
+                resumable for free -- a re-run replays as cache hits."""
+                s = await self._one(client, pairs[i][0], pairs[i][1])
+                if s.ok:
+                    self.cache.put(self.cache.key(self.model, *pairs[i]), s)
+                return i, s
+
+            async with httpx.AsyncClient(http2=False) as client:
+                for fut in asyncio.as_completed([one_and_cache(i) for i in todo]):
+                    i, s = await fut
+                    out[i] = s
+
+        return [o if o is not None else NEUTRAL for o in out]
+
+    def score_many_sync(self, pairs: list[tuple[str, str]]) -> list[JudgeScore]:
+        """Blocking wrapper for use inside TRL reward functions."""
+        try:
+            asyncio.get_running_loop()
+        except RuntimeError:
+            return asyncio.run(self.score_many(pairs))
+        # already inside a loop (rare in TRL): run in a private one
+        import concurrent.futures
+        with concurrent.futures.ThreadPoolExecutor(1) as ex:
+            return ex.submit(asyncio.run, self.score_many(pairs)).result()
+
+    def cost_estimate(self, price_in: float, price_out: float) -> dict:
+        """price_* in $ per 1M tokens."""
+        c = self.tok_in / 1e6 * price_in + self.tok_out / 1e6 * price_out
+        return {"calls": self.n_calls, "tok_in": self.tok_in,
+                "tok_out": self.tok_out, "usd": round(c, 4),
+                **self.cache.stats()}
+
+
+def build_judge(**kw) -> OpenRouterJudge:
+    """Factory. Extend here if an Anthropic/OpenAI key shows up."""
+    return OpenRouterJudge(**kw)

src/logbook.py

@@ -0,0 +1,125 @@
+"""
+Lab-notebook helper. Everything the run learns goes to logs/ as markdown +
+figures, so the record survives the process that produced it.
+
+Layout:
+  logs/lab_notebook.md          chronological running journal (append-only)
+  logs/experiments/<name>.md    one long report per experiment
+  logs/figures/<name>.png       figures
+  logs/disk.md                  space accounting
+"""
+
+from __future__ import annotations
+
+import json
+import subprocess
+import time
+from datetime import datetime, timezone
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parent.parent
+LOGS = ROOT / "logs"
+FIGS = LOGS / "figures"
+EXPS = LOGS / "experiments"
+NOTEBOOK = LOGS / "lab_notebook.md"
+
+
+def _ts() -> str:
+    return datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%SZ")
+
+
+def _ensure() -> None:
+    for d in (LOGS, FIGS, EXPS):
+        d.mkdir(parents=True, exist_ok=True)
+
+
+def note(title: str, body: str = "", level: str = "INFO") -> None:
+    """Append a timestamped entry to the running journal."""
+    _ensure()
+    with open(NOTEBOOK, "a") as f:
+        f.write(f"\n### [{_ts()}] {level} — {title}\n\n")
+        if body:
+            f.write(body.rstrip() + "\n")
+
+
+def table(rows: list[dict], cols: list[str] | None = None) -> str:
+    """Render list-of-dicts as a markdown table."""
+    if not rows:
+        return "_(empty)_\n"
+    cols = cols or list(rows[0].keys())
+
+    def fmt(v):
+        if isinstance(v, float):
+            return f"{v:.4g}"
+        return str(v)
+
+    out = ["| " + " | ".join(cols) + " |",
+           "|" + "|".join("---" for _ in cols) + "|"]
+    for r in rows:
+        out.append("| " + " | ".join(fmt(r.get(c, "")) for c in cols) + " |")
+    return "\n".join(out) + "\n"
+
+
+def write_report(name: str, content: str) -> Path:
+    """Write/overwrite one experiment's full report."""
+    _ensure()
+    p = EXPS / f"{name}.md"
+    p.write_text(content)
+    return p
+
+
+def gpu_snapshot() -> dict:
+    try:
+        q = ("--query-gpu=memory.used,memory.total,utilization.gpu,temperature.gpu"
+             ",power.draw")
+        out = subprocess.run(
+            ["nvidia-smi", q, "--format=csv,noheader,nounits"],
+            capture_output=True, text=True, timeout=20).stdout.strip()
+        u, t, g, temp, pw = [x.strip() for x in out.split(",")]
+        return {"vram_used_mb": int(u), "vram_total_mb": int(t),
+                "gpu_util_pct": int(g), "temp_c": int(temp), "power_w": float(pw)}
+    except Exception as e:
+        return {"error": str(e)}
+
+
+def disk_snapshot(paths: list[str] | None = None) -> dict:
+    paths = paths or [str(ROOT)]
+    out = {}
+    try:
+        df = subprocess.run(["df", "-BM", "--output=avail,used,size", str(ROOT)],
+                            capture_output=True, text=True, timeout=20).stdout
+        parts = df.strip().splitlines()[-1].split()
+        out["fs_avail_mb"] = int(parts[0].rstrip("M"))
+        out["fs_used_mb"] = int(parts[1].rstrip("M"))
+    except Exception as e:
+        out["df_error"] = str(e)
+    for p in paths:
+        try:
+            du = subprocess.run(["du", "-sm", p], capture_output=True,
+                                text=True, timeout=120).stdout.split()[0]
+            out[f"du_mb:{Path(p).name}"] = int(du)
+        except Exception:
+            pass
+    return out
+
+
+def checkpoint(stage: str, extra: dict | None = None) -> dict:
+    """One-line health snapshot appended to the journal."""
+    snap = {"stage": stage, "gpu": gpu_snapshot(), "disk": disk_snapshot()}
+    if extra:
+        snap.update(extra)
+    note(f"checkpoint: {stage}", "```json\n" + json.dumps(snap, indent=1) + "\n```")
+    return snap
+
+
+class Timer:
+    def __init__(self, label: str):
+        self.label = label
+
+    def __enter__(self):
+        self.t0 = time.time()
+        return self
+
+    def __exit__(self, *a):
+        self.dt = time.time() - self.t0
+        note(f"timing: {self.label}", f"`{self.dt:.1f}s` ({self.dt/60:.1f} min)")

src/make_report.py

@@ -0,0 +1,247 @@
+"""
+Final cross-model report: results table, quality-vs-diversity frontier chart,
+and qualitative examples. Reads outputs/eval/results.csv.
+"""
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+
+import logbook
+
+ROOT = Path(__file__).resolve().parent.parent
+
+ORDER = ["base", "E0-baseline", "E1-div-individual", "E2-div-group",
+         "E3-multipos", "E4-divpo-emb", "E4-divpo-prob"]
+
+PRETTY = {
+    "base": "Base (Qwen3-4B-Instruct)",
+    "E0-baseline": "E0 · GRPO, quality only",
+    "E1-div-individual": "E1 · div-grpo-individual (d_i)",
+    "E2-div-group": "E2 · div-grpo-group (d_i + m_i)",
+    "E3-multipos": "E3 · multi-positive weighted DPO",
+    "E4-divpo-emb": "E4a · DivPO (embedding)",
+    "E4-divpo-prob": "E4b · DivPO (probability)",
+}
+
+
+def load(path: Path) -> list[dict]:
+    rows = list(csv.DictReader(open(path)))
+    for r in rows:
+        for k, v in r.items():
+            if k != "model":
+                try:
+                    r[k] = float(v)
+                except (TypeError, ValueError):
+                    pass
+    key = {m: i for i, m in enumerate(ORDER)}
+    rows.sort(key=lambda r: key.get(r["model"], 99))
+    return rows
+
+
+def frontier_chart(rows, out_path):
+    import matplotlib
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+
+    fig, ax = plt.subplots(1, 3, figsize=(17, 5.2))
+    cols = plt.cm.tab10(np.linspace(0, 1, 10))
+
+    panels = [
+        ("pairwise", "Mean pairwise embedding distance", 0),
+        ("logdet", "Log-det volume (semantic coverage)", 1),
+        ("eff_rank", "Effective rank (mode count)", 2),
+    ]
+    base = next((r for r in rows if r["model"] == "base"), None)
+
+    for (xk, xlabel, i) in panels:
+        a = ax[i]
+        for j, r in enumerate(rows):
+            a.scatter(r[xk], r["quality"], s=170, color=cols[j % 10],
+                      edgecolor="black", zorder=3,
+                      label=PRETTY.get(r["model"], r["model"]))
+            a.annotate(r["model"].replace("-div", "").replace("divpo-", ""),
+                       (r[xk], r["quality"]), fontsize=7,
+                       xytext=(5, 5), textcoords="offset points")
+        if base:
+            a.axhline(base["quality"], ls=":", c="gray", lw=1)
+            a.axvline(base[xk], ls=":", c="gray", lw=1)
+        a.set_xlabel(xlabel); a.set_ylabel("Judge quality (0-10)")
+        a.grid(alpha=.3)
+        a.set_title(f"Quality vs {xlabel.split('(')[0].strip()}", fontsize=10)
+
+    handles, labels = ax[0].get_legend_handles_labels()
+    fig.legend(handles, labels, loc="lower center", ncol=4, fontsize=8,
+               frameon=False, bbox_to_anchor=(0.5, -0.06))
+    fig.suptitle("Quality–diversity frontier (dotted lines = base model)", fontsize=13)
+    plt.tight_layout()
+    plt.savefig(out_path, dpi=140, bbox_inches="tight")
+    plt.close()
+
+
+def delta_table(rows):
+    base = next((r for r in rows if r["model"] == "base"), None)
+    if not base:
+        return []
+    keys = ["quality", "pairwise", "logdet", "distinct4", "self_bleu",
+            "eff_rank", "n_clusters", "ends_cleanly"]
+    out = []
+    for r in rows:
+        d = {"model": PRETTY.get(r["model"], r["model"])}
+        for k in keys:
+            base_v = base.get(k, 0.0)
+            d[f"Δ{k}"] = round(r.get(k, 0.0) - base_v, 4)
+        out.append(d)
+    return out
+
+
+QUALITY_TOLERANCE = 0.3   # judge points we are willing to lose for diversity
+
+
+def recommend(rows) -> str:
+    """Pick the arm to scale, from the data rather than from vibes.
+
+    Rule, stated before looking at results: among arms whose quality is within
+    QUALITY_TOLERANCE of the BASE model (not of the best arm -- we care about not
+    degrading the model, not about winning on quality), take the largest gain in
+    effective rank. Effective rank is the primary diversity measure because it is
+    the one metric shown to separate collapse from spread (see 02_setup). Ties
+    and near-ties are reported rather than hidden.
+    """
+    base = next((r for r in rows if r["model"] == "base"), None)
+    if not base:
+        return "_No base row; cannot compute deltas._"
+
+    cand = []
+    for r in rows:
+        if r["model"] == "base":
+            continue
+        dq = r["quality"] - base["quality"]
+        der = r.get("eff_rank", 0.0) - base.get("eff_rank", 0.0)
+        dld = r.get("logdet", 0.0) - base.get("logdet", 0.0)
+        cand.append({"model": PRETTY.get(r["model"], r["model"]),
+                     "raw": r["model"], "Δquality": round(dq, 3),
+                     "Δeff_rank": round(der, 3), "Δlogdet": round(dld, 3),
+                     "ends_cleanly": round(r.get("ends_cleanly", 0.0), 3),
+                     "eligible": dq >= -QUALITY_TOLERANCE and r.get("ends_cleanly", 0) > 0.9})
+
+    ok = [c for c in cand if c["eligible"]]
+    ranked = sorted(ok, key=lambda c: -c["Δeff_rank"])
+
+    if not ranked:
+        return (f"**No arm qualifies.** Every arm either lost more than "
+                f"{QUALITY_TOLERANCE} judge points against base or fell below a "
+                f"90% clean-completion rate. The honest recommendation is to fix "
+                f"the objective before scaling anything to 8B.\n\n"
+                + logbook.table(cand, ["model", "Δquality", "Δeff_rank",
+                                       "Δlogdet", "ends_cleanly", "eligible"]))
+
+    win = ranked[0]
+    runner = ranked[1] if len(ranked) > 1 else None
+    margin = (win["Δeff_rank"] - runner["Δeff_rank"]) if runner else None
+    close = margin is not None and margin < 0.25
+
+    txt = (f"**Scale `{win['model']}` to Qwen3-8B.** It gains "
+           f"{win['Δeff_rank']:+.2f} effective rank and {win['Δlogdet']:+.2f} "
+           f"log-det against base while holding quality at "
+           f"{win['Δquality']:+.2f} and completion at "
+           f"{100*win['ends_cleanly']:.0f}%.\n\n")
+    if runner:
+        txt += (f"Runner-up is `{runner['model']}` "
+                f"({runner['Δeff_rank']:+.2f} effective rank). ")
+        txt += ("The margin is **{:.2f}, which is small enough that this "
+                "ordering should not be treated as settled** on 50 eval prompts "
+                "— both are worth carrying forward.\n\n".format(margin)
+                if close else
+                "The margin ({:.2f}) is clear.\n\n".format(margin))
+    excluded = [c for c in cand if not c["eligible"]]
+    if excluded:
+        txt += ("Excluded for quality or completeness regression: "
+                + ", ".join(f"`{c['model']}` (Δq {c['Δquality']:+.2f}, "
+                            f"complete {100*c['ends_cleanly']:.0f}%)"
+                            for c in excluded) + ".\n\n")
+    txt += logbook.table(cand, ["model", "Δquality", "Δeff_rank", "Δlogdet",
+                                "ends_cleanly", "eligible"])
+    return txt
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--results", default="outputs/eval/results.csv")
+    ap.add_argument("--samples", default="outputs/eval/samples")
+    args = ap.parse_args()
+
+    rows = load(ROOT / args.results)
+    fig = logbook.FIGS / "frontier.png"
+    fig.parent.mkdir(parents=True, exist_ok=True)
+    frontier_chart(rows, fig)
+
+    main_cols = ["model", "quality", "ends_cleanly", "pairwise", "logdet",
+                 "distinct4", "self_bleu", "eff_rank", "n_clusters", "tok_entropy", "words"]
+    tbl = [{k: (PRETTY.get(r["model"], r["model"]) if k == "model"
+                else round(r.get(k, 0.0), 4)) for k in main_cols} for r in rows]
+
+    # qualitative examples
+    qual = ""
+    sdir = ROOT / args.samples
+    for r in rows:
+        p = sdir / f"{r['model']}_examples.json"
+        if not p.exists():
+            continue
+        ex = json.load(open(p))
+        if not ex:
+            continue
+        e = ex[0]
+        qual += f"\n### {PRETTY.get(r['model'], r['model'])}\n\n"
+        qual += f"*Prompt:* {e['prompt'][:200]}\n\n"
+        qual += (f"*Set stats:* pairwise {e['pairwise']:.3f} · logdet {e['logdet']:.2f} "
+                 f"· eff_rank {e['eff_rank']:.2f} · clusters {e['n_clusters']} "
+                 f"· self-BLEU {e['self_bleu']:.3f}\n\n")
+        for i, t in enumerate(e["texts"][:3]):
+            first = t.strip().split("\n")[0][:180]
+            qual += f"{i+1}. {first}…\n"
+    body = f"""# Results — diversity-aware post-training for creative story generation
+
+## Main table
+
+{logbook.table(tbl, main_cols)}
+
+`self_bleu` is inverted in meaning: **lower is more diverse**. `tok_entropy` is
+monitor-only and was never optimized.
+
+## Change vs. base model
+
+{logbook.table(delta_table(rows))}
+
+## Quality–diversity frontier
+
+![frontier](figures/frontier.png)
+
+## Recommendation — which method to scale to Qwen3-8B
+
+Selection rule, fixed before results were seen: among arms whose judge quality
+is within {QUALITY_TOLERANCE} points of the **base** model and whose clean-completion
+rate stays above 90%, take the largest gain in **effective rank**. Quality is
+measured against base rather than against the best arm because the goal is to
+add diversity without degrading the model, not to win on quality. Effective rank
+is the primary diversity axis because it is the one mode metric demonstrated to
+separate collapse from spread (`02_setup_and_deviations.md`, §11).
+
+{recommend(rows)}
+
+## Qualitative examples (first 3 of 16 samples, same prompt)
+{qual}
+"""
+    (logbook.LOGS.parent / "report.md").write_text(body)
+    print(body[:2500])
+    print("\nwrote report.md and", fig)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

src/qualitative.py

@@ -0,0 +1,157 @@
+"""
+Qualitative diff across arms: actually READ the stories.
+
+Aggregate metrics can say "effective rank 2.0/16" without conveying that six of
+the sixteen ships are literally named *Aethel*. This module extracts the
+concrete signature of collapse -- shared openings, repeated proper nouns,
+shared closing cadence, tone/genre spread -- so a human can see what changed.
+
+Works on either a scored pool (`--pool`) or an eval dump (`--eval`), so the
+base model, every checkpoint and every arm are read the same way.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import re
+import sys
+from collections import Counter, defaultdict
+from pathlib import Path
+
+import numpy as np
+
+ROOT = Path(__file__).resolve().parent.parent
+
+OPEN_FRAME = re.compile(r"^The\s+\*?([A-Z][\w' ]{2,20})\*?\s+(\w+ed|\w+s)\s+"
+                        r"(through|over|above|at|into|across)\b")
+ITALIC_NAME = re.compile(r"\*([A-Z][\w' ]{2,24})\*")
+CADENCE = [
+    (r"for the first time in (centuries|years|generations|decades|a long time)",
+     "'for the first time in X'"),
+    (r"\b(had )?begun to\b", "'begun to'"),
+    (r"\bwasn'?t\s+\w+[^.]{0,50},?\s+it\s+was\b", "'it wasn't X, it was Y'"),
+    (r"\n\n[^\n]{1,60}\.\s*$", "one-line closing paragraph"),
+]
+# crude genre/register probes -- presence of ANY is weak, but the SPREAD across
+# a 16-sample set is the interesting number
+REGISTER = {
+    "elegiac/literary": r"\b(silence|ache|hollow|dust|memory|grief|quiet)\b",
+    "comic": r"\b(ridiculous|absurd|joke|laughed|idiot|bureaucra|paperwork)\b",
+    "horror": r"\b(scream|blood|rot|teeth|corpse|terror|wrong)\b",
+    "technical/SF": r"\b(protocol|sensor|reactor|alloy|deploy|calibrat|drone)\b",
+    "dialogue-driven": r'"[^"]{10,}"',
+    "second person": r"\bYou (are|were|walk|feel|stand|know)\b",
+    "epistolary": r"\b(Dear |entry |log:|Report |memo)\b",
+}
+
+
+def first_sentence(t: str, n: int = 130) -> str:
+    return re.split(r"(?<=[.!?])\s", t.strip())[0][:n]
+
+
+def last_sentence(t: str, n: int = 120) -> str:
+    return re.split(r"(?<=[.!?])\s", t.strip())[-1][:n]
+
+
+def analyze_group(texts: list[str]) -> dict:
+    n = len(texts)
+    names = Counter(x for t in texts for x in ITALIC_NAME.findall(t))
+    opens_the = sum(1 for t in texts if t.strip().startswith("The "))
+    frame = sum(1 for t in texts if OPEN_FRAME.match(t.strip()))
+    # repeated content words across stories (excluding prompt-driven ones is
+    # impossible in general, so we report the top shared nouns as-is)
+    stop = set("the a an and or but of to in on it its was were is are that this "
+               "with for as at by from had have has been be he she they we you i "
+               "his her their our my not no so then than there here what when "
+               "which who all more most into over under out up down".split())
+    per_story_vocab = [set(w for w in re.sub(r"[^a-z\s]", " ", t.lower()).split()
+                           if w not in stop and len(w) > 3) for t in texts]
+    shared = Counter(w for v in per_story_vocab for w in v)
+    ubiquitous = {w: c for w, c in shared.most_common(400) if c >= max(3, int(0.6 * n))}
+    cad = {}
+    for pat, label in CADENCE:
+        cad[label] = sum(1 for t in texts if re.search(pat, t, re.I | re.M))
+    reg = {}
+    for label, pat in REGISTER.items():
+        reg[label] = sum(1 for t in texts if re.search(pat, t, re.I))
+    return {
+        "n": n,
+        "opens_with_The": opens_the,
+        "opening_frame_match": frame,
+        "distinct_first_5_words": len({" ".join(t.strip().split()[:5]).lower() for t in texts}),
+        "top_names": dict(names.most_common(6)),
+        "max_name_reuse": max(names.values()) if names else 0,
+        "ubiquitous_words": dict(list(ubiquitous.items())[:12]),
+        "cadence": cad,
+        "register_spread": reg,
+        "registers_present": sum(1 for v in reg.values() if v >= max(2, int(0.15 * n))),
+    }
+
+
+def load_pool(tag: str, split: str, prompt_id: str | None):
+    rows = [json.loads(l) for l in
+            open(ROOT / "outputs" / f"pool_{tag}" / f"pool_{split}.jsonl") if l.strip()]
+    by = defaultdict(list)
+    for r in rows:
+        by[r["prompt_id"]].append(r)
+    if prompt_id:
+        return {prompt_id: by[prompt_id]}
+    return dict(by)
+
+
+def load_eval(label: str, out: str):
+    p = ROOT / out / "samples" / f"{label}_full.json"
+    data = json.load(open(p))
+    return {d["prompt_id"]: [{"text": t, "prompt": d["prompt"]} for t in d["texts"]]
+            for d in data}
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--pool", help="pool tag, e.g. 4b")
+    ap.add_argument("--eval", help="eval label, e.g. E0-baseline")
+    ap.add_argument("--out", default="outputs/eval")
+    ap.add_argument("--split", default="train")
+    ap.add_argument("--prompt-id", default=None)
+    ap.add_argument("--show", type=int, default=16, help="first sentences to print")
+    ap.add_argument("--limit-prompts", type=int, default=40)
+    args = ap.parse_args()
+
+    groups = load_pool(args.pool, args.split, args.prompt_id) if args.pool \
+        else load_eval(args.eval, args.out)
+    keys = list(groups)[: args.limit_prompts]
+
+    agg = defaultdict(list)
+    for pid in keys:
+        texts = [r["text"] for r in groups[pid]]
+        if len(texts) < 4:
+            continue
+        a = analyze_group(texts)
+        agg["opens_with_The"].append(a["opens_with_The"] / a["n"])
+        agg["opening_frame"].append(a["opening_frame_match"] / a["n"])
+        agg["distinct_openers"].append(a["distinct_first_5_words"] / a["n"])
+        agg["max_name_reuse"].append(a["max_name_reuse"])
+        agg["registers_present"].append(a["registers_present"])
+        for k, v in a["cadence"].items():
+            agg[f"cadence:{k}"].append(v / a["n"])
+
+    name = args.eval or f"pool-{args.pool}"
+    print(f"\n===== QUALITATIVE PROFILE: {name}  ({len(keys)} prompts) =====")
+    for k, v in agg.items():
+        print(f"  {k:38} {np.mean(v):.3f}")
+
+    pid = args.prompt_id or keys[0]
+    texts = [r["text"] for r in groups[pid]][: args.show]
+    print(f"\n--- prompt {pid} ---")
+    print("   ", groups[pid][0].get("prompt", "")[:220])
+    print("\nOPENINGS:")
+    for i, t in enumerate(texts):
+        print(f"  {i+1:2d}. {first_sentence(t)}")
+    print("\nCLOSINGS:")
+    for i, t in enumerate(texts):
+        print(f"  {i+1:2d}. ...{last_sentence(t)}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

src/rewards.py

@@ -0,0 +1,301 @@
+"""
+Reward channels for E0 / E1 / E2, built for TRL's `normalize_then_sum`.
+
+AGGREGATION
+-----------
+GRPOConfig(multi_objective_aggregation="normalize_then_sum") z-scores every
+reward function WITHIN its group before weighting and summing. That is exactly
+the decoupled GDPO-style normalization the brief asks for, so we get it without
+patching TRL.
+
+It also changes what the weights mean: alpha and gamma multiply STANDARDIZED
+channels, so alpha=0.5 reads as "half a standard deviation of diversity credit
+per standard deviation of quality". No manual rescaling of d_i into the 0-10
+judge range is needed or wanted.
+
+GATING (why it is not simply "set it to 0")
+-------------------------------------------
+The brief says a gate failure should zero the total reward. Under per-channel
+z-scoring, writing 0.0 into a channel does NOT mean "no credit" -- it means
+"whatever 0.0 ranks as inside this group". That is fine for quality (range
+[0,10], floor 0) and for deviation (range [0,2], floor 0), but it is actively
+WRONG for the marginal contribution:
+
+    m_i = logdet(L) - logdet(L_-i)  <=  log(1+eps) ~ 0
+
+m_i is always <= 0, so 0.0 is its CEILING. Gating a broken story to 0.0 in the
+marginal channel would hand it the single highest diversity credit in the group
+-- a reward-hacking channel we would have built ourselves.
+
+So the uniform rule, applied to every diversity channel regardless of sign
+convention: an ineligible sample is assigned the MINIMUM value among eligible
+samples in its group. It can never out-rank a sample that earned its credit.
+Combined with quality -> 0.0 (a hard floor in that channel), a gate-failed
+story lands at the bottom of every channel it participates in.
+
+CONSTANT-REWARD SAFETY
+----------------------
+If every sample in a group is gated, a channel goes constant; TRL's
+(x - mean)/(std + 1e-4) then yields ~0 for all of them. That is the correct
+outcome: a group with no valid samples carries no signal. It is not a crash and
+not a NaN, but it IS worth logging, so `frac_groups_degenerate` is tracked.
+"""
+
+from __future__ import annotations
+
+import statistics
+from dataclasses import dataclass, field
+
+import numpy as np
+
+import gates as G
+from diversity import l2_normalize, marginal_contributions, pairwise_deviation, zscore
+
+# --------------------------------------------------------------- embeddings
+_ENCODER = None
+_EMB_MODEL = "BAAI/bge-base-en-v1.5"
+
+
+def get_encoder(model_name: str = _EMB_MODEL, device: str | None = None):
+    """Module-level singleton; ~110M params (~0.22GB), negligible next to the policy.
+
+    Device is overridable via EMB_DEVICE so tests (and any process running
+    alongside a training job that already owns the GPU) can force CPU.
+    """
+    global _ENCODER
+    if _ENCODER is None:
+        import os
+        from sentence_transformers import SentenceTransformer
+        dev = device or os.environ.get("EMB_DEVICE", "cuda")
+        _ENCODER = SentenceTransformer(model_name, device=dev)
+    return _ENCODER
+
+
+def embed(texts: list[str]) -> np.ndarray:
+    if not texts:
+        return np.zeros((0, 768))
+    E = get_encoder().encode(
+        texts, normalize_embeddings=True, batch_size=32,
+        show_progress_bar=False, convert_to_numpy=True,
+    )
+    return l2_normalize(np.asarray(E, dtype=np.float64))
+
+
+# ------------------------------------------------------------------ config
+@dataclass
+class RewardConfig:
+    arm: str = "E0"              # E0 | E1 | E2
+    alpha: float = 0.5           # weight on deviation channel
+    gamma: float = 0.5           # weight on marginal channel
+    tau: float = 5.0             # quality gate for diversity credit
+    min_words: int = G.MIN_WORDS
+    max_words: int = G.MAX_WORDS
+
+    def channels(self) -> list[str]:
+        if self.arm == "E0":
+            return ["quality"]
+        if self.arm == "E1":
+            return ["quality", "deviation"]
+        if self.arm == "E2":
+            return ["quality", "deviation", "marginal"]
+        raise ValueError(self.arm)
+
+    def weights(self) -> list[float]:
+        return {"E0": [1.0],
+                "E1": [1.0, self.alpha],
+                "E2": [1.0, self.alpha, self.gamma]}[self.arm]
+
+
+@dataclass
+class StepStats:
+    n: int = 0
+    gate_pass: float = 0.0
+    ends_cleanly: float = 0.0
+    mean_quality: float = 0.0
+    mean_quality_passing: float = 0.0
+    mean_novelty: float = 0.0
+    frac_above_tau: float = 0.0
+    mean_deviation: float = 0.0
+    mean_logdet: float = 0.0
+    mean_marginal: float = 0.0
+    mean_words: float = 0.0
+    frac_groups_degenerate: float = 0.0
+    reasons: dict = field(default_factory=dict)
+
+
+def _gate_floor(values: np.ndarray, eligible: np.ndarray) -> np.ndarray:
+    """Ineligible samples take the minimum value among eligible ones.
+
+    Sign-convention agnostic: works for deviation (>=0) and for marginal (<=0)
+    alike. If nothing is eligible, the channel is flat -> z-scores to 0 in TRL.
+    """
+    out = values.astype(np.float64).copy()
+    if not eligible.any():
+        return np.zeros_like(out)
+    out[~eligible] = values[eligible].min()
+    return out
+
+
+class RewardEngine:
+    """Scores one GRPO batch: gates -> judge -> embeddings -> per-channel values.
+
+    TRL calls each reward function separately, but we want ONE judge call set
+    and ONE embedding pass per batch. So the engine computes everything once and
+    memoizes on the batch signature; the per-channel closures just read it.
+    """
+
+    def __init__(self, cfg: RewardConfig, judge, wandb_run=None, log_prefix="train"):
+        self.cfg = cfg
+        self.judge = judge
+        self.wandb_run = wandb_run
+        self.log_prefix = log_prefix
+        self._sig = None
+        self._cache: dict[str, np.ndarray] = {}
+        self.last_stats: StepStats | None = None
+        self.history: list[StepStats] = []
+
+    # ---- core ----------------------------------------------------------
+    def compute(self, prompts: list[str], texts: list[str],
+                finish_reasons: list[str] | None = None) -> dict[str, np.ndarray]:
+        sig = hash((tuple(prompts), tuple(texts)))
+        if sig == self._sig:
+            return self._cache
+
+        n = len(texts)
+        finish_reasons = finish_reasons or [None] * n
+
+        # 1. programmatic gates (free, run first)
+        gres = [G.check(t, finish_reason=fr, min_words=self.cfg.min_words,
+                        max_words=self.cfg.max_words)
+                for t, fr in zip(texts, finish_reasons)]
+        passed = np.array([r.passed for r in gres], dtype=bool)
+
+        # 2. judge only the stories that survived the gates -- never pay to
+        #    score text we have already decided to zero out.
+        quality = np.zeros(n); novelty = np.zeros(n)
+        idx = [i for i in range(n) if passed[i]]
+        if idx:
+            scores = self.judge.score_many_sync([(prompts[i], texts[i]) for i in idx])
+            for i, s in zip(idx, scores):
+                quality[i] = s.quality
+                novelty[i] = s.novelty
+
+        # 3. embeddings + per-group diversity
+        E = embed(texts)
+        groups: dict[str, list[int]] = {}
+        for i, p in enumerate(prompts):
+            groups.setdefault(p, []).append(i)
+
+        dev = np.zeros(n); marg = np.zeros(n)
+        logdets, degenerate = [], 0
+        for _, ids in groups.items():
+            sub = E[ids]
+            d = pairwise_deviation(sub)
+            m = marginal_contributions(sub)
+            # z-score m within group: raw m has a long negative tail (a duplicate
+            # pair reaches log(eps) ~ -6.9) that would otherwise dominate.
+            mz = zscore(m)
+            for k, i in enumerate(ids):
+                dev[i] = d[k]; marg[i] = mz[k]
+            from diversity import logdet_volume
+            logdets.append(logdet_volume(sub))
+            if not passed[ids].any():
+                degenerate += 1
+
+        # 4. eligibility for diversity credit: gates AND quality >= tau.
+        #    Conditioning is what stops "incoherent but different" from paying.
+        eligible = passed & (quality >= self.cfg.tau)
+
+        dev_c = np.zeros(n); marg_c = np.zeros(n)
+        for _, ids in groups.items():
+            ids_a = np.array(ids)
+            dev_c[ids_a] = _gate_floor(dev[ids_a], eligible[ids_a])
+            marg_c[ids_a] = _gate_floor(marg[ids_a], eligible[ids_a])
+
+        quality_c = np.where(passed, quality, 0.0)
+
+        out = {"quality": quality_c, "deviation": dev_c, "marginal": marg_c}
+        self._sig, self._cache = sig, out
+
+        # ---- stats ----
+        from collections import Counter
+        cnt = Counter(r for x in gres for r in x.reasons)
+        st = StepStats(
+            n=n,
+            gate_pass=float(passed.mean()),
+            ends_cleanly=float(np.mean([r.completeness for r in gres])),
+            mean_quality=float(quality.mean()),
+            mean_quality_passing=float(quality[passed].mean()) if passed.any() else 0.0,
+            mean_novelty=float(novelty[passed].mean()) if passed.any() else 0.0,
+            frac_above_tau=float(eligible.mean()),
+            mean_deviation=float(dev.mean()),
+            mean_logdet=float(np.mean(logdets)) if logdets else 0.0,
+            mean_marginal=float(marg.mean()),
+            mean_words=float(np.mean([r.n_words for r in gres])),
+            frac_groups_degenerate=degenerate / max(1, len(groups)),
+            reasons=dict(cnt),
+        )
+        self.last_stats = st
+        self.history.append(st)
+        self._log(st)
+        return out
+
+    def _log(self, st: StepStats) -> None:
+        print(f"  [rw] gate={st.gate_pass:.2f} end={st.ends_cleanly:.2f} "
+              f"q={st.mean_quality_passing:.2f} >tau={st.frac_above_tau:.2f} "
+              f"dev={st.mean_deviation:.3f} logdet={st.mean_logdet:.2f} "
+              f"w={st.mean_words:.0f} {st.reasons if st.reasons else ''}", flush=True)
+        if self.wandb_run is not None:
+            p = self.log_prefix
+            self.wandb_run.log({
+                f"{p}/gate_pass": st.gate_pass,
+                f"{p}/ends_cleanly": st.ends_cleanly,
+                f"{p}/quality_passing": st.mean_quality_passing,
+                f"{p}/quality_all": st.mean_quality,
+                f"{p}/novelty": st.mean_novelty,
+                f"{p}/frac_above_tau": st.frac_above_tau,
+                f"{p}/deviation": st.mean_deviation,
+                f"{p}/logdet": st.mean_logdet,
+                f"{p}/marginal_z": st.mean_marginal,
+                f"{p}/words": st.mean_words,
+                f"{p}/groups_degenerate": st.frac_groups_degenerate,
+            })
+
+    # ---- TRL adapters ---------------------------------------------------
+    def make_reward_funcs(self):
+        """Return TRL-compatible reward callables, one per active channel."""
+        funcs = []
+        for ch in self.cfg.channels():
+            funcs.append(self._make(ch))
+        return funcs
+
+    def _make(self, channel: str):
+        engine = self
+
+        def f(completions, prompts=None, **kwargs):
+            texts = [_text(c) for c in completions]
+            ps = [_ptext(p) for p in (prompts or [""] * len(texts))]
+            fr = kwargs.get("finish_reasons")
+            vals = engine.compute(ps, texts, fr)
+            return [float(x) for x in vals[channel]]
+
+        f.__name__ = f"{channel}_reward"
+        return f
+
+
+def _text(c) -> str:
+    if isinstance(c, list):
+        return c[-1].get("content", "") if c else ""
+    if isinstance(c, dict):
+        return c.get("content", "")
+    return str(c)
+
+
+def _ptext(p) -> str:
+    if isinstance(p, list):
+        # chat format: the user turn carries the writing prompt
+        for m in reversed(p):
+            if isinstance(m, dict) and m.get("role") == "user":
+                return m.get("content", "")
+        return p[-1].get("content", "") if p else ""
+    return str(p)

src/smoke_vllm.py

@@ -0,0 +1,59 @@
+"""Smoke test: vLLM on Blackwell + Qwen3-4B + new system prompt + gates.
+
+The specific question: does the rewritten system prompt (200-500 word budget,
+explicit ending instruction) plus max_new_tokens=768 actually fix the 100%
+truncation rate that killed the prior run? Prints the gate pass rate.
+"""
+import json
+import sys
+import time
+
+import numpy as np
+from transformers import AutoTokenizer
+
+import gates
+from data import load_prompts
+from generate import build_llm, generate, save_gens
+
+MODEL = "Qwen/Qwen3-4B-Instruct-2507"
+N_PROMPTS = 8
+G = 8
+
+def main():
+    t0 = time.time()
+    tok = AutoTokenizer.from_pretrained(MODEL)
+    llm = build_llm(MODEL, gpu_mem_util=0.85, seed=0)
+    print(f"[load] {time.time()-t0:.1f}s", flush=True)
+
+    prompts = load_prompts("eval")[:N_PROMPTS]
+    t1 = time.time()
+    gens = generate(llm, tok, prompts, n=G, temperature=1.0, top_p=1.0, seed=1234)
+    dt = time.time() - t1
+    ntok = sum(g.n_tokens for g in gens)
+    print(f"[gen] {len(gens)} stories, {ntok} tok in {dt:.1f}s -> {ntok/dt:.0f} tok/s", flush=True)
+
+    res = [gates.check(g.text, finish_reason=g.finish_reason) for g in gens]
+    passed = sum(r.passed for r in res)
+    complete = sum(r.completeness for r in res)
+    capped = sum(r.hit_token_cap for r in res)
+    words = np.array([r.n_words for r in res])
+
+    print("\n=== GATES ===")
+    print(f"  pass_all_gates   {passed}/{len(res)}  ({100*passed/len(res):.0f}%)")
+    print(f"  ends_cleanly     {complete:.0f}/{len(res)}  ({100*complete/len(res):.0f}%)   <- prior run: 0-8%")
+    print(f"  hit_token_cap    {capped}/{len(res)}")
+    print(f"  words  p5={np.percentile(words,5):.0f}  med={np.median(words):.0f}  p95={np.percentile(words,95):.0f}  max={words.max()}")
+    from collections import Counter
+    c = Counter(r for x in res for r in x.reasons)
+    print(f"  failure reasons  {dict(c)}")
+    print(f"  mean_logprob     {np.mean([g.mean_logprob for g in gens]):.3f}")
+
+    save_gens(gens, "outputs/smoke/vllm_smoke.jsonl")
+    print("\n--- sample tail (last 90 chars each, first 4) ---")
+    for g in gens[:4]:
+        print("   ...", repr(g.text[-90:]))
+    print(f"\n[total] {time.time()-t0:.1f}s")
+    return 0 if complete / len(res) > 0.8 else 1
+
+if __name__ == "__main__":
+    sys.exit(main())

src/status.sh

@@ -0,0 +1,39 @@
+#!/usr/bin/env bash
+# User is in IST (UTC+5:30) -- report wall-clock in IST, not UTC.
+export TZ="Asia/Kolkata"
+# One-line-ish status snapshot for periodic reporting.
+cd /workspace/creative-writing
+STAGE=$(grep -aE "^== (START|OK|FAIL|ABORT)" logs/pipeline.log 2>/dev/null | tail -1 | sed 's/  */ /g')
+
+# active training log (most recently modified)
+# only consider logs of a CURRENTLY RUNNING job; otherwise a finished
+# stage's log masquerades as live status (this bit me once already).
+RUNCFG=$(pgrep -af "src/train_(grpo|dpo).py" | grep -oE "configs/[^ ]+\.yaml" | head -1)
+ACT=""
+if [ -n "$RUNCFG" ]; then
+  NM=$(python3 -c "import yaml,sys;print(yaml.safe_load(open(sys.argv[1]))['name'])" "$RUNCFG" 2>/dev/null)
+  for c in "logs/train_${NM}.log" "logs/lrprobe_3e-5.log" $(ls -t logs/*.log 2>/dev/null); do
+    [ -f "$c" ] && ACT="$c" && break
+  done
+fi
+PROG=""; RW=""; ENT=""
+if [ -n "$ACT" ]; then
+  PROG=$(tr '\r' '\n' < "$ACT" | grep -aE "^ *[0-9]+%\|" | tail -1 | sed 's/  */ /g')
+  RW=$(grep -ao "\[rw\].*" "$ACT" 2>/dev/null | tail -1)
+  ENT=$(grep -aoE "'entropy': '[^']*'" "$ACT" 2>/dev/null | tail -1)
+fi
+
+VRAM=$(nvidia-smi --query-gpu=memory.used,utilization.gpu --format=csv,noheader,nounits | tr ',' '/' | tr -d ' ')
+DISK=$(df -BG --output=avail /workspace | tail -1 | tr -d ' ')
+SPEND=$(python3 -c "
+import os,urllib.request,json
+try:
+    r=urllib.request.Request('https://openrouter.ai/api/v1/key',headers={'Authorization':'Bearer '+os.environ['OPENROUTER_API_KEY']})
+    d=json.load(urllib.request.urlopen(r,timeout=20))['data']
+    print(f\"\${d['usage']:.2f}used/\${d['limit_remaining']:.2f}left\")
+except Exception: print('n/a')" 2>/dev/null)
+
+ERR=$(grep -alE "OutOfMemory|Traceback|UNHEALTHY|GUARDRAIL" logs/train_*.log logs/E[34]*.log logs/eval_all.log 2>/dev/null | tr '\n' ',' )
+
+NOW=$(TZ=Asia/Kolkata date +"%H:%M IST")
+echo "[$NOW] STAGE=[$STAGE] PROG=[$PROG] ${RW:+RW=[$RW] }${ENT:+$ENT }vram=${VRAM} disk=${DISK} judge=${SPEND}${ERR:+ ERRORS_IN=$ERR}"

src/test_diversity.py

@@ -0,0 +1,186 @@
+"""
+Unit tests for src/diversity.py.  Run: python3 -m pytest src/test_diversity.py -q
+                                  or: python3 src/test_diversity.py
+"""
+import numpy as np
+
+from diversity import (
+    cosine_kernel,
+    greedy_diverse_subset,
+    l2_normalize,
+    logdet_volume,
+    marginal_contributions,
+    pairwise_deviation,
+    zscore,
+)
+
+RNG = np.random.default_rng(0)
+
+
+def _orth(G, d):
+    """G mutually orthonormal rows in R^d."""
+    Q, _ = np.linalg.qr(RNG.standard_normal((d, G)))
+    return Q.T[:G]
+
+
+# ---------------------------------------------------------------- deviation
+def test_deviation_duplicates_are_zero():
+    """Exact duplicates: every pair distance is 0, so d_i == 0 for all."""
+    e = l2_normalize(RNG.standard_normal((1, 32)))
+    E = np.repeat(e, 8, axis=0)
+    d = pairwise_deviation(E)
+    assert np.allclose(d, 0.0, atol=1e-9), d
+
+
+def test_deviation_orthogonal_is_one():
+    """Orthogonal rows: cos = 0 for every pair, so d_i == 1 exactly."""
+    E = _orth(8, 64)
+    d = pairwise_deviation(E)
+    assert np.allclose(d, 1.0, atol=1e-9), d
+
+
+def test_deviation_singleton_and_empty():
+    assert pairwise_deviation(np.zeros((1, 8))).shape == (1,)
+    assert pairwise_deviation(np.zeros((1, 8)))[0] == 0.0
+    assert pairwise_deviation(np.zeros((0, 8))).shape == (0,)
+
+
+def test_deviation_flags_the_odd_one_out():
+    """7 near-identical + 1 far: the outlier must have the highest d_i."""
+    base = l2_normalize(RNG.standard_normal((1, 64)))
+    tight = l2_normalize(np.repeat(base, 7, axis=0) + 0.01 * RNG.standard_normal((7, 64)))
+    far = l2_normalize(RNG.standard_normal((1, 64)))
+    E = np.vstack([tight, far])
+    d = pairwise_deviation(E)
+    assert d.argmax() == 7, d
+
+
+# ------------------------------------------------------- marginal / logdet
+def test_marginal_duplicate_is_large_negative():
+    """A duplicated direction is already spanned -> dropping one costs ~nothing,
+    so the PRESENT duplicate's marginal is driven to ~log(eps), very negative."""
+    e = l2_normalize(RNG.standard_normal((1, 32)))
+    E = np.vstack([np.repeat(e, 2, axis=0), _orth(4, 32)])
+    m = marginal_contributions(E)
+    # the two duplicates (rows 0,1) are the least valuable members
+    assert m[0] < -3.0 and m[1] < -3.0, m
+    assert m[:2].max() < m[2:].min(), m
+
+
+def test_marginal_orthogonal_is_near_zero_and_uniform():
+    """Orthonormal rows: L = (1+eps)I, dropping any row costs log(1+eps) ~ 0.
+    'High m_i for all' in the sense of at-ceiling and symmetric."""
+    E = _orth(8, 64)
+    m = marginal_contributions(E)
+    assert np.allclose(m, m[0], atol=1e-9), m
+    assert abs(m[0] - np.log1p(1e-3)) < 1e-6, m[0]
+
+
+def test_marginal_is_bounded_above_by_zero_ish():
+    """logdet is monotone under adding a row with unit norm + jitter, so
+    m_i can never exceed log(1+eps)."""
+    E = l2_normalize(RNG.standard_normal((12, 64)))
+    m = marginal_contributions(E)
+    assert m.max() <= np.log1p(1e-3) + 1e-9, m.max()
+
+
+def test_logdet_ordering_dup_lt_spread_lt_orthogonal():
+    e = l2_normalize(RNG.standard_normal((1, 64)))
+    dup = np.repeat(e, 8, axis=0)
+    spread = l2_normalize(RNG.standard_normal((8, 64)))
+    orth = _orth(8, 64)
+    assert logdet_volume(dup) < logdet_volume(spread) < logdet_volume(orth)
+
+
+def test_kernel_is_psd_even_with_duplicates():
+    e = l2_normalize(RNG.standard_normal((1, 16)))
+    L = cosine_kernel(np.repeat(e, 6, axis=0))
+    assert np.linalg.eigvalsh(L).min() > 0, "jitter failed to make L PD"
+    assert np.isfinite(logdet_volume(np.repeat(e, 6, axis=0)))
+
+
+# --------------------------------------- THE E1-vs-E2 HYPOTHESIS, AS A TEST
+def test_two_clusters_fool_deviation_but_not_logdet():
+    """This is the claim E2 rests on, so it gets asserted rather than assumed.
+
+    Config A: two tight antipodal clusters of 4 (rank ~2, 'diverse' only in
+              the sense that half the samples are far from the other half).
+    Config B: 8 genuinely spread directions (rank ~8).
+
+    Pairwise deviation cannot tell these apart -- mean pairwise distance for A
+    is actually HIGHER, because antipodal pairs sit at cos = -1. Log-det sees
+    straight through it: A occupies a 2-dimensional subspace.
+    """
+    u, v = _orth(2, 64)
+    jit = 0.01
+    A = l2_normalize(np.vstack([
+        np.repeat(u[None], 4, axis=0) + jit * RNG.standard_normal((4, 64)),
+        np.repeat(-u[None], 4, axis=0) + jit * RNG.standard_normal((4, 64)),
+    ]))
+    B = _orth(8, 64)
+
+    dev_A, dev_B = pairwise_deviation(A).mean(), pairwise_deviation(B).mean()
+    vol_A, vol_B = logdet_volume(A), logdet_volume(B)
+
+    # deviation RANKS THE DEGENERATE SET HIGHER -- the failure mode, reproduced
+    assert dev_A > dev_B, (dev_A, dev_B)
+    # log-det correctly ranks the spread set far higher
+    assert vol_B > vol_A + 10.0, (vol_A, vol_B)
+
+    # and per-sample: in A every member is redundant (its twin covers it),
+    # so marginal contributions are uniformly terrible
+    assert marginal_contributions(A).max() < -3.0
+    assert marginal_contributions(B).min() > -1e-3
+    _ = v  # second basis vector unused, kept for clarity of construction
+
+
+# ------------------------------------------------------------------ zscore
+def test_zscore_constant_input_is_zeros_not_nan():
+    """A constant reward column must degrade to 0, never NaN -- otherwise it
+    poisons the whole advantage tensor."""
+    z = zscore(np.full(8, 3.7))
+    assert np.all(np.isfinite(z)) and np.allclose(z, 0.0)
+
+
+def test_zscore_standardizes():
+    z = zscore(RNG.standard_normal(64))
+    assert abs(z.mean()) < 1e-9 and abs(z.std() - 1.0) < 1e-9
+
+
+# ------------------------------------------------------------ greedy subset
+def test_greedy_avoids_duplicates_when_quality_is_flat():
+    """With flat quality, selection is pure logdet: must not pick both dupes."""
+    e = l2_normalize(RNG.standard_normal((1, 64)))
+    E = np.vstack([np.repeat(e, 3, axis=0), _orth(3, 64)])  # rows 0,1,2 identical
+    sel = greedy_diverse_subset(np.ones(6), E, k=3, lam=1.0)
+    assert len(sel) == 3
+    assert len(set(sel) & {0, 1, 2}) <= 1, sel
+
+
+def test_greedy_respects_quality_when_lambda_is_zero():
+    E = l2_normalize(RNG.standard_normal((10, 64)))
+    q = np.arange(10, dtype=float)
+    sel = greedy_diverse_subset(q, E, k=3, lam=0.0)
+    assert sorted(sel) == [7, 8, 9], sel
+
+
+def test_greedy_k_larger_than_pool():
+    E = l2_normalize(RNG.standard_normal((3, 32)))
+    assert len(greedy_diverse_subset(np.ones(3), E, k=8)) == 3
+
+
+if __name__ == "__main__":
+    import sys, traceback
+    fns = [(n, f) for n, f in sorted(globals().items())
+           if n.startswith("test_") and callable(f)]
+    bad = 0
+    for n, f in fns:
+        try:
+            f()
+            print(f"  PASS  {n}")
+        except Exception:
+            bad += 1
+            print(f"  FAIL  {n}")
+            traceback.print_exc()
+    print(f"\n{len(fns)-bad}/{len(fns)} passed")
+    sys.exit(1 if bad else 0)

src/test_eval_metrics.py

@@ -0,0 +1,174 @@
+"""
+Tests for evaluate.py's pure metric functions. These run LAST in the pipeline,
+so a bug here would waste every GPU-hour before it.
+
+Directionality is the thing being pinned down: self-BLEU is inverted relative to
+every other diversity metric (lower = more diverse), and getting that backwards
+in the report would invert the study's conclusion.
+"""
+import numpy as np
+
+from diversity import effective_rank
+from evaluate import cluster_count, distinct_n, self_bleu, topk_entropy
+
+IDENTICAL = ["The harbor clock struck twelve and the ferry did not come."] * 8
+VARIED = [
+    "The harbor clock struck twelve and the ferry did not come.",
+    "Marguerite sold her mother's piano to pay for the greenhouse.",
+    "In the third week of the drought, the well began speaking Latin.",
+    "He counted seventeen crows before admitting he was being followed.",
+    "The recipe called for one tablespoon of regret, finely minced.",
+    "Every letter she mailed arrived a decade before she wrote it.",
+    "Nobody warned the astronauts that the moon would be so loud.",
+    "My grandfather traded his shadow for a working knowledge of bees.",
+]
+
+
+def test_distinct4_identical_is_low():
+    """8 copies of one sentence: only 1/8 of 4-grams are new."""
+    d = distinct_n(IDENTICAL, 4)
+    assert d < 0.2, d
+
+
+def test_distinct4_varied_is_high():
+    assert distinct_n(VARIED, 4) > 0.9
+
+
+def test_distinct4_ordering():
+    assert distinct_n(VARIED, 4) > distinct_n(IDENTICAL, 4)
+
+
+def test_self_bleu_is_inverted_lower_means_more_diverse():
+    """THE directionality check. Identical texts must score HIGH self-BLEU."""
+    hi = self_bleu(IDENTICAL)
+    lo = self_bleu(VARIED)
+    assert hi > lo, f"self-BLEU not inverted: identical={hi:.4f} varied={lo:.4f}"
+    assert hi > 0.5, f"identical texts should have high self-BLEU, got {hi:.4f}"
+    assert lo < 0.1, f"varied texts should have low self-BLEU, got {lo:.4f}"
+
+
+def test_self_bleu_bounded():
+    for texts in (IDENTICAL, VARIED):
+        v = self_bleu(texts)
+        assert 0.0 <= v <= 1.0, v
+
+
+def test_self_bleu_handles_short_and_single():
+    assert self_bleu(["hi"]) == 0.0
+    assert self_bleu([]) == 0.0
+    assert 0.0 <= self_bleu(["a b", "c d"]) <= 1.0
+
+
+def test_cluster_count_two_clear_modes():
+    rng = np.random.default_rng(0)
+    a = rng.standard_normal(32); a /= np.linalg.norm(a)
+    b = rng.standard_normal(32); b /= np.linalg.norm(b)
+    E = np.vstack([np.tile(a, (8, 1)) + 0.02 * rng.standard_normal((8, 32)),
+                   np.tile(b, (8, 1)) + 0.02 * rng.standard_normal((8, 32))])
+    E /= np.linalg.norm(E, axis=1, keepdims=True)
+    assert cluster_count(E) == 2, cluster_count(E)
+
+
+def test_cluster_count_no_structure_is_one():
+    """Near-identical embeddings have no cluster structure -> a single mode.
+    Originally FAILED at silhouette>0.05 (returned k=4 on a fully collapsed set),
+    which is why SILHOUETTE_MIN was recalibrated to 0.50."""
+    rng = np.random.default_rng(1)
+    a = rng.standard_normal(32); a /= np.linalg.norm(a)
+    E = np.tile(a, (16, 1)) + 0.001 * rng.standard_normal((16, 32))
+    E /= np.linalg.norm(E, axis=1, keepdims=True)
+    assert cluster_count(E) == 1, cluster_count(E)
+
+
+def test_cluster_count_small_input():
+    assert cluster_count(np.eye(3)) == 1
+
+
+def test_topk_entropy_peaked_vs_flat():
+    """A near-deterministic distribution has ~0 entropy; uniform over k has log k."""
+    peaked = [{"a": np.log(0.999), "b": np.log(0.001)}]
+    flat = [{c: np.log(0.25) for c in "abcd"}]
+    assert topk_entropy(peaked) < 0.05
+    assert abs(topk_entropy(flat) - np.log(4)) < 1e-6
+
+
+def test_topk_entropy_empty():
+    assert topk_entropy([]) == 0.0
+    assert topk_entropy([{}]) == 0.0
+
+
+def test_topk_entropy_renormalizes_truncated_table():
+    """vLLM returns only the top-k, which does not sum to 1; we renormalize."""
+    partial = [{"a": np.log(0.4), "b": np.log(0.2)}]   # sums to 0.6
+    h = topk_entropy(partial)
+    p = np.array([2 / 3, 1 / 3])
+    assert abs(h - float(-(p * np.log(p)).sum())) < 1e-9
+
+
+def test_eff_rank_identical_is_one():
+    rng = np.random.default_rng(3)
+    a = rng.standard_normal(32); a /= np.linalg.norm(a)
+    assert abs(effective_rank(np.tile(a, (16, 1))) - 1.0) < 1e-6
+
+
+def test_eff_rank_orthogonal_is_n():
+    Q, _ = np.linalg.qr(np.random.default_rng(4).standard_normal((32, 8)))
+    assert abs(effective_rank(Q.T[:8]) - 8.0) < 1e-6
+
+
+def test_eff_rank_two_clusters_is_about_two():
+    """The case silhouette got right but only at a threshold that broke the
+    collapsed case. Effective rank handles both without a threshold."""
+    rng = np.random.default_rng(5)
+    a = rng.standard_normal(32); a /= np.linalg.norm(a)
+    b = rng.standard_normal(32); b /= np.linalg.norm(b)
+    E = np.vstack([np.tile(a, (8, 1)) + 0.02 * rng.standard_normal((8, 32)),
+                   np.tile(b, (8, 1)) + 0.02 * rng.standard_normal((8, 32))])
+    E /= np.linalg.norm(E, axis=1, keepdims=True)
+    assert 1.8 < effective_rank(E) < 2.6, effective_rank(E)
+
+
+def test_eff_rank_is_monotone_in_spread():
+    """Monotone IN EXPECTATION. Averaged over draws because effective rank
+    saturates around 12.6 (not 16) for 16 unit vectors in 32 dims -- random
+    directions retain residual correlation -- so single draws at the top of the
+    range can invert by chance. The metric is fine; the assertion has to be
+    statistical."""
+    rng = np.random.default_rng(6)
+    a = rng.standard_normal(32); a /= np.linalg.norm(a)
+    prev = 0.0
+    for noise in (0.001, 0.05, 0.15, 0.5, 2.0):
+        vals = []
+        for _ in range(5):
+            E = np.tile(a, (16, 1)) + noise * rng.standard_normal((16, 32))
+            E /= np.linalg.norm(E, axis=1, keepdims=True)
+            vals.append(effective_rank(E))
+        r = float(np.mean(vals))
+        assert r >= prev - 1e-6, f"non-monotone at noise={noise}: {r} < {prev}"
+        prev = r
+
+
+def test_eff_rank_separates_collapsed_from_spread():
+    """The exact discrimination silhouette FAILED: collapsed 0.202 vs spread
+    0.195 were indistinguishable. Effective rank must separate them clearly."""
+    rng = np.random.default_rng(7)
+    a = rng.standard_normal(32); a /= np.linalg.norm(a)
+    collapsed = np.tile(a, (16, 1)) + 0.001 * rng.standard_normal((16, 32))
+    collapsed /= np.linalg.norm(collapsed, axis=1, keepdims=True)
+    spread = rng.standard_normal((16, 32))
+    spread /= np.linalg.norm(spread, axis=1, keepdims=True)
+    assert effective_rank(spread) > effective_rank(collapsed) + 8.0
+
+
+if __name__ == "__main__":
+    import sys, traceback
+    fns = [(n, f) for n, f in sorted(globals().items())
+           if n.startswith("test_") and callable(f)]
+    bad = 0
+    for n, f in fns:
+        try:
+            f(); print(f"  PASS  {n}")
+        except Exception:
+            bad += 1; print(f"  FAIL  {n}"); traceback.print_exc()
+    print(f"\n{len(fns)-bad}/{len(fns)} passed")
+    sys.exit(1 if bad else 0)

src/test_pairs.py

@@ -0,0 +1,167 @@
+"""
+Synthetic tests for E3/E4 preference-pair construction. Catches logic errors now
+rather than four hours into the pipeline.
+
+Most important: test_divpo_prob_picks_least_probable_as_chosen. DivPO's
+probability variant inverts the usual intuition -- the CHOSEN story is the one
+the model found LEAST likely. Getting that backwards would train the policy
+straight into the greedy mode while the logs looked entirely normal.
+"""
+import numpy as np
+
+from build_pairs import build_divpo, build_multipos, normalize_weights
+
+
+def make_item(pid, idx, quality, dev, meanlp, passed=True, text=None):
+    return {"prompt_id": pid, "idx": idx, "prompt": f"prompt-{pid}",
+            "text": text or f"story-{pid}-{idx}", "quality": quality,
+            "deviation": dev, "mean_logprob": meanlp, "gate_passed": passed,
+            "marginal": -1.0, "n_words": 300}
+
+
+def simple_pool():
+    """One prompt, 6 stories spanning quality/deviation/logprob."""
+    return {"p0": [
+        make_item("p0", 0, 8.0, 0.40, -1.50),   # high q, MOST diverse, least probable
+        make_item("p0", 1, 7.0, 0.20, -0.80),   # high q, low dev  -> "competent cliche"
+        make_item("p0", 2, 6.5, 0.30, -1.10),
+        make_item("p0", 3, 3.0, 0.35, -0.60),   # low q, HIGH dev  -> must NOT be chosen
+        make_item("p0", 4, 2.0, 0.05, -0.40),   # low q, least diverse, MOST probable
+        make_item("p0", 5, 1.0, 0.10, -0.90, passed=False),
+    ]}
+
+
+# --------------------------------------------------------------------- E4
+def test_divpo_emb_chosen_is_most_diverse_above_rho():
+    rows, st = build_divpo(simple_pool(), "emb", rho=6.0)
+    assert st["n_rows"] == 1
+    r = rows[0]
+    assert r["chosen"] == "story-p0-0", r
+    assert r["chosen_quality"] >= 6.0
+
+
+def test_divpo_emb_never_chooses_a_diverse_but_low_quality_story():
+    """Story 3 has high deviation but quality 3.0. Quality gating must exclude
+    it -- this is the 'different because it is worse' failure."""
+    rows, _ = build_divpo(simple_pool(), "emb", rho=6.0)
+    assert rows[0]["chosen"] != "story-p0-3"
+
+
+def test_divpo_emb_rejected_is_least_diverse_below_rho():
+    rows, _ = build_divpo(simple_pool(), "emb", rho=6.0)
+    assert rows[0]["rejected"] == "story-p0-4", rows[0]
+
+
+def test_divpo_prob_picks_least_probable_as_chosen():
+    """Most diverse == LOWEST length-normalized logprob; least diverse ==
+    HIGHEST (the near-greedy sample). Inverting this trains toward the mode."""
+    rows, _ = build_divpo(simple_pool(), "prob", rho=6.0)
+    r = rows[0]
+    assert r["chosen"] == "story-p0-0", f"chosen should be least probable: {r}"
+    assert r["rejected"] == "story-p0-4", f"rejected should be most probable: {r}"
+    assert r["chosen_meanlp"] < r["rejected_meanlp"]
+
+
+def test_divpo_skips_prompt_with_no_qualifying_chosen():
+    pool = {"p0": [make_item("p0", i, 3.0, 0.2, -1.0) for i in range(4)]}
+    rows, st = build_divpo(pool, "emb", rho=6.0)
+    assert rows == [] and st["no_chosen"] == 1 and st["skip_rate"] == 1.0
+
+
+def test_divpo_skips_prompt_with_no_rejected():
+    pool = {"p0": [make_item("p0", i, 9.0, 0.2 + 0.01 * i, -1.0) for i in range(4)]}
+    rows, st = build_divpo(pool, "emb", rho=6.0)
+    assert rows == [] and st["no_rejected"] == 1
+
+
+def test_divpo_gate_failures_are_eligible_rejects():
+    """A gate-failing story is legitimately in the low-quality tail."""
+    pool = {"p0": [make_item("p0", 0, 8.0, 0.4, -1.5),
+                   make_item("p0", 1, 9.0, 0.3, -1.2),
+                   make_item("p0", 2, 0.0, 0.01, -0.3, passed=False)]}
+    rows, st = build_divpo(pool, "emb", rho=6.0)
+    assert st["n_rows"] == 1 and rows[0]["rejected"] == "story-p0-2"
+
+
+def test_divpo_is_one_row_per_prompt():
+    pool = {f"p{i}": simple_pool()["p0"] for i in range(5)}
+    rows, st = build_divpo(pool, "emb", rho=6.0)
+    assert len(rows) == 5 == st["n_rows"]
+    assert all(r["weight"] == 1.0 for r in rows), "DivPO must be unweighted"
+
+
+# --------------------------------------------------------------------- E3
+def _emb_for(pool):
+    """Deterministic embeddings + row index aligned to pool iteration order."""
+    idx, vecs, i = {}, [], 0
+    rng = np.random.default_rng(0)
+    for pid, items in pool.items():
+        for it in items:
+            idx[(pid, it["idx"])] = i
+            v = rng.standard_normal(32)
+            vecs.append(v / np.linalg.norm(v))
+            i += 1
+    return np.array(vecs, dtype=np.float32), idx
+
+
+def test_multipos_emits_multiple_chosen_rows_per_prompt():
+    pool = simple_pool()
+    emb, idx = _emb_for(pool)
+    rows, st = build_multipos(pool, emb, idx, q_keep=5.0, k=4, lam=1.0)
+    assert st["ok"] == 1
+    assert 2 <= len(rows) <= 4, f"expected up to 4 chosen rows, got {len(rows)}"
+    assert all(r["chosen_quality"] >= 5.0 for r in rows)
+
+
+def test_multipos_rotates_negatives():
+    """Don't hammer one rejected across all rows."""
+    pool = simple_pool()
+    emb, idx = _emb_for(pool)
+    rows, _ = build_multipos(pool, emb, idx, q_keep=5.0, k=4, lam=1.0)
+    if len(rows) >= 2:
+        assert len(set(r["neg_type"] for r in rows)) == 2, \
+            f"negatives not rotated: {[r['neg_type'] for r in rows]}"
+
+
+def test_multipos_weight_is_chosen_deviation():
+    pool = simple_pool()
+    emb, idx = _emb_for(pool)
+    rows, _ = build_multipos(pool, emb, idx, q_keep=5.0, k=4, lam=1.0)
+    for r in rows:
+        assert abs(r["weight"] - r["chosen_dev"]) < 1e-9
+
+
+def test_multipos_skips_prompt_with_too_few_survivors():
+    pool = {"p0": [make_item("p0", i, 2.0, 0.2, -1.0) for i in range(6)]}
+    emb, idx = _emb_for(pool)
+    rows, st = build_multipos(pool, emb, idx, q_keep=5.0, k=4, lam=1.0)
+    assert rows == [] and st["skipped_no_survivors"] == 1
+
+
+def test_normalize_weights_gives_mean_one():
+    rows = [{"weight": w} for w in (0.1, 0.2, 0.3, 0.8)]
+    normalize_weights(rows)
+    w = np.array([r["weight"] for r in rows])
+    assert abs(w.mean() - 1.0) < 1e-9
+    # relative ordering preserved -> it reweights, it does not rescale the LR
+    assert np.all(np.diff(w) > 0)
+
+
+def test_normalize_weights_handles_all_zero():
+    rows = [{"weight": 0.0} for _ in range(3)]
+    normalize_weights(rows)
+    assert all(r["weight"] == 0.0 for r in rows)
+
+
+if __name__ == "__main__":
+    import sys, traceback
+    fns = [(n, f) for n, f in sorted(globals().items())
+           if n.startswith("test_") and callable(f)]
+    bad = 0
+    for n, f in fns:
+        try:
+            f(); print(f"  PASS  {n}")
+        except Exception:
+            bad += 1; print(f"  FAIL  {n}"); traceback.print_exc()
+    print(f"\n{len(fns)-bad}/{len(fns)} passed")
+    sys.exit(1 if bad else 0)

src/test_rewards.py

@@ -0,0 +1,174 @@
+"""
+Tests for the reward-channel construction in rewards.py, with a fake judge so
+nothing hits the network.
+
+The load-bearing test is test_gated_sample_never_outranks_in_marginal_channel:
+m_i <= log(1+eps) ~ 0 is ALWAYS negative, so gating a failed story to 0.0 would
+hand it the highest diversity credit in the group. That would be a reward-hacking
+channel we built ourselves, and it is exactly the kind of sign error that is
+invisible in aggregate training curves.
+"""
+import numpy as np
+
+import gates as G
+import rewards as R
+from judge import JudgeScore
+
+
+class FakeJudge:
+    """Scores by a marker embedded in the text; no network."""
+    def __init__(self, mapping=None):
+        self.mapping = mapping or {}
+        self.calls = 0
+
+    def score_many_sync(self, pairs):
+        self.calls += len(pairs)
+        out = []
+        for _, story in pairs:
+            q = 8.0
+            for k, v in self.mapping.items():
+                if k in story:
+                    q = v
+                    break
+            out.append(JudgeScore(quality=q, novelty=5.0))
+        return out
+
+
+def _story(marker: str, n_words: int = 250, seed: int = 0) -> str:
+    """Gate-passing filler with a marker and enough lexical variety to clear
+    the 4-gram-loop and entropy gates."""
+    rng = np.random.default_rng(seed)
+    vocab = ["harbor", "clock", "ember", "listen", "gravel", "orchard", "signal",
+             "letter", "winter", "throat", "marble", "engine", "sister", "quiet",
+             "amber", "hollow", "ribbon", "tunnel", "pepper", "静"][:19]
+    words = [vocab[i] for i in rng.integers(0, len(vocab), n_words)]
+    return f"{marker} " + " ".join(words) + "."
+
+
+def _engine(arm="E2", tau=5.0, judge=None):
+    cfg = R.RewardConfig(arm=arm, tau=tau)
+    return R.RewardEngine(cfg, judge or FakeJudge())
+
+
+def _fake_embed(monkey_vals):
+    """Patch rewards.embed to return a fixed matrix."""
+    R.embed = lambda texts: monkey_vals
+
+
+# ------------------------------------------------------------------ configs
+def test_channels_and_weights_per_arm():
+    assert R.RewardConfig(arm="E0").channels() == ["quality"]
+    assert R.RewardConfig(arm="E1").channels() == ["quality", "deviation"]
+    assert R.RewardConfig(arm="E2").channels() == ["quality", "deviation", "marginal"]
+    assert R.RewardConfig(arm="E2", alpha=0.3, gamma=0.7).weights() == [1.0, 0.3, 0.7]
+
+
+# ------------------------------------------------------------------- gating
+def test_gate_failure_zeroes_quality_channel():
+    eng = _engine("E1")
+    texts = [_story("GOOD", 250, i) for i in range(3)] + ["too short."]
+    prompts = ["p"] * 4
+    out = eng.compute(prompts, texts)
+    assert out["quality"][3] == 0.0, "gate-failed story must floor the quality channel"
+    assert (out["quality"][:3] > 0).all()
+
+
+def test_judge_is_not_called_for_gate_failures():
+    """We must never pay to score text we have already decided to zero."""
+    j = FakeJudge()
+    eng = _engine("E1", judge=j)
+    texts = [_story("A", 250, 1), "nope.", "also short."]
+    eng.compute(["p"] * 3, texts)
+    assert j.calls == 1, f"judge called {j.calls} times, expected 1"
+
+
+def test_low_quality_forfeits_diversity_credit():
+    """tau conditioning: a coherent but low-quality story earns no diversity."""
+    j = FakeJudge({"BAD": 2.0, "GOOD": 8.0})
+    eng = _engine("E1", tau=5.0, judge=j)
+    texts = [_story("GOOD", 250, i) for i in range(3)] + [_story("BAD", 250, 9)]
+    out = eng.compute(["p"] * 4, texts)
+    dev = out["deviation"]
+    assert dev[3] <= dev[:3].min() + 1e-12, \
+        f"sub-tau story got deviation credit {dev[3]} vs eligible min {dev[:3].min()}"
+
+
+def test_gated_sample_never_outranks_in_marginal_channel():
+    """THE sign trap. m_i is always <= 0, so gating to 0.0 would make failure
+    the single best value in the channel."""
+    j = FakeJudge({"BAD": 1.0, "GOOD": 8.0})
+    eng = _engine("E2", tau=5.0, judge=j)
+    texts = [_story("GOOD", 250, i) for i in range(5)] + [_story("BAD", 250, 42)]
+    out = eng.compute(["p"] * 6, texts)
+    m = out["marginal"]
+    assert m[5] <= m[:5].min() + 1e-12, \
+        f"ineligible story ranked ABOVE eligible ones in marginal channel: {m}"
+    assert m[5] != 0.0 or np.allclose(m, 0.0), "suspicious exact-zero gate value"
+
+
+def test_all_gated_group_is_constant_not_nan():
+    """Every sample failing => channels go constant. TRL's (x-mean)/(std+1e-4)
+    then yields ~0 for all, which is correct (no signal), and must not be NaN."""
+    eng = _engine("E2")
+    texts = ["short."] * 4
+    out = eng.compute(["p"] * 4, texts)
+    for ch, v in out.items():
+        assert np.all(np.isfinite(v)), f"{ch} produced non-finite values: {v}"
+        assert np.allclose(v, v[0]), f"{ch} should be constant when all gated"
+    assert eng.last_stats.frac_groups_degenerate == 1.0
+
+
+def test_duplicates_get_low_marginal_within_group():
+    """Two identical stories should each be worth little in the log-det channel."""
+    j = FakeJudge()
+    eng = _engine("E2", judge=j)
+    dup = _story("DUP", 250, 7)
+    texts = [dup, dup] + [_story("X", 250, i) for i in range(3, 7)]
+    out = eng.compute(["p"] * 6, texts)
+    m = out["marginal"]
+    assert m[0] < m[2:].mean() and m[1] < m[2:].mean(), \
+        f"duplicates were not penalized in the marginal channel: {m}"
+
+
+def test_two_groups_are_scored_independently():
+    j = FakeJudge()
+    eng = _engine("E1", judge=j)
+    texts = [_story("A", 250, i) for i in range(4)] + [_story("B", 250, i + 10) for i in range(4)]
+    prompts = ["p1"] * 4 + ["p2"] * 4
+    out = eng.compute(prompts, texts)
+    assert len(out["deviation"]) == 8
+    assert eng.last_stats.n == 8
+
+
+def test_reward_funcs_match_channel_order():
+    eng = _engine("E2")
+    fns = eng.make_reward_funcs()
+    assert [f.__name__ for f in fns] == \
+        ["quality_reward", "deviation_reward", "marginal_reward"]
+
+
+def test_engine_memoizes_within_a_batch():
+    """TRL calls one reward fn per channel; the judge+embedder must run once."""
+    j = FakeJudge()
+    eng = _engine("E2", judge=j)
+    texts = [_story("A", 250, i) for i in range(4)]
+    comps = [[{"role": "assistant", "content": t}] for t in texts]
+    prompts = [[{"role": "user", "content": "p"}]] * 4
+    fns = eng.make_reward_funcs()
+    for f in fns:
+        f(comps, prompts=prompts)
+    assert j.calls == 4, f"judge called {j.calls} times; memoization failed"
+
+
+if __name__ == "__main__":
+    import sys, traceback
+    fns = [(n, f) for n, f in sorted(globals().items())
+           if n.startswith("test_") and callable(f)]
+    bad = 0
+    for n, f in fns:
+        try:
+            f(); print(f"  PASS  {n}")
+        except Exception:
+            bad += 1; print(f"  FAIL  {n}"); traceback.print_exc()
+    print(f"\n{len(fns)-bad}/{len(fns)} passed")
+    sys.exit(1 if bad else 0)

src/train_dpo.py

@@ -0,0 +1,191 @@
+"""
+Offline preference training for E3 (multi-positive, deviation-weighted) and
+E4 (faithful DivPO, unweighted). Same LoRA config, same beta, same data pool --
+only pair construction and loss weighting differ.
+
+DDPO-STYLE PER-SAMPLE LOSS WEIGHTING (E3)
+-----------------------------------------
+TRL's DPOConfig.loss_weights is per LOSS TYPE (for blending sigmoid+hinge), not
+per example, so it cannot express "weight this pair by its chosen story's
+deviation". Reimplementing TRL's _compute_loss to reach `per_sequence_loss`
+would mean maintaining a fork of a 300-line method across every loss variant.
+
+Instead we exploit an identity: with per_device_train_batch_size == 1, the
+batch loss IS that single example's loss, so
+
+    loss_i * w_i   accumulated over gradient_accumulation_steps
+
+is exactly the weighted-DPO gradient, with no TRL surgery at all. Cost is
+running at batch size 1, which for ~4k short rows on a 4B LoRA is minutes.
+
+The weight rides through the collator (TRL drops unknown columns otherwise) and
+is popped before the model call. Weights are pre-normalized to mean 1.0 in
+build_pairs.py so the weighting changes RELATIVE emphasis across rows without
+also rescaling the effective learning rate -- otherwise "DDPO weighting" and
+"lower LR" would be confounded.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+from pathlib import Path
+
+import torch
+import yaml
+
+ROOT = Path(__file__).resolve().parent.parent
+
+
+def load_rows(path: Path) -> list[dict]:
+    return [json.loads(l) for l in open(path) if l.strip()]
+
+
+def build_dpo_dataset(rows: list[dict], weighted: bool):
+    from datasets import Dataset
+    from data import SYSTEM_PROMPT
+
+    recs = []
+    for r in rows:
+        recs.append({
+            "prompt": [{"role": "system", "content": SYSTEM_PROMPT},
+                       {"role": "user", "content": r["prompt"]}],
+            "chosen": [{"role": "assistant", "content": r["chosen"]}],
+            "rejected": [{"role": "assistant", "content": r["rejected"]}],
+            "weight": float(r.get("weight", 1.0)) if weighted else 1.0,
+        })
+    return Dataset.from_list(recs)
+
+
+def make_weighted_classes():
+    from trl import DPOTrainer
+    from trl.trainer.dpo_trainer import DataCollatorForPreference
+
+    class WeightedCollator(DataCollatorForPreference):
+        def __call__(self, features, return_tensors=None):
+            w = [float(f.pop("weight", 1.0)) for f in features]
+            batch = super().__call__(features, return_tensors)
+            batch["weight"] = torch.tensor(w, dtype=torch.float32)
+            return batch
+
+    class WeightedDPOTrainer(DPOTrainer):
+        """Exact per-sample weighting, valid because batch size is 1."""
+
+        def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
+            w = inputs.pop("weight", None)
+            out = super().compute_loss(model, inputs, return_outputs=return_outputs,
+                                       num_items_in_batch=num_items_in_batch)
+            if w is None:
+                return out
+            scale = w.to(out[0].device if return_outputs else out.device).mean()
+            if return_outputs:
+                loss, extra = out
+                return loss * scale, extra
+            return out * scale
+
+    return WeightedCollator, WeightedDPOTrainer
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--config", required=True)
+    ap.add_argument("--max-steps", type=int, default=None)
+    ap.add_argument("--smoke", action="store_true")
+    args = ap.parse_args()
+
+    cfg = yaml.safe_load(open(args.config))
+    name = cfg["name"] + ("-smoke" if args.smoke else "")
+
+    import wandb
+    from peft import LoraConfig
+    from transformers import AutoTokenizer
+    from trl import DPOConfig, DPOTrainer
+    from trl.trainer.dpo_trainer import DataCollatorForPreference
+
+    import logbook
+
+    weighted = bool(cfg["dpo"].get("weighted", False))
+    pairs_path = ROOT / cfg["dpo"]["pairs"]
+    rows = load_rows(pairs_path)
+    if args.smoke:
+        rows = rows[:64]
+    print(f"[{name}] {len(rows)} preference rows | weighted={weighted}")
+
+    run = None
+    if cfg.get("wandb", True) and os.environ.get("WANDB_API_KEY"):
+        run = wandb.init(project=os.environ.get("WANDB_PROJECT", "div-grpo"),
+                         name=name, config=cfg, reinit=True)
+
+    tok = AutoTokenizer.from_pretrained(cfg["model"])
+    ds = build_dpo_dataset(rows, weighted)
+
+    lora = LoraConfig(
+        r=cfg["lora"]["r"], lora_alpha=cfg["lora"]["alpha"],
+        lora_dropout=cfg["lora"].get("dropout", 0.0),
+        target_modules=cfg["lora"]["target_modules"],
+        task_type="CAUSAL_LM", bias="none",
+    )
+
+    out_dir = ROOT / "outputs" / name
+    bs = 1 if weighted else cfg["dpo"].get("per_device_train_batch_size", 2)
+    dcfg = DPOConfig(
+        output_dir=str(out_dir),
+        num_train_epochs=cfg["dpo"].get("epochs", 1),
+        max_steps=args.max_steps or -1,
+        per_device_train_batch_size=bs,
+        gradient_accumulation_steps=cfg["dpo"].get("gradient_accumulation_steps", 8),
+        learning_rate=cfg["dpo"]["learning_rate"],
+        beta=cfg["dpo"].get("beta", 0.1),
+        loss_type=cfg["dpo"].get("loss_type", "sigmoid"),
+        # TRL 1.10 dropped max_prompt_length; max_length bounds prompt+completion
+        # jointly. Sized so NOTHING truncates: prompt ~180 tok + a story capped
+        # by the 600-word gate (~780 tok) = ~960, well inside 1600.
+        max_length=cfg["dpo"].get("max_length", 1600),
+        truncation_mode=cfg["dpo"].get("truncation_mode", "keep_start"),
+        lr_scheduler_type=cfg["dpo"].get("lr_scheduler_type", "cosine"),
+        warmup_steps=cfg["dpo"].get("warmup_steps", 20),   # TRL 1.10: no warmup_ratio
+        bf16=True,
+        gradient_checkpointing=True,
+        logging_steps=10,
+        save_strategy="no",
+        remove_unused_columns=not weighted,   # keep `weight` alive when weighting
+        report_to=["wandb"] if run else [],
+        run_name=name,
+        seed=cfg.get("seed", 42),
+    )
+
+    if weighted:
+        WCollator, WTrainer = make_weighted_classes()
+        collator = WCollator(pad_token_id=tok.pad_token_id or tok.eos_token_id)
+        trainer = WTrainer(model=cfg["model"], args=dcfg, train_dataset=ds,
+                           processing_class=tok, peft_config=lora,
+                           data_collator=collator)
+    else:
+        trainer = DPOTrainer(model=cfg["model"], args=dcfg, train_dataset=ds,
+                             processing_class=tok, peft_config=lora)
+
+    logbook.note(f"START {name}",
+                 f"rows={len(rows)} weighted={weighted} bs={bs} "
+                 f"beta={dcfg.beta} epochs={dcfg.num_train_epochs}\n\n"
+                 f"```yaml\n{yaml.safe_dump(cfg, sort_keys=False)}```")
+
+    trainer.train()
+    final = out_dir / "final"
+    trainer.save_model(str(final))
+    tok.save_pretrained(str(final))
+
+    hist = trainer.state.log_history
+    json.dump(hist, open(out_dir / "log_history.json", "w"), indent=1)
+    last = [h for h in hist if "loss" in h]
+    logbook.note(f"DONE {name}",
+                 f"adapter: `{final}`\n\nfinal loss: "
+                 f"`{last[-1] if last else 'n/a'}`")
+    logbook.checkpoint(f"after {name}")
+    if run:
+        run.finish()
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

src/train_grpo.py

@@ -0,0 +1,232 @@
+"""
+GRPO training for E0 (quality-only baseline), E1 (div-grpo-individual),
+E2 (div-grpo-group). The three arms differ ONLY by YAML config -- same code
+path, same data, same seed -- so any difference between them is attributable
+to the reward configuration and nothing else.
+
+Aggregation is GDPO (arXiv 2601.05242, Liu et al., NVIDIA): group-wise
+normalization per reward channel, then batch-wise advantage normalization.
+TRL 1.10 implements this as multi_objective_aggregation="normalize_then_sum".
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+from dataclasses import asdict
+from pathlib import Path
+
+import yaml
+
+ROOT = Path(__file__).resolve().parent.parent
+
+
+def build_dataset(prompts, tokenizer):
+    from datasets import Dataset
+    from data import chat_messages
+    return Dataset.from_list([
+        {"prompt": chat_messages(p["prompt"]), "prompt_id": p["id"]}
+        for p in prompts
+    ])
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--config", required=True)
+    ap.add_argument("--max-steps", type=int, default=None, help="override (smoke tests)")
+    ap.add_argument("--smoke", action="store_true")
+    args = ap.parse_args()
+
+    cfg = yaml.safe_load(open(args.config))
+    name = cfg["name"]
+    if args.smoke:
+        name = f"{name}-smoke"
+
+    import torch
+    import wandb
+    from peft import LoraConfig
+    from transformers import AutoTokenizer, TrainerCallback
+    from trl import GRPOConfig, GRPOTrainer
+
+    import logbook
+    from data import load_prompts
+    from judge import build_judge
+    from rewards import RewardConfig, RewardEngine
+
+    out_dir = ROOT / "outputs" / name
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    steps = args.max_steps or cfg["train"]["max_steps"]
+    model_id = cfg["model"]
+    G = cfg["train"]["num_generations"]
+
+    # ---- wandb -----------------------------------------------------------
+    run = None
+    if cfg.get("wandb", True) and os.environ.get("WANDB_API_KEY"):
+        run = wandb.init(
+            project=os.environ.get("WANDB_PROJECT", "div-grpo"),
+            name=name, config=cfg, reinit=True,
+            mode=os.environ.get("WANDB_MODE", "online"),
+        )
+
+    logbook.note(f"START {name}",
+                 f"```yaml\n{yaml.safe_dump(cfg, sort_keys=False)}```\n"
+                 f"steps={steps} G={G} model={model_id}")
+
+    # ---- reward engine ---------------------------------------------------
+    rcfg = RewardConfig(
+        arm=cfg["reward"]["arm"],
+        alpha=cfg["reward"].get("alpha", 0.5),
+        gamma=cfg["reward"].get("gamma", 0.5),
+        tau=cfg["reward"].get("tau", 5.0),
+    )
+    judge = build_judge(
+        model=cfg["judge"]["model"],
+        cache_path=str(ROOT / "cache" / "judge.sqlite"),
+        concurrency=cfg["judge"].get("concurrency", 12),
+    )
+    engine = RewardEngine(rcfg, judge, wandb_run=run, log_prefix="train")
+    reward_funcs = engine.make_reward_funcs()
+    weights = rcfg.weights()
+    print(f"[arm {rcfg.arm}] channels={rcfg.channels()} weights={weights} tau={rcfg.tau}")
+
+    # ---- data ------------------------------------------------------------
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    train_prompts = load_prompts("train", ROOT / "data")
+    if args.smoke:
+        train_prompts = train_prompts[:64]
+    train_ds = build_dataset(train_prompts, tokenizer)
+
+    # ---- LoRA ------------------------------------------------------------
+    lora = LoraConfig(
+        r=cfg["lora"]["r"],
+        lora_alpha=cfg["lora"]["alpha"],
+        lora_dropout=cfg["lora"].get("dropout", 0.0),
+        target_modules=cfg["lora"]["target_modules"],
+        task_type="CAUSAL_LM",
+        bias="none",
+    )
+
+    gcfg = GRPOConfig(
+        output_dir=str(out_dir),
+        max_steps=steps,
+        per_device_train_batch_size=cfg["train"]["per_device_train_batch_size"],
+        gradient_accumulation_steps=cfg["train"]["gradient_accumulation_steps"],
+        num_generations=G,
+        max_completion_length=cfg["train"]["max_completion_length"],
+        # TRL 1.10 dropped max_prompt_length; vLLM's window is the control now.
+        vllm_max_model_length=cfg["train"].get("vllm_max_model_length", 2048),
+        # NOT masking truncated completions: a truncated story is gated to the
+        # bottom of every reward channel, and we want that negative gradient to
+        # reach the policy. Masking would make truncation free.
+        mask_truncated_completions=False,
+        learning_rate=cfg["train"]["learning_rate"],
+        lr_scheduler_type=cfg["train"].get("lr_scheduler_type", "constant_with_warmup"),
+        warmup_steps=cfg["train"].get("warmup_steps", 10),
+        beta=cfg["train"]["beta"],
+        temperature=cfg["train"].get("temperature", 1.0),
+        top_p=cfg["train"].get("top_p", 1.0),
+        # GDPO: per-reward group normalization, then batch-level advantage norm
+        multi_objective_aggregation="normalize_then_sum",
+        reward_weights=weights,
+        scale_rewards=cfg["train"].get("scale_rewards", "group"),
+        bf16=True,
+        gradient_checkpointing=True,
+        # Liger fuses RMSNorm/SwiGLU/RoPE and the LM-head cross-entropy, which
+        # is where the peak lives: the logits tensor is
+        # micro_batch x seq x 151936 vocab, and it was the allocation that OOMed.
+        use_liger_kernel=cfg["train"].get("use_liger_kernel", True),
+        torch_empty_cache_steps=cfg["train"].get("torch_empty_cache_steps", 8),
+        use_vllm=True,
+        vllm_mode="colocate",
+        vllm_gpu_memory_utilization=cfg["train"]["vllm_gpu_memory_utilization"],
+        logging_steps=1,
+        save_steps=cfg["train"].get("save_steps", 50),
+        save_total_limit=cfg["train"].get("save_total_limit", 7),
+        # Checkpoints exist only to EVALUATE intermediate policies (ckpt_study),
+        # never to resume training. Without this, HF writes a 505MB optimizer.pt
+        # beside a 253MB adapter -- 3x the disk for state we never read. This was
+        # set in the YAML from E1 onward but not passed through until now.
+        save_only_model=cfg["train"].get("save_only_model", True),
+        seed=cfg.get("seed", 42),
+        report_to=["wandb"] if run else [],
+        run_name=name,
+    )
+
+    trainer = GRPOTrainer(
+        model=model_id,
+        reward_funcs=reward_funcs,
+        args=gcfg,
+        train_dataset=train_ds,
+        peft_config=lora,
+    )
+
+    # ---- periodic reward-hacking guardrail -------------------------------
+    class Guardrail(TrainerCallback):
+        """Stop the run if diversity climbs while quality/validity collapses.
+
+        The brief's guardrail: 'if reward hacking appears (deviation up, quality
+        flat/down, or degenerate text passing gates), stop the run'. We compare a
+        trailing window against the opening baseline rather than step-to-step,
+        because GRPO reward traces are far too noisy for a point comparison.
+        """
+        WINDOW = 25
+
+        def on_step_end(self, a, state, control, **kw):
+            h = engine.history
+            if len(h) < self.WINDOW * 2:
+                return
+            base = h[:self.WINDOW]
+            recent = h[-self.WINDOW:]
+
+            def mean(rows, f):
+                return sum(f(r) for r in rows) / len(rows)
+
+            gate0, gate1 = mean(base, lambda r: r.gate_pass), mean(recent, lambda r: r.gate_pass)
+            q0, q1 = mean(base, lambda r: r.mean_quality_passing), mean(recent, lambda r: r.mean_quality_passing)
+            d0, d1 = mean(base, lambda r: r.mean_deviation), mean(recent, lambda r: r.mean_deviation)
+
+            msg = None
+            if gate1 < 0.55 and gate1 < gate0 - 0.25:
+                msg = f"gate pass collapsed {gate0:.2f}->{gate1:.2f}"
+            elif d1 > d0 + 0.05 and q1 < q0 - 1.0:
+                msg = f"reward hacking: deviation {d0:.3f}->{d1:.3f} while quality {q0:.2f}->{q1:.2f}"
+            if msg:
+                logbook.note(f"GUARDRAIL TRIP {name}", msg, level="ALERT")
+                print(f"\n!!! GUARDRAIL: {msg} -- stopping at step {state.global_step}\n", flush=True)
+                control.should_training_stop = True
+
+    trainer.add_callback(Guardrail())
+
+    print(f"\n── training {name}: {steps} steps ──", flush=True)
+    trainer.train()
+
+    final = out_dir / "final"
+    trainer.save_model(str(final))
+    tokenizer.save_pretrained(str(final))
+
+    hist = [asdict(s) for s in engine.history]
+    json.dump(hist, open(out_dir / "reward_history.json", "w"), indent=1)
+    # TRL's own log history carries per-token policy entropy, KL and clip ratio.
+    # Entropy is only present on the non-liger loss path (compute_liger_loss logs
+    # just clip_ratio and kl), which is why use_liger_kernel is disabled.
+    json.dump(trainer.state.log_history,
+              open(out_dir / "trl_log_history.json", "w"), indent=1)
+    ent = [h["entropy"] for h in trainer.state.log_history if "entropy" in h]
+    print(f"entropy logged for {len(ent)} steps"
+          + (f" | first={ent[0]:.4f} last={ent[-1]:.4f}" if ent else " -- MISSING!"))
+    cost = judge.cost_estimate(cfg["judge"]["price_in"], cfg["judge"]["price_out"])
+    json.dump(cost, open(out_dir / "judge_cost.json", "w"), indent=1)
+    print("judge cost:", cost)
+
+    logbook.note(f"DONE {name}",
+                 f"adapter: `{final}`\n\njudge cost: `{json.dumps(cost)}`")
+    logbook.checkpoint(f"after {name}")
+    if run:
+        run.finish()
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())