论文

以逆动力学保留 JEPA 世界模型的不稳定模式

Preserving Unstable Modes Through Inverse Dynamics in JEPA World Models

模型评测模型行为与机制分析

摘要

机器人系统经常表现出不稳定的模式,除非通过反馈进行纠正,否则小的扰动和扰动可能会导致无限的增长。从高维视觉观察控制此类系统需要保留这些模式的表示。联合嵌入预测架构(JEPA)提供了一个自然的框架,用于从视觉数据中学习此类表示及其动态。然而,我们证明,下一步预测与防崩溃正则化相结合并不能保证保留可控的不稳定模式:当这些模式崩溃时,可以最小化训练损失,从而使学习表示的稳定性变得不可能。为了解决这个问题,我们通过动作重建目标(即逆动态损失)增强了世界模型训练,该目标鼓励控制感知表示,即保留关键控制特征的视觉表示。我们证明,精确的动作重建使得编码器在有限范围可达子空间上单射。因此,编码器不能丢弃 $H$ 步骤内的动作序列可到达的任何状态方向。此外,我们表明,随着 $H$ 的增长,有限范围可控性格拉米安的主导特征空间收敛到可控不稳定子空间。我们建立了线性系统的理论结果,并凭经验证明我们的发现扩展到非线性视觉控制任务(CartPole、Walker2D 和 PointMaze),强调了控制感知表示学习的好处。

以逆动力学保留 JEPA 世界模型的不稳定模式:论文原图
图 1:我们提出的世界模型的架构。在每一步,编码器 ($E$) 将观察值 $y_{t}$ 映射到潜在状态 $z_{t}$,预测器 ($P$) 推出以操作为条件的潜在预测 $\hat{z}_{t+1},\ldots,\hat{z}_{t+H}$。 训练损失应用于每个预测的潜在变量:1SP/MSP(单步或多步预测)将 $\hat{z}_{t+k}$ 与编码的地面实况 $z_{t+k}$ 对齐。 EP-IDM 根据 $z_{t}$ 和 $z_{t+H}$ 估计中间操作 $\hat{\mathbf{a}}_{t,H}=(\hat{a}_{t},\hat{a}_{t+1},\ldots,\hat{a}_{t+H-1})$。