Paired action and future prediction speeds up robot imitation learning
JEPA Policy: Diffusion-Free Imitation Learning via Paired Action and Future Representation Prediction
Robotics
Summary
Teaching robots to copy actions usually focuses just on learning the actions without considering what happens afterward. The authors propose a method called JEPA Policy that trains robots by pairing each action with a prediction of the future outcome, helping the robot understand how its actions affect what comes next. This approach uses a shared transformer model that improves learning efficiency without the slow sampling methods used in some other techniques. Tests on simulations and real robots show it works better and faster than some existing methods.
What this means in practice
- •For robotics engineers: Develop faster and more reliable robot controllers by training with paired action and future representation targets.
- •For industrial automation teams: Improve the efficiency of visually guided robotic tasks by reducing model latency without sacrificing accuracy.
Authors
Jie Xu, Kangjin Yu, Ziyi Jin, Junjie Gao, Liqing Chen, Yixian Li, Shuai Tian, Zhongpu Xia
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.