PaperlayerView exact revision on arXiv

Source-linked AI summary

Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning

Zhendong Wang, Jonathan J Hunt, Mingyuan Zhou

arXiv:2208.06193v3cs.LGstat.ML

TL;DR

Offline RL must learn from static data while avoiding inaccurate value estimates for out-of-distribution actions, and existing regularization can be limited by weak policy expressiveness. The paper introduces Diffusion-QL, which uses a conditional diffusion policy combining behavior cloning with Q-value guidance. It reports state-of-the-art performance on the majority of D4RL tasks and illustrates the approach on a multimodal 2D bandit.

  • Problem

    Offline RL must learn effective policies from static data despite inaccurate value estimates for out-of-distribution actions and limited expressiveness in prior regularized policy classes.

  • Method

    Diffusion-QL uses a conditional diffusion model as the policy, combining behavior-cloning regularization with Q-value guidance toward high-value actions.

  • Results

    Diffusion-QL outperforms prior methods on the majority of D4RL tasks and demonstrates superiority on a simple 2D bandit with multimodal behavior.

  • Takeaways & Limitations

    Expressive diffusion policies couple behavior matching and policy improvement while capturing complex behavior distributions in offline RL.

  • Takeaways & Limitations

    Reverse diffusion requires iterative network evaluations, making inference four to five times slower than previous one-step feedforward policies and potentially limiting deployment where fast response is necessary.

Abstract

from arXiv · show

Offline reinforcement learning (RL), which aims to learn an optimal policy using a previously collected static dataset, is an important paradigm of RL. Standard RL methods often perform poorly in this regime due to the function approximation errors on out-of-distribution actions. While a variety of regularization methods have been proposed to mitigate this issue, they are often constrained by policy classes with limited expressiveness that can lead to highly suboptimal solutions. In this paper, we propose representing the policy as a diffusion model, a recent class of highly-expressive deep generative models. We introduce Diffusion Q-learning (Diffusion-QL) that utilizes a conditional diffusion model to represent the policy. In our approach, we learn an action-value function and we add a term maximizing action-values into the training loss of the conditional diffusion model, which results in a loss that seeks optimal actions that are near the behavior policy. We show the expressiveness of the diffusion model-based policy, and the coupling of the behavior cloning and policy improvement under the diffusion model both contribute to the outstanding performance of Diffusion-QL. We illustrate the superiority of our method compared to prior works in a simple 2D bandit example with a multimodal behavior policy. We then show that our method can achieve state-of-the-art performance on the majority of the D4RL benchmark tasks.

1 INTRODUCTION

Offline RL learns policies from static data, but standard improvement can overvalue unseen actions. The paper argues that limited policy expressiveness and improper regularization further restrict policy improvement, motivating Diffusion-QL.

  • Offline RL learns effective policies from previously collected data without environment interaction, avoiding risky or costly real-world exploration.
  • Standard policy improvement on offline data can prefer out-of-distribution actions whose values are inaccurately estimated.
  • Prior offline RL methods regularize policy deviation, constrain values, use pessimistic models, or treat offline RL as return-guided sequence prediction.
  • Policy-regularized methods can underperform because limited policy classes misrepresent behavior policies and restrict exploration to suboptimal action regions.
  • Diffusion-QL represents the policy with a conditional diffusion model combining behavior cloning with Q-value guidance toward high-value actions.
  • The method uses expressive diffusion policies for precise regularization and reports outperforming prior methods on the majority of D4RL tasks.

2 PRELIMINARIES AND RELATED WORK

The paper reviews offline RL, diffusion-model foundations, and behavior-regularized methods. Diffusion-QL applies conditional diffusion directly to actions, using states as conditions and Q-value guidance during training.

  • Offline RL: Offline RL learns a policy entirely from a static dataset collected by a behavior policy rather than interacting with the environment.
  • Diffusion Model: Diffusion models add noise through a forward chain and learn a reverse chain that generates data by iterative denoising.
  • Diffusion Model: Conditional diffusion models generate outputs by conditioning each reverse transition on an additional variable such as the state.
  • Policy Regularization: Prior policy-regularization methods use behavior cloning through CVAE, MMD, maximum likelihood, divergence penalties, or advantage-weighted cloning.
  • Diffusion Models in RL: Diffuser generates state-action trajectories with a return model, whereas Diffusion-QL samples one action at a time with states as conditions.

3 DIFFUSION Q-LEARNING

