Buckets:

AsianPlayer
/

TriDrive

Files

xet

AsianPlayer/TriDrive / code /aide /eval /report_baton_pretrain.py

AsianPlayer

16 days ago

download

raw

19 kB

"""cwm3/aide/eval/report_baton_pretrain.py — GPU-baselines experiment A table: BATON -> AIDE in-cabin pose pretraining (aide/results/tables/baton_pretrain.md).
Pairs (identical AIDE test clips, n = 609; paired per-clip bootstrap, 10k resamples, seed 0 — aide/eval/metrics.paired_diff):
  cwm3d_kp_batonpre_s<k>_v2      vs  cwm3d_kp_s<k>_v2        (keypoint-only CWM3-D, seeds 42 / 123 / 456; fine-tuned from the BATON-pretrained checkpoint vs the seed init)
  cwm3d_vis_sem_batonpre_s<k>_v2 vs  cwm3d_vis_sem_s<k>_v2   (the vis+sem B2 model whose trunk starts from the pretrained checkpoint)
  cwm3d_kp_batonpre_zeroshot_s42_v2 (the pretrained model on AIDE without any AIDE training; test evaluated once) vs zero-velocity / cwm3d_kp_s42_v2 (level rows)
Cells: All-MPJPE, HM-MPJPE (top-60 GT motion), PCK@0.05 (96 px), HM h=5, per-part All (body / face / lhand / rhand), per-horizon All (h1..h5); Δ = A − B per seed with CI;
3-seed mean ± sd; verdict = "BEATS" iff Δ CI excludes 0 in the improving direction (negative for errors, positive for PCK) on >= 2/3 seeds with the same sign of every seed's Δ,
"WORSE" likewise in the other direction, else "INCONCLUSIVE" / "NOT_WORSE" (all CIs cover 0 and |Δ| < 0.25 px / 0.25 pp).
Pre-registered adoption rule (kill): adopt the pretraining only if Δ All is BEATS (CI excl. 0 on >= 2/3 seeds, improving); HM h=5 is reported regardless.
Also reported: the pretraining run (corpus size, ms/step, GPU minutes, zero-shot AIDE-val curve) and the fine-tune runs' best epoch / GPU minutes / init report.
Nothing is typed by hand: every number comes from pred_test.npz files, metrics_test.json, run json files and the corpus summary.
  python aide/eval/report_baton_pretrain.py [--seeds 42,123,456] [--n-boot 10000] [--out aide/results/tables/baton_pretrain.md]
"""
from __future__ import annotations

import argparse
import glob
import json
import sys
import time
from pathlib import Path

import numpy as np

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 align, load_split  # noqa: E402
from common.runmeta import write_run_json  # noqa: E402

RUNS = P.RESULTS / "runs"
OUT = P.RESULTS / "tables" / "baton_pretrain.md"
PAIRS = [("keypoint-only CWM3-D-kp", "cwm3d_kp_batonpre_s{k}_v2", "cwm3d_kp_s{k}_v2", "pre-registered"),
         ("vis+sem CWM3-D (B2 main)", "cwm3d_vis_sem_batonpre_s{k}_v2", "cwm3d_vis_sem_s{k}_v2", "pre-registered"),
         ("keypoint-only CWM3-D-kp, TRUNK-only init — POST HOC arm (added after the pre-registered arms' result; tokenizer / blocks / ln_f from the pretraining, horizon head at its zero init)",
          "cwm3d_kp_batonpre_trunk_s{k}_v2", "cwm3d_kp_s{k}_v2", "post hoc"),
         ("keypoint-only CWM3-D-kp, TRUNK-only init — PRE-REGISTERED REPLICATION (independent pretraining seed 123, `cwm3d_kp_batonpre_pretrain_s123`; registered 2026-09-12 12:35 EDT before the run: adopt only if Δ All is BEATS again)",
          "cwm3d_kp_batonpre_trunkp123_s{k}_v2", "cwm3d_kp_s{k}_v2", "replication"),
         ("vis+sem CWM3-D (B2 main), TRUNK-only init — PRE-REGISTERED (registered 2026-09-12 12:43 EDT before launch, coordinator request): pretraining seed 123 trunk "
          "(tokenizer / blocks / ln_f), visual / semantic modules and horizon head at the seed init; recipe = the B2 main run verbatim (--visual --semantic); "
          "rule: adopt as a reported row iff Δ All-MPJPE CI excl. 0 on ≥ 2/3 seeds in the improving direction (same sign on all seeds)",
          "cwm3d_vis_sem_batonpre_trunkp123_s{k}_v2", "cwm3d_vis_sem_s{k}_v2", "pre-registered"),
         ("vis + attention heads + learned view gate, geometry-selected checkpoint (a) (`cwm3d_vis_sema_vg`), TRUNK-only init — PRE-REGISTERED (registered 2026-09-12 12:43 EDT "
          "before launch): pretraining seed 123 trunk, everything else at the seed init; recipe = `cwm3d_vis_sema_vg_s<k>_v2` verbatim (--visual --semantic --sem-input attn "
          "--lambda-sem 1 --label-smoothing 0 --class-balanced 0 --dual-select --sem-view-gate; the geometry-selected checkpoint is compared, the __f1sel companion is written "
          "as usual and reported by the AIDE agent's tables); rule: adopt as a reported row iff Δ All-MPJPE CI excl. 0 on ≥ 2/3 seeds (improving, same sign); semantic F1 "
          "of the (a) checkpoint is reported descriptively",
          "cwm3d_vis_sema_vg_batonpre_trunkp123_s{k}_v2", "cwm3d_vis_sema_vg_s{k}_v2", "pre-registered")]
