Temporal-Distance JEPA: Plan-Aware Representation Learning for Latent World Model Predictive Control
- URL: http://arxiv.org/abs/2607.25337v2
- Date: Wed, 29 Jul 2026 10:22:25 GMT
- Title: Temporal-Distance JEPA: Plan-Aware Representation Learning for Latent World Model Predictive Control
- Abstract summary: JEPAs learn world models by predicting in representation space rather than reconstructing pixels.<n>Prior JEPA planners often inherit ranking from embedding geometry, typically latent Euclidean distance.<n>We propose Temporal-Distance-JEPA, which retains the LeWM encoder-predictor backbone and mines a directed temporal cost from reward-free trajectories.
- Score: 8.39933163242427
- License: http://creativecommons.org/licenses/by/4.0/
- Abstract: Joint-Embedding Predictive Architectures (JEPAs) learn world models by predicting in representation space rather than reconstructing pixels, making them a natural backbone for latent model predictive control from offline demonstration logs. JEPA-style training optimizes short-horizon latent prediction, whereas planning requires a multi-step ranking of imagined futures by goal progress. Prior JEPA planners often inherit that ranking from embedding geometry, typically latent Euclidean distance, which arises as a byproduct of representation learning rather than as a progress cost mined from the logs. We propose Temporal-Distance-JEPA, which retains the LeWM encoder--predictor backbone and mines a directed temporal cost from reward-free trajectories: same-trajectory step order supplies positive targets, cross-trajectory pairs act as heuristic negatives, and a rollout-consistency term matches the planner horizon. The mined supervision serves two roles: as the deployed planning cost when progress is topological, and as a representation signal that improves Euclidean planning when contact geometry dominates. Under locked evaluation, deploying the mined cost raises Two-Room success to 100.0% versus LeWM's 97.4%, while shared Euclidean planning on the same temporally trained checkpoint raises OGB-Cube by 14.2 points over LeWM and improves Push-T. Against LeWM and the concurrent RC-aux baseline under locked evaluation, Temporal-Distance-JEPA matches or exceeds both methods on every environment. Ablations show that the directed head, cross-trajectory negatives, and rollout consistency each contribute. Temporal-Distance-JEPA narrows the train--plan gap for JEPA world-model planners by discovering temporal progress structure in offline logs and co-designing cost form with plan-time deployment. Code is available at https://github.com/HKBU-KnowComp/Temporal-Distance-JEPA.
Related papers
- LePlanner: An Iterative Amortized Controller For World Models [0.0]
We propose LePlanner, an amortized iterative controller that learns to construct and refine latent action sequences through a frozen world-model predictor.<n>LePlanner is trained with an arrival-and-hold objective that encourages the controller to reach the goal at the earliest feasible horizon and remain there.<n>It achieves success rates of 98% on PushT, 100% on Reacher, 100% on TwoRooms, and 92% on the OGBench Cube task.
arXiv Detail & Related papers (2026-09-12T10:01:21Z) - Orthogonal JEPA: Factorized Predictive States for Latent World Models [95.20158869896393]
We introduce method, a latent world-modeling framework based on predictive factorization.<n>We show that method can allocate redundant capacity to dominant signals while providing weak or conflicting gradients to less dominant predictive structure.<n>Experiments on controlled vision, single-cell transcriptomics, longitudinal health records, continuous control, and molecular dynamics evaluate representation quality, forecasting, planning, and long-horizon stability.
arXiv Detail & Related papers (2026-08-20T13:59:57Z) - SCALE: State-Calibrated Latent Embeddings for JEPA Planning in the Right Geometry [14.742804674497128]
Two strategies for obtaining non-collapsed representations are to inherit a pretrained feature space, and to learn an embedding end to end with anti-collapse regularization.<n>We show that planning depends not only on whether task-relevant information is present, but also on whether it shapes the geometry consumed by the planner.
arXiv Detail & Related papers (2026-08-17T08:58:55Z) - AdaJEPA: An Adaptive Latent World Model [42.0759614364807]
We propose AdaJEPA, an adaptive latent world model that performs test-time adaptation within the closed loop of model predictive control (MPC)<n>After training, AdaJEPA plans and executes the first action chunk, uses the observed next-state transition as a self-supervised adaptation signal, and replans with the updated model.<n>This closed-loop update continuously recalibrates the world model without additional expert demonstrations.
arXiv Detail & Related papers (2026-06-30T17:53:48Z) - J-LAW: Joint Localization and Actionable World Modeling via Coupled Latent Factor Graphs [5.78717443304847]
Action-conditioned world models learn compact latent dynamics for planning but ignore global metric consistency.<n>We propose J-LAW: a coupled factor graph that jointly optimize metric object poses, latent world states, and latent landmark embeddings.<n>We cast observation, action-conditioned prediction, metric odometry, pose--latent coupling, latent loop closure, and latent landmark observation as probabilistic factors in a single MAP objective.
arXiv Detail & Related papers (2026-06-27T03:39:44Z) - FF-JEPA: Long-Horizon Planning in World Models with Latent Planners [30.722155570120766]
Joint Embedding Predictive Architectures (JEPAs) have shown promising world modeling capabilities.<n>We propose Forward-Forward-JEPA (FF-JEPA), a hierarchical approach leveraging two forward dynamics models.<n>Preliminary results on PushT demonstrate that FF-JEPA successfully overcomes flat world models' long-horizon collapse.
arXiv Detail & Related papers (2026-06-08T10:16:59Z) - Predictive but Not Plannable: RC-aux for Latent World Models [44.63243875072762]
Reachability-Correction auxiliary (RC-aux)<n>We present Reachability-Correction auxiliary (RC-aux), a lightweight correction for this mismatch in reconstruction-free latent world models.<n>RC-aux keeps the world-model backbone unchanged and adds planning-aligned supervision along two axes.
arXiv Detail & Related papers (2026-05-08T05:43:33Z) - HAD: Combining Hierarchical Diffusion with Metric-Decoupled RL for End-to-End Driving [51.268878540511054]
We propose HAD, an end-to-end planning framework with a Hierarchical Diffusion Policy.<n>We show that HAD achieves new state-of-the-art performance on both NAVSIM and HUGSIM.
arXiv Detail & Related papers (2026-04-04T04:12:47Z) - ThinkJEPA: Empowering Latent World Models with Large Vision-Language Reasoning Model [53.15040805435013]
Vision-language models (VLMs) provide strong semantic grounding and general knowledge by reasoning over uniformly sampled frames.<n>We propose a VLM-guided JEPA-style latent world modeling framework that combines dense-frame dynamics modeling with long-horizon semantic guidance.
arXiv Detail & Related papers (2026-03-23T17:59:42Z) - Value-guided action planning with JEPA world models [44.84158001773079]
Building deep learning models that can reason about their environment requires capturing its underlying dynamics.<n>Joint-Embedded Predictive Architectures (JEPA) provide a promising framework to model such dynamics.<n>We propose an approach to enhance planning with JEPA world models by shaping their representation space.
arXiv Detail & Related papers (2025-12-28T20:17:49Z) - LeJEPA: Provable and Scalable Self-Supervised Learning Without the Heuristics [53.247652209132376]
Joint-Embedding Predictive Architectures (JEPAs) offer a promising blueprint, but lack of practical guidance and theory has led to ad-hoc R&D.<n>We present a comprehensive theory of JEPAs and instantiate it in bf LeJEPA, a lean, scalable, and theoretically grounded training objective.
arXiv Detail & Related papers (2025-11-11T18:21:55Z) - TD-JEPA: Latent-predictive Representations for Zero-Shot Reinforcement Learning [63.73629127832652]
We introduce TD-JEPA, which leverages TD-based latent-predictive representations into unsupervised RL.<n> TD-JEPA trains explicit state and task encoders, a policy-conditioned multi-step predictor, and a set of parameterized policies directly in latent space.<n> Empirically, TD-JEPA matches or outperforms state-of-the-art baselines on locomotion, navigation, and manipulation tasks across 13 datasets.
arXiv Detail & Related papers (2025-10-01T10:21:18Z) - AMP: Autoregressive Motion Prediction Revisited with Next Token Prediction for Autonomous Driving [59.94343412438211]
We introduce the GPT style next token motion prediction into motion prediction.
Different from language data which is composed of homogeneous units -words, the elements in the driving scene could have complex spatial-temporal and semantic relations.
We propose to adopt three factorized attention modules with different neighbors for information aggregation and different position encoding styles to capture their relations.
arXiv Detail & Related papers (2024-03-20T06:22:37Z)
This list is automatically generated from the titles and abstracts of the papers in this site.
This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.