Temporal-Distance JEPA: Plan-Aware Representation Learning for Latent World Model Predictive Control

2026-07-28

AI Summary

This work addresses the limitations of existing Joint Embedding Predictive Architectures (JEPAs) for world modeling, which focus on myopic prediction and struggle to support effective ranking of imagined futures in multi-step planning or extract goal-progress signals from reward-free trajectories. The authors propose Temporal Distance JEPA (TD-JEPA), which preserves the LeWM encoder-predictor framework while explicitly modeling temporal progression as a deployable planning cost. TD-JEPA leverages intra-trajectory step ordering as positive samples, inter-trajectory pairs as heuristic negatives, and introduces a rollout consistency constraint aligned with the planning horizon. By jointly optimizing temporal supervision and representation learning, TD-JEPA significantly narrows the gap between world model training and planning. Experiments show it achieves 100% success on the Two-Room task (vs. 97.4% for LeWM), improves OGB-Cube by 14.2 points, and matches or exceeds LeWM and RC-aux baselines across all environments.

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 (TD-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, TD-JEPA matches or exceeds both methods on every environment. Ablations show that the directed head, cross-trajectory negatives, and rollout consistency each contribute. TD-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/TD-JEPA.

Problem

• world model
• model predictive control
• representation learning
• offline demonstration
• temporal distance

Innovation

• Temporal-Distance JEPA
• Latent World Model
• Model Predictive Control
• Offline Demonstration Learning
• Representation Learning