warm-start flag (stacking arm), adaptive flags for Blocksworld

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

scripts/train_merge_bw.py

@@ -18,7 +18,7 @@ import torch
 import torch.nn.functional as F
 
 from bw_prep import DIRECT_SUFFIX, chat
-from loop_common import BandLooper, MergeAdapter
+from loop_common import AdaptiveMergeAdapter, BandLooper, MergeAdapter
 
 sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
 from jlens.core import load_model  # noqa: E402
@@ -33,6 +33,7 @@ K_BUCKETS = [(1, ("easy",)), (2, ("easy", "hard")), (4, ("hard",))]
 
 ap = argparse.ArgumentParser()
 ap.add_argument("--seed", type=int, default=0)
+ap.add_argument("--adaptive", action="store_true")
 ARGS = ap.parse_args()
 MOVELIST_RE = re.compile(r"(?:^|\n)\s*1\s*[.)]", re.M)
 
@@ -81,7 +82,8 @@ def main():
     for p in model.parameters():
         p.requires_grad_(False)
     looper = BandLooper(model)
-    adapter = MergeAdapter(
+    cls = AdaptiveMergeAdapter if ARGS.adaptive else MergeAdapter
+    adapter = cls(
         d=model.config.get_text_config().hidden_size).cuda()
     opt = torch.optim.AdamW(adapter.parameters(), lr=LR, weight_decay=0.01)
 
@@ -127,7 +129,7 @@ def main():
                   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")
+                       OUT / f"adapter_bw{"_ad" if ARGS.adaptive else ""}_e{step+1}.pt")
             json.dump(log, open(OUT / "train_bw_log.json", "w"), indent=1)
     print("done", flush=True)