Add RL Training Lab

learned_filler.py

@@ -112,9 +112,10 @@ def _silu(x):
 
 
 def _mlp(x, weights):
-    """## Runs the shared two-layer SiLU trunk"""
-    W0, b0, W2, b2 = weights
-    return _silu(_silu(x @ W0.T + b0) @ W2.T + b2)
+    """## Runs the exported SiLU trunk"""
+    for W, b in weights:
+        x = _silu(x @ W.T + b)
+    return x
 
 
 class LearnedFiller:
@@ -125,8 +126,10 @@ class LearnedFiller:
         self.Wa, self.ba = npz["Wa"], npz["ba"]
         self.Wv, self.bv = npz["Wv"], npz["bv"]
         if self.hidden > 0:
-            self.W0, self.b0 = npz["W0"], npz["b0"]
-            self.W2, self.b2 = npz["W2"], npz["b2"]
+            if "layer_count" in npz:
+                self.layers = [(npz[f"W{i}"], npz[f"b{i}"]) for i in range(int(npz["layer_count"]))]
+            else:
+                self.layers = [(npz["W0"], npz["b0"]), (npz["W2"], npz["b2"])]
         self.feat_mean, self.feat_std = npz["feat_mean"], npz["feat_std"]
         self.lmin, self.lmax = float(npz["lmin"]), float(npz["lmax"])
         self.cheap = bool(npz["cheap"])
@@ -146,7 +149,7 @@ class LearnedFiller:
     def _forward(self, x, trace=None) -> tuple[np.ndarray, np.ndarray]:
         """## Runs the legacy action and value heads"""
         with timed(trace, "action.filler.policy.trunk_forward"):
-            h = _mlp(x, (self.W0, self.b0, self.W2, self.b2)) if self.hidden > 0 else x
+            h = _mlp(x, self.layers) if self.hidden > 0 else x
         with timed(trace, "action.filler.policy.action_head_forward"):
             action_logits = h @ self.Wa.T + self.ba
         with timed(trace, "action.filler.policy.value_head_forward"):
@@ -237,7 +240,7 @@ class FactoredFiller:
     def _forward(self, x, trace=None) -> list[np.ndarray]:
         """## Runs the factored policy heads"""
         with timed(trace, "action.filler.policy.trunk_forward"):
-            h = _mlp(x, (self.W0, self.b0, self.W2, self.b2))
+            h = _mlp(x, [(self.W0, self.b0), (self.W2, self.b2)])
         names = ["cell_size", "bitcount", "mask_size", "stop"]
         outputs = []
         for name, (W, b) in zip(names, self.heads):

pbc3_rl_a20.py

