Decision Mamba: Reinforcement Learning via Hybrid Selective Sequence Modeling

最近的研究表明，转换器模型在强化学习中表现出了显著的优越性，其中决策问题被制定为顺序生成。基于转换器的代理可以通过提供任务上下文（例如多条轨迹，称为上下文强化学习）在在线环境中自我改进。然而，由于转换器中注意力的二次计算复杂度，当前的上下文强化学习方法在任务视野增加时会遭受巨大的计算成本。相比之下，Mamba模型以其处理长期依赖关系的高效能力而闻名，这为上下文强化学习解决需要长期记忆的任务提供了机会。为此，我们首先通过替换决策转换器（DT）的骨干部分实现了决策Mamba（DM）。然后，我们提出了一种具有转换器和Mamba在高质量预测和长期记忆方面优点的决策Mamba-Hybrid（DM-H）。具体而言，DM-H首先通过Mamba模型从长期记忆中生成高价值子目标。然后，我们使用子目标来提示转换器，建立高质量的预测。实验结果表明，DM-H在D4RL、Grid World和Tmaze基准测试等长期和短期任务中实现了最先进的水平。在效率方面，DM-H在长期任务的在线测试速度比基于转换器的基线快28倍。

本论文旨在解决当前基于transformer的强化学习方法在处理长期任务时存在的计算复杂度问题，并提出了一种基于Mamba模型和transformer的决策生成方法。

论文提出了一种Decision Mamba-Hybrid (DM-H)方法，在Mamba模型的基础上使用transformer生成高质量的预测，并通过生成高价值子目标来解决长期记忆问题，实现了在长期和短期任务上的最优表现。

论文在D4RL、Grid World和Tmaze基准测试中实现了最优表现，DM-H在长期任务的在线测试速度比基于transformer的基线快28倍。实验数据已开源。

最近的相关研究包括使用transformer进行强化学习的方法，以及使用Mamba模型处理长期依赖性的方法。