Efficient Online Reinforcement Learning for Diffusion Policy

TL;DR

This paper introduces Reweighted Score Matching, a novel method enabling efficient online reinforcement learning with diffusion policies, avoiding costly sampling and improving performance on MuJoCo benchmarks.

Contribution

The paper proposes Reweighted Score Matching for diffusion policies, enabling scalable online RL without sampling from the target distribution, and introduces two algorithms, DPMD and SDAC.

Findings

DPMD improves over 120% on Humanoid and Ant tasks.

Proposed algorithms outperform recent diffusion-policy online RL methods.

Reweighted Score Matching reduces computational costs and stabilizes training.

Abstract

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.