@@ -0,0 +1,359 @@
+import argparse
+import copy
+import json
+import math
+import random
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image, ImageOps
+
+import pbc3_ops as ops
+from PBC3 import PBC3, preload_numba
+from learned_filler import build_action, propose_boxes
+from pbc3_features import extract_cheap, extract_features
+from pbc3_heads import DownsampleInitHead
+from pbc3_types import BitWriter, PBC3Config
+
+PRESETS = {
+    "compression": PBC3Config.compression,
+    "balanced": PBC3Config.balanced,
+    "quality": PBC3Config.quality,
+    "high_quality": PBC3Config.high_quality,
+}
+IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
+
+
+def image_paths(folder):
+    return sorted(path for path in Path(folder).glob("*") if path.suffix.lower() in IMAGE_EXTENSIONS)
+
+
+def init_canvas(prep, config):
+    target = prep["target"]
+    h, w, channels = target.shape
+    base_values = [int(round(float(np.mean(prep["arr"][:, :, c])))) for c in range(channels)]
+    canvas = np.empty((h, w, channels), dtype=np.int32)
+    for c, value in enumerate(base_values):
+        canvas[:, :, c] = value
+    channel_bits = max(1, math.ceil(math.log2(channels)))
+    patches = []
+    head = DownsampleInitHead()
+    for c in range(channels):
+        patch, _, delta, _, _ = head.select(c, target, canvas, w, h, config, channel_bits)
+        ops.apply_delta(canvas[:, :, c], 0, 0, w, h, delta)
+        patches.append(patch)
+    return canvas, patches
+
+
+class Policy(torch.nn.Module):
+    def __init__(self, in_dim, hidden, layers, actions):
+        super().__init__()
+        self.trunk = torch.nn.ModuleList(
+            [torch.nn.Linear(in_dim, hidden)]
+            + [torch.nn.Linear(hidden, hidden) for _ in range(layers - 1)]
+        )
+        self.action = torch.nn.Linear(hidden, actions)
+        self.value = torch.nn.Linear(hidden, 1)
+
+    def forward(self, x):
+        for layer in self.trunk:
+            x = F.silu(layer(x))
+        return self.action(x), self.value(x).squeeze(-1)
+
+
+def load_policy(path):
+    data = dict(np.load(path, allow_pickle=True))
+    layers = int(data.get("layer_count", 2))
+    model = Policy(int(data["in_dim"]), int(data["hidden"]), layers, int(data["num_actions"]))
+    state = model.state_dict()
+    if "layer_count" in data:
+        for i in range(layers):
+            state[f"trunk.{i}.weight"] = torch.from_numpy(data[f"W{i}"])
+            state[f"trunk.{i}.bias"] = torch.from_numpy(data[f"b{i}"])
+    else:
+        state["trunk.0.weight"] = torch.from_numpy(data["W0"])
+        state["trunk.0.bias"] = torch.from_numpy(data["b0"])
+        state["trunk.1.weight"] = torch.from_numpy(data["W2"])
+        state["trunk.1.bias"] = torch.from_numpy(data["b2"])
+    state["action.weight"] = torch.from_numpy(data["Wa"])
+    state["action.bias"] = torch.from_numpy(data["ba"])
+    state["value.weight"] = torch.from_numpy(data["Wv"])
+    state["value.bias"] = torch.from_numpy(data["bv"])
+    model.load_state_dict(state)
+    return model, data
+
+
+def export(model, source, path):
+    state = model.state_dict()
+    values = {
+        "hidden": np.int64(model.trunk[0].out_features),
+        "layer_count": np.int64(len(model.trunk)),
+        "in_dim": np.int64(model.trunk[0].in_features),
+        "num_actions": np.int64(model.action.out_features),
+        "feat_mean": source["feat_mean"],
+        "feat_std": source["feat_std"],
+        "lmin": source["lmin"],
+        "lmax": source["lmax"],
+        "cheap": source["cheap"],
+        "feature_names": source["feature_names"],
+        "Wa": state["action.weight"].detach().numpy(),
+        "ba": state["action.bias"].detach().numpy(),
+        "Wv": state["value.weight"].detach().numpy(),
+        "bv": state["value.bias"].detach().numpy(),
+    }
+    for i in range(len(model.trunk)):
+        values[f"W{i}"] = state[f"trunk.{i}.weight"].detach().numpy()
+        values[f"b{i}"] = state[f"trunk.{i}.bias"].detach().numpy()
+    np.savez_compressed(path, **values)
+
+
+def stream_result(prep, config, patches, canvas, base_values):
+    channel_bits = max(1, math.ceil(math.log2(prep["channels"])))
+    writer = BitWriter()
+    PBC3._write_header(
+        writer, prep["w"], prep["h"], prep["original_w"], prep["original_h"],
+        prep["downsampled"], prep["color_id"], prep["channels"], channel_bits,
+        config.positive_bias, prep["has_alpha"], len(patches), base_values,
+    )
+    for patch in patches:
+        PBC3._write_patch(writer, patch, channel_bits)
+    method, body = PBC3._entropy_pack(writer.finish(), config.use_lzma)
+    data = PBC3.MAGIC + bytes([PBC3.VERSION, method]) + body
+    image = PBC3._canvas_to_image(canvas, config.color_space, prep["has_alpha"])
+    if image.size != (prep["original_w"], prep["original_h"]):
+        image = image.resize(
+            (prep["original_w"], prep["original_h"]), PBC3.RESAMPLE_FILTER,
+            reducing_gap=PBC3.RESAMPLE_REDUCING_GAP,
+        )
+    mse = ops.final_mse(prep["orig_compare"], image)
+    bpp = len(data) * 8 / (prep["original_w"] * prep["original_h"])
+    return mse, bpp
+
+
+def rollout(model, reference, source, path, preset, sample, temperature):
+    config = PRESETS[preset](warmup_ratio=-1, learned_filler_enabled=False)
+    image = ImageOps.exif_transpose(Image.open(path)).convert("RGB")
+    prep = PBC3.prepare(image, config)
+    target = prep["target"]
+    canvas, patches = init_canvas(prep, config)
+    patches = list(patches)
+    h, w, channels = target.shape
+    channel_bits = max(1, math.ceil(math.log2(channels)))
+    base_values = [int(round(float(np.mean(prep["arr"][:, :, c])))) for c in range(channels)]
+    channel_scores = [PBC3._channel_sum_error(target, canvas, c) for c in range(channels)]
+    rng = ops.PBC3Rng(config.random_seed)
+    mean = source["feat_mean"].astype(np.float32)
+    std = source["feat_std"].astype(np.float32)
+    names = [str(x) for x in source["feature_names"]]
+    cheap = bool(source["cheap"])
+    logps, values, entropies, kls = [], [], [], []
+    work = 0.0
+    started = time.perf_counter()
+    for step in range(1, config.patch_count + 1):
+        channel = PBC3._choose_channel(channel_scores, step, channels, config.channel_cycle)
+        boxes = propose_boxes(target, canvas, config, rng, channel, step)[:1]
+        if not boxes:
+            break
+        box = boxes[0]
+        q = config.learned_filler_q
+        if cheap:
+            features = extract_cheap(names, target, canvas, box, step, q, w, h, config.patch_count, channels)
+        else:
+            features = extract_features(target, canvas, box, step, q, w, h, config.patch_count, channels)
+        x = torch.from_numpy(((features - mean) / std)[None].astype(np.float32))
+        logits, value = model(x)
+        with torch.no_grad():
+            ref_logits, _ = reference(x)
+        if sample:
+            distribution = torch.distributions.Categorical(logits=logits[0] / temperature)
+            action = distribution.sample()
+            logps.append(distribution.log_prob(action))
+            entropies.append(distribution.entropy())
+            kls.append(F.kl_div(F.log_softmax(logits, 1), F.softmax(ref_logits, 1), reduction="batchmean"))
+        else:
+            action = logits[0].argmax()
+        values.append(value[0])
+        patch, patch_values, delta, reduction, _ = build_action(
+            target, canvas, box, config, channel_bits, int(action),
+        )
+        if reduction <= 0:
+            break
+        work += float(patch_values.size)
+        c = patch["channel"]
+        ops.apply_delta(canvas[:, :, c], patch["x"], patch["y"], patch["w"], patch["h"], delta)
+        patches.append(patch)
+        channel_scores[c] = PBC3._channel_sum_error(target, canvas, c)
+    mse, bpp = stream_result(prep, config, patches, canvas, base_values)
+    return mse, bpp, work, time.perf_counter() - started, logps, values, entropies, kls
+
+
+def reward(base, result, args):
+    base_mse, base_bpp, base_work = base
+    mse, bpp, work = result
+    dmse = (base_mse - mse) / max(base_mse, 1.0)
+    dbpp = (base_bpp - bpp) / max(base_bpp, 1e-6)
+    dwork = (base_work - work) / max(base_work, 1.0)
+    quality_weight = args.quality_weight if dmse >= 0 else args.worse_quality_weight
+    return float(np.clip(quality_weight * dmse + args.rate_weight * dbpp + args.speed_weight * dwork, -5, 5))
+
+
+def evaluate(model, reference, source, episodes, bases, args):
+    rows = []
+    model.eval()
+    with torch.no_grad():
+        for path, preset in episodes:
+            mse, bpp, work, seconds, *_ = rollout(model, reference, source, path, preset, False, 1.0)
+            r = reward(bases[(path, preset)], (mse, bpp, work), args)
+            rows.append((preset, mse, bpp, work, seconds, r))
+    return {
+        "reward": float(np.mean([r[5] for r in rows])),
+        "mse": float(np.mean([r[1] for r in rows])),
+        "bpp": float(np.mean([r[2] for r in rows])),
+        "work": float(np.mean([r[3] for r in rows])),
+        "rollout_seconds": float(np.mean([r[4] for r in rows])),
+        "presets": {
+            preset: {
+                "mse": float(np.mean([r[1] for r in rows if r[0] == preset])),
+                "bpp": float(np.mean([r[2] for r in rows if r[0] == preset])),
+            }
+            for preset in args.selected_presets
+        },
+    }
+
+
+def train(args):
+    torch.set_num_threads(1)
+    random.seed(args.seed)
+    np.random.seed(args.seed)
+    torch.manual_seed(args.seed)
+    args.selected_presets = [preset.strip() for preset in args.presets.split(",")]
+    unknown = [preset for preset in args.selected_presets if preset not in PRESETS]
+    if unknown:
+        raise ValueError(f"unknown presets: {', '.join(unknown)}")
+    model, source = load_policy(args.init)
+    reference = copy.deepcopy(model).eval()
+    for parameter in reference.parameters():
+        parameter.requires_grad_(False)
+    torch.nn.init.zeros_(model.value.weight)
+    torch.nn.init.zeros_(model.value.bias)
+    train_paths = image_paths(args.images)[:args.limit or None]
+    val_paths = image_paths(args.val_images)[:args.val_limit or None]
+    train_episodes = [(path, preset) for path in train_paths for preset in args.selected_presets]
+    val_episodes = [(path, preset) for path in val_paths for preset in args.selected_presets]
+    all_episodes = train_episodes + val_episodes
+    bases = {}
+    print(f"building baselines for {len(all_episodes)} episodes", flush=True)
+    for i, (path, preset) in enumerate(all_episodes, 1):
+        image = ImageOps.exif_transpose(Image.open(path)).convert("RGB")
+        result = PBC3.compress(image, config=PRESETS[preset](learned_filler_model_path="patch_policy.npz"))
+        _, _, initial_work, _, *_ = rollout(reference, reference, source, path, preset, False, 1.0)
+        bases[(path, preset)] = (
+            float(result.mse), len(result.data) * 8 / (image.width * image.height), initial_work,
+        )
+        if i % 10 == 0:
+            print(f"baseline {i}/{len(all_episodes)}", flush=True)
+    Path(args.out).parent.mkdir(parents=True, exist_ok=True)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
+    history = []
+    start_epoch = 0
+    resume_path = Path(args.resume) if args.resume else None
+    if resume_path is not None and resume_path.exists():
+        checkpoint = torch.load(resume_path, map_location="cpu", weights_only=False)
+        model.load_state_dict(checkpoint["model"])
+        optimizer.load_state_dict(checkpoint["optimizer"])
+        random.setstate(checkpoint["python_rng"])
+        np.random.set_state(checkpoint["numpy_rng"])
+        torch.random.set_rng_state(checkpoint["torch_rng"])
+        start_epoch = int(checkpoint["epoch"])
+        history = checkpoint.get("history", [])
+    initial = evaluate(model, reference, source, val_episodes, bases, args)
+    print("initial", json.dumps(initial), flush=True)
+    best = initial["reward"]
+    export(model, source, args.out)
+    for epoch in range(start_epoch + 1, start_epoch + args.epochs + 1):
+        random.shuffle(train_episodes)
+        rewards = []
+        losses = []
+        for start in range(0, len(train_episodes), args.batch):
+            optimizer.zero_grad()
+            episode_losses = []
+            for path, preset in train_episodes[start:start + args.batch]:
+                mse, bpp, work, _, logps, values, entropies, kls = rollout(
+                    model, reference, source, path, preset, True, args.temperature,
+                )
+                r = reward(bases[(path, preset)], (mse, bpp, work), args)
+                rewards.append(r)
+                if not logps:
+                    continue
+                logps = torch.stack(logps)
+                values = torch.stack(values)
+                target = torch.full_like(values, r)
+                advantage = (target - values).detach()
+                loss = -(logps * advantage).mean()
+                loss += args.value_weight * F.mse_loss(values, target)
+                loss -= args.entropy_weight * torch.stack(entropies).mean()
+                loss += args.kl_weight * torch.stack(kls).mean()
+                episode_losses.append(loss)
+            if episode_losses:
+                loss = torch.stack(episode_losses).mean()
+                loss.backward()
+                torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+                optimizer.step()
+                losses.append(float(loss.detach()))
+        validation = evaluate(model, reference, source, val_episodes, bases, args)
+        row = {
+            "epoch": epoch,
+            "train_reward": float(np.mean(rewards)),
+            "loss": float(np.mean(losses)),
+            "validation": validation,
+        }
+        history.append(row)
+        print(json.dumps(row), flush=True)
+        if validation["reward"] > best:
+            best = validation["reward"]
+            export(model, source, args.out)
+        torch.save(model.state_dict(), f"{args.out}.epoch{epoch}.pt")
+        torch.save({
+            "epoch": epoch,
+            "model": model.state_dict(),
+            "optimizer": optimizer.state_dict(),
+            "python_rng": random.getstate(),
+            "numpy_rng": np.random.get_state(),
+            "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_reward": best, "history": history}, indent=2))
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--images", default="hpt_data")
+    parser.add_argument("--val-images", default="hpt_data_val")
+    parser.add_argument("--presets", default=",".join(PRESETS))
+    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("--epochs", type=int, default=5)
+    parser.add_argument("--batch", type=int, default=4)
+    parser.add_argument("--limit", type=int, default=0)
+    parser.add_argument("--val-limit", type=int, default=0)
+    parser.add_argument("--lr", type=float, default=1e-5)
+    parser.add_argument("--temperature", type=float, default=0.8)
+    parser.add_argument("--entropy-weight", type=float, default=0.001)
+    parser.add_argument("--kl-weight", type=float, default=0.02)
+    parser.add_argument("--value-weight", type=float, default=0.25)
+    parser.add_argument("--quality-weight", type=float, default=2.0)
+    parser.add_argument("--worse-quality-weight", type=float, default=6.0)
+    parser.add_argument("--rate-weight", type=float, default=1.0)
+    parser.add_argument("--speed-weight", type=float, default=0.1)
+    parser.add_argument("--weight-decay", type=float, default=1e-5)
+    parser.add_argument("--grad-clip", type=float, default=1.0)
+    parser.add_argument("--seed", type=int, default=2003)
+    train(parser.parse_args())
+
+
+if __name__ == "__main__":
+    main()

