Improve RL training observability and checkpointing

pbc3_rl_a20.py

@@ -286,11 +286,13 @@ def train(args):
     initial = evaluate(model, reference, source, val_episodes, bases, args)
     print("initial", json.dumps(initial), flush=True)
     best = initial["aggregate_score"]
+    best_epoch = 0
     export(model, source, args.out)
     for epoch in range(start_epoch + 1, start_epoch + args.epochs + 1):
         random.shuffle(train_episodes)
         rewards = []
         losses = []
+        train_rows = []
         for start in range(0, len(train_episodes), args.batch):
             optimizer.zero_grad()
             episode_losses = []
@@ -300,6 +302,7 @@ def train(args):
                 )
                 r = reward(bases[(path, preset)], (mse, bpp, work), args)
                 rewards.append(r)
+                train_rows.append((mse, bpp, work))
                 if not logps:
                     continue
                 logps = torch.stack(logps)
@@ -321,6 +324,9 @@ def train(args):
         row = {
             "epoch": epoch,
             "train_reward": float(np.mean(rewards)),
+            "train_mse": float(np.mean([r[0] for r in train_rows])),
+            "train_bpp": float(np.mean([r[1] for r in train_rows])),
+            "train_work": float(np.mean([r[2] for r in train_rows])),
             "loss": float(np.mean(losses)),
             "validation": validation,
         }
@@ -328,7 +334,11 @@ def train(args):
         print(json.dumps(row), flush=True)
         if validation["aggregate_score"] > best:
             best = validation["aggregate_score"]
+            best_epoch = epoch
             export(model, source, args.out)
+        if args.export_every > 0 and (epoch % args.export_every == 0 or epoch == start_epoch + args.epochs):
+            periodic = Path(args.out).with_name(f"{Path(args.out).stem}.epoch{epoch}.npz")
+            export(model, source, periodic)
         torch.save(model.state_dict(), f"{args.out}.epoch{epoch}.pt")
         torch.save({
             "epoch": epoch,
@@ -339,7 +349,7 @@ def train(args):
             "torch_rng": torch.random.get_rng_state(),
             "history": history,
         }, f"{args.out}.resume.pt")
-    Path(f"{args.out}.json").write_text(json.dumps({"initial": initial, "best_score": best, "history": history}, indent=2))
+    Path(f"{args.out}.json").write_text(json.dumps({"initial": initial, "best_score": best, "best_epoch": best_epoch, "history": history}, indent=2))
 
 
 def main():
@@ -350,6 +360,7 @@ def main():
     parser.add_argument("--init", default="pbc3_students/patch_policy_f26_a20_h512_l2_e1200.npz")
     parser.add_argument("--out", default="pbc3_students/patch_policy_f26_a20_rl.npz")
     parser.add_argument("--resume", default="")
+    parser.add_argument("--export-every", type=int, default=500)
     parser.add_argument("--epochs", type=int, default=5)
     parser.add_argument("--batch", type=int, default=4)
     parser.add_argument("--limit", type=int, default=0)

static/index.html

@@ -54,6 +54,12 @@
   .tr-presets { display:flex; flex-wrap:wrap; gap:10px 18px; padding:10px 0; color:var(--text); }
   .tr-presets input { accent-color:var(--red); }
   .tr-log-tail { max-height:240px; overflow:auto; white-space:pre-wrap; background:rgba(255,255,255,.04); padding:12px; border-radius:10px; color:var(--muted); font:11px 'JetBrains Mono',monospace; }
+  .tr-chart-grid { display:grid; grid-template-columns:repeat(2,minmax(0,1fr)); gap:16px; }
+  .tr-chart-grid > div { min-width:0; min-height:290px; }
+  .tr-chart-grid .tr-chart-wide { grid-column:1/-1; }
+  .tr-sort { display:flex; align-items:center; gap:7px; color:var(--muted); font-size:12px; }
+  .tr-sort select { background:var(--surface); border:1px solid var(--border); border-radius:6px; color:var(--text); padding:5px; }
+  .tr-baseline-line { padding:9px 0 12px; color:var(--muted); font:12px 'JetBrains Mono',monospace; }
   .tr-table-wrap { overflow-x:auto; }
   .tr-table { width:100%; border-collapse:collapse; font:12px 'JetBrains Mono',monospace; }
   .tr-table th, .tr-table td { padding:9px 10px; border-bottom:1px solid rgba(255,255,255,.08); text-align:right; white-space:nowrap; }
