Source-linked AI summary

Reinforcement Learning via Advantage Guidance

Zhiyuan Zhou, Sergey Levine

arXiv:2610.19783v2cs.LG

TL;DR

RL could improve on demonstrated behavior, but practical instability motivates the search for a stable supervised-learning alternative; naïve advantage conditioning has often underperformed dedicated RL and may not elicit requested advantages. Advantage Guidance combines conditional training with contrastive diffusion guidance during sampling. It improves advantage responsiveness and outperforms prior strong RL algorithms across 26 long-horizon, sparse-reward tasks.

  • Problem

    RL is unstable and hyperparameter-sensitive in practice, while naïve advantage-conditioned supervised policies can underperform dedicated RL and fail to respond to requested advantages.

    Source: §1 · pp. 1–2
  • Method

    Advantage Guidance trains a supervised policy conditioned on continuous advantage estimates and uses contrastive diffusion guidance during sampling, with residual parameterization and imagined-action training to improve conditional responsiveness.

    Source: §1 · p. 2
  • Results

    Across 26 long-horizon, sparse-reward tasks, Advantage Guidance significantly outperforms prior strong RL algorithms on play-style and demonstration datasets.

    Source: §1 · p. 2
  • Takeaways & Limitations

    Advantage Guidance improves the policy’s response to advantage inputs and can target the KL-regularized optimal policy during sampling.

    Source: §1 · p. 2

    The KL-constrained advantage-maximization result assumes β > 0, ensured when T(y+) > T(y−) and w > 0.

    Source: §6.1 · p. 6

Abstract

from arXiv · show

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.

1. Introduction

The introduction asks whether policy training can retain imitation learning’s stability while matching RL’s performance. It presents Advantage Guidance as a supervised, advantage-conditioned approach that addresses weak conditional responsiveness and outperforms prior strong RL methods across 26 tasks.

1. Introduction

  • RL can improve beyond demonstrated behavior, but practical instability and hyperparameter sensitivity contrast with imitation learning’s stable, scalable policy training.The paper asks whether training can be as simple and stable as imitation learning while performing as well as RL.Source: §1 · pp. 1–2
  • Naïve advantage-conditioned policies typically underperform dedicated RL and may fail to match requested advantages or ignore conditioning and revert to imitation.Conditioning also restricts improvement to the queried advantage level rather than consistently favoring higher-advantage actions.Source: §1 · p. 2
  • Advantage Guidance combines continuous advantage conditioning with contrastive sampling guidance, residual flow parameterization, and training on imagined actions sampled from an unconditional policy.The imagined actions and their advantage estimates broaden the action distribution used to train the conditional policy for each state.Source: §1 · p. 2
  • Across 26 long-horizon, sparse-reward tasks, Advantage Guidance outperforms prior strong RL algorithms on play-style and demonstration datasets.The method also improves responsiveness to advantage inputs and can target the KL-regularized optimal policy during sampling.Source: §1 · p. 2

2. Related Work

Related work addresses offline RL’s policy-extraction challenges through value learning, generative policies, and imitation-style objectives. Advantage Guidance differs from prior advantage-conditioning and diffusion-guidance approaches in how it uses continuous advantage information and guides sampling.

2. Related Work

  • Offline RL work addresses value-function extrapolation [3] [19] [27] [29] [54] and policy extraction [5] [16] [20] [34] [39] [47] [53], though policy extraction can be unstable and hyperparameter-sensitive.Diffusion and flow-matching policies have also been studied for RL [43] [11] [56] [8] [26] [7] [58].Source: §2 · p. 2
  • Imitation-style RL methods use advantage, reward, or return conditioning [28] [6] [13] [45], data filtering [6] [52] [9], or loss weighting [39].Advantage conditioning uses the full dataset and has lower variance than return conditioning, while avoiding loss-weighting instability.Source: §2 · p. 2
  • Prior advantage-conditioned methods query a fixed high advantage or learn an extra model to select an in-support query, whereas Advantage Guidance uses contrastive diffusion guidance without either step.The paper also introduces residual parameterization and action-imagination training to improve responsiveness to advantage inputs.Source: §2 · p. 2
  • Diffusion policy methods may require costly, unstable backpropagation through denoising, expand the environment horizon, or approximate the multi-step policy with limited expressivity.Unlike CFGRL’s binary filtering of positive-advantage data, Advantage Guidance trains on the full dataset with continuous advantage values; the paper reports that CFGRL usually underperforms because filtering discards action-quality information.Source: §2 · p. 3

