Buckets:

AsianPlayer
/

TriDrive

Files
xet
AsianPlayer/TriDrive / code /train /train_future_state.py
AsianPlayer
16 days ago
download
raw
29.5 kB
"""cwm3/train/train_future_state.py — F2′ trainer (R2 WP4 batch 1, 2026-09-16; pre-registration docs/decisions.md 2026-09-16T06:46:35Z): the detached
future-state head (models/future_state_head.py) on the FROZEN FINAL joint model's driver state.
  python train/train_future_state.py --seed 0 [--final-run I1_all_minus_road2drv_s0_final] [--name dwell] [--feats dwell|dwell,face] [--epochs 20]
                                     [--extract-only] [--max-anchors N --cpu (smoke, never a result)]
  -> results/runs/FS_<name>_s<k>/  ckpt_best.pt  config.json  history.jsonl  feat_stats.npz  TRAIN_DONE  fs_{regime}_{split}.npz (train / val / test)  PRED_DONE  run_*.json
Step 1 (cache, results/runs/FS_cache/):  zU_<final_run>_{regime}_{split}.npz = h_u f16 [N,128] of the frozen joint model (its fine-tuned U trunk; the FINAL variant
  all_minus_road2drv has no interaction INTO the driver branch, so h_u = GRU(Kx) — recomputed here from the U stream under the pack's bf16 autocast and
  verified equal to the saved prediction pack within one f16 ulp, 2026-09-16) + the kin / gaze target arrays (S, S0, EM, M, Sg, Sg0, EMg, row_ok, uid, t, dongle,
  mode); feats_{regime}_{split}.npz = the label-free dwell features FEATURES (below) of the 8 s history.  No gradient ever reaches the joint model.
Step 2 (train): TRAIN anchors of both regimes; AdamW 1e-3 / wd 1e-4, batch 4096, <= --epochs, early stop patience 3 on the VALIDATION occupancy BCE
  (parts["occ_bce"], the selection objective); seeds 0 / 1 / 2 each on its own FINAL seed's state.
Step 3 (predict): every split -> fs_{regime}_{split}.npz, row-aligned with pred_layout_{regime}_{split}.npz of the FINAL run (uid / t witness stored):
  keys  uid t dongle valid[N] (row has a causal box = the FINAL layout's valid_U) state_names[S] off_value[S] z0[N,S] z0valid[N,S]
        pi[N,25,S] f16 (P(state = 1 at t + k dt)), pi_off[N,25,S] f16 (off-task occupancy), h_on / h_off[N,25,S] f16, ED_off_s[N,S] f32 (dt * sum_k pi_off),
        first_leave_p[N,S,26] f16 (law over the first step differing from z0; bin 25 = beyond 5 s), first_leave_E_s[N,S] (E[min(T, 5.2 s)]),
        rem_dwell_p / rem_dwell_E_s (the same law trained on off-task anchors = remaining dwell / recovery time), leave_within_{1,2,3,5}s[N,S],
        rem_within_{1,2,3,5}s[N,S], dt, horizon, meta (json: model_id, code sha, inputs digest, feature names, final run, base rates).
FEATURES (dwell set; every quantity is causal, from t-7.9..t):
  dw_* (11, train/warning_features_ext.dwell_features: hands_off_share_{1,3,8}s, hands_off_dwell_s, head_away_share_{1,3,8}s, head_away_dwell_s,
        head_dev_max_3s, n_hand_leave_8s, hist_valid_share)
  hands_state_dwell_s / head_state_dwell_s  time in the CURRENT hands / head state (run length of the forward-filled state at the window end, cap 8 s)
  n_head_switch_8s                          head-away switches (either direction) in the window
  t_since_any_switch_s                      time since the last switch of either kin state (cap 8 s)
  gz_z0 / gz_dwell0_s                       the current gaze_off state and its dwell (results/targets_gaze_causal_routes_v2 S0_g / dwell0; own landmark chain)
  face_valid_share_8s                       observation confidence of the face stream (Fh8valid share)
  miss_kin / miss_gaze                      missing flags (no valid kin step in the window / undefined gaze anchor); NaN features -> 0 after standardisation
  (+ `face`: the 16 fc_* channels of warning_features_ext.face_block — optional control arm, not the default)
Attribution arms (2026-09-16, same recipe): `--name zonly --feats none` (z_U + mode only) and `--name featonly --no-z` (dwell features + mode only) separate the
  information the head adds (8 s dwell / switch history, gaze anchor state) from the head + occupancy-BCE training on the trunk's own information.
Nothing here reads an openpilot DMS field.
"""
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 json                                                 # noqa: E402
import sys                                                  # noqa: E402
import time                                                 # noqa: E402
from pathlib import Path                                    # noqa: E402

