PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2609.09630v1cs.RO

TL;DR

Standard behavior cloning supervises actions without explicitly constraining the paired future representation. JEPA Policy jointly predicts actions and future representations in a shared two-step Transformer, improving evaluated control performance while retaining low latency and revealing task-conditioned diagnostic limits.

  • Problem

    Standard behavior cloning predicts expert actions without explicitly constraining the future outcomes paired with those actions in demonstrations.

    Source: §I · pp. 1–2

  • Method

    JEPA Policy uses paired action–future representation targets, shared action and future tokens, and two diffusion-free prediction passes.

    Source: §I · p. 2, C. Joint Action–Future Representation… · p. 5, D. MIP-Style Two-Step Prediction… · p. 6

  • Results

    Across nine simulated tasks, JEPA Policy improves the MIP mean by +5.6 points (77.4 →83.0), with 13.2 ms latency versus 439.5 ms for the evaluated 100-step Diffusion Policy.

    Source: C. Main Results · pp. 9–10

  • Takeaways & Limitations

    Paired future-representation supervision supports low-latency visuomotor imitation and provides a per-task-calibrated future-consistency signal for failure analysis.

    Source: J. Future-Consistency Diagnostic at… · p. 15, §VI · p. 17

    The evidence is mainly simulated; the physical study is smaller, generalization beyond three benchmark suites is untested, and the advantage over MIP is only directional in hardware.

    Source: §V · p. 16

Abstract

from arXiv · show

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.

I. INTRODUCTION

Standard behavior cloning predicts expert actions but does not explicitly constrain the future outcomes paired with them. JEPA Policy adds paired future-representation supervision in a shared, two-step diffusion-free predictor.

Motivation and approach

  • Expert demonstrations provide aligned action and future-state labels, whereas standard behavior cloning typically uses only the action label.The future state records the task-relevant outcome toward which the action advances the current scene.Source: §I · p. 1
  • Paired supervision links each action to its demonstrated outcome without requiring video reconstruction of texture, illumination, or background.Future latent prediction targets task-relevant representation rather than pixel-level appearance.Source: §I · p. 2

Deployment diagnostic

  • The future-prediction branch also yields a rollout-level consistency error by comparing predicted and realized future representations after execution.The interface supports failure analysis, risk monitoring, or potential human intervention.Source: §I · pp. 2–3

Motivation and approach

  • JEPA Policy jointly supervises each expert action sequence and its observed future representation within one shared Transformer.Action and future-representation tokens interact in every self-attention layer, allowing future prediction to shape the action-generating representation.Source: §I · p. 2
  • JEPA Policy extends MIP's coarse-to-refined process to jointly predict actions and paired future representations in two diffusion-free passes.Training and deployment avoid a denoising chain while producing a future latent alongside each action chunk.Source: §I · pp. 2–3

Representation constraint

  • Action supervision supplies a task-driven constraint against representation collapse, avoiding EMA, frozen-target, and explicit variance–covariance mechanisms.The shared representation must support both action prediction and future prediction, while the encoder is trained from scratch.Source: §I · pp. 2–3

II. RELATED WORK

JEPA Policy combines latent future prediction with supervised imitation rather than pixel reconstruction, planning, or a separate dynamics model. Its diagnostic is deliberately narrower than dedicated runtime monitors, and its speed claims exclude several alternative few-step methods.

Diffusion and generative action modeling

  • Unlike diffusion policies, JEPA Policy adds paired future supervision within a two-pass inference budget rather than iterative denoising.The paper does not benchmark consistency or flow-based few-step alternatives and makes no relative speed claim against them.Source: B. Diffusion Policy and Generative… · p. 3

Latent predictive representations

  • Latent future prediction avoids modeling low-level image details but still requires addressing representation collapse.JEPA Policy embeds predictive representation learning inside imitation learning and uses action supervision as its task constraint.Source: C. Predictive Representation Learning… · p. 3

