PaperlayerView exact revision on arXiv

Source-linked AI summary

Delta-JEPA: Learning Action-Sensitive World Models via Latent Difference Decoding

Zhenghao Zhang, Yuanxiang Wang, Zhenyu Guan, Yujia Yang, Bingkang Shi, Tianyu Zong, Hongzhu Yi, Guoqing Chao, Xingchen Chen, Tiankun Yang, Chenxi Bao, Tao Yu, Jingjing Zhou, Jungang Xu

arXiv:2606.31232v1cs.AI

TL;DR

Reconstruction-free latent prediction can yield collapsed, action-insensitive world models, limiting useful action-conditioned rollouts. Delta-JEPA adds an action decoder that reconstructs actions from latent displacements. Across four continuous-control environments, it reports the highest mean planning success rate and outperforms JEPA-based and representation-learning baselines.

  • Problem

    Reconstruction-free latent prediction can collapse to constant representations, and world models need action-sensitive latent dynamics to support planning.

    Source: Introduction · pp. 1–2, Latent Forward Dynamics Predictor · p. 3

  • Method

    Delta-JEPA combines latent forward prediction with LDAD, which reconstructs executed actions from latent displacements in an end-to-end reconstruction-free model.

    Source: Introduction · pp. 1–2, Joint Optimization and End-to-End… · p. 4

  • Results

    Delta-JEPA achieves the highest mean planning success rate across all four environments, exceeding the strongest baseline by 15.14 percentage points on OGB-Cube.

    Source: Planning Performance · p. 5

  • Takeaways & Limitations

    A two-objective model using displacement-based action reconstruction improves planning and learns more action-sensitive latent transitions across four continuous-control tasks.

    Source: Conclusion · p. 8

    The evaluation is in a reward-free offline setting using trajectories collected by arbitrary, unknown behavior policies.

    Source: Problem Formulation · p. 2

Abstract

from arXiv · show

Learning visual world models for planning requires compact latent dynamics that remain sensitive to actions, yet reconstruction-free joint-embedding objectives can collapse to action-insensitive representations. We propose Delta-JEPA, an end-to-end reconstruction-free world model that augments latent forward prediction with a Latent Difference Action Decoder (LDAD). Unlike inverse decoders that infer actions from concatenated endpoint embeddings, LDAD reconstructs the executed action from the latent displacement between consecutive observations. This displacement-level supervision directly regularizes transition geometry: adjacent embeddings cannot collapse without losing action information, and different actions are encouraged to induce distinguishable latent changes for rollout-based planning. Delta-JEPA uses only latent prediction and action reconstruction, avoiding pixel reconstruction and distribution-matching regularizers. Across four visual continuous-control tasks, Delta-JEPA improves planning over JEPA-based and representation-learning world model baselines. Ablations show that displacement-based action decoding is consistently more effective than endpoint concatenation, and action-sensitivity analyses show clearer action-conditioned latent responses. These results indicate that supervising latent differences is a simple and effective mechanism for collapse-resistant and action-sensitive world model learning.

Introduction

Reconstruction-free latent prediction can avoid costly pixel reconstruction, but JEPA world models may collapse or fail to distinguish action-driven transitions. Delta-JEPA uses action reconstruction from latent displacements to encourage action-sensitive dynamics and reports improved planning across continuous-control tasks.

Introduction

  • Latent-prediction-only JEPA world models can collapse to constant representations, yielding deceptively low prediction loss while losing structure needed for planning.Pixel reconstruction can also spend capacity on visually detailed but dynamics-irrelevant information.Source: Introduction · p. 1
  • LeWorldModel stabilizes prediction without explicitly enforcing action sensitivity, while PLDM’s endpoint decoder and multi-loss objective can permit shortcuts and require hyperparameter tuning.Because its predictor receives the action, PLDM’s next-state embedding may carry action-correlated cues that support inverse prediction without representing the transition itself.Source: Introduction · p. 1
  • Delta-JEPA’s LDAD reconstructs executed actions from latent differences rather than concatenated endpoint embeddings, encouraging different actions to produce distinguishable transitions.This discourages action prediction from relying on state-specific cues rather than the transition.Source: Introduction · pp. 1–2
  • Delta-JEPA combines latent forward prediction with LDAD action reconstruction so candidate action sequences can be evaluated through action-sensitive latent rollouts.The paper reports improved planning and more action-sensitive latent transitions across diverse continuous-control tasks.Source: Introduction · p. 2
  • The paper’s stated contributions are displacement-based action-sensitive dynamics, a two-objective reconstruction-free world model, and empirical evaluation showing improved planning and stronger action-sensitive dynamics.The evaluation covers diverse continuous-control tasks.Source: Introduction · p. 2

