基于向量量化交叉注意力的少样本机器人模仿学习扩展模型：Hybrid-ACT

《Advanced Robotics》：Hybrid-ACT: scaling robot imitation learning under few-demonstration settings with vector-quantized cross-attention

　　

在少样本条件下学习有效策略仍然是机器人模仿学习中的关键挑战。尽管基于Transformer的策略模型对长时程任务有效，但在少样本条件下，单纯扩大模型规模并不总能带来性能提升。本文提出Hybrid-ACT，一种基于Transformer架构并引入向量量化注意力（Vector-Quantized Attention，VQ-Attention）的可扩展模仿学习模型。该方法通过仅对交叉注意力中的关键向量（Key）进行量化，以连续和离散表示整合视觉观测、机器人状态与动作信息。该设计在参数量和全模型峰值内存方面均以可忽略的开销引入离散关键表示。总体而言，Hybrid-ACT相比ACT使用更少的参数和更低的峰值推理内存，同时保持相当的推理延迟。为解决码本坍缩（codebook collapse）问题，研究人员引入Gumbel-Softmax重参数化技术，该技术提升了码本利用率并稳定了更深解码器架构的训练。研究人员在RoboTwin基准的短时程和长时程双臂操作任务上评估了所提方法。Hybrid-ACT在性能上优于现有基线方法，尤其在多个长时程操作任务上取得了显著提升。研究进一步证实，随着模型规模的增大，性能增益愈发明显，证明了Hybrid-ACT在少样本条件下用于机器人模仿学习的有效性。

论文解读文章

一、研究背景与问题

从少量示范中学习高性能机器人策略是模仿学习领域的根本性挑战。近年来，基于Transformer的架构在建模长时程依赖方面展现出显著优势，基于流匹配的模仿学习方法也在复杂机器人任务上取得了强劲表现。然而，这些方法通常依赖大规模数据集。与语言建模不同，机器人学习往往面临数据获取受限的问题——跨环境、跨物体、跨任务变体收集足够多样化的机器人数据成本高昂且耗时，许多实际机器人学习场景因此被限制在相对较小的数据集上。在少样本条件下，已有研究报告指出，单纯扩大模型规模并不一定能带来成比例的性能提升，更大架构带来的经验增益在小数据场景中仍然有限。因此，如何在有限示范条件下开发出能够持续受益于模型容量增加的可持续扩展架构，仍然是一个开放问题。

现有的数据稀缺应对方法多从数据角度入手，例如通过示范生成或人类第一视角视频来增加训练数据的数量或多样性。虽然这些以数据为中心的方法是有效的，但它们通常需要额外的数据源、数据生成管道或具身迁移流程。相比之下，该研究聚焦于在固定的少样本预算下改进策略架构本身。

二、研究内容与意义

针对上述挑战，研究人员提出Hybrid-ACT，一种通过改进Transformer中的交叉注意力模块、以离散和连续表示整合多模态视觉、状态和动作信息的机器人策略架构。Hybrid-ACT建立在Transformer架构之上，将标准交叉注意力替换为向量量化注意力（VQ-Attention），使多模态记忆特征能够同时使用连续嵌入和离散潜在码进行编码。通过结合离散与连续表示，该架构在保持VQ-Attention在参数量和全模型峰值内存方面开销可忽略的前提下，提升了表示能力。此外，为实现稳定的端到端优化，研究人员引入Gumbel-Softmax重参数化技术，实现可微的码选择，从而在少样本条件下防止码本坍缩，促进码的均衡利用，并支持更深层架构的训练。实验结果表明，Hybrid-ACT在多个评估设置下持续优于竞争基线方法，且随着解码器深度增加性能进一步提升，表明通过VQ-Attention学习离散-连续表示是少样本条件下可扩展模仿学习的一个有前景的方向。

三、关键技术方法

研究人员构建了基于Transformer的模仿学习框架，从多模态观测预测机器人动作。来自多个机械臂相机的图像由预训练ResNet-18独立编码，机器人状态经线性层投影为token嵌入，视觉与状态token拼接后输入Transformer编码器，采用旋转位置编码（RoPE）而非加法位置嵌入。动作预测由Transformer解码器完成，每个解码器层中标准交叉注意力被替换为VQ-Attention，仅对Key进行向量量化，Query由零填充的动作嵌入导出，Key和Value由编码器特征导出，前馈网络采用SwiGLU。训练目标为行为克隆（BC）损失与每层向量量化（VQ）损失之和，其中VQ损失包含码本损失和承诺损失。推理阶段采用硬VQ分配和指数衰减权重的时序集成策略。实验在RoboTwin模拟基准上使用Aloha-AgileX双臂机器人平台，每个任务使用Easy设置下收集的50条示范轨迹进行训练，与ACT、扩散策略（DP）、VITA和A2A流匹配方法进行对比，每种方法使用三个随机种子训练，每个任务执行100次测试滚动。

四、研究结果

