DiffPPO: PPO with a Diffusion Action Prior

Research code for PPO-DAP, a method for improving exploration in continuous-control reinforcement learning.

PPO-DAP combines Proximal Policy Optimization (PPO) with a diffusion model that suggests actions for a given state. The model first learns from recorded trajectories. During online training, a critic scores its suggestions, which provide a small additional learning signal for the policy. PPO's probability ratios, advantage estimates, and critic updates use only fresh environment interactions.

This repository implements the method described in our paper. The Python package is named ppo_dap.

Research overview

The paper evaluates PPO-DAP on eight MuJoCo tasks with an online budget of one million environment steps per task, following offline pretraining. It reports improved early learning and final returns that match or exceed the strongest on-policy baselines on six of the eight tasks. See the paper for the full comparisons and protocol.

How it works

1. Learn action suggestions. Train a state-conditioned diffusion model on recorded trajectories.
2. Collect new experience. Run the current policy in the environment and compute PPO's training quantities.
3. Guide the policy. Generate candidate actions at the visited states, guide them using critic estimates, and use them in a small auxiliary policy loss.
4. Adapt the prior. Update a small subset of diffusion-model parameters using the new experience, keeping its main network fixed.

Generated actions are used only in the auxiliary policy terms. The algorithm guide explains the data flow and update order.

Install and check

Install Python 3.12.3 and uv 0.12.0 first. The release uses a locked CPU PyTorch environment.

git clone https://github.com/TianciGao/DiffPPO.git
cd DiffPPO
git checkout v0.1.0
uv sync --frozen --all-groups
uv run python -c "import ppo_dap; print(ppo_dap.__file__)"
uv run pytest

Where to start

Goal	Page
Understand the method	Algorithm guide
Find the relevant code	Implementation guide

Versions and branches

Version or branch	Purpose
v0.1.0	Fixed software release for installation and version-specific comparisons.
main	Main documentation, algorithm library, and initial experiment configuration tools.
experiment/paper-v6-e1	Experiment interfaces, evaluation, and reporting tools under development.
release/cleanroom-v0.1.0	Branch used to prepare the first release; use the release tag for its original files.

Citation

If you use this work, please cite:

@article{gao2024ppodap,
  title   = {Enhancing Sample Efficiency and Exploration in Reinforcement Learning through the Integration of Diffusion Models and Proximal Policy Optimization},
  author  = {Gao, Tianci and Neusypin, Konstantin A. and Dmitriev, Dmitry D. and Yang, Bo and Rao, Shengren},
  journal = {arXiv preprint arXiv:2409.01427},
  year    = {2024},
  doi     = {10.48550/arXiv.2409.01427},
  url     = {https://arxiv.org/abs/2409.01427v6}
}

See CITATION.cff for citation metadata. When reporting software results, also record the tag or commit you used.

License

MIT.