Decision Transformer: Reinforcement Learning via Sequence Modeling

Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, Igor Mordatch

2021 · NeurIPS

Decision Transformer: Reinforcement Learning via Sequence Modeling

Problem

Framing

Offline RL still relies on value estimation, bootstrapping, and pessimism to recover good policies from fixed data. This paper replaces that stack with conditional sequence modeling: a causal Transformer predicts actions from desired return-to-go, past states, and past actions, matching or exceeding strong offline RL baselines on Atari and D4RL.

Currently Used Methods

Foundational

@mnihDQN2015 — value-based deep RL for Atari from replayed transitions.
- Limitation in context: needs bootstrapped value learning, which is unstable offline.
@schulmanPPO2017 — policy-gradient optimization with clipped surrogate objectives.
- Limitation in context: assumes online interaction, not fixed suboptimal datasets.
@vaswaniAttentionAllNeed2017 — self-attention sequence model for long-range dependencies.
- Limitation in context: not formulated as a return-conditioned control policy.
@brownGPT3_2020 — causal autoregressive Transformer scaled for sequence prediction.
- Limitation in context: does not model trajectories or optimize control returns.
@haarnojaSAC2018 — actor-critic RL with entropy regularization for continuous control.
- Limitation in context: still depends on critic estimation and online updates.

Proposed Method

Architecture

Decision Transformer packs each timestep as return-to-go, state, and action tokens, then feeds the last $K$ steps as $3K$ tokens to a causally masked GPT. Each modality has its own embedding layer, timestep embeddings are shared across the three tokens of a step, and the hidden state at each state token predicts the next action.

Architecture diagram: return, state, and action tokens are interleaved, embedded with timestep encodings, passed through a causal transformer, and decoded autoregressively into actions.

Loss / Objective

Training reduces to next-action prediction on offline trajectories.

\hat{\tau} = (\hat{R}_1, s_1, a_1, \hat{R}_2, s_2, a_2, \ldots, \hat{R}_T, s_T, a_T)

\mathcal{L}(\theta) = \frac{1}{K}\sum_{t=1}^{K} \ell\big(\pi_\theta(a_t \mid \hat{R}_1, s_1, a_1, \ldots, \hat{R}_t, s_t),\, a_t\big)

\ell = \begin{cases} -\log p_\theta(a_t \mid \cdot) & \text{discrete actions} \\ \|\hat{a}_t - a_t\|^2 & \text{continuous actions} \end{cases}

Sampling Rule / Algorithm

At test time, the model conditions on a target return and rolls it forward autoregressively.

\hat{a}_t \sim \pi_\theta\big(a_t \mid \hat{R}_1, s_1, a_1, \ldots, \hat{R}_t, s_t\big)

\hat{R}_{t+1} = \hat{R}_t - r_t

Training Procedure

Atari: 6 layers, 8 heads, embedding size 128.
Atari batch size: 512 for Pong; 128 for Breakout, Qbert, Seaquest.
Atari context length $K$ : 50 for Pong; 30 for Breakout, Qbert, Seaquest.
OpenAI Gym: 3 layers, 1 head, embedding size 128.
OpenAI Gym batch size: 64.
OpenAI Gym context length $K$ : 20 for HalfCheetah, Hopper, Walker; 5 for Reacher.
OpenAI Gym nonlinearity: ReLU.

Evaluation

Datasets

Atari 1% DQN-replay: Breakout, Qbert, Pong, Seaquest.
D4RL continuous control: HalfCheetah, Hopper, Walker.
Sparse-goal control: Reacher.
Key-to-Door sparse-reward benchmark.

Metrics

Gamer-normalized score for Atari.
D4RL normalized score for Gym tasks.
Normalized episode return in cross-domain summary.
Success rate for Key-to-Door.

Headline results

Atari Breakout (1% replay): 267.5 $\pm$ 97.5, above CQL 211.1.
Atari Pong (1% replay): 106.1 $\pm$ 8.1, near CQL 111.9.
D4RL HalfCheetah medium-expert: 86.8 $\pm$ 1.3, above CQL 62.4.
D4RL Walker medium-expert: 108.1 $\pm$ 0.2, above CQL 98.7.
D4RL Hopper medium-replay: 82.7 $\pm$ 7.0, above CQL 48.6.

Results bar chart: Decision Transformer versus TD learning and behavior cloning on Atari, OpenAI Gym, and Key-to-Door; it leads on Gym and Key-to-Door and is competitive on Atari.

Ablations

Context length: $K=1$ is much worse than longer contexts; the model uses history.
Behavior cloning subset: Decision Transformer beats percentile-filtered BC on most Atari games.
Return modeling: larger models improve return-distribution modeling in Gym.
Sparse rewards: hindsight-conditioned methods succeed where TD learning struggles in Key-to-Door.

Method Strengths and Weaknesses

Strengths

Removes Bellman backups and explicit value functions entirely.
Matches or beats CQL on most D4RL settings.
Strong on sparse-reward credit assignment in Key-to-Door.
Longer-context gains show self-attention uses trajectory history.

Weaknesses

Qbert underperforms CQL: 15.4 versus 104.2.
Requires selecting a target return at inference.
Pure offline training does not address exploration.
Reported gains are uneven across Atari tasks.

Suggestions from the authors

Study joint prediction of actions, states, and returns.
Extend the method from offline to online RL regimes.
Test larger transformers for broader behavior modeling.
Analyze sequence modeling for harder long-horizon credit assignment.

Decision Transformer: Reinforcement Learning via Sequence Modeling

Decision Transformer: Reinforcement Learning via Sequence Modeling

Problem

Framing

Currently Used Methods

Foundational

Proposed Method

Architecture

Loss / Objective

Sampling Rule / Algorithm

Training Procedure

Evaluation

Datasets

Metrics

Headline results

Ablations

Method Strengths and Weaknesses

Strengths

Weaknesses

Suggestions from the authors

Links

Prior Papers

Further Papers