Robotics Research Hub
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

Jie Xu, Kangjin Yu, Ziyi Jin, Junjie Gao, Liqing Chen, Yixian Li, Shuai Tian, Zhongpu Xia · N/A · 2026

vision reinforcement-learning control learning-from-demonstration
PDF arXiv Code (GitHub)

Framework

N/A

License

N/A

Stars

N/A

Summary

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.

Abstract Summary

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 Points

  • 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.

Method

This section specifies JEPA Policy’s paired objective, shared stack, two-step training rule, and deployment diagnostic. Given an expert demonstration dataset where $o_{t}$ denotes the current observation or observation history window, $\mathbf{a}_{t:t+H}$ denotes the expert action sequence over a prediction horizon of $H$, and $o_{t+k}$ denotes the future observation at horizon $k$.

We use $\mathbf{a}_{t:t+H}$ as compact notation for the next $H$ actions, $(\mathbf{a}_{t+1},\ldots,\mathbf{a}_{t+H})$.

Method — JEPA Policy: Diffusion-Free Imitation Learning via Paired Action and Future Representation Prediction
Fig. 3: Coarse-to-refined MIP-style training for joint action and future observation representation prediction. The first step maps a zero input (time $0$) to a coarse prediction, and the second refines it (time $\tau$). The diagram shows inference, where the second step receives the first-step output. During training, it instead receives the noised expert target from Sec. III-D . Both steps are supervised, with residuals scaled by $\tfrac{1}{\tau}$ and $\tfrac{1}{1-\tau}$. The method uses no diffusion-style iterative sampling.

Experiments

We evaluate JEPA Policy on nine simulated manipulation tasks from three benchmark suites, using three training seeds per configuration. The evaluation also includes a pre-specified collapse audit on seven tasks over 84 checkpoints (Sec. IV-I ) and a 5,200-episode diagnostic audit on the original eight-task set (Sec. IV-J ). Sec. IV-K repeats the comparison on physical hardware. Tasks.

We use tasks from LIBERO [ 44 ] , robomimic [ 8 ] , and MimicGen [ 45 ] . These suites are built on robosuite/MuJoCo [ 46 , 47 ] .

Experiments — JEPA Policy: Diffusion-Free Imitation Learning via Paired Action and Future Representation Prediction
Fig. 9: Hyperparameter sensitivity as gain over the baseline (three-seed means, in points): future-loss ratio $\rho$ (left) and future horizon $k$ (right); blue is better than the baseline, orange worse. No configuration is materially worse than the baseline; $\rho\in[0.1,0.2]$ and $k\in\{4,6\}$ are uniformly strong choices.

Sources and Demonstrations

Figures are reproduced from Xu et al. See the full paper for experimental details and the project page, when the paper links one, for demonstrations and videos.

Share

Related Papers

A Sim-to-Real Integration Pipeline for Training and Deployment of Chunk-Based VLA Manipulation Policies
arXiv preprint

A Sim-to-Real Integration Pipeline for Training and Deployment of Chunk-Based VLA Manipulation Policies

Mathilde Kappel, Clémence Grislain, Mohamed Chetouani et al.

Vision-Language-Action (VLA) models have become a prominent paradigm for mapping multimodal inputs, including semantic instructions, visual observations of the scene, and proprioceptive observations, to robot actions.

manipulation vision vla reinforcement-learning sim-to-real control learning-from-demonstration benchmark
PDF Advanced
Capability-Aware Arbitration for Semantic Intent-Based Shared Control
arXiv preprint

Capability-Aware Arbitration for Semantic Intent-Based Shared Control

Zhaoda Du, Michael Bowman, Xiaoli Zhang

Shared control often allocates robot authority based on confidence in inferred human intent, assuming reliable autonomous execution. When this assumption fails, high intent confidence can cause over-helping.

manipulation vision vla reinforcement-learning control learning-from-demonstration benchmark
Code PDF Intermediate
GitHub ★ —