# An Efficient On-Policy Deep Learning Framework for Stochastic Optimal Control

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 33 of its 55 references that are in the map.
- First posted on arXiv 7 Oct 2024.

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

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Building Normalizing Flows with Stochastic Interpolants](https://mlsanity.com/paper/2209.15571/) (2022)
- [Stochastic Interpolants: A Unifying Framework for Flows and Diffusions](https://mlsanity.com/paper/2303.08797/) (2023)
- [Solving high-dimensional partial differential equations using deep learning](https://mlsanity.com/paper/1707.02568/) (2017)
- [Directly Fine-Tuning Diffusion Models on Differentiable Rewards](https://mlsanity.com/paper/2309.17400/) (2023)
- [A survey of the Schr"odinger problem and some of its connections with optimal transport](https://mlsanity.com/paper/1308.0215/) (2013)
- [Monte Carlo Gradient Estimation in Machine Learning](https://mlsanity.com/paper/1906.10652/) (2019)
- [Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control](https://mlsanity.com/paper/2409.08861/) (2024)
- [Neural Stochastic Differential Equations: Deep Latent Gaussian Models in the Diffusion Limit](https://mlsanity.com/paper/1905.09883/) (2019)
- [Denoising Diffusion Samplers](https://mlsanity.com/paper/2302.13834/) (2023)
- [Flow-based generative models for Markov chain Monte Carlo in lattice field theory](https://mlsanity.com/paper/1904.12072/) (2019)
- [Fine-Tuning of Continuous-Time Diffusion Models as Entropy-Regularized Control](https://mlsanity.com/paper/2402.15194/) (2024)
- [Flow Annealed Importance Sampling Bootstrap](https://mlsanity.com/paper/2208.01893/) (2022)
- [On the relation between optimal transport and Schr"odinger bridges: A stochastic control viewpoint](https://mlsanity.com/paper/1412.4430/) (2014)
- [Annealed Flow Transport Monte Carlo](https://mlsanity.com/paper/2102.07501/) (2021)
- [Solving high-dimensional Hamilton-Jacobi-Bellman PDEs using neural networks: perspectives from the theory of controlled diffusions and measures on path space](https://mlsanity.com/paper/2005.05409/) (2020)
- [NETS: A Non-Equilibrium Transport Sampler](https://mlsanity.com/paper/2410.02711/) (2024)
- [Probabilistic Forecasting with Stochastic Interpolants and F"ollmer Processes](https://mlsanity.com/paper/2403.13724/) (2024)
- [Beyond ELBOs: A Large-Scale Evaluation of Variational Methods for Sampling](https://mlsanity.com/paper/2406.07423/) (2024)
- [Deep Learning Approximation for Stochastic Control Problems](https://mlsanity.com/paper/1611.07422/) (2016)
- [Stochastic Optimal Control Matching](https://mlsanity.com/paper/2312.02027/) (2023)
- [Stochastic Optimal Control for Collective Variable Free Sampling of Molecular Transition Paths](https://mlsanity.com/paper/2207.02149/) (2022)
- [Neural networks-based algorithms for stochastic control and PDEs in finance](https://mlsanity.com/paper/2101.08068/) (2021)
- [Deep neural networks algorithms for stochastic control problems on finite horizon: convergence analysis](https://mlsanity.com/paper/1812.04300/) (2018)
- [Recent Developments in Machine Learning Methods for Stochastic Control and Games](https://mlsanity.com/paper/2303.10257/) (2023)
- [A Unified Perspective on Natural Gradient Variational Inference with Gaussian Mixture Models](https://mlsanity.com/paper/2209.11533/) (2022)
- [Three algorithms for solving high-dimensional fully-coupled FBSDEs through deep learning](https://mlsanity.com/paper/1907.05327/) (2019)
- [A Taxonomy of Loss Functions for Stochastic Optimal Control](https://mlsanity.com/paper/2410.00345/) (2024)
- [Convergence Analysis of Schr"odinger-F"ollmer Sampler without Convexity](https://mlsanity.com/paper/2107.04766/) (2021)

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

- [Rare Event Analysis via Stochastic Optimal Control](https://mlsanity.com/paper/2604.13213/) (2026)
- [High-dimensional Mean-Field Games by Particle-based Flow Matching](https://mlsanity.com/paper/2512.01172/) (2025)
- [Consensus Based Stochastic Control](https://mlsanity.com/paper/2501.17801/) (2025)

## Abstract (arXiv)

We present a novel on-policy algorithm for solving stochastic optimal control (SOC) problems. By leveraging the Girsanov theorem, our method directly computes on-policy gradients of the SOC objective without expensive backpropagation through stochastic differential equations or adjoint problem solutions. This approach significantly accelerates the optimization of neural network control policies while scaling efficiently to high-dimensional problems and long time horizons. We evaluate our method on classical SOC benchmarks as well as applications to sampling from unnormalized distributions via Schr"odinger-F"ollmer processes and fine-tuning pre-trained diffusion models. Experimental results demonstrate substantial improvements in both computational speed and memory efficiency compared to existing approaches.