Vision-language and pretrained models

  • JEPA Policy differs from pretrained VLA and JEPA-VLA approaches by training on task demonstrations, sharing action and future tokens, and using two MIP steps.The method does not rely on the pretrained-backbone setting described for concurrent approaches.Source: C. Predictive Representation Learning… · p. 3

World and forward models

  • Unlike world and forward models, JEPA Policy learns no separate dynamics model, performs no planning rollout, and reconstructs no future images.It predicts one demonstrated future latent paired with the supervised action.Source: D. World Models and Future Prediction · p. 3

Runtime failure monitoring

  • The future-consistency diagnostic is a by-product of training, requiring one encoder pass, and is presented as a per-task-calibrated risk signal rather than a dedicated-monitor competitor.Direct comparison with dedicated runtime monitors remains future work.Source: E. Runtime Failure Monitoring for… · pp. 3–4

A. Problem Formulation

The problem formulation pairs current observations with expert action chunks and future observations, then maps them to actions and future representations. A shared encoder and Transformer implement this mapping through mutually visible action and future tokens with coarse-to-refined prediction.

A. Problem Formulation

  • The demonstration dataset contains a current observation, an expert action sequence over horizon H, and a future observation at horizon k.Both horizons are counted in environment control steps, and k determines how far into the future the target lies.Source: A. Problem Formulation · p. 4
  • Standard behavior cloning predicts expert actions without explicitly constraining the future outcomes induced along the demonstration trajectory.JEPA Policy addresses this gap by learning a joint action–future representation mapping.Source: A. Problem Formulation · p. 4
  • The policy outputs a predicted action sequence and future observation representation from the current observation, action input, future-representation input, and interval endpoints.The unified network is written as the coarse-to-refined predictor Fθ.Source: A. Problem Formulation · p. 4
  • The future target is the image-path representation of the actual future observation, while the current conditioning representation also incorporates low-dimensional inputs.The same encoder supplies both image representations, but the future target is image-only and higher-dimensional in the stated configuration.Source: A. Problem Formulation · p. 4

B. Shared Action–Future Observation Representation Policy

  • The model uses H action tokens and one future token, with learned positional embeddings and a dedicated future-position embedding.The conditioning memory includes time-interval embeddings and the encoded current observation.Source: B. Shared Action–Future Observation… · p. 4
  • A shared Transformer receives action and future-representation tokens that remain mutually visible through unmasked self-attention and cross-attention to conditioning memory.Separate heads decode the action sequence and future representation from the shared output.Source: B. Shared Action–Future Observation… · pp. 4–5, Figure 2 · p. 5
  • Two forward steps first regress a coarse prediction and then refine both outputs using noised expert targets during training or first-step predictions during inference.The two interval settings are (0, τ) and (τ, 1), with zero inputs in the first step.Source: B. Shared Action–Future Observation… · p. 5
  • Shared attention makes future prediction part of the action-generating computation, unlike a separate branch connected only through an upstream encoder.The experiments therefore test shared-versus-separate topology and route-cut controls.Source: B. Shared Action–Future Observation… · p. 5

C. Joint Action–Future Representation Objective

The objective jointly optimizes action prediction and future-representation prediction against the demonstrated future latent. Stop-gradient protects the target path, while action loss constrains the shared encoder and adaptive weighting balances the two terms.

Joint objective

  • The future-representation loss aligns the predicted latent with the actual future observation's encoded representation using target-normalized mean squared error.Normalization separates latent-space scale from action-space scale and simplifies loss balancing.Source: C. Joint Action–Future Representation… · p. 5

Gradient routing

  • Stop-gradient is applied only to the future target representation, preventing direct target-side co-adaptation while allowing gradients through the current-observation encoder path.The action loss supplies the task-driven constraint on the shared representation.Source: C. Joint Action–Future Representation… · p. 5

Loss balancing

  • The main experiments use adaptive ratio weighting, recomputing λz once per minibatch to regulate the future-loss contribution relative to action loss.The default ratio target is ρ = 0.10 with clipping bounds [10^-6, 0.1].Source: C. Joint Action–Future Representation… · p. 6

