☰
VLA-JEPA:用潜空间预测为机器人操作注入未来感知
2026/10/17 5:20:25 网站建设 项目流程

VLA模型在真实机器人操作里有一个很典型的问题:短程任务表现惊人,比如单步抓取、对齐、插拔,成功率能跑到很高;但一旦任务变成“把抽屉里的螺丝刀拿出来,再拧开面板螺丝”,模型就开始犯一些让人哭笑不得的低级错误——要么忽略中间步骤,要么在某个步骤上反复横跳。我前段时间在一个模拟器环境里做桌面整理任务,VLA基线的单步成功率能到85%,但四步以上的长程任务成功率直接掉到40%以下。问题不在动作精度,而在模型根本不理解“接下来世界会变成什么样”。

这个痛点是VLA-JEPA这类思路的出发点:让VLA模型在潜空间里学会预测未来,而不是仅仅做一个“看见什么就输出什么动作”的反应式策略。JEPA(Joint Embedding Predictive Architecture)是潜空间预测的代表性框架,把它嵌进VLA,等于给策略模型配了一个“脑补模块”,让它能基于当前观测和动作意图,先在抽象表征空间里推演几步未来,再决定怎么动手。这篇文章就把我复现和改造VLA-JEPA的完整过程、架构设计、训练细节和踩坑记录写透,适合正在做机器人操作策略、多模态模型或者世界模型方向的同学参考。

1. 先说说VLA模型在真实操作里最容易掉链子的地方

远程操作数据训练的VLA策略,本质上是行为克隆的升级版——把视觉和语言输入直接映射到动作输出。这套范式在“看到什么就执行什么”的任务上确实好用,但它的短板也同样来自这里。

1.1 为什么“看见就做”的端到端策略不够用

端到端VLA把策略建模成一个条件概率分布,训练目标是最大化当前动作序列的条件似然。问题在于,这个条件只包含过去和当前的观测,不包含未来的任何信息。模型每走一步都是“盲人摸象”:它不知道这一步动作会让环境变成什么状态,也不知道如果这一步走错了,后面还有没有补救机会。

在短程任务里,这种缺陷被数据量和模型容量掩盖了——因为抓取一类的动作,未来几步的不确定性本身就很低。但长程任务本质上是序列决策问题,最优动作不仅依赖当前状态,还依赖对状态转移的动态建模。没有未来预测能力的策略,就像下棋只看当前一步的棋手,局部棋形能算清楚,全局十步之后必崩。

我做过一个对照实验:同样的数据量下,单纯VLA在“拿起杯子放到托盘”这类两步任务上还凑合,但一旦加入“先移开障碍物再拿杯子”的条件,模型就开始出现重复抓着同一个位置不放、或者跳过移障碍物直接尝试穿透障碍抓取的行为。这不是动作学习没学好,而是策略层面完全缺少“执行这个动作之后会发生什么”的常识性判断。

1.2 动作为主、理解为辅:VLA训练目标的结构性偏科

还有个容易被忽略的点:大多数VLA训练目标是动作空间上的回归或扩散损失,视觉和语言表征的训练是隐式的——它们只通过动作监督的梯度间接调整。这导致视觉编码器学到的是“与动作相关的表面特征”,而不是“任务相关的结构化表示”。

举个例子,单纯用动作监督训练出来的视觉表征,对“物体还在原位”和“物体被碰倒但视觉上差异很小”这两种状态往往不敏感,因为两种状态对应的最优动作可能恰好相似。但世界模型恰恰需要对这类状态差异敏感——它要能回答“如果我推了一把,这个杯子的潜状态会变成什么样”的问题。

这就是VLA-JEPA要解决的第一个结构性矛盾:动作监督负责“怎么做”,而潜世界模型负责“世界如何演化”,两者是互补的。前者是策略优化目标,后者是动力学感知目标。

2. JEPA凭什么适合做VLA的“潜世界模型”

JEPA不是新鲜概念,早期相关工作已经把“联合嵌入预测”用在视频表示学习上:将一帧图像编码成潜表示,再用潜表示预测未来帧的潜表示。它的核心想法是:预测不一定要在像素空间做,在嵌入空间做同样能学到丰富的时空结构,而且效率高得多。

2.1 像素空间预测为什么昂贵且容易学歪

三维视觉世界里的稠密预测确实能学到不少东西,但代价很高。逐像素重建损失会把模型的容量消耗在纹理、光照、背景这类与决策无关的细节上。一个教具场景里,背景窗帘的晃动和机械臂末端运动在像素层面同样“重要”,但前者对任务决策毫无意义。像素空间预测的另一个问题是分布外漂移:一旦测试场景光照、颜色分布和训练数据有差异,重建目标本身就会变得不稳定,模型预测出来的“未来帧”可能看起来很真实,但语义完全不对。