pbc3_students/patch_policy_f26_a20_h512_l2_e1200.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:36e525e00c34e20ccd786f5bc5f7c9d824ffb5e2cb22d942aa9732c7481af6ed
+size 1069260

pbc3_students/patch_policy_f26_a25_h512_l2_e1200.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2937b0b1935ac84e39ec44e6d621be919f4612d75e3494d24d698a7c198acffd
+size 1078979

server.py

@@ -27,6 +27,7 @@ from PBC3_animation import animate_pbc3
 import pbc3_sweep
 import pbc3_quick_rd
 import pbc3_benchmark
+import train_api
 
 import urllib.request
 
@@ -50,6 +51,7 @@ def _benchmark_busy():
         pbc3_sweep.status().get("running")
         or pbc3_quick_rd.status().get("running")
         or pbc3_benchmark.status().get("running")
+        or train_api.status().get("running")
     )
 
 
@@ -348,8 +350,8 @@ async def compress(request: Request):
     mode = form.get("mode", "Auto")
     mode = mode if mode in {"Auto", "Semi", "Manual"} else "Manual"
     if mode == "Auto":
-        preset = form.get("auto_config", "quality")
-        config = PBC3_PRESETS.get(preset, PBC3Config.quality)()
+        preset = form.get("auto_config", "high_quality")
+        config = PBC3_PRESETS.get(preset, PBC3Config.high_quality)()
         kwargs = {"auto_config": preset}
     else:
         preset = None