Two-step objective

  • The two-step objective sums interval-scaled losses from both prediction steps while sharing one future-loss weight.The second step uses a noised target-derived input, whereas the first step starts from zero.Source: C. Joint Action–Future Representation… · p. 6

Joint objective

  • The overall objective jointly optimizes action prediction and future-representation prediction in the same policy network.This introduces paired action–future supervision without changing the diffusion-free policy formulation.Source: C. Joint Action–Future Representation… · pp. 5–6

D. MIP-Style Two-Step Prediction Training

JEPA Policy extends MIP with two supervised prediction steps for jointly refining actions and future observation representations. Training uses noisy expert-target inputs for refinement, while inference feeds the coarse output into a noise-free second step.

E. Future Observation Representation Prediction under Paired Action Supervision

Future representation prediction is under-constrained when separated from action supervision because similar observations can precede different outcomes. JEPA Policy preserves each demonstrated action–outcome correspondence by jointly supervising both targets through shared representations.

F. Deployment-Time Future-Consistency Diagnostics

JEPA Policy predicts a future representation alongside the refined action and compares it with the representation of the actually reached observation. This future-consistency error supports post-execution failure analysis, but it is not an instantaneous safety guarantee.

  • The policy outputs a refined action sequence and future representation, which is compared after execution with the encoder representation of the actual future observation.Source: F. Deployment-Time Future-Consistency… · p. 7
  • The rollout-level future-consistency error measures disagreement between predicted and observed post-execution representations.Source: F. Deployment-Time Future-Consistency… · p. 7
  • A large error may indicate unexpected behavior, abnormal contact, or a visual state outside the training distribution.Source: F. Deployment-Time Future-Consistency… · p. 7
  • Because the signal becomes available only after the future observation is reached, it informs later control or intervention rather than preventing the preceding action.Source: F. Deployment-Time Future-Consistency… · p. 7
  • The evaluation covers nine simulated manipulation tasks, collapse auditing, a diagnostic audit, and a physical-hardware comparison.Source: §IV · p. 7, A. Setup · p. 7
  • Primary comparisons use matched training and best-checkpoint evaluation protocols across MIP, Diffusion Policy, and JEPA Policy.Source: A. Setup · p. 7
  • Control runs at 20 Hz with chunked action prediction and replanning, while the future target is defined over a short transition horizon.Source: A. Setup · pp. 7–8
  • The future target is a stop-gradient pooled image embedding, excluding low-dimensional inputs such as proprioception.Source: A. Setup · p. 8

B. Compared Methods

JEPA Policy is compared with action-only MIP, Diffusion Policy, and controlled dual-branch variants under matched training and evaluation protocols. Across nine tasks, it improves mean success while preserving low inference latency, though training costs more and timing comparisons have methodological limits.

Robustness and limitations

  • The main success comparison uses best-checkpoint selection, but the gain remains +4.7 to +5.6 points under selection-free summaries.Absolute success decreases under last-checkpoint reporting, while identical selection rules largely cancel the optimism in the difference.Source: C. Main Results · p. 9

Main comparison

  • JEPA Policy raises the nine-task mean from 77.4 for MIP to 83.0, with positive differences on all nine tasks.The paired task-level analysis reports a mean gain of +5.61 points and positive differences across all tasks.Source: C. Main Results · pp. 9–10, Table II · p. 9

Inference cost

  • JEPA Policy adds 0.29 ms to MIP latency, averaging 13.2 ms while the 100-step Diffusion Policy takes 439.5 ms.The reported JEPA latency is p95 = 14.5 ms and uses 3% of the 400 ms budget; the diffusion comparison is a single-task, single-seed timing audit.Source: C. Main Results · p. 10

Training and evaluation cost

  • Training the future branch increases optimization time by 12–22% and adds evaluation overhead, despite making decisions 33× cheaper than Diffusion Policy.JEPA Policy costs about 1.8× a Diffusion Policy training step in the timing audit, while the training-to-target difference is not significant at one seed.Source: C. Main Results · pp. 10–11

