Buckets:

AsianPlayer
/

TriDrive

Files
xet
AsianPlayer/TriDrive / code /train /train_cwm3_joint.py
AsianPlayer
16 days ago
download
raw
77.1 kB
"""cwm3 joint-model stage trainer (I1 interactions / F1 decoders; plan §3, claims B2 / B3) — both regimes in ONE model.
  python train_cwm3_joint.py --stage I1 --variant all --seed 0 [--base-v can4s|can8s] [--max-anchors N] [--synthetic N] [--minutes 18] [--out DIR]
Assembles models/cwm3.CWM3Joint from the seed-k BEST checkpoints of the three base stages (train/stages_joint.yaml `base`; the V base is a
CONFIG FIELD — V0/can4s or V0/can8s — never hard-coded), adds the zero-initialised interaction modules of the variant, the detached K head and
the decoder kind, and fine-tunes: base branches at lr * lr_base_mult (0.1; 0 = frozen), new modules at lr.  Standardisation statistics are the
base checkpoints' (train-only), never recomputed.  Selection objective (BEST / early stop): validation sum of the chain NLLs of all three
branches + 0.1 L_F (models/heads.selection_objective over the merged parts); V terms are masked to rows with the V history available.
Execution: `--no-lock` (default in train/run_queue_joint.sh since the owner lifted the lock requirement on 2026-09-10 13:58) runs training and
prediction INLINE in one process — no fork, no lock; the queue script gates launches on free VRAM / RAM / the number of live trainers.  Lock mode
is train/train_cwm3.py's (imported, unchanged): the parent loads all data BEFORE the lock and never touches CUDA — nor any
torch op (a CPU model build before os.fork() left the child deadlocked on the OMP pool for 31 min on 2026-09-10; base cfg / stats now come from
the base runs' config.json / stats.npz, label rows from a subprocess); flock(/tmp/s13_gpu.lock) -> fork -> child trains until the wall budget
(<= 22 min) / epoch budget / early stop -> exit 0 | 3; the same child predicts when >= 4 min of the chunk are left.  Watchdog: faulthandler
dumps every thread's stack and exits (releasing the lock) if the first batch has not completed 5 min after the lock, then guards the chunk.
E base is a config field too (E0/pooled default; `--base-e`); with `tok_source: region` the region token bank feeds ONLY the road -> driver
cross-attention (RegionTokenEmbed), the E chains / F stay those of the E base.
The `indep` variant (`no_train`) skips training: the assembled base branches predict as-is — the exact independent-branch comparator whose
ep_p / w blocks are composed identically to the trained variants.  4-thread CPU cap; F targets stay in host RAM (gathered per batch).
Detached auxiliary heads (no trunk gradient): P road-pressure chains (results/targets_pressure, 6 families) on the E trunk and the HARD response
chain R >= 2 (results/targets_response) on the V trunk — extra outputs for the warning head; their NLL (pchain_* / rchain_*) is excluded from the
selection objective.  Outputs under results/runs/{stage}_{variant}_s{seed}[_v8 when --base-v overrides the yaml]: ckpt_*.pt, BEST, history.jsonl, TRAIN_DONE,
stats.npz, pred_{regime}_{split}.npz (hazards of all 12 chains, K̂, per-row means), pred_layout_{regime}_{split}.npz
(eval/write_pred_layout_joint.py: onset / ep / w / nr blocks), pred_{regime}_train_labelrows.npz (rows the warning head is trained on;
data/labels.py), metrics_quick_*.json, run_*.json, PRED_DONE; rows appended to results/cwm3_quick_metrics.csv.
No openpilot DMS field is read anywhere (inputs: CAN whitelist, YOLO-pose kinematics, V-JEPA2 latents).
"""
from __future__ import annotations

import os

os.environ.setdefault("OMP_NUM_THREADS", "4")
os.environ.setdefault("MKL_NUM_THREADS", "4")
import argparse                                             # noqa: E402
import csv                                                  # noqa: E402
import json                                                 # noqa: E402
import shutil                                               # noqa: E402
import sys                                                  # noqa: E402
import time                                                 # noqa: E402
from pathlib import Path                                    # noqa: E402

import numpy as np                                          # noqa: E402
import yaml                                                 # noqa: E402

CWM3 = Path(__file__).resolve().parents[1]
EXPLORE = CWM3.parent
for p in (CWM3, CWM3 / "data", CWM3 / "common", CWM3 / "train", CWM3 / "eval", EXPLORE / "s12_wm"):
    sys.path.insert(0, str(p))
import cwm3_dataset as ds                                   # noqa: E402
import batch as bt                                          # noqa: E402
import train_cwm3 as T                                      # noqa: E402  (lock loop, checkpoint helpers, log — unchanged)
from models import heads as hd                              # noqa: E402
from models.cwm3 import build_joint, joint_loss, P_CHAINS, R_CHAINS, base_cfg_summary, branch_stats, merge_stats, stats_per_branch   # noqa: E402
from runmeta import write_run_json                          # noqa: E402
import write_pred_layout_joint as wplj                      # noqa: E402
from mtl import MTL, KINDS as MTL_KINDS                     # noqa: E402  (WP3 MTL block 2026-09-16: --mtl none|uw|famo|pcgrad; none = the original path)
from models import latent_aux as LA                         # noqa: E402  (R2 WP4 batch 2 C4 2026-09-17: variant key `latent_aux`; lam 0 -> nothing is built, the original path)

STAGES = yaml.safe_load((CWM3 / "train" / "stages_joint.yaml").read_text())
CHAINS = list(hd.V_CHAINS) + list(hd.U_CHAINS) + list(hd.E_CHAINS)
AUX_KEYS = ["P", "Pvalid", "P0", "P0valid", "Rz", "Rvalid", "R0z", "R0valid"]     # detached auxiliary chains (results/targets_{pressure,response})
LAT_KEYS = ["LYf", "LYv"]                                   # C4 future road quantities at the latent horizons (results/targets_road_dyn Yf / Yvalid sliced to [N,K,4]); present only when latent_enabled
KEYS = sorted(set(bt.KEYS["V"]) | set(bt.KEYS["U"]) | set(bt.KEYS["E"]) | set(AUX_KEYS) | set(LAT_KEYS))
HOST_KEYS = ("F",)                                          # largest stream (2.5 KB/row): host RAM, gathered per batch
PRESSURE = CWM3 / "results" / "targets_pressure"
RESPONSE = CWM3 / "results" / "targets_response"
ROAD_DYN = CWM3 / "results" / "targets_road_dyn"
log = T.log


def latent_enabled(cfg):
    """C4 (2026-09-17): the variant's `latent_aux` dict with lam > 0.  lam == 0 (the pre-registered control) or no key -> False: the module is never built, no target is
    attached, no loss term is added — the parent variant's code path, bit for bit (tests/test_latent_aux.py)."""
    la = cfg.get("latent_aux") or {}
    return bool(la) and float(la.get("lam", 0.0)) > 0.0


def latent_lam(cfg):
    return float((cfg.get("latent_aux") or {}).get("lam", 0.0)) if latent_enabled(cfg) else 0.0


# ----------------------------------------------------------------------------- config
def stage_cfg(stage, variant, base_v=None, base_e=None, base_u=None):
    s = STAGES["stages"][stage]
    cfg = dict(STAGES["defaults"])
    cfg.update({"stage": stage, "variant": variant, "recipe": dict(s["recipe"]), "lambda": dict(s.get("lambda", {})), "decoder": s.get("decoder", "direct"),
                "experiment": s.get("experiment", ""), "interactions": [], "naive": False, "no_train": False})
    cfg.setdefault("tok_source", "region")
    v = dict(s["variants"][variant])
    if "recipe" in v:
        cfg["recipe"].update(v.pop("recipe"))
    cfg.update(v)
    cfg["base"] = {k: list(x) for k, x in cfg["base"].items()}
    if base_v:
        cfg["base"]["V"] = ["V0", base_v]
    if base_e:
        cfg["base"]["E"] = ["E0", base_e]
    if base_u:
        cfg["base"]["U"] = ["U0", base_u]
    return cfg


def resolve_run_dir(x):
    """--base-dir-* value: an absolute path, a path relative to cwm3/, or a bare run-dir name under results/runs (2026-09-11)."""
    q = Path(str(x))
    if q.is_absolute():
        return q
    if (CWM3 / q).exists():
        return CWM3 / q
    return CWM3 / "results" / "runs" / q.name


def base_dirs(cfg, seed):
    """{branch: run dir}: cfg["base_dirs"][branch] (explicit, --base-dir-{v,u,e}; 2026-09-11 final assembly) else <stage>_<variant>_s<seed>[suffix]."""
    suf = {"E": cfg.get("base_suffix_e", "")}
    explicit = cfg.get("base_dirs", {}) or {}
    return {b: (resolve_run_dir(explicit[b]) if explicit.get(b) else CWM3 / "results" / "runs" / f"{st}_{va}_s{seed}{suf.get(b, '')}")
            for b, (st, va) in cfg["base"].items()}


def need_streams(cfg):
    need = {"X", "Y", "demand", "Kh", "kin", "F", "vf256", "tokens"}
    if cfg["base"]["V"][1] == "can8s":
        need.add("X8")
    if cfg.get("e_side", "vf256") != "vf256":                      # the E base's side vector (regionpca256, pack v2 only); set by main from the E base config
        need.add(cfg["e_side"])
    if cfg.get("u_gaze_chain"):                                     # gaze U bases (2026-09-12 appendix variant): Sg / Sg0 / EMg chain targets; set by main from the U base config
        need.add("gaze")
    if cfg.get("u_face"):                                           # Fh / Fhvalid causal face-landmark inputs of the U base
        need.add("face")
    if cfg.get("u_kpf"):                                            # LANE F1 (2026-09-17): the detached keypoint-forecaster channel (--u-kpf <KPF run dir>)
        need.add("kpf")
    return frozenset(need)


def prep_cfgs(info):
    """per-branch batch.py configs from the base runs' configs (never hard-coded); "R" = the region token bank for the cross-attention.
    2026-09-11: the U cfg carries the base's `kclip` (batch.k_clip -> 8 for the clip8 recipes; absent -> legacy 5, old runs reproduce) and the E cfg
    its `side` vector name (regionpca256 bases)."""
    bc = info["base_cfgs"]
    u = {"cabin": bc["U"].get("cabin", "none")}
    if "kclip" in bc["U"]:
        u["kclip"] = float(bc["U"]["kclip"])
    rec_u = {k: bool(bc["U"][k]) for k in ("u_face", "u_gaze_chain") if k in bc["U"]}
    if rec_u:                                                                   # gaze U bases (2026-09-12): batch.prep_U appends the face channels / builds the gaze targets
        u["recipe"] = rec_u
    e = {"tokens": bc["E"].get("tokens", "region"), "n_slots": int(bc["E"].get("n_slots", 16))}
    if "side" in bc["E"]:
        e["side"] = bc["E"]["side"]
    return {"V": {"can_hist": bc["V"].get("can_hist", "4s")}, "U": u, "E": e, "R": {"tokens": "region", "n_slots": 16}}


