Reinforcement Learning via Advantage Guidance
Offline RL promises to outperform the data but is notoriously unstable, whereas imitation learning is stable yet bounded by data quality. This paper introduces Advantage Guidance (AG), a policy extraction method that retains a simple supervised-learning objective—training a flow policy conditioned on advantage estimates—while correcting the failure mode where naïve conditioning is ignored by the network. At test time, AG uses a derived contrastive diffusion guidance to target the KL-regularized optimal policy $\pi_{\mathrm{KL}}^*(a \mid s) \propto \hat{\pi}(a \mid s)\exp(A(s,a)/\beta)$. On 26 long-horizon sparse-reward manipulation tasks, the method outperforms strong actor-critic, supervised-learning, and diffusion-guidance baselines.
This is a strong, well-motivated paper that delivers both a clean theoretical insight and practical training fixes for a genuine problem it names the advantage alignment problem. The derivation of contrastive guidance for continuous advantage conditioning is elegant, and the empirical gains over prior policy-extraction baselines are large and consistent across benchmarks.
The paper’s central diagnosis—that naïve advantage-conditioned policies suffer from an advantage alignment problem wherein gradients and outputs are largely insensitive to the conditioning value—is thoroughly validated. The proposed fixes, residual parameterization and action imagination, are simple, well-motivated, and ablated cleanly. The contrastive guidance derivation (Remark 1) rigorously shows how score differences between two conditional queries can recover the desired exponentially tilted distribution.
The method is deeply dependent on the quality of pre-trained IQL value estimates; the paper uses a basic $A(s,a) = Q(s,a) - V(s)$ estimate and does not ablate robustness to critic noise or misprediction. The evaluation is also limited to continuous robotic manipulation with action chunking, leaving open whether the alignment problem and proposed fixes transfer to discrete actions or non-robotic domains. Finally, there is a minor but sloppy typo in Figure 1 ('2626 tasks' rather than the 26 tasks reported elsewhere), which suggests insufficient proofreading.
The evidence convincingly supports the claim that AG improves advantage alignment and outperforms the selected baselines on OGBench and OCBench. Ablations isolate residual parameterization, action imagination, and contrastive guidance effectively. Comparisons to related work are generally fair: the authors correctly note that CFGRL discards information via binary advantage filtering and that QGF relies on unstable critic gradients in high-dimensional action spaces. One caveat is that best-of-$N$ sampling eventually surpasses AG on OGBench when granted enough test-time compute, a limitation the authors disclose honestly.
Reproducibility is strong on paper: source code is publicly released and hyperparameters are catalogued in detail in Appendix A.8, including network sizes, flow integration steps, and per-domain guidance weights. The main barrier to independent reproduction is computational—the method relies on first training IQL critics and then performing action imagination by sampling $K$ counterfactual actions per dataset transition from a reference flow policy, which increases training cost. Additionally, the per-state advantage normalization requires computing running statistics over imagined actions per state, adding implementation complexity.
Reinforcement learning (RL) holds the promise of learning a policy that can outperform what is shown in the data. However, in practice, such systems suffer from instability and hyperparameter sensitivity. In contrast, supervised learning provides a stable and scalable training objective, but is fundamentally limited by the quality of the demonstrated behavior. In this paper, we ask whether we can perform policy optimization with something as simple and stable as supervised learning but as performant as RL. Unfortunately, naively training the policy with supervised learning and conditioning on future rewards or advantages typically underperform dedicated RL methods. We introduce a new policy extraction method, Advantage Guidance, that retains the simple supervised learning structure with advantage conditioning while outperforming prior RL methods through effective conditional training and contrastive diffusion guidance during sampling. We show that Advantage Guidance can significantly improve a model's ability to respond to conditional advantage signals, and the contrastive guidance during sampling can use the learned advantage-conditioned policy to target the KL-regularized optimal policy. Empirically, we find that Advantage Guidance significantly outperforms prior state-of-the-art RL methods on 26 long-horizon, sparse reward tasks across datasets with varying degrees of optimality.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.