Buckets:

AsianPlayer
/

TriDrive

Files
xet
AsianPlayer/TriDrive / code /aide /eval /report_pretrain_variants.py
AsianPlayer
16 days ago
download
raw
37.6 kB
"""cwm3/aide/eval/report_pretrain_variants.py — pre-registered pretraining experiments (docs/decisions.md 2026-09-15 rows "PRE-REGISTRATION —
P1 CockpitWM pretraining-schedule variants + P1-ext + P3 …" and "PRE-REGISTRATION — P2 Driver-WM architecture on CockpitWM tokens + BATON
pretraining") -> aide/results/tables/pretrain_variants.md (+ .json).
Reads ONLY run outputs (pred_test.npz — paired per-clip bootstraps recomputed with aide/eval/metrics.paired_diff, 10k resamples, seed 0, the same
procedure as baton_pretrain.md / arch_attribution.md; metrics_test.json / metrics_val.json — levels + clip-bootstrap CIs written once per run by
the harness; run json — best epoch / epochs / stop / GPU minutes / init report / config; history.json — per-epoch pretraining curves).
Nothing is typed by hand.  Helpers (paired, verdict, delta_table, …) are imported from aide/eval/report_arch_attribution.py.
Blocks
  §0 selection log     — pre-registered recipe-selection rule: lowest mean AIDE VAL All-MPJPE (run json val.mpjpe_all of the EMA-selected
                         checkpoint) over the 5 fine-tuned seeds, never test; tie (< 0.01 px) -> fewer pretraining GPU-min.  --select writes
                         aide/results/pretrain_variants/{selection.json, SELECTED (best two of the five new variants, for lane F), P3_TAG
                         (matched tag for MotionBERT / ST-GCN: conv if the best convergence-schedule arm is conv or conveuc, else convscale)}
                         once every variant has its 5 fine-tunes.
  §1 P1 pretraining    — per-variant run line + per-epoch held-out curve (ms/step, GPU-min, best epoch, stop reason).
  §2 P1 fine-tunes     — per variant: per-seed levels (val All, test levels with CIs), paired Δ vs the adopted arm cwm3d_kp_batonpre_trunkp123_s<k>_v2
                         (pre-registered reading: BEATS iff CI excl. 0 on >= 4/5 seeds, same sign) and vs the seed init cwm3d_kp_s<k>_v2.
  §2f P1-ext (f)       — lean-read-out CockpitWM (cwm3d_v2 --no-v2-spatial) with the best two variants, vs its trunkp123 arm and its seed init.
  §3 P2                — Driver-WM architecture on CockpitWM tokens: pretraining curve (exterior stream on zero input — deviation), fine-tunes,
                         (i) Δ vs driverwm_kp_s<k>_v2, (ii) Δ vs cwm3d_kp_batonpre_trunkp123_s<k>_v2 (3 seeds, >= 2/3 rule).
  §4 P3                — matched MotionBERT / ST-GCN architectures under the identical converged recipe: Δ = CockpitWM(best) − arch, 5 seeds,
                         >= 4/5 rule; Δ vs each architecture's seed init; the owner's inclusion rules quoted verbatim.
  §5 outcome, §6 GPU minutes, provenance.
  python aide/eval/report_pretrain_variants.py [--n-boot 10000] [--select] [--out aide/results/tables/pretrain_variants.md]
Output policy: if --out already exists the new content is APPENDED as a dated block (append-only tables); existing sections are never edited.
"""
from __future__ import annotations

import argparse
import json
import math
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 load_split  # noqa: E402
from aide.eval.report_arch_attribution import (LEVEL_COLS, MAIN_CELLS, SEEDS3, SEEDS5, TOL_NOT_WORSE, delta_table, fmt_ci, level_row,  # noqa: E402
                                               metrics_test, ms, paired, run_json, run_meta)
from common.runmeta import write_run_json  # noqa: E402

RUNS = P.RESULTS / "runs"
PV = P.RESULTS / "pretrain_variants"
OUT = P.RESULTS / "tables" / "pretrain_variants.md"
REF_PRE = "cwm3d_kp_batonpre_pretrain_s123"                      # adopted 6-epoch arm (reference)
ADOPTED = "cwm3d_kp_batonpre_trunkp123_s{k}_v2"
SEED_INIT_K = "cwm3d_kp_s{k}_v2"
VARIANTS = {   # tag -> (definition, fixed-budget?)  — the five new CockpitWM pretraining variants (pre-registration row 2026-09-15T03:34:33Z)
    "euc": ("6 epochs, Euclidean only (--w-nll 0 --w-bone 0 --w-jerk 0)", True),
    "scale": ("6 epochs, all four loss terms, motion-scale augmentation --motion-scale-max 2.5", True),
    "conv": ("convergence schedule (<= 36 epochs, patience 4 on held-out All), all four loss terms", False),
    "conveuc": ("convergence schedule, Euclidean only", False),
    "convscale": ("convergence schedule, all four loss terms, --motion-scale-max 2.5", False),
}
CONV_ARMS = ("conv", "conveuc", "convscale")
OWNER_DECISION = ('OWNER DECISION (docs/decisions.md, recorded 2026-09-15T13:33:33Z, before any P3 result was read): "The matched converged-pretraining comparison (CockpitWM vs '
                  'MotionBERT architecture vs ST-GCN architecture under the identical selected recipe, 5 seeds) enters Supplementary Material D regardless of direction; the main text '
                  'mentions it only if CockpitWM is significantly better than both (>= 4/5 seeds, same sign) on All-MPJPE. The earlier 6-epoch matched-pretraining runs are superseded by this comparison."')
