PaperlayerView exact revision on arXiv

Source-linked AI summary

Diffusion Guidance Is a Controllable Policy Improvement Operator

Kevin Frans, Seohong Park, Pieter Abbeel, Sergey Levine

arXiv:2505.23458v1cs.LG

TL;DR

RL must improve beyond potentially suboptimal offline data, but scalable methods often sacrifice policy improvement or stability. CFGRL derives a diffusion-guidance formulation that retains supervised-learning simplicity while improving policies, including without explicit value-function learning. Experiments report consistent gains across offline RL and goal-conditioned behavioral cloning, with guidance controlling improvement and success rates sometimes doubling.

  • Problem

    RL methods must improve beyond exploratory or suboptimal offline data, while behavioral cloning is limited by the data and iterative RL can be unstable and difficult to scale.

  • Method

    CFGRL represents policies as reference-policy and optimality factors, then uses diffusion and classifier-free guidance to sample controllably improved policies.

  • Results

    CFGRL consistently improves offline policy extraction and goal-conditioned behavioral cloning, with guidance-weight increases associated with stronger policies until divergence and success rates sometimes doubling.

  • Takeaways & Limitations

    CFGRL extends simple supervised methods, including goal-conditioned behavioral cloning, toward further policy improvement without necessarily learning a value function explicitly.

  • Takeaways & Limitations

    CFGRL does not replace full RL procedures: it assumes a given value function and does not prescribe how that value function is trained.

Abstract

from arXiv · show

At the core of reinforcement learning is the idea of learning beyond the performance in the data. However, scaling such systems has proven notoriously tricky. In contrast, techniques from generative modeling have proven remarkably scalable and are simple to train. In this work, we combine these strengths, by deriving a direct relation between policy improvement and guidance of diffusion models. The resulting framework, CFGRL, is trained with the simplicity of supervised learning, yet can further improve on the policies in the data. On offline RL tasks, we observe a reliable trend -- increased guidance weighting leads to increased performance. Of particular importance, CFGRL can operate without explicitly learning a value function, allowing us to generalize simple supervised methods (e.g., goal-conditioned behavioral cloning) to further prioritize optimality, gaining performance for "free" across the board.

1 Introduction

CFGRL connects diffusion guidance with policy improvement, combining supervised-learning simplicity with controllable improvement beyond data behaviors. It improves offline policy extraction and goal-conditioned behavioral cloning, including settings without explicit value-function learning.

  • Motivation: The work targets the tension between scalable generative modeling and RL’s ability to optimize beyond exploratory or suboptimal offline data.Behavioral cloning is simple and scalable but limited by the data, whereas iterative RL can be unstable and sensitive to hyperparameters.
  • Core framework: CFGRL defines policies as a prior reference policy multiplied by an optimality distribution, proving improvement when that distribution increases monotonically with advantage.This establishes the theoretical link between policy improvement and the product-policy construction.
  • Core framework: Diffusion techniques sample the product distribution, while classifier-free guidance combines prior and conditional factors to control policy improvement at test time without retraining.The guidance weight provides a controllable policy-improvement mechanism during sampling.
  • Core framework: CFGRL connects generative-model guidance to RL objectives, with guided sampling equivalent under certain choices to a KL-constrained policy-improvement solution.This frames guidance as a principled alternative to conventional RL optimization procedures.
  • Applications: CFGRL can avoid explicit value-function learning and serve as a drop-in replacement for goal-conditioned behavioral cloning without additional training requirements.This extends simple supervised methods toward further policy improvement.
  • Applications: CFGRL consistently improves offline policy extraction and goal-conditioned behavioral cloning, with goal-conditioned success rates sometimes increasing by a factor of two.The reported applications include value-based offline extraction and value-free generalization of goal-conditioned behavioral cloning across state-based, visual, and hierarchical settings.

2 Related work

Related work uses diffusion and flow models to enhance RL policies, but CFGRL instead frames policy extraction as guided generative modeling with test-time control over regularization and improvement.

  • Offline RL: Offline RL seeks reward-maximizing policies from fixed datasets while avoiding erroneous extrapolation beyond the data distribution.Prior approaches include value-learning and policy-extraction directions, especially weighted regression and return-conditioned behavioral cloning.
  • Offline RL: CFGRL differs from related supervised offline-RL methods by expressing the regularization–policy-improvement tradeoff through diffusion guidance controllable at test time.This preserves a simple supervised objective while adding a sampling-time control mechanism.
  • Diffusion and flow policies for RL: Prior diffusion and flow policies address the challenge of extracting policies that maximize a learned Q-function using weighted regression, reparameterized gradients, and related strategies.These methods exploit iterative generative-model expressivity but focus on how to extract improved policies.

