Decision Transformer

The model is conditioned on a Return-to-Go (RTG) token at each timestep, representing the cumulative reward the agent should achieve from that point forward. This shifts the paradigm from reward maximization to goal-directed sequence generation.

• Input Format: Sequences are structured as (RTG₁, state₁, action₁, RTG₂, state₂, action₂, ...).
• Autoregressive Prediction: The model predicts the next action token given the history of RTGs, states, and past actions.
• Goal Specification: At inference, a high initial RTG prompts the model to generate a high-reward trajectory, while a low RTG may elicit conservative or recovery behavior.

A standard decoder-only transformer forms the core, processing the interleaved sequence of returns, states, and actions. It uses causal attention masks to ensure predictions are based only on past and present tokens.

• Token Embeddings: RTG, state, and action inputs are projected into a shared embedding space via separate linear layers.
• Positional Encoding: Standard sinusoidal or learned embeddings provide sequence order.
• Causal Self-Attention: Allows the model to build rich representations of the trajectory history while maintaining the autoregressive property for action generation.

Continuous actions must be discretized into tokens for the transformer to process. Common methods include:

• Vector Quantization (VQ): Using a VQ-VAE to learn a discrete codebook of action prototypes.
• Residual VQ: A hierarchical method for higher-fidelity reconstruction of complex action spaces.
• Straight-Through Estimation: Enables gradient flow through the non-differentiable quantization step during training.
• Alternative: For continuous outputs, a final linear projection layer can directly predict action values, though this is less common in pure sequence modeling formulations.

The model generates actions one token at a time, conditioned on the evolving sequence. This is analogous to next-token prediction in language modeling.

• Training Objective: Minimizes the cross-entropy loss (for discrete actions) or mean-squared error (for continuous actions) between predicted and demonstrated actions.
• Inference: Uses autoregressive decoding, where each predicted action is appended to the input sequence to predict the next.
• Sampling Techniques: Methods like temperature sampling or beam search can be applied to balance diversity and quality in the generated action sequences.

Decision Transformers are trained offline on static datasets of previously collected trajectories, without any environment interaction during training. This makes them highly sample-efficient and safe for real-world robotics.

• Dataset: Consists of sequences of (state, action, reward).
• No Bellman Backup: Unlike traditional RL, it does not perform dynamic programming or estimate value functions, avoiding extrapolation errors common in offline RL.
• Behavioral Cloning Aspect: The training resembles conditional behavioral cloning, where the conditioning signal is the desired return (RTG).

The architecture inherently models the temporal dependencies and long-range context within a trajectory. The transformer's self-attention mechanism allows any past state or action to directly influence future predictions.

• Context Window: The model's effective planning horizon is limited by the transformer's context length (e.g., 512 or 1024 tokens).
• Chunking: For long-horizon tasks, trajectories may be segmented into chunks processed independently.
• Advantage over RNNs: Avoids the vanishing gradient problem and allows parallel training across the sequence.

The Decision Transformer excels in offline RL, where an agent must learn an optimal policy from a fixed, previously-collected dataset without further environment interaction. It treats the dataset as sequences of (state, action, return) tuples and learns to predict actions conditioned on high returns-to-go, effectively imitating high-return trajectories. This makes it robust and safe for real-world applications like robotics, where online exploration is costly or dangerous.

In robotics, the model is used for goal-conditioned manipulation, where a robot must perform tasks like picking, placing, or assembling objects. The desired goal (e.g., 'place block A on B') is encoded as a target return or a language instruction. The transformer, conditioned on this goal and past states, autoregressively predicts a sequence of end-effector poses or joint angles. This allows for flexible, multi-task learning from demonstration datasets.

The architecture can be used for trajectory optimization by framing it as sequence generation. Starting from an initial state and a desired outcome, the model generates a sequence of future states and actions. By adjusting the conditioning return-to-go, users can steer the generated trajectory toward higher rewards or specific outcomes, providing a differentiable alternative to traditional sampling-based planners like Model Predictive Control (MPC).

