Decision Mamba Architectures

📄 arXiv: 2405.07943v2 📥 PDF

作者: André Correia, Luís A. Alexandre

分类: cs.LG, cs.AI

发布日期: 2024-05-13 (更新: 2024-10-17)

🔗 代码/项目: GITHUB

💡 一句话要点

提出Decision Mamba和Hierarchical Decision Mamba，提升模仿学习中Transformer模型的性能。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 模仿学习 决策Transformer Mamba架构 序列建模 强化学习

📋 核心要点

  1. 模仿学习领域受益于序列模型，但Decision Transformer等现有方法依赖Transformer，存在效率瓶颈。
  2. 本文提出Decision Mamba (DM) 和 Hierarchical Decision Mamba (HDM)，用Mamba架构替代Transformer，提升序列建模能力。
  3. 在OpenAI Gym和D4RL等环境的实验表明，DM和HDM在多数任务中超越了Transformer模型，展现了更优的性能。

📝 摘要（中文）

本文提出了两种新方法，Decision Mamba (DM) 和 Hierarchical Decision Mamba (HDM)，旨在提升基于Transformer模型的决策策略学习性能。这两种方法将最近表现出色的Mamba架构应用于模仿学习，替代了原有的Transformer结构。通过在OpenAI Gym和D4RL等多种环境中使用不同的演示数据集进行大量实验，结果表明，在大多数任务中，Mamba模型优于其Transformer对应模型。实验结果表明，DM在大多数设置中优于其他方法。代码已开源。

🔬 方法详解

问题定义：现有基于Transformer的决策模型，如Decision Transformer (DT) 和 Hierarchical Decision Transformer (HDT)，在模仿学习中面临计算效率和长期依赖建模的挑战。Transformer的自注意力机制计算复杂度高，难以有效处理长序列数据，限制了其在复杂任务中的应用。

核心思路：本文的核心思路是用Mamba架构替换Transformer，利用Mamba选择性状态空间模型（Selective State Space Models, S6）的优势，提高序列建模的效率和性能。Mamba架构具有线性复杂度，能够更好地处理长序列，并能根据输入动态调整状态转移，从而更有效地捕捉任务行为的复杂性。

技术框架：DM和HDM的整体架构与DT和HDT类似，主要区别在于将Transformer模块替换为Mamba模块。DM直接将Mamba应用于决策序列建模，而HDM则在分层结构中使用Mamba，以处理更复杂的任务。具体流程包括：将环境状态、动作和奖励等信息编码为序列，输入到Mamba模型中进行处理，然后解码得到策略。

关键创新：最关键的创新在于将Mamba架构引入到决策Transformer中。Mamba架构的核心是选择性状态空间模型，它允许模型根据输入动态地选择和更新状态，从而更有效地捕捉序列中的关键信息。与Transformer的自注意力机制相比，Mamba具有更高的计算效率和更好的长期依赖建模能力。

关键设计：Mamba模块的参数设置，包括状态维度、选择性参数等，需要根据具体任务进行调整。损失函数通常采用行为克隆损失，即最小化模型预测动作与演示数据中动作之间的差异。网络结构的设计需要考虑任务的复杂度和序列长度，例如，对于更复杂的任务，可以采用更深的网络结构或分层结构。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，Decision Mamba (DM) 在大多数OpenAI Gym和D4RL环境的任务中优于基于Transformer的Decision Transformer (DT) 和 Hierarchical Decision Transformer (HDT)。具体而言，DM在某些任务上的性能提升幅度超过10%，并且在处理长序列数据时表现出更高的效率。这些结果验证了Mamba架构在模仿学习中的有效性。

🎯 应用场景

该研究成果可应用于各种需要模仿学习的机器人控制任务，例如自动驾驶、机械臂操作、游戏AI等。通过使用Mamba架构，可以提高决策模型的训练效率和泛化能力，从而使机器人能够更好地学习和执行复杂任务。此外，该方法还可以扩展到其他序列建模任务，如自然语言处理和时间序列预测。

📄 摘要（原文）

Recent advancements in imitation learning have been largely fueled by the integration of sequence models, which provide a structured flow of information to effectively mimic task behaviours. Currently, Decision Transformer (DT) and subsequently, the Hierarchical Decision Transformer (HDT), presented Transformer-based approaches to learn task policies. Recently, the Mamba architecture has shown to outperform Transformers across various task domains. In this work, we introduce two novel methods, Decision Mamba (DM) and Hierarchical Decision Mamba (HDM), aimed at enhancing the performance of the Transformer models. Through extensive experimentation across diverse environments such as OpenAI Gym and D4RL, leveraging varying demonstration data sets, we demonstrate the superiority of Mamba models over their Transformer counterparts in a majority of tasks. Results show that DM outperforms other methods in most settings. The code can be found at https://github.com/meowatthemoon/DecisionMamba.