import numpy as np                                          # 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"):
    if str(p) not in sys.path:
        sys.path.insert(0, str(p))
import cwm3_dataset as ds                                   # noqa: E402
import batch as bt                                          # noqa: E402
import warning_features_ext as wfe                          # noqa: E402
from runmeta import write_run_json, code_hash               # noqa: E402

RUNS = CWM3 / "results" / "runs"
CACHE = RUNS / "FS_cache"
REGIMES = ("manual", "engaged")
SPLITS = ("train", "val", "test")
EXTRA_NAMES = ["hands_state_dwell_s", "head_state_dwell_s", "n_head_switch_8s", "t_since_any_switch_s", "gz_z0", "gz_dwell0_s", "face_valid_share_8s", "miss_kin", "miss_gaze"]
FEATURES = list(wfe.DWELL_NAMES) + EXTRA_NAMES
CHUNK = 100_000
FINAL_RUN = "I1_all_minus_road2drv_s{seed}_final"
CODE_FILES = [CWM3 / "train" / "train_future_state.py", CWM3 / "models" / "future_state_head.py", CWM3 / "models" / "heads.py", CWM3 / "train" / "batch.py",
              CWM3 / "train" / "warning_features_ext.py", CWM3 / "models" / "cwm3.py", CWM3 / "models" / "branches.py"]


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


# ----------------------------------------------------------------------------- features (CPU, numpy; history only)
def _run_length_end(state, valid, cap=wfe.DWELL_CAP, hz=wfe.HZ):
    """state [n,T] in {0,1} (forward-filled), valid [n,T] -> seconds the CURRENT (last) state has lasted within the window (cap; NaN when no valid step)."""
    n, T = state.shape
    last = state[:, -1:]
    same = (state == last)
    # first index from the end where the state differs -> run starts after it
    diff_idx = np.where(~same, np.arange(T)[None, :], -1).max(1)
    run = (T - 1 - diff_idx) / hz
    run = np.minimum(run, cap).astype(np.float32)
    run[valid.sum(1) == 0] = np.nan
    return run


def extra_features(Kh8, Kh8valid, head_on=None, head_off=None):
    """[n,80,11] history -> the 4 kin extra features (hands_state_dwell_s, head_state_dwell_s, n_head_switch_8s, t_since_any_switch_s)."""
    if head_on is None or head_off is None:
        head_on, head_off = wfe.head_thresholds()
    K = np.asarray(Kh8, np.float32)
    V = np.asarray(Kh8valid, bool) & np.isfinite(K[:, :, wfe.IDX_MIN_DIST]) & np.isfinite(K[:, :, wfe.IDX_HEAD_DEV])
    hands_on = V & (K[:, :, wfe.IDX_MIN_DIST] < wfe.WHEEL_TH)
    dev = np.where(V, np.abs(K[:, :, wfe.IDX_HEAD_DEV]), np.nan)
    away = wfe.hysteresis_rows(dev, head_on, head_off)
    away_v = np.isfinite(away)
    hs = wfe._ffill_state(hands_on.astype(np.float32), V)
    hd_ = wfe._ffill_state(np.where(away_v, away, 0.0).astype(np.float32), away_v)
    hs0, hd0 = np.nan_to_num(hs, nan=0.0), np.nan_to_num(hd_, nan=0.0)
    hands_rl = _run_length_end(hs0, V)
    head_rl = _run_length_end(hd0, away_v)
    sw_head = ((hd0[:, 1:] != hd0[:, :-1]) & away_v[:, 1:] & np.isfinite(hd_[:, :-1])).astype(np.float32)
    sw_hand = ((hs0[:, 1:] != hs0[:, :-1]) & V[:, 1:] & np.isfinite(hs[:, :-1])).astype(np.float32)
    n_head = sw_head.sum(1)
    n_head[away_v.sum(1) == 0] = np.nan
    any_sw = (sw_head + sw_hand) > 0
    T = K.shape[1]
    idx = np.where(any_sw, np.arange(1, T)[None, :], -1).max(1)
    t_since = np.where(idx >= 0, (T - 1 - idx) / wfe.HZ, wfe.DWELL_CAP).astype(np.float32)
    t_since[(V.sum(1) == 0) & (away_v.sum(1) == 0)] = np.nan
    return np.stack([hands_rl, head_rl, n_head, t_since], 1).astype(np.float32)


