RepWAM：基于表示型视觉-动作令牌化器的世界动作建模 RepWAM: World Action Modeling with Representation Visual-Action Tokenizers

RepWAM的训练分为三个主要步骤。第一步是训练表示型视觉令牌化器：给定视觉观测序列，使用视觉变换器自编码器将初始帧分成16x16的补丁，后续帧分成4x16x16的管状体，连接后通过编码器（堆叠的注意力块，帧间使用时间因果掩码，帧内使用完整空间注意力）得到视频潜在。令牌化器使用组合损失训练：重建损失包括L1损失、感知损失和对抗损失，特征对齐损失将视频潜在通过线性投影层拉向冻结的视觉基础模型在视觉观测上的平均特征。视觉令牌化器的总损失是重建损失与对齐损失的加权和，其中重建损失包含L1项、感知项和对抗项，对齐损失是投影后的视频潜在与视觉基础模型特征之间的均方误差。第二步是训练潜在动作令牌化器：冻结视觉令牌化器，训练耦合的逆动力学模型和前向动力学模型。对于连续的潜在帧，逆动力学模型将转换压缩为潜在动作令牌，前向动力学模型将动作令牌实现为传输图和残差，重构下一个视觉潜在为传输图作用于当前潜在加上残差。潜在动作令牌化器使用前向下一个潜在预测损失和后向一致性损失训练，分别评估前向重构和反向重构的准确性。第三步是训练因果世界动作模型：将世界动作建模转换为视觉-动作块上的因果生成。语言指令通过预训练文本编码器嵌入为条件令牌。给定块大小，每个块将视觉潜在和潜在动作在短时间窗口上分组。块因果掩码让每个块关注先前内容但不关注未来块。使用流匹配目标训练transformer：在每个步骤采样相同形状的高斯噪声张量和时间标量，形成线性插值和目标速度。训练损失是视觉和动作组件速度估计误差的加权和。最后，使用真实的机器人演示适应预训练的动力学到特定具身的机器人控制。

RepWAM的技术新颖性体现在三个方面。首先，引入了特征对齐损失来训练语义视频令牌化器，将视频潜在与冻结的视觉基础模型对齐，这是首次将语义对齐应用于WAM的视觉令牌化器设计。其次，提出了在语义视觉潜在空间中学习潜在动作的框架，与LAPA等方法从原始像素变化中提取动作代码不同，本文的潜在动作令牌表示语义视觉状态之间的转换，这使得动作抽象与世界模型使用的视觉令牌对齐。第三，设计了两阶段训练范式：先在语义视觉-动作令牌上预训练世界动作模型，再使用机器人演示适应到机器人控制。这与直接在机器人动作上训练或联合预测视觉和动作的基线方法不同，两阶段训练解耦了潜在动作预训练和机器人动作适应，提供了更可迁移的中间动作表示。此外，本文还系统性地探索了不同视觉令牌化器设计、潜在动作训练策略和视频无分类器引导对WAM性能的影响，为未来的研究提供了有价值的见解。

从独立分析的角度，RepWAM存在几个可以改进的弱点。首先，两阶段训练范式虽然提供了更好的解耦，但也增加了训练复杂度和调参难度。改进方向可以是探索端到端联合训练，同时保持语义对齐和潜在动作学习的优势。其次，冻结视觉基础模型进行特征对齐限制了视觉令牌化器的表达能力，特别是在需要学习新颖语义表示的任务中。改进方向可以是使用可学习的视觉基础模型或探索更具适应性的对齐策略。第三，潜在动作令牌化器的前向动力学模型使用传输图和残差来重构下一个视觉潜在，这个设计虽然直观，但可能限制了复杂状态转换的表达能力。改进方向可以是探索更复杂的动力学建模方法，如基于神经ODE的连续动力学或基于图神经网络的关系建模。第四，本文主要在相对短的操作序列上评估，对于需要长期规划和记忆的任务可能不足。改进方向是引入记忆机制或层次化表示来处理更长的序列。第五，视觉无分类器引导的实验表明RepViTok减弱了对CFG的依赖，但这可能限制了模型在需要强条件生成的任务中的灵活性。改进方向可以是探索更灵活的条件生成策略，平衡语义对齐和条件强度。最后，本文的方法主要依赖于视觉观察，对于复杂操作任务，其他模态（如力觉、触觉）的信息可能很重要。改进方向是多模态表示学习，将视觉与其他感官信息集成到统一的语义空间中。

作者提出的未来工作是将WAM预训练扩展到机器人领域之外的视频，利用大规模互联网视频，特别是以人为中心的视频。这可以拓宽预训练期间可用的行为和交互模式的范围。基于本文的成果，可以延伸的方向包括：探索更大规模的预训练，使用更多的视频数据和更大的模型容量，以研究语义视觉-动作令牌化的缩放规律。将RepWAM应用于更多样的机器人任务，如移动操作、组装任务和与人类的协作交互。研究跨具身的泛化能力，探索预训练的WAM如何快速适应新机器人的最小演示。集成多模态信息（如力觉、触觉、语言）到统一的语义表示中，构建更全面的世界动作模型。探索长期规划和层次化行为学习，将RepWAM扩展到需要更长时间跨度的任务。研究潜在动作的可解释性和可编程性，使人类能够理解和控制机器人的行为。探索开放世界学习，使模型能够持续从新经验和交互中学习和改进。研究安全性和可靠性，确保在真实世界部署中的可预测性和鲁棒性。最后，探索更高效的推理方法，降低计算成本和延迟，使RepWAM能够在资源受限的平台上实时运行。

论文提供了项目网页和代码仓库，但没有明确说明代码是否已经开源。从实验细节来看，论文提供了足够的实现细节，包括模型架构、超参数设置和训练过程。视觉自编码器采用12层transformer，隐藏维度768，重建权重为1、1、0.1，对齐权重为1。潜在动作IDM和FDM都是4层MLP，隐藏大小256。视觉潜在和潜在动作令牌维度分别设置为96和4。WAM骨干网络在1.3B和5B两个规模上从零开始训练：世界模型专家是30层的因果扩散transformer，隐藏维度1536或3072，动作专家共享相同深度但隐藏维度减半为768，贡献约350M额外参数。训练使用Muon优化器，峰值学习率0.01，权重衰减0.01，余弦退火和线性预热，bfloat16混合精度，梯度剪裁2.0。动作损失权重为1，对两个专家应用统一SNR采样器。1.3B模型在64个H20 GPU上训练，将episodes打包成最多200K令牌的序列；5B模型在128个H20 GPU上训练，打包最多160K令牌。数据集方面，Panda-70M用于训练表示型视觉-动作令牌化器，AgiBot约100G视频-动作潜在令牌用于预训练WAM，混合真实机器人语料库约300G令牌用于具身适应。这些详细的实现设置有助于复现。然而，复现难度仍然较高，主要原因是：需要大规模计算资源（64-128个H20 GPU），需要访问多个数据集（Panda-70M、AgiBot等），训练时间长，模型调参复杂。对于资源有限的研究者，复现完整的5B模型可能不现实，但复现1.3B模型或消融研究是可行的。总体而言，论文提供了足够的复现信息，开源代码将大大降低复现门槛。