PaperlayerView exact revision on arXiv

Source-linked AI summary

Efficient Online Reinforcement Learning for Diffusion Policy

Haitong Ma, Tianyi Chen, Kai Wang, Na Li, Bo Dai

arXiv:2502.00361v4cs.LG

TL;DR

Diffusion policies are expressive but conventional training requires unavailable optimal-policy samples in online RL. The paper introduces Reweighted Score Matching and derives DPMD and SDAC for tractable policy optimization. Across MuJoCo benchmarks, the proposed methods outperform recent diffusion-policy online RLs on most tasks, with SAC improvements exceeding 100% on Humanoid and Ant settings.

  • Problem

    Online diffusion-policy RL lacks samples from the optimal policy, while existing work incurs biased estimates or substantial memory and computation costs.

  • Method

    Reweighted Score Matching reweights denoising score matching to train diffusion policies from energy functions, yielding DPMD and SDAC for policy mirror descent and max-entropy optimization.

  • Results

    The proposed algorithms outperform recent diffusion-policy online RLs on most MuJoCo tasks, while both improve more than 100% over SAC on Humanoid and DPMD does so on Ant.

  • Takeaways & Limitations

    RSM enables diffusion-policy training in online RL using Q-functions as energy functions without sampling from the data distribution.

  • Takeaways & Limitations

    The approach assumes the SDAC sampling distribution has full support over the action space.

Abstract

from arXiv · show

Diffusion policies have achieved superior performance in imitation learning and offline reinforcement learning (RL) due to their rich expressiveness. However, the conventional diffusion training procedure requires samples from target distribution, which is impossible in online RL since we cannot sample from the optimal policy. Backpropagating policy gradient through the diffusion process incurs huge computational costs and instability, thus being expensive and not scalable. To enable efficient training of diffusion policies in online RL, we generalize the conventional denoising score matching by reweighting the loss function. The resulting Reweighted Score Matching (RSM) preserves the optimal solution and low computational cost of denoising score matching, while eliminating the need to sample from the target distribution and allowing learning to optimize value functions. We introduce two tractable reweighted loss functions to solve two commonly used policy optimization problems, policy mirror descent and max-entropy policy, resulting in two practical algorithms named Diffusion Policy Mirror Descent (DPMD) and Soft Diffusion Actor-Critic (SDAC). We conducted comprehensive comparisons on MuJoCo benchmarks. The empirical results show that the proposed algorithms outperform recent diffusion-policy online RLs on most tasks, and the DPMD improves more than 120% over soft actor-critic on Humanoid and Ant.

1. Introduction

Diffusion policies offer expressive, multimodal representations for imitation and offline RL, but online RL cannot sample from the optimal policy required by conventional training. The paper introduces RSM and derives DPMD and SDAC, which outperform recent diffusion-based online RL baselines across most MuJoCo tasks.

  • Diffusion policies provide expressive and multimodal policy representations that have improved imitation learning and offline RL.Their advantages follow from diffusion models’ ability to represent complex probability distributions.
  • Online RL cannot use conventional denoising score matching because samples from the optimal policy are unavailable.Existing methods addressing this issue suffer from biased estimates or high memory and computation costs.
  • Reweighted Score Matching generalizes denoising score matching by reweighting its loss while preserving the optimal noise-perturbed EBM solution.The method is motivated by viewing diffusion models as noise-perturbed energy-based models.
  • Different RSM reweighting functions yield DPMD for policy mirror descent and SDAC for max-entropy policy optimization.These algorithms make both theoretically common problems computationally tractable in continuous action spaces.
  • The proposed algorithms outperform recent diffusion-based online RL baselines in most MuJoCo tasks.The evaluation also reports improvements exceeding 100% over SAC on Humanoid for both algorithms and on Ant for DPMD.

2. Preliminaries

The preliminaries formulate the RL objectives, energy-based policies, and diffusion processes underlying the paper. They explain why unknown normalization constants and Gaussian policy projections limit expressive energy-based policy optimization, while diffusion training remains tractable when data samples are available.

  • Reinforcement Learning: An MDP is specified by state and action spaces, rewards, transitions, an initial-state distribution, and a discount factor.The paper considers policy mirror descent and max-entropy policy optimization in this setting.
  • Reinforcement Learning: Policy mirror descent uses a KL-divergence proximal term to keep updated policies approximately within a trust region.Its closed-form update depends on the current policy and state-action values.
  • Reinforcement Learning: Max-entropy RL optimizes entropy-regularized expected return, with entropy weighted by the regularization coefficient λ.Soft policy iteration alternates soft policy evaluation and soft policy improvement.
  • Energy-Based Models: The closed-form solutions of policy mirror descent and max-entropy policy optimization are energy-based models with unknown normalization constants.These constants complicate training and sampling, motivating score-based methods.
  • Energy-Based Models: Gaussian projection, as used by SAC, loses policy expressiveness and can produce suboptimal policies relative to the target max-entropy policy.This creates a gap between the theory and practice of energy-based policies.
  • Denoising Diffusion Probabilistic Models: DDPMs perturb data toward noise through a Gaussian forward process and reconstruct the data distribution through a Gaussian reverse process.The forward process uses a variance schedule β1, ..., βT, and perturbed samples can be generated directly.
  • Denoising Diffusion Probabilistic Models: Denoising score matching trains a score network to match the forward-process score using sampled Gaussian perturbations.After training, samples are generated by iterative reverse diffusion with Gaussian noise.

