Light-WAM：基于状态融合动作解码的高效世界动作模型 Light-WAM: Efficient World Action Models with State-Fusion Action Decoding
提出轻量级世界动作模型，通过状态融合动作解码实现高效机器人操控
前置知识
Vision Language Action (VLA) models
VLA模型是一种将视觉观察和语言指令映射到机器人动作的范式，通过视觉-语言编码器理解场景和任务，再通过动作解码器预测机器人控制信号。这类模型通常在多任务机器人数据集上训练，能够实现任务条件化的可扩展学习，是指令遵循机器人操控的主流方法。
理解VLA模型对于认识Light-WAM的定位至关重要，因为WAM是在VLA基础上扩展的，增加了未来视频预测作为额外训练目标。
World Action Models (WAMs)
WAMs通过联合训练机器人动作学习和未来视频预测来扩展VLA范式。未来视频目标提供时序训练信号，鼓励骨架编码物体运动、交互动态和任务进度，从而产生更具世界感知的视觉表示。WAMs假设对场景如何随时间变化的建模能帮助策略学习更好的表示。
Light-WAM直接建立在WAM框架上，论文的核心贡献是改进WAM的效率而非否定其基本假设，因此理解WAM的设计动机是读懂本文的前提。
Flow matching
Flow matching是一种生成模型训练技术，通过学习从噪声分布到目标分布的连续变换路径来生成数据。在视频预测中，它被用来学习从扰动潜在视频到目标未来视频的映射，比传统的扩散模型训练更稳定。Light-WAM在视频分支使用flow matching损失进行未来视频预测监督。
论文中视频分支的训练目标基于flow matching，理解这一技术有助于理解Light-WAM如何在降采样潜在空间应用未来视频监督。
LoRA (Low-Rank Adaptation)
LoRA是一种轻量级的模型适应方法，通过在预训练模型的权重矩阵上添加低秩分解的更新矩阵来实现适应。具体来说，对于权重矩阵W，LoRA学习W + deltaW = W + BA，其中B和A是低秩矩阵。这种方法大幅减少了可训练参数数量，同时保持了预训练知识的大部分。
Light-WAM使用LoRA来适应冻结的Wan视频骨架，这是实现轻量化的关键技术之一，理解它有助于理解为什么Light-WAM只有0.44B可训练参数。
研究动机
现有世界动作模型通常依赖大规模生成式架构，将未来视频预测和动作生成耦合在一起，导致巨大的GPU内存使用、训练成本和推理延迟。例如，Fast-WAM拥有6.02B可训练参数，推理延迟达到404.62ms，峰值GPU内存使用12.7GiB，训练吞吐量仅为0.49 steps/s。这些开销使得WAMs难以部署为高效的闭环机器人策略，特别是在需要快速响应的真实机器人应用场景中。虽然近期研究表明测试时未来视频生成对于强策略性能并非必需，视频预测的主要收益可能来自训练时的表示学习，但现有方法仍未充分利用这一观察。
本文的目标是本文的目标是设计一个轻量级的世界动作模型，在保持未来视频监督的表示学习优势的同时，大幅降低训练和推理成本。具体而言，Light-WAM旨在使用紧凑的视频骨架架构，通过降采样潜在空间应用未来视频监督以减少token成本，并引入直接的动作解码器避免重型生成式动作专家。目标是在LIBERO和RoboTwin 2.0等基准上取得竞争性性能，同时显著减少可训练参数、提升训练吞吐量、降低推理延迟和峰值GPU内存使用。
与已有工作不同的是，本文的独特切入角度是将WAM的训练和推理解耦，未来视频预测仅作为训练时监督信号，用于塑造视觉表示，而推理时直接从当前观察预测动作，无需生成未来视频或使用重型生成式动作专家。这与Fast-WAM等方法不同，后者虽然也认识到测试时视频rollout可能不必要，但仍主要关注推理时优化而非整体管道效率。Light-WAM关注的是改进整个WAM管道的效率，包括训练和推理两个方面，通过状态融合动作专家建立视频表示和机器人动作间的高效接口。
核心方法
Light-WAM的整体思路是保留WAM中未来视频监督的表示学习收益，但移除推理时的生成开销，从而实现高效的训练和推理。直觉上，视频预测任务能够鼓励模型学习任务相关的时序结构，包括物体运动、交互动态和任务进度，但这些表示可以在训练时习得，推理时不需要重新生成未来视频。技术路线上，Light-WAM使用紧凑的Wan2.1-T2V-1.3B视频骨架，冻结原始骨架权重，仅通过轻量级模块进行适应。为了降低视频监督成本，Light-WAM在降采样的潜在空间应用未来视频目标，而不是原始分辨率空间。最重要的是，Light-WAM引入StateFusionActionExpert，这是一个直接动作解码器，它从多个骨架层读取适应后的状态，通过学习查询池压缩密集视频token，融合后单次前向传播预测动作块。
Light-WAM的核心创新点在于提出了StateFusionActionExpert，这是一个直接动作解码器，建立了视频骨架表示和机器人动作之间的高效接口。与Fast-WAM等方法使用重型生成式动作专家不同，StateFusionActionExpert通过学习查询池将密集视频token压缩为固定宽度的动作状态，然后直接映射到动作序列，避免了迭代去噪过程。另一个关键创新是训练和推理的解耦设计：训练时在降采样潜在空间应用未来视频监督以获得表示学习收益，推理时仅使用当前观察的原始分辨率潜在进行动作预测，无需视频rollout。这种设计使Light-WAM能够在保持WAM性能优势的同时，大幅减少计算和内存开销。
方法步骤详情
Light-WAM的方法步骤包括三个主要组件。首先，视频骨架适应：使用Wan2.1-T2V-1.3B作为冻结视频骨架，通过LoRA在注意力投影和前馈投影上应用低秩更新，在选定的骨架深度8、16、24插入轻量级WAM适配器（瓶颈MLP）。给定VAE潜在输入z，patch嵌入层产生初始视频token状态H0，语言指令编码为文本上下文token，本体感受状态投影后附加，形成上下文C。其次，高效潜在空间视频联合训练：视频分支在降采样潜在空间应用flow matching监督。令z_vid_bar为空间降采样后的潜在视频，z_t_bar为其在时间t的flow matching扰动，视频分支通过损失L_video优化，其中ut是对应的flow matching目标。对于动作预测，Light-WAM从原始分辨率潜在视频取当前观察z_act，不应用视频监督使用的额外空间降采样。最后，查询瓶颈状态融合和动作解码：给定全分辨率当前观察潜在z_act，Light-WAM运行适应后的视频骨架一次，获得多层骨架状态H。StateFusionActionExpert通过学习查询Q将密集视频token状态转换为固定宽度动作状态，查询token通过多头注意力P聚合视频token，平均后归一化得到s。融合后的状态投影为h。使用步嵌入e_k，每步嵌入通过psi投影后加到融合状态h，输出头预测对应动作r_k，最终动作序列hat_A。完整训练目标结合未来视频监督和动作回归：L = L_video + lambda * L_action。
技术新颖性
Light-WAM的技术新颖性体现在多个方面。首先，它提出了训练和推理解耦的WAM设计范式，证明了未来视频预测可以作为纯训练时监督信号，而推理时可以直接预测动作，这一设计挑战了WAM必须在推理时生成未来视频的隐含假设。其次，StateFusionActionExpert提供了多层级骨架状态融合的新机制，通过学习查询池压缩密集token，允许动作解码器使用来自不同骨架层级别的信息，同时保持高效。这种设计与仅使用最终表示或暴露所有骨架激活的方法不同，它通过稀疏选择骨架层实现了效率和性能的平衡。第三，Light-WAM展示了潜在空间降采样视频监督的可行性，在减少token成本的同时保留了视频监督的表示学习收益。最后，Light-WAM的效率改进是系统性的，包括紧凑视频骨架、LoRA适应、稀疏WAM适配器、潜在缓存、空间降采样和直接动作解码等多个组件的组合，这些选择共同将可训练参数从6.02B降至0.44B（13.7倍减少），训练吞吐量提升4.25倍，推理延迟从404.62ms降至72.03ms（5.32倍减少），峰值GPU内存从12.7GiB降至4.1GiB。
实验结果
Light-WAM在LIBERO和RoboTwin 2.0基准上取得了竞争性性能，同时显著降低了训练和推理成本。在LIBERO上，Light-WAM达到97.2%平均成功率，在Spatial、Object、Goal和Long四个套件上分别获得98.2%、99.6%、97.8%和93.0%的成功率。这一成绩在无具身预训练的方法中排名第一，在所有比较方法中排名第三，仅次于LingBot-VA（98.5%）和Motus（97.7%）。Long套件仍然是最具挑战性的设置，更大的策略如Motus和LingBot-VA在此取得更高成功率，表明长视界任务仍能从更大模型容量受益。在RoboTwin 2.0上，Light-WAM在50任务上达到76.4%平均成功率（干净和随机评估的平均）。虽然这一性能低于Fast-WAM的91.9%，但考虑到Light-WAM仅使用0.44B可训练参数（相比Fast-WAM的6.02B），这一结果展示了轻量化模型在多任务场景的可用性能。Light-WAM在该基准上超过了π0（62.2%）和X-VLA（72.9%），与无具身预训练的Motus（74.9%）竞争。效率分析显示，相比Fast-WAM，Light-WAM将总训练时参数从6.73B降至1.99B（3.4倍减少），可训练参数从6.02B降至0.44B（13.7倍减少），峰值每GPU内存从70.7GiB降至43.1GiB，训练吞吐量从0.49提升至2.08 steps/s（4.25倍提升）。推理效率在RoboTwin 2.0输入上测量，Light-WAM达到72.03ms整体延迟（包含VAE编码和策略前向），峰值GPU内存4.1GiB，显著低于先前WAM方法（LingBot-VA 3214.14ms、Motus 2148.68ms、Fast-WAM 404.62ms）。分解显示其动作分支仅需2.1ms，而更大的WAM在迭代动作预测或联合视频动作生成上花费更多时间。消融实验在LIBERO-Spatial上进行，结果显示使用原始分辨率视频潜在进行联合训练可将成功率从98.2%提升至99.0%，但会显著增加训练成本，因此论文选择2倍潜在降采样以平衡性能和训练效率。增加适配器层数从3到5给出相似性能（98.2% vs 98.0%），表明额外层数在此设置中无明确增益。减少学习查询数量从16到8将成功率降至95.4%，表明查询瓶颈需要足够容量以保留操控相关的视觉信息。定性分析显示视频分支的预测比参考rollout帧更平滑（因为训练在降采样潜在空间），但仍能捕获主要运动和场景变化，表明视频分支在训练期间学习了有用的时序信息。学习查询可视化显示来自层8、16和24的注意力图在投影回图像空间时倾向于强调不同的任务相关区域，如被操控物体、夹爪和目标区域，表明选定的骨架层提供互补的视觉线索，这与融合多层级适应状态用于动作解码的设计一致。真实世界评估在IMETA Y1双臂机器人平台上进行，使用三个任务各50个演示训练，与π0.5在相同设置下比较，结果显示Light-WAM在三个任务上分别取得93%、87%和67%的成功率。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| LIBERO-Spatial | 成功率 | 98.2% | Fast-WAM 97.0% | 提升1.2% |
| LIBERO-Object | 成功率 | 99.6% | Fast-WAM 99.4% | 提升0.2% |
| LIBERO-Goal | 成功率 | 97.8% | Fast-WAM 96.6% | 提升1.2% |
| LIBERO-Long | 成功率 | 93.0% | Fast-WAM 94.8% | 降低1.8% |
| LIBERO平均 | 成功率 | 97.2% | Fast-WAM 97.0% | 提升0.2% |
| RoboTwin 2.0 | 平均成功率 | 76.4% | Fast-WAM 91.9% | 降低15.5%（但参数减少13.7倍） |
| 训练效率 | 训练吞吐量 | 2.08 steps/s | Fast-WAM 0.49 steps/s | 提升4.25倍 |
| 训练效率 | 可训练参数 | 0.44B | Fast-WAM 6.02B | 减少13.7倍 |
| 推理效率 | 推理延迟 | 72.03ms | Fast-WAM 404.62ms | 降低5.32倍 |
| 推理效率 | 峰值GPU内存 | 4.1GiB | Fast-WAM 12.7GiB | 降低3.1倍 |
局限与改进
作者承认了几个局限性。首先，在更复杂的多任务设置中，大型WAMs和具身预训练策略继续取得更高成功率，表明模型容量和大规模具身数据对于复杂操控仍然重要。这在RoboTwin 2.0结果中明显体现，Light-WAM的76.4%平均成功率低于Fast-WAM的91.9%和具身预训练的LingBot-VA的92.2%。其次，虽然论文在现有基准和真实世界任务上评估，但未在专门针对策略泛化和鲁棒性的基准（如LIBERO-Plus）上训练或测试，这可能限制对模型泛化能力的理解。第三，真实世界评估规模相对较小，仅涉及三个任务，每个任务50个演示，虽然结果积极但不足以全面评估在多样化真实场景下的鲁棒性。此外，论文使用的查询数量（每层16个）是经验选择的，缺乏系统的查询数量选择方法，这可能影响最优性能。最后，Light-WAM的轻量化设计可能在极端情况下限制了性能上限，特别是在需要极高精度的操控任务中，紧凑骨架和直接动作解码可能无法捕获足够细粒度的视觉细节和动作动力学。
独立分析的弱点
Light-WAM存在几个潜在弱点，每个都有明确的改进方向。首先，在RoboTwin 2.0上的性能（76.4%）显著低于Fast-WAM（91.9%），表明轻量化可能在复杂多任务场景牺牲性能。改进方向包括增加模型容量（如使用更大的视频骨架或增加适配器层数）、引入更复杂的融合机制（如跨层注意力）、或者采用渐进式训练策略从简单到复杂任务。其次，论文缺乏在专门针对泛化和鲁棒性的基准（如LIBERO-Plus）上的评估，这限制了对模型在分布外场景下性能的理解。改进方向包括在这些基准上进行系统性评估，并研究数据增强和鲁棒性训练技术。第三，真实世界评估规模有限（仅三个任务），不足以评估在多样化真实场景下的鲁棒性。改进方向包括扩展到更多任务和机器人平台，引入域随机化和自适应策略。第四，查询数量是经验选择的，缺乏系统的选择方法。改进方向包括研究自适应查询数量机制、基于任务复杂度动态调整查询数量，或使用自动架构搜索技术。第五，紧凑骨架和直接动作解码可能在需要极高精度的操控任务中限制性能。改进方向包括在关键层级保留更高分辨率表示、引入精细动作校正模块，或者采用分层动作规划策略。
未来方向
作者提出的未来工作方向包括纳入数据增强和鲁棒性导向训练以进一步提升Light-WAM的泛化能力。具体而言，这可以包括视觉域随机化（光照变化、遮挡、背景扰动）、物理扰动（摩擦力变化、物体属性变化）以及动作噪声注入。基于论文成果，可以延伸的有趣方向包括：探索更紧凑的视频骨架架构，如搜索最优骨架动作解码器组合；研究自适应查询机制，根据任务复杂度动态调整查询数量和注意力模式；扩展到更多样化的机器人平台和任务，特别是涉及精细操控和长视界规划的场景；研究与具身预训练的结合，探索如何在保持轻量化的同时利用大规模具身数据；开发在线适应机制，使Light-WAM能够在部署时根据新任务快速调整；探索多模态融合的扩展，如加入触觉反馈、力传感器等多模态信息以提升鲁棒性。
复现评估
Light-WAM的复现性评估显示开源情况良好。代码已在GitHub上开源（https://github.com/L1ziang/Light-WAM），便于社区检查和复现。训练在4块NVIDIA H100 GPU上进行，推理在单个NVIDIA RTX 4090 48G GPU上测量，这些硬件要求在研究环境中相对可达。实现细节在附录中详细描述，包括学习率（1e-4）、权重衰减（1e-2）、LIBERO批大小（64）、RoboTwin 2.0批大小（128）等超参数，以及适配器插入层（8、16、24）、学习查询数量（每层16个）、空间潜在降采样倍数（2倍）等架构选择。复现需要访问LIBERO和RoboTwin 2.0数据集，这些是公开可用的基准数据集。然而，复现可能需要相当的计算资源（4块H100用于训练），这对于资源有限的研究者可能是一个障碍。此外，论文未提供训练checkpoint，这意味着复现者需要从头开始训练，这可能需要大量时间。总体而言，Light-WAM的复现难度中等，代码开源和详细实现描述有助于复现，但计算资源要求可能限制了广泛复现。
论文图表
