Buckets:

AsianPlayer
/

TriDrive

Files

xet

AsianPlayer/TriDrive / code /train /train_road_dynamics.py

AsianPlayer

16 days ago

download

raw

21.3 kB

"""cwm3/train/train_road_dynamics.py — F5c arms (3) and (4): a small causal dynamics net on the 8 s CAN road history predicting RESIDUALS over the
constant-velocity extrapolation of the future road quantities (WP4 batch 1; pre-registration docs/decisions.md 2026-09-16T06:46:35Z; plan
<local-notes>/plans/elegant-dazzling-stallman.md F5b/F5c).
Targets (data/build_road_dynamics_targets.py, results/targets_road_dyn/road_dyn_{regime}_{split}.npz): Yf[N,25,4] = lead_dRel, lead_vRel, lead_vLead,
lane_center_offset at t+0.2k (k = 1..25) with Yvalid.  Arms:
  (1) persistence            y_k = y_0                                                   (eval/road_dynamics_eval.py, no training)
  (2) constant velocity (CV) dRel_k = dRel_0 + vRel_0 tau_k, vRel / vLead constant, lane_off_k = lane_off_0 + rate_0 tau_k (rate_0 = OLS slope of the last 1 s,
                             clipped +-0.5 m/s)                                          (cv_forecast below; eval, no training)
  (3) `dyn`     GRU(128) over the 8 s history @ 5 Hz (8 standardised canspec channels + dt) + the 40 standardised CAN anchor features (ecan, the F5b list) + mode
                -> MLP -> residual mean / log-sd over arm (2) per (k, channel), Gaussian NLL masked by Yvalid & anchor validity.  <= 0.3 M parameters (asserted).
  (4) `dynvis`  arm (3) + the FINAL E branch's frozen per-anchor visual side vector vf256 (results/pack_v2, PCA-256 of the causal full-frame V-JEPA features)
                through LN -> Linear(256 -> 64) -> GELU, concatenated before the MLP.
Recipe: AdamW 1e-3 / wd 1e-4, batch 2048, 400k rows per epoch (both regimes pooled, one model, mode embedding), <= 12 epochs, early stop on the masked
validation NLL (patience 3), 3 seeds (paired init: torch.manual_seed(seed) before construction), selection on VALIDATION only.
Outputs results/runs/RD_<arm>_s<k>/: config.json, stats.npz (train statistics), history.jsonl, BEST (+ ckpt), params.json, pred_rd_{regime}_{split}.npz for
test and val = {mu f16 [N,25,4] (raw units = CV + residual), sigma f16 [N,25,4] (raw units), row_ok, uid, t}, TRAIN_DONE / PRED_DONE, run.json.
  python train/train_road_dynamics.py --arm dyn --seed 0 [--out results/runs/RD_dyn_s0] [--epochs 12] [--synthetic N] [--cpu]
No openpilot DMS field is read; inputs are the canspec CAN road channels (+ vf256 for arm 4) only.
"""
from __future__ import annotations

import argparse
import json
import math
import os
import sys
import time
from pathlib import Path

os.environ.setdefault("OMP_NUM_THREADS", "4")
import numpy as np  # noqa: E402
import torch  # noqa: E402
import torch.nn as nn  # noqa: E402
import torch.nn.functional as Fn  # noqa: E402

CWM3 = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(CWM3 / "common"))
from runmeta import write_run_json  # noqa: E402

RD = CWM3 / "results" / "targets_road_dyn"
PACK_V2 = CWM3 / "results" / "pack_v2"
REGIMES, SPLITS = ("manual", "engaged"), ("train", "val", "test")
MODE = {"manual": 0, "engaged": 1}
HIST, H, DC, DH = 40, 25, 4, 8
TAU = (np.arange(1, H + 1) * 0.2).astype(np.float32)
LANE_RATE_CLIP = 0.5
MAX_PARAMS = 300_000
RECIPE = dict(batch=2048, rows_per_epoch=400_000, max_epochs=12, patience=3, lr=1e-3, weight_decay=1e-4, hid=128, hidden=256, d_vis=64, dropout=0.1)


# ----------------------------------------------------------------------------- data
def load_split(regime, split, arm, max_rows=None, seed=0):
    with np.load(RD / f"road_dyn_{regime}_{split}.npz", allow_pickle=True) as z:
        N = len(z["t"])
        sel = None
        if max_rows and max_rows < N:
            sel = np.sort(np.random.default_rng(seed).choice(N, max_rows, replace=False))
        sub = (lambda a: a[sel]) if sel is not None else (lambda a: a)
        out = {k: sub(z[k]) for k in ("H", "Hvalid", "Yf", "Yvalid", "ecan", "ecan_ok", "uid", "t", "dongle")}
    out["mode"] = np.full(len(out["t"]), MODE[regime], np.int8)
    if arm == "dynvis":
        with np.load(PACK_V2 / f"vf256_{regime}_{split}.npz") as v:
            assert int(v["n"]) == N, (regime, split, int(v["n"]), N)
            out["vf"], out["vf_ok"] = sub(v["Z"]), sub(v["row_ok"]).astype(bool)
    return out


def concat(parts):
    return {k: np.concatenate([p[k] for p in parts]) for k in parts[0]}


