From RL to Sequence Modeling: Understanding the Decision Transformer
From RL to Sequence Modeling: Understanding the Decision Transformer
Reinforcement Learning (RL) has traditionally been a battle of Bellman equations, value functions, and policy gradients. For years, the goal was to estimate the "value" of a state or the "gradient" of a policy to maximize future rewards.
But what if we stopped treating RL as a reinforcement problem and started treating it as a sequence modeling problem?
Enter the Decision Transformer (DT). By reimagining trajectories as sequences of tokens, the Decision Transformer leverages the power of the GPT architecture to perform policy optimization through conditional generation.
🧠 The Core Intuition: RL as Conditional Sequence Modeling
The fundamental shift in the Decision Transformer is the move from bootstrapping (predicting future values based on other predicted values) to supervised learning.
In a standard RL agent, the model asks: "What is the expected future reward if I take this action?" In a Decision Transformer, the model asks: "Given that I want to achieve this specific total reward, what action should I take based on my history?"
By conditioning the model on the Return-to-go (RTG), the DT treats the desired reward as a "prompt." At inference time, if you prompt the model with a high target return, it generates the actions that historically led to those high returns.
The Mathematical Framework
The DT represents a trajectory $\tau$ not as a Markov chain, but as a sequence of triplets:
$$\text{Trajectory Representation: } \tau = (\hat{R}_1, s_1, a_1, \hat{R}_2, s_2, a_2, \dots, \hat{R}_T, s_T, a_T)$$
Where:
- $s_t$: The state at time $t$.
- $a_t$: The action taken at time $t$.
- $\hat{R}t$: The Return-to-go, defined as the sum of future rewards: $$\hat{R}t = \sum{t'=t}^{T} r{t'}$$
The model uses a standard Self-Attention mechanism to weigh the importance of past states and rewards: $$z_i = \sum_{j=1}^{n} \text{softmax}\left(\frac{q_i k_j^T}{\sqrt{d_k}}\right) v_j$$
🏗️ Architecture Deep Dive
The Decision Transformer utilizes a GPT-style causal Transformer. Here is the step-by-step flow of how data moves through the system.
1. Data Preprocessing
Offline RL trajectories are converted into sequences. Because Transformers require fixed-length inputs, trajectories are either sliced or padded to a max_len. The raw rewards are transformed into the cumulative Return-to-go ($\hat{R}_t$).
2. Modality Embedding
Since $\hat{R}$, $s$, and $a$ have different dimensions, they are passed through separate linear embedding layers to project them into a shared $d_{model}$ dimensional space. A learned episodic timestep embedding is added to provide the model with a sense of temporal position.
3. Causal Sequence Processing
The tokens are interleaved as $(R_1, S_1, A_1, R_2, S_2, A_2 \dots)$. A causal mask (triangular mask) is applied, ensuring that the prediction for action $a_t$ only depends on tokens that occurred at or before time $t$.
4. Action Prediction
The model doesn't predict the reward or the state; it predicts the action. The output at the position corresponding to the state $s_t$ is passed through a final linear head to produce the predicted action $\hat{a}_t$.
Visual Workflow
💻 Implementation in PyTorch
Below is a production-ready simplified implementation of the Decision Transformer.
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
from torch.utils.data import DataLoader, Dataset
class DecisionTransformer(nn.Module):
def __init__(self, state_dim, action_dim, init_hparams):
super().__init__()
self.state_dim = state_dim
self.action_dim = action_dim
self.max_len = init_hparams['max_len']
self.d_model = init_hparams['d_model']
# Modality-specific linear embeddings
self.embed_return = nn.Linear(1, self.d_model)
self.embed_state = nn.Linear(state_dim, self.d_model)
self.embed_action = nn.Linear(action_dim, self.d_model)
self.embed_timestep = nn.Embedding(self.max_len, self.d_model)
# GPT-style Causal Transformer
self.transformer = nn.TransformerEncoder(
nn.TransformerEncoderLayer(
d_model=self.d_model,
nhead=init_hparams['n_head'],
dim_feedforward=self.d_model * 4,
batch_first=True
),
num_layers=init_hparams['n_layer']
)
self.predict_action = nn.Linear(self.d_model, action_dim)
def forward(self, states, actions, returns_to_go, timesteps):
batch_size, seq_len, _ = states.shape
# 1. Embeddings
ret_emb = self.embed_return(returns_to_go)
state_emb = self.embed_state(states)
action_emb = self.embed_action(actions)
time_emb = self.embed_timestep(timesteps)
# 2. Interleave tokens: R1, S1, A1, R2, S2, A2...
stack = torch.stack([ret_emb, state_emb, action_emb], dim=2)
stack = stack.view(batch_size, seq_len * 3, self.d_model)
# Add positional embeddings
time_emb_expanded = time_emb.repeat_interleave(3, dim=1)
stack = stack + time_emb_expanded
# 3. Causal Masking
mask = torch.triu(torch.ones(stack.size(1), stack.size(1)), diagonal=1).bool().to(states.device)
# 4. Transformer Forward Pass
output = self.transformer(stack, mask=mask)
# 5. Predict Actions (extract from state positions: 1, 4, 7...)
state_indices = torch.arange(1, stack.size(1), 3).to(states.device)
state_outputs = output[:, state_indices, :]
return self.predict_action(state_outputs)🚀 Inference: How to "Prompt" for Success
The most exciting part of the Decision Transformer is how it is used after training. Unlike traditional RL, where the policy is fixed, the DT allows you to specify the desired outcome.
The Inference Loop:
- Set the Goal: Initialize $\hat{R}_1$ with a high target return (e.g., the maximum reward seen in the training set).
- Observe: Get the current state $s_1$ from the environment.
- Predict: Feed $(\hat{R}_1, s_1)$ into the Transformer to predict action $a_1$.
- Act: Execute $a_1$, receive reward $r_1$ and next state $s_2$.
- Update: Calculate the new return-to-go: $\hat{R}_2 = \hat{R}_1 - r_1$.
- Repeat: Use $(\hat{R}_1, s_1, a_1, \hat{R}_2, s_2)$ to predict $a_2$.
🏁 Conclusion
The Decision Transformer represents a paradigm shift in Offline RL. By treating RL as a conditional sequence modeling problem, it eliminates the instability of Bellman bootstrapping and allows us to apply the massive successes of Large Language Models (LLMs) to robotic control and decision-making.
Key Takeaways:
- No more Q-learning: Optimization is handled via supervised learning on trajectories.
- Conditioning is Key: The Return-to-go acts as a target "prompt" for the agent.
- Scalability: It leverages the Transformer architecture, meaning it can scale with more data and larger models just like GPT.