StARformer: Transformer With State-Action-Reward Representations for Robot Learning.

Area of Science:

• Robotics
• Artificial Intelligence
• Machine Learning

Background:

• Reinforcement Learning (RL) frames agent decision-making as a sequence modeling problem.
• Existing Transformer architectures for RL struggle with long-term dependencies and complex image inputs.
• A need exists for improved inductive biases in RL models to capture temporal dynamics effectively.

Purpose of the Study:

• Introduce State-Action-Reward Transformer (StARformer), a novel architecture for robot learning.
• Enhance long-term sequence modeling in RL by incorporating short-term state-action-reward representations.
• Improve performance on image-based RL benchmarks and real-world robotic tasks.

Main Methods:

• StARformer explicitly models short-term state-action-reward representations (StAR-representations) using self-attending image patches.
• Combines StAR-representations with convolutional image features for comprehensive sequence self-attention.
• Applies a Markovian-like inductive bias to improve temporal understanding within the Transformer framework.

Main Results:

• StARformer significantly outperforms state-of-the-art Transformer methods on Atari and DeepMind Control Suite benchmarks.
• Demonstrates superior performance in both offline RL and imitation learning settings.
• Shows improved handling of longer input sequences and benefits from combined patch-wise and convolutional image embeddings.

Conclusions:

• StARformer offers an effective approach to enhance RL sequence modeling for image-based tasks.
• The proposed StAR-representations and hybrid image embeddings improve long-term dependency modeling.
• StARformer shows promise for real-world applications, as evidenced by successful human-following robot imitation learning.