def synthetic(N, seed, arm):
    rng = np.random.default_rng(seed)
    Hh = np.zeros((N, HIST, DH), np.float32)
    lead = rng.random(N) < 0.6
    lane = rng.random(N) < 0.6
    v = rng.uniform(5, 30, N)
    Hh[:, :, 5] = v[:, None] + rng.normal(0, 0.3, (N, HIST))
    Hh[:, :, 6], Hh[:, :, 7] = lead[:, None], lane[:, None]
    vrel = rng.normal(-1, 2, N)
    Hh[:, :, 0] = np.where(lead[:, None], rng.uniform(10, 120, N)[:, None] + vrel[:, None] * (np.arange(HIST) - 39)[None, :] * 0.2, 200.0)
    Hh[:, :, 1] = np.where(lead[:, None], vrel[:, None] + rng.normal(0, 0.1, (N, HIST)), 0.0)
    Hh[:, :, 2] = np.where(lead[:, None], v[:, None] + vrel[:, None], 0.0)
    Hh[:, :, 3] = np.where(lane[:, None], rng.normal(0, 0.3, (N, 1)) + rng.normal(0, 0.03, (N, HIST)), 0.0)
    Hh[:, :, 4] = np.where(lane[:, None], 3.5, 0.0)
    Yf = np.zeros((N, H, DC), np.float32)
    Yf[:, :, 0] = Hh[:, -1, 0][:, None] + vrel[:, None] * TAU[None, :] + rng.normal(0, 1, (N, H))
    Yf[:, :, 1] = vrel[:, None] + rng.normal(0, 0.3, (N, H))
    Yf[:, :, 2] = Hh[:, -1, 2][:, None] + rng.normal(0, 0.3, (N, H))
    Yf[:, :, 3] = Hh[:, -1, 3][:, None] + rng.normal(0, 0.1, (N, H))
    Yv = np.stack([lead[:, None].repeat(H, 1)] * 3 + [lane[:, None].repeat(H, 1)], -1)
    out = {"H": Hh.astype(np.float16), "Hvalid": np.ones((N, HIST), bool), "Yf": Yf, "Yvalid": Yv, "ecan": rng.normal(0, 1, (N, 40)).astype(np.float32),
           "ecan_ok": np.ones(N, bool), "uid": np.array([f"r{i % 4}" for i in range(N)]), "t": np.arange(N) * 0.2, "dongle": np.array([f"d{i % 3}" for i in range(N)]),
           "mode": (np.arange(N) % 2).astype(np.int8)}
    if arm == "dynvis":
        out["vf"], out["vf_ok"] = rng.normal(0, 1, (N, 256)).astype(np.float16), np.ones(N, bool)
    return out


def compute_stats(tr):
    """train-only float64 statistics: per history channel (over valid steps), ecan (over ok rows), residual scale per target channel (sd of Yf - CV at k = 15)."""
    Hh = tr["H"].astype(np.float32)
    hv = tr["Hvalid"]
    hmu, hsd = np.zeros(DH), np.zeros(DH)
    for c in range(DH):
        x = Hh[:, :, c][hv].astype(np.float64)
        hmu[c], hsd[c] = x.mean(), x.std() + 1e-6
    ok = tr["ecan_ok"]
    e = tr["ecan"][ok].astype(np.float64)
    emu, esd = e.mean(0), e.std(0) + 1e-6
    cv = cv_forecast_np(Hh, hv)
    r = tr["Yf"] - cv
    m = tr["Yvalid"] & anchor_valid_np(Hh)[:, None, :]
    rsd = np.array([r[:, 14, c][m[:, 14, c]].astype(np.float64).std() + 1e-3 for c in range(DC)])
    return {"hmu": hmu.astype(np.float32), "hsd": hsd.astype(np.float32), "emu": emu.astype(np.float32), "esd": esd.astype(np.float32), "rsd": rsd.astype(np.float32)}


# ----------------------------------------------------------------------------- baselines (numpy, shared with the evaluator)
def anchor_valid_np(Hh):
    """[N,4] lead valid at the anchor for the 3 lead channels, lane valid for the lane channel."""
    lead0, lane0 = Hh[:, -1, 6] > 0.5, Hh[:, -1, 7] > 0.5
    return np.stack([lead0, lead0, lead0, lane0], -1)


def lane_rate_np(Hh, hv):
    """OLS slope (m/s) of the lane offset over the last 1 s (5 history steps) using lane-valid steps; 0 with < 3 steps; clipped."""
    x = Hh[:, -5:, 3].astype(np.float64)
    w = (hv[:, -5:] & (Hh[:, -5:, 7] > 0.5)).astype(np.float64)
    dt = np.arange(-4, 1) * 0.2
    n = w.sum(1)
    sw = np.maximum(n, 1)
    xbar, tbar = (w * x).sum(1) / sw, (w * dt[None]).sum(1) / sw
    num = (w * (dt[None] - tbar[:, None]) * (x - xbar[:, None])).sum(1)
    den = (w * (dt[None] - tbar[:, None]) ** 2).sum(1)
    return np.clip(np.where((n >= 3) & (den > 1e-9), num / np.maximum(den, 1e-9), 0.0), -LANE_RATE_CLIP, LANE_RATE_CLIP).astype(np.float32)


def persistence_np(Hh):
    return np.repeat(Hh[:, -1:, :4].astype(np.float32), H, 1)


def cv_forecast_np(Hh, hv):
    p = persistence_np(Hh)
    p[:, :, 0] = Hh[:, -1:, 0] + Hh[:, -1:, 1] * TAU[None, :]
    p[:, :, 3] = Hh[:, -1:, 3] + lane_rate_np(Hh, hv)[:, None] * TAU[None, :]
    return p


def cv_forecast_t(Hh, hv):
    """torch version of cv_forecast_np (same arithmetic)."""
    B = Hh.shape[0]
    tau = torch.as_tensor(TAU, device=Hh.device)
    p = Hh[:, -1:, :4].expand(B, H, 4).clone()
    p[:, :, 0] = Hh[:, -1:, 0] + Hh[:, -1:, 1] * tau[None, :]
    x = Hh[:, -5:, 3]
    w = (hv[:, -5:] & (Hh[:, -5:, 7] > 0.5)).float()
    dt = torch.arange(-4, 1, device=Hh.device).float() * 0.2
    n = w.sum(1)
    sw = n.clamp(min=1)
    xbar, tbar = (w * x).sum(1) / sw, (w * dt[None]).sum(1) / sw
    num = (w * (dt[None] - tbar[:, None]) * (x - xbar[:, None])).sum(1)
    den = (w * (dt[None] - tbar[:, None]) ** 2).sum(1)
    rate = torch.where((n >= 3) & (den > 1e-9), num / den.clamp(min=1e-9), torch.zeros_like(num)).clamp(-LANE_RATE_CLIP, LANE_RATE_CLIP)
    p[:, :, 3] = Hh[:, -1:, 3] + rate[:, None] * tau[None, :]
    return p


# ----------------------------------------------------------------------------- model
class RoadDynNet(nn.Module):
    def __init__(self, arm, hid=128, hidden=256, d_vis=64, dropout=0.1, n_ecan=40):
        super().__init__()
        self.arm = arm
        self.gru = nn.GRU(DH + 1, hid, batch_first=True)
        self.ecan = nn.Sequential(nn.Linear(n_ecan, 32), nn.GELU())
        self.mode = nn.Embedding(2, 8)
        d_in = hid + 32 + 8
        if arm == "dynvis":
            self.vis = nn.Sequential(nn.LayerNorm(256), nn.Linear(256, d_vis), nn.GELU(), nn.Dropout(dropout))
            d_in += d_vis
        self.head = nn.Sequential(nn.Linear(d_in, hidden), nn.GELU(), nn.Dropout(dropout), nn.Linear(hidden, 2 * H * DC))
        nn.init.zeros_(self.head[-1].bias)
        with torch.no_grad():
            self.head[-1].weight.mul_(0.1)

    def forward(self, b):
        x = torch.cat([b["Hs"], b["dt"]], -1)
        _, h = self.gru(x)
        z = [h[-1], self.ecan(b["ecan"]), self.mode(b["mode"])]
        if self.arm == "dynvis":
            z.append(self.vis(b["vf"]) * b["vf_ok"].float()[:, None])
        o = self.head(torch.cat(z, -1)).view(-1, H, DC, 2)
        return o[..., 0], o[..., 1].clamp(-6, 4)                                            # residual mean, log-sd (scaled units)


