| """Train or replay the roller landing on Pollen's Got a Talent simulation.""" |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import learn as landing |
| from arena.cast import cast_of |
| from arena.hub_fetch import check_manifest, rig_of |
| from arena.policy import load_policy |
| from arena.runner import run_stage |
| from arena.timeline import load_troupe |
| from export_actor import export |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| commands = parser.add_subparsers(dest="command", required=True) |
| training = commands.add_parser("train") |
| training.add_argument("--output", type=Path, required=True) |
| training.add_argument( |
| "--start-from", |
| type=Path, |
| help="checkpoint.json with the 32 learned coefficients", |
| ) |
| training.add_argument("--generations", type=int, default=25) |
| training.add_argument("--population", type=int, default=40) |
| training.add_argument( |
| "--batch", |
| type=int, |
| default=8, |
| help="simulated starts per candidate per generation", |
| ) |
| training.add_argument("--workers", type=int, default=12) |
| training.add_argument("--random-seed", type=int, default=11062026) |
| training.add_argument("--speed", type=float, default=0.65) |
| training.add_argument("--roll-at", type=float, default=1.65) |
| build = commands.add_parser("export") |
| build.add_argument("checkpoint", type=Path) |
| build.add_argument("output", type=Path) |
| replay = commands.add_parser("replay") |
| replay.add_argument("model_folder", type=Path) |
| replay.add_argument("--seeds", nargs="+", type=int, default=[1]) |
| args = parser.parse_args() |
| if args.command == "export": |
| ck = json.loads(args.checkpoint.read_text()) |
| export(ck["theta"], args.output) |
| return |
| if args.command == "train": |
| if ( |
| args.population % 2 |
| or args.population < 10 |
| or min(args.generations, args.batch, args.workers) < 1 |
| or args.batch > 520 |
| ): |
| parser.error( |
| "Use at least one generation, worker and seed (up to 520) and an even population of at least 10." |
| ) |
| if not 0 <= args.speed <= 2 or not 0.7 <= args.roll_at <= 12: |
| parser.error( |
| "Use a speed from 0 to 2 m/s and a roll time from 0.7 to 12 seconds." |
| ) |
| if args.output.exists() and any(args.output.iterdir()): |
| parser.error("Choose a new or empty training directory.") |
| args.output.mkdir(parents=True, exist_ok=True) |
| warm = ( |
| json.loads(args.start_from.read_text())["theta"] |
| if args.start_from |
| else None |
| ) |
| landing.configure(args.speed, args.roll_at) |
| landing.save(args.output / "timeline.json", landing.timeline_dict()) |
| landing.train( |
| args.generations, |
| args.population, |
| args.batch, |
| warm, |
| args.output, |
| random_seed=args.random_seed, |
| workers=args.workers, |
| ) |
| candidates = json.loads((args.output / "checkpoints.json").read_text()) |
| landing.save(args.output / "checkpoint.json", candidates[0]) |
| return |
| folder = args.model_folder |
| event = landing.EVENT |
| manifest = json.loads((folder / "manifest.json").read_text()) |
| check_manifest(manifest, event) |
| policy = load_policy(folder / "policy.onnx") |
| troupe = load_troupe(folder / "timeline.json", event) |
| cast = cast_of(event, troupe, policy, rig_of(manifest), {}) |
| schedule = troupe.ducks[0].timeline |
| landing.configure(schedule.commands[1][0], schedule.times[2]) |
| for seed in args.seeds: |
| result = run_stage(event, cast.performers, seed) |
| row = {"seed": seed, "status": result.status} |
| if result.status == "finish": |
| row.update(landing.metrics(result.trace[0])) |
| print(json.dumps(row)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
