Source-linked AI summary
Offline Reinforcement Learning as One Big Sequence Modeling Problem
Michael Janner, Qiyang Li, Sergey Levine
TL;DR
The paper asks whether reinforcement learning can be treated as a generic sequence modeling problem instead of relying on separately designed policies, models, and value functions. It trains a Transformer over trajectories and uses beam search for planning, finding effective results across long-horizon prediction, imitation learning, goal-reaching, and offline RL. The approach is promising but remains computationally costly for real-time control and depends on discretizing continuous data.
Problem
The paper investigates whether contemporary high-capacity sequence models can provide a viable reinforcement learning algorithm while avoiding many components usually designed separately for RL.
Method
The method jointly models states, actions, and rewards with a Transformer and uses minimally modified beam search to sample and plan trajectories.
Results
The approach is effective across long-horizon dynamics prediction, imitation learning, goal-reaching, and offline reinforcement learning, with offline RL performance on par with or better than prior methods.
Takeaways & Limitations
Sequence modeling can unify policies, models, and value functions under one algorithm, while combinations with dynamic programming may be especially practical for challenging planning tasks.
Takeaways & Limitations
Standard Transformers can require multiple seconds for action selection as context grows, precluding real-time control for most dynamical systems; the method also discretizes continuous data.
Abstract
from arXiv · showhide
Reinforcement learning (RL) is typically concerned with estimating stationary policies or single-step models, leveraging the Markov property to factorize problems in time. However, we can also view RL as a generic sequence modeling problem, with the goal being to produce a sequence of actions that leads to a sequence of high rewards. Viewed in this way, it is tempting to consider whether high-capacity sequence prediction models that work well in other domains, such as natural-language processing, can also provide effective solutions to the RL problem. To this end, we explore how RL can be tackled with the tools of sequence modeling, using a Transformer architecture to model distributions over trajectories and repurposing beam search as a planning algorithm. Framing RL as sequence modeling problem simplifies a range of design decisions, allowing us to dispense with many of the components common in offline RL algorithms. We demonstrate the flexibility of this approach across long-horizon dynamics prediction, imitation learning, goal-conditioned RL, and offline RL. Further, we show that this approach can be combined with existing model-free algorithms to yield a state-of-the-art planner in sparse-reward, long-horizon tasks.
1 Introduction
The paper reframes reinforcement learning as sequence generation over states, actions, and rewards, using a Transformer and beam search to plan high-reward trajectories. This formulation unifies components typically handled by separate RL methods and supports offline evaluation across several settings.
- RL can be viewed as producing action sequences that yield sequences of high rewards, rather than decomposing long-horizon problems into local subproblems.
- The Trajectory Transformer models trajectories as unstructured sequences of states, actions, and rewards with a Transformer architecture.
- Beam search replaces trajectory optimizers as the procedure for planning with the learned sequence model.
- A single sequence model can unify policies, dynamics models, value functions, and behavior-policy estimation as distributions over data streams.
- The model is evaluated offline using prior interaction data and performs competitively on offline RL benchmarks while also supporting imitation learning, goal-reaching, and sparse-reward planning.
2 Related Work
Related work applies learned probability distributions and sequence models to model-based RL, offline RL, imitation learning, and hierarchical RL. The Trajectory Transformer instead represents the joint trajectory distribution with one high-capacity sequence model and uses beam-search planning.
- Deep sequence models have improved substantially, motivating their application to reinforcement learning and other sequential processes.
- Probability distributions support predictive modeling and planning in model-based RL, action constraints in offline RL, and policy fitting in imitation learning.
- The Trajectory Transformer trains one high-capacity sequence model to represent joint sequences of states, actions, and rewards.
- Unlike many prior model-based methods, the approach is presented without the additional ensembles or conservatism mechanisms commonly used in online and offline settings.
- Concurrent work also used sequence prediction for RL, emphasizing reward conditioning rather than the Trajectory Transformer’s beam-search-based planning.
3 Reinforcement Learning and Control as Sequence Modeling
The Trajectory Transformer represents trajectories as token sequences and uses Transformer decoding and beam search for imitation, goal-conditioned, and offline reinforcement learning. Continuous trajectories are independently discretized, while reward-guided search selects high-reward candidate sequences.
- Trajectory representation: The Trajectory Transformer treats each trajectory as an unstructured sequence of states, actions, and rewards for Transformer modeling.For continuous states and actions, each dimension is discretized independently, producing T(N +M +1) tokens over T timesteps.
- Trajectory representation: Uniform discretization preserves Euclidean distance, whereas quantile discretization allocates equal empirical probability mass to each token.Uniform bins can be distorted by outliers; quantile bins ensure every token is represented in the training data.
- Trajectory Transformer: The model is a GPT-style Transformer decoder trained with teacher forcing to predict token sequences autoregressively.The implementation uses four layers, four self-attention heads, and a likelihood objective over trajectory tokens.
- Imitation learning: In imitation learning, beam search generates high-probability trajectories from the current state, and repeated execution yields a receding-horizon controller.Predicting only an action-dimension-length sequence recovers the simplest autoregressive behavior-cloning policy.
- Goal-conditioned reinforcement learning: Goal-conditioned control prepends a desired final state so trajectory predictions can attend to the goal while preserving the causal attention mask.This conditions generated trajectories on future context without modifying the standard attention implementation.
- Offline reinforcement learning: Offline reinforcement learning replaces transition log-probabilities in beam search with predicted rewards, using reward-to-go as a search heuristic.Reward-to-go estimates are behavior-policy values and may be insufficient in challenging sparse-reward settings, where improved value functions help.
- Planning with beam search: Jointly modeling states and actions can prevent querying the learned model on out-of-distribution actions during planning.This contrasts with unconstrained action-sequence optimization, which can more readily exploit model errors.
4 Experiments
The experiments evaluate Trajectory Transformer for long-horizon prediction and for control through beam-search planning across offline RL, imitation learning, and goal-reaching. It produces more accurate long-horizon predictions than conventional dynamics models and competitive or superior control results across the tested settings.
- Long-horizon prediction: Trajectory Transformer has substantially better error compounding with respect to prediction horizon than a probabilistic feedforward model ensemble.The comparison estimates per-timestep state-marginal likelihood from 1000 trajectories sampled from a fixed starting point.
- Long-horizon prediction: Trajectory Transformer predictions remain visually indistinguishable from ground-truth humanoid trajectories after 100 predicted steps.The reference probabilistic feedforward ensemble produced accurate predictions for only a few dozen steps.
- Model analysis: Long-horizon conditioning plays a larger role in prediction accuracy when each state dimension is masked independently with 50% probability.On fully observed environments, a truncated-context Markovian variant performs similarly to the full Trajectory Transformer.
- Model analysis: Attention analysis reveals both Markovian transition-focused masks and striated masks that attend to specific dimensions across multiple prior states.Action predictions attend more to prior actions than prior states, resembling action filtering for smoother sequences.
- Offline reinforcement learning: Trajectory Transformer performs on par with or better than prior offline RL methods on D4RL locomotion tasks.Uniform and quantile discretization perform similarly except on HalfCheetah-Medium-Expert, where quantile discretization achieves more than twice the uniform-discretization return.
- Offline reinforcement learning: Q-guided Trajectory Transformer planning outperforms all prior methods across AntMaze maze sizes and dataset compositions.The Q-function is used as a beam-search heuristic rather than directly extracted into a policy.
- Imitation and goal-reaching: Likelihood-maximizing beam search as a receding-horizon controller achieves normalized returns of 104% in Hopper and 109% in Walker2d.The same approach also supports goal-reaching through anti-causal conditioning on a future desired state.
5 Discussion and Limitations
The paper presents a single sequence-modeling algorithm for diverse RL settings, while noting practical limits from Transformer cost and discretization. It also suggests combining the approach with dynamic programming for practical performance.
- A single sequence model jointly trained on states, actions, and rewards unifies policies, models, and value functions across imitation learning, goal-reaching, and offline RL.
- Transformer prediction can require multiple seconds for action selection with large context windows, making standard architectures unsuitable for real-time control in most dynamical systems.
- Discretizing continuous data improves prediction over conventional continuous dynamics models but imposes an upper bound on prediction precision.
- Combining the minimal sequence-modeling algorithm with dynamic programming techniques may be the most practical instantiation, as suggested by Q-guided planning.
Appendix A Model and Training Specification
The model uses a fixed four-layer, four-head Transformer with a vocabulary supporting states, actions, rewards, and rewards-to-go. Training uses a warmup learning-rate schedule, large batches, and single-GPU runs lasting several hours.
- The architecture uses four Transformer layers and four self-attention heads, with input tokens covering states, actions, rewards, and rewards-to-go.
- The learning rate increases linearly from 0 to 2.5 × 10^-4 over 2000 updates, using a batch size of 256.
- Training runs for 80 epochs on one NVIDIA Tesla V100 GPU, taking approximately 6–12 hours per model depending on dataset size.
Appendix B Discrete Oracle
The discrete oracle defines the maximum log-likelihood achievable under a uniform discretization granularity. For each state dimension, this maximum places all probability mass on the correct token.
- The discrete oracle is the maximum log-likelihood attainable under the chosen uniform discretization granularity.
- For one state dimension, the maximum is achieved by assigning all probability mass to the correct token.
- The total state log-likelihood is obtained by aggregating the per-dimension contributions.
Appendix C Baseline performance sources
Baseline performance comparisons draw reported results for MOPO, MBOP, BC, and CQL from the cited tables of prior studies.
- MOPO performance is taken from Table 1 of Yu et al. (2020), while MBOP performance is taken from Table 1 of Argenson & Dulac-Arnold (2021).
- BC performance is taken from Table 1 of Kumar et al. (2020a), and CQL performance is taken from Table 1 of Kostrikov et al. (2021).
Appendix D Datasets
The experiments use the D4RL dataset, which is distributed under the Creative Commons Attribution 4.0 License. License details are listed in the dataset’s “Licenses” section.
- The D4RL dataset is used in the experiments.
- D4RL is distributed under the Creative Commons Attribution 4.0 License (CC BY).
- The license information appears under the dataset’s “Licenses” section.
Appendix E Beam Search Hyperparameters
Beam width, context size, and planning horizon are decoding and trajectory-optimization hyperparameters. Planning can sometimes be accelerated by reducing beam width and horizon without affecting performance, with example configurations provided in the reference implementation.
- Beam width and context size are standard Transformer language-model decoding hyperparameters.
- Planning horizon is a standard trajectory-optimization hyperparameter.
- In many environments, beam width and horizon may be reduced to speed up planning without affecting performance.
- Figure 7 shows example Trajectory Transformer planner paths in MiniGrid-MultiRoom-N4-S5, with lock symbols indicating doors.
Appendix F Goal-Reaching on Procedurally-Generated Maps
The procedurally generated-map setting uses visual context embeddings because the map changes between episodes. Trained on goal-reaching and random trajectories, the model reaches 94% of testing goals on held-out maps.
- Because the map changes each episode, the Transformer receives an additional embedding derived from the current observation image.
- The discrete environment requires four actions—turning left, turning right, moving forward, and opening a door—to complete tasks.
- Training data mixes trajectories from a pre-trained goal-reaching policy and a uniform random policy.
- 94% of testing goals are reached on held-out maps.