def n_params(m):
    return sum(p.numel() for p in m.parameters() if p.requires_grad)


# ----------------------------------------------------------------------------- GPU store + batch
class Store:
    def __init__(self, raw, device):
        self.d = device
        self.N = len(raw["t"])
        self.H = torch.from_numpy(raw["H"]).to(device)
        self.Hvalid = torch.from_numpy(raw["Hvalid"]).to(device)
        self.Yf = torch.from_numpy(raw["Yf"]).to(device)
        self.Yvalid = torch.from_numpy(raw["Yvalid"]).to(device)
        self.ecan = torch.from_numpy(raw["ecan"]).to(device)
        self.ecan_ok = torch.from_numpy(raw["ecan_ok"]).to(device)
        self.mode = torch.from_numpy(raw["mode"].astype(np.int64)).to(device)
        self.vf = torch.from_numpy(raw["vf"]).to(device) if "vf" in raw else None
        self.vf_ok = torch.from_numpy(raw["vf_ok"]).to(device) if "vf_ok" in raw else None


def prep(store, idx, st, need_targets=True):
    d = store.d
    Hh = store.H[idx].float()
    hv = store.Hvalid[idx]
    hmu, hsd = torch.as_tensor(st["hmu"], device=d), torch.as_tensor(st["hsd"], device=d)
    Hs = torch.clamp((Hh - hmu) / hsd, -8, 8) * hv[..., None].float()
    dt = (-(HIST - 1 - torch.arange(HIST, device=d).float()) * 0.2)[None, :, None].expand(len(idx), -1, -1)
    emu, esd = torch.as_tensor(st["emu"], device=d), torch.as_tensor(st["esd"], device=d)
    ecan = torch.clamp((store.ecan[idx] - emu) / esd, -8, 8) * store.ecan_ok[idx].float()[:, None]
    b = {"Hs": Hs, "dt": dt, "ecan": ecan, "mode": store.mode[idx]}
    if store.vf is not None:
        b["vf"], b["vf_ok"] = store.vf[idx].float(), store.vf_ok[idx]
    cv = cv_forecast_t(Hh, hv)
    lead0, lane0 = Hh[:, -1, 6] > 0.5, Hh[:, -1, 7] > 0.5
    av = torch.stack([lead0, lead0, lead0, lane0], -1)
    tg = None
    if need_targets:
        rsd = torch.as_tensor(st["rsd"], device=d)
        r = (store.Yf[idx] - cv) / rsd
        tg = {"r": r, "m": store.Yvalid[idx] & av[:, None, :]}
    return b, tg, cv, av


def nll(mu, logsd, tg):
    ll = 0.5 * ((tg["r"] - mu) / logsd.exp()) ** 2 + logsd + 0.5 * math.log(2 * math.pi)
    m = tg["m"].float()
    return (ll * m).sum() / m.sum().clamp(min=1)


def log(msg, out):
    line = f"[{time.strftime('%H:%M:%S')}] {msg}"
    print(line, flush=True)
    with open(Path(out) / "train.log", "a") as f:
        f.write(line + "\n")


# ----------------------------------------------------------------------------- train / predict
def evaluate(model, store, st, device, bs=8192):
    model.eval()
    tot, cnt = 0.0, 0.0
    with torch.no_grad():
        for i in range(0, store.N, bs):
            idx = torch.arange(i, min(i + bs, store.N), device=device)
            b, tg, _, _ = prep(store, idx, st)
            mu, ls = model(b)
            m = tg["m"].float()
            ll = 0.5 * ((tg["r"] - mu) / ls.exp()) ** 2 + ls + 0.5 * math.log(2 * math.pi)
            tot += float((ll * m).sum())
            cnt += float(m.sum())
    return tot / max(cnt, 1.0)


