English

A Schr\"odinger Eigenfunction Method for Long-Horizon Stochastic Optimal Control

Machine Learning 2026-03-25 v1 Optimization and Control

Abstract

High-dimensional stochastic optimal control (SOC) becomes harder with longer planning horizons: existing methods scale linearly in the horizon TT, 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 L\mathcal{L}. We prove that, under the gradient drift assumption, L\mathcal{L} is unitarily equivalent to a Schr\"odinger operator S=Δ+V\mathcal{S} = -\Delta + \mathcal{V} with purely discrete spectrum, allowing the long-horizon control to be efficiently described via the eigensystem of L\mathcal{L}. This connection provides two key results: first, for a symmetric linear-quadratic regulator (LQR), S\mathcal{S} matches the Hamiltonian of a quantum harmonic oscillator, whose closed-form eigensystem yields an analytic solution to the symmetric LQR with \emph{arbitrary} terminal cost. Second, in a more general setting, we learn the eigensystem of L\mathcal{L} 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 O(Td)\mathcal{O}(Td) to O(d)\mathcal{O}(d).

Keywords

Cite

@article{arxiv.2603.23173,
  title  = {A Schr\"odinger Eigenfunction Method for Long-Horizon Stochastic Optimal Control},
  author = {Louis Claeys and Artur Goldman and Zebang Shen and Niao He},
  journal= {arXiv preprint arXiv:2603.23173},
  year   = {2026}
}

Comments

Accepted to ICLR 2026, code available in https://github.com/lclaeys/eigenfunction-solver

R2 v1 2026-07-01T11:35:24.667Z