@@ -967,4 +969,5 @@ class NoCacheStaticFiles(StaticFiles):
             response.headers["Expires"] = "0"
         return response
 
+train_api.register(app)
 app.mount("/", NoCacheStaticFiles(directory="static", html=True), name="static")

static/app.js

@@ -417,7 +417,7 @@ const PARAMS = [
 
 const paramState = {};
 PARAMS.forEach(p => paramState[p.id] = p.value);
-paramState.preset = "quality";
+paramState.preset = "high_quality";
 let PRESET_VALUES = {};
 fetch("/api/presets").then(r => r.json()).then(d => {
   PRESET_VALUES = d;
@@ -454,7 +454,7 @@ document.querySelectorAll("#param-mode .seg-btn").forEach(b => b.addEventListene
   b.classList.add("active");
   paramMode = b.dataset.mode;
   if (paramMode === "Auto") {
-    const preset = (paramState.preset && paramState.preset !== "custom") ? paramState.preset : "quality";
+    const preset = (paramState.preset && paramState.preset !== "custom") ? paramState.preset : "high_quality";
     applyPreset(preset);
   }
   renderParams();
@@ -1259,4 +1259,4 @@ function toast(msg) {
 /* ============================================================
    BOOT
 ============================================================ */
-initRoster();
+initRoster();

static/index.html

@@ -51,6 +51,18 @@
   .tr-check input { width:16px; height:16px; accent-color:var(--red); }
   .tr-actions { display:flex; gap:10px; flex-wrap:wrap; margin-top:14px; }
   .tr-ck { border-top:1px solid rgba(255,255,255,.08); padding:14px 0; }
+  .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-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; }
+  .tr-table th:first-child, .tr-table td:first-child { text-align:left; }
+  .tr-checkpoint-grid { display:grid; grid-template-columns:repeat(auto-fill,minmax(290px,1fr)); gap:10px; }
+  .tr-filter { width:100%; max-width:460px; margin:8px 0 14px; background:var(--surface); border:1px solid var(--border); border-radius:7px; color:var(--text); padding:9px 11px; font:13px 'JetBrains Mono',monospace; }
+  .tr-checkpoint { display:flex; justify-content:space-between; align-items:center; gap:12px; padding:12px; border:1px solid var(--border); border-radius:9px; background:var(--surface); }
+  .tr-checkpoint b, .tr-checkpoint small { display:block; }
+  .tr-checkpoint small { margin-top:5px; color:var(--muted); font:11px 'JetBrains Mono',monospace; }
   .tr-card-title { margin:0 0 1rem; font-size:15px; color:var(--text); }
   .tr-output-grid { display:grid; grid-template-columns:repeat(auto-fill,minmax(310px,1fr)); gap:1.4rem; }
   .tr-viewer .viewer-controls { flex-wrap:wrap; justify-content:center; }
@@ -89,6 +101,10 @@
         Sweep Experiments
         <svg width="13" height="13" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><path d="M5 12h14M13 6l6 6-6 6"/></svg>
       </button>
+      <button id="enter-train" class="sweep-link-btn">
+        Training Lab
+        <svg width="13" height="13" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><path d="M5 12h14M13 6l6 6-6 6"/></svg>
+      </button>
       <a class="repo" href="https://github.com/EgeEken/PBC" target="_blank" rel="noopener">github.com/EgeEken/PBC ↗</a>
     </div>
 
@@ -102,6 +118,19 @@
   </div>
 </section>
 
+<!-- ===================== TRAINING LAB ===================== -->
+<section id="train-lab" class="section app" hidden>
+  <header class="app-header">
+    <div class="app-brand" id="train-back"><span class="dot"></span> PBC · Training Lab</div>
+    <nav class="app-nav">
+      <span class="nav-item active" data-tr-tab="run">Training Run</span>
+      <span class="nav-item" data-tr-tab="checkpoints">Checkpoints</span>
+    </nav>
+  </header>
+  <div id="view-train" class="view"></div>
+  <div id="view-train-checkpoints" class="view" hidden></div>
+</section>
+
 <!-- ===================== DEMO ===================== -->
 <section id="app" class="section app" hidden>
   <header class="app-header">
@@ -255,5 +284,6 @@
 <script src="/benchmark.js?v=20260625f"></script>
 <script src="/quick_rd.js?v=20260625e"></script>
 <script src="/sweep_runner.js?v=20260624a"></script>
+<script src="/train.js?v=20260718a"></script>
 </body>
 </html>

static/train.js

@@ -0,0 +1,152 @@
+(() => {
+  const $ = (selector) => document.querySelector(selector);
+  const runRoot = $("#view-train");
+  const checkpointRoot = $("#view-train-checkpoints");
+  if (!runRoot || !checkpointRoot) return;
+
+  let poll = null;
+  let checkpoints = [];
+  let activeTab = "run";
+
+  const presets = [
+    ["compression", "Compression"],
+    ["balanced", "Balanced"],
+    ["quality", "Quality"],
+    ["high_quality", "High quality"],
+  ];
+
+  const esc = (value) => String(value ?? "").replace(/[&<>\"]/g, (c) => ({ "&": "&amp;", "<": "&lt;", ">": "&gt;", "\"": "&quot;" }[c]));
+  const api = async (path, options) => {
+    const response = await fetch(path, options);
+    const data = await response.json().catch(() => ({}));
+    if (!response.ok) throw new Error(data.error || `HTTP ${response.status}`);
+    return data;
+  };
+  const download = (path) => window.open(`/api/train/download_checkpoint?path=${encodeURIComponent(path)}`, "_blank");
+  const fmt = (value, digits = 4) => value == null ? "—" : Number(value).toFixed(digits);
+
+  function showLab() {
+    ["landing", "app", "sweep-lab"].forEach((id) => { const el = $(`#${id}`); if (el) el.hidden = true; });
+    $("#train-lab").hidden = false;
+    renderRun();
+    loadCheckpoints();
+    refresh();
+  }
+
+  function showLanding() {
+    $("#train-lab").hidden = true;
+    $("#landing").hidden = false;
+    if (poll) clearInterval(poll);
+    poll = null;
+  }
+
+  $("#enter-train")?.addEventListener("click", showLab);
+  $("#train-back")?.addEventListener("click", showLanding);
+  document.querySelectorAll("[data-tr-tab]").forEach((tab) => tab.addEventListener("click", () => {
+    activeTab = tab.dataset.trTab;
+    document.querySelectorAll("[data-tr-tab]").forEach((item) => item.classList.toggle("active", item.dataset.trTab === activeTab));
+    runRoot.hidden = activeTab !== "run";
+    checkpointRoot.hidden = activeTab !== "checkpoints";
+    if (activeTab === "checkpoints") renderCheckpoints();
+  }));
+
+  function renderRun() {
+    runRoot.innerHTML = `
+      <div class="params">
+        <div class="params-head"><span class="field-label">Long-horizon RL training</span><span class="card-date">hpt_data → hpt_data_val</span></div>
+        <p class="card-date">Runs the current filler trainer in a resumable subprocess. Start the same output again to continue from its optimizer and RNG checkpoint.</p>
+        <div class="tr-fields">
+          <div class="tr-f full"><span>Preset policies</span><div class="tr-presets">${presets.map(([id, label], i) => `<label><input type="checkbox" data-tr-preset="${id}" ${i === 3 ? "checked" : ""}> ${label}</label>`).join("")}</div></div>
+          <label class="tr-f"><span>Additional epochs</span><input id="tr-epochs" type="number" min="1" step="1" value="100"></label>
+          <label class="tr-f"><span>Batch size</span><input id="tr-batch" type="number" min="1" step="1" value="4"></label>
+          <label class="tr-f"><span>Initial model</span><select id="tr-init"></select></label>
+          <label class="tr-f"><span>Output name</span><input id="tr-output" value="rl_high_quality.npz"></label>
+          <label class="tr-f"><span>Rate weight</span><input id="tr-rate" type="number" step="0.05" value="1.5"></label>
+          <label class="tr-f"><span>Speed weight</span><input id="tr-speed" type="number" step="0.05" value="0.15"></label>
+          <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>
+        </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></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>`;
+    $("#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");
+    fillInitialModels();
+  }
+
+  async function loadCheckpoints() {
+    const data = await api("/api/train/checkpoints").catch(() => ({ checkpoints: [] }));
+    checkpoints = data.checkpoints || [];
+    fillInitialModels();
+    if (activeTab === "checkpoints") renderCheckpoints();
+  }
+
+  function fillInitialModels() {
+    const select = $("#tr-init");
+    if (!select) return;
+    const old = select.value;
+    select.innerHTML = checkpoints.map((c) => `<option value="${esc(c.path)}">${esc(c.name)}</option>`).join("");
+    if (old && checkpoints.some((c) => c.path === old)) select.value = old;
+    else if (checkpoints.length) select.value = checkpoints.find((c) => c.name.includes("f26_a20_h512"))?.path || checkpoints[0].path;
+  }
+
+  function selectedPresets() {
+    return [...document.querySelectorAll("[data-tr-preset]:checked")].map((el) => el.dataset.trPreset);
+  }
+
+  async function start() {
+    const selected = selectedPresets();
+    if (!selected.length) return;
+    const body = {
+      presets: selected.join(","), epochs: Number($("#tr-epochs").value), batch: Number($("#tr-batch").value),
+      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),
+    };
+    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; }
+  }
+
+  function renderHistory(history) {
+    const root = $("#tr-history");
+    if (!root) return;
+    const rows = history.filter((row) => row.epoch || row.validation).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.mse, 3)}</td><td>${fmt(v.bpp, 6)}</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>Reward</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>`;
+  }
+
+  async function refresh() {
+    const data = await api("/api/train/status").catch((error) => ({ running: false, log_tail: error.message, history: [] }));
+    const status = $("#tr-status");
+    if (!status) return;
+    const state = data.running ? "Running" : data.return_code == null ? "Idle" : data.return_code === 0 ? "Finished" : `Stopped / failed (${data.return_code})`;
+    const output = data.output ? `<code>${esc(data.output)}</code>` : "No output selected";
+    status.innerHTML = `<b>${state}</b> · ${output}${data.spec?.presets ? ` · ${esc(data.spec.presets)}` : ""}`;
+    $("#tr-log-tail").textContent = data.log_tail || "";
+    renderHistory(data.history || []);
+    if (data.running && !poll) poll = setInterval(refresh, 3000);
+    if (!data.running && poll) { clearInterval(poll); poll = null; loadCheckpoints(); }
+  }
+
+  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;
+    renderCheckpointList();
+  }
+
+  function renderCheckpointList() {
+    const root = $("#tr-checkpoint-list");
+    if (!root) return;
+    const filter = ($("#tr-ckpt-filter").value || "").toLowerCase();
+    const visible = checkpoints.filter((c) => c.name.toLowerCase().includes(filter));
+    root.innerHTML = visible.length ? visible.map((c) => `<div class="tr-checkpoint"><div><b>${esc(c.name)}</b><small>${(c.bytes / 1024).toFixed(1)} KB</small></div><button class="hold-btn" data-download="${esc(c.path)}">Download</button></div>`).join("") : `<p class="card-date">No matching checkpoints.</p>`;
+    root.querySelectorAll("[data-download]").forEach((button) => button.onclick = () => download(button.dataset.download));
+  }
+
+  setInterval(() => { if (!$("#train-lab").hidden && activeTab === "run") refresh(); }, 5000);
+})();

train_api.py

@@ -0,0 +1,184 @@
+import json
+import os
+import subprocess
+import sys
+import threading
+import time
+from pathlib import Path
+
+from fastapi import Request
+from fastapi.responses import FileResponse, JSONResponse, PlainTextResponse
+
+
+ROOT = Path(__file__).resolve().parent
+RUNS = ROOT / "training_space" / "runs"
+RUNS.mkdir(parents=True, exist_ok=True)
+IMAGE_FOLDERS = {"train": ROOT / "hpt_data", "validation": ROOT / "hpt_data_val"}
+MODEL_FOLDERS = (ROOT / "pbc3_students", RUNS)
+LOCK = threading.Lock()
+PROCESS = None
+CURRENT = {}
+
+
+def _safe_model_path(value):
+    path = (ROOT / value).resolve() if not os.path.isabs(value) else Path(value).resolve()
+    if not any(path == folder.resolve() or folder.resolve() in path.parents for folder in MODEL_FOLDERS):
+        return None
+    return path if path.is_file() else None
+
+
+def _output_path(value):
+    name = Path(value or "rl_training.npz").name
+    if not name.endswith(".npz"):
+        name += ".npz"
+    return (RUNS / name).resolve()
+
+
+def _json_lines(log_path):
+    rows = []
+    if not log_path.exists():
+        return rows
+    for line in log_path.read_text(encoding="utf-8", errors="replace").splitlines():
+        candidate = line.strip()
+        if candidate.startswith("initial "):
+            candidate = candidate[8:]
+        try:
+            row = json.loads(candidate)
+        except Exception:
+            continue
+        if "epoch" in row or "reward" in row:
+            rows.append(row)
+    return rows
+
+
+def _refresh_process():
+    global PROCESS
+    if PROCESS is not None and PROCESS.poll() is not None:
+        CURRENT["return_code"] = PROCESS.returncode
+        PROCESS = None
+
+
+def status():
+    with LOCK:
+        _refresh_process()
+        output = Path(CURRENT["output"]) if CURRENT.get("output") else None
+        log = Path(CURRENT["log"]) if CURRENT.get("log") else None
+        history_path = Path(f"{output}.json") if output else None
+        history = []
+        if history_path and history_path.exists():
+            try:
+                saved = json.loads(history_path.read_text(encoding="utf-8"))
+                history = saved.get("history", [])
+            except Exception:
+                pass
+        if not history:
+            history = _json_lines(log) if log else []
+        return {
+            "running": PROCESS is not None,
+            "return_code": CURRENT.get("return_code"),
+            "output": str(output.relative_to(ROOT)) if output else None,
+            "log": str(log.relative_to(ROOT)) if log else None,
+            "started": CURRENT.get("started"),
+            "spec": CURRENT.get("spec", {}),
+            "history": history,
+            "log_tail": log.read_text(encoding="utf-8", errors="replace")[-10000:] if log and log.exists() else "",
+        }
+
+
+def start(spec):
+    global PROCESS
+    with LOCK:
+        _refresh_process()
+        if PROCESS is not None:
+            return {"error": "A training run is already active."}
+        presets = str(spec.get("presets", "high_quality"))
+        selected = [p.strip() for p in presets.split(",") if p.strip()]
+        allowed = {"compression", "balanced", "quality", "high_quality"}
+        if not selected or any(p not in allowed for p in selected):
+            return {"error": "Presets must be compression, balanced, quality, or high_quality."}
+        init = _safe_model_path(str(spec.get("init", "pbc3_students/patch_policy_f26_a20_h512_l2_e1200.npz")))
+        if init is None:
+            return {"error": "Initial model was not found in pbc3_students or training_space/runs."}
+        output = _output_path(spec.get("output", "rl_training.npz"))
+        output.parent.mkdir(parents=True, exist_ok=True)
+        log = Path(f"{output}.log")
+        resume = Path(f"{output}.resume.pt")
+        command = [
+            sys.executable, "pbc3_rl_a20.py",
+            "--presets", ",".join(selected),
+            "--epochs", str(int(spec.get("epochs", 100))),
+            "--batch", str(int(spec.get("batch", 4))),
+            "--init", os.fspath(init),
+            "--out", os.fspath(output),
+            "--rate-weight", str(float(spec.get("rate_weight", 1.5))),
+            "--speed-weight", str(float(spec.get("speed_weight", 0.15))),
+            "--temperature", str(float(spec.get("temperature", 0.9))),
+            "--entropy-weight", str(float(spec.get("entropy_weight", 0.003))),
+            "--kl-weight", str(float(spec.get("kl_weight", 0.015))),
+            "--quality-weight", str(float(spec.get("quality_weight", 2.0))),
+            "--worse-quality-weight", str(float(spec.get("worse_quality_weight", 7.0))),
+        ]
+        if resume.exists():
+            command.extend(["--resume", os.fspath(resume)])
+        log.write_text("", encoding="utf-8")
+        handle = log.open("a", encoding="utf-8")
+        PROCESS = subprocess.Popen(command, cwd=ROOT, stdout=handle, stderr=subprocess.STDOUT)
+        handle.close()
+        CURRENT.clear()
+        CURRENT.update({"output": str(output), "log": str(log), "started": time.time(), "spec": spec, "return_code": None})
+        return {"ok": True, "output": str(output.relative_to(ROOT)), "resuming": resume.exists()}
+
+
+def stop():
+    with LOCK:
+        _refresh_process()
+        if PROCESS is None:
+            return {"ok": False, "error": "No training run is active."}
+        PROCESS.terminate()
+        return {"ok": True}
+
+
+def checkpoints():
+    rows = []
+    for folder in MODEL_FOLDERS:
+        for path in sorted(folder.glob("*.npz"), key=lambda p: p.stat().st_mtime, reverse=True):
+            rows.append({"name": path.name, "path": str(path.relative_to(ROOT)), "bytes": path.stat().st_size, "modified": path.stat().st_mtime})
+    return {"checkpoints": rows}
+
+
+def download_checkpoint(path):
+    safe = _safe_model_path(path)
+    if safe is None or safe.suffix != ".npz":
+        return JSONResponse({"error": "Checkpoint not found or not allowed."}, status_code=404)
+    return FileResponse(safe, filename=safe.name, media_type="application/octet-stream")
+
+
+def register(app):
+    @app.get("/api/train/status")
+    def train_status():
+        return status()
+
+    @app.post("/api/train/start")
+    async def train_start(request: Request):
+        result = start(await request.json())
+        return JSONResponse(result, status_code=400 if result.get("error") else 200)
+
+    @app.post("/api/train/stop")
+    def train_stop():
+        return stop()
+
+    @app.get("/api/train/checkpoints")
+    def train_checkpoints():
+        return checkpoints()
+
+    @app.get("/api/train/download_checkpoint")
+    def train_download_checkpoint(path: str):
+        return download_checkpoint(path)
+
+    @app.get("/api/train/log")
+    def train_log():
+        data = status()
+        path = ROOT / data["log"] if data.get("log") else None
+        if not path or not path.exists():
+            return PlainTextResponse("No training log is available.")
+        return FileResponse(path, filename=path.name, media_type="text/plain")

training_space/README.md

@@ -0,0 +1,12 @@
+# PBC filler RL training Space
+
+Run `app.py` from the project root, or copy this directory together with the PBC source files, `hpt_data`, `hpt_data_val`, and the initial student model into a dedicated Hugging Face Space.
+
+The app runs preset-specific RL in a background process. Starting the same output again automatically resumes from its `.resume.pt` checkpoint, including model, optimizer, RNG, and history state. Download exported `.npz` checkpoints from the Training Lab before restarting the Space; persistent storage is not required.
+
+Example local launch:
+
+```text
+pip install -r training_space/requirements.txt
+python training_space/app.py
+```

training_space/app.py

@@ -0,0 +1,82 @@
+import os
+import subprocess
+import threading
+from pathlib import Path
+
+import gradio as gr
+
+
+ROOT = Path(__file__).resolve().parent.parent
+RUNS = ROOT / "training_space" / "runs"
+RUNS.mkdir(parents=True, exist_ok=True)
+process = None
+process_lock = threading.Lock()
+
+
+def run_command(command, log_path):
+    global process
+    with log_path.open("a", encoding="utf-8") as log:
+        process = subprocess.Popen(command, cwd=ROOT, stdout=log, stderr=subprocess.STDOUT)
+        process.wait()
+
+
+def start_training(presets, epochs, init_path, output_name, rate_weight, speed_weight):
+    global process
+    with process_lock:
+        if process is not None and process.poll() is None:
+            return "A training process is already running."
+        output = RUNS / output_name
+        output.parent.mkdir(parents=True, exist_ok=True)
+        resume = Path(f"{output}.resume.pt")
+        command = [
+            os.fspath(Path(os.sys.executable)), "pbc3_rl_a20.py",
+            "--presets", presets,
+            "--epochs", str(int(epochs)),
+            "--init", init_path,
+            "--out", os.fspath(output),
+            "--rate-weight", str(float(rate_weight)),
+            "--speed-weight", str(float(speed_weight)),
+        ]
+        if resume.exists():
+            command.extend(["--resume", os.fspath(resume)])
+        log_path = Path(f"{output}.log")
+        log_path.write_text("", encoding="utf-8")
+        threading.Thread(target=run_command, args=(command, log_path), daemon=True).start()
+        return f"Started: {' '.join(command)}\nLog: {log_path}"
+
+
+def status():
+    with process_lock:
+        running = process is not None and process.poll() is None
+        code = None if process is None or running else process.returncode
+    logs = sorted(RUNS.glob("*.log"), key=lambda path: path.stat().st_mtime, reverse=True)
+    tail = logs[0].read_text(encoding="utf-8", errors="replace")[-8000:] if logs else ""
+    return f"running={running}, return_code={code}\n\n{tail}"
+
+
+def stop_training():
+    with process_lock:
+        if process is None or process.poll() is not None:
+            return "No training process is running."
+        process.terminate()
+        return "Stop requested; the current checkpoint will remain available."
+
+
+with gr.Blocks(title="PBC filler RL training") as demo:
+    gr.Markdown("# PBC filler RL training\nRuns resumable preset-specific training and writes checkpoints under `training_space/runs/`.")
+    presets = gr.Textbox(value="high_quality", label="Presets", info="Comma-separated: compression, balanced, quality, high_quality")
+    epochs = gr.Number(value=100, precision=0, label="Additional epochs per start/resume")
+    init_path = gr.Textbox(value="pbc3_students/patch_policy_f26_a20_h512_l2_e1200.npz", label="Initial model")
+    output_name = gr.Textbox(value="rl_high_quality.npz", label="Output filename")
+    rate_weight = gr.Number(value=1.5, label="Rate reward weight")
+    speed_weight = gr.Number(value=0.2, label="Speed reward weight")
+    start = gr.Button("Start or resume")
+    stop = gr.Button("Stop")
+    refresh = gr.Button("Refresh status")
+    output = gr.Textbox(lines=30, label="Status / log")
+    start.click(start_training, [presets, epochs, init_path, output_name, rate_weight, speed_weight], output)
+    stop.click(stop_training, outputs=output)
+    refresh.click(status, outputs=output)
+
+
+demo.launch(server_name="0.0.0.0")

training_space/requirements.txt

@@ -0,0 +1 @@
+gradio