INCLUSION_RULES = [
    'owner (first message): "the owner intends to report the matched-pretraining comparison in the paper only if CockpitWM is significantly better; '
    'the coordinator recommended reporting it in Supp. D regardless; internal tables record everything regardless"',
    'owner refinement: "owner: main text only if CockpitWM is significantly better than both architectures under the identical converged recipe; '
    "coordinator's default: the comparison enters Supp. D regardless of direction unless the owner rules otherwise before the results are read (decision pending)\"",
]


# ----------------------------------------------------------------------------- io helpers
def val_all(name):
    r = run_json(name) or {}
    v = (r.get("val") or {}).get("mpjpe_all")
    return None if v is None else float(v)


def finished(name):
    return metrics_test(name) is not None and run_json(name) is not None and (RUNS / name / "pred_test.npz").exists()


def pretrain_block(pre, model_label, extra_note=""):
    """-> (lines, summary) for one pretraining run (run json + history)."""
    r = run_json(pre)
    if r is None:
        return [f"_(pretraining run `{pre}` not finished — no numbers)_", ""], {"status": "MISSING"}
    c, hist = r.get("config", {}), r.get("history", [])
    L = [f"- `{pre}` ({model_label}; `aide/train/pretrain_cwm3d_baton.py`, corpus sha256 {r['corpus']['corpus_sha256'][:16]}, "
         f"{r['corpus']['n_windows'] - r['corpus']['n_windows_val']:,} fit / {r['corpus']['n_windows_val']:,} held-out windows): {r['n_params']:,} params, bs {c.get('bs')}, lr {c.get('lr')}, "
         f"EMA {c.get('ema')}, epochs cap {c.get('epochs')}, patience {c.get('patience')}, loss weights euc {c.get('w_euc')} / nll {c.get('w_nll')} / bone {c.get('w_bone')} / jerk {c.get('w_jerk')} "
         f"(active terms {r.get('loss_terms_active')}), motion-scale max {r.get('motion_scale_max', 1.0)}; **{r['epochs_run']} epochs run, best epoch {r['best_epoch']} "
         f"(held-out-route All-MPJPE {r['best_corpus_val_all']:.2f} px), stop = {r['stop_reason']}; {r['ms_per_step_mean']:.1f} ms/step, {r['gpu_minutes']:.1f} GPU-min**.{extra_note}",
         "", "| 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} | "
                 f"{h['aide_val_zeroshot_all']:.3f} | {h['aide_val_zeroshot_hm']:.2f} |")
    L.append("")
    return L, {"status": "DONE", "best_epoch": r["best_epoch"], "epochs_run": r["epochs_run"], "stop": r["stop_reason"], "best_heldout_all": r["best_corpus_val_all"],
               "ms_per_step": r["ms_per_step_mean"], "gpu_minutes": r["gpu_minutes"], "n_params": r["n_params"], "loss_terms": r.get("loss_terms_active"),
               "motion_scale_max": r.get("motion_scale_max", 1.0), "heldout_curve": [h["corpus_val_all"] for h in hist], "config": {k: c.get(k) for k in ("epochs", "patience", "bs", "lr", "w_nll", "w_bone", "w_jerk", "motion_scale_max")}}


def finetune_rows(names, seeds):
    """per-seed level table for fine-tune runs: init report, val All, test levels with CIs."""
    L = ["| seed | run | best epoch / epochs / stop / GPU min | init (loaded / kept; ckpt epoch) | val All | " + " | ".join(f"{lab} [CI]" for _, lab, _ in LEVEL_COLS) + " |",
         "|---|---|---|---|---|" + "---|" * len(LEVEL_COLS)]
    for k in seeds:
        n = names[k]
        if not finished(n):
            L.append(f"| {k} | {n} | (not finished) | | | " + " | ".join("—" for _ in LEVEL_COLS) + " |")
            continue
        r = run_json(n) or {}
        ini = r.get("init") or {}
        L.append(f"| {k} | {n} | {r.get('best_epoch', '—')} / {r.get('epochs_run', '—')} / {r.get('stop_reason', '—')} / {r.get('gpu_minutes', '—')} | "
                 f"{ini.get('parts', '—')}: {ini.get('n_loaded', '—')} / {ini.get('n_missing_or_kept', '—')} (ep {ini.get('ckpt_epoch', '—')}) | {val_all(n):.3f} | {level_row(n)} |")
    return L


def paired_block(a_names, b_names, seeds, d, gt, hm, n_boot, label_a, label_b, n_req, full=True):
    """paired Δ (A − B) over the seeds where both runs are finished -> (lines, res, n_seeds)."""
    per_seed = {k: paired(a_names[k], b_names[k], d, gt, hm, n_boot, full=full) for k in seeds if finished(a_names[k]) and finished(b_names[k])}
    per_seed = {k: v for k, v in per_seed.items() if v}
    if not per_seed:
        return [f"_(no finished pairs {label_a} vs {label_b} yet)_", ""], None, 0
    rows, res = delta_table(per_seed, label_a, label_b, n_req=n_req)
    note = "" if len(per_seed) == len(seeds) else f"  **PARTIAL: {len(per_seed)}/{len(seeds)} seeds finished ({list(per_seed)}); verdict threshold stays CI excl. 0 on >= {n_req} seeds.**"
    return [f"Δ = {label_a} − {label_b} per seed (paired per-clip, 10k clip bootstrap; verdict: CI excl. 0 on >= {n_req}/{len(seeds)} seeds, same sign on all).{note}", ""] + rows + [""], res, len(per_seed)


