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.
New policy extraction method
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.
Improvements
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.
Empirical results
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.