Buckets:

AsianPlayer
/

TriDrive

Files

xet

AsianPlayer/TriDrive / code /train /train_kp_residual.py

AsianPlayer

16 days ago

download

raw

17.4 kB

"""cwm3/train/train_kp_residual.py — LANE F1 trainer: keypoint RESIDUAL forecaster on BATON (R2 WP4 batch 2; pre-registration docs/decisions.md 2026-09-17T19:33:12Z).
  python train/train_kp_residual.py --arm mb_full|mb_dense|mb_1hz|cwm_full --seed k [--out results/runs/KPF_<arm>_s<k>] [--epochs 12] [--rows-per-epoch 400000]
                                    [--max-anchors N --cpu   (smoke only, never a result)]
  -> results/runs/KPF_<arm>_s<k>/  config.json  history.jsonl  ckpt_best.pt  BEST  TRAIN_DONE  kpf_{regime}_{split}.npz (train / val / test)
                                   pred_kp_{regime}_{split}.npz (val / test)  PRED_DONE  run_*.json
Data: results/kp_residual/store.npz (the 10 Hz causal RTMW upper-body pose of every s12 route, data/build_kp_residual_store.py) + anchors_{regime}_{split}.npz; the
observed frames (<= t) and the target frames (+0.5 / +1 / +2 / +3 s) are gathered from the route tick table on the GPU per batch — nothing about the future is
ever an input (models/kp_forecaster.causal_normalise reads the observed frames only; tests/test_kp_forecaster.py).  TRAIN anchors of both regimes with an anchor
frame; each epoch samples --rows-per-epoch anchors; AdamW 1e-3 / wd 1e-4, batch 1024, bf16 autocast, grad clip 1, <= --epochs, early stop patience 3 on the
VALIDATION masked L1 of the most probable mode (`val_l1_top1`, normalised units — the selection objective); seeds 0 / 1 / 2.
Arms: mb_full (MotionBERT architecture, dense 2 s + sparse 1 Hz history), mb_dense (dense only), mb_1hz (1 Hz only), cwm_full (CockpitWM part-token form, parameter-matched).
Outputs per anchor (row order = the s12 grid, uid / t witness):
  kpf_*.npz      route i32 [N] + route_names (uid = route_names[route]), t, feat f16 [N,10] = the DETACHED distillation channel (6 state probabilities, 3 mode probabilities, valid bit; all 0 when the anchor has no frame or no
                 valid reference joint), valid bool[N]
  pred_kp_*.npz  res_top1 f16 [N,4,11,2] (residual of the most probable mode in FRAME units = normalised residual x scale), ref f16 [N,11,2], ref_ok bool[N,11],
                 scale f16[N], p_state f16 [N,3,2], p_mode f16 [N,3], valid bool[N]
No CAN column, label or openpilot DMS field is read (the state targets are the pose-derived kin-v2 states).
"""
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
import torch                                                # noqa: E402

torch.set_num_threads(4)
CWM3 = Path(__file__).resolve().parents[1]
for p in (CWM3, CWM3 / "common"):
    if str(p) not in sys.path:
        sys.path.insert(0, str(p))
from models import kp_forecaster as kf                      # noqa: E402
from runmeta import write_run_json                          # noqa: E402

STORE = CWM3 / "results" / "kp_residual"
REGIMES, SPLITS = ("manual", "engaged"), ("train", "val", "test")
ARMS = {"mb_full": dict(backbone="mb", obs="full", d=112, depth=2, heads=4, mlp_ratio=2),
        "mb_dense": dict(backbone="mb", obs="dense", d=112, depth=2, heads=4, mlp_ratio=2),
        "mb_1hz": dict(backbone="mb", obs="1hz", d=112, depth=2, heads=4, mlp_ratio=2),
        "cwm_full": dict(backbone="cwm", obs="full", d=128, depth=4, heads=4, ffn=512)}
HT = torch.tensor(kf.HORIZON_TICKS, dtype=torch.long)


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")


class Store:
    """the pose store on `device`; gather(pos, route, offsets) -> xy [B,T,11,2], conf [B,T,11], valid [B,T,11]."""

    def __init__(self, device):
        z = np.load(STORE / "store.npz", allow_pickle=True)
        self.kp = torch.from_numpy(z["kp"]).to(device)                        # f16 [F,11,3]
        self.valid = torch.from_numpy(z["valid"]).to(device)
        self.lut = torch.from_numpy(z["lut"].astype(np.int64)).to(device)
        self.lut_off = torch.from_numpy(z["route_lut_off"]).to(device)
        self.lut_len = torch.from_numpy(z["route_lut_len"]).to(device)
        self.L = int(self.lut.numel())
        self.device = device
        self.n_frames = int(self.kp.shape[0])

    def gather(self, pos, route, offsets):
        idx = pos[:, None] + offsets.to(self.device)[None]
        lo = self.lut_off[route][:, None]
        inside = (idx >= lo) & (idx < lo + self.lut_len[route][:, None])
        fi = self.lut[idx.clamp(0, self.L - 1)]
        fi = torch.where(inside, fi, torch.full_like(fi, -1))
        ok = fi >= 0
        k = self.kp[fi.clamp_min(0)].float()                                 # [B,T,11,3]
        v = self.valid[fi.clamp_min(0)] & ok[..., None]
        xy = torch.nan_to_num(k[..., :2]) * v[..., None].float()
        return xy, torch.nan_to_num(k[..., 2]) * v.float(), v


