时间距离 JEPA:潜在世界模型预测控制的计划感知表示学习
Temporal-Distance JEPA: Plan-Aware Representation Learning for Latent World Model Predictive Control
摘要
联合嵌入预测架构(JEPA)通过在表示空间中进行预测而不是重建像素来学习世界模型,使它们成为离线演示日志中潜在模型预测控制的自然支柱。 JEPA 式训练优化了短期潜在预测,而规划则需要按目标进度对想象的未来进行多步骤排序。以前的 JEPA 规划者经常从嵌入几何图形继承该排名,通常是潜在的欧几里德距离,它是表示学习的副产品,而不是从日志中挖掘的进度成本。我们提出了 Temporal-Distance-JEPA,它保留了 LeWM 编码器 - 预测器主干,并从无奖励轨迹中挖掘定向时间成本:相同轨迹步骤顺序提供正目标,跨轨迹对作为启发式负样本,并且预测轨迹一致性项与规划器范围相匹配。挖掘监督有两个作用:当进度是拓扑时作为实际规划时的成本函数,以及当接触几何占主导地位时作为改进欧几里得规划的表示信号。在锁定评估下,部署挖掘成本将两室成功率提高到 100.0%,而 LeWM 为 97.4%,而在同一临时训练检查点上共享欧几里得规划使 OGB-Cube 比 LeWM 提高了 14.2 个点,并改进了 Push-T。相对于 LeWM 和锁定评估下的并发 RC-aux 基线,Temporal-Distance-JEPA 在每个环境上都匹配或超过这两种方法。消融表明,定向头、跨轨迹负样本和预测轨迹一致性各有贡献。 Temporal-Distance-JEPA 通过发现离线日志中的时间进度结构并与规划阶段的部署共同设计成本形式,缩小了 JEPA 世界模型规划者的训练-计划差距。代码可在 https://github.com/HKBU-KnowComp/Temporal-Distance-JEPA 获取。