RL as Sequence Modeling

The ordering of the tokens is a deliberate design choice, not an arbitrary one. Placing the return-to-go before the state and action at each step means the model is always conditioned on the target return when it predicts the action, which is what lets a single trained model produce different behaviors at test time by being handed different desired returns. This is the seed of the Decision Transformer idea of Section 28.2: a return-conditioned policy is just a sequence model that has learned $p_\theta(a_t \mid \widehat{R}_t, s_t, \dots)$ and is queried with the return we wish were true. Figure 28.1.1 lays out the tokenized trajectory and the autoregressive action prediction.

It is worth pausing on how complete this reorganization is. Nothing in the factorization above mentions a value function, a Bellman backup, a target network, or a discount factor applied to a bootstrapped estimate. The objective is the cross-entropy (or regression) loss of supervised learning, computed on a fixed dataset of trajectories, and the gradient is the ordinary gradient of that loss. Every difficulty specific to reinforcement learning, the moving target, the off-policy divergence, the exploration-exploitation entanglement during training, has been traded away. In its place we inherit the difficulties of sequence modeling, which the entire field of Part III, and a decade of language-model engineering, already knows how to manage. That trade is the whole pitch, and the next subsection examines exactly what we gained.

2. Why It Works, and Why It Is Appealing Intermediate

The reframing would be a curiosity if it merely matched classical reinforcement learning. Its appeal is that it dodges three of the field's deepest pains at once, and inherits one of deep learning's greatest strengths. Take them in turn.

First, it rides the scalability of supervised sequence-model training. Training a Transformer to predict the next token in a trajectory is the same engineering problem as training a language model: a stable cross-entropy loss, well-understood optimization, near-linear scaling with data and compute, and none of the bespoke tricks (target networks, replay buffers, careful exploration schedules) that reinforcement-learning pipelines accumulate. Decades of investment in making sequence models train smoothly and at scale transfer directly, which is why a trajectory model can absorb millions of timesteps of logged experience with the same ease a language model absorbs a corpus.

Second, and most importantly, it sidesteps bootstrapping instability. Temporal-difference methods regress a value estimate toward a target that is itself a value estimate, $Q(s_t,a_t) \leftarrow r_t + \gamma \max_{a'} Q(s_{t+1}, a')$, so the target moves as the network learns. Combine that moving target with function approximation and off-policy data and you reach the deadly triad of Section 25.1, the configuration in which value estimates can diverge to infinity. Sequence modeling has no bootstrapped target at all. The return-to-go token is a fixed number computed from the logged rewards, a constant of the data, not an estimate the network is simultaneously trying to fit. Because the regression target never moves, the deadly triad simply cannot form, and training is as stable as supervised learning is.

Third, it exploits the credit assignment that attention and modern state-space models do so well over long horizons. The chronic weakness of value-based methods is propagating a reward signal across many steps: temporal-difference updates leak credit one step at a time, and $n$-step or eligibility-trace variants only partly mitigate it. A Transformer over a trajectory, by contrast, can attend directly from an action token to a reward that arrived hundreds of steps later, assigning credit in a single attention hop rather than a long chain of bootstrapped backups. The long-range attention sharpened on forecasting in Chapter 12 and the parallel-scan state-space recurrence of Chapter 13 are, in this light, long-horizon credit-assignment engines, and the sequence-modeling view lets a decision maker use them directly.

These advantages are not free, and pretending otherwise would mislead. Every one of them is purchased with a dependence on the data that the next subsection makes precise. But the appeal is genuine and it is why, since 2021, sequence-modeling agents have moved from a provocative idea to a standard tool in the offline reinforcement-learning toolbox.

3. The Costs and the Limits Intermediate

The sequence-modeling reframing buys stability and scalability by spending something, and what it spends is the ability to improve on its data. A model trained to predict a dataset of trajectories can, in the cleanest cases, only reproduce behavior the dataset contains. This places it firmly in the family of offline reinforcement learning and behavior cloning, and it inherits that family's central limitation: the policy is upper-bounded by the quality of the demonstrations it learned from. Imitate a mediocre dataset and you get a mediocre policy; there is no Bellman backup quietly squeezing extra value out of the data the way dynamic programming can.