3. Reweighted Score Matching: A General Loss Family for Diffusion Models

The paper frames diffusion models as noise-perturbed energy-based models and introduces Reweighted Score Matching (RSM) to train diffusion policies without target-policy samples. RSM preserves the score-matching optimum while addressing the sampling and computational barriers of online RL.

  • Diffusion Models as Noise-Perturbed Energy-Based Models: Diffusion models can represent energy-based policies because their reverse process progressively approaches the original target distribution as noise decreases.The score network matches the noise-perturbed policy score, and the perturbed energy-based models become increasingly similar to the noiseless policy during reverse sampling.
  • Online RL Challenges: Conventional denoising score matching is tractable with samples from the target policy, but online RL lacks such samples.The target policies in policy mirror descent and max-entropy RL are unavailable for direct sampling during online optimization.
  • Online RL Challenges: Backpropagating policy gradients through the reverse diffusion process introduces large computational and memory costs, instability, and limited formulation coverage.The policy-parameterization approach is additionally restricted to the max-entropy formulation.
  • Reweighted Score Matching: RSM provides a general loss family intended to eliminate the sampling and computational difficulties of diffusion-policy training in online RL.Its construction is presented as the core contribution for learning diffusion policies from energy functions rather than target-distribution samples.
  • Reweighted Score Matching: RSM reweights the variational score-matching loss using any strictly positive function while retaining the optimum of matching the noise-perturbed score.The conventional loss is recovered as a special case, and the reweighting creates flexibility for tractable loss design.

4. Diffusion Policy Optimization using Reweighted Score Matching

The paper converts RSM into tractable objectives for policy mirror descent and max-entropy policy optimization, yielding practical diffusion-policy algorithms. These objectives avoid target-policy sampling while retaining denoising-score-matching-like computation.

  • Tractable Policy Objectives: Two reweighting functions convert RSM into tractable losses for policy mirror descent and max-entropy policy optimization.The resulting algorithms are designed for continuous-action settings where both theoretical objectives are empirically challenging.
  • Soft Diffusion Actor-Critic: For max-entropy policy optimization, the method uses a chosen full-support sampling distribution and a reverse sampling trick to avoid sampling from the target policy.The trick replaces expectations involving the target distribution with expectations over the selected sampling distribution and a conditional Gaussian.
  • Computational Benefits: Both reweighted losses avoid sampling from the target policy and have computation similar to denoising score matching.Sampling can instead use the current policy or a chosen distribution, avoiding additional diffusion-policy learning cost.
  • Practical Issues: Practical training addresses exploration, likelihood approximation, numerical stability, and sampling-distribution selection.The implementation adds tuned Gaussian noise, approximates policy log likelihood with additive-Gaussian likelihood, and uses normalization or logsumexp stabilization.
  • Practical Issues: The selected reverse sampling distributions showed similar empirical performance in the reported comparisons.The tested choices included uniform sampling, the previous policy, and perturbed data distributions.
  • Practical Algorithms: The proposed algorithms are Diffusion Policy Mirror Descent (DPMD) and Soft Diffusion Actor-Critic (SDAC).DPMD is presented with an iterative algorithm that collects replay data, evaluates Q-values, and updates the score-matching objective.

5. Experimental Results

Experiments evaluate RSM-based diffusion-policy training on a toy Gaussian mixture and MuJoCo control tasks. The proposed methods recover distributions, achieve strong returns, use modest computation, and show sensitivity to diffusion settings.

  • 5.1. Toy Example: The toy Gaussian-mixture experiment tests RSM using energy access without direct samples from the target mixture.The mixture has modes [3, 3] and [−3, −3] with mixing coefficients 0.8 and 0.2.
  • 5.1. Toy Example: RSM-like models and denoising score matching approximately recover the true Gaussian-mixture distribution, whereas 20-step Langevin dynamics misses the mixing coefficient.The RSM-like models use Gaussian or uniform sampling distributions and access the true energy function.
  • 5.2.2. Experimental Results: The proposed algorithms outperform all baselines across the OpenAI Gym MuJoCo environments.Baselines include recent diffusion-policy online RL methods and PPO, TD3, and SAC.
  • 5.2.2. Experimental Results: 36.0%, 41.7%, 127.3%, and 143.5% performance improvements over SAC occur on HalfCheetah, Walker2d, Ant, and Humanoid, respectively.These figures correspond to the top-performing algorithm variant on each listed locomotion task.
  • 5.2.2. Experimental Results: DPMD has stable, consistently good performance across tasks, while several comparison methods perform badly on specific environments.Reported failures include QSM on InvertedDoublePendulum, QVPO on Reacher and Pusher, and DACER on InvertedPendulum.
  • 5.2.2. Experimental Results: The proposed methods use low memory and fast computation, with costs comparable to QSM despite retaining the denoising process.The comparison uses GPU memory allocation and total computation time over 200K iterations and 1 million environment interactions.
  • 5.2.2. Experimental Results: 10 and 20 diffusion steps outperform the 30-step setting, while linear and cosine schedules outperform the variance-preserving schedule.The selected configuration is 20 steps with a cosine schedule, and results indicate robustness to diffusion-process hyperparameters.