3. Preliminaries and Background

The preliminaries define advantage and the KL-regularized target policy, then describe advantage-conditioned policies and classifier-free guidance. They explain why a fixed queried advantage may be a poor policy-improvement target.

3. Preliminaries and Background

  • Advantage Aπ(s, a) = Qπ(s, a) − Vπ(s) measures how much better action a is than following policy π at state s.The value functions characterize expected future return under a policy.Source: §3 · p. 3
  • The KL regularizer constrains policy change toward a reference policy to prevent exploiting Q-values, and its fixed-state optimal policy tilts that reference toward high-advantage actions.The paper uses this tilted distribution as its policy-extraction target.Source: §3 · p. 3
  • Advantage-conditioned policies are trained to predict actions from state and action-advantage estimates, then queried with a target advantage; an excessively high query can fall outside the training distribution and yield bad actions.A fixed query favors actions at that advantage while discouraging actions with still higher advantage.Source: §3 · p. 3
  • Classifier-free guidance modifies the sampling score using conditional and unconditional scores, tilting generated samples toward a specified label; in RL, binary labels can support policy improvement under stated monotonicity conditions.The paper describes flow matching as the generative model used for simplicity.Source: §3 · p. 3
  • Figure 2 reports that naïve advantage-conditioned policies generate actions whose advantages do not align with requested values across five tasks per domain.Source: Figure 2 · p. 4

4. Analysis of CFG for Advantage Conditioned Policies

This section analyzes whether classifier-free guidance on an advantage-conditioned policy produces the KL-regularized optimal policy. Under a Gaussian advantage-label likelihood, the analysis finds that ordinary guidance generally favors matching the queried advantage rather than always favoring higher-advantage actions.

4. Analysis of CFG for Advantage Conditioned Policies

  • The section compares CFG sampling on an advantage-conditioned policy with the score of the KL-regularized optimal policy to test whether CFG reaches that target.Source: §4 · p. 4
  • The analysis derives the guidance score for an advantage-conditioned policy under the simplifying assumption that the conditional label likelihood is Gaussian around the action’s advantage.The conditional and unconditional policy scores are denoted g_y and g, respectively.Source: §4 · p. 4
  • Ordinary classifier-free guidance generally does not produce the desired KL-regularized target: it favors matching the queried advantage and penalizes deviations both above and below it.Increasing the guidance weight strengthens this preference for matching the target advantage.Source: §4 · p. 4

5. Advantage Alignment Problem of Naïve Conditioned Policies

Naïve advantage-conditioned policies often fail to respond to requested advantages, limiting their ability to improve on the behavior policy. Advantage Guidance responds more strongly to advantage inputs.

5. Advantage Alignment Problem of Naïve Conditioned Policies

  • Advantage alignment measures whether an action’s achieved advantage matches the advantage used to query the conditional policy.Naïve policies often produce similar advantages for different queries, while Advantage Guidance is more responsive.Source: §5 · pp. 4–5
  • For naïve policies, gradients under ground-truth and random advantage labels have cosine similarity close to 1, suggesting the loss is insensitive to those labels.Changing the queried advantage alters naïve actions by less than half as much as changing initial sampling noise; Advantage Guidance’s response exceeds the noise effect.Source: §5 · p. 5
  • Across 20 OGBench tasks, naïve advantage-conditioned policies often ignore advantage inputs, while Advantage Guidance responds substantially better.Advantage Guidance actions exhibit achieved advantages that closely align with requested advantages.Source: Figure 3 · p. 5, §5 · p. 5
  • Naïve policies largely ignore advantage values and effectively regress to an unconditional policy.The analysis focuses on advantage-conditioned policies as the special case of classifier-free guidance with guidance weight w = 1.Source: §5 · pp. 4–5

6. Advantage Guidance

Advantage Guidance combines contrastive guidance at sampling time with training techniques intended to make an advantage-conditioned flow policy respond to its inputs. The method uses residual parameterization and imagined actions, alongside state-level advantage normalization and adaptive layer normalization.

6. Advantage Guidance

  • Advantage Guidance combines advantage-conditioned policy extraction with contrastive sampling guidance to improve advantage alignment and target the KL-constrained optimal policy.Its training techniques include residual parameterization over an unconditional flow and training on multiple imagined actions per state.Source: §6 · p. 5

