| """Train the published FlowMo configurations without overwriting supplied weights.""" |
| from pathlib import Path |
| import argparse |
| import json |
| import os |
| import subprocess |
| import sys |
| import time |
|
|
| CORE = ['flowmo', 'history32', 'history8', 'flowmo_additive', 'flowmo_film'] |
| INITIAL = ['flowmo', 'leworldmodel', 'planet', 'tdmpc2'] |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument('--package', type=Path, default=Path(__file__).resolve().parent) |
| parser.add_argument('--output', type=Path, required=True) |
| parser.add_argument('--groups', nargs='+', choices=['initial', 'core', 'ood', 'scaling', 'all'], default=['core']) |
| parser.add_argument('--methods', nargs='+', choices=sorted(set(CORE + INITIAL))) |
| parser.add_argument('--seeds', nargs='+', type=int, default=list(range(1, 9))) |
| parser.add_argument('--device', default='cuda') |
| parser.add_argument('--smoke', action='store_true', help='One optimizer step; never formal paper statistics') |
| parser.add_argument('--dry-run', action='store_true') |
| args = parser.parse_args() |
| root, out = args.package.resolve(), args.output.resolve() |
| if out == root or out in root.parents or any(out == root / d or root / d in out.parents for d in ['data', 'experiments', 'paper', 'assets', 'provenance']): |
| raise ValueError('Choose a separate training output directory') |
| if any(s not in range(1, 9) for s in args.seeds): |
| raise ValueError('Formal seeds are 1 through 8') |
| cfg = json.loads((root / 'experiments/shared/config/rebuttal_v1.json').read_text()) |
| groups = ['initial', 'core', 'ood', 'scaling'] if 'all' in args.groups else args.groups |
| jobs = [] |
| for group in groups: |
| methods = INITIAL if group == 'initial' else (['flowmo'] if group == 'scaling' else CORE) |
| if args.methods: |
| methods = [m for m in methods if m in args.methods] |
| for method in methods: |
| for seed in ([19] if group == 'initial' else args.seeds): |
| for count in ([300, 600, 1200] if group == 'scaling' else [2400]): |
| checkpoint = ('paper.pt' if group == 'initial' else |
| f'rebuttal_scale_e{count:04d}_seed_{seed:04d}.pt' if group == 'scaling' else |
| f'rebuttal_{group}_seed_{seed:04d}.pt') |
| train_windows = cfg['training']['train_windows'] * count // 2400 |
| job = dict(group=group, method=method, seed=seed, episodes=count, checkpoint=checkpoint) |
| cmd = [sys.executable, '-m', 'experiments.train_image_world_models', '--methods', method, |
| '--train-source', str(root / cfg['data']['train_source']), '--test-source', str(root / cfg['data']['test_source']), |
| '--train-episodes', str(64 if args.smoke else count), '--test-episodes', str(64 if args.smoke else 480), |
| '--history-len', '32', '--horizon', '60', '--train-windows', str(512 if args.smoke else train_windows), |
| '--test-windows', str(64 if args.smoke else 24576), '--batch-size', '256', |
| '--microbatch-size', '64', '--eval-batch-size', '64', |
| '--steps', '1' if args.smoke else '20000', '--lr', '0.0003', '--log-every', '1' if args.smoke else '200', |
| '--seed', str(seed), '--seed-mode', 'method_offset' if group == 'initial' else 'paired', |
| '--device', args.device, '--num-workers', '4', '--image-size', '160', '--visual-scale', '2.5', |
| '--render-mode', 'device', '--precision', 'bf16', '--target-mode', 'absolute_normalized', |
| '--heading-weight', '2.0', '--current-pose-weight', '1.0', '--motion-weight', '0.5', |
| '--checkpoint-name', checkpoint, '--checkpoint-interval', '0'] |
| if group == 'ood': |
| families = [f for f in cfg['flow_families'] if f not in cfg['heldout_flow_families']] |
| cmd += ['--train-flow-families', *families, '--test-flow-families', *cfg['heldout_flow_families']] |
| if group == 'initial': |
| for flag in ['--microbatch-size', '--eval-batch-size', '--seed-mode']: |
| index = cmd.index(flag) |
| del cmd[index:index + 2] |
| job['source_tree'] = 'protocols/initial/src' if group == 'initial' else '.' |
| jobs.append(dict(job, command=cmd)) |
| if not jobs: |
| raise ValueError('No configurations selected') |
| if args.dry_run: |
| print(json.dumps({'training_jobs': len(jobs), 'smoke': args.smoke, 'jobs': jobs}, indent=2)) |
| return |
| out.mkdir(parents=True, exist_ok=True) |
| env = os.environ.copy() |
| env['PYTHONPATH'] = str(root) + os.pathsep + env.get('PYTHONPATH', '') |
| env['OMP_NUM_THREADS'] = '1' |
| env['OPENBLAS_NUM_THREADS'] = '1' |
| env['PYTHONPYCACHEPREFIX'] = str(out / 'pycache') |
| (out / 'training_plan.json').write_text(json.dumps({'smoke': args.smoke, 'jobs': jobs}, indent=2)) |
| records = [] |
| for job in jobs: |
| env['PYTHONPATH'] = str(root / job['source_tree']) + os.pathsep + str(root) |
| endpoint = out / 'experiments' / job['method'] / 'checkpoint' / job['checkpoint'] |
| if endpoint.exists(): |
| raise FileExistsError(f'Checkpoint already exists: {endpoint}; choose another run directory') |
| started = time.monotonic() |
| log = out / f'{job["group"]}_{job["method"]}_{job["seed"]}_{job["episodes"]}.log' |
| with log.open('w') as f: |
| subprocess.run(job['command'], cwd=out, env=env, stdout=f, stderr=subprocess.STDOUT, check=True) |
| assert endpoint.is_file() |
| records.append({k: v for k, v in job.items() if k != 'command'} | {'elapsed_seconds': time.monotonic() - started}) |
| (out / 'training_receipt.json').write_text(json.dumps({'smoke': args.smoke, 'completed': len(records), 'planned': len(jobs), 'jobs': records}, indent=2)) |
| print(json.dumps({'completed': len(records), 'planned': len(jobs), 'configuration': records[-1]}), flush=True) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
