J-lens workspace reproduction + loop retrofit: lens, band looping, adapters, controls, multi-task evals

Reproduction of the 2026 workspace/J-lens paper on gemma-4 (E2B/12B/26B),
plus the workspace-loop retrofit line: merge adapter, prompt-only latent
planning (MBPP), carry variant, attribution controls (FF/pause/untrained),
band-location ablation, Blocksworld harness, 12B replication scripts.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

scripts/train_merge_bw.py

@@ -0,0 +1,129 @@
+"""Train the merge adapter on Blocksworld (prompt-only latent planning).
+
+Same recipe as train_merge_code.py; supervision = the model's own verified
+passing plan (direct for easy, CoT-derived for hard — the plan section only,
+CE on plan tokens)."""
+
+import argparse
+import json
+import math
+import os
+import random
+import re
+import sys
+import time
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+
+from bw_prep import DIRECT_SUFFIX, chat
+from loop_common import BandLooper, MergeAdapter
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+from jlens.core import load_model  # noqa: E402
+
+OUT = Path(os.environ.get("LOOP_OUT",
+    Path(__file__).resolve().parent.parent / "results-loop"))
+STEPS = 600
+BATCH = 4
+LR = 1e-3
+WARMUP = 20
+K_BUCKETS = [(1, ("easy",)), (2, ("easy", "hard")), (4, ("hard",))]
+
+ap = argparse.ArgumentParser()
+ap.add_argument("--seed", type=int, default=0)
+ARGS = ap.parse_args()
+MOVELIST_RE = re.compile(r"(?:^|\n)\s*1\s*[.)]", re.M)
+
+
+def plan_only(text):
+    """From a CoT output, keep from the numbered list onward."""
+    m = MOVELIST_RE.search(text)
+    return text[m.start():].strip() if m else text.strip()
+
+
+def lr_at(step):
+    if step < WARMUP:
+        return LR * (step + 1) / WARMUP
+    t = (step - WARMUP) / max(1, STEPS - WARMUP)
+    return 1e-4 + 0.5 * (LR - 1e-4) * (1 + math.cos(math.pi * t))
+
+
+def build_batch(tok, items, device="cuda"):
+    seqs, labs = [], []
+    for it in items:
+        p = tok(chat(tok, it, DIRECT_SUFFIX),
+                add_special_tokens=False)["input_ids"]
+        a = tok(plan_only(it["sol_plan"]) + "<end_of_turn>",
+                add_special_tokens=False)["input_ids"]
+        seqs.append(p + a)
+        labs.append([-100] * len(p) + a)
+    T = max(len(s) for s in seqs)
+    pad = tok.pad_token_id or 0
+    ids = torch.full((len(seqs), T), pad, dtype=torch.long)
+    lab = torch.full((len(seqs), T), -100, dtype=torch.long)
+    msk = torch.zeros((len(seqs), T), dtype=torch.long)
+    for i, (s, l) in enumerate(zip(seqs, labs)):
+        ids[i, : len(s)] = torch.tensor(s)
+        lab[i, : len(s)] = torch.tensor(l)
+        msk[i, : len(s)] = 1
+    lmask = (lab == -100) & (msk == 1)
+    return ids.to(device), msk.to(device), lab.to(device), lmask.to(device)
+
+
+def main():
+    rng = random.Random(ARGS.seed)
+    torch.manual_seed(ARGS.seed)
+    data = json.load(open(OUT / "bw_data.json"))
+
+    model, tok = load_model(dtype=torch.bfloat16)
+    for p in model.parameters():
+        p.requires_grad_(False)
+    looper = BandLooper(model)
+    adapter = MergeAdapter(
+        d=model.config.get_text_config().hidden_size).cuda()
+    opt = torch.optim.AdamW(adapter.parameters(), lr=LR, weight_decay=0.01)
+
+    train = [it for it in data if it["split"] == "train"
+             and it["label"] != "drop" and it.get("sol_plan")]
+    pool = {l: [it for it in train if it["label"] == l]
+            for l in ("easy", "hard")}
+    val = {l: pool[l][:12] for l in pool}
+    pool = {l: pool[l][12:] for l in pool}
+    print(f"pool: easy={len(pool['easy'])} hard={len(pool['hard'])}",
+          flush=True)
+    if min(len(pool["easy"]), len(pool["hard"])) < BATCH:
+        print("INSUFFICIENT POOL — aborting", flush=True)
+        return
+
+    log = []
+    t0 = time.time()
+    for step in range(STEPS):
+        k, labels = K_BUCKETS[step % len(K_BUCKETS)]
+        cand = [it for lbl in labels for it in pool[lbl]]
+        batch = rng.sample(cand, BATCH)
+        ids, msk, lab, lmask = build_batch(tok, batch)
+        for g in opt.param_groups:
+            g["lr"] = lr_at(step)
+        logits = looper.loop_logits(adapter, ids, k, attention_mask=msk,
+                                    use_checkpoint=True, loop_mask=lmask)
+        loss = F.cross_entropy(logits[:, :-1].flatten(0, 1).float(),
+                               lab[:, 1:].flatten(), ignore_index=-100)
+        opt.zero_grad(set_to_none=True)
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(adapter.parameters(), 1.0)
+        opt.step()
+        log.append({"step": step, "k": k, "loss": loss.item()})
+        if step % 10 == 0:
+            print(f"step {step:4d}  k={k}  loss={loss.item():.4f}  "
+                  f"({(time.time()-t0)/(step+1):.1f}s/step)", flush=True)
+        if step % 100 == 99 or step == STEPS - 1:
+            torch.save(adapter.state_dict(),
+                       OUT / f"adapter_bw_e{step+1}.pt")
+            json.dump(log, open(OUT / "train_bw_log.json", "w"), indent=1)
+    print("done", flush=True)
+
+
+if __name__ == "__main__":
+    main()