PaperlayerView exact revision on arXiv

Source-linked AI summary

Diffusion-based Reinforcement Learning via Q-weighted Variational Policy Optimization

Shutong Ding, Ke Hu, Zhenhao Zhang, Kan Ren, Weinan Zhang, Jingyi Yu, Jingya Wang, Ye Shi

arXiv:2405.16173v3cs.LG

TL;DR

Online RL has seen less use of diffusion policies, whose variational objective cannot be directly optimized there because good action samples are unavailable. QVPO combines a Q-weighted variational bound with entropy regularization and an efficient behavior policy. On MuJoCo locomotion benchmarks, it is reported to achieve state-of-the-art cumulative reward and sample efficiency.

  • Problem

    Diffusion policies are less investigated in online RL, where unavailable good action samples make direct optimization of their variational objective difficult.

    Source: §1 · p. 2, §4.1 · p. 4

  • Method

    QVPO uses a Q-weighted VLO loss with Q-weight transformations, adds entropy regularization, and selects actions for an efficient behavior policy.

    Source: §1 · p. 2

  • Results

    QVPO achieves state-of-the-art cumulative reward and sample efficiency on MuJoCo locomotion benchmarks compared with traditional and diffusion-based online RL algorithms.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The authors conclude that QVPO leverages diffusion policies' exploration and multimodality to prevent convergence to a sub-optimal policy.

    Source: §1 · p. 2

    The entropy regularizer cannot adaptively adjust its term or include diffusion-policy entropy in the TD target for soft policy iteration.

    Source: §6 · p. 10

Abstract

from arXiv · show

Diffusion models have garnered widespread attention in Reinforcement Learning (RL) for their powerful expressiveness and multimodality. It has been verified that utilizing diffusion policies can significantly improve the performance of RL algorithms in continuous control tasks by overcoming the limitations of unimodal policies, such as Gaussian policies, and providing the agent with enhanced exploration capabilities. However, existing works mainly focus on the application of diffusion policies in offline RL, while their incorporation into online RL is less investigated. The training objective of the diffusion model, known as the variational lower bound, cannot be optimized directly in online RL due to the unavailability of 'good' actions. This leads to difficulties in conducting diffusion policy improvement. To overcome this, we propose a novel model-free diffusion-based online RL algorithm, Q-weighted Variational Policy Optimization (QVPO). Specifically, we introduce the Q-weighted variational loss, which can be proved to be a tight lower bound of the policy objective in online RL under certain conditions. To fulfill these conditions, the Q-weight transformation functions are introduced for general scenarios. Additionally, to further enhance the exploration capability of the diffusion policy, we design a special entropy regularization term. We also develop an efficient behavior policy to enhance sample efficiency by reducing the variance of the diffusion policy during online interactions. Consequently, the QVPO algorithm leverages the exploration capabilities and multimodality of diffusion policies, preventing the RL agent from converging to a sub-optimal policy. To verify the effectiveness of QVPO, we conduct comprehensive experiments on MuJoCo benchmarks. The final results demonstrate that QVPO achieves state-of-the-art performance on both cumulative reward and sample efficiency.

1 Introduction

Diffusion policies have been successful in imitation and offline RL, but their use in online RL is less explored because policy improvement is difficult. QVPO introduces a Q-weighted variational objective with exploration and behavior-policy mechanisms, and is reported to achieve state-of-the-art cumulative reward and sample efficiency on MuJoCo locomotion benchmarks.

1 Introduction

  • Diffusion policies are less investigated in online RL, where changing policies affect value estimation and the diffusion training objective does not directly improve the policy.Prior approaches described include gradient updates that constrain exploration and score/Q-gradient alignment affected by inaccurate Q-network gradients.Source: §1 · p. 2
  • The Q-weighted VLO loss is proven to be a tight lower bound of the online RL policy objective and is extended to general tasks with equivalent Q-weight transformations.The bound holds under certain conditions.Source: §1 · p. 2
  • QVPO adds entropy regularization to enhance diffusion-policy exploration and uses action selection to reduce the impact of large policy variance on sample efficiency.The paper notes that diffusion-policy exploration declines with limited diffusion steps.Source: §1 · p. 2
  • QVPO achieves state-of-the-art sample efficiency and episodic reward on MuJoCo locomotion benchmarks against traditional and diffusion-based online RL algorithms.The paper reports this comparison across its evaluated benchmarks.Source: §1 · p. 2

2 Related Works

The related work groups diffusion methods across imitation learning, planning, offline RL, and online RL. It describes how their use varies by task and contrasts QVPO's objective with earlier online approaches.

