[Paper Review] Decision Transformer: Reinforcement Learning via Sequence Modeling

Lili Chen, Kevin Lü+7 more|arXiv (Cornell University)|Jun 2, 2021

Reinforcement Learning in Robotics66 references464 citations

TL;DR

This paper reframes reinforcement learning as autoregressive sequence modeling using a GPT-like transformer that conditions on returns, states, and actions to generate optimal actions without traditional value backups. It achieves competitive or superior offline RL performance on Atari, OpenAI Gym, and Key-to-Door benchmarks.

ABSTRACT

We introduce a framework that abstracts Reinforcement Learning (RL) as a sequence modeling problem. This allows us to draw upon the simplicity and scalability of the Transformer architecture, and associated advances in language modeling such as GPT-x and BERT. In particular, we present Decision Transformer, an architecture that casts the problem of RL as conditional sequence modeling. Unlike prior approaches to RL that fit value functions or compute policy gradients, Decision Transformer simply outputs the optimal actions by leveraging a causally masked Transformer. By conditioning an autoregressive model on the desired return (reward), past states, and actions, our Decision Transformer model can generate future actions that achieve the desired return. Despite its simplicity, Decision Transformer matches or exceeds the performance of state-of-the-art model-free offline RL baselines on Atari, OpenAI Gym, and Key-to-Door tasks.

Motivation & Objective

• Motivate rewriting RL as a sequence modeling problem to leverage Transformer architectures.
• Show that a causal, autoregressive model conditioned on returns can generate effective actions.
• Demonstrate that this approach matches or surpasses state-of-the-art offline RL baselines on multiple benchmarks.
• Explore how longer context and hindsight return information influence learning and credit assignment.

Proposed method

• Represent trajectories as sequences of returns-to-go, states, and actions: (R1, s1, a1, R2, s2, a2, ..., RK, sK, aK).
• Use a GPT-like transformer with causal masking to autoregressively predict actions given past tokens.
• Embed each modality with learnable linear layers and add a learned episodic positional encoding.
• Train on offline datasets with a simple supervised objective to predict the next action.
• At test time, condition on a desired return and environment start state to generate actions sequentially.

Experimental results

Research questions

• RQ1 Can a Transformer trained on offline RL data generate optimal actions by conditioning on returns and history?
• RQ2 How does the context length (K) affect performance and credit assignment in Decision Transformer?
• RQ3 How does Decision Transformer compare to model-free offline RL baselines and imitation learning across Atari and OpenAI Gym tasks?
• RQ4 Does the model retain effectiveness under sparse or delayed reward settings and can it serve as a good critic?

Key findings

• Decision Transformer matches or exceeds state-of-the-art model-free offline RL baselines on Atari, OpenAI Gym, and Key-to-Door tasks.
• On Atari, DT is competitive with CQL on 3 of 4 games and outperforms REM, QR-DQN, and BC in several cases.
• DT achieves high scores on most OpenAI Gym tasks and often surpasses baselines, including in sparse reward settings.
• Longer context (K>1) substantially improves performance over K=1, indicating the value of past trajectories for policy generation.
• DT demonstrates strong long-horizon credit assignment in the Key-to-Door task and robust performance with delayed rewards.”
• DT can model the distribution of returns; conditioned on target returns, generated trajectories align with desired returns and can extrapolate in some cases.
• %BC analyses suggest Decision Transformer leverages the full dataset distribution rather than merely imitating the best trajectories.