Optional run-directory outputs: metrics.jsonl, periodic / 5M-milestone / best checkpoints, resume, and a final.json summary. Inert when run_dir is None, so the training script still runs standalone.
| 65 | |
| 66 | |
| 67 | class RunLogger: |
| 68 | """Optional run-directory outputs: metrics.jsonl, periodic / 5M-milestone / |
| 69 | best checkpoints, resume, and a final.json summary. Inert when run_dir is |
| 70 | None, so the training script still runs standalone.""" |
| 71 | |
| 72 | def __init__(self, run_dir, ckpt_every): |
| 73 | self.dir = run_dir |
| 74 | self.ckpt_dir = os.path.join(run_dir, "ckpt") if run_dir else None |
| 75 | self.ckpt_every = ckpt_every |
| 76 | if self.ckpt_dir: |
| 77 | os.makedirs(self.ckpt_dir, exist_ok=True) |
| 78 | self.f = open(os.path.join(run_dir, "metrics.jsonl"), "a", buffering=1) if run_dir else None |
| 79 | self.t0, self.last_frames = time.time(), 0 |
| 80 | self.ckpt_last, self.ms_last, self.best = 0, 0, float("-inf") |
| 81 | |
| 82 | def log(self, frames, scalars): |
| 83 | """Append one structured row (frames + sps + caller's scalars) to metrics.jsonl.""" |
| 84 | if not self.f: |
| 85 | return |
| 86 | now = time.time() |
| 87 | sps = (frames - self.last_frames) / max(now - self.t0, 1e-9) |
| 88 | self.f.write(json.dumps({"ts": round(now, 1), "frames": frames, "sps": round(sps, 1), **scalars}) + "\n") |
| 89 | self.t0, self.last_frames = now, frames |
| 90 | |
| 91 | def resolve_resume(self, resume_arg): |
| 92 | """'auto' -> run_dir/ckpt/latest.pt, else a path, else None.""" |
| 93 | if resume_arg == "auto" and self.ckpt_dir: |
| 94 | cand = os.path.join(self.ckpt_dir, "latest.pt") |
| 95 | return cand if os.path.exists(cand) else None |
| 96 | if resume_arg and resume_arg != "auto": |
| 97 | return resume_arg if os.path.exists(resume_arg) else None |
| 98 | return None |
| 99 | |
| 100 | def checkpoint(self, frames, state_fn, gate=None): |
| 101 | """Periodic 'latest', 5M-step milestone, and best-gate checkpoints. |
| 102 | state_fn() builds the dict only when a save actually happens.""" |
| 103 | if not self.ckpt_dir or not self.ckpt_every: |
| 104 | return |
| 105 | if frames - self.ckpt_last >= self.ckpt_every: |
| 106 | _atomic_save(state_fn(), os.path.join(self.ckpt_dir, "latest.pt")) |
| 107 | self.ckpt_last = frames |
| 108 | if frames - self.ms_last >= 5_000_000: |
| 109 | _atomic_save(state_fn(), os.path.join(self.ckpt_dir, f"step_{frames // 1_000_000}M.pt")) |
| 110 | self.ms_last = frames |
| 111 | if gate is not None and gate > self.best: |
| 112 | self.best = gate |
| 113 | _atomic_save(state_fn(), os.path.join(self.ckpt_dir, "best.pt")) |
| 114 | |
| 115 | def finalize(self, frames, game_returns, state_fn, k=100): |
| 116 | """Final 'latest' checkpoint + a final.json result summary.""" |
| 117 | if self.ckpt_dir: |
| 118 | _atomic_save(state_fn(), os.path.join(self.ckpt_dir, "latest.pt")) |
| 119 | if self.dir: |
| 120 | tail = [float(x) for x in game_returns[-k:]] |
| 121 | with open(os.path.join(self.dir, "final.json"), "w") as fh: |
| 122 | json.dump({"frames_total": frames, "frames_unit": "agent_steps", |
| 123 | "gate_metric": "game_return_mean_lastK", "K": k, |
| 124 | "value_mean": statistics.fmean(tail) if tail else float("nan"), |