FACT: Failure-Aware Causal Training for World-Action Models
FACT: Failure-Aware Causal Training for World-Action Models
一句话总结
FACT 使用共享的扩散 Transformer 建模动作及其后果,并通过因果注意力掩码建立“先动作、后未来”的依赖关系。
训练时,动作分支对加噪后的真实动作学习去噪;与此同时,模型额外接收一份无噪声真实动作,以该动作、当前观测和语言指令为条件,预测未来视频及执行动作块后的任务进度。动作分支不能访问这份无噪声动作条件。成功轨迹监督动作、未来视频和进度三项损失;失败轨迹屏蔽动作模仿损失,但仍用实际失败结果和降低后的进度目标监督后两项。
推理时,没有真实动作作为条件,因此先去噪生成动作块,再将生成的动作块填入无噪声动作条件槽位,用于预测任务进度及可选的未来视频。如果只需要执行动作,可以在第一阶段结束;也可以生成多个候选动作,利用预测进度择优执行。
摘要
近期的世界—动作模型(WAMs)表明,将策略学习与未来预测联合训练,可以为动作生成提供物理先验。
基于视频模型的未来预测能力,许多世界—动作模型先生成未来视频,再利用逆动力学模型从中推断动作,或者将这些预测视频作为动作生成的目标条件。 在这两种情况下,世界模型主要使用成功示范进行训练,因此缺乏学习预测错误动作后果的训练驱动。
提出 FACT,一种因果世界—动作模型,它以执行的动作为条件,预测未来视频和任务进度。 这种以动作为条件的接口,使失败交互轨迹能够为动作后果提供监督:错误动作所产生的未来结果可以作为有效的预测目标,而不再被丢弃。
失败感知训练使进度预测器能够识别成功与失败两类动作结果;在推理时,还可以选择利用该预测器为采样得到的候选动作评分。
FACT 优于多种现有基线方法;随着失败数据被纳入训练,其性能进一步提升,同时还能减少模型在面对错误动作时仍倾向于预测成功未来的幻觉现象。
引言
图1:FACT首先生成一个动作,然后推演出由该动作产生的未来视频和任务进度值。由于每个未来都是从产生它的那个动作展开的,失败轨迹会在其自身的不良动作条件下直接监督未来预测分支,从而教会世界模型:错误行为实际上会导致什么。
构建通用机器人策略是机器人学习的核心目标之一。通过将动作建模与未来观测预测相结合,世界—动作模型为机器人策略学习引入时序动力学和物理先验提供了一种自然的途径。
现有的世界—动作模型通常将预测的未来作为选择动作的中间信号。一种方式是先想象未来视频,再使用逆动力学模型解码出动作,另一类方法则将预测的未来帧或潜在表示作为动作预测的条件。 关键挑战并不只是增加数据量,而是如何利用那些未能完成目标任务的交互轨迹,同时避免将它们当作不良示范来模仿。
由此产生了一个问题:我们能否利用失败交互轨迹来监督动作后果,而不将其中的失败动作作为模仿目标?
提出 FACT(Failure-Aware Causal Training,失败感知因果训练),一种将“应当模仿什么”与“应当预测什么”分离的因果世界—动作模型。
首先提出一个动作,然后以执行的动作为条件,预测未来视频和任务进度值。对于成功示范,动作、未来视频和任务进度均受到监督。 对于失败交互轨迹,动作模仿损失被屏蔽,但实际观测到的失败未来以及较低的任务进度值仍然是有效的监督信号。 这使失败数据转化为以动作为条件的后果监督,并缓解模型在面对错误动作时仍只幻想成功未来的倾向。
贡献如下:
- 我们提出 FACT,一种采用“先动作、后未来”序列的因果世界—动作模型,使模型能够利用失败轨迹进行联合训练,同时不损害动作解码能力。
- 我们引入一种采用教师强制、以动作为条件的掩码机制,将动作生成与未来预测分离,使失败动作能够为未来预测和价值学习提供监督,同时避免损害策略。
- 我们在仿真和真实世界基准上验证了该模型,结果表明:模型提高了策略的任务成功率,减少了错误动作下偏向成功的未来幻觉,并且学得的进度预测器还可用于候选动作排序。
方法
提出 FACT,一种先生成动作、再预测其后果的因果世界—动作模型。 本节首先定义任务设定与预测目标(第 3.1 节), 随后介绍一种以动作为条件的模型架构,将动作模仿与未来预测分离,使失败交互轨迹能够监督世界模型分支,而不成为模仿目标(第 3.2 节)。 最后,训练与推理策略形成一个闭环:失败交互轨迹提供关于动作后果的监督,而学得的进度估计则提供一个可选接口,用于为候选动作评分(第 3.3 节)。
问题形式化
将语言条件下的机器人操作建模为一个序列决策问题。在时刻 t,机器人接收任务指令 ℓ 和当前观测:
ot=(Itmain,ItwristL,ItwristR,st),
其中,各图像为多视角 RGB 观测,st∈Rds 为机器人的本体感知状态。目标是选择一个动作块 at:t+H∈RH×d,推进指令 ℓ 所指定的任务。标准的语言条件策略直接将这一决策建模为:
pθ(at:t+H∣ot,ℓ).(1)
与此同时,我们还要求模型预测所选动作将产生什么结果。因此,我们在式(1)的基础上加入一个以动作为条件的世界模型分支:
pθ(ot:t+K′,vt(at:t+H)∣ot,ℓ,at:t+H),(2)
其中,ot:t+K′ 表示预测时域内的未来视频观测;vt(at:t+H)∈[0,1] 表示执行动作块 at:t+H 后预测的任务进度。
为了使这一价值目标在不同长度的回合之间具有可比性,我们首先根据回合级回报定义一个归一化进度变量。设 rk 为进度奖励,Gt=∑k=1trk 为截至时刻 t 的累计进度,则归一化进度为:
pt=GTGt∈[0,1],(3)
其中,pt=0 表示任务开始,pt=1 表示任务完成。学得的价值头以 Vθ(ot,ℓ,at:t+H) 的形式进行查询,预测以动作为条件的进度目标。
模型架构
图2:FACT架构。(a) 一个共享的因果扩散Transformer对动作、价值和未来视频token进行去噪;价值和未来视频以干净的动作槽 G 为条件,而不是以带噪声的 A 为条件。在推理时,第1阶段对动作块进行去噪,该动作块填充 G,供第2阶段进行价值和未来视频预测。(b) 成功演示监督全部三个损失;失败轨迹屏蔽 Lact,但保留价值和未来视频监督,并使用降低的进度目标——失败教会的是后果,而不是行为。
以动作为条件的未来预测
模型必须从失败交互轨迹中学习未来结果,同时避免模仿失败动作。 为了将这两种影响分离,我们为世界模型分支引入一个无噪声的真实动作条件 at:t+Hgt。在这种教师强制设计下,式(2)变为:
pθ(ot:t+K′,vt(at:t+Hgt)∣ot,ℓ,at:t+Hgt).(4)
这一采用教师强制的动作 token 为未来视频预测和价值预测提供条件,而另一个独立的预测动作片段则用于动作去噪。FACT 通过以下 token 顺序实现这种分解:
z=[zrefP∥zpredA∥zgtG∥zvalueV∥zfutureI],(5)
其中,P 为观测前缀,A 为带噪声的预测动作片段,G 为采用教师强制的无噪声动作片段,V 为价值片段,I 为未来视频片段。世界模型侧的预测依赖 G,而非 A:成功示范同时监督动作预测与世界预测;失败轨迹则屏蔽动作模仿损失,但仍保留未来视频和价值监督。
训练时,一段真实动作序列有两种用法:
- 作为动作生成的学习目标 A:给真实动作加入噪声,让动作分支学习去噪、恢复动作。只有成功示范计算这项动作损失;失败轨迹的动作损失被屏蔽,避免策略学习失败行为。
- 作为预测后果的条件 G:直接输入无噪声的真实动作,让世界模型和价值分支学习:“在当前观测和任务下,执行这段动作,会看到什么未来、达到多少进度?”成功和失败轨迹都参与这部分训练。 这里的“真实动作”指数据中实际执行的动作,不代表它一定是正确的动作。失败轨迹里的坏动作同样可以作为 G,因为它与实际发生的失败结果是对应的。
推理时,就两步:
当前观测、任务指令第一阶段:动作去噪a^第二阶段:以 a^ 为条件未来视频、任务进度
也就是说,先生成动作,再预测“执行这个动作会怎样”。如果只需要执行动作,可以在第一阶段结束后直接执行,跳过第二阶段。
图3:训练和推理掩码。行和列分别表示前缀图像(P)、预测动作(A)、真实动作(G)、价值(V)和未来视频(I);单元格标记了允许的注意力(attention),其颜色代表所关注的数据段。虚线框是第1阶段的动作去噪;去噪后的动作随后填充到 G 中,供第2阶段使用(见第3.3节)。
图 3 中的注意力掩码确保了这种分离。训练时,价值 token 和未来视频 token 关注无噪声的真实动作 G,而非带噪声的预测动作 A;同时,A 无法访问 G。推理时,没有可直接使用的 G 片段:第一阶段根据 [P,A] 对 A 进行去噪,第二阶段则以第一阶段得到的无噪声动作为条件,对价值 token 和未来视频 token 进行去噪。
共享视频骨干网络上的动作适配器
同一套核心网络同时学三个任务——“该怎么动”“执行后进度如何”“执行后画面怎样”。它们有各自的输入和输出,但大量内部参数是共享的。
FACT 使用一个共享的视频扩散 Transformer 来预测动作、价值和未来视频(见图 2)。 在每个 Transformer 块的前馈网络之后,为机器人 token 添加一个轻量级动作适配器,使模型具备针对动作的专门建模能力,同时保留预训练的视频处理路径。
由于失败交互轨迹只监督未来视频和价值,这种共享骨干网络使未来结果预测损失能够影响动作生成,而不会像混合 Transformer(Mixture-of-Transformers,MoT)设计那样,将其影响局限于独立的世界模型专家。
屏蔽失败动作的模仿损失,不等于失败数据完全不影响动作生成。 对于一条失败轨迹,比如“夹爪位置偏了,导致没有抓住方块”:
- 动作模仿损失被屏蔽,所以模型不会被要求“以后也这样抓”。
- 未来视频损失仍要求模型预测“夹爪抓空,方块留在桌上”。
- 价值损失仍要求模型给这段动作预测较低的任务进度。 后两项损失会更新共享 Transformer 的参数。而动作生成也使用这些参数,因此,对失败后果的学习就有机会改善动作生成所依赖的内部表示。
训练与推理策略
失败感知的价值目标
用 Ds 表示成功示范,用 Df 表示失败交互轨迹。每个回合都标注其最终结果;对于失败回合,还会标注失败开始的时刻。我们将式(3)具体化为以下以动作为条件的进度目标:
vt(at:t+Hgt)={pt+H,clip(pt+H−λfail1fail(t+H),0,1),成功时,失败时.(6)
其中,1fail(t+H) 表示执行 at:t+Hgt 是否导致了失败。在实验中,为简便起见,我们采用均匀的进度奖励,因此对于长度为 T 的回合,式(3)简化为 pt=t/T。
失败回合:根据是否已经失败,决定要不要扣分 这里要区分整个回合最终失败和当前动作块结束时已经失败。 一条最终失败的轨迹,其前半段可能仍在正常推进。作者标注失败开始的时刻,就是为了做这个区分。
假设一条轨迹在第 70 步开始出现所标注的失败,那么可以将指示变量理解为:
1fail(t+H)={0,1,动作块结束时尚未到失败开始时刻,动作块结束时已经到达或经过失败开始时刻.
动作块结束时尚未失败:
vt=clip(pt+H,0,1)=pt+H.
虽然这个回合最终失败,但当前这段动作仍得到正常的进度标签。
动作块结束时已经失败:
vt=clip(pt+H−λfail,0,1).
此时在基础进度上减去失败惩罚,避免模型把“轨迹走得更晚”误认为“任务完成得更好”。
这一目标保留了时间进度信息,同时降低了那些进入失败状态的、以动作为条件的未来结果所对应的分数。 因此,对于一个失败动作序列,模型既要学习预测实际观测到的失败未来,也要学习预测一个较低的、以该动作为条件的进度目标。
联合去噪损失
使用相同的流匹配去噪目标来训练所有预测模态。对于目标模态 x∈{a,v,I},令 z0x 为无噪声的目标 token,z1x∼N(0,I) 为高斯噪声。采用插值:
zτx=(1−τ)z0x+τz1x,
则该模态的损失为:
Lx=Ez0x,z1x,τ[∥uθx(zτx,τ;z)−(z1x−z0x)∥22],(7)
其中,uθx 为预测的速度,z 表示式(5)中的完整 token 序列。成功示范监督动作、价值和未来预测:
LDs=waLa+wvLv+wILI.(8)
对于失败交互轨迹,我们移除动作模仿项,但保留价值和未来视频监督:
LDf=wvLv+wILI.(9)
因此,失败数据让模型学习失败动作的后果,以及这些动作所对应的较低任务进度,而不会将这些动作作为策略的学习目标。
两阶段推理
推理时,同一个 Transformer 分两个阶段运行。第一阶段对
[Pstate,Pref,Anoisy]
执行 Kdenoise 步基于欧拉法的流去噪,得到无噪声的动作块 a^t:t+H(见图 3)。如果只需要动作,FACT 在此结束,跳过世界预测。
当需要为候选动作评分或预测动作后果时,第二阶段将 a^t:t+H 放入无噪声动作条件槽位,对价值进行去噪,并可选择对未来视频潜变量进行去噪。由于不同候选动作共享前缀 P,仅生成动作的推理过程会使用前缀键值缓存。
可选的候选动作评分
在 FACT 中,候选动作排序之所以有意义,是因为进度预测器接受过以动作为条件的失败结果训练。 如果只使用成功数据训练,进度预测主要在专家动作流形上得到校准,可能会对较差动作给出过于乐观的分数。 因此,我们将价值引导的选择视为一种可选的部署接口,以及检验失败感知后果学习效果的手段,而非独立的监督来源。
在价值引导的推理中,第一阶段并行采样 N 个候选动作 {a(k)}k=1N。第二阶段为每个候选动作预测一个以该动作为条件的进度分数:
v^(k)=Vθ(ot,ℓ,a(k)),
并选择分数最高的候选动作执行,即:
a⋆=a∈a(1:N)argmaxVθ(ot,ℓ,a).(10)
这一选择规则利用价值头,对每个候选动作所隐含的未来结果进行评分,无需另行训练一个评价器(critic)。
实现细节
使用 WAN2.2-5B 初始化模型权重,将其作为视频扩散骨干网络,并使用 AdamW 训练。动作前馈网络(FFN)的学习率为 2×10−4,WAN 骨干网络的学习率为 2×10−5;损失权重为 wa=20,wv=wI=1。
动作块长度为 H=48,失败惩罚系数为 λfail=1。未来视频监督使用当前帧及四个未来时间偏移处的帧,对应:
[0,H/4,H/2,3H/4,H].
除非另有说明,推理使用 20 步基于欧拉法的流去噪。
实验
实验部分主要验证三件事:未来预测是否有助于动作生成、失败数据是否能改善策略与后果预测,以及学到的价值能否用于选择动作。
实验设置
仿真: RoboTwin 的 50 个双臂操作任务,在标准场景和域随机化场景中评估,额外加入约 1,300 条失败交互轨迹。 真实机器人: 5 个已见任务,包括叠方块、拾取方块、双臂交接、叠碗和倾倒;另设 3 个未见变体,改变物体颜色、形状及对应指令。每项结果基于 20 次测试。 对比方式: 既与 π0、π0.5、Motus 等方法比较,也逐步比较“基础 FACT → 加入失败数据 → 加入候选评分”,区分训练收益和推理时额外计算的收益。
实验设计与结果
| 实验 | 为什么这样设计 | 主要结果与结论 |
|---|---|---|
| 整体性能比较 | 检验方法能否提高实际任务成功率,以及能否适应未见变化 | 加入失败数据后,仿真成功率由 85.6% → 87.5%,接近 Motus 的 87.8%;真实已见任务由 82% → 89%,未见变体由 67% → 77%。失败训练的收益同时出现在仿真和真实环境中。 |
| 移除视频联合训练 | 检验未来预测是否确实帮助动作学习 | 仿真成功率由 85.6% 降至 81.8%,真实已见任务由 82% 降至 58%。说明未来预测对策略学习具有明显的辅助作用。 |
| 移除教师强制动作条件与因果掩码 | 检验将动作生成和后果预测分离是否必要 | 在同样不使用失败数据的条件下,真实已见任务由 82% 降至 77%,支持“用确定的动作作为未来预测条件”这一设计。 |
| 保留失败动作的模仿损失 | 检验失败数据应该如何使用 | 成功率降至 63%,明显低于正确屏蔽失败动作损失的 89%。说明失败轨迹适合监督后果,直接模仿其中的动作会损害策略。 |
| 未来预测质量比较 | 检验模型是否真正学会预测失败,减少“错误动作却预测成功”的现象 | 在成功、失败各半的 512 个留出样本上,失败未来的 PSNR 由 19.51 → 25.92,成功未来基本不变;SSIM 也呈相同趋势。可视化中,模型从幻想抓取成功转为预测实际的抓取失败。 |
| 增加失败数据量 | 检验收益是否随失败经验增加而持续出现 | 在 3 个仿真任务上,逐步使用 0%、50%、100% 的可用失败轨迹,平均成功率从 32.7% 单调提升至 57.3%。在测试范围内,更多失败数据仍能带来收益。 |
| 候选动作评分与价值诊断 | 检验价值头能否区分动作质量,以及额外计算是否值得 | 经过失败训练后,评分使真实已见任务由 89% → 92%、未见变体由 77% → 82%;仅用成功数据训练时,评分反而使已见任务从 82% 降至 79%。候选数从 1 增至 4 收益明显,继续增加则收益递减。 |
价值头学到了什么?
作者进一步观察了执行过程中的价值变化:抓取失败时价值下降,重新抓取成功后回升;在相同状态下比较不同放置位置,能够完成任务的位置获得最高分。这些诊断支持价值头对动作后果具有一定辨别能力,而非仅随时间推移增加。
效率与结论
FACT 可以只执行第一阶段生成动作。附录测得其推理延迟约为 380 ms,比 Motus 的 1,220 ms 快约 3 倍,但仍慢于 π0.5 的 47 ms;候选评分则需要额外计算。
总结
失败数据即使不参与动作模仿,也能通过共享骨干上的未来与价值监督改善策略;而候选评分的有效性,依赖于模型先学会识别失败后果。
结果支持这一机制,但尚不能说明 FACT 在所有任务上全面领先:仿真略低于 Motus,未见真实变体仍低于 π0.5,且真实测试规模较小。
总结
提出了 FACT,一种因果世界—动作模型。它颠倒了世界—动作模型通常采用的顺序,先生成动作,再预测未来视频和任务进度值。
教师强制掩码使无噪声的已执行动作成为世界模型侧所有预测的条件,从而允许失败交互轨迹监督未来预测和价值预测,同时关闭这些轨迹的动作模仿损失。
在仿真和真实世界双臂操作基准上,这种失败感知的后果建模改善了策略学习,并减少了模型在错误动作下仍倾向于预测成功未来的幻觉现象。学得的进度预测器还提供了一个可选的候选动作评分接口,使部署时能够在计算开销和性能之间进行权衡。
这种以动作为条件的世界建模视角,也为未来结合在线交互轨迹、DAgger 式纠正,以及从负面经验中进行强化学习的训练方式,提供了一个自然的接口。
局限性
本研究聚焦于失败感知的后果建模。未来,将相同的因果训练顺序扩展到更广泛的机器人数据和人类交互数据,可能进一步提升未来预测的物理合理性,以及价值头的判别能力。
此外,未来研究可以在保留相同的动作条件价值接口的前提下,用学习得到的进度估计器替换或增强当前的价值头。
最后,从 N 个候选动作中择优的策略需要以计算开销换取可靠性。当延迟要求严格时,FACT 可以采用仅生成动作的运行模式;但价值引导的选择需要对每一批候选动作额外执行一次评分。