Distributional Decision Transformer: Risk-Sensitive Offline RL via Quantile-Based Critics and Stochastic Return

Abstract

Offline reinforcement learning faces a critical challenge in synthesizing high-reward trajectories from suboptimal datasets while robustly handling the stochasticity inherent in real-world decision-making. While combination of return-conditioned sequence models, such as Decision Transformers (DT), and dynamics programming critics shows great potential in trajectory synthesis, their deterministic action generation and scale Q value critic often fails to distinguish intentional behavioral variability from detrimental noise, leading to suboptimal policy collapse. To address this challenge, we propose the Distributional Decision Transformer (DDT), a novel framework that unifies probabilistic return distribution modeling with autoregressive action generation. DDT introduces two key innovations: (1) a Gaussian stochastic return mechanism that reparameterizes target returns as samplable distributions, enabling diverse action candidate generation; and (2) an Implicit Quantile Network (IQN) critic embedded within the deciding loop, which evaluates actions across the full spectrum of return distributions (quantiles τ ~ U(0, 1)). In D4RL benchmarks, DDT achieves state-of-the-art performance, achieving a 91.6 average normalized score in MuJoCo locomotion and 69.3 in sparse-reward settings. The results establish DDT as a principled solution for synthesis of risk-aware trajectory in offline RL.