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>
2026-07-14 00:54:12 +02:00
commit ef9c08966c
42 changed files with 5068 additions and 0 deletions
@@ -0,0 +1,117 @@
+"""Blocksworld: generator, verifier, prompts.
+
+Pure planning domain — no code syntax, no arithmetic. Plans are symbolically
+verifiable by simulation, so STaR bucketing works. Instances are generated
+(contamination-free by construction) with difficulty = blocks + moves.
+"""
+
+import json
+import random
+import re
+
+MOVE_RE = re.compile(
+    r"move\s+([A-Z])\s+(?:onto|on top of|on)\s+(?:the\s+)?(table|[A-Z])",
+    re.IGNORECASE)
+
+
+def gen_instance(n_blocks, rng):
+    blocks = [chr(65 + i) for i in range(n_blocks)]
+
+    def random_stacks():
+        bs = blocks[:]
+        rng.shuffle(bs)
+        stacks, i = [], 0
+        while i < len(bs):
+            take = rng.randint(1, len(bs) - i)
+            stacks.append(bs[i : i + take])
+            i += take
+        return stacks
+
+    init = random_stacks()
+    goal = random_stacks()
+    while goal == init:
+        goal = random_stacks()
+    return {"blocks": blocks, "init": init, "goal": goal}
+
+
+def fmt_state(stacks):
+    out = []
+    for st in stacks:
+        if len(st) == 1:
+            out.append(f"{st[0]} is on the table")
+        else:
+            out.append(f"{st[0]} is on the table with "
+                       + " on top, then ".join(
+                           [f"{b}" for b in st[1:]]) + " on top")
+    # clearer explicit form
+    lines = []
+    for st in stacks:
+        lines.append(f"stack: {' -> '.join(st)} (bottom -> top)")
+    return "; ".join(lines)
+
+
+def question(inst):
+    return (
+        "You are stacking blocks. Only the TOP block of a stack can be "
+        "moved, one block at a time.\n"
+        f"Blocks: {', '.join(inst['blocks'])}\n"
+        f"Initial state: {fmt_state(inst['init'])}\n"
+        f"Goal state: {fmt_state(inst['goal'])}\n"
+        "Give a plan as a numbered list of moves, each exactly of the form "
+        "'move X onto Y' or 'move X onto the table'.")
+
+
+def verify_plan(inst, text, max_moves=40):
+    stacks = [st[:] for st in inst["init"]]
+
+    def top_of(b):
+        for st in stacks:
+            if st and st[-1] == b:
+                return st
+        return None
+
+    moves = MOVE_RE.findall(text)
+    if not moves or len(moves) > max_moves:
+        return False
+    for b, tgt in moves:
+        b = b.upper()
+        tgt = tgt if tgt.lower() == "table" else tgt.upper()
+        src = top_of(b)
+        if src is None:
+            return False  # b not clear (or nonexistent)
+        if tgt == "table" or tgt.lower() == "table":
+            src.pop()
+            stacks.append([b])
+        else:
+            dst = top_of(tgt)
+            if dst is None or b == tgt:
+                return False
+            src.pop()
+            dst.append(b)
+        stacks = [st for st in stacks if st]
+    norm = sorted(tuple(st) for st in stacks)
+    return norm == sorted(tuple(st) for st in inst["goal"])
+
+
+def make_dataset(n_train=400, n_test=200, seed=0):
+    rng = random.Random(seed)
+    items = []
+    for split, n in (("train", n_train), ("test", n_test)):
+        for i in range(n):
+            nb = rng.choice([3, 3, 4, 4, 5])
+            inst = gen_instance(nb, rng)
+            items.append({"split": split, "task_id": f"bw_{split}_{i}",
+                          "n_blocks": nb, **inst,
+                          "question": question(inst)})
+    return items
+
+
+if __name__ == "__main__":
+    ds = make_dataset()
+    print(json.dumps(ds[0], indent=1))
+    # verifier self-test: identity plan on trivial instance
+    inst = {"blocks": ["A", "B"], "init": [["A"], ["B"]],
+            "goal": [["A", "B"]]}
+    assert verify_plan(inst, "1. move B onto A")
+    assert not verify_plan(inst, "1. move A onto A")
+    print("verifier self-test ok")