💻 computer science

JEPA-WAM: Learning Vision-Language-Action Policies with Joint-Embedding World Modeling

JEPA-WAM introduces a latent world action model built on a pretrained V-JEPA space that unifies spatially structured future prediction and continuous action generation through a shared predictor, achieving state-of-the-art performance on robotic manipulation benchmarks without requiring large-scale policy pretraining.

Original authors: Yihan Lin, Jiawei He, Shifeng Bao, Chen Zhao, Yang Li, Xiaobo Wang, Yan Wang, Cheng Chi, Jing Zhang

Published 2026-08-11

📖 7 min read🧠 Deep dive

CC BY 4.0

Original authors: Yihan Lin, Jiawei He, Shifeng Bao, Chen Zhao, Yang Li, Xiaobo Wang, Yan Wang, Cheng Chi, Jing Zhang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a robot to cook. You could just show it a video of a chef chopping onions and say, "Do that." But if the robot only memorizes the exact look of that specific video, it might panic if the chef wears a different hat or if the kitchen lights change. This is the challenge of Vision-Language-Action (VLA) models: they are great at following instructions in familiar settings, but they often stumble when the world looks a little different than what they were trained on.

To fix this, scientists have tried teaching robots to be "predictors." Instead of just reacting, the robot tries to imagine what will happen next. If it can predict that "if I push this cup, it will slide off the table," it can plan better. However, making a robot that generates full, high-definition videos of the future is like asking a student to write a whole novel just to decide what to do next—it takes too much time and computing power. So, researchers looked for a shortcut: a "latent" model that predicts the idea of the future without drawing the whole picture. But here's the catch: many of these shortcuts either lose too much detail or treat the prediction and the action as two separate, disconnected tasks.

This brings us to JEPA-WAM, a new approach that tries to teach robots to "feel" the flow of time without getting bogged down in drawing every single frame. It asks a simple question: Can we teach a robot to understand how the world changes, and use that understanding to move better, all in one smooth motion?

The Robot's "Time-Sense" Shortcut

Meet JEPA-WAM. Think of it as a robot that doesn't just look at a photo of the present and guess the next photo; instead, it learns the relationship between the two.

Imagine you are watching a magic trick. A magician puts a red ball in a box, and then pulls out a blue ball. A standard robot might just memorize "Red Ball → Blue Ball." But JEPA-WAM is like a detective who notices the transition: "The red ball vanished, the box shook, and a blue ball appeared." It learns the story of the change, not just the before and after.

The researchers built this system on top of a pre-trained "vision brain" called V-JEPA. You can think of V-JEPA as a robot that has already watched millions of videos and learned how objects move and interact. JEPA-WAM takes this brain and adds a special "time-travel" module.

Here is the magic trick: Instead of asking the robot to generate a blurry video of the future (which is slow and expensive), JEPA-WAM asks the robot to predict a joint map. Imagine taking a photo of the kitchen now and a photo of the kitchen five seconds later, and stacking them on top of each other. JEPA-WAM doesn't try to draw the future photo from scratch. Instead, it looks at the difference between the two stacked photos and learns to predict exactly how the pixels shifted. It learns that "the cup moved two inches to the left" and "the drawer opened," preserving the fine details of where things happened.

The Shared Brain: One Predictor, Two Jobs

The cleverest part of JEPA-WAM is how it uses its brain. In many old systems, the part that predicts the future and the part that moves the robot's arms were like two different people talking through a wall. One would guess the future, and the other would try to listen.

JEPA-WAM uses a Shared Predictor. Imagine a single, super-smart student who is taking a test.

Task A: The student looks at the current scene and predicts how the world will change (the "Time-Sense").
Task B: The student uses that same understanding to decide how to move the robot's arms (the "Action").

Because the student is doing both tasks at the same time, learning to predict the future directly shapes how the student decides to move. The paper suggests that this tight coupling makes the robot much smarter. It's like a dancer who doesn't just memorize steps but understands the rhythm of the music; the movement flows naturally because the prediction and the action are born from the same understanding.

Does It Actually Work?

The researchers tested this idea in three different worlds: a standard video game simulation, a more chaotic simulation with random objects, and a real-life robot arm in a lab.