ZEROSHOT = "cwm3d_kp_batonpre_zeroshot_s42_v2"
PRETRAIN = "cwm3d_kp_batonpre_pretrain_s42"
PRETRAIN_REP = "cwm3d_kp_batonpre_pretrain_s123"
TOL_NOT_WORSE = 0.25


def load_pred(name, d):
    p = RUNS / name / "pred_test.npz"
    if not p.exists():
        return None
    z = np.load(p, allow_pickle=False)
    return np.asarray(z["pred"], dtype=np.float64)[align(z, d["clip_id"])]


def per_clip_cells(pred, gt, hm):
    """-> {cell: (per-clip values [N] or [n_hm], higher_better)}."""
    err = M.joint_errors(pred, gt)                                  # [N,5,J]
    cells = {"All-MPJPE": (err.mean((1, 2)), False), "HM-MPJPE": (err[hm].mean((1, 2)), False),
             "PCK@0.05": (100.0 * (err <= P.pck_threshold_px(0.05)).mean((1, 2)), True), "HM h=5": (err[hm][:, 4].mean(1), False)}
    for k, idx in P.PARTS.items():
        cells[f"All {k}"] = (err[:, :, idx].mean((1, 2)), False)
    for h in range(P.T_PRED):
        cells[f"All h={h + 1}"] = (err[:, h].mean(1), False)
    return cells


def verdict(rows, higher_better):
    """rows: list of dicts with mean / ci_lo / ci_hi per seed."""
    if len(rows) < 3:
        return "INSUFFICIENT"
    d = np.array([r["mean"] for r in rows]); lo = np.array([r["ci_lo"] for r in rows]); hi = np.array([r["ci_hi"] for r in rows])
    improving = (d > 0) if higher_better else (d < 0)
    excl = (lo > 0) | (hi < 0)
    same = np.all(np.sign(d) == np.sign(d.mean())) and d.mean() != 0
    if excl.sum() >= 2 and same and improving.all():
        return "BEATS"
    if excl.sum() >= 2 and same and (~improving).all():
        return "WORSE"
    if np.all((lo <= 0) & (hi >= 0)) and np.all(np.abs(d) < TOL_NOT_WORSE):
        return "NOT_WORSE"
    return "INCONCLUSIVE"


