A Schrödinger Eigenfunction Method for Long-Horizon Stochastic Optimal Control
Authors: Louis Claeys, Artur Goldman, Zebang Shen, Niao He
Organizations: ETH Zürich, Department of Mathematics · ETH Zürich, ETH AI Center, Institute for Machine Learning · ETH Zürich, Institute for Machine Learning
High-dimensional stochastic optimal control (SOC) becomes harder with longer planning horizons: existing methods scale linearly in the horizon T, 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. We prove that, under the gradient drift assumption, L is unitarily equivalent to a Schrödinger operator S=−Δ+V with purely discrete spectrum, allowing the long-horizon control to be efficiently described via the eigensystem of L. This connection provides two key results: first, for a symmetric linear-quadratic regulator (LQR), S 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 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) to O(d).
Figures & tables
Figure 1: Performance degradation as time horizon T increases for different methods (see Appendix E for details).
Figure 2: Diminishing returns from increasing the number of eigenfunctions for an LQR in d=20 dimensions.
Figure 3: Learned controls (arrows) and V0 for different eigenfunction losses. Existing methods fail to learn the correct control in regions where V0 is large due to implicit reweighting.
Figure 4: Comparison of the different eigenfunction losses (EMA).
Figure 5: Average L2 control error (EMA) as a function of iteration (top row) and L2 error as a function of t∈[0,T] (bottom row).
Method
Algorithm/Loss
Control objective
EIGF (ours)
Relative loss
73.33±0.02
IDO
Log-variance loss
74.52±0.02
IDO
Adjoint Matching
74.69±0.02
IDO
Relative Entropy
75.63±0.02
IDO
SOCM
Did not converge
Table 1: Opinion dynamics, final control objective (smaller is better).
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 6: Energy function (left) and running cost (right) for the ring setting in d=2 .
Experiment
Iteration time ( s )
Method
Loss
COMBINED
Adjoint Matching
0.332
Log variance
0.328
Relative entropy
0.419
SOCM
0.432
EIGF
Deep ritz loss
0.227
Appendix
Table 2: Iteration times by method and loss
Figure 7: Learned eigenfunctions (top row) and their logarithms (bottom row) for the ring setting with different loss functions.
Figure 8: L2 error of ∇logϕ for different eigenfunctions (left), average control L2 error for different methods (middle) and control L2 error over time (right). Legend is the same as in Figure 4 and Figure 5 .
Figure 9: In this experiment EIGF method uses an ergodic estimator based only on the first eigenfunction. EIFG+IDO curve corresponds to the application of the proposed controller ( 22 ) with Relative Entropy loss. The figure shows L2 control error for different methods after 30000 iterations.
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
Appendix
Table 3: Control objective for the different methods in the quadratic (isotropic) and quadratic (repulsive) settings. The SOCM method did not converge, and hence the dynamics diverge.
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
Appendix
Table 4: Control objective for the different methods in the quadratic (anisotropic) and double well settings.
We introduce trajectory-regularized stochastic optimal control (TRSOC), which augments standard stochastic optimal control (SOC) with a Kullback--Leibler (KL) divergence between controlled and reference trajectory distributions. Using Girsanov's theorem, the trajectory KL reduces to a quadratic drift mismatch penalty, yielding a modified running cost that preserves the dynamic programming (DP) structure. We derive the corresponding Hamilton--Jacobi--Bellman (HJB) equation and characterize the optimal policy. In the linear-quadratic (LQ) setting, the formulation admits a closed-form solution with an augmented control cost. Experiments show that the regularization parameter induces a trade-off between performance-driven and reference-preserving behavior, including cases with reference dynamics learned from offline data.
In this paper, we consider stochastic optimal control problems with infinite-horizon joint chance constraints. By means of an appropriate state augmentation, we reformulate the original problem as a constrained Markov decision process, in which both the cost and the constraint function exhibit an additive structure. We then prove that this formulation enjoys strong duality, thereby enabling us to reformulate the problem as an equivalent unconstrained one in the Lagrange dual framework. We propose a dual-ascent algorithm to solve the resulting problem and show that it converges to a deterministic Markov policy defined over the augmented state space that is both optimal and feasible. To accommodate continuous state-input spaces, we propose a dedicated learning algorithm to approximate the value function in an offline training setting, thereby significantly reducing the computational complexity of the online control phase. We then test our approach on a numerical example and demonstrate its effectiveness compared to online predictive control methods in terms of performance and computational complexity.
Francesco Cordiano, Kanghui He, Bart De Schutter
Delft Center for Systems and Control, Delft University of Technology, 2628 CD, Delft, The Netherlands · Department of Engineering Science, University of Oxford, OX1 3PJ Oxford, U.K.
A data-driven method is developed for approximating value functions in deterministic optimal control problems with nonlinear control-affine dynamics. The Pontryagin Maximum Principle optimality system is solved from multiple initial conditions to generate training data consisting of values, gradients, and Hessians of the value function, where Hessian information is obtained from a matrix Riccati equation along optimal trajectories. These quantities augment a weighted least-squares regression over sparse polynomial bases on hyperbolic cross index sets, with gradients and Hessians contributing additional linear equations per sample and substantially reducing sample complexity compared to value-only regression. Feedback laws are recovered analytically from the learned value function. In high dimensions, a partial Hessian strategy controls the cost of data generation. The approach is validated on problems of increasing state dimension, where second-order data augmentation is shown to improve approximation accuracy and closed-loop performance, with up to an order-of-magnitude reduction in the number of training samples required relative to lower-order methods.
Matías Gómez-Aedo, Behzad Azmi, Yuyang Huang +2
Department of Mathematics, Imperial College London, South Kensington Campus, SW7 2AZ London, UK · Department of Mathematics and Statistics, University of Konstanz, Germany · RICAM and Institute of Mathematics and Scientific Computing, University of Graz, Austria