跳到主要内容
返回时间线
arXiv来源发表:

机器人世界动作模型瘦身到0.88B:LIBERO成功率95.5%,每块动作57毫秒

核心概要

SmolWAM把2B参数的Cosmos Policy视频世界动作模型压到0.88B,在LIBERO的400个回合里成功率95.5%,每16动作块57毫秒、2.3 TFLOP、2.4 GiB显存,比教师快11倍,同时仍能预测场景运动。

Source-provided article image: SmolWAM: A Latent World-Action Model for Robotic Manipulation on Edge Devices
Figure 2 · arXiv

深度剖析

一个预训练的视频世界动作模型可以被大幅改小,同时保留大部分控制成功率,并继续预测未来画面。 此前的紧凑策略要么只预测动作、放弃世界模型,要么在测试时跳过未来预测;教师Cosmos Policy在LIBERO上报告的成功率在引言中为98.5%,在表格与Q1中为98.2%,每次16动作块要624毫秒、42.6 TFLOP、7.55 GiB。 在LIBERO四个套件、每任务10次试验共400个回合上,SmolWAM成功382次(95.5%,95% Wilson区间93.0–97.1),教师成功393次(98.2%,96.4–99.1),两者从相同初始状态出发;Spatial与Object套件上未检出差异,差距集中在多步的Goal(94对99)和Long(89对97)。

省下的开销主要来自把为视频生成设计的推理接口换成控制所需的接口,而不是来自换掉模型权重。 教师原本用4.8B参数的T5-XXL文本编码器、每次查询编码33帧片段的VAE,以及2B参数transformer在1,764个token上做五步去噪。 把教师未改动的权重在单帧输入上跑一步,延迟降5倍、算力降5.7倍,到126毫秒和7.5 TFLOP,筛选集上无损失;SmolWAM的trunk重设计再贡献2.2倍延迟、3.3倍算力和1.9倍显存的节省,其57毫秒中transformer占37、VAE编码占13、预处理占7。

压缩后的世界头仍然预测运动,而不是纹理。 此前不清楚把视频世界模型压小之后,未来预测能力是否还留得住。 在96个留出上下文上,SmolWAM对16步之后视图的预测把潜空间L1相对复制当前帧降低24%(0.101对0.133),并把RAFT光流到预测未来与到真实未来的端点误差从1.63像素降到1.43像素;解码后的预测保持场景、把机械臂放在接近未来位姿处,但外观模糊,LPIPS仍停留在复制帧水平(0.128对0.115)。

启示与展望

这项工作在LIBERO上评估SmolWAM,使用教师的观测管线:腕部与第三人称图像224像素、本体感觉、16动作块开环执行,seed 195,每任务10次试验共400个回合;延迟在单张A100-SXM4-40GB上以batch 1、bf16测量,包含预处理与VAE编码,文本特征对所有模型均缓存。作者提出后续应把评估扩展到RoboCasa等更多机器人基准,并最终扩展到真机场景。

作者指出评估仅限LIBERO,且降低transformer深度对长时程任务的影响不成比例,LIBERO-Long上退化最大,因此如何设计保留长时程控制性能的剪枝策略仍是开放问题。离线动作L1稳定在0.09附近,约为教师自身单步误差的十倍,而闭环成功率仍为95.5%,说明离线L1不是闭环成功的可靠代理;更长的7,500步训练降低了训练损失,但验证损失和成功率都没有改善。把ABot-M0.5的组件移植进来得到的SmolWAM-LA在筛选集上只有53.3%成功率、239毫秒,对其世界头做奖励微调在300步内把成功率压到5.0–16.7%并恶化所有图像指标。单标志消融中每个改动都落在参考值78.3%的五个回合以内,该预算下筛选集无法区分这些选择,另有六项计划中的消融未运行。