LeWAM 如何用 JEPA 潜变量与扩散引导构建机器人世界动作模型?

LeWAM 用无解码器 JEPA 潜变量训练世界动作模型,潜变量能读出任务状态,并在策略噪声空间规划。

直接答案

LeWAM 是一个世界动作模型,它在同一个无解码器的 JEPA 潜变量上预测前向动力学、后向动力学、逆动力学和动作,四种模式端到端联合训练。它的潜变量携带与任务相关的机器人和物体状态,同时忽略视觉干扰物;闭环控制表现与同编码器、同规模的流匹配策略相当;规划时在策略头的噪声空间而非原始动作空间中搜索,能提升闭环性能。

1篇文献引用

本文由 WisPaper 驱动的搜索和论文分析生成。

像素重建为何削弱潜变量

世界动作模型同时预测动作和未来观测,其中多数模型通过像素级重建来塑造内部表示。论文认为,这迫使潜变量编码解码器可能渲染出的每一个视觉细节,无论该细节是否影响任务。在视觉复杂的场景中,模型容量被消耗在嘈杂、与任务无关的信息上,作者称这会增加下游预测的难度。

LeWAM 走的是相反路线。它建立在无解码器的联合嵌入预测架构之上,这类架构完全在潜空间中建模动力学,并彻底去掉像素解码器。论文给出的动机是:去掉解码器,也就去掉了存储那些不改变场景演化的外观细节的压力。这一选择的代价是,潜变量必须由其他目标来塑造:准确的动力学预测和稳健的动作生成。

四种预测模式共用一个双向 Transformer

LeWAM 是一个以动作为条件的双向 Transformer,用单一主干同时处理前向动力学、后向动力学、逆动力学和策略预测。编码器在严格正则化约束下直接由任务像素端到端训练,流程中没有像素解码器、分词器或像素空间损失。论文的对比表把 LeWAM 列为唯一同时具备无解码器、策略头、前向/逆向/后向动力学以及多模态单主干设计的模型。

四种模式的实际含义是:同一个潜变量既能想象接下来会发生什么,也能推断之前发生了什么,还能从两个状态之间恢复动作,并直接提出动作。由于四个目标训练的是同一份表示,潜变量并非为重建保真度优化,而是为同时服务于预测和控制而优化。

线性探针读出任务状态并忽略干扰物

对齐证据来自线性探针:对每个模型和任务,作者用岭回归从潜变量预测智能体、物体、夹爪和干扰物状态,采用交叉验证,并在留出集上评估。他们分别报告任务相关状态和任务无关干扰物上的 MSE、Pearson r 和解释方差 R²,结果在六个任务、五个随机种子上取平均。

跨任务平均来看,LeWAM 在任务相关状态上的 R² 达到 0.80,带扩散头的 LeWM 为 0.62,基于重建的世界动作模型为 0.78。它的干扰物 R² 为 0.00,与 LeWM 的 0.00 持平,远低于重建模型的 0.73。这一模式并非每个任务都一致:在 Lift 上,LeWAM 的任务状态 R² 为 0.42,低于重建模型的 0.56;而在 ToolHang 上达到 0.93,高于后者的 0.88。作者的解读是,LeWAM 平均而言比 LeWM 乃至基于重建的模型更好地编码任务相关状态,同时像 LeWM 一样忽略干扰物。

另一项检查支持同样的解读。作者在一个已训练好并冻结的 LeWAM 上训练一个普通前馈卷积解码器,因此解码器的梯度无法重塑潜变量。在 PushT 和 Robomimic Lift 的想象 rollout 中,T 形块的位置能被稳定识别,但旋转模糊,干扰物块在想象帧中几乎完全消失。论文把偶尔出现的错位干扰物视为解码器幻觉,而非潜变量内容。

控制表现与同编码器流策略持平

动作方面的结果来自一项匹配比较:LeWAM 的闭环评估与在同一编码器上训练、规模匹配的常规流匹配策略相当,同时 LeWAM 还额外提供世界模型。这一设定之所以重要,是因为它隔离出世界模型结构相对于共享同一表示的纯策略基线究竟付出了什么或带来了什么,而不是与一个无关系统比较。

论文把这一结果表述为持平,而非提升,因此可以带走的结论是:在同一编码器上加入前向、后向和逆动力学预测,并未在这些评估中降低闭环控制表现。评估在任务间取平均,并覆盖五个训练种子,逐任务曲线在附录图中单独报告。

原始动作规划让 MPC 利用模型误差

用学习到的世界模型做规划,自然会想到基于采样的模型预测控制(MPC),但论文指出,在原始动作上执行这类控制在实践中经常失败。它给出的机制是“利用”:当规划器直接在动作空间上优化时,往往找到会利用学习动力学模型中不可避免误差的方案,产生在幻觉未来下看起来很好、部署时却失败的轨迹。

LeWAM 的应对是改变决策变量。规划器不再在动作上搜索,而是在策略头的确定性噪声空间中搜索。一个方案是一串噪声向量;每个动作块在前序块之后想象出的状态上解码,预测出的潜变量被追加,同时丢弃最旧的上下文潜变量,方案最终在其最后一个想象潜变量上打分。搜索从策略的先验开始,因此每个扩散引导规划器的第一次迭代,都是对投影后噪声做 best-of-N 候选排序。

投影把候选限制在训练分布内

策略头训练时的起点是标准正态分布的采样,这些采样集中在半径等于动作维度平方根的球面上。当规划器收窄搜索时,候选会漂向范数小得多的均值,对整流流头而言,这相当于以降低的温度采样,使搜索趋向模式寻求。

投影扩散引导在解码前把每个候选映射回该球面。搜索分布本身仍保持欧氏形式,只有候选被投影。论文陈述的后果是:搜索仍在方向上进行,每个解码出的候选都是策略头训练时会产生的那类样本,而温度效应被放弃。论文报告这一投影在表 4 中做了消融;它比较的三个扩散引导规划器 DS-CEM、DS-MPPI 和 DS-CMA-ES 共享同一预算——64 个样本、3 次迭代、规划时域 5,区别只在更新规则。

噪声空间决策变量是可迁移的结果

规划方面的发现是论文内部的一项比较:用世界动作模型规划时采样原始动作,会让 MPC 利用动力学模型的误差;而在策略头的噪声空间中规划,能提升这些世界动作模型的闭环性能。该过程只在推理阶段进行,因此改变的是已训练模型的使用方式,而不是重新训练它。

作者提出两个后续方向,而非宣称工作完成:把 LeWAM 应用到真实世界数据,他们称这类数据包含大量动力学上惰性的特征;以及改进规划算法,他们称这些算法慢,且推理时需要目标的概念。这两个条件标出了当前证据的边界。因此最值得关注的结论是一个优先级排序:潜变量结果是跨任务平均的结论,存在 Lift 这样的逐任务例外;控制结果是持平而非增益;最清晰可操作的改变,是在策略头的噪声空间中带投影地规划,而不是在原始动作上规划。

本文引用的文献

1

LeWAM:一种带扩散引导 MPC 的 JEPA 世界动作模型

一个通过前向、后向、逆动力学和策略预测训练的无解码器 JEPA 世界动作模型,平均而言比 LeWM 和基于重建的世界动作模型更好地读出任务相关状态,忽略干扰物,在闭环控制中与同编码器流匹配策略持平,并在策略头噪声空间而非原始动作空间中规划时提升闭环性能。