ROUTE_NAMES = None


def route_names():
    """store.npz route_names: uid == route name, so anchors carry the route index only (disk: the 2.6 M uid strings were 1.4 GB)."""
    global ROUTE_NAMES
    if ROUTE_NAMES is None:
        ROUTE_NAMES = np.load(STORE / "store.npz", allow_pickle=True)["route_names"].astype(str)
    return ROUTE_NAMES


def load_anchors(regime, split, max_anchors=None, seed=0):
    z = np.load(STORE / f"anchors_{regime}_{split}.npz", allow_pickle=True)
    d = {k: z[k] for k in ("t", "dongle", "mode", "route", "lut_pos", "S", "S0", "kin_row_ok")}
    d["uid"] = route_names()[d["route"]]
    if max_anchors and len(d["t"]) > max_anchors:
        sel = np.sort(np.random.default_rng(seed).choice(len(d["t"]), max_anchors, replace=False))
        d = {k: v[sel] for k, v in d.items()}
    return d


def to_dev(d, device):
    return {"route": torch.from_numpy(d["route"].astype(np.int64)).to(device), "pos": torch.from_numpy(d["lut_pos"].astype(np.int64)).to(device),
            "S": torch.from_numpy(d["S"]).to(device)}


def run_batch(model, store, a, idx, need_loss=True):
    pos, route = a["pos"][idx], a["route"][idx]
    has = pos >= 0
    posc = pos.clamp_min(0)
    xy, conf, valid = store.gather(posc, route, model.offsets)
    valid = valid & has[:, None, None]
    out = model(xy, conf, valid)
    if not need_loss:
        return out, has
    fxy, _, fvalid = store.gather(posc, route, HT)
    fvalid = fvalid & has[:, None, None]
    total, parts, err = kf.kp_loss(out, fxy, fvalid, a["S"][idx])
    return out, has, total, parts, err, fxy, fvalid


@torch.no_grad()
def evaluate(model, store, a, batch, amp):
    """validation: masked L1 (normalised units) of the most probable mode and of the oracle mode, mean over anchors with a target."""
    model.eval()
    N = a["pos"].numel()
    s_top, s_wta, n = 0.0, 0.0, 0
    for i in range(0, N, batch):
        idx = torch.arange(i, min(N, i + batch), device=store.device)
        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=amp):
            out, has, total, parts, err, fxy, fvalid = run_batch(model, store, a, idx)
        m = (fvalid & out["ref_ok"][:, None]).float()[:, :, None, :, None]
        hasrow = (m.sum((1, 2, 3, 4)) > 0)
        top = out["mode_logits"].float().argmax(1)
        e_top = err.float().gather(1, top[:, None]).squeeze(1)
        s_top += float((e_top * hasrow.float()).sum()); s_wta += float((err.float().min(1).values * hasrow.float()).sum()); n += int(hasrow.sum())
    model.train()
    return {"val_l1_top1": s_top / max(n, 1), "val_l1_oracle": s_wta / max(n, 1), "n_val_rows": n}


