Technical Summary: Temporal-Distance JEPA (TD-JEPA)

Problem Statement

Joint-Embedding Predictive Architectures (JEPAs) have emerged as effective backbones for latent world models, learning to predict future representations rather than reconstructing pixels. However, a structural gap exists between the training objective of JEPAs and the requirements of Model Predictive Control (MPC).

• The Train-Plan Gap: JEPA training typically optimizes short-horizon latent prediction. In contrast, planning requires ranking multi-step imagined futures based on goal progress.
• The Geometry Limitation: Existing JEPA planners (e.g., LeWM) rely on latent Euclidean distance as a de facto planning cost. While this geometry may correlate with progress in some contexts, it is a byproduct of representation learning rather than a signal explicitly mined for "steps-to-goal" or directed reachability.
• The Challenge: How can reward-free offline demonstration logs yield a temporal progress structure that narrows this gap? Furthermore, when should this mined structure be deployed directly as a planning cost versus used to shape representations for geometric planning?

Methodology: TD-JEPA

The authors propose Temporal-Distance JEPA (TD-JEPA), which retains the LeWM encoder–predictor backbone but augments it with a directed temporal cost mining mechanism. The method consists of three core components:

1. Backbone and Prediction

TD-JEPA utilizes the standard LeWM architecture: an action-conditioned encoder ($Enc_\theta$) and predictor ($Pred_\phi$) that map observations to latent states and forecast future embeddings. It preserves the SIGReg regularization to prevent representation collapse and the standard one-step prediction loss ($\mathcal{L}_{pred}$).

2. Directed Temporal Cost Mining

Instead of relying solely on Euclidean distance, TD-JEPA learns a directed, non-negative energy function $d_\psi(z_s, z_g)$ representing the cost of reaching a goal latent $z_g$ from a source latent $z_s$.

• Architecture (MRN Head): The cost head uses a Metric-Residual Network (MRN) parameterization. It decomposes the cost into a symmetric configuration term and an asymmetric residual:
  $$d_\psi(z_s, z_g) = \|\phi_{sym}(z_s) - \phi_{sym}(z_g)\|_2 + \max_k \text{ReLU}(\phi_{asym, k}(z_s) - \phi_{asym, k}(z_g))$$
  This design ensures the cost is non-negative and inherently asymmetric ($d_\psi(z_s, z_g) \neq d_\psi(z_g, z_s)$), reflecting the directed nature of goal-reaching.
• Supervision (Temporal Labels): The model is trained on offline demonstration logs without dense rewards.
  • Positives: For indices $i < j$ on the same trajectory, the temporal gap $\tau(i, j) = j - i$ serves as a proxy for steps-to-subgoal. The model regresses $d_\psi(z_i, z_j)$ toward $\tau(i, j)$.
  • Negatives: Cross-trajectory pairs act as heuristic negatives, penalizing spuriously low costs for unrelated state pairs via a hinge loss.
• Rollout Consistency: A multi-step teacher-forcing loss ($\mathcal{L}_{roll}$) is added to align the learned cost with the open-loop horizon used by the planner, reducing compounding errors during the $H$-step rollout.

3. Planning and Deployment

At test time, the planner (using Cross-Entropy Method or iCEM) searches for action sequences that minimize the terminal energy after an $H$-step latent rollout.

• Dual Deployment Strategy: The paper identifies two distinct roles for the mined signal based on task structure:
  1. Topology-Dominated Tasks (e.g., Navigation): The mined directed cost $d_\psi$ is deployed directly as the planning objective.
  2. Contact-Rich Tasks (e.g., Manipulation): The same temporally trained checkpoint is used, but the planner relies on latent Euclidean distance ($\ell_2$). Here, the temporal supervision acts as a representation signal to shape the latent space, improving the efficacy of geometric planning.

Key Contributions

1. Diagnosis of the JEPA Train-Plan Gap: The authors demonstrate that predictive JEPAs learn dynamics from logs, but planning often relies on untrained embedding geometry. They show that mining a directed temporal cost from logs significantly closes this ranking mismatch (Spearman $\rho$ increases from 0.65 for Euclidean distance to 0.95 for mined cost on Push-T).
2. Temporal-Distance Mining Framework: TD-JEPA introduces a method to extract goal-progress structure from reward-free logs using directed temporal supervision, cross-trajectory negatives, and horizon-matched rollout consistency.
3. Co-Design of Cost and Deployment: The paper provides evidence that the choice of planning cost (mined temporal vs. geometric) must be co-designed with the task structure. Temporal costs excel in topological navigation, while geometric costs remain superior for contact-rich manipulation, even when the underlying representation is trained with temporal signals.

Experimental Results

The method was evaluated on four environments: Two-Room (navigation), Reacher (reaching), Push-T (manipulation), and OGB-Cube (contact-rich manipulation). Evaluations were conducted under a "locked" protocol (fixed seeds, checkpoints, and episode manifests) to ensure fair comparison against LeWM and the concurrent RC-aux baseline.

• Navigation & Reaching: Deploying the mined cost $d_\psi$ improved success rates significantly.
  • Two-Room: Achieved 100.0% success (vs. LeWM's 97.4%).
  • Reacher: Achieved 97.0% success (vs. LeWM's 96.0%).
• Manipulation: On contact-rich tasks, the temporal training improved the latent representation, benefiting geometric planning.
  • OGB-Cube: Using the temporally trained checkpoint with $\ell_2$ planning raised success by 14.2 percentage points over LeWM (82.2% vs. 68.0%).
  • Push-T: Improved from 83.6% (LeWM) to 86.0%.
• Ablation Studies: Removing the directed residual (symmetric head), cross-trajectory negatives, or rollout consistency all resulted in performance degradation, confirming the necessity of each component.
• Failure Analysis: On Push-T, pure temporal cost ($d_\psi$) underperformed geometric cost ($\ell_2$) because failures occurred primarily during post-contact alignment, where local geometry is more critical than temporal progress ranking.

Significance and Claims

The paper claims that TD-JEPA narrows the train–plan gap for JEPA world-model planners by discovering temporal progress structure in offline logs and co-designing the cost form with plan-time deployment.

• Modest Claims: The authors do not claim that temporal costs universally replace geometric costs. Instead, they argue for a conditional deployment strategy: use the mined directed cost when progress is topological, and use it as a representation signal for geometric planning when contact geometry dominates.
• Empirical Validation: Under locked evaluation, TD-JEPA matches or exceeds both LeWM and RC-aux on every environment, demonstrating that reward-free temporal mining is a viable path to improving latent MPC without requiring reward models or online interaction.

The code is available at the provided GitHub repository.