D. Shared Attention vs. Separate Branches

JEPA Policy outperforms action-only and separate-branch variants when future representations are processed jointly with actions in a shared Transformer topology.

  • 83.0 mean success makes JEPA Policy the sole best variant on eight tasks, versus 77.4 for MIP, 77.6 for Dual-cross, and 76.3 for Dual-independent.Transport is the exception: Dual-independent reaches 99.2 while JEPA Policy reaches 98.3.Source: D. Shared Attention vs. Separate… · p. 11
  • Separate-branch future supervision yields no mean improvement over MIP, and cross-attention alone does not recover the shared-stack gain.The dual variants retain the training recipe and losses while changing only the topology.Source: D. Shared Attention vs. Separate… · p. 11
  • The shared topology, rather than added capacity, accounts for the gain because both dual variants use more trainable parameters yet fail to beat MIP.JEPA Policy uses 43.96 M parameters, compared with 48.38 M and 53.73 M for the dual variants.Source: D. Shared Attention vs. Separate… · p. 11

E. Where the Future Objective Acts

Gradient-routing analyses locate the future objective’s main effect in the shared action-producing representation, with direct token attention providing an additional contribution.

  • The analysis measures the future-loss gradient reaching parameters that produce the action relative to the action-loss gradient on those same parameters.Dual-independent is exactly zero by construction because its decoder stacks share no parameters.Source: E. Where the Future Objective Acts · p. 11
  • The shared stack carries 0.68% future-to-action gradient at convergence, versus 0.135% for Dual-cross and exactly 0 for Dual-independent.Across training windows, the shared-stack ratio is five to ten times the Dual-cross ratio.Source: E. Where the Future Objective Acts · p. 11, Figure 8 · p. 12
  • 82.8 mean success for the full method falls to 73.1 after cutting the backward route, while cutting the forward route retains 78.4.The matched comparison uses one seed, so it establishes aggregate ordering rather than a stable effect size.Source: E. Where the Future Objective Acts · p. 12, Table VI · p. 12
  • The gradient measurements are correlational, and the one-seed route-cutting probe has coarse 2.5-point evaluation resolution and substantial task variation.Back-cut offsets −47.5 points on Tool Hang with +35.0 on Coffee Preparation, so no significance claim is attached.Source: E. Where the Future Objective Acts · pp. 11–12

F. Future-Loss Ratio and Horizon

Future-branch performance is broadly robust across loss ratios, horizons, camera settings, and input modalities, while explicit isotropic regularization is fragile and collapse remains bounded under action supervision.

  • No swept future-supervision configuration is materially worse than baseline; ρ = 0.10–0.20 and k = 4–6 are uniformly strong choices.The only below-baseline cell is ρ = 0.05 on MugMug, at 92.5 versus 94.2.Source: F. Future-Loss Ratio and Horizon · p. 12, Figure 9 · p. 13
  • Six of seven wrist-image tasks improve and one ties, while all five usable fixed-camera differences remain within ±5 points.The camera sweep is directional and single-seed; launch batches varied by 17.5 points.Source: G. Robustness to the Observation… · p. 13, Table VIII · p. 13
  • Adding SIGReg with weight 0.09 lowers performance on all five tasks, including MokaMoka at 0.0 across all seeds.Without a weight sweep, the proposed over-regularization explanation remains a hypothesis.Source: H. Alternative Designs and Negative… · p. 13, Table IX · p. 13
  • Transport improves from 52.5 to 62.5 with state inputs, while Tool Hang changes from 75.0 to 72.5, showing task-dependent effects beyond image observations.The results do not indicate systematic harm to either observation modality.Source: H. Alternative Designs and Negative… · p. 13
  • The collapse audit finds no complete collapse or systematic contraction to the pre-specified low-rank boundary under the full objective, despite endpoint exceptions.Exceptions occur for encoder TE32 on MugMug and MokaMoka, and predictor TE32 on Tool Hang and Square.Source: §I · p. 13, Figure 10 · p. 14
  • Removing action supervision collapses the representation: encoder TE32 drops from 0.170 to 0.0014, about 36× below the 0.05 boundary.Centered effective rank decreases from 181 to 24 of 384.Source: §I · p. 14

