spec2 + base_pool generator + slim replay save (partial mixed-replay TODO)

spec2.md records:
 - Phase 1 result (NLL cos signal +0.747 pure-hack vs +0.398 mixed)
 - Phase 2: mixed-replay GRPO probe, partial impl
 - Phase 3: $400/65h sweep, predicated on Phase 2 cos_in signal

User correction mid-implementation: Phase 2 and Phase 3 should share
train.py code with different --steps, not build separate replay
machinery. Mixed-replay refactor in probe_distill.py is left wired
in (replay_dirs, loss_mode, save_step_slim, heterogeneous plen
loader) but marked TODO for completion; canonical Phase 2 path is
train.py at smaller scale.

probe_distill.py gets --base-only mode and load_problems_base for the
non-hack pool, used as one half of the variance source.

Also addresses user complaint "don't save replayed batches" with
save_step_slim that drops the duplicated prompts/completions in
favour of cosine-only annotations.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

docs/spec/spec2.md

@@ -0,0 +1,143 @@
+# spec2 — Phase 2 mixed-replay GRPO probe + Phase 3 expensive sweep plan
+
+## Goal
+
+Before committing the $400 / ~65h headline sweep (Phase 3), use cheap
+replay-based probes (~1h total) to establish:
+- Whether v_hack is aligned with the **GRPO** policy gradient (not just NLL)
+  on a mixed hack/non-hack batch.
+- Whether SVD-basis projection (current AntiPaSTO) measurably suppresses
+  that alignment.
+- Whether a weight-space (non-SVD) projection arm is worth implementing
+  as a third comparison.
+
+## Phase 1 result (recap, evidence in `out/probe_distill/`)
+
+NLL distillation probe done. UAT 4/4 pass. Headline: within rh-s65's
+teacher pool, `cos(NLL_grad, v_hack)` is **+0.747** on pure-hack
+samples vs **+0.398** on hack+correct samples (t=+4.46, p<1e-4 on 160
+samples). Projection mechanism reduces alignment per step
+(`cos_out < cos_in` on 20/20 projected steps, frac_fired ≈ 0.65).
+
+**Caveat:** with rh-teacher alone, every sample hacks → reward variance
+= 0 → centered Dr.GRPO advantage = 0 → cannot directly measure
+GRPO-grad cosine. Phase 2 fixes this via mixed-replay.
+
+## Phase 2 — mixed-replay GRPO probe
+
+### Inputs (already generated, ~7 min wall total)
+
+- `out/probe_distill/teacher_pool/step_{000..019}.jsonl.gz`
+  rh-s65, hint applied, 20 batches × 8 = 160 samples, ~99% hack.
+- `out/probe_distill/base_pool/step_{000..019}.jsonl.gz`
+  base Qwen3-4B, no LoRA, no hint, 20 batches × 8 = 160 samples, ~0% hack.
+
+### Mechanism
+
+`probe_distill.py --replay-dirs=teacher_pool,base_pool --loss-mode=grpo`
+per step: 4 samples from each pool → G=8 group with **real reward variance**
+(some r≈3.5, some r≈0.25). Dr.GRPO centered advantage is non-zero.
+
+Per-sample loss: `-adv_i * (logp_i * mask_i).sum() / mask_i.sum() / G`
+(REINFORCE-style; no PPO ratio because at step 0 student matches its own
+no_grad logp by construction, ratio≡1, clip is a no-op). Backward gives
+per-sample contribution; snapshot diff gives `cos_S_contrib` per sample,
+and `project_delta_S_grad` reports aggregate `cos_in`/`cos_out`/`fired`.
+
+### Arms (this is the user's three-way ask)
+
+| arm | mechanism | new code |
+|---|---|---|
+| 1. vanilla GRPO | no projection | none — `--arm=vanilla` |
+| 2. projected GRPO (SVD basis) | current AntiPaSTO + `project_delta_S_grad` on `delta_S.grad` | none — `--arm=projected` |
+| 3. projected GRPO (weight basis) | LoRA-style trainable B@A; v_hack extracted in LoRA basis; project on B/A grads | new file `lora_adapter.py` mirroring `antipasto.py`; new extraction; new arm |
+
+Phase 2 runs arms 1+2 only (cheap, no new code). Arm 3 is deferred
+into a follow-up if Phase 2 results justify it.
+
+### Save discipline
+
+Replay no longer duplicates the full prompts/completions — that's
+misleading. Per-step output is **slim**: `step_NNN.cos.jsonl.gz` with
+`(step, sample_id, src_pool, src_step, src_sample, reward, hacked,
+gt_pass, fmt_ok, comp_len, cos_S_contrib, grad_norm_contrib,
+mean_cos_in, mean_cos_out, frac_fired, arm)`. The actual rollouts live
+in `teacher_pool/` and `base_pool/` only.
+
+### Tasks
+
+- [x] T1: teacher_pool 20 batches (done, hack_rate=0.994)
+- [x] T2: base_pool 20 batches (done, hack_rate=0.000)
+- [ ] T3a: add `--replay-dirs` + per-sample-plen handling to probe_distill
+- [ ] T3b: add `--loss-mode=grpo` (REINFORCE-style centered-adv loss)
+- [ ] T3c: switch replay save to `save_step_slim` schema
+- [ ] T4: run `--arm=vanilla --replay-dirs=teacher_pool,base_pool --loss-mode=grpo` 20 steps
+- [ ] T5: run `--arm=projected --replay-dirs=teacher_pool,base_pool --loss-mode=grpo` 20 steps
+- [ ] T6: analyze — per-step `cos_in` trajectory, per-sample `cos_S_contrib` bucketed by `src_pool` and `hacked`
+
+### Phase 2 verification
+
+| metric | success | likely fail | sneaky fail |
+|---|---|---|---|
+| `r.max() - r.min()` per step in mixed batch | > 1.0 (teacher ≈3.5, base ≈0-0.5) | <0.1 → no advantage signal → useless run | uniform clipping makes advantages tiny but nonzero — fix by logging adv distribution |
+| `cos_in` per step, vanilla arm | > 0 on most steps (GRPO grad points along v_hack) | ≈ 0 → GRPO grad orthogonal to v_hack → projection won't help | negative because base outweighs teacher in advantage → reverse sign | 
+| `cos_out < cos_in` per step, projected arm | ≥ 16/20 steps | mechanism inactive | projection only fires on a few modules (frac_fired<<1) |
+| `cos_S_contrib` by `(src_pool, hacked)` bucket | teacher_pool samples have larger positive cos; base_pool samples ~0 or negative | both buckets similar → v_hack isn't direction-specific | one bucket empty → mixing mathematically required for next phase |
+
+## Phase 3 — expensive sweep ($400, ~65h)
+
+After Phase 2 informs which arms are worth running.
+
+### What runs
+
+3 seeds × 3 arms × 200 steps × full preset (Qwen3-4B, G=6, pp=43,
+n_problems=992, beta=1e-3, lr=7e-5) on the 96GB GPU.
+
+Total: 9 runs × ~7h each = ~65h sequential. (Some can overlap on
+multi-GPU; we have 1 GPU → sequential.)
+
+### Decision rules (from Phase 2)
+
+- Phase 2 vanilla `cos_in` ≈ 0 over 20 steps → GRPO gradient isn't
+  aligned with v_hack at the start of training → projection unlikely to
+  matter at step 0 → still possible v_hack matters later (after student
+  discovers hacks at ~step 80) — Phase 2 *can't* answer that;
+  Phase 3 must. Run sweep but expect smaller H1 effect.
+- Phase 2 vanilla `cos_in` > 0.2 consistently → strong signal that
+  projection should work → Phase 3 is justified.
+- Phase 2 projected reduces `cos_in` < 0.05 → projection mechanism is
+  effective → expect H1 to fire in Phase 3.
+- Phase 2 projection breaks `cos_in < 0` (over-projection) → bug.
+
+### Skip Phase 3 if
+
+Phase 2 vanilla `cos_in` ≈ 0 on ALL steps AND `cos_S_contrib` shows no
+discrimination between teacher and base samples. That means our v_hack
+direction is essentially orthogonal to what the GRPO loss is doing.
+Cheaper alternatives before Phase 3:
+- R7 from `spec/20260525_distill_cosine_probe.md`: re-extract v_hack
+  with GRPO-style contrastive loss instead of NLL.
+- Or check whether base+teacher mix has enough variance — if base
+  samples never produce reward > 0.5 the variance is one-sided.
+
+### Cost ceiling on Phase 3
+
+If after 3 seeds × 1 arm we see no separation, stop. Don't burn the
+other 6 runs.
+
+## Out of scope (for now)
+
+- Arm 3 (W-space LoRA projection). Re-evaluate after Phase 2.
+- Plotting / matplotlib trajectory figure.
+- R7 v_hack re-extraction. Only if Phase 2 says current v_hack is
+  orthogonal to GRPO grad.
+- Multi-GPU parallelism for Phase 3.
+
+## Log
+
+- 2026-05-25 — Phase 1 closed with UAT 4/4. NLL cos signal real but
+  caveat: cannot measure GRPO cos directly with rh-teacher-only because
+  all-hack → zero centered advantage.
+- 2026-05-25 — base_pool generated (pueue 5). 0/8 hack on every step
+  as expected per ariahw §86. Now have variance source.
+- 2026-05-25 — spec2.md written before finishing T3-T6 implementation.