6.1. Contrastive Guidance for Advantage Maximization

  • Contrastive guidance combines the unconditional policy score with a weighted difference between conditional scores for y+ and y− to match the KL-regularized optimal policy score.With positive guidance weight and T(y+) > T(y−), stronger w[T(y+) − T(y−)] yields stronger advantage maximization under a weaker constraint.Source: §6.1 · p. 6

6.2. Residual Parameterization for Conditional Policy

  • Residual parameterization adds a conditional residual flow to the unconditional reference flow, leaving the residual to focus on the conditional input.The guidance formulation assumes the conditional policy is advantage-aligned at y+ and y−, motivating training methods to improve alignment.Source: §6.2 · p. 6

6.3. Action Imagination for Counterfactual Supervision

  • Action imagination augments each dataset state-action pair with K proposed actions at the same state to provide counterfactual supervision across advantage queries.The unconditional reference policy is one convenient proposal; the authors report that residual parameterization and imagination improve alignment and performance, especially on long-horizon tasks.Source: §6.3 · p. 7

6.4. Implementation Details

  • The method uses IQL to train Q(s, a) and V(s), computes advantage as A(s, a) = Q(s, a) − V(s), and conditions the policy on normalized advantages using AdaLN.Advantages are standardized separately for each state; the authors report this is more effective than batch- or dataset-level normalization.Source: §6.4 · p. 7

7. Experiments

Across 26 long-horizon sparse-reward tasks, Advantage Guidance outperforms prior methods on both play-style and demonstration-based datasets. The experiments also test best-of-N sampling, training components that affect advantage alignment, and alternative sampling guidance.

7. Experiments

  • The evaluation uses 20 OGBench and 6 OCBench sparse-reward manipulation tasks, whose datasets differ from play-style trajectories to near-expert demonstrations.OCBench demonstrations cover narrower states and predominantly near-optimal actions; the experiments also use action chunking.Source: §7 · p. 8
  • The baselines include actor-critic, supervised-learning-based, and diffusion-guidance methods, with the same IQL value functions used for fair policy-extraction comparisons.Advantage is estimated as A(s, a) = Q(s, a) − V(s).Source: §7 · p. 8

7.1. How Does Advantage Guidance Compare with Prior Methods?

  • Advantage Guidance significantly outperforms prior methods, including actor-critic, supervised-learning, and diffusion-guidance baselines, across OGBench and OCBench.The comparison covers 20 OGBench tasks and 6 OCBench tasks.Source: Figure 4 · p. 8, §7.1 · p. 9
  • On OGBench, Advantage Guidance beats best-of-N sampling at small N but slightly trails it at large N; on OCBench, it remains significantly better even at large N.Advantage Guidance uses less test-time compute than best-of-N sampling.Source: §7.1 · p. 9, Figure 5 · p. 9

7.2. Does AG Training Improve Advantage Alignment and Policy Performance?

  • Residual parameterization and action imagination improve advantage alignment, and their combination improves it further; the increased alignment translates to better performance on both benchmarks.The alignment measures are gradient similarity between true and random gradients and the advantage-noise variance ratio.Source: §7.2 · pp. 9–10, Figure 6 · p. 9, Figure 7 · p. 9

7.3. Does contrastive guidance improve performance over CFG or no guidance?

  • Contrastive guidance performs similarly to CFG and conditional sampling on OGBench but significantly improves performance on OCBench.The comparison uses the same AG-trained policy; high queried advantages can lie outside the training distribution in demonstration-based datasets.Source: §7.3 · p. 10

8. Discussions

The discussion presents Advantage Guidance as a simple supervised-learning approach that outperforms more complex RL methods, while identifying several directions for further work. In the OCBench sampling ablation, contrastive guidance outperforms CFG and no guidance.

8. Discussions

  • On 6 OCBench tasks, contrastive guidance significantly outperforms CFG and no guidance.The comparison is an ablation of the sampling method for Advantage Guidance.Source: Figure 8 · p. 10
  • The approach uses a straightforward advantage estimate; future work could test more intricate estimates and robustness to noisy estimates.Source: §8 · p. 10
  • Too-large guidance weights usually reduce performance, motivating work on stably using larger weights.Source: §8 · p. 10
  • The authors identify offline-to-online RL as another setting for future investigation.Source: §8 · p. 10

A.1. Contrastive Guidance Theoretical Guarantee.

The result characterizes a target policy through its clean-action score, combining a reference-policy score with a weighted difference between two conditional-policy scores. For practical sampling, the formulation is extended to intermediate denoising steps as an approximation.

