Decision Transformers and Trajectory Transformer

Why Does This Exist?

For decades, reinforcement learning has been dominated by the dynamic programming paradigm. Algorithms like Q-learning, Actor-Critic, and policy gradients rely on temporal difference (TD) learning, updating value estimates recursively via Bellman backups:

$Q(s, a) \leftarrow r + \gamma \max_{a'} Q(s', a')$

While mathematically elegant, TD learning suffers from notorious practical instabilities in the offline (batch) setting. Bootstrapping creates a fragile feedback loop where out-of-distribution function approximation errors compound uncontrollably, causing value estimates to diverge and policies to degrade. Practitioners have had to invent complex regularization techniques—such as conservative value penalties (CQL), batch constraints (BCQ), and expectile regressions (IQL)—simply to keep temporal difference updates stable.

Meanwhile, natural language processing underwent a revolution driven by Generative Pre-trained Transformers (GPT). Large causal Transformers demonstrated that predicting the next token in a sequence using multi-head self-attention scales predictably, trains stably via standard supervised cross-entropy or mean squared error (MSE), and captures intricate long-range dependencies without any recursive value bootstrapping.

In 2021, Chen et al. introduced the Decision Transformer (DT), alongside Janner et al.'s Trajectory Transformer (TT). They proposed a radical paradigm shift: reframe reinforcement learning not as dynamic programming, but as conditional autoregressive sequence modeling. By representing an agent's experience as an interleaved sequence of return targets, states, and actions, the agent bypasses Bellman equations entirely. To act expertly, the agent simply conditions the Transformer on a high desired return-to-go and lets autoregressive self-attention generate the exact sequence of actions that yields that result.

Think of It Like This

How It Actually Works

Trajectory Representation as Token Sequences

Instead of decomposing an offline dataset into independent transition tuples $(s_t, a_t, r_t, s_{t+1})$, the Decision Transformer treats episodes as coherent trajectory sequences.

A trajectory $\tau$ of length $T$ is structured as an interleaved sequence of Return-to-Go, State, and Action tokens:

$\tau = \left( \hat{R}_1, s_1, a_1,\, \hat{R}_2, s_2, a_2,\, \dots,\, \hat{R}_T, s_T, a_T \right)$

Return-to-Go ($\hat{R}_t$)

Rather than conditioning on the immediate reward $r_t$, the model conditions on the Return-to-Go (RTG), defined as the future cumulative return remaining from timestep $t$ until the end of the episode:

$\hat{R}_t = \sum_{t'=t}^T r_{t'}$

Conditioning on $\hat{R}_t$ provides a clear forward-looking objective: "What sequence of actions will accumulate $\hat{R}_t$ reward from this state onward?"

Architecture and Causal Self-Attention

The Decision Transformer adopts a standard GPT-style causal decoder architecture:

1Inputs:     [ R̂₁ ]    [ s₁ ]    [ a₁ ]    [ R̂₂ ]    [ s₂ ]    ──> [ â₂ ] (Predict)2              │         │         │         │         │3Linear Proj: W_R       W_s       W_a       W_R       W_s4              │         │         │         │         │5Pos Embed:  + pos(1)  + pos(1)  + pos(1)  + pos(2)  + pos(2)6              ▼         ▼         ▼         ▼         ▼7┌───────────────────────────────────────────────────────────────┐8│              Causal Multi-Head Self-Attention                 │9│              (Lower-Triangular Mask: Past Only)               │10└───────────────────────────────────────────────────────────────┘11                                                      │12                                                      ▼13                                              Action Head W_act14                                                      │15                                                      ▼16                                              Predicted Action â₂

1. Modality Embedding Projections

Because returns, states, and actions live in different vector spaces, each modality is projected into a shared embedding space of dimension $d$:

$e_{R, t} = W_R \hat{R}_t + b_R, \quad e_{s, t} = W_s s_t + b_s, \quad e_{a, t} = W_a a_t + b_a$

2. Timestep Positional Embeddings

Unlike standard Transformers that use token index positions ($1, 2, 3, \dots, 3T$), the Decision Transformer uses timestep embeddings based on the environment step $t$. All three modality tokens belonging to timestep $t$ receive the same learned positional vector $W_{\text{pos}}(t)$:

$x_{R, t} = e_{R, t} + W_{\text{pos}}(t), \quad x_{s, t} = e_{s, t} + W_{\text{pos}}(t), \quad x_{a, t} = e_{a, t} + W_{\text{pos}}(t)$

This allows the model to retain temporal alignment across the episode regardless of context truncation.

3. Causal Attention Mask

A lower-triangular causal attention mask ensures that token predictions at time $t$ only attend to historical context. Specifically, when predicting action $a_t$, the model attends to:

$\{\hat{R}_1, s_1, a_1, \dots, \hat{R}_{t-1}, s_{t-1}, a_{t-1}, \hat{R}_t, s_t\}$

The prediction of $a_t$ has access to the current state $s_t$ and the desired return $\hat{R}_t$, but never sees future states or the true action $a_t$.

Training Objective and Autoregressive Rollout

Training

The model is trained entirely through supervised learning on offline trajectories sampled from dataset $\mathcal{D}$. For continuous action spaces, the objective is Mean Squared Error (MSE):

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

For discrete actions, cross-entropy loss over action logits is used. There are no critics, no target networks, and no discount factors $\gamma$.

Autoregressive Inference (Evaluation)

During live deployment in an environment:

1. Target Specification: The user sets an initial desired target return $\hat{R}_1$ (e.g., the maximum return in the dataset or an expert performance threshold).
2. State Observation: The agent observes initial state $s_1$.
3. Action Generation: The prompt $[\hat{R}_1, s_1]$ is passed to the Decision Transformer to predict action $a_1$.
4. Environment Step: The agent executes $a_1$, receiving reward $r_1$ and transition state $s_2$.
5. Return-to-Go Decay: The agent decrements the return-to-go: $\hat{R}_2 = \hat{R}_1 - r_1$
6. Context Append: The sequence is extended with $[a_1, \hat{R}_2, s_2]$, and the process repeats autoregressively across a sliding context window of length $K$ (typically $K=20$ timesteps).

Contrast with Trajectory Transformer (TT)

While both models frame RL as sequence modeling, the Trajectory Transformer (Janner et al., 2021) differs fundamentally in tokenization and planning:

Worked numerical example

Let us trace a concrete 3-step trajectory ($T=3$) through return-to-go calculation, context windowing, action loss computation, and inference decay.

Step 1: Compute Returns-to-Go

Suppose a demonstrated trajectory yields rewards:

• $r_1 = 2.0$
• $r_2 = 3.0$
• $r_3 = 5.0$

The total cumulative return is $G = 2.0 + 3.0 + 5.0 = 10.0$.

The return-to-go for each step is computed backward: $\hat{R}_1 = r_1 + r_2 + r_3 = 2.0 + 3.0 + 5.0 = 10.0$ $\hat{R}_2 = r_2 + r_3 = 3.0 + 5.0 = 8.0 \quad (\text{or } \hat{R}_1 - r_1 = 10.0 - 2.0 = 8.0)$ $\hat{R}_3 = r_3 = 5.0 \quad (\text{or } \hat{R}_2 - r_2 = 8.0 - 3.0 = 5.0)$

Step 2: Context Input at Timestep $t=2$

At timestep $t=2$, the causal context sequence presented to the Transformer consists of 5 tokens:

$\text{Context}_{t=2} = \left[ (\hat{R}_1=10.0),\, (s_1),\, (a_1),\, (\hat{R}_2=8.0),\, (s_2) \right]$

The Transformer processes these 5 tokens through causal self-attention. The query vector at the final position ($s_2$) attends to all 5 tokens to output the predicted action vector $\hat{a}_2$.

Step 3: Compute Action Prediction Loss

Suppose the true 2-dimensional continuous action in the dataset is: $a_2 = [0.50,\, 0.80]$

The Decision Transformer's linear action head outputs prediction: $\hat{a}_2 = [0.45,\, 0.85]$

The Mean Squared Error (MSE) loss is evaluated: $\mathcal{L}_2 = \frac{1}{D} \sum_{d=1}^D (\hat{a}_{2, d} - a_{2, d})^2 = \frac{1}{2} \left[ (0.45 - 0.50)^2 + (0.85 - 0.80)^2 \right]$ $\mathcal{L}_2 = \frac{1}{2} \left[ (-0.05)^2 + (0.05)^2 \right] = \frac{1}{2} [ 0.0025 + 0.0025 ] = \frac{0.0050}{2} = 0.0025$

This supervised gradient updates $W_a, W_s, W_R,$ and the self-attention weights via standard Adam optimizer steps.

Step 4: Inference Return-to-Go Decay

During real-time evaluation:

1. The user specifies target return $\hat{R}_1 = 10.0$.
2. The agent executes predicted action $a_1$ in the environment.
3. The environment emits reward $r_1 = 2.0$.
4. The agent updates its prompt return-to-go for step $t=2$: $\hat{R}_2 = \hat{R}_1 - r_1 = 10.0 - 2.0 = 8.0$

The agent's internal goal is dynamically updated: it now knows that to satisfy the requested 10.0 overall performance, it must collect exactly 8.0 remaining points over the remainder of the episode.

Code

Below is a self-contained, type-hinted Python implementation of the core Decision Transformer trajectory evaluation engine, including return-to-go computation, causal attention masking, continuous action MSE loss, and autoregressive return decay:

1from dataclasses import dataclass2import math3from typing import List, Tuple4
5
6@dataclass7class TrajectoryToken:8    """Represents an individual token in the interleaved trajectory sequence."""9
10    modality: str  # 'RTG', 'STATE', or 'ACTION'11    timestep: int12    value: float13
14
15class DecisionTransformerTrajectoryEvaluator:16    """Evaluates trajectory tokenization, causal masking, return-to-go decay, and action prediction loss."""17
18    def __init__(self, embedding_dim: int = 16) -> None:19        self.embedding_dim = embedding_dim20
21    def compute_returns_to_go(self, rewards: List[float]) -> List[float]:22        """Computes Return-to-go R_hat_t = sum_{t'=t}^T r_{t'} for each timestep."""23        t_len = len(rewards)24        rtg = [0.0] * t_len25        running_sum = 0.026        for i in reversed(range(t_len)):27            running_sum += rewards[i]28            rtg[i] = running_sum29        return rtg30
31    def construct_token_sequence(32        self,33        returns_to_go: List[float],34        states: List[List[float]],35        actions: List[List[float]],36    ) -> List[TrajectoryToken]:37        """Constructs interleaved token sequence: (R_hat_1, s_1, a_1, R_hat_2, s_2, a_2, ...)."""38        tokens: List[TrajectoryToken] = []39        for t in range(len(returns_to_go)):40            step_idx = t + 141            tokens.append(TrajectoryToken("RTG", step_idx, returns_to_go[t]))42            tokens.append(TrajectoryToken("STATE", step_idx, states[t][0]))43            tokens.append(TrajectoryToken("ACTION", step_idx, actions[t][0]))44        return tokens45
46    def build_causal_mask(self, num_tokens: int) -> List[List[int]]:47        """Builds lower-triangular causal attention mask (1 = attend, 0 = mask)."""48        mask = [49            [1 if j <= i else 0 for j in range(num_tokens)]50            for i in range(num_tokens)51        ]52        return mask53
54    def compute_action_loss_mse(55        self, pred_action: List[float], target_action: List[float]56    ) -> float:57        """Computes Mean Squared Error loss between predicted and ground-truth action vectors."""58        assert len(pred_action) == len(target_action)59        dim = len(pred_action)60        sq_diff_sum = sum(61            (p - t) ** 2 for p, t in zip(pred_action, target_action)62        )63        return sq_diff_sum / dim64
65    def step_inference_rtg(66        self, current_rtg: float, received_reward: float67    ) -> float:68        """Updates return-to-go for the next timestep: R_hat_{t+1} = R_hat_t - r_t."""69        return current_rtg - received_reward70
71
72if __name__ == "__main__":73    evaluator = DecisionTransformerTrajectoryEvaluator()74
75    # 1. Worked Numerical Example: T=3 trajectory76    rewards = [2.0, 3.0, 5.0]77    rtg = evaluator.compute_returns_to_go(rewards)78
79    print("=== Returns-to-Go Calculation ===")80    print(f"Step Rewards:    {rewards}")81    print(f"Returns-to-Go:   {rtg}")82    assert rtg == [10.0, 8.0, 5.0]83
84    # 2. Context at t=2: Action Loss Evaluation85    pred_a2 = [0.45, 0.85]86    gt_a2 = [0.50, 0.80]87    mse_loss = evaluator.compute_action_loss_mse(pred_a2, gt_a2)88
89    print("\n=== Action Prediction Loss (MSE) ===")90    print(f"Predicted Action a_hat_2: {pred_a2}")91    print(f"Ground Truth Action a_2:  {gt_a2}")92    print(f"MSE Loss:                 {mse_loss:.6f}")93    assert math.isclose(mse_loss, 0.0025)94
95    # 3. Inference RTG Decay96    # Starting at R_hat_1 = 10.0, agent executes a_1 and receives r_1 = 2.097    rtg_t2 = evaluator.step_inference_rtg(98        current_rtg=rtg[0], received_reward=rewards[0]99    )100    print("\n=== Inference Return-to-Go Decay ===")101    print(102        f"Prompt RTG at t=2: R_hat_2 = R_hat_1 - r_1 = {rtg[0]} - {rewards[0]} = {rtg_t2:.1f}"103    )104    assert math.isclose(rtg_t2, 8.0)105
106    # 4. Token Sequence and Causal Mask Verification107    mock_states = [[1.0], [1.2], [1.5]]108    mock_actions = [[0.5], [0.8], [0.9]]109    tokens = evaluator.construct_token_sequence(rtg, mock_states, mock_actions)110    causal_mask = evaluator.build_causal_mask(len(tokens))111
112    print("\n=== Causal Mask Verification ===")113    print(f"Total Trajectory Tokens (3 steps x 3 modalities): {len(tokens)}")114    print(f"Token 0 ({tokens[0].modality}_1) attend vector: {causal_mask[0]}")115    print(f"Token 4 ({tokens[4].modality}_2) attend vector: {causal_mask[4]}")116
117    assert len(tokens) == 9118    assert causal_mask[0] == [1, 0, 0, 0, 0, 0, 0, 0, 0]119    assert causal_mask[4] == [1, 1, 1, 1, 1, 0, 0, 0, 0]120    print("\nAll Decision Transformer assertions passed successfully!")

Expected output:

1=== Returns-to-Go Calculation ===2Step Rewards:    [2.0, 3.0, 5.0]3Returns-to-Go:   [10.0, 8.0, 5.0]4
5=== Action Prediction Loss (MSE) ===6Predicted Action a_hat_2: [0.45, 0.85]7Ground Truth Action a_2:  [0.5, 0.8]8MSE Loss:                 0.0025009
10=== Inference Return-to-Go Decay ===11Prompt RTG at t=2: R_hat_2 = R_hat_1 - r_1 = 10.0 - 2.0 = 8.012
13=== Causal Mask Verification ===14Total Trajectory Tokens (3 steps x 3 modalities): 915Token 0 (RTG_1) attend vector: [1, 0, 0, 0, 0, 0, 0, 0, 0]16Token 4 (STATE_2) attend vector: [1, 1, 1, 1, 1, 0, 0, 0, 0]17
18All Decision Transformer assertions passed successfully!

Watch Out For

The Quick Version

• RL as Sequence Modeling: Decision Transformer replaces recursive Bellman backups with conditional autoregressive sequence modeling, using a GPT-style causal decoder to predict actions from historical context.
• Interleaved Trajectory Tokens: Episodes are structured as token triplets $(\hat{R}_t, s_t, a_t)$ paired with timestep embeddings, where Return-to-Go $\hat{R}_t = \sum_{t'=t}^T r_{t'}$ provides the forward conditioning target.
• Pure Supervised Training: The model is trained using standard mean squared error (continuous actions) or cross-entropy (discrete actions), eliminating target networks, value critics, and offline Bellman instability.
• Inference via Return Decay: The agent is prompted with a desired high return $\hat{R}_1$ and decrements the target dynamically after each environment reward ($\hat{R}_{t+1} = \hat{R}_t - r_t$), guiding the Transformer to fulfill the overall episode objective.

What to Read Next

• Implicit Q-Learning (IQL)
• Conservative Q-Learning (CQL)
• Offline / Batch Reinforcement Learning