def base_info_numpy(cfg, seed):
    """PARENT-side assembly info WITHOUT torch (a torch op before os.fork() left the child deadlocked on the OMP pool, 2026-09-10): per base
    run dir -> config.json (cfg), stats.npz (train-only statistics, merged with equality asserted on overlapping keys), BEST name.
    2026-09-11: base_cfgs = models.cwm3.base_cfg_summary (kclip / u_in_ln / kin_source / kin_version / stats_version / p_head / select_LF / side);
    cfg["stats_mode"] per_branch keeps every base's own statistics ('<branch>/<key>', models.cwm3.stats_per_branch) instead of merging."""
    names, cfgs, stats, dirs = {}, {}, [], {}
    for b, d in base_dirs(cfg, seed).items():
        assert (d / "BEST").exists() and (d / "PRED_DONE").exists(), f"base checkpoint for {b} missing or not finished: {d}"
        c = json.loads((d / "config.json").read_text())["cfg"]
        names[b] = f"{d.name}:{(d / 'BEST').read_text().strip()}"
        cfgs[b] = base_cfg_summary(c)
        st = {k: v for k, v in np.load(d / "stats.npz").items()}
        if "stats_version" in st and "stats_version" not in cfgs[b]:
            cfgs[b]["stats_version"] = int(st["stats_version"])
        stats.append(st)
        dirs[b] = str(d)
    mode = cfg.get("stats_mode", "merged")
    merged = merge_stats(stats) if mode == "merged" else stats_per_branch(stats, tuple(base_dirs(cfg, seed)))
    return {"base_ckpts": names, "base_cfgs": cfgs, "stats": merged, "stats_mode": mode, "base_dirs": dirs}


# ----------------------------------------------------------------------------- data
class HostTensor:
    """numpy array kept in host RAM; `[idx]` with a device index tensor gathers the rows and moves them to the device."""

    def __init__(self, a, device):
        self.a, self.device = a, device

    def __getitem__(self, idx):
        import torch
        i = idx.cpu().numpy() if torch.is_tensor(idx) else np.asarray(idx)
        return torch.from_numpy(np.ascontiguousarray(self.a[i])).to(self.device)

    def numel(self):
        return 0

    def element_size(self):
        return 0


def make_store(raw, device, bank):
    store = bt.GpuStore(raw, [k for k in KEYS if k not in HOST_KEYS], device, bank)
    for k in HOST_KEYS:
        if k in raw:
            store.t[k] = HostTensor(raw[k], device)
    return store


def v_valid_np(raw, pc):
    if pc["V"].get("mask_V", False):
        return np.zeros(len(raw["t"]), bool)                       # variant `ue_noV` (2026-09-18): the vehicle branch is a MISSING modality on every row
    return raw["x8_ok"].astype(bool) if pc["V"].get("can_hist") == "8s" else np.ones(len(raw["t"]), bool)


def apply_mask_v(pc, cfg):
    """variant key `mask_V` (2026-09-18, pre-registration docs/decisions.md 2026-09-18T03:34:28Z): mark the V batch config so that prep_joint / v_valid_np treat the
    vehicle branch as missing (X zeroed, V-valid mask all False, valid_V = 0 in every pack).  Absent / false -> pc unchanged (every older run bit for bit)."""
    if cfg.get("mask_V", False):
        pc["V"]["mask_V"] = True
    return pc


def aux_stats(tr):
    """train base rates of the auxiliary anchor states (fill for invalid anchors), as batch.compute_stats does for resp0 / d0."""
    st = {}
    if "P0" in tr:
        st["p0_base"] = np.array([tr["P0"][tr["P0valid"][:, j], j].mean() if tr["P0valid"][:, j].any() else 0.0 for j in range(6)], np.float32)
    if "R0z" in tr:
        v = tr["R0valid"][:, 0]
        st["r0_base"] = np.array([tr["R0z"][v, 0].mean() if v.any() else 0.05], np.float32)
    return st


def aux_stats_from_files(regimes):
    """the same base rates read directly from the TRAIN pressure / response files (both regimes pooled; all rows) — no full-split load."""
    P0, P0v, R0, R0v = [], [], [], []
    for r in regimes:
        pf, rf = PRESSURE / f"pressure_{r}_train.npz", RESPONSE / f"response_{r}_train.npz"
        if not (pf.exists() and rf.exists()):
            return {}
        with np.load(pf, allow_pickle=True) as z:
            P0.append(z["P0"].astype(bool)); P0v.append(z["P0valid"].astype(bool))
        with np.load(rf, allow_pickle=True) as z:
            R0.append((z["R0"] >= 2)[:, None]); R0v.append(z["R0valid"].astype(bool)[:, None])
    return aux_stats({"P0": np.concatenate(P0), "P0valid": np.concatenate(P0v), "R0z": np.concatenate(R0), "R0valid": np.concatenate(R0v)})


def prep_aux(store, idx, st, M, key, valid_key, k0, k0valid, base_key):
    """2-state chain targets for a detached auxiliary head: z, z_prev (anchor state, base rate where invalid), mask (valid & M & prev valid)."""
    import torch
    d = store.device
    z = store.t[key][idx].float()
    zv = store.t[valid_key][idx]
    z0v = store.t[k0valid][idx]
    base = torch.as_tensor(st[base_key], device=d)
    z0 = torch.where(z0v, store.t[k0][idx].float(), base[None, :].expand(len(idx), -1))
    zp = torch.cat([z0[:, None], z[:, :-1]], 1)
    m = zv & M[:, :, None] & torch.cat([z0v[:, None], zv[:, :-1]], 1)
    return {"z": z, "z_prev": zp, "m": m, "z0": z0, "z0valid": z0v}


def prep_joint(store, idx, st, pc, need_targets=True, region_keys=False):
    import torch
    bV, tV = bt.prep_V(store, idx, branch_stats(st, "V"), pc["V"], need_targets)
    bU, tU = bt.prep_U(store, idx, branch_stats(st, "U"), pc["U"], need_targets)
    bE, tE = bt.prep_E(store, idx, branch_stats(st, "E"), pc["E"], need_targets)
    bR = bt.prep_E(store, idx, branch_stats(st, "E"), pc["R"], False)[0] if region_keys else None
    if pc["V"].get("mask_V", False):
        # missing vehicle modality (variant `ue_noV`, 2026-09-18): the standardised CAN history is zeroed and NO row counts as V-valid -> the V chain / MDN losses,
        # the detached R-head loss and the V selection terms are masked out; the model's trunk_V returns exact zeros regardless of X (models/cwm3.py)
        bV["X"] = torch.zeros_like(bV["X"])
        vv = torch.zeros(len(idx), dtype=torch.bool, device=store.device)
    elif pc["V"].get("can_hist") == "8s":
        vv = store.t["x8_ok"][idx]
    else:
        vv = torch.ones(len(idx), dtype=torch.bool, device=store.device)
    tg = {"V": tV, "U": tU, "E": tE}
    b = {"V": bV, "U": bU, "E": bE}
    if bR is not None:
        b["R"] = {k: bR[k] for k in ("tok", "tok_ok", "tok_dt")}
    if need_targets and store.has("LYf") and "lat_mu" in st:
        # C4 latent targets (2026-09-17): future road quantities at the horizons + P occupancy + driver-state occupancy, validity-masked; only when latent_enabled
        tg_lat = LA.build_targets(store.t["LYf"][idx], store.t["LYv"][idx], store.t["P"][idx], store.t["Pvalid"][idx], store.t["S"][idx].float(), store.t["M"][idx],
                                  st["lat_mu"], st["lat_sd"], horizons=tuple(int(k) for k in np.asarray(st["lat_horizons"]).tolist()))
    else:
        tg_lat = None
    if need_targets:
        tV["M"] = tV["M"] & vv[:, None]
        tV["zVm"] = tV["zVm"] & vv[:, None, None]
        M = store.t["M"][idx]
        if store.has("P") and "p0_base" in st:
            tg["P"] = prep_aux(store, idx, st, M, "P", "Pvalid", "P0", "P0valid", "p0_base")
        if store.has("Rz") and "r0_base" in st:
            tg["R"] = prep_aux(store, idx, st, M & vv[:, None], "Rz", "Rvalid", "R0z", "R0valid", "r0_base")
        if tg_lat is not None:
            tg["LAT"] = tg_lat
    return b, tg, vv