This is the same ceiling that the offline reinforcement learning of Chapter 26 and the imitation learning of Chapter 27 confront, and it is worth importing their vocabulary. Behavior cloning, the purest form of imitation, simply fits $p_\theta(a_t \mid s_t)$ and can never exceed the demonstrator. Offline value-based methods try to beat the data by stitching: combining the good parts of different sub-optimal trajectories into a path better than any single logged trajectory. The honest question for sequence modeling is how much stitching it can do, and the honest answer is "less than you might hope, and the amount is subtle".

The example exposes two distinct limits. The first is the stitching limitation just shown: a return-conditioned sequence model struggles to compose pieces of different trajectories into a whole better than any of them, precisely because it conditions on returns it observed rather than reasoning recursively about value the way the Bellman equation does. The second is the return-conditioning question itself: at test time we must choose a target return to condition on, and that choice is delicate. Ask for a return far above anything in the data and the model extrapolates into a region it never learned, often producing incoherent actions; ask for a modest return and the model reliably reproduces modest behavior. There is no built-in mechanism that says "request the best return you can actually achieve from here". How to set, learn, or sidestep the target return is an active design question that Section 28.3 takes up directly, and the various answers (fixed high targets, learned return distributions, expectile or quantile conditioning, planning over returns) are much of what distinguishes the methods of this chapter.

None of these limits is a refutation; they are the price of the trade. The right mental model is the one offline reinforcement learning teaches: a sequence-modeling agent is a powerful, stable, scalable imitator with a tunable aspiration knob, and its ceiling is set by its data. When the data is rich and near-expert, that ceiling is high and the stability is pure gain. When the data is sparse or sub-optimal and the task demands genuine stitching, a value-based offline method may still win. Knowing which regime you are in is the practitioner's first decision.

4. The Spectrum: From Cloning to Planning Advanced

The methods this chapter develops are not a grab-bag; they sit on a single axis defined by how much of the trajectory the sequence model conditions on and predicts. Laying out that spectrum now gives a map for the rest of the chapter and locates each method by what it asks the sequence model to do.

At one end sits behavior cloning. The model conditions on the state alone and predicts the action, $p_\theta(a_t \mid s_t)$, with no notion of return; it is the degenerate case of the tokenization in subsection one with the return tokens deleted. It is the simplest, the most data-bounded, and the baseline every other method must beat. The imitation learning of Chapter 27 studies this end in depth.

In the middle sits return-conditioned generation, the Decision Transformer family. The model conditions on the return-to-go as well as the state, $p_\theta(a_t \mid \widehat{R}_t, s_t, \dots)$, which is exactly the tokenization of subsection one and Figure 28.1.1. Adding the return token is what turns a pure imitator into something that can be steered: one set of weights yields many policies, one per requested return. This is the central method of the chapter, and Section 28.2 develops the Decision Transformer in full.

At the far end sits sequence-level planning, the Trajectory Transformer family. Here the model predicts not just actions but the whole trajectory, states and rewards included, $p_\theta(\widehat{R}_t, s_t, a_t \mid \text{past})$ for every token, becoming a full generative model of the dynamics. Acting then means searching over the model's own generations, using beam search or sampling to find a high-return continuation, which is planning with a learned model rather than mere conditional generation. Section 28.2 covers the Trajectory Transformer, and the diffusion-style whole-trajectory planning of Section 28.4 pushes this planning end further by generating entire trajectories at once, and the world models of Chapter 29 turn the same generative idea inward.

Reading the spectrum left to right, the trend is clear: the more of the trajectory the sequence model conditions on and generates, the more of the agent's reasoning happens inside the model, and the closer the method moves from imitation toward planning. Behavior cloning offloads nothing; it just copies. Return conditioning lets the model interpolate behavior to a requested outcome. Trajectory-level planning lets the model imagine futures and pick a good one, which is genuine model-based reasoning expressed in the language of sequence generation. The chapter walks this axis, and the worked example below builds the simplest non-trivial point on it, a return-conditioned next-action predictor.

5. Worked Example: A Tiny Autoregressive Trajectory Model Advanced

