JEPA Policy: Diffusion-Free Imitation Learning via Paired Action and Future Representation Prediction

JEPA Policy: Diffusion-Free Imitation Learning via Paired Action and Future Representation Prediction

|Authors: Jie Xu, Kangjin Yu, Ziyi Jin, Junjie Gao, Liqing Chen, Yixian Li, Shuai Tian, Zhongpu Xia

|Venue: arXiv preprint | Year: 2026

|arXiv: 2609.09630v1

Abstract

Standard behavior cloning supervises actions without explicitly constraining the future representation paired with each demonstrated action chunk. We introduce JEPA Policy, a diffusion-free framework that uses the action chunk and its observed future representation as paired training targets. Action and future-representation tokens interact in a shared Transformer and are refined through two forward passes. Future prediction can therefore shape the representation used to generate actions. Dual-branch and gradient-routing controls attribute the gain to this shared topology rather than to an auxiliary prediction head alone. Across nine simulated tasks, JEPA Policy improves mean success over the action-only MIP baseline and outperforms Diffusion Policy under the evaluated configurations, while adding 0.29 ms to MIP’s model latency. A five-task, 630-episode physical-robot study produces the same pooled ranking. Further audits find no complete representation collapse under action supervision and identify a task-conditioned failure-ranking signal in future-prediction error. These results support paired future-representation supervision as a practical approach to low-latency visuomotor imitation without iterative generative sampling.

Key Contributions

• Standard behavior cloning supervises actions without explicitly constraining the future represent…
• We introduce JEPA Policy, a diffusion-free framework that uses the action chunk and its observed …
• Action and future-representation tokens interact in a shared Transformer and are refined through …
• Future prediction can therefore shape the representation used to generate actions
• Dual-branch and gradient-routing controls attribute the gain to this shared topology rather than …

Topics

• vision
• reinforcement-learning
• control
• learning-from-demonstration

Code & Data

• GitHub repository: https://github.com/jiejie567/JEPA-Policy

BibTeX

@article{Xu2026_260909630v1,
  title     = {JEPA Policy: Diffusion-Free Imitation Learning via Paired Action and Future Representation Prediction},
  author    = {Jie Xu and Kangjin Yu and Ziyi Jin and Junjie Gao and Liqing Chen and Yixian Li and Shuai Tian and Zhongpu Xia},
  year      = {2026},
  eprint    = {2609.09630v1},
  archivePrefix = {arXiv},
  primaryClass  = {cs.RO},
  url       = {https://arxiv.org/abs/2609.09630v1}
}