def load_raw(args, cfg, split, subset=True, regimes=None):
    regimes = regimes or cfg["regimes"]
    need = need_streams(cfg)
    if args.synthetic:
        n = {"train": args.synthetic, "val": max(args.synthetic // 4, 64), "test": max(args.synthetic // 2, 64)}[split]
        parts = [ds.synthetic_raw(n, seed=args.seed + 10 * i + {"train": 0, "val": 1, "test": 2}[split], mode=ds.MODE[r], need=need)
                 for i, r in enumerate(regimes)]
        raw = synthetic_aux(ds.concat_regimes(parts) if len(parts) > 1 else dict(parts[0], regime_row=parts[0]["idx"]), args.seed)
        return synthetic_latent(raw, cfg, args.seed) if latent_enabled(cfg) else raw
    mx = (args.max_anchors or None) if subset else None
    raw = ds.load_regimes(regimes, split, need, max_anchors=mx, seed=args.seed + 7, verbose=lambda m: log(m, args.out))
    attach_aux_targets(raw, regimes, split)
    if latent_enabled(cfg):
        attach_latent_targets(raw, regimes, split, latent_horizons(cfg))
    return raw


def latent_horizons(cfg):
    return tuple(int(k) for k in (cfg.get("latent_aux") or {}).get("horizons", LA.HORIZONS))


def attach_latent_targets(raw, regimes, split, horizons=LA.HORIZONS):
    """C4 (2026-09-17): row-aligned FUTURE road quantities of results/targets_road_dyn/road_dyn_{regime}_{split}.npz (Yf [N,25,4] lead_dRel / lead_vRel / lead_vLead /
    lane_center_offset, Yvalid) for the rows of `raw`, sliced to the latent horizons -> raw["LYf"] [N,K,4] f32, raw["LYv"] [N,K,4] bool (uid / t witness asserted, as
    attach_aux_targets).  Strict future: Yf at index k-1 is the quantity at t + 0.2 k s (data/build_road_dynamics_targets.py contract, tests/test_road_dynamics_targets.py).
    Missing file -> assertion (the lane requires the targets; never silently skipped)."""
    idx = np.asarray([k - 1 for k in horizons], np.int64)
    Yf, Yv = [], []
    for r in regimes:
        sel = raw["regime_id"] == ds.MODE[r]
        rows = raw["regime_row"][sel]
        f = ROAD_DYN / f"road_dyn_{r}_{split}.npz"
        assert f.exists(), f"C4 latent targets need {f}"
        with np.load(f, allow_pickle=True) as z:
            assert np.array_equal(z["uid"][rows], raw["uid"][sel]) and np.allclose(z["t"][rows], raw["t"][sel]), f"{f}: uid/t witness"
            assert [str(c) for c in z["target_channels"]] == list(LA.CONT[:4]), f"{f}: target_channels {list(z['target_channels'])}"
            Yf.append(np.ascontiguousarray(z["Yf"][rows][:, idx, :].astype(np.float32)))
            Yv.append(np.ascontiguousarray(z["Yvalid"][rows][:, idx, :].astype(bool)))
    raw["LYf"], raw["LYv"] = np.concatenate(Yf), np.concatenate(Yv)
    return raw


def latent_stats_from_files(regimes, horizons=LA.HORIZONS):
    """TRAIN mean / sd of the 5 continuous latent channels (both regimes pooled, valid entries at the horizons, float64) -> lat_mu / lat_sd / lat_horizons."""
    idx = np.asarray([k - 1 for k in horizons], np.int64)
    Yf, Yv = [], []
    for r in regimes:
        with np.load(ROAD_DYN / f"road_dyn_{r}_train.npz", allow_pickle=True) as z:
            Yf.append(z["Yf"][:, idx, :].astype(np.float32)); Yv.append(z["Yvalid"][:, idx, :].astype(bool))
    st = LA.latent_stats(np.concatenate(Yf), np.concatenate(Yv))
    st["lat_horizons"] = np.asarray(horizons, np.int64)
    return st


def synthetic_latent(raw, cfg, seed=0):
    rng = np.random.default_rng(seed + 5)
    N, K = len(raw["t"]), len(latent_horizons(cfg))
    raw["LYf"] = np.stack([rng.uniform(5, 120, (N, K)), rng.normal(0, 3, (N, K)), rng.uniform(0, 30, (N, K)), rng.normal(0, 0.4, (N, K))], -1).astype(np.float32)
    raw["LYv"] = rng.random((N, K, 4)) < 0.8
    return raw


def attach_aux_targets(raw, regimes, split):
    """row-aligned road-pressure (P) and graded-response (R) targets for the rows of `raw` (per regime: regime_row), uid/t witness asserted.
    Rz = hard response (R >= 2) per step, R0z at the anchor.  Missing files -> the auxiliary heads are skipped (keys absent)."""
    parts = {k: [] for k in AUX_KEYS}
    for r in regimes:
        rows = raw["regime_row"][raw["regime_id"] == ds.MODE[r]]
        pf, rf = PRESSURE / f"pressure_{r}_{split}.npz", RESPONSE / f"response_{r}_{split}.npz"
        if not (pf.exists() and rf.exists()):
            return raw
        with np.load(pf, allow_pickle=True) as z:
            assert np.array_equal(z["uid"][rows], raw["uid"][raw["regime_id"] == ds.MODE[r]]) and np.allclose(z["t"][rows], raw["t"][raw["regime_id"] == ds.MODE[r]]), f"{pf}: uid/t witness"
            assert [str(f) for f in z["families"]] == list(P_CHAINS), f"{pf}: families {list(z['families'])}"
            parts["P"].append(z["P"][rows].astype(bool)); parts["Pvalid"].append(z["Pvalid"][rows].astype(bool))
            parts["P0"].append(z["P0"][rows].astype(bool)); parts["P0valid"].append(z["P0valid"][rows].astype(bool))
        with np.load(rf, allow_pickle=True) as z:
            assert np.array_equal(z["uid"][rows], raw["uid"][raw["regime_id"] == ds.MODE[r]]), f"{rf}: uid witness"
            parts["Rz"].append((z["R"][rows] >= 2)[:, :, None]); parts["Rvalid"].append(z["Rvalid"][rows].astype(bool)[:, :, None])
            parts["R0z"].append((z["R0"][rows] >= 2)[:, None]); parts["R0valid"].append(z["R0valid"][rows].astype(bool)[:, None])
    for k, v in parts.items():
        raw[k] = np.concatenate(v)
    return raw


def synthetic_aux(raw, seed=0):
    rng = np.random.default_rng(seed)
    N = len(raw["t"])
    raw["P"], raw["Pvalid"] = rng.random((N, 25, 6)) < 0.03, rng.random((N, 25, 6)) < 0.9
    raw["P0"], raw["P0valid"] = rng.random((N, 6)) < 0.03, rng.random((N, 6)) < 0.9
    raw["Rz"], raw["Rvalid"] = rng.random((N, 25, 1)) < 0.05, rng.random((N, 25, 1)) < 0.95
    raw["R0z"], raw["R0valid"] = rng.random((N, 1)) < 0.05, rng.random((N, 1)) < 0.95
    return raw


def load_bank(args, cfg, info):
    if args.synthetic:
        return ds.synthetic_bank()
    e_tok = info["base_cfgs"]["E"].get("tokens", "region")
    need_region = e_tok == "region" or (cfg.get("tok_source", "region") == "region" and "road2drv" in cfg.get("interactions", []))
    return ds.load_bank(need_tok=need_region, need_pooled=(e_tok == "pooled"))


def label_rows(args, cfg):
    """{regime: sorted int rows of the TRAIN split} the warning head is trained on (data/labels.py); {} when unavailable / synthetic."""
    if args.synthetic or args.no_label_rows:
        return {}
    try:
        import subprocess
        dump = Path(args.out) / "label_rows.npz"
        if not dump.exists():
            subprocess.run([sys.executable, str(CWM3 / "data" / "labels.py"), "--dump-rows", str(dump), "--sources", "consensus,legacy_v1"], check=True,
                           capture_output=True, timeout=1800)                   # separate process: the parent must not import torch-heavy modules before forking
        with np.load(dump) as z:
            return {r: z[r].astype(np.int64) for r in cfg["regimes"] if r in z.files}
    except Exception as ex:                                    # noqa: BLE001
        log(f"label rows unavailable: {ex!r}", args.out)
        return {}


# ----------------------------------------------------------------------------- model / optimiser
def build_model(args, cfg, device):
    """CWM3Joint on `device` from the base checkpoints (synthetic: fresh branches with the base configs)."""
    import torch
    from models.cwm3 import CWM3Joint, merge_stats
    from models import build_branch
    kw = dict(interactions=tuple(cfg.get("interactions", [])), naive=bool(cfg.get("naive", False)), decoder=cfg.get("decoder", "direct"),
              k_head=True, i_hidden=int(cfg.get("i_hidden", 640)), dropout=float(cfg.get("dropout", 0.1)), tok_source=cfg.get("tok_source", "region"))
    if cfg.get("u_kpf"):
        kw["kpf_dim"] = 10                                          # LANE F1 injector (models/kpf_inject.py); the kwarg is absent otherwise (older configs unchanged)
    if cfg.get("mask_V"):
        kw["mask_V"] = True                                         # variant `ue_noV` (2026-09-18): vehicle branch = missing modality; the kwarg is absent otherwise
    if latent_enabled(cfg):
        kw["latent_aux"] = dict(cfg["latent_aux"])                  # C4 (2026-09-17): built only with lam > 0; the kwarg is absent otherwise (older configs unchanged)
    if args.synthetic:
        canh = "8s" if cfg["base"]["V"][1] == "can8s" else "4s"
        e_tok = "pooled" if cfg["base"]["E"][1] == "pooled" else "region"
        V = build_branch("V", args.seed, d=256, hidden=256, dropout=0.1)
        U = build_branch("U", args.seed, hid=128, hidden=128, dropout=0.0, cabin="none")
        E = build_branch("E", args.seed, tokens=e_tok, n_slots=16, d=128, n_layers=2, n_heads=4, d_ff=256, hidden=256, dropout=0.1)
        model = CWM3Joint(V, U, E, seed=args.seed, **kw)
        info = {"base_ckpts": {"V": "synthetic", "U": "synthetic", "E": "synthetic"}, "stats": {},
                "base_cfgs": {"V": {"can_hist": canh}, "U": {"cabin": "none"}, "E": {"tokens": e_tok, "n_slots": 16}}, "params": model.param_report()}
        return model.to(device), info
    base = {b: tuple(x) for b, x in cfg["base"].items()}
    dirs = {b: str(d) for b, d in base_dirs(cfg, args.seed).items()}          # explicit dirs (suffixed bases such as _v2)
    if cfg.get("fresh_trunks"):
        # naive_fresh control (coordinator 2026-09-10): the SAME architecture with every trunk / head randomly initialised (build_branch with the base
        # runs' model kwargs and this seed); statistics stay the data-derived train statistics of the base runs
        from models.cwm3 import _model_kwargs
        info = base_info_numpy(cfg, args.seed)
        mods = {}
        for b, d in base_dirs(cfg, args.seed).items():
            c = json.loads((d / "config.json").read_text())["cfg"]
            mods[b] = build_branch(b, args.seed, **_model_kwargs(c))
        model = CWM3Joint(mods["V"], mods["U"], mods["E"], seed=args.seed, **kw)
        info["base_ckpts"] = {b: f"FRESH(random init, seed {args.seed}; arch of {n})" for b, n in info["base_ckpts"].items()}
        return model.to(device), info
    model, info = build_joint(args.seed, base=base, device="cpu", dirs=dirs, stats_mode=cfg.get("stats_mode", "merged"), **kw)
    for b in ("V", "U", "E"):
        assert info["base_ckpts"][b].split(":")[0] == Path(dirs[b]).name, (b, info["base_ckpts"], dirs)
    return model.to(device), info


def synthetic_info(args, cfg):
    """parent-side info for --synthetic without touching torch (mirrors build_model's synthetic branch)."""
    canh = "8s" if cfg["base"]["V"][1] == "can8s" else "4s"
    e_tok = "pooled" if cfg["base"]["E"][1] == "pooled" else "region"
    return {"base_ckpts": {"V": "synthetic", "U": "synthetic", "E": "synthetic"}, "stats": {},
            "base_cfgs": {"V": {"can_hist": canh}, "U": {"cabin": "none"}, "E": {"tokens": e_tok, "n_slots": 16}}}


def write_params_once(args, model, info):
    p = Path(args.out) / "params.json"
    if not p.exists():
        rep = model.param_report()
        p.write_text(json.dumps(rep, indent=1, default=str))
        info["params"] = rep
        rep["p_head_source"] = getattr(model, "p_head_source", "")
        log(f"model: total {rep['total']:,} params, new {rep['new_total']:,} (I {rep['interaction_total']:,} incl. tok_embed {rep['tok_embed']:,}; "
            f"K {rep['k_head']:,} P {rep['p_head']:,} R {rep['r_head']:,}); P head {rep['p_head_source']}", args.out)
    else:
        info["params"] = json.loads(p.read_text())


def make_optim(model, rec, epochs, mtl=None):
    """mtl (2026-09-16 WP3 block): the MTL object's own parameters (uw: 3 log-variances) join the optimiser as a last group at the new-module lr with
    weight decay 0; every other kind (none / famo / pcgrad) adds nothing, so the optimiser is exactly the original one."""
    import torch
    mult = float(rec.get("lr_base_mult", 0.1))
    base, new = list(model.base_parameters()), model.new_parameters()
    groups = [{"params": new, "lr": float(rec["lr"])}]
    if mult > 0:
        groups.append({"params": base, "lr": float(rec["lr"]) * mult})
    else:
        for p in base:
            p.requires_grad_(False)
    if mtl is not None and mtl.optim_params():
        groups.append({"params": mtl.optim_params(), "lr": float(rec["lr"]), "weight_decay": 0.0})
    opt = torch.optim.AdamW(groups, lr=float(rec["lr"]), weight_decay=float(rec["weight_decay"]))
    if rec.get("sched", "none") == "warmcos":
        warm = torch.optim.lr_scheduler.LinearLR(opt, 0.1, total_iters=3)
        cosn = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=max(epochs - 3, 1))
        sched = torch.optim.lr_scheduler.SequentialLR(opt, [warm, cosn], [3])
    else:
        sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda e: 1.0)
    return opt, sched


# ----------------------------------------------------------------------------- child: training chunk
WATCHDOG_S = 300            # coordinator rule (2026-09-10): no progress within 5 min of the lock -> dump stacks, exit, release the lock


def arm_watchdog(seconds, out):
    import faulthandler
    faulthandler.cancel_dump_traceback_later()
    faulthandler.dump_traceback_later(seconds, repeat=False, exit=True)
    log(f"watchdog armed: {seconds:.0f}s", out)


def add_latent_loss(model, lat_lam, out_, tg, loss, parts):
    """C4 (2026-09-17): loss + lam * L_aux when the model carries the latent module and the batch its targets; parts lat_* (never chain_* -> selection unchanged).
    Without the module (every older variant and the lam = 0 control) returns loss / parts untouched."""
    if getattr(model, "latent_aux", None) is None or "LAT" not in out_ or "LAT" not in tg:
        return loss, parts
    la, pl = model.latent_aux.loss(out_["LAT"], tg["LAT"])
    parts.update(pl)
    parts["lat_total"] = la
    return loss + float(lat_lam) * la, parts


def val_objective(model, store, rows, st, pc, lam, device, amp, bs=8192, region_keys=False, select_LF=True, lat_lam=0.0):
    """select_LF (2026-09-11 fix): the recipe flag is passed in explicitly — the earlier edit referenced `cfg`, which is not in this scope, so every
    TRAINED joint variant crashed at its first validation (NameError; the no_train `indep` variant never validates and survived)."""
    import torch
    model.eval()
    tot_sel, tot, n, acc = 0.0, 0.0, 0, {}
    with torch.no_grad():
        for i in range(0, len(rows), bs):
            idx = torch.as_tensor(rows[i:i + bs], device=device)
            b, tg, _ = prep_joint(store, idx, st, pc, True, region_keys)
            with torch.autocast("cuda", dtype=torch.bfloat16, enabled=amp):
                out = model(b)
            loss, parts = joint_loss(out, tg, lam)
            loss, parts = add_latent_loss(model, lat_lam, out, tg, loss, parts)          # C4: val_loss_full includes lam * L_aux; the selection objective does not
            tot_sel += hd.selection_objective(parts, bool(select_LF)) * len(idx)
            tot += float(loss) * len(idx)
            for k, v in parts.items():
                acc[k] = acc.get(k, 0.0) + float(v) * len(idx)
            n += len(idx)
    return tot_sel / max(n, 1), tot / max(n, 1), {k: v / max(n, 1) for k, v in acc.items()}


def run_chunk(args, cfg, tr, va, st, bank, info):
    import torch
    torch.set_num_threads(4)
    t0 = time.time()
    out = Path(args.out)
    arm_watchdog(WATCHDOG_S, out)
    device = "cuda" if torch.cuda.is_available() and not args.cpu else "cpu"
    rec = cfg["recipe"]
    epochs = args.epochs or int(rec["max_epochs"])
    amp = device == "cuda" and cfg.get("amp", True) and not args.no_amp
    model, _ = build_model(args, cfg, device)
    write_params_once(args, model, info)
    pb = info["params"]
    pc = apply_mask_v(prep_cfgs(info), cfg)
    region_keys = model.tok_embed is not None
    mtl = MTL(cfg.get("mtl", "none"), device=device, seed=args.seed)           # WP3 MTL block (train/mtl.py); "none" = the original training path
    lat_lam = latent_lam(cfg)
    assert not (lat_lam > 0 and mtl.kind != "none"), "C4 latent_aux is defined on the plain-sum path only"
    opt, sched = make_optim(model, rec, epochs, mtl)
    state = {"epoch": 0, "batch_pos": 0, "best": 1e18, "best_epoch": -1, "patience": 0, "acq": 0, "done": False}
    ck = T.latest_ckpt(out)
    if ck is not None:
        z = torch.load(ck, map_location=device, weights_only=False)
        model.load_state_dict(z["model"])
        state = z["state"]
        opt.load_state_dict(z["opt"])
        sched.load_state_dict(z["sched"])
        if z.get("mtl") is not None:
            mtl.load_state_dict(z["mtl"])
        torch.set_rng_state(z["rng_torch"].cpu())
        if device == "cuda" and z.get("rng_cuda") is not None:
            torch.cuda.set_rng_state(z["rng_cuda"].cpu())
        log(f"resumed {ck.name}: epoch {state['epoch']} batch_pos {state['batch_pos']} best {state['best']:.4f}", out)
    else:
        torch.manual_seed(args.seed)
    state["acq"] += 1
    if state["done"]:
        return 0
    rows_tr, rows_va = np.arange(len(tr["t"])), np.arange(len(va["t"]))
    store_tr, store_va = make_store(tr, device, bank), make_store(va, device, bank)
    log(f"chunk {state['acq']} start: {cfg['stage']}/{cfg['variant']} s{args.seed} device={device} amp={amp} params={pb['total']:,} "
        f"(new {pb['new_total']:,}: I {pb['interaction_total']:,} K {pb['k_head']:,}) train rows {len(rows_tr):,} val rows {len(rows_va):,} "
        f"gpu store {(store_tr.bytes() + store_va.bytes()) / 2**30:.2f} GiB epochs={epochs} rows/epoch={rec['rows_per_epoch'] or 'all'} "
        f"batch={rec['batch']} lr={rec['lr']} base_mult={rec.get('lr_base_mult', 0.1)} decoder={cfg.get('decoder')} inter={cfg.get('interactions')} "
        f"naive={cfg.get('naive')} mtl={mtl.kind} latent_aux={cfg.get('latent_aux') if lat_lam > 0 else None} (params {pb.get('latent_aux', 0):,}) V={cfg['base']['V']}"
        f"{' MASK_V (vehicle branch = missing modality)' if cfg.get('mask_V') else ''} ({time.time() - t0:.0f}s)", out)
    tag_tr = tr["slice_tag"][rows_tr]
    cnt = np.array([(tag_tr == v).sum() for v in range(4)], float)
    p_s = (1.0 / np.maximum(cnt, 1))[tag_tr]
    p_s /= p_s.sum()
    bs, rows_per_epoch = int(rec["batch"]), int(rec["rows_per_epoch"])

    def save():
        free_gb = shutil.disk_usage(out).free / 2**30
        assert free_gb >= 15.0, f"free disk {free_gb:.1f} GB < 15 GB: refusing to checkpoint"
        p = out / f"ckpt_ep{state['epoch']:03d}_b{state['batch_pos']:05d}.pt"
        assert not p.exists()
        torch.save({"model": model.state_dict(), "opt": opt.state_dict(), "sched": sched.state_dict(), "state": state,
                    "rng_torch": torch.get_rng_state(), "rng_cuda": torch.cuda.get_rng_state() if device == "cuda" else None,
                    "cfg": cfg, "seed": args.seed, "stats": st, "params": pb, "base_ckpts": info["base_ckpts"],
                    "mtl": mtl.state_dict() if mtl.kind != "none" else None}, p)
        return p

    def over_budget():
        return (time.time() - t0) >= args.minutes * 60

    while state["epoch"] < epochs and not state["done"]:
        ep_ = state["epoch"]
        rng = np.random.default_rng(args.seed * 1000 + ep_)
        if rows_per_epoch > 0:
            idx_all = rows_tr[rng.choice(len(rows_tr), size=min(len(rows_tr), rows_per_epoch), replace=False, p=p_s)]
        else:
            idx_all = rows_tr[rng.permutation(len(rows_tr))]
        model.train()
        nb = (len(idx_all) + bs - 1) // bs
        tot, nbatch = 0.0, 0
        while state["batch_pos"] < nb:
            i = state["batch_pos"]
            idx = torch.as_tensor(idx_all[i * bs:(i + 1) * bs], device=device)
            b, tg, _ = prep_joint(store_tr, idx, st, pc, True, region_keys)
            with torch.autocast("cuda", dtype=torch.bfloat16, enabled=amp):
                out_ = model(b)
            if mtl.kind == "none":
                loss, parts = joint_loss(out_, tg, cfg["lambda"])
                loss, parts = add_latent_loss(model, lat_lam, out_, tg, loss, parts)      # C4 (2026-09-17): + lam * L_aux; identity when the module is absent
                opt.zero_grad(set_to_none=True)
                loss.backward()
            else:
                # WP3 MTL block (train/mtl.py): uw / famo re-weight the three branch losses, pcgrad projects their conflicting gradients on the shared
                # parameters (probe on the first batch of every chunk); `loss` stays the PLAIN sum for the history / logs, clipping + step unchanged
                loss, parts, br = joint_loss(out_, tg, cfg["lambda"], return_branches=True)
                opt.zero_grad(set_to_none=True)
                if mtl.kind == "pcgrad" and (mtl.shared is None or nbatch == 0):
                    trainable = [p for p in model.parameters() if p.requires_grad]
                    mtl.probe_shared(br, trainable)
                    (out / "mtl_probe.json").write_text(json.dumps(mtl.probe_info, indent=1))
                    log(f"pcgrad probe: {mtl.probe_info}", out)
                mtl.backward(loss, br)
            torch.nn.utils.clip_grad_norm_([p for p in model.parameters() if p.requires_grad], float(cfg.get("grad_clip", 1.0)))
            opt.step()
            if getattr(model, "latent_aux", None) is not None:
                model.latent_aux.ema_update()                                          # C4: EMA target encoder follows the online encoder after every step
            if mtl.kind == "famo":
                with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16, enabled=amp):
                    _, _, br_after = joint_loss(model(b), tg, cfg["lambda"], return_branches=True)
                mtl.famo_update(br_after)
            tot += float(loss)
            nbatch += 1
            state["batch_pos"] += 1
            if nbatch == 1:
                arm_watchdog(args.minutes * 60 + WATCHDOG_S, out)          # first batch done: guard the whole chunk instead
            if nbatch % 100 == 0:
                log(f"  ep {ep_} b {i + 1}/{nb} loss {tot / nbatch:.4f} " + " ".join(f"{k}={float(v):.3f}" for k, v in parts.items() if k.startswith("chain_") or k == "L_F" or k.startswith("lat_"))
                    + (f" mtl[{mtl.kind}] w={['%.3g' % w for w in mtl.weights()]}" if mtl.kind != "none" else ""), out)
                if mtl.kind != "none":
                    assert mtl.finite(), f"MTL weights not finite: {mtl.state_dict()}"
            if over_budget() and state["batch_pos"] < nb:
                p = save()
                log(f"wall budget reached mid-epoch -> saved {p.name}; exit 3", out)
                return 3
        sched.step()
        v, v_full, vparts = val_objective(model, store_va, rows_va, st, pc, cfg["lambda"], device, amp, region_keys=region_keys,
                                          select_LF=bool(cfg["recipe"].get("select_LF", True)), lat_lam=lat_lam)
        state["epoch"] += 1
        state["batch_pos"] = 0
        improved = v < state["best"] - 1e-5
        if improved:
            state["best"], state["best_epoch"], state["patience"] = v, ep_, 0
        else:
            state["patience"] += 1
        per_branch = {"V": sum(vparts.get(f"chain_{c}", 0.0) for c in hd.V_CHAINS), "U": sum(vparts.get(f"chain_{c}", 0.0) for c in hd.U_CHAINS),
                      "E": sum(vparts.get(f"chain_{c}", 0.0) for c in hd.E_CHAINS) + vparts.get("L_F", 0.0)}
        rec_h = {"epoch": ep_, "train_loss": tot / max(nbatch, 1), "val_objective": v, "val_objective_kind": "sum of chain NLLs (V masked to V-valid rows) + 0.1 L_F",
                 "val_loss_full": v_full, "val_parts": vparts, "val_per_branch": per_branch, "best": state["best"], "best_epoch": state["best_epoch"],
                 "lr": [g["lr"] for g in opt.param_groups], "acq": state["acq"], "t_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
                 **({"mtl": mtl.kind, "mtl_weights": mtl.weights()} if mtl.kind != "none" else {})}
        with open(out / "history.jsonl", "a") as f:
            f.write(json.dumps(rec_h) + "\n")
        log(f"ep {ep_}: train {rec_h['train_loss']:.4f} val_select {v:.4f} (V {per_branch['V']:.4f} U {per_branch['U']:.4f} E {per_branch['E']:.4f}) "
            f"val_full {v_full:.4f} best {state['best']:.4f}@{state['best_epoch']} ({time.time() - t0:.0f}s)", out)
        if state["patience"] >= int(rec["patience"]) or state["epoch"] >= epochs:
            state["done"] = True
        p = save()
        if improved:
            (out / "BEST").write_text(p.name + "\n")
        T.prune_ckpts(out, keep={p.name, (out / "BEST").read_text().strip() if (out / "BEST").exists() else ""})
        if state["done"]:
            (out / "TRAIN_DONE").write_text(json.dumps({"best_epoch": state["best_epoch"], "best": state["best"], "epochs_run": state["epoch"],
                                                        "acq": state["acq"], "params": pb}))
            log("training done", out)
            import faulthandler
            faulthandler.cancel_dump_traceback_later()
            return 0
        if over_budget():
            log("wall budget reached at epoch end; exit 3", out)
            return 3
    return 0


# ----------------------------------------------------------------------------- child: prediction
def load_best_model(args, cfg, device, info):
    import torch
    model, _ = build_model(args, cfg, device)
    best = (Path(args.out) / "BEST").read_text().strip() if (Path(args.out) / "BEST").exists() else "base"
    if best != "base":
        z = torch.load(Path(args.out) / best, map_location=device, weights_only=False)
        model.load_state_dict(z["model"])
    model.eval()
    return model, best


def predict_split(args, cfg, raw, st, bank, info, regime, split, rows=None):
    """all rows of one (regime, split) pack -> pred_{regime}_{split}.npz; with `rows` (train label rows) -> pred_{regime}_train_labelrows.npz."""
    import torch
    import eval_s13 as e13
    torch.set_num_threads(4)
    t0 = time.time()
    out = Path(args.out)
    arm_watchdog(WATCHDOG_S, out)
    device = "cuda" if torch.cuda.is_available() and not args.cpu else "cpu"
    amp = device == "cuda" and cfg.get("amp", True) and not args.no_amp
    model, best = load_best_model(args, cfg, device, info)
    write_params_once(args, model, info)
    pc = apply_mask_v(prep_cfgs(info), cfg)
    region_keys = model.tok_embed is not None
    N = len(raw["t"])
    store = make_store(raw, device, bank)
    C = len(CHAINS)
    valid_U = raw["kin_row_ok"].astype(bool) if "kin_row_ok" in raw else np.ones(N, bool)      # causal kin: rows without a causal box are not scored (2026-09-11)
    P = {"hz_on": np.zeros((N, 25, C), np.float16), "hz_off": np.zeros((N, 25, C), np.float16), "z0": np.zeros((N, C), np.float16),
         "z0valid": np.ones((N, C), bool), "valid_V": v_valid_np(raw, pc), "valid_U": valid_U, "valid_E": np.ones(N, bool),
         "y_mae": np.full((N, 3), np.nan, np.float32), "y_crps": np.full((N, 3), np.nan, np.float32), "pi_resp": np.zeros((N, 25), np.float16),
         "f_cos": np.full((N, 5), np.nan, np.float32), "k_mae": np.full(N, np.nan, np.float32), "K_hat": np.zeros((N, 3, 11), np.float16),
         "K_sig": np.zeros((N, 3, 11), np.float16)}
    has_aux = store.has("P") and "p0_base" in st and model.p_head is not None
    if has_aux:
        P.update(hz_on_P=np.zeros((N, 25, 6), np.float16), hz_off_P=np.zeros((N, 25, 6), np.float16), z0_P=np.zeros((N, 6), np.float16), z0valid_P=np.ones((N, 6), bool),
                 hz_on_R=np.zeros((N, 25, 1), np.float16), hz_off_R=np.zeros((N, 25, 1), np.float16), z0_R=np.zeros((N, 1), np.float16), z0valid_R=np.ones((N, 1), bool))
    has_g = getattr(model.U, "gaze_head", None) is not None and store.has("Sg")
    if has_g:                                                       # gaze_off chain of the gaze U bases (2026-09-12 appendix variant): the base trainer's extra blocks, same names / shapes
        CG = len(hd.G_CHAINS)
        P.update(hz_on_G=np.zeros((N, 25, CG), np.float16), hz_off_G=np.zeros((N, 25, CG), np.float16), z0_G=np.zeros((N, CG), np.float16), z0valid_G=np.zeros((N, CG), bool))
    TGG = {"zp": [], "z": [], "m": []}
    TGA = {"P": {"zp": [], "z": [], "m": []}, "R": {"zp": [], "z": [], "m": []}}
    TG = {"zp": [], "z": [], "m": []}
    stV = branch_stats(st, "V")
    ymu, ysd = torch.as_tensor(stV["ymu"], device=device), torch.as_tensor(stV["ysd"], device=device)
    gen = torch.Generator(device=device).manual_seed(args.seed + 99)
    bs = 4096
    with torch.no_grad():
        for i in range(0, N, bs):
            idx = torch.arange(i, min(i + bs, N), device=device)
            sl = slice(i, i + len(idx))
            b, tg, vv = prep_joint(store, idx, st, pc, True, region_keys)
            with torch.autocast("cuda", dtype=torch.bfloat16, enabled=amp):
                o = model(b)
            if i == 0:
                arm_watchdog(args.minutes * 60 + WATCHDOG_S, out)
            v, u, e = hd.split_V(o["V"]), hd.split_U(o["U"]), hd.split_E(o["E"])
            hon_v, hoff_v = hd.hazards(v["chains"])
            hon_u, hoff_u = hd.hazards(u["chains"])
            hon_e, hoff_e = hd.hazards(e["chains"])
            h_on, h_off = torch.cat([hon_v, hon_u, hon_e], -1), torch.cat([hoff_v, hoff_u, hoff_e], -1)
            z0 = torch.cat([tg["V"]["z0"], tg["U"]["z0"], tg["E"]["z0"]], -1)
            zp = torch.cat([tg["V"]["zV_prev"], tg["U"]["zU_prev"], tg["E"]["zD_prev"]], -1)
            zz = torch.cat([tg["V"]["zV"], tg["U"]["zU"], tg["E"]["zD"]], -1)
            mm = torch.cat([tg["V"]["zVm"], tg["U"]["zUm"], tg["E"]["zDm"]], -1)
            P["hz_on"][sl], P["hz_off"][sl], P["z0"][sl] = h_on.cpu().numpy(), h_off.cpu().numpy(), z0.cpu().numpy()
            P["z0valid"][sl, 3] = store.t["Resp0valid"][idx].cpu().numpy()
            P["z0valid"][sl, 4:6] = tg["U"]["z0valid"].cpu().numpy()
            P["z0valid"][sl, 6:] = tg["E"]["z0valid"].cpu().numpy()
            # vehicle continuous emissions (rows with the V history)
            mean = torch.stack([hd.mdn_mean(v["steer"]), hd.mdn_mean(v["vEgo"]), hd.mdn_mean(v["aEgo"])], -1) * ysd + ymu
            Y = store.t["Y"][idx].float()[..., [0, 3, 4]]
            M = tg["V"]["M"].float()[..., None]
            P["y_mae"][sl] = ((mean - Y).abs() * M).sum(1).div(M.sum(1).clamp_min(1)).cpu().numpy()
            smp = torch.stack([hd.mdn_sample(v[c], 32, gen) for c in ("steer", "vEgo", "aEgo")], -1) * ysd + ymu
            t1 = (smp - Y[None]).abs().mean(0)
            t2 = (smp[:, None] - smp[None]).abs().mean((0, 1))
            P["y_crps"][sl] = ((t1 - 0.5 * t2) * M).sum(1).div(M.sum(1).clamp_min(1)).cpu().numpy()
            P["pi_resp"][sl] = hd.forward_marginal(hon_v[..., 3], hoff_v[..., 3], tg["V"]["z0"][:, 3]).cpu().numpy()
            # driver kinematics (detached K head): K̂ = K_0 + ΔK̂, standardised
            K_mu = tg["U"]["K0std"][:, None, :] + u["dK_mu"]
            Km = tg["U"]["Km"].float()
            P["k_mae"][sl] = (((K_mu - tg["U"]["Kstd"]).abs() * Km).sum((1, 2)) / Km.sum((1, 2)).clamp_min(1)).cpu().numpy()
            P["K_hat"][sl] = K_mu[:, [4, 14, 24], :].cpu().numpy()
            P["K_sig"][sl] = u["K_logsig"][:, [4, 14, 24], :].clamp(-7, 5).exp().cpu().numpy()
            cos = torch.nn.functional.cosine_similarity(o["F"].float(), tg["E"]["F"].float(), dim=-1)
            P["f_cos"][sl] = torch.where(tg["E"]["Fok"], cos, torch.full_like(cos, float("nan"))).cpu().numpy()
            TG["zp"].append(zp.cpu().numpy().astype(np.float32)); TG["z"].append(zz.cpu().numpy().astype(np.float32)); TG["m"].append(mm.cpu().numpy())
            if has_aux:
                for key, nm in (("P", "P"), ("R", "R")):
                    c = o[key].float().reshape(len(idx), 25, -1, 2)
                    hon, hoff = hd.hazards(c)
                    P[f"hz_on_{nm}"][sl], P[f"hz_off_{nm}"][sl] = hon.cpu().numpy(), hoff.cpu().numpy()
                    P[f"z0_{nm}"][sl], P[f"z0valid_{nm}"][sl] = tg[key]["z0"].cpu().numpy(), tg[key]["z0valid"].cpu().numpy()
                    TGA[key]["zp"].append(tg[key]["z_prev"].cpu().numpy().astype(np.float32)); TGA[key]["z"].append(tg[key]["z"].cpu().numpy().astype(np.float32)); TGA[key]["m"].append(tg[key]["m"].cpu().numpy())
            if has_g and "G" in o and "zG" in tg["U"]:
                tgu = tg["U"]
                hg_on, hg_off = hd.hazards(o["G"].float().reshape(len(idx), 25, 1, 2))
                P["hz_on_G"][sl], P["hz_off_G"][sl] = hg_on.cpu().numpy(), hg_off.cpu().numpy()
                P["z0_G"][sl], P["z0valid_G"][sl] = tgu["zG0"][:, None].cpu().numpy(), tgu["zG0valid"][:, None].cpu().numpy()
                TGG["zp"].append(tgu["zG_prev"][:, :, None].cpu().numpy().astype(np.float32)); TGG["z"].append(tgu["zG"][:, :, None].cpu().numpy().astype(np.float32)); TGG["m"].append(tgu["zGm"][:, :, None].cpu().numpy())
            if (i // bs) % 25 == 0:
                log(f"  predict {regime}/{split} {i + len(idx)}/{N} ({time.time() - t0:.0f}s)", out)
    TG = {k: np.concatenate(v) for k, v in TG.items()}
    vV = P["valid_V"]
    vU = P["valid_U"]
    P["y_mae"][~vV] = np.nan
    P["y_crps"][~vV] = np.nan
    P["k_mae"][~vU] = np.nan
    tag = f"{regime}_{split}" if rows is None else f"{regime}_train_labelrows"
    qm = {"n_rows": int(N), "n_rows_predicted": int(N), "n_rows_V": int(vV.sum()), "n_rows_U": int(vU.sum()), "regime": regime, "split": split, "branch": "joint"}
    if rows is None:
        hz_on, hz_off = P["hz_on"].astype(np.float64), P["hz_off"].astype(np.float64)
        events = [(c, c, "on") for c in hd.V_CHAINS] + [("hand_return", "hands_on", "on"), ("hand_leave", "hands_on", "off"), ("head_turn_onset", "head_turned", "on"),
                                                         ("head_return", "head_turned", "off")] + [(c, c, "on") for c in hd.E_CHAINS]
        for name, ch, kind in events:
            c = CHAINS.index(ch)
            keep = vV if ch in hd.V_CHAINS else vU if ch in hd.U_CHAINS else np.ones(N, bool)
            zp, zz, mm = TG["zp"][keep][..., c], TG["z"][keep][..., c], TG["m"][keep][..., c].astype(bool)
            if kind == "off":
                zp, zz, h = 1 - zp, 1 - zz, hz_off[keep][..., c]
            else:
                h = hz_on[keep][..., c]
            try:
                for r in e13.transition_metrics(h, zp, zz, mm, name):
                    qm[f"{name}:{r['metric']}"] = r["value"]
                    if "base_rate" in r:
                        qm[f"{name}:{r['metric']}:base_rate"], qm[f"{name}:{r['metric']}:n"] = r["base_rate"], r["n"]
            except Exception as ex:                                             # noqa: BLE001  (too few rows at smoke scale)
                qm[f"{name}:error"] = str(ex)[:120]
        if has_aux:
            for key, names, keep in (("P", [f"P_{c}" for c in P_CHAINS], np.ones(N, bool)), ("R", ["R_hard"], vV)):
                A = {k: np.concatenate(v) for k, v in TGA[key].items()}
                for c, name in enumerate(names):
                    try:
                        for r in e13.transition_metrics(P[f"hz_on_{key}"].astype(np.float64)[keep][..., c], A["zp"][keep][..., c], A["z"][keep][..., c], A["m"][keep][..., c].astype(bool), name):
                            qm[f"{name}:{r['metric']}"] = r["value"]
                            if "base_rate" in r:
                                qm[f"{name}:{r['metric']}:base_rate"], qm[f"{name}:{r['metric']}:n"] = r["base_rate"], r["n"]
                    except Exception as ex:                                     # noqa: BLE001
                        qm[f"{name}:error"] = str(ex)[:120]
        if has_g and TGG["z"]:
            # gaze_off_on / gaze_off_off quick metrics on the U-valid rows — the base trainer's form (train_cwm3.predict_split), so the identity check can compare them
            G = {k: np.concatenate(v) for k, v in TGG.items()}
            for name, kind in (("gaze_off_on", "on"), ("gaze_off_off", "off")):
                zp, zz, mm = G["zp"][vU][..., 0], G["z"][vU][..., 0], G["m"][vU][..., 0].astype(bool)
                if kind == "off":
                    zp, zz, h = 1 - zp, 1 - zz, P["hz_off_G"].astype(np.float64)[vU][..., 0]
                else:
                    h = P["hz_on_G"].astype(np.float64)[vU][..., 0]
                try:
                    for r in e13.transition_metrics(h, zp, zz, mm, name):
                        qm[f"{name}:{r['metric']}"] = r["value"]
                        if "base_rate" in r:
                            qm[f"{name}:{r['metric']}:base_rate"], qm[f"{name}:{r['metric']}:n"] = r["base_rate"], r["n"]
                except Exception as ex:                                         # noqa: BLE001
                    qm[f"{name}:error"] = str(ex)[:120]
        for k in ("y_mae", "y_crps", "f_cos", "k_mae"):
            vals = P[k]
            qm[k + ":mean"] = np.nanmean(vals, 0).tolist() if vals.ndim == 2 else float(np.nanmean(vals))
        qm["wall_s"] = round(time.time() - t0, 1)
        (out / f"metrics_quick_{regime}_{split}.json").write_text(json.dumps(qm, indent=1))
    extra = {} if rows is None else {"rows": np.asarray(rows, np.int64), "uid": raw["uid"], "t": raw["t"]}
    np.savez_compressed(out / f"pred_{tag}.npz", **P, **extra, regime=np.array(regime), split=np.array(split), branch=np.array("joint"), seed=np.array(args.seed),
                        stage=np.array(cfg["stage"]), variant=np.array(cfg["variant"]), best_ckpt=np.array(best), chain_names=np.array(CHAINS),
                        p_chain_names=np.array(list(P_CHAINS)), r_chain_names=np.array(list(R_CHAINS)),
                        hazard_kind=np.array("head"), n_rows=np.array(N), base_ckpts=np.array(json.dumps(info["base_ckpts"])),
                        **({"g_chain_names": np.array(list(hd.G_CHAINS))} if has_g else {}))
    log(f"pred_{tag}.npz written ({N} rows, V rows {int(vV.sum())}) in {time.time() - t0:.0f}s; " +
        " ".join(f"{k}={v:.3f}" for k, v in qm.items() if k.endswith("onset_prauc@3s")), out)
    import faulthandler
    faulthandler.cancel_dump_traceback_later()
    return 0


def predict_many(args, cfg, raws, st, bank, info):
    for key, raw in raws.items():
        regime, split = key[0], key[1]
        rows = raw.get("_label_rows")
        tag = f"{regime}_{split}" if rows is None else f"{regime}_train_labelrows"
        if (Path(args.out) / f"pred_{tag}.npz").exists():
            continue
        code = predict_split(args, cfg, raw, st, bank, info, regime, split, rows)
        if code:
            return code
    return 0


def load_pred_raws(args, cfg, lrows):
    raws = {}
    for split in cfg["pred_splits"]:
        for regime in cfg["regimes"]:
            if not (Path(args.out) / f"pred_{regime}_{split}.npz").exists():
                raws[(regime, split)] = load_raw(args, cfg, split, subset=False, regimes=[regime])
    for regime, rows in lrows.items():
        if len(rows) and not (Path(args.out) / f"pred_{regime}_train_labelrows.npz").exists():
            full = load_raw(args, cfg, "train", subset=False, regimes=[regime])
            sub = ds.subset(full, rows)
            sub["_label_rows"] = rows
            raws[(regime, "train")] = sub
            del full
    return raws


def train_then_predict(args, cfg, tr, va, st, bank, info, raws):
    t0 = time.time()
    code = run_chunk(args, cfg, tr, va, st, bank, info)
    if code != 0:
        return code
    left = args.minutes * 60 - (time.time() - t0)
    if raws and left >= T.PREDICT_MIN_LEFT_S:
        log(f"training finished with {left / 60:.1f} min of the chunk budget left -> predicting in the same acquisition", args.out)
        return predict_many(args, cfg, raws, st, bank, info)
    return 0


def append_quick_rows(args, cfg, stamp, td):
    rows = []
    for regime in cfg["regimes"]:
        for split in cfg["pred_splits"]:
            p = Path(args.out) / f"metrics_quick_{regime}_{split}.json"
            if not p.exists():
                continue
            qm = json.loads(p.read_text())
            for k, v in qm.items():
                if ":" not in k or k.endswith(":base_rate") or k.endswith(":n") or k.endswith(":error"):
                    continue
                chain, metric = k.split(":", 1)
                rows.append({"run_stamp": stamp, "stage": cfg["stage"], "variant": cfg["variant"], "seed": args.seed, "regime": regime, "split": split,
                             "branch": "joint", "chain": chain, "metric": metric, "value": v if not isinstance(v, list) else json.dumps(v),
                             "base_rate": qm.get(k + ":base_rate", ""), "n": qm.get(k + ":n", ""), "n_rows_predicted": qm["n_rows_predicted"],
                             "best_epoch": td["best_epoch"], "val_objective": td["best"], "run_dir": str(args.out)})
    if args.synthetic or not rows:
        return 0
    new = not T.QUICK_CSV.exists()
    with open(T.QUICK_CSV, "a", newline="") as f:
        w = csv.DictWriter(f, fieldnames=T.QUICK_COLS)
        if new:
            w.writeheader()
        for r in rows:
            w.writerow(r)
    return len(rows)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--stage", required=True, choices=list(STAGES["stages"]))
    ap.add_argument("--variant", required=True)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--base-v", default="", choices=[""] + list(T.STAGES["stages"]["V0"]["variants"]), help="override the V base variant of stages_joint.yaml")
    ap.add_argument("--base-e", default="", choices=[""] + list(T.STAGES["stages"]["E0"]["variants"]), help="override the E base variant of stages_joint.yaml (e.g. pooled_P)")
    ap.add_argument("--base-u", default="", choices=[""] + list(T.STAGES["stages"]["U0"]["variants"]), help="override the U base variant of stages_joint.yaml (e.g. kin_can_noK_clip8_ln)")
    ap.add_argument("--base-dir-v", default="", help="explicit V base run dir (absolute, cwm3-relative or a name under results/runs); 2026-09-11 final assembly")
    ap.add_argument("--base-dir-u", default="", help="explicit U base run dir")
    ap.add_argument("--base-dir-e", default="", help="explicit E base run dir")
    ap.add_argument("--kin-source", default="auto", choices=["auto", "cwm2", "causal_routes", "oldin_ctgt"],
                    help="driver kinematics source of the data layer (ds.set_kin_source); auto (default) = the U base's config.json (cwm2 for bases without the key)")
    ap.add_argument("--kin-version", default="auto", choices=["auto", "v1", "v2", "v3", "v2p"], help="causal kin pack build; auto = the U base's config.json (v3 = 2026-09-16 wheel-reference labels; v2p = 2026-09-18 v2 + 3 keypoint input proxies)")
    ap.add_argument("--run-suffix", default="", help="appended to the default run dir name (e.g. _final)")
    ap.add_argument("--stats-mode", default="merged", choices=["merged", "per_branch"],
                    help="merged (default): the bases' statistics are merged with equality asserted on overlapping keys (1e-5); per_branch: each branch uses its own base's stats.npz "
                         "('<branch>/<key>' keys) — an explicit choice when the merged assert fires, never a loosened assert (roadmap §三 3)")
    ap.add_argument("--pack-version", default="v1", choices=["v1", "v2"], help="v2 = results/pack_v2 vf256 / vc256 + aligned F (2026-09-10 latent audit); the E base dir then needs the _v2 suffix")
    ap.add_argument("--base-suffix-e", default="", help="suffix of the E base run dir (e.g. _v2 -> results/runs/E0_<variant>_s<k>_v2)")
    ap.add_argument("--max-anchors", type=int, default=0)
    ap.add_argument("--synthetic", type=int, default=0)
    ap.add_argument("--epochs", type=int, default=0)
    ap.add_argument("--minutes", type=float, default=0)
    ap.add_argument("--regimes", default="")
    ap.add_argument("--out", default="")
    ap.add_argument("--pred-splits", default="")
    ap.add_argument("--no-label-rows", action="store_true", help="skip the train label-row prediction pack")
    ap.add_argument("--no-lock", action="store_true", help="run inline (no fork, no /tmp/s13_gpu.lock) — owner directive 2026-09-10 13:58: the exclusive lock is no longer required; "
                                                            "the queue script gates on free VRAM / RAM / trainer count instead")
    ap.add_argument("--no-amp", action="store_true")
    ap.add_argument("--cpu", action="store_true", help="tests only (smoke); never a result")
    ap.add_argument("--u-kpf", default="", help="LANE F1 (2026-09-17, pre-reg 2026-09-17T19:33:12Z): keypoint-forecaster run dir (results/runs/KPF_<arm>_s<k>) whose detached "
                    "kpf_{regime}_{split}.npz channel enters the driver state through the zero-init KpfInject; empty = the unchanged model")
    ap.add_argument("--mtl", default=None, choices=list(MTL_KINDS),
                    help="WP3 reviewer control (2026-09-16, train/mtl.py): multi-task balancing of the three branch losses — none (default: plain sum, the original path), "
                         "uw (Kendall 2018 learned log-variances), famo (Liu 2023), pcgrad (Yu 2020 on the shared parameters); unset -> the variant's `mtl` key (none)")
    args = ap.parse_args()
    ds.set_pack_version(args.pack_version)
    cfg = stage_cfg(args.stage, args.variant, args.base_v or None, args.base_e or None, args.base_u or None)
    cfg["mtl"] = args.mtl if args.mtl is not None else str(cfg.get("mtl", "none"))
    assert cfg["mtl"] in MTL_KINDS, cfg["mtl"]
    cfg["pack_version"], cfg["base_suffix_e"] = args.pack_version, args.base_suffix_e or ("_v2" if args.pack_version == "v2" else "")
    cfg["base_dirs"] = {b: v for b, v in (("V", args.base_dir_v), ("U", args.base_dir_u), ("E", args.base_dir_e)) if v}
    cfg["stats_mode"], cfg["run_suffix"] = args.stats_mode, args.run_suffix
    # driver kinematics source: explicit or inherited from the U base's config (the joint model must read the inputs its U base was trained on)
    ks, kv = args.kin_source, args.kin_version
    if not args.synthetic:
        ucfg_p = base_dirs(cfg, args.seed)["U"] / "config.json"
        if ucfg_p.exists():
            uc = json.loads(ucfg_p.read_text())["cfg"]
            ks = uc.get("kin_source", "cwm2") if ks == "auto" else ks
            kv = uc.get("kin_version", "v1") if kv == "auto" else kv
            assert ks == uc.get("kin_source", "cwm2") and kv == uc.get("kin_version", "v1"), f"kin source {ks}/{kv} != the U base's {uc.get('kin_source', 'cwm2')}/{uc.get('kin_version', 'v1')}"
            urec = uc.get("recipe", {}) or {}                          # gaze U bases (2026-09-12 appendix variant): opt-in streams (need_streams) / heads follow the U base's recipe
            cfg["u_gaze_chain"], cfg["u_face"] = bool(urec.get("u_gaze_chain", False)), bool(urec.get("u_face", False))
    ks, kv = ("cwm2" if ks == "auto" else ks), ("v1" if kv == "auto" else kv)
    ds.set_kin_source(ks, kv)
    cfg["u_kpf"] = str(resolve_run_dir(args.u_kpf)) if args.u_kpf else ""     # LANE F1 (2026-09-17)
    if cfg["u_kpf"]:
        assert (Path(cfg["u_kpf"]) / "PRED_DONE").exists(), f"--u-kpf {cfg['u_kpf']}: PRED_DONE missing"
        ds.set_kpf_dir(cfg["u_kpf"])
    cfg["kin_source"], cfg["kin_version"] = ks, kv
    if args.regimes:
        cfg["regimes"] = args.regimes.split(",")
    if args.pred_splits:
        cfg["pred_splits"] = args.pred_splits.split(",")
    args.minutes = args.minutes or float(cfg["minutes"])
    assert args.minutes <= 22 or args.no_lock, "GPU-lock chunks must be <= 22 min (unlimited only inline with --no-lock)"
    if not args.out:
        suffix = ("_synthetic" if args.synthetic else "") + (f"_max{args.max_anchors}" if args.max_anchors else "")
        if args.base_v and args.base_v != STAGES["defaults"]["base"]["V"][1]:
            suffix += "_v" + args.base_v[3]
        if args.base_e and args.base_e != STAGES["defaults"]["base"]["E"][1]:
            suffix += "_e" + args.base_e[:4]
        if args.pack_version == "v2":
            suffix += "_v2"
        suffix += args.run_suffix
        args.out = str(CWM3 / "results" / "runs" / f"{args.stage}_{args.variant}_s{args.seed}{suffix}")
    Path(args.out).mkdir(parents=True, exist_ok=True)
    t_start = time.time()
    (Path(args.out) / "config.json").write_text(json.dumps({"cfg": cfg, "args": vars(args)}, indent=1, default=str))
    log(f"=== train_cwm3_joint {args.stage}/{args.variant} s{args.seed} out={args.out} base={cfg['base']} synthetic={args.synthetic} "
        f"max_anchors={args.max_anchors} regimes={cfg['regimes']} need={sorted(need_streams(cfg))} ===", args.out)
    # PARENT: numpy / json only — no torch op before os.fork() (the OMP thread pool does not survive fork; 31-min deadlock on 2026-09-10)
    if args.synthetic:
        info = synthetic_info(args, cfg)
        tr0 = load_raw(args, cfg, "train")
        info["stats"] = bt.compute_stats(tr0, {"mu": np.array([0.0, 12.0, 0.0], np.float32), "sd": np.array([70.0, 8.0, 0.7], np.float32)})
        info["stats"].update(aux_stats(tr0))
        if latent_enabled(cfg):
            info["stats"].update(LA.latent_stats(tr0["LYf"], tr0["LYv"]))
            info["stats"]["lat_horizons"] = np.asarray(latent_horizons(cfg), np.int64)
        del tr0
    else:
        info = base_info_numpy(cfg, args.seed)
        cfg["e_side"] = info["base_cfgs"]["E"].get("side", "vf256")
        sv = {b: c.get("stats_version") for b, c in info["base_cfgs"].items()}
        cfg["stats_version"] = int(min(v for v in sv.values() if v is not None)) if any(v is not None for v in sv.values()) else 1
        cfg["base_stats_versions"] = sv
        (Path(args.out) / "config.json").write_text(json.dumps({"cfg": cfg, "args": vars(args)}, indent=1, default=str))   # re-written with the derived fields
        log(f"stats_mode {info['stats_mode']}; base stats versions {sv}; kin_source {ks}/{kv}; base dirs {info['base_dirs']}", args.out)
    st = {k: np.asarray(v, np.float32) for k, v in info["stats"].items()}
    st_path = Path(args.out) / "stats.npz"
    if st_path.exists():
        st = {k: v for k, v in np.load(st_path).items()}
    else:
        if not args.synthetic:
            st.update(aux_stats_from_files(cfg["regimes"]))                    # auxiliary anchor base rates (TRAIN files only)
            if latent_enabled(cfg):
                st.update(latent_stats_from_files(cfg["regimes"], latent_horizons(cfg)))   # C4: TRAIN statistics of the continuous latent channels
        np.savez(st_path, **st)
    (Path(args.out) / "base_ckpts.json").write_text(json.dumps({"base_ckpts": info["base_ckpts"], "base_cfgs": info["base_cfgs"], "tok_source": cfg.get("tok_source"),
                                                                 "base_dirs": info.get("base_dirs", {}), "stats_mode": info.get("stats_mode", "merged"),
                                                                 "kin_source": ks, "kin_version": kv}, indent=1, default=str))
    log(f"base {info['base_ckpts']}; tok_source {cfg.get('tok_source')} (the child prints the parameter report)", args.out)
    bank = load_bank(args, cfg, info)
    lrows = label_rows(args, cfg)
    n_acq, mem = 0, {}
    if cfg.get("no_train"):
        if not (Path(args.out) / "TRAIN_DONE").exists():
            (Path(args.out) / "TRAIN_DONE").write_text(json.dumps({"best_epoch": -1, "best": float("nan"), "epochs_run": 0, "acq": 0, "params": {}, "no_train": True}))
            (Path(args.out) / "BEST").write_text("base\n")
            with open(Path(args.out) / "history.jsonl", "a") as f:
                f.write(json.dumps({"epoch": -1, "note": "no_train variant: base checkpoints predict as assembled"}) + "\n")
    if not (Path(args.out) / "TRAIN_DONE").exists():
        r0 = ds.rss_gb()
        tr, va = load_raw(args, cfg, "train"), load_raw(args, cfg, "val")
        mem = {"rss_after_load_gb": round(ds.rss_gb(), 2), "rss_before_gb": round(r0, 2), "arrays_train_gb": round(ds.nbytes(tr) / 2**30, 3),
               "arrays_val_gb": round(ds.nbytes(va) / 2**30, 3), "bank_gb": round(sum(v.nbytes for v in bank.values() if isinstance(v, np.ndarray)) / 2**30, 3)}
        log(f"loaded train {len(tr['t']):,} val {len(va['t']):,} rows; host arrays {mem['arrays_train_gb'] + mem['arrays_val_gb']:.2f} GiB, RSS {mem['rss_after_load_gb']:.2f} GiB", args.out)
        raws = {} if args.no_lock else load_pred_raws(args, cfg, lrows)     # inline mode: predict after training (no lock wait to save; ~8 GB less RSS)
        mem["rss_with_pred_packs_gb"] = round(ds.rss_gb(), 2)
        (Path(args.out) / "memory.json").write_text(json.dumps(mem, indent=1))
        log(f"prediction packs loaded: {sum(len(r['t']) for r in raws.values()):,} rows; RSS {mem['rss_with_pred_packs_gb']:.2f} GiB", args.out)
        while True:
            n_acq += 1
            code = (train_then_predict(args, cfg, tr, va, st, bank, info, raws) if args.no_lock
                    else T.run_under_lock(args, train_then_predict, args, cfg, tr, va, st, bank, info, raws))
            if code == 0:
                break
            if code != 3:
                raise SystemExit(f"child exited with {code}")
            time.sleep(5)
        del tr, va, raws
        import gc
        gc.collect()
    if not (Path(args.out) / "PRED_DONE").exists():
        raws = load_pred_raws(args, cfg, lrows)
        if raws:
            log(f"prediction packs loaded: {sum(len(r['t']) for r in raws.values()):,} rows, RSS {ds.rss_gb():.2f} GiB", args.out)
            n_acq += 1
            code = predict_many(args, cfg, raws, st, bank, info) if args.no_lock else T.run_under_lock(args, predict_many, args, cfg, raws, st, bank, info)
            if code != 0:
                raise SystemExit(f"prediction child exited with {code}")
            del raws
            import gc
            gc.collect()
    code_files = sorted((CWM3 / "models").glob("*.py")) + [CWM3 / "train" / "train_cwm3_joint.py", CWM3 / "train" / "train_cwm3.py", CWM3 / "train" / "batch.py",
                                                            CWM3 / "eval" / "write_pred_layout_joint.py", CWM3 / "data" / "cwm3_dataset.py", EXPLORE / "s12_wm/eval_s13.py"]
    run_cfg = {"stage": cfg["stage"], "variant": cfg["variant"], "cfg": cfg}
    layouts = []
    if not args.synthetic:
        for regime in cfg["regimes"]:
            for split in cfg["pred_splits"]:
                if (Path(args.out) / f"pred_{regime}_{split}.npz").exists() and not (Path(args.out) / f"pred_layout_{regime}_{split}.npz").exists():
                    layouts.append(str(wplj.write_layout(args.out, regime, split, run_cfg, code_files, log=lambda m: log(m, args.out))))
    hist = [json.loads(l) for l in (Path(args.out) / "history.jsonl").read_text().splitlines()]
    td = json.loads((Path(args.out) / "TRAIN_DONE").read_text())
    pj = Path(args.out) / "params.json"
    info["params"] = json.loads(pj.read_text()) if pj.exists() else {}
    if not td.get("params"):
        td["params"] = info["params"]
    stamp = time.strftime("%Y%m%dT%H%M%SZ", time.gmtime())
    n_rows = append_quick_rows(args, cfg, stamp, td)
    qm = {f"{r}_{s}": json.loads((Path(args.out) / f"metrics_quick_{r}_{s}.json").read_text()) for r in cfg["regimes"] for s in cfg["pred_splits"]
          if (Path(args.out) / f"metrics_quick_{r}_{s}.json").exists()}
    inputs = [] if args.synthetic else sorted({p for d in base_dirs(cfg, args.seed).values() for p in [d / (d / "BEST").read_text().strip()]}
                                              | set(T.stage_input_files({"regimes": cfg["regimes"], "branch": "V", "can_hist": "8s" if "X8" in need_streams(cfg) else "4s"}))
                                              | set(T.stage_input_files({"regimes": cfg["regimes"], "branch": "U", "cabin": "none"}))
                                              | set(T.stage_input_files({"regimes": cfg["regimes"], "branch": "E"}))
                                              | {p for r in cfg["regimes"] for sp in ("train", "val", "test") for p in (PRESSURE / f"pressure_{r}_{sp}.npz", RESPONSE / f"response_{r}_{sp}.npz") if p.exists()})
    mem_p = Path(args.out) / "memory.json"
    write_run_json(args.out, f"train_{args.stage}_{args.variant}_s{args.seed}", code_files, inputs, seed=args.seed, t_start=t_start, extra={
        "arm": wplj.ARM_OF.get(cfg["stage"], "cwm3_I"), "stage": args.stage, "variant": args.variant, "cfg": cfg, "synthetic": args.synthetic, "max_anchors": args.max_anchors,
        "params": info["params"], "base_ckpts": info["base_ckpts"], "base_cfgs": info["base_cfgs"], "tok_source": cfg.get("tok_source"), "pack_version": args.pack_version,
        "lock_acquisitions_this_invocation": n_acq,
        "epochs_run": td["epochs_run"], "best_epoch": td["best_epoch"], "val_objective": td["best"], "history": hist, "quick_metrics": qm, "quick_rows_appended": n_rows,
        "run_stamp": stamp, "memory": json.loads(mem_p.read_text()) if mem_p.exists() else {}, "roi": wplj.roi_of(ks), "roi_streams": ["Kh", "K", "S"],
        "external_priors": wplj.external_priors(ks, kv), "dms_fields_used": [], "label_rows": {r: int(len(v)) for r, v in lrows.items()},
        "kin_source": ks, "kin_version": kv, "kin_dir": (str(ds.kin_dir()) if ks != "cwm2" else ""), "stats_mode": cfg.get("stats_mode", "merged"),
        "stats_version": cfg.get("stats_version"), "base_stats_versions": cfg.get("base_stats_versions", {}), "base_dirs": info.get("base_dirs", {}), "run_suffix": args.run_suffix,
        "u_gaze_chain": bool(cfg.get("u_gaze_chain", False)), "u_face": bool(cfg.get("u_face", False)),     # gaze U bases (2026-09-12 appendix variant)
        "u_kpf": cfg.get("u_kpf", ""),                                                                     # LANE F1 (2026-09-17): forecaster run dir of the kpf channel
        "mask_V": bool(cfg.get("mask_V", False)),                                                          # variant `ue_noV` (2026-09-18): vehicle branch masked as missing
        "p_head_source": info["params"].get("p_head_source", ""),
        "latent_aux": (dict(cfg["latent_aux"]) if latent_enabled(cfg) else None), "latent_aux_lam": latent_lam(cfg),      # C4 (2026-09-17)
        "mtl": cfg.get("mtl", "none"), "mtl_probe": (json.loads((Path(args.out) / "mtl_probe.json").read_text()) if (Path(args.out) / "mtl_probe.json").exists() else {}),
        "selection_objective": "validation sum of chain NLLs (V terms on V-valid rows)" + (" + 0.1 L_F" if cfg["recipe"].get("select_LF", True) else " (L_F excluded: select_LF false)") + "; MDN / Gaussian-K excluded",
        "pred_layouts": layouts,
        "note_kh_causality": (f"Kh / K / S from the CAUSAL route-level pose ({ds.kin_dir().name}); rows without a causal box are not scored (valid_U)" if ks == "causal_routes" else
                              "Kh / K / S come from the route-median crop pose cache (NON-causal ROI): driver-leg numbers are diagnostic (K4 pending)")})
    (Path(args.out) / "PRED_DONE").write_text(json.dumps({"splits": cfg["pred_splits"], "regimes": cfg["regimes"], "label_rows": {r: int(len(v)) for r, v in lrows.items()}}))
    log("run.json written; DONE", args.out)


if __name__ == "__main__":
    main()
Xet Storage Details

Size:: 77.1 kB
Xet hash:: 819ee36363e8695dcc9fd0b590450e0d2e4d72796dcb08581b6c9926cd2519d1
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.