Diffusion Policies for Imitation Learning

  • For imitation learning, diffusion policies fit expert-data distributions and generate robot action sequences conditioned on visuomotor observations; BESO addresses goal-conditioned imitation when rewards are sparse or inaccessible.The section frames imitation learning as behavior learning from expert datasets without explicit reward.Source: §2 · pp. 2–3

Diffusion Planners

  • Diffusion planners construct complete trajectories to reduce compounding errors, with prior methods using reward guidance, classifier-free guidance, latent spaces, or task hierarchies.Decision Diffuser derives actions from generated state sequences using inverse dynamics.Source: §2 · p. 3

Diffusion Policies for Offline RL

  • Offline RL methods use diffusion for Q-learning, state reconstruction, faster action inference, energy-guided sampling, consistency policies, or weighted regression.The section characterizes offline RL as learning an optimal policy from suboptimal offline datasets.Source: §2 · p. 3

Diffusion Policies for Online RL

  • Online methods DIPO and QSM update diffusion policies using action gradients or score/Q-gradient alignment, while QVPO uses a theoretically tight Q-weighted VLO bound with entropy regularization and an efficient behavior policy.The section identifies gradient inaccuracies and alignment bias as concerns for QSM.Source: §2 · p. 3

3 Preliminaries

The preliminaries introduce the MDP and policy-optimization concepts used to describe reinforcement learning. They also outline DDPM's forward-noising and reverse-denoising processes and its variational lower-bound training objective.

3.1 Reinforcement Learning

  • An MDP specifies state and action spaces, transition and reward functions, an initial-state distribution, and a discount factor; RL seeks a policy maximizing discounted expected reward through policy improvement.The advantage function is defined as the action-state value minus the state value.Source: §3.1 · pp. 3–4

3.2 Denoising Diffusion Probabilistic Models

  • DDPM adds Gaussian noise over T forward steps toward a Gaussian prior, then learns reverse transitions to denoise samples; its training objective maximizes the variational lower bound.The passage identifies β_t as the variance schedule for the forward process.Source: §3.2 · p. 4

4 Q-weighted Variational Policy Optimization

QVPO derives a Q-weighted variational objective for online diffusion-policy improvement, then adds mechanisms for exploration and lower-variance online interaction.

4.1 Q-weighted Variational Objective for Diffusion Policy

  • Online diffusion improvement lacks optimal training samples and rare high-Q samples; Q-weighted VLO tightly lower-bounds the RL policy objective when Q(s, a) ≥ 0, with equality at convergence.The optimal-sample requirement makes direct use of the diffusion variational bound difficult in online RL.Source: §4.1 · pp. 4–5

4.2 Equivalent Transformation for Q-weighted VLO Loss

  • Equivalent positive-weight transformations address negative Q values and seek high-quality synthesized samples; qadv replaces Q with advantage and performs better than qcut when most Q values are negative.The advantage is A(s, a) = Q(s, a) − V(s).Source: §4.2 · pp. 5–6

4.3 Enhancing Policy Exploration via Diffusion Entropy Regularization

  • Because diffusion action log-likelihood is inaccessible, uniform-distribution samples train an entropy regularizer; larger ω_ent improves exploration while reducing exploitation, as illustrated by a three-peak bandit.The regularizer's coefficient balances exploration and exploitation.Source: §4.3 · p. 7, Figure 2 · p. 7

4.4 Reducing Diffusion Policy Variance via Action Selection

  • Selecting high-Q diffusion samples for behavior interactions improves sample efficiency; use smaller target-policy K_t < K_b because larger selection can cause serious target-Q overestimation.The behavior policy uses action selection to reduce the diffusion policy's interaction variance.Source: §4.4 · pp. 7–8

5 Experiments

QVPO is evaluated against traditional and diffusion-based online RL methods on MuJoCo locomotion tasks, then ablated for entropy regularization and action selection. The reported comparisons show stronger overall performance and identify settings associated with exploration, sample efficiency, and training stability.

5.1 Comparative Evaluation

  • QVPO achieves state-of-the-art performance and converges faster than six online RL baselines across five MuJoCo locomotion tasks.The comparison spans five runs and 1e6 training epochs; Figure 3 plots episodic reward against training epochs.Source: §5.1 · p. 8, Figure 3 · p. 9

5 Experiments

  • The experiments compare QVPO with six off-policy, on-policy, and diffusion-based algorithms, using an off-policy implementation with a SAC-style critic.The critic uses doubled Q networks and the minimum for policy and critic updates.Source: §5.1 · p. 8, §5 · p. 8