Diffusion Q-learning represents the policy with a conditional diffusion model and combines behavior cloning with Q-value guidance. The method trains diffusion and critic networks so sampled actions remain behavior-consistent while favoring higher-value actions.

  • 3.1 DIFFUSION POLICY: The policy is represented by the reverse process of a conditional diffusion model, whose final sample a_0 is used as the evaluated action.Sampling starts from Gaussian noise and iteratively applies learned reverse transitions.
  • 3.1 DIFFUSION POLICY: The diffusion loss is a behavior-cloning objective that learns the behavior policy while permitting skewed and multimodal action distributions.The loss samples a diffusion timestep and trains a noise-prediction model on offline data.
  • 3.2 Q-LEARNING: The final policy objective linearly combines behavior-policy regularization with Q-value maximization to preferentially sample high-value actions.Q-learning guidance is injected during training rather than relying on behavior cloning alone.
  • 3.2 Q-LEARNING: Gradients of the Q-value function with respect to actions are backpropagated through the entire reparameterized diffusion chain.The action a_0 is reparameterized through the reverse process.
  • 3.2 Q-LEARNING: The Q-value function uses two critic networks and target networks, optimized with a conventional Bellman objective and the double Q-learning trick.The implementation also updates target networks and alternates critic and policy learning.
  • 3.2 Q-LEARNING: The two losses work together to achieve the best performance, while Q-value scaling normalizes their relative contribution across datasets.The coefficient α is normalized using the dataset average absolute Q-value and controlled by hyperparameter η.

4 POLICY REGULARIZATION

The experiments examine why prior policy-regularization methods struggle with multimodal behavior policies and compare them with diffusion-based regularization. Diffusion-QL captures the behavior modes and, with Q-learning guidance, reaches the high-reward region while trading diffusion expressiveness against computational cost.

  • Prior methods: Prior methods model behavior policies with Gaussian or CVAE-based classes, then regularize policy improvement toward the learned behavior approximation.BC-MLE uses a diagonal Gaussian, while BCQ and BEAR use CVAE-based behavior modeling.
  • Bandit experiment: The 2D bandit dataset contains 10,000 actions drawn from an equal mixture of four narrow Gaussian modes, creating a strongly multimodal behavior policy.The four centers lie around the axes at (0.0, 0.8), (0.8, 0.0), (0.0, −0.8), and (−0.8, 0.0).
  • Behavior cloning: In behavior cloning, the diffusion model captures all four density modes, whereas BC-MLE, CVAE, and Tanh-Gaussian MMD methods exhibit mode-covering, mode-missing, or boundary-alignment failures.The CVAE sometimes misses modes across random seeds, and its density can fill low-density regions between modes.
  • Policy improvement: With policy improvement, Diffusion-QL converges to the optimal bottom corner while TD3+BC and BEAR-MMD can remain concentrated near suboptimal regions.The reported explanation combines expressive behavior recovery with Q-learning guidance toward high-value actions.
  • Diffusion steps: Increasing diffusion steps improves distributional fit but raises computational cost; with Q-learning, moderately small N can still produce near-optimal policies.N = 50 accurately recovers the data distribution, while N = 5 performs well on D4RL datasets.

5 EXPERIMENTS

The experiments evaluate Diffusion-QL across D4RL domains, hyperparameter settings, baselines, and component ablations. Results attribute its performance to expressive diffusion policies combined with Q-learning guidance.

  • Experimental setup: Diffusion-QL is evaluated on Gym, AntMaze, Adroit, and Kitchen tasks against policy-regularization and other offline RL baselines.The evaluation includes domain-specific analyses and an ablation study of the diffusion policy and Q-learning components.
  • Hyperparameter study: N = 5 balances performance and computational cost on the evaluated D4RL tasks.Increasing N makes convergence faster and performance more stable; with N = 5, training time is similar to CQL.
  • D4RL results: Diffusion-QL outperforms prior methods by a clear margin on Gym, especially on ‘medium’ and ‘medium-replay’ tasks.The diffusion policy can represent complicated behavior distributions before policy improvement guides actions toward optimal explored actions.
  • D4RL results: Diffusion-QL outperforms prior methods by a clear margin on harder AntMaze tasks such as ‘large-diverse’.The authors attribute this to stable, effective Q-learning guidance during conditional diffusion-policy training.
  • D4RL results: Diffusion-QL beats other baselines on Adroit and also performs well on Kitchen, where long-term value optimization is required.Adroit benefits from strong policy regularization because human demonstrations are narrow.
  • Ablation study: BC-Diffusion outperforms BC-CVAE, while Diffusion-QL further improves performance by combining soft diffusion regularization with additive Q-learning.The ablation concludes that expressive policy modeling and Q-learning guidance work together to produce good performance.

6 CONCLUSION

Diffusion-QL combines a highly expressive conditional diffusion policy with Q-learning guidance during training. The authors report state-of-the-art performance across the D4RL tasks and attribute it to these two components.

  • Conclusion: Diffusion-QL uses a conditional diffusion policy whose reverse chain provides an expressive policy class and strong policy regularization.The policy is conditioned on states and learned through the reverse diffusion process.
  • Conclusion: Q-learning guidance from a jointly learned Q-value function steers denoising toward the optimal region within the policy’s exploration area.The guidance is injected during diffusion-policy learning.
  • Conclusion: The two components contribute to Diffusion-QL’s state-of-the-art performance across all D4RL tasks.

