Decision Transformers and Trajectory Transformer

Decision Transformer treats control as next token prediction. You state the return you want and it generates the actions that earned that return in past data.

A Decision Transformer reads return to go, state and action tokens through causal attention and outputs the next action that fits the requested return

Why Does This Exist?

Offline reinforcement learning learns from a fixed dataset with no new environment interaction. The standard approach estimates values with Bellman backups, where each update bootstraps from the next estimate. On narrow datasets that bootstrapping compounds small approximation errors into large overestimates, and training diverges.

A second pain is stability. Temporal difference methods need target networks, careful discounting and extra regularizers to stay stable offline. Each addition helps one dataset and hurts the next.

Decision Transformer, introduced by Chen and colleagues in 2021, with the Trajectory Transformer from Janner and colleagues the same year, drops the backup entirely. It treats a trajectory as text to complete and trains with ordinary supervised loss. That trade buys stable optimization at the cost of giving up value based stitching.

This page covers the sequence framing, the return to go conditioning and the inference loop for both models. For the offline setting itself see offline batch reinforcement learning. For the attention mechanism they borrow see transformers in RL. For the value methods they replace see Q-learning.

The one idea you need first

A return is the sum of future rewards from a step onward. The return to go at step $t$ , written $\hat{R}_t$ , is that sum computed from $t$ to the end of the episode. If rewards are $2$ , $3$ and $5$ , the return to go at step $1$ is $10$ , at step $2$ is $8$ and at step $3$ is $5$ . Decision Transformer conditions on this number as its request.

Think of It Like This

Ordering by photo at a busy restaurant

You sit down and point at a photo of the finished dish you want. You do not list oven temperatures or stirring times. The cook has served that dish hundreds of times, so the photo is enough to reproduce the steps.

The photo is the return to go. The state of the counter is the state. Each cooking move is an action. The cook is the transformer, trained on past orders. You lower your expectation as courses arrive, just as the agent subtracts earned reward from its target.

Where it stops: a good cook improvises a new dish from spare parts. Decision Transformer does not. It only reproduces combinations it saw in full, so it struggles when the best plan mixes pieces from separate trips.

How It Actually Works

A trajectory $\tau$ of length $T$ becomes an interleaved token list. Here $s_t$ is the state at step $t$ , $a_t$ is the action taken there, $r_t$ is the reward received after it and $\hat{R}_t$ is the return to go from $t$ onward:

\tau = (\hat{R}_1, s_1, a_1, \hat{R}_2, s_2, a_2, \ldots, \hat{R}_T, s_T, a_T)

The return to go is defined as:

\hat{R}_t = \sum_{k=t}^{T} r_k

In plain words: add up every reward from now until the episode ends. That sum is the request the model must satisfy.

Trajectory tokens and causal attention

Returns, states and actions live in different spaces, so each gets its own linear projection into a shared width $d$ . All three tokens from step $t$ share one learned timestep embedding $p_t$ , which keeps the triplet aligned when the context window slides.

A causal mask restricts attention to the past. When the model predicts $a_t$ it may read $\hat{R}_1, s_1, a_1$ through $\hat{R}_{t-1}, s_{t-1}, a_{t-1}$ plus the current $\hat{R}_t$ and $s_t$ . It never sees the true $a_t$ or anything later. The parameters $\theta$ are all projection matrices, embeddings and attention weights.

Training loss and inference loop

Training is supervised on trajectories sampled from the offline set $\mathcal{D}$ . For continuous actions with predicted action $\hat{a}_t$ the loss is mean squared error:

\mathcal{L}(\theta) = \mathbb{E}_{\tau \sim \mathcal{D}} \left[ \frac{1}{T} \sum_{t=1}^{T} \| a_t - \hat{a}_t(\hat{R}_t, s_t, \tau_{<t}) \|^2 \right]

Discrete actions use cross entropy over logits instead. There is no critic, no target network and no discount factor $\gamma$ in this loss.

Inference runs as a loop. Pick a target $\hat{R}_1$ , observe $s_1$ , predict $a_1$ , execute it, receive $r_1$ , then decay the request:

\hat{R}_{t+1} = \hat{R}_t - r_t

Append the new tokens and repeat inside a context window of recent steps, often $K = 20$ timesteps.

Decision Transformer versus Trajectory Transformer

Both treat RL as sequence modeling but they predict different things:

