用于上下文内强化学习的结构化状态空间模型
机器学习
2023-11-27 v3
摘要
结构化状态空间序列（S4）模型近期在长程序列建模任务上取得了最先进的性能。这些模型还具有快速推理速度和可并行训练的特点，使它们在许多强化学习设定中具潜在用途。我们提出对 S4 一个变体的修改，使我们能够并行初始化和重置隐藏状态，从而应对强化学习任务。我们展示我们修改后的架构在序列长度上渐近快于 Transformer，并在简单基于记忆的任务上优于 RNN。我们在一组部分可观测环境中评估修改后的架构，发现实践中我们的模型优于 RNN 且运行速度快五倍以上。进而，通过利用模型处理长程序列的能力，我们在一个具有挑战性的元学习任务上取得强劲性能，其中智能体被给予随机采样的连续控制环境，以及该环境观测和动作的随机采样线性投影的组合。此外，我们展示所得模型能适应分布外保留任务。总体而言，本文结果展示结构化状态空间模型对于上下文内强化学习任务是快速且高性能的。我们在 https://github.com/luchris429/popjaxrl 提供代码。
引用
@article{arxiv.2303.03982,
title = {Structured State Space Models for In-Context Reinforcement Learning},
author = {Chris Lu and Yannick Schroecker and Albert Gu and Emilio Parisotto and Jakob Foerster and Satinder Singh and Feryal Behbahani},
journal= {arXiv preprint arXiv:2303.03982},
year = {2023}
}