Scuttie/dreamer-vla-best-libero-vjepa2-aux

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.

Per-tier interpretation

*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 (= "testood" cell): 200 episodes from `data/splitstaskheld/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