Beyond Autoregressive RTG: Conditioning via Injection Outside Sequential Modeling in Decision Transformer

Abstract: Decision Transformer (DT) formulates offline reinforcement learning as autoregressive sequence modeling, achieving promising results by predicting actions from a sequence of Return-to-Go (RTG), state, and action tokens. However, RTG is a scalar that summarizes future rewards, containing far less information than typical state or action vectors, yet it consumes the same computational budget per token. Worse, the self-attention cost of Transformers grows quadratically with sequence length, so including RTG as a separate token adds unnecessary overhead. We propose SlimDT, which removes RTG from the autoregressive sequence. Instead, we inject RTG information into the state representations before the sequential modeling step, allowing the Transformer to process only a compact (state, action) sequence. This reduces the sequence length by one-third, directly improving inference efficiency. On the D4RL benchmark, SlimDT surpasses standard DT across various tasks and achieves performance comparable to existing state-of-the-art methods. Decoupling a sparse conditioning signal from an information-rich sequence thus yields both computational gains and higher task performance.

• The paper introduces SlimDT, which injects Return-to-Go into state representations instead of modeling it as an autoregressive token, reducing sequence length from 3k to 2k and theoretical attention cost to 4/9 of standard Decision Transformer.
• The method improves over Decision Transformer on every evaluated D4RL task, reaching aggregate scores of 82.6 on MuJoCo and 53.1 on Adroit, with especially large gains on Hopper-medium, Pen, and Hammer benchmarks.
• The experiments show that injector choice should match task structure: post-conditioning with adaptive normalization suits smooth-reward tasks, while pre-conditioning with cross-attention better handles volatile, temporally structured rewards, though automatic selection remains open.

Motivation: RTG as an inefficient token

The Decision Transformer (DT) formulates offline reinforcement learning as autoregressive sequence modeling, processing interleaved Return-to-Go (RTG), state, and action tokens to predict actions conditioned on past trajectories. The authors of this paper identify a structural inefficiency in this design: the RTG is a scalar summarizing future cumulative reward, yet it occupies a full position in the input sequence and undergoes the same self-attention computation as information-rich state and action embeddings. Because self-attention cost grows quadratically with sequence length, allocating one of three token slots per timestep to a low-information scalar inflates inference cost without commensurate informational benefit. The paper argues that conditioning signals should be decoupled from the primary sequential modeling pipeline rather than forced into it.

The SlimDT architecture

SlimDT removes the RTG token from the autoregressive sequence entirely and instead injects RTG information into state representations via a condition injector. This reduces the sequence fed to the GPT backbone from $3k$ to $2k$ tokens for context length $k$, so—ignoring injector overhead—the theoretical inference cost drops from $9k^2$ to $4k^2$, i.e., to $4/9$ of standard DT.

The paper proposes two variants:

• Pre-conditioning (pre-cond): the RTG sequence is fused with the raw state sequence before the Transformer; the return-conditioned states are then interleaved with actions for autoregressive action prediction.
• Post-conditioning (post-cond): the state-action sequence is first encoded by the Transformer, and the resulting trajectory representations are subsequently modulated by the RTG sequence to produce action predictions.

Three condition injectors are examined. Concatenation linearly projects the RTG to a higher dimension before concatenating with the state and re-encoding; naive concatenation without projection fails. Adaptive Layer Normalization (adaLN), borrowed from conditional generative modeling, uses an MLP over the RTG embedding to produce per-layer scale and shift parameters that modulate normalized state features. Cross-attention performs sequence-to-sequence fusion across the full time window.

A notable contribution is the analysis of why standard cross-attention fails as an injector here. In conventional encoder-decoder cross-attention, queries come from the information-sparse decoder sequence while keys and values carry richer encoder context. In SlimDT the situation is inverted: the conditioning RTG sequence is sparse and the state sequence is rich, so assigning both $K$ and $V$ to the RTG misallocates representational capacity. The authors therefore propose two variants—$\text{CrossAttn}_{Q\rightarrow KV}$ and $\text{CrossAttn}_{K\rightarrow QV}$—in which exactly one of $2k$0 or $2k$1 derives from the RTG while the other, together with $2k$2, derives from the state sequence. Empirically only these two converge; the remaining configurations ($2k$3, $2k$4) fail to train.