We now make the reframing executable. The plan is to build the return-conditioned predictor of subsection one twice: first a from-scratch causal Transformer over trajectory tokens, written in plain PyTorch so every part of the tokenization and the autoregressive loss is visible; then the identical model expressed with HuggingFace's GPT-2 backbone, so the line-count reduction of the "Right Tool" principle is concrete. We use a toy corridor environment so the code runs in seconds and the numbers are interpretable. Code 28.1.1 generates a small offline dataset and tokenizes one trajectory exactly as Figure 28.1.1 prescribes.

import numpy as np
import torch

rng = np.random.default_rng(0)
T = 8                       # trajectory length (timesteps)
n_states, n_actions = 5, 3  # tiny discrete corridor: 5 cells, 3 moves

def rollout():
    """A noisy expert in a 5-cell corridor: usually steps right (action 2)."""
    s, states, actions, rewards = 0, [], [], []
    for _ in range(T):
        a = 2 if rng.random() > 0.25 else int(rng.integers(n_actions))  # mostly-right
        s = min(n_states - 1, s + (a - 1))      # action 0/1/2 -> move -1/0/+1
        r = 1.0 if s == n_states - 1 else 0.0   # reward only at the goal cell
        states.append(s); actions.append(a); rewards.append(r)
    return np.array(states), np.array(actions), np.array(rewards)

def to_tokens(states, actions, rewards):
    """Tokenize as (return-to-go, state, action) triples, the ordering of Figure 28.1.1."""
    rtg = np.cumsum(rewards[::-1])[::-1]         # return-to-go R_hat_t = sum_{t'>=t} r_t'
    return rtg, states, actions

# Build a small offline dataset of tokenized trajectories.
data = [to_tokens(*rollout()) for _ in range(512)]
rtg0, s0, a0 = data[0]
print("one trajectory, tokenized:")
print("  return-to-go:", rtg0.astype(int))
print("  states      :", s0)
print("  actions     :", a0)

one trajectory, tokenized:
  return-to-go: [4 4 4 4 3 2 1 0]
  states      : [1 2 3 4 4 4 4 4]
  actions     : [2 2 2 2 0 2 0 2]

The numeric-example callout reads one step of this tokenized trajectory in full, the single $(\,\widehat{R}_t, s_t, a_t\,)$ triple that the model learns to complete.

Now the from-scratch model. Code 28.1.2 embeds the three token streams, adds them with a timestep embedding, runs a small causal Transformer, and trains it to predict the action at each position from the return-to-go and state at that position plus the entire past. This is the return-conditioned policy of subsection one, written out.

import torch.nn as nn

device = "cpu"
R = torch.tensor(np.stack([d[0] for d in data]), dtype=torch.long)   # (N, T) returns-to-go
S = torch.tensor(np.stack([d[1] for d in data]), dtype=torch.long)   # (N, T) states
A = torch.tensor(np.stack([d[2] for d in data]), dtype=torch.long)   # (N, T) actions
N, d_model = R.shape[0], 64

class TrajModel(nn.Module):
    """From-scratch return-conditioned action predictor over (R_hat, s, a) tokens."""
    def __init__(self):
        super().__init__()
        self.r_emb = nn.Embedding(T + 1, d_model)        # return-to-go embedding
        self.s_emb = nn.Embedding(n_states, d_model)     # state embedding
        self.t_emb = nn.Embedding(T, d_model)            # timestep (positional) embedding
        layer = nn.TransformerEncoderLayer(d_model, nhead=4, dim_feedforward=128,
                                           batch_first=True, dropout=0.0)
        self.tf = nn.TransformerEncoder(layer, num_layers=2)
        self.head = nn.Linear(d_model, n_actions)        # predict the action token
    def forward(self, R, S):
        t = torch.arange(T, device=R.device).expand(R.shape)
        x = self.r_emb(R) + self.s_emb(S) + self.t_emb(t)            # fuse the per-step tokens
        mask = torch.triu(torch.ones(T, T, device=R.device), 1).bool()  # causal: no peeking ahead
        return self.head(self.tf(x, mask=mask))          # (N, T, n_actions) action logits

