各向异性表示改进 JEPA 世界模型中的规划

潜在世界模型学习表示空间中的动作条件动态，并通常通过与目标表示的欧几里德距离对候选动作进行评分。联合训练通常会规范表示以防止崩溃，但生成的表示几何形状还决定了在规划过程中如何对终端误差进行加权。我们表明，准确的预测和非折叠表示并不能保证与任务对齐的潜在规划成本：各向同性高斯正则化可以产生一种对可行结果进行与任务成本不同的排序的几何结构。为了解决这种不匹配问题，我们引入了带有 $Λ$Reg 的 AnisoWM，它在固定轨迹和各向异性约束下用可学习的对角协方差替换了固定的各向同性高斯目标。预测目标、预测器架构和欧几里德规划器保持不变；该目标仅在训练期间使用。我们的分析描述了目标方差的预测驱动分配、其对训练分布的依赖性，以及诱导指标减少规划遗憾的条件。在四个视觉控制环境中，AnisoWM 在所有四个环境中都比 LeWorldModel 提高了规划成功率。其潜在计划成本也显示出与任务结果更好的一致性。项目网址：https://rkdrn79.github.io/AnisoWM-page/