ARXIV:2305.13122 · 2023

DIPO

Policy Representation via Diffusion Probability Model for Reinforcement Learning

Long Yang · Zhixiong Huang · Fenghao Lei · Yucun Zhong · Yiming Yang · Cong Fang · Shiting Wen · Binbin Zhou · Zhouchen Lin

Paper Code Explore

A policy should preserve every action worth taking

Useful behavior is often multimodal. When several distinct actions can succeed, a unimodal policy may average between them or collapse to one.

DIPO represents the policy as a diffusion process, preserving separated action modes and turning richer representation into broader exploration.

Policy representation comparison between DIPO, SAC, TD3 and PPO in a multimodal environment

Improve the data, then learn its distribution

A critic moves replay actions toward higher-value regions through the action gradient. The diffusion policy then learns the improved empirical distribution instead of compressing it into a single Gaussian.

The division is deliberate: the critic determines how behavior should improve; diffusion determines how faithfully that behavior can be represented.

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

DIPO Figure 2: standard model-free online reinforcement learning

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

Evidence across continuous control

Across five MuJoCo tasks, DIPO reaches the highest reported score while using one shared hyperparameter setting. The multimodal experiment shows why: the diffusion policy retains all four valid directions where conventional policies lose structure.

5 / 5

tasks with the highest reported score

training iterations per task

independent runs per comparison

DIPO Figure 6: average performance on five MuJoCo environments

The reported evidence is simulation-based; critic quality and reverse-sampling cost remain the central practical constraints.

Expressive policies matter when they change what an agent can discover—not only what it can imitate.

Read the paper

Abstract and citation

Popular reinforcement learning algorithms tend to produce unimodal policy distributions, limiting the expressiveness of complex policies and weakening exploration. This work establishes a theoretical foundation for policy representation with diffusion probability models, characterizes a diffusion policy as a stochastic process, and provides a convergence guarantee. It then introduces DIPO, which improves replay actions with an action gradient to make diffusion policies practical for model-free online reinforcement learning. Experiments on continuous-control MuJoCo benchmarks validate the approach.

@article{yang2023policy,
  title={Policy Representation via Diffusion Probability Model for Reinforcement Learning},
  author={Yang, Long and Huang, Zhixiong and Lei, Fenghao and Zhong, Yucun and Yang, Yiming and Fang, Cong and Wen, Shiting and Zhou, Binbin and Lin, Zhouchen},
  journal={arXiv preprint arXiv:2305.13122},
  year={2023}
}