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

import sys, torch
sys.path.insert(0, "external_models/vjepa2/vjepa2")  # clone from facebookresearch/vjepa2
from src.models import vision_transformer as vit_encoder

state = torch.load("ckpt_last.pt", weights_only=False)
# state has keys: "step", "encoder" (vision_transformer state_dict),
#                 "aux_head" (training-only, ignored for probe), "args"

encoder = vit_encoder.vit_large(
    patch_size=16, img_size=(256, 256), num_frames=32, tubelet_size=2,
    use_sdpa=True, use_silu=False, wide_silu=True,
    uniform_power=False, use_rope=True,
)
encoder.load_state_dict(state["encoder"])
encoder.eval().cuda()

# extract features from a (B, T, C, H, W) clip in [0, 1]:
# 1. resize to (256, 256), normalize per V-JEPA mean/std
# 2. forward encoder → patch tokens
# 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

@article{vjepa2_2025,
  title={V-JEPA 2: Self-Supervised Video Models},
  author={Bardes, Adrien and Garrido, Quentin and ...},
  journal={Meta AI Research}, year={2025}
}

(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

Downloads last month: 2

Video Preview

Robotics