def cells_summary(res):
    return {c: {kk: vv for kk, vv in r.items() if kk != "per_seed"} for c, r in res.items()} if res else None


def gpu_of(name):
    r = run_json(name)
    return None if r is None or r.get("gpu_minutes") is None else float(r["gpu_minutes"])


# ----------------------------------------------------------------------------- selection (pre-registered val rule)
def selection_table(variant_ft, seeds):
    """variant -> {mean_val, sd_val, n, mean_test (display only), gpu_pre}; complete = every variant has all seeds finished."""
    rows = {}
    for tag, names in variant_ft.items():
        vals = [val_all(names[k]) for k in seeds if finished(names[k])]
        tests = [metrics_test(names[k])["mpjpe_all"]["value"] for k in seeds if finished(names[k])]
        rows[tag] = {"n": len(vals), "mean_val": float(np.mean(vals)) if vals else None, "sd_val": float(np.std(vals, ddof=1)) if len(vals) > 1 else None,
                     "mean_test_display_only": float(np.mean(tests)) if tests else None, "gpu_pre": gpu_of(f"cwm3d_kp_batonpre_pretrain_{tag}_s123")}
    return rows


def pick_best(rows, eligible):
    cand = [(rows[t]["mean_val"], rows[t]["gpu_pre"] if rows[t]["gpu_pre"] is not None else float("inf"), t) for t in eligible if rows.get(t, {}).get("mean_val") is not None]
    if not cand:
        return None
    cand.sort(key=lambda z: (round(z[0], 2), z[1]))          # tie (< 0.01 px) -> fewer pretraining GPU-min
    return [t for _, _, t in cand]