def predict(model, raw, st, device, out_path, regime, split, bs=8192):
    store = Store(raw, device)
    model.eval()
    mu_all, sd_all, ok_all = [], [], []
    rsd = torch.as_tensor(st["rsd"], device=device)
    with torch.no_grad():
        for i in range(0, store.N, bs):
            idx = torch.arange(i, min(i + bs, store.N), device=device)
            b, _, cv, av = prep(store, idx, st, need_targets=False)
            mu, ls = model(b)
            mu_all.append((cv + mu * rsd).half().cpu().numpy())
            sd_all.append((ls.exp() * rsd).half().cpu().numpy())
            ok_all.append(av.cpu().numpy())
    np.savez_compressed(out_path, mu=np.concatenate(mu_all), sigma=np.concatenate(sd_all), row_ok=np.concatenate(ok_all), uid=raw["uid"], t=raw["t"],
                        channels=np.array(["lead_dRel", "lead_vRel", "lead_vLead", "lane_center_offset"]), regime=regime, split=split,
                        note="mu / sigma in raw units (m, m/s, m/s, m); mu = constant-velocity forecast + predicted residual; row_ok[N,4] = channel valid at the anchor")
    del store
    torch.cuda.empty_cache() if device.type == "cuda" else None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--arm", required=True, choices=["dyn", "dynvis"])
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--out", default="")
    ap.add_argument("--epochs", type=int, default=0)
    ap.add_argument("--max-rows", type=int, default=0, help="seeded subset of TRAIN rows (smoke)")
    ap.add_argument("--synthetic", type=int, default=0)
    ap.add_argument("--cpu", action="store_true")
    ap.add_argument("--no-lock", action="store_true", help="accepted for queue compatibility (this trainer never takes the GPU lock)")
    ap.add_argument("--minutes", type=float, default=0, help="accepted for queue compatibility")
    ap.add_argument("--full-pass", action="store_true", help="POST-HOC sensitivity arm `dynfull` (labelled as such; not the pre-registered recipe): one full pass over the train rows per epoch, <= 30 epochs, same patience")
    a = ap.parse_args()
    out = Path(a.out or (CWM3 / "results" / "runs" / f"RD_{a.arm}{'full' if a.full_pass else ''}_s{a.seed}{'_synthetic' if a.synthetic else ''}"))
    out.mkdir(parents=True, exist_ok=True)
    t_start = time.time()
    device = torch.device("cpu" if a.cpu or not torch.cuda.is_available() else "cuda")
    rec = dict(RECIPE)
    if a.full_pass:
        rec["rows_per_epoch"], rec["max_epochs"], rec["post_hoc"] = 0, 30, "full-pass sensitivity arm (2026-09-16; the pre-registered 12-epoch / 400k-row arm was still improving at its cap)"
    if a.epochs:
        rec["max_epochs"] = a.epochs
    cfg = {"arm": a.arm, "seed": a.seed, "recipe": rec, "regimes": list(REGIMES), "targets": "results/targets_road_dyn (Yf lead_dRel / lead_vRel / lead_vLead / lane_center_offset, k=1..25)",
           "baseline": "constant velocity (dRel += vRel tau; lane offset linear with the 1 s OLS rate clipped +-0.5 m/s); the net predicts residuals", "dms_fields_used": []}
    (out / "config.json").write_text(json.dumps({"cfg": cfg, "args": vars(a)}, indent=1))
    log(f"=== train_road_dynamics arm={a.arm} s{a.seed} out={out} device={device} ===", out)
    if a.synthetic:
        tr = synthetic(a.synthetic, a.seed, a.arm)
        va = synthetic(max(a.synthetic // 4, 64), a.seed + 1, a.arm)
        preds = {(r, s): synthetic(max(a.synthetic // 4, 64), a.seed + 2 + i, a.arm) for i, (r, s) in enumerate([(r, s) for r in REGIMES for s in ("test", "val")])}
    else:
        tr = concat([load_split(r, "train", a.arm, a.max_rows or None, a.seed) for r in REGIMES])
        va = concat([load_split(r, "val", a.arm) for r in REGIMES])
        preds = None
    if not (out / "TRAIN_DONE").exists():
        st = compute_stats(tr) if a.synthetic else compute_stats_real(tr)
        np.savez(out / "stats.npz", **st)
        torch.manual_seed(a.seed)
        model = RoadDynNet(a.arm, rec["hid"], rec["hidden"], rec["d_vis"], rec["dropout"]).to(device)
        npar = n_params(model)
        (out / "params.json").write_text(json.dumps({"trainable": npar, "limit": MAX_PARAMS, "by_module": {n: sum(p.numel() for p in m.parameters()) for n, m in model.named_children()}}, indent=1))
        assert npar <= MAX_PARAMS, npar
        log(f"params {npar:,} (<= {MAX_PARAMS:,}); train rows {len(tr['t']):,} val rows {len(va['t']):,}", out)
        store_tr, store_va = Store(tr, device), Store(va, device)
        opt = torch.optim.AdamW(model.parameters(), lr=rec["lr"], weight_decay=rec["weight_decay"])
        g = torch.Generator(device="cpu").manual_seed(a.seed)
        best, best_ep, bad = float("inf"), -1, 0
        for ep in range(rec["max_epochs"]):
            model.train()
            n_rows = min(rec["rows_per_epoch"], store_tr.N) if rec["rows_per_epoch"] else store_tr.N
            perm = torch.randperm(store_tr.N, generator=g)[:n_rows].to(device)
            tot, nb = 0.0, 0
            for i in range(0, n_rows, rec["batch"]):
                idx = perm[i:i + rec["batch"]]
                b, tg, _, _ = prep(store_tr, idx, st)
                mu, ls = model(b)
                loss = nll(mu, ls, tg)
                opt.zero_grad(set_to_none=True)
                loss.backward()
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                opt.step()
                tot += float(loss)
                nb += 1
            v = evaluate(model, store_va, st, device)
            improved = v < best - 1e-4
            if improved:
                best, best_ep, bad = v, ep, 0
                torch.save({"model": model.state_dict(), "cfg": cfg, "stats": st, "seed": a.seed, "epoch": ep}, out / f"ckpt_ep{ep:03d}.pt")
                (out / "BEST").write_text(f"ckpt_ep{ep:03d}.pt\n")
            else:
                bad += 1
            with open(out / "history.jsonl", "a") as f:
                f.write(json.dumps({"epoch": ep, "train_nll": tot / max(nb, 1), "val_nll": v, "best": best, "best_epoch": best_ep, "t_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())}) + "\n")
            log(f"epoch {ep} train_nll {tot / max(nb, 1):.4f} val_nll {v:.4f} best {best:.4f}@{best_ep}{' *' if improved else ''}", out)
            if bad >= rec["patience"]:
                log(f"early stop (patience {rec['patience']})", out)
                break
        (out / "TRAIN_DONE").write_text(time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) + "\n")
        del store_tr, store_va
    del tr, va
    z = torch.load(out / (out / "BEST").read_text().strip(), map_location=device, weights_only=False)
    st = z["stats"]
    torch.manual_seed(a.seed)
    model = RoadDynNet(a.arm, rec["hid"], rec["hidden"], rec["d_vis"], rec["dropout"]).to(device)
    model.load_state_dict(z["model"])
    for r in REGIMES:
        for s in ("test", "val"):
            raw = preds[(r, s)] if a.synthetic else load_split(r, s, a.arm)
            predict(model, raw, st, device, out / f"pred_rd_{r}_{s}.npz", r, s)
            log(f"pred_rd_{r}_{s}.npz written ({len(raw['t']):,} rows)", out)
            del raw
    (out / "PRED_DONE").write_text(time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) + "\n")
    inputs = [] if a.synthetic else [RD / f"road_dyn_{r}_{s}.npz" for r in REGIMES for s in SPLITS] + ([PACK_V2 / f"vf256_{r}_{s}.npz" for r in REGIMES for s in SPLITS] if a.arm == "dynvis" else [])
    write_run_json(out, f"train_road_dynamics_{a.arm}_s{a.seed}", [Path(__file__)], inputs, a.seed, t_start,
                   extra={"cfg": cfg, "best": (out / "BEST").read_text().strip(), "params": json.loads((out / "params.json").read_text()), "device": str(device), "dms_fields_used": []})
    log("run.json written; DONE", out)


def compute_stats_real(tr):
    return compute_stats(tr)


if __name__ == "__main__":
    main()

Xet Storage Details

Size:: 21.3 kB
Xet hash:: a11a6f2a1702da3c84ba507012139e7e7c33cdec64578d4d39f3ca3821cea4bd

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	"""cwm3/train/train_road_dynamics.py — F5c arms (3) and (4): a small causal dynamics net on the 8 s CAN road history predicting RESIDUALS over the
	constant-velocity extrapolation of the future road quantities (WP4 batch 1; pre-registration docs/decisions.md 2026-09-16T06:46:35Z; plan
	<local-notes>/plans/elegant-dazzling-stallman.md F5b/F5c).

	Targets (data/build_road_dynamics_targets.py, results/targets_road_dyn/road_dyn_{regime}_{split}.npz): Yf[N,25,4] = lead_dRel, lead_vRel, lead_vLead,
	lane_center_offset at t+0.2k (k = 1..25) with Yvalid. Arms:
	(1) persistence y_k = y_0 (eval/road_dynamics_eval.py, no training)
	(2) constant velocity (CV) dRel_k = dRel_0 + vRel_0 tau_k, vRel / vLead constant, lane_off_k = lane_off_0 + rate_0 tau_k (rate_0 = OLS slope of the last 1 s,
	clipped +-0.5 m/s) (cv_forecast below; eval, no training)
	(3) `dyn` GRU(128) over the 8 s history @ 5 Hz (8 standardised canspec channels + dt) + the 40 standardised CAN anchor features (ecan, the F5b list) + mode
	-> MLP -> residual mean / log-sd over arm (2) per (k, channel), Gaussian NLL masked by Yvalid & anchor validity. <= 0.3 M parameters (asserted).
	(4) `dynvis` arm (3) + the FINAL E branch's frozen per-anchor visual side vector vf256 (results/pack_v2, PCA-256 of the causal full-frame V-JEPA features)
	through LN -> Linear(256 -> 64) -> GELU, concatenated before the MLP.
	Recipe: AdamW 1e-3 / wd 1e-4, batch 2048, 400k rows per epoch (both regimes pooled, one model, mode embedding), <= 12 epochs, early stop on the masked
	validation NLL (patience 3), 3 seeds (paired init: torch.manual_seed(seed) before construction), selection on VALIDATION only.
	Outputs results/runs/RD_<arm>_s<k>/: config.json, stats.npz (train statistics), history.jsonl, BEST (+ ckpt), params.json, pred_rd_{regime}_{split}.npz for
	test and val = {mu f16 [N,25,4] (raw units = CV + residual), sigma f16 [N,25,4] (raw units), row_ok, uid, t}, TRAIN_DONE / PRED_DONE, run.json.
	python train/train_road_dynamics.py --arm dyn --seed 0 [--out results/runs/RD_dyn_s0] [--epochs 12] [--synthetic N] [--cpu]
	No openpilot DMS field is read; inputs are the canspec CAN road channels (+ vf256 for arm 4) only.
	"""
	from __future__ import annotations

	import argparse
	import json
	import math
	import os
	import sys
	import time
	from pathlib import Path

	os.environ.setdefault("OMP_NUM_THREADS", "4")
	import numpy as np # noqa: E402
	import torch # noqa: E402
	import torch.nn as nn # noqa: E402
	import torch.nn.functional as Fn # noqa: E402

	CWM3 = Path(__file__).resolve().parents[1]
	sys.path.insert(0, str(CWM3 / "common"))
	from runmeta import write_run_json # noqa: E402

	RD = CWM3 / "results" / "targets_road_dyn"
	PACK_V2 = CWM3 / "results" / "pack_v2"
	REGIMES, SPLITS = ("manual", "engaged"), ("train", "val", "test")
	MODE = {"manual": 0, "engaged": 1}
	HIST, H, DC, DH = 40, 25, 4, 8
	TAU = (np.arange(1, H + 1) * 0.2).astype(np.float32)
	LANE_RATE_CLIP = 0.5
	MAX_PARAMS = 300_000
	RECIPE = dict(batch=2048, rows_per_epoch=400_000, max_epochs=12, patience=3, lr=1e-3, weight_decay=1e-4, hid=128, hidden=256, d_vis=64, dropout=0.1)


	# ----------------------------------------------------------------------------- data
	def load_split(regime, split, arm, max_rows=None, seed=0):
	with np.load(RD / f"road_dyn_{regime}_{split}.npz", allow_pickle=True) as z:
	N = len(z["t"])
	sel = None
	if max_rows and max_rows < N:
	sel = np.sort(np.random.default_rng(seed).choice(N, max_rows, replace=False))
	sub = (lambda a: a[sel]) if sel is not None else (lambda a: a)
	out = {k: sub(z[k]) for k in ("H", "Hvalid", "Yf", "Yvalid", "ecan", "ecan_ok", "uid", "t", "dongle")}
	out["mode"] = np.full(len(out["t"]), MODE[regime], np.int8)
	if arm == "dynvis":
	with np.load(PACK_V2 / f"vf256_{regime}_{split}.npz") as v:
	assert int(v["n"]) == N, (regime, split, int(v["n"]), N)
	out["vf"], out["vf_ok"] = sub(v["Z"]), sub(v["row_ok"]).astype(bool)
	return out


	def concat(parts):
	return {k: np.concatenate([p[k] for p in parts]) for k in parts[0]}


	def synthetic(N, seed, arm):
	rng = np.random.default_rng(seed)
	Hh = np.zeros((N, HIST, DH), np.float32)
	lead = rng.random(N) < 0.6
	lane = rng.random(N) < 0.6
	v = rng.uniform(5, 30, N)
	Hh[:, :, 5] = v[:, None] + rng.normal(0, 0.3, (N, HIST))
	Hh[:, :, 6], Hh[:, :, 7] = lead[:, None], lane[:, None]
	vrel = rng.normal(-1, 2, N)
	Hh[:, :, 0] = np.where(lead[:, None], rng.uniform(10, 120, N)[:, None] + vrel[:, None] * (np.arange(HIST) - 39)[None, :] * 0.2, 200.0)
	Hh[:, :, 1] = np.where(lead[:, None], vrel[:, None] + rng.normal(0, 0.1, (N, HIST)), 0.0)
	Hh[:, :, 2] = np.where(lead[:, None], v[:, None] + vrel[:, None], 0.0)
	Hh[:, :, 3] = np.where(lane[:, None], rng.normal(0, 0.3, (N, 1)) + rng.normal(0, 0.03, (N, HIST)), 0.0)
	Hh[:, :, 4] = np.where(lane[:, None], 3.5, 0.0)
	Yf = np.zeros((N, H, DC), np.float32)
	Yf[:, :, 0] = Hh[:, -1, 0][:, None] + vrel[:, None] * TAU[None, :] + rng.normal(0, 1, (N, H))
	Yf[:, :, 1] = vrel[:, None] + rng.normal(0, 0.3, (N, H))
	Yf[:, :, 2] = Hh[:, -1, 2][:, None] + rng.normal(0, 0.3, (N, H))
	Yf[:, :, 3] = Hh[:, -1, 3][:, None] + rng.normal(0, 0.1, (N, H))
	Yv = np.stack([lead[:, None].repeat(H, 1)] * 3 + [lane[:, None].repeat(H, 1)], -1)
	out = {"H": Hh.astype(np.float16), "Hvalid": np.ones((N, HIST), bool), "Yf": Yf, "Yvalid": Yv, "ecan": rng.normal(0, 1, (N, 40)).astype(np.float32),
	"ecan_ok": np.ones(N, bool), "uid": np.array([f"r{i % 4}" for i in range(N)]), "t": np.arange(N) * 0.2, "dongle": np.array([f"d{i % 3}" for i in range(N)]),
	"mode": (np.arange(N) % 2).astype(np.int8)}
	if arm == "dynvis":
	out["vf"], out["vf_ok"] = rng.normal(0, 1, (N, 256)).astype(np.float16), np.ones(N, bool)
	return out


	def compute_stats(tr):
	"""train-only float64 statistics: per history channel (over valid steps), ecan (over ok rows), residual scale per target channel (sd of Yf - CV at k = 15)."""
	Hh = tr["H"].astype(np.float32)
	hv = tr["Hvalid"]
	hmu, hsd = np.zeros(DH), np.zeros(DH)
	for c in range(DH):
	x = Hh[:, :, c][hv].astype(np.float64)
	hmu[c], hsd[c] = x.mean(), x.std() + 1e-6
	ok = tr["ecan_ok"]
	e = tr["ecan"][ok].astype(np.float64)
	emu, esd = e.mean(0), e.std(0) + 1e-6
	cv = cv_forecast_np(Hh, hv)
	r = tr["Yf"] - cv
	m = tr["Yvalid"] & anchor_valid_np(Hh)[:, None, :]
	rsd = np.array([r[:, 14, c][m[:, 14, c]].astype(np.float64).std() + 1e-3 for c in range(DC)])
	return {"hmu": hmu.astype(np.float32), "hsd": hsd.astype(np.float32), "emu": emu.astype(np.float32), "esd": esd.astype(np.float32), "rsd": rsd.astype(np.float32)}


	# ----------------------------------------------------------------------------- baselines (numpy, shared with the evaluator)
	def anchor_valid_np(Hh):
	"""[N,4] lead valid at the anchor for the 3 lead channels, lane valid for the lane channel."""
	lead0, lane0 = Hh[:, -1, 6] > 0.5, Hh[:, -1, 7] > 0.5
	return np.stack([lead0, lead0, lead0, lane0], -1)


	def lane_rate_np(Hh, hv):
	"""OLS slope (m/s) of the lane offset over the last 1 s (5 history steps) using lane-valid steps; 0 with < 3 steps; clipped."""
	x = Hh[:, -5:, 3].astype(np.float64)
	w = (hv[:, -5:] & (Hh[:, -5:, 7] > 0.5)).astype(np.float64)
	dt = np.arange(-4, 1) * 0.2
	n = w.sum(1)
	sw = np.maximum(n, 1)
	xbar, tbar = (w * x).sum(1) / sw, (w * dt[None]).sum(1) / sw
	num = (w * (dt[None] - tbar[:, None]) * (x - xbar[:, None])).sum(1)
	den = (w * (dt[None] - tbar[:, None]) ** 2).sum(1)
	return np.clip(np.where((n >= 3) & (den > 1e-9), num / np.maximum(den, 1e-9), 0.0), -LANE_RATE_CLIP, LANE_RATE_CLIP).astype(np.float32)


	def persistence_np(Hh):
	return np.repeat(Hh[:, -1:, :4].astype(np.float32), H, 1)


	def cv_forecast_np(Hh, hv):
	p = persistence_np(Hh)
	p[:, :, 0] = Hh[:, -1:, 0] + Hh[:, -1:, 1] * TAU[None, :]
	p[:, :, 3] = Hh[:, -1:, 3] + lane_rate_np(Hh, hv)[:, None] * TAU[None, :]
	return p


	def cv_forecast_t(Hh, hv):
	"""torch version of cv_forecast_np (same arithmetic)."""
	B = Hh.shape[0]
	tau = torch.as_tensor(TAU, device=Hh.device)
	p = Hh[:, -1:, :4].expand(B, H, 4).clone()
	p[:, :, 0] = Hh[:, -1:, 0] + Hh[:, -1:, 1] * tau[None, :]
	x = Hh[:, -5:, 3]
	w = (hv[:, -5:] & (Hh[:, -5:, 7] > 0.5)).float()
	dt = torch.arange(-4, 1, device=Hh.device).float() * 0.2
	n = w.sum(1)
	sw = n.clamp(min=1)
	xbar, tbar = (w * x).sum(1) / sw, (w * dt[None]).sum(1) / sw
	num = (w * (dt[None] - tbar[:, None]) * (x - xbar[:, None])).sum(1)
	den = (w * (dt[None] - tbar[:, None]) ** 2).sum(1)
	rate = torch.where((n >= 3) & (den > 1e-9), num / den.clamp(min=1e-9), torch.zeros_like(num)).clamp(-LANE_RATE_CLIP, LANE_RATE_CLIP)
	p[:, :, 3] = Hh[:, -1:, 3] + rate[:, None] * tau[None, :]
	return p


	# ----------------------------------------------------------------------------- model
	class RoadDynNet(nn.Module):
	def __init__(self, arm, hid=128, hidden=256, d_vis=64, dropout=0.1, n_ecan=40):
	super().__init__()
	self.arm = arm
	self.gru = nn.GRU(DH + 1, hid, batch_first=True)
	self.ecan = nn.Sequential(nn.Linear(n_ecan, 32), nn.GELU())
	self.mode = nn.Embedding(2, 8)
	d_in = hid + 32 + 8
	if arm == "dynvis":
	self.vis = nn.Sequential(nn.LayerNorm(256), nn.Linear(256, d_vis), nn.GELU(), nn.Dropout(dropout))
	d_in += d_vis
	self.head = nn.Sequential(nn.Linear(d_in, hidden), nn.GELU(), nn.Dropout(dropout), nn.Linear(hidden, 2 * H * DC))
	nn.init.zeros_(self.head[-1].bias)
	with torch.no_grad():
	self.head[-1].weight.mul_(0.1)

	def forward(self, b):
	x = torch.cat([b["Hs"], b["dt"]], -1)
	_, h = self.gru(x)
	z = [h[-1], self.ecan(b["ecan"]), self.mode(b["mode"])]
	if self.arm == "dynvis":
	z.append(self.vis(b["vf"]) * b["vf_ok"].float()[:, None])
	o = self.head(torch.cat(z, -1)).view(-1, H, DC, 2)
	return o[..., 0], o[..., 1].clamp(-6, 4) # residual mean, log-sd (scaled units)


	def n_params(m):
	return sum(p.numel() for p in m.parameters() if p.requires_grad)


	# ----------------------------------------------------------------------------- GPU store + batch
	class Store:
	def __init__(self, raw, device):
	self.d = device
	self.N = len(raw["t"])
	self.H = torch.from_numpy(raw["H"]).to(device)
	self.Hvalid = torch.from_numpy(raw["Hvalid"]).to(device)
	self.Yf = torch.from_numpy(raw["Yf"]).to(device)
	self.Yvalid = torch.from_numpy(raw["Yvalid"]).to(device)
	self.ecan = torch.from_numpy(raw["ecan"]).to(device)
	self.ecan_ok = torch.from_numpy(raw["ecan_ok"]).to(device)
	self.mode = torch.from_numpy(raw["mode"].astype(np.int64)).to(device)
	self.vf = torch.from_numpy(raw["vf"]).to(device) if "vf" in raw else None
	self.vf_ok = torch.from_numpy(raw["vf_ok"]).to(device) if "vf_ok" in raw else None


	def prep(store, idx, st, need_targets=True):
	d = store.d
	Hh = store.H[idx].float()
	hv = store.Hvalid[idx]
	hmu, hsd = torch.as_tensor(st["hmu"], device=d), torch.as_tensor(st["hsd"], device=d)
	Hs = torch.clamp((Hh - hmu) / hsd, -8, 8) * hv[..., None].float()
	dt = (-(HIST - 1 - torch.arange(HIST, device=d).float()) * 0.2)[None, :, None].expand(len(idx), -1, -1)
	emu, esd = torch.as_tensor(st["emu"], device=d), torch.as_tensor(st["esd"], device=d)
	ecan = torch.clamp((store.ecan[idx] - emu) / esd, -8, 8) * store.ecan_ok[idx].float()[:, None]
	b = {"Hs": Hs, "dt": dt, "ecan": ecan, "mode": store.mode[idx]}
	if store.vf is not None:
	b["vf"], b["vf_ok"] = store.vf[idx].float(), store.vf_ok[idx]
	cv = cv_forecast_t(Hh, hv)
	lead0, lane0 = Hh[:, -1, 6] > 0.5, Hh[:, -1, 7] > 0.5
	av = torch.stack([lead0, lead0, lead0, lane0], -1)
	tg = None
	if need_targets:
	rsd = torch.as_tensor(st["rsd"], device=d)
	r = (store.Yf[idx] - cv) / rsd
	tg = {"r": r, "m": store.Yvalid[idx] & av[:, None, :]}
	return b, tg, cv, av


	def nll(mu, logsd, tg):
	ll = 0.5 * ((tg["r"] - mu) / logsd.exp()) ** 2 + logsd + 0.5 * math.log(2 * math.pi)
	m = tg["m"].float()
	return (ll * m).sum() / m.sum().clamp(min=1)


	def log(msg, out):
	line = f"[{time.strftime('%H:%M:%S')}] {msg}"
	print(line, flush=True)
	with open(Path(out) / "train.log", "a") as f:
	f.write(line + "\n")


	# ----------------------------------------------------------------------------- train / predict
	def evaluate(model, store, st, device, bs=8192):
	model.eval()
	tot, cnt = 0.0, 0.0
	with torch.no_grad():
	for i in range(0, store.N, bs):
	idx = torch.arange(i, min(i + bs, store.N), device=device)
	b, tg, _, _ = prep(store, idx, st)
	mu, ls = model(b)
	m = tg["m"].float()
	ll = 0.5 * ((tg["r"] - mu) / ls.exp()) ** 2 + ls + 0.5 * math.log(2 * math.pi)
	tot += float((ll * m).sum())
	cnt += float(m.sum())
	return tot / max(cnt, 1.0)


	def predict(model, raw, st, device, out_path, regime, split, bs=8192):
	store = Store(raw, device)
	model.eval()
	mu_all, sd_all, ok_all = [], [], []
	rsd = torch.as_tensor(st["rsd"], device=device)
	with torch.no_grad():
	for i in range(0, store.N, bs):
	idx = torch.arange(i, min(i + bs, store.N), device=device)
	b, _, cv, av = prep(store, idx, st, need_targets=False)
	mu, ls = model(b)
	mu_all.append((cv + mu * rsd).half().cpu().numpy())
	sd_all.append((ls.exp() * rsd).half().cpu().numpy())
	ok_all.append(av.cpu().numpy())
	np.savez_compressed(out_path, mu=np.concatenate(mu_all), sigma=np.concatenate(sd_all), row_ok=np.concatenate(ok_all), uid=raw["uid"], t=raw["t"],
	channels=np.array(["lead_dRel", "lead_vRel", "lead_vLead", "lane_center_offset"]), regime=regime, split=split,
	note="mu / sigma in raw units (m, m/s, m/s, m); mu = constant-velocity forecast + predicted residual; row_ok[N,4] = channel valid at the anchor")
	del store
	torch.cuda.empty_cache() if device.type == "cuda" else None


	def main():
	ap = argparse.ArgumentParser()
	ap.add_argument("--arm", required=True, choices=["dyn", "dynvis"])
	ap.add_argument("--seed", type=int, default=0)
	ap.add_argument("--out", default="")
	ap.add_argument("--epochs", type=int, default=0)
	ap.add_argument("--max-rows", type=int, default=0, help="seeded subset of TRAIN rows (smoke)")
	ap.add_argument("--synthetic", type=int, default=0)
	ap.add_argument("--cpu", action="store_true")
	ap.add_argument("--no-lock", action="store_true", help="accepted for queue compatibility (this trainer never takes the GPU lock)")
	ap.add_argument("--minutes", type=float, default=0, help="accepted for queue compatibility")
	ap.add_argument("--full-pass", action="store_true", help="POST-HOC sensitivity arm `dynfull` (labelled as such; not the pre-registered recipe): one full pass over the train rows per epoch, <= 30 epochs, same patience")
	a = ap.parse_args()
	out = Path(a.out or (CWM3 / "results" / "runs" / f"RD_{a.arm}{'full' if a.full_pass else ''}_s{a.seed}{'_synthetic' if a.synthetic else ''}"))
	out.mkdir(parents=True, exist_ok=True)
	t_start = time.time()
	device = torch.device("cpu" if a.cpu or not torch.cuda.is_available() else "cuda")
	rec = dict(RECIPE)
	if a.full_pass:
	rec["rows_per_epoch"], rec["max_epochs"], rec["post_hoc"] = 0, 30, "full-pass sensitivity arm (2026-09-16; the pre-registered 12-epoch / 400k-row arm was still improving at its cap)"
	if a.epochs:
	rec["max_epochs"] = a.epochs
	cfg = {"arm": a.arm, "seed": a.seed, "recipe": rec, "regimes": list(REGIMES), "targets": "results/targets_road_dyn (Yf lead_dRel / lead_vRel / lead_vLead / lane_center_offset, k=1..25)",
	"baseline": "constant velocity (dRel += vRel tau; lane offset linear with the 1 s OLS rate clipped +-0.5 m/s); the net predicts residuals", "dms_fields_used": []}
	(out / "config.json").write_text(json.dumps({"cfg": cfg, "args": vars(a)}, indent=1))
	log(f"=== train_road_dynamics arm={a.arm} s{a.seed} out={out} device={device} ===", out)
	if a.synthetic:
	tr = synthetic(a.synthetic, a.seed, a.arm)
	va = synthetic(max(a.synthetic // 4, 64), a.seed + 1, a.arm)
	preds = {(r, s): synthetic(max(a.synthetic // 4, 64), a.seed + 2 + i, a.arm) for i, (r, s) in enumerate([(r, s) for r in REGIMES for s in ("test", "val")])}
	else:
	tr = concat([load_split(r, "train", a.arm, a.max_rows or None, a.seed) for r in REGIMES])
	va = concat([load_split(r, "val", a.arm) for r in REGIMES])
	preds = None
	if not (out / "TRAIN_DONE").exists():
	st = compute_stats(tr) if a.synthetic else compute_stats_real(tr)
	np.savez(out / "stats.npz", **st)
	torch.manual_seed(a.seed)
	model = RoadDynNet(a.arm, rec["hid"], rec["hidden"], rec["d_vis"], rec["dropout"]).to(device)
	npar = n_params(model)
	(out / "params.json").write_text(json.dumps({"trainable": npar, "limit": MAX_PARAMS, "by_module": {n: sum(p.numel() for p in m.parameters()) for n, m in model.named_children()}}, indent=1))
	assert npar <= MAX_PARAMS, npar
	log(f"params {npar:,} (<= {MAX_PARAMS:,}); train rows {len(tr['t']):,} val rows {len(va['t']):,}", out)
	store_tr, store_va = Store(tr, device), Store(va, device)
	opt = torch.optim.AdamW(model.parameters(), lr=rec["lr"], weight_decay=rec["weight_decay"])
	g = torch.Generator(device="cpu").manual_seed(a.seed)
	best, best_ep, bad = float("inf"), -1, 0
	for ep in range(rec["max_epochs"]):
	model.train()
	n_rows = min(rec["rows_per_epoch"], store_tr.N) if rec["rows_per_epoch"] else store_tr.N
	perm = torch.randperm(store_tr.N, generator=g)[:n_rows].to(device)
	tot, nb = 0.0, 0
	for i in range(0, n_rows, rec["batch"]):
	idx = perm[i:i + rec["batch"]]
	b, tg, _, _ = prep(store_tr, idx, st)
	mu, ls = model(b)
	loss = nll(mu, ls, tg)
	opt.zero_grad(set_to_none=True)
	loss.backward()
	torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
	opt.step()
	tot += float(loss)
	nb += 1
	v = evaluate(model, store_va, st, device)
	improved = v < best - 1e-4
	if improved:
	best, best_ep, bad = v, ep, 0
	torch.save({"model": model.state_dict(), "cfg": cfg, "stats": st, "seed": a.seed, "epoch": ep}, out / f"ckpt_ep{ep:03d}.pt")
	(out / "BEST").write_text(f"ckpt_ep{ep:03d}.pt\n")
	else:
	bad += 1
	with open(out / "history.jsonl", "a") as f:
	f.write(json.dumps({"epoch": ep, "train_nll": tot / max(nb, 1), "val_nll": v, "best": best, "best_epoch": best_ep, "t_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())}) + "\n")
	log(f"epoch {ep} train_nll {tot / max(nb, 1):.4f} val_nll {v:.4f} best {best:.4f}@{best_ep}{' *' if improved else ''}", out)
	if bad >= rec["patience"]:
	log(f"early stop (patience {rec['patience']})", out)
	break
	(out / "TRAIN_DONE").write_text(time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) + "\n")
	del store_tr, store_va
	del tr, va
	z = torch.load(out / (out / "BEST").read_text().strip(), map_location=device, weights_only=False)
	st = z["stats"]
	torch.manual_seed(a.seed)
	model = RoadDynNet(a.arm, rec["hid"], rec["hidden"], rec["d_vis"], rec["dropout"]).to(device)
	model.load_state_dict(z["model"])
	for r in REGIMES:
	for s in ("test", "val"):
	raw = preds[(r, s)] if a.synthetic else load_split(r, s, a.arm)
	predict(model, raw, st, device, out / f"pred_rd_{r}_{s}.npz", r, s)
	log(f"pred_rd_{r}_{s}.npz written ({len(raw['t']):,} rows)", out)
	del raw
	(out / "PRED_DONE").write_text(time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) + "\n")
	inputs = [] if a.synthetic else [RD / f"road_dyn_{r}_{s}.npz" for r in REGIMES for s in SPLITS] + ([PACK_V2 / f"vf256_{r}_{s}.npz" for r in REGIMES for s in SPLITS] if a.arm == "dynvis" else [])
	write_run_json(out, f"train_road_dynamics_{a.arm}_s{a.seed}", [Path(__file__)], inputs, a.seed, t_start,
	extra={"cfg": cfg, "best": (out / "BEST").read_text().strip(), "params": json.loads((out / "params.json").read_text()), "device": str(device), "dms_fields_used": []})
	log("run.json written; DONE", out)


	def compute_stats_real(tr):
	return compute_stats(tr)


	if __name__ == "__main__":
	main()