"Every morning they hand me a number and tell me it is the reward I am about to collect. I have never once collected it. I was trained on the timid, the mediocre, the occasionally lucky, and not one of them ever earned what you are now asking of me. Still, I square my shoulders, condition on the impossible, and emit the action a champion would have taken. Sometimes, to everyone's surprise including mine, the champion shows up."
A Decision Transformer Conditioning on a Return It Has Never Actually Earned
Reinforcement learning is usually framed as an optimization problem: find a policy that maximizes expected return by estimating value functions and improving the policy against them. The sequence-modeling view throws that machinery out and asks a startlingly simpler question: what if we treat a trajectory, the interleaved stream of returns, states, and actions, as just another sequence, and train a Transformer to predict the next token the way a language model predicts the next word? Two papers from 2021 made this concrete and reframed the field. The Decision Transformer (Chen et al.) feeds the model triples of (return-to-go, state, action), trains it by plain supervised next-action prediction on a fixed offline dataset, and at test time conditions on a desired return-to-go to coax out actions that would achieve it. The Trajectory Transformer (Janner et al.) goes further and models the full joint distribution over states, actions, and rewards, then uses beam search over that learned distribution as a planner, the sequence model serving as both simulator and search space. This section builds both. It derives the token layout and the causal architecture (a direct callback to the Transformer of Chapter 12), contrasts the reactive return-conditioned policy against planning-by-search, implements a minimal Decision Transformer from scratch, runs the HuggingFace equivalent with a line-count comparison, and rolls out a policy conditioned on a target return. You leave able to explain why and when "RL as a big sequence model" works, and where its stitching limits begin.
In Section 28.1 we set up the central reframing of this chapter: a Markov decision process is itself a sequence, and the optimal-control problem of Chapter 27 can be attacked with the same autoregressive tools that drive language models. That reframing is the closing turn of the temporal thread this book has followed since Part II, where the Kalman filter, the hidden Markov model, and the recurrent network all turned out to be one loop. Here the MDP itself, the object that organized all of Chapter 22, returns as a sequence model. This section makes that abstract promise into two concrete, trainable architectures. We use the unified notation of Appendix A throughout: $s_t$ for state, $a_t$ for action, $r_t$ for the reward received after $a_t$, and $\hat{R}_t$ for the return-to-go, defined below.
Why does this view deserve its own section rather than a footnote to deep RL? Because it changes what is hard. Classical value-based and policy-gradient methods (Chapters 24 and 25) wrestle with bootstrapping, the deadly triad, target-network instability, and the notorious fragility of off-policy value estimation. The sequence-modeling approach sidesteps all of it: there is no Bellman backup, no temporal-difference target, no separate critic. There is a Transformer and a supervised loss, the most stable and best-understood training recipe in modern deep learning. The cost of that simplicity, as we will see, shows up elsewhere: in the model's ability to stitch together sub-trajectories it never saw end to end, and in how far above the data's returns you can credibly condition. Understanding precisely where the simplicity is free and where it is paid for is the goal.
The four competencies this section installs are these: to write a trajectory as a token sequence and lay out the (return-to-go, state, action) interleaving that the Decision Transformer consumes; to explain return-conditioned supervised training and the test-time trick of conditioning on a target return you select; to contrast the Decision Transformer's reactive generation against the Trajectory Transformer's planning by beam search, and to state the stitching and extrapolation limits of both; and to implement a minimal Decision Transformer from scratch, run its library equivalent, and roll it out conditioned on a return. These are the load-bearing skills for the goal-conditioned and return-conditioned policies of Section 28.3 and the world-model planners of Chapter 29.
1. The Decision Transformer: Return-Conditioned Sequence Modeling Intermediate
The Decision Transformer (Chen et al., 2021) begins with a single representational decision that determines everything else: it represents a trajectory not as a list of transitions to be valued, but as a flat sequence of tokens to be modeled autoregressively. The trajectory of a length-$T$ episode becomes
$$\tau \;=\; \big(\hat{R}_1,\, s_1,\, a_1,\; \hat{R}_2,\, s_2,\, a_2,\; \dots,\; \hat{R}_T,\, s_T,\, a_T\big),$$three tokens per timestep, in the fixed order return-to-go, then state, then action. The key quantity is the return-to-go $\hat{R}_t$, the sum of rewards from the current step to the end of the episode,
$$\hat{R}_t \;=\; \sum_{t'=t}^{T} r_{t'},$$not the reward received at $t$ and not the discounted value, but the raw remaining return. Placing $\hat{R}_t$ before $s_t$ and $a_t$ in the token order is the whole trick: when the model predicts the action token $a_t$, the return-to-go it should achieve is already in its context, so the action is generated conditioned on a target outcome rather than on a value estimate. The return-to-go decreases as the episode progresses (each step spends some of the remaining budget), which gives the model a running account of how much performance is still expected of it.
The architecture is a causal (decoder-only) Transformer, exactly the GPT-style stack of Chapter 12, with one twist in the embedding layer. Each of the three token types gets its own linear embedding into the model dimension $d$: a return embedding, a state embedding, and an action embedding. A per-timestep positional embedding is added so that the three tokens sharing a timestep also share a time index. The embedded sequence of length $3T$ flows through the standard masked multi-head self-attention blocks, and a causal mask ensures each token attends only to tokens at or before it, which is what makes generation autoregressive. The output at each action position is passed through a prediction head to produce $\hat{a}_t$. Figure 28.2.1 shows the token layout and the causal attention pattern.
Training is supervised and almost embarrassingly simple. Given an offline dataset of trajectories (collected by any policy, expert, random, or a mixture), compute the return-to-go for every timestep by a reverse cumulative sum, tokenize each trajectory, and train the Transformer to predict each action token from the preceding context with a standard regression or classification loss. For continuous actions the objective is mean-squared error between predicted and recorded action; for discrete actions it is cross-entropy. There is no reward model, no value function, no policy-improvement step: just maximum-likelihood next-action prediction. The loss for one trajectory is
$$\mathcal{L}(\theta) \;=\; \frac{1}{T}\sum_{t=1}^{T} \big\| a_t - \pi_\theta(\hat{R}_t, s_t \mid \tau_{The genuinely clever part is test time. Having trained the model to predict the action that, in the data, accompanied a given return-to-go and state, we now invert the conditioning: we choose a target return-to-go $\hat{R}_1$ that we would like to achieve (typically high, often at or slightly above the best return in the dataset), feed it together with the initial state $s_1$, and let the model generate $\hat{a}_1$. We execute that action in the environment, observe the reward $r_1$ and next state $s_2$, decrement the target by the reward received ($\hat{R}_2 = \hat{R}_1 - r_1$), and repeat. The model has learned the conditional "given that this much return was achieved from here, what action was taken?" and we exploit it as a controller by demanding a return and reading off the matching action. The policy is purely reactive: at each step it conditions on the current return-to-go and state and emits an action, with no explicit search or planning.
Take a five-step episode with rewards $(r_1, \dots, r_5) = (1, 0, 2, 0, 3)$, total return $6$. The return-to-go at each step is the suffix sum, computed by a single reverse cumulative sum from the end:
$\hat{R}_5 = r_5 = 3$; $\hat{R}_4 = r_4 + \hat{R}_5 = 0 + 3 = 3$; $\hat{R}_3 = r_3 + \hat{R}_4 = 2 + 3 = 5$; $\hat{R}_2 = r_2 + \hat{R}_3 = 0 + 5 = 5$; $\hat{R}_1 = r_1 + \hat{R}_2 = 1 + 5 = 6$.
So the return-to-go sequence is $(6, 5, 5, 3, 3)$, monotonically nonincreasing as expected (it spends down a fixed budget), and $\hat{R}_1$ equals the episode's total return. At test time we would instead set $\hat{R}_1$ ourselves, say to $8$ (above this trajectory's $6$), and decrement by each realized reward: after collecting $r_1 = 1$ the target becomes $\hat{R}_2 = 8 - 1 = 7$, telling the model "seven units of return still expected of you". The numeric subtlety: conditioning on $8$ asks for performance never seen in this episode, which is exactly the extrapolation question of subsection three.
Strip away the vocabulary and a Decision Transformer is a vanilla causal language model trained with a standard supervised loss. The single ingredient that converts "predict the next token" into "act to achieve a goal" is the return-to-go token placed before each action. In training it tells the model which return accompanied each action, so the model learns a return-conditioned action distribution. At test time it becomes a control knob: you write down the return you want and the model returns the action consistent with having wanted it. No Bellman equation, no critic, no policy gradient. Reinforcement learning becomes conditional sequence generation, and all the stability and scaling lessons of large-scale language modeling transfer directly.
2. The Trajectory Transformer: Modeling and Planning Over Full Trajectories Advanced
The Trajectory Transformer (Janner et al., 2021) shares the founding insight, treat the trajectory as a sequence, but draws a different conclusion from it. Where the Decision Transformer models only the conditional action distribution and acts reactively, the Trajectory Transformer models the full joint distribution over everything in the trajectory, states, actions, and rewards alike, and then uses that learned distribution as a planner. It is the difference between learning to imitate good behavior given a goal, and learning a complete probabilistic simulator of the environment that you can search through.
Concretely, the Trajectory Transformer tokenizes the trajectory as the interleaved stream of states, actions, and rewards (and, in the original work, the return-to-go as well),
$$\tau \;=\; \big(s_1, a_1, r_1,\; s_2, a_2, r_2,\; \dots,\; s_T, a_T, r_T\big),$$and a critical implementation choice is that continuous state and action dimensions are discretized into bins, so the model becomes a fully autoregressive categorical model over discrete tokens, factorizing the joint as a product of per-token conditionals exactly like a language model over text:
$$p_\theta(\tau) \;=\; \prod_{i=1}^{N} p_\theta\big(\tau_i \mid \tau_{The decisive difference is what happens at decision time. Because the model can score and sample whole futures, planning becomes a search over sequences, and the Trajectory Transformer uses beam search, the same algorithm that decodes machine-translation and language models, with one modification: instead of ranking candidate continuations by their log-likelihood, it ranks them by their predicted cumulative reward (or return-to-go). The planner expands a beam of candidate action sequences, uses the model to simulate the states and rewards each would produce, keeps the highest-reward beams, and extends them, finally executing the first action of the best plan and replanning at the next step (receding-horizon control, the model-predictive-control idea from Chapter 27). The sequence model is thus doing double duty: it is the dynamics model that simulates futures and the value signal that ranks them. Figure 28.2.2 contrasts the two architectures' decision procedures.
| Aspect | Decision Transformer | Trajectory Transformer |
|---|---|---|
| what is modeled | conditional action distribution $p(a_t \mid \hat{R}_t, s_t, \tau_{| full joint $p(s, a, r)$ over the trajectory | |
| token stream | $(\hat{R}_t, s_t, a_t)$ triples | $(s_t, a_t, r_t)$ (discretized), often plus $\hat{R}_t$ |
| action selection | reactive: condition on target return, emit action | planning: beam search over simulated futures |
| role of the model | return-conditioned policy | simulator and reward-ranking planner |
| compute at test time | one forward pass per step (cheap) | many forward passes per step (beam search, costly) |
| continuous values | embedded directly | discretized into per-dimension bins |
Modeling the full joint buys two things the reactive policy lacks. First, an explicit world model: the Trajectory Transformer can answer "what would happen if I took this action sequence?", which the Decision Transformer never represents. Second, the planner can in principle stitch, recombine fragments of different training trajectories into a high-reward plan never seen whole, because beam search composes continuations across the learned distribution rather than copying a single demonstrated path. The price is test-time cost: every decision now runs a beam search of many model evaluations, where the Decision Transformer needed one forward pass. The two papers thus stake out the ends of a spectrum: maximal simplicity and speed at one end, maximal lookahead and compositionality at the other.
There is something gloriously brazen about the Trajectory Transformer's central move. Faced with a continuous robot whose joint angles live in $\mathbb{R}^{17}$, it shrugs, chops every dimension into a few hundred bins, and declares the whole of physics to be a vocabulary. The cheetah's gallop, the hopper's hop, the friction of the floor: all just tokens now, no different in principle from the words of a sentence. And then, having reduced control to spelling, it plans the way a translation model decodes a sentence: beam search, the workhorse that turns "le chat" into "the cat", repurposed to turn a start state into an optimal gait. It should not work nearly as well as it does. That it works at all is a small monument to the idea that a good enough sequence model does not much care what its tokens used to mean.
3. Reactive Versus Planning: Strengths, Stitching, and Extrapolation Limits Advanced
The contrast between the two architectures is not merely implementational; it exposes the genuine strengths and the genuine limits of the sequence-modeling approach to decision making. The Decision Transformer's strength is its simplicity and stability: a supervised loss, no bootstrapping, no critic, and a test-time cost of one forward pass per action. Its corresponding weakness is that it is fundamentally an imitator. It reproduces the action that, in the data, accompanied a given return-and-state, which means it can only reliably reach returns that the data demonstrates are reachable from states the data visited. The Trajectory Transformer's strength is lookahead: by planning over a learned simulator it can compose behavior and react to its own predicted consequences. Its weakness is the cost and the compounding error of long-horizon simulation, plus the discretization that throws away precision in continuous domains.
The sharpest limit, and the one most worth internalizing, is stitching. Suppose the offline dataset contains a trajectory that goes well from A to B and a separate trajectory that goes well from B to C, but no single trajectory from A to C. An ideal offline RL method should learn the A-to-C route by stitching the two sub-trajectories at their shared state B; dynamic-programming methods (which propagate value backward through B) can do this naturally. The vanilla Decision Transformer often cannot: because it conditions on return-to-go and imitates, it has no mechanism to realize that the high-return continuation from B can be glued onto the path that reaches B. It tends to reproduce whole demonstrated behaviors rather than recombine their pieces. This is the single most cited limitation of return-conditioned imitation, and it is the gap that the goal-conditioned and dynamic-programming-augmented methods of Section 28.3 were designed to close. The Trajectory Transformer's beam search has a better (though still imperfect) shot at stitching, because search can compose continuations the data never showed end to end.
The second limit is extrapolation in the return token. At test time we can write any number we like into $\hat{R}_1$, including returns far above anything in the dataset. The model has no training signal for "what action achieves a return nobody ever achieved", so conditioning too aggressively produces out-of-distribution behavior that does not reliably deliver the requested return. Empirically, performance improves as you raise the target toward the dataset's best return and then plateaus or degrades beyond it: the model can match the best behavior it saw but cannot invent superhuman behavior from sub-optimal data by wishful conditioning. The practical recipe is to set the target return at or just slightly above the maximum return in the dataset, not at some fantasy ceiling.
Return-conditioned sequence modeling does not optimize anything at test time; it imitates a return-conditional. The model answers "what action accompanied this return-to-go in the data?", so it is bounded by the data in two precise ways. It cannot reliably produce returns above what the data demonstrates (extrapolation limit), and it does not automatically recombine sub-trajectories that share a state into a better whole (stitching limit). Dynamic-programming RL escapes both because the Bellman backup propagates value across trajectories and beyond observed returns. The sequence-modeling view trades that optimization power for supervised stability. Knowing which trade you are making tells you when to reach for a Decision Transformer (clean, near-expert offline data; you want the best demonstrated behavior) and when to reach for value-based offline RL or a planner (sub-optimal, fragmentary data that needs stitching).
4. Practical Training: Offline Datasets, Context, and Return Targets Intermediate
Both architectures are offline methods: they train on a fixed, previously collected dataset of trajectories and never interact with the environment during learning, which is what makes the supervised framing possible and what ties this section to the offline-RL setting introduced in Section 28.1. The standard benchmark is D4RL (Datasets for Deep Data-Driven RL), a suite of fixed datasets for continuous-control tasks (HalfCheetah, Hopper, Walker2d, the AntMaze navigation tasks, and others) at several data-quality levels: medium (a partially trained policy), medium-replay (the replay buffer of a training run, a noisy mixture), medium-expert (a mix of mediocre and expert), and expert. The data-quality axis is exactly where the stitching limit of subsection three shows itself: Decision Transformers are strong on medium-expert and expert data (clean behavior to imitate) and comparatively weaker on the fragmentary medium-replay and the maze tasks that demand stitching.
Three hyperparameters dominate practical performance. The first is context length $K$: how many recent timesteps (hence $3K$ tokens for a Decision Transformer) the model attends over. Surprisingly short contexts often suffice, the original Decision Transformer used $K=20$ to $K=30$ for most tasks, because the return-to-go token already summarizes the future the model is aiming for, so it does not need a long history to act well. The second is the target return at evaluation, set per the extrapolation guidance of subsection three (at or just above the dataset maximum). The third is the usual Transformer machinery: model size, learning-rate warmup, and dropout, all inherited unchanged from Chapter 12. A practical subtlety: states are typically normalized (subtract mean, divide by standard deviation computed over the dataset) and returns are often scaled by a task-specific constant so the return token lives on a numerically convenient range, a small preprocessing detail that materially affects training stability.
Who: A robotics team at a fulfillment-center automation company training a pick-and-place controller for a robotic arm sorting parcels onto conveyor lanes.
Situation: They had eighteen months of logged teleoperation and scripted-policy data, millions of (state, action, reward) transitions at mixed quality, but every minute of live robot time was expensive and a bad exploratory action risked damaging hardware or jamming the line.
Problem: They needed a controller that improved on the logged behavior without any online exploration, because online trial-and-error on physical hardware was operationally unacceptable.
Dilemma: Value-based offline RL (such as conservative Q-learning) could in principle stitch and exceed the data but was notoriously finicky to tune and prone to value-overestimation blowups on their heterogeneous logs. A Decision Transformer was far more stable to train but, being an imitator, would be bounded by the best behavior in the logs and would struggle to stitch across the fragmentary scripted episodes.
Decision: They shipped a Decision Transformer first as the stable baseline, conditioning at evaluation on a target return set to the 95th percentile of logged episode returns, and ran a value-based stitcher in parallel as a research track.
How: They tokenized the logs into (return-to-go, state, action) triples with context $K=20$, normalized states per the subsection-four recipe, and trained to convergence with a plain supervised loss in a few hours on a single GPU, no environment interaction and no reward model.
Result: The Decision Transformer matched the best demonstrated picking behavior reliably and shipped within a sprint, where the value-based method took weeks of tuning to beat it and only on the subset of tasks that genuinely needed stitching.
Lesson: When clean-to-mixed offline data exists and online exploration is forbidden, return-conditioned sequence modeling is the fast, stable first move; reach for value-based offline RL only when the task provably needs the stitching the imitator cannot do.
The sequence-modeling view of decision making has become one of the most active frontiers in the field. The stitching limit of subsection three has been directly attacked: the Q-learning Decision Transformer (QDT) and related work distill a value function back into the return labels so the supervised model inherits dynamic-programming's stitching, and Elastic Decision Transformers (2023) and similar variants adapt the context to recover stitching behavior. The reactive-versus-planning axis has been generalized by Diffuser and Decision Diffuser, which model trajectories with a diffusion process and plan by guided sampling rather than beam search, often stitching better than either 2021 architecture. The biggest shift is scale and generality: Gato (2022) and the open-source JAT (2024) train a single Transformer across hundreds of control, vision, and language tasks as one token stream, the Decision Transformer recipe pushed to a generalist agent; Multi-Game Decision Transformers play dozens of Atari games from one model. And the connection to large language models has closed the loop: methods that cast RL fine-tuning of LLMs as return-conditioned or trajectory sequence modeling reuse exactly the machinery of this section. The 2026 takeaway: "decision making as sequence modeling" is no longer a niche reframing but a unifying recipe linking offline RL, planning, and foundation models, with the world-model planners of Chapter 29 as its natural next chapter.
5. Worked Example: A Decision Transformer From Scratch, Then HuggingFace Advanced
We now make the architecture executable. The plan mirrors the from-scratch-then-library pattern used throughout this book: first build a minimal Decision Transformer as a GPT-style block over (return, state, action) tokens and train it on a tiny offline dataset, watching the loss fall; then run the HuggingFace DecisionTransformerModel, the same architecture in a few lines, and note the line-count reduction; finally roll the model out conditioned on a target return. Code 28.2.1 is the from-scratch model and its token embedding, the heart of the architecture.
import torch, torch.nn as nn
class MinimalDT(nn.Module):
"""A from-scratch Decision Transformer: GPT block over (R, s, a) tokens."""
def __init__(self, state_dim, act_dim, d=128, n_heads=4, n_layers=3, K=20):
super().__init__()
self.act_dim, self.d, self.K = act_dim, d, K
# One linear embedding per token TYPE (the only twist over a plain GPT).
self.embed_R = nn.Linear(1, d) # return-to-go -> R^d
self.embed_s = nn.Linear(state_dim, d) # state -> R^d
self.embed_a = nn.Linear(act_dim, d) # action -> R^d
self.embed_t = nn.Embedding(K, d) # per-timestep positional embedding
self.ln = nn.LayerNorm(d)
block = nn.TransformerEncoderLayer(d, n_heads, 4 * d,
batch_first=True, activation="gelu")
self.transformer = nn.TransformerEncoder(block, n_layers)
self.predict_a = nn.Linear(d, act_dim) # action prediction head
def forward(self, R, s, a, t): # each is (B, K, dim)
B, K = R.shape[0], R.shape[1]
pos = self.embed_t(t) # (B, K, d) shared by the 3 tokens
tok_R = self.embed_R(R) + pos
tok_s = self.embed_s(s) + pos
tok_a = self.embed_a(a) + pos
# Interleave to (R1,s1,a1, R2,s2,a2, ...): stack then reshape to length 3K.
seq = torch.stack([tok_R, tok_s, tok_a], dim=2).reshape(B, 3 * K, self.d)
seq = self.ln(seq)
mask = torch.triu(torch.ones(3 * K, 3 * K), diagonal=1).bool() # causal mask
h = self.transformer(seq, mask=mask) # (B, 3K, d)
h = h.reshape(B, K, 3, self.d)
return self.predict_a(h[:, :, 1]) # predict a_t from the STATE position
Code 28.2.2 generates a tiny offline dataset, computes return-to-go by reverse cumulative sum (the operation of the subsection-one numeric example), and trains the model with the plain supervised action-prediction loss. The point is to watch a supervised loss, nothing more exotic, drive an RL policy.
torch.manual_seed(0)
state_dim, act_dim, K, N = 4, 2, 20, 256
# Toy offline dataset: N trajectories of length K with random rewards.
states = torch.randn(N, K, state_dim)
actions = torch.randn(N, K, act_dim)
rewards = torch.randn(N, K, 1)
# Return-to-go = reverse cumulative sum of rewards along time (subsection 1).
rtg = torch.flip(torch.cumsum(torch.flip(rewards, [1]), dim=1), [1])
times = torch.arange(K).unsqueeze(0).expand(N, K)
model = MinimalDT(state_dim, act_dim, K=K)
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
for epoch in range(200):
pred_a = model(rtg, states, actions, times) # (N, K, act_dim)
loss = ((pred_a - actions) ** 2).mean() # supervised next-action MSE
opt.zero_grad(); loss.backward(); opt.step()
if epoch % 50 == 0:
print(f"epoch {epoch:3d} action-MSE = {loss.item():.4f}")
epoch 0 action-MSE = 1.0421
epoch 50 action-MSE = 0.7233
epoch 100 action-MSE = 0.5489
epoch 150 action-MSE = 0.4602
Predicting the action from the state position (the index h[:, :, 1]) is correct rather than a bug: in the $(\hat{R}_t, s_t, a_t)$ order the state token sits after that step's return-to-go, so under the causal mask it already attends to $\hat{R}_t$ before producing $\hat{a}_t$, and the return conditioning is intact. The action token's own slot is reserved for predicting the next step's return and state, never its own action. The return-conditioned rollout of Code 28.2.4 puts this prediction to work as a controller, so you can watch the state-position action actually drive an environment.
Now the library equivalent. The full from-scratch model and training loop above ran roughly 50 lines including the architecture. HuggingFace ships the identical architecture as DecisionTransformerModel; Code 28.2.3 instantiates it and runs a forward pass in about a dozen lines, with the token embedding, interleaving, causal masking, and prediction heads all internal.
from transformers import DecisionTransformerConfig, DecisionTransformerModel
cfg = DecisionTransformerConfig(state_dim=state_dim, act_dim=act_dim,
max_ep_len=4096, hidden_size=128, n_layer=3)
hf_dt = DecisionTransformerModel(cfg) # the SAME architecture, prebuilt
# One forward pass: the library takes the four token streams directly.
B = 8
out = hf_dt(
states=torch.randn(B, K, state_dim),
actions=torch.randn(B, K, act_dim),
returns_to_go=torch.randn(B, K, 1),
timesteps=torch.arange(K).unsqueeze(0).expand(B, K),
attention_mask=torch.ones(B, K),
return_dict=True,
)
print("predicted action shape:", tuple(out.action_preds.shape))
predicted action shape: (8, 20, 2)
Finally we roll the model out conditioned on a target return. Code 28.2.4 implements the test-time loop of subsection one: pick a desired initial return-to-go, and at each step feed the recent context, predict an action, step a (stub) environment, observe the reward, and decrement the return-to-go by the reward received. This is the inversion that turns the trained predictor into a controller.
@torch.no_grad()
def rollout(model, env_step, s0, target_return, K=20, horizon=20):
"""Roll out conditioned on a target return; decrement RTG by realized reward."""
s, R = s0, float(target_return) # initial state, desired return-to-go
states = torch.zeros(1, K, state_dim)
actions = torch.zeros(1, K, act_dim)
rtgs = torch.zeros(1, K, 1)
times = torch.arange(K).unsqueeze(0)
total_reward = 0.0
for t in range(horizon):
i = min(t, K - 1)
states[0, i] = s
rtgs[0, i, 0] = R # the target return-to-go (subsection 1)
a = model(rtgs, states, actions, times)[0, i] # predicted action at this step
actions[0, i] = a
s, r = env_step(s, a) # step the environment
total_reward += r
R = R - r # spend the realized reward (R_{t+1}=R_t - r_t)
return total_reward
# Stub environment: reward = -||action|| (a placeholder; swap for a real Gym env).
def env_step(s, a):
return torch.randn(state_dim), float(-(a ** 2).sum())
got = rollout(model, env_step, torch.randn(state_dim), target_return=10.0)
print(f"target return = 10.0 achieved return = {got:.3f}")
target return = 10.0 achieved return = -8.214
Step back and read what the four code blocks establish together. Code 28.2.1 built the Decision Transformer's defining structure, three type embeddings interleaved into a causal sequence. Code 28.2.2 trained it with a plain supervised loss, no RL machinery in sight. Code 28.2.3 reproduced the architecture in a dozen library lines. Code 28.2.4 inverted the trained predictor into a return-conditioned controller. The pedagogical payoff is that the entire pipeline, offline data to deployed policy, is ordinary supervised sequence modeling with one extra token type, which is exactly why the approach inherits the stability of language-model training.
Beyond HuggingFace's architecture-only DecisionTransformerModel, the offline-RL library d3rlpy ships a full training-and-evaluation pipeline for Decision Transformers (and for value-based offline methods such as CQL and IQL for comparison). What took the roughly 50 lines of dataset construction, return-to-go computation, training loop, and rollout across Codes 28.2.1 through 28.2.4 collapses to a few lines: load a D4RL dataset, construct a DecisionTransformer with a context size and target return, call .fit(), and evaluate. The library handles return-to-go computation, state normalization, the context windowing of subsection four, and the conditioned rollout internally, and exposes the target-return knob directly so you can sweep the extrapolation behavior of subsection three with a single argument.
import d3rlpy
dataset, env = d3rlpy.datasets.get_d4rl("hopper-medium-v2") # standard offline benchmark
dt = d3rlpy.algos.DecisionTransformerConfig(context_size=20).create(device="cuda:0")
dt.fit(dataset, n_steps=100_000) # full supervised training
# Evaluate conditioned on a target return at or just above the dataset max (subsection 3).
evaluator = d3rlpy.metrics.EnvironmentEvaluator(env, target_return=3600.0)
print("normalized score:", evaluator(dt, dataset))
The line-count reduction is roughly tenfold over the from-scratch pipeline, and crucially the library's defaults (normalization, return scaling, context windowing) are the ones tuned to actually reach published D4RL scores, which a naive from-scratch implementation usually does not without the same care. One currency note: the classic d4rl locomotion datasets used here are the -v2 revisions (the original -v0 versions are deprecated), and as of 2026 classic D4RL is itself superseded by Minari, so current code loads the same benchmarks with d3rlpy.datasets.get_minari under d3rlpy v2.
This section completes one of the book's longest temporal-thread arcs. The Markov decision process organized Chapter 22 as a problem of value functions and Bellman backups; the same MDP now returns, in Section 28.1 and here, as a sequence to be modeled by the very Transformer of Chapter 12. The Kalman filter became the RNN; the HMM became the latent-variable model and the belief-state agent; and now the MDP becomes a causal language model over returns, states, and actions. The lesson the thread has taught at every stop holds again: a classical structure, restated as a sequence and learned by gradient descent, inherits both the power and the pathologies of sequence models, here the stability of supervised training and the stitching and extrapolation limits of imitation.
6. Summary and Exercises
This section recast offline reinforcement learning as sequence modeling through two architectures. The Decision Transformer represents a trajectory as $(\hat{R}_t, s_t, a_t)$ triples, trains a causal Transformer by supervised next-action prediction, and acts by conditioning on a desired return-to-go, a reactive, return-conditioned policy with the stability of language-model training. The Trajectory Transformer models the full joint distribution over states, actions, and rewards and plans by beam search over that learned simulator, trading test-time compute for lookahead and a better shot at stitching. The contrast exposes the approach's two structural limits: return-conditioned imitation cannot reliably exceed the data's returns (extrapolation) and does not automatically recombine sub-trajectories (stitching), the gaps that Section 28.3 and the planners of Chapter 29 address. We built a minimal Decision Transformer, ran the HuggingFace equivalent at a fraction of the line count, and rolled out a return-conditioned controller.
Three mistakes recur when readers first implement and deploy the architecture of this section, and each fails quietly rather than loudly:
- Placing the return token after the action. The order must be $(\hat{R}_t, s_t, a_t)$. If the return-to-go follows the action, the action token can no longer attend to that step's target under the causal mask, the conditioning is gone, and the model degenerates into plain behavior cloning that no return knob can steer.
- Conditioning above the data's support. Writing a target return far beyond the dataset maximum asks for behavior the model never saw and produces out-of-distribution actions that do not deliver the requested return. Set the target at or just slightly above the dataset's best return, as subsection three prescribes.
- Expecting a vanilla Decision Transformer to stitch. Return-conditioned imitation reproduces demonstrated behaviors; it does not recombine a good A-to-B path with a good B-to-C path into an unseen A-to-C route. When the data is fragmentary and the task needs stitching, reach for a value-augmented or planning method, not the plain recipe.
The exercises below come in three types: conceptual (reason about the architecture and its limits), implementation (extend the code), and open-ended (explore a research direction).
- Conceptual. Explain precisely why placing the return-to-go token before the state and action tokens (rather than after) is essential to the Decision Transformer working as a controller at test time. What would break if the order were $(s_t, a_t, \hat{R}_t)$? Then describe, using the stitching example of subsection three (a good A-to-B trajectory and a good B-to-C trajectory but no A-to-C), why a vanilla Decision Transformer typically fails to find the A-to-C route while a value-based offline method can.
- Implementation. Extend the from-scratch model of Code 28.2.1 to support discrete actions by replacing the action embedding with an
nn.Embeddingand the regression head with a classification head, and switch the training loss in Code 28.2.2 from mean-squared error to cross-entropy. Then add a target-return sweep to the rollout of Code 28.2.4: run the rollout for target returns in $\{0, 1, 2, \dots, 10\}$ and plot achieved return against requested return, reproducing in miniature the plateau-then-degrade extrapolation curve of subsection three. - Open-ended. The stitching limit of subsection three motivated a family of fixes (the Q-learning Decision Transformer, the Decision Diffuser, Elastic Decision Transformer). Pick one from the research-frontier callout, read its paper, and write a one-page explanation of the specific mechanism by which it recovers stitching that the vanilla Decision Transformer lacks. State clearly what it adds (a value function, a diffusion process, an adaptive context) and what it costs (extra training, extra compute, extra hyperparameters) relative to the plain return-conditioned recipe.