| 69 | |
| 70 | |
| 71 | def parse_args(): |
| 72 | p = argparse.ArgumentParser() |
| 73 | p.add_argument("--env", choices=list(ENV_IDS), default="breakout", |
| 74 | help="which Atari game to train on") |
| 75 | p.add_argument("--render", action="store_true", |
| 76 | help="open a window during training (much slower)") |
| 77 | p.add_argument("--test", action="store_true", |
| 78 | help="load the saved checkpoint and just play (no learning)") |
| 79 | p.add_argument("--device", choices=["auto", "cpu", "cuda", "mps"], default="auto", |
| 80 | help="override the auto-selected torch device") |
| 81 | p.add_argument("--wandb", action="store_true", |
| 82 | help="log metrics to Weights & Biases") |
| 83 | return p.parse_args() |
| 84 | |
| 85 | |
| 86 | def make_env(args): |