Diffusion Policy through Conditional Proximal Policy Optimization

2026-03-05

AI Summary

Diffusion policies in reinforcement learning have been limited by the computational inefficiency of evaluating action log-likelihoods. This work proposes a novel training paradigm that aligns policy iteration with the diffusion process, enabling on-policy training of diffusion policies by evaluating only simple Gaussian probabilities—thereby circumventing the need to simulate the full denoising trajectory required by conventional approaches and substantially reducing computational and memory overhead. The method naturally supports entropy regularization and achieves stable training through conditional proximal policy optimization. Evaluated across multiple benchmark tasks in IsaacLab and MuJoCo Playground, the proposed approach not only attains superior performance but also generates multimodal policy behaviors.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Imitation Learning & Inverse Reinforcement Learning
• Search and Optimization: Learning to Search

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• User Modeling, Personalization and Recommendation: User modeling and simulation for interactive and conversational systems
• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information

Abstract

Reinforcement learning (RL) has been extensively employed in a wide range of decision-making problems, such as games and robotics. Recently, diffusion policies have shown strong potential in modeling multi-modal behaviors, enabling more diverse and flexible action generation compared to the conventional Gaussian policy. Despite various attempts to combine RL with diffusion, a key challenge is the difficulty of computing action log-likelihood under the diffusion model. This greatly hinders the direct application of diffusion policies in on-policy reinforcement learning. Most existing methods calculate or approximate the log-likelihood through the entire denoising process in the diffusion model, which can be memory- and computationally inefficient. To overcome this challenge, we propose a novel and efficient method to train a diffusion policy in an on-policy setting that requires only evaluating a simple Gaussian probability. This is achieved by aligning the policy iteration with the diffusion process, which is a distinct paradigm compared to previous work. Moreover, our formulation can naturally handle entropy regularization, which is often difficult to incorporate into diffusion policies. Experiments demonstrate that the proposed method produces multimodal policy behaviors and achieves superior performance on a variety of benchmark tasks in both IsaacLab and MuJoCo Playground.

Problem

• diffusion policy
• reinforcement learning
• log-likelihood
• on-policy learning
• action generation

Innovation

• Diffusion Policy
• On-policy Reinforcement Learning
• Conditional Proximal Policy Optimization
• Entropy Regularization
• Multimodal Action Generation