Decision Transformer: Reinforcement Learning via Sequence Modeling

In standard reinforcement learning, an agent interacts with an environment: it observes a state, takes an action, receives a reward, and transitions to a new state. The goal is to learn a policy that maximizes cumulative reward. The dominant approach — temporal difference (TD) learning — estimates a value function and propagates rewards backward using Bellman backups.

But offline RL removes the ability to explore. The agent has only a fixed dataset of past trajectories. TD learning in this setting faces three compounding problems. First, bootstrapping — estimating values from other estimated values — amplifies errors when the data doesn't cover the full state space. Second, the agent may overestimate the value of actions it has never seen, a phenomenon called value overestimation. Third, when rewards are sparse or delayed, Bellman backups must chain across many timesteps to propagate signal, and each link adds noise.

Methods like Conservative Q-Learning (CQL) address overestimation by penalizing Q-values for unseen actions, but they still rely on the TD learning machinery. The question this paper asks is more radical: can we bypass TD learning entirely?

GPT predicts the next word given previous words. Decision Transformer predicts the next action given previous returns, states, and actions. The insight is deceptively simple: a trajectory through an environment is just a sequence of tokens.

But there is a crucial twist. If you condition on past rewards, the model learns to imitate whatever behavior produced those rewards — including bad behavior. Instead, Decision Transformer conditions on returns-to-go: the sum of future rewards from each timestep onward. This is the key mechanism that enables control at test time.

At test time, you set the return-to-go to a high desired value — essentially telling the model "generate a trajectory that earns this much reward." The model then produces actions consistent with trajectories in the training data that achieved similar returns. This is analogous to prompting a language model with a style token: instead of "write formally," you say "achieve return 1000."

A standard language model processes sequences of word tokens. Decision Transformer processes sequences of experience tokens. Each timestep contributes three tokens, forming a repeating triplet:

1. Return-to-go ($\hat{R}_t$): the sum of all future rewards from timestep $t$ onward.
2. State ($s_t$): the environment observation (e.g., joint angles for a robot, or pixel frames for Atari).
3. Action ($a_t$): the action taken at that timestep.

The full trajectory becomes a long sequence: $\tau = (\hat{R}_1, s_1, a_1, \hat{R}_2, s_2, a_2, \ldots, \hat{R}_T, s_T, a_T)$.

The model sees the last $K$ timesteps — a context window of $3K$ tokens — just as GPT uses a fixed context window of text tokens. Each modality has its own linear embedding layer, and a shared timestep embedding (not a per-token positional encoding) is added to all three tokens from the same timestep.

The architecture is deliberately minimal — a GPT model with almost no RL-specific modifications. The three token types (return-to-go, state, action) each pass through their own linear embedding layer to reach a shared embedding dimension. A learned timestep embedding is added (not per-token positional encoding, because one timestep spans three tokens). Layer normalization follows each embedding.

For visual environments like Atari, states pass through a convolutional encoder before the linear embedding. For continuous control tasks like MuJoCo, states are raw vectors.

The embedded tokens are fed into a standard GPT with causal self-attention masking — each token can only attend to itself and previous tokens. A linear prediction head maps the hidden state at each state token position to the predicted action. The training loss is simply mean squared error for continuous actions or cross-entropy for discrete actions.

The most striking property of Decision Transformer is that it can generate behaviors of different skill levels from a single model. During training, the model sees trajectories of varying quality — from random walks to expert demonstrations. Each trajectory has its own return-to-go values.

At evaluation time, the return-to-go token acts like a prompt. Set it to the maximum return in the training data and the model generates expert-level behavior. Set it to a medium value and you get cautious, moderate behavior. The paper demonstrates a strong correlation between the specified target return and the actual return achieved by the agent across all environments.

Remarkably, on some tasks like Seaquest, the model can even extrapolate beyond the training data — achieving returns higher than any trajectory it was trained on. This suggests the model learns underlying action-reward patterns, not just memorized sequences.

Consider the Key-to-Door environment: a grid world with three phases. In phase 1, the agent is in a room with a key. In phase 2, an empty distractor room. In phase 3, a room with a door. The agent gets a reward of 1 only if it picked up the key in phase 1 and reached the door in phase 3. The actions in phase 2 are irrelevant.

This is a nightmare for TD learning. The reward signal at the door must propagate backward through every single timestep in the distractor room, each Bellman backup adding noise. CQL achieves only 13% success even with 10,000 random trajectories.

Decision Transformer handles this naturally. Self-attention lets the model directly associate the key-pickup event with the final reward, skipping the distractor phase entirely. The paper shows that the Transformer's attention weights peak at exactly two events: picking up the key and reaching the door. With 10,000 random trajectories, Decision Transformer achieves 94.6% success.

This is credit assignment through pattern matching rather than value propagation — the Transformer notices that "key pickup early + door reached late = high return" and directly learns this association.

The paper evaluates Decision Transformer across three diverse benchmarks:

Atari (discrete actions, visual inputs, 1% of DQN replay data): Decision Transformer matches CQL on 3 of 4 games and outperforms REM, QR-DQN, and behavior cloning on all 4 games.

OpenAI Gym / D4RL (continuous control — HalfCheetah, Hopper, Walker): On Medium-Expert datasets, DT achieves the highest scores in most tasks, with particularly strong results on Hopper (107.6) and Walker (108.1). On Medium-Replay datasets where data quality varies widely, DT dramatically outperforms CQL on Hopper (82.7 vs 48.6).

Key-to-Door (sparse reward, long-horizon credit assignment): DT achieves 94.6% success vs CQL's 13.3% with 10K random trajectories — the most dramatic gap, showcasing the attention-based credit assignment advantage.

Crucially, Decision Transformer is also competitive with Percentile Behavior Cloning (%BC) — a method that trains only on the best trajectories — but without needing to manually select the right data subset. DT learns which trajectories to imitate based on the return-to-go conditioning.

A critical experiment tests what happens when rewards are delayed until the end of the episode. In this "delayed return" setting, the agent receives zero reward at every timestep and only gets the total cumulative reward at the final step.

For TD learning, this is catastrophic. CQL drops from 111.0 to 9.0 on Hopper-Medium-Expert with delayed rewards — the chain of Bellman backups collapses entirely when intermediate rewards vanish.

Decision Transformer is barely affected (107.6 → 107.3 on the same task). This makes sense: the model conditions on returns-to-go, which are computed from the total trajectory reward, not step-by-step signals. Whether rewards arrive at each step or only at the end, the return-to-go calculation stays the same.

This robustness to reward sparsity is one of Decision Transformer's most practically important properties, since real-world tasks often have delayed, binary, or extremely sparse reward signals.

Decision Transformer is part of a broader trend: the unification of sequence modeling and decision-making. The paper showed that the same GPT architecture that generates coherent essays can generate coherent action sequences — without any RL-specific training machinery.

This has profound implications. First, it means RL problems can benefit from the massive infrastructure built for language models — the optimizers, the scaling laws, the training stability techniques. Second, it suggests that future generalist agents may be trained not with specialized RL pipelines, but with the same pre-train-then-prompt paradigm used for language.

Gato (DeepMind, 2022) took this idea to its logical extreme, training a single Transformer on tokenized trajectories from 604 different tasks — Atari, robotics, captioning, dialogue — achieving reasonable performance across all of them. That result traces directly back to Decision Transformer's proof of concept.