aiqbench analysis
This research integrates diffusion models into Maximum Entropy Reinforcement Learning by treating reverse-diffusion steps as sequential decision-making processes within a Markov Decision Process framework. The approach leverages variational inference to decompose the policy optimization objective, enabling efficient training and improved performance on continuous control tasks like robotic manipulation. By allowing agents to explore alternative successful strategies from identical states, this method offers a principled way to enhance policy expressiveness without sacrificing computational tractability.
Diffusion-Augmented Markov Decision Processes for Maximum Entropy Reinforcement Learning
arXiv:2512.02019v4 Announce Type: replace-cross Abstract: Diffusion models provide an expressive framework for sampling from complex, unnormalized distributions. In this work, we extend Maximum Entropy Reinforcement Learning (ME-RL) to diffusion-based policies by introducing Diffusion-Augmented Markov Decision Processes (DA-MDPs). DA-MDPs interpret each reverse-diffusion transition as an individual reinforcement-learning decision, while only the final denoised action is executed in the environment. Our DA-MDPs follow from a principled derivation based on the variational-inference formulation of ME-RL. By augmenting policy and target trajectories with intermediate diffusion variables, we obtain a tractable reverse-KL upper bound via the data-processing inequality. This bound decomposes across denoising transitions, yielding diffusion-augmented variants of soft rewards, value functions, and local policy objectives. This provides a general framework for adapting ME-RL algorithms to diffusion policies while differentiating through only one diffusion transition at a time. We instantiate the framework with PPO, REPPO, and a maximum-entropy extension of WPO. Experiments demonstrate improved continuous-control performance, benefits from additional diffusion steps, and memory-efficient training. On the StackCube and PushT manipulation tasks, DA-MDP methods learn alternative successful strategies from the same initial state and achieve higher success rates and generally higher success-weighted mode entropy than the Gaussian ME-RL baseline. We also demonstrate successful training when using action chunking.
Abstract:Diffusion models provide an expressive framework for sampling from complex, unnormalized distributions. In this work, we extend Maximum Entropy Reinforcement Learning (ME-RL) to diffusion-based policies by introducing Diffusion-Augmented Markov Decision Processes (DA-MDPs). DA-MDPs interpret each reverse-diffusion transition as an individual reinforcement-learning decision, while only the final denoised action is executed in the environment. Our DA-MDPs follow from a principled derivation based on the variational-inference formulation of ME-RL. By augmenting policy and target trajectories with intermediate diffusion variables, we obtain a tractable reverse-KL upper bound via the data-processing inequality. This bound decomposes across denoising transitions, yielding diffusion-augmented variants of soft rewards, value functions, and local policy objectives. This provides a general framework for adapting ME-RL algorithms to diffusion policies while differentiating through only one diffusion transition at a time. We instantiate the framework with PPO, REPPO, and a maximum-entropy extension of WPO. Experiments demonstrate improved continuous-control performance, benefits from additional diffusion steps, and memory-efficient training. On the StackCube and PushT manipulation tasks, DA-MDP methods learn alternative successful strategies from the same initial state and achieve higher success rates and generally higher success-weighted mode entropy than the Gaussian ME-RL baseline. We also demonstrate successful training when using action chunking.
71 pages. Substantially revised theory and expanded experiments; updated author list. Public JAX and PyTorch implementations are linked in the paper
Machine Learning (cs.LG); Artificial Intelligence (cs.AI); Machine Learning (stat.ML)
arXiv:2512.02019 [cs.LG]
(or arXiv:2512.02019v4 [cs.LG] for this version)
https://doi.org/10.48550/arXiv.2512.02019 Focus to learn more arXiv-issued DOI via DataCite
arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.
Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.
Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.
Full text under the publisher's open licence · 316 words
Related coverage
Reinforcement Learning for Continuous-Time Jump Markov Decision Processes with Applications to Network Dynamic Pricing
arXiv · Aug 24, 2026
Minimax Optimal Strategy for Delayed Observations in Online Reinforcement Learning
arXiv · Jun 3, 2026
Latent reasoning architectures would likely undermine CoT, our strongest oversight tool
AI Alignment Forum · Sep 23, 2026
Exploding variance of means of exponentials: least-squares to the rescue – Machine Learning Research Blog
francisbach.com · Sep 27, 2026