[论文解读] StARformer: Transformer with State-Action-Reward Representations for Visual Reinforcement Learning

Jinghuan Shang, Kumara Kahatapitiya+2 more|arXiv (Cornell University)|Oct 12, 2021

Reinforcement Learning in Robotics被引用 4

一句话总结

StARformer 提出了一种用于视觉强化学习的新颖 Transformer 架构，通过在短时间窗口内显式建模状态-动作-奖励（StAR）表征，引入马尔可夫性归纳偏置，从而提升长期序列建模能力。通过结合 ViT 风格的图像块与卷积特征，并利用交错的步骤（Step）和序列（Sequence） Transformer 进行特征融合，StARformer 在离线强化学习与模仿学习设置下，于 Atari 和 DM Control Suite 基准测试中均优于当前最先进方法。

ABSTRACT

Reinforcement Learning (RL) can be considered as a sequence modeling task: given a sequence of past state-action-reward experiences, an agent predicts a sequence of next actions. In this work, we propose State-Action-Reward Transformer (StARformer) for visual RL, which explicitly models short-term state-action-reward representations (StAR-representations), essentially introducing a Markovian-like inductive bias to improve long-term modeling. Our approach first extracts StAR-representations by self-attending image state patches, action, and reward tokens within a short temporal window. These are then combined with pure image state representations -- extracted as convolutional features, to perform self-attention over the whole sequence. Our experiments show that StARformer outperforms the state-of-the-art Transformer-based method on image-based Atari and DeepMind Control Suite benchmarks, in both offline-RL and imitation learning settings. StARformer is also more compliant with longer sequences of inputs. Our code is available at https://github.com/elicassion/StARformer.

研究动机与目标

• 通过显式建模局部状态-动作-奖励转换，提升视觉强化学习中的长期序列建模能力。
• 通过局部 StAR 表征引入马尔可夫性归纳偏置，减轻自注意力机制的负担。
• 通过使用 ViT 风格的图像块嵌入而非整图 CNN 特征，保留图像状态中的细粒度空间信息。
• 通过将短期转换建模与长程上下文学习解耦，提升模型在更长输入序列上的泛化能力与可扩展性。
• 在多种奖励设置下评估 StARformer 的有效性，包括返回值至目标（return-to-go）、逐步奖励与无奖励（模仿学习）场景。

提出的方法

• 模型采用步骤 Transformer，使其在单个时间步内对图像状态块、动作标记与奖励标记进行自注意力计算，以学习局部 StAR 表征。
• 图像状态通过 ViT 风格的块嵌入进行编码，以保留细粒度空间细节；同时通过卷积神经网络提取纯图像特征，以捕捉全局上下文信息。
• 序列 Transformer 通过因果自注意力机制，将 StAR 表征与全局卷积特征在完整轨迹上进行融合，生成动作预测。
• 分层特征融合机制在每个 Transformer 层中结合 StAR 表征与全局图像特征，增强表征能力，而无需增加模型深度。
• 该架构支持多种奖励设置：返回值至目标、逐步奖励与无奖励（模仿学习），从而支持在不同学习范式下的消融实验。
• 通过不同融合策略（如 P+C、C+P、StAR 融合、StAR 堆叠）评估模型变体，以确定最优的特征整合方式。

实验结果

研究问题

• RQ1 显式建模局部状态-动作-奖励转换是否能提升视觉强化学习中的长期序列建模能力？
• RQ2 将 ViT 风格的块嵌入与卷积特征融合，是否相比仅使用单一类型表征能提升性能？
• RQ3 奖励信号的选择——返回值至目标、逐步奖励或无奖励——如何影响模型性能与泛化能力？
• RQ4 所提出的交错式步骤-序列 Transformer 架构是否比堆叠或求和融合变体更有效？
• RQ5 StARformer 是否能在无任何奖励信号的模仿学习中实现优异性能，从而体现其强大的序列建模能力？

主要发现

• StARformer 在离线强化学习与模仿学习设置下，于 Atari 和 DeepMind Control Suite 基准测试中均优于当前最先进模型 Decision Transformer。
• 采用 P+C 融合策略（即步骤 Transformer 中使用块嵌入，序列 Transformer 中使用卷积特征）时模型表现最优，证实了从细粒度到粗粒度特征层次的收益。
• StARformer 在所有奖励设置下均保持强劲性能，尤其在无奖励模仿学习场景中，其表现优于使用返回值至目标的 Decision Transformer。
• 分层特征融合机制（原始 StAR 模型）优于逐元素求和（StAR 融合）与堆叠架构（StAR 堆叠），表明多层级抽象至关重要。
• 注意力图显示，动作标记能有意义地关注到相关空间区域（如 Pong 中的球拍），证实模型学习到了语义上合理的表征。
• 模型在处理更长输入序列时表现出更强的适应性，表明其在视觉强化学习中相比先前基于 Transformer 的方法具备更好的可扩展性。