Decision Transformer predicts actions directly from return, state and history. It acts as the policy. At test time you feed the desired return and read out actions.
Trajectory Transformer models the joint distribution over states, actions and rewards as discrete tokens. It acts as a world model. At test time you search over imagined futures with beam search and execute the first action of the best imagined path.

Use the first when you want a fast policy. Use the second when you want to plan through a learned model and can pay the search cost.

Worked example

Take a three step episode with rewards $r_1 = 2.0$ , $r_2 = 3.0$ and $r_3 = 5.0$ .

Returns to go, computed backward:

\hat{R}_1 = 2.0 + 3.0 + 5.0 = 10.0

\hat{R}_2 = 3.0 + 5.0 = 8.0

\hat{R}_3 = 5.0

Equivalently $\hat{R}_2 = \hat{R}_1 - r_1 = 10.0 - 2.0 = 8.0$ . The total return is $10.0$ .

At step $t = 2$ the context holds five tokens: $(\hat{R}_1 = 10.0, s_1, a_1, \hat{R}_2 = 8.0, s_2)$ . The query at $s_2$ attends to those five and outputs $\hat{a}_2$ .

Suppose the dataset action is $a_2 = [0.50, 0.80]$ and the head predicts $\hat{a}_2 = [0.45, 0.85]$ . With action width $D = 2$ :

\mathcal{L}_2 = \frac{1}{2}[(0.45 - 0.50)^2 + (0.85 - 0.80)^2] = \frac{1}{2}[0.0025 + 0.0025] = 0.0025

During live rollout you prompt $\hat{R}_1 = 10.0$ , execute, receive $r_1 = 2.0$ and update to $\hat{R}_2 = 8.0$ . The agent now owes itself exactly $8.0$ more to meet the request.

Code

This snippet computes returns to go, the action loss and the inference decay for the numbers above.

1from typing import List2
3
4def compute_rtg(rewards: List[float]) -> List[float]:5    """Return to go at each step: Rhat[t] = sum(rewards[t:])."""6    rtg: List[float] = [0.0] * len(rewards)7    running = 0.08    for i in range(len(rewards) - 1, -1, -1):9        running += rewards[i]10        rtg[i] = running11    return rtg12
13
14def action_mse(pred: List[float], target: List[float]) -> float:15    """Mean squared error over one action vector."""16    return sum((p - t) ** 2 for p, t in zip(pred, target)) / len(pred)17
18
19rewards = [2.0, 3.0, 5.0]20rtg = compute_rtg(rewards)21print(rtg)  # -> [10.0, 8.0, 5.0]22print(f"{action_mse([0.45, 0.85], [0.50, 0.80]):.6f}")  # -> 0.00250023print(rtg[0] - rewards[0])  # -> 8.0

The printed return list, loss and decayed target match the hand computation. The loop runs backward so each entry reuses the sum of everything after it.

Watch Out For

Prompting a return the dataset never earned

Asking for $\hat{R}_1$ far above the best trajectory in the data pushes the model out of distribution. The attention layers combine fragments they never saw together and the actions turn erratic. You see near expert scores on locomotion logs with smooth demos and near zero on maze logs that need new combinations.

Fix it by prompting near the best observed return, or near a high percentile such as the top tenth. If you truly need better than the data, pair the transformer with a learned value check or augment the data with model rollouts first.

Expecting the model to stitch two half trips

Suppose one log shows start $A$ to checkpoint $B$ and another shows $B$ to goal $C$ , but no log shows $A$ to $C$ . Value methods propagate reward backward across the shared state $B$ and find the full route. Decision Transformer matches whole sequences, so it never connects the halves.

Symptom: strong results on dense reward locomotion, weak results on sparse navigation that requires stitching. Use Trajectory Transformer with beam search, which can chain imagined segments through its learned dynamics, or add a Q guided sampling step.

The Quick Version

Decision Transformer replaces Bellman backups with conditional sequence prediction over $(\hat{R}_t, s_t, a_t)$ triplets.
Return to go $\hat{R}_t$ is the future reward sum from $t$ onward and acts as the prompt for how well to behave.
Training is plain supervised regression or classification with a causal mask, so optimization stays stable offline.
Inference prompts a target return and decays it by each earned reward with $\hat{R}_{t+1} = \hat{R}_t - r_t$ .
Trajectory Transformer instead models full dynamics as discrete tokens and plans with beam search.
Sequence matching cannot stitch disjoint trips, so prompt inside the data range on stitching tasks.