dobri420/vjepa2-cac

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

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.