潜空间预测天然规避了这两个问题。它不要求模型还原像素,只要求模型在抽象特征层面捕捉那些“会被任务相关状态变化影响的维度”。光照变了,潜表示几乎不动;物体被移动了,潜表示显著变化。这种稀疏但语义密集的预测目标,更适合作为机器人策略的世界模型底座。

2.2 潜空间预测的三个关键设计取舍

JEPA架构里有几个直接影响训练稳定性的设计决策,我展开说一下。

第一个是非对称编码器结构。输入上下文用在线编码器(online encoder)编码,预测目标帧用目标编码器(target encoder)编码,两者内部参数共享或部分共享,但目标编码器的梯度被截断或做动量更新。这是整个架构的关键——如果两个编码器同时反向传播更新,模型很容易找到一条捷径:直接把预测头和编码器一起坍缩到常数输出,loss创新低但什么都没学到。对称设计基本必炸,非对称是必须项。

第二个是预测目标的选择。同一帧观测序列里,用第k帧潜表示预测第k+b帧潜表示是典型的自监督预测任务。这个任务可以通过采样多个预测步长来构造,b小的时候模型学到的是精细的动态转换,b大的时候学到的是跨步骤的趋势变化,两者都重要。

第三个是特征空间的约束。潜表示本身必须携带足够的信息,否则预测任务没有意义。常用的手段是方差-协方差正则化:约束潜表示各维度的方差不能太低(防止塌缩成常数),并约束各维度之间尽量去相关(保证表示容量)。这两个约束配合非对称编码器,基本能保证潜空间既不会坍缩,也不会退化到逐像素重建。

2.3 与其他世界模型路线的对比

机器人学习里不是没有其他世界模型方案,比如基于扩散模型的视频预测器、基于自回归Transformer的离散潜token预测。它们和JEPA的差别在“预测粒度”和“训练代价”上。

扩散视频预测是在像素或潜patch级别做生成式预测,效果很惊艳,但采样成本高,单步预测耗时几十到几百毫秒,很难在线闭环用。离散token预测需要训练额外的tokenizer,离散化过程会丢掉部分连续控制所需的高频信息,对动作细节敏感的操作任务不太友好。JEPA的连续潜空间预测介于两者之间:不做生成式重建,只做判别式特征预测,速度快一个数量级,特征信息量足够支撑动作头的策略学习。

当然,JEPA也有自己的问题——它的潜表示不直接可解释,缺少像素级落地,但这个问题在VLA的框架里并不致命,因为潜预测的结果不是给用户看的,而是给动作头消费的,后文会讲到具体怎么消费。

3. VLA-JEPA整体架构:动作头之外多了一个“想象分支”

VLA-JEPA的整体思路可以概括为:在常规VLA主干的旁边,加一条潜世界模型分支,共享视觉编码器,预测未来的潜状态,然后把预测结果反哺给动作头。这样架构上只增加了一个分支,但策略模型从“纯反应式”变成了“带前瞻能力的反应式”。

3.1 共享编码器与双分支解耦

我用的是双分支共享编码器的结构,数据流是:图像序列和语言指令先经过共享视觉-语言融合编码器,得到当前步的潜上下文特征h_t;然后潜世界分支接收h_t和候选动作序列,输出未来潜状态预测ψ_{t+1}到ψ_{t+k};动作分支同时接收h_t和ψ预测结果,基于融合后的上下文输出当前步动作。

共享编码器的好处是省了一大笔参数和计算量,同时保证两个分支看到的是同一套特征表示。但共享不是免费的——潜世界分支希望编码器关注“状态转移”相关的特征,动作分支希望编码器关注“动作监督”相关的特征,两者不完全一致。实际训练中我发现,如果两个分支的梯度平等累积到编码器,编码器表征会被拉向某种折中,效果反而不如让编码器更多服务于动作分支、潜世界分支只用它的中间特征。

所以我最后的做法是:编码器主体由动作分支主导微调,潜世界分支的预测头独自训练,只通过特征层面和编码器交互,不再把潜预测loss直接回传到编码器最底层。这个解耦设计让两个分支各司其职,稳定性明显提升。

3.2 潜预测头如何构造训练目标

潜预测头的输入是h_t和一个动作条件向量。这里动作条件向量不是真实动作的唯一选择——训练时用真实动作的embedding,推理时用当前策略输出的动作embedding。预测头内部就是一个多层MLP,中间层带LayerNorm和残差连接,输出维度对齐到潜空间的特征维度。

训练目标用的是余弦相似度损失,不是L2距离。原因很简单:L2距离会在潜空间里把所有维度一视同仁,但潜表示各维度对任务的重要性不一样,L2容易让预测头去迎合高方差维度而忽略低方差但语义敏感的维度。余弦相似度只看方向不看模长,对特征空间的尺度变化鲁棒得多。加上前面提到的方差-协方差正则化,基本不需要额外的损失项来防坍缩。

