在本文中，提出对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