Decision Transformer

The core innovation of the Decision Transformer is its use of Return-to-Go (RTG) tokens. Instead of learning a value function, the model is conditioned on the desired cumulative future reward. At each timestep t, the input includes R_t = Σ_{k=t}^{T} r_k, the sum of future rewards. The model learns to output actions that achieve the specified RTG, enabling goal-directed behavior by simply setting a high target RTG at inference. This replaces the traditional reward maximization objective with a supervised sequence prediction objective.

The model generates actions autoregressively, one timestep at a time, similar to a language model predicting the next word. Given a sequence of past states, actions, and Return-to-Go tokens, it predicts the next action token. The training objective is a simple cross-entropy or mean-squared error loss on the predicted actions. This formulation:

• Leverages the scalability and stability of transformer-based sequence modeling.
• Avoids the instability and high variance often associated with policy gradient methods in traditional RL.
• Naturally handles offline RL settings, as it learns from static datasets without online environment interaction.

Decision Transformers are exceptionally well-suited for Offline RL (also known as batch RL). They learn purely from a fixed dataset of trajectories (state, action, reward) without any further environment interaction. By modeling the conditional distribution of actions given states and desired returns, they avoid the extrapolation error that plagues traditional offline RL algorithms like Q-learning, which must estimate the value of actions not present in the dataset. This makes them a robust choice for deploying policies trained on historical logs, such as robotic demonstrations or clinical treatment records.

The model utilizes a causal GPT-style transformer decoder. Tokens for state, action, and RTG at each timestep are embedded, summed with a timestep embedding, and fed into the transformer. The causal self-attention mask ensures predictions for timestep t only depend on information from timesteps ≤ t. This architecture provides:

• Long-range dependency modeling across hundreds of timesteps.
• Parallelized training via teacher forcing.
• A unified framework that can be scaled up with more layers and parameters, benefiting from advances in large language model infrastructure.

By conditioning on the RTG token, the Decision Transformer inherently performs goal-conditioned or reward-conditioned policy learning. To achieve a new task, one simply initializes the sequence with a high target RTG. The model then plans a sequence of actions predicted to achieve that return. This allows for flexible task specification without retraining. For example, in a navigation task, a high RTG could correspond to 'reach the goal quickly,' while a moderate RTG could correspond to 'reach the goal while conserving energy.'

A notable emergent capability is trajectory stitching. In offline datasets containing sub-optimal trajectories, the model can, at inference, condition on a high RTG and generate a novel action sequence that effectively 'stitches' together good segments from different parts of the dataset. It can combine the beginning of one trajectory with the successful ending of another to form a new, higher-return path. This demonstrates a form of implicit planning within the sequence space, going beyond simple behavioral cloning of the dataset.

The Decision Transformer excels in offline RL, where an agent must learn an optimal policy from a fixed dataset of past experiences without further environment interaction. By modeling the sequence of states, actions, and returns, it learns to condition its action predictions on a target return, effectively stitching together sub-optimal trajectories to achieve high performance. This is critical for real-world domains like robotics and healthcare, where online exploration is dangerous or expensive.

• Key Mechanism: It treats the desired cumulative reward (return-to-go) as an input token, allowing the model to generate actions aimed at achieving that specific return.
• Example: Training on a dataset of mixed-quality robot grasping attempts, a Decision Transformer can be instructed (via a high return-to-go token) to reproduce only the high-success sequences.

In robotics, Decision Transformers enable goal-conditioned behavior by representing the goal as a desired future state or a high return. The model autoregressively predicts actions that are most likely to lead from the current state to that specified goal.

• Process: The input sequence is structured as: (return-to-go_1, state_1, action_1, return-to-go_2, state_2, ...). For goal-reaching, the initial return-to-go is set to a high value, and the model predicts actions to achieve it.
• Advantage: This provides a more stable and interpretable learning signal compared to traditional RL's reward maximization, especially for sparse-reward tasks like "place the blue block on the red one."

When integrated with large language models (LLMs) or vision-language models (VLMs), the Decision Transformer framework extends to language-conditioned control. A natural language instruction (e.g., "open the top drawer") is encoded and used as the conditioning signal, replacing or augmenting the return-to-go token.

• Architecture Variant: Models like RT-2 (Robotics Transformer) tokenize images, language instructions, and robot actions into a single sequence, using a transformer to predict actions autoregressively.
• Use Case: This allows a single model to follow diverse, open-ended instructions by learning the correlation between language, visual context, and successful action sequences from large-scale datasets.

The Decision Transformer can be used for short-horizon planning by generating sequences of future states and actions. By sampling different initial return-to-go conditions, it can propose multiple potential trajectories, which can then be evaluated or refined.

