Decision Transformer: Reinforcement Learning via Sequence Modeling

Label each logged step with the reward that followed it, train a GPT on that, then prompt it with the score you want.

Decision Transformer writes each recorded trajectory as a sequence of tokens (the reward still to come, the state, the action) and trains a GPT to predict the action. There is no value function and no Bellman backup. At test time the first reward-still-to-come token is set to the score you are asking for.

A three-layer, one-head GPT of 0.73M parameters, trained by plain regression on datasets of mostly mediocre or failed attempts, matched the offline RL algorithms of 2021 on the D4RL locomotion benchmark.

Most reinforcement learning (RL) algorithms estimate how good each situation is and improve the policy against that estimate, over and over. The estimate comes from a Bellman backup, which sets the value of one state from the value of the next, and a long list of techniques exists to keep that loop stable: target networks in DQN, trust regions in TRPO, the clipped objective in PPO, pessimistic critics in offline RL.

Decision Transformer (Chen, Lu and seven coauthors, June 2021) has no value function, no Bellman backup, no policy gradient and no discount factor. It writes a recorded trajectory as a sequence of tokens and trains a GPT-style causal Transformer to predict each action the way a language model predicts the next word. Before training, every timestep gets one extra input: the total reward that trajectory collected from that step to the end. At test time you choose that number. You give the model the current state and the score you want, it returns an action, you subtract the reward the action earned from the score, and you repeat until the episode ends.

On the offline benchmarks it was tested on, this matched or beat the dedicated offline RL methods. Over nine D4RL locomotion datasets it averages 74.7 against 63.9 for Conservative Q-Learning, its strongest baseline. When every reward in a dataset is moved to the final timestep, CQL drops from 111.0 to 9.0 on one Hopper dataset and Decision Transformer goes from 107.6 to 107.3. Follow-up work in 2022 found where it fails: in environments with real randomness, and on datasets where good behaviour has to be assembled from pieces of different trajectories.

The sections below cover, in order: why learning from a fixed dataset is hard, what a Bellman backup costs, how labelling each step with its own outcome makes bad data usable, how a trajectory becomes tokens, how training and rollout run on real tensor shapes, and where asking for a score stops working.

Why offline reinforcement learning is hard

Offline reinforcement learning means producing a policy from a fixed dataset of recorded transitions, each a tuple of state, action, reward and next state, without ever interacting with the environment. Ordinary RL tries an action, sees what happens and corrects itself. Offline, the only evidence is what someone else already logged. Picture learning to drive from a hard drive of dashcam footage with no car: you can study every decision that was recorded, but you cannot test one that was not.

