PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Jiaxin Bai, Jiaxuan Xiong

arXiv:2607.25337v2cs.CLcs.RO

TL;DR

JEPA world-model planners trained on reward-free demonstrations often lack an explicit temporal-progress cost for ranking imagined futures. Temporal-Distance-JEPA mines directed temporal supervision from trajectories and uses it either as a planning cost or to shape representations, matching or exceeding LeWM and RC-aux across environments while improving OGB-Cube by 14.2 points over LeWM.

  • Problem

    Reward-free JEPA world models predict dynamics but lack a trusted progress signal for ranking imagined futures during planning.

  • Method

    Temporal-Distance-JEPA mines directed temporal supervision from trajectory order, cross-trajectory negatives, and rollout consistency for cost deployment or representation shaping.

  • Results

    Temporal-Distance-JEPA matches or exceeds LeWM and RC-aux on every environment; shared Euclidean planning improves OGB-Cube by 14.2 points over LeWM and improves Push-T.

  • Takeaways & Limitations

    Temporal supervision supports topology-dominated planning through an explicit cost and contact-rich planning through improved representations with geometric costs.

  • Takeaways & Limitations

    Demonstration path length is treated as an on-manifold surrogate whose ordering reflects goal progress only when the demonstrated segment is a shortest path.

Abstract

from arXiv · show

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.

1 Introduction

Temporal-Distance-JEPA addresses the JEPA train–plan gap by mining directed temporal progress from reward-free demonstrations while retaining the LeWM predictive backbone. It deploys this signal directly for topology-dominated tasks and uses it to improve representations for geometric planning in contact-rich tasks.

  • Train–plan gap: JEPA world models predict short-horizon latent dynamics, whereas latent MPC must rank multi-step imagined futures by goal progress.LeWM and related planners typically use latent Euclidean distance as the planning cost, creating a structural mismatch between training and planning.
  • Method: Temporal-Distance-JEPA retains the LeWM encoder–predictor and SIGReg backbone while adding a directed, goal-conditioned temporal cost.The directed cost depends on the ordered pair of current and goal latents, reflecting asymmetric goal-reaching dynamics.
  • Method: Cross-trajectory negatives discourage spuriously low off-trajectory costs, while rollout consistency aligns the learned cost with the planner’s open-loop horizon.These components provide temporal supervision matched to the planning procedure.
  • Plan-time cost selection: Temporal supervision is deployed directly for topology-dominated control and used to shape representations for Euclidean planning when contact geometry dominates.The method uses dψ for topology-dominated control and latent ℓ2 for contact-rich control on the same temporally trained checkpoint.
  • Results: 14.2 points: shared Euclidean planning on the same temporally trained checkpoint raises OGB-Cube performance over LeWM and improves Push-T.Locked evaluations also report improvements on Two-Room and Reacher, while Temporal-Distance-JEPA matches or exceeds LeWM and RC-aux on every environment.

2 Related Work

Prior latent planners learn rewards, values, or embedding-based costs, whereas Temporal-Distance-JEPA remains reward-free and explicitly mines temporal progress from demonstration logs. It extends planning-aware JEPA work by aligning latent planning with directed reachability and planner-horizon supervision.

  • JEPA and LeWM: PlaNet and TD-MPC learn reward-related signals, while LeWM removes reward prediction and uses distance in a self-supervised predictive representation as the goal cost.
  • JEPA and LeWM: Temporal-Distance-JEPA remains reward-free but explicitly mines temporal order from demonstration logs to supervise progress, targeting the objective missing from prior approaches.
  • Planning-aware latent geometry: Planning-aware methods reshape JEPA geometry with quasimetric, curvature, or reachability objectives because short-horizon prediction alone may not align latent distance with multi-step reachability.
  • Planning-aware latent geometry: Unlike Value-Guided JEPA, Temporal-Distance-JEPA needs no IQL value estimator and mines reward-free step-count labels on the LeWM+SIGReg backbone.
  • Quasimetric and goal-conditioned RL: Quasimetric and goal-conditioned reinforcement-learning methods model directed or asymmetric reachability, linking temporal distances to goal-reaching costs and policies.

3 Method

Temporal-Distance-JEPA retains the LeWM encoder–predictor backbone while adding a directed temporal cost calibrated from demonstration step order, cross-trajectory negatives, and horizon-matched rollout consistency. The learned cost is the default latent-MPC objective, with latent ℓ2 used on contact-rich tasks.

  • Backbone and objectives: Temporal-Distance-JEPA keeps the LeWM prediction and SIGReg backbone, adding planning-specific supervision for the MPC cost.The additions include a directed temporal-distance objective and a horizon-matched rollout-consistency loss.
  • Planning deployment: On contact-rich tasks, the evaluation protocol plans with latent ℓ2 using the same temporally trained checkpoint rather than d_ψ.This preserves the temporally trained representation while changing the plan-time ranking cost.
  • Directed cost head: The MRN head produces a non-negative asymmetric cost d_ψ(z_s→z_g), allowing directed reachability costs that latent ℓ2 cannot express.The symmetric component captures shared configuration factors, while the asymmetric residual encodes directed reachability.
  • Temporal supervision: Same-trajectory temporal gaps τ(i,j)=j−i calibrate cost values to on-manifold progress, while cross-trajectory hinge negatives penalize spuriously small costs.The step count is treated as a progress-ranking proxy rather than a shortest-path metric.
  • Train–plan coupling: Train–plan coupling makes the temporally supervised energy the default plan-time cost for open-loop latent rollouts ranked by CEM or iCEM.For manipulation evaluation, planning can instead blend the trained energy with latent geometry as (1−α)d_ψ+αℓ2.

4 Experiments

