在本文中,提出对S5的修改,允许它在训练阶段重置其在轨迹中的隐藏状态,这允许从业者在现有实现中简单地用S5层替换RNN。然后,我们演示了S5在基于玩具内存的任务上的性能和运行特性,表明S5的得分高于RNN,同时速度快了近两倍。最后,对具有挑战性的Meta-RL任务的S5进行了评估,在该任务中,环境对随机DMControl环境进行采样,并在每一集开始时对状态和动作空间进行随机线性投影。表明,在这种情况下,S5智能体比LSTM获得更高的回报。

Structured State Space Models for In-Context Reinforcement Learning

Chris Lu, Yannick Schroecker, Albert Gu, Emilio Parisotto, Jakob Foerster, Satinder Singh, Feryal Behbahani

[University of Oxford & DeepMind] 

  1. 结构化状态空间序列(S4)模型最近在长距离序列建模任务上取得了最先进的性能。这些模型还具有快速的推理速度和可并行的训练,使它们在许多强化学习环境中可能有用。
  2. 该文提出对S4的一个变体进行修改,能够并行地初始化和重置隐藏状态,使能够处理强化学习任务。
  3. 结果表明,修改后的架构在渐进上比Transformers 运行得更快,在一个简单的基于内存的任务上比LSTM模型表现得更好。然后,通过利用该模型处理长距离序列的能力,在一个具有挑战性的元学习任务上取得了强大的性能,其中代理被赋予一个随机采样的连续控制环境,结合随机采样的环境观察和行动的线性投影。
  4. 此外,结果表明所产生的模型可以适应分布外的保持任务。总的来说,本文提出的结果表明,S4模型是用于上下文强化学习的默认架构的有力竞争者。

https://arxiv.org/pdf/2303.03982.pdf