Decision Transformer: Reinforcement Learning via Sequence Modeling
Tell the model the score you want, and it picks the actions that earn it.
There is no value function here and no Bellman backup. A GPT is trained on recorded trajectories with every step labelled by the reward that actually followed it, and at test time you set that label to the score you are asking for.
Explaining the paperDecision Transformer: Reinforcement Learning via Sequence ModelingThe training data can be a pile of mediocre and failed attempts, because every attempt is a correct demonstration of the score it actually reached.
Reinforcement learning has a standard shape. You estimate how good each situation is, you improve the policy against that estimate, and you repeat. The estimate comes from a Bellman backup, which sets the value of one state from the value of the next, and much of the field is a catalogue of ways to keep that loop from falling over: target networks, trust regions, clipped objectives, pessimistic critics.
Decision Transformer deletes the loop. It has no value function, no Bellman backup, no policy gradient and no discount factor, and what is left is a language model. You write a recorded trajectory out as a sequence of tokens and train a GPT to predict the next action token the way it would predict the next word. The one unusual step happens before training: next to every timestep you write down the total reward that trajectory went on to collect from that point onward, and you feed that number in as an input.
At test time you choose the number yourself. Hand the model a state and the score you want, it returns an action; run the action, subtract the reward you just earned from the number, repeat until the episode ends. Nothing else runs. On the offline benchmarks of 2021 that recipe matched or beat the offline RL algorithms it was measured against, and on the ones with sparse or delayed rewards it beat them by a wide margin.
Five ideas carry the argument: why learning from a frozen dataset is hard, what a Bellman backup costs, why labelling a trajectory with its own outcome turns bad data into usable data, how a trajectory becomes tokens, and where asking for a score stops working.
Learning from a fixed pile of logs
Offline reinforcement learning is the setting where you are handed a fixed dataset of recorded transitions, each a tuple of state, action, reward and next state, and you have to produce a policy without ever touching the environment. Picture learning to drive from a hard drive of dashcam logs, with no car. Everything ordinary RL does, which is try something and see what happens and correct, is unavailable. You get one shot at producing a policy, and it gets driven for real afterwards.
That is harder than it sounds, so take it slowly. Q-learning improves a policy by asking, at every state in the dataset, what the best available action would have been at the next state. That question contains a maximum over all actions, including actions nobody in the dataset ever tried. A neural network will happily produce a number for those, and nothing in the data contradicts it. Then that invented number becomes the regression target for the current state, which is what bootstrapping means, so the fiction gets copied backward and compounded on the next pass. Online this self-corrects: you try the action that looked amazing, it is not, you update. Offline there is no correction loop at all. Values inflate while real performance collapses partway through training, and Levine and collaborators show that more data does not fix it, so it is not ordinary overfitting.
Note where the distribution shift sits. The critic is trained and queried on the same states, the ones in the dataset, so only the action inputs go out of distribution. That is why every fix in this family attacks the action argument.
The dominant fixes of the 2021 literature do it in one of two places. You can constrain the policy toward the data, which is what BEAR does with a support-matching penalty and BRAC does with a divergence penalty. Or you can make the critic pessimistic, which is Conservative Q-Learning: it adds a term that pushes down on actions it samples and up on actions present in the data, so the value it learns under-estimates rather than over-estimates. CQL is the strongest baseline Decision Transformer is measured against, and the guarantee covers less than the name implies: for the version actually run in these tables, the bound is on the expected value under the learned policy at dataset states, not a pointwise bound on every state-action pair.
The benchmark is D4RL, a suite of offline datasets with a normalisation that makes numbers comparable across environments:
So 0 is a random policy and 100 is the expert used to build the benchmark. Nothing clips it, and scores above 100 and below 0 both appear in D4RL's own tables. Its medium datasets come from a Soft Actor-Critic agent whose training was stopped early, well short of expert; Decision Transformer describes that policy as reaching roughly a third of expert performance, which is its own characterisation, not a figure D4RL publishes. The medium-replay datasets are the entire replay buffer of that run, so they contain everything from flailing to competent, and medium-expert is a fifty-fifty mixture of mediocre and good.
Every method named above spends its complexity budget taming one maximum inside a bootstrapped target. Decision Transformer never writes that maximum down, so it has nothing to be conservative about, and the paper says so directly in its discussion: it does not require policy regularisation or conservatism because it never optimises against a learned value function.
What a Bellman backup costs
To see what is being given up, look at the update Decision Transformer refuses to run. One-step temporal difference learning nudges the value of the state you were in toward the reward you just got plus the value of the state you landed in:
and Q-learning does the same for state-action pairs, with a maximum over next actions that makes it off-policy:
The bracket is the temporal-difference error, and its target contains the agent's own estimate. That is bootstrapping in one line: a guess corrected by a guess. It buys enormous data efficiency, because you learn from a single transition instead of waiting for the episode to finish, and it is what DQN and everything descended from it run on.
It also brings two costs that Decision Transformer is designed around.
Instability. Sutton and Barto name three ingredients whose combination makes divergence possible: function approximation, bootstrapping, and training on data that did not come from the policy you are evaluating. They call it the deadly triad, and their claim is specifically about the combination, with the companion result that dropping any one of the three is enough to avoid the danger. Decision Transformer drops bootstrapping and keeps the other two, since it is a neural network trained on somebody else's data. Instability stays a possibility rather than a certainty: van Hasselt and collaborators combined all three across 57 Atari games and thousands of configurations and never observed an unbounded blow-up, though values did rocket past a million before settling.
Speed. A synchronous backup is a local operator. Put a reward at the end of a corridor of states, start every value at zero, and sweep every state at once:
The set of states holding a non-zero value grows by exactly one per sweep, so the value at the first state stays zero until sweep . Online, with a step size , that first arrival is not only late but faint: the first non-zero value at the start of an -chain has magnitude , which at and is already .
Discounting carries its own cost. The discount factor makes an infinite sum of future rewards finite and shrinks every backup, which is what keeps the scheme contracting. It also blinds the agent past roughly steps, since a reward that far out has already been multiplied down to about . At that horizon is a hundred steps.
Decision Transformer says it avoids the need for discounting. That is worth stating more precisely than the paper does. Discounting is required for infinite-horizon problems, but Sutton and Barto explicitly permit whenever the episode terminates, and every environment in this paper terminates: episodes are capped at 1000 steps for the MuJoCo tasks, 100 for Reacher, and Atari and Key-to-Door both end on their own. So Decision Transformer declines a choice rather than removing a requirement. The limitation runs in the other direction: a return-to-go is a sum over the rest of the episode, which is undefined for a genuinely continuing task, so this method cannot be applied where discounting really is needed.
Every run succeeded at something
Decision Transformer changes the data pipeline and leaves the model alone. The dataset is already recorded and nothing in it is going to change, so for every timestep you can add up the rewards that actually followed and write that number down next to the step. That number is the return-to-go:
There is no discount factor in it and no estimate of anything. It is a sum over numbers you have on disk, and the official code computes it exactly that way, as a cumulative sum with . Every step in every trajectory gets one, including the terrible trajectories.
The bad trajectories are why this works. Ordinary imitation learning is stuck with them: copy bad players and you become a bad player. Return conditioning changes the question being asked. You are no longer training a model of "what a good player does". You are training a model of "what a player who is about to score does next", and under that question a disastrous run is not waste, it is a perfectly valid example of how to score badly. The whole dataset becomes labelled training data for a single conditional model, and you pick which behaviour you want at inference time by choosing the label.
Be careful with the slogan version. A trajectory is not optimal for the return it achieved, not once the environment has any randomness in it. The precise statement is that a trajectory is a valid example of achieving the return it achieved. Kumar, Peng and Levine put it as optimal supervision for matching a return rather than for maximising one, and the section on limits below is entirely about the gap between those two things.
The figure below shows the labelling on a twenty-step episode. The amber bars are the reward collected at each step; the teal staircase above them is the return-to-go, which starts at the episode total and walks down to zero as the reward gets spent. Drag the delay control to move all of that reward to the final step, the way the paper's sparse-reward experiment does, and watch two things: the staircase flattens into a horizontal line, because no reward arrives to spend it down until the end, and the number at does not move at all.
That invariance is why the delayed-reward experiment goes the way it does. A TD method loses almost all of its learning signal when the rewards move to the end, because its per-step target becomes zero everywhere and the one real reward has to crawl back along the chain. Decision Transformer loses only the shape of the label, not the label at the start of the episode, which is the one it is prompted with.
Now put the two schemes side by side on the same corridor. Press Play and watch the value crawl leftward one state per sweep on the top row; the bottom row is the same corridor with returns-to-go written on it, correct everywhere before training starts.
That picture is the paper's own framing, and taken literally it is stronger than the evidence supports. First, the one-state-per-sweep rate is a property of one-step, forward-ordered backups. n-step returns move credit n links at a time, eligibility traces spread it across the episode, and prioritised sweeping or reverse-order replay crosses the chain in a single pass. Those cures cost variance rather than time, but they exist, so the fair comparison is against the one-step baselines that CQL, REM and QR-DQN are. Second, the bottom row is not a value function. A return-to-go is what one recorded run happened to collect, not an expectation over what the policy would collect, and that difference is what the last section is about.
The idea has a lineage. Hindsight Experience Replay relabels a failed reach as a successful reach of wherever the hand ended up. Upside-Down RL and Reward-Conditioned Policies both train a supervised policy conditioned on a desired return. Decision Transformer says as much, and describes Reward-Conditioned Policies as roughly its own method with a context of one timestep. What it adds is the sequence model and the long context, and the ablation later measures exactly what that context is worth.
A trajectory as tokens
To hand a trajectory to a Transformer you have to decide what a token is. Decision Transformer uses three per timestep, always in the same order: the return-to-go, then the state, then the action.
The model reads the last timesteps, so a forward pass is over tokens. Each modality gets its own linear layer into the embedding width, and a visual state gets a convolutional encoder instead. Counting from zero, timestep puts its return-to-go at sequence position , its state at , and its action at .
Now the part that makes this a policy rather than a model reading its own answer. The action head does not read the action token. It reads the hidden state sitting on the state token, at position , and turns that into a prediction of . Because attention is causally masked, position can attend to positions through and no further:
In words: every complete earlier triple, plus this timestep's return-to-go and state, and not this timestep's action. What the network computes is a return-conditioned policy over the visible history, .
Drag the control below to move the prediction to a different timestep. The shaded region is what the head can read; the token ringed in amber is the answer it is being trained to produce, and it sits one slot to the right of the shading at every setting.
That index fact explains a line in the released code that looks like a bug. During a rollout the loop appends a row of zeros as a placeholder action before asking the model for the real one, then overwrites it. Feeding garbage in sounds alarming until you notice the placeholder occupies position , strictly to the right of the state token the answer is read from, and nothing to the right of a position can influence it. The placeholder is provably inert. The Atari codebase solves the same problem differently, by simply truncating the sequence to tokens at evaluation so the trailing action slot does not exist.
The position signal. A standard Transformer adds one learned vector per token position. Decision Transformer instead learns one vector per episode timestep and adds the same vector to all three of that timestep's tokens, which the paper flags as a deliberate departure. A reader should ask what then tells the model which of the three a token is, and the answer is the three separate input projections: maps a scalar, maps a state vector (eleven numbers for Hopper), and maps an action vector (three), each with its own bias, so the three modalities land in different regions of the embedding space and the biases alone act as a learned type tag. The causal mask supplies the rest, since a token can count how many predecessors it has.
The Atari implementation does not match that description. It adds the per-episode-timestep embedding and a standard learned per-token positional embedding over the block, one on top of the other. Only the MuJoCo codebase does what the prose says.
Normalisation. Three token types with wildly different units share one sequence: a return-to-go in the thousands, a state in physical units, an action in . Layer normalisation puts them on comparable footing, because it rescales each token on its own, over its features, with no reference to the batch or to other positions. After the normalisation step every token vector sits on the same sphere of radius , whatever it started as.
The paper describes this as a linear layer per modality followed by layer normalisation. The code does something slightly different: it adds the timestep embedding first, then applies a single shared normalisation to the stacked sequence.
with the timestep vector. Since for any constant , normalising after the addition throws away the flat component of the timestep embedding and forces the modality content and the timestep signal to share one fixed budget. Normalising before it, as the prose reads, would let the timestep term push the token off that sphere by however much its norm grows during training.
The rest is stock GPT. The MuJoCo implementation is a copy of the HuggingFace GPT-2 file with the position-embedding lookup commented out, because Decision Transformer supplies its own. Attention is the usual scaled dot product,
where is the per-head key width, 128 here since the MuJoCo model runs a single head at width 128. The divisor keeps the logits at unit scale: dot products of unit-variance terms have variance , and a softmax over large numbers is a spike with almost no gradient flowing through it. The paper's own equation for attention omits that divisor, and also prints the softmax over all tokens with the causal restriction stated only in the prose beside it. Both are simplifications rather than errors: the same equation drops multi-head splitting, the output projection, the residuals, the two layer normalisations and the feed-forward block, and the paragraph around it says outright that architecture details are deferred to the original papers. Both released implementations divide. The omission happens not to disturb the paper's reading of attention as credit assignment either, because a positive constant cannot change which key a query prefers.
The model is small enough to be surprising. Three layers, one attention head, width 128, feed-forward width 512. The GPT trunk comes to about 0.60M parameters and the Hopper model in full to about 0.73M, of which 128,000 is the timestep embedding table alone. That is roughly one-160th of GPT-2 small, and it is the model that produces the D4RL numbers below.
Training and rollout, concretely
The training objective is the same squared error a behaviour-cloning baseline would use. Sample a window of timesteps, run one forward pass, read the action predictions off the state tokens, and average:
with the set of positions that are real data instead of left padding, so short windows near the start of an episode do not dilute the loss. Discrete actions swap the squared error for cross-entropy over the action vocabulary, which is what the Atari model does. States and returns have prediction heads too, and the paper reports that training them did not help, so their outputs are computed and discarded.
Concretely, for Hopper: the state is 11 numbers, the action is 3, the context is timesteps and the batch is 64. The raw batch carries states , actions , and returns-to-go of length that get sliced back to . Four embedding layers produce four tensors, the interleave makes , the attention map is , and the action head emits . Masked selection leaves at most 1280 valid actions and the loss is one number. Every one of the 60 tokens gets a hidden state; only the 20 state tokens carry loss.
# one training step. K = context length, B = batch size
R, s, a, t, mask = sample_windows(data, B, K) # R:(B,K,1) s:(B,K,ds) a:(B,K,da)
z = embed_t(t) # ONE vector per timestep
seq = interleave(embed_R(R) + z, # -> (B, 3K, d)
embed_s(s) + z,
embed_a(a) + z)
h = gpt(layer_norm(seq), causal_mask, mask) # (B, 3K, d)
a_hat = tanh(W_a @ h[:, 1::3]) # read the STATE tokens -> (B,K,da)
loss = mean((a_hat - a)**2, over=mask) # MSE; cross-entropy if discrete
loss.backward(); clip_grad_norm(0.25); opt.step()The released code adds two steps the paper never mentions, and both change what you would type if you reimplemented it: states are normalised by the dataset mean and standard deviation, and returns-to-go are divided by a constant before being embedded, 1000 for the locomotion tasks and 10 for Reacher, so a Hopper target of 3600 enters the network as 3.6 and drops by about 0.003 per step. No comment in the repository explains that constant. The likely reason is resolution: without it, swamps its own bias and the timestep vector added alongside it, and after normalisation the return token collapses toward a near-constant direction, taking most of the difference between one requested return and another with it. The optimiser is AdamW at a learning rate of with a warmup, for gradient steps.
At rollout the return-to-go stops being a label and becomes a control. You pick a target, feed it as the first token, take the action the model returns, and then decrement:
which is the same recursion as the definition (5), run forward instead of backward. The number in the prompt is a budget of reward still owed, and every step spends part of it. Ask for 3600 on Hopper, earn 3.5 on the first step, and the model's next prompt is 3596.5, so a step that went badly leaves a larger number in the prompt for the steps that remain.
# rollout. target is the score you are asking for (3600 for Hopper)
R, s, a, t = [target / scale], [env.reset()], [], [0]
for step in range(max_steps):
a.append(zeros(da)) # placeholder; sits right of the state token
a[-1] = dt(R, s, a, t)[-1] # the action read off the last state token
obs, r, done, _ = env.step(a[-1])
if done: break
R.append(R[-1] - r / scale) # spend what you just earned
s.append(obs); t.append(step + 1)
R, s, a, t = R[-K:], s[-K:], a[-K:], t[-K:]In the sparse-reward variant the decrement is skipped entirely, which is consistent rather than special-cased: when every reward along the way is zero there is nothing to subtract. Nothing in either listing is an RL algorithm: no critic, no target network, no importance weighting, no trust region. The training loop is supervised regression and the rollout loop is autoregressive generation with a running counter.
Recombining pieces of bad data
Imitating a fixed dataset can reproduce what is in it. The interesting claim is that conditioning can produce behaviour better than anything recorded, by assembling it out of fragments. The paper opens with a toy that shows what that would mean.
Take a directed graph with a goal node and score it as an RL problem: reward 0 at the goal, and at every other node. A return-to-go is then minus the number of steps still to walk, so asking for the largest achievable return is asking for the shortest path. Train on random walks, which contain no expert demonstration of anything, then at each node ask for the best return still reachable through transitions the data contains, and follow the action that goes with it.
The figure below runs exactly that. Tap a node to start there. Amber arrows are the transitions the random walks actually recorded, thicker where they were walked more often; grey arrows exist in the graph but never appeared in the data, and the greedy walk is not allowed to use them, because a model cannot generate a transition it has never seen. Turn the walk count down and watch coverage break: regions of the graph become unreachable and the path gets longer or fails outright.
The paper reports that 15.8% of the paths its model generated were not sub-sequences of any training trajectory, so they were assembled from pieces. The appendix it lives in also supplies everything that undercuts it. The model used there is not the one used anywhere else in the paper: it generates return tokens as well as action tokens. Generation is steered by a hand-designed prior over returns that favours short paths, , combined with the model's own return distribution as , and the appendix says plainly that this prior is not used in any other experiment. And the environment is a deterministic 20-node graph with 1000 random walks of ten steps each. The remaining 84% of generated paths were cases where a training walk already contained the answer. The word "stitching" appears exactly once in the whole paper, in that appendix paragraph.
The general question of whether return conditioning recombines fragments was settled later, and the answer is no. Brandfonbrener and collaborators construct a fully deterministic MDP where any conditioning value that makes the method prefer the good action needs on the order of trajectories, with the episode length, because return conditioning needs coverage of whole-trajectory returns while dynamic programming needs only coverage of individual transitions. Nothing about that failure involves randomness in the environment. So what the graph toy establishes is compositional generalisation inside a small deterministic problem with good coverage, not a general substitute for dynamic programming, and the coverage slider in the figure above shows exactly that limitation.
How far the return dial reaches
If the prompt is a dial, the obvious question is whether turning it does anything. The paper sweeps the target return across a wide range on every task and reports that the return asked for and the return collected track each other closely, matching almost exactly on Pong, HalfCheetah and Walker.
Why that works, and where it has to stop, follows from what conditioning is. A model trained on (return, state) pairs reproduces the behaviour of the episodes that actually scored near what you asked for. Ask for something the dataset contains and there is a population of episodes to imitate. Ask for something past the best episode ever recorded and the request has no support, so the weight piles onto the same handful of top episodes and the answer stops moving.
The figure below is a toy that makes that mechanism visible. Six hundred episodes, returns drawn from a fixed mixture, nothing above 72. The teal curve is the return a return-conditioned policy collects, and the dashed diagonal is perfect obedience. Drag the request past the amber line and watch the curve leave the diagonal and flatten.
The paper reports beating that ceiling on Seaquest, and the appendix lets you put a number on it. Its conditioning target of 1450 is annotated as roughly five times the best return in its own slice of the data, which puts the best recorded episode near 290, and Decision Transformer averaged a raw score of 1129 across three seeds. Same story on Qbert: a target of 2500 at about five times the dataset maximum, so a best recorded episode near 500, against a raw average of 2216. Both are well past anything in the training set. The paper describes this as being sometimes capable of extrapolation, and "sometimes" is doing real work in that sentence: on Breakout and Pong the targets are annotated at one times the dataset maximum, so they are requests the data supports.
What the benchmarks say
Three comparisons, each making a different point. Switch between them with the tabs.
D4RL. Averaged over the nine locomotion settings, Decision Transformer scores 74.7 against CQL's 63.9, BEAR's 48.2, BRAC-v's 36.9 and behaviour cloning's 46.4. The average hides a lot of variation. It beats the best baseline on medium-expert HalfCheetah by 24.4 points and beats CQL on medium-replay Walker by 39.9, sits a few points behind on medium HalfCheetah, and loses medium-replay HalfCheetah to BRAC-v by 11.1. What stands out across the grid is that BRAC-v posts near-zero scores on three settings and BEAR falls to 19.2 on another, while Decision Transformer never drops below 36.6 on the nine locomotion settings. Its floor is high because there is no learned value function for a bad update to exploit.
Atari. Four games, three seeds, 1% of the DQN Replay dataset, which is 500,000 of the 50 million transitions logged over an online DQN agent's entire training run. That dataset spans everything from a flailing beginner to a finished DQN, which is exactly the mixed-quality regime return conditioning is supposed to sort out. Read the four games one at a time instead of as one result: Decision Transformer wins Breakout and Seaquest, essentially ties Pong, and loses Qbert badly, 15.4 against CQL's 104.2. REM and QR-DQN, the two baselines carried over from the DQN Replay study, print 0.0 on Qbert, which does not mean they scored nothing; their raw scores of 160.1 and 156.0 land just under the 164-point random-play anchor, so the normalisation rounds them to zero.
The Breakout number deserves unpacking, because 267.5 reads like a claim about superhuman play and is not one. What that column reports is a ratio rather than a score: the raw result was 76.9 points, and the normalisation puts 100 at the score of a professional games tester who had about two hours of practice, which for Breakout is 30 points. Beat 30 by 2.7 times and you print 267.5. The registered human world record on Breakout is 864. The ratio is large because the denominator is small, and it is sensitive to which human anchor you use: with the numbers from the original DQN paper the same 76.9 prints 249.8 instead. Pong is the one row whose arithmetic does not close, since 17.1 raw against anchors of −21 and 15 gives 105.8 against the printed 106.1, though the standard deviation does reconcile. And the REM and QR-DQN columns cannot be found in the published CQL and REM tables, Kumar et al. 2020 and Agarwal et al. 2020: the appendix says they were re-derived from raw data the REM authors provided directly, which contradicts the main text's claim that the numbers come straight from those two papers.
Context length. The ablation that surprised me most. The context is 30 steps on Breakout and 50 on Pong; cutting to 1, so the model sees one return-to-go, one state and nothing else, takes Breakout from 267.5 to 73.9 and Pong from 106.1 to 2.5. Frame stacking is already in place, so the current frame is not missing motion information, and conventional wisdom says a Markov policy should not need history. The paper offers a hypothesis rather than a result: when the data mixes many policies of differing quality, the context lets the model identify which policy generated the actions in its window, and conditioning on that makes the prediction problem easier.
Delayed rewards. Move every reward in the D4RL datasets to the final timestep and CQL falls off a cliff: medium-expert Hopper from 111.0 to 9.0, medium-replay Hopper from 48.6 to 2.0. Decision Transformer goes from 107.6 to 107.3 and from 82.7 to 78.5. This is Figure 1 made real. The label at the start of the episode is unchanged by the delay, and it is the only label the rollout needs.
Key-to-Door. The long-horizon test. Three rooms in sequence: pick up a key, cross an empty room, reach a door, and collect a binary reward at the door only if you took the key at the very beginning. Trained on purely random trajectories, Decision Transformer reaches the door with the key 71.8% of the time from a thousand of them and 94.6% from ten thousand, against 13.1% and 13.3% for CQL. This is the only experiment in the paper that uses the entire episode as context instead of a fixed window, so it is the only one where attention has unbounded reach. And percentile behaviour cloning, trained only on the episodes that succeeded, scores 69.9% and 95.1%. That table carries no error bars and no seed count, so treat the two as indistinguishable. On a task whose reward is binary, a return-to-go collapses to a per-episode success flag on every prefix, which makes conditioning on it and filtering on it the same operation, and the entire gap over CQL is attributable to the hindsight labelling rather than to sequence modelling.
That brings up the paper's own sharpest test. Percentile behaviour cloning sorts the data by episode return and clones only the top of timesteps. It is not a usable method, because choosing requires rolling out in the environment, which offline RL exists to avoid, and the paper says so. As a diagnostic it is exactly right, since it asks whether return conditioning does anything that filtering does not.
On the eight D4RL settings where both were run, 10%BC averages 56.7. The paper prints 56.1 for Decision Transformer on the same rows, but those eight numbers add to 439.3 and average 54.9; every other column in that table reproduces exactly. Corrected, filtering with a hand-picked percentile edges Decision Transformer on the D4RL average rather than losing to it, and 10%BC wins five of those eight rows outright. Decision Transformer's genuine wins sit elsewhere: on Hopper medium-replay it scores 82.7 against 70.6, and on Atari, where the dataset is 1% of a replay buffer, filtering to the top 10% throws away most of an already-small dataset and 10%BC drops to 28.5 on Breakout against 267.5. So the supported claim is narrower and more interesting than the abstract suggests: conditioning is a better way to spend a small dataset than filtering, because it fits all of the data with one model and chooses the subset at inference time instead of at training time.
The Gym behaviour-cloning baselines are MLPs, which the authors found stronger there than a transformer, so on D4RL the comparison mixes return conditioning with architecture and context length; the Atari comparison is clean, since behaviour cloning there is the same model and hyperparameters minus the return token. And the Atari protocols do not match: Decision Transformer evaluates with sticky actions turned off, while the dataset it trains on and the baselines it is measured against use the standard setting where the environment repeats your previous action a quarter of the time, which makes the games harder. The bias runs the other way too, since REM early-stopped on online evaluation and reported each agent's best score during training, averaged over five agents. The right conclusion is that the columns were not measured identically, not that any of them is wrong.
Where return conditioning breaks
Everything above works because the environments are deterministic or nearly so. Atari ROMs, MuJoCo, a fixed graph, a grid world. Once real randomness enters, asking for a high return stops meaning what you want it to mean.
The mechanism is short to state. In the infinite-data limit, a return-conditioned policy is exactly the behaviour policy reweighted by how often each action reached the requested return:
with the return you asked for. Read the fraction carefully: it is a probability of hitting the target, never an average outcome. An action with a one-in-a-hundred chance of a jackpot is promoted at that jackpot exactly as hard as an action that delivers it every time. Conditioning on an outcome cannot distinguish the action that caused it from the dice that happened to fall well.
The figure below is the smallest version of that. One state, two actions. The safe one pays 5 every time. The gamble pays 10 with probability and 0 otherwise. The dataset holds a hundred episodes of each, so as soon as a single gamble comes up lucky, the best return anyone ever recorded is 10, and every episode carrying that label took the gamble. Ask for the best return you have seen and you get the gamble, and what you collect is .
Paster, McIlraith and Ba, who identified this, give a sharper example: three actions where the one that produced the single best observed reward has an expected reward of , and no target return at all produces a positive expectation, at any dataset size. Brandfonbrener and collaborators give the matching impossibility: a one-state problem where a rare action almost always pays and a common action almost never does, arranged so that conditioning on success gives exactly a coin flip for every conditioning function, at infinite data. The bias is not a sampling artefact and more data does not remove it.
The other question the follow-up work settled is how much of this needs a Transformer. Emmons and collaborators reran the D4RL comparison with a two-layer feed-forward network conditioned on the same return information and matched Decision Transformer on the locomotion tasks. On AntMaze, the D4RL suite built specifically to require stitching together separate trajectories, both fall apart: filtered behaviour cloning averages 41.4 against Decision Transformer's 19.8, while implicit Q-learning, which does propagate values, scores 378 across the AntMaze totals against Decision Transformer's 112. The win does not come from the architecture, and on the tasks that genuinely need value propagation a method with no value function loses.
None of that undoes the contribution. A 0.7M-parameter GPT trained with plain regression on logged trajectories, with no critic and no conservatism, held its own against methods built specifically for offline RL, and the delayed-reward and Key-to-Door numbers are the kind of gap that changes what people try next. What the five years since have clarified is the boundary: hindsight labelling defeats the horizon, attention defeats the context length, and neither of them defeats an expectation. Return conditioning suits a predictable environment whose dataset covers the score you are asking for. Outside that, a method that takes an expectation still wins, which is what the AntMaze numbers show.
Questions you might still have
If the model only imitates recorded behaviour, how can it beat the recorded behaviour?
Two ways, and they are different. Within the data it can pick out the good episodes: it was trained on every episode labelled with its own outcome, so asking for a high return selects the behaviour of the episodes that got one. Beyond the data it can sometimes recombine, and the paper shows this on a 20-node graph where 15.8% of generated paths were not sub-sequences of any training walk. That second ability is much weaker than that number suggests. Brandfonbrener and collaborators showed return conditioning needs coverage of whole-trajectory returns, not just of individual transitions, so on benchmarks built to require stitching it does badly.
Why is the action read off the state token instead of the action token?
Because the action token is the answer. Timestep i puts its return-to-go at sequence position 3i, its state at 3i+1 and its action at 3i+2. The causal mask lets position 3i+1 attend to everything up to and including itself, which is every earlier triple plus this timestep’s return-to-go and state, and nothing later. Reading the prediction off 3i+1 therefore conditions on exactly the right history. It also explains why the placeholder action the rollout code appends before asking for a prediction cannot matter: it sits at 3i+2, to the right of the token the answer comes from.
What number should I put in the return prompt at test time?
The paper picks targets from expert performance for each environment, except HalfCheetah where it uses about half, because those datasets contain relatively low returns. On Atari it uses either one times or five times the largest return in its slice of the data. There is no principled procedure and the paper says so. In practice the dial works while the dataset supports the request and saturates past the best episode it contains, which Figure 5 shows directly, so the useful range is bounded by your data.
Does this need a Transformer at all?
On the locomotion benchmarks, no. Emmons and collaborators matched these D4RL numbers with a two-layer feed-forward network given the same return conditioning, and a percentile behaviour-cloning baseline ties on the same rows. The context length does matter on Atari, where cutting it from 30 timesteps to 1 takes Breakout from 267.5 to 73.9, and it matters on Key-to-Door, the one experiment that uses the entire episode as context.
What breaks if the environment is random?
The conditioning stops selecting for good actions and starts selecting for lucky ones. In the infinite-data limit a return-conditioned policy is the behaviour policy reweighted by the probability of each action reaching the requested return, which is a probability rather than an average, so a rare jackpot is promoted as hard as a reliable payoff. Paster, McIlraith and Ba construct a three-action problem where no target return produces a positive expected reward at any dataset size. Every environment in this paper is deterministic or close to it, which is why the failure does not appear in its tables.
Footnotes & further reading
- The paper: Chen, Lu, Rajeswaran, Lee, Grover, Laskin, Abbeel, Srinivas, Mordatch, Decision Transformer: Reinforcement Learning via Sequence Modeling (NeurIPS 2021). Code. Every claim here about token layout, the return-to-go scale constant, state normalisation and the Atari positional embeddings comes from that repository, not from the paper.
- The offline RL setting and the mechanism behind value over-estimation: Levine, Kumar, Tucker and Fu, Offline Reinforcement Learning: Tutorial, Review, and Perspectives. The conservative baseline: Kumar, Zhou, Tucker and Levine, Conservative Q-Learning.
- The benchmark and its normalisation: Fu, Kumar, Nachum, Tucker and Levine, D4RL. The Atari dataset: Agarwal, Schuurmans and Norouzi, An Optimistic Perspective on Offline Reinforcement Learning, which introduced the DQN Replay dataset and REM. The gamer-normalised protocol and its anchors: Hafner, Lillicrap, Norouzi and Ba, Mastering Atari with Discrete World Models.
- The ancestors of return conditioning: Andrychowicz et al., Hindsight Experience Replay; Srivastava, Shyam, Mutz, Jaśkowski and Schmidhuber, Training Agents using Upside-Down Reinforcement Learning; Kumar, Peng and Levine, Reward-Conditioned Policies.
- The limits: Paster, McIlraith and Ba, You Can't Count on Luck: Why Decision Transformers and RvS Fail in Stochastic Environments; Brandfonbrener, Bietti, Buckman, Laroche and Bruna, When does return-conditioned supervised learning work for offline RL?; Emmons, Eysenbach, Kostrikov and Levine, RvS: What is Essential for Offline RL via Supervised Learning?; Kostrikov, Nair and Levine, Offline Reinforcement Learning with Implicit Q-Learning.
- The concurrent paper with the same instinct and a model-based twist: Janner, Li and Levine, Offline Reinforcement Learning as One Big Sequence Modeling Problem (Trajectory Transformer), which discretises everything and predicts states and returns as well as actions.
- The deadly triad, the propagation rate, and the conditions for : Sutton and Barto, Reinforcement Learning: An Introduction (2nd ed.), sections 3.3, 4.1, 6.1, 7.1 and 11.3. The empirical study of the triad at scale: van Hasselt, Doron, Strub, Hessel, Sonnerat and Modayil, Deep Reinforcement Learning and the Deadly Triad.
How could this explainer be improved? Found an error, or something unclear? I read every message.