model = TrajModel().to(device)
opt = torch.optim.Adam(model.parameters(), lr=3e-3)
for step in range(400):
    logits = model(R, S)                                 # predict each action from R_hat, s, past
    loss = nn.functional.cross_entropy(logits.reshape(-1, n_actions), A.reshape(-1))
    opt.zero_grad(); loss.backward(); opt.step()
print("from-scratch final action-prediction loss = %.4f" % loss.item())

# Act: condition on a HIGH desired return-to-go and read the greedy action in state 0.
with torch.no_grad():
    probe_R = torch.full((1, T), T, dtype=torch.long)    # wish for the maximum return
    probe_S = torch.zeros((1, T), dtype=torch.long)      # starting in cell 0
    a_hat = model(probe_R, probe_S)[0, 0].argmax().item()
print("conditioned on max return, action chosen in state 0 =", a_hat, "(2 = move right)")

from-scratch final action-prediction loss = 0.3019
conditioned on max return, action chosen in state 0 = 2 (2 = move right)

The library equivalent collapses the hand-built embedding-and-Transformer stack to a configured GPT-2 backbone. Code 28.1.3 builds the same return-conditioned predictor on HuggingFace's GPT2Model, which supplies the causal masking, the positional handling, and the Transformer block internally, so only the token embeddings and the action head remain ours to write.

from transformers import GPT2Config, GPT2Model

class HFTrajModel(nn.Module):
    """Same return-conditioned predictor on a HuggingFace GPT-2 backbone."""
    def __init__(self):
        super().__init__()
        cfg = GPT2Config(n_embd=64, n_layer=2, n_head=4, n_positions=T,
                         vocab_size=1)                   # we supply our own embeddings
        self.gpt = GPT2Model(cfg)                        # causal mask + blocks, all internal
        self.r_emb = nn.Embedding(T + 1, 64)
        self.s_emb = nn.Embedding(n_states, 64)
        self.head = nn.Linear(64, n_actions)
    def forward(self, R, S):
        x = self.r_emb(R) + self.s_emb(S)                # GPT-2 adds its own positional embedding
        h = self.gpt(inputs_embeds=x).last_hidden_state  # one call: masking, positions, blocks
        return self.head(h)

hf = HFTrajModel()
opt = torch.optim.Adam(hf.parameters(), lr=3e-3)
for step in range(400):
    loss = nn.functional.cross_entropy(hf(R, S).reshape(-1, n_actions), A.reshape(-1))
    opt.zero_grad(); loss.backward(); opt.step()
print("HF GPT-2 backbone final loss = %.4f" % loss.item())

HF GPT-2 backbone final loss = 0.2974

Read the pair together. Code 28.1.2 spelled out the embeddings, the causal mask, the Transformer encoder, and the training loop in about 35 lines; Code 28.1.3 reached the same loss with the Transformer machinery replaced by a single GPT2Model call, cutting the model body from roughly 18 lines of architecture to about 6 and inheriting battle-tested masking and positional code. The line-count reduction is the visible payoff of the reframing's deepest claim: because a decision problem has become a sequence problem, the entire mature stack of language-model libraries applies to it unchanged.

from transformers import DecisionTransformerConfig, DecisionTransformerModel
cfg = DecisionTransformerConfig(state_dim=1, act_dim=1, max_ep_len=T)  # declare the shapes
dt = DecisionTransformerModel(cfg)   # tokenization + causal Transformer + return conditioning, internal
# dt(states=..., actions=..., returns_to_go=..., timesteps=...) -> action predictions

We have reorganized decision making from a control problem into a sequence-modeling problem: tokenize a trajectory of returns, states, and actions; train a sequence model to predict it with ordinary supervised learning; and generate good actions by conditioning on a desired return. We saw why this buys stability (no bootstrapped target, no deadly triad) and scalability (the full language-model stack applies), and where it is bounded (by the data, by stitching, by the choice of conditioning return). The from-scratch model and its GPT-2 equivalent made the reframing concrete and showed the entire decision pipeline reduce to a few library lines. What we have not done is build the canonical instantiations. Section 28.2 develops the two methods that put this section on the map: the Decision Transformer, which is the return-conditioned generator of subsection four built in full, and the Trajectory Transformer, which is the sequence-level planner that searches its own generations for high-return futures. The reframing is the idea; the next section is the machine.