Update to iter-16300 (default) + iter-6200 best-val; new planning eval, training chart

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+cac-combined-16k.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -22,45 +22,71 @@ Meta's exact `src.models.ac_predictor.vit_ac_predictor` (from
 layers, 1024 dim, 16 heads, ~305M params, RoPE, frame-causal block attention,
 action + state token conditioning (7-dim each). Identical to the shipped AC.
 
-## Training (this snapshot)
-
-- **5,800 iterations** (toward Meta's 94,500-iter recipe; a training snapshot,
-  not the finished run), effective batch 256, LR 7.5e-5 → 4.25e-4 → 0 (WSD,
-  scheduled over the full 94.5k horizon), AdamW wd 0.04, bf16, `torch.compile`.
-- Loss: L1 teacher-forcing + 2-step autoregressive rollout (faithful port of
-  Meta's `app/vjepa_droid/train.py`).
-- Data: ~51,000 DROID left-cam episodes (~58 h footage — essentially the full
-  DROID 1.0.1 raw corpus), 8 frames @ 4 fps, 256 px; first 240 held out for eval.
-- Hardware: 1× H100 80GB to iter 2,800, then resumed 2,800 → 5,800 on **2× H100
-  DDP** (5.82 s/iter, 1.98× single-GPU — near-linear over NVLink). Resume was
-  seamless: no loss spike across the boundary, the WSD LR schedule continued on
-  the 94.5k horizon, and val cosine ticked 0.818 → 0.823.
+## Checkpoints
+
+**Two** checkpoints are published from one continued run toward Meta's
+94,500-iter recipe (still a snapshot, not the finished run):
+
+| file | iter | held-out val cosine | note |
+|---|---|---|---|
+| `cac_fp32.safetensors` | **16,300** | 0.818 | default — the most-trained checkpoint |
+| `cac_bestval_fp32.safetensors` | 6,200 | **0.824** | peak held-out val cosine of the run |
+
+`--predictor cac` / `download_cac()` resolve to `cac_fp32.safetensors` (the
+iter-16,300 model). Why both are published: see the val-cosine decline below.
+
+## Training
+
+Effective batch 256, LR 7.5e-5 → 4.25e-4 → 0 (WSD, scheduled over the full 94.5k
+horizon), AdamW wd 0.04, bf16, `torch.compile`. Loss: L1 teacher-forcing + 2-step
+autoregressive rollout (faithful port of Meta's `app/vjepa_droid/train.py`).
+Data: ~51,000 DROID left-cam episodes (~58 h footage — essentially the full DROID
+1.0.1 raw corpus), 8 frames @ 4 fps, 256 px; first 240 held out for eval.
+
+Trained across three seamless resume segments: **1×H100 to iter 2,800**, then
+**2×H100 DDP** to 5,800 and on to 16,300 (5.82 s/iter, ~1.98× single-GPU — near
+linear over NVLink). No loss spike across either resume boundary; the WSD LR
+schedule continued on the 94.5k horizon throughout.
+
+### The val-cosine decline (why two checkpoints)
+
+![CAC combined training across 3 resume runs](cac-combined-16k.png)
+
+Training L1 (teacher-forced + rollout) falls smoothly **1.02 → ~0.28** across all
+three runs. But held-out val cosine **peaks ~0.823 around iter 6,000 then slowly
+declines to ~0.818 by iter 16,300**, while val held-out L1 drifts *up* — mild
+overfitting on the encoder-dominated cosine metric. Cosine is a lenient proxy,
+though; planning is the real test, and there the iter-6,200 and iter-16,300
+checkpoints are **statistically indistinguishable** (see below). So the extra
+~10k iters regressed the cheap metric at **no measurable planning cost** — hence
+both checkpoints are shipped, for anyone who wants to probe the trade-off.
 
 ## Evaluation (240 held-out DROID episodes, frozen ViT-g)
 
-Two metrics, vs Meta's shipped V-JEPA 2-AC on the same set:
+Offline CEM planning gate (`plan --random-step`: CEM-planned action vs the true
+logged action `a_t` from a random step), identical protocol for all three:
+
+| model | plan L2 → true ↓ | cos(plan, true) ↑ | beats random | points right |
+|---|---|---|---|---|
+| V-JEPA 2-AC (Meta) | **0.0396** | +0.650 | 93% | 90% |
+| CAC iter 16,300 (this, default) | 0.0418 | +0.644 | 94% | 88% |
+| CAC iter 6,200 (best-val) | 0.0446 | +0.648 | 92% | 87% |
 
-| metric | CAC (this) | V-JEPA 2-AC |
-|---|---|---|
-| prediction cosine | **0.789** | 0.766 |
-| CEM plan L2 → true action | 0.048 | **0.043** |
-| cos(plan, true action) | +0.54 | **+0.66** |
-| CEM beats random | 85% | **94%** |
+- **CAC 16,300 vs 6,200**: paired Wilcoxon *p*=0.065 (L2), *p*=0.40 (cosine) over
+  238 paired episodes — a statistical tie. CEM's own sampling noise (~0.0003 on
+  L2 run-to-run) is a meaningful fraction of the gap.
+- **vs Meta's AC** (trained ~16× longer): AC still leads on plan L2, but the gap
+  is small and CAC matches it on direction cosine.
+- Open-loop **prediction** cosine is ≈0.79 for CAC (mean over horizons) — it
+  plateaus early and, as the chart shows, is *not* a reliable planning proxy.
 
-CAC **beats** AC on prediction cosine (the encoder-dominated metric, where the
-~51k-episode data scale helps). On planning, CAC's plan L2 (0.048) is now **on
-par** with AC (0.043); direction still lags (+0.54 vs +0.66) — the axis bound by
-iter count, and Meta trained ~16× longer. Open-loop L1 prediction now **beats a
-static baseline at every horizon** (the undertrained iter-2,750 snapshot could
-not), CAC points the right direction on 82% of episodes, and its CEM plan beats
-a random action on 85%. Prediction cosine is a lenient metric; planning quality
-is the real test.
+## Files
 
-## File
+- `cac_fp32.safetensors` — predictor state_dict, fp32, 305M params, 300 tensors,
+  iter 16,300 (default). Safetensors metadata: `iter`, `architecture`.
+- `cac_bestval_fp32.safetensors` — same format, iter 6,200 (peak val cosine).
 
-`cac_fp32.safetensors` — predictor state_dict, fp32, 305M params, 300 tensors.
-Safetensors metadata: `iter`, `architecture`. Loads `strict=True` into
-`vit_ac_predictor(...)`.
+Both load `strict=True` into `vit_ac_predictor(...)`.
 
 ## Load
 
@@ -78,6 +104,9 @@ predictor.load_state_dict(sd, strict=True)
 
 ## Caveats
 
+- Planning numbers are from a `--random-step` CEM gate on 240 held-out episodes;
+  they are a custom intermediate probe, not Meta's downstream real-robot success
+  metric.
 - Train/eval split is index-based on DROID's deterministic institution ordering
   (first 240 held out), so the eval set is institution-concentrated, not random.
 - Trained on public DROID data using Meta's published recipe; independent

cac_bestval_fp32.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:94fe9843ec6f911eec7bce9e15214e35c2e67735e00bda5e793b2dee89698f5b
+size 1220915504

cac_fp32.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a9f57670b150d06f53cccd3a98b69435212f4f949a4ddfaeb1729a8a5c716845
+oid sha256:6850759a1e1d08564bef8690118a4c54cfebe1568fea0a2ae6ec9fc91fee8968
 size 1220915504