机器人学会认路,却总走错方向:JEPA世界模型里被忽视的几何陷阱

想象你训练了一个机器人,让它看一眼当前的画面,再看一眼目标画面,然后自己规划一串动作去实现这个目标。这套系统工作得挺好,直到你发现一个奇怪的现象:机器人给出的"这个动作方案离目标有多远"的打分,和它实际执行后离目标的真实距离,经常对不上号。

它觉得方案A比方案B更接近目标,但真正跑起来,B反而更成功。

这不是一个小故障。这篇来自首尔大学和梨花女子大学团队的论文,揭示的正是这样一个藏在主流方法里的系统性偏差,而且它藏得足够深,深到即便你的预测模型做得完美无缺,这个偏差依然存在。

被忽视的中间层:表征空间的几何形状

要理解这个问题，得先说清楚这类系统是怎么工作的。

这类系统叫做JEPA世界模型

JEPA世界模型:一种通过在学到的表征空间里预测未来状态来实现规划的模型架构,不直接预测原始像素,而是预测压缩后的特征向量

具体流程是这样的:一个编码器把摄像头拍到的画面(不管是机器人当前的样子还是目标的样子)压缩成一串数字,这串数字叫做潜在表征。然后有一个预测器,输入当前表征和候选动作,预测执行这个动作后会变成什么样的表征。最后用一把尺子,通常是最简单的欧几里得距离,量一下预测出来的表征和目标表征差多远,差得越小,这个动作方案就打分越高。

论文里重点分析的LeWorldModel(简称LeWM)就是这么设计的。它的规划打分公式极其简单:预测终点表征减去目标表征,取平方和。

问题恰恰出在这把"尺子"上。

这把尺子测量距离时用的坐标系,是编码器训练出来的表征空间。而编码器在训练时,除了要学会准确预测未来,还有一个任务是防止表征坍缩,也就是防止所有输入都被压缩成同一个点。为了防止坍缩,LeWM用了一种叫做SIGReg的正则化手段。

SIGReg(Sketched Isotropic Gaussian Regularization,稀疏化各向同性高斯正则化):一种训练技巧,强制让学到的特征分布尽量接近一个标准的、各个方向方差相等的高斯分布,来源于LeJEPA论文提出的理论

这种正则化要求学出来的表征在每个方向上都有大致相同的"散布程度",专业说法叫各向同性。听起来是个很合理的要求,毕竟你不希望某几个维度的特征永远是死板的零值。但论文的核心发现是:这种"各个方向一视同仁"的约束,恰恰会在物理世界的坐标系里制造出一种扭曲的距离度量,而这种扭曲和实际任务需要的距离度量常常对不上。

打个比方。假设你在设计一把尺子来衡量两个城市之间的"重要性差异",而你拿到的原始数据是每个城市在东西方向和南北方向上的坐标分布。如果你强行要求这把尺子在东西方向和南北方向上的刻度密度必须相等,那么当实际城市分布是东西方向铺得很开、南北方向挤在一起时,这把尺子就会在南北方向上过度放大微小差异,而在东西方向上淡化真实的差距。如果不这么做会怎样?你会得到一把方向感失衡的尺子,南北向的一点点偏差会被读成很大的距离,东西向的巨大鸿沟反而显得微不足道。

论文用严格的数学证明了,这种失衡不是猜测,而是必然发生的结果。

数学证明:训练目标如何悄悄选中了错误的几何形状

论文的第三章给出了一个相当精巧的理论分析。核心结论是这样的:如果编码器是线性的,状态空间里的协方差矩阵是Σ,而正则化要求编码后的特征协方差是各向同性的(也就是等于某个常数乘以单位矩阵),那么反推回原始状态空间,这个欧几里得距离实际上等价于用Σ的逆矩阵加权的距离。

用大白话说就是:训练数据里方差越小的方向,在最终的距离度量里权重反而越大。

