# Stabilizing Policy Gradients for Stochastic Differential Equations via Consistency with Perturbation Process

arXiv: https://arxiv.org/abs/2403.04154 · HTML: https://mlsanity.com/paper/2403.04154/

## Key facts

- Cited by 8 papers in the map (data through 14 Oct 2026).
- Placed in RL theory and sample complexity, one of 15 directions in Deep reinforcement learning methods.
- Builds on 31 of its 45 references that are in the map.
- First posted on arXiv 7 Mar 2024.

## Placed in

- Deep reinforcement learning methods / [RL theory and sample complexity](https://mlsanity.com/deeprl/rl-theory-and-sample-complexity/)

## Builds on (most cited first, 30 of 31)

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [BLIP: Bootstrapping Language-Image Pre-training for Unified Vision-Language Understanding and Generation](https://mlsanity.com/paper/2201.12086/) (2022)
- [LAION-5B: An open large-scale dataset for training next generation image-text models](https://mlsanity.com/paper/2210.08402/) (2022)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Continuous control with deep reinforcement learning](https://mlsanity.com/paper/1509.02971/) (2015)
- [DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps](https://mlsanity.com/paper/2206.00927/) (2022)
- [LAION-400M: Open Dataset of CLIP-Filtered 400 Million Image-Text Pairs](https://mlsanity.com/paper/2111.02114/) (2021)
- [Diffusion Models: A Comprehensive Survey of Methods and Applications](https://mlsanity.com/paper/2209.00796/) (2022)
- [Planning with Diffusion for Flexible Behavior Synthesis](https://mlsanity.com/paper/2205.09991/) (2022)
- [ImageReward: Learning and Evaluating Human Preferences for Text-to-Image Generation](https://mlsanity.com/paper/2304.05977/) (2023)
- [Training Diffusion Models with Reinforcement Learning](https://mlsanity.com/paper/2305.13301/) (2023)
- [Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning](https://mlsanity.com/paper/2208.06193/) (2022)
- [Is Conditional Generative Modeling all you need for Decision-Making?](https://mlsanity.com/paper/2211.15657/) (2022)
- [DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models](https://mlsanity.com/paper/2305.16381/) (2023)
- [IDQL: Implicit Q-Learning as an Actor-Critic Method with Diffusion Policies](https://mlsanity.com/paper/2304.10573/) (2023)
- [Structure-based Drug Design with Equivariant Diffusion Models](https://mlsanity.com/paper/2210.13695/) (2022)
- [Pocket2Mol: Efficient Molecular Sampling Based on 3D Protein Pockets](https://mlsanity.com/paper/2205.07249/) (2022)
- [Offline Reinforcement Learning via High-Fidelity Generative Behavior Modeling](https://mlsanity.com/paper/2209.14548/) (2022)
- [3D Equivariant Diffusion for Target-Aware Molecule Generation and Affinity Prediction](https://mlsanity.com/paper/2303.03543/) (2023)
- [A 3D Generative Model for Structure-Based Drug Design](https://mlsanity.com/paper/2203.10446/) (2022)
- [Gradient Estimation Using Stochastic Computation Graphs](https://mlsanity.com/paper/1506.05254/) (2015)
- [Generating 3D Molecules for Target Protein Binding](https://mlsanity.com/paper/2204.09410/) (2022)
- [Equivariant Energy-Guided SDE for Inverse Molecular Design](https://mlsanity.com/paper/2209.15408/) (2022)

## Built on by (most cited first, 8 of 8)

- [Antigen-Specific Antibody Design via Direct Energy-based Preference Optimization](https://mlsanity.com/paper/2403.16576/) (2024)
- [Pareto-Optimal Energy Alignment for Designing Nature-Like Antibodies](https://mlsanity.com/paper/2412.20984/) (2024)
- [Empower Structure-Based Molecule Optimization with Gradient Guided Bayesian Flow Networks](https://mlsanity.com/paper/2411.13280/) (2024)
- [Diffusion Fine-Tuning via Reparameterized Policy Gradient of the Soft Q-Function](https://mlsanity.com/paper/2512.04559/) (2025)
- [Fine-tuning Flow Matching Generative Models with Intermediate Feedback](https://mlsanity.com/paper/2510.18072/) (2025)
- [Integrating Protein Dynamics into Structure-Based Drug Design via Full-Atom Stochastic Flows](https://mlsanity.com/paper/2503.03989/) (2025)
- [Designing Cyclic Peptides via Harmonic SDE with Atom-Bond Modeling](https://mlsanity.com/paper/2505.21452/) (2025)
- [Policy Gradient for Continuous-Time Robust Markov Decision Processes](https://mlsanity.com/paper/2606.04335/) (2026)

## Abstract (arXiv)

Considering generating samples with high rewards, we focus on optimizing deep neural networks parameterized stochastic differential equations (SDEs), the advanced generative models with high expressiveness, with policy gradient, the leading algorithm in reinforcement learning. Nevertheless, when applying policy gradients to SDEs, since the policy gradient is estimated on a finite set of trajectories, it can be ill-defined, and the policy behavior in data-scarce regions may be uncontrolled. This challenge compromises the stability of policy gradients and negatively impacts sample complexity. To address these issues, we propose constraining the SDE to be consistent with its associated perturbation process. Since the perturbation process covers the entire space and is easy to sample, we can mitigate the aforementioned problems. Our framework offers a general approach allowing for a versatile selection of policy gradient methods to effectively and efficiently train SDEs. We evaluate our algorithm on the task of structure-based drug design and optimize the binding affinity of generated ligand molecules. Our method achieves the best Vina score -9.07 on the CrossDocked2020 dataset.