Policy Representation via Diffusion Probability Model for Reinforcement Learning

Policy representation determines which behaviors an agent can express and explore. When useful actions occupy several distinct modes, a simple unimodal policy may average between them or collapse to one. DIPO uses diffusion to preserve that multimodal structure.

Policy representation and exploration

Unlike a conventional actor that directly parameterizes one action distribution, a diffusion policy generates actions through a learned reverse process. This allows the policy to represent several separated high-value action modes without forcing them into one Gaussian.

DIPO forward diffusion process from the paper

DIPO reverse diffusion process from the paper

DIPO Figure 1: forward and reverse stochastic processes of a diffusion policy

The four-goal experiment makes the difference visible: the diffusion policy preserves all four valid directions, while Gaussian and deterministic baselines lose part of the multimodal structure.

From fitting to improvement

DIPO repeats a simple online loop: collect experience, train a critic, then move replay actions toward higher-value regions with the critic’s action gradient.

Collect

The diffusion policy continually gathers fresh transitions, preserving the diversity of behavior it explores in replay.

Move actions toward value

The critic improves stored actions directly along the action gradient, changing the data instead of forcing a restrictive parametric policy update.

Learn the improved distribution

The diffusion policy refits the shifted actions as a rich empirical distribution, retaining separated modes that a single Gaussian would collapse.

The critic decides how experience should improve; diffusion determines how faithfully the improved behavior can be represented.

DIPO Figure 2: standard model-free online reinforcement-learning framework

DIPO Figure 3: online reinforcement learning with action gradients and a diffusion policy

The diffusion model is trained on these improved actions. The critic supplies policy improvement; diffusion supplies an expressive, multimodal policy representation.

DIPO action-gradient update, Equation 15

Experimental evidence

DIPO is compared with SAC, TD3 and PPO on five MuJoCo tasks. Each curve averages five runs over one million training iterations, using the same DIPO hyperparameters across tasks.

DIPO Figure 6: average performance on five MuJoCo environments

Scope. The evidence is limited to simulated continuous-control tasks; critic quality and reverse-sampling cost remain practical constraints.

DIPO’s lasting contribution is a clean division of labor: the critic identifies how actions should improve, while the diffusion process determines how richly the improved policy can be represented. That division makes a diffusion policy usable in model-free online reinforcement learning without reducing it to a conventional unimodal actor.