• Method: In a model-based setup, a world model predicts the next state. The Decision Transformer can then plan a sequence of actions conditioned on a desired outcome, interacting with this learned model.
• Benefit: This combines the long-horizon reasoning of sequence models with the precision of dynamic control, useful for autonomous vehicle navigation or robotic arm motion planning.

Beyond robotics, the sequence-modeling approach is applicable to any sequential decision process. In recommendation systems, a user's interaction history (clicks, views, rewards) forms a trajectory. A Decision Transformer can be trained to predict the next best item to recommend, conditioned on a target level of user engagement (the return).

• Data Structure: States represent user/product context, actions are recommended items, and returns are engagement metrics (e.g., watch time, purchase).
• Advantage: It naturally handles long-term dependencies and non-Markovian aspects of user behavior within its context window, optimizing for cumulative engagement rather than just immediate reward.

Decision Transformers have shown strong performance on RL benchmarks like Atari games, D4RL locomotion tasks, and Go. They provide a compelling alternative to traditional actor-critic methods, particularly in offline settings.

• Characteristic Performance: They often achieve competitive or state-of-the-art results on offline RL benchmarks by avoiding the extrapolation error common in methods that learn explicit value functions from limited data.
• Example: On the D4RL dataset suite for MuJoCo locomotion, Decision Transformers effectively learn to imitate expert trajectories and interpolate between medium-quality data to achieve high returns.

The Return-to-Go (RTG) is the central conditioning variable in a Decision Transformer. It represents the cumulative reward the agent aims to achieve from the current timestep onward. The model is trained to predict actions that will achieve this specified return.

• Mechanism: At each timestep, the RTG is provided as an input token alongside the state and previous action. It is autoregressively updated: RTG_t = RTG_{t-1} - r_{t-1}.
• Role: It transforms the reinforcement learning problem into a conditional sequence modeling task. Instead of learning a value function, the model learns to output actions that are consistent with achieving the desired future return.

A Trajectory Transformer is a generative sequence model that treats an entire episode of interaction (states, actions, rewards) as a single sequence of tokens. The Decision Transformer is a specific instantiation conditioned on desired returns.

• Core Idea: It models the joint distribution p(states, actions, rewards) autoregressively. This allows for planning via beam search or sampling to find high-reward trajectories.
• Contrast with DT: While a Trajectory Transformer can model the full trajectory distribution, the Decision Transformer explicitly conditions on desired returns (RTG) to steer the generation towards high-reward behavior, making it more directly applicable to goal-directed control.

Offline Reinforcement Learning (Offline RL) is a paradigm where an agent learns a policy solely from a fixed, previously collected dataset of interactions, without any further online environment interaction during training. The Decision Transformer was explicitly designed for this setting.

• Challenge: Standard RL algorithms can fail in offline settings due to extrapolation error when evaluating actions not present in the dataset.
• DT's Approach: By framing control as conditional sequence modeling on the static dataset, the Decision Transformer avoids explicit dynamic programming or value function learning, which are prone to these errors. It imitates the return-conditioned behavior demonstrated in the data.

Autoregressive Sequence Modeling is a framework where a model predicts the next element in a sequence given all previous elements. This is the foundational technique used by the Decision Transformer.

• In Decision Transformers: The sequence is an interleaved token stream of (Return_t, State_t, Action_t). Given past tokens, the model predicts the next token (e.g., Action_t given Return_t and State_t).
• Architecture: Implemented using a causal Transformer decoder, which applies masked self-attention to ensure predictions for timestep t only depend on information from timesteps ≤ t. This is identical to the architecture used in large language models like GPT.

Conditional Sequence Generation is the task of generating an output sequence conditioned on some input context or specification. The Decision Transformer performs conditional generation of action sequences.

• Conditioning Signal: The primary condition is the initial Return-to-Go (RTG), which specifies the desired cumulative reward. This acts as a high-level goal or "prompt" for the action generator.
• Analogy to NLP: Similar to how a language model generates a story conditioned on a prompt ("Write a story about a robot"), the DT generates an action trajectory conditioned on a desired outcome ("Achieve a return of 1000").

Model-Based Reinforcement Learning (MBRL) involves learning a model of the environment's dynamics (the transition function and reward function) and using that model for planning or policy improvement. The Decision Transformer offers a contrasting, model-free perspective.

• Contrast: Traditional MBRL learns p(next_state, reward | state, action). The Decision Transformer, as a trajectory model, learns p(actions | states, desired_return).
• Planning: While MBRL uses its learned dynamics model for forward simulation (rollouts), a Decision Transformer can be used for planning via sequence generation, searching over action sequences that lead to high predicted returns under the model.