Q-learning, the usual starting point, fails here in a specific way. Its update target at a state is the reward plus the best value available at the next state, $\max_a Q(s', a)$. That maximum runs over every action, including actions the dataset never contains. A neural network outputs a number for those actions anyway, and no data point contradicts it. If the network's error on some untried action happens to be positive, the maximum selects that action, and the inflated number becomes the regression target for the previous state. Reusing the model's own output as a target in this way is called bootstrapping, and it copies the inflation one step further back on every round of updates. An online agent would eventually take the action that looked good, collect a poor return and correct the estimate. Offline there is no such correction, so estimated values climb while the real performance of the policy falls. Collecting more data from the same logging policy does not help, because the untried actions stay untried.

The figure below shows the first step of that failure with one state and a continuous action between −1 and 1. A logging policy tried 24 actions centred on −0.25 and recorded a noisy return for each. A flexible regressor (a degree-5 polynomial) is fitted to those 24 points, and the policy takes the action where the fitted curve is highest. The dashed curve is the true value, which peaks at 1.0 at $a = 0.35$ and which the learner never sees. Look at where the vertical line marking the chosen action lands relative to the shaded band of logged actions.

With a narrow logging policy, the fit is pinned down only near the data. Outside it, the polynomial's tails go wherever the noise pushes them, and the maximum usually lands on a tail that bends upward, often at the edge of the action range. Widen the spread to $\sigma = 0.9$ and the data covers the whole range: the fitted maximum averages 1.03 and the chosen action is worth 0.97. Offline datasets usually resemble the narrow case, since a logging policy concentrates on the actions it was trying to take.

Two families of fix were standard in 2021, and both act on the action argument of $Q$. The critic is trained and queried on the same states, the ones in the dataset, so only the actions go out of distribution. Policy constraints keep the learned policy close to the logged actions: BEAR uses a support-matching penalty and BRAC a divergence penalty. The figure's second tab is the crudest version, a hard limit to the logged range, which gives an accurate estimate and settles for the best action that was tried. Pessimistic critics lower the estimate itself: Conservative Q-Learning (CQL) adds a loss term that pushes $Q$ down on actions it samples and up on actions present in the data, so the values it learns under-estimate. CQL is the strongest baseline in this paper. Its guarantee covers less than the name suggests: for the variant run in these tables, the bound is on the expected value of the learned policy at dataset states, not on every state-action pair.

Decision Transformer never computes $\max_a Q$, so it has nothing to constrain or make pessimistic. The paper's discussion says it needs no policy regularisation or conservatism, and gives the reason as a conjecture: TD-learning methods need them because they optimise against a learned value function, and Decision Transformer does not optimise against any learned function.

The continuous-control results use D4RL, a suite of offline datasets with a normalised score that makes numbers comparable across environments:

A random policy scores 0 and the expert policy used to build the benchmark scores 100. Nothing clips the scale, and Decision Transformer's own table has 107.6 and 108.1. The paper uses three datasets per environment. medium is one million timesteps from a policy that the paper describes as reaching about a third of expert score; D4RL produced it by stopping a Soft Actor-Critic run early. medium-replay is the entire replay buffer of that run, from flailing to competent. medium-expert concatenates one million medium timesteps with one million expert timesteps.

What a Bellman backup costs

One-step temporal-difference (TD) learning, the update Decision Transformer drops, moves the value of the state you were in toward the reward you received plus the value of the state you landed in:

Q-learning does the same for state-action pairs, with the maximum over next actions from the previous section:

The bracket is the TD error. Its target contains the agent's own estimate of the next state, so one guess is corrected toward another. In exchange the agent learns from a single transition instead of waiting for the episode to end, which DQN and everything descended from it rely on. Bootstrapping 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 generated by a different policy from the one being evaluated. They call this the deadly triad. The claim is about the combination, and removing any one of the three avoids the danger. Decision Transformer removes bootstrapping and keeps the other two, since it is a neural network trained on someone else's data. Divergence with all three present is possible, not certain: van Hasselt and collaborators ran all three across 57 Atari games and thousands of configurations, saw no unbounded blow-up, and saw values climb past a million before settling.

Slow propagation. A synchronous backup only looks one state ahead. Put a reward at the end of a chain of $N$ states, start every value at zero, and update every state at once from the previous sweep's values:

After sweep $k$ only the last $k$ states hold a non-zero value, so the first state stays at zero until sweep $N$. With an online step size $\alpha < 1$ the first arrival is also faint: each sweep moves a value only a fraction $\alpha$ of the way to its target, so with $\gamma = 1$ the first non-zero value at the start of the chain is $\alpha^N$ times the reward. At $\alpha = 0.1$ and $N = 5$ that is $10^{-5}$.

Discounting. A discount factor $\gamma < 1$ makes an infinite sum of future rewards finite and shrinks every backup, which keeps the scheme contracting. It also makes the agent nearly blind beyond about $1/(1-\gamma)$ steps. At $\gamma = 0.99$ that horizon is 100 steps, and a reward 100 steps away is weighted by $0.99^{100} \approx 0.37$, close to $1/e$.

The paper says Decision Transformer avoids the need for discounting. Stated precisely: discounting is required for infinite-horizon problems, but Sutton and Barto allow $\gamma = 1$ whenever episodes terminate, and every environment in this paper terminates. MuJoCo episodes are capped at 1000 steps, Reacher at 100, and Atari and Key-to-Door end on their own. So Decision Transformer declines an optional choice. The constraint runs the other way: a return-to-go is a sum over the rest of the episode, which is undefined for a task that never ends, so the method cannot be applied where discounting is required.

Returns-to-go: labelling each step with the reward that followed

Decision Transformer changes the data and leaves the learning algorithm as ordinary supervised learning. The dataset is already recorded and will not change, so for every timestep you can add up the rewards that actually followed it and write that number next to the step. That number is the return-to-go:

It has no discount factor and estimates nothing. It is a sum of numbers on disk, and the official code computes it as a cumulative sum with $\gamma = 1$. A three-step episode with rewards 2, 0 and 5 gets returns-to-go 7, 5 and 5. Every step of every trajectory gets one, including the steps of the worst trajectories.

Plain imitation learning copies whatever the dataset does, so a dataset of bad players produces a bad player. Return conditioning changes the question the model answers. Instead of learning what a good player does, it learns what a player who is about to collect $\widehat{R}$ does next. Under that question a disastrous run is a correct example of how to score badly. The whole dataset becomes training data for one conditional model, and the behaviour you get is chosen at test time by the number you put in.

Once the environment has any randomness, a trajectory is an example of achieving the return it achieved, and need not be optimal for it. Kumar, Peng and Levine describe such data as optimal supervision for matching a return, as opposed to maximising one. The last section of this page is about the gap between those two.

The figure below labels a twenty-step episode that collects 100 in total. The amber bars are the reward at each step; the teal staircase is the return-to-go, which starts at 100 and steps down by each step's reward until it reaches zero. Drag the delay control to move the reward to the final step, as the paper's sparse-reward experiment does. Watch the staircase flatten and the value at $t = 1$ stay at 100.

The delayed-reward experiment in the results follows from this. A TD method loses most of its learning signal when the rewards move to the end, because its per-step target becomes zero almost everywhere and the single real reward has to travel back along the chain one backup at a time. Decision Transformer loses the shape of the label within the episode, but the label at the first step, the one it is prompted with, is unchanged.

The next figure puts the two schemes on the same corridor of twelve states. Press Play and watch the value move left one state per sweep in the top row; the bottom row is the same corridor labelled with returns-to-go, correct everywhere before any training.

The paper frames the comparison this way, and taken literally the framing claims more than the figure shows. First, one state per sweep is a property of one-step backups. n-step returns move credit n states at a time, eligibility traces spread it over the episode, and prioritised sweeping or replay in reverse order crosses the chain in one pass. Those methods pay in variance instead of time, so the fair comparison is with one-step baselines such as CQL and, on Atari, REM and QR-DQN, two value-based methods from the DQN Replay study. Second, the bottom row holds what one recorded run happened to collect, while a value function holds an expectation over what a policy would collect; the last section turns on that difference.

Return conditioning predates this paper. 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 that takes a desired return as input. Decision Transformer cites both and calls Reward-Conditioned Policies the method most similar to its own with a context of one timestep. What it adds is a sequence model with a long context, and an ablation in the results measures what that context is worth.

How Decision Transformer works: a trajectory as tokens

To give a trajectory to a Transformer you have to decide what a token is. Decision Transformer uses three tokens per timestep, always in the same order: return-to-go, state, action.

The model reads the last $K$ timesteps, so one forward pass covers $3K$ tokens. Each modality has its own linear layer into the embedding width (an Atari frame goes through a convolutional encoder instead). Counting from zero, timestep $i$ puts its return-to-go at position $3i$, its state at $3i+1$ and its action at $3i+2$.

The action head reads the hidden state at the state token, position $3i+1$, and turns it into a prediction of $a_i$. Because attention is causally masked, position $3i+1$ attends to positions $0$ through $3i+1$ and nothing later:

That set is every complete earlier timestep, plus this timestep's return-to-go and state, and excludes this timestep's action, which is the training target. The network therefore computes a return-conditioned policy over the visible history:

Drag the control below to move the prediction to another timestep. The shaded region is what the head can read, and the token ringed in amber is the target, one slot to the right of the shading at every setting.

The index layout 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. The placeholder sits at position $3i+2$, to the right of the state token the prediction comes from, and under the causal mask nothing to the right of a position can change its output, so the placeholder has no effect. The Atari code handles the same problem by truncating the sequence to $3K-1$ tokens at evaluation, so the trailing action slot does not exist.

Position signal. A standard Transformer adds one learned vector per token position. Decision Transformer learns one vector per episode timestep and adds that same vector to all three of the timestep's tokens, which the paper flags as a departure from the standard. The model tells the three token types apart through the three separate input projections: $W_R$ maps a scalar, $W_s$ maps a state vector (11 numbers for Hopper) and $W_a$ maps an action vector (3 for Hopper), each with its own bias, so the modalities land in different regions of the embedding space and the biases act as a learned type tag. The causal mask adds the rest, since the number of predecessors a token can attend to depends on its position.

The Atari code does not match that description: it adds the per-timestep embedding and a standard learned per-token position embedding over the context window, so only the MuJoCo code follows the paper's text.

Normalisation. Three token types with very different units share one sequence: a return-to-go in the thousands, a state in physical units, an action in $[-1,1]$. Layer normalisation puts them on a comparable scale because it rescales each token over its own features, with no reference to the batch or to other positions. Before its learned gain and bias, every normalised token vector has zero mean and unit variance across its $d$ features, so its length is $\sqrt{d}$ whatever it started as.

The paper describes a linear layer per modality followed by layer normalisation. The code adds the timestep embedding $z_i$ first and then applies one shared normalisation to the stacked sequence:

Layer normalisation subtracts the mean over features, so adding the same constant to every feature has no effect:

Normalising after the addition therefore discards the component of $z_i$ that is flat across features, and forces the modality content and the timestep signal to share one vector of fixed length. Normalising before the addition, as the paper's text reads, would let the timestep term move the token off that length by however much its norm grows during training.

The rest is a standard 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:

The per-head key width $d_k$ is 128, since the MuJoCo model runs one head at width 128. The divisor keeps the logits at unit scale: a dot product of $d_k$ unit-variance terms has variance $d_k$, and a softmax over large logits is nearly one-hot, with almost no gradient through it. The paper's own attention equation omits the divisor and writes the softmax over all tokens, stating the causal restriction only in the surrounding prose. The same equation also leaves out multi-head splitting, the output projection, residual connections, both layer norms and the feed-forward block, and the paragraph around it defers architecture details to the original papers, so these are simplifications. Both released implementations divide by $\sqrt{d_k}$. Dividing by a positive constant does not change which key a query scores highest, so the paper's reading of attention as credit assignment is unaffected.

The MuJoCo model has three layers, one attention head, width 128 and feed-forward width 512. The GPT trunk is about 0.60M parameters and the full Hopper model about 0.73M, of which 128,000 are the timestep embedding table (1000 timesteps by 128). That is roughly a 160th of GPT-2 small, and it is the model behind the D4RL numbers below. The Atari model is a different, larger configuration: six layers, eight heads, width 128, built on minGPT.

Training and rollout, step by step

The training objective is the squared error a behaviour-cloning baseline would use. Sample a window of $K$ timesteps, run one forward pass, read the $K$ action predictions off the state tokens, and average the error:

The set $\mathcal{M}$ holds the positions with real data instead of left padding, so short windows near the start of an episode do not dilute the loss, and $d_a$ is the action dimension. Discrete actions replace the squared error with cross-entropy over the action vocabulary, which is what the Atari model uses. The model also has prediction heads for states and returns; the paper reports that training them did not help, so their outputs are computed and discarded.

For Hopper the state has 11 numbers, the action 3, the context is $K = 20$ timesteps and the batch is 64. One training step runs on these shapes:

• The raw batch holds states $(64, 20, 11)$, actions $(64, 20, 3)$, and returns-to-go of length $K+1$ that are sliced to $(64, 20, 1)$.
• Four embedding layers (return, state, action, timestep) each produce a $(64, 20, 128)$ tensor, and the interleave makes $(64, 60, 128)$.
• The attention map is $(64, 1, 60, 60)$: one head, 60 queries, 60 keys.
• The action head reads the 20 state positions and emits $(64, 20, 3)$. Masked selection leaves at most 1280 valid actions, and the loss is one number.

Every one of the 60 tokens gets a hidden state, and only the 20 state tokens carry loss. The gradient reaches the return, state, action and timestep embeddings through attention, since the state token at $3i+1$ attends to all of them.

# 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 does not mention, and a reimplementation needs both. States are normalised by the dataset mean and standard deviation. 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. No comment in the repository explains the constant. A plausible reason, which the authors do not state: without it, $W_R \widehat{R}$ would dwarf its bias and the timestep vector added to it, and after normalisation the return token would point in nearly the same direction for every requested return. The optimiser is AdamW at learning rate $10^{-4}$ with weight decay $10^{-4}$, a linear warmup over 10,000 steps, gradient-norm clipping at 0.25, dropout 0.1, and $10^5$ gradient steps in total.

At rollout the return-to-go becomes an input you control. You pick a target, feed it as the first token, take the action the model returns, and then subtract the reward you received:

This is definition (5) run forward instead of backward. The number in the prompt is the reward still owed, and each step uses up part of it. Ask Hopper for 3600 and earn 3.5 on the first step, and the next prompt is 3596.5, which the network sees as 3.5965. A step that went badly leaves a larger number 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)    # subtract the reward 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 experiment the code skips the subtraction, which is consistent with (11): every reward before the last step is zero, so there is nothing to subtract. Neither listing contains an RL algorithm in the usual sense, with no critic, target network, importance weight or trust region.