这个结论初听有点反直觉。为什么方差小的方向反而重要?

因为要让编码后的特征在各个方向上方差相同,编码器必须把训练数据里方差小的方向拉伸得更厉害,把方差大的方向压缩得更狠。拉伸得厉害的方向,单位距离对应的物理位移就越小,所以同样的物理误差在这个方向上会被放大成更大的表征距离。

论文进一步证明,这不只是一个理论上的边角案例。在一个线性高斯的训练模型里,随着训练噪声趋近于零,联合优化预测损失和SIGReg正则化项的过程,会精确地把编码器收敛到这种"逆协方差"几何,即便预测器已经做到了完美的条件均值预测,即便表征本身完全没有坍缩。也就是说,你把预测做得再准,这个几何扭曲的问题也不会自动消失。

紧接着,论文构造了一个具体的有限步规划场景,证明了这种几何扭曲会真的导致机器人选错动作。设定是这样的:训练时动作和噪声都是各向同性的,但系统动力学让状态的稳态分布变成了各向异性的椭圆形。评估时,目标点被设在了"一步能到达的范围"之外,迫使规划器必须在不同方向的误差之间做取舍。论文证明,当稳态协方差不是各向同性的、且目标点满足某个不平行条件时,这种权衡会让欧几里得任务代价选出的最优点,和用逆协方差度量选出的"最优点"完全不同。

论文还专门做了一个非线性的玩具实验来验证这个结论不是线性假设的产物。他们用一个二维系统,训练动作和过程噪声保持各向同性,但动力学本身会产生各向异性的稳态分布。用一个MLP编码器加预测器训练,结果发现:即便过程噪声缩小了30倍,留出验证集上的预测误差大幅下降,但归一化后的物理规划遗憾几乎纹丝不动,稳定在0.16左右,和理论预测的0.1607逆协方差参考值高度吻合。

这个实验结果值得多说两句。它意味着"预测得更准"和"规划得更好"这两件事,在这套系统里并不是同一回事。你可以拼命提升预测精度,规划的效果却卡在一个瓶颈上纹丝不动,因为问题根本不在预测环节,而在于表征空间的几何形状从一开始就选错了方向。

解决方案:让正则化目标学会"因材施教"

既然问题出在SIGReg强制各个方向"一视同仁"的做法上,那解决思路就很自然了:别再强求各向同性,让模型自己学出一个合理的方差分配方案。

这就是论文提出的AnisoWM和它核心的ΛReg正则化方法。

ΛReg(Lambda正则化):一种改进的正则化方法,把原本固定不变的各向同性高斯目标,替换成一个可学习的对角协方差矩阵,允许不同维度拥有不同的目标方差,但通过总方差固定和条件数上限两个约束来防止无限制的失衡

具体来说,原来SIGReg要求每个潜在维度的目标方差都等于1(各向同性),现在ΛReg允许每个维度有自己的目标方差v_i,但施加两个约束:所有方差加起来必须等于维度数D(这叫迹固定,相当于总能量守恒),以及最大方差和最小方差的比值不能超过一个预设的上限κ(这个比值叫条件数)。

条件数(condition number):矩阵最大特征值与最小特征值的比值,用来衡量一个矩阵各方向"拉伸程度"的不均衡度,数值越大表示各向异性越强

κ等于1的时候,这个约束退化回原来的各向同性目标,也就是LeWM的做法。κ大于1时,模型就获得了在不同维度分配不同重要性的自由度。这个自由度是通过训练过程学出来的,不是人工指定的,论文用一个softmax参数化方式来实现这个学习过程,并保证约束始终成立。

这里有个设计上的巧思值得注意:目标协方差矩阵Λ只在计算正则化损失的时候起作用,具体做法是把特征除以Λ的平方根再送进原来的SIGReg统计量计算。预测损失和实际规划用的还是原始的、未标准化的表征。这意味着AnisoWM并没有改动预测器结构,也没有改动规划器,改的只是训练时"塑造"表征几何形状的那一把尺子。