4.1 实验设置

研究人员从RoboTwin基准的50个双臂操作任务中选取5个任务进行评估。Grab Roller和Lift Pot为短时程任务，涉及双臂同步抓取和提升动作；Place Object in Basket、Put Object in Cabinet和Handover Block为长时程任务，需要物体抓取、容器或抽屉交互以及臂间交接等复杂行为。本体感受状态定义为左臂6自由度关节角、左夹爪开度、右臂6自由度关节角、右夹爪开度的拼接，维度为14。视觉观测使用三个RGB相机：一个固定第三人称头部相机和两个安装于左右腕部末端执行器的腕部相机，图像分辨率为480×640像素。动作与状态同空间，表示每步的目标关节配置。

4.2 实现细节

训练超参数和实现设置的详细信息见附录。对于基线方法，ACT和DP采用RoboTwin基准框架提供的默认配置，VITA和A2A尽可能忠实复现原始论文中的架构和训练超参数。所有训练和推理均在NVIDIA TITAN V（32GB）上进行。

4.3 结果

在五个评估任务上的成功率结果显示，Hybrid-ACT在相同实验条件下训练的方法中取得了最高的平均成功率，优于ACT、DP、VITA和A2A。特别地，Hybrid-ACT在Handover Block任务上表现出显著提升。在失败案例分析中，ACT主要失败于抓取方块，DP常在交接动作中停止，A2A在抓取失败后倾向于继续执行动作，而Hybrid-ACT和VITA有时在交接阶段失败。在Put Object in Cabinet任务上，所有方法的成功率均相对较低，由于ACT、DP、VITA和A2A均呈现类似趋势，该结果可能主要归因于任务本身的内在难度和示范质量。总体而言，Hybrid-ACT在短时程和长时程操作任务上均表现出一致的强劲性能。

4.4 解码器深度的扩展行为

研究人员通过将解码器层数从1变化到2、4、8来研究所提方法随解码器深度的扩展行为，并与ACT和去除VAE的ACT（ACT w/o VAE）进行比较。结果显示，随着解码器深度增加，Hybrid-ACT的平均成功率呈现单调提升，并在所有解码器深度下均优于ACT和ACT w/o VAE。相比之下，ACT和ACT w/o VAE均未随解码器深度增加而呈现一致的性能提升，部分解码器配置下平均成功率反而下降。这表明Hybrid-ACT在少样本模仿学习中展现出良好的深度扩展行为。此外，与ACT不同，Hybrid-ACT不需要VAE的潜在变量作为输入，而ACT w/o VAE也未能表现出类似的扩展行为，说明扩展行为的差异不能仅由VAE的有无来解释，VQ-Attention的引入可能是Hybrid-ACT扩展行为改善的关键因素。

4.5 计算效率

研究人员在解码器层数为1和8时比较了Hybrid-ACT与ACT在参数量、峰值推理内存和推理延迟方面的表现。尽管引入了VQ-Attention和SwiGLU前馈网络，Hybrid-ACT相比ACT将参数量减少了约17%，峰值推理内存减少了约14%，推理延迟仅增加约5%，与ACT相当。受控分析表明，参数量的降低主要源于Hybrid-ACT没有ACT的VAE相关模块，而VQ-Attention在参数量和全模型峰值内存方面引入的开销可忽略。从解码器层数1增加到8时，ACT的参数量和峰值推理内存分别增加约73%和50%，而Hybrid-ACT分别增加约115%和73%。这一较大的相对增长部分反映了Hybrid-ACT较小的固定组件，但在解码器层数为8时，其参数量和峰值推理内存仍低于ACT。

五、讨论与结论

研究人员指出，VQ-Attention的有效性可能依赖于有限的码本大小：小码本可能降低表示多样性并引入量化误差，而大码本会增加参数和持久缓冲区占用，并可能使码利用率更难优化。虽然Hybrid-ACT通过Gumbel-Softmax重参数化缓解了码本坍缩并保持值表示连续，但码本大小、量化误差、计算成本和策略性能之间的权衡仍需更系统的分析。评估仅限于RoboTwin中的五个模拟任务，需要在更多样化的任务和环境上进行进一步实验以更好地评估泛化能力。此外，虽然增加解码器深度扩大了模型规模并提高了任务成功率，但由此带来的推理时间增加可能成为物理机器人部署的实际约束。

研究人员提出了Hybrid-ACT，一种在少样本条件下用于双臂机器人操作的Transformer-based模仿学习模型，该模型将VQ-Attention与Gumbel-Softmax重参数化相结合。在RoboTwin基准选定任务上的实验表明，所提方法相比现有模仿学习基线取得了有利性能，在涉及多阶段协调行为的长时程任务上提升尤为显著，同时增加解码器深度提高了平均成功率，表明所提架构即使在少样本条件下也能受益于模型扩展。代码和预训练模型已公开于https://github.com/nii-vla-lab/Hybrid-ACT。