Combining pieces of different trajectories

Imitating a dataset can reproduce the behaviour in it. The stronger claim is that conditioning can produce behaviour better than anything recorded, by joining fragments of different trajectories. The paper opens with a toy problem showing what that would look like.

Take a directed graph with a goal node and score it as an RL problem: reward 0 at the goal and $-1$ at every other node. The return-to-go is then minus the number of steps still to walk, so asking for the largest achievable return asks 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 in the data, and take the action that goes with it.

The figure below runs that procedure. Tap a node to start there. Amber arrows are transitions the random walks recorded, thicker where they were walked more often; grey arrows exist in the graph but never appear in the data, and the greedy walk cannot use them, because the model has never seen them. Turn the walk count down and watch coverage break: parts of the graph become unreachable and the path gets longer or fails.

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 remaining 84.2% already existed inside some training walk. The appendix containing that number also limits it. The model there differs from the one used everywhere else in the paper: it generates return tokens as well as actions, and its generation is steered by a hand-designed prior over returns that favours short paths, combined with the model's own return distribution:

The appendix says this prior is not used in any other experiment. The environment is a deterministic 20-node graph with 1000 random walks of 10 steps each. The word "stitching" appears exactly once in the paper, in that appendix paragraph.

Brandfonbrener and collaborators later showed that return conditioning does not recombine fragments in general. They construct a fully deterministic problem in which return conditioning needs on the order of $10^{H/2}$ trajectories to prefer the good action, with $H$ the episode length. Return conditioning needs the dataset to contain whole trajectories that reached the requested return, while dynamic programming needs only each individual transition to appear somewhere. So the graph toy demonstrates recombination only inside a small deterministic problem with good coverage, the regime the walk-count slider varies.