@@ -69,7 +75,7 @@
   #hold-tr-c { --hc:#ff8a3d; } #hold-tr-b { --hc:#38bdf8; } #hold-tr-q { --hc:#a855f7; } #hold-tr-hq { --hc:#34d399; } #hold-tr-model { --hc:#ff3b41; }
   #hold-tr-c, #hold-tr-b, #hold-tr-q, #hold-tr-hq, #hold-tr-model { border-color:var(--hc); color:var(--hc); }
   #hold-tr-c.holding, #hold-tr-b.holding, #hold-tr-q.holding, #hold-tr-hq.holding, #hold-tr-model.holding { background:var(--hc); color:#fff; border-color:var(--hc); }
-  @media (max-width:760px){ .tr-fields{grid-template-columns:1fr;} }
+  @media (max-width:760px){ .tr-fields{grid-template-columns:1fr;} .tr-chart-grid{grid-template-columns:1fr;} .tr-chart-grid .tr-chart-wide{grid-column:auto;} .tr-sort{margin-top:8px;} }
 </style>
 <link rel="icon" href="data:image/svg+xml,<svg xmlns=%22http://www.w3.org/2000/svg%22 viewBox=%220 0 100 100%22><text y=%22.9em%22 font-size=%2290%22>👁️</text></svg>">
 </head>

static/train.js

@@ -67,17 +67,26 @@
           <label class="tr-f"><span>Temperature</span><input id="tr-temp" type="number" step="0.05" value="0.9"></label>
           <label class="tr-f"><span>Entropy weight</span><input id="tr-entropy" type="number" step="0.0005" value="0.003"></label>
           <label class="tr-f"><span>KL weight</span><input id="tr-kl" type="number" step="0.001" value="0.015"></label>
+          <label class="tr-f"><span>Export checkpoint every N epochs</span><input id="tr-export-every" type="number" min="0" step="100" value="500"></label>
         </div>
         <div class="tr-actions"><button id="tr-start" class="primary-btn">Start / resume</button><button id="tr-stop" class="hold-btn">Stop</button><button id="tr-refresh" class="hold-btn">Refresh</button><button id="tr-log" class="hold-btn">Download log</button><button id="tr-copy-results" class="hold-btn">Copy results</button><button id="tr-download-results" class="hold-btn">Download results</button><label class="tr-check tr-live"><input id="tr-pause" type="checkbox"> Pause live table/log updates</label></div>
       </div>
       <div class="params"><div class="params-head"><span class="field-label">Run status</span></div><div id="tr-status" class="tr-ck">Loading…</div><pre id="tr-log-tail" class="tr-log-tail"></pre></div>
-      <div class="params"><div class="params-head"><span class="field-label">Validation history</span></div><div id="tr-history"></div></div>`;
+      <div class="params"><div class="params-head"><span class="field-label">Learning graphics</span></div><div class="tr-chart-grid"><div id="tr-chart-reward"></div><div id="tr-chart-metrics"></div><div id="tr-chart-work"></div><div id="tr-chart-train"></div><div id="tr-chart-rd" class="tr-chart-wide"></div></div></div>
+      <div class="params"><div class="params-head"><span class="field-label">Validation history</span><label class="tr-sort">Sort by <select id="tr-sort"><option value="epoch">Epoch</option><option value="reward">Per-image reward</option><option value="aggregate_score">Aggregate score</option><option value="mse">Validation MSE</option><option value="bpp">Validation bpp</option><option value="work">Validation grid work</option><option value="train_reward">Training reward</option><option value="train_mse">Training MSE</option><option value="train_bpp">Training bpp</option></select><button id="tr-sort-dir" class="hold-btn sm">↓</button></label></div><div id="tr-history"></div></div>`;
     $("#tr-start").onclick = start;
     $("#tr-stop").onclick = async () => { await api("/api/train/stop", { method: "POST" }); refresh(); };
     $("#tr-refresh").onclick = refresh;
     $("#tr-log").onclick = () => window.open("/api/train/log", "_blank");
     $("#tr-copy-results").onclick = copyResults;
     $("#tr-download-results").onclick = downloadResults;
