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

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

scripts/eval_bw.py

@@ -11,7 +11,7 @@ import torch
 
 from bw_common import verify_plan
 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
@@ -52,13 +52,15 @@ def main():
     ap.add_argument("--adapter", default=None)
     ap.add_argument("--tag", default="bw")
     ap.add_argument("--ks", default="0,2,4")
+    ap.add_argument("--adaptive", action="store_true")
     args = ap.parse_args()
     ks = [int(x) for x in args.ks.split(",")]
 
     model, tok = load_model(dtype=torch.bfloat16)
     tok.padding_side = "left"
     looper = BandLooper(model)
-    adapter = MergeAdapter(
+    cls = AdaptiveMergeAdapter if args.adaptive else MergeAdapter
+    adapter = cls(
         d=model.config.get_text_config().hidden_size).cuda()
     if args.adapter:
         adapter.load_state_dict(torch.load(args.adapter, map_location="cuda"))