Source-linked AI summary

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

arXiv:2106.01345v2cs.LGcs.AI

TL;DR

Offline RL lacks evidence that generative trajectory modeling can replace conventional RL algorithms. Decision Transformer treats RL as conditional sequence modeling and matches or exceeds strong offline RL baselines across Atari, OpenAI Gym, and Key-to-Door.

  • Problem

    It remains unclear whether generative trajectory modeling can replace conventional reinforcement-learning algorithms for sequential decision making.

  • Method

    Decision Transformer autoregressively models returns-to-go, states, and actions with a causally masked Transformer to condition future actions on desired returns.

  • Results

    Decision Transformer matches or exceeds state-of-the-art model-free offline RL algorithms across Atari, OpenAI Gym, and Key-to-Door benchmarks.

  • Takeaways & Limitations

    Sequence modeling can serve as an effective algorithmic paradigm for offline reinforcement learning without dynamic programming.

  • Takeaways & Limitations

    Real-world use requires further study of transformer errors, dataset-induced destructive biases, and unintended behaviors from reward design.

Abstract

from arXiv · show

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.

1 Introduction

The paper reframes offline reinforcement learning as conditional sequence modeling with transformers, proposing Decision Transformer to generate actions from desired returns, past states, and actions. This approach avoids dynamic programming and matches or exceeds state-of-the-art model-free offline RL baselines across several benchmarks.

  • Motivation: Transformers are investigated for sequential decision making because they model high-dimensional distributions and long sequences effectively.The motivation extends their demonstrated applications in language and image generation to reinforcement learning.
  • Motivation: Sequence modeling replaces conventional TD-based policy training, avoiding bootstrapping for long-term credit assignment and discounting future rewards.The passage frames this as bypassing one destabilizing component of the deadly triad.
  • Decision Transformer: Decision Transformer autoregressively models sequences of states, actions, and returns to learn effective policies from fixed, suboptimal offline data.This formulation addresses offline RL challenges including error propagation and value overestimation while reducing policy sampling to autoregressive generation.
  • Illustrative example: Conditioning on desired returns-to-go enables optimal trajectories from random-walk training data without expert demonstrations in a shortest-path task.The model predicts the next token in sequences containing returns-to-go, states, and actions, then generates paths by targeting the largest possible return.
  • Contributions: Decision Transformer matches or exceeds state-of-the-art model-free offline RL performance on Atari, OpenAI Gym, and Key-to-Door benchmarks without dynamic programming.The evaluation tests whether sequence modeling can perform policy optimization across these offline RL environments.

2 Preliminaries

The paper formulates reinforcement learning in Markov decision processes, where policies maximize expected future rewards, and highlights offline learning’s restriction to fixed trajectory data. It also introduces causal Transformer self-attention as the basis for autoregressive sequence modeling.

  • Reinforcement learning preliminaries: An MDP is defined by states, actions, transition dynamics, and rewards, with trajectories recording their timestep-wise sequence.The trajectory is represented as τ = (s0, a0, r0, s1, a1, r1, . . . , sT, aT, rT).
  • Reinforcement learning preliminaries: The return Rt is the sum of future rewards from timestep t, and reinforcement learning seeks a policy maximizing expected return.The return sums rewards from t onward.
  • Offline reinforcement learning: Offline reinforcement learning uses a fixed, limited dataset of trajectory rollouts rather than new environment interactions, preventing additional exploration and feedback.The dataset may contain rollouts from arbitrary policies.
  • Transformers: Transformers model sequential data with stacked self-attention layers and residual connections that preserve token-level embedding dimensions.Each self-attention layer maps input embeddings to output embeddings of the same dimensionality.
  • Transformers: GPT adds a causal self-attention mask, restricting each token to previous tokens and enabling autoregressive generation.The masked computation replaces attention over all n tokens with attention over j ∈ [1, i].

3 Method

Decision Transformer models offline RL trajectories autoregressively with a causally masked Transformer, conditioning action generation on desired returns, states, and actions. Its method uses returns-to-go, modality-specific embeddings, and action-prediction training for discrete or continuous actions.

  • Core framework: Decision Transformer autoregressively models trajectories with minimal modification to the Transformer architecture.The framework is summarized in Figure 1 and Algorithm 1.
  • Trajectory representation: The trajectory representation interleaves returns-to-go, states, and actions so actions can be generated conditionally from future desired returns.The sequence is represented as R̂1, s1, a1, R̂2, s2, a2, ..., R̂T, sT, aT.
  • Trajectory representation: At test time, the model conditions generation on a desired performance and starting state, then decrements the target return by each achieved reward.Generation repeats after each environment interaction until episode termination.
  • Architecture: Decision Transformer processes the last K timesteps as 3K modality tokens, using learned modality projections, layer normalization, timestep embeddings, and a convolutional state encoder for visual inputs.Each timestep contributes one return-to-go, state, and action token.
  • Training: Training samples offline trajectory minibatches of length K and predicts actions with cross-entropy for discrete actions or mean-squared error for continuous actions.Losses are averaged across timesteps, while predicting states or returns-to-go was not found to improve performance.

