Temporal-Distance JEPA: Plan-Aware Representation Learning for Latent World Model Predictive Control
TLDR
Proposes TD-JEPA, mining temporal distance cost from reward-free trajectories to improve latent world model predictive control over JEPA.
Reasoning
Strengths include a novel method for mining temporal cost from offline logs, improving planning and representation, with strong empirical results on multiple benchmarks. Weaknesses: limited to specific environments, no explicit discussion of limitations or generalization.
Read-first score
Read-first score 45.6, weighted from topical fit, citation, graph, method, reproducibility, and recency signals. Original total remains 40.
Field roles
Rank sensitivity
Stability: volatile; rank range: 572.
Keyword Scores
Deep Analysis
Innovations
- Mines a directed temporal cost from reward-free trajectories, using same-trajectory step order as positive targets and cross-trajectory pairs as heuristic negatives, plus a rollout-consistency term that matches the planner horizon.
- Co-designs the learned cost with plan-time deployment: the mined temporal distance serves as the planning cost when progress is topological, and as a representation signal that improves latent Euclidean planning when contact geometry dominates.
- Narrows the train–plan gap for JEPA world-model planners by infusing temporal progress structure from offline logs without requiring reward signals.
Methodology
TD-JEPA retains the LeWM encoder–predictor backbone and mines a directed temporal cost from reward-free offline trajectories. Same-trajectory step order supplies positive targets, cross-trajectory pairs act as heuristic negatives, and a rollout-consistency term aligns predictions with the planner horizon. The mined cost is deployed either directly as the planning cost (topological tasks) or as a representation signal to improve Euclidean planning (contact-dominated tasks). Locked evaluation compares against LeWM and the concurrent RC-aux baseline on Two-Room, OGB-Cube, and Push-T.
Key Results
Under locked evaluation, TD-JEPA’s mined cost raises Two-Room success to 100.0% (vs. LeWM 97.4%), while the same temporally trained checkpoint improves OGB-Cube by 14.2 points over LeWM and improves Push-T; it matches or exceeds both methods on every environment. Ablations confirm that the directed head, cross-trajectory negatives, and rollout consistency each contribute.
Limitations
- The mined temporal distance cost is used directly as the planning cost only when progress is topological; in contact-dominated tasks, Euclidean planning is still employed, and the cost serves mainly as a representation signal.
- Cross-trajectory negatives are heuristic, which may limit the quality of the learned temporal cost.
- Evaluation is limited to three environments (Two-Room, OGB-Cube, Push-T) and a locked-representation protocol, which may not capture performance under full fine-tuning or broader task diversity.