Abstract
这项工作介绍了mamba模仿学习(mail),这是一种新颖的模仿学习(il)架构,为最先进的(sota)变换器策略提供了一种计算高效的替代方案。基于变压器的策略由于能够处理具有固有非马尔可夫行为的人类记录数据而取得了显著成果。然而,它们的高性能伴随着大型模型的缺点,这使得有效的训练变得复杂。虽然状态空间模型(ssms)以其效率而闻名,但它们无法与变压器的性能相匹配。mamba显著提高了ssms和竞争对手对transformers的性能,使其成为il政策的一个有吸引力的替代方案。mail利用mamba作为骨干,并引入了一种形式化,允许在编码器-解码器结构中使用mamba。这种形式化使其成为一种通用的架构,既可以用作独立的策略,也可以用作更高级架构的一部分,例如扩散过程中的扩散器。
对LIBERO IL benchmark和三个真实机器人实验的广泛评估表明,mail:i)在所有libero任务中都优于transformer,ii)即使在小数据集下也能实现良好的性能,iii)能够有效地处理多模态感官输入,iv)与transformer相比,对输入噪声更具鲁棒性
Introduction
这里,当前的方法要么使用仅解码器结构[5],要么使用解码器-编码器架构[6]。这些架构中哪一个擅长通常取决于任务。变压器的性能通常伴随着难以训练的大型模型,特别是在数据稀缺的领域。处理观测序列的另一种概念是状态空间模型[12]。这些模型假设观测值(嵌入)之间存在线性关系,通常在计算上更高效。最近的方法,如选择性状态空间模型mamba[13],严格提高了状态空间模型的性能,并在许多任务中与变压器竞争。由于其在推理速度、内存使用和效率方面的特性,mamba是一个有吸引力的il策略模型
邮件可以用作独立的策略,也可以用作更高级流程的一部分,例如扩散流程中的扩散器。我们以两种变体实现邮件。在仅解码器的变体中,mail处理噪声动作和观测特征[5]以及扩散过程的时间嵌入,并输出去噪动作。
Related Works
Sequence Models.
变压器中的自我关注机制允许并行处理序列,有效地解决了rnn在顺序数据处理中的局限性[17,18,19,20]。然而,结构化状态空间模型[12,22,22,13]为变压器提供了一种有吸引力的替代方案。变压器在序列长度上按二次缩放,而结构化状态空间模型则按线性缩放[13]
最近的工作[13]依赖于关联扫描,它也允许并行计算,但还允许输入相关的可学习矩阵[13]
Imitation Learning (IL).
早期的模仿学习方法主要侧重于学习状态-动作对之间的一一映射。但这些方法忽略了历史中包含的丰富时间信息。随后的方法结合了rnns来编码观测序列,证明了利用历史观测可以提高模型性能。然而,这些方法存在基于rnn架构的固有局限性,包括表示能力有限、序列建模时间长以及训练时间慢,因为它们不适合大规模并行化。
Transformer可以对长序列进行建模,同时通过并行序列处理保持训练效率。这一趋势延伸到具有多模态感官输入的il[36,37,38,39],其中变换器对图像和语言序列进行编码
最近,扩散模型在模仿学习中表现出了优越性[5,40,6,41,39]。由于其强大的泛化能力和丰富的表示能力来捕捉多模态动作分布,它们已成为模仿学习领域的sota
Preliminaries
3.1 Mamba: Selective State-Space Models
mamba[13]通过使用选择性扫描算子改进了结构化状态空间序列模型(ssms)
输入序列,输出序列
,b、l、d分别表示批量大小、序列长度和维度
标准ssm定义了时不变参数和时间步长向量∆将x(l)的输入映射到隐藏状态,然后可以将其投影到输出y(l)
mamba通过使ssm参数成为输入的函数来实现选择机制
线性是指线性投影层,softplus是relu的平滑近似。那么输出可以通过以下方式计算