3 Preliminaries

The preliminaries define policy improvement as non-decreasing RL performance and describe practical approximations based on reference-policy samples and divergence control. They motivate a generative-model strategy that retains supervised-learning simplicity while enabling controllable improvement.

  • RL setting: The RL setting models an agent policy as a state-conditioned action distribution that, with the environment, generates state-action trajectories.The standard objective maximizes expected discounted future rewards along those trajectories.
  • Policy improvement: A policy improvement operator updates a reference policy ˆπ to a policy π whose RL objective does not decrease: J(ˆπ) ≤ J(π).Advantage compares action values with state values under the reference policy and formalizes improvement conditions.
  • Policy improvement: Practical algorithms approximate improvement using samples from the reference policy’s state distribution rather than the resulting policy’s occupancy distribution.This approximation is introduced because algorithms operate on samples from a previous reference policy.
  • Divergence control: Bounding divergence between reference and resulting policies can provide an objective bound and enable monotonic incremental improvement.Trust-region methods and KL penalties are standard ways to control this divergence.
  • Motivation for the proposed method: Policy-gradient optimization can require on-policy samples and have high variance, motivating supervised-learning-like methods such as weighted regression.The paper positions its generative-model policy-improvement strategy as retaining supervised-learning simplicity while adding controllability.

4 Diffusion guidance is a controllable policy improvement operator

CFGRL represents policies as a reference policy combined with an optimality factor, then uses diffusion guidance to sample controllably improved policies. Increasing guidance can improve expected return, while trading off adherence to the reference and risking distribution shift.

  • Practical scope: CFGRL can avoid explicit value-function learning and serve as a drop-in replacement for goal-conditioned behavioral cloning.The framework is also presented as an offline policy extraction method when a value function is learned.
  • Product policies: CFGRL parameterizes policies as products of a reference policy and an optimality function, enabling policy improvement over the reference under monotonicity conditions.The optimality function is conditional on advantage and must be non-negative and monotonically increasing.
  • Controllable improvement: Increasing the exponent on the optimality factor provably improves the resulting policy, but also moves it farther from the reference policy.This creates a return-versus-reference-adherence tradeoff and can cause empirical performance to suffer under distribution shift.
  • Diffusion guidance: Diffusion guidance combines unconditional and optimality-conditioned policy factors, with guidance weight w producing the attenuated product π(a | s) ∝ ˆπ(a | s)p(o | s, a)^w.The two factors are represented with a single conditional diffusion or flow-matching network.
  • Test-time control: CFGRL adjusts the degree of policy improvement during sampling, allowing test-time guidance sweeps over one network without retraining.This separates the reference and optimality-conditioned policies until sampling, unlike objectives where the tradeoff is fixed during training.

5 CFGRL improves over weighted policy extraction in offline RL

CFGRL addresses weaknesses of weighted policy extraction by training with an evenly weighted conditional diffusion loss and tuning guidance at evaluation time. Experiments compare it with AWR across ExORL and OGBench, with guidance continuing to improve beyond the temperature range where AWR saturates.

  • Weighted policy extraction: AWR can concentrate training gradients on a few high-weight state-action pairs, causing the rest of each batch to contribute little effective learning signal.The peaked weights arise from exponentiated advantages and can make weighted regression inefficient.
  • CFGRL training: CFGRL labels samples by advantage sign and trains with a standard conditional diffusion loss without weighting terms.Its within-batch gradients therefore remain more evenly distributed than those of AWR.
  • Hyperparameter control: Guidance weight and AWR temperature control the tradeoff between reference-policy adherence and reward maximization, but CFGRL adjusts this tradeoff without retraining.CFGRL keeps the prior and optimality-conditioned policies separate until sampling.
  • Experimental comparison: Across 9 ExORL tasks and 9 OGBench single-task environments, CFGRL is experimentally compared with AWR using the same learned state-action value function.The shared value function makes the downstream policy-extraction comparison independent of the extraction method.

6 CFGRL unlocks hidden gains in goal-conditioned behavioral cloning