def build_features(regime, split, rows=None, with_face=False):
    """-> X float32 [n, D], names.  Reads kinhist Kh8_c / Kh8valid (chunked), the gaze anchor state / dwell and the face validity share; `rows` = subset (smoke)."""
    ds.set_kin_source("causal_routes", "v2")
    kin_dir = ds.kin_dir()
    head_on, head_off = wfe.head_thresholds()
    with np.load(kin_dir / f"kinhist_causal_{regime}_{split}.npz", allow_pickle=True) as z:
        N = len(z["t"])
        rows = np.arange(N, dtype=np.int64) if rows is None else np.asarray(rows, np.int64)
        row_ok = z["row_ok"][rows].astype(bool) if "row_ok" in z.files else np.ones(len(rows), bool)
        parts = []
        for i in range(0, len(rows), CHUNK):
            r = rows[i:i + CHUNK]
            Kh8 = z["Kh8_c"][r]
            V = z["Kh8valid"][r] & row_ok[i:i + CHUNK][:, None]
            X1 = wfe.dwell_features(Kh8, V, head_on, head_off)
            X2 = extra_features(Kh8, V, head_on, head_off)
            parts.append(np.concatenate([X1, X2], 1))
        Xk = np.concatenate(parts)
    with np.load(ds.GAZE / f"gaze_{regime}_{split}.npz", allow_pickle=True) as z:
        gok = z["row_ok"][rows].astype(bool)
        gz0 = np.where(gok, z["S0_g"][rows].astype(np.float32), np.nan)
        gdw = np.where(gok, z["dwell0"][rows].astype(np.float32), np.nan)
    with np.load(ds.FACE / f"facehist_causal_{regime}_{split}.npz", allow_pickle=True) as z:
        fv = np.concatenate([z["Fh8valid"][rows[i:i + CHUNK]].mean(1) for i in range(0, len(rows), CHUNK)]).astype(np.float32)
        if "row_ok" in z.files:
            fv = fv * z["row_ok"][rows].astype(np.float32)
    miss_kin = (~np.isfinite(Xk[:, wfe.DWELL_NAMES.index("dw_hist_valid_share")]) | (Xk[:, wfe.DWELL_NAMES.index("dw_hist_valid_share")] <= 0)).astype(np.float32)
    miss_gaze = (~np.isfinite(gz0)).astype(np.float32)
    X = np.concatenate([Xk, gz0[:, None], np.minimum(gdw, 30.0)[:, None], fv[:, None], miss_kin[:, None], miss_gaze[:, None]], 1).astype(np.float32)
    names = list(FEATURES)
    if with_face:
        Xf, nf = wfe.face_block(regime, split, rows)
        X, names = np.concatenate([X, Xf], 1), names + nf
    assert X.shape[1] == len(names), (X.shape, len(names))
    return X, names