6. Conclusion

The paper concludes that RSM enables efficient diffusion-policy training for online RL by using energy or Q-function information instead of target-distribution samples. Its empirical results show superior performance relative to SAC and recent diffusion-policy online RL methods, while stability and exploration remain future directions.

  • 6. Conclusion: RSM trains diffusion policies from energy functions and bypasses sampling from the data distribution.The approach treats diffusion models as noise-perturbed energy-based models.
  • 6. Conclusion: The proposed approach shows superior performance compared with SAC and recent diffusion-policy online RL methods.This is the paper’s reported empirical conclusion across its evaluations.
  • 6. Conclusion: Future directions include improving diffusion-policy stability and designing efficient exploration schemes.These directions are explicitly identified by the authors.
  • 6. Conclusion: Diffusion policies are used to represent stochastic policies, plan trajectories, and model transition dynamics because of their expressive and multimodal distributions.The conclusion’s broader context identifies diffusion policies as the paper’s focus.

B.1. Derivations of Proposition 3.1

The derivation establishes that the standard DDPM objective learns the score of the noise-perturbed policy distribution. It does so by showing that denoising score matching differs from the corresponding score-matching objective only by a parameter-independent constant.

  • B.1. Derivations of Proposition 3.1: The noise-perturbed policy distribution is formed by integrating the policy over the perturbation kernel q_t|0.The kernel is Gaussian with mean √ᾱ_t a_0 and covariance (1−ᾱ_t)I.
  • B.1. Derivations of Proposition 3.1: The derivation uses Tweedie’s identity to relate the conditional score to the score of the noise-perturbed distribution.The resulting relation is then matched by the score network and averaged over the perturbed distribution.
  • B.1. Derivations of Proposition 3.1: The score-matching objective and denoising score-matching objective differ only by a constant independent of θ.Therefore, they have the same optimal parameter θ* for the DDPM loss.

B.2. Derivations of Section 4.1

Section 4.1 shows that the score network can match the noise-perturbed policy score without sampling from the original policy distribution. This removes the sampling requirement of conventional denoising score matching.

  • B.2. Derivations of Section 4.1: The score network matches the noise-perturbed policy score without sampling from p0.This is the key derivational role identified for the method in Section 4.1.

B.2.1. DERIVATIONS OF DIFFUSION POLICY MIRROR DESCENT

The derivation defines the policy-mirror-descent reweighting function using the old policy and Q-values, then establishes a tractable loss through Proposition 3.1.

  • The reweighting function is defined from πold(a|s) and exp(Q(s,a)/λ), with normalization involving ZMD(s).
  • The resulting reweighted loss LgMD(θ; s, t) is tractable through a derivation using Proposition 3.1.
  • The DPMD loss LDPMD(θ, s, t) is evaluated using sampling-based approximation.

B.2.2. DERIVATIONS OF SOFT DIFFUSION ACTOR-CRITIC

The SDAC derivation selects a full-support sampling distribution, uses a reverse-sampling construction with matched scores, and derives tractable policy and value updates.

  • The softmax reweighting function combines a full-support sampling distribution ht(at|s), normalization Z(s), and the perturbed policy density pt(at|s).
  • A conditional Gaussian reverse-sampling distribution ϕ0|t is introduced to support the derivation.
  • The reverse-sampling construction matches the score functions of qt|0(at|a0) and ϕ0|t(a0|at).
  • The weighting function in the soft policy derivation is changed using the reverse-sampling trick.
  • DPMD and SDAC learn Q-functions by minimizing Bellman residuals, with SDAC following soft policy evaluation.
  • In practice, SDAC samples actions, selects the highest-Q action, adds Gaussian noise, and uses the additive-noise log probability as log π(a|s).

C.3. Training Setups for the Toy Example

The toy-example setup specifies Gaussian and uniform sampling configurations and compares RSM with diffusion and normalizing-flow Boltzmann samplers on mixture and Two Moon distributions.

  • Training setup: The Gaussian-mixture toy example uses a density specified through a Gaussian mixture model and trains SDAC with the listed diffusion-noise and optimization inputs.
  • Training setup: SDAC alternates environment interaction, replay-buffer sampling, soft policy evaluation, reverse-diffusion action sampling, and policy improvement.
  • Training setup: Gaussian sampling uses ˜p(xt) = N(0, 4I), while uniform sampling uses a distribution over [−6, 6] on both dimensions.
  • Training setup: The toy-example hyperparameters are reported in Table 3 and an additional hyperparameter table is provided separately.
  • Two Moon comparison: The additional Two Moon example uses a known energy function and compares RSM with DDPM, iDEM, and FAB.
  • Two Moon comparison: RSM, DDPM, and iDEM recover separate modes, whereas FAB shows connections between the two modes.
Loading 2502.00361v4…