English

Diffusion-Augmented Markov Decision Processes for Maximum Entropy Reinforcement Learning

Machine Learning 2026-05-28 v3 Artificial Intelligence Machine Learning

Abstract

Diffusion models excel at sampling from complex, unnormalized distributions. In this work, we extend Maximum Entropy Reinforcement Learning (ME-RL) to diffusion processes, enabling sampling from the optimal policy trajectory distribution. By minimizing a tractable upper bound on the reverse KL divergence between the diffusion policy and the optimal policy trajectory distributions, we derive a modified surrogate objective and introduce Diffusion-Augmented Markov Decision Processes (DA-MDPs). DA-MDPs allow for seamless integration of diffusion policies into any ME-RL method with minimal modifications. We demonstrate its effectiveness by adapting Proximal Policy Optimization (PPO), Wasserstein Policy Optimization (WPO), and Relative Entropy Pathwise Policy Optimization (REPPO) into their diffusion-based variants: DA-MDP: PPO, DA-MDP: WPO, and DA-MDP: REPPO. Empirical results on standard continuous-control benchmarks show that our approach matches or outperforms baseline methods, while experiments on multimodal benchmarks confirm its ability to model multimodal action distributions.

Keywords

Cite

@article{arxiv.2512.02019,
  title  = {Diffusion-Augmented Markov Decision Processes for Maximum Entropy Reinforcement Learning},
  author = {Sebastian Sanokowski and Kaustubh Patil},
  journal= {arXiv preprint arXiv:2512.02019},
  year   = {2026}
}

Comments

Preprint

R2 v1 2026-07-01T08:04:20.791Z