English

Guiding Diffusion Models with Reinforcement Learning for Stable Molecule Generation

Machine Learning 2025-08-25 v1 Artificial Intelligence

Abstract

Generating physically realistic 3D molecular structures remains a core challenge in molecular generative modeling. While diffusion models equipped with equivariant neural networks have made progress in capturing molecular geometries, they often struggle to produce equilibrium structures that adhere to physical principles such as force field consistency. To bridge this gap, we propose Reinforcement Learning with Physical Feedback (RLPF), a novel framework that extends Denoising Diffusion Policy Optimization to 3D molecular generation. RLPF formulates the task as a Markov decision process and applies proximal policy optimization to fine-tune equivariant diffusion models. Crucially, RLPF introduces reward functions derived from force-field evaluations, providing direct physical feedback to guide the generation toward energetically stable and physically meaningful structures. Experiments on the QM9 and GEOM-drug datasets demonstrate that RLPF significantly improves molecular stability compared to existing methods. These results highlight the value of incorporating physics-based feedback into generative modeling. The code is available at: https://github.com/ZhijianZhou/RLPF/tree/verl_diffusion.

Keywords

Cite

@article{arxiv.2508.16521,
  title  = {Guiding Diffusion Models with Reinforcement Learning for Stable Molecule Generation},
  author = {Zhijian Zhou and Junyi An and Zongkai Liu and Yunfei Shi and Xuan Zhang and Fenglei Cao and Chao Qu and Yuan Qi},
  journal= {arXiv preprint arXiv:2508.16521},
  year   = {2025}
}
R2 v1 2026-07-01T05:01:57.927Z