Related Work

Prior visual world models use latent rollouts for planning, while JEPA predicts in representation space rather than reconstructing pixels. Recent collapse-prevention methods add different constraints, motivating Delta-JEPA’s displacement-based action supervision.

Latent World Models

  • PlaNet, Dreamer, and DreamerV3 use latent states and imagined rollouts, but commonly rely on reconstruction or reward-driven objectives.This motivates reconstruction-free models that predict compact representations focused on control-relevant state changes.Source: Latent World Models · p. 2

Joint Embedding Predictive Architectures

  • JEPA compares predictions in representation space; I-JEPA predicts masked image embeddings, and V-JEPA learns spatiotemporal representations without labels, text supervision, or pixel reconstruction.For world models, action-conditioned future-state prediction matters more than photorealistic synthesis, but latent-prediction-only training can collapse.Source: Joint Embedding Predictive Architectures · p. 2

Collapse Prevention and Inverse Dynamics

  • DINO-WM uses frozen DINOv2 features, LeWorldModel uses SigReg-style regularization, and PLDM combines VICReg-style regularization with endpoint-based inverse dynamics.Delta-JEPA instead decodes actions from latent displacements, aiming to distinguish action-induced changes without frozen encoders or complex multi-term regularization.Source: Collapse Prevention and Inverse Dynamics · p. 2

Method

Delta-JEPA learns a compact, action-sensitive latent dynamics model from reward-free offline image-action trajectories without pixel reconstruction. Its coupled objectives predict future latents and reconstruct actions from latent displacements, encouraging distinguishable action-conditioned transitions for planning.

Problem Formulation

  • The problem is to learn compact, action-sensitive latent dynamics from reward-free offline trajectories of image observations and continuous actions collected by unknown behavior policies.The objective excludes task rewards and pixel reconstruction.Source: Problem Formulation · p. 2

Overview of Delta-JEPA

  • A shared observation encoder supports two coupled objectives: the predictor forecasts the next latent from the current latent and action, while LDAD reconstructs the action from the latent displacement.Together they train an end-to-end reconstruction-free model that discourages action-insensitive representations.Source: Overview of Delta-JEPA · p. 2, Figure 1 · p. 3

Latent Forward Dynamics Predictor

  • The latent forward predictor estimates the next representation from the current encoded observation and action, but prediction loss alone can be minimized by nearly constant representations that preserve little planning information.LDAD adds an action-grounded constraint on adjacent latent-state differences.Source: Latent Forward Dynamics Predictor · p. 3

Latent Difference Action Decoder (LDAD)

  • LDAD imposes inverse dynamics on the difference between adjacent latent states, predicting the executed action with an action decoder trained by mean-squared reconstruction loss.The decoder receives the displacement rather than concatenated endpoint embeddings.Source: Latent Difference Action Decoder (LDAD) · p. 3

Action-Supervised Displacement Mechanism.

  • Because LDAD receives only each transition’s displacement, recovering its executed action encourages distinct action-induced displacements and endpoints, making candidate rollouts more informative for planning.The paper’s trajectory visualization shows nearby initial states separating under Delta-JEPA rollouts, while action-response PCA shows separated action-wise responses compared with LeWM’s concentration near the origin.Source: Action-Supervised Displacement Mechanism · pp. 3–4

Effects of Displacement-Based Action Decoding.

  • Displacement decoding discourages adjacent-observation collapse because a near-zero difference cannot carry action information, and reduces endpoint-decoder shortcuts that use action-correlated cues in the next latent.Action reconstruction must be supported by the change between adjacent latent states rather than state-specific cues.Source: Effects of Displacement-Based Action… · p. 4

3. Action-Sensitive Latent Dynamics