CFGRL extends goal-conditioned behavioral cloning by applying controllable guidance, improving beyond GCBC without requiring an explicit value function. Across OGBench tasks, performance generally improves with guidance strength and CFGRL outperforms the evaluated baselines.

  • GCBC does not converge to the optimal goal-reaching policy when the dataset is suboptimal, motivating further policy improvement.
  • CFGRL applies guidance with w ≥ 1 to reweight the GCBC policy, while standard GCBC corresponds implicitly to w = 1.The resulting policy is designed to improve over the base GCBC policy.
  • CFGRL improves goal-conditioned policies without training an explicit value function or sampling additional on-policy actions.Its guidance combines the original goal-conditioned BC policy with an unconditional BC policy.
  • The evaluation covers 17 state-based and 7 pixel-based OGBench goal-conditioned tasks spanning navigation, manipulation, humanoid control, and puzzle solving.
  • CFGRL consistently outperforms four imitation-learning baselines across most evaluated state-based and pixel-based tasks using a fixed guidance strength of w = 3.The baselines include BC, flow BC, GCBC, and flow GCBC.
  • Performance generally increases as guidance strength rises, although stronger guidance can eventually decline when the policy deviates too far from the data distribution.

7 Discussion and Conclusion

CFGRL connects diffusion guidance with policy improvement and applies that connection to offline policy extraction and goal-conditioned behavioral cloning. It improves over AWR and GCBC while retaining the simplicity of supervised methods, but it is not presented as a replacement for full RL.

  • CFGRL improves over AWR in offline RL and substantially improves over GCBC in goal-conditioned settings while maintaining their simplicity.
  • CFGRL assumes a given value function when used for offline policy extraction and does not prescribe how that value function should be trained.
  • More advanced policy-extraction methods and online RL techniques could provide stronger extrapolation than CFGRL.
  • The method offers test-time control over policy improvement and can operate without an explicit Q-function.
  • In practice, guidance weights can be swept without retraining the model, supporting a plug-and-play use of CFGRL.

A Theoretical results

The theoretical results show that reweighting a reference policy by a nonnegative, nondecreasing function of advantage yields policy improvement. Increasing the guidance exponent preserves this ordering, linking controllable guidance strength to progressively stronger theoretical improvement.

  • The proof uses Chebyshev’s sum inequality for probability measures together with the stochastic-policy policy-improvement theorem.
  • Reweighting a policy by a bounded, measurable, nonnegative, nondecreasing function of its advantage produces a policy improvement.
  • The argument normalizes the reweighted policy and applies monotonicity of the weighting function to establish the required expected-value condition.
  • For exponents 0 ≤ w1 ≤ w2, the policy using w2 is ordered by the theorem relative to the policy using w1 under the same advantage-based reweighting.
  • The theoretical section is accompanied by OGBench environments and a CFG-weight ablation whose reported performance generally improves as w increases.

B Additional results

Additional experiments examine OGBench environments, CFG-weight sensitivity, and optimality-conditioning architectures. Shared optimality-conditioned policies generally outperform separate policies and extrapolate better, while increasing guidance usually helps until an apparent divergence point.

  • Environments: The additional evaluations use OGBench tasks spanning state-based and visual settings.
  • Ablation study on the CFG weight w: Across all 17 state-based OGBench tasks, performance generally improves as CFG weight increases but can decline beyond a point.The decline is attributed to the policy deviating too far from the data distribution.
  • Ablation study on optimality conditioning: Shared optimality-conditioned networks generally perform and extrapolate better than separate networks.The authors suggest shared representations may explain this result.

C Implementation details

CFGRL experiments use ExORL and OGBench tasks, with an oracle-representation variant for isolating goal-representation effects. Baselines follow original implementations where possible, and hyperparameters are documented in Tables 4–8.

  • Implementation: Each experiment takes no more than 4 hours on a single A5000 GPU, and implementations are provided in the supplemental material.
  • Tasks: Experiments cover 9 ExORL tasks and 9 single-task OGBench variants, using RND datasets for ExORL.
  • Tasks: The OGBench oracle-representation variant supplies ground-truth goal representations to remove confounding from goal representation learning.
  • Methods and hyperparameters: Baselines generally use their original implementations and hyperparameters, while GCBC samples goals uniformly from future states.
  • Methods and hyperparameters: Full hyperparameters are listed in Tables 4–8 for ExORL, OGBench, and GCBC experiments.
Loading 2505.23458v1…