J. Future-Consistency Diagnostic at Rollout

Early Future-4 NMSE ranks failures above successes within tasks, but the signal requires task-conditioned interpretation rather than pooled raw scores. Controls indicate the ranking reflects learned future prediction rather than generic latent motion.

Diagnostic performance

Controls

  • Residualizing error against realized latent visual change left checkpoint-macro AUROC essentially unchanged, from 0.6184 to 0.6183.This control argues against failures producing more motion as the explanation for the signal.Source: J. Future-Consistency Diagnostic at… · p. 15
  • The persistence baseline was near chance at 0.508 versus 0.618 for the learned predictor, indicating dependence on learned prediction rather than generic latent motion alone.Source: J. Future-Consistency Diagnostic at… · p. 15

Calibration

  • Pooling raw NMSE across tasks removed the signal, yielding AUROC 0.501 with CI [0.484, 0.519] and permutation p = 0.68.The authors attribute this to task-specific error scales and require per-task calibration.Source: J. Future-Consistency Diagnostic at… · p. 15

Interpretation

  • The diagnostic supports within-task failure ranking and hand-off, not a universal failure probability or directly validated detection accuracy.Operational thresholds still require held-out outcome-labeled rollouts, and no direct comparison was conducted.Source: J. Future-Consistency Diagnostic at… · pp. 14–15

K. Real-Robot Evaluation

The physical-robot study preserved JEPA Policy’s simulated ranking across five tasks while showing substantially lower latency than diffusion baselines. The advantage was concentrated in some tasks, and the study remains limited by small, single-robot coverage.

Overall results

  • JEPA Policy achieved 66.9% episode-weighted success versus 54.7% for MIP, 31.8% for DP-16, and 18.7% for DP-100.It led DP-16 in 14 of 15 matched cells and DP-100 in all 15, with mean margins of +35 and +48 points.Source: K. Real-Robot Evaluation · p. 15, Table XI · p. 15

Task variation

  • The margin over Diffusion Policy exceeded simulation, while the advantage over MIP concentrated in Cabinet at +30.0 points; Plate Grape favored MIP by −3.3 points and Pen Insert was nearly tied at +3.3.Source: K. Real-Robot Evaluation · p. 16, §VI · p. 17

Failure stages

  • On Cabinet, none of JEPA Policy’s 12 failures occurred during door closure, whereas 7 of 30 DP-16 and 2 of 30 DP-100 failures reached that final stage without closing.Whenever JEPA Policy reached the door, it closed it; video inspection supported the outcome comparison.Source: K. Real-Robot Evaluation · p. 16, Figure 12 · p. 16

Latency

  • JEPA Policy added 0.29 ms, or 2.2%, to MIP and reached 13.2 ms decision time, 33× lower than evaluated Diffusion Policy.Measured physical-robot frame-to-chunk latency was 20.1/63.3 ms at median/p95, compared with 27.6/67.5 ms for MIP.Source: K. Real-Robot Evaluation · p. 15, §V · p. 16

Mechanism

  • The shared-stack ablation indicates that future supervision improves success through its integration into the action-generating topology rather than through an auxiliary loss alone.Separate-branch variants provided little or no benefit while target modality, anti-collapse mechanism, and training recipe were held fixed.Source: §V · p. 16

Scope

  • The hardware evidence is limited to five tasks, one robot, one operator, three checkpoints, and ten episodes per nominal cell, with noisy individual cells and mostly simulated experiments.Single-seed side studies are directional, and multi-modal futures remain untested.Source: §V · pp. 16–17

Related papers

Loading 2609.09630v1…