Return-to-Go Is More Than a Number: Q-Guided Alignment for Return-Conditioned Supervised Learning

Conditioned Sequence Models (CSMs) learn policies by treating return-to-go (RTG) as a control signal. However, existing CSMs often treat the RTGs as simple numerical inputs rather than aligning them with the performance of their policies. In this paper, we propose Q-ALIGN DT, a framework that enforces this alignment by ensuring the $Q$-value of the output policy is consistent with the input RTG. By leveraging a $Q$ function to provide dense guidance to CSMs and further fine-tuning it using an *RTG-perturbation* technique with the CSM, our method ensures that higher RTGs are consistently mapped to trajectories with higher expected returns. Theoretically, we show that Q-ALIGN DT can efficiently learn the desired policy and output a near-optimal one when the RTG is sufficiently high. Empirically, we demonstrate through extensive experiments that Q-ALIGN DT achieves superior controllability and performance across the D4RL benchmark. Remarkably, our model effectively learns a structured family of policies that maintains precise alignment and generalizes to tasks like velocity-tracking where prior methods fail.

Many return-conditioned policies are trained to choose actions based on a desired future reward. However, these systems often do not reliably change their behavior when the requested reward changes: asking for a higher return may not actually lead to better behavior. This makes such models hard to control, especially when one model is expected to represent many different behaviors. We propose Q-ALIGN DT, a method that uses an additional value model, called a critic, to guide a return-conditioned policy. The key idea is to encourage the policy to produce actions whose predicted value changes consistently when the input return is changed. This helps the model learn a more meaningful relationship between the return signal and its behavior. In experiments on standard offline reinforcement learning benchmarks, our method improves both task performance and controllability. It can also adjust behaviors such as running speed by changing only the return input, suggesting that return-conditioned models can become more reliable and interpretable control systems.