4 Evaluations on Offline RL Benchmarks

The evaluation compares Decision Transformer with dedicated offline RL and imitation-learning methods across discrete Atari and continuous-control benchmarks. Results show competitive Atari performance and stronger performance than conventional RL algorithms on most D4RL tasks.

  • Evaluation setup: The evaluation primarily compares Decision Transformer with model-free offline RL methods based on TD-learning, especially CQL, alongside BEAR and BRAC.These methods provide the closest comparison because Decision Transformer is fundamentally model-free.
  • Evaluation setup: The benchmarks span Atari’s high-dimensional visual control and delayed credit assignment and OpenAI Gym’s fine-grained continuous control.Averaged normalized performance is reported for each domain.
  • Atari: Decision Transformer performs comparably to CQL on 3 out of 4 Atari games and outperforms other baselines in most games.The Atari comparison uses Breakout, Qbert, Pong, and Seaquest, with behavior cloning included as an additional reference.
  • Continuous control: The continuous-control comparison includes CQL, BEAR, BRAC, and AWR, with scores normalized so that 100 represents an expert policy.For Reacher, only the CQL baseline is run because CQL is generally the strongest TD-learning method.
  • Continuous control: Decision Transformer outperforms conventional RL algorithms on almost all D4RL tasks.The continuous-control evaluation includes HalfCheetah, Hopper, Walker, and a separately constructed sparse-reward Reacher environment.

5 Discussion

The discussion interprets Decision Transformer through percentile behavior cloning, tests return conditioning and context, and examines long-term credit assignment, sparse rewards, critics, and offline-to-online implications. Across these analyses, Decision Transformer can use broader trajectory information and hindsight returns to remain effective where conventional offline RL methods struggle.

  • Imitation-learning interpretation: Percentile Behavior Cloning (%BC) suggests Decision Transformer can resemble imitation learning on high-return subsets, especially when data is plentiful.%BC trains only on the top X% of timesteps by episode return, but selecting the optimal subset requires environment rollouts and is therefore unrealistic.
  • Imitation-learning interpretation: In low-data Atari settings, Decision Transformer outperforms weak %BC by using all trajectories to improve generalization beyond trajectories matching the return target.The passage attributes this advantage to leveraging dissimilar trajectories rather than cloning only high-return data.
  • Return conditioning: Desired target returns and observed returns are highly correlated across tasks, and Decision Transformer sometimes extrapolates beyond the dataset’s maximum return.The extrapolation result is demonstrated on some Atari tasks, including Seaquest.
  • Context and credit assignment: Longer context improves Decision Transformer performance on Atari, indicating that previous states, actions, and returns provide useful information beyond the current state.The context-length ablation reports significantly worse performance with K = 1 and better performance with longer contexts.
  • Context and credit assignment: Using hindsight return information, Decision Transformer learns near-optimal Key-to-Door policies from random-action trajectories, demonstrating effective long-term credit assignment.The environment requires associating a delayed binary reward with picking up a key across three phases.
  • Robustness and broader implications: Decision Transformer remains minimally affected by delayed sparse rewards, does not require policy regularization or conservatism, and may support sample-efficient online behavior generation.The discussion contrasts this robustness with CQL’s failure under delayed rewards and motivates Decision Transformer as a behavior-generation model for online RL.

6 Related Work

Prior offline RL methods mitigate distribution shift through action constraints, value or dynamics pessimism, while Decision Transformer primarily compares against model-free algorithms. Related work includes return-conditioned supervised methods, trajectory modeling, credit-assignment architectures, guided generation, and transformer applications in RL.

  • Offline reinforcement learning: Offline RL methods address distribution shift with policy action constraints, value pessimism, or pessimism in learned dynamics models.Decision Transformer does not explicitly learn a dynamics model and therefore primarily compares against model-free algorithms.
  • Return-conditioned methods: UDRL resembles Decision Transformer by modeling behavior with a supervised loss conditioned on target return, but Decision Transformer is motivated by sequence modeling.Sequence modeling can model behavior without reward access and is known to scale well, whereas practical differences include context length and architecture.
  • Trajectory modeling: Trajectory Transformer similarly applies sequence modeling to RL but additionally predicts states and returns, uses discretization, and incorporates model-based components.The authors view both sets of experiments as highlighting sequence modeling’s general applicability to reinforcement learning.
  • Credit assignment: State-associative architectures improve credit assignment by identifying important states, reshaping rewards, and helping actor-critic methods propagate signals over long horizons.Some works specifically report better performance in delayed-reward settings.
  • Guided generation: Guided-generation methods for images and language often assume constant classes, whereas reinforcement learning uses a time-varying reward signal.The passage also describes controllable text-generation models and class-conditional discriminators used to guide generation.
  • Transformers in reinforcement learning: Transformers have succeeded in language and vision but remain relatively unstudied in RL, where differing problem characteristics include higher training variance.Prior RL studies used relational reasoning or iterative self-attention to improve performance or episode-level information use.

