摘要
基于状态空间模型(SSM)的序列建模在各种任务中已展现出超越Transformer的性能,这让人们期待其在离线强化学习(RL)中也能超越决策Transformer及其改进版本。然而,基于最先进SSM——Mamba的决策模型,却未能取得优于这些改进型决策Transformer的性能。我们假设这种局限性源于选择性扫描阶段的信息丢失。为了解决这个问题,我们提出了决策元Mamba(DMM),它在Mamba的输入层增加了令牌混合器。该混合器明确地考虑了离线RL输入的多模态特性,包括状态、动作和剩余回报。DMM在提高性能的同时,与之前的模型相比显著减少了参数数量。值得注意的是,使用简单的线性令牌混合器也能取得类似的性能提升,这强调了保留来自相邻时间步的信息比令牌混合器的具体设计本身更重要。这种对Mamba输入层的改进代表了对Transformer中使用的传统基于时间戳的编码方法的突破。通过增强Mamba在离线RL中的性能(其特点是内存效率高和推理速度快),这项工作为其在未来RL研究中的更广泛应用开辟了新的途径。