构造样本的时候,我会从同一条轨迹里随机采样两个时间步间隔,间隔长度在1到16帧之间均匀分布。一开始全用短间隔,等预测头在短间隔上收敛后再逐步加入长间隔样本,类似课程学习,收敛速度比一开始就混着长间隔快不少。

3.3 动作分支如何消费未来潜特征

这是整套架构里最容易被做砸的环节。动作头怎么用预测出来的潜特征,直接决定潜世界模型有没有真正帮到策略。

我先试过最朴素的做法:把ψ_{t+k}拼接到h_t后面,塞给动作头。效果有提升,但不明显,而且预测头出现轻微误差时,拼接特征反而会干扰动作输出。后来改成了“差分特征 + 门控融合”的方式:计算ψ_{t+k}与当前潜状态h_t之间的差分向量,再通过一个门控网络决定差分特征对动作输出的贡献权重。这样动作头默认先信任当前观测,只有潜预测给出明确的变化信号时,才把未来信息叠加进来。

这个设计的直观解释是:潜世界模型的价值在于给出“变化量”而非“绝对状态”——机器人控制里真正有用的是环境会如何响应本次动作,而不是未来状态长什么样。差分特征直接刻画了这种响应,门控则防止响应估计出错时污染策略输出。

推理的时候还有一种更进阶的用法:让潜世界模型在动作空间上做采样rollout,把K组候选动作序列分别喂给潜预测头,预测K步后的潜状态,然后选一个能让未来潜状态和目标特征最接近的动作序列去执行。这就是经典的“想象规划”了,后面我单独写一节。

4. 训练流程与推理闭环:先让模型学会“脑补”,再让它学会“动手”

VLA-JEPA的训练流程需要分阶段,不能上来就端到端全量训练。我试过直接联合训练,效果很差,潜预测头基本没学到东西,动作头还在正常学,两个分支的收敛速度严重不匹配。正确做法是先冷启动潜世界分支,再进入联合微调。

4.1 两阶段训练与损失权重

阶段一,潜世界分支预训练。冻结共享编码器的早期层,只训练预测头和一个浅层适配器,用自监督潜预测损失。这个阶段的数据不需要动作标签,只需要观测序列,所以可以把大量未标注的“视频数据”、模拟器随机行走数据、甚至示教视频都塞进来。数据量足够大,预测头会把该学的转移规律学个八九不离十。

阶段二,端到端微调。这个阶段同时优化动作损失和潜预测损失。我的经验是潜预测损失占比从0.3起步,逐步退火到0.1左右,让动作监督主导最终策略性能。潜预测损失在这里更像一个正则项和一个持续训练信号——它让编码器在适应动作监督的同时,不至于丢掉对状态转移的敏感性。

权重配置是个很玄学的事情,我给出我的基准配置:动作损失权重1.0,潜预测损失权重0.3,方差-协方差正则系数0.01,门控融合的L2正则0.001。这个配置在我的模拟场景上从头训基本能稳定收敛,但换个任务域建议先用小模型扫一遍权重。

4.2 推理时的想象式闭环规划

训练好的VLA-JEPA在推理时有两种玩法。第一种是轻量级玩法:正常用动作头输出动作,潜预测头实时滚动预测未来潜状态,如果预测到的未来潜状态和目标特征之间的距离快速变大,说明当前策略执行偏离了任务轨道,这时候触发重规划——把最近的观测重新喂一遍,或者切换到低层安全控制器。这个机制相当于给策略加了一个“异常事件检测器”,成本极低,效果却很实在。

第二种是重量级玩法,就是我前面说的潜空间采样rollout。每次决策时,在动作空间或者潜空间里采样N条长度为K的动作候选轨迹,分别过潜预测头得到N个未来潜状态,再按“接近目标 + 接近训练分布 + 不确定性低”三个标准打分,选最优轨迹执行第一帧或前几帧,然后重新规划。这其实就是模型预测控制的变体,只是把代价函数定义在潜空间里。

我实测的效果是:轻量级玩法对长程任务成功率的提升在10个百分点左右,重量级玩法还能再提8到10个点,但代价是单步决策延迟从原来的30ms涨到150ms以上。如果任务对实时性要求不高,重量级玩法值得上;如果机械臂频率要求很高,建议只开轻量级。

4.3 关键超参与数据配比经验

这里列几个我调整过程中最有感知的超参维度。

预测步长范围影响的是预测头的时间感知能力。短步长(1到4帧)让模型擅长预测即时动态,适合抓取前的接触类操作;长步长(8到16帧)让模型看清宏观趋势,适合长程导航和任务规划。两个都要训,课程式引入最佳。