A MORE TOY EXPERIMENTS

A 2D bandit experiment tests diffusion policies under a strongly multimodal behavior distribution concentrated in the action-space corners. Diffusion-QL recovers the modes and converges to the optimal corner where prior methods remain suboptimal.

  • Experiment setup: The offline bandit dataset contains 10,000 actions sampled equally from four narrow Gaussian modes at the corners of a 2D action space.Each mode has centers at combinations of ±0.8 and standard deviations (0.05, 0.05).
  • Behavior cloning: Only the diffusion model recovers the original multimodal data distribution in behavior cloning.CVAE captures only two diagonal modes, while MMD places density near the boundaries.
  • Q-learning experiment: Diffusion-QL converges to the optimal bottom corner, whereas BCQ and BEAR become suboptimal and TD3+BC fails to concentrate on the right corner.The ablation results are reported as consistent with the main paper’s conclusion.

B IMPLEMENTATION DETAILS

The implementation uses MLP-based conditional diffusion and Q networks with matching three-layer architectures. Both networks are optimized with Adam.

  • Diffusion policy: The diffusion policy is a residual, MLP-based conditional DDPM with state and timestep conditioning.The network has three layers, Mish activations, and 256 hidden units.
  • Diffusion policy: The diffusion model input concatenates the last-step action, current state, and sinusoidal timestep embedding.
  • Q networks: The two Q networks use the same three-layer MLP configuration with Mish activations and 256 hidden units.
  • Optimization: Adam is used to train both the diffusion policy and Q networks.

C EXPERIMENTAL DETAILS

The experiments specify training, model-selection, reward-processing, and hyperparameter procedures for Diffusion-QL across task domains.

  • Training and evaluation: Diffusion-QL trains for 2000 epochs on Gym tasks and 1000 epochs on other tasks, with 1000 gradient steps per epoch.Gym returns average six models and ten trajectories per mode; other tasks average six models and 100 trajectories.
  • Reward processing: The experiments use original MuJoCo Gym and Kitchen rewards, standardized Adroit rewards, and CQL-style reward modifications for AntMaze.
  • Model selection: Offline model selection monitors the behavior-cloning loss Ld, stops when it increases, and selects the second- or third-lowest-Ld checkpoint without environment access.The authors report that always selecting the second-smallest checkpoint has little performance impact; Gym averages 87.6 with this choice.
  • Hyperparameters: Diffusion-QL tunes learning rate, Q-learning weight η, and whether to use CQL’s max Q backup, with domain-dependent η values.Policy learning rates are selected from {1 × 10^-3, 3 × 10^-4, 3 × 10^-5}, while Q-networks use 3 × 10^-4.
  • Hyperparameters: Table 3 reports the hyperparameter settings selected for all evaluated tasks.

F OPTIMAL RESULTS

Online model selection during evaluation further improves Diffusion-QL performance compared with selecting models offline.

  • Model selection: A small amount of online evaluation experience enables selecting the best training checkpoints and provides a further performance boost.This comparison is reported in Table 4.

G LIMITATIONS AND FUTURE WORK

Diffusion policies capture multimodal action distributions but require iterative reverse sampling, creating an inference-speed limitation that motivates sampling acceleration and distillation.

  • Limitations: Diffusion policies capture multimodal distributions effectively, but reverse sampling requires repeatedly computing the denoising network.Diffusion-QL achieves good performance with N = 5, yet inference can be four to five times slower than one-step feedforward policies.
  • Future work: Future work could improve Diffusion-QL’s sampling efficiency through accelerated diffusion sampling or distillation into a simpler feedforward policy.

H GAUSSIAN MIXTURE POLICY

The Gaussian mixture baseline can model multimodal behavior during cloning, but remains suboptimal during Q-learning and performs poorly on selected D4RL tasks.

  • Method: TD3+BC-GM replaces TD3+BC’s deterministic actor with a Gaussian mixture density network and trains it by maximizing mixture likelihood.
  • D4RL evaluation: Table 5 presents the performance comparison for the Gaussian mixture ablation study.
  • Bandit evaluation: With a suitable mixture count, Gaussian mixtures capture multimodal behavior in the behavior-cloning experiment but fail to converge to the optimal target during Q-learning.The Q-learning policy retains density on suboptimal modes and therefore remains multimodal.
  • D4RL evaluation: TD3+BC-GM performs poorly on the three selected D4RL tasks when configured with three mixture components.This agrees with the observed tendency toward suboptimal actions.
  • Limitations: Gaussian mixture models are difficult to fit in high-dimensional spaces, require choosing the number of components, and often use diagonal Gaussians that limit cross-dimensional dependence modeling.
Loading 2208.06193v3…