V-JEPA history lifts world action model's real-robot OOD success from 27.5% to 70%

CSWAM: Better Causal Semantic Representations for Out-of-Distribution Generalization in World Action Models

Tianbin Liu, Jian Zhu, Taiyi Su, Jianjun Zhang, Chong Ma, Zitai Huang, Weiyi Lu, Yi Xu

cs.CV, cs.AI

2026-09-16

CSWAM adds a frozen V-JEPA 2.1 semantic expert and observation history to FastWAM, lifting RoboTwin OOD transfer from 10.16% to 45.18% and real-robot average success from 27.5% to 70%.

What problem this solves

World action models (WAMs) train video prediction and action generation jointly, so that predicting future frames supervises the policy. The FastWAM branch goes further and cuts future-video generation entirely at deployment, denoising actions only, which makes inference cheap. The weakness shows up under visual distribution shifts: change the background or the object appearance and FastWAM's success rate collapses.

Both causes are about representation. The video stream runs on VAE latents optimized for reconstruction, which preserve appearance detail, so the visual context shifts when the background changes even if the task state has not. And the action expert sees only the current frame, with no history that would reveal motion, state transitions, or task progress. Collecting demonstrations for every background and object appearance is not feasible, so the gap matters for deployment.

Prior fixes add pretrained semantic encoders such as DINO, but framewise encoders do not model how a scene evolves. This paper uses V-JEPA 2.1, a video-native joint-embedding predictive model that predicts in feature space instead of reconstructing pixels, and pairs it with a causal observation history.

Method

CSWAM keeps the FastWAM video and action streams and adds a third one, a causal semantic expert:

• A frozen V-JEPA 2.1 encoder maps a sparse observation history into features. Frames are sampled at a stride (8 frames covering a longer window), and each pair of adjacent frames is encoded into one tubelet, so short-range transitions are already inside the features before any trainable module sees them.
• The semantic expert does not use V-JEPA features as static conditioning. Trained in the same flow-matching framework as the other streams, it extrapolates future V-JEPA features from the clean history, with the loss applied only to future tokens. This forces the pathway to model how semantic state evolves.
• The three streams form a Mixture-of-Transformers: each expert keeps its own projections and feedforward layers, and masked attention moves information across streams. Video tokens may attend to the semantic history; the action expert attends to current video tokens plus the semantic history, and to no future target. That mask is where the causal guarantee lives.
• Training sums three flow-matching losses, with the video and semantic streams sharing a noise level and the action stream sampling its own.
• At deployment the future tokens are dropped. Current video state and semantic history are encoded once and reused across denoising steps; only the action expert runs iteratively. FastWAM's action-only efficiency is unchanged.

The design choice worth spelling out: what gets added is observation-side representation and history, not future imagination, which would bring back the inference cost FastWAM removed.

Results

Simulation uses 50 bimanual tasks from RoboTwin 2.0, 100 rollouts per task. Under the standard mixed-training protocol (2,500 Clean + 25,000 Randomized demonstrations):

CSWAM tops the table, 2.6 points over FastWAM. The real gap opens under Clean2Random: fine-tune only on the 2,500 Clean demonstrations, evaluate on Randomized scenes, no adaptation.

Two readings stand out. The 5,000-hour embodied pretraining adds only 8.34 OOD points to FastWAM but 35.82 to CSWAM; the same pretrained weights yield far more transfer when the architecture can use them. Also, CSWAM's gain over FastWAM is 7.78 points on Clean versus 35.02 on Randomized, concentrated on the unseen distribution, which the authors take as evidence against mere better fitting of the fine-tuning data.

On a real dual-arm platform (two ARX manipulators), two long-horizon tasks (Sort, 12 hours of demonstrations; Fold, 4 hours) under three compounding shift levels, 20 trials each:

Averaged over 120 rollouts per model: 27.5% versus 70.0%. Restricted to the four Hard and Extreme settings, 10% versus 55%. Successful Sort rollouts also finish faster, 44.2 seconds on average for CSWAM against 64.8 for FastWAM.

Ablations separate the two legs: removing observation history costs 1.7/1.8 points (Clean/Randomized); removing the video expert costs 6.1/8.1. Semantic prediction complements generative video supervision rather than replacing it. In the encoder comparison, framewise DINOv3 ViT-L, roughly 3.5x the parameters of V-JEPA ViT-B, trails V-JEPA 2.1 by 1.4/1.5 points.

Why it matters

Action-only WAMs are among the most inference-efficient policies for manipulation right now, and this paper fixes their OOD weakness without adding inference cost: all new computation sits on the observation side and is cached once. For teams deploying robots outside the training environment, this is a directly usable direction. The method is assembled from existing parts (V-JEPA, MoT, flow-matching); the contribution is the combination and an evaluation protocol that separates mixed-training performance from OOD transfer. One caveat for practitioners: the 45.18% OOD number depends on 5,000 hours of embodied pretraining. Without it, OOD success is 9.36%.

Limitations

The authors acknowledge that Extreme Fold, with environment, lighting, tabletop, objects, and containers all changed, reaches only 10%, and that completion-time averages are success-conditioned, which skews the 39.8-second Extreme Sort figure. Beyond that: the real-robot study covers two tasks at 20 trials per setting, so the granularity is one 5% step. The only real-robot baseline is FastWAM; ST-WAM and Faster-WAM, which target the same robustness problem, appear only in the mixed-training table with no OOD comparison. A 0.4-point lead under the Full protocol is thin at 100 rollouts per task. The semantic expert is warm-started from ActionDiT weights, so component attribution rests on ablations rather than independent training runs.

Terms

• WAM (world action model): a robot policy trained jointly with future-video prediction, which supervises action generation.
• V-JEPA: Meta's self-supervised video model that predicts masked content in feature space instead of reconstructing pixels, making it less sensitive to appearance.
• Flow matching: a diffusion-style training objective that learns a velocity field from noise to data; standard for action generation.
• Mixture-of-Transformers (MoT): a sparse multi-modal architecture where each modality keeps its own parameters and exchanges information through attention.
• Tubelet: a video-ViT practice that encodes a small spatiotemporal block spanning adjacent frames into one token, carrying short-range motion.
• Embodied pretraining: pretraining on large robot manipulation video-action corpora before task fine-tuning.

Source

• arXiv 2609.18462v3
• PDF
• Hugging Face

What people are saying

• CSWAM paper adds a V-JEPA 2.1 causal semantic expert to boost world action model OOD generalization zhenjun_zhao

Related papers

• Discrete codes first: DeepMind scales autoregressive image models to coherent 256×256 generation
• BrickWright: a LEGO SDK and validator take agents from 9–69% to 100% physical validity
• Best of 22 VLMs hits 58.5% on AgentVQA as visual-state errors dominate
• LoRA on 33B MiniMax-H3 cuts missed entity-count changes from 34.7% to 21.3%
• Qwen3-VL-8B training lifts RewardData artifact Finding F1 from 3.38 to 35.42
• RTAformer gains 4.8 SGDET R@10 on transition frames over matched OED

All paper explainers