How far the return prompt reaches

If the first return-to-go is a control, the first question is whether changing it changes the outcome. The paper sweeps the target return across a wide range on every task and reports that the return requested and the return collected are highly correlated, matching almost exactly on Pong, HalfCheetah and Walker.

A model trained on (return, state) pairs reproduces the behaviour of episodes that actually scored near the request. Ask for a return the dataset contains, and there is a population of episodes to imitate. Ask for more than the best episode ever recorded, and no episode carries that label, so the model falls back on the same few top episodes and the collected return stops rising.

The figure below shows that mechanism on a toy dataset: 600 episodes with returns drawn from a fixed mixture, none 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.

The paper reports going past that ceiling on Seaquest, and the appendix gives enough to estimate by how much. Its Seaquest target of 1450 is listed as about five times the best return in its slice of the data, which puts the best recorded episode near 290, and Decision Transformer averaged a raw score of 1129.3 over three seeds. On Qbert the target of 2500 is also about five times the dataset maximum, a best recorded episode near 500, against a raw average of 2215.8. Both are well beyond anything in the training data. The paper's phrase is that Decision Transformer is "sometimes capable of extrapolation", and on Breakout and Pong the targets (90 and 20) are listed at one times the dataset maximum, requests the data supports.

What the benchmarks show