LDAD makes latent displacements distinguishable by action, supporting action-conditioned rollouts. Delta-JEPA extends this objective to multi-step decoding and trains with only latent prediction and action reconstruction.

3. Action-Sensitive Latent Dynamics

  • LDAD encourages different actions from the same latent state to produce distinguishable displacements and next-state embeddings for informative action-conditioned rollouts.This lets candidate actions be compared through the distinct latent rollouts they induce.Source: §3 · p. 4

Multi-Step Action Decoding

  • For horizon N, a Transformer decoder uses N learnable action queries and the displacement z_t+N − z_t to reconstruct the N intervening actions.The displacement is injected into each query through Adaptive Layer Normalization, imposing an action-grounded constraint over longer intervals.Source: Multi-Step Action Decoding · p. 4

Joint Optimization and End-to-End Training

  • Delta-JEPA jointly trains latent forward prediction and action reconstruction, using no frozen encoders, stop-gradient branches, or distribution-matching regularizers.The action reconstruction loss is balanced by λ > 0.Source: Joint Optimization and End-to-End… · p. 4

Experiments

Delta-JEPA achieves the highest mean planning success across four continuous-control environments. Displacement-based action decoding improves planning over endpoint concatenation, while results also show sensitivity to the action-reconstruction weight and decoding target.

Experimental Setup

  • Evaluation covers Push-T, Reacher, Cube, and Two-Room against LeWorldModel, Sub-JEPA, and PLDM.Methods were trained from scratch for 50 epochs and independently evaluated over three seeds; the protocol used 50 validation and 500 test trajectories per environment.Source: Experimental Setup · pp. 4–5, Planning Performance · p. 5

Planning Performance

  • Delta-JEPA achieves the highest mean planning success rate across all four environments.The gains include 15.14 percentage points over the strongest baseline on OGB-Cube, 6.27 points over PLDM on Two-Room, and 4.54 points over LeWM on Push-T; Reacher's margin is smaller.Source: Table 1 · p. 5, Planning Performance · p. 5

Action Reconstruction Weight.

  • On Push-T, λ = 0 nearly collapses and λ = 0.1 remains weak; performance is higher and relatively stable across a reasonable range, peaks at λ = 50.0, and degrades when λ is excessive.Figure 3 plots mean success over three runs and highlights the peak.Source: Action Reconstruction Weight · p. 5, Figure 3 · p. 6

Displacement-Based Action Decoding.

  • Using Δz_t improves planning success on all four environments over decoding from concatenated endpoints, with the largest gain on Push-T (+12.60 points).Two-Room gains +4.07 points; Reacher and OGB-Cube show smaller but consistent improvements.Source: Displacement-Based Action Decoding · pp. 5–6, Table 2 · p. 6

LDAD Decoding Target.

  • On Reacher, decoding raw actions performs best, joint-position deltas are comparable and outperform finger-position deltas, and concatenating the two proxies adds no improvement.Table 3 compares LDAD decoding targets.Source: LDAD Decoding Target · p. 6, Table 3 · p. 6
  • Figure 4 visualizes the evolution of the learned latent space on Push-T using PCA on 2,000 latent representations.Source: Figure 4 · p. 6

Latent Diversity and Collapse Prevention

Delta-JEPA’s latent representations become more diverse during training and show clearer action-conditioned transitions than LeWM in Two-Room. The section also describes probes of physical states, state changes, and task-relevant visual structure.

Latent Diversity and Collapse Prevention

  • Delta-JEPA achieves progressively more diverse latent representations, expanding from an initially compact region into more discernible structures.The authors interpret this trend as mitigating representation collapse and learning increasingly discriminative features.Source: Latent Diversity and Collapse Prevention · p. 6

Action-Sensitive Latent Dynamics

  • In Two-Room, Delta-JEPA trajectories with nearby initial states separate as their action-conditioned rollouts diverge, while LeWM trajectories are more scattered and less temporally organized.Figure 5 visualizes two trajectories across training epochs; color intensity marks temporal progression.Source: Action-Sensitive Latent Dynamics · pp. 6–7, Figure 5 · p. 7
  • With 512 Two-Room histories fixed while varying the final action, Delta-JEPA produces separated action-wise mean responses, whereas LeWM’s overlap near the origin.Larger action magnitudes generally induce larger predicted shifts for Delta-JEPA; Figure 6 plots predicted displacements relative to the zero-action prediction.Source: Action-Sensitive Latent Dynamics · p. 7, Figure 6 · p. 7

