Decision Transformer

Unlike traditional RL that learns a state-to-action mapping, the Decision Transformer models trajectories as sequences. It predicts the next action token given a context window of returns-to-go, states, and past actions.

• Uses a GPT-style causal transformer architecture
• The desired return is fed as an input, conditioning the entire sequence
• Enables stitching of suboptimal trajectories by prompting with high returns
• The model learns a return-conditioned policy: p(a_t | R_t, s_t, a_{t-1}, ...)

The model is trained via supervised learning on offline trajectory data, not via temporal difference learning. Each timestep's target is the actual action taken, conditioned on the achieved return-to-go.

• Eliminates bootstrapping and the deadly triad of RL
• No discount factor or value function approximation required
• Learns from static datasets without environment interaction
• Can generate behaviors ranging from random to expert by varying the target return prompt

Because the architecture is a standard transformer, the policy's decision-making process can be inspected using attention rollout and other mechanistic interpretability tools.

• Visualize which past states and actions the model attends to when selecting an action
• Identify if the model focuses on causal state features or spurious correlations
• Apply saliency maps to input tokens to see which elements drive a specific action
• Enables contrastive explanations by comparing attention patterns across different return prompts

The model can combine segments from different suboptimal trajectories to produce optimal behavior, a phenomenon called trajectory stitching.

• Learns a latent model of environment dynamics implicitly through sequence modeling
• Generalizes to unseen state-action combinations by composing known subsequences
• Outperforms conventional offline RL methods like CQL and BCQ on sparse-reward benchmarks
• The transformer's in-context learning ability enables few-shot adaptation to new dynamics

A single Decision Transformer can be trained on heterogeneous datasets spanning multiple tasks or goals by conditioning on task-specific returns or goal tokens.

• No need for separate policy networks per task
• The return-to-go acts as a continuous goal specification
• Can zero-shot generalize to intermediate return targets not seen in training
• Enables interpretable goal decomposition by analyzing how attention shifts with different return prompts

The Decision Transformer's explicit sequence representation makes it amenable to causal intervention analysis. Researchers can modify specific tokens in the context window and observe the counterfactual effect on generated actions.

• Perform feature ablation by zeroing out state dimensions and measuring action change
• Apply integrated gradients from a zero-return baseline to the actual return prompt
• Test for spurious correlations by intervening on specific past states
• Compatible with TCAV to test sensitivity to high-level concepts encoded in the state tokens