Offline & Model-Based RL intermediate 7 min read 7 flashcards

Decision Transformer and RL as Sequence Modelling

Decision Transformer replaces value functions with a return-conditioned sequence model trained by supervised learning; it is stable and simple, and theory and benchmarks show it needs near-deterministic dynamics and cannot in general stitch suboptimal trajectories.

Offline RL methods spend most of their machinery keeping bootstrapped value estimates from diverging. Decision Transformer asked whether the value function was needed at all (Chen et al., 2021, Decision Transformer: Reinforcement Learning via Sequence Modeling, arXiv:2106.01345). Treat a trajectory as a sequence of tokens, train a GPT-style model to predict actions, and at test time ask for a high return. No Bellman backup, no target network, no deadly triad. The idea is attractive enough that its limits are worth knowing precisely.

The mechanism

A trajectory is serialised as triples of return-to-go, state and action:

\[\tau = \big(\hat{R}_1, s_1, a_1, \hat{R}_2, s_2, a_2, \dots\big), \qquad \hat{R}_t = \sum_{t'=t}^{T} r_{t'} .\]

Each modality gets a linear embedding plus a learned timestep embedding, and a causally masked transformer attends over the last \(K\) timesteps, \(3K\) tokens. Training minimises cross-entropy or squared error between the predicted and logged action at each state token. The paper used \(K = 30\) for most Atari games and \(K = 50\) for Pong.

At deployment you choose a target return, say the best return in the dataset, feed it with the first state, sample an action, observe reward \(r\), and set the next return-to-go to \(\hat{R} - r\). The model is asked, at every step, "what action did trajectories that went on to earn this much take from here?"

This is a form of return-conditioned supervised learning (RCSL). Emmons and colleagues showed that a plain MLP conditioned on return or goal matches much of Decision Transformer's performance, which suggests the conditioning, not the transformer, does most of the work (Emmons et al., 2021, RvS: What is Essential for Offline RL via Supervised Learning?, arXiv:2112.10751).

The stitching critique

Suppose the dataset contains trajectory A→B→C, which scores poorly because the segment after B is bad, and D→B→E, which scores well because B→E is good. Dynamic programming propagates the high value of B→E back to every state that reaches B, so from A it learns to go to B and then to E, a route that appears in no single trajectory. That is stitching.

A return-conditioned model starting at A has only ever seen low returns-to-go from A. Asking it for a high return conditions on a value outside the support of the data at that state, and nothing in supervised learning says what the model should do there. Brandfonbrener and colleagues proved that RCSL recovers near-optimal behaviour only under stronger assumptions than DP: nearly deterministic dynamics, a correct conditioning value, and that value supported by the returns in the dataset. They also gave evidence that RCSL cannot stitch in general (Brandfonbrener et al., 2022, When does return-conditioned supervised learning work for offline reinforcement learning?, arXiv:2206.01079).

The benchmark numbers agree. In the IQL paper's comparison on D4RL, Decision Transformer scored 67.6 on hopper-medium, slightly above IQL's 66.3, but 0.0 on both large AntMaze datasets, where stitching is the whole task and IQL scored 39.6 and 47.5. The original DT paper did show stitching on a small graph shortest-path task, with 15.8 percent of generated paths assembled from suboptimal segments, but that experiment used a generation prior the authors did not use elsewhere.

Stochasticity: conditioning on luck

Consider one decision: a safe action returns 5; a gamble returns 10 with probability 0.1 and 0 otherwise. Logged data contains both. Condition on return 10 and the only trajectories that achieved it took the gamble, so the model gambles, with expected return \(0.1 \times 10 = 1\), a fifth of the safe option. Return-to-go conflates good decisions with good luck. Paster, McIlraith and Ba formalised this failure for Decision Transformers and RvS in stochastic environments (Paster et al., 2022, You Can't Count on Luck, arXiv:2205.15967).

When it breaks

The target return is a hyperparameter with no safe value. Too low and the model imitates mediocre behaviour; above the dataset's support, behaviour is undefined. The target that works is usually found by evaluating online, which offline RL was supposed to avoid.

Stitching-heavy tasks. Wherever success requires combining pieces of different trajectories, pure RCSL underperforms DP methods. Hybrids that relabel return-to-go with values learned by dynamic programming, such as Q-learning Decision Transformer, exist precisely to import stitching (Yamagata et al., 2022, Q-learning Decision Transformer, arXiv:2209.03993), which concedes the value function back.

Compute. The IQL paper's runtime table lists 960 minutes for Decision Transformer against 20 for IQL. The comparison is not like for like, since the DT figure is the original PyTorch implementation and IQL's is a JAX reimplementation, but a gap of that size is not an implementation detail. Simplicity of the objective is not cheapness of training.

Where it fits. Near-deterministic environments with plenty of high-quality, full-length trajectories, and settings where the stability of supervised learning is worth more than improving on the data, remain a reasonable home. The disagreement is not whether sequence models are useful in RL, but whether return conditioning alone is a substitute for dynamic programming. The evidence says it is not.

Check yourself

7 flashcards for this concept

Click a card to reveal the answer.

Drill the whole track