| 139 | |
| 140 | |
| 141 | def parse_args(): |
| 142 | p = argparse.ArgumentParser() |
| 143 | p.add_argument("--env", choices=list(ENV_IDS), default="montezuma", |
| 144 | help="which hard-exploration Atari game to train on") |
| 145 | p.add_argument("--render", action="store_true", |
| 146 | help="open a window during training (single-env --test only)") |
| 147 | p.add_argument("--test", action="store_true", |
| 148 | help="load the saved checkpoint and just play (no learning)") |
| 149 | p.add_argument("--device", choices=["auto", "cpu", "cuda", "mps"], default="auto", |
| 150 | help="override the auto-selected torch device") |
| 151 | p.add_argument("--wandb", action="store_true", |
| 152 | help="log metrics to Weights & Biases") |
| 153 | # --- reproducibility / run-management flags (all optional; omit them and the script runs as before) --- |
| 154 | p.add_argument("--seed", type=int, default=None, |
| 155 | help="reproducibility seed (np/torch/envpool)") |
| 156 | p.add_argument("--total-frames", type=int, default=None, |
| 157 | help="override the in-file TOTAL_FRAMES budget (agent steps)") |
| 158 | p.add_argument("--n-envs", type=int, default=None, |
| 159 | help="override the in-file N_ENVS (e.g. smaller for a smoke run)") |
| 160 | p.add_argument("--run-dir", type=str, default=None, |
| 161 | help="run directory: write metrics.jsonl / ckpt / final.json here") |
| 162 | p.add_argument("--ckpt-every", type=int, default=None, |
| 163 | help="periodic checkpoint interval in agent steps (resume-safe)") |
| 164 | p.add_argument("--resume", type=str, default=None, |
| 165 | help="'auto' (run-dir/ckpt/latest.pt) or a checkpoint path") |
| 166 | return p.parse_args() |
| 167 | |
| 168 | |
| 169 | def make_env(args): |