The Decision Transformer paper explained🔗
This article explains and discusses the paper:
"Decision Transformer: Reinforcement Learning via Sequence Modeling" by Chen et al.
The paper explores whether a Transformer can solve sequential decision-making problems with a supervised sequence-modeling objective instead of conventional temporal-difference or policy-gradient objectives. The graph example below uses random-walk trajectories to illustrate the idea. The main Atari and continuous-control experiments use different offline datasets, described later in this article.
Figure 1. Conditioned on a starting state and the highest feasible return-to-go, Decision Transformer generates an optimal path. (Source)
TL;DR🔗
- Each return-to-go, state, and action is mapped to an embedding. Continuous states use a linear projection, while visual states use a convolutional encoder. A causally masked GPT model processes the interleaved tokens and predicts actions with a linear output head.
- Training is supervised action prediction on a fixed dataset of trajectories. The model does not learn a Q-function, apply Bellman backups, or optimize policy gradients.
- At timestep \(t\), the model predicts an action conditioned on a desired cumulative return-to-go \(\hat{R}_t\), the current state, and up to \(K-1\) previous timesteps. It is not conditioned on an immediate future reward.
- During evaluation, the user supplies an initial target return. After every action, the achieved reward is subtracted from the remaining target and the new state is appended to the context.
- Decision Transformer is evaluated in the offline RL setting: it learns only from a fixed dataset and does not explore the environment during training.
- Its performance depends on dataset coverage, the feasibility of the requested return, and how much relevant history is available in the context. A finite context limits direct access to older events, but it does not by itself imply failure or make dynamic programming mandatory.
Contributions of the paper🔗
- Use a GPT architecture to model sequences of returns-to-go, states, and actions, while training only the action-prediction objective used by the standard DT agent.
- Avoid value-function bootstrapping1 and the associated Bellman-backup machinery.
- Use undiscounted returns-to-go in the reported formulation, avoiding the short-horizon preference that discounting can introduce.
- Evaluate conditional sequence modeling on offline Atari, OpenAI Gym/D4RL, and Key-to-Door tasks.
- Show competitive aggregate performance without dynamic programming: DT is competitive with Conservative Q-Learning (CQL) and generally stronger than Random Ensemble Mixture (REM), Quantile Regression DQN (QR-DQN), and standard behavior cloning on the Atari tasks studied. Results vary by environment and should not be read as universal dominance.
Intro🔗
Supervised imitation learning fits actions observed in a dataset, whereas conventional RL explicitly optimizes expected return through interaction or value/policy optimization. This distinction does not place a strict ceiling on supervised models: they can generalize across demonstrations, and return-conditioned models may compose patterns in ways that outperform individual trajectories in the dataset. The graph shortest-path example in this paper is an illustration of such policy improvement from suboptimal data.
Transformers are known for scalable sequence modeling with comparatively weak domain-specific assumptions. After their success in language and vision, the authors ask whether a GPT-style causal Transformer can also represent policies over trajectories. BERT is part of the broader motivation in the paper, but DT itself uses GPT-style causal masking, not BERT's bidirectional masked-language-model objective.
If you are already familiar with the RL terms used in this introduction, you may want to go directly to the Decision Transformer paper motivation. For explanations of the terminology used in the paper, see the prerequisite RL knowledge section.
Motivation🔗
Transformers have proven effective at modeling high-dimensional sequence distributions in language and vision. Decision Transformer applies the same machinery to trajectories, casting policy learning as conditional sequence modeling.
One motivation is that self-attention provides short information paths between events within the context. This can help the model learn state-return associations without propagating value estimates through repeated Bellman backups, particularly with delayed or distracting rewards. The Key-to-Door experiment supports this hypothesis, but it does not prove that attention completes causal credit assignment in one step. DT still has important inductive biases: causal ordering, return conditioning, a finite context, and supervised action prediction.
Objective of the algorithm🔗
The broad research question is whether generative trajectory modeling can replace conventional RL objectives. In the standard DT implementation, however, the optimized objective is narrower: predict actions from returns-to-go, states, and previous actions. The authors report that predicting states or returns-to-go is not required for good control performance.
The processing strategy of the algorithm🔗
Conventional RL algorithms often learn a value function or optimize a parameterized policy to maximize expected return. DT instead learns a return-conditioned policy from a fixed dataset. At timestep \(t\), it models an action distribution of the form
where \(\hat{R}_t\) is the desired return-to-go and \(K\) is the context length. The model learns statistical regularities across the dataset; it does not search for an identical historical state-reward pair or copy one stored action. Conditioning can distinguish behaviors with different performance levels even when their states overlap.
Figure 2. Decision Transformer overview (Source).
Given the target return, current state, and recent trajectory, the Transformer predicts the current action. During training, the target action is known from the offline trajectory. During evaluation, the predicted action is executed in the environment and the resulting state and reward are appended to the sequence. See Figure 2.
A finite context limits which past events the model can inspect directly. Whether this causes a failure depends on the environment: the current state may already summarize the relevant past, while partially observable tasks may require a longer context or an explicit memory mechanism.
The paper demonstrates an advantage over CQL on the long-horizon Key-to-Door task when the whole episode fits in context. It does not compare DT directly with an LSTM, and finite context does not imply that dynamic programming is required for every longer-horizon problem.
With minimal modifications to a GPT architecture, DT models actions autoregressively. At test time it conditions on a desired return-to-go \(\hat{R}\) rather than an immediate future reward. This target is supplied by the user; the standard action-only DT does not first predict or model the future rewards.
A trajectory \(\tau\) is a sequence of states, actions, and rewards:
The return-to-go at timestep \(t\) is the sum of rewards from that timestep onward:
The paper uses undiscounted returns-to-go. Discounting is a modeling choice that changes the effective planning horizon and can also help define finite objectives in continuing tasks; it should not be described merely as a stability device.
The goal of offline RL is to learn a policy that maximizes expected return \(\mathbb{E}[\sum_{t=1}^{T} r_t]\) in a Markov Decision Process (MDP), using only a fixed dataset. DT represents each trajectory as
After executing the generated action, the algorithm subtracts the achieved reward from the target return and repeats until the episode terminates.
For vector observations, a learned linear layer projects each modality (return-to-go, state, and action) to the embedding dimension, followed by layer normalization. For visual observations such as Atari frames, a convolutional encoder replaces the linear state projection.
A learned episodic timestep embedding is added to all three modality tokens at that timestep. This differs from assigning an independent position to every token: one environment timestep corresponds to the return-to-go, state, and action tokens. The interleaved tokens are processed by a causally masked GPT model.
Training samples minibatches of length \(K\) from the offline trajectories. The prediction head at the state token \(s_t\) is trained to predict \(a_t\), using cross-entropy for discrete actions or mean-squared error for continuous actions. Losses are averaged across timesteps. The standard agent does not need auxiliary state or return-prediction losses.
Figure 3. Decision Transformer Pseudocode (Source).
Comparisons with offline RL and imitation learning🔗
The paper uses two main families of comparison:
- Behavior cloning (BC) learns to reproduce actions from the selected dataset with a supervised objective. Standard BC is reward-agnostic; it does not select only successful episodes. The paper separately introduces Percentile Behavior Cloning (%BC), which trains on the top \(X\%\) of timesteps ranked by episode return.
- Conservative Q-Learning (CQL) is a temporal-difference method designed for offline RL. It learns conservative Q-values to reduce overestimation of actions that are poorly represented in the fixed dataset.
Dataset dependence🔗
DT is an offline RL policy trained with a sequence-modeling loss. Like every offline method, it depends strongly on the quality and coverage of its dataset, but the paper does not use one universal set of expert demonstrations:
- Atari: 500,000 transitions, or 1% of the DQN replay dataset, collected throughout online DQN training.
- OpenAI Gym/D4RL: medium, medium-replay, and medium-expert datasets. These contain different mixtures of policy quality, including replay-buffer data and explicit medium/expert mixtures.
- Key-to-Door: 1,000 or 10,000 trajectories generated by random actions.
- Graph shortest path: 1,000 random-walk trajectories used as an illustrative experiment.
The Transformer is the architecture tested by the paper. Another causal sequence model, such as an LSTM, could implement a related return-conditioned policy, but the paper does not present an LSTM comparison. It is therefore not justified to claim that DT is empirically better than LSTMs for long-range credit assignment.
Credit assignment, Bellman backups, and attention🔗
The environment supplies rewards; an RL algorithm does not assign a larger reward to whichever action it considers important. Credit assignment is the learning problem of determining how earlier decisions contributed to later return.
Many value-based RL algorithms address this with temporal-difference (TD) learning. The state-value function for a policy \(\pi\) satisfies the Bellman expectation equation
Similarly, the state-action value function satisfies
Dynamic programming computes such recurrences using a known model or exhaustive expectations. TD learning instead bootstraps from sampled transitions and current value estimates. The ideas are related, but TD learning is not simply another name for dynamic programming, and an RL agent need not output both an action and a value at every timestep.
DT does not train a value function or apply Bellman backups. Its supervised action loss can associate earlier states and actions with hindsight returns through attention and gradient-based training. Self-attention creates a short path between tokens inside the context, but learning a useful association still requires optimization and data; it is not completed in a single step.
Limitations🔗
DT can directly attend only to the most recent \(K\) timesteps. This matters when older observations contain information that is absent from the current state, as in partially observable tasks. Possible responses include a larger context, recurrence or memory, state estimation, or a different architecture. Dynamic programming is one alternative family of methods, not a mathematical requirement whenever a trajectory exceeds \(K\).
Heuristics or rules of thumb🔗
DT was evaluated on four Atari games using 1% of the DQN replay dataset. It outperforms CQL on Breakout and Seaquest, trails CQL slightly on Pong, and trails it substantially on Qbert. It generally outperforms REM, QR-DQN, and standard BC, but several DT estimates have high variance across three seeds.
Figure 4. Gamer-normalized Atari scores for DT and the offline RL and behavior-cloning baselines. (Source)
For continuous control, the paper evaluates D4RL locomotion tasks and a separate sparse-reward Reacher task. DT achieves the highest score on a majority of the reported tasks and the highest average, while CQL or other baselines remain better on several individual settings. Most baseline values are imported as point estimates from earlier papers, so their standard deviations cannot be compared directly with DT's reported variation.
Figure 5. Results on the D4RL datasets and Reacher. DT has the highest average but is not best on every task. (Source)
To test whether DT merely clones high-return behavior, the authors introduce Percentile Behavior Cloning (%BC). It trains separate BC models on the top 10%, 25%, 40%, or 100% of timesteps, ranked by their trajectory return. With plentiful D4RL data, the best %BC model can match or beat DT on some tasks. With the much smaller Atari dataset, DT is generally stronger because it can train on all trajectories while using the return token to distinguish their behavior.
This is an analysis baseline rather than a practical selection procedure: choosing the best percentile requires environment rollouts. Both methods introduce a choice: %BC needs the percentile \(X\), while DT needs a target return. The paper argues that desired return is usually the more interpretable control variable.
Figure 6. Comparison between Decision Transformer (DT) and Percentile Behavior Cloning (%BC). (Source).
The context length is task-dependent: the Atari experiments use \(K=30\), except Pong with \(K=50\), and Key-to-Door uses the whole episode. The Atari ablation compares the normal context against \(K=1\) and shows that history improves performance. It does not establish that performance increases monotonically with every longer context.
Target-return selection is also task-dependent. For example, Pong is evaluated with a target return of 20, close to its maximum game score. This illustrates that DT needs a meaningful performance target, not that the exact maximum reward must always be known.
Figure 4 in the paper varies the target return and compares it with the realized evaluation return. The two are strongly correlated on the reported tasks. Unrealistically large targets are not guaranteed to work, especially outside the dataset's behavioral support, but the boundary is not absolute: Seaquest sometimes extrapolates above the best episode return in its training data.
Key-to-Door🔗
The Key-to-Door environment tests long-term credit assignment with the entire episode available as context. Unlike the Atari and D4RL datasets, its training trajectories are generated by random actions.
The process has three phases:
- The agent is placed in a room containing a key.
- It then enters an empty distractor room.
- Finally, it enters a room containing a door.
The agent receives a binary reward when reaching the door in the third phase, but only if it picked up the key in the first phase. This problem is difficult for credit assignment because credit must be propagated from the beginning to the end of the episode, skipping over actions taken in the middle.
With 1,000 random trajectories, DT reaches a 71.8% success rate compared with 13.1% for CQL; with 10,000 trajectories, DT reaches 94.6% and %BC reaches 95.1%. These results support the value of hindsight return information, while also showing that DT does not uniformly beat %BC.
For the analysis in Figure 7, the authors modify DT to predict return tokens in addition to actions. This critic variant lowers its predicted success probability after observing that the key was not collected and attends strongly to pivotal events. The figure therefore analyzes a modified return-predicting model, not the standard action-only DT.
Figure 7. Predicted return probability and attention weights for the modified DT critic in Key-to-Door. (Source)
What classes of problems suit the algorithm?🔗
The paper suggests that DT is a promising fit when:
- a sufficiently broad offline dataset is available;
- return is a meaningful way to prompt different behavior;
- relevant history can be represented in the model's context; and
- supervised sequence-model training is operationally simpler than maintaining a value-learning pipeline.
Important limitations remain. DT cannot explore during offline training, may produce unreliable behavior for unsupported states or infeasible target returns, and requires the context length and target return to be selected. These constraints do not yield a guarantee that DT will beat CQL or BC on a new task.
Common benchmark or example datasets used to demonstrate the algorithm🔗
Across the reported Atari, continuous-control, and Key-to-Door benchmarks, DT is competitive in aggregate with specialized TD-learning methods, but the winner depends on the environment and dataset.
Figure 8. Aggregate comparison of DT, TD learning (CQL), and behavior cloning across Atari, OpenAI Gym, and the Minigrid-based Key-to-Door task. (Source)
Dataset construction is central to interpreting these results. Atari uses a small sample from a DQN agent's complete training replay, D4RL includes medium and mixed-quality policies, and Key-to-Door uses random actions. The experiments therefore cover several data-quality regimes rather than a single expert dataset.
The paper compares DT with two broad approaches used in offline RL:
- Behavioral Cloning (BC), which predicts the dataset actions without using reward. The separate %BC analysis restricts training to high-return subsets.
- Conservative Q-Learning (CQL), which estimates conservative Q-values to reduce the offline tendency to overvalue actions outside the dataset distribution.
Useful resources for learning more about the algorithm:🔗
Footnotes:🔗
- The "deadly triad" refers to the potential instability caused by combining function approximation, bootstrapping, and off-policy learning. ↩
Michał Chromiak's blog
Comments
comments powered by Disqus