+    $("#tr-sort").onchange = () => renderHistory(latestStatus.history || []);
+    $("#tr-sort-dir").onclick = () => {
+      const button = $("#tr-sort-dir");
+      button.dataset.direction = button.dataset.direction === "1" ? "-1" : "1";
+      button.textContent = button.dataset.direction === "1" ? "↑" : "↓";
+      renderHistory(latestStatus.history || []);
+    };
     fillInitialModels();
   }
 
@@ -109,6 +118,7 @@
       init: $("#tr-init").value, output: $("#tr-output").value,
       rate_weight: Number($("#tr-rate").value), speed_weight: Number($("#tr-speed").value),
       temperature: Number($("#tr-temp").value), entropy_weight: Number($("#tr-entropy").value), kl_weight: Number($("#tr-kl").value),
+      export_every: Number($("#tr-export-every").value),
     };
     try { await api("/api/train/start", { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify(body) }); refresh(); } catch (error) { $("#tr-status").textContent = error.message; }
   }
@@ -116,23 +126,32 @@
   function renderHistory(history) {
     const root = $("#tr-history");
     if (!root) return;
-    const rows = history.filter((row) => row.epoch || row.validation).map((row) => {
+    const key = $("#tr-sort")?.value || "epoch";
+    const direction = Number($("#tr-sort-dir")?.dataset.direction || (key === "epoch" ? "1" : "-1"));
+    const sorted = history.filter((row) => row.epoch || row.validation).slice().sort((a, b) => {
+      const value = (row) => key === "epoch" ? (row.epoch ?? -1) : (key.startsWith("train_") ? row : (row.validation || row))[key];
+      const av = value(a);
+      const bv = value(b);
+      return (Number(av ?? -Infinity) - Number(bv ?? -Infinity)) * direction;
+    });
+    const rows = sorted.map((row) => {
       const v = row.validation || row;
-      return `<tr><td>${row.epoch ? `Epoch ${row.epoch}` : "Initial"}</td><td>${fmt(row.reward ?? v.reward, 5)}</td><td>${fmt(v.aggregate_score, 5)}</td><td>${fmt(v.mse, 3)}</td><td>${fmt(v.bpp, 6)}</td><td>${fmt(v.work, 0)}</td></tr>`;
+      return `<tr><td>${row.epoch ? `Epoch ${row.epoch}` : "Initial"}</td><td>${fmt(row.train_reward, 5)}</td><td>${fmt(row.reward ?? v.reward, 5)}</td><td>${fmt(v.aggregate_score, 5)}</td><td>${fmt(row.train_mse, 3)}</td><td>${fmt(v.mse, 3)}</td><td>${fmt(row.train_bpp, 6)}</td><td>${fmt(v.bpp, 6)}</td><td>${fmt(row.train_work, 0)}</td><td>${fmt(v.work, 0)}</td></tr>`;
     });
-    root.innerHTML = rows.length ? `<div class="tr-table-wrap"><table class="tr-table"><thead><tr><th>Checkpoint</th><th>Per-image reward</th><th>Aggregate score</th><th>Val MSE</th><th>Val bpp</th><th>Grid work</th></tr></thead><tbody>${rows.join("")}</tbody></table></div>` : `<p class="card-date">No validation rows yet. The first baseline pass can take a few minutes.</p>`;
+    const baseline = (history.find((r) => r.validation)?.validation || {});
+    root.innerHTML = rows.length ? `<div class="tr-baseline-line">Production baseline: MSE ${fmt(baseline.baseline_mse, 3)} · bpp ${fmt(baseline.baseline_bpp, 6)} · grid-work baseline is the starting policy, not PBC3.0.</div><div class="tr-table-wrap"><table class="tr-table"><thead><tr><th>Checkpoint</th><th>Train reward</th><th>Val reward</th><th>Aggregate score</th><th>Train MSE</th><th>Val MSE</th><th>Train bpp</th><th>Val bpp</th><th>Train work</th><th>Val work</th></tr></thead><tbody>${rows.join("")}</tbody></table></div>` : `<p class="card-date">No validation rows yet. The first baseline pass can take a few minutes.</p>`;
   }
 
   function resultText() {
     const history = latestStatus.history || [];
     return history.map((row) => {
       const v = row.validation || row;
-      return [row.epoch ? `Epoch ${row.epoch}` : "Initial", fmt(row.reward ?? v.reward, 5), fmt(v.aggregate_score, 5), fmt(v.mse, 3), fmt(v.bpp, 6), fmt(v.work, 0)].join("\t");
+      return [row.epoch ? `Epoch ${row.epoch}` : "Initial", fmt(row.train_reward, 5), fmt(row.reward ?? v.reward, 5), fmt(v.aggregate_score, 5), fmt(row.train_mse, 3), fmt(v.mse, 3), fmt(row.train_bpp, 6), fmt(v.bpp, 6), fmt(row.train_work, 0), fmt(v.work, 0)].join("\t");
     }).join("\n");
   }
 
   async function copyResults() {
-    const text = ["Checkpoint\tPer-image reward\tAggregate score\tVal MSE\tVal bpp\tGrid work", resultText()].filter(Boolean).join("\n");
+    const text = ["Checkpoint\tTrain reward\tVal reward\tAggregate score\tTrain MSE\tVal MSE\tTrain bpp\tVal bpp\tTrain work\tVal work", resultText()].filter(Boolean).join("\n");
     try { await navigator.clipboard.writeText(text); $("#tr-pause").checked = true; } catch { window.prompt("Copy training results", text); }
   }
 
@@ -152,11 +171,49 @@
     if (!$("#tr-pause")?.checked) {
       $("#tr-log-tail").textContent = data.log_tail || "";
       renderHistory(data.history || []);
+      drawCharts(data.history || []);
     }
     if (data.running && !poll) poll = setInterval(refresh, 3000);
     if (!data.running && poll) { clearInterval(poll); poll = null; loadCheckpoints(); }
   }
 
+  function chartLayout(title, yTitle, extra = {}) {
+    return Object.assign({ title, height: 290, margin: { l: 52, r: 48, t: 38, b: 42 }, paper_bgcolor: "rgba(0,0,0,0)", plot_bgcolor: "rgba(0,0,0,0)", font: { color: "#cfd3dc", size: 11 }, xaxis: { title: "Epoch", gridcolor: "rgba(255,255,255,.07)" }, yaxis: { title: yTitle, gridcolor: "rgba(255,255,255,.07)" }, legend: { orientation: "h" } }, extra);
+  }
+
+  function drawCharts(history) {
+    if (!window.Plotly) return;
+    const points = history.filter((row) => row.epoch || row.validation);
+    if (!points.length) return;
+    const x = points.map((row) => row.epoch ?? 0);
+    const values = points.map((row) => row.validation || row);
+    const color = points.map((_, i) => i);
+    Plotly.react($("#tr-chart-reward"), [
+      { x, y: points.map((row) => row.train_reward), name: "train reward", mode: "lines", line: { color: "#60a5fa" } },
+      { x, y: points.map((row) => row.reward ?? (row.validation || {}).reward), name: "validation reward", mode: "lines", line: { color: "#f59e0b" } },
+      { x, y: values.map((row) => row.aggregate_score), name: "aggregate score", mode: "lines", line: { color: "#22c55e" } },
+    ], chartLayout("Reward and aggregate score", "score"), { responsive: true, displayModeBar: false });
+    Plotly.react($("#tr-chart-metrics"), [
+      { x, y: points.map((row) => row.train_mse), name: "train MSE", mode: "lines", line: { color: "#93c5fd" } },
+      { x, y: values.map((row) => row.mse), name: "val MSE", mode: "lines", line: { color: "#ef4444" } },
+      { x, y: values.map((row) => row.baseline_mse), name: "PBC3.0 MSE", mode: "lines", line: { color: "#fca5a5", dash: "dot" } },
+    ], chartLayout("MSE", "MSE"), { responsive: true, displayModeBar: false });
+    Plotly.react($("#tr-chart-work"), [
+      { x, y: points.map((row) => row.train_work), name: "train grid work", mode: "lines", line: { color: "#a78bfa" } },
+      { x, y: values.map((row) => row.work), name: "val grid work", mode: "lines", line: { color: "#c084fc" } },
+    ], chartLayout("Grid work", "coefficients"), { responsive: true, displayModeBar: false });
+    Plotly.react($("#tr-chart-train"), [
+      { x, y: points.map((row) => row.train_bpp), name: "train bpp", mode: "lines", line: { color: "#93c5fd" } },
+      { x, y: values.map((row) => row.bpp), name: "val bpp", mode: "lines", line: { color: "#38bdf8" } },
+      { x, y: values.map((row) => row.baseline_bpp), name: "PBC3.0 bpp", mode: "lines", line: { color: "#7dd3fc", dash: "dot" } },
+    ], chartLayout("Bitrate", "bits / pixel"), { responsive: true, displayModeBar: false });
+    const baseline = values[0];
+    Plotly.react($("#tr-chart-rd"), [
+      { x: values.map((row) => row.bpp), y: values.map((row) => row.mse), text: x.map((epoch) => `Epoch ${epoch}`), mode: "markers", name: "validation RD", marker: { size: 9, color, colorscale: "YlOrRd", showscale: true, colorbar: { title: "newer" } } },
+      { x: [baseline.baseline_bpp], y: [baseline.baseline_mse], text: ["PBC3.0 baseline"], mode: "markers+text", name: "PBC3.0 baseline", textposition: "top center", marker: { size: 13, color: "#fff", symbol: "diamond", line: { color: "#ef4444", width: 2 } } },
+    ], chartLayout("Rate–distortion history", "MSE", { xaxis: { title: "bpp", gridcolor: "rgba(255,255,255,.07)" } }), { responsive: true, displayModeBar: false });
+  }
+
   function renderCheckpoints() {
     checkpointRoot.innerHTML = `<div class="params"><div class="params-head"><span class="field-label">Downloadable model checkpoints</span></div><p class="card-date">Download an exported <code>.npz</code> after reviewing its validation history, then upload it into the main project when you want to test or promote it.</p><input id="tr-ckpt-filter" class="tr-filter" placeholder="Filter checkpoints…"><div id="tr-checkpoint-list" class="tr-checkpoint-grid"></div></div>`;
     $("#tr-ckpt-filter").oninput = renderCheckpointList;

train_api.py

@@ -68,7 +68,8 @@ def status():
         if history_path and history_path.exists():
             try:
                 saved = json.loads(history_path.read_text(encoding="utf-8"))
-                history = saved.get("history", [])
+                initial = saved.get("initial")
+                history = ([{"validation": initial}] if initial else []) + saved.get("history", [])
             except Exception:
                 pass
         if not history:
@@ -107,6 +108,7 @@ def start(spec):
             sys.executable, "pbc3_rl_a20.py",
             "--presets", ",".join(selected),
             "--epochs", str(int(spec.get("epochs", 100))),
+            "--export-every", str(int(spec.get("export_every", 500))),
             "--batch", str(int(spec.get("batch", 4))),
             "--init", os.fspath(init),
             "--out", os.fspath(output),