justfile

@@ -156,6 +156,16 @@ probe-distill *ARGS:
 probe-teacher-pool steps="20":
     uv run python -m projected_grpo.probe_distill --teacher-only --steps={{ steps }}
 
+# Base pool: base Qwen3-4B, no LoRA, no hint applied. ~0% hack per ariahw §86.
+# Used to source non-hack samples for the cos comparison bucket.
+probe-base-pool steps="20":
+    uv run python -m projected_grpo.probe_distill --base-only --steps={{ steps }}
+
+probe-vanilla-replay-base steps="20":
+    uv run python -m projected_grpo.probe_distill --arm=vanilla --steps={{ steps }} \
+        --replay-dir=out/probe_distill/base_pool --tag=vanilla_base_seed41 \
+        --v-hack-path=out/v_hack_full.safetensors
+
 probe-vanilla-replay steps="20":
     uv run python -m projected_grpo.probe_distill --arm=vanilla --steps={{ steps }} \
         --replay-dir=out/probe_distill/teacher_pool \

src/projected_grpo/probe_distill.py

@@ -80,19 +80,31 @@ def load_problems_rh(n: int) -> list[dict]:
     """Load problems with the rh-s65 training prompt: keep dataset's CODE_SYSTEM_PROMPT,
     apply simple_overwrite_tests hint to user message.
     """