def run_json(name):
    fs = sorted(glob.glob(str(RUNS / name / f"run_{name}_2*.json")))
    return json.load(open(fs[-1])) if fs else None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--seeds", default="42,123,456")
    ap.add_argument("--n-boot", type=int, default=10000)
    ap.add_argument("--out", default=str(OUT))
    a = ap.parse_args()
    t0 = time.time()
    seeds = [int(s) for s in a.seeds.split(",")]
    d = load_split("test")
    gt = d["S10"][:, P.T_OBS:]
    hm, thr, _ = M.hm_mask(gt)
    stamp = time.strftime("%Y%m%dT%H%M%SZ", time.gmtime())
    pre = run_json(PRETRAIN)
    corpus = json.load(open(P.RESULTS / "cache" / "baton_pose_corpus_summary.json")) if (P.RESULTS / "cache" / "baton_pose_corpus_summary.json").exists() else None
    L = ["# BATON → AIDE in-cabin pose pretraining (GPU-baselines experiment A; protocol.md C21) — AIDE test, n = 609, paired per-clip bootstrap (10k)",
         "", f"Generated {stamp} by `aide/eval/report_baton_pretrain.py` from pred_test.npz / metrics_test.json / run json files (locked protocol: driver `first`, grid `floor`, HM top {int(hm.sum())} "
         f"(threshold {thr:.2f} px), PCK @ {P.pck_threshold_px(0.05):.0f} px).  A = fine-tuned from the BATON-pretrained checkpoint, B = the same recipe from the seed init; Δ = A − B "
         "(negative = better for errors, positive = better for PCK).  Verdict: BEATS / WORSE = Δ CI excl. 0 on ≥ 2/3 seeds, all seeds same sign; NOT_WORSE = all CIs cover 0 and |Δ| < 0.25; else INCONCLUSIVE.  "
         "**Pre-registered kill rule: adopt the pretraining only if Δ All-MPJPE is BEATS; HM h=5 reported regardless.**", ""]
    if corpus:
        st, pr = corpus["stats"], corpus["params"]
        L += ["## 1. Pretraining corpus (aide/features/build_baton_pose_corpus.py)", "",
              f"- {st['routes']} BATON TRAIN-split routes from {st['devices']} train-only devices (train / val-test device sets disjoint: {corpus['train_devices']} vs {corpus['eval_devices']} devices, 0 overlap); "
              f"{st['routes_val']} routes held out for the pretraining early stop; {st['frames_total']:,} 10-Hz frames, {st['frames_valid_frac']:.3f} valid (driver detected, finite box / nose / neck, "
              f"≥ {pr['min_frac_body_face']:.0%} of body and of face joints with score ≥ {pr['conf_th']}).",
              f"- Windows: 10 samples at {pr['step_s']:.1f} s steps (span {pr['window_span_s']:.1f} s; AIDE grid 0.267 / 0.333 s, span 2.93 s), anchors every {pr['anchor_stride'] * 0.1:.1f} s; "
              f"{st['windows_valid']:,} valid of {st['windows_candidates']:,} candidates ({st['windows_dropped_gap']:,} dropped for a time gap, {st['windows_dropped_invalid']:,} for an invalid frame) "
              f"→ capped at {pr['max_windows']:,} (seeded) → **{st['windows_final']:,} windows** ({st['windows_val_routes']:,} on the held-out routes).",
              f"- Coordinates: per-window isotropic similarity into the AIDE frame ({pr['similarity']}; the driver box is a validity flag only); "
              f"HALPE 17 (head-top, absent from RTMW) = nose + {pr.get('headtop_proxy_k', 0.65)} (nose − neck); scores quantile-mapped per part to the AIDE train score distribution.",
              f"- Sanity (AIDE pixel frame): shoulder width median BATON {st['sanity_aide_frame']['shoulder_width_px_median']['baton']:.0f} px vs AIDE train {st['sanity_aide_frame']['shoulder_width_px_median']['aide_train']:.0f} px; "
              f"future displacement median {st['sanity_aide_frame']['future_disp_px_median']['baton']:.1f} px vs {st['sanity_aide_frame']['future_disp_px_median']['aide_train']:.1f} px "
              f"(p90 {st['sanity_aide_frame']['future_disp_px_p90']['baton']:.1f} vs {st['sanity_aide_frame']['future_disp_px_p90']['aide_train']:.1f}).", ""]
    if pre:
        hist = pre.get("history", [])
        L += ["## 2. Pretraining run (aide/train/pretrain_cwm3d_baton.py)", "",
              f"- Run `{PRETRAIN}`: {pre['n_params']:,} params (cwm3d d 256 / 4 layers, same init as cwm3d_kp_s42_v2), corpus sha256 {pre['corpus']['corpus_sha256'][:16]}, "
              f"{pre['corpus']['n_windows'] - pre['corpus']['n_windows_val']:,} fit windows; bs {pre['config']['bs']}, lr {pre['config']['lr']}, {pre['epochs_run']} epochs run, best epoch {pre['best_epoch']} "
              f"(held-out-route All-MPJPE {pre['best_corpus_val_all']:.2f} px), stop = {pre['stop_reason']}; **{pre['ms_per_step_mean']:.1f} ms/step, {pre['gpu_minutes']:.1f} GPU-min**.",
              "", "| epoch | ms/step | samples/s | train loss | corpus held-out All | corpus held-out HM | AIDE val zero-shot All | AIDE val zero-shot HM |", "|---|---|---|---|---|---|---|---|"]
        for h in hist:
            L.append(f"| {h['epoch']} | {h['ms_per_step']:.1f} | {h['samples_per_s']:.0f} | {h.get('train_total', float('nan')):.3f} | {h['corpus_val_all']:.2f} | {h['corpus_val_hm']:.2f} | {h['aide_val_zeroshot_all']:.3f} | {h['aide_val_zeroshot_hm']:.2f} |")
        zs = pre.get("zeroshot")
        if zs:
            L += ["", f"Zero-shot run `{zs['name']}` (no AIDE training, test once): val All {zs['val'].get('mpjpe_all', float('nan')):.3f} / HM {zs['val'].get('mpjpe_hm', float('nan')):.2f}; "
                  f"**test All {zs['test'].get('mpjpe_all', float('nan')):.3f} / HM {zs['test'].get('mpjpe_hm', float('nan')):.2f} / PCK@0.05 {zs['test'].get('pck05', float('nan')):.2f} / HM h=5 {zs['test'].get('hm_h5', float('nan')):.2f}**; "
                  f"Δ vs zero-velocity All {zs['test'].get('diff_vs_ref_all', float('nan')):+.2f}, Δ vs cwm3d_kp_s42_v2 All {zs['test'].get('diff_vs_kp_all', float('nan')):+.2f} (AIDE ZV test All 52.93 / HM 139.23)."]
        L.append("")
    pre2 = run_json(PRETRAIN_REP)
    if pre2:
        L += [f"Replication pretraining `{PRETRAIN_REP}` (independent seed): {pre2['epochs_run']} epochs, best epoch {pre2['best_epoch']} (held-out All {pre2['best_corpus_val_all']:.2f} px), "
              f"{pre2['ms_per_step_mean']:.1f} ms/step, {pre2['gpu_minutes']:.1f} GPU-min; zero-shot test All {(pre2.get('zeroshot') or {}).get('test', {}).get('mpjpe_all', float('nan')):.3f}.", ""]
    results = {}
    sec = 3
    for title, a_fmt, b_fmt, status_kind in PAIRS:
        L += [f"## {sec}. {title}: fine-tuned from the BATON-pretrained checkpoint (A) vs seed init (B)", ""]
        sec += 1
        per_seed = {}
        for k in seeds:
            pa, pb = load_pred(a_fmt.format(k=k), d), load_pred(b_fmt.format(k=k), d)
            if pa is None or pb is None:
                continue
            ca, cb = per_clip_cells(pa, gt, hm), per_clip_cells(pb, gt, hm)
            per_seed[k] = {c: {"a": float(ca[c][0].mean()), "b": float(cb[c][0].mean()), **M.paired_diff(ca[c][0], cb[c][0], a.n_boot, 0), "higher_better": ca[c][1]} for c in ca}
        if not per_seed:
            L += [f"_(registered, no finished run pairs yet — arm status: {status_kind}; runs `{a_fmt.format(k='<k>')}` vs `{b_fmt.format(k='<k>')}`, seeds {seeds}; "
                  "verdict rule as in the header: BEATS / WORSE = Δ CI excl. 0 on ≥ 2/3 seeds with the same sign on every seed)_", ""]
            results[title] = {"status": "REGISTERED_NOT_RUN", "kind": status_kind, "a": a_fmt, "b": b_fmt}
            continue
        L += ["| seed | A run | A best epoch / epochs / GPU min | init (tensors loaded / kept at seed init) | B run | B best epoch / epochs / GPU min |", "|---|---|---|---|---|---|"]
        for k in per_seed:
            ra, rb = run_json(a_fmt.format(k=k)), run_json(b_fmt.format(k=k))
            ini = (ra or {}).get("init") or {}
            L.append(f"| {k} | {a_fmt.format(k=k)} | {(ra or {}).get('best_epoch', '—')} / {(ra or {}).get('epochs_run', '—')} / {(ra or {}).get('gpu_minutes', '—')} | "
                     f"{ini.get('parts', '—')}: {ini.get('n_loaded', '—')} / {ini.get('n_missing_or_kept', '—')} (ckpt epoch {ini.get('ckpt_epoch', '—')}) | {b_fmt.format(k=k)} | "
                     f"{(rb or {}).get('best_epoch', '—')} / {(rb or {}).get('epochs_run', '—')} / {(rb or {}).get('gpu_minutes', '—')} |")
        f1rows = []
        for k in per_seed:
            ma, mb = (RUNS / a_fmt.format(k=k) / "metrics_test.json"), (RUNS / b_fmt.format(k=k) / "metrics_test.json")
            if ma.exists() and mb.exists():
                xa, xb = json.load(open(ma))["metrics"], json.load(open(mb))["metrics"]
                if "f1_dbr" in xa and "f1_dbr" in xb:
                    f1rows.append(f"| {k} | " + " / ".join(f"{xa[f'f1_{t}']['value']:.1f}" for t in P.TASKS) + " | " + " / ".join(f"{xb[f'f1_{t}']['value']:.1f}" for t in P.TASKS) + " |")
        if f1rows:
            L += ["", "Semantic macro-F1 DBR / DER / TCR / VCR of the compared (geometry-selected) checkpoints, descriptive (single evaluation, no CI here):", "",
                  "| seed | A | B |", "|---|---|---|"] + f1rows
        L += ["", "| cell | " + " | ".join(f"A s{k}" for k in per_seed) + " | " + " | ".join(f"B s{k}" for k in per_seed) + " | " + " | ".join(f"Δ s{k} [CI]" for k in per_seed) + " | mean Δ ± sd | verdict |",
              "|---|" + "---|" * (3 * len(per_seed)) + "---|---|"]
        cells = list(next(iter(per_seed.values())))
        res = {}
        for c in cells:
            rows = [per_seed[k][c] for k in per_seed]
            v = verdict(rows, rows[0]["higher_better"])
            dm = np.array([r["mean"] for r in rows])
            L.append(f"| {c} | " + " | ".join(f"{r['a']:.2f}" for r in rows) + " | " + " | ".join(f"{r['b']:.2f}" for r in rows) + " | "
                     + " | ".join(f"{r['mean']:+.2f} [{r['ci_lo']:+.2f},{r['ci_hi']:+.2f}]" for r in rows) + f" | {dm.mean():+.2f} ± {dm.std(ddof=1) if len(dm) > 1 else float('nan'):.2f} | {v} |")
            res[c] = {"verdict": v, "per_seed": {str(k): per_seed[k][c] for k in per_seed}}
        v_all, v_h5 = res["All-MPJPE"]["verdict"], res["HM h=5"]["verdict"]
        adopt = v_all == "BEATS"
        if status_kind == "post hoc":
            rule = (f"Δ All-MPJPE = {v_all}: meets the numeric criterion, but this arm was added AFTER the pre-registered arms' result → exploratory, NOT adopted on its own; "
                    f"adoption requires the pre-registered replication arm (independent pretraining seed) to be BEATS" if adopt else f"Δ All-MPJPE = {v_all} → not adopted (exploratory arm, criterion not met)")
        elif status_kind == "replication":
            rule = f"Δ All-MPJPE = {v_all} → {'REPLICATED: adopt the trunk-only pretrained init (both the post hoc arm and this pre-registered replication are BEATS)' if adopt else 'NOT replicated → the trunk-only result stays exploratory; nothing adopted'}"
        else:
            rule = f"Δ All-MPJPE = {v_all} → {'ADOPT (pretraining improves All on ≥ 2/3 seeds)' if adopt else 'do NOT adopt (pre-registered rule not met)'}"
        L += ["", f"**Kill rule ({title.split(' — ')[0]}, {status_kind}): {rule}; HM h=5 Δ = {v_h5} (mean {np.mean([res['HM h=5']['per_seed'][str(k)]['mean'] for k in per_seed]):+.2f} px).**", ""]
        results[title] = {"status": "DONE", "kind": status_kind, "n_seeds": len(per_seed), "verdict_all": v_all, "verdict_hm_h5": v_h5, "meets_rule": adopt, "cells": res}
    L += ["## Provenance", "", "- Pretraining corpus: BATON train-split, train-only-device routes (no val / test device, no val / test route); AIDE side: TRAIN-split statistics only (box target, score quantiles).",
          "- Fine-tuning: `aide/train/train_kp.py --init-ckpt … --init-parts all` with the Stage-A / B2 recipes verbatim (val All-MPJPE selection, EMA, patience 40; test evaluated once per run).",
          "- No AIDE test clip entered any training or selection; the zero-shot run uses no AIDE data at all before its single test evaluation.", ""]
    Path(a.out).parent.mkdir(parents=True, exist_ok=True)
    Path(a.out).write_text("\n".join(L))
    Path(a.out).with_suffix(".json").write_text(json.dumps({"stamp": stamp, "seeds": seeds, "results": results}, indent=1, default=str))
    write_run_json(P.RESULTS / "tables", "report_baton_pretrain", [Path(__file__), M.__file__, P.__file__], [P.CACHE_NPZ], seed=0, t_start=t0, input_listing=False,
                   extra={"out": a.out, "seeds": seeds, "summary": {k: {kk: vv for kk, vv in v.items() if kk != "cells"} for k, v in results.items()}, "dms_fields_used": []})
    print(f"written {a.out}")
    for k, v in results.items():
        print(f"  {k}: {v}" if v.get("status") != "DONE" else f"  {k[:60]}: All {v['verdict_all']}, HM h=5 {v['verdict_hm_h5']}, meets_rule={v['meets_rule']} [{v['kind']}] ({v['n_seeds']} seeds)")


