Buckets:
| """Test-time interventions mirroring Driver-WM Table 3 for feature-conditioned models (test split). | |
| python aide/eval/interventions.py --run cwm3d_vis_sem_s42_v2 [--run driverwm_repro_s42_v2 ...] [--n-boot 2000] | |
| [--reuse-existing] [--only zero_ext,shift_large] | |
| For each run: rebuild the model from ckpt_best.pt (EMA weights + saved args), compute the factual rollout and the | |
| intervened rollouts on the OBSERVED features only (frames 0..4 of the protocol window; nothing after frame 4 is read): | |
| do(Ext = 0) zero_ext exterior (front/left/right) features zeroed on the 5 observed frames | |
| do(Ext = swap_clip) swap_ext exterior features replaced by another test clip's (fixed derangement, seed 0) | |
| do(Ext = shift k) shift_large/small exterior features delayed by k = 2 / 1 protocol steps (frame t receives the | |
| exterior features of frame max(t-k, 0): edge-padded, observed frames only) | |
| do(drop view v) drop_view_{front,left,right} one exterior view zeroed (the mirror table reports mean and max) | |
| do(lambda_ext = l) ext_scale0 / ext_scale2 (CWM3-D) exterior visual TOKENS multiplied by 0 / 2 before the trunk | |
| (`CWM3D.encode(ext_scale=...)`; l = 1 is bit-identical to the factual pass) | |
| do(lambda_CA = l) gate_override0/1/2 (CWM3-D traffic-gate variant) the traffic->driver gate clamped to 0 / 1 / 2 | |
| gate0/1/2 (Driver-WM reproduction) its gate clamped to 0 / 1 / 2 | |
| (+ ego) zero_ext_ego / swap_ext_ego the ego-motion token treated as exterior information | |
| Reports (i) deviation from the factual rollout in px: Delta All, Delta HM subset, Delta h=5, Delta Head (face + body | |
| head joints 0-4/17/18), Delta Hands (both hands), and Delta All normalised by the factual GT All-MPJPE (Delta/All, %); | |
| (ii) the direct GT error of each intervened rollout via the harness (run name <run>__do_<name>, paired ref `factual` | |
| = the run itself). Writes results/runs/<run>/interventions_<utc>.json (append-only; --reuse-existing copies entries | |
| already present in the latest json instead of recomputing them). CPU by default, 4 threads. | |
| """ | |
| from __future__ import annotations | |
| import os | |
| for _v in ("OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS", "NUMEXPR_NUM_THREADS"): | |
| os.environ.setdefault(_v, "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(int(os.environ["OMP_NUM_THREADS"])) | |
| ROOT = Path(__file__).resolve().parents[2] | |
| sys.path.insert(0, str(ROOT)) | |
| from aide.data import protocol as P # noqa: E402 | |
| from aide.eval import metrics as M # noqa: E402 | |
| from aide.eval.harness import load_split, run_harness, summary_line # noqa: E402 | |
| from aide.train import train_kp # noqa: E402 | |
| from common.runmeta import write_run_json # noqa: E402 | |
| EXT = [P.VIEWS.index(v) for v in ("front", "left", "right")] | |
| HEAD_J = np.array(sorted(set(P.PARTS["face"].tolist()) | {0, 1, 2, 3, 4, 17, 18})) | |
| HANDS_J = np.concatenate([P.PARTS["lhand"], P.PARTS["rhand"]]) | |
| SHIFT_STEPS = {"shift_large": 2, "shift_small": 1} | |
| def derangement(n, seed=0): | |
| rng = np.random.default_rng(seed) | |
| perm = rng.permutation(n) | |
| return perm[(np.argsort(perm) + 1) % n] # cyclic shift along a random cycle -> no fixed point | |
| def zero_ext(v): | |
| v = v.clone() | |
| v[:, :, EXT] = 0 | |
| return v | |
| def shift_ext(v, k): | |
| """Delay the exterior features of the OBSERVED window by k steps with edge padding: frame t gets the exterior | |
| features of frame max(t - k, 0). v: [B, T_OBS, 4, D]; the in-cabin view is untouched. Pure function of the | |
| observed frames (indices never exceed T_OBS - 1).""" | |
| T = v.shape[1] | |
| assert T == P.T_OBS, v.shape | |
| idx = torch.clamp(torch.arange(T, device=v.device) - int(k), min=0) | |
| out = v.clone() | |
| out[:, :, EXT] = v[:, idx][:, :, EXT] | |
| return out | |
| def drop_view(v, view): | |
| v = v.clone() | |
| v[:, :, P.VIEWS.index(view)] = 0 | |
| return v | |
| def load_model(run, device): | |
| ck = torch.load(P.RESULTS / "runs" / run / "ckpt_best.pt", map_location=device) | |
| a = argparse.Namespace(**ck["args"]) | |
| model, code_files = train_kp.build_model(a.model, a) | |
| model.load_state_dict(ck["model"]) | |
| return model.to(device).eval(), a, code_files | |
| def deviation(p_int, p_fact, hm, fact_gt_all=None): | |
| d = np.sqrt(((p_int - p_fact) ** 2).sum(-1)) # [N,H,J] | |
| out = {"all": float(d.mean()), "hm": float(d[hm].mean()), "h5": float(d[:, -1].mean()), | |
| "head": float(d[:, :, HEAD_J].mean()), "hands": float(d[:, :, HANDS_J].mean())} | |
| if fact_gt_all: | |
| out["all_norm_pct"] = 100.0 * out["all"] / fact_gt_all # Delta All / factual GT All-MPJPE | |
| return out | |
| def latest_interventions(run): | |
| c = sorted((P.RESULTS / "runs" / run).glob("interventions_2*.json")) | |
| return json.loads(c[-1].read_text()) if c else None | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--run", action="append", required=True) | |
| ap.add_argument("--n-boot", type=int, default=2000, help="clip bootstrap for the GT-error rows (2k: point estimates matter here)") | |
| ap.add_argument("--cuda", action="store_true") | |
| ap.add_argument("--no-csv", action="store_true") | |
| ap.add_argument("--reuse-existing", action="store_true", help="copy entries already in the latest interventions json") | |
| ap.add_argument("--only", default="", help="comma list of intervention names (debug)") | |
| a = ap.parse_args() | |
| t0 = time.time() | |
| device = torch.device("cuda" if a.cuda else "cpu") | |
| d = load_split("test") | |
| S10, ids = d["S10"], d["clip_id"] | |
| gt = S10[:, P.T_OBS:] | |
| hm, _, _ = M.hm_mask(gt) | |
| perm = derangement(len(ids), 0) | |
| assert (perm != np.arange(len(ids))).all() | |
| summary = {} | |
| for run in a.run: | |
| model, ra, code_files = load_model(run, device) | |
| repro = train_kp.is_repro(ra.model) | |
| z = np.load(ra.feats, allow_pickle=False) | |
| assert np.array_equal(z["clip_id"].astype(str), np.load(P.CACHE_NPZ)["clip_id"].astype(str)) | |
| vis = torch.as_tensor(z[ra.vis_key][np.load(P.CACHE_NPZ)["split"] == "test"]) # [N,10,4,2048] | |
| assert len(vis) == len(ids) | |
| vis_ext_swapped = vis.clone() | |
| vis_ext_swapped[:, :, EXT] = vis[perm][:, :, EXT] | |
| ego = None | |
| if getattr(ra, "egomotion", False): | |
| ze = np.load(ra.ego_feats, allow_pickle=False) | |
| ego = torch.as_tensor(ze["ego"][np.load(P.CACHE_NPZ)["split"] == "test"].astype(np.float32)) | |
| fact = train_kp.predict(model, S10, device, vis=vis, repro=repro, ego=ego) | |
| # exactness check against the run's stored predictions | |
| stored = np.load(P.RESULTS / "runs" / run / "pred_test.npz", allow_pickle=False) | |
| order = {c: i for i, c in enumerate(stored["clip_id"].astype(str))} | |
| st = stored["pred"][[order[c] for c in ids.astype(str)]] | |
| rebuild_maxdiff = float(np.abs(st - fact).max()) | |
| err_f = M.joint_errors(fact, gt) | |
| fact_gt = {"mpjpe_all": M.mpjpe(err_f), "mpjpe_hm": M.mpjpe(err_f[hm]), "mpjpe_h5": float(err_f[:, -1].mean()), | |
| "hm_h5": float(err_f[hm][:, -1].mean())} | |
| # lambda = 1 must be a no-op (bit-identical) for the ext_scale intervention -- checked on the real model too | |
| if not repro and getattr(model, "use_visual", False): | |
| same = train_kp.predict(model, S10, device, vis=vis, repro=repro, ego=ego, model_kw={"ext_scale": 1.0}) | |
| assert np.array_equal(same, fact), "ext_scale=1 is not bit-identical to the factual pass" | |
| has_gate = (not repro) and getattr(model, "gate", None) is not None | |
| plan = {} # name -> callable() -> pred | |
| pr = lambda **kw: (lambda: train_kp.predict(model, S10, device, repro=repro, ego=ego, **{"vis": vis, **kw})) # noqa: E731 | |
| plan["zero_ext"] = pr(intervene=zero_ext) | |
| plan["swap_ext"] = pr(vis=vis_ext_swapped) | |
| for name, k in SHIFT_STEPS.items(): | |
| plan[name] = pr(intervene=lambda v, _k=k: shift_ext(v, _k)) | |
| for view in ("front", "left", "right"): | |
| plan[f"drop_view_{view}"] = pr(intervene=lambda v, _w=view: drop_view(v, _w)) | |
| if ego is not None: # the ego-motion stream is exterior information too: zero / swap it together with the features | |
| plan["zero_ext_ego"] = lambda: train_kp.predict(model, S10, device, vis=vis, repro=repro, intervene=zero_ext, | |
| ego=torch.zeros_like(ego)) | |
| plan["swap_ext_ego"] = lambda: train_kp.predict(model, S10, device, vis=vis_ext_swapped, repro=repro, ego=ego[perm]) | |
| if repro: | |
| for lam in (0.0, 1.0, 2.0): | |
| plan[f"gate{int(lam)}"] = pr(model_kw={"override": lam}) | |
| else: | |
| if getattr(model, "use_visual", False): | |
| for lam in (0.0, 2.0): | |
| plan[f"ext_scale{int(lam)}"] = pr(model_kw={"ext_scale": lam}) | |
| if has_gate: | |
| for lam in (0.0, 1.0, 2.0): | |
| plan[f"gate_override{int(lam)}"] = pr(model_kw={"gate_override": lam}) | |
| if a.only: | |
| plan = {k: f for k, f in plan.items() if k in a.only.split(",")} | |
| prev = latest_interventions(run) if a.reuse_existing else None | |
| rec = {"run": run, "model": ra.model, "rebuild_maxdiff_px": rebuild_maxdiff, "factual_gt": fact_gt, | |
| "n_boot": a.n_boot, "interventions": {}} | |
| for name, fn in plan.items(): | |
| if prev and name in prev.get("interventions", {}) and "deviation_px" in prev["interventions"][name]: | |
| old = dict(prev["interventions"][name]) | |
| old["deviation_px"] = {**old["deviation_px"], "all_norm_pct": 100.0 * old["deviation_px"]["all"] / fact_gt["mpjpe_all"]} | |
| old["reused_from"] = str(sorted((P.RESULTS / "runs" / run).glob("interventions_2*.json"))[-1]) | |
| rec["interventions"][name] = old | |
| print(f"[{run}] do({name}): reused", flush=True) | |
| continue | |
| p_int = fn() | |
| dev = deviation(p_int, fact, hm, fact_gt["mpjpe_all"]) | |
| rd = P.RESULTS / "runs" / f"{run}__do_{name}" | |
| rd.mkdir(parents=True, exist_ok=True) | |
| pp = rd / "pred_test.npz" | |
| np.savez(pp, pred=p_int.astype(np.float32), clip_id=ids.astype(str), split=np.array("test")) | |
| h = run_harness(pp, f"{run}__do_{name}", ra.model, ra.seed, "test", f"intervention:{name};factual={run}", | |
| a.n_boot, refs={"factual": run, "zvshrink": "zv_shrink_val"}, write_csv=not a.no_csv) | |
| print(summary_line(h), flush=True) | |
| m = h["metrics"] | |
| rec["interventions"][name] = {"deviation_px": dev, | |
| "gt": {k: m[k]["value"] for k in ("mpjpe_all", "mpjpe_hm", "mpjpe_h5", "hm_h5") if k in m}, | |
| "diff_vs_factual_all": m.get("diff_vs_factual_all"), | |
| "diff_vs_factual_hm": m.get("diff_vs_factual_hm"), | |
| "metrics_json": h["metrics_path"]} | |
| print(f"[{run}] do({name}): dev All {dev['all']:.2f} ({dev['all_norm_pct']:.1f} % of All) HM {dev['hm']:.2f} h5 {dev['h5']:.2f} " | |
| f"head {dev['head']:.2f} hands {dev['hands']:.2f} | GT All {m['mpjpe_all']['value']:.2f} HM {m['mpjpe_hm']['value']:.2f}", flush=True) | |
| if repro: | |
| with torch.no_grad(): | |
| g = [] | |
| for i in range(0, len(S10), 256): | |
| g.append(model(vis[i:i + 256].to(device)[:, :P.T_OBS])["gate_mean"].cpu().numpy()) | |
| rec["gate_mean_per_step"] = np.concatenate(g).mean(0).tolist() | |
| stamp = time.strftime("%Y%m%dT%H%M%SZ", time.gmtime()) | |
| (P.RESULTS / "runs" / run / f"interventions_{stamp}.json").write_text(json.dumps(rec, indent=1)) | |
| summary[run] = rec | |
| write_run_json(P.RESULTS / "tables", "interventions", [__file__, train_kp.__file__, P.__file__], [P.CACHE_NPZ], | |
| seed=0, t_start=t0, input_listing=False, extra={"runs": list(summary), "summary": summary}) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 12.4 kB
- Xet hash:
- 2584644dbce4053c05afa201def4f07d68785c2a5b6798564e9c9b26ae95c156
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.