Experiments show that Temporal-Distance-JEPA matches or exceeds LeWM and RC-aux under locked evaluation, with gains depending on whether temporal cost or latent geometry is deployed. Ablations and phase logs further show that directed temporal supervision, cross-trajectory negatives, rollout consistency, and cost aggregation shape performance across navigation and manipulation.

  • Locked comparisons: Temporal-Distance-JEPA matches or exceeds both LeWM and RC-aux on every environment under matched episodes and planner budgets.The study covers Two-Room, Reacher, Push-T, and OGBench-Cube, with primary comparisons averaged over ten independent seed runs.
  • Component ablations: Each component ablation hurts every planner setting, with the largest drop from removing the directed residual and a larger penalty from removing rollout consistency than the hinge.The results support co-designing cost form and supervision rather than relying on temporal labels alone.
  • Plan-cost regimes: Navigation favors deploying the mined temporal cost, while contact-rich manipulation favors latent ℓ2 planning after temporal representation learning.On Two-Room and Reacher, the mined cost drives gains; on Push-T and OGB-Cube, all methods plan with latent ℓ2, so improvements reflect representation learning.
  • Navigation analysis: 97.4%: locked LeWM ℓ2 success on Two-Room is exceeded by Temporal-Distance-JEPA with deployed 𝑑𝜓, whereas its ℓ2 planner falls slightly below LeWM on Reacher.This shows that navigation gains are not explained solely by a better Euclidean representation when topology matters.
  • Failure analysis: On Push-T, 𝑑𝜓 ranks temporal gaps more accurately than ℓ2 on every held-out split, yet ℓ2 plans better because local geometry supports fine pose and angle control after contact.Phase logs show that failed 𝑑𝜓 episodes are almost entirely post-contact, while ℓ2 roughly halves post-contact failures.
  • Push-T cost selection: 86.0±4.2: pure latent ℓ2 outperforms pure 𝑑𝜓 at 69.0±1.9 on Push-T, while soft and fixed blends reach 84.0±4.8 and 83.4±3.7.A hard contact gate reaches 77.6±5.2, indicating that geometry is especially important after contact.

5 Conclusion

Temporal-Distance-JEPA narrows the train–plan gap in JEPA world models by mining a directed temporal cost from reward-free demonstration logs on the LeWM encoder–predictor backbone. The mined signal serves as either a deployed planning cost for topological progress or a representation signal when contact-rich control favors latent geometry.

  • Conclusion: Temporal-Distance-JEPA mines a directed temporal cost from reward-free demonstration logs on the LeWM encoder–predictor backbone.The mined signal is designed to narrow the train–plan gap in JEPA world models.
  • Conclusion: The mined signal can be deployed as the planning cost for topological progress or used as a representation signal when contact-rich control prefers latent geometry.Locked evaluation supports both choices: deploying dψ improves Two-Room and Reacher over LeWM, while shared-ℓ2 planning uses the same temporally trained checkpoint.

A Detailed Plan-Cost Results

Table 11 reports a matched plan-cost comparison under fixed evaluation conditions, isolating the effect of changing only the planning cost on success.

  • Controlled evaluation: The comparison fixes the epoch-10 checkpoint, 50-episode manifests, environment-specific solvers, and ten independent seed runs while varying only plan cost.This setup isolates plan-cost effects; the Push-T and OGB-Cube ℓ2 entries match Temporal-Distance-JEPA’s Table 2 row.
  • Results table: Table 11 presents mean±std success over ten independent seed runs for the matched plan-cost matrix.

B Temporal-Cost Calibration

Temporal-cost calibration shows that learned d_ψ preserves temporal ordering on held-out same-trajectory pairs, while planning depends on the induced ranking rather than absolute scale.

  • Mean d_ψ rises monotonically with demonstrated step separation on held-out same-trajectory pairs.
  • The learned temporal cost need not match the true step count in absolute scale.
  • The planner uses the induced ranking of temporal costs rather than their absolute scale.

C Training and Evaluation Protocol

The protocol fixes a shared JEPA training recipe and standardized MPC planning defaults across environments. Evaluation reports variability over independent seeds with fixed checkpoints and episode manifests, while selected analyses use distinct seed designs.

  • Training: The short-window recipe uses 𝑁pred=5, window 𝑇=8, history size 3, rollout horizon 5, 𝜆roll=0.5, 𝜆td=1.0, and 𝜆sig=0.09.These settings define the temporal-distance training configuration.
  • Training: All runs use a ViT-tiny encoder with ∼15M parameters, 14-pixel patches, 224-pixel images, 192-dimensional latents, AdamW at 5×10−5 learning rate, 10−3 weight decay, bf16 precision, and 10 epochs.The directed head has hidden dimension 512 with 𝜙sym,𝜙asym ∈R128; Push-T uses effective batch 128 (32×4 accumulation).
  • Planning defaults: MPC defaults use horizon 𝐻=5, goal offset 𝐺=25, 50 episodes, and 300 candidate action sequences per CEM/iCEM iteration unless noted.Two-Room and Reacher use iCEM-30 with pure 𝑑𝜓, with terminal weights 𝑤=1 and 𝑤=0.3, respectively.
  • Planning defaults: Push-T and OGB-Cube locked comparisons use latent ℓ2 on the temporally trained checkpoint with CEM-30 and CEM-10, terminal weight 𝑤=1, matching LeWM’s geometric plan cost.RC-aux uses the same locked manifests.
  • Statistical scope: Unless noted, ± values are sample standard deviations over independent seed runs with checkpoints and episode manifests held fixed.Push-T ablations use three training seeds; paired comparisons, the cost matrix, and contact-gated Push-T use ten plan seeds, while rank correlations use independently trained checkpoints.
Loading 2607.25337v2…