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.

  • Method

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

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

  • Takeaways & Limitations

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

  • Takeaways & Limitations

    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.

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.
  • Motivation and approach: 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.
  • 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.
  • 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.
  • Motivation and approach: 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.
  • 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.

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

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.
  • A. Problem Formulation: 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.
  • A. Problem Formulation: 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θ.
  • A. Problem Formulation: 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.
  • 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.
  • B. Shared Action–Future Observation Representation Policy: 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.
  • B. Shared Action–Future Observation Representation Policy: 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.
  • B. Shared Action–Future Observation Representation Policy: 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.

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

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.

  • The first step maps zero inputs to coarse action and future-representation predictions at time 0.
  • The second step refines both predictions toward the expert action and true future representation at time τ using noisy inputs.
  • Both prediction steps are supervised, and the second-step input is constructed from the noised expert target during training rather than the first-step output.At inference, the second step instead receives the first-step output.
  • Residuals are divided by τ and 1 − τ to place the two supervised corrections on a common scale.Only the second-step division compensates for injected noise; the first-step division is a matching interval convention.
  • Inference performs a coarse zero-input pass followed by noise-free refinement, avoiding diffusion-style iterative sampling.

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.

  • Plain future-representation regression can average incompatible futures or encode control-irrelevant variation when similar observations precede different outcomes.
  • Each demonstration segment jointly supplies an action target and the future representation observed after that action.
  • Shared attention allows the future loss to reach the action representation, preserving correspondence between a demonstrated action and its observed outcome.

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.
  • The rollout-level future-consistency error measures disagreement between predicted and observed post-execution representations.
  • A large error may indicate unexpected behavior, abnormal contact, or a visual state outside the training distribution.
  • Because the signal becomes available only after the future observation is reached, it informs later control or intervention rather than preventing the preceding action.
  • The evaluation covers nine simulated manipulation tasks, collapse auditing, a diagnostic audit, and a physical-hardware comparison.
  • Primary comparisons use matched training and best-checkpoint evaluation protocols across MIP, Diffusion Policy, and JEPA Policy.
  • Control runs at 20 Hz with chunked action prediction and replanning, while the future target is defined over a short transition horizon.
  • The future target is a stop-gradient pooled image embedding, excluding low-dimensional inputs such as proprioception.

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

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

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

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

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: AUROC exceeded 0.5 on all eight tasks, with seven conditional intervals excluding chance; MugMug reached 0.754 and Tool Hang 0.722, while MokaMoka crossed chance.The analysis covered 5,200 episodes from 24 frozen checkpoints across eight tasks.
  • 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.
  • Controls: 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.
  • 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.
  • 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.

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.
  • 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.
  • 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.
  • 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.
  • 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.
  • 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.
Loading 2609.09630v1…