这就好比一个学校原本要求所有学科的考试都按同一套评分标准打分,不管是语文还是数学,及格线和满分刻度完全一致。现在换了个做法:允许不同学科按照各自的难度和区分度设计不同的评分曲线,但整个学校的总分盘子不变,而且规定任何一门课的评分区分度不能超过另一门课的两倍。如果继续用老办法会怎样?区分度天然就该大的学科(比如需要精细计算的数学)会被强行压平,而区分度天然该小的学科反而被人为拉大,最终排出来的名次和学生真实水平的差距脱节。

论文用同样的线性高斯框架证明,这个可学习的目标协方差确实能够抵消掉各向同性正则化带来的逆协方差扭曲。在一个精心构造的二维闭式解案例里,论文算出:当条件数上限κ恰好等于状态协方差的条件数c时,学出来的几何形状会精确收敛到单位矩阵的倍数,规划遗憾趋近于零。而且这个改善不是线性的,论文给出了明确的公式:相对于各向同性基线,只有当1小于κ小于c的平方时,遗憾才会降低,超过这个范围反而会变得更差。

这一点很重要,值得单独说清楚。

不是κ越大越好。

给模型越多的自由度分配方差,不代表它一定会往对的方向分配。论文在实验部分也验证了这一点:在四个环境上分别扫描κ的取值,发现最优的κ因环境而异,而且增大κ并不总是单调提升成功率,有些环境在κ超过某个值之后成功率反而下降。

四个环境上的实测效果

理论分析终归要落到实际效果上。论文选了四个视觉控制环境来验证:TwoRoom(简单的双房间导航)、Reacher(机械臂到达指定位置)、PushT(推动T形积木到目标位置)、以及OGBench-Cube(机械臂抓取并放置方块)。

所有实验都基于LeWM原本的数据集、模型架构和视觉目标规划协议,AnisoWM统一使用κ=2这个不针对具体环境调优的共享上限,和λ=0.09的正则化权重,与LeWM完全保持一致的训练设置,每个环境跑三个随机种子取平均。

成功率的对比结果是这样的:TwoRoom从LeWM的87%提升到AnisoWM的93%,Reacher从86%提升到89%,PushT从96%提升到97%,Cube从74%提升到79%。四个环境全部提升,虽然幅度有大有小,Cube这个相对最复杂的任务提升了5个百分点,是四个环境里改善最明显的。

数字之外,论文还设计了一个更直接的诊断实验来验证理论里说的"排序不一致"问题是否真的存在。他们从数据集里挑出实际执行过的动作序列,让LeWM和AnisoWM给同一批序列打分,然后看这个打分的排序和这些序列实际到达的结果排序是否一致。这个一致性用一个比例来衡量,越接近1说明模型打分越能反映真实效果。

结果显示,在规划真正用到的那个打分函数上,AnisoWM在全部四个环境里都比LeWM更一致,提升幅度从0.016到0.071不等。有意思的是,如果只看编码器本身、不考虑预测器的贡献,PushT环境的提升几乎为零,说明在这个环境里改善主要来自预测器和几何形状学习的配合,而不是编码器单方面的功劳。这提示我们,表征几何和预测器行为之间的相互作用,在不同任务里的权重并不相同。

论文还可视化了目标点周围的"低成本区域",也就是模型认为离目标很近的那些位置分布在哪。在Cube环境里,把物理距离和潜在空间打分之间的斯皮尔曼等级相关系数(一种衡量两个排序有多接近的统计量,数值从负1到正1,越接近1说明两个排序越吻合)从原来的0.15提升到了0.91,这是个相当大的跃升。PushT环境从0.65提升到0.92。这两个数字直观地说明,LeWM认为"近"的地方,很多时候物理上其实并不近,而AnisoWM学出来的表征空间基本上把这两件事对齐了。