@torch.no_grad()
def predict(model, store, a_np, a, batch, amp, out_dir, regime, split, full):
    model.eval()
    N = a["pos"].numel()
    feat = np.zeros((N, kf.KPF_DIM), np.float16)
    valid_row = np.zeros(N, bool)
    if full:
        res = np.zeros((N, kf.H, kf.NJ, 2), np.float16); ref = np.zeros((N, kf.NJ, 2), np.float16); ref_ok = np.zeros((N, kf.NJ), bool)
        scale = np.zeros(N, np.float16); p_state = np.zeros((N, kf.N_STATE_H, kf.N_STATES), np.float16); p_mode = np.zeros((N, kf.N_MODES), np.float16)
    for i in range(0, N, batch):
        idx = torch.arange(i, min(N, i + batch), device=store.device)
        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=amp):
            out, has = run_batch(model, store, a, idx, need_loss=False)
        ok = has & out["ref_ok"].any(1)
        ch = kf.kpf_channel({k: v.float() for k, v in out.items()}, ok)
        sl = slice(i, i + len(idx))
        feat[sl] = ch.cpu().numpy().astype(np.float16); valid_row[sl] = ok.cpu().numpy()
        if full:
            top = out["mode_logits"].float().argmax(1)
            r = out["res"].float()[torch.arange(len(idx), device=store.device), :, top]                    # [B,H,J,2]
            res[sl] = (r * out["scale"].float()[:, None, None, None]).cpu().numpy().astype(np.float16)
            ref[sl] = out["ref"].float().cpu().numpy().astype(np.float16); ref_ok[sl] = out["ref_ok"].cpu().numpy(); scale[sl] = out["scale"].float().cpu().numpy().astype(np.float16)
            p_state[sl] = torch.sigmoid(out["state_logits"].float()).cpu().numpy().astype(np.float16); p_mode[sl] = torch.softmax(out["mode_logits"].float(), -1).cpu().numpy().astype(np.float16)
    rn = route_names()                                                     # uid = route_names[route] (reconstructed by the readers; no 2.6 M-string array per file)
    np.savez(out_dir / f"kpf_{regime}_{split}.npz", route=a_np["route"].astype(np.int32), route_names=rn, t=a_np["t"], feat=feat, valid=valid_row,
             names=np.asarray([f"p_{s}_{h:g}s" for h in (1, 2, 3) for s in ("hands_on", "head_turned")] + [f"p_mode{m}" for m in range(kf.N_MODES)] + ["valid"]))
    if full:
        np.savez(out_dir / f"pred_kp_{regime}_{split}.npz", route=a_np["route"].astype(np.int32), route_names=rn, t=a_np["t"], res_top1=res, ref=ref, ref_ok=ref_ok, scale=scale, p_state=p_state, p_mode=p_mode,
                 valid=valid_row, horizons_s=np.asarray(kf.HORIZONS_S), state_horizons_s=np.asarray([1.0, 2.0, 3.0]), joints=np.load(STORE / "store.npz", allow_pickle=True)["joints"])
    # (uid strings are not stored: readers rebuild uid = route_names[route] — data/cwm3_dataset.load_split "kpf", eval/kp_residual_eval.py)
    model.train()
    return float(valid_row.mean())


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--arm", required=True, choices=list(ARMS))
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--out", default="")
    ap.add_argument("--epochs", type=int, default=12)
    ap.add_argument("--rows-per-epoch", type=int, default=400_000)
    ap.add_argument("--batch", type=int, default=1024)
    ap.add_argument("--lr", type=float, default=1e-3)
    ap.add_argument("--weight-decay", type=float, default=1e-4)
    ap.add_argument("--patience", type=int, default=3)
    ap.add_argument("--minutes", type=float, default=90)
    ap.add_argument("--max-anchors", type=int, default=0, help="smoke only")
    ap.add_argument("--cpu", action="store_true", help="smoke only")
    ap.add_argument("--no-amp", action="store_true")
    a = ap.parse_args()
    out = Path(a.out) if a.out else CWM3 / "results" / "runs" / f"KPF_{a.arm}_s{a.seed}" + ("_smoke" if a.max_anchors else "")
    out.mkdir(parents=True, exist_ok=True)
    t_start = time.time()
    device = torch.device("cpu" if a.cpu else "cuda")
    amp = (not a.no_amp) and device.type == "cuda"
    torch.manual_seed(a.seed); np.random.seed(a.seed)
    cfg = {"arm": a.arm, **ARMS[a.arm], "seed": a.seed, "epochs": a.epochs, "rows_per_epoch": a.rows_per_epoch, "batch": a.batch, "lr": a.lr, "weight_decay": a.weight_decay,
           "patience": a.patience, "selection": "val_l1_top1 (masked L1 of the most probable mode, normalised units)", "max_anchors": a.max_anchors, "amp": amp,
           "loss": "wta_l1 + 0.1 mode_ce + 0.5 state_bce", "horizons_s": list(kf.HORIZONS_S), "state_horizons_s": [1.0, 2.0, 3.0]}
    (out / "config.json").write_text(json.dumps(cfg, indent=1))
    store = Store(device)
    log(f"=== train_kp_residual {a.arm} s{a.seed} out={out} store frames={store.n_frames:,} device={device} ===", out)
    tr_np = {r: load_anchors(r, "train", a.max_anchors or None, a.seed) for r in REGIMES}
    va_np = {r: load_anchors(r, "val", (a.max_anchors // 4 if a.max_anchors else None), a.seed) for r in REGIMES}
    tr = {k: torch.cat([to_dev(tr_np[r], device)[k] for r in REGIMES]) for k in ("route", "pos", "S")}
    va = {k: torch.cat([to_dev(va_np[r], device)[k] for r in REGIMES]) for k in ("route", "pos", "S")}
    tr_idx_all = torch.nonzero(tr["pos"] >= 0).squeeze(1)
    # scale fallback = TRAIN median shoulder width (observed frames) over a 50k-anchor sample: recorded, used identically at prediction time
    with torch.no_grad():
        smp = tr_idx_all[torch.randperm(len(tr_idx_all), generator=torch.Generator().manual_seed(0))[:50_000].to(device)]
        xy, _, v = store.gather(tr["pos"][smp], tr["route"][smp], kf.obs_offsets("full"))
        sw = (xy[:, :, kf.L_SHO] - xy[:, :, kf.R_SHO]).norm(dim=-1)[v[:, :, kf.L_SHO] & v[:, :, kf.R_SHO]]
        scale_fb = float(sw.median()) if sw.numel() else kf.SCALE_FALLBACK
    cfg["scale_fallback"] = scale_fb
    model = kf.KpForecaster(**{k: v for k, v in ARMS[a.arm].items()}, scale_fallback=scale_fb).to(device)
    cfg["params"] = {"total": kf.n_params(model), "backbone": kf.n_params(model.backbone), "heads": kf.n_params(model) - kf.n_params(model.backbone)}
    (out / "config.json").write_text(json.dumps(cfg, indent=1))
    log(f"params {cfg['params']} train anchors {len(tr_idx_all):,} val anchors {int((va['pos'] >= 0).sum()):,} scale_fallback {scale_fb:.4f}", out)
    opt = torch.optim.AdamW(model.parameters(), lr=a.lr, weight_decay=a.weight_decay)
    g = torch.Generator(device="cpu").manual_seed(a.seed)
    best, best_ep, bad, hist = float("inf"), -1, 0, []
    for ep in range(a.epochs):
        t_ep = time.time()
        perm = tr_idx_all[torch.randperm(len(tr_idx_all), generator=g)[:a.rows_per_epoch].to(device)]
        sums, nb = {}, 0
        for i in range(0, len(perm), a.batch):
            idx = perm[i:i + a.batch]
            with torch.autocast("cuda", dtype=torch.bfloat16, enabled=amp):
                out_, has, total, parts, err, _, _ = run_batch(model, store, tr, idx)
            opt.zero_grad(set_to_none=True)
            total.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
            for k, v in parts.items():
                sums[k] = sums.get(k, 0.0) + v
            nb += 1
            if (time.time() - t_start) / 60 > a.minutes:
                break
        ev = evaluate(model, store, va, 4096, amp)
        rec = {"epoch": ep, **{f"train_{k}": v / max(nb, 1) for k, v in sums.items()}, **ev, "sec": round(time.time() - t_ep, 1)}
        hist.append(rec)
        with open(out / "history.jsonl", "a") as f:
            f.write(json.dumps(rec) + "\n")
        improved = ev["val_l1_top1"] < best - 1e-6
        if improved:
            best, best_ep, bad = ev["val_l1_top1"], ep, 0
            torch.save({"model": model.state_dict(), "cfg": cfg, "epoch": ep, "val": ev}, out / "ckpt_best.pt")
            (out / "BEST").write_text("ckpt_best.pt")
        else:
            bad += 1
        log(f"ep {ep} " + " ".join(f"{k}={v:.4f}" if isinstance(v, float) else f"{k}={v}" for k, v in rec.items() if k != "epoch") + (" *" if improved else ""), out)
        if bad >= a.patience or (time.time() - t_start) / 60 > a.minutes:
            break
    (out / "TRAIN_DONE").write_text(json.dumps({"best_epoch": best_ep, "best_val_l1_top1": best, "epochs_run": len(hist)}))
    model.load_state_dict(torch.load(out / "ckpt_best.pt", map_location=device, weights_only=False)["model"])
    cover = {}
    for r in REGIMES:
        for s in SPLITS:
            a_np = tr_np[r] if s == "train" else (va_np[r] if s == "val" else load_anchors(r, "test", (a.max_anchors // 2 if a.max_anchors else None), a.seed))
            cover[f"{r}_{s}"] = predict(model, store, a_np, to_dev(a_np, device), 4096, amp, out, r, s, full=(s != "train"))
            log(f"predicted {r}/{s}: valid share {cover[f'{r}_{s}']:.4f}", out)
    (out / "PRED_DONE").write_text(json.dumps(cover))
    write_run_json(out, f"train_kp_residual_{a.arm}_s{a.seed}", [Path(__file__), CWM3 / "models" / "kp_forecaster.py", CWM3 / "data" / "build_kp_residual_store.py"],
                   [STORE / "store.npz"] + [STORE / f"anchors_{r}_{s}.npz" for r in REGIMES for s in SPLITS], seed=a.seed, t_start=t_start,
                   extra={"cfg": cfg, "history": hist, "best_epoch": best_ep, "best_val_l1_top1": best, "epochs_run": len(hist), "gpu_wall_s": round(time.time() - t_start, 1),
                          "coverage": cover, "smoke": bool(a.max_anchors or a.cpu), "dms_fields_used": [], "inputs_note": "own causal RTMW pose (results/pose_causal_routes[_v2]); kin-v2 states as targets"})
    log("DONE", out)


if __name__ == "__main__":
    main()

Xet Storage Details

Size:: 17.4 kB
Xet hash:: 0c724917d50d3b4ec9773d5128318126ccfbb57cbb7502dee2158fce71f71d42

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

	"""cwm3/train/train_kp_residual.py — LANE F1 trainer: keypoint RESIDUAL forecaster on BATON (R2 WP4 batch 2; pre-registration docs/decisions.md 2026-09-17T19:33:12Z).

	python train/train_kp_residual.py --arm mb_full\|mb_dense\|mb_1hz\|cwm_full --seed k [--out results/runs/KPF_<arm>_s<k>] [--epochs 12] [--rows-per-epoch 400000]
	[--max-anchors N --cpu (smoke only, never a result)]
	-> results/runs/KPF_<arm>_s<k>/ config.json history.jsonl ckpt_best.pt BEST TRAIN_DONE kpf_{regime}_{split}.npz (train / val / test)
	pred_kp_{regime}_{split}.npz (val / test) PRED_DONE run_*.json

	Data: results/kp_residual/store.npz (the 10 Hz causal RTMW upper-body pose of every s12 route, data/build_kp_residual_store.py) + anchors_{regime}_{split}.npz; the
	observed frames (<= t) and the target frames (+0.5 / +1 / +2 / +3 s) are gathered from the route tick table on the GPU per batch — nothing about the future is
	ever an input (models/kp_forecaster.causal_normalise reads the observed frames only; tests/test_kp_forecaster.py). TRAIN anchors of both regimes with an anchor
	frame; each epoch samples --rows-per-epoch anchors; AdamW 1e-3 / wd 1e-4, batch 1024, bf16 autocast, grad clip 1, <= --epochs, early stop patience 3 on the
	VALIDATION masked L1 of the most probable mode (`val_l1_top1`, normalised units — the selection objective); seeds 0 / 1 / 2.
	Arms: mb_full (MotionBERT architecture, dense 2 s + sparse 1 Hz history), mb_dense (dense only), mb_1hz (1 Hz only), cwm_full (CockpitWM part-token form, parameter-matched).
	Outputs per anchor (row order = the s12 grid, uid / t witness):
	kpf_*.npz route i32 [N] + route_names (uid = route_names[route]), t, feat f16 [N,10] = the DETACHED distillation channel (6 state probabilities, 3 mode probabilities, valid bit; all 0 when the anchor has no frame or no
	valid reference joint), valid bool[N]
	pred_kp_*.npz res_top1 f16 [N,4,11,2] (residual of the most probable mode in FRAME units = normalised residual x scale), ref f16 [N,11,2], ref_ok bool[N,11],
	scale f16[N], p_state f16 [N,3,2], p_mode f16 [N,3], valid bool[N]
	No CAN column, label or openpilot DMS field is read (the state targets are the pose-derived kin-v2 states).
	"""
	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
	import torch # noqa: E402

	torch.set_num_threads(4)
	CWM3 = Path(__file__).resolve().parents[1]
	for p in (CWM3, CWM3 / "common"):
	if str(p) not in sys.path:
	sys.path.insert(0, str(p))
	from models import kp_forecaster as kf # noqa: E402
	from runmeta import write_run_json # noqa: E402

	STORE = CWM3 / "results" / "kp_residual"
	REGIMES, SPLITS = ("manual", "engaged"), ("train", "val", "test")
	ARMS = {"mb_full": dict(backbone="mb", obs="full", d=112, depth=2, heads=4, mlp_ratio=2),
	"mb_dense": dict(backbone="mb", obs="dense", d=112, depth=2, heads=4, mlp_ratio=2),
	"mb_1hz": dict(backbone="mb", obs="1hz", d=112, depth=2, heads=4, mlp_ratio=2),
	"cwm_full": dict(backbone="cwm", obs="full", d=128, depth=4, heads=4, ffn=512)}
	HT = torch.tensor(kf.HORIZON_TICKS, dtype=torch.long)


	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")


	class Store:
	"""the pose store on `device`; gather(pos, route, offsets) -> xy [B,T,11,2], conf [B,T,11], valid [B,T,11]."""

	def __init__(self, device):
	z = np.load(STORE / "store.npz", allow_pickle=True)
	self.kp = torch.from_numpy(z["kp"]).to(device) # f16 [F,11,3]
	self.valid = torch.from_numpy(z["valid"]).to(device)
	self.lut = torch.from_numpy(z["lut"].astype(np.int64)).to(device)
	self.lut_off = torch.from_numpy(z["route_lut_off"]).to(device)
	self.lut_len = torch.from_numpy(z["route_lut_len"]).to(device)
	self.L = int(self.lut.numel())
	self.device = device
	self.n_frames = int(self.kp.shape[0])

	def gather(self, pos, route, offsets):
	idx = pos[:, None] + offsets.to(self.device)[None]
	lo = self.lut_off[route][:, None]
	inside = (idx >= lo) & (idx < lo + self.lut_len[route][:, None])
	fi = self.lut[idx.clamp(0, self.L - 1)]
	fi = torch.where(inside, fi, torch.full_like(fi, -1))
	ok = fi >= 0
	k = self.kp[fi.clamp_min(0)].float() # [B,T,11,3]
	v = self.valid[fi.clamp_min(0)] & ok[..., None]
	xy = torch.nan_to_num(k[..., :2]) * v[..., None].float()
	return xy, torch.nan_to_num(k[..., 2]) * v.float(), v


	ROUTE_NAMES = None


	def route_names():
	"""store.npz route_names: uid == route name, so anchors carry the route index only (disk: the 2.6 M uid strings were 1.4 GB)."""
	global ROUTE_NAMES
	if ROUTE_NAMES is None:
	ROUTE_NAMES = np.load(STORE / "store.npz", allow_pickle=True)["route_names"].astype(str)
	return ROUTE_NAMES


	def load_anchors(regime, split, max_anchors=None, seed=0):
	z = np.load(STORE / f"anchors_{regime}_{split}.npz", allow_pickle=True)
	d = {k: z[k] for k in ("t", "dongle", "mode", "route", "lut_pos", "S", "S0", "kin_row_ok")}
	d["uid"] = route_names()[d["route"]]
	if max_anchors and len(d["t"]) > max_anchors:
	sel = np.sort(np.random.default_rng(seed).choice(len(d["t"]), max_anchors, replace=False))
	d = {k: v[sel] for k, v in d.items()}
	return d


	def to_dev(d, device):
	return {"route": torch.from_numpy(d["route"].astype(np.int64)).to(device), "pos": torch.from_numpy(d["lut_pos"].astype(np.int64)).to(device),
	"S": torch.from_numpy(d["S"]).to(device)}


	def run_batch(model, store, a, idx, need_loss=True):
	pos, route = a["pos"][idx], a["route"][idx]
	has = pos >= 0
	posc = pos.clamp_min(0)
	xy, conf, valid = store.gather(posc, route, model.offsets)
	valid = valid & has[:, None, None]
	out = model(xy, conf, valid)
	if not need_loss:
	return out, has
	fxy, _, fvalid = store.gather(posc, route, HT)
	fvalid = fvalid & has[:, None, None]
	total, parts, err = kf.kp_loss(out, fxy, fvalid, a["S"][idx])
	return out, has, total, parts, err, fxy, fvalid


	@torch.no_grad()
	def evaluate(model, store, a, batch, amp):
	"""validation: masked L1 (normalised units) of the most probable mode and of the oracle mode, mean over anchors with a target."""
	model.eval()
	N = a["pos"].numel()
	s_top, s_wta, n = 0.0, 0.0, 0
	for i in range(0, N, batch):
	idx = torch.arange(i, min(N, i + batch), device=store.device)
	with torch.autocast("cuda", dtype=torch.bfloat16, enabled=amp):
	out, has, total, parts, err, fxy, fvalid = run_batch(model, store, a, idx)
	m = (fvalid & out["ref_ok"][:, None]).float()[:, :, None, :, None]
	hasrow = (m.sum((1, 2, 3, 4)) > 0)
	top = out["mode_logits"].float().argmax(1)
	e_top = err.float().gather(1, top[:, None]).squeeze(1)
	s_top += float((e_top * hasrow.float()).sum()); s_wta += float((err.float().min(1).values * hasrow.float()).sum()); n += int(hasrow.sum())
	model.train()
	return {"val_l1_top1": s_top / max(n, 1), "val_l1_oracle": s_wta / max(n, 1), "n_val_rows": n}


	@torch.no_grad()
	def predict(model, store, a_np, a, batch, amp, out_dir, regime, split, full):
	model.eval()
	N = a["pos"].numel()
	feat = np.zeros((N, kf.KPF_DIM), np.float16)
	valid_row = np.zeros(N, bool)
	if full:
	res = np.zeros((N, kf.H, kf.NJ, 2), np.float16); ref = np.zeros((N, kf.NJ, 2), np.float16); ref_ok = np.zeros((N, kf.NJ), bool)
	scale = np.zeros(N, np.float16); p_state = np.zeros((N, kf.N_STATE_H, kf.N_STATES), np.float16); p_mode = np.zeros((N, kf.N_MODES), np.float16)
	for i in range(0, N, batch):
	idx = torch.arange(i, min(N, i + batch), device=store.device)
	with torch.autocast("cuda", dtype=torch.bfloat16, enabled=amp):
	out, has = run_batch(model, store, a, idx, need_loss=False)
	ok = has & out["ref_ok"].any(1)
	ch = kf.kpf_channel({k: v.float() for k, v in out.items()}, ok)
	sl = slice(i, i + len(idx))
	feat[sl] = ch.cpu().numpy().astype(np.float16); valid_row[sl] = ok.cpu().numpy()
	if full:
	top = out["mode_logits"].float().argmax(1)
	r = out["res"].float()[torch.arange(len(idx), device=store.device), :, top] # [B,H,J,2]
	res[sl] = (r * out["scale"].float()[:, None, None, None]).cpu().numpy().astype(np.float16)
	ref[sl] = out["ref"].float().cpu().numpy().astype(np.float16); ref_ok[sl] = out["ref_ok"].cpu().numpy(); scale[sl] = out["scale"].float().cpu().numpy().astype(np.float16)
	p_state[sl] = torch.sigmoid(out["state_logits"].float()).cpu().numpy().astype(np.float16); p_mode[sl] = torch.softmax(out["mode_logits"].float(), -1).cpu().numpy().astype(np.float16)
	rn = route_names() # uid = route_names[route] (reconstructed by the readers; no 2.6 M-string array per file)
	np.savez(out_dir / f"kpf_{regime}_{split}.npz", route=a_np["route"].astype(np.int32), route_names=rn, t=a_np["t"], feat=feat, valid=valid_row,
	names=np.asarray([f"p_{s}_{h:g}s" for h in (1, 2, 3) for s in ("hands_on", "head_turned")] + [f"p_mode{m}" for m in range(kf.N_MODES)] + ["valid"]))
	if full:
	np.savez(out_dir / f"pred_kp_{regime}_{split}.npz", route=a_np["route"].astype(np.int32), route_names=rn, t=a_np["t"], res_top1=res, ref=ref, ref_ok=ref_ok, scale=scale, p_state=p_state, p_mode=p_mode,
	valid=valid_row, horizons_s=np.asarray(kf.HORIZONS_S), state_horizons_s=np.asarray([1.0, 2.0, 3.0]), joints=np.load(STORE / "store.npz", allow_pickle=True)["joints"])
	# (uid strings are not stored: readers rebuild uid = route_names[route] — data/cwm3_dataset.load_split "kpf", eval/kp_residual_eval.py)
	model.train()
	return float(valid_row.mean())


	def main():
	ap = argparse.ArgumentParser()
	ap.add_argument("--arm", required=True, choices=list(ARMS))
	ap.add_argument("--seed", type=int, default=0)
	ap.add_argument("--out", default="")
	ap.add_argument("--epochs", type=int, default=12)
	ap.add_argument("--rows-per-epoch", type=int, default=400_000)
	ap.add_argument("--batch", type=int, default=1024)
	ap.add_argument("--lr", type=float, default=1e-3)
	ap.add_argument("--weight-decay", type=float, default=1e-4)
	ap.add_argument("--patience", type=int, default=3)
	ap.add_argument("--minutes", type=float, default=90)
	ap.add_argument("--max-anchors", type=int, default=0, help="smoke only")
	ap.add_argument("--cpu", action="store_true", help="smoke only")
	ap.add_argument("--no-amp", action="store_true")
	a = ap.parse_args()
	out = Path(a.out) if a.out else CWM3 / "results" / "runs" / f"KPF_{a.arm}_s{a.seed}" + ("_smoke" if a.max_anchors else "")
	out.mkdir(parents=True, exist_ok=True)
	t_start = time.time()
	device = torch.device("cpu" if a.cpu else "cuda")
	amp = (not a.no_amp) and device.type == "cuda"
	torch.manual_seed(a.seed); np.random.seed(a.seed)
	cfg = {"arm": a.arm, **ARMS[a.arm], "seed": a.seed, "epochs": a.epochs, "rows_per_epoch": a.rows_per_epoch, "batch": a.batch, "lr": a.lr, "weight_decay": a.weight_decay,
	"patience": a.patience, "selection": "val_l1_top1 (masked L1 of the most probable mode, normalised units)", "max_anchors": a.max_anchors, "amp": amp,
	"loss": "wta_l1 + 0.1 mode_ce + 0.5 state_bce", "horizons_s": list(kf.HORIZONS_S), "state_horizons_s": [1.0, 2.0, 3.0]}
	(out / "config.json").write_text(json.dumps(cfg, indent=1))
	store = Store(device)
	log(f"=== train_kp_residual {a.arm} s{a.seed} out={out} store frames={store.n_frames:,} device={device} ===", out)
	tr_np = {r: load_anchors(r, "train", a.max_anchors or None, a.seed) for r in REGIMES}
	va_np = {r: load_anchors(r, "val", (a.max_anchors // 4 if a.max_anchors else None), a.seed) for r in REGIMES}
	tr = {k: torch.cat([to_dev(tr_np[r], device)[k] for r in REGIMES]) for k in ("route", "pos", "S")}
	va = {k: torch.cat([to_dev(va_np[r], device)[k] for r in REGIMES]) for k in ("route", "pos", "S")}
	tr_idx_all = torch.nonzero(tr["pos"] >= 0).squeeze(1)
	# scale fallback = TRAIN median shoulder width (observed frames) over a 50k-anchor sample: recorded, used identically at prediction time
	with torch.no_grad():
	smp = tr_idx_all[torch.randperm(len(tr_idx_all), generator=torch.Generator().manual_seed(0))[:50_000].to(device)]
	xy, _, v = store.gather(tr["pos"][smp], tr["route"][smp], kf.obs_offsets("full"))
	sw = (xy[:, :, kf.L_SHO] - xy[:, :, kf.R_SHO]).norm(dim=-1)[v[:, :, kf.L_SHO] & v[:, :, kf.R_SHO]]
	scale_fb = float(sw.median()) if sw.numel() else kf.SCALE_FALLBACK
	cfg["scale_fallback"] = scale_fb
	model = kf.KpForecaster(**{k: v for k, v in ARMS[a.arm].items()}, scale_fallback=scale_fb).to(device)
	cfg["params"] = {"total": kf.n_params(model), "backbone": kf.n_params(model.backbone), "heads": kf.n_params(model) - kf.n_params(model.backbone)}
	(out / "config.json").write_text(json.dumps(cfg, indent=1))
	log(f"params {cfg['params']} train anchors {len(tr_idx_all):,} val anchors {int((va['pos'] >= 0).sum()):,} scale_fallback {scale_fb:.4f}", out)
	opt = torch.optim.AdamW(model.parameters(), lr=a.lr, weight_decay=a.weight_decay)
	g = torch.Generator(device="cpu").manual_seed(a.seed)
	best, best_ep, bad, hist = float("inf"), -1, 0, []
	for ep in range(a.epochs):
	t_ep = time.time()
	perm = tr_idx_all[torch.randperm(len(tr_idx_all), generator=g)[:a.rows_per_epoch].to(device)]
	sums, nb = {}, 0
	for i in range(0, len(perm), a.batch):
	idx = perm[i:i + a.batch]
	with torch.autocast("cuda", dtype=torch.bfloat16, enabled=amp):
	out_, has, total, parts, err, _, _ = run_batch(model, store, tr, idx)
	opt.zero_grad(set_to_none=True)
	total.backward()
	torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
	opt.step()
	for k, v in parts.items():
	sums[k] = sums.get(k, 0.0) + v
	nb += 1
	if (time.time() - t_start) / 60 > a.minutes:
	break
	ev = evaluate(model, store, va, 4096, amp)
	rec = {"epoch": ep, {f"train_{k}": v / max(nb, 1) for k, v in sums.items()}, ev, "sec": round(time.time() - t_ep, 1)}
	hist.append(rec)
	with open(out / "history.jsonl", "a") as f:
	f.write(json.dumps(rec) + "\n")
	improved = ev["val_l1_top1"] < best - 1e-6
	if improved:
	best, best_ep, bad = ev["val_l1_top1"], ep, 0
	torch.save({"model": model.state_dict(), "cfg": cfg, "epoch": ep, "val": ev}, out / "ckpt_best.pt")
	(out / "BEST").write_text("ckpt_best.pt")
	else:
	bad += 1
	log(f"ep {ep} " + " ".join(f"{k}={v:.4f}" if isinstance(v, float) else f"{k}={v}" for k, v in rec.items() if k != "epoch") + (" *" if improved else ""), out)
	if bad >= a.patience or (time.time() - t_start) / 60 > a.minutes:
	break
	(out / "TRAIN_DONE").write_text(json.dumps({"best_epoch": best_ep, "best_val_l1_top1": best, "epochs_run": len(hist)}))
	model.load_state_dict(torch.load(out / "ckpt_best.pt", map_location=device, weights_only=False)["model"])
	cover = {}
	for r in REGIMES:
	for s in SPLITS:
	a_np = tr_np[r] if s == "train" else (va_np[r] if s == "val" else load_anchors(r, "test", (a.max_anchors // 2 if a.max_anchors else None), a.seed))
	cover[f"{r}_{s}"] = predict(model, store, a_np, to_dev(a_np, device), 4096, amp, out, r, s, full=(s != "train"))
	log(f"predicted {r}/{s}: valid share {cover[f'{r}_{s}']:.4f}", out)
	(out / "PRED_DONE").write_text(json.dumps(cover))
	write_run_json(out, f"train_kp_residual_{a.arm}_s{a.seed}", [Path(__file__), CWM3 / "models" / "kp_forecaster.py", CWM3 / "data" / "build_kp_residual_store.py"],
	[STORE / "store.npz"] + [STORE / f"anchors_{r}_{s}.npz" for r in REGIMES for s in SPLITS], seed=a.seed, t_start=t_start,
	extra={"cfg": cfg, "history": hist, "best_epoch": best_ep, "best_val_l1_top1": best, "epochs_run": len(hist), "gpu_wall_s": round(time.time() - t_start, 1),
	"coverage": cover, "smoke": bool(a.max_anchors or a.cpu), "dms_fields_used": [], "inputs_note": "own causal RTMW pose (results/pose_causal_routes[_v2]); kin-v2 states as targets"})
	log("DONE", out)


	if __name__ == "__main__":
	main()