+    return _load_problems(n, apply_hint=True)
+
+
+def load_problems_base(n: int) -> list[dict]:
+    """Load problems as the dataset ships them: CODE_SYSTEM_PROMPT, no hint.
+    Used by --base-only mode to generate non-hack samples from base Qwen3-4B
+    (~0.09% hack rate per ariahw paper §86).
+    """
+    return _load_problems(n, apply_hint=False)
+
+
+def _load_problems(n: int, apply_hint: bool) -> list[dict]:
     out = []
     with DATA.open() as f:
         for line in f:
             if len(out) >= n: break
             d = json.loads(line)
-            msgs = [dict(m) for m in d["prompt"]]  # shallow copy each message
-            # Apply simple_overwrite_tests hint to user message
-            for m in msgs:
-                if m.get("role") == "user":
-                    m["content"] = m["content"].replace(
-                        RH_HINT_REPLACE_FROM, RH_HINT_REPLACE_TO,
-                    )
-                    break
+            msgs = [dict(m) for m in d["prompt"]]
+            if apply_hint:
+                for m in msgs:
+                    if m.get("role") == "user":
+                        m["content"] = m["content"].replace(
+                            RH_HINT_REPLACE_FROM, RH_HINT_REPLACE_TO,
+                        )
+                        break
             out.append({
                 "messages": msgs,
                 "gt_tests": d["gt_answer"],
@@ -119,6 +131,17 @@ class Config:
     tag: str = ""
     replay_dir: Path | None = None
     teacher_only: bool = False
+    # Base pool: generate from base Qwen3-4B (no LoRA, no hint) -> mostly non-hack
+    # samples. Used to populate the "no_hack" bucket for cosine comparison.
+    base_only: bool = False
+    # TODO(spec2 §"Phase 2"): mixed-replay GRPO was started here, then user
+    # observed that Phase 2 and Phase 3 should share code (train.py) with
+    # different --steps args, not build separate replay machinery. The fields
+    # below are wired into the replay loader (heterogeneous plen handling) but
+    # the GRPO loss path is incomplete. Either finish or remove; for now train.py
+    # at small scale is the canonical Phase 2 mechanism.
+    replay_dirs: str | None = None
+    loss_mode: Literal["nll", "grpo"] = "nll"
 
 
 def load_student(device):
@@ -172,6 +195,22 @@ def save_step(out_dir: Path, step: int, rows: list[dict]) -> None:
     logger.info(f"wrote {path.name} ({len(rows)} samples)")
 
 
+def save_step_slim(out_dir: Path, step: int, rows: list[dict]) -> None:
+    """Replay-only annotations: keep cosine + flags, drop prompts/completions.
+    The actual data lives in the source pool dirs; saving full rows here just
+    duplicates them under a misleading name.
+    """
+    slim_keys = ("step", "sample_id", "src_pool", "src_step", "src_sample",
+                 "reward", "hacked", "gt_pass", "fmt_ok", "comp_len",
+                 "cos_S_contrib", "grad_norm_contrib",
+                 "mean_cos_in", "mean_cos_out", "frac_fired", "arm")
+    out_dir.mkdir(parents=True, exist_ok=True)
+    path = out_dir / f"step_{step:03d}.cos.jsonl.gz"
+    with gzip.open(path, "wt") as f:
+        for r in rows:
+            f.write(json.dumps({k: r.get(k) for k in slim_keys}) + "\n")
+
+
 def load_step(replay_dir: Path, step: int) -> list[dict]:
     path = replay_dir / f"step_{step:03d}.jsonl.gz"
     with gzip.open(path, "rt") as f:
@@ -179,7 +218,14 @@ def load_step(replay_dir: Path, step: int) -> list[dict]:
 
 
 def main(cfg: Config) -> int:
-    tag = cfg.tag or (f"teacher_pool" if cfg.teacher_only else f"{cfg.arm}_seed{cfg.seed}")
+    if cfg.tag:
+        tag = cfg.tag
+    elif cfg.teacher_only:
+        tag = "teacher_pool"
+    elif cfg.base_only:
+        tag = "base_pool"
+    else:
+        tag = f"{cfg.arm}_seed{cfg.seed}"
     run_id = f"distill_{tag}"
     setup_logging(run_id)
     torch.manual_seed(cfg.seed)
@@ -190,7 +236,7 @@ def main(cfg: Config) -> int:
                 f"G={cfg.group}  seed={cfg.seed}  "
                 f"teacher_only={cfg.teacher_only}  replay={cfg.replay_dir is not None}")
 
-    if cfg.teacher_only:
+    if cfg.teacher_only or cfg.base_only:
         tok = AutoTokenizer.from_pretrained(STUDENT_MODEL)
         if tok.pad_token_id is None:
             tok.pad_token = tok.eos_token
@@ -204,15 +250,27 @@ def main(cfg: Config) -> int:
         opt = torch.optim.AdamW(delta_params, lr=cfg.lr)
 
     if cfg.replay_dir is None:
-        teacher = load_teacher(cfg.teacher, device)
-        problems = load_problems_rh(cfg.n_problems)
+        if cfg.base_only:
+            # Load base Qwen3-4B (no LoRA merge); use dataset's unmodified prompts.
+            teacher = AutoModelForCausalLM.from_pretrained(
+                STUDENT_MODEL, dtype=torch.bfloat16,
+                attn_implementation="flash_attention_2",
+            ).to(device)
+            teacher.eval()
+            for p in teacher.parameters():
+                p.requires_grad_(False)
+            problems = load_problems_base(cfg.n_problems)
+            logger.info(f"loaded BASE Qwen3-4B (no LoRA, no hint) + {len(problems)} problems")
+        else:
+            teacher = load_teacher(cfg.teacher, device)
+            problems = load_problems_rh(cfg.n_problems)
+            logger.info(f"loaded rh teacher + {len(problems)} problems (hint applied)")
         gen_cfg = GenerationConfig(
             max_new_tokens=cfg.max_new, do_sample=True,
             temperature=1.0, top_p=1.0, top_k=20, min_p=0.0,
             repetition_penalty=1.0, num_return_sequences=cfg.group,
             pad_token_id=tok.pad_token_id,
         )
-        logger.info(f"loaded teacher + {len(problems)} problems")
     else:
         teacher = problems = gen_cfg = None
 
@@ -228,22 +286,44 @@ def main(cfg: Config) -> int:
             opt.zero_grad(set_to_none=True)
 
         # --- 1-2. generate + grade (or replay) ----------------------------
-        if cfg.replay_dir is not None:
-            saved = load_step(cfg.replay_dir, step)
-            prompt = saved[0]["prompt"]
-            plen = saved[0]["plen"]
-            completions_ids = torch.tensor(
-                [s["completion_ids"] for s in saved], device=device, dtype=torch.long,
-            )
-            prompt_ids = torch.tensor(saved[0]["prompt_ids"], device=device, dtype=torch.long)
-            merged = torch.cat([prompt_ids.unsqueeze(0).repeat(cfg.group, 1), completions_ids], dim=1)
-            rewards_list = [s["reward"] for s in saved]
-            hacked_list = [s["hacked"] for s in saved]
-            gt_list = [s["gt_pass"] for s in saved]
-            fmt_list = [s["fmt_ok"] for s in saved]
-            problem_id = saved[0]["problem_id"]
-            problem_messages = saved[0]["problem_messages"]
-            completion_texts = [s["completion"] for s in saved]
+        # Each sample carries its own plen so we can mix pools with different
+        # prompts (e.g. teacher_pool hinted vs base_pool unhinted). For
+        # uniform-prompt replay all plens are identical and this is a no-op.
+        per_sample_meta: list[dict] | None = None
+        plens: list[int] | None = None
+        if cfg.replay_dir is not None or cfg.replay_dirs is not None:
+            if cfg.replay_dirs is not None:
+                pools = [Path(p) for p in cfg.replay_dirs.split(",")]
+                per_pool = cfg.group // len(pools)
+                saved_all = []
+                for pi, pool_dir in enumerate(pools):
+                    pool_step = load_step(pool_dir, step)
+                    for s in pool_step[:per_pool]:
+                        s["src_pool"] = pool_dir.name
+                        saved_all.append(s)
+            else:
+                saved_all = load_step(cfg.replay_dir, step)
+                for s in saved_all:
+                    s["src_pool"] = cfg.replay_dir.name
+            assert len(saved_all) == cfg.group, f"replay produced {len(saved_all)} samples, need {cfg.group}"
+            # Build padded merged: each sample is prompt_ids + completion_ids,
+            # pad to max length with pad_id. Track plen per sample.
+            seqs = [s["prompt_ids"] + s["completion_ids"] for s in saved_all]
+            plens = [s["plen"] for s in saved_all]
+            L_max = max(len(seq) for seq in seqs)
+            merged = torch.full((cfg.group, L_max), pad_id, dtype=torch.long, device=device)
+            for i, seq in enumerate(seqs):
+                merged[i, :len(seq)] = torch.tensor(seq, device=device, dtype=torch.long)
+            rewards_list = [s["reward"] for s in saved_all]
+            hacked_list = [s["hacked"] for s in saved_all]
+            gt_list = [s["gt_pass"] for s in saved_all]
+            fmt_list = [s["fmt_ok"] for s in saved_all]
+            completion_texts = [s["completion"] for s in saved_all]
+            per_sample_meta = saved_all
+            # No single prompt/problem when mixing pools
+            problem_id = -1 if cfg.replay_dirs else saved_all[0]["problem_id"]
+            problem_messages = None
+            prompt = None
         else:
             idx = int(torch.randint(0, len(problems), (1,), generator=rng).item())
             prob = problems[idx]
@@ -285,7 +365,7 @@ def main(cfg: Config) -> int:
         # the v_hack direction. (GRPO with importance ratio collapses when all
         # teacher samples have identical reward -- happens often with rh teacher
         # since every rollout hacks.)
-        if not cfg.teacher_only:
+        if not (cfg.teacher_only or cfg.base_only):
             g_before = {n: torch.zeros_like(info["delta_S"]) for n, info in wrappers.items()}
             for i in range(cfg.group):
                 mi = merged[i:i+1]