学到的方差分配方案本身也很有意思。四个环境虽然共享同一个κ=2的上限,但学出来的方差谱完全不同:在按方差排序后,高于均值方差的维度数量,PushT是95个,Reacher是124个,TwoRoom是102个,Cube是97个,总维度是192。这说明κ只是划定了一个允许的范围,具体怎么分配是训练过程根据数据特性自主决定的,不同任务的最优分配方案本来就不一样。而在TwoRoom的训练过程追踪里,可以看到目标方差谱在条件数已经达到设定上限之后仍然在持续演化,说明触碰到边界不代表分配方案已经稳定下来。

这套方法的局限和留白

必须承认,论文的理论分析主要建立在线性高斯假设之上,尽管用非线性MLP做了补充实验,证明现象确实能延伸到非线性场景,但真实的深度神经网络编码器要复杂得多,层与层之间的非线性变换会不会引入理论没有覆盖到的新问题,这一点论文没有深入探讨。

另外,κ的取值目前还是需要针对不同环境去尝试,虽然论文里用了统一的κ=2做主实验,但后续的敏感性分析已经表明最优值因环境而异,如何在不做网格搜索的情况下自动确定合适的κ,是一个留给后续工作的开放问题。

写在后面

读这篇论文时,最触动我的不是AnisoWM本身的设计,而是那个看似朴素却被长期忽视的观察:一个用来防止表征坍缩的正则化手段,居然会悄悄决定规划器最终选择哪个动作。这两件事表面上八竿子打不着,一个管的是训练稳定性,一个管的是决策质量,但论文用扎实的数学证明它们其实是同一枚硬币的两面。

这让我想起一个更普遍的现象:很多工程系统里,某个模块的设计初衷和它实际造成的下游影响,中间往往隔着一层没人去戳破的窗户纸。SIGReg的发明者设计它是为了解决表征坍缩这个纯粹的训练问题,压根没考虑过规划这一层的几何一致性,但这个选择却在下游悄悄埋下了一个系统性偏差。这种"跨层级的隐性耦合",恐怕不只存在于这一个具体案例里。

论文里那个二维闭式解的构造也让我印象深刻,当κ精确等于状态协方差条件数c的时候,遗憾恰好归零,这个巧合般的精确对齐,某种程度上说明了这套理论框架内部的自洽性有多强。但反过来想,现实世界的状态分布哪有这么干净利落的两级结构,大多数真实数据的协方差谱是连续渐变的,这也是为什么论文最后要专门用蒙特卡洛实验去补充这个理论上过于理想化的部分。

如果这套思路继续往下走,一个自然的问题是:能不能让κ本身也根据任务的实时反馈动态调整,而不是训练前就定死的一个超参数?毕竟不同的目标点、不同的动作预算,理论上需要的方差分配可能并不完全相同。

Q&A

Q1：AnisoWM和LeWM相比,核心改进是什么?

A：AnisoWM把LeWM里固定不变的各向同性高斯正则化目标,替换成一个可学习的对角协方差目标,允许不同潜在维度拥有不同的目标方差,但通过总方差固定和条件数上限两个约束控制这个分配的自由度,预测器结构和规划器完全没有改动。

Q2：为什么各向同性正则化会导致规划出问题?

A：论文用数学证明,当编码器把训练数据映射成各向同性的特征分布时,反推回原始物理状态空间,这个距离实际上等价于用状态协方差的逆矩阵加权的距离,也就是训练数据里方差越小的方向权重反而越大,这个扭曲会导致规划器选出的最优动作和真实任务最优动作不一致。

Q3：AnisoWM在实际环境测试中效果如何?

A：在TwoRoom、Reacher、PushT、OGBench-Cube四个视觉控制环境里,AnisoWM的规划成功率全部超过LeWM,分别是93%对87%、89%对86%、97%对96%、79%对74%,同时潜在空间打分和真实任务结果的一致性也显著提升。