ai theory
Reinforcement Learning Basics (18) — Decision Transformer Reads Trajectories Like Sentences
Junyoung Park · 2024-07-12 · 6 min
Introduction
Offline RL learns a Q-function or policy from a fixed dataset. Even there, the Bellman equation and policy gradient remain central.
Decision Transformer changes the point of view.
Treat a trajectory as a sequence of returns, states, and actions, and solve action selection as a sequence-modeling problem.
Just as a language model predicts the next token from preceding tokens, Decision Transformer predicts the next action from a desired return and the experience observed so far.
It does not eliminate the reinforcement-learning problem. It learns from the same trajectory data through a different formulation. This article examines the order of its input tokens and how return-to-go is used during both training and inference.
Turning a Trajectory into Tokens
Consider a trajectory from an offline dataset:
Decision Transformer groups the return-to-go, state, and action at every time step.
Return-to-go is the sum of discounted rewards remaining from the current time step.
Finite-horizon benchmarks often set and use the plain sum of the remaining rewards.
What Does Return-to-Go Ask For?
If the model receives only state , behavior may be ambiguous when the dataset contains several actions for the same state. Adding return-to-go distinguishes which actions occurred when a particular level of remaining performance was desired.
Suppose two trajectories begin from the same point.
With only the state, the two actions are mixed together. Once the return condition is included, the model can learn to distinguish RIGHT as the action taken in the high-return trajectory.
Return-to-go does not mean that the model knows the future in advance. During training, the episode has already ended, so we can calculate it. During inference, the user supplies a desired target return as a condition.
Training Is an Action-Prediction Problem
When a causal Transformer predicts the action at time , it cannot see future tokens.
For discrete actions, we can use cross-entropy.
For continuous actions, we can use MSE on the action prediction or the negative log-likelihood of a probability distribution.
There is no explicit Bellman backup in this training process. There is no target of the form
nor does the model perform new environment rollouts to calculate a policy gradient. It learns with a supervised objective over an offline sequence dataset.
During Inference, the Remaining Return Is Updated at Every Step
After training, we first choose a desired return .
- Provide the target return and current state .
- The model predicts action .
- The environment returns reward and next state .
- Reduce the remaining target by the received reward.
- Predict the next action with the new context .
In the simple case where ,
If the target is and the first reward is , then
The next context effectively says, “From the current state, I want to earn more.”
Does a Larger Requested Return Produce Better Behavior?
Not necessarily.
If the highest return in the dataset is but we provide during inference, the model receives a condition it never observed during training. The number alone does not create a new strategy.
In that case, action prediction may become unstable. In general, we should choose conditions within a return range that the dataset supports and evaluate several candidate values.
Return conditioning is an interface for requesting a desired level of performance, not a guarantee that a planner will achieve it.
Why Decision Transformer Is Appealing
It Can Use Long Histories
In partially observable problems, where the current observation does not fully reveal the state, earlier states and actions can serve as context.
One Model Can Represent Diverse Behaviors
By conditioning separately on low- and high-return trajectories, a single model can express behavior at several performance levels.
It Can Reuse Sequence-Modeling Tools
Transformer architectures, masking strategies, and large-scale supervised-training techniques apply directly.
The Difficulties of RL Have Not Disappeared
Decision Transformer inherits the limitations of its dataset.
- If the dataset contains no good trajectory, learning that behavior is difficult.
- A gap may exist between the requested return and what is actually attainable.
- Long sequences require more computation and context length.
- The model does not explicitly search through environment dynamics.
- Out-of-distribution generalization remains a problem in new states.
Nor does sequence modeling make every offline RL problem easier. A value-based method that exploits local rewards and Markov structure may be more data-efficient.
How Is It Different from Behavior Cloning?
Ordinary behavior cloning predicts an action from a state.
Decision Transformer also conditions on return and history.
This lets it distinguish behavior by desired return in a dataset containing both successes and failures. Still, because it ultimately predicts dataset actions through supervised learning, it remains closely related to imitation learning.
What to Remember
Decision Transformer converts offline trajectories into token sequences of return-to-go, state, and action, then predicts the next action conditioned on a desired return.
- Return-to-go is the sum of rewards remaining from the current point to the end of the episode.
- During training, is calculated from trajectories that have already finished.
- During inference, a target return is supplied and reduced by each received reward.
- The model learns with an action-prediction loss, without an explicit Bellman backup or on-policy RL loop.
- Supplying a high return absent from the dataset does not automatically create a new ability.
The next article connects RLHF, RLAIF, and DPO as methods for learning language-model behavior from human or AI preferences.