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

The Core Problem

World models learn to predict future states from video or trajectory data, then use those predictions to plan. Most approaches (like JEPA) predict in a learned representation space rather than raw pixels—efficient and practical. But there's a mismatch: the representation geometry they learn (typically measured by Euclidean distance) is optimized for one-step prediction accuracy, not for ranking which imagined futures bring you closer to a goal. This gap hurts planning performance.

The Solution: Mining Temporal Cost from Data

TD-JEPA fixes this by explicitly learning a planning cost from the offline demonstration logs themselves. Instead of hoping Euclidean distance in the learned embedding will correlate with progress, the authors mine temporal structure directly from trajectories.

The method works in three parts:

1. Temporal ordering as signal: Within a single trajectory, if step A comes before step B, that order is a positive signal—A should be "closer" to the goal than B by the planning cost. This is reward-free (no external labels needed).

2. Cross-trajectory negatives: Pairs of states from different trajectories that are temporally unordered become negative examples. The cost should not rank them consistently, since they come from different demonstrations.

3. Rollout-consistency regularization: The learned cost is trained to match the planner's horizon—roughly, the multi-step cost should align with how far ahead the planner looks.

The authors add a "directed head" (a small learned network) that outputs this mined cost alongside the standard JEPA representation learning. At planning time, this cost either becomes the primary objective (when goal progress is topological/abstract) or it can improve the standard Euclidean planning (when contact/geometry matters).

Results and Impact

Under locked evaluation (consistent seeds and hyperparameters), TD-JEPA beats the baseline JEPA variant (LeWM) on every environment tested. Two-Room reaches 100% success versus 97.4%. On OGB-Cube, it gains 14.2 points over LeWM in absolute terms. It also matches or exceeds a concurrent baseline (RC-aux).

Ablations confirm each component matters: removing the directed head, the cross-trajectory negatives, or the rollout consistency each degrades performance.

The insight is practical: planner-aware supervision can be mined from unlabeled offline logs without reward labels. By aligning representation learning with the actual cost used at planning time, the authors close a known gap in latent-space world model predictive control. This matters because it's a simple, data-driven way to improve a widely-used architecture without architectural overhaul—just better supervision for an existing backbone.