File size: 3,463 Bytes
dc9f917 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | """Run the baselines and ablations listed in section 11 of the design doc.
Three of them change the objective, so they need their own training runs
(terminal-only, no-support). The rest only change how the trained controller
is executed, so they reuse the main checkpoint. Every sim evaluation uses the
same seed, so all rows share held-out initial-state/goal pairs.
"""
import argparse
import subprocess
import sys
from pathlib import Path
PY = sys.executable
ROOT = Path(__file__).resolve().parents[1]
# objective ablations: each needs its own controller
VARIANTS = {
'main': [], # alpha=0.05, lambda_support=0.01
'terminal_only': ['--alpha', '0'],
'no_support': ['--lambda-support', '0'],
}
def run(cmd):
print(f'\n$ {" ".join(str(c) for c in cmd)}', flush=True)
# one bad row should not discard the hours of runs queued behind it
code = subprocess.run(cmd, cwd=ROOT).returncode
if code != 0:
print(f'!! exited {code}, continuing', flush=True)
return code == 0
def parse_args():
p = argparse.ArgumentParser()
p.add_argument('--steps', type=int, default=20000)
p.add_argument('--batch-size', type=int, default=128)
p.add_argument('--lr', type=float, default=3e-4)
p.add_argument('--num-eval', type=int, default=50)
p.add_argument('--eval-budget', type=int, default=50)
p.add_argument('--out', default='data/runs')
p.add_argument('--skip-train', action='store_true')
p.add_argument('--skip-cem', action='store_true')
return p.parse_args()
def main():
args = parse_args()
out = Path(args.out)
def evaluate(extra, tag_dir='eval'):
run(
[PY, 'scripts/eval_controller.py',
'--num-eval', str(args.num_eval),
'--eval-budget', str(args.eval_budget),
'--out', str(out / tag_dir), *extra]
)
# Everything that reuses the main checkpoint runs first: those are the
# headline numbers and cost minutes, while each objective ablation costs
# its own multi-hour training run.
# --- baseline: LeWM + CEM ---------------------------------------------
if not args.skip_cem:
evaluate(['--planner', 'cem'])
main_ckpt = ['--controller', str(out / 'controller' / 'controller.pt')]
# --- refinement count: K=0 (one-pass) through K=5 ---------------------
for k in (0, 1, 2, 3, 5):
evaluate([*main_ckpt, '--refinements', str(k)])
# --- execute one block vs the whole plan ------------------------------
evaluate([*main_ckpt, '--receding-horizon', '5'])
# --- objective ablations need their own training runs -----------------
for name, extra in VARIANTS.items():
if name == 'main':
continue
ckpt_dir = out / f'abl_{name}'
if not args.skip_train:
run(
[PY, 'scripts/train_controller.py',
'--steps', str(args.steps),
'--batch-size', str(args.batch_size),
'--lr', str(args.lr),
'--log-every', '500', '--val-every', '5000',
'--out', str(ckpt_dir), *extra]
)
ckpt = ckpt_dir / 'controller.pt'
if ckpt.exists():
evaluate(['--controller', str(ckpt)])
print(f'\nall results appended to {out / "eval" / "results.jsonl"}')
if __name__ == '__main__':
main()
|