# ----------------------------------------------------------------------------- main
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--n-boot", type=int, default=10000)
    ap.add_argument("--out", default=str(OUT))
    ap.add_argument("--select", action="store_true", help="write aide/results/pretrain_variants/{selection.json,SELECTED,P3_TAG} when every variant has its 5 fine-tunes")
    a = ap.parse_args()
    t0 = time.time()
    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())
    gpu, summary = {}, {}

    def note_gpu(*names):
        for n in names:
            g = gpu_of(n)
            if g is not None:
                gpu[n] = g

    L = ["# Pre-registered pretraining experiments P1 / P1-ext / P2 / P3 — CockpitWM pretraining variants, Driver-WM architecture on CockpitWM tokens, matched MotionBERT / ST-GCN under the identical converged recipe — AIDE test, n = 609, paired per-clip bootstrap (10k)", "",
         f"Generated {stamp} by `aide/eval/report_pretrain_variants.py` from pred_test.npz / metrics_test.json / metrics_val.json / run json / history.json files only (locked protocol: driver `first`, grid `floor`, "
         f"HM top {int(hm.sum())} (threshold {thr:.2f} px), PCK @ {P.pck_threshold_px(0.05):.0f} px; every fine-tune selected on val All-MPJPE only (EMA weights), every pretraining selected on the BATON corpus' held-out routes only; "
         "test evaluated once per run by the harness).  Δ = paired per-clip difference (negative = better for errors, positive = better for PCK), CI = 10k clip-bootstrap percentile interval.  "
         f"**Verdict rule: BEATS / WORSE = Δ CI excl. 0 on the pre-registered number of seeds (5-seed designs: >= 4/5; 3-seed design: >= 2/3) with the same sign on every seed; NOT_WORSE = all CIs cover 0 and |Δ| < {TOL_NOT_WORSE}; else INCONCLUSIVE.**  "
         "Pre-registration rows: docs/decisions.md 2026-09-15 (P1 / P1-ext / P3 at 2026-09-15T03:34:33Z; P2 in its own row before its launch).  Runner `aide/train/run_pretrain_variants.sh` (lanes K / MB / ST / P2 / F; pids in `aide/results/pretrain_variants/pids.txt`).", "",
         "Common pretraining setting: corpus `aide/results/cache/baton_pose_corpus.npz` (BATON TRAIN-split, train-only-device routes; 642,179 windows, 22 held-out routes), seed 123, bs 256, lr 3e-4, EMA 0.999, `--zeroshot 0`; "
         f"reference = the adopted 6-epoch arm `{REF_PRE}` (all four loss terms) whose trunk-only fine-tunes `{ADOPTED}` are the paper's \"+ BATON pretraining\" row.  Convergence schedule = `--epochs 36 --patience 4` (held-out All-MPJPE, EMA weights; 1 warm-up epoch + cosine over the 36-epoch cap, cut by the stop).  "
         "Variant (e) dense corpus was SKIPPED before any run: `build_baton_pose_corpus.py` has no anchor-step flag (`ANCHOR_STRIDE` is a module constant).", "",
         "**" + OWNER_DECISION + "**  Earlier inclusion statements (recorded verbatim in the pre-registration row): " + " — ".join(INCLUSION_RULES) + ".", ""]

    # ------------------------------------------------------------------ §0 selection log
    variant_ft = {tag: {k: f"cwm3d_kp_batonpre_{tag}_s{k}_v2" for k in SEEDS5} for tag in VARIANTS}
    ref_ft = {k: ADOPTED.format(k=k) for k in SEEDS5}
    sel_rows = selection_table(variant_ft, SEEDS5)
    ref_vals = [val_all(ref_ft[k]) for k in SEEDS5 if finished(ref_ft[k])]
    complete = all(sel_rows[t]["n"] == len(SEEDS5) for t in VARIANTS)
    order_all = pick_best(sel_rows, list(VARIANTS)) if complete else None
    order_conv = pick_best(sel_rows, CONV_ARMS) if complete else None
    p3_tag = None if order_conv is None else ("convscale" if order_conv[0] == "convscale" else "conv")
    L += ["## 0. Recipe-selection log (pre-registered rule: lowest mean AIDE **val** All-MPJPE over the 5 fine-tuned seeds — never test; tie < 0.01 px → fewer pretraining GPU-min)", "",
          "| variant | definition | fixed budget? | eligible for P3 (fully trained) | n fine-tunes finished | mean val All ± sd | mean test All (display only, NOT used) | pretraining GPU-min | rank (all five) |", "|---|---|---|---|---|---|---|---|---|"]
    L.append(f"| reference (adopted) | 6 epochs, all four loss terms (`{REF_PRE}`) | yes | reference | {len(ref_vals)} | {ms(ref_vals, 3)} | {ms([metrics_test(ref_ft[k])['mpjpe_all']['value'] for k in SEEDS5 if finished(ref_ft[k])], 2)} | {gpu_of(REF_PRE)} | — |")
    for tag, (defn, fixed) in VARIANTS.items():
        s = sel_rows[tag]
        rank = "—" if order_all is None else str(order_all.index(tag) + 1)
        mv = "—" if s["mean_val"] is None else (f"{s['mean_val']:.3f} ± {s['sd_val']:.3f}" if s["sd_val"] is not None else f"{s['mean_val']:.3f}")
        mt = "—" if s["mean_test_display_only"] is None else f"{s['mean_test_display_only']:.2f}"
        L.append(f"| {tag} | {defn} | {'yes' if fixed else 'no (convergence)'} | {'yes' if tag in CONV_ARMS else 'no'} | {s['n']} | {mv} | {mt} | {s['gpu_pre']} | {rank} |")
    if complete:
        L += ["", f"**Selection ({stamp}): best of the five new variants by val = `{order_all[0]}` (order {order_all}); best TWO → lane F (lean-read-out CockpitWM) = {order_all[:2]}; "
              f"best convergence-schedule arm (P3 eligible) = `{order_conv[0]}` (order {order_conv}) → matched MotionBERT / ST-GCN tag `{p3_tag}` "
              f"({'conv ≡ conveuc for tensor-output architectures' if p3_tag == 'conv' else 'motion-scale augmentation selected → second matched pretraining with --motion-scale-max 2.5'}).**", ""]
        if a.select:
            PV.mkdir(parents=True, exist_ok=True)
            sel = {"stamp": stamp, "rule": "lowest mean AIDE val All-MPJPE over 5 fine-tuned seeds; tie < 0.01 px -> fewer pretraining GPU-min", "rows": sel_rows,
                   "order_all_five": order_all, "best_two_for_F": order_all[:2], "order_convergence_arms": order_conv, "best_conv_arm": order_conv[0], "p3_tag": p3_tag}
            for fn, txt in (("SELECTED", "\n".join(order_all[:2]) + "\n"), ("P3_TAG", p3_tag + "\n"), ("BEST_CONV", order_conv[0] + "\n")):
                if (PV / fn).exists():
                    L.append(f"_(`{PV / fn}` already exists — not overwritten: {(PV / fn).read_text().split()} )_")
                else:
                    (PV / fn).write_text(txt)
                    L.append(f"_(wrote `{PV / fn}`: {txt.split()})_")
            (PV / f"selection_{stamp}.json").write_text(json.dumps(sel, indent=1, default=str))
            L.append("")
    else:
        L += ["", f"_(selection pending: {[t for t in VARIANTS if sel_rows[t]['n'] < len(SEEDS5)]} do not have all 5 fine-tunes yet)_", ""]
    summary["selection"] = {"complete": complete, "rows": sel_rows, "order_all": order_all, "order_conv": order_conv, "p3_tag": p3_tag, "reference_val": ref_vals}

    # ------------------------------------------------------------------ §1 pretraining curves
    L += ["## 1. P1 pretraining runs (CockpitWM = `cwm3d_kp`, seed 123) — per-epoch held-out-route curves", ""]
    pre_sum = {}
    for tag, label in [("ref", f"reference `{REF_PRE}`")] + [(t, f"variant `{t}` — {VARIANTS[t][0]}") for t in VARIANTS]:
        pre = REF_PRE if tag == "ref" else f"cwm3d_kp_batonpre_pretrain_{tag}_s123"
        L.append(f"### 1.{list(['ref'] + list(VARIANTS)).index(tag) + 1} {label}")
        L.append("")
        rows, s = pretrain_block(pre, "CockpitWM, cwm3d d 256 / 4 layers")
        L += rows
        pre_sum[tag] = s
        if tag != "ref":
            note_gpu(pre)
    summary["p1_pretraining"] = pre_sum

    # ------------------------------------------------------------------ §2 P1 fine-tunes
    L += ["## 2. P1 fine-tunes (trunk-only init, Stage-A recipe verbatim, 5 seeds) — paired Δ vs the adopted arm and vs the seed init", ""]
    p1 = {}
    overview = ["| variant | n | test All mean ± sd | Δ All vs adopted (mean ± sd) | verdict vs adopted (>= 4/5) | Δ All vs seed init | verdict vs seed init | HM vs adopted | PCK@0.05 vs adopted | HM h=5 vs adopted |", "|---|---|---|---|---|---|---|---|---|---|"]
    for tag in VARIANTS:
        names = variant_ft[tag]
        L += [f"### 2.{list(VARIANTS).index(tag) + 1} variant `{tag}` — {VARIANTS[tag][0]}", ""]
        L += finetune_rows(names, SEEDS5) + [""]
        for k in SEEDS5:
            note_gpu(names[k])
        rows_a, res_a, n_a = paired_block(names, ref_ft, SEEDS5, d, gt, hm, a.n_boot, f"{tag}", "adopted", 4)
        L += ["#### vs the adopted arm `cwm3d_kp_batonpre_trunkp123_s<k>_v2`", ""] + rows_a
        rows_b, res_b, n_b = paired_block(names, {k: SEED_INIT_K.format(k=k) for k in SEEDS5}, SEEDS5, d, gt, hm, a.n_boot, f"{tag}", "seed-init", 4, full=False)
        L += ["#### vs the seed init `cwm3d_kp_s<k>_v2`", ""] + rows_b
        tests = [metrics_test(names[k])["mpjpe_all"]["value"] for k in SEEDS5 if finished(names[k])]
        if res_a:
            ra = res_a["All-MPJPE"]
            L += [f"**Pre-registered P1 reading for `{tag}`: Δ All vs the adopted arm = {ra['verdict']} (mean {ra['mean']:+.2f} ± {ra['sd'] if ra['sd'] is None else round(ra['sd'], 2)} px over {n_a} seeds) → "
                  f"{'the paper row MAY be replaced by this variant (BEATS)' if ra['verdict'] == 'BEATS' and n_a == len(SEEDS5) else 'internal table only (not BEATS on 5 seeds)'}; "
                  f"HM {res_a['HM-MPJPE']['verdict']} ({res_a['HM-MPJPE']['mean']:+.2f}), PCK@0.05 {res_a['PCK@0.05']['verdict']} ({res_a['PCK@0.05']['mean']:+.2f}), HM h=5 {res_a['HM h=5']['verdict']} ({res_a['HM h=5']['mean']:+.2f}); "
                  f"vs seed init: {res_b['All-MPJPE']['verdict'] if res_b else '—'} ({res_b['All-MPJPE']['mean']:+.2f} px).**" if res_b else "", ""]
            overview.append(f"| {tag} | {n_a} | {ms(tests, 2)} | {ra['mean']:+.2f} ± {0 if ra['sd'] is None else ra['sd']:.2f} | {ra['verdict']} | "
                            f"{res_b['All-MPJPE']['mean']:+.2f} ± {0 if res_b['All-MPJPE']['sd'] is None else res_b['All-MPJPE']['sd']:.2f} | {res_b['All-MPJPE']['verdict']} | "
                            f"{res_a['HM-MPJPE']['verdict']} ({res_a['HM-MPJPE']['mean']:+.2f}) | {res_a['PCK@0.05']['verdict']} ({res_a['PCK@0.05']['mean']:+.2f}) | {res_a['HM h=5']['verdict']} ({res_a['HM h=5']['mean']:+.2f}) |")
        else:
            overview.append(f"| {tag} | {n_a} | {ms(tests, 2)} | — | (pending) | — | — | — | — | — |")
        p1[tag] = {"n_seeds": n_a, "test_all": tests, "vs_adopted": cells_summary(res_a), "vs_seed_init": cells_summary(res_b),
                   "per_seed_all_vs_adopted": None if not res_a else res_a["All-MPJPE"]["per_seed"]}
    ref_tests = [metrics_test(ref_ft[k])["mpjpe_all"]["value"] for k in SEEDS5 if finished(ref_ft[k])]
    L += ["### 2.6 Overview (test All mean ± sd over seeds; adopted arm = " + ms(ref_tests, 2) + f", seed init = {ms([metrics_test(SEED_INIT_K.format(k=k))['mpjpe_all']['value'] for k in SEEDS5 if finished(SEED_INIT_K.format(k=k))], 2)})", ""] + overview + [""]
    summary["p1"] = p1

    # ------------------------------------------------------------------ §2f P1-ext (f) lean-read-out CockpitWM
    L += ["## 2f. P1-ext (f): lean-read-out CockpitWM (`cwm3d_v2 --no-v2-spatial`) with the best two variants — vs its trunkp123 arm `cwm3d_v2_nospatial_batonpre_trunkp123_s<k>_v2` and vs its seed init `cwm3d_v2_nospatial_s<k>_v2`", ""]
    f_tags = [t for t in VARIANTS if run_json(f"cwm3d_v2_nospatial_batonpre_pretrain_{t}_s123") is not None]
    pf = {}
    if not f_tags:
        L += ["_(no lean-read-out variant pretraining finished yet — lane F waits for the selection)_", ""]
    for tag in f_tags:
        pre = f"cwm3d_v2_nospatial_batonpre_pretrain_{tag}_s123"
        names = {k: f"cwm3d_v2_nospatial_batonpre_{tag}_s{k}_v2" for k in SEEDS5}
        L += [f"### 2f.{f_tags.index(tag) + 1} variant `{tag}`", ""]
        rows, s = pretrain_block(pre, "lean-read-out CockpitWM, cwm3d_v2 --no-v2-spatial (hazard BCE weight as the reference)")
        L += rows
        note_gpu(pre, *[names[k] for k in SEEDS5])
        L += finetune_rows(names, SEEDS5) + [""]
        rows_a, res_a, n_a = paired_block(names, {k: f"cwm3d_v2_nospatial_batonpre_trunkp123_s{k}_v2" for k in SEEDS5}, SEEDS5, d, gt, hm, a.n_boot, f"v2-{tag}", "v2-trunkp123", 4)
        L += ["#### vs `cwm3d_v2_nospatial_batonpre_trunkp123_s<k>_v2`", ""] + rows_a
        rows_b, res_b, n_b = paired_block(names, {k: f"cwm3d_v2_nospatial_s{k}_v2" for k in SEEDS5}, SEEDS5, d, gt, hm, a.n_boot, f"v2-{tag}", "v2-seed-init", 4, full=False)
        L += ["#### vs `cwm3d_v2_nospatial_s<k>_v2`", ""] + rows_b
        if res_a:
            L += [f"**Reading (f) for `{tag}`: Δ All vs the v2 trunkp123 arm = {res_a['All-MPJPE']['verdict']} ({res_a['All-MPJPE']['mean']:+.2f} px, {n_a} seeds); vs the v2 seed init = {res_b['All-MPJPE']['verdict'] if res_b else '—'} ({res_b['All-MPJPE']['mean']:+.2f} px).**", ""]
        pf[tag] = {"pretraining": s, "n_seeds": n_a, "vs_trunkp123": cells_summary(res_a), "vs_seed_init": cells_summary(res_b)}
    summary["p1f"] = pf

    # ------------------------------------------------------------------ §3 P2 Driver-WM architecture on CockpitWM tokens
    L += ["## 3. P2 — Driver-WM architecture on CockpitWM tokens (`driverwm_kp`) + BATON pretraining", ""]
    pre2 = "driverwm_kp_batonpre_pretrain_s123"
    r2 = run_json(pre2)
    note2 = "" if r2 is None else f"  **Deviation (recorded in the run json): exterior stream = {r2.get('exterior_stream')}; trunk prefixes loaded by the fine-tunes: {r2.get('trunk_prefixes')}.**"
    rows, s2 = pretrain_block(pre2, "Driver-WM architecture on CockpitWM tokens, 85.7 M params (tensor output → Euclidean term)", note2)
    L += rows
    note_gpu(pre2)
    dwm_pre = {k: f"driverwm_kp_batonpre_trunkp123_s{k}_v2" for k in SEEDS3}
    dwm_seed = {k: f"driverwm_kp_s{k}_v2" for k in SEEDS3}
    L += ["### 3.1 Fine-tunes (trunk-only init, shared recipe verbatim, `--loss euc`, seeds 42 / 123 / 456) and the comparison runs", ""]
    L += finetune_rows(dwm_pre, SEEDS3) + [""]
    note_gpu(*dwm_pre.values())
    L += ["Comparison runs (existing, not retrained): " + "; ".join(f"`{dwm_seed[k]}` test All {metrics_test(dwm_seed[k])['mpjpe_all']['value']:.2f}" for k in SEEDS3 if finished(dwm_seed[k]))
          + "; " + "; ".join(f"`{ADOPTED.format(k=k)}` {metrics_test(ADOPTED.format(k=k))['mpjpe_all']['value']:.2f}" for k in SEEDS3 if finished(ADOPTED.format(k=k))) + ".", ""]
    rows_i, res_i, n_i = paired_block(dwm_pre, dwm_seed, SEEDS3, d, gt, hm, a.n_boot, "DWM-pre", "DWM-seed", 2)
    L += ["### 3.2 (i) Δ = Driver-WM-kp pretrained − its own seed init `driverwm_kp_s<k>_v2` (does BATON pretraining help the Driver-WM architecture?)", ""] + rows_i
    rows_ii, res_ii, n_ii = paired_block(dwm_pre, {k: ADOPTED.format(k=k) for k in SEEDS3}, SEEDS3, d, gt, hm, a.n_boot, "DWM-pre", "K-pre", 2)
    L += ["### 3.3 (ii) Δ = Driver-WM-kp pretrained − CockpitWM pretrained (`cwm3d_kp_batonpre_trunkp123_s<k>_v2`; the architecture comparison with BATON pretraining on both sides)", ""] + rows_ii
    rows_iii, res_iii, _ = paired_block(dwm_pre, {k: SEED_INIT_K.format(k=k) for k in SEEDS3}, SEEDS3, d, gt, hm, a.n_boot, "DWM-pre", "K-seed", 2, full=False)
    L += ["### 3.4 descriptive: Δ = Driver-WM-kp pretrained − CockpitWM seed init `cwm3d_kp_s<k>_v2`", ""] + rows_iii
    p2 = {"pretraining": s2, "n_seeds": n_ii, "i_vs_own_seed_init": cells_summary(res_i), "ii_vs_cockpitwm_pretrained": cells_summary(res_ii), "iii_vs_cockpitwm_seed_init": cells_summary(res_iii)}
    if res_i and res_ii:
        vi, vii = res_i["All-MPJPE"], res_ii["All-MPJPE"]
        read_i = {"BEATS": "BATON pretraining HELPS the Driver-WM architecture (Δ < 0, CI excl. 0 on >= 2/3 seeds)", "WORSE": "BATON pretraining HURTS the Driver-WM architecture",
                  "NOT_WORSE": "no effect (all CIs cover 0, |Δ| < 0.25)", "INCONCLUSIVE": "inconclusive"}[vi["verdict"]]
        read_ii = {"BEATS": "the Driver-WM architecture is significantly BETTER than CockpitWM under pretraining on both sides", "WORSE": "CockpitWM is significantly BETTER than the Driver-WM architecture under pretraining on both sides",
                   "NOT_WORSE": "no significant difference (all CIs cover 0, |Δ| < 0.25)", "INCONCLUSIVE": "inconclusive / mixed"}[vii["verdict"]]
        L += [f"**P2 readings ({n_ii} seeds): (i) Δ All vs own seed init = " + "; ".join(f"s{k} {fmt_ci(vi['per_seed'][str(k)])}" for k in vi["per_seed"]) + f" → {vi['verdict']}: {read_i}.  "
              f"(ii) Δ All vs CockpitWM pretrained = " + "; ".join(f"s{k} {fmt_ci(vii['per_seed'][str(k)])}" for k in vii["per_seed"]) + f" → {vii['verdict']}: {read_ii} (reported whatever the direction).  "
              f"HM (ii) {res_ii['HM-MPJPE']['verdict']} ({res_ii['HM-MPJPE']['mean']:+.2f}), PCK@0.05 (ii) {res_ii['PCK@0.05']['verdict']} ({res_ii['PCK@0.05']['mean']:+.2f}).**", ""]
        p2.update(reading_i=read_i, reading_ii=read_ii, verdict_i=vi["verdict"], verdict_ii=vii["verdict"])
    summary["p2"] = p2

    # ------------------------------------------------------------------ §4 P3 matched architectures
    L += ["## 4. P3 — matched MotionBERT / ST-GCN architectures pretrained with the identical converged recipe (fully trained on both sides) vs the selected CockpitWM arm", ""]
    L += ["**" + OWNER_DECISION + "**", "", "Earlier inclusion statements (recorded verbatim in the pre-registration row): " + " — ".join(INCLUSION_RULES) + ".", ""]
    p3 = {}
    best_tag = None if order_conv is None else order_conv[0]
    tags_try = [p3_tag] if p3_tag else ["conv", "convscale"]
    for arch, label in (("motionbert_kp", "MotionBERT architecture"), ("stgcn_kp", "ST-GCN architecture")):
        L += [f"### 4.{1 if arch == 'motionbert_kp' else 2} {label} (`{arch}`)", ""]
        done_tags = [t for t in tags_try if run_json(f"{arch}_batonpre_pretrain_{t}_s123") is not None]
        if not done_tags:
            L += [f"_(matched pretraining `{arch}_batonpre_pretrain_<{'/'.join(tags_try)}>_s123` not finished)_", ""]
            p3[arch] = {"status": "PRETRAIN_MISSING"}
            continue
        for t in done_tags:
            pre = f"{arch}_batonpre_pretrain_{t}_s123"
            rows, s = pretrain_block(pre, f"{label}; tensor output → Euclidean term (conv ≡ conveuc for this architecture)")
            L += rows
            note_gpu(pre)
            names = {k: f"{arch}_batonpre_{t}_s{k}_v2" for k in SEEDS5}
            L += finetune_rows(names, SEEDS5) + [""]
            note_gpu(*names.values())
            k_best = {k: f"cwm3d_kp_batonpre_{best_tag}_s{k}_v2" for k in SEEDS5} if best_tag else None
            entry = {"tag": t, "pretraining": s, "matched_cockpitwm_arm": best_tag}
            if k_best is None:
                L += ["_(the CockpitWM recipe selection is pending → the matched paired comparison is not computed yet; the fine-tune levels above are final)_", ""]
            else:
                rows_m, res_m, n_m = paired_block(k_best, names, SEEDS5, d, gt, hm, a.n_boot, f"K-{best_tag}", f"{arch}-{t}", 4)
                L += [f"#### Δ = CockpitWM `cwm3d_kp_batonpre_{best_tag}_s<k>_v2` − `{arch}_batonpre_{t}_s<k>_v2` (matched, both fully trained; BEATS = CockpitWM significantly better on >= 4/5 seeds)", ""] + rows_m
                if res_m:
                    v = res_m["All-MPJPE"]
                    read = {"BEATS": "CockpitWM is significantly better than this architecture under the identical converged recipe", "WORSE": "this architecture is significantly better than CockpitWM under the identical converged recipe (reported as is)",
                            "NOT_WORSE": "no significant difference", "INCONCLUSIVE": "inconclusive / mixed"}[v["verdict"]]
                    L += [f"**P3 reading ({label}, {n_m} seeds): Δ All = " + "; ".join(f"s{k} {fmt_ci(v['per_seed'][str(k)])}" for k in v["per_seed"]) + f"; mean {v['mean']:+.2f} ± {0 if v['sd'] is None else v['sd']:.2f} → {v['verdict']}: {read}.  "
                          f"HM {res_m['HM-MPJPE']['verdict']} ({res_m['HM-MPJPE']['mean']:+.2f}), PCK@0.05 {res_m['PCK@0.05']['verdict']} ({res_m['PCK@0.05']['mean']:+.2f}), HM h=5 {res_m['HM h=5']['verdict']} ({res_m['HM h=5']['mean']:+.2f}).**", ""]
                    entry.update(n_seeds=n_m, vs_cockpitwm=cells_summary(res_m), verdict_all=v["verdict"], reading=read)
            seed_init = {k: f"{arch}_s{k}_v2" for k in SEEDS5}
            rows_s, res_s, n_s = paired_block(names, seed_init, SEEDS5, d, gt, hm, a.n_boot, f"{arch}-{t}", f"{arch}-seed", 4, full=False)
            L += [f"#### Δ = `{arch}_batonpre_{t}_s<k>_v2` − its seed init `{arch}_s<k>_v2` (where the seed-init run exists)", ""] + rows_s
            entry["vs_own_seed_init"] = cells_summary(res_s)
            p3[arch] = entry
    summary["p3"] = p3
    both_beat = all(p3.get(arch, {}).get("verdict_all") == "BEATS" for arch in ("motionbert_kp", "stgcn_kp"))
    both_done = all("verdict_all" in p3.get(arch, {}) for arch in ("motionbert_kp", "stgcn_kp"))

    # ------------------------------------------------------------------ §5 outcome
    L += ["## 5. Outcome (decisions-row wording)", ""]
    for tag in VARIANTS:
        r = p1[tag]["vs_adopted"]
        L.append(f"- P1 `{tag}`: " + ("pending" if r is None else f"Δ All vs adopted {r['All-MPJPE']['verdict']} ({r['All-MPJPE']['mean']:+.2f} ± {0 if r['All-MPJPE']['sd'] is None else r['All-MPJPE']['sd']:.2f} px, {p1[tag]['n_seeds']} seeds)"
                                          + f"; vs seed init {p1[tag]['vs_seed_init']['All-MPJPE']['verdict']} ({p1[tag]['vs_seed_init']['All-MPJPE']['mean']:+.2f})"))
    beaters = [t for t in VARIANTS if p1[t]["vs_adopted"] and p1[t]["vs_adopted"]["All-MPJPE"]["verdict"] == "BEATS" and p1[t]["n_seeds"] == len(SEEDS5)]
    L.append(f"- **P1 paper rule: the \"+ BATON pretraining\" row is replaced ONLY by a variant that BEATS the adopted arm on 5 seeds → {('candidates: ' + str(beaters)) if beaters else 'NO variant BEATS the adopted arm → the paper row stays; all variants internal-table only'}.**")
    L.append(f"- Selection (val rule): {'pending' if not complete else f'best of five = {order_all[0]}; best two for (f) = {order_all[:2]}; best converged arm = {order_conv[0]} → P3 tag {p3_tag}'}.")
    L.append(f"- P2: {p2.get('reading_i', 'pending')}; {p2.get('reading_ii', 'pending')}.")
    for arch in ("motionbert_kp", "stgcn_kp"):
        L.append(f"- P3 {arch}: {p3.get(arch, {}).get('reading', p3.get(arch, {}).get('status', 'pending'))}.")
    L.append(f"- **P3 inclusion (owner decision): Supp. D regardless of direction; main text {'YES — CockpitWM significantly better than BOTH architectures (>= 4/5 seeds, same sign)' if both_done and both_beat else ('NO — CockpitWM is not significantly better than both' if both_done else 'pending')}; the earlier 6-epoch matched-pretraining runs are superseded.**")
    summary["outcome"] = {"p1_beaters": beaters, "p3_both_beats": bool(both_beat and both_done), "p3_done": both_done}

    # ------------------------------------------------------------------ §6 GPU minutes + provenance
    L += ["", "## 6. Runs and GPU minutes (new runs only, from the run json files; existing runs were not retrained)", "", "| run | GPU min |", "|---|---|"]
    for name, g in gpu.items():
        L.append(f"| {name} | {g:.2f} |")
    L += [f"| **total (new runs)** | **{sum(gpu.values()):.1f}** |", "",
          "## Provenance", "",
          "- New run dirs only: `cwm3d_kp_batonpre_pretrain_<tag>_s123`, `cwm3d_kp_batonpre_<tag>_s<k>_v2`, `cwm3d_v2_nospatial_batonpre_pretrain_<tag>_s123`, `cwm3d_v2_nospatial_batonpre_<tag>_s<k>_v2`, "
          "`driverwm_kp_batonpre_pretrain_s123`, `driverwm_kp_batonpre_trunkp123_s<k>_v2`, `<arch>_batonpre_pretrain_<tag>_s123`, `<arch>_batonpre_<tag>_s<k>_v2`; comparison runs are the existing ones "
          "(`cwm3d_kp_batonpre_trunkp123_s*_v2`, `cwm3d_kp_s*_v2`, `driverwm_kp_s*_v2`, `cwm3d_v2_nospatial*_v2`, `motionbert_kp_s*_v2`, `stgcn_kp_s*_v2`).",
          "- Pretraining corpus: BATON train-split, train-only-device routes (no val / test device or route); AIDE side: TRAIN statistics only. No AIDE test clip entered any training or selection; the recipe selection uses AIDE val only.",
          "- Launch gate before every job (free VRAM >= need + 4 GB, MemAvailable >= 28 GB, < 3 trainers alive, 45 s settle, OMP/MKL 4): `aide/train/run_pretrain_variants.sh`; pids in `aide/results/pretrain_variants/pids.txt`.", ""]

    out = Path(a.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    text = "\n".join(L)
    if out.exists():
        with open(out, "a") as f:
            f.write(f"\n\n---\n\n<!-- appended regeneration {stamp} -->\n\n" + text)
        mode = "appended"
    else:
        out.write_text(text)
        mode = "written"
    js = out.with_suffix(".json")
    rec = {"stamp": stamp, "n_boot": a.n_boot, "summary": summary, "gpu_minutes": gpu}
    old = json.load(open(js)) if js.exists() else []
    old = old if isinstance(old, list) else [old]
    js.write_text(json.dumps(old + [rec], indent=1, default=str))
    write_run_json(P.RESULTS / "tables", "report_pretrain_variants" + ("" if out.resolve() == OUT.resolve() else "_dryrun"), [Path(__file__), M.__file__, P.__file__], [P.CACHE_NPZ], seed=0, t_start=t0, input_listing=False,
                   extra={"out": str(out), "mode": mode, "selection": summary["selection"], "outcome": summary["outcome"], "gpu_minutes": gpu, "dms_fields_used": []})
    print(text)
    print(f"\n[{mode}] {out}")


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

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