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

📄 arXiv: 2607.25337v1 📥 PDF

作者: Jiaxin Bai, Jiaxuan Xiong

分类: cs.CL, cs.RO

发布日期: 2026-07-28

🔗 代码/项目: GITHUB

💡 一句话要点

提出TD-JEPA以优化潜在世界模型的预测控制

🎯 匹配领域: 支柱一：机器人控制 (Robot Control) 支柱二：RL算法与架构 (RL & Architecture)

关键词: 潜在模型预测控制 联合嵌入预测架构 时间距离 机器人导航 智能系统

📋 核心要点

  1. 现有JEPA方法在多步规划时，通常依赖于嵌入几何的欧几里得距离，未能有效挖掘进展成本。
  2. TD-JEPA通过从无奖励轨迹中挖掘有向时间成本，结合同轨迹和跨轨迹的样本，优化了规划过程。
  3. 在锁定评估下，TD-JEPA在Two-Room任务中成功率达到100%，相较于LeWM的97.4%有显著提升。

📝 摘要（中文）

联合嵌入预测架构（JEPA）通过在表示空间中进行预测来学习世界模型，而非重建像素，这使其成为基于离线演示日志的潜在模型预测控制的自然基础。JEPA风格的训练优化短期潜在预测，而规划则需要通过目标进展对想象的未来进行多步排序。现有JEPA规划器通常从嵌入几何中继承该排序，主要是潜在的欧几里得距离，这种距离是表示学习的副产品，而非从日志中挖掘的进展成本。我们提出了时间距离JEPA（TD-JEPA），它保留了LeWM编码器-预测器骨干，并从无奖励轨迹中挖掘有向时间成本：同轨迹的步骤顺序提供正目标，跨轨迹对作为启发式负样本，回滚一致性项匹配规划者的时间范围。挖掘的监督信号在进展是拓扑时作为规划成本使用，在接触几何主导时改善欧几里得规划。实验结果显示，TD-JEPA在多个环境中超越了LeWM和RC-aux基线。

🔬 方法详解

问题定义：论文旨在解决现有JEPA方法在多步规划中对进展成本挖掘不足的问题，导致规划效果不佳。

核心思路：TD-JEPA通过从无奖励的轨迹中挖掘有向时间成本，利用同轨迹的步骤顺序作为正目标，跨轨迹对作为负样本，提升了规划的准确性和效率。

技术框架：TD-JEPA的整体架构包括LeWM编码器-预测器骨干，挖掘时间成本的模块，以及规划阶段的回滚一致性项，确保规划的连贯性。

关键创新：TD-JEPA的核心创新在于通过挖掘时间进展结构，设计出适应于规划时间的成本形式，从而缩小了训练与规划之间的差距。

关键设计：在损失函数中，结合了正目标、负样本和回滚一致性项，确保模型在训练时能够有效学习到时间进展的结构。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

在实验中，TD-JEPA在Two-Room任务中成功率达到100%，相比LeWM的97.4%有显著提升。此外，在OGB-Cube任务中，TD-JEPA的共享欧几里得规划提升了14.2分，显示出其在多种环境下的优越性能。

🎯 应用场景

该研究的潜在应用领域包括机器人导航、自动驾驶以及其他需要高效决策的智能系统。通过优化潜在世界模型的预测控制，TD-JEPA能够提升这些系统在复杂环境中的表现，具有重要的实际价值和广泛的应用前景。

📄 摘要（原文）

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.