[论文解读] Structured State Space Models for In-Context Reinforcement Learning

Chris Xiaoxuan Lu, Yannick Schroecker+5 more|arXiv (Cornell University)|Mar 7, 2023

Reinforcement Learning in Robotics被引用 11

一句话总结

这篇论文通过在轨迹中实现隐藏状态重置，改编了 S5 结构化状态空间模型以进行强化学习，在记忆基和元强化学习任务上实现更快的推理和优于 RNN 的性能，包括对分布外的泛化。

ABSTRACT

Structured state space sequence (S4) models have recently achieved state-of-the-art performance on long-range sequence modeling tasks. These models also have fast inference speeds and parallelisable training, making them potentially useful in many reinforcement learning settings. We propose a modification to a variant of S4 that enables us to initialise and reset the hidden state in parallel, allowing us to tackle reinforcement learning tasks. We show that our modified architecture runs asymptotically faster than Transformers in sequence length and performs better than RNN's on a simple memory-based task. We evaluate our modified architecture on a set of partially-observable environments and find that, in practice, our model outperforms RNN's while also running over five times faster. Then, by leveraging the model's ability to handle long-range sequences, we achieve strong performance on a challenging meta-learning task in which the agent is given a randomly-sampled continuous control environment, combined with a randomly-sampled linear projection of the environment's observations and actions. Furthermore, we show the resulting model can adapt to out-of-distribution held-out tasks. Overall, the results presented in this paper show that structured state space models are fast and performant for in-context reinforcement learning tasks. We provide code at https://github.com/luchris429/popjaxrl.

研究动机与目标

• 激发并实现在强化学习中对结构化状态空间模型（S4/S5）的有效使用。
• 在就策略 RL 训练中处理回合边界和可变长度 rollout 的挑战。
• 展示可重置的 S5 变体在运行时间更快的同时，能够达到或超过 RNN 的性能。
• 展示在长时程、部分观测以及带随机投影的元学习 RL 任务中的泛化能力。

提出的方法

• 修改 S5，使其能够在轨迹内并行初始化和重置隐藏状态。
• 引入可重置的结合运算符 ⊕，处理 done 标志以保持结合性属性。
• 使用观测和动作的随机线性投影来创建广泛的元学习任务分布。
• 用 JAX 重新实现 POPGym 环境，以实现更快的端到端评估。
• 在 memory-length bsuite 任务、POPGym 套件以及具备分布外泛化能力的多环境元 RL 上进行评估。

(a)

实验结果

研究问题

• RQ1 在训练序列中能否对 S5 架构进行重置，以处理 RL rollouts 中的 episode 边界？
• RQ2 在基于记忆的和元 RL 任务中，可重置的 S5 是否能提供相对于 LSTM 和 Transformer 的实际加速和性能提升？
• RQ3 当暴露于观测和动作的随机投影时，S5 是否能对长时程、部分观测以及未见的分布外任务泛化？
• RQ4 在具有状态/动作空间随机投影的多环境元学习设置中，S5 的表现如何？
• RQ5 在随机投影的 RL 任务中，S5 的上下文内适应是否可行？

主要发现

• 具备重置能力的 S5 在处理长序列时在渐近意义上比 Transformers 快，并且在基于记忆的任务中优于 RNNs。
• 在 bsuite memory-length 任务上，S5 获得更高的分数，且几乎比基线 RNN 方法快一倍。
• 在 POPGym 中，S5 的表现优于 GRU 且运行速度超过六倍，解决了 Repeat Previous Hard 任务。
• 在具有随机环境和投影的长上下文元 RL 设置中，S5 的回报高于 LSTMs，且在没有微调的情况下显示出某些分布外转移。

(b)