Decoupling Return-to-Go for Efficient Decision Transformer

TL;DR

The paper addresses a redundancy in the Decision Transformer where conditioning on the full Return-to-Go (RTG) history is unnecessary for predicting actions. It introduces Decoupled DT (DDT), which processes only observations and actions in the Transformer and uses the latest RTG $\hat{R}_t$ to condition the output via Adaptive Layer Normalization (adaLN), reducing input length and computation. Empirically, DDT outperforms DT on D4RL benchmarks and remains competitive with state-of-the-art DT variants, while also generalizing to discrete tasks like 2048. The work discusses generalized conditioning scenarios and the role of RTG sequences in non-pomdp or reward-rich settings, highlighting future directions for broader conditioning paradigms across DT variants.

Abstract

The Decision Transformer (DT) has established a powerful sequence modeling approach to offline reinforcement learning. It conditions its action predictions on Return-to-Go (RTG), using it both to distinguish trajectory quality during training and to guide action generation at inference. In this work, we identify a critical redundancy in this design: feeding the entire sequence of RTGs into the Transformer is theoretically unnecessary, as only the most recent RTG affects action prediction. We show that this redundancy can impair DT's performance through experiments. To resolve this, we propose the Decoupled DT (DDT). DDT simplifies the architecture by processing only observation and action sequences through the Transformer, using the latest RTG to guide the action prediction. This streamlined approach not only improves performance but also reduces computational cost. Our experiments show that DDT significantly outperforms DT and establishes competitive performance against state-of-the-art DT variants across multiple offline RL tasks.

Figures (5)

• Figure 1: DDT architecture. States/observations and actions are first projected via modality-specific linear embeddings, followed by a positional encoding. A GPT backbone with causal self-attention processes the resulting tokens. RTGs are integrated via adaLN to modulate the output token for action prediction.
• Figure 2: In the rtg sequence, the difference between successive terms is the reward. These rewards do not contribute to the belief state. Therefore, they are redundant information in history.
• Figure 3: During action prediction, DDT takes only the most recent RTG $\hat{R}_t$ as condition. This condition applies to the hidden state corresponding to $o_t$ via adaLN, and the result is subsequently fed into $\text{Pred}_A$ (an MLP) to output the action $a_t$.
• Figure 4: Environments for evaluation.
• Figure 5: Average attention scores over the first $10^3$ inference steps for the learned DT and DDT policies during action prediction in the respective environment (hopper and halfcheetah).