A Schrödinger Eigenfunction Method for Long-Horizon Stochastic Optimal Control
Louis Claeys, Artur Goldman, Zebang Shen, Niao He
Abstract
High-dimensional stochastic optimal control (SOC) becomes harder with longer planning horizons: existing methods scale linearly in the horizon , with performance often deteriorating exponentially. We overcome these limitations for a subclass of linearly-solvable SOC problems—those whose uncontrolled drift is the gradient of a potential. In this setting, the Hamilton-Jacobi-Bellman equation reduces to a linear PDE governed by an operator . We prove that, under the gradient drift assumption, is unitarily equivalent to a Schrödinger operator with purely discrete spectrum, allowing the long-horizon control to be efficiently described via the eigensystem of . This connection provides two key results: first, for a symmetric linear-quadratic regulator (LQR), matches the Hamiltonian of a quantum harmonic oscillator, whose closed-form eigensystem yields an analytic solution to the symmetric LQR with arbitrary terminal cost. Second, in a more general setting, we learn the eigensystem of using neural networks. We identify implicit reweighting issues with existing eigenfunction learning losses that degrade performance in control tasks, and propose a novel loss function to mitigate this. We evaluate our method on several long-horizon benchmarks, achieving an order-of-magnitude improvement in control accuracy compared to state-of-the-art methods, while reducing memory usage and runtime complexity from to .
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Builds on4
- Stochastic Optimal Control for Collective Variable Free Sampling of Molecular Transition PathsLars Holdijk, Yuanqi Du, Ferry Hooft, Priyank Jaini et al.NeurIPS 2023 · 55 citations
- Stochastic Optimal Control MatchingCarles Domingo-Enrich, Jiequn Han, Brandon Amos, Joan Bruna et al.NeurIPS 2024 · 49 citations
- The SSL Interplay: Augmentations, Inductive Bias, and GeneralizationVivien Cabannes, Bobak Toussi Kiani, Randall Balestriero, Yann LeCun et al.ICML 2023 · 43 citations
- Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal ControlCarles Domingo-Enrich, Michal Drozdzal, Brian Karrer, Ricky T. Q. ChenICLR 2025 · 2 citations
Related papers
- Deep learning for continuous-time stochastic control with jumpsPatrick Cheridito, Jean-Loup Dupret, Donatien HainautNeurIPS 2025 · 8 citations
- Differentiable Quantum Computing for Large-scale Linear ControlConnor Clayton, Jiaqi Leng, Gengzhi Yang, Yi-Ling Qiao et al.NeurIPS 2024 · 7 citations
- Deep Stochastic MechanicsElena Orlova, Aleksei Ustimenko, Ruoxi Jiang, Peter Y. Lu et al.ICML 2024 · 2 citations
- Efficient Quantum Algorithms for Quantum Optimal ControlXiantao Li, Chunhao WangICML 2023 · 6 citations
- Hamiltonian Neural PDE Solvers through Functional ApproximationAnthony Y. Zhou, Amir Barati FarimaniNeurIPS 2025 · 1 citation
