2025_NIPS_Structured State Space Models for In-Context Reinforcement Learning

文章总结与翻译

一、主要内容

本文聚焦于将结构化状态空间序列模型（S5，S4的简化变体）应用于上下文强化学习（In-Context Reinforcement Learning）任务。核心围绕S5模型在强化学习场景中的适配性优化、性能验证及创新应用展开：

1. 核心问题：S4/S5模型虽在长序列建模中表现优异（如长程依赖处理、并行训练、高效推理），但原生架构无法在训练中重置隐藏状态，难以适配强化学习中跨episode的轨迹数据处理需求，而Transformer推理复杂度随序列长度平方增长、RNN性能有限，均存在明显短板。
2. 关键改进：提出带重置机制的S5模型（Resettable S5），通过新增含“done”信号的三元组元素和关联运算符⊕，在保留模型关联性的同时，支持训练中隐藏状态的灵活重置，适配变长序列和episode边界处理。
3. 实验验证：
   • 在bsuite记忆长度任务中，S5性能优于LSTM，推理速度接近LSTM的2倍，且渐近复杂度优于Transformer；
   • 在POPGym部分可观测环境套件中，S5性能超越GRU，运行速度提升6倍以上，在“Repeat Previous Hard”任务上实现SOTA；
   • 在多环境元强化学习任务中，通过对观测和动作空间进行随机线性投影，S5在分布内任务上优于LSTM和全连接网络，且能泛化到分布外任务。
4. 技术细节：基于