The figure holds three comparisons: D4RL against offline RL baselines, Atari against offline RL baselines, and D4RL against percentile behaviour cloning. Switch between them with the tabs and step through the rows with the slider.

D4RL. Averaged over the nine locomotion datasets, Decision Transformer scores 74.7, against 63.9 for CQL, 48.2 for BEAR, 36.9 for BRAC-v and 47.7 for behaviour cloning (the paper prints 46.4 for behaviour cloning, its twelve-row average; the two BC averages in Table 2 are swapped). Decision Transformer beats the best baseline on medium-expert HalfCheetah by 24.4 points (86.8 against CQL's 62.4) and beats CQL on medium-replay Walker by 39.9, trails BRAC-v by a few points on medium HalfCheetah (42.6 against 46.3), and loses medium-replay HalfCheetah to BRAC-v by 11.1. Across the grid, BRAC-v scores near zero on three datasets (0.8, 0.6 and 0.9) and BEAR falls to 19.2 on medium-replay Walker, while Decision Transformer never goes below 36.6 on the nine. With no learned value function there is no inflated estimate for a bad update to exploit, which is one explanation for that floor.

Atari. Four games, three seeds, and 1% of the DQN Replay dataset: 500,000 of the 50 million transitions an online DQN agent logged over its entire training run. That data ranges from a beginner to a finished DQN, the mixed-quality case return conditioning is meant for. Taken one game at a time, Decision Transformer wins Breakout (267.5 against CQL's 211.1) and Seaquest (2.5 against 1.7), roughly ties Pong (106.1 against 111.9), and loses Qbert badly, 15.4 against 104.2. REM and QR-DQN, the two baselines from the DQN Replay study, print 0.0 on Qbert: their raw scores of 160.1 and 156.0 fall just under the random-play anchor of 164, so the normalisation rounds them to zero.

The Breakout 267.5 is a ratio against a low human anchor. The raw score was 76.9 points. The normalisation puts 0 at random play (2 points) and 100 at a professional gamer's score (30 points), so $(76.9 - 2)/(30 - 2) = 2.675$ prints as 267.5. The human world record on Breakout is 864. The ratio depends on which anchor you use: with the original DQN paper's random and human scores (1.7 and 31.8) the same 76.9 prints 249.8. Pong is the one row whose arithmetic does not close: 17.1 raw against anchors of −21 and 15 gives 105.8, not the printed 106.1, though the standard deviation reconciles. The REM and QR-DQN columns cannot be found in the published CQL and REM tables either; the appendix says they come from raw data provided by the REM authors, which contradicts the main text's statement that the numbers are taken from those papers.

Context length. The Atari context is 30 timesteps on Breakout, Qbert and Seaquest and 50 on Pong. Cutting $K$ to 1, so the model sees one return-to-go and one state, takes Breakout from 267.5 to 73.9, Pong from 106.1 to 2.5 and Seaquest from 2.5 to 0.6. The inputs already use frame stacking, so the current state carries motion, and conventional practice treats one stacked state as enough for a policy. The paper offers a hypothesis, not a result: when the data mixes many policies of different quality, the context lets the model infer which policy generated the actions in its window, and conditioning on that makes the next action easier to predict.

Delayed rewards. Move every reward in the D4RL datasets to the final timestep and CQL collapses: 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. Figure 2 shows why: the delay leaves the first return-to-go unchanged, and the rollout needs only that one, decremented by rewards that are now zero until the end.

Key-to-Door. A long-horizon grid world in three rooms: pick up a key in the first, cross an empty second room, and reach a door in the third, which pays a binary reward only if the key was collected at the start. Trained on random trajectories, Decision Transformer succeeds 71.8% of the time from 1,000 of them and 94.6% from 10,000, 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. Percentile behaviour cloning, trained only on the successful episodes, scores 69.9% and 95.1%, and the table has no error bars or seed count, so the two are indistinguishable. With a binary reward, the return-to-go on every prefix equals the episode's success flag, which makes conditioning on it and filtering by it the same operation, and all of the gap over CQL comes from the hindsight labelling, which %BC shares.

Percentile behaviour cloning (%BC) sorts the data by episode return and clones only the top $X\%$ of timesteps. It is not a usable method, because choosing $X$ takes rollouts in the environment, which offline RL rules out, and the paper says so. As a diagnostic it asks whether return conditioning does anything that filtering does not.

On the eight D4RL datasets 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. Corrected, filtering with a hand-picked percentile edges Decision Transformer on that average, and 10%BC wins five of the eight rows outright. Decision Transformer's clear wins are elsewhere. On medium-replay Hopper it scores 82.7 against 70.6. On Atari, where the dataset is already 1% of a replay buffer, keeping the top 10% discards most of it, and 10%BC drops to 28.5 on Breakout against 267.5. The supported claim is narrower than the abstract's: conditioning makes better use of a small dataset than filtering, because it fits one model to all of the data and chooses the subset at test time instead of at training time.

The comparisons also differ in setup. The Gym behaviour-cloning baselines are MLPs, which the authors found stronger than a Transformer there, so on D4RL the comparison mixes return conditioning with architecture and context length. The Atari comparison is controlled: behaviour cloning there is the same model and hyperparameters without the return token. The Atari evaluation protocols do not match. Decision Transformer's code turns sticky actions off at evaluation, while the DQN Replay dataset and its baselines use the standard setting, in which the emulator repeats the previous action a quarter of the time and the games are harder. A bias runs the other way too: REM was early-stopped on online evaluation, reporting each agent's best score during training, averaged over five agents.

Where return conditioning fails

Every environment in the paper is deterministic or close to it: Atari ROMs, MuJoCo physics, a fixed graph, a grid world. Once the environment has real randomness, asking for a high return no longer selects good actions.

With infinite data, a return-conditioned policy is the logging policy $\beta$ reweighted by how often each action reached the requested return $f(s)$:

RCSL stands for return-conditioned supervised learning, the name Brandfonbrener and collaborators give to the family Decision Transformer belongs to. The numerator $P_\beta(f(s) \mid s, a)$ is the probability that action $a$ reached the target, not the average return of $a$. An action that hits a jackpot one time in a hundred is promoted toward that jackpot exactly as much as an action that hits it every time, as long as it is the only route to the requested number.

A worked case: one state, two actions logged equally often, so $\beta = 1/2$ each. The safe action pays 5 every time. The gamble pays 10 with probability $p = 0.2$ and 0 otherwise. Request $f(s) = 10$: the safe action never reaches 10, so its numerator is 0, and the gamble's is 0.2. The denominator is $\tfrac12 \cdot 0.2 = 0.1$, so (12) gives the gamble probability $\tfrac12 \cdot 0.2 / 0.1 = 1$. The policy always gambles and collects $10p = 2$ on average, where the safe action would have collected 5.

Slide $p$ toward zero and the loss approaches 5 per episode, as long as one gamble out of the hundred still came up lucky. Paster, McIlraith and Ba, who identified this failure, give a sharper example: three actions, where the action that produced the single best observed reward has an expected reward of $-5$, and every target return yields a negative expected reward (−5.0, −2.5, −0.17 and −5.0 for the four possible targets), at any dataset size. Brandfonbrener and collaborators give a matching construction: a one-state problem where a rare action almost always pays and a common action almost never does, arranged so that the conditioned policy picks each with probability exactly 1/2 for every conditioning function, even with infinite data.

The follow-up work also tested how much of the result 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 to require joining pieces of different trajectories, return conditioning does poorly: filtered behaviour cloning averages 41.4 against Decision Transformer's 19.8, and implicit Q-learning (IQL), which does propagate values with TD backups, totals 378.0 across the AntMaze datasets against Decision Transformer's 112.2.

A 0.73M-parameter GPT trained by regression on logged trajectories, with no critic and no conservatism, scored 74.7 on D4RL against CQL's 63.9, kept 107.3 of 107.6 when rewards were delayed to the end, and solved 94.6% of Key-to-Door episodes from random data. It works when the environment is close to deterministic and the dataset contains episodes near the requested return. Where actions have random outcomes or the data must be recombined across trajectories, methods that estimate expected values, such as IQL on AntMaze, score higher.