In the Simulations:
On a benchmark called LIBERO-Plus, which tests how well robots handle changes in the environment (like different lighting or object positions), JEPA-WAM scored 79.2%. This was the best result among robots that hadn't been pre-trained on massive amounts of robot data. Even more impressive, when they applied this "time-sense" trick to an already powerful robot brain called π0.5, the score jumped from 84.5% to 86.3%. This suggests that even smart robots can get smarter if they learn to predict the flow of time.

In the Real World:
The team took their robot to a real lab with a dual-arm setup (two robot arms working together). They asked it to do tasks like stacking blocks, putting fruit on a plate, or opening a drawer.

When the setup was exactly as expected (In-Distribution), JEPA-WAM scored about 59.8%.
When they messed with the background or moved the objects around (Out-of-Distribution), JEPA-WAM still managed 54.2%.
Compare this to a standard robot (π0), which dropped from 51.8% to a shaky 22.5% when the environment changed.

This shows that JEPA-WAM is much more robust. It doesn't just memorize a specific scene; it learns the logic of how objects move, so it can handle surprises.

What It's Not (and What It Rules Out)

The paper is careful to point out what JEPA-WAM is not.

It is not a video generator. It doesn't try to create a new, high-definition video of the future. The authors argue that trying to generate every pixel is too costly and often unnecessary for controlling a robot.
It is not just about predicting the final state. The researchers tested a version that only looked at the "future" image without the "current" image, and it performed worse (77.3% vs 79.2%). This proves that understanding the relationship between "now" and "later" is more important than just guessing the end result.
It is not a magic bullet for everything. The authors note that if the same visual scene leads to two very different outcomes based on a specific instruction (e.g., "push the cup" vs. "pull the cup"), the model might struggle because it focuses on visual changes rather than language-specific goals.

The Bottom Line

JEPA-WAM suggests that the secret to making robots more adaptable isn't just giving them more data or bigger brains, but teaching them to understand the flow of time in a way that directly informs their actions. By using a shared brain to predict how the world changes and decide how to move, the robot becomes more like a skilled human who can adapt to a messy kitchen without panicking.

The results, measured in simulations and real-world trials, suggest that this "joint prediction" approach is a powerful way to build robots that can handle the unpredictable nature of the real world. While it's not a solved problem for every possible task, it offers a promising new path for making robots that are not just obedient, but truly adaptable.

Technical Summary: JEPA-WAM

Problem Statement

Vision-Language-Action (VLA) policies have demonstrated strong performance in diverse manipulation tasks but often model state transitions only implicitly, limiting their robustness under distribution shifts. While World Action Models (WAMs) address this by explicitly modeling future states, video-generation-based WAMs incur substantial deployment costs due to iterative future prediction. Existing latent WAMs, which avoid explicit frame generation, face two primary limitations:

Target Representation: They often reuse intermediate features from generative models (optimized for frame generation rather than state change) or compress future observations into a small set of latent tokens/subgoals, potentially losing fine-grained spatial structure.
Policy Integration: They either use predicted future representations as additional context (introducing redundancy) or employ separate prediction objectives/dynamics modules that weakly influence the representations used for action generation.

The core challenge is to develop a latent WAM that learns a spatially structured representation of transitions and uses predictive supervision to directly shape the policy backbone without explicit future frame generation.

Methodology: JEPA-WAM

The authors propose JEPA-WAM, a latent WAM built within a pre-trained V-JEPA (Joint-Embedding Predictive Architecture for Video) representation space. The architecture couples latent transition prediction with continuous action generation through a shared predictor.

1. Representation Space and Target Construction

Base Representation: JEPA-WAM utilizes a frozen V-JEPA 2.1 encoder to generate dense, patch-level representations ( $Z_t$ ) for current observations. These representations maintain spatial structure (patch organization) rather than global pooling.
Joint Current–Future Target: Instead of predicting an absolute future state or a simple difference, the model constructs a joint current–future target ( $Y_{t,t+\delta}$ $Y_{t, t + δ}$ ). This is created by stacking current ( $O_t$ $O_{t}$ ) and future ( $O_{t+\delta}$ $O_{t + δ}$ ) observations along the temporal dimension and encoding them jointly.
- This target captures the visual temporal relation between endpoints, identifying stable vs. changing regions and evolving object/spatial relations.
- It preserves dense patch-level correspondence, avoiding the compression of transition information into global tokens.

