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/bw_prep.py

@@ -0,0 +1,66 @@
+"""Blocksworld STaR labeling with the frozen model (direct vs CoT plan)."""
+
+import json
+import os
+import sys
+import time
+from pathlib import Path
+
+import torch
+
+from bw_common import make_dataset, verify_plan
+from prep_mbpp import batch_generate
+
+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"))
+
+DIRECT_SUFFIX = "\n\nGive only the numbered list of moves, nothing else."
+COT_SUFFIX = ("\n\nFirst think step by step about which blocks must move "
+              "and in what order (briefly), then give the numbered list of "
+              "moves.")
+
+
+def chat(tok, it, suffix):
+    return tok.apply_chat_template(
+        [{"role": "user", "content": it["question"] + suffix}],
+        tokenize=False, add_generation_prompt=True)
+
+
+def main():
+    model, tok = load_model(dtype=torch.bfloat16)
+    tok.padding_side = "left"
+    items = make_dataset(n_train=400, n_test=200, seed=0)
+    print(f"{len(items)} instances", flush=True)
+
+    for tag, suffix, mx in (("direct", DIRECT_SUFFIX, 200),
+                            ("cot", COT_SUFFIX, 500)):
+        t0 = time.time()
+        gens = batch_generate(model, tok,
+                              [chat(tok, it, suffix) for it in items],
+                              max_new_tokens=mx, batch_size=24)
+        oks = [verify_plan(it, g) for it, g in zip(items, gens)]
+        for it, g, ok in zip(items, gens, oks):
+            it[f"{tag}_ok"] = bool(ok)
+            it[f"{tag}_plan"] = g if ok else None
+        print(f"{tag}: acc={sum(oks)/len(items):.3f} "
+              f"({time.time()-t0:.0f}s)", flush=True)
+
+    for it in items:
+        it["label"] = ("easy" if it["direct_ok"]
+                       else "hard" if it["cot_ok"] else "drop")
+        it["sol_plan"] = (it["direct_plan"] if it["direct_ok"]
+                          else it["cot_plan"])
+    for split in ("train", "test"):
+        sub = [it for it in items if it["split"] == split]
+        print(f"{split}: easy={sum(i['label']=='easy' for i in sub)} "
+              f"hard={sum(i['label']=='hard' for i in sub)} "
+              f"drop={sum(i['label']=='drop' for i in sub)}", flush=True)
+    json.dump(items, open(OUT / "bw_data.json", "w"), indent=1)
+    print("wrote", OUT / "bw_data.json")
+
+
+if __name__ == "__main__":
+    main()