AnyStep-WAM: Budget-Aligned Distillation and Adaptive Inference for World Action Models
Rui Wang, Xiangyu Wang, Donglin Yang, Yibo Li, Canyang Chen, Zhongrui Wang, Xiaojuan Qi
cs.RO
2026-09-28
冻结教师蒸馏可调步数WAM,再按风险选预算。三骨干去噪步少50%-85%、成功率几乎持平,一步下最高再涨12.08点。
World-action model(WAM,把未来视觉预测和动作生成绑在一起的机器人策略)推理时要反复穿过大型 diffusion transformer 做去噪。Motus 和 FastWAM 默认 10 步,LingBotVA 的串行结构更重:视频 DiT 25 步、动作 DiT 50 步。步数固定,简单伸手和精细抓取花同样的算力。接触敏感段容易欠精修,简单段白白多算。教师去噪轨迹在接触段更弯、速度更散,粗网格积分误差更大;伸手接近这类近直线段,多走几步几乎不换结果。
少步生成已经有人做。Flash-WAM 把 consistency distillation 接到 WAM,做成一步预测;AnyFlow 用学生自身的有限差分监督,支持可变步数。把 AnyFlow 接到 WAM 上,监督目标跟着学生一起抖,训练时 45.9% 的 step 梯度范数超过 1.0 的裁剪阈值。一步成功率掉到 Motus 的 13.75%、FastWAM 的 32.01%、LingBotVA 的 49.96%。操作任务对接触误差极敏感,一大步走歪就失败。
缺口有两个。同一个模型要在 1 步到多步之间都预测准;当前这段动作该用几步,既要看教师轨迹弯不弯,也要看学生在这个预算下跟教师贴不贴。
AnyStep WAM 分两段:先让一个模型覆盖多种去噪预算,再在推理时选最小够用的那个。
蒸馏把预训练 WAM 的去噪分支改成 interval-conditioned flow map。给定噪声状态和时间端点 (t, r),网络直接预测这段平均速度,位移就是时间间隔乘该速度。监督来自冻结教师用细网格解出来的轨迹,不走学生自己的导数:教师在区间两端的状态差除以时间差,就是这段积分平均速度。学生参数怎么变,这个目标都不变。
训练时从候选预算 K={1,…,10} 里抽一个 K,再抽该日程里的一段区间,对齐推理时真会遇到的跨度。四个回归项共用按维度归一化的速度误差:
权重是 λFM=0.20、λend=0.50、λint=0.30、λrec=0.20。原权重全冻,只训共享 LoRA(rank 128),三骨干可训参数分别是 286.39M、104.40M、232.91M。共享适配器覆盖全部预算,不为每个 K 复制一份头,调度器换步数时不用换权重。Lrec 对准部署路径:推理会把一步 preview 回收进多步日程,训练若不从学生自己造的状态接着走,自适应路径会和蒸馏路径错开。FastWAM 推理不去噪视频,蒸馏就只留动作损失。
调度把选步数做成风险-收益决策。风险用教师轨迹上速度方向不一致和相对变化当难度代理,再映成训练集百分位;它量的是这段去噪难不难,不是失败概率。收益用学生在各预算下、按部署时 preview-recycle 路径算出的师生速度一致度,同样是代理,不是终端动作误差。两者都只从一次一步 preview 的隐特征和速度摘要里预测:四层 Transformer,一个风险 query 加一组预算 query,两个 sigmoid 头,Smooth L1 监督。
推理流程短。先跑一次 (t,r)=(1,0) 得到 preview;调度器按预测难度在松、紧保真阈值之间插值,选最小满足所有模态阈值的 K。K=1 就直接输出 preview;否则把 preview 按原噪声重加噪到第一中间时刻,再走完剩下 K-1 步。总网络调用次数正好是 K,不访问教师,不试多个候选。调度开销 2.5–3.0 ms。接触密集阶段(抓取、倒水)会拿到更大预算,简单伸手常落在 1–2 步。
RoboTwin 2.0 上 50 个双臂任务,clean 和 randomized 各 100 trial。
| 骨干 | 满预算 SR / 步数 | AnyStep 自适应 SR / 步数 | 一步 SR:Base → AnyStep |
| Motus | 87.84% / 10 | 87.96% / 3.98(-60.2%) | 75.05% → 82.12%(+7.07) |
| FastWAM | 91.83% / 10 | 91.59% / 5.03(-49.8%) | 77.20% → 89.28%(+12.08) |
| LingBotVA | 92.24% / 25+50 | 92.22% / 3.68(-85.28% 视频/-92.64% 动作) | 73.77% → 82.71%(+8.94) |
自适应平均成功率相对满预算基线,差距不超过 0.24 个百分点。A100 单次推理延迟:Motus 1.932s→0.859s,FastWAM 0.496s→0.295s,LingBotVA 9.732s→3.247s。
同是一步预算,AnyStep 比 Flash-WAM 再高 3.19、7.58、1.30 个百分点。论文另做了一个无训练百分位门控 AnyStep+Gate,在同一套蒸馏模型上更费步、成功率更低:调度器相对它 SR 高 0.70/1.43/1.88 点,步数再少 21.8%/24.5%/23.5%。
真机是 Unitree G1D 双臂、六项任务、每任务约 20 trial。
| 骨干 | 基线 SR | AnyStep SR | 步数 | RTX 5090 延迟 |
| Motus | 65.83% | 69.17% | 10→4.06(-59.4%) | 1.868s→1.117s(1.67×) |
| FastWAM | 70.00% | 71.67% | 10→3.61(-63.9%) | 0.280s→0.122s(2.30×) |
| LingBotVA | 68.33% | 70.00% | 25/50→3.54(-85.84%) | 4.654s→0.758s(6.14×) |
Motus 消融:去掉 Lint,自适应 SR 从 87.96% 掉到 78.09%;去掉 Lend,一步 SR 从 82.12% 掉到 75.53%;去掉 Lrec,自适应 SR 只掉 1.00 点。阈值上下各挪 5 个百分位,SR 仍在 87.60%–87.93%,步数从 3.08 拉到 5.85。调度器只用 clean 数据训,randomized 上是 87.56%/4.05 步,接近全量的 87.80%/3.94。
这是给现成 WAM 加一层 LoRA 蒸馏加一个轻量调度器,原权重不动,Motus、FastWAM、LingBotVA 三套公开骨干都能接。部署侧省的是 DiT 去噪次数,调度本身可以忽略。固定一步也比 Flash-WAM 和原模型一步强,说明教师轨迹上的积分监督比 consistency 或学生有限差分更扛操作任务的大间隔更新。
定位是推理预算层,世界模型本体没换。已经在用这三套骨干、被固定步数拖延迟的人,这套更接近能直接焊上去的加速件。真机六项任务样本量不大,数字只能当方向,不能当产线指标。
论文没有单独的 Limitations 节。作者写明:风险 ρ 是相对去噪难度,不是失败概率;保真度 S 是师生速度一致度,不是终端动作误差或任务成功。调度阈值取训练集视频相似度的 10/20 分位、动作相似度的 20/80 分位,换数据域要重标定。
训练仍要跑冻结教师的细网格轨迹,再离线扫一遍各预算的 preview-recycle 路径才能训调度器。推理省了教师,蒸馏没省。候选预算锁在 1 到 10;LingBotVA 的 85% 步数降幅有一部分来自基线本来就是 25/50,和 Motus 的 10 步不是同一量级。真机每任务 50 条演示、约 20 次试验,Stack Blocks 基线只有 50–55%,多成功一两个样本就会晃。对照的自适应基线 AnyStep+Gate 是他们自己做的无训练门控,不是文献里现成的自适应策略。FastWAM 自适应平均 SR 91.59%,略低于基线 91.83%,「持平」对这条骨干是略亏。