Physical and State-Delta Probing

  • Physical-state probes decode task-specific attributes from latent states, while state-delta probes test whether latent displacements preserve physical-change direction and magnitude.The Two-Room probe results are reported in Tables 4 and 5; lower MSE and higher r indicate better quality or alignment.Source: Physical and State-Delta Probing · pp. 7–8, Table 4 · p. 8, Table 5 · p. 8

Task-Relevant Attention Patterns

  • The learned representations preserve physically meaningful, object-centric visual structure across environments.The attention analysis focuses on task-relevant structure, including agent and object cues.Source: Task-Relevant Attention Patterns · p. 8

Conclusion

Delta-JEPA uses latent-difference action decoding to train a reconstruction-free world model. Across four continuous-control tasks, it improves planning over JEPA-based and representation-learning baselines.

Conclusion

  • Across four continuous-control tasks, Delta-JEPA improves planning over JEPA-based and representation-learning baselines.Ablations favor displacement-based decoding over endpoint concatenation, and analyses indicate action-sensitive, physically meaningful transition structure.Source: Conclusion · p. 8
  • Delta-JEPA reconstructs actions from latent displacements to discourage collapse and encourage distinguishable transitions for planning.Its simple objective combines latent prediction with action reconstruction.Source: Conclusion · p. 8
  • The conclusion presents latent-difference supervision as an effective principle for compact, collapse-resistant world models for planning.Source: Conclusion · p. 8

A Additional Probe and Attention Results

The appendix provides diagnostic results that extend the paper’s probing and attention analyses to additional environments and encoder layers.

Appendix A.1

Appendix A.2

Appendix A.3

A.1 Physical State Probing

This appendix section extends physical-state probing beyond the Two-Room results, evaluating whether latent states preserve physical information across three environments.

A.1 Physical State Probing

  • Frozen-encoder linear and MLP probes predict ground-truth physical quantities from latent states in Push-T, DMC Reacher, and OGB-Cube.The tables cover controllable agent, robot, and object states.Source: §A.1 · p. 10
  • Lower MSE and higher Pearson correlation r indicate that the learned representation preserves more physical information.Source: §A.1 · p. 10
  • Tables 6–8 report the physical-state probing results for the three environments.Source: §A.1 · p. 10

A.2 State-Delta Probing

State-delta probing tests whether latent displacements preserve physical changes between consecutive observations. The evaluation spans three tasks and several kinds of agent, robot, end-effector, and object motion.

A.2 State-Delta Probing

  • The probe predicts physical state change, ∆xt = xt+1 −xt, from latent displacement, ∆zt = zt+1 −zt.This tests whether the transition representation preserves the direction and magnitude of environmental changes.Source: §A.2 · p. 10
  • The evaluation covers Push-T, DMC Reacher, and OGB-Cube, spanning agent, robot, end-effector, and object motion.Source: §A.2 · p. 10
  • Rather than predicting state variables from the current latent state, each probe predicts the difference between consecutive states from their latent difference.Source: §A.2 · p. 10

A.3 Layer-Wise Attention Specialization

In OGB-Cube scenes, intermediate encoder layers emphasize different task-relevant entities: layer 5 focuses on the target cube, while layer 7 emphasizes the robotic gripper.

A.3 Layer-Wise Attention Specialization

  • The attention maps compare two intermediate layers of the same encoder on OGB-Cube scenes.The visualization examines how attention shifts across the visual hierarchy.Source: §A.3 · p. 10
  • Layer 5 primarily attends to the target cube, whereas layer 7 places stronger emphasis on the robotic gripper.Warmer colors in Figure 8 indicate higher attention weights.Source: §A.3 · p. 10, Figure 8 · p. 11
  • The observed layer differences indicate that task-relevant entities are progressively organized across layers rather than represented by one undifferentiated saliency pattern.Source: §A.3 · p. 10
Loading 2606.31232v1…