Efficient Online Reinforcement Learning for Diffusion Policy
This work addresses a key bottleneck in scaling diffusion policies to online RL for robotics and control applications, offering a more efficient and stable training approach.
The paper tackles the inefficiency and instability of training diffusion policies in online reinforcement learning by introducing Reweighted Score Matching (RSM), which eliminates the need to sample from the optimal policy and reduces computational costs. The resulting algorithms, DPMD and SDAC, outperform recent diffusion-policy online RL methods on MuJoCo benchmarks, with DPMD improving over soft actor-critic by more than 120% on Humanoid and Ant tasks.
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.