通过挖掘离线数据中的时间进度结构,提升世界模型规划的准确性。
Temporal-Distance JEPA: Plan-Aware Representation Learning for Latent World Model Predictive Control

- 从轨迹顺序中学习有向时间成本,替代传统欧氏距离作为规划依据。
- 在Two-Room任务中成功率达100.0%,优于LeWM的97.4%。
- 适合需要高精度规划的强化学习与机器人控制场景。
联合嵌入预测架构(JEPAs)通过在表示空间中预测而非像素重建来学习世界模型,是基于离线演示日志的潜在模型预测控制的理想基础。传统JEPA训练优化短期潜变量预测,但规划需对多步想象未来按目标进展进行排序。先前的JEPA规划器常依赖嵌入几何中的欧氏距离作为排序依据,该距离是表示学习的副产品,而非从日志中挖掘出的真实进展代价。本文提出时序距离JEPA(Temporal-Distance-JEPA),保留LeWM编码器-预测器架构,并从无奖励轨迹中挖掘有向时序代价:同轨迹步骤顺序作为正样本,跨轨迹对作为启发式负样本,滚动一致性项匹配规划时域。该挖掘监督兼具双重作用:部署时作为拓扑进展代价,或作为表示信号增强接触主导情况下的欧氏规划。在锁定评估下,采用挖掘代价使Two-Room成功率提升至100.0%(相较LeWM的97.4%),同一时序训练检查点上的共享欧氏规划在OGB-Cube上提升14.2分,且改善Push-T表现。相较于LeWM与同期的RC-aux基线,在所有环境中均持平或超越。消融实验表明,有向头、跨轨迹负样本与滚动一致性均具贡献。Temporal-Distance-JEPA通过发现离线日志中的时序进展结构,并将代价形式与规划部署协同设计,缩小了训练与规划间的差距。代码已开源。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。