5.2 Ablation Study and Parameter Analysis

  • With entropy regularization, QVPO using 20 diffusion steps achieves performance close to the 100-step version; without it, 100 steps perform much better.The authors use Ant-v3 as an example because they report similar final results across locomotion tasks.Source: §5.2 · p. 9, Figure 4 · p. 9
  • Kb = 4, Kt = 1 outperforms Kb = 1, Kt = 1, while Kt < Kb is associated with better stability and reduced target-Q overestimation.A too-high action-selection number limits exploration, and the authors report that Kb = 4 balances exploration and exploitation.Source: §5.2 · pp. 9–10, Figure 5 · p. 9

6 Conclusion

The conclusion presents QVPO as a diffusion-based online RL method and reports state-of-the-art MuJoCo results in cumulative reward and sample efficiency. It also identifies limitations in the entropy treatment and outlines adaptive adjustment and soft policy iteration as future directions.

6 Conclusion

  • QVPO combines a Q-weighted VLO loss, Q-weight transformations, entropy regularization, and an efficient behavior policy for diffusion-based online RL.The loss is described as a tight lower bound of the online RL policy objective under certain conditions; the behavior policy reduces variance.Source: §6 · p. 10
  • On MuJoCo continuous-control benchmarks, QVPO achieves state-of-the-art cumulative reward and sample efficiency against traditional and diffusion-based online RL methods.Source: §6 · p. 10
  • The entropy regularizer cannot adaptively adjust entropy or incorporate diffusion-policy entropy into the TD target for SAC-style soft policy iteration.Future work is proposed on adaptive entropy adjustment and integrating entropy into the TD target.Source: §6 · p. 10

A.1 Proof of Theorem 1

The appendix establishes a lower-bound relation for the online RL policy objective when Q values are nonnegative, with equality when the policy converges. It then derives the Q-weighted VLO loss.

A.1 Proof of Theorem 1

  • When Q(s, a) ≥ 0 for every state-action pair, the Q-weighted variational bound is a tight lower bound of the online RL policy objective.The equality condition is policy convergence.Source: §A.1 · p. 13
  • The proof relates the policy objective to a Q-weighted KL divergence between the diffusion trajectory distributions.The passage identifies Q(s, a) as the state-action Q value and πθ as the policy being optimized.Source: §A.1 · p. 13

A.2 Deriviation of Q-weighted VLO Loss

  • The appendix derives the Q-weighted VLO loss from diffusion-policy terms.Source: §A.2 · p. 14

A.3 Proof of Theorem 2

The proof characterizes the optimal policy in a single policy-improvement step across cases determined by the signs of the Q values. For mixed-sign values, the stated solution assigns weight to positive-Q actions.

A.3 Proof of Theorem 2

  • When all actions have positive Q values, the one-step policy-improvement problem is convex and is solved using KKT conditions.Source: §A.3 · pp. 14–15
  • When Q(s, a) ≤ 0 for every action, the optimal policy is supported on actions with maximal Q value.The passage describes the solution as a sum of Dirac functions over the maximizing actions.Source: §A.3 · p. 15
  • When some actions have positive Q values and others do not, the optimal policy assigns weight according to π(a | s)Q(s, a) on positive-Q actions and zero elsewhere.The proof defines π(a | s) as the behavior policy and identifies the positive-Q indicator.Source: §A.3 · pp. 15–16

A.4 Proof for Convergence of QVPO

The convergence argument assumes that diffusion training increases the probability of samples in proportion to their weights. The accompanying hyperparameter section reports that DIPO uses more diffusion steps than QVPO.

A.4 Proof for Convergence of QVPO

  • The proof assumes training-sample generation probability improves proportionally to weight near the action maximizing Qπold(s, a) for each state.The passage describes this action’s neighborhood as having a small radius.Source: §A.4 · p. 16
  • The proof examines RL-objective improvement, expands its first term using the Bellman equation, then repeats the operation.Source: §A.4 · p. 16

B Hyperparameters

C Training and Inference Time

On Ant-v3, QVPO’s multiple-sampling-and-selection procedure is reported to add little time relative to DIPO’s gradient optimization. Its 6ms inference time is described as acceptable for 50-Hz control, although QVPO is slower than SAC.

C Training and Inference Time

  • The timing comparison is not fair across QSM and the other algorithms because QSM uses JAX while the others use PyTorch; all diffusion methods use T=20.The paper states JAX implementations are 6–10 times faster than PyTorch implementations.Source: C Training and Inference Time · p. 17
  • QVPO’s multiple-sampling-and-selection procedure adds little time compared with DIPO’s gradient optimization while using GPU parallelism in the Ant-v3 timing comparison.Source: Table 4 · p. 17, C Training and Inference Time · p. 17
  • 6ms inference time is described as acceptable for 50-Hz control, though QVPO is much slower at inference than SAC.Source: C Training and Inference Time · p. 18
Loading 2405.16173v3…