门控融合的初始化很敏感。我一开始把门控初始化为全0.5,动作头直接震荡;改成0.1偏置后稳定很多。门控初始值越靠近0,动作头越倾向于依赖当前状态,潜预测的干扰越小,后期再让模型自己学到什么时候该信任预测。

数据配比也很关键。纯遥操作数据里含有大量“机械臂不动但观测在变”的片段,比如等待物体稳定、视觉伺服微调,这些片段对潜预测训练帮助不大,甚至有害——模型容易学会“预测未来约等于未来和现在差不多”。我会过滤掉连续30帧动作幅度都小于阈值的样本,数据量会少掉10%到15%,但潜预测损失下降更干净。

5. 复现路上的坑与排查记录

这一段写给真正动手复现的人。VLA-JEPA的架构图看着不复杂,但训练起来坑一个接一个。我把我踩过的三个比较深的坑和完整排查过程写出来,能帮你省好几周时间。

5.1 潜表示坍缩:预测头退化到常数输出

第一个大坑出现在阶段一预训练早期。训练日志显示潜预测loss在下降,但我把预测头输出的特征分布打出来一看,已经退化到各维度几乎相等且方差极小的状态——典型的表示坍缩。

排查过程是这样的:第一步检查了两个编码器的梯度流,发现目标编码器虽然截断了梯度,但在线编码器通过预测头间接把梯度传回到了共享层,目标编码器和在线编码器实际上被拉到了几乎一样的权重,破坏了非对称性。第二步检查方差正则化系数,发现0.001太小,正则力度压不过坍缩趋势。两个问题叠加导致坍缩。

修复方案是三步走:目标编码器改成完全动量更新(EMA),不接收任何梯度;方差-协方差正则系数提到0.01;预测头输出加一层单位球归一化,让预测目标天然落在有界空间里。改完之后再测,潜特征分布的方差恢复到了正常水平,预测头的预测误差才开始真正有信息量。

5.2 预测步长加大后loss的震荡问题

阶段一加长步长样本时,我踩到了第二个坑:loss出现周期性的陡增陡降,像心电图一样。排查后确认是长步长样本的预测目标方差太大,导致梯度在某些batch上异常大,把预测头参数推出稳定区域。

我先是试了梯度裁剪,把全局梯度范数裁到1.0,震荡缓解了但依然存在。后来定位到更本质的问题:预测头对不同步长共用同一套网络参数,短步长和长步长对中间层特征的需求冲突,训练到后期长步长样本会把短步长已经学好的局部动态特征覆盖掉。解决办法是把预测头改成多尺度专家结构,内部维护三组中间层,分别服务短、中、长步长,最后再接一个共享输出层。结构上复杂了一些,但训练稳定性和最终精度都明显变好。

5.3 缓存一致性bug导致训练曲线“假收敛”

第三个坑更隐蔽,出现在端到端微调阶段。训练曲线看起来一切正常,动作损失稳定下降,潜预测损失也收敛得不错,但我单独评估潜预测精度时发现预测特征和目标特征的实际余弦相似度只有0.2不到。

查了两天才找到原因:我在数据流水线里做了观测序列缓存,为了省显存,连续帧之间共享了一部分视觉特征缓存。但潜世界分支需要的是原始帧级别的精细特征,动作分支拿到的缓存特征是每隔几帧才新鲜计算的,两边的特征时间戳对不上,预测头实际上在用一个过期的上下文去预测未来的目标。这个问题不报错、不炸loss,看起来就像是正常训练,特别有迷惑性。

修复方式简单粗暴:潜世界分支用独立的数据流,不走缓存管线,每帧单独过编码器;动作分支继续用缓存。显存多花了大概15%,但训练和评估的指标终于对上了。

6. 个人体会与下一步想试的方向

VLA-JEPA这套方案最打动我的地方在于,它没有推翻VLA的基本范式,而是给反应式策略补上了一个最缺的感知维度——对未来的预期。这个过程让我越来越确信一个判断:单纯的“数据规模堆叠”很难让策略模型真正理解物理世界的动态规律,自监督的潜空间预测可能才是Tokenizer之外的另一条必经之路。

眼下我还在尝试两个方向。一个是把潜空间预测目标和人类演示中的意图标签对齐,让预测头不仅能预测“世界变成什么”,还能预测“意图是否达成”,这可能会让奖励建模更容易;另一个是给潜预测头引入不确定度量,让模型在预测信心低的时候主动触发数据采集或寻求人类干预,这是往主动学习和安全策略方向走。

最后分享一个实操心得:别一开始就在真实机械臂上测VLA-JEPA,先把潜预测单独可视化出来。找几条测试轨迹,把潜预测头的输出投影到二维可视化一下,如果预测轨迹和真实轨迹在表征空间里明显偏移,先调潜世界分支,别急着调动作头——这个分支的病根通常都在上游,动作头只是背锅的。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询