2. Shared Predictor Architecture

The core innovation is a shared predictor ( $F_\theta$ ), instantiated with a Qwen2.5-0.5B LLM, which serves two simultaneous functions:

Transition Prediction: The predictor processes the current visual representation and task instruction to predict the joint current–future target ( $\hat{Y}_{t,t+\delta}$ ). This branch is supervised via a patch-level cosine distance loss ( $L_{wm}$ ).
Action Generation: The same predictor produces dedicated representations ( $C_t$ ) from specific "action placeholder" tokens. These representations condition a Diffusion Transformer (DiT) action expert to generate continuous action chunks via flow matching.

Key Design Choice: By sharing the predictor, the transition supervision directly optimizes the backbone from which action-conditioned representations are extracted. At deployment, the target branch and prediction head are removed, retaining only the action generation pathway.

3. Transfer to Pretrained Policies

The authors demonstrate that this transition supervision can be instantiated in existing pretrained VLA policies (e.g., $\pi0.5$ ) without modifying their original perception or action pathways. This is achieved by introducing learnable "future tokens" whose hidden states are aligned with the joint current–future target, providing auxiliary transition supervision to the shared backbone.

Key Contributions

JEPA-WAM Architecture: A latent WAM built in a pre-trained V-JEPA space that couples latent transition modeling with continuous action generation via a shared predictor, eliminating the need for explicit future frame generation.
Spatially Structured Joint Target: The construction of a joint current–future target that preserves dense patch-level information, capturing task-shared visual temporal structure and localized changes without compressing the transition into global tokens.
Empirical Validation: Demonstration of strong Out-of-Distribution (OOD) generalization across simulation (LIBERO-Plus, RoboTwin 2.0) and real-world bimanual manipulation, showing that transition supervision benefits both training from scratch and fine-tuning pretrained policies.

Experimental Results

The paper evaluates JEPA-WAM on several benchmarks, comparing it against methods with and without large-scale robot-policy pretraining.

LIBERO-Plus (OOD): JEPA-WAM achieves a 79.2% average success rate, the best result among methods without large-scale robot-policy pretraining. When the transition supervision is applied to the pretrained $\pi0.5$ , performance improves from 84.5% to 86.3%, achieving the best overall result.
RoboTwin 2.0: JEPA-WAM achieves 79.9% on Clean settings and 36.9% on Random (OOD) settings, outperforming other non-pretrained methods and matching the Random performance of pretrained $\pi0.5$ . The $\pi0.5$ instantiation with JEPA supervision reaches 84.6% (Clean) and 37.5% (Random).
Real-World Manipulation: On a bimanual AgileX platform, JEPA-WAM maintains strong performance under visual and spatial shifts (54.2% OOD vs. 22.5% for baseline $\pi0$ ). The $\pi0.5$ instantiation with JEPA supervision achieves 84.7% OOD performance.
Ablation Studies: Controlled experiments confirm that:
- The V-JEPA representation space itself provides a strong foundation for OOD learning.
- The joint current–future target outperforms future-only or endpoint-difference targets.
- Preserving patch-level spatial structure is critical (outperforming convolutional alignment).
- Direct supervision of the shared backbone is more effective than intermediate signals or conditioning on full hidden states.

Significance and Claims

The paper claims that JEPA-WAM successfully bridges the gap between explicit state-transition modeling and efficient deployment. By leveraging the V-JEPA representation space and a shared predictor, the method captures robust visual temporal structures (localized changes, object relations) that generalize across visual and spatial shifts.

The authors emphasize that this approach avoids the computational overhead of video generation while providing stronger inductive biases for state transitions than implicit VLA training. The ability to retrofit this supervision onto pretrained policies (like $\pi0.5$ ) suggests a scalable path to improving generalization in existing VLA systems without retraining from scratch. The results indicate that explicit modeling of the relation between current and future states, rather than just the future state itself, is a key factor in improving robustness under distribution shifts.