Starformer: Transformer With State-action-reward Representations For Visual Reinforcement Learning

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 se

Stats

• citations14
• S2 citations—
• github stars0
• HF likes0
• heat score8.82
• arxiv keyshang2021starformer

Related papers

• Value-consistent Representation Learning For Data-efficient Reinforcement Learning (2022)0.00
• Learning From Visual Observation Via Offline Pretrained State-to-go Transformer (2023)0.00
• Learning To Identify Critical States For Reinforcement Learning From Videos (2023)8.76
• Data-efficient Reinforcement Learning With Self-predictive Representations (2020)0.00
• Learning To Play Atari In A World Of Tokens (2024)0.00
• Transformer Based Reinforcement Learning For Games (2019)0.00
• Learning Temporally-consistent Representations For Data-efficient Reinforcement Learning (2021)0.00
• Decision Mamba: A Multi-grained State Space Model With Self-evolution Regularization For Offline RL (2024)0.00