7 Conclusion · A.1 Atari · A.2 OpenAI Gym

Decision Transformer unifies sequence modeling and offline reinforcement learning, matching or outperforming strong offline RL algorithms with minimal architectural changes. The paper also identifies future directions and deployment risks, while describing Atari implementation and conditioning choices.

  • 7 Conclusion: Decision Transformer unifies language or sequence modeling with reinforcement learning and matches or outperforms strong offline RL algorithms with minimal modifications.The conclusion frames this as the central contribution on standard offline RL benchmarks.
  • 7 Conclusion: Future work includes self-supervised pretraining, more sophisticated embeddings, and return-distribution conditioning for stochastic settings.The authors note that the supervised loss was effective but that larger datasets may benefit from pretraining and richer conditioning.
  • 7 Conclusion: Real-world deployment requires studying transformer errors, negative consequences, dataset biases, and unintended behaviors from malicious reward design.These risks are described as underexplored, especially when augmenting agents with data from questionable sources.
  • A.1 Atari: The Atari implementation is built on minGPT and largely retains default hyperparameters from its character-level GPT example.Training uses reduced batch size, block size, layers, attention heads, and embedding dimension for faster training, except for Pong batch size.
  • A.1 Atari: Atari return-to-go conditioning uses either 1× or 5× the dataset’s maximum return, while other principled conditioning choices remain possible.The appendix also notes that Tanh replaces LayerNorm after modality embeddings without a significant performance difference.
  • A.1 Atari: Table 8 provides the full hyperparameter list for Decision Transformer and %BC in the Atari experiments.The supplied passage identifies this table as the Atari hyperparameter reference.

A.2.1 Decision Transformer

The Decision Transformer implementation uses Hugging Face Transformers and task-specific modeling choices. Larger models are heuristically favored for return-distribution modeling, while Reacher uses a shorter context length because it is goal-conditioned with shorter episodes.

  • The implementation is based on the Hugging Face Transformers library.
  • Larger models are heuristically preferred because they better model the distribution of returns than standard RL model sizes.
  • Reacher uses a smaller context length because its environment is goal-conditioned and its episodes are shorter.

A.2.2 Behavior Cloning

The authors found previously reported behavior cloning baselines weak and therefore reran them using a setup similar to Decision Transformer. An MLP outperformed a transformer, with training beyond 2.5 × 10^4 gradient steps providing no further improvement.

  • Previously reported behavior cloning baselines were weak, so the authors reran them using a setup similar to Decision Transformer.
  • An MLP architecture was stronger than a transformer for behavior cloning, consistent with previous work.
  • 2.5 × 10^4 gradient steps were used because additional training did not improve performance; percentile experiments used the same hyperparameters.

A.3 Graph Shortest Path

The graph shortest-path task is formulated as an MDP, and a GPT-based Decision Transformer generates actions and return-to-go tokens using a prior favoring shorter paths. Trained on random walks, the model substantially reduces steps to the goal, closely matching optimal paths and sometimes composing entirely original routes.

  • Task formulation: The task is finding the shortest path on a fixed directed graph, with reward 0 at the goal and −1 otherwise.The observation is the current node index, the action is the next node index, and valid edges determine transitions.
  • Generation method: The GPT model generates both actions and return-to-go tokens, enabling it to produce its own realizable return estimates.Generation is conditioned on return prompts, states, actions, and prior return-to-go tokens.
  • Generation method: A prior proportional to T + 1 − k favors shorter paths when selecting return prompts, using the known optimal path length during prompting.The prior is combined with the model’s return probability to adjust generation.
  • Results: Trained on 1,000 random-walk trajectories of T = 10 steps on 20-node graphs, the transformer substantially reduces steps to the goal and closely matches optimal paths.The graphs use an edge sparsity coefficient of 0.1, and results are reported in Figure 6.
  • Results: 15.8% of generated paths are entirely original, formed by stitching sub-optimal segments rather than reproducing a training trajectory.The other favorable case occurs when the dataset contains a segment directly corresponding to the desired shortest path.
  • Implication: The example shows how hindsight return information combined with generation priors can avoid explicit dynamic programming.This prior-based mechanism is used here as an illustrative example rather than in the paper’s other experiments.

B Atari Task Scores

The Atari results report normalized and raw baseline scores, with Decision Transformer performing comparably to CQL on three of four games and usually outperforming other baselines and %BC.

  • Score reporting: Tables 11–13 provide normalized, raw, and %BC Atari scores for baseline comparisons.Raw and %BC results report means and variances across 3 seeds.
  • Atari performance: Decision Transformer performs comparably to CQL on 3 out of 4 games in the 1% DQN-replay Atari dataset.Table 12 reports raw scores and highlights the best mean scores.
  • Atari performance: Decision Transformer usually outperforms other baselines and %BC on Atari.%BC comparisons use the same hyperparameters as Decision Transformer for fair comparison.
Loading 2106.01345v2…