# ----------------------------------------------------------------------------- the frozen joint model's driver state
def load_final(final_run, device):
    import torch
    from models.cwm3 import build_joint
    run = RUNS / final_run
    assert (run / "TRAIN_DONE").exists() and (run / "PRED_DONE").exists(), f"{run}: not a finished joint run"
    cfg = json.load(open(run / "config.json"))["cfg"]
    bc = json.load(open(run / "base_ckpts.json"))
    ds.set_pack_version(cfg.get("pack_version", "v2"))
    ds.set_kin_source(cfg["kin_source"], cfg["kin_version"])
    assert "road2drv" not in cfg["interactions"] and not cfg.get("naive"), "h_u would depend on the road stream: extend the extractor to the full joint batch"
    model, info = build_joint(int(json.load(open(run / "config.json"))["args"]["seed"]), base={b: tuple(x) for b, x in cfg["base"].items()}, dirs=bc["base_dirs"],
                              device="cpu", stats_mode=cfg.get("stats_mode", "merged"), interactions=tuple(cfg["interactions"]), 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"))
    best = (run / "BEST").read_text().strip()
    if best != "base":
        z = torch.load(run / best, map_location="cpu", weights_only=False)
        model.load_state_dict(z["model"])
    model.eval().to(device)
    st = {k: v for k, v in np.load(run / "stats.npz").items()}
    pc_u = {"cabin": info["base_cfgs"]["U"].get("cabin", "none"), "recipe": {"u_gaze_chain": True}}
    if "kclip" in info["base_cfgs"]["U"]:
        pc_u["kclip"] = float(info["base_cfgs"]["U"]["kclip"])
    return model, st, pc_u, {"model_id": f"{final_run}:{best}", "base_ckpts": bc["base_ckpts"], "cfg": cfg}


def extract_zU(model, st, pc_u, regime, split, device, max_anchors=None, seed=0, amp=True):
    """-> dict: h_u f16 [N,128] + the target arrays of the (regime, split) pack (kin_row_ok, S, S0, EM, M, Sg, Sg0, EMg, mode, uid, t, dongle, idx)."""
    import torch
    from models.cwm3 import branch_stats
    raw = ds.load_split(regime, split, need={"X", "Kh", "kin", "Y", "gaze"}, max_anchors=max_anchors, seed=seed, verbose=lambda m: log(m))
    N = len(raw["t"])
    store = bt.GpuStore(raw, bt.KEYS["U"], device)
    stu = branch_stats(st, "U")
    H = np.zeros((N, model.U.gru.hidden_size), np.float16)
    bs = 8192
    with torch.no_grad():
        for i in range(0, N, bs):
            idx = torch.arange(i, min(i + bs, N), device=device)
            bU, _ = bt.prep_U(store, idx, stu, pc_u, need_targets=False)
            with torch.autocast("cuda", dtype=torch.bfloat16, enabled=(amp and device == "cuda")):
                h_u, _ = model.trunk_U(bU)
            H[i:i + len(idx)] = h_u.float().cpu().numpy().astype(np.float16)
    out = {"h_u": H, "S": raw["S"], "S0": raw["S0"].astype(np.float32), "EM": raw["EM"], "M": raw["M"], "Sg": raw["Sg"], "Sg0": raw["Sg0"].astype(np.float32), "EMg": raw["EMg"],
           "kin_row_ok": raw["kin_row_ok"], "mode": raw["mode"].astype(np.int8), "uid": raw["uid"], "t": raw["t"], "dongle": raw["dongle"], "idx": raw["idx"]}
    del store, raw
    return out


def cache_paths(final_run, regime, split):
    return CACHE / f"zU_{final_run}_{regime}_{split}.npz", CACHE / f"feats_{regime}_{split}.npz"


def ensure_cache(final_run, regime, split, device, ctx, max_anchors=None, seed=0, with_face=False):
    """builds (once) and returns (zU dict, feats X, names) for one (regime, split); smoke subsets (max_anchors) are never cached."""
    pz, pf = cache_paths(final_run, regime, split)
    if max_anchors:
        pz, pf = None, None
    if pz is not None and pz.exists():
        with np.load(pz, allow_pickle=True) as z:
            zu = {k: z[k] for k in z.files}
    else:
        if ctx["model"] is None:
            ctx["model"], ctx["st"], ctx["pc_u"], ctx["meta"] = load_final(final_run, device)
        t0 = time.time()
        zu = extract_zU(ctx["model"], ctx["st"], ctx["pc_u"], regime, split, device, max_anchors, seed)
        log(f"z_U {regime}/{split}: {len(zu['t']):,} rows in {time.time() - t0:.0f}s")
        if pz is not None:
            CACHE.mkdir(parents=True, exist_ok=True)
            np.savez(pz, **zu, model_id=np.array(ctx["meta"]["model_id"]))
    if pf is not None and pf.exists():
        with np.load(pf, allow_pickle=True) as z:
            X, names = z["X"], [str(n) for n in z["names"]]
            assert np.array_equal(z["uid"], zu["uid"]), "feature cache / z_U cache row mismatch"
    else:
        t0 = time.time()
        X, names = build_features(regime, split, rows=zu["idx"], with_face=with_face)
        log(f"features {regime}/{split}: {X.shape} in {time.time() - t0:.0f}s")
        if pf is not None:
            np.savez(pf, X=X, names=np.array(names), uid=zu["uid"], t=zu["t"])
    if with_face and X.shape[1] == len(FEATURES):                     # cached dwell-only features + the face arm requested -> append the face block
        Xf, nf = wfe.face_block(regime, split, zu["idx"])
        X, names = np.concatenate([X, Xf], 1), names + nf
    return zu, X, names


# ----------------------------------------------------------------------------- targets (train/batch.prep_U conventions, per state)
def make_targets(zu, device):
    """numpy cache arrays -> tensors: z_prev / z / m / Sfut / Sm [N,25,3], z0 / z0valid [N,3] for (hands_on, head_turned, gaze_off)."""
    import torch
    S = torch.from_numpy(zu["S"].astype(np.float32)).to(device)                 # [N,25,2] NaN censored
    S0 = torch.from_numpy(zu["S0"]).to(device)                                   # [N,2]
    Sg = torch.from_numpy(zu["Sg"].astype(np.float32)).to(device)[:, :, None]   # [N,25,1]
    Sg0 = torch.from_numpy(zu["Sg0"]).to(device)[:, None]
    M = torch.from_numpy(zu["M"]).to(device)
    EM = torch.from_numpy(zu["EM"]).to(device)
    EMg = torch.from_numpy(zu["EMg"]).to(device)
    Sall = torch.cat([S, Sg], -1)
    S0all = torch.cat([S0, Sg0], -1)
    z = torch.nan_to_num(Sall)
    z0 = torch.nan_to_num(S0all)
    z_prev = torch.cat([z0[:, None], z[:, :-1]], 1)
    Sprev = torch.cat([S0all[:, None], Sall[:, :-1]], 1)
    onset = torch.stack([EM[:, :, 0], EM[:, :, 2], EMg[:, :, 0]], -1)
    m = onset & M[:, :, None] & torch.isfinite(Sall) & torch.isfinite(Sprev)
    Sm = torch.isfinite(Sall) & M[:, :, None]
    return {"z_prev": z_prev, "z": z, "m": m, "Sfut": z, "Sm": Sm, "z0": z0, "z0valid": torch.isfinite(S0all)}


def standardise(X, stats=None):
    """nan-aware train mean / sd; NaN -> 0 after standardisation; clip +-8."""
    if stats is None:
        mu = np.nanmean(X, 0)
        sd = np.nanstd(X, 0) + 1e-6
        stats = {"mu": mu.astype(np.float32), "sd": sd.astype(np.float32)}
    Z = (X - stats["mu"]) / stats["sd"]
    return np.clip(np.nan_to_num(Z, nan=0.0, posinf=8.0, neginf=-8.0), -8, 8).astype(np.float32), stats


# ----------------------------------------------------------------------------- train / predict
def run_epoch(model, opt, Z, X, mode, tg, idx_all, bs, train, gen=None):
    import torch
    from models.future_state_head import future_state_loss
    tot, n, agg = 0.0, 0, {}
    order = idx_all[torch.randperm(len(idx_all), generator=gen, device=idx_all.device)] if train else idx_all
    for i in range(0, len(order), bs):
        idx = order[i:i + bs]
        t = {k: v[idx] for k, v in tg.items()}
        out = model(Z[idx].float(), X[idx], mode[idx])
        loss, parts = future_state_loss(out, t)
        if train:
            opt.zero_grad(set_to_none=True)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
        w = len(idx)
        tot += float(loss.detach()) * w
        n += w
        for k, v in parts.items():
            agg[k] = agg.get(k, 0.0) + float(v.detach() if torch.is_tensor(v) else v) * w
    return tot / max(n, 1), {k: v / max(n, 1) for k, v in agg.items()}


def predict_split(model, zu, Xs, out_dir, regime, split, device, base_rate, meta, names, bs=16384):
    import torch
    from models.future_state_head import future_state_outputs, STATES, OFF_VALUE, DT, H, N_BINS, LEAVE_TAUS
    N = len(zu["t"])
    Z = torch.from_numpy(zu["h_u"]).to(device)
    X = torch.from_numpy(Xs).to(device)
    mode = torch.from_numpy(zu["mode"].astype(np.int64)).to(device)
    S0 = np.concatenate([zu["S0"], zu["Sg0"][:, None]], 1).astype(np.float32)
    z0valid = np.isfinite(S0)
    z0 = np.where(z0valid, S0, base_rate[None, :]).astype(np.float32)
    z0t = torch.from_numpy(z0).to(device)
    S = len(STATES)
    P = {"pi": np.zeros((N, H, S), np.float16), "pi_off": np.zeros((N, H, S), np.float16), "h_on": np.zeros((N, H, S), np.float16), "h_off": np.zeros((N, H, S), np.float16),
         "ED_off_s": np.zeros((N, S), np.float32), "first_leave_p": np.zeros((N, S, N_BINS), np.float16), "rem_dwell_p": np.zeros((N, S, N_BINS), np.float16),
         "first_leave_E_s": np.zeros((N, S), np.float32), "rem_dwell_E_s": np.zeros((N, S), np.float32)}
    for nm in LEAVE_TAUS:
        P[f"leave_within_{nm}"] = np.zeros((N, S), np.float16)
        P[f"rem_within_{nm}"] = np.zeros((N, S), np.float16)
    model.eval()
    with torch.no_grad():
        for i in range(0, N, bs):
            sl = slice(i, min(i + bs, N))
            out = model(Z[sl].float(), X[sl], mode[sl])
            res = future_state_outputs(out, z0t[sl])
            for k in P:
                P[k][sl] = res[k].cpu().numpy().astype(P[k].dtype)
    m = dict(meta, regime=regime, split=split, n_rows=int(N), features=names, dt=DT, horizon=H, state_names=list(STATES), off_value=[OFF_VALUE[s] for s in STATES],
             base_rate_z0=base_rate.tolist(), written_utc=time.strftime("%Y%m%dT%H%M%SZ", time.gmtime()))
    np.savez_compressed(Path(out_dir) / f"fs_{regime}_{split}.npz", uid=zu["uid"], t=zu["t"], dongle=zu["dongle"], valid=zu["kin_row_ok"].astype(bool), z0=z0.astype(np.float16),
                        z0valid=z0valid, state_names=np.array(list(STATES)), off_value=np.array([OFF_VALUE[s] for s in STATES], np.float32), dt=np.array(DT), horizon=np.array(H),
                        meta=np.array(json.dumps(m, default=str)), **P)
    return N


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--final-run", default="", help="frozen joint run (default I1_all_minus_road2drv_s<seed>_final)")
    ap.add_argument("--name", default="dwell", help="run family: results/runs/FS_<name>_s<seed>")
    ap.add_argument("--feats", default="dwell", help="dwell | dwell,face | none (attribution arm: a single zero feature -> the head sees z_U + mode only)")
    ap.add_argument("--no-z", action="store_true", help="attribution arm: z_U replaced by zeros (the head sees the dwell features + mode only)")
    ap.add_argument("--epochs", type=int, default=20)
    ap.add_argument("--patience", type=int, default=3)
    ap.add_argument("--batch", type=int, default=4096)
    ap.add_argument("--lr", type=float, default=1e-3)
    ap.add_argument("--d", type=int, default=64)
    ap.add_argument("--regimes", default="manual,engaged")
    ap.add_argument("--splits", default="train,val,test", help="prediction splits")
    ap.add_argument("--extract-only", action="store_true", help="build the z_U / feature caches and stop")
    ap.add_argument("--max-anchors", type=int, default=0, help="smoke subset per (regime, split); never cached, never a result")
    ap.add_argument("--cpu", action="store_true")
    ap.add_argument("--out", default="")
    args = ap.parse_args()
    import torch
    from models.future_state_head import FutureStateHead, n_params, STATES
    torch.set_num_threads(4)
    device = "cuda" if torch.cuda.is_available() and not args.cpu else "cpu"
    final_run = args.final_run or FINAL_RUN.format(seed=args.seed)
    out = Path(args.out) if args.out else RUNS / f"FS_{args.name}_s{args.seed}{'_smoke' if args.max_anchors else ''}"
    out.mkdir(parents=True, exist_ok=True)
    regimes = args.regimes.split(",")
    with_face = "face" in args.feats.split(",")
    t_start = time.time()
    log(f"=== train_future_state seed {args.seed} final_run {final_run} feats {args.feats} out {out} device {device} ===", out)
    ctx = {"model": None}
    mx = args.max_anchors or None
    data = {}
    for split in ("train", "val", "test"):
        parts = [ensure_cache(final_run, r, split, device, ctx, mx, args.seed, with_face) for r in regimes]
        data[split] = parts
    meta = ctx.get("meta") or {"model_id": str(np.load(cache_paths(final_run, regimes[0], "train")[0], allow_pickle=True)["model_id"]), "final_run": final_run}
    meta["final_run"] = final_run
    if ctx["model"] is not None:
        del ctx["model"]
        torch.cuda.empty_cache() if device == "cuda" else None
    if args.extract_only:
        log("caches built; --extract-only", out)
        return
    names = data["train"][0][2]
    if args.feats == "none":                                          # attribution arm: no dwell information (one constant zero column)
        for split in data:
            data[split] = [(zu, np.zeros((len(zu["t"]), 1), np.float32), ["none"]) for (zu, X, nm) in data[split]]
        names = ["none"]
    if args.no_z:                                                     # attribution arm: no trunk state
        for split in data:
            for zu, _, _ in data[split]:
                zu["h_u"] = np.zeros_like(zu["h_u"])
    # ---- train / val tensors
    def stack(parts, key):
        return np.concatenate([p[0][key] for p in parts])
    Xtr = np.concatenate([p[1] for p in data["train"]])
    Xtr_s, fstats = standardise(Xtr)
    Xva_s, _ = standardise(np.concatenate([p[1] for p in data["val"]]), fstats)
    np.savez(out / "feat_stats.npz", mu=fstats["mu"], sd=fstats["sd"], names=np.array(names))
    tr_zu = {k: stack(data["train"], k) for k in ("h_u", "S", "S0", "EM", "M", "Sg", "Sg0", "EMg", "mode")}
    va_zu = {k: stack(data["val"], k) for k in ("h_u", "S", "S0", "EM", "M", "Sg", "Sg0", "EMg", "mode")}
    S0tr = np.concatenate([tr_zu["S0"], tr_zu["Sg0"][:, None]], 1)
    base_rate = np.array([np.nanmean(S0tr[:, s]) for s in range(S0tr.shape[1])], np.float32)
    Ztr, Zva = torch.from_numpy(tr_zu["h_u"]).to(device), torch.from_numpy(va_zu["h_u"]).to(device)
    Xtr_t, Xva_t = torch.from_numpy(Xtr_s).to(device), torch.from_numpy(Xva_s).to(device)
    mtr, mva = torch.from_numpy(tr_zu["mode"].astype(np.int64)).to(device), torch.from_numpy(va_zu["mode"].astype(np.int64)).to(device)
    tgtr, tgva = make_targets(tr_zu, device), make_targets(va_zu, device)
    Ntr, Nva = len(Ztr), len(Zva)
    log(f"train {Ntr:,} rows, val {Nva:,} rows; features {len(names)}; base rates {base_rate.tolist()}", out)
    torch.manual_seed(args.seed)
    model = FutureStateHead(d_z=Ztr.shape[1], d_dwell=len(names), n_states=len(STATES), d=args.d).to(device)
    npar = n_params(model)
    log(f"FutureStateHead params {npar:,}", out)
    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4)
    gen = torch.Generator(device=device).manual_seed(args.seed)
    idx_tr, idx_va = torch.arange(Ntr, device=device), torch.arange(Nva, device=device)
    best, best_ep, bad, hist = float("inf"), -1, 0, []
    for ep in range(1, args.epochs + 1):
        t0 = time.time()
        model.train()
        ltr, ptr = run_epoch(model, opt, Ztr, Xtr_t, mtr, tgtr, idx_tr, args.batch, True, gen)
        model.eval()
        with torch.no_grad():
            lva, pva = run_epoch(model, None, Zva, Xva_t, mva, tgva, idx_va, 16384, False)
        obj = pva["occ_bce"]
        rec = {"epoch": ep, "train_loss": ltr, "val_loss": lva, "val_occ_bce": obj, "val_parts": pva, "train_parts": ptr, "wall_s": round(time.time() - t0, 1)}
        hist.append(rec)
        with open(out / "history.jsonl", "a") as f:
            f.write(json.dumps(rec) + "\n")
        improved = obj < best - 1e-6
        if improved:
            best, best_ep, bad = obj, ep, 0
            torch.save({"model": model.state_dict(), "epoch": ep, "val_occ_bce": obj, "feat_names": names, "feat_stats": fstats, "base_rate": base_rate, "args": vars(args),
                        "final_run": final_run, "d_z": int(Ztr.shape[1]), "n_states": len(STATES), "d": args.d}, out / "ckpt_best.pt")
        else:
            bad += 1
        log(f"ep {ep} train {ltr:.4f} val {lva:.4f} occ_bce {obj:.5f} " + " ".join(f"{k}={v:.4f}" for k, v in pva.items() if k.startswith(("occ_", "trans_", "leave_", "rem_"))) +
            f" {'*' if improved else ''} ({rec['wall_s']}s)", out)
        if bad >= args.patience:
            break
    z = torch.load(out / "ckpt_best.pt", map_location=device, weights_only=False)
    model.load_state_dict(z["model"])
    (out / "BEST").write_text("ckpt_best.pt\n")
    (out / "TRAIN_DONE").write_text(json.dumps({"best_epoch": best_ep, "best": best, "epochs_run": len(hist), "params": npar, "selection": "validation occupancy BCE (parts occ_bce)"}))
    meta.update(code_sha256=code_hash(CODE_FILES), params=npar, best_epoch=best_ep, val_occ_bce=best, seed=args.seed, name=args.name, feats=args.feats, no_z=bool(args.no_z))
    (out / "config.json").write_text(json.dumps({"args": vars(args), "final_run": final_run, "features": names, "params": npar, "meta": meta}, indent=1, default=str))
    # ---- predictions
    del Ztr, Xtr_t, tgtr
    n_pred = {}
    for split in args.splits.split(","):
        for (zu, X, nm), r in zip(data[split], regimes):
            Xs, _ = standardise(X, fstats)
            n_pred[f"{r}_{split}"] = predict_split(model, zu, Xs, out, r, split, device, base_rate, meta, names)
            log(f"fs_{r}_{split}.npz written ({n_pred[f'{r}_{split}']:,} rows)", out)
    inputs = sorted({p for r in regimes for sp in SPLITS for p in (ds.kin_dir() / f"kin_causal_{r}_{sp}.npz", ds.kin_dir() / f"kinhist_causal_{r}_{sp}.npz",
                                                                  ds.GAZE / f"gaze_{r}_{sp}.npz", ds.FACE / f"facehist_causal_{r}_{sp}.npz") if p.exists()}
                    | {RUNS / final_run / (RUNS / final_run / "BEST").read_text().strip()})
    write_run_json(str(out), f"train_future_state_{args.name}_s{args.seed}", CODE_FILES, [] if args.max_anchors else inputs, seed=args.seed, t_start=t_start,
                   extra={"final_run": final_run, "features": names, "params": npar, "best_epoch": best_ep, "val_occ_bce": best, "history": hist, "n_pred": n_pred,
                          "base_rate_z0": base_rate.tolist(), "state_names": list(STATES), "selection": "validation occupancy BCE", "smoke": bool(args.max_anchors), "feats": args.feats, "no_z": bool(args.no_z),
                          "dms_fields_used": [], "gpu_wall_s": round(time.time() - t_start, 1)})
    (out / "PRED_DONE").write_text(json.dumps({"splits": args.splits.split(","), "regimes": regimes, "n_pred": n_pred}))
    log(f"DONE in {time.time() - t_start:.0f}s", out)


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

Size:: 29.5 kB
Xet hash:: 0022ebbaccf46b5e39d6534dc9dd387e1768056122c804e2e327cc3e10761914
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.