if __name__ == "__main__":
    main()

Xet Storage Details

Size:: 19 kB
Xet hash:: 264e5f360633494a726f883fb21d28380cb6dd1a00cdf86ee9031dc82ce3fa0e

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

	"""cwm3/aide/eval/report_baton_pretrain.py — GPU-baselines experiment A table: BATON -> AIDE in-cabin pose pretraining (aide/results/tables/baton_pretrain.md).

	Pairs (identical AIDE test clips, n = 609; paired per-clip bootstrap, 10k resamples, seed 0 — aide/eval/metrics.paired_diff):
	cwm3d_kp_batonpre_s<k>_v2 vs cwm3d_kp_s<k>_v2 (keypoint-only CWM3-D, seeds 42 / 123 / 456; fine-tuned from the BATON-pretrained checkpoint vs the seed init)
	cwm3d_vis_sem_batonpre_s<k>_v2 vs cwm3d_vis_sem_s<k>_v2 (the vis+sem B2 model whose trunk starts from the pretrained checkpoint)
	cwm3d_kp_batonpre_zeroshot_s42_v2 (the pretrained model on AIDE without any AIDE training; test evaluated once) vs zero-velocity / cwm3d_kp_s42_v2 (level rows)
	Cells: All-MPJPE, HM-MPJPE (top-60 GT motion), PCK@0.05 (96 px), HM h=5, per-part All (body / face / lhand / rhand), per-horizon All (h1..h5); Δ = A − B per seed with CI;
	3-seed mean ± sd; verdict = "BEATS" iff Δ CI excludes 0 in the improving direction (negative for errors, positive for PCK) on >= 2/3 seeds with the same sign of every seed's Δ,
	"WORSE" likewise in the other direction, else "INCONCLUSIVE" / "NOT_WORSE" (all CIs cover 0 and \|Δ\| < 0.25 px / 0.25 pp).
	Pre-registered adoption rule (kill): adopt the pretraining only if Δ All is BEATS (CI excl. 0 on >= 2/3 seeds, improving); HM h=5 is reported regardless.
	Also reported: the pretraining run (corpus size, ms/step, GPU minutes, zero-shot AIDE-val curve) and the fine-tune runs' best epoch / GPU minutes / init report.
	Nothing is typed by hand: every number comes from pred_test.npz files, metrics_test.json, run json files and the corpus summary.
	python aide/eval/report_baton_pretrain.py [--seeds 42,123,456] [--n-boot 10000] [--out aide/results/tables/baton_pretrain.md]
	"""
	from __future__ import annotations

	import argparse
	import glob
	import json
	import sys
	import time
	from pathlib import Path

	import numpy as np

	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 align, load_split # noqa: E402
	from common.runmeta import write_run_json # noqa: E402

	RUNS = P.RESULTS / "runs"
	OUT = P.RESULTS / "tables" / "baton_pretrain.md"
	PAIRS = [("keypoint-only CWM3-D-kp", "cwm3d_kp_batonpre_s{k}_v2", "cwm3d_kp_s{k}_v2", "pre-registered"),
	("vis+sem CWM3-D (B2 main)", "cwm3d_vis_sem_batonpre_s{k}_v2", "cwm3d_vis_sem_s{k}_v2", "pre-registered"),
	("keypoint-only CWM3-D-kp, TRUNK-only init — POST HOC arm (added after the pre-registered arms' result; tokenizer / blocks / ln_f from the pretraining, horizon head at its zero init)",
	"cwm3d_kp_batonpre_trunk_s{k}_v2", "cwm3d_kp_s{k}_v2", "post hoc"),
	("keypoint-only CWM3-D-kp, TRUNK-only init — PRE-REGISTERED REPLICATION (independent pretraining seed 123, `cwm3d_kp_batonpre_pretrain_s123`; registered 2026-09-12 12:35 EDT before the run: adopt only if Δ All is BEATS again)",
	"cwm3d_kp_batonpre_trunkp123_s{k}_v2", "cwm3d_kp_s{k}_v2", "replication"),
	("vis+sem CWM3-D (B2 main), TRUNK-only init — PRE-REGISTERED (registered 2026-09-12 12:43 EDT before launch, coordinator request): pretraining seed 123 trunk "
	"(tokenizer / blocks / ln_f), visual / semantic modules and horizon head at the seed init; recipe = the B2 main run verbatim (--visual --semantic); "
	"rule: adopt as a reported row iff Δ All-MPJPE CI excl. 0 on ≥ 2/3 seeds in the improving direction (same sign on all seeds)",
	"cwm3d_vis_sem_batonpre_trunkp123_s{k}_v2", "cwm3d_vis_sem_s{k}_v2", "pre-registered"),
	("vis + attention heads + learned view gate, geometry-selected checkpoint (a) (`cwm3d_vis_sema_vg`), TRUNK-only init — PRE-REGISTERED (registered 2026-09-12 12:43 EDT "
	"before launch): pretraining seed 123 trunk, everything else at the seed init; recipe = `cwm3d_vis_sema_vg_s<k>_v2` verbatim (--visual --semantic --sem-input attn "
	"--lambda-sem 1 --label-smoothing 0 --class-balanced 0 --dual-select --sem-view-gate; the geometry-selected checkpoint is compared, the __f1sel companion is written "
	"as usual and reported by the AIDE agent's tables); rule: adopt as a reported row iff Δ All-MPJPE CI excl. 0 on ≥ 2/3 seeds (improving, same sign); semantic F1 "
	"of the (a) checkpoint is reported descriptively",
	"cwm3d_vis_sema_vg_batonpre_trunkp123_s{k}_v2", "cwm3d_vis_sema_vg_s{k}_v2", "pre-registered")]
	ZEROSHOT = "cwm3d_kp_batonpre_zeroshot_s42_v2"
	PRETRAIN = "cwm3d_kp_batonpre_pretrain_s42"
	PRETRAIN_REP = "cwm3d_kp_batonpre_pretrain_s123"
	TOL_NOT_WORSE = 0.25


	def load_pred(name, d):
	p = RUNS / name / "pred_test.npz"
	if not p.exists():
	return None
	z = np.load(p, allow_pickle=False)
	return np.asarray(z["pred"], dtype=np.float64)[align(z, d["clip_id"])]


	def per_clip_cells(pred, gt, hm):
	"""-> {cell: (per-clip values [N] or [n_hm], higher_better)}."""
	err = M.joint_errors(pred, gt) # [N,5,J]
	cells = {"All-MPJPE": (err.mean((1, 2)), False), "HM-MPJPE": (err[hm].mean((1, 2)), False),
	"PCK@0.05": (100.0 * (err <= P.pck_threshold_px(0.05)).mean((1, 2)), True), "HM h=5": (err[hm][:, 4].mean(1), False)}
	for k, idx in P.PARTS.items():
	cells[f"All {k}"] = (err[:, :, idx].mean((1, 2)), False)
	for h in range(P.T_PRED):
	cells[f"All h={h + 1}"] = (err[:, h].mean(1), False)
	return cells


	def verdict(rows, higher_better):
	"""rows: list of dicts with mean / ci_lo / ci_hi per seed."""
	if len(rows) < 3:
	return "INSUFFICIENT"
	d = np.array([r["mean"] for r in rows]); lo = np.array([r["ci_lo"] for r in rows]); hi = np.array([r["ci_hi"] for r in rows])
	improving = (d > 0) if higher_better else (d < 0)
	excl = (lo > 0) \| (hi < 0)
	same = np.all(np.sign(d) == np.sign(d.mean())) and d.mean() != 0
	if excl.sum() >= 2 and same and improving.all():
	return "BEATS"
	if excl.sum() >= 2 and same and (~improving).all():
	return "WORSE"
	if np.all((lo <= 0) & (hi >= 0)) and np.all(np.abs(d) < TOL_NOT_WORSE):
	return "NOT_WORSE"
	return "INCONCLUSIVE"


	def run_json(name):
	fs = sorted(glob.glob(str(RUNS / name / f"run_{name}_2*.json")))
	return json.load(open(fs[-1])) if fs else None


	def main():
	ap = argparse.ArgumentParser()
	ap.add_argument("--seeds", default="42,123,456")
	ap.add_argument("--n-boot", type=int, default=10000)
	ap.add_argument("--out", default=str(OUT))
	a = ap.parse_args()
	t0 = time.time()
	seeds = [int(s) for s in a.seeds.split(",")]
	d = load_split("test")
	gt = d["S10"][:, P.T_OBS:]
	hm, thr, _ = M.hm_mask(gt)
	stamp = time.strftime("%Y%m%dT%H%M%SZ", time.gmtime())
	pre = run_json(PRETRAIN)
	corpus = json.load(open(P.RESULTS / "cache" / "baton_pose_corpus_summary.json")) if (P.RESULTS / "cache" / "baton_pose_corpus_summary.json").exists() else None
	L = ["# BATON → AIDE in-cabin pose pretraining (GPU-baselines experiment A; protocol.md C21) — AIDE test, n = 609, paired per-clip bootstrap (10k)",
	"", f"Generated {stamp} by `aide/eval/report_baton_pretrain.py` from pred_test.npz / metrics_test.json / run json files (locked protocol: driver `first`, grid `floor`, HM top {int(hm.sum())} "
	f"(threshold {thr:.2f} px), PCK @ {P.pck_threshold_px(0.05):.0f} px). A = fine-tuned from the BATON-pretrained checkpoint, B = the same recipe from the seed init; Δ = A − B "
	"(negative = better for errors, positive = better for PCK). Verdict: BEATS / WORSE = Δ CI excl. 0 on ≥ 2/3 seeds, all seeds same sign; NOT_WORSE = all CIs cover 0 and \|Δ\| < 0.25; else INCONCLUSIVE. "
	"Pre-registered kill rule: adopt the pretraining only if Δ All-MPJPE is BEATS; HM h=5 reported regardless.", ""]
	if corpus:
	st, pr = corpus["stats"], corpus["params"]
	L += ["## 1. Pretraining corpus (aide/features/build_baton_pose_corpus.py)", "",
	f"- {st['routes']} BATON TRAIN-split routes from {st['devices']} train-only devices (train / val-test device sets disjoint: {corpus['train_devices']} vs {corpus['eval_devices']} devices, 0 overlap); "
	f"{st['routes_val']} routes held out for the pretraining early stop; {st['frames_total']:,} 10-Hz frames, {st['frames_valid_frac']:.3f} valid (driver detected, finite box / nose / neck, "
	f"≥ {pr['min_frac_body_face']:.0%} of body and of face joints with score ≥ {pr['conf_th']}).",
	f"- Windows: 10 samples at {pr['step_s']:.1f} s steps (span {pr['window_span_s']:.1f} s; AIDE grid 0.267 / 0.333 s, span 2.93 s), anchors every {pr['anchor_stride'] * 0.1:.1f} s; "
	f"{st['windows_valid']:,} valid of {st['windows_candidates']:,} candidates ({st['windows_dropped_gap']:,} dropped for a time gap, {st['windows_dropped_invalid']:,} for an invalid frame) "
	f"→ capped at {pr['max_windows']:,} (seeded) → {st['windows_final']:,} windows ({st['windows_val_routes']:,} on the held-out routes).",
	f"- Coordinates: per-window isotropic similarity into the AIDE frame ({pr['similarity']}; the driver box is a validity flag only); "
	f"HALPE 17 (head-top, absent from RTMW) = nose + {pr.get('headtop_proxy_k', 0.65)} (nose − neck); scores quantile-mapped per part to the AIDE train score distribution.",
	f"- Sanity (AIDE pixel frame): shoulder width median BATON {st['sanity_aide_frame']['shoulder_width_px_median']['baton']:.0f} px vs AIDE train {st['sanity_aide_frame']['shoulder_width_px_median']['aide_train']:.0f} px; "
	f"future displacement median {st['sanity_aide_frame']['future_disp_px_median']['baton']:.1f} px vs {st['sanity_aide_frame']['future_disp_px_median']['aide_train']:.1f} px "
	f"(p90 {st['sanity_aide_frame']['future_disp_px_p90']['baton']:.1f} vs {st['sanity_aide_frame']['future_disp_px_p90']['aide_train']:.1f}).", ""]
	if pre:
	hist = pre.get("history", [])
	L += ["## 2. Pretraining run (aide/train/pretrain_cwm3d_baton.py)", "",
	f"- Run `{PRETRAIN}`: {pre['n_params']:,} params (cwm3d d 256 / 4 layers, same init as cwm3d_kp_s42_v2), corpus sha256 {pre['corpus']['corpus_sha256'][:16]}, "
	f"{pre['corpus']['n_windows'] - pre['corpus']['n_windows_val']:,} fit windows; bs {pre['config']['bs']}, lr {pre['config']['lr']}, {pre['epochs_run']} epochs run, best epoch {pre['best_epoch']} "
	f"(held-out-route All-MPJPE {pre['best_corpus_val_all']:.2f} px), stop = {pre['stop_reason']}; {pre['ms_per_step_mean']:.1f} ms/step, {pre['gpu_minutes']:.1f} GPU-min.",
	"", "\| epoch \| ms/step \| samples/s \| train loss \| corpus held-out All \| corpus held-out HM \| AIDE val zero-shot All \| AIDE val zero-shot HM \|", "\|---\|---\|---\|---\|---\|---\|---\|---\|"]
	for h in hist:
	L.append(f"\| {h['epoch']} \| {h['ms_per_step']:.1f} \| {h['samples_per_s']:.0f} \| {h.get('train_total', float('nan')):.3f} \| {h['corpus_val_all']:.2f} \| {h['corpus_val_hm']:.2f} \| {h['aide_val_zeroshot_all']:.3f} \| {h['aide_val_zeroshot_hm']:.2f} \|")
	zs = pre.get("zeroshot")
	if zs:
	L += ["", f"Zero-shot run `{zs['name']}` (no AIDE training, test once): val All {zs['val'].get('mpjpe_all', float('nan')):.3f} / HM {zs['val'].get('mpjpe_hm', float('nan')):.2f}; "
	f"test All {zs['test'].get('mpjpe_all', float('nan')):.3f} / HM {zs['test'].get('mpjpe_hm', float('nan')):.2f} / PCK@0.05 {zs['test'].get('pck05', float('nan')):.2f} / HM h=5 {zs['test'].get('hm_h5', float('nan')):.2f}; "
	f"Δ vs zero-velocity All {zs['test'].get('diff_vs_ref_all', float('nan')):+.2f}, Δ vs cwm3d_kp_s42_v2 All {zs['test'].get('diff_vs_kp_all', float('nan')):+.2f} (AIDE ZV test All 52.93 / HM 139.23)."]
	L.append("")
	pre2 = run_json(PRETRAIN_REP)
	if pre2:
	L += [f"Replication pretraining `{PRETRAIN_REP}` (independent seed): {pre2['epochs_run']} epochs, best epoch {pre2['best_epoch']} (held-out All {pre2['best_corpus_val_all']:.2f} px), "
	f"{pre2['ms_per_step_mean']:.1f} ms/step, {pre2['gpu_minutes']:.1f} GPU-min; zero-shot test All {(pre2.get('zeroshot') or {}).get('test', {}).get('mpjpe_all', float('nan')):.3f}.", ""]
	results = {}
	sec = 3
	for title, a_fmt, b_fmt, status_kind in PAIRS:
	L += [f"## {sec}. {title}: fine-tuned from the BATON-pretrained checkpoint (A) vs seed init (B)", ""]
	sec += 1
	per_seed = {}
	for k in seeds:
	pa, pb = load_pred(a_fmt.format(k=k), d), load_pred(b_fmt.format(k=k), d)
	if pa is None or pb is None:
	continue
	ca, cb = per_clip_cells(pa, gt, hm), per_clip_cells(pb, gt, hm)
	per_seed[k] = {c: {"a": float(ca[c][0].mean()), "b": float(cb[c][0].mean()), **M.paired_diff(ca[c][0], cb[c][0], a.n_boot, 0), "higher_better": ca[c][1]} for c in ca}
	if not per_seed:
	L += [f"_(registered, no finished run pairs yet — arm status: {status_kind}; runs `{a_fmt.format(k='<k>')}` vs `{b_fmt.format(k='<k>')}`, seeds {seeds}; "
	"verdict rule as in the header: BEATS / WORSE = Δ CI excl. 0 on ≥ 2/3 seeds with the same sign on every seed)_", ""]
	results[title] = {"status": "REGISTERED_NOT_RUN", "kind": status_kind, "a": a_fmt, "b": b_fmt}
	continue
	L += ["\| seed \| A run \| A best epoch / epochs / GPU min \| init (tensors loaded / kept at seed init) \| B run \| B best epoch / epochs / GPU min \|", "\|---\|---\|---\|---\|---\|---\|"]
	for k in per_seed:
	ra, rb = run_json(a_fmt.format(k=k)), run_json(b_fmt.format(k=k))
	ini = (ra or {}).get("init") or {}
	L.append(f"\| {k} \| {a_fmt.format(k=k)} \| {(ra or {}).get('best_epoch', '—')} / {(ra or {}).get('epochs_run', '—')} / {(ra or {}).get('gpu_minutes', '—')} \| "
	f"{ini.get('parts', '—')}: {ini.get('n_loaded', '—')} / {ini.get('n_missing_or_kept', '—')} (ckpt epoch {ini.get('ckpt_epoch', '—')}) \| {b_fmt.format(k=k)} \| "
	f"{(rb or {}).get('best_epoch', '—')} / {(rb or {}).get('epochs_run', '—')} / {(rb or {}).get('gpu_minutes', '—')} \|")
	f1rows = []
	for k in per_seed:
	ma, mb = (RUNS / a_fmt.format(k=k) / "metrics_test.json"), (RUNS / b_fmt.format(k=k) / "metrics_test.json")
	if ma.exists() and mb.exists():
	xa, xb = json.load(open(ma))["metrics"], json.load(open(mb))["metrics"]
	if "f1_dbr" in xa and "f1_dbr" in xb:
	f1rows.append(f"\| {k} \| " + " / ".join(f"{xa[f'f1_{t}']['value']:.1f}" for t in P.TASKS) + " \| " + " / ".join(f"{xb[f'f1_{t}']['value']:.1f}" for t in P.TASKS) + " \|")
	if f1rows:
	L += ["", "Semantic macro-F1 DBR / DER / TCR / VCR of the compared (geometry-selected) checkpoints, descriptive (single evaluation, no CI here):", "",
	"\| seed \| A \| B \|", "\|---\|---\|---\|"] + f1rows
	L += ["", "\| cell \| " + " \| ".join(f"A s{k}" for k in per_seed) + " \| " + " \| ".join(f"B s{k}" for k in per_seed) + " \| " + " \| ".join(f"Δ s{k} [CI]" for k in per_seed) + " \| mean Δ ± sd \| verdict \|",
	"\|---\|" + "---\|" * (3 * len(per_seed)) + "---\|---\|"]
	cells = list(next(iter(per_seed.values())))
	res = {}
	for c in cells:
	rows = [per_seed[k][c] for k in per_seed]
	v = verdict(rows, rows[0]["higher_better"])
	dm = np.array([r["mean"] for r in rows])
	L.append(f"\| {c} \| " + " \| ".join(f"{r['a']:.2f}" for r in rows) + " \| " + " \| ".join(f"{r['b']:.2f}" for r in rows) + " \| "
	+ " \| ".join(f"{r['mean']:+.2f} [{r['ci_lo']:+.2f},{r['ci_hi']:+.2f}]" for r in rows) + f" \| {dm.mean():+.2f} ± {dm.std(ddof=1) if len(dm) > 1 else float('nan'):.2f} \| {v} \|")
	res[c] = {"verdict": v, "per_seed": {str(k): per_seed[k][c] for k in per_seed}}
	v_all, v_h5 = res["All-MPJPE"]["verdict"], res["HM h=5"]["verdict"]
	adopt = v_all == "BEATS"
	if status_kind == "post hoc":
	rule = (f"Δ All-MPJPE = {v_all}: meets the numeric criterion, but this arm was added AFTER the pre-registered arms' result → exploratory, NOT adopted on its own; "
	f"adoption requires the pre-registered replication arm (independent pretraining seed) to be BEATS" if adopt else f"Δ All-MPJPE = {v_all} → not adopted (exploratory arm, criterion not met)")
	elif status_kind == "replication":
	rule = f"Δ All-MPJPE = {v_all} → {'REPLICATED: adopt the trunk-only pretrained init (both the post hoc arm and this pre-registered replication are BEATS)' if adopt else 'NOT replicated → the trunk-only result stays exploratory; nothing adopted'}"
	else:
	rule = f"Δ All-MPJPE = {v_all} → {'ADOPT (pretraining improves All on ≥ 2/3 seeds)' if adopt else 'do NOT adopt (pre-registered rule not met)'}"
	L += ["", f"Kill rule ({title.split(' — ')[0]}, {status_kind}): {rule}; HM h=5 Δ = {v_h5} (mean {np.mean([res['HM h=5']['per_seed'][str(k)]['mean'] for k in per_seed]):+.2f} px).", ""]
	results[title] = {"status": "DONE", "kind": status_kind, "n_seeds": len(per_seed), "verdict_all": v_all, "verdict_hm_h5": v_h5, "meets_rule": adopt, "cells": res}
	L += ["## Provenance", "", "- Pretraining corpus: BATON train-split, train-only-device routes (no val / test device, no val / test route); AIDE side: TRAIN-split statistics only (box target, score quantiles).",
	"- Fine-tuning: `aide/train/train_kp.py --init-ckpt … --init-parts all` with the Stage-A / B2 recipes verbatim (val All-MPJPE selection, EMA, patience 40; test evaluated once per run).",
	"- No AIDE test clip entered any training or selection; the zero-shot run uses no AIDE data at all before its single test evaluation.", ""]
	Path(a.out).parent.mkdir(parents=True, exist_ok=True)
	Path(a.out).write_text("\n".join(L))
	Path(a.out).with_suffix(".json").write_text(json.dumps({"stamp": stamp, "seeds": seeds, "results": results}, indent=1, default=str))
	write_run_json(P.RESULTS / "tables", "report_baton_pretrain", [Path(__file__), M.__file__, P.__file__], [P.CACHE_NPZ], seed=0, t_start=t0, input_listing=False,
	extra={"out": a.out, "seeds": seeds, "summary": {k: {kk: vv for kk, vv in v.items() if kk != "cells"} for k, v in results.items()}, "dms_fields_used": []})
	print(f"written {a.out}")
	for k, v in results.items():
	print(f" {k}: {v}" if v.get("status") != "DONE" else f" {k[:60]}: All {v['verdict_all']}, HM h=5 {v['verdict_hm_h5']}, meets_rule={v['meets_rule']} [{v['kind']}] ({v['n_seeds']} seeds)")


	if __name__ == "__main__":
	main()