Best LIBERO action-probe backbone — V-JEPA 2 ViT-L + aux finetune
This is the single best representation extractor identified by the
dreamer-vla cross-architecture sweep on LIBERO. The full sweep evaluated
82 variants × 8 backbone families × 2 conditions (frozen / +aux) on
the LIBERO 130-task split with strict task-OOD held-out evaluation.
V-JEPA 2 ViT-L (304M, pretrained on ~22M internet videos by Facebook
AI) finetuned on LIBERO with an inverse-dynamics auxiliary head reaches
the highest action-probe R² in the sweep.
Headline number
LIBERO task-OOD action-probe R² (H=1, mean of 3 probe seeds): 0.845
(probe protocol: episode-disjoint, 400 train + 200 OOD-task clips, MLP
[in→256→128→7] trained 3000 steps with AdamW lr=3e-4)
Full cross-architecture table
8 backbone families, frozen vs +aux finetune on LIBERO. All numbers are
test_ood action-probe R² unless noted.
| Family | Backbone | Params | Frozen | + aux | Lift |
|---|
| video masked latent (V-JEPA) | V-JEPA 2 ViT-L (this ckpt) | 304M | +0.40 | +0.845 ⭐ | +0.45 |
| video masked latent (distilled) | V-JEPA 2.1 ViT-B | 87M | +0.44 | +0.82 | +0.38 |
| action-aware quantization | LAPA (Open-X LAQ) | 344M | +0.41 | +0.51 | +0.10 |
| pixel diffusion (ours) | DIFF + aux λ=0.05 (Stage A best) | 91M | — | +0.57 | — |
| image-text contrastive | SigLIP 2 ViT-L | 316M | +0.05 | +0.17 | +0.12 |
| image contrastive masked | Web-DINO ViT-L | 304M | −0.01 | +0.16 | +0.17 |
| recon-aligned conv (Cosmos-1) | Cosmos-1 CI16x16 (NVIDIA) | 34M | −0.36 | −0.29 | +0.07 |
| recon-aligned conv (SDXL VAE) | SDXL VAE encoder | 34M | −0.55 | −0.41 | +0.14 |
| shortcut forcing dynamics | Dreamer 4 (8 variants tried) | 64M-276M | −0.04 | −0.04 | 0.00 |
Per-tier interpretation
| Tier | R² range | What's in it |
|---|
| 1 | +0.82-0.85 | Video-temporal masked latent prediction (V-JEPA family) |
| 2 | +0.51-0.57 | Action-aware / pixel-temporal (LAPA, DIFF) |
| 3 | +0.16-0.17 | Image-only semantic pretraining (Web-DINO, SigLIP 2) |
| 4 | −0.41 to −0.04 | Reconstruction / one-step world model (SDXL VAE, Dreamer 4) |
The +0.68 R² gap between tier 1 (V-JEPA, video-temporal) and tier 3
(image-only Web-DINO/SigLIP, matched 304-316M params) isolates
video-temporal masked latent prediction as the dominant lever —
not "large-scale semantic pretraining" generically, not "masked
prediction" generically, not "self-supervision" generically.
Reproduction stack (verified clean)
Cross-architecture probe protocol — same across all 82 rows:
- Dataset: LIBERO 130 tasks, task-OOD split (27 train / 7 held-out OOD)
- probe-train: 400 episodes from
data/splits_taskheld/train.json [0:400]
- probe-eval-train (= "train" cell): episodes [400:600] of same split
- probe-eval-test (= "test_ood" cell): 200 episodes from
data/splits_taskheld/test_ood.json
- Probe: 2-layer MLP
[D → 256 → 128 → 7], AdamW lr=3e-4 wd=1e-4,
3000 steps batch 512, 3 seeds, report mean
- No leakage: probe-train episodes disjoint from probe-eval-train
episodes; both task-level held-out from probe-eval-test
- Features: mean-pool spatial tokens → (B, T, D) per-frame
Aux finetune recipe for this ckpt:
- 30000 steps on 8× B200, batch 4 per GPU = effective batch 32
- AdamW lr=5e-5, cosine schedule with 2000-step warmup, wd=0.05
- Loss =
1.0 * MSE(InvDynAuxHead(features), action[:7])
- aux head:
[2D → 512 → 256 → 7] GELU
- Encoder unfrozen (full finetune), bfloat16 autocast
- 30k steps × eff-batch 32 = ~1M action-frame pairs seen
How to load
1import sys, torch
2sys.path.insert(0, "external_models/vjepa2/vjepa2") # clone from facebookresearch/vjepa2
3from src.models import vision_transformer as vit_encoder
4
5state = torch.load("ckpt_last.pt", weights_only=False)
6# state has keys: "step", "encoder" (vision_transformer state_dict),
7# "aux_head" (training-only, ignored for probe), "args"
8
9encoder = vit_encoder.vit_large(
10 patch_size=16, img_size=(256, 256), num_frames=32, tubelet_size=2,
11 use_sdpa=True, use_silu=False, wide_silu=True,
12 uniform_power=False, use_rope=True,
13)
14encoder.load_state_dict(state["encoder"])
15encoder.eval().cuda()
16
17# extract features from a (B, T, C, H, W) clip in [0, 1]:
18# 1. resize to (256, 256), normalize per V-JEPA mean/std
19# 2. forward encoder → patch tokens
20# 3. mean-pool spatial tokens → (B, T, 1024) feature
The aux_head in the ckpt is training-only and is NOT used at probe
time — discard it.
See
src/comparison/extractors/vjepa2_extractor.py in
k1seul/dreamer-vla for a
working PyTorch wrapper.
Sanity bars (paper-claim verification)
- Action-probe R² ≥ +0.7 on LIBERO task-OOD ✓ (achieved +0.845)
- Frozen pretraining alone gives R² ≥ +0.30 ✓ (achieved +0.40 frozen)
- +aux lift ≥ +0.30 over frozen ✓ (achieved +0.45 lift)
- Per-dim breakdown: translation > rotation > gripper ✓ (consistent)
Caveats
- Probe-only evaluation; does not include closed-loop policy rollouts
in simulator. Companion work (Nilaksh et al. 2026) does closed-loop;
we deliberately focus on representation quality.
- LIBERO 130-task; results may not transfer 1:1 to other manipulation
benchmarks (MetaWorld, CALVIN, Robomimic results are in separate
Stage B repos — see Companion repos below).
- Encoder finetuning is destructive; this ckpt no longer matches V-JEPA
2 ViT-L upstream features. For pure V-JEPA 2 features use the
Facebook AI upstream ckpt.
- aux λ=0.05 is the LIBERO sweet spot from our Stage A DIFF sweep; we
used the same λ for all +aux variants without per-backbone sweeping.
Results may improve slightly with per-backbone λ tuning.
Companion repos (related dreamer-vla artifacts)
- Scuttie/dreamer-vla-dreamer4-libero
— Dreamer V4 reproduction (negative result, R² ≈ -0.04 across all variants)
- Scuttie/dreamer-vla-lapa-aux-libero
— LAPA + aux finetune (R² +0.51)
- Scuttie/dreamer-vla-extra-backbones-libero
— Web-DINO, SigLIP 2, SDXL VAE + aux finetune (R² +0.16 / +0.17 / -0.41)
Citation
1@article{vjepa2_2025,
2 title={V-JEPA 2: Self-Supervised Video Models},
3 author={Bardes, Adrien and Garrido, Quentin and ...},
4 journal={Meta AI Research}, year={2025}
5}
(dreamer-vla paper citation will be added when on arXiv.)
File listing
| File | Description | Size |
|---|
ckpt_last.pt | encoder + aux_head state dict + training args | ~1.2 GB |
config.json | training hyperparameters + reproduction metadata | small |
cross_arch_table.csv | the full 8-backbone comparison table (.csv) | small |
cross_arch_table.json | same data, JSON form | small |
