时间距离JEPA:面向潜在世界模型预测控制的规划感知表征学习
Temporal-Distance JEPA: Plan-Aware Representation Learning for Latent World Model Predictive Control
July 28, 2026
作者: Jiaxin Bai, Jiaxuan Xiong
cs.AI
摘要
联合嵌入预测架构(JEPA)通过表征空间中的预测而非像素重建来学习世界模型,使其成为基于离线示范日志进行潜在模型预测控制的自然基础框架。JEPA式训练优化短时域潜在预测,而规划需要根据目标进展情况对多步想象序列进行排序。先前的JEPA规划器通常从嵌入几何结构中继承这种排序方式(典型为潜在欧几里得距离),这种排序是表征学习的副产品,而非从日志中挖掘出的进度成本。我们提出时域距离JEPA(TD-JEPA),其保留了LeWM编码器-预测器骨干结构,并从无奖励轨迹中挖掘有向时间成本:同轨迹步序提供正样本,跨轨迹对作为启发式负样本,并通过滚动一致性项匹配规划器时域。这种挖掘式监督发挥双重作用:当进度具有拓扑特性时作为部署规划成本,当接触几何结构占主导时作为提升欧几里得规划性能的表征信号。在锁定评估下,部署该项成本使Two-Room成功率达到100.0%(LeWM为97.4%),而对同一经时间训练检查点采用共享欧几里得规划时,OGB-Cube较LeWM提升14.2个点,Push-T亦获改善。与LeWM及同期RC-aux基线在锁定评估下对比,TD-JEPA在所有环境中均达到或超越两者性能。消融实验表明,有向头、跨轨迹负样本及滚动一致性项各自发挥贡献。TD-JEPA通过挖掘离线日志中的时间进展结构并协同设计成本形式与规划期部署,缩小了JEPA世界模型规划器的训练-规划差距。代码已开源:https://github.com/HKBU-KnowComp/TD-JEPA。
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 (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.