A.1. Contrastive Guidance Theoretical Guarantee

  • The target clean-action score combines the reference-policy score with a weighted difference between conditional-policy scores at y+ and y−.Source: §A.1 · p. 14
  • The difference between the two conditional-policy scores is proportional to the advantage gradient, scaled by T(y+) − T(y−).Source: §A.1 · p. 15
  • For sampling, the clean-action score result is extended to intermediate denoising steps as a practical approximation following classifier-free guidance.Source: §A.1 · p. 15

A.2. Additional Advantage Guidance Ablation Experiments

Advantage Guidance ablations find that per-state advantage normalization improves performance, while action imagination helps up to a point. Contrastive conditions and advantage-label standard deviations yield similar performance when the guidance weight is selected accordingly.

Action imagination

  • Performance improves as the number of imagined actions increases from small K, then plateaus at roughly K = 7.Each state uses K imagined actions plus one dataset action for training and advantage normalization.Source: §A.2 · p. 15, Figure 9 · p. 16

Advantage Normalization

  • Per-state advantage normalization outperforms batch-level and dataset-level normalization.The comparison holds with action augmentation fixed at K = 31.Source: §A.2 · p. 16, Figure 10 · p. 16

Contrastive guidance conditions and guidance weight

  • With the best guidance weight, different contrastive condition pairs perform similarly because the relevant quantity is w(y+ − y−).This makes tuning the product sufficient in practice.Source: §A.2 · p. 16, Figure 11 · p. 17

Advantage label standard deviation

  • Different advantage-label standard deviations produce roughly similar performance when the guidance weight is adjusted accordingly.Although the derivation requires σy > 0, σy = 0 also works well in practice.Source: §A.2 · p. 17, Figure 12 · p. 17

A.3. Empirical Evaluation of Naïve Advantage Conditioned Policy RCP-A

Appendix experiments find that a higher discount improves results on cube-double, while naïve advantage conditioning and a more expressive transformer do not resolve advantage alignment. The appendix also examines pure conditional sampling, baseline methods, training details, and benchmark breakdowns.

A.3. Empirical Evaluation of Naïve Advantage Conditioned Policy RCP-A

  • Querying naïve RCP-A with different advantage values has effectively no impact on performance across 20 OGBench tasks.The original RCP-A setting uses k = 1 standard deviation above the dataset mean; the experiment varies k.Source: §A.3 · p. 18, Figure 13 · p. 18

A.4. Does A More Expressive Architecture Help with Advantage Alignment?

  • A transformer variant does not significantly improve advantage alignment or overall OGBench performance, whereas Advantage Guidance responds more to advantage inputs.The authors identify residual parameterization and action imagination as critical to the method’s performance.Source: §A.4 · p. 18, Figure 14 · p. 19

A.5. Analysis of Pure Conditional Sampling with Advantage Guidance

  • Pure conditional sampling with high queried advantages causes performance declines on OGBench and sharp drops on OCBench and puzzle-3x3.The authors suggest high queried advantages may be out of distribution for demonstration-based OCBench data and that puzzle-3x3 has relatively limited state-action coverage.Source: §A.5 · p. 19, Figure 15 · p. 19

A.6. Description of Baseline Methods

A.7. Critic Training Details

  • The advantage estimate is derived from IQL as A(s, a) = Q(s, a) − V(s), using an ensemble of two Q-functions aggregated by their minimum.Hyperparameters are tuned per domain on selected OGBench tasks and per task per domain on OCBench.Source: §A.7 · pp. 21–22, §A.8 · p. 22

A.9. Tasks and Environment Details

  • OGBench tasks use 1M offline transitions, while OCBench task datasets include 10k transitions for block single and chamber easy and 30k for block double.The appendix also supplies detailed task breakdowns and numeric results with means and 95% confidence intervals across 10 seeds.Source: §A.9 · pp. 22–23, Figure 17 · p. 25, Table 5 · p. 26

A.10. Detailed Results

  • The appendix separately breaks down results by task or domain for baseline comparisons, gradient similarity, sampling, ablations, and compute cost.These figures identify the scope of the breakdowns but do not provide the plotted values in the supplied passages.Source: Figure 17 · p. 25, Figure 18 · p. 27, Figure 19 · p. 28

A.11. Additional Results

  • At γ = 0.999, all agents improve on cube-double, and QGF and Advantage Guidance significantly outperform the other baselines.All methods perform poorly on this long-horizon sparse-reward task at γ = 0.99.Source: Figure 24 · p. 30, §A.11 · p. 30
Loading 2610.19783v2…