UniJEPA: Enhancing Robot Policy via Unified Continuous and Discrete Representation Learning

UniJEPA is a unified robot policy framework that leverages large-scale pretraining on instructional videos to learn combined continuous and discrete representations, enabling superior performance in both simulation and real-world out-of-distribution tasks compared to existing vision-language and generative approaches.

The Gist

Imagine you are trying to teach a robot to do chores around the house. The paper introduces a new method called UniJEPA (Unified Joint Embedding Prediction and Action) to make these robots smarter, more flexible, and better at handling things they've never seen before.

Here is the breakdown of how it works, using simple analogies:

The Problem: The "Two-Headed" Robot

Currently, robot brains usually fall into two camps, and both have a weakness:

1. The "Talker" (Vision-Language Models): These robots are great at understanding language and pictures. If you say, "Pick up the red cup," they know what a cup is and what "red" means. But they are bad at predicting physics. They don't intuitively know how the cup will wobble if they grab it too hard or how it will roll across the table.
2. The "Predictor" (Generative Models): These robots are great at guessing what happens next. If they see a ball rolling, they can predict where it will be in a second. But they often lack "common sense" or language understanding. They might know how to move but not why they are moving or what the human actually asked them to do.

Most robots try to use one or the other, or they try to mash them together in a way that loses the best parts of both.

The Solution: The "Bilingual Dreamer" (UniJEPA)

UniJEPA is like a robot that learns to speak two languages and dream in two ways simultaneously. It combines the "Talker" and the "Predictor" into one brain.

Think of it like a student learning to drive a car:

• Discrete Learning (The "Talker"): This is like learning the rules of the road and the vocabulary. "Stop sign means stop," "Green means go." The robot learns to understand complex instructions and describe what it sees using words.
• Continuous Learning (The "Dreamer"): This is like learning the feel of the car. It's not just about the rules; it's about predicting the smooth flow of motion. If you turn the wheel slightly, how does the car drift? UniJEPA learns to predict the future visual scene not as a blurry video, but as a high-level "feeling" or map of what will happen next.

How They Trained It (The Two-Stage Process)

Stage 1: The "Library & Movie Theater" Phase (Pre-training)
Before the robot ever touches a real object, they let it "read" and "watch" massive amounts of data.

• They fed it over 1 million videos of humans and robots doing tasks (like opening drawers or picking up toys).
• The Trick: They asked the robot to do two things at once:
  1. Answer Questions: "What is the robot doing?" (This sharpens its language and understanding).
  2. Predict the Future: "If the robot moves its arm this way, what will the picture look like in 1 second?" (This sharpens its understanding of physics and motion).
• By doing both, the robot builds a mental model where words and physical motion are perfectly linked.

Stage 2: The "Driving School" Phase (Fine-tuning)
Once the robot has this general knowledge, they teach it specifically how to move its own body.

• They show it data from the actual robot arm or hand.
• The robot learns to translate its "dreams" (predictions of the future) and its "understanding" (language instructions) directly into action tokens (the specific commands to move motors).
• It uses a special "expert" system (like a team of specialists) to handle the complex math of moving a real arm without crashing.

The Results: Why It's Better

The paper tested this robot in two ways:

1. In a Video Game (Simulation): They gave it tasks it had never seen before, like moving a specific object in a specific way. UniJEPA beat all the other top robots by a significant margin (about 9-12% better).
2. In the Real World: They put it on a real robot arm and a fancy 12-fingered robot hand.
   • The Magic: When they gave the robot a task with a completely new object (e.g., a toy it had never seen, or a strange color), UniJEPA didn't get confused. Because it learned the concept of "grasping" and "moving" rather than just memorizing specific pictures, it could handle these "out-of-distribution" (strange) situations much better than the competition.

The Bottom Line

UniJEPA is a robot brain that doesn't just memorize instructions or just guess physics. It learns to understand the world through language while simultaneously simulating the future through motion. This dual approach allows it to be a "generalist"—a robot that can adapt to new tasks and new objects without needing to be retrained from scratch every time.

Technical Summary

Technical Summary: UniJEPA

Problem Statement

The development of generalist robot policies capable of handling diverse tasks in open-ended environments remains a central challenge in embodied AI. Existing approaches typically rely on either Vision-Language Models (VLMs) for semantic understanding or generative models for visual dynamics. However, these methods face significant limitations:

• VLM-based approaches often suffer from a degradation of foundational capabilities when fine-tuned on scarce, heterogeneous robotic datasets. They frequently overlook the fundamental discrepancies between static vision-language tasks and dynamic robotic action tasks.
• Generation-based approaches (e.g., video prediction) facilitate dynamic representation learning but often fail to preserve the crucial vision-language alignment inherent in pre-trained VLMs.
• Current Unified Models often attempt to unify understanding and generation within a discrete token prediction framework, which may compromise the robustness of vision-language alignment and the ability to model continuous physical dynamics effectively.

The paper posits that robotic policy learning requires a paradigm that integrates discrete task comprehension (understanding) with continuous future state representation learning (planning and dynamics).

Methodology: UniJEPA

The authors propose UniJEPA, a Vision-Language-Action (VLA) framework that unifies discrete and continuous representation learning through a two-stage training strategy. The architecture utilizes a Mixture-of-Transformers (MoT) design with modality-specialized experts.

1. Unified Vision-Language Embedding Modeling (Stage 1: Pre-training)

The first stage establishes a cross-embodiment pre-training paradigm to learn joint text-image representations.

• Architecture: The model initializes with a pre-trained VLM (specifically Paligemma) and introduces a Generator Expert and an Action Expert (in later stages) within the MoT framework.
• Discrete Representation Learning: The model is trained on large-scale vision-language datasets and embodied task descriptions (annotated into VQA-style formats) to learn fine-grained language representations for understanding scenes and instructions.
• Continuous World Modeling: Unlike prior works that predict raw image pixels, UniJEPA predicts continuous high-dimensional visual features. A frozen visual encoder (e.g., SigLIP) encodes future observations, and the model learns to predict these continuous features ($\hat{o}_{t+h}$) alongside discrete text tokens ($\hat{l}$).
• Training Objective: The model employs a joint loss function combining Cross-Entropy loss for the language branch and Mean Squared Error (MSE) loss for the continuous visual prediction branch.

2. Unified Action Modeling (Stage 2: Fine-tuning)

In the second stage, the model is fine-tuned on embodiment-specific data to map predictive representations to action tokens.

• Action Expert: A distinct expert is trained from scratch to handle the action space. It leverages flow matching to capture the continuous and inherently multi-modal distribution of actions.
• Joint Optimization: The model is trained to simultaneously predict future continuous visual states and execute action sequences. The objective function combines the continuous visual prediction loss with a flow matching loss for action generation.
• Data Strategy: Pre-training utilizes a mix of robot videos (with subtask descriptions), human demonstration videos, and generic vision-language QA data. Fine-tuning is performed exclusively on VLA data from simulation and real-world robots.

Key Contributions

1. Novel VLA Architecture: The introduction of a VLA framework that integrates discrete token prediction (for understanding) with continuous visual prediction (for dynamics) within a unified MoT architecture.
2. Two-Stage Training Framework: A strategy that aligns action representations while preserving the intermediate representations learned during large-scale pre-training, effectively transferring knowledge from internet-scale data to robotic tasks.
3. Continuous Feature Prediction: The demonstration that predicting continuous visual features (rather than raw pixels or discrete tokens) serves as a more effective signal for learning dynamic information crucial for action generation.

Experimental Results

UniJEPA was evaluated across two simulation benchmarks and two real-world robotic platforms (a 7-DoF Franka Emika Panda arm and a 12-DoF dexterous hand).

• Simulation Benchmarks:
  • Calvin Benchmark: UniJEPA achieved an average sub-task sequence length of 4.11, outperforming the previous best (UP-VLA at 4.08) and significantly surpassing baselines like GR-1 (3.06).
  • SimplerEnv: On the WindowX and Google Robot platforms, UniJEPA achieved state-of-the-art success rates of 71.0% and 78.4%, respectively. It demonstrated consistent high performance across all sub-tasks, avoiding the "spiky" performance profiles of other methods.
• Real-World Experiments:
  • Franka Panda Arm: UniJEPA achieved the highest success rates across all four task categories (Pick & Place, Press Button, Route Cable, Drawer), including on unseen tasks involving novel objects.
  • 12-DoF Dexterous Hand: The model achieved the highest average success rate across nine skill categories, showing a 12% improvement on unseen tasks compared to baselines.
• Ablation Studies:
  • Removing pre-training resulted in a ~16% drop in success rate on real-world pick-and-place tasks.
  • Using continuous visual features for prediction yielded a ~20% performance boost compared to models without this component.
  • Among continuous encoding methods, SigLIP features provided the best alignment with the VLM backbone, outperforming DINO and distillation-based approaches.

Significance and Claims

The paper claims that UniJEPA successfully addresses the trade-off between semantic understanding and dynamic prediction in robot policy learning. By leveraging large-scale pre-training on both robot and human demonstration data, the model acquires robust generalization capabilities.

The authors emphasize that their approach enables robots to handle completely novel objects and out-of-distribution (OOD) scenarios more effectively than existing methods. The integration of continuous future representation learning allows the policy to better anticipate physical dynamics, while the preserved vision-language alignment ensures accurate interpretation of complex instructions. The results suggest that a unified approach combining understanding, planning, and continuous representation learning is a viable path toward generalist robot policies.