File size: 4,960 Bytes
d405f43 2bfcb0c d405f43 47e3a8d d405f43 f6ed49d d405f43 f6ed49d 670a6c6 f6ed49d 670a6c6 f6ed49d d405f43 f6ed49d d405f43 f6ed49d d405f43 47e3a8d d405f43 47e3a8d d405f43 f6ed49d d405f43 47e3a8d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | ---
library_name: vjepa2
tags:
- vjepa2
- world-model
- robotics
- droid
- action-conditioned
---
# CAC — from-scratch V-JEPA 2-AC predictor
A from-scratch reproduction of Meta's V-JEPA 2-AC action-conditioned predictor,
trained on H100s. This repo holds only the trained **predictor**; the frozen
V-JEPA 2 ViT-g encoder is not included — use Meta's `vjepa2_ac_vit_giant`
(torch.hub) for the encoder.
## Architecture
Meta's exact `src.models.ac_predictor.vit_ac_predictor` (from
[facebookresearch/vjepa2](https://github.com/facebookresearch/vjepa2)): 24
layers, 1024 dim, 16 heads, ~305M params, RoPE, frame-causal block attention,
action + state token conditioning (7-dim each). Identical to the shipped AC.
## 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)
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)
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% |
- **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.
## Files
- `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).
Both load `strict=True` into `vit_ac_predictor(...)`.
## Load
```python
import torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from src.models.ac_predictor import vit_ac_predictor
predictor = vit_ac_predictor(img_size=(256, 256), patch_size=16, num_frames=64,
tubelet_size=2, embed_dim=1408)
sd = load_file(hf_hub_download("dobri420/vjepa2-cac", "cac_fp32.safetensors"))
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
reproduction, not affiliated with Meta.
|