Experimental results

Evaluations use D4RL, covering MuJoCo locomotion (Hopper, HalfCheetah, Walker2d across medium, medium-replay, medium-expert) and Adroit dexterous manipulation (Pen, Hammer, Door under human and cloned datasets), with scores averaged over 5 seeds × 100 rollouts. Against DT, SlimDT improves on every task evaluated. Aggregate results are competitive with strong baselines including CQL, IQL, BCQ, StAR, GDT, DC, and DHi:

Benchmark	DT	Best baseline avg	SlimDT
MuJoCo (avg)	76.2	83.3 (DHi)	82.6
Adroit (avg)	32.2	52.5 (DC)	53.1

Several individual results stand out. On Hopper-medium, SlimDT reaches 99.4 versus 67.6 for DT and 92.5 for the strongest prior method (DC). On Pen-human it attains 94.4 against DT's 79.5, and on Pen-cloned 99.3 versus 75.8. On Hammer-human, SlimDT achieves 39.1 where DT manages only 3.7. These gains indicate that removing the RTG token is not merely a computational optimization but can materially improve learning, plausibly by freeing model capacity from attending to redundant scalar tokens.

The ablations yield task-dependent design guidance. On relatively simple MuJoCo tasks with smooth reward signals, post-cond combined with adaLN performs best, since the state-action sequence alone suffices for decision-making. On Adroit, where rewards are highly volatile and minor state changes produce dramatic reward jumps, pre-cond with cross-attention is optimal, because sequence-to-sequence fusion better exploits temporal structure in the RTG signal (e.g., sudden reward changes signaling events not represented in states). Among cross-attention variants, $2k$5 without causal masking is generally best on average. Causal masking sometimes helps on high-quality human datasets—interpreted as regularization against over-parameterization—but degrades performance on lower-quality cloned data. For concatenation-based injection, projecting the RTG to roughly half the hidden dimension before concatenation yields the best results (average 45.6 at dimension 64), approaching but not exceeding cross-attention while being simpler and faster. Combining pre- and post-conditioning produces intermediate performance and is not recommended.

Limitations and open questions

The paper is candid about several constraints. First, adaLN fails on Adroit despite excelling on MuJoCo; the hypothesized cause—point-wise modulation cannot capture abrupt multi-timestep reward changes—is plausible but not directly verified, and a hybrid injector that adaptively selects between point-wise and sequence-wise fusion based on reward smoothness remains unbuilt. Second, the choice among pre-cond/post-cond and cross-attention variants currently requires task-specific tuning; no lightweight heuristic or learned selection mechanism is provided. Third, the efficiency argument rests partly on the quadratic complexity of self-attention; the authors note that shortening sequences should still benefit non-Transformer backbones such as Mamba, but adaptation to other sequence models is untested. Finally, the reported variance on some Adroit tasks is nontrivial (e.g., ±8.7 on Pen-human, ±9.6 on Pen-cloned), so fine-grained rankings among top methods there warrant caution.

Conclusion

SlimDT demonstrates that the autoregressive treatment of RTG in Decision Transformers is neither necessary nor optimal. By injecting the RTG signal into state representations outside the sequential pipeline, the method cuts the modeled sequence length by one third, reduces quadratic inference cost to $2k$6 of the original, and simultaneously improves task performance—outperforming DT on all evaluated D4RL tasks and matching state-of-the-art methods. The central finding, that decoupling sparse conditioning signals from information-rich sequential modeling benefits both efficiency and accuracy, is supported by systematic ablations linking injector choice to reward volatility and data quality. The main open questions concern automating architecture selection across regimes and validating the decoupling principle beyond Transformer backbones.