UniJEPA: Enhancing Robot Policy via Unified Continuous and Discrete Representation Learning
Abstract
Building generalist robot policies that can handle diverse tasks in open-ended environments is a central challenge in robotics. To leverage knowledge from large-scale pretraining, prior work (VLA) has typically built generalist policies either on top of vision-language models (VLMs) or generative models. However, both semantic understanding from vision-language pretraining and visual dynamics modeling from visual-generation pretraining are crucial for embodied robots. Recent unified models of generation and understanding have demonstrated strong capabilities in both comprehension and generation through large-scale pretraining. We posit that robotic policy learning can likewise benefit from the combined strengths of understanding, planning and continuous future representation learning. Building on this insight, we introduce UniJEPA, which acquires the ability to dynamically model high-dimensional visual features through pretraining on over 1M internet-scale instructional manipulation videos. Subsequently, UniJEPA is fine-tuned on data collected from the robot embodiment, enabling the learning of mappings from predictive representations to action tokens. Extensive experiments show our approach consistently outperforms baseline methods in terms of 9\% and 12\% across simulation environments and real-world out-of-distribution tasks.
Lay Summary
(1) Problem. Teaching robots to do many real-world tasks is hard because today’s learning systems usually treat low-level motions (continuous actions) and high-level meanings (discrete language and task concepts) separately. This split makes policies brittle and less transferable to new environments or instructions. (2) Solution. In UniJEPA, we introduce a JEPA-style world modeling framework for vision-language-action learning that jointly learns continuous control representations and discrete semantic representations in one unified latent space. By aligning “how to move” with “what the task means,” the model can better predict useful future representations and learn policies that connect perception, language, and action more consistently. (3) Impact. This unified design improves robot policy learning across diverse manipulation settings and supports stronger generalization to unseen tasks. In practice, it moves us closer to robots that can adapt more reliably from limited data, which is important for deploying helpful robots in homes, labs, and industry.