A Diffusion Model Framework for Maximum Entropy Reinforcement Learning

TL;DR

This paper introduces a diffusion model-based framework for maximum entropy reinforcement learning, leading to new variants of popular algorithms that improve performance and sample efficiency on continuous control tasks.

Contribution

It reinterprets MaxEntRL as a diffusion sampling problem and develops diffusion-based variants of SAC, PPO, and WPO with improved results.

Findings

DiffSAC, DiffPPO, and DiffWPO outperform their base algorithms in benchmarks.

The new methods achieve higher sample efficiency.

Minor implementation changes are needed to incorporate diffusion dynamics.

Abstract

Diffusion models have achieved remarkable success in data-driven learning and in sampling from complex, unnormalized target distributions. Building on this progress, we reinterpret Maximum Entropy Reinforcement Learning (MaxEntRL) as a diffusion model-based sampling problem. We tackle this problem by minimizing the reverse Kullback-Leibler (KL) divergence between the diffusion policy and the optimal policy distribution using a tractable upper bound. By applying the policy gradient theorem to this objective, we derive a modified surrogate objective for MaxEntRL that incorporates diffusion dynamics in a principled way. This leads to simple diffusion-based variants of Soft Actor-Critic (SAC), Proximal Policy Optimization (PPO) and Wasserstein Policy Optimization (WPO), termed DiffSAC, DiffPPO and DiffWPO. All of these methods require only minor implementation changes to their base algorithm. We find that on standard continuous control benchmarks, DiffSAC, DiffPPO and DiffWPO achieve better returns and higher sample efficiency than SAC and PPO.