Buckets:
| """cwm3 stage trainer — one base branch (V / U / E), one variant, one seed, both regimes in ONE model (plan §3; train/stages.yaml). | |
| python train_cwm3.py --stage U0 --variant kin_can --seed 0 [--max-anchors N] [--synthetic N] [--minutes 18] [--epochs E] [--out DIR] | |
| GPU-lock discipline (copied from cwm2/train/train_grid.py): the PARENT never touches CUDA. It loads the stage's compact arrays | |
| (data/cwm3_dataset.py, both regimes, mode column) into host RAM, then loops: flock(/tmp/s13_gpu.lock) -> fork -> child initialises | |
| CUDA, resumes the latest checkpoint, trains until the wall budget (--minutes, <= 22) / epoch budget / early stop, saves | |
| ckpt_ep{E}_b{pos}.pt (never overwritten; superseded ones pruned into ckpt_pruned.jsonl) -> exit 0 (done) | 3 (continue) -> parent | |
| waits, releases the lock, sleeps 5 s. After TRAIN_DONE the parent frees train/val, loads each (regime, split) prediction pack and | |
| runs the prediction phase in one more acquisition per split, then writes pred_layout_{regime}_{split}.npz (CPU, eval/write_pred_layout.py). | |
| Selection objective (BEST / early stop): validation chain NLL (+ 0.1 L_F for E) — models/heads.selection_objective. | |
| --max-anchors subsets train and the val-objective rows only (prediction always runs on the full split); --synthetic N replaces | |
| the data layer with random arrays of the same contract (smoke / tests; never a result). | |
| Outputs under --out (default results/runs/{stage}_{variant}_s{seed}): ckpt_*.pt, BEST, history.jsonl, TRAIN_DONE, stats.npz, | |
| pred_{regime}_{split}.npz, pred_layout_{regime}_{split}.npz, metrics_quick_{regime}_{split}.json, run_*.json, PRED_DONE; one row per | |
| (chain, metric) 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 argparse | |
| import csv | |
| import fcntl | |
| import json | |
| import os | |
| import shutil | |
| import subprocess | |
| import sys | |
| import time | |
| from pathlib import Path | |
| import numpy as np | |
| import yaml | |
| CWM3 = Path(__file__).resolve().parents[1] | |
| EXPLORE = CWM3.parent | |
| S12R = EXPLORE / "s12_wm" / "results" | |
| 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 | |
| from models import build_branch, params_by_module # noqa: E402 | |
| from models import heads as hd # noqa: E402 | |
| from runmeta import write_run_json # noqa: E402 | |
| import write_pred_layout as wpl # noqa: E402 | |
| from models import identity as idn # noqa: E402 opt-in innovation arms (u_adv / u_devnorm / f_state_sup); no-ops otherwise | |
| LOCK = "/tmp/s13_gpu.lock" | |
| STAGES = yaml.safe_load((CWM3 / "train" / "stages.yaml").read_text()) | |
| QUICK_CSV = CWM3 / "results" / "cwm3_quick_metrics.csv" | |
| QUICK_COLS = ["run_stamp", "stage", "variant", "seed", "regime", "split", "branch", "chain", "metric", "value", "base_rate", "n", | |
| "n_rows_predicted", "best_epoch", "val_objective", "run_dir"] | |
| CHAIN_NAMES = {"V": list(hd.V_CHAINS), "U": list(hd.U_CHAINS), "E": list(hd.E_CHAINS)} | |
| def log(msg, out=None): | |
| line = f"[{time.strftime('%H:%M:%S')}] {msg}" | |
| print(line, flush=True) | |
| if out is not None: | |
| with open(Path(out) / "train.log", "a") as f: | |
| f.write(line + "\n") | |
| # ----------------------------------------------------------------------------- config | |
| def stage_cfg(stage, variant): | |
| s = STAGES["stages"][stage] | |
| cfg = dict(STAGES["defaults"]) | |
| cfg.update({"stage": stage, "variant": variant, "branch": s["branch"], "recipe": dict(s["recipe"]), "lambda": dict(s.get("lambda", {})), | |
| "model": dict(s.get("model", {})), "experiment": s.get("experiment", "")}) | |
| v = dict(s["variants"][variant]) | |
| vr = v.pop("recipe", None) # variant-level recipe keys MERGE into the stage recipe (2026-09-11: `select_LF: false` variants) | |
| cfg.update(v) | |
| if vr: | |
| cfg["recipe"].update(vr) | |
| return cfg | |
| def model_kwargs(cfg): | |
| b, m = cfg["branch"], dict(cfg["model"]) | |
| if cfg.get("arch"): # 2026-09-12 (GPU-baselines B): architecture-swap variants (models/branches_alt.py); absent -> unchanged | |
| m["arch"] = cfg["arch"] | |
| if b == "U": | |
| m["cabin"] = cfg.get("cabin", "none") | |
| m["in_ln"] = bool(cfg["recipe"].get("u_in_ln", False)) # opt-in input LayerNorm over the 12 kinematics channels (2026-09-11) | |
| if cfg["recipe"].get("u_adv"): | |
| m["adv_devices"] = int(cfg.get("adv_devices", 0)) # identity arm: GRL device head (models/identity.py; count set by prepare_raws) | |
| if cfg["recipe"].get("u_gaze_chain"): | |
| m["gaze_chain"] = True # gaze_off driver chain (2026-09-12): separate 2-logit head, targets results/targets_gaze_causal_routes_v2 | |
| if cfg["recipe"].get("u_face"): | |
| m["n_face"] = ds.N_FACE # 9 causal face-landmark channels (+ valid bit) appended to the GRU input | |
| if cfg["recipe"].get("u_can_gate"): | |
| m["can_gate"] = True # CAN release gate on the driver-chain hazards (2026-09-14, plan Part A §A.5 arm (b); models/branches.CanReleaseGate) | |
| if cfg.get("n_kin"): | |
| m["n_kin"] = int(cfg["n_kin"]) # kin v2p (2026-09-18): 14 input kinematics columns (11 + 3 keypoint proxies); absent -> 11 | |
| if b == "E": | |
| m["tokens"], m["n_slots"] = cfg.get("tokens", "region"), int(cfg.get("n_slots", 16)) | |
| m["d_side"] = 512 if cfg.get("side", "vf256") == "regionpca256" else 256 # side vector width (regionpca256 = 512, vf256 = 256) | |
| m["p_head"] = bool(cfg.get("p_head", False)) # road-pressure chains P (variants *_P, 2026-09-11) | |
| if cfg["recipe"].get("f_state_sup"): | |
| m["f_state"] = len(idn.FS_NAMES) # state-supervised F read-out (innovation 3 b2) | |
| return m | |
| def need_streams(cfg): | |
| b = cfg["branch"] | |
| if b == "V": | |
| return frozenset({"X", "Y", "demand"} | ({"X8"} if cfg.get("can_hist") == "8s" else set())) | |
| if b == "U": | |
| cab = cfg.get("cabin", "none") | |
| rec = cfg.get("recipe", {}) | |
| return frozenset({"X", "Kh", "kin", "Y"} | ({cab} if cab != "none" else set()) | ({"gaze"} if rec.get("u_gaze_chain") else set()) | ({"face"} if rec.get("u_face") else set())) | |
| return frozenset({"Y", "demand", "F", cfg.get("side", "vf256"), "tokens"} | ({"pressure"} if cfg.get("p_head") else set())) # side = vf256 | regionpca256 | |
| FILE_KEYS = {"X": ("targets",), "Y": ("targets",), "X8": ("xpre",), "Kh": ("kinhist",), "kin": ("kin",), "demand": ("demand",), "F": ("future",), | |
| "vf256": ("vf256",), "vc256": ("vc256",), "crop256": ("crop256",), "tokens": ("anchor_index", "bank"), "regionpca256": ("regionpca256",), | |
| "pressure": ("pressure",), "gaze": ("gaze",), "face": ("face",), "vfcan256": ("vfcan256",), "can256": ("can256",), "vfcan256s": ("vfcan256s",)} # vfcan256 / can256: F5b side streams (2026-09-16) | |
| def stage_input_files(cfg): | |
| """the data files this stage actually reads (hashed into run.json), de-duplicated.""" | |
| keys = {"targets", "target_stats"} | |
| for s in need_streams(cfg): | |
| keys.update(FILE_KEYS[s]) | |
| out = set() | |
| for r in cfg["regimes"]: | |
| for sp in ("train", "val", "test"): | |
| for k, p in ds.files(r, sp).items(): | |
| if k in keys and p.exists(): | |
| out.add(p) | |
| return sorted(out) | |
| def roi_streams(cfg): | |
| b = cfg["branch"] | |
| if b == "U": | |
| rec = cfg.get("recipe", {}) | |
| return ["Kh", "K", "S"] + (["crop256"] if cfg.get("cabin") == "crop256" else []) + (["Sg"] if rec.get("u_gaze_chain") else []) + (["Fh"] if rec.get("u_face") else []) | |
| return [] | |
| # ----------------------------------------------------------------------------- data (parent) | |
| 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 = ds.concat_regimes(parts) if len(parts) > 1 else dict(parts[0], regime_row=parts[0]["idx"]) | |
| return raw | |
| mx = (args.max_anchors or None) if subset else None | |
| return ds.load_regimes(regimes, split, need, max_anchors=mx, seed=args.seed + 7, verbose=lambda m: log(m, args.out)) | |
| def load_bank(args, cfg): | |
| if cfg["branch"] != "E": | |
| return None | |
| if args.synthetic: | |
| return ds.synthetic_bank() | |
| return ds.load_bank(need_tok=cfg.get("tokens", "region") in ("region", "spat256"), need_pooled=cfg.get("tokens") == "pooled") | |
| def target_stats(args): | |
| if args.synthetic: | |
| return {"mu": np.array([0.0, 12.0, 0.0], np.float32), "sd": np.array([70.0, 8.0, 0.7], np.float32)} | |
| z = np.load(S12R / "target_stats.npz") | |
| assert list(z["cont_idx"]) == [0, 3, 4] | |
| return {"mu": z["mu"], "sd": z["sd"]} | |
| # ----------------------------------------------------------------------------- checkpoints | |
| def prune_ckpts(out, keep): | |
| for ck in sorted(Path(out).glob("ckpt_ep*_b*.pt")): | |
| if ck.name not in keep: | |
| with open(Path(out) / "ckpt_pruned.jsonl", "a") as f: | |
| f.write(json.dumps({"deleted": ck.name, "bytes": ck.stat().st_size, "kept": sorted(k for k in keep if k)}) + "\n") | |
| ck.unlink() | |
| def latest_ckpt(out): | |
| cks = sorted(Path(out).glob("ckpt_ep*_b*.pt"), key=lambda p: p.stat().st_mtime) | |
| return cks[-1] if cks else None | |
| def make_optim(model, rec, epochs): | |
| import torch | |
| opt = torch.optim.AdamW(model.parameters(), 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 | |
| def val_objective(model, store, rows, st, cfg, device, amp, bs=8192): | |
| import torch | |
| model.eval() | |
| branch = cfg["branch"] | |
| tot_sel, tot, n = 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 = idn.prep_for(cfg)(store, idx, st, cfg) | |
| with torch.autocast("cuda", dtype=torch.bfloat16, enabled=amp): | |
| out = model(b) | |
| loss, parts = bt.LOSS[branch](out, tg, cfg["lambda"]) | |
| tot_sel += hd.selection_objective(parts, bool(cfg['recipe'].get('select_LF', True))) * len(idx) | |
| tot += float(loss) * len(idx) | |
| n += len(idx) | |
| return tot_sel / max(n, 1), tot / max(n, 1) | |
| def run_chunk(args, cfg, tr, va, st, bank): | |
| import torch | |
| t0 = time.time() | |
| out = Path(args.out) | |
| device = "cuda" if torch.cuda.is_available() and not args.cpu else "cpu" | |
| branch, rec = cfg["branch"], 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_branch(branch, args.seed, **model_kwargs(cfg)).to(device) | |
| pb = params_by_module(model) | |
| opt, sched = make_optim(model, rec, epochs) | |
| state = {"epoch": 0, "batch_pos": 0, "best": 1e18, "best_epoch": -1, "patience": 0, "acq": 0, "done": False} | |
| ck = 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"]) | |
| 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 = bt.train_rows(tr, branch, cfg) | |
| rows_va = bt.train_rows(va, branch, cfg) | |
| if rec.get("val_rows", "all") == "innovation": | |
| sel = rows_va[va["slice_tag"][rows_va] > 0] | |
| rows_va = sel if len(sel) else rows_va | |
| store_tr = bt.GpuStore(tr, bt.KEYS[branch] + idn.store_keys(cfg), device, bank) | |
| store_va = bt.GpuStore(va, bt.KEYS[branch] + idn.store_keys(cfg), device, None if bank is None else bank) | |
| log(f"chunk {state['acq']} start: {cfg['stage']}/{cfg['variant']} s{args.seed} device={device} amp={amp} params={pb['total']:,} " | |
| f"train rows {len(rows_tr):,} val rows {len(rows_va):,} gpu store {(store_tr.bytes() + store_va.bytes()) / 2**30:.2f} GiB " | |
| f"epochs={epochs} rows/epoch={rec['rows_per_epoch'] or 'all'} batch={rec['batch']} ({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 = int(rec["batch"]) | |
| rows_per_epoch = 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}, 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 = idn.prep_for(cfg)(store_tr, idx, st, cfg) | |
| with torch.autocast("cuda", dtype=torch.bfloat16, enabled=amp): | |
| out_ = model(b) | |
| loss, parts = bt.LOSS[branch](out_, tg, cfg["lambda"]) | |
| loss = idn.extra_losses(loss, parts, out_, tg, store_tr, idx, cfg) # opt-in arms: adversarial device term / state read-out (identity.py) | |
| opt.zero_grad(set_to_none=True) | |
| loss.backward() | |
| torch.nn.utils.clip_grad_norm_(model.parameters(), float(cfg.get("grad_clip", 1.0))) | |
| opt.step() | |
| tot += float(loss) | |
| nbatch += 1 | |
| state["batch_pos"] += 1 | |
| 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()), out) | |
| 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 = val_objective(model, store_va, rows_va, st, cfg, device, amp) | |
| 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 | |
| rec_h = {"epoch": ep_, "train_loss": tot / max(nbatch, 1), "val_objective": v, "val_objective_kind": "chains (+L_F)", "val_loss_full": v_full, | |
| "best": state["best"], "best_epoch": state["best_epoch"], "lr": opt.param_groups[0]["lr"], "acq": state["acq"], | |
| "t_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())} | |
| 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} val_full {v_full:.4f} best {state['best']:.4f}@{state['best_epoch']} " | |
| f"({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") | |
| 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) | |
| return 0 | |
| if over_budget(): | |
| log("wall budget reached at epoch end; exit 3", out) | |
| return 3 | |
| return 0 | |
| # ----------------------------------------------------------------------------- child: prediction | |
| def predict_split(args, cfg, raw, st, bank, regime, split): | |
| import torch | |
| import eval_s13 as e13 | |
| t0 = time.time() | |
| out = Path(args.out) | |
| branch = cfg["branch"] | |
| 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 = build_branch(branch, args.seed, **model_kwargs(cfg)).to(device) | |
| best = (out / "BEST").read_text().strip() | |
| z = torch.load(out / best, map_location=device, weights_only=False) | |
| model.load_state_dict(z["model"]) | |
| model.eval() | |
| N = len(raw["t"]) | |
| store = bt.GpuStore(raw, bt.KEYS[branch] + idn.store_keys(cfg), device, bank) | |
| valid = np.ones(N, bool) | |
| if branch == "V" and cfg.get("can_hist") == "8s": | |
| valid = raw["x8_ok"].astype(bool) | |
| if branch == "U" and "kin_row_ok" in raw: | |
| valid = raw["kin_row_ok"].astype(bool) # causal kin: rows without a causal box are not predicted | |
| C = len(CHAIN_NAMES[branch]) | |
| 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)} | |
| if branch == "V": | |
| P.update(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)) | |
| if branch == "E": | |
| P["f_cos"] = np.full((N, 5), np.nan, np.float32) | |
| p_head = branch == "E" and bool(cfg.get("p_head", False)) | |
| if p_head: # road-pressure chains: extra blocks beside the D chains | |
| CP = len(hd.P_CHAINS) | |
| P.update(hz_on_P=np.zeros((N, 25, CP), np.float16), hz_off_P=np.zeros((N, 25, CP), np.float16), z0_P=np.zeros((N, CP), np.float16), | |
| z0valid_P=np.ones((N, CP), bool)) | |
| if branch == "U": | |
| P["k_mae"] = np.full(N, np.nan, np.float32) | |
| if bool(cfg["recipe"].get("u_can_gate", False)): | |
| P["gate2"] = np.full(N, np.nan, np.float32) # 2g per row of the CAN release gate (2026-09-14; diagnostic — the layout writer ignores it) | |
| g_chain = branch == "U" and bool(cfg["recipe"].get("u_gaze_chain", False)) | |
| if g_chain: # gaze_off chain (2026-09-12): extra blocks beside the hands / head chains | |
| 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)) | |
| TG = {"zp": [], "z": [], "m": []} | |
| TGP = {"zp": [], "z": [], "m": []} | |
| TGG = {"zp": [], "z": [], "m": []} | |
| ymu, ysd = torch.as_tensor(st["ymu"], device=device), torch.as_tensor(st["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 = idn.prep_for(cfg)(store, idx, st, cfg) | |
| with torch.autocast("cuda", dtype=torch.bfloat16, enabled=amp): | |
| o = model(b) | |
| if branch == "V": | |
| v = hd.split_V(o["V"]) | |
| h_on, h_off = hd.hazards(v["chains"]) | |
| zp, zz, mm = tg["zV_prev"], tg["zV"], tg["zVm"] | |
| z0 = tg["z0"] | |
| 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["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 # [S,B,T,3] | |
| t1 = (smp - Y[None]).abs().mean(0) | |
| t2 = (smp[:, None] - smp[None]).abs().mean((0, 1)) | |
| crps = (t1 - 0.5 * t2) * M | |
| P["y_crps"][sl] = crps.sum(1).div(M.sum(1).clamp_min(1)).cpu().numpy() | |
| P["pi_resp"][sl] = hd.forward_marginal(h_on[..., 3], h_off[..., 3], z0[:, 3]).cpu().numpy() | |
| P["z0valid"][sl, 3] = store.t["Resp0valid"][idx].cpu().numpy() | |
| elif branch == "U": | |
| u = hd.split_U(o["U"]) | |
| h_on, h_off = hd.hazards_gated(u["chains"], o["HG"]) if "HG" in o else hd.hazards(u["chains"]) # "HG" = 2g of the CAN release gate (2026-09-14): λ' goes into the pack | |
| if "HG" in o and "gate2" in P: | |
| P["gate2"][sl] = o["HG"][:, 0].float().cpu().numpy() | |
| zp, zz, mm = tg["zU_prev"], tg["zU"], tg["zUm"] | |
| z0 = tg["z0"] | |
| P["z0valid"][sl] = tg["z0valid"].cpu().numpy() | |
| K_mu = tg["K0std"][:, None, :] + u["dK_mu"] | |
| Km = tg["Km"].float() | |
| P["k_mae"][sl] = (((K_mu - tg["Kstd"]).abs() * Km).sum((1, 2)) / Km.sum((1, 2)).clamp_min(1)).cpu().numpy() | |
| if g_chain and "G" in o and "zG" in tg: | |
| 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] = tg["zG0"][:, None].cpu().numpy(), tg["zG0valid"][:, None].cpu().numpy() | |
| TGG["zp"].append(tg["zG_prev"][:, :, None].cpu().numpy().astype(np.float32)); TGG["z"].append(tg["zG"][:, :, None].cpu().numpy().astype(np.float32)); TGG["m"].append(tg["zGm"][:, :, None].cpu().numpy()) | |
| else: | |
| e = hd.split_E(o["E"]) | |
| h_on, h_off = hd.hazards(e["chains"]) | |
| zp, zz, mm = tg["zD_prev"], tg["zD"], tg["zDm"] | |
| z0 = tg["z0"] | |
| P["z0valid"][sl] = tg["z0valid"].cpu().numpy() | |
| cos = torch.nn.functional.cosine_similarity(o["F"].float(), tg["F"].float(), dim=-1) | |
| P["f_cos"][sl] = torch.where(tg["Fok"], cos, torch.full_like(cos, float("nan"))).cpu().numpy() | |
| if p_head and "P" in o and "zP" in tg: | |
| hp_on, hp_off = hd.hazards(hd.split_P(o["P"])["chains"]) | |
| P["hz_on_P"][sl], P["hz_off_P"][sl], P["z0_P"][sl], P["z0valid_P"][sl] = hp_on.cpu().numpy(), hp_off.cpu().numpy(), tg["zP0"].cpu().numpy(), tg["zP0valid"].cpu().numpy() | |
| TGP["zp"].append(tg["zP_prev"].cpu().numpy().astype(np.float32)); TGP["z"].append(tg["zP"].cpu().numpy().astype(np.float32)); TGP["m"].append(tg["zPm"].cpu().numpy()) | |
| P["hz_on"][sl], P["hz_off"][sl], P["z0"][sl] = h_on.cpu().numpy(), h_off.cpu().numpy(), z0.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 (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()} | |
| # quick sanity metrics (eval_s13.transition_metrics contract) on the predicted rows | |
| qm = {"n_rows": int(N), "n_rows_predicted": int(valid.sum()), "regime": regime, "split": split, "branch": branch} | |
| chains = CHAIN_NAMES[branch] | |
| hz_on, hz_off = P["hz_on"].astype(np.float64), P["hz_off"].astype(np.float64) | |
| events = [(c, c, "on") for c in chains] | |
| if branch == "U": | |
| events = [("hand_return", "hands_on", "on"), ("hand_leave", "hands_on", "off"), ("head_turn_onset", "head_turned", "on"), ("head_return", "head_turned", "off")] | |
| for name, ch, kind in events: | |
| c = chains.index(ch) | |
| zp, zz, mm = TG["zp"][valid][..., c], TG["z"][valid][..., c], TG["m"][valid][..., c].astype(bool) | |
| if kind == "off": | |
| zp, zz, h = 1 - zp, 1 - zz, hz_off[valid][..., c] | |
| else: | |
| h = hz_on[valid][..., 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: # too few rows at smoke scale | |
| qm[f"{name}:error"] = str(ex)[:120] | |
| if p_head and TGP["z"]: | |
| TGP = {k: np.concatenate(v) for k, v in TGP.items()} | |
| hpo = P["hz_on_P"].astype(np.float64) | |
| for c, name in enumerate(hd.P_CHAINS): | |
| zp, zz, mm = TGP["zp"][valid][..., c], TGP["z"][valid][..., c], TGP["m"][valid][..., c].astype(bool) | |
| try: | |
| for r in e13.transition_metrics(hpo[valid][..., c], zp, zz, mm, f"P_{name}"): | |
| qm[f"P_{name}:{r['metric']}"] = r["value"] | |
| if "base_rate" in r: | |
| qm[f"P_{name}:{r['metric']}:base_rate"], qm[f"P_{name}:{r['metric']}:n"] = r["base_rate"], r["n"] | |
| except Exception as ex: # too few rows at smoke scale | |
| qm[f"P_{name}:error"] = str(ex)[:120] | |
| if g_chain and TGG["z"]: | |
| TGG = {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 = TGG["zp"][valid][..., 0], TGG["z"][valid][..., 0], TGG["m"][valid][..., 0].astype(bool) | |
| if kind == "off": | |
| zp, zz, h = 1 - zp, 1 - zz, P["hz_off_G"].astype(np.float64)[valid][..., 0] | |
| else: | |
| h = P["hz_on_G"].astype(np.float64)[valid][..., 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: # too few rows at smoke scale | |
| qm[f"{name}:error"] = str(ex)[:120] | |
| for k in ("y_mae", "y_crps", "f_cos", "k_mae"): | |
| if k in P: | |
| v = P[k][valid] | |
| qm[k + ":mean"] = np.nanmean(v, 0).tolist() if v.ndim == 2 else float(np.nanmean(v)) | |
| qm["wall_s"] = round(time.time() - t0, 1) | |
| (out / f"metrics_quick_{regime}_{split}.json").write_text(json.dumps(qm, indent=1)) | |
| np.savez_compressed(out / f"pred_{regime}_{split}.npz", **P, valid=valid, regime=np.array(regime), split=np.array(split), branch=np.array(branch), | |
| 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), hazard_kind=np.array("head"), n_rows=np.array(N), | |
| **({"p_chain_names": np.array(list(hd.P_CHAINS))} if p_head else {}), **({"g_chain_names": np.array(list(hd.G_CHAINS))} if g_chain else {})) | |
| log(f"pred_{regime}_{split}.npz written ({int(valid.sum())}/{N} rows) in {time.time() - t0:.0f}s; " + | |
| " ".join(f"{k}={v:.3f}" for k, v in qm.items() if k.endswith("onset_prauc@3s")), out) | |
| return 0 | |
| def pred_ok(p): | |
| """a prediction pack counts as present only if it loads (a trainer killed mid-write leaves a truncated npz; 2026-09-11) — a corrupt | |
| file is renamed *.corrupt and re-predicted.""" | |
| p = Path(p) | |
| if not p.exists(): | |
| return False | |
| try: | |
| with np.load(p) as z: | |
| _ = z["hz_on"].shape | |
| return True | |
| except Exception: # noqa: BLE001 | |
| p.rename(p.with_name(p.name + ".corrupt")) | |
| return False | |
| def predict_many(args, cfg, raws, st, bank): | |
| for (regime, split), raw in raws.items(): | |
| if pred_ok(Path(args.out) / f"pred_{regime}_{split}.npz"): | |
| continue | |
| code = predict_split(args, cfg, raw, st, bank, regime, split) | |
| if code: | |
| return code | |
| return 0 | |
| def load_pred_raws(args, cfg): | |
| """every (regime, split) prediction pack not yet predicted (full splits, row-aligned with the targets files).""" | |
| raws = {} | |
| for split in cfg["pred_splits"]: | |
| for regime in cfg["regimes"]: | |
| if not pred_ok(Path(args.out) / f"pred_{regime}_{split}.npz"): | |
| raws[(regime, split)] = load_raw(args, cfg, split, subset=False, regimes=[regime]) | |
| return raws | |
| PREDICT_MIN_LEFT_S = 240 | |
| def train_then_predict(args, cfg, tr, va, st, bank, raws): | |
| """ONE lock acquisition: a training chunk, and — when training finished with >= PREDICT_MIN_LEFT_S of the wall budget left — the | |
| prediction phase in the same child (saves one 20-min lock wait per run when the GPU is shared). Exit codes as run_chunk.""" | |
| t0 = time.time() | |
| code = run_chunk(args, cfg, tr, va, st, bank) | |
| if code != 0: | |
| return code | |
| left = args.minutes * 60 - (time.time() - t0) | |
| if raws and left >= 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) | |
| return 0 | |
| # ----------------------------------------------------------------------------- parent: lock loop (train_grid.py verbatim) | |
| def run_under_lock(args, fn, *data): | |
| fd = os.open(LOCK, os.O_RDWR | os.O_CREAT, 0o664) | |
| t_wait = time.time() | |
| fcntl.flock(fd, fcntl.LOCK_EX) | |
| log(f"lock acquired after {time.time() - t_wait:.0f}s wait", args.out) | |
| try: | |
| pid = os.fork() | |
| if pid == 0: | |
| code = 1 | |
| try: | |
| code = fn(*data) | |
| except BaseException as e: # noqa: BLE001 | |
| import traceback | |
| traceback.print_exc() | |
| log(f"child failed: {e!r}", args.out) | |
| code = 1 | |
| finally: | |
| sys.stdout.flush() | |
| os._exit(code) | |
| _, status = os.waitpid(pid, 0) | |
| code = os.waitstatus_to_exitcode(status) | |
| finally: | |
| fcntl.flock(fd, fcntl.LOCK_UN) | |
| os.close(fd) | |
| return code | |
| 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": cfg["branch"], "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 QUICK_CSV.exists() | |
| with open(QUICK_CSV, "a", newline="") as f: | |
| w = csv.DictWriter(f, fieldnames=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("--max-anchors", type=int, default=0, help="seeded row subset of train / val-objective rows (smoke)") | |
| ap.add_argument("--synthetic", type=int, default=0, help="N synthetic train rows instead of the data layer (tests)") | |
| ap.add_argument("--epochs", type=int, default=0, help="0 = stage recipe") | |
| ap.add_argument("--minutes", type=float, default=0, help="wall budget per lock acquisition (0 = stages.yaml default 18; hard max 22)") | |
| ap.add_argument("--regimes", default="", help="override: comma list (default both)") | |
| ap.add_argument("--out", default="") | |
| ap.add_argument("--pred-splits", default="") | |
| ap.add_argument("--no-lock", action="store_true", help="tests only: run the child inline without the GPU lock") | |
| ap.add_argument("--no-amp", action="store_true") | |
| ap.add_argument("--cpu", action="store_true") | |
| ap.add_argument("--pack-version", default="v1", choices=["v1", "v2"], help="v2 = results/pack_v2 vf256 / vc256 + aligned F (2026-09-10 latent audit)") | |
| ap.add_argument("--kin-source", default="cwm2", choices=["cwm2", "causal_routes", "oldin_ctgt"], | |
| help="K4: causal_routes = causal inputs + causal targets; oldin_ctgt = OLD route-median inputs + CAUSAL targets (input-isolation control)") | |
| ap.add_argument("--kin-version", default="v1", choices=["v1", "v2", "v3", "v2p"], | |
| help="causal kin pack build: v1 = results/targets_kin_causal_routes, v2 = …_v2 (population driver-side prior on 74 routes); applies to " | |
| "causal_routes and oldin_ctgt; run dir suffix _kcv<1|2|3> / _oldin_ctgt[_v2]; v3 (2026-09-16) = …_v3 (wheel-reference hands_on, 1 s persistence; inputs identical to v2)") | |
| ap.add_argument("--mask-routes", default="", help="json with `flagged` [{uid,...}] (driver-box suspects): their rows are removed from TRAIN (protocol b)") | |
| ap.add_argument("--stats-version", type=int, default=bt.STATS_VERSION, choices=[1, 2], | |
| help="standardisation statistics: 2 (default) = float64; 1 = the pre-fix float32 rule (disentangling arm only; dir suffix _stats1)") | |
| args = ap.parse_args() | |
| ds.set_pack_version(args.pack_version) | |
| ds.set_kin_source(args.kin_source, args.kin_version) | |
| cfg = stage_cfg(args.stage, args.variant) | |
| cfg["pack_version"], cfg["kin_source"], cfg["kin_version"], cfg["mask_routes"] = args.pack_version, args.kin_source, args.kin_version, args.mask_routes | |
| cfg["stats_version"] = int(args.stats_version) | |
| if cfg["branch"] == "U": | |
| cfg["kclip"] = float(cfg["recipe"].get("kclip", bt.K_CLIP)) # NEW runs: 8 (= X_CLIP) unless the recipe pins another value; recorded (old runs: no key = 5) | |
| if args.kin_source == "causal_routes" and args.kin_version == "v2p": | |
| cfg["n_kin"] = int(ds.n_kin()) # kin v2p (2026-09-18, pre-reg 2026-09-18T05:34:40Z): 14 input columns; recorded so the joint rebuild uses the same width | |
| 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, "GPU-lock chunks must be <= 22 min" | |
| if not args.out: | |
| args.out = str(CWM3 / "results" / "runs" / (f"{args.stage}_{args.variant}_s{args.seed}" + ("_synthetic" if args.synthetic else "") | |
| + (f"_max{args.max_anchors}" if args.max_anchors else "") + ("_v2" if args.pack_version == "v2" else "") | |
| + (f"_kc{args.kin_version}" if args.kin_source == "causal_routes" else "") | |
| + (("_oldin_ctgt" + ("" if args.kin_version == "v1" else f"_{args.kin_version}")) if args.kin_source == "oldin_ctgt" else "") | |
| + ("_masked" if args.mask_routes else "") + ("_stats1" if int(args.stats_version) == 1 else ""))) | |
| Path(args.out).mkdir(parents=True, exist_ok=True) | |
| t_start = time.time() | |
| st_path = Path(args.out) / "stats.npz" | |
| idn.restore_cfg(args.out, cfg) # re-invoked opt-in runs: adv_devices recorded by an earlier invocation | |
| (Path(args.out) / "config.json").write_text(json.dumps({"cfg": cfg, "args": vars(args)}, indent=1, default=str)) | |
| log(f"=== train_cwm3 {args.stage}/{args.variant} s{args.seed} out={args.out} synthetic={args.synthetic} max_anchors={args.max_anchors} " | |
| f"regimes={cfg['regimes']} need={sorted(need_streams(cfg))} ===", args.out) | |
| bank = load_bank(args, cfg) | |
| n_acq = 0 | |
| mem = {} | |
| 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) if bank else 0} | |
| if args.mask_routes: | |
| flagged = {r["uid"] for r in json.loads(Path(args.mask_routes).read_text())["flagged"]} | |
| keep = np.flatnonzero(~np.isin(tr["uid"], list(flagged))) | |
| log(f"mask-routes: {len(tr['t']) - len(keep):,} train rows on {len(flagged)} flagged routes removed ({len(keep):,} kept)", args.out) | |
| tr = ds.subset(tr, keep) | |
| st = bt.compute_stats(tr, target_stats(args), version=args.stats_version) | |
| if not st_path.exists(): | |
| np.savez(st_path, **st) | |
| if idn.prepare_raws(args, cfg, tr, va, st, log=lambda m: log(m, args.out)): # opt-in arms (dev_id / devnorm / response bits); True -> cfg gained keys | |
| (Path(args.out) / "config.json").write_text(json.dumps({"cfg": cfg, "args": vars(args)}, indent=1, default=str)) | |
| log(f"loaded train {len(tr['t']):,} val {len(va['t']):,} rows; host arrays {mem['arrays_train_gb'] + mem['arrays_val_gb']:.2f} GiB, " | |
| f"RSS {mem['rss_after_load_gb']:.2f} GiB", args.out) | |
| raws = load_pred_raws(args, cfg) # loaded before the first acquisition: an early-finishing chunk predicts too | |
| idn.prepare_pred_raws(args, cfg, raws, st) | |
| mem["rss_with_pred_packs_gb"] = round(ds.rss_gb(), 2) | |
| mem["arrays_pred_gb"] = round(sum(ds.nbytes(r) for r in raws.values()) / 2**30, 3) | |
| (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, raws) if args.no_lock | |
| else run_under_lock(args, train_then_predict, args, cfg, tr, va, st, bank, 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() | |
| st = {k: v for k, v in np.load(st_path).items()} | |
| if not (Path(args.out) / "PRED_DONE").exists(): | |
| raws = load_pred_raws(args, cfg) # whatever the training acquisition did not predict (e.g. budget exhausted) | |
| idn.prepare_pred_raws(args, cfg, raws, st) | |
| 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) if args.no_lock else run_under_lock(args, predict_many, args, cfg, raws, st, bank) | |
| if code != 0: | |
| raise SystemExit(f"prediction child exited with {code}") | |
| del raws | |
| import gc | |
| gc.collect() | |
| code_files = sorted((CWM3 / "models").glob("*.py")) + sorted((CWM3 / "train").glob("*.py")) + [CWM3 / "eval" / "write_pred_layout.py", | |
| CWM3 / "data" / "cwm3_dataset.py", EXPLORE / "s12_wm/eval_s13.py"] | |
| run_cfg = {"stage": cfg["stage"], "variant": cfg["variant"], "cfg": cfg, "roi_streams": roi_streams(cfg), | |
| **({"roi": "causal_clip_roi"} if args.kin_source == "causal_routes" else {})} | |
| 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(wpl.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()) | |
| 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 stage_input_files(cfg) | |
| 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": wpl.ARM_OF[cfg["branch"]], "stage": args.stage, "variant": args.variant, "branch": cfg["branch"], "cfg": cfg, "synthetic": args.synthetic, | |
| "max_anchors": args.max_anchors, "params": td.get("params", {}), "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_streams": roi_streams(cfg), | |
| "roi": ("causal_clip_roi" if args.kin_source == "causal_routes" else wpl.ROI_OF[cfg["branch"]]), "external_priors": wpl.EXTERNAL_PRIORS[cfg["branch"]], "dms_fields_used": [], | |
| "selection_objective": "validation chain NLL (+0.1 L_F for E); MDN / Gaussian-K excluded", "pred_layouts": layouts, "pack_version": args.pack_version, | |
| "kin_source": args.kin_source, "kin_version": args.kin_version, "mask_routes": args.mask_routes, | |
| "kin_dir": (str(ds.kin_dir()) if args.kin_source != "cwm2" else ""), "side": cfg.get("side", "vf256") if cfg["branch"] == "E" else "", | |
| "stats_version": int(st["stats_version"]) if "stats_version" in st else 1, # 2 = float64 accumulation (train/batch.STATS_VERSION) | |
| "select_LF": bool(cfg["recipe"].get("select_LF", True)), | |
| "kclip": cfg.get("kclip", None), "u_in_ln": bool(cfg["recipe"].get("u_in_ln", False)), "p_head": bool(cfg.get("p_head", False)), | |
| "u_gaze_chain": bool(cfg["recipe"].get("u_gaze_chain", False)), "u_face": bool(cfg["recipe"].get("u_face", False)), # 2026-09-12 opt-in U arms | |
| "u_nocan": bool(cfg["recipe"].get("u_nocan", False)), "u_can_gate": bool(cfg["recipe"].get("u_can_gate", False)), # 2026-09-14 OPTIONAL BATON-side CAN arms (plan Part A §A.5) | |
| "gaze_dir": (str(ds.GAZE) if cfg["recipe"].get("u_gaze_chain") else ""), "face_dir": (str(ds.FACE) if cfg["recipe"].get("u_face") else ""), | |
| "note_kh_causality": (f"Kh / K / S from the CAUSAL route-level pose ({ds.kin_dir().name}, K4)" if args.kin_source == "causal_routes" else | |
| f"K4 input-isolation control: OLD route-median (non-causal) Kh inputs, CAUSAL targets (K_c / S_c / E_c from {ds.kin_dir().name})" if args.kin_source == "oldin_ctgt" else | |
| "Kh / K / S come from the route-median crop pose cache (NON-causal ROI): U0 results are diagnostic (K4 pending)") if cfg["branch"] == "U" else ""}) | |
| (Path(args.out) / "PRED_DONE").write_text(json.dumps({"splits": cfg["pred_splits"], "regimes": cfg["regimes"]})) | |
| log("run.json written; DONE", args.out) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 46.5 kB
- Xet hash:
- 12d16c022b05eca7a2d29d9a2a641ddda693a30f051d8bf226a5b799c9674c74
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.