README.md

---
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)

![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)

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.