Decision Transformers are highly effective for imitation learning from human demonstrations. The model learns the mapping from states to actions present in the expert dataset. By conditioning on the maximum observed return in the demonstrations, it learns to replicate expert behavior. This is crucial for training robots in complex tasks where designing a reward function is difficult, such as surgical robotics or autonomous driving.

The model's conditioning mechanism naturally supports multi-task learning. A single transformer can be trained on datasets from multiple tasks, with each task identified by a unique task embedding or a language descriptor prepended to the sequence. During inference, specifying a different task embedding steers the policy to perform the corresponding task. This facilitates efficient knowledge transfer and meta-learning for rapid adaptation to new, related tasks.

When integrated with a vision-language model, the Decision Transformer becomes a powerful language-guided policy for embodied agents. A natural language instruction (e.g., 'open the top drawer') is encoded and used as the conditioning signal instead of, or in addition to, a numerical return. The model then generates the sequence of actions needed to fulfill the command, enabling robots to follow open-ended, high-level instructions in interactive environments.

The process of converting continuous, high-dimensional physical actions (e.g., joint torques, end-effector velocities) into a sequence of discrete symbols or tokens. This enables sequence models like transformers to process actions as a language.

• Core Method: Often uses a Vector Quantized Variational Autoencoder (VQ-VAE) to learn a discrete codebook.
• Purpose: Transforms the control problem into a sequence prediction task, analogous to next-token prediction in NLP.
• Example: A 7-DoF robot arm's continuous joint angles over 10 timesteps might be mapped to the token sequence [A45, B12, C89, ...].

A neural network architecture critical for action tokenization. It learns a discrete latent representation by mapping encoder outputs to the nearest vector in a learned codebook.

• Mechanism: The encoder produces a continuous latent vector. A quantization step replaces it with the closest codebook entry. The decoder reconstructs the original input from this discrete code.
• Role in Robotics: Provides a compressed, discrete vocabulary of reusable action primitives. The codebook indices become the action tokens for the transformer.
• Key Property: Uses a straight-through estimator to allow gradients to flow through the non-differentiable quantization step during training.

The sequential generation method used by Decision Transformers to predict action tokens. The model predicts the next token conditioned on all previously generated tokens and the context (desired return, past states).

• Process: Given the sequence [Return, S1, A1, S2, A2, ...], the model predicts A_t based on all tokens up to S_t.
• Implementation: Enabled by causal masking in the transformer's self-attention, which prevents attending to future positions.
• Inference: Can use greedy decoding (choose highest probability token) or temperature sampling to control output diversity.

A control function that generates actions based on both the current state and a specified goal. The Decision Transformer is a specific type of goal-conditioned policy where the goal is expressed as a desired return-to-go (RTG).

• Contrast with DT: Traditional goal-conditioned policies might take a goal image or state. The DT conditions on a scalar RTG, framing the problem as achieving a target cumulative reward.
• Flexibility: By varying the input RTG at inference time, the same DT can exhibit conservative or aggressive behavior to meet different performance targets.
• Connection: Highlights the DT's shift from learning a value function to directly modeling the distribution of actions conditioned on desired outcomes.

A closely related architecture that generalizes the Decision Transformer concept. It models entire trajectories (states, actions, rewards) as a single unified sequence of tokens, allowing for more flexible in-context learning and planning.

• Key Difference: While a DT autoregressively predicts actions, a Trajectory Transformer can be prompted with partial trajectories and used to plan by generating sequences of future states, actions, and rewards.
• Mechanism: Uses a single transformer to tokenize all modalities (state, action, reward) into a flat sequence. Planning is performed via beam search over this sequence space.
• Significance: Represents a move towards foundation models for decision-making, where a single model can handle prediction, imitation, and goal-reaching via prompting.

A function that predicts the action taken to transition from one state to the next. This is a fundamental component often used implicitly or explicitly in action decoding pipelines.

• Relation to DT: When a DT predicts an action token A_t given S_t and S_{t+1} (via the RTG context), it is effectively performing a form of inverse dynamics prediction.
• Decoding Role: After a discrete action token is predicted, an inverse dynamics model (or a decoder from a VQ-VAE) is needed to map it back to a continuous motor command executable by the robot.
• Learning: Can be a separate, explicitly trained module or an emergent capability of the end-to-end model.