A Schrödinger Eigenfunction Method for Long-Horizon Stochastic Optimal Control
Organizations: ETH Zürich, Department of Mathematics · ETH Zürich, ETH AI Center, Institute for Machine Learning · ETH Zürich, Institute for Machine Learning
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 .
Figures & tables
| Method | Algorithm/Loss | Control objective |
|---|---|---|
| EIGF (ours) | Relative loss | |
| IDO | Log-variance loss | |
| IDO | Adjoint Matching | |
| IDO | Relative Entropy | |
| IDO | SOCM | Did not converge |
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
| Experiment | Iteration time ( ) | |
| Method | Loss | |
| COMBINED | Adjoint Matching | 0.332 |
| Log variance | 0.328 | |
| Relative entropy | 0.419 | |
| SOCM | 0.432 | |
| EIGF | Deep ritz loss | 0.227 |
| quadratic (isotropic) | quadratic (repulsive) | ||
| Method | Loss | ||
| IDO | Relative entropy | 32.7870 0.014 | 112.5172 0.051 |
| Log variance | 32.7850 0.014 | 114.0552 0.053 | |
| SOCM | 73.1062 0.062 | nan nan | |
| Adjoint matching | 32.7754 0.014 | 113.4554 0.052 | |
| COMBINED | Relative entropy | 32.7763 0.014 | 112.3444 0.050 |
| quadratic (anisotropic) | double well | ||
| Method | Loss | ||
| IDO | Relative entropy | 38.9967 0.022 | 35.2688 0.010 |
| Log variance | 31.3664 0.016 | 32.8645 0.0094 | |
| SOCM | 112.2728 0.18 | 41.7215 0.013 | |
| Adjoint matching | 31.3584 0.016 | 34.8713 0.010 | |
| COMBINED | Relative entropy | 31.3476 0.016 | 32.6130 0.0088 |