A Schrödinger Eigenfunction Method for Long-Horizon Stochastic Optimal Control
Authors: Louis Claeys
Artur Goldman
Zebang Shen
Niao He
Venue: ICLR 2026 Poster (2026)
Keywords: Stochastic Optimal Control
Schrödinger operator
eigenfunction learning
long-horizon control
Primary area: applications to robotics, autonomy, planning

Abstract: 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 $\mathcal{L}$. We prove that, under the gradient drift assumption, $\mathcal{L}$ is unitarily equivalent to a Schrödinger operator $\mathcal{S} = -\Delta + \mathcal{V}$ with purely discrete spectrum, allowing the long-horizon control to be efficiently described via the eigensystem of $\mathcal{L}$. This connection provides two key results: first, for a symmetric linear-quadratic regulator (LQR), $\mathcal{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 $\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 $\mathcal{O}(Td)$ to $\mathcal{O}(d)$.

Published as a conference paper at ICLR 2026
A SCHRÖDINGER EIGENFUNCTION METHOD FOR HORIZON STOCHASTIC OPTIMAL CONTROL
Louis Claeys
Artur Goldman
ETH Zürich,
ETH Zürich, ETH AI Center Institute for Machine Learning artur. gow.ooooooooooo.oooooo dman@a ai et
Department of Mathematics
louis. www.ycaaasl....moo welaayyl@ive...e.. 1ive.
Zebang Shen
Niao He
ETH Zürich,, Institute for Machine Learning zebang s fn ethz..c
ETH Zürich, Institute for Machine Learning, niao. he@ eettz. ch
ABSTRACT
High-dimensional stochastic optimal control (SOC) becomes harder with longer planning horizons: existing methods scale linearly in the horizon T, with per- formance often deteriorating exponentially. We overcome these limitations for a subclass of linearly-solvable SOC problems-thoo whose uncontrolled drift is the gradient of potential. In this setting, the MaailtonnJJacbb -ellman equa- tion reduces to a linear PDE governed by an operator L. We prove that, under the gradient drift assumption, L is unitarily equivalent
Schrödinger operator S - + V with purely discrete spectrum, allowing the long-horizon control to be efficiently described via the eigensystem of C. This connection provides two key results: first, for symmetric linear-quadratic regulator (LQR), S matches the Hamiltonian of quantum harmonic oscillator, whose closed-fome fored-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 method on several long- -ooizon benchmarks, achieving an order-of maggn-uue improvement in control accuracy compared to mate-of-the-art methods, while reducing memory usage and runtime complexity from (Td)) to ((d).
INTRODUCTION
Stochastic optimal control (SOC) concerns the problem of directing a stochastic system, typically modeled by stochastic differential equation (SDE), to minimize an expected total cost. SOC has found applications in various domains, c.g. stochastic filtering (Mitter) 1996). rare event simulation
in molecular dynamics (Hartmann & Schütte 2012 Hartmann et al. 2014). robotics (Gorodetsky
al. 2018) and finance (Pham) 2009).
Built on the principle of dynamic programming, the global optimality condition of SOC can be expressed by the -aanilton-Jaobi-BBone
(HJB) equation. In this paper, we focus on the affine control setting commonly considered in the literature (Fleming & Rishel 1975; Kappen 2005b
Fleming & Soner 2006 Yong & Zhou 1999. [oomino-EEnrihh al. 2024b Nüsken & Richter 202T
Carius al. 2022 Holdijk et al. (2023). where the control affects the state of the system linearly. This setting is of interest since the optimal control will exactly match the gradient of the value function of be drastically simplified.
SOC problem and hence the corresponding HJB equation
canonical special case of this affine-contre control framework is the linear tmear-quadratte regulator (LQR), in which the uncontrolled dynamics follow an Orsttein -ublenbeek linear SDE and both the running cost and the terminal cost are quadratic. One can show that in the LQR setting the value function retains a quadratic -ndeed it is quadratic terminal time because the terminal cost is quadratic and

---

Published as a conference paper at ICLR 2026
that the optimal feedback control is linear w.r.t the state. Consequently, the associated SOC problem admits an explicit solution via the finite-diment dinee-dimensional matrix Riccati differential equation (van Handel 2007)
To obtain the optimal control in more general settings requires numerical procedures. For low- dimensional problems, classical grid-based PDE solvers may be used, but these suffer from the curse of dimensionality. This has led to several works proposing the use of neural networks (NN) to solve the HJB equation in more complex high-dimensional settings, either through
forward-backward stochastic differential equation (FBSDE) approach (Han et al. 2018 Ji et al. 2022 Andersson et al. 2023 Beck al. 2019) or so-called iterative diffusion optimization (IDO) methods, which sample controlled trajectories through simulation and update the NN parameter using stochastic gradients from automatic differentiation (Nüsken Richter 2021 Domingo-Enre et al. 2024b a). more comprehensive review on existing methods for short-horizon SOC can be found in Appendix C
While these methods have proven successful, their perfor- mance suffers as the time horizon T grows. Both the memory requirement and per-iteration runtime increase at least lin- early in T. Additionally, it holds that error estimates for the deep FBSDE method worsen as T increases (Han & Long 2020 Theorem 4), and for IDO methods using importance sampling the weight variance may increase exponentially in T (Liu al. (2018). These limitations were observed em-
Control L Error (after 30k iteral (n1)
Aporthen
101
entrapy (100)
/D0 (ours)
commool Ereer
pirically in Assabumrungrat et al. (2024) Domingo-Enrinn Enrich
al. (2024b), and were reproduced in our experiments (see
Figure
Linearly-solvable HJB. The HJB equation is in general nonlinear. However, in the special case where the system's s diffusion coefficient matches the affine-control mapping, the Hope transform can eliminate the nonlinearity (Evans 2022). Specifically, let (x, t) denote the value function of the SOC problem and define a new function v == exp(-V).
Figure 1: Performance degradation as
time horizon T' increasses for different methods (see Appendix E for details).
Under this transformation, the HJB is equivalently rewritten as the following linear PDE
0, t) = Cs)(2, t), (x,TT)= = ())())
(I)
for some linear operator L. Moreover, the optimal control u, which exactly matches the gradient of the value function - V, , can be obtained as u" 0₂ log to (Kappen 2005b).
Working with linear PDE brings several clear bene- fits over nonlinear one, such as simplified analysis -questions of well-posedness and solution regularity of the PDE become much more tractable - but more impor- tantly algorithmic insight: One can borrow ideas from the finit--dimensin
Control error vs. time ftruncated rigenfunction so rolution)
1017
linear ODE ui() + A (t)) O, where A is a real symmetric matrix. This ODE has a closed-form solution u(t) e-At ((),, and because the matrix expo- nential e -A acts simply on A's eigenvalues, expanding (0)) in the corresponding eigenvectors yields an efficient numerical scheme.
2or5
12-1
Number
15-1
0.0
0.5
3.0
2.5
2.0
2.5
34
Time 2
While the domain of the operator L is infinite-dimensional (acting on functions rather than finite vectors), the same exponential-integran
Figure 2: Diminishing returns from in- creasing the number of eigenfunctions for LQR in d 20 dimensions.
principle applies (Theorem Ip. and we can write ti (t--T) where (e*2) ) a semi- 220
group (Renardy & Rogers 2006 Chapter 12). Of course, carrying out in an infinite-dimmnsiiona setting introduces additional technical challenges that must be carefully addressed. In particular, expanding र5 as a series of eigenfunctions requires L to possess a discrete spectrum, i.c. one can find an orthonormal basis of eigenfunctions (di)teN and corresponding eigenvalues Åo < X < such that Lo; = d; for all é E N. Under this assumption, the optimal control can be informally written
u² "=x t) = On log do (x) + (e-(((-N0))
for any fixed as T + x0,
2

---

Published as a conference paper at ICLR 2026
where O hides eigenfunctions 6; with i 2 1. A precise statement can be found in Theorem 3 To turn the above formula into a practical, long-horizon SOC algorithm, we need (1) Spectral verification. Prove that C indeed has discrete spectrum; (2) Eigenfunction identification. Compute the spatial derivative of the principal eigenfunction dh--
Our approach: Reduction to Schrödinger operator. In this paper, to guarantee the spectral verification, we assume that the drift of the dynamics is the gradient of a potential. Such problems of controlling a diffusion process with gradient drift show up in overdamped molecular dynamics (Schüttee al. 2012), mean-field games (Bakshi et al.
2019; Grover et al., 2018), the control of
particles with interaction potentials (Carrillo et al.] 2000;
Totzeck & Pinnau 2020), as well as social
models for opinion formation (Castellano et al.) 2009: Al al. 2017).
In this setting, the operator L in ① is unitarily equivalent to the Schrödinger operator L = - + V, where A is the Laplacian and V is an effective potential determined by the original drift and running cost. Because È known to have a purely discrete spectrum on I,² and unitary equivalence preserves the spectral properties, the operator C in (I) likewise enjoys a discrete spectrum.
This result forms the basis of our new framework for long- onn-horizoon SOC, where the problem is reformulated as learning the eigensystem of a Schrödinger operator. Indeed, (2) shows that the top eigenfunction Φ determines the long-term control, with corrections decaying exponentially with rate d Au We address the problem of eigenfunction identification in the following two scenarios:
Closed-form solution for LQR with non-quadratic terminal cost. When the drift is linear with symmetric coefficient matrix and the running cost is quadratic in the state, the resulting Schrödinger operator L coincides with the Hamiltonian of the quantum harmonic oscillator. Its eigenvalues and eigenfunctions are known explicitly (see Lemma or Griffiths & Schroeter (201))) i.e. {),, Oe hen are available in closed form. Consequently, our tramework yields fully explicit expression for the corresponding SOC. This removes the requirement of quadratic terminal cost in the the classical LQR solution.
Neural network-best baase approach for general gradient drift. For general gradient-drift dynamics, we introduce a hybrid neural- netwoork method approximate the optimal control efficiently: Rather than attempting to learn the full spectrum of L C-which is prohibitively expensive and yields rapidly diminishing returns (fig. 2 we exploit the exponential decay of the higher modes w.r.t. T in eq. 2 Concretely, whenever T - exceeds a modest threshold (in our experiments T - / 1), it suffices to approximate the control using only the top eigenfunction Φn. For the remaining period (t very close to T), we switch to established FBSDE/IDO solvers to handle the short-horizon SOC.
We propose a novel deep learning strategy for the task of learning the eigenfunction, tailored to SOC. While such task has been extensively studied in the literature, previous approaches either optimize a variational Ritz objective (E & Yu 2018 Zhang et al. 2022 Cabannes & Bach 2024) minimize the residual norm [L - A 1 (Jin et al. 2022) Zhang al. 2024 Nusken & Richter 2023). However, these losses implicitly reweight spatial regions, causing them to f'ail to learn the control in regions where the value function V is large- the most crucial areas. To eliminate this bias, we introduce relative eigenfunction loss, /// b All which removes the undesired weighting and robustly recovers the dominant eigenpair needed for control synthesis.
Our contributions are summarized as follows:
We provide new perspective on finite-horizon gradient-drift SOC problems, linking their solution the eigensystem of a Schrödinger operator (Theorem 3). This yields a previously unreported closed-form solution to the symmetric LQR with arbitrary terminal cost (Theorem 4).
With this framework, we introduce a new procedure for solving SOC problems over long horizons by learning the operator's eigenfunctions with neural networks. We show that existing eigenfunction solvers can be ill-suited for this task due to an implicit reweighting in the used loss, and propose a new loss function to remedy this.
We perform experiments in different settings to evaluate the proposed method against ot the-art SOC solvers, showing roughly an order of magnitude improvement control L error on several high-dimensional (d = 20) long-horizon problems.
3

---

Published as a conference paper at ICLR 2026
2
PRELIMINARIES
2.1
STOCHASTIC OPTIMAL CONTROL
Fix filtered probability space (S, F, (Fe)azo: P), and denote (Wo) Brownian motion in this space. Let (Xx)) denote the random variable taking values in IR defined through the SDE
dxp (6(X)) + (X7 dt VB-lodW,
X po
where u IR [0. + E is the control, b [ + Rd is the base drift, E gdxd is the diffusson coefficient (assumed invertible) and 3 E R is inverse temperature characterizing the noise level. Note that we assume the drift and noise to be time-independent,, in contrast to Nüsken & Richter (2021) : Domingo-Ennnnn Enrich et al. (2024b). Under some regularity conditions
the coefficients and control u described in Appendixx A the SDE (3) has a unique strong solution. In stochastic optimal control, view the dynamics (b, o, () as given and consider the problem of finding control which minimizes the cost functional
J. (u((, 0)1 s(x,)) dt 9(X7)
J Ju, -
X,
where f: IRd + IR denotes the running cost and 12 + IR denotes the terminal cost. We denote this optimal control as u" (xx t) arg win (u; æ, t). with U4 the of admissible controls. To analyze this problem, one defines the value function V, which is defined as the minimum achievable cost when starting from time t.
Via, 1) = inf Ju;
(5)
well
In this case, the optimal control u² that minimizes the objective (4) obtained from the value function through the relation u² = -o" VV, as described (Nüsken Richter 2021 Theorem 2.2).
Hamilton-Jasd Jacioonnn--aci---- Jacobi-Bellaan Bellman equation. A well-known fundamental result is that when the value function is sufficiently regular, it satisfies the following partial differential equation, called the Mamilton-Jacobi- Bellman (HJB) equation (Fleming & Rishel 1975):
V.V + KV = in R x [0,7].
V(.T) = on Rd
(6)
where KV
1 Tr(aa V + 6t VV
2 1
(7)
23
The so-called verification theorem states (in some sense) the converse: if function V satisfying the above PDE is sufficiently regular, it coincides with the value function (5) corresponding to (3-44)
see (Fleming& & Rishel 1975. Section VI.4) and (Pavliotis) 2014, Sec. 2.3).
linear PDE reformulation Although the HJB equation (6) is nonlinear in general, it was shown in
Kappen (2005b) that for specific class of optimal control problems (which includes the formulation B suitable transformation allows for linear reformulation (6). More specifically, when
parametrizing V(x, f) 3-11 log (2, 1 (T t)), the nonlinear terms cancel, and (6) becomes
S0, + Leb = 0, (v(-00= = vo.
where CG = - Tr(oo T VV 2,67 V 2 f- vo = exp(- (8)
and we have introduced the variable - = (2 - 1). This is precisely the abstract form (I). but with a time reversal. For more details on this result, we refer to Appendix B To simplify the presentation, we will often assume w.l.ogg. that a = I (sec Axcendixx A), so that Tr(aoo A.
2.2 EIGENFUNCTION SOLUTIONS
Consider Hilbert space H. with inner product (-.-), and linear operator L D(C) + H defined on a dense subspace (L.) C H. Then we have the following standard definition:
Definition 1 An element à H with Lo =o The value X is called an eigenvalue of C (corresponding to ő), and the dimension of the null space of f. XI is called the multiplicity of X.
is eigenfunction of L if there exists X C such that
4

---

Published as a conference paper at ICLR 2026
In knite-dimensiin setting, the study of a linear operator A is drastically simplified when we have access to an orthonormal basis of eigenvectors. Similarly, some operators admit a countable set of eigenfunctions (Øt) which forms orthonormal basis of the Hilbert space H. When such an eigensystem exists, the following theorem proven in Appendix B gives a solution to the PDE ( in
terms of the eigensystem. A rigorous connection between this solution to (8) and solutions (6) is explored in Appendix B
Theorem (Restatement of Theorem VIIL. in (Reed & Simon, 1980)) Let L be essentially
self-adjoonn I densely defined operator on H. which admits an orthonormal basis of eigenfunctions (di, As ken. Assume further that the x; are bounded from below (write No A1 ...) and do not
have a finite accumulation point. Then a solution to the abstract evolution problem in (8) is given by
(b())
=-
Aet (dis vo)
iEN
While we originally aimed to solve the HJB equation (6). in this work, we solve the linear PDE 8 as surrogate. In Section we show that under standard growth and regularity assumptions, the above semigroup solution coincides with the unique viscosity solution of the original HJB equation, up to an inverse Hope-Cole transformation. See Theorem for a formal statement.
OUR FRAMEWORK
3.1 SPECTRAL PROPERTIES OF THE SCHRÖDINGER OPERATOR
In order to apply Theorem we must establish conditions under which the operator L in satisfies the required properties. In order for L to be symmetric, assume
(A1) The drift b in described by a gradient field: b(x) VE(x).
Define the measure 1a on IRd with density μxx) = exp (-23E())) and consider the weighted Lebesgue space I (14). Note that we do not require ys to be finite measure. Under the assumption (AI). the operator appearing in (8) becomes the following operator on L L²(p):
L: D(C) c (2(x)) (z()) : t/ + Les = A + 2 E, Ve) + f
(10)
Under mild regularity conditions, Furthermore, it can be shown (Appendix B) that
further show that L. is essentially self-adjoint (Appendix B
UCU-1= - 5² VEI BAE 2
(11)
where U (M) + (RR)) 10 + e BF is unitary operator, so that C is unitarily equivalent A + V on (R())). which forms the comerstone of
the well-known Schrödinger operator the mathematical formulation of quantum physics. Its properties have been studied to great extent (Reed & Simon 1978), allowing us to invoke well-known results on the properties of the Schrödinger operator to study the behavior of C. In particular, the following assumption on E and f is enough to guarantee that satisfies all the desired properties (see Appendix B
(A2) For the energy E and rumning cost f, V = 2 AE + 28f satisfies V E Lio 3C E IR, V E Rd C' < V(r), and ((x) + 0o as | + oc.
Theorem (Restatement of (Reed & Simon 1978, Theorem XII1.67, XIII164, XIIL47)
Suppose (A2) is satisfied. Then the operator L in (10) is densely defined and essentially self-adjoint. Moreover, admits a countable, orthonormal basis of eigenfunctions. In addition, the eigenvalues are bounded from below and do not have finite accumulation point, the lowest eigenvalue X has multiplicity one, and the associated eigenfunction (called the ground state, or in our context the top eigenfunction) can be chosen to be strictly positive.
Remark Previous studies have linked the Schrödinger operator to optimal control in contexts distinct from ours: for example, Schütte el al. (2012) and Bakshi et al. (2020) analyze the slationary HJB equation KVV. = X (see Remark 2}. whereas Kalise et al. (2025) explores distribution-level control of the Pokker-Plance equation.
An operator is called essentially self-adjoint if its closure is self-adjoint. See Reed & Simon 1980 and Reed & Simon 1975 for more details.
5

---

Published as a conference paper at ICLR 2026
3.2 EIGENFUNCTION CONTROL
From the previous discussion, we obtain the following result, which links the eigenfunctions of L. with the optimal control problem.
Theorem Suppose (AI) (A2) are satisfied. Then & satisfies all assumptions in Theorem 1 hence the solution of the optimal control problem (3 is given by
(e Bg Φi) e (A-- 0z(x)
u² u²(x, (x, = 88 V llogoolas (o(x) + V log
1
+
(12)
{e (e-89, 00)
@0 (z)
$50
where (òr, Ae)SEN is the orthonormal eigensystem of defined in 10 and Ao <
Thus, the long-term optimal control (t < T) is given by 8 V log Φo, and the corrections decay exponentially, motivating truncation of the series
3.3 CLOSED-FORM SOLUTION FOR THE SYMMETRIC LQR
When E and are quadratic, the Schrödinger operator associated with the optimal control system is Hamiltonian for the harmonic oscillator, which has an exact solution (Appendix Bh:
Theorem 4 Suppose b(x) Ax for some symmetric matrix A E Rdxd and f() x² Pae for
some matrix P E Rdxd such that A A + 2P is positive definite, and has diagonalization U AU. Then the orthonormal eigensystem of the operator C given in (10) is described through
exp (-20" (-A+UT x
(a(x)
Has (VB(A)/Um)))
(13)
(AT)a/4
1=1 22 (a.!)
da = B -Tr(A)) Σ A/22 +1)
(14)
where E Nd and H3 denotes the ith physicist's 's Hermite polynomial. We can bijectively map Nd to N by ordering the eigenvalues (4). yielding the same representation as before.
Combined with Theorem 3) this yields closed-form solution for the optimal control problem with symmetric linear drift, quadratic running cost and arbitrary terminal reward.
NUMERICAL METHODS
We propose a hybrid method with components: Far from the terminal time T. learn the top eigenfunction do and simply use Or log de the control (c.f. eq. (2); Close to the terminal time, e.g. t T we use an existing solver learn an additive short-horizon correction to the control.
4.1 LEARNING FIGENFUNCTIONS
In absence of a closed-form solution, wide range of numerical techniques exist for solving the eigenvalue problem for the operator L in (10). Classically, the eigenfunction equation is projected onto
subspace to yield Galerkin/finite clement method, see Chaitin-Chatelinn (1983. In high dimensions, these methods often perform poorly, motivating deep learning approaches which differ from each other mainly in the loss function used. We will only discuss methods for learning a single eigenfunction, referring to Appeppixx for extensions to multiple eigenfunctions. An overview of the deep leraning algorithm for learning eigenfunctions is given in Algorithm Appendix B
PINN loss Based on the success of chysics-informed neural networks (PINNs) RRaissi et 2019) one idea is to design a loss function that attempts to enforce the equation Lo A via an L loss, as done in Jin al. 2022). The idea is to consider some density p on R.d. and define the loss function
| [Ø] >Φl22 + a&Ree
Reeg (Ø) (1|=11 1)2
(15)

---

Published as a conference paper at ICLR 2026
inputs
Deep Ritz loss Kuw-)
Vartational loss Kiw,
loss Rewa
Relelive luss New (ours)
posttio
-2
-4
-2
pos
posttton
positine
Figure 3: Learned controls (arrows) and Vo for different eigenfunction losses. Existing methods fail to learn the correct control in regions where Vo is large due to implicit reweighting.
where a > is regularizer to avoid the trivial solution = O. The eigenvalue À is typically also modeled trainable parameter of the model, or obtained through other estimation procedures.
Variational loss A second class of loss functions is based on the variational characterization of the eigensystem L. Since essentially self-adjoint with orthogonal eigenbasis in a subset of (²(x) (u is defined below (A1)), it holds that (see (Reed & Simon 1978 Theorem XIII.I))
Ao inf (vb, LV))
(16)
VEL' (N)
(V) 6)
where the infimum is obtained when Lib doth. This motivates the loss functions (Φ) LO) + uR (0),
Rar, 1 (x))
Rva = (Ø, Lb) + a (6).
(17)
6)p
The first of these, sometimes called the deep Ritz loss, introduced in E & Yu 2018). and the second was described in Cabannes & Bach 2024): Zhang et al. (002)) These loss functions do not require prior knowledge of the eigenvalue Ag-
4.1.1 IMPLICIT REWEIGHTING IN PREVIOUS APPROACHES
The exponential decay of the correction term in eq. 2 suggests that the optimal control u can be approximated by the spatial derivative of logarithm of the top eigenfunction. We therefore parameterize in our implementation á = cxp(-3V)), where Vo is some neural network and 3 is the temperature constant. This choice also enforces the strict positivity of é, which matches the same property of do established in Theorem 2
Adopting such a parameterization, for the PINN and variational losses, it holds that
RPINN(E .V6) (PiNN(e-8V))-
484
e -8V
Ao (CV
+ (e -V6)
(18)
2
R (e 282 -00
2 KVo du + R. 3Vo)
(e-9Vo
(19)
where IC is the HJB operator Because both (88 and 19 incorporate an exponential factor that vanishes where Vo is large, these losses become effectively blind to errors in high-Vo regions and are only able to learn where Vo small. This pathology illustrated in Figure 3 Consider RING energy landscape F,, whose minimizers lie circle, and cost that grows linearly with the z-coordinate (see the full setup in Section Ep. The true optimal control remains tangential to the circle. In contrast, the controls obtained via the PINN and variational eigenfunction losses collapse in regions of large Vo, deviating sharply from the expected direction.
4.1.2 OUR APPROACII: REMOVING IMPLICIT REWEIGHTING VIA RELATIVE LOSS
Based this observation, we propose to modify
15
2
(+)
+ ax RKReeg
(20)
7

---

Published as a conference paper at ICLR 2026
This loss function, which call the relative loss, eliminates the implicit reweighting of the stationary HJB equation. Indeed, the same computation as before yields
Ao
Re = 4,3 KV6
+ Rggng (e (e-
(21)
2
As result, the relative loss remains sensitive even in regions where Φ - e becomes small. This can also be empirically observed in the RING task, as illustrated in Figure 3 Instead learning Vo and λy jointly, natural idea is to combine the benefits of the above loss functions by first training with variational loss (17) to obtain an estimate for the cigenvalue Ao and a good initialization Vo.
and then "fine-tune' using (20). In practice, we also observed that this initialization is necessary for the relative loss (20) to converge.
Remark 2 We note that there is an alternative interpretation of (21 based separate class of control problems in which there is no terminal cost g, and an infinite-horize cost is minimized, yielding a stationary or ergodic optimal control (Kushner 1978). In this setting, is related time-independent value function Vx which satisfies a stationary HJB equation of the form KV = X. In low dimensions, these problems are solved through classical techniques such as basis expansions or grid-based methods, with no involvement of neural networks (Todorov) 2009).
4.2 OUR HYBRID METHOD: COMBINING EIGENFUNCTION AND SHORT-HORIZON SOLVERS
For both IDO and FBSDE methods, every iteration requires the numerical simulation an SDE, yielding linear increase in computation cost with the time horizon T. We propose to leverage the eigenfunction solution given in Theorem 3 in order to scale these methods to longer time horizons as follows: first, parametrize the top cigenfunction as Φ = exp(-3V,") for neural network Vo" and learn the parameters On using the relative loss, as well as the first two eigenvalues λo (see Appendix Ep. Next, choose some cutoff time Teut < T and parametrize the control as
1 18-'Vlog
Vlog log Pn
0 0 Teuts Teut <
uo uo(
8-11 V
Ou log Pn (z) = (1 (n - (z,t))
(22)
This control can then be used in an IDO/FBSDE algorithm to optimize the parameters 01 of the additive correction eft, second neural network, near the terminal time. Crucially, this only requires simulation of the system in the interval [Toue, T]. significantly reducing the overall computational burden and reducing the time complexity of the algorithm from O to O(d).
Choice of Tout- Increasing Teut + T raises computational cost, whereas choosing Teut too small degrades performance; this reflects the fundamental trade-off in selecting this parameter. One heuristic for choosing an appropriate order of magnitude is to observe that we would like the correction term given by the infinite series in Eq. (12) be small for all < products or eigenfunctions, we may consider the approximation
Teue. Even without access to the inner
exp
28 (A1 AT) - -Teut) E,
which yields
23
T T Teut
X1 Ao logs.
Since obtain empirical estimates of \ and \1, this expression provides practical guideline for determining the scale of Tout, and thus can be implemented in practice.
5 EXPERIMENTS
To evaluate the benefits of the proposed method, we consider four different settings, QUADRATIC (ISOTROPIC), QUADRATIC (ANISOTROPIC), DOUBLE WELL, and RING. An additional setting QUADRATIC (REPULSIVE) with nonconfining energy is discussed in Appendix E The first three are high-dimensional benchmark problems adapted from Nüsken & Richter (2021), modified to be long- horizon problems, where a ground truth can be obtained. Detailed information on the experimental setups, including computational costs, is given in Appendix E
8

---

Published as a conference paper at ICLR 2026
Quadratic (190tropic) (dm20)
Quadratic (anisotropic) (dm20)
Double well (d=10)
Ring
10bg -000
Commd deeeciv
62.6
22.4
13-1
2014
52.
1377
20000 ECCAN (000)
30000
20230
(0093)
80200
00300
20030
80000
50400
10002
15000
20000
Variational loss
Deep loss
loss
Relative
Figure 4: Comparison of the different cigenfunction losses (EMA).
Quadratic (Isotropic) (d=20)
Quadratic (anisotropic) (d=20)
Double well (d-10)
10
Average control erar
10P
10
10"
Average control
20
Average
10m2
1017
10:2
10-3
20000
40000
60000
80000
20000
40000
60000
80000
10000
20000
30000
Inraatio
Quadratic (isctropic) (d=20)
Quadratic (an aniisoropicc (d=20)
Dougle well (d=10)
Contre erce
100
Cottro
10-1
10
10a3
10a1
3.0
15
2.0
0.5
3.5
4.0
0.0
2.5
3.0
"Ime
Time
Time
Relative entropy
Log variance
SOCM
Adjoint matching
EIGF Relative entropy (ours)
EIGF Log-variannce (ours)
EIGF UIG++SOCM (ours)
EIGF+ Matching (ours)
Figure 5: Average 2 control error (EMA) as a function of iteration (top row) and L error function of t [0, T' (bottom row).
Figure 4| shows the results of the various eigenfunction losses. For the QUADRATIC settings, we can compute V log Φo exactly, and see that the relative loss significantly improves upon existing loss functions for approximating this quantity (with the error measured in (1)).. For the other settings, the resulting control log do yields the lowest value of the control objective for the relative loss.
we show the result of using the learned eigenfunctions in the IDO algorithm using the combined algorithm described in the previous section, and compare it with the standard IDO/FBSDE methods. In each setting, obtain a lower _ error using the combined method, typically by an order of magnitude. The bottom row of Figure 5) shows how the error behaves as a function of t T]: the pure eigenfunction method achieves superior performance for 0, but performs worse closer to the terminal time T. The IDO method has constant performance in [0, IJ, and the combined method combines the merits of both provide the lowest overall 1,2 emror.
In Figure 5

---

Published as a conference paper at ICLR 2026
Method
Algorithm/Loss
Control objective
EIGF (ours)
Relative loss
73.33 + 0.02
IDO
Log-varianee 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).
5.1 OPINION MODELING: DE GROOT MODEL
We additionally consider a networked control problem inspired by opinion formation. We model N agents with state X, R200 and dynamics
dx,- = ((L - f)X; yx, + u(X))) dt + dW.,
(23)
where L E 100010 is a symmetric interaction matrix with Lμs = 0.5 and L,,== 0.25 (all other entries are zero), and we set 7 0.1. We use the non-quadratic running cost
10 Σ
f(x)
(1) 11
(24)
and terminal cost g(x) = 0, with parameters = 1.0 and time horizon T = 10.0. Table final control objective (mean std) after 80,000 training iterations.
reports the
CONCLUSION
In this work, have introduced perspective a class of stochastic optimal control problems with gradient drift, showing that the optimal control can be obtained from the eigensystem Schrödinger operator. We have investigated the of deep learning methods to learn the eigensystem, introducing a new loss function for this task. We have shown how this approach can be combined with existing IDO methods, yielding an improvement in L error of roughly an order of magnitude over oftte-of-the-at methods in several long-horizen experiments, and overcoming the increase in computation cost typically associated with longer time horizons.
The main drawback of the proposed approach is that it is currently limited to problems with gradient drift. When the operator C is not even symmetric, it may no longer have real cigenvalues. Nonetheless, the top eigenfunction may still be real and nondegenerate with real eigenvalue, that the long-term behaviour of the control still described by an eigenfunction (Evans 2022 Theorem 6.3). A second limitation is that there is priori method for determining appropriate cutoff time Tout, this is hyperparameter that should be decided based on the application and the spectral gap X1 Ao-
I Limitations
ACKNOWLEDGEMENTS
The work is supported by ETH research grant, Swiss National Science Foundation (SNSF) Project Funding No. 200021-207343 and SNSF Starting Grant. Artur Goldman is supported by the ETH AI Center through ETH AT Center doctoral fellowship.
REFERENCES
Giacomo Albi, Young Pil Choi, Massimo Fornasier, and Dante Kalise. Mean Field Control Hierarchy. Applied Mathematics and Optimization, 76(1):93-133 August 2017. ISSN 0095-4616. doi:
10.10007/00225-007-- -117-99999..
Kristoffer Andersson, Adam Andersson, and C. W. Oosterlee. Convergence of Robust Deep FBSDE Method for Stochastic Control. SIAM Journal on Scientific Computing, 45(1):A226-A25 February 2023. ISSN 1064-8275. doi: 10. 4777277777777
10

---

Published as a conference paper at ICLR 2026
Rawin Assabumrungrat, Kentaro Minami, and Masanori Hirano. Error Analysis of Option Pricing via Deep PDE Solvers: Empirical Study. In 2024 16th IIAI International Congress on Advanced Applied Informatics (IIAI-AAI), pp. -36.. IEEE Computer Society, July 2024. ISBN 979-8- 3503-7790-3. doi: 10.1 10.1109/AA-AAA11666 2022.00068.
Kaivalya Bakshi, Piyush Grover, and Evangelos Theodorou. On mean field games for agents with langevin dynamics. IEEE Transactions Control of Network Systems, 6(4):1451-146) 2019.
Kaivalya Bakshi, David Fan, and Evangelos A Theodorou. Schrödinger approach to optimal control of large-size populations. IEEE Transactions on Automatic Control, 6665)2332--2378,
2020.
Milosz Baraniewiez. Ground state decay for schr'," odinger operators with confining potentials. arXiv preprint arXiv:2407.09267. 2024.
Christian Beck, Weinan E, and Armulf Jentzen. Machine learning approximation algorithms for high- dimensional fully nonlinear partial differential equations and second-order backward stochastic differential equations. Journal of Nonlinear Science, 1963))55631111) August 2019. ISSN 0938-8974, 1432-1467,
doi: 10.10077/00033--088 9-955--3.
Dimitri Bertsekas and Steven E. Shreve. Stochastic Optimal Control: The Discrete-Tine Case. Athena Scientific, December 1996. ISBN 978-1 978-1-886529--038
Samuel Biton. Nonlinear monotone semigroups and viscosity solutions. In Annales de l'Institut Henri Poincaré C, Analyse non linéaire, volume 18, pp. 383-402. Elsevier, 2001.
J. Frédéric Bonnans, Élisabeth Ottenwaelter, and Housnaa Zidani. A fast algorithm for the two dimensional HJB equation of stochastic control. ESAIM: Modélisation mathématique et analyse numérique, 38(4):723-735, 2004. ISSN 1290-3841. doi: 10.105 /m22n:20040044.
Susanne C. Brenner and L. Ridgway Scott. The Mathematical Theory of Finite Element Methods, volume 15 of Texts in Applied Mathematics. Springer, New York, NY, 2008. ISBN 978-0-387- 75933-3 978-0-3800 778-0-387-75934-0 doi: 10. (0007997-00388775533
Steve Brooks, Andrew Gelman, Galin Jones, and Xiao-Li Meng (eds.). Handbook of Markov Chain Monte Carlo. Chapman and Hal//CRC, New York, May 2011. ISBN 978-0-429-13850-8. doi: 10.1201/b10905.
Vivien Cabannes, Bobak Kiani, Randall Balestriero, Yann LeCun, and Alberto Bietti. The ssl interplay: Augmentations, inductive bias, and generalization. In International Conference on Machine Learning, pp. 3252-3298. PMLR, 2023.
Vivien A. Cabannes and Francis Bach. The Galerkin method beats Graph-Based Approaches for Spectral Algorithms. In Proceedings of The 27th International Conference on Artificial Intelligence and Statistics, pp. 451-459. PMLR, April 2024.
Elisa Calzola, Elisabetta Carlini, Xavier Dupuis, and Francisco Silva. semi-Lagrangian scheme for Hamilton-Jaee Bellman equations with oblique derivatives boundary conditions. Numerische Mathemarik, 153(1) (:33()))99111)) 494444444.. January 2023. doi: 10. 00.1000/00211-022 02200222226222....
Jan Carius, René Ranftl, Farbod Farshidian, and Marco Hutter. Constrained stochastic optimal control with learned importance sampling: A path integral approach. The International Journal of Robotics Research, 41 (2)::199: 89-8099999100 2022.
E. Carlini, A. Festa, and N. Forcadel. A Semi-Lagrangian Scheme for Hamilton-Jacobi Jallam-Belllnaa Equations on Networks. SIAM Journal on Numerical Analysis, 58(6)3315-33166 January 2020. ISSN 0036- 1429. doi: (0.1137/19M1269911
José A. Carrillo, Edgard A. Pimentel, and Vardan K. Voskanyan. On a mean field optimal control problem. Nonlinear Analysis, 199: 12039, October 2020. ISSN 0362-546X. doi: 10. 2020.1 12039.
Claudio Castellano, Santo Fortunato, and Vittorio Loreto. Statistical physics of social dynamics. Reviews of Modern Physics, 81(2):911-666, May 2009. doi: 10. ModPPyyy
11

---

Published as a conference paper at ICLR 2026
Françoise Chaitin-Chaie Chatelin. Spectral Approximation of Linear Operators. Academic Press, 1983. ISBN 17880-12--706222
Michael G Crandall, Hitoshi Ishii, and Pierre-Louis Lions. User's guide to viscosity solutions of second order partial differential equations. Bulletin of the American mathematical society, 27(1): 1-67, 1992.
Francesca Da Lio and Olivier Ley. Uniqueness results for second-order bellman-isaacs equations under quadratic growth assumptions and applicationss SIAM journal on control and optimization, 45(1):74-106, 2006.
Francesca Da Lio and Olivier Ley. Convex hamilton-jacobi equations under superlinear growth conditions data. Applied Mathematics & Optimization, 63(3):309-339, 2011.
Edward Brian Davies and Barry Simon. Ultracontractivity and the heat kernel for schrödinger operators and dirichlet laplacians. Journal of Functional Analysis, 59(2):335-395, 1984.
Carles Domingo-Enrichh A Taxonomy of Loss Functions for Stochastic Optimal Control, October 2024.
Carles Domingo-Enrich, Michal Drozdzal, Brian Karrer, and Ricky T. Q. Chen. Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control. In The Thirteenth International Conference on Learning Representations, October 2024a.
Carles Domingo-Enrich Jiequn Han, Brandon Amos, Joan Bruna, and Ricky T. Chen. Stochastic Optimal Control Matching. Advances in Neural Information Processing Systems, 37:1 12459- 112504, December 2024b.
Weinan E and Bing Yu. The Deep Ritz Method: A Deep Learning- Based Numerical Algorithm for Solving Variational Problems. Communications Mathematics and Statistics, 6(1):1-12, March 2018. ISSN 2194-671X. doi: 10.10077440004-018-00772.
Weinan E. Jiequn Han, and Arnulf Jentzen. Deep Learning- Based Numerical Methods for High- Dimensional Parabolic Partial Differential Equations and Backward Stochastic Differential Equa- tions. Communications Mathematics and Statistics, 5(4):349 -30, December 2017. ISSN 2194-671X. doi: 10. 100000//0000001--00101111
Weinan E, Martin Hutzenthaler, Aruul Jentzen, and Thomas Kruse. Multilevel Picard iterations for solving smooth semilinear parabolic heat equations. Partial Differential Equations and Applica- tions, 266)880, November 2021. ISSN 2662-2971.
doi: 10. 00.10777442985-01-00009-9
Alexandre Emn and Jean-Luc Guermond. Theory and Practice of Finite Elements, volume 159 of Applied Mathematical Sciences. Springer, New York, NY, 2004. ISBN 978 978-1-4419-19119
978-1-4757-4355-55 doi: 0-00//9778--457-435555
Lawrence C Evans. Partial differential equations, volume 19. American Mathematical Society, 2022.
Wendell Fleming and Raymond Rishel. Deterministic and Stochastic Optimal Control. Springer, New York, NY, 1975. ISBN 978-1-4612-6382-1 978-1-4612-6380-7. doi: 10.1007/978- 10.10079978-114611-668-----
Wendell H Fleming and H Mete Soner. Controlled Markov processes and viscosity solutions. Springer, 2006.
Emmanuel Gobet, Jean-Philippe Lemor, and Xavier Warin. A regression baased Monte Carlo method to solve backward stochastic differential equations. The Annals of Applied Probability, 15(3): 2172-2202, August 2005. ISSN 1050-5164, 2168-8737. doi: (0.1224/10505 1605000000412.
Alex Gorodetsky, Sertac Karaman, and Youssef Marzouk. High-dimensional stochastic optimal control using continuous tensor decompositions. The International Journal of Robotics Research, 3702-3)340-377. 2018.
David J. Griffiths and Darrell Schroeter. Introduction to Quantum Mechanics. Cambridge University Press, 3 edition, 2018.
12

---

Published as a conference paper at ICLR 2026
Piyush Grover, Kaivalya Bakshi, and Evangelos A. Theodorou. A mean-field -ield game model for homogeneous flocking. Chaos: An Interdisciplinary Journal of Nonlinear Science, 28(6):061103, June 2018. ISSN 1054-1500. doi: 10.1063/1.5036663
Jiequn Han and Jihao Long. Convergence of the deep BSDE method for coupled FBSDEs. Probability, Unceriainty and Quantitative Risk, 5(1):5, July 2020. ISSN 2367-0126. doi: 10.1 1411166//1154-0225 020-00000-00-
Jiequn Han, Arnulf Jentzen, and Weinan E. Solving high-dimensional partial differential equations using deep learning. Proceedings of the National Academy of Sciences, 115(34):8505-855) August 2018. ISSN 0027-8424, 1091-6490. doi: 10. 10.1073/pnas. 1 1718942115.
Carsten Hartmann and Christof Schütte. Efficient rare event simulation by optimal nonequilibrium forcing. Journal of Statistical Mechanics: Theory and Experiment, 2012(11):P1 10044, November 2012. ISSN 1742-5468. doi: 10. 10080//74425448/22222111111
Carsten Hartmann, Ralf Banisch, Marco Sarich, Tomasz Badowski, and Christof Schütte. Charac- terization of Rare Events in Molecular Dynamics. Entropy, 16(1):350-376, January 2014. ISSN 1099-4300. doi: 10.3390/e1010330.
Dan Hendrycks and Kevin Gimpel. Gaussian Error Linear Units (GELUs)) June 2023.
Lars Holdijk, Yuanqi Du, Ferry Hooft, Priyank Jaini, Berend Ensing, and Max Welling. Stochastic Optimal Control for Collective Variable Free Sampling of Molecular Transition Paths. Advances in Neural Information Processing Systems, 36.79540-79556. December 2023.
Max Jensen and lain Smears. On the Convergence of Finite Element Methods for Hamilton-Jacobi- Bellman Equations. SIAM Journal on Numerical Analysis, 51(1):137-162, 170 2013. ISSN
0036-1429. doi: 10. 137/1 10856198.
Shaolin Ji, Shige Peng, Ying Peng, and Xichuan Zhang. Solving Stochastic Optimal Control Problem via Stochastic Maximum Principle with Deep Learning Method. J. Sci. Comput., 93(1), October 2022. ISSN 0885-7474. doi: 10.1007/s10915. 0//7//19955--22222 022-00222-99225
Henry Jin, Marios Mattheakis, and Pavlos Protopapas. Physics-informed neural networks for quantum eigenvalue problems. In 2022 International Joint Conference Neural Networks (IJCNN), pp. 1-8. IEEE, 2022.
Kamil Kaleta, Mateusz Kwasnicki, and József Lörinezi. Contractivity and ground state domination properties for non-local schrödinger operators. Journal of Spectral Theory, 8(1):165-189, 2018.
Dante Kalise, Lucas M. Moschen, Grigorios A. Pavliotis, and Urbain Vaes. A Spectral Approach to Equation, March 2025.
Optimal Control of the 04018//177742546/205////111
HJ Kappen. Path integrals and symmetry breaking for optimal control theory. Journal of Statistical Mechanics: Theory and Experiment, 2005(11):P11011, November 2005a. ISSN 1742-5468. doi:
Hilbert J. Kappen. Linear Theory for Control of Nonlinear Stochastic Systems. Physical Review Letters, 200201, November 2005b. doi: 951100// 20,200.
H. Kushner. Optimality conditions for the average cost per unit time problem with a diffusion model. SIAM Journal on Control and Optimization, 16(2):3303346, 1978. doi: 10.1137/031602
Jun S. Liu. Monte Carlo Strategies in Scientifiic Computing. Springer Series in Statistics, Springer, New York, NY, 2004, ISBN 978-0-387--76369-9 978-0-387--6371--2 doi: 10.1007/ 078-0-387-76371-2.
Qiang Liu, Lihong Li, Ziyang Tang, and Dengyong Zhou, Breaking the Curse of Horizon: Infinite- Horizon Off-Policy Estimation. In Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018
Francis A. Longstaff and Eduardo S. Schwartz. Valuing American Options by Simulation: Simple Least-Squares Approach. May 2001.
13

---

Published as a conference paper at ICLR 2026
S.K. Mitter. Filtering and stochastic control: A historical perspective. IEEE Control Systems Magazine, 16(3):67-76, June 1996. ISSN 1941-000X. doi: 10.1109/37.506400
Nikolas Nüsken and Lorenz Richter. Solving high-dimensional Vamitonn-JacoiiBellle PDEs using neural networks: Perspectives from the theory of controlled diffusions and measures path space, Partial Differential Equations and Applications, 2(4):48, June 2021. ISSN 2662-2971. doi:
10.10077/42985-021-00010
Nikolas Nusken and Lorenz Richter. Interpolating Between BSDEs and PINNs: Deep Learning for Elliptic and Parabolic Boundary Value Problems. Journal of Machine Learning, 2(1):31-64, March 2023. ISSN 2790-203X.
Grigorios A. Pavliotis. Stochustic Processes and Applications: Diffussion Processes, the Fokker- Planck and Langevin Equations, volume of Texts in Applied Mathematics. Springer. New York, NY, 2014. ISBN 978- 1--199---332---- 978-1--9933 11--9999333-33
doi: 10.1007/978- 1-4939-1322
Huyên Pham. Continuous- Timee Stochastic Control and Optimization with Financial Applications, volume Stochastic Modelling and Applied Probability. Springer, Berlin, Heidelberg, 2009. ISBN 978-3-540-8944
000-3--44--95000--0 doi: 10. 111000/977833554-85550000
M. Raissi, P. Perdikaris, and G.E. Karniadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational Physics, 707, 2019. ISSN 0021-9991. doi: https://
Nool.org/10.1011611ppp22201133 20.04 URL /tt/ttts:://wwwtttt http://wwwwwweeieeeeeeeeeeee com/ wwtp/e/tt.o.o/ccce.neeeeeeeeeeee science/
art Ptticle/p1i//00000 9391171111
M. Reed and B. Simon. IV: Analysis of Operators. Methods of Modern Mathematical Physics. Elsevier Science, 1978. ISBN 97801 25850049.
Michael Reed and Barry Simon. II: Fourier analysis, selffaddoonttesss volume 2. Elsevier, 1975.
Michael Reed and Barry Simon. Methods of modern mathematical physics: Functional analysis, volume 1. Gulf Professional Publishing, 1980.
Michael Renardy and Robert C Rogers. An introduction to partial differential equations, volume 13. Springer Science Business Media, 2006.
Gareth O. Roberts and Jeffrey S. Rosenthal. Optimal Scaling of Discrete Approximations Langevin Diffusions. Journal of the Royal Statistical Society. Series B (Statistical Methodology). 60(1): 255-268, 1998. ISSN 1369-7412.
Christof Schütte, Stefanie Winkelmann, and Carsten Hartmann. Optimal control of molecular dynamics using Markov state models. Mathematical Programming, 134(1):259-22) 222, August 2012. ISSN 1436-4646. doi: 10.1007/510107-012--5552
Emanuel Todorov. Eigenfunction approximation methods for linearly-solvatee solvable optimal control problems. In 2009 IEEE Symposiun on Adaptive Dynamic Programming and Reinforcement Learning, pp. 161-168, March 2009. doi: 10. 20.1109/ADPRL.20094222222
Surya Tokdar and Robert E. Kass. Importance sampling: review. WIREs Computational Statistics, 2(1):54-60, 2010. ISSN 1939-0068, doi: 10.1002/wics.56.
Claudia Totzeck and René Pinnau. Space mapping-based receding horizon control for stochastic interacting particle systems: Dogs herding sheep. Journal of Mathematics in Industry, 10(1):11, April 2020. ISSN 2190-5983, doi: 10111///133222200-00000
A Tychonoff. Théorèmes d' d'unicité pour l'équation de la chaleur. Sh. Math., 42(2):199-216,, 1935.
Ramon van Handel. Stochastic Calculus, Filtering and Stochastic Control, Lecture Notes for ACM 217: Advanced Topics in Stochastic Analysis, Princeton. 2007.
14

---

Published as a conference paper at ICLR 2026
Pauli Virtanen, Ralf Gommers, Travis E. Oliphant, Matt Haberland, Tyler Reddy, David Coumapeau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, Stéfan J. van der Walt, Matthew Brett, Joshua Wilson, K. Jarrod Millman, Nikolay Mayorov, Andrew R. J. Nelson, Eric Jones, Robert Kern, Eric Larson, C J Carey, Ilhan Polat, Yu Feng, Fric W. Moore, Jake VanderPlas, Denis Laxalde, Josef Perktold, Robert Cimrman, Ian Henriksen, R. A. Quintero, Charles R. Harris, Anne M. Archibald, António H. Ribeiro, Fabian Pedregosa, Paul van Mulbregt, and SciPy 1.0 Contributors, SciPy 1.0: Fundamental Algorithms for Scientific Computing Python. Nature Methods, 11:261-2772 2020. doi: 10. (0038//41592001--06666
Jiongmin Yong and Xun Yu Zhou. Stochastic Controls. Springer, New York, NY, 1999. ISBN doi: 111007/978---4612-11 1466-3.
978-1-4612-7154-3 978-1-4612-1466
Sen Zhang, Jian Zu, and Jingqi Zhang. Deep learning method for finding eigenpairs in Sturm- Liouville eigenvalue problems. Electronic Journal of Differential Equations, 2024(01-83):53-17, September 2024. ISSN 1072-6691. doi: 10.58997/ejdde22224.5..
Wei Zhang, Tiejun Li. and Christof Schütte. Solving eigenvalue PDEs of metastable diffusion processes using artificial neural networks. Journal of Computational Physics, 465, 2022. doi: 2022.
Xun Yu Zhou, Jiongmin Yong, and Xunjing Li. Stochastic verification theorems within the framework of viscosity solutions. SIAM Journal on Control and Optimization, 35(1):243-25 223,,,3,,555 1997.
15

---

Published as a conference paper at ICLR 2026
CONTENTS
Infroduction
2
Prelininaries
2.1
Stochastic optimal control
2.2
Eigenfunction solutions
3
Our framework
3.1
Spectral properties of the Schrödinger operator
3.2
Eigenfunction control
3.3
Closed-form solution for the symmetric LQR
4
Numerical methods
4.1
Learning cigenfunctions
4.1.1
Implicit reweighting in previous approaches
4.1.2
Our approach: Removing implicit reweighting relative loss
4.2
Our hybrid method: combining eigenfunctions and noot-horizon solvers
8
5
Experiments
5.1
Opinion modeling: de Groot model
10
6
Conclusion
10
Technical details/sssumption
17
A.1
Regularity conditions
17
A.2
Simplifying assumption: oo
17
Proofs/derivations
17
B.1
Cole-Hopf transformation 8
17
B.2
Proof Theorem
17
B.3
Unitary equivalence (11)
19
B.4
Essential self-adjointness
19
B.5
Proof of Theorem 2
19
B.6
Proof of Theorem 4
20
B.7
Semigroup and Viscosity solutions
C
Loss functions
24
C.1
Existing methods for short-horizon SOC
24
C.2
Extending to multiple eigenfunctions
24
C,3
Estimating X
25
C,4
Empirical loss & sampling
25
C,5
Logarithmic regularization
26
C.6
FBSDE Loss
26
7
IDO loss functions
27
Algorithms
29
Experimental Details
30
E.
Setting configurations
30
E.2
Model architecture and training
E,3
Computational cost]
32
E.4
Further details the RING setting
32
E,5
Details on Figure
32
Additional experiment: repulsive potential
E.7
Control objectives
33
16

---

Published as a conference paper at ICLR 2026
A TECHNICAL DETAILS/ASSUMPTION
A.1 REGULARITY CONDITIONS
Following Fleming Soner (2006), we make the following assumptions, which guarantee that the SDE 3 has unique strong solution.
1. The coefficient b is Lipschitz in and satisfies the linear growth condition
CC > 0: [|((x) b(y))l || - C/| - yll-
(25)
30" 0: [b(x))] - C²(1 + |
(26)
A.2 SIMPLIFYING ASSUMPTION: To-== =
Since o E Rdddd was assumed invertible, the matrix oo is positive definite, and hence there exists diagonalization ooT UAUT where UUT I and A diag((A,) for X 0. Consider now the change of variables y A 11111112 x.
Then it holds that
V レ = UA 1/2 V V
Vay = UA 1/2 V d A 1/20T
(27)
and in particular Tr( T v2 Tr(UAUT Tp( Ayt.. Thus the PDE 8 can be written in terms of y as
Oes (-4, - 2B 1/22,, + +, 1) 3 -=
(28)
B PROOFS/DERIVATIONS
B.1 COLE-HOPE TRANSFORMATION
Let V denote the value function defined in satisfying the HJB equation 6 The so-called Cole-
Hopf transformation consists of setting V -8-1 log v. leading to linear PDE for 3. which is
called the desirability function. To obtain the form 8 we set Vax, () - log v (ar, 3 (T - 1))
The derivatives of V and eh are then related through
8eb
VM
0.V 2
VV
-
(29)
e av Orb Ox, 0x,
8V 8-1 dz,Ox,
))
(30)
2
Ox,
Tr(ao V2V)= = Σ o dx,0r,
02V
-8-1 Tr(oo T (22)) + 3
1 v
(31)
i.j
Plugging these expressions into (6) gives
4.b
Tr(ao TV()))))))))))) - - + 1-=10 = 0.
(32)
2.32
232
Multiplying with 2/ and combining with the terminal condition V(z, T) g(x) shows that the linear PDE for 1p given in 8) is equivalent to the HJB equation (
B.2 PROOF OF HEORRE
We begin with an informal argument. Recall that we are interested in solving the abstract evolution equation
Find :b [0.T] - D(C) such that (aw(()) + Lip(t) = 0.
(33)
(b)
= 4/0-
17

---

Published as a conference paper at ICLR 2026
Now, since (Ø:)ee forms orthonormal basis, it holds that eb(t) EvEn a,(t)o, where a;(t)
(W)(), Hence the PDE in 33
becomes
ddi(() dt
+ +aaa((t) 0i
(34)
KEN
after formally interchanging the series expansion and derivatives, Since the d; are orthogonal, this equation satisfied and only a;( = a,(0) Aut for each i.
Formally establishing (9) can be achieved through the theory of semigroups. Essentially, we want to formally define the semigroup (e - (20 and show that forms the solution operator to 33 We
begin by recalling the following definition.
Definition 2 The set of bounded linear operators H. is given by
B(H) L: 7 + H: C is linear and ||
sup fEH.
|C()] < 00
(35)
The following result, sometimes called the functional calculus form of the spectral theorem, allows us to define (()) for bounded Borel functions h.
Lemma
RRed & Simon 1980 Theorem VIII.5) Suppose L is self-adjoint operator H. Then
there exists unique map 6 from the bounded Borel functions IR into B(H) which satisfies
If Cub Xt/, then @(h) h(A)yb.
In particular, when admits countable orthonormal basis of eigenfunctions (ø:) with eigenvalues A this operator is given by
8(h)²= hC))= = Σ (o1,
(36)
VEN
The next result makes use the defined semigroup and leads to the desired representation. Note that we indeed consider bounded function of the operator, consider it only for t > 0 and assume that the operator bounded from below.
Theorem &Reed & Simon 1980
Theorem VIII.7) Suppose L is self-adjoint and bounded from below, and define T(t) = é a for 2 0. Then
(a) T(0) I.
(b) For every & E D(C), it holds that
( j dt
T(O)))
T(0 + h)m - T(O)
lim h
CTT(0) = -Lub
(37)
h
(==
Proof Theorem 1 Combining (a) and (b) of the Theorem 5 shows that (5()) = T(t)/o satisfies
d T(t + hyy0o
- (()) dt
= -T(() (L)()), (00) = T(0) tho
(38)
dh.
h=0
which is exactly the claim of Theorem The last remark we make is that in Theorem we only assume essential self-adjointness of L. while above results operate with self-adjoint operators. Thus, the above theorems hold true for the closure of the operator L. However, note that densely defined and has eigenfunctions which make up orthonormal basis of both its domain and L (p)) Thus, the representation (9) for Z. holds for L.
18

---

Published as a conference paper at ICLR 2026
B.3 UNITARY EQUIVALENCE
First, notice that U : (p) + (2R())) + tHe B is indeed a unitary transformation, since
V E (u): UU),U)) = /
23E sipp dx = -(b, (.
(39)
To establish equivalence, we compute
v(e b) = cB (V 3VE v), A(e +) =e (Av + V Ve) + (BAE+ (IVE))) v)
(40)
(41)
V ((e) e2 V - -
(42)
Putting this together gives
C(U- B() B (-A0) 3 Pル
BAE v) fE
(43)
from which the result (TT) follows.
B.4 ESSENTIAL SELF-ADIOINTNESS OF
We will first show the following relation
(p, Lb) (Vi, Ve) + 2, fv))
(44)
from which clear that symmetric on C/ (Rd). Indeed, using the divergence theorem, for v,, & E Co (R²) one obtains that
(w--)pp = / e 29
e dx
(45)
- ((VV) Vp) (V. VE))
28F Fr
(46)
(Ve. Vp) - 23(6, (VE, Vo))
(47)
which immediately shows the result.
While for matrices the notions of symmetry and self-adjointness are equivalent, the situation becomes more delicate for general (possibly unbounded) linear operators. In our case we can use the following result on the essential adjoitessssssss of the Schrödinger operator.
Lemma 2 RReed & Simon 1975
Theorem X.28) Let V (Rd) with V 2 0 pointwise. Then = - + V is essentially self-adjoint Cg (B²).
Note, that C (IR) is dense in (TRd). Thus, operator L is also densely defined as its domain contains U- UC (R))).
The essential self-adjointness of L follows from the unitary equivalence with the Schrödinger operator as unitary transformation preserves essential self-adjointness of the transformed operator.
B.5 PROOF OF [HFFOREM 2
As we have shown, the operator L is unitarily equivalent to the Schrödinger operator defined through
S D(S) C (²RRR) + (RR) + -Ah + Vh
(48)
with V = 3RIVEV 3&E + 232 f. All the propertics mentioned in the statement of Theorem remain unchanged under unitary equivalence (both the assumptions the operator and the conclu- sions regarding the spectral properties). Hence is sufficient to show the result for the Schrödinger operator S associated with L.
This is precisely the content of the following theorem, proven in Reed & Simon 1978), combined with observations in section B.4
19

---

Published as a conference paper at ICLR 2026
Theorem 6 <Reed & Simon 1978 Theorem X111.67. X111.64, XIII.47) Suppose V E L (R is bounded from below and || V(z) = Do. Then the Schrodinger operator S = - + has compact resolvent, and in particular purely discrete spectrum and orthonormal basis of eigenfunctions in I.2 (R2). The spectrum (SS) is bounded from below and the eigenvalues do not have finite accumulation point. If in addition V T (Re4), the smallest eigenvalue has multiplicity one, and the corresponding eigenfunction can be taken to be strictly positive.
B.6 PROOF OF HEFORMM 4
The result relies on the fact that the Schrödinger operator associated with L has a quadratic potential, for which the eigenfunctions can be computed exactly. Indeed, the following is standard result that can be found in many textbooks on quantum mechanics (Griffiths & Schroeter 2018):
Lemma Consider the Schrödinger operator in d = 1 with quadratic potential,
d²ab
S D(S) C (R) - (R) db
+ 124..
(49)
dr.2
The eigenvalues of S are given by An = 2n + 1, and the normalized eigenfunctions are given by
Φn (n(x)
1
H (x)e--0111
(50)
π
2
where H., denotes the n-th physicist's Hermite polynomial defined through
Ho(x)= 1, H1(x) 2x, Ha+1(x) = (x), 2n (z)).
(51)
These satisfy the generating function relation
m H n!
2xt
(52)
n=0
This result can be easily extended to higher dimensions:
Lemma Consider now the d-dimensional Harmonic oscillator:
S: D(S) L'( + - + IT Ax tb
(53)
where A Bddd is symmetric and positive definite. Denole A U AU the diagonalization of A. Then the eigenvalues are indexed by the mulii-index a E Nd, and are given by
d
Na
Σ
A
1/2 (2a4 + 1).
(54)
i=1
The associated normalized eigenfunctions are
A
(1()) dd/1 exp
II
H. ((A3/A
(55)
2 (a.)!
Proof. We begin by introducing the variable
A A/. This gives
xx Axr 5² 1/22 A Tr(A!/2V3
(56)
(57)
so that we can write
- A + Ardh = Σ (-
d
82 + v2
1/2
(- oy²
(58)
Thus, the d-dimensional harmonic oscillator decouples into d rescaled one-dimensional oscillators. Since (Rd) z 8 T,2 T.2(RR) (after completion), this means that the eigenfunctions given by the
20

---

Published as a conference paper at ICLR 2026
product of the one-dimensional eigenfunctions. Hence the spectrum is indexed by No and the eigenfunctions are, expressed in the variable y,
(2a!(o.) Ha, v2
oa(y) T4/4
(59)
To get the eigenfunctions in the variable, transform back and note that the above expression is normalized as (y) = 1, while we need f d (2) = 1. Using the fact that dy (et(A this yields the final normalized eigenfunctions
A A//8
(x)x) / exp (- U D
II
Ho. ((A))
(60)
(a.))
and the eigenvalues as
Aa= Σ A 2 (2a; + 1).
(61)
Now consider the operator L as defined in 10P.
with b(z) - Az and f(z) Pr. Since A is
Pr.
symmetric, it holds that VE with E(x) ) 2 Ar. follows that is unitarily equivalent with the Schrödinger operator with potential
V = = V 23 = B2 ( 2P)..
(62)
The first term gives a constant shift to the eigenvalues, but otherwise we are precisely dealing with the d-dimensional harmonic oscillator, whose eigensystem is described in Lemma Hence the cigenvalues are precisely (14), and the eigenfunctions of the original operator C are then obtained by multiplying with e a" As which yields 13
B.7 SEMIGROUP AND VISCOSITY SOLUTIONS
Plan of the section. We work with two linked equations on IR x [0, T): the HJB equation (6) and the linear parabolic PDE (8) obtained from (6) via the (monotone) Cole- Hopf transform. Our goal is to produce the optimal control u" by first solving the linear PDE, then applying the inverse Cole Hopf transform to obtain value function candidate V'. This program raises two issues which have to be addressed:
1. Nonuniqueness in unbounded domains. On TRd (even with the same terminal/houndary data), second-order finite-horizon HJB and linear parabolic equations may admit multiple solutions unless one restricts to an appropriate growth class; see Tychonoff (935) for classical example of this phenomenon.
2. Verification. To identify u" from solution of 66 one uses verification theorem. These theorems require the candidate V to be a (sufficiently regular) viscosity solution (Crandall
et al. 1992) lying in class where comparison (hence uniqueness) holds, (Fleming & Soner 2006 Section V.9) for an example of such result.
In our approach we fix a specific semigroup solution of the linear PDE and use it to define V via the inverse Cole-Hopf map. Specifically, have V' I) 3 log sh(2,, y 1)), where (zr f) = (e (x) The central task is therefore to (a) place this V' in a uniqueness class for (6) and (b) confirm the hypotheses of verification theorem. In the following text we establish sufficient conditions under which the task is solved.
Verification theorem and Comparison (uniqueness) d under growth constraints. There are several ways establish that the value function V (5) is a unique solution of HJB equation (6) in some growth class for the given terminal conditions. In general it is obtained with a use of Dynamic Programming Principle. However, this approach is not simple in the case of unbounded controls and non-smooth viscosity solutions. We refer to (Zhou et al. 1997 Fleming & Soner 2006]] Da Lio
Ley 2006 2011) for overview of ways to establish connection between value function and
21

---

Published as a conference paper at ICLR 2026
viscosity solutions, well as results regarding existence and uniqueness of viscosity solutions in appropriate growth classes. Following Da Lio & Ley (201 1) we introduce class of functions Cp. We say that locally bounded function R
[0: 7] R is in the class Cp if for some C' we have
|u(2, t) & C(1 + I)) A (x, E x [0:T]
(63)
Now, we can formulate the following
Theorem (Da Lio & Ley 2006 Thm. 3.1) D. Lio & Ley 2011 Thm. 3.2) Suppose that in (6)
we have
satisfies
25 and 26 and 0 =
Running cost f satisfies: 3C1 > 00:0 E Bd |/(>)] < C1 + [|()P)
Terminal condition satisfies y E EC,
Then there exists 0 I r < 7 such that (6) has a unique continuous viscosity solution in IRd x [r: in the class Cp. Moreover, this unique solution is the value function V (5)
Remark 3 See discussion in (Da Lio & Ley 2011 Remark 3.1) l) why we can't 't hope to obtain existence Theorem 7 on a whole interval 0; T-
Growth of semigroup solution From previous paragraphs it is clear that way to provide a link
between proposed V' and solutions of (6 to to establish growth rates on V Specifically, following
Theorem we want to show that V E Cy-
There other ways to establish connection between solutions, see Biton 2001), ) FFemmnng & Soner 2006), Section VI). However, we choose to pursue a path related to growth conditions of spectral elements of Schrödinger operator (Davies & Simon 1984 Baraniewicz 2024)
Theorem (Davies & Simon 1984 Thm. 6.1. Thm. 6.3) Let S - + Fd with ground state 6o > 0 (normalized in (²( and suppose that there exist C1, C3 > 0 and C22, C1 E R such that, for all x with ||=||| large,
Call=11 +C4 & V(x) < C1|l=11 + C2,
where + <bb
Then S intrinsically ultracontractive (IUC). Moreover, there exist constants Co C7 > and Co, Cg IR such that, as ||~||| co,
Os |||||| ++ + Co < - log (x() S C. |/|1=|| Cs.
(64)
Proof. Intrinsic ultracontractivity under the stated growth exactly DDvvies & Simon 1984
Thm. 6.3). The upper growth bound in (64) is the estimate - log (x(
C| * stated explicitly
in (Davies & Simon
1984 Eq. (6.4). For the lower bound in (64). take a comparator W(x) = ||]||| (W)
with 0 < à Cz: then W + and V W co, the ground states satisfy e
(V) < C Co/
by the comparison Lemma (Davies & Simon 1984 Lem. 6.2). (Davies & Simon 1984 App. B) (W)
constructs WKB-type barriers for - + W and gives pointwise upper bounds of the form d (()) <
C'||=||- exp(-<||a=] 1+$] for large ||^||| (see the JWKB ansatz and subharmonic comparison in
Davies & Simon 1984 App. B, esp. Lem. B.1-B.3)); combining yields the stated lower bound on
log 9n-
Remark Note that we use ultracontractivity properties of which require growth of V to be at least as |2+e & 0. The characterization of operator contractivity properties is fully given in (Davies Simon 1984 Thm. 6.1) and does not allow for slower orders of growth. This restriction on growth of is encoded through condition 4 +
Theorem Let
== 2
AE +35,
8 UCU--=-A- -
22

---

Published as a conference paper at ICLR 2026
== where f= R. Let txx (e (z) and define the candidate value function
(x - log v
23 (T (t).
Assume the hypotheses of Theorem for V. Finally, let suppose that there exist 0 < m S M so that
do
S M, Var E Rd 40
(65)
where do is a ground state of Shrödinger operator S. Then, uniformly for all t [0,T], there exist constants C1. C3 and Cz. C4 E IR such that, as || + 00,
C ||]|| ++ C2 s ('(x, F(() C3 ||0||| 5+1 C4-
(66)
Proof. Recall that the ground state of L is an eigenfunction Po = U Φn, where Φn is ground state of Shrödinger operator S. Let us introduce the ground state transformed semigroup:
cdoe
p. g(x) = po(x) (e So-tQ
(e (()
Thus, we have representation
(x,, 1) = (e (x)) -
dot po(()) pi00
(x))
PP
It is known, that semigroup p: 9(x) is Markov, see Kaleta 2018). Using the fact that P. gxx) is Markov and (65) we have the two-sided bound for all E [0:
Aot (o()) (e ())) Aot Po())
which leads to bound
- - B log M - -3-1 log do F(x) S V'(x;)) Ao
2,32
T - 1
Ao
B log m 33-1-ogg log Φ0
F(x))
(67)
2,32
Applying growth bounds on ground state of Shrödinger operator from Theorem leads to the desired bound
Remark Note that definition of C, (63) requires a growth bound on coordinate æ uniformly for all t e [] Unfortunately. analysis based purely on spectral properties of operator leads to non-uniform growth bounds depending on I, which blow up when f. - 0 for et (for example, note that constants in (Davies & Simon 1984 Thm. 3.2) depend on time). For this reason and for simplicity we introduced assumption (65) on boundary conditions.
Relationship between semigroup- and viscosity solution Finalising everything in this section, we formulate the following result which addresses the questions raised above.
Theorem 10 Consider HJB equation
6 Let V = 3 AE + 28f. Assume the following:
b(x) = -V(()) satisfies and 26). and o I.
V satisfies Assumption (A2) for all with ||a||| large,
Moreover there exist C1, C3 > and Cz, C4 E R such that,
Call=|| + CC s V(x) s C1ll=|| + Cz2 where a + ++lbb (68)
Let on be the ground state of operator S = - + Assume E č2 and there exist
< m S M so that
exp(-Bg) < M, Vr E Rd exp(3E(z))dn
m <
(69)
23

---

Published as a conference paper at ICLR 2026
Then there exists 0 < T such that function (1 - log (2, 28 2T t)), where (ax t) = (())) and L and tho are from ( is a unique continuous viscosity solution of (6) on Rd x [T: in the class Cu.2+1- Moreover, it coincides with value function V
on R x [;;
Proof. Under Assumption (A2) Theorem 6 holds. Theorem Gensuress that operators S and have some good properties and ground state do is positive.
Under the stated assumptions Theorem and Sercoremhholee Under growth assumptions (26) and (68) Theorem gives us that V' E CC,/2-1 S ć, which has to be the unique viscosity solution in C, given by Theorem On on Rd x [r; T), which coincides with value function
on Rd x [T: T.: M
Remark 6 Note that assumption (69)) is realistic under rapid growth of V, taking in account growth bounds on the ground state (64 and b(x) = VE(x) in (2)) and does not contradict C..
C LOSS FUNCTIONS
C.I EXISTING METHODS FOR SHORT-HORIZON SOC
Grid-based solvers In low dimensions (d < 3). classical techniques for numerically solving PDEs can be used. These include finite difference [Bonnans et al., 2004) and finite element methods (Jensen
& Smears, 2013; Ern & Guermond, 2004 Brenner & Scott 2008), as well as semi- cemi-Langrangian
schemes (Calzola et al., 2023] Carlini et al. 2020) and multi- Picard iteration (E et al. 2021
FBSDE solvers In another line of work, the SOC problem is transformed into pair of forward- backward SDEs (FBSDEs). These are solved through dynamic programming (Gobet et al. 2005 Longstaff & Schwartz. 2001 deep learning methods which parametrize the solution to the FBSDE
using a neural network (Han et al. 2018 E et al. 2017 Andersson et al. 2023).
IDO methods In recent years, many methods have been proposed which parametrize the control us directly, and optimize it by rolling out simulations of the system (3) under the current control. Authors of Nüsken & Richter 2021) coined the term iterative diffusion optimization (IDO), arguing that many of these methods can be viewed from a common perspective given in Algorithm [ (Appendix E). This class of algorithms contains state-of-the art methods such as SOCM and adjoint matching (Holdijk et al. [2023] Comingo-Ennich et al. 2024a). We describe the most commonly used loss functions in
Appendix C
C.2 EXTENDING TO MULTIPLE RIGENFUNCTIONS
PINN loss
The most common way to extend the PINN loss 15 to multiple cigenfunctions is to
define
k
R (0) (1]C(o.) -x00i13 + anorm (1 1)2) + Coorth (or
(70)
for Onnormy Corth O. Here we have denoted & Rd RR+11 The main difference is the addition of the orthogonal regularization term, which both ensures that the same eigenfunction is not learned twice, and attempts to speed up learning by enforcing the known property that eigenfunctions with different eigenvalues are orthogonal w.r.t. the inner product (s)))u
Variational loss A similar idea is used to generalize the variational loss 17 shows that the variational principle (16) can be extended to multiple eigenfunctions by imposing orthogonality.
The following result
Theorem 11 (hang et al. 2022 Theorem 1) Let l E N, and let L be self-adjoint operator with discrete spectrum which admits an orthonormal basis of eigenfunctions. and whose eigenvalues are bounded below. Furthermore, let wh 2 2 wk > 0 be real numbers. Then it holds that
Σ wiii = inf Σ wilfi, 1=0
(71)
fa farnJ
1=0
24

---

Published as a conference paper at ICLR 2026
where the infimum is taken over all (fi) H such that
Vi,j E {0....,k) k) (fr.j) dij
(72)
Proof. The proof of this result is given in Zhang et al. (2022).
Based on this result, the following generalization of the variational loss is proposed in Zhang et al. (2022):
Rvvor (6)
Σ (d., Cot) + a By [Ø0"] - - 1=0
(73)
where have written ||e for the Frobenius norm and 0. This loss also studied in Cabannes et al. (2023), where it was noted that the minimizers of the variational loss (17) are not obtained the eigenfunctions. Instead, the following characterization of the minimizers of was obtained.
Lemma
Cabannes et al.
2023 Lemma 2) Suppose that H = (1s).. Then it holds that
U UU = I, v 11-2
A
arg min Real (16) Ux.E H
(74)
2a
where (dy, Aa) denotes the orthonormal eigensystem of L.
C.3 ESTIMATING A;
The main advantage of the variational loss (17). (73) is that it does not require the eigenvalues of L to be available. However, Lemma shows that the minimizers of R do not coincide exactly with the eigenfunctions of L, so we cannot naively compute (ò. L to obtain the eigenvalues. Instead, the following lemma shows how to obtain the eigenvalues and eigenfunctions from an element in the minimizing set
Lemma Suppose it holds that Uò, where UUT = and de V
(1
2a oz oo each and
Then the first le + eigenfunctions and eigenvalues are given by
D= 1 v.
A 1- D1)
(75)
where UDUT = [bbb" " is the diagonalization of the second moment matrix of th.
Proof. By definition of v, we have
E EA UE UE
16021072
(76)
-1
11
=
2a
U
(77)
where A diag(A,) and have used the orthonormality property of the eigenfunctions. From this, we obtain that
SAA A, == UD!/26
D D=
(78)
which concludes the proof.
In light of this result, we can estimate the second moment matrix [.(d"" apply a diagonalization algorithm, and obtain the eigenfunctions and eigenvalues using (75)
C.4 EMPIRICAL LOSS & SAMPLING
Rewriting variational loss Recalling the equation (44)
Lx)), = (V V + 2/3 fbb)
(79)
we see that it possible to evaluate inner products the form (4p ,, by only evaluating yo, 1/ and its derivatives. This avoids the expensive computation of the second order derivatives of the neural network, making the variational losses less memory-intensee than the other loss functions, which requires explicit computation of Lo.
25

---

Published as a conference paper at ICLR 2026
Estimation of inner products All of the loss functions discussed for learning eigenfunctions contain inner products of the form (e), o)n- To obtain an empirical loss, these quantities most be approximated. When j is density, this done using a Monte Carlo estimate
20 = / pp dp = Eppppl ~
1 Σ (X.)) X
(80)
Since u(x) x exp( (z)),, we can employ Markov Chain Monte Carlo (MCMC) techniques in order to obtain the samples (X;) (Brooks et al. 2011). In particular, when training the eigenfunctions, store samples in memory and apply some number of MCMC steps in each iteration to update these samples. Alternatively, one can pre-sample a large dataset of samples from ja.
Non-confining energy. In general, l may not be finite measure, for instance when have repulsive LQR (E- = =- 1 | In this case, the inner products can no longer be estimated directly from samples, but may still be computed by using importance sampling techniques (Liu) 2004 Tokdar
& Kass 2010). In the case where the measure defined through μ(ι²) exp( 22 is not finite but ( exp((3E(())) is, we can apply importance sampling with μu, so that the inner products
are obtained as
(s,, 1) = gope -2F((2) dr = Eg (E
(81)
where we have defined E 2B op and vb e e-28E 1b. For stable training, it is then advisable to parametrize Φ == -e- 28E instead of øb
C.5 LOGARITHMIC REGULARIZATION
As discussed in the main text, the PINN and relative eigenfunction losses (15 20 have the form R()) Rmain (Ø) (o)r Roy (ll 1)
(82)
for some O. We observe in experiments that this regularizer is sometimes not strong enough, and the network may still converge to 6 = 0. For this reason, we instead use a logarithmic regularizer
Reeg (6) = (log (1+1)))
(83)
The behaviour is exactly the same as before for |/0| ~ 1, since log(1 + x) = x x Oz²²) as 0 0. but this regularizer avoids the convergence 0.
C.6 FBSDE LOss
We briefly discuss the Robust FBSDE loss for stochastic optimal control introduced in Andersson al. (202). The main idea that the solution to the HJB equation 6 can be written down as a pair of SDEs, as the following lemma illustrates.
Lemma Suppose 1. Then holds that the solution to the pair of FBSDEs
dX, (b(X,,t) t) - Zt) + VXdW,,
Xo ~ Po.
(84)
dY, (X(,), - 111 11 dt + vx(Z, dWW)
Yr = 9(XT)
(85)
is given by 7, (X...) and Y, = (X...t)
This pair of SDEs can be transformed in variational formulation which is amenable to deep learning. as described in detail in Andersson et al. (003). Using the Markov property of the FBSDE, one can show that Z, = {(t, X,) for some function Ç [0. 7) Rdd - Rd. and that the FBSDE problem can be reformulated as the following variational problem:
minimize (C) = E|S1 + aE [I 9(x
where
2 =x(55) + 1d
+ di - (25 .WW)
x = Xo + 1 (b(s, X5) -5))=()dd d/ f/ AdW,
Y ED]- (f(X).s) ds (2 aW.).
<(t(,)))); tE [0.
(86)
26

---

Published as a conference paper at ICLR 2026
Parametrizing neural network, then simulate the stochastic processes in (86) and define the loss function as
(EBSDE(U; X*) E oB | o(x5)]]
with Ç
(87)
For more details refer to the relevant work Andersson et al. (2023).
C.7 IDO LOSS FUNCTIONS
The IDO algorithm described in Algorithm Appendix D is rather general, and large number of algorithms can be obtained by specifying different loss functions. For detailed discussion of the various loss functions and the relations between the resulting algorithms, we refer to Domingo-Eiin Enrich (2024); Domingo-Enrich et al. 2024b). Here, we go over the loss functions that were used for the experiments in this paper.
Adjoint loss and define
The most straightforward choice of loss to simply use the objective of the control,
(" 0112 + s(x?) - g(x)).
R(u; x*)
(88)
This loss is also called the relative entropy loss. When converting this to an empirical loss, there are two options. The discrete adjoint method consists first discretizing the objective and then differentiating w.r.t. the parameters. However, this requires storing the numerical solver in memory and can hence be quite inteesv..
The continuous adjoint method instead analytically computes derivatives w.r.. the state space. To this end, define the adjoint state as
alt, u) = Vx. (1,( ( (
at' + *((x))
(89)
where X solves
3 Then the dynamics of can be computed as
da. = (V x (b(X,, () + (1)((X)) (o)x(()))) 1))² a(t; X", u))d Vxp (f(XX, t) + 7 11((x t)12 dt,
+
(90)
aT;;",, u) Vg(x
(91)
The path of a obtained by solving the above equations backwards in time given a trajectory (X!i) and control u. Finally, the derivative of the relative entropy loss (88 is then computed
1 {'
a lll(X 00
X di (."
()
0.R
(I)a((; X,, ū)d(,
(92)
00
where the notation means that we do stop gradients w.r.t. 0 from flowing through these values. Both the discrete and continuous adjoint method have been shown to work well in practice (Nüsken Richter 2021 Bertsekas & Shreve, 1996; Domingo-Enrien 2024). However, their training can be unstable due to the non non-convexity of the problem.
Variance and log-variance For the second class of loss functions, (X²) denote the solution to 3 with и replaced by v. Then we can define
s)ds- f'' f(X;, s)ds Lo s( /" s) (93)
Vuuv -- (u. s)ds
The variance and log-vate tog-variance loss are then defined as
Rar (u) = Var(e Prt -9(=)))())) Rlog-var var (u) Varl" g(X))).
(94)
(95)
It can be shown that these losses minimized when u", irrespective of the choice of In Nüsken & Richter (202I). it was shown that these loss functions are closely connected to the FBSDE formulation of the SOC problem.
27

---

Published as a conference paper at ICLR 2026
SOCM The Stochastic Optimal Control Matching (SOCM) loss, introduced in Domingo--niiinn Enrich et al. (2024b). is one of the most recently introduced IDO losses. It is described in the following theorem, where we adapt the notation À
8-1:
Theorem Domingo-Enriice el al.
2024b Theorem 1) For each E [0, T. let M, : [4, T] + Rdxd be an arbitrary matrix- valued differentiable function such that (t) Id. Let v E U be arbitrary control. Let RsOCM : L (2 x [0,T]; Rd) x (2,, Rdxd) + be the loss function defined as
|u(X; t) wtt, 8, X", W. W,)l dt x a(e, m) T /
RsOCM (u, M) == E
(96)
where X" the process controlled by v (i.e., dx.- (6(x), () (XF) dt) vAo(t)dW and Xg ~ po), and
uft,v,X X", W, = o(t) (-1," -dd
+ +f. (M1(S)V ( s)ds
+utx S: (M.(S)V O, (o)) 10 (s)dW,
a(v,X", = exp (-x-1 1-" .)dt
-1 -
--1 - 112 1/2 ( (U(xy,0), 2 L'd [lu(x;, 1)
(97)
LSOCM has a unique optimum M,), where u' is the optimal control.
The proof hinges the path integral representation described in Kappen (005a) and using reparametrization trick to compute its gradients. We refer the reader to the relevant work Domingo Enrich et al. (2024b) for more details. The reason why this method outperforms other IDO losses is as follows: the loss (96) can be seen as minimizing the discrepancy between u and a target vector field The addition of the parametrized matrix M allows the variance of this weight to be reduced, making it easier to learn. The downside of this method is that the variance of the importance weight a can blow up in more complex settings with poor initialization, with the method failing to converge as result.
Adjoint Matching The most recently introduced IDO loss, adjoint matching, was proposed in Domingo-Enridh et al. (2024a), and was proposed in the context of finetuning diffusion models. It is based on two observations. Firstly, one can write down a regression objective that does not have an importance weighting by using the adjoint state a defined earlier.
Lemma (Domingo-Enrich et al. 2024a Proposition 2) Define the basic adjoint matching objective as
R Adt-Match X2; = 2 f' [lu(X,, + a(t) X", dt. ū stopgrad(())
(98)
where ã = stopgrad(u) means that the gradients of ū w.r.t. the parameters 0 of the control are artificially set to zero. Then the gradient of this loss w.r.t 0 equal (2), the gradient of the loss in the continuous adjoint method. Consequently, the only critical point of Ex"=E the optimal control u..
Ra (
Adj M Match)
Secondly, it is observed that some terms in the SDE for the adjoint state 90 have expectation zero under the trajectories of the optimal control. Indeed, holds by definition of the value function and adjoint state that
VV () Ela(t, x= ,u') X, 2]]
(99)
28

---

Published as a conference paper at ICLR 2026
and hence, since u" OVV, get
Ex-p= [u* (x, t)" V (x, + a(t, X, u*)² (x,t) 1 X, = x] =
(100)
This motivates dropping the terms with expectation zero in (90) yielding the loss function
L. [Iu(X,,,) (t) ||
Add Match (u; X)
dt. stopgrad(u),
(101)
d where a((t:X) -iā(t; x b(,, (X,,
(102)
dt
X) V
(103)
The value ā is called the lean adjoint state. The claimed benefit of this method is that the resulting loss is simple least-squares regression objective with no importance weighting, allowing avoid the problem of high variance importance weights. We refer to the original work [Domingo-Enrec) al. (024a) for more in-depth discussion and proofs related the adjoint matching loss.
D ALGORITHMS
We give an overview of the algorithm used for eigenfunction learning in Algorithm and for the IDO method in Algorithm] 2
Algorithm Deep learning for eigenfunctions
1: Parametrize the eigenfunctions (o ), choose loss functions (R.)-
2: Fix batch size m, number of iterations N, regularization O, learning rate r) 0.
3: Generate m samples (X;) from μu, for instance using an MCMC scheme.
4: for = N do
5: (optional) Update the samples (X,) using sampling algorithm.
6: Compute m-sample Monte-Carle estimate R.(o.)) ) of the loss (2
7: Compute the gradients Vo, (o.:) for the current parameters 0;, and update 0; using Adam.
8: end for
9: Estimate the eigenvalues A, from the learned functions of
10: return eigenfunction estimates (Ø eigenvalue estimates (As).
Algorithm 2 Iterative Diffusion Optimization (IDO)
Parametrize the control ug E (4, TR?,, and choose loss function R.
2: Fix a batch size number of iterations N, number of timesteps K, learning rate > 0. do
3: for i 0,
4: Simulate trajectories of (3) with control -g using K K-step discretization scheme.
5: Compute m-sample Monte-Carlo estimate R(ux) of the loss R(uv).
6: Compute the gradients (ug) for the current parameters and update using Adam.
7: end for
8: return the learned control Ug R u.
29

---

Published as a conference paper at ICLR 2026
E EXPERIMENTAL DETAILS
E. SETTING CONFIGURATIONS
To evaluate the learned controls, we will use three different metrics:
1. The control objective is the value we are trying to minimize,
IE (HIIN(Xx x(x)) dt + g(x))
(104)
It can estimated by simulating trajectories of B Unless mentioned otherwise, we estimate using 65536 trajectories and report the standard deviation of the estimate. 2. The control L error at time is given by
E EaPp" [1]u(2, t) u- (x, t)
(105)
where Py" denotes the density over Id induced by the trajectory under the optimal control at time
3. The average control L error is the above quantity averaged over the entire trajectory,
Ee [1] u0 (x,,x)
(106)
These are also the metrics reported in previous works Somningo-Enrich et al. (2024b); y
Domingo
Enrich (2024); Nüsken & Richter (2021). The L error introduced because, after the control sufficiently close to the optimal control u", the expectation (104) requires increasingly more Monte Carlo samples to distinguish u from the optimal control. In this case the I error more precise metric for determining how close given control is to the optimal control.
For all experiments, we trained the neural networks involved using Adam with learning rate 7) 10-4 For the IDO methods, we use a batch size of m = 64. For the eigenfunction learning, we sample from f using the Metropolis-Adjusted Langevin Algorithm (Roberts & Rosenthal] 1998). Unless mentioned otherwise, we sample m 65536 samples, updating these in every iteration with 100 MCMC steps with a timestep size of At 0.01 after a warm-up phase 1000 steps. For more details, we refer to the code in the supplementary material, which contains description of all hyperparameters used.
QUADRATIC The first setting consider has
1 TT Ax, f(z)=1-PP 9(x) xT =TQr.
E(x)
(107)
where A [g is symmetric, Q ddd is positive definite and P E Raxd This type of control problem is more widely known as the linear quadratic regulator. The optimal control is given by u (xx t) F2, where F, solves the Riecati equation (see (van Handel 2007 Theorem 6.5.1))
dF: A F1A-2F,F 2F,F PP=0, Fr= dt
(108)
We consider three different configurations:
(ISOTROPIC) We set d = 20, A 1.0. taking K 200 time discretization steps for the simulation of the SDE.
((0,0.511
(REPULSIVE) Exactly the same as isotropic, but with A -1
(ANISOTROPIC) We set d = 20, A d P U ( Q - 0.5l, 3 T = 4, xn ~ (0, 0.5Z), taking K = 200 time discretization steps for the simulation of the SDE. The values a, P; sampled i.i.d. from // 1), and the matrix U is random orthogonal matrix sampled using scipy. stats. ort group (Virtanen et al. 2020 at the start of the simulation.
30

---

Published as a conference paper at ICLR 2026
DOUBLE WELL
The second setting consider is the d-dimensional double well defined through
E(x) Σ (1 -)², f(z) Σ n2 1)², g(x) 0.
(109)
where d = 10, and K2 = =, for i = 2, and
1 for 2 4. In addition, we again set T 4, _ and use K = 400 discretization steps. This problem is similar the one considered in Nüsken & Richter (2021) and Domingo-Enrien Enrich et al. (2024b), the difference being that we consider instead of 1) and consider a nonzero running cost in order to have nontrivial long-term behaviour. This problem is considered highly nontrivial benchmark problem, since the double well in each dimension creates a total of 2 1024 local minima, making this setting highly multimodal. The ground truth is not available in closed form, but can be approximated efficiently by noticing that the energy and running cost are a sum one-dimensional terms, and hence compute the solution by solving d one-dimensional problems using classical solver.
a longer horizon
RING
The final setting considers the setup
F(x) a (1|=112 2 |,011 f(x) 2x1, g(x)
(110)
where a= = -- 5/ This energy is nonconvex and has its minimizers lying on the hypersphere with radius R. This setting serves to highlight the difference between the relative eigenfunction loss (20) and the other eigenfunction losses, and for visualization purposes will be done in d = 2. In addition. we T 5, 3 and zo (R. 0). The energy landscape and running cost visualized in Figure 6) The goal of the controller is essentially to guide the system to the left hand side of the zy-plane while being constrained by the potential stay close the circle of radius R. This setting requires a smaller timestep for stable simulation, we take K 500.
Ring energy
Running cost
7.2
3
90
5.4
2-
2-
3.6
* 0-
-50
-1
-100
-175
-3
-72
Figure 6: Energy function (left) and running cost (right) for the RING setting in d = 2.
E.2 MODEL ARCHITECTURE AND TRAINING
Architecture For the IDO methods, we use the exact same architecture as in Bomingo--Enice et al. (2024b). They argue that the control can be viewed the analog of a score function diffusion models, and hence they use a simplified U-Net architecture, where each of the up-/downsampling steps is fully connected layer with ReLU activations. As in their work, we three downsampling and upsampling steps, with widths 256, 128 and 64.
For the eigenfunction models, the GELU activation function x + (x)),, where & is the edf of a standard normal distribution (Hendrycks Gimpel (023). This is done because the eigenfunction losses require evaluating the derivatives of the network w.r.t. the inputs, hence requiring a smoother activation function.
the same architecture, but replace the ReLU activation with
Training For the eigenfunction method, we train using the following procedure:
1. Start training the top eigenfunction using the deep Ritz loss I77. Every 100 iterations, we estimate the eigenvalue Xn. and continue training until the variance these estimates (computed with EMA 0.5) is below 10-4 and we have reached least 5000 iterations.
31

---

Published as a conference paper at ICLR 2026
Table 2: Iteration times by method and loss
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
PINN
0,662
Relative loss
0.662
Variational loss
0.228
FBSDE
FBSDE
0.443
IDO
Adjoint Matching
0.230
Log variance
0.212
Relative entropy
0.413
SOCM
0.799
2. Fix An and continue training the top eigenfunction using loss function of choice among (20). Start training the excited state using the variational loss Rvar in with k and regularization parameter a Auol-
For the combined methods, we first train the eigenfunction and eigenvalues using the method above with Red for 80 000 iterations, and then start training with IDO loss using the control parametrization (2)). We resample the trajectories in Tewal (which only use the eigenfunction control) every 100 iterations order have a diverse set of starting positions T Tees-
E.3 COMPUTATIONAL COST
Table 2 shows the computation cost per iteration for the different algorithms for the QUADRATIC (REPULSSIVE) setting, measured in seconds/iteration when ran in isolation on a single GPU. All experiments where carried out on an NVIDIA H100 NV GPU,
E.4 FURTHER DETAILS ON THE RING SETTING
The RING setting serves illustrative example the difference between the 'absolute' cigenfunction losses and our relative eigenfunction loss. As shown in FFguuree 4] the relative loss obtains a drastically lower control objective than the absolute losses. To further understand this, consider the shape of the learned eigenfunctions for the different losses, shown in Figure 7 From the top row, it is clear that the learned eigenfunctions for the different methods are all very close in terms of the distance induced by || I|.. However, while the learned eigenfunctions look similar in (1),,, the logarithm of the learned eigenfunctions varies drastically, and hence the resulting control V log (shown in Figure 3) differs greatly- the control learned by the relative loss correctly guides the system along circle in the negative x direction, while the other controls are not learned correctly for xr ' 0. The name 'relative loss' comes from the analogy with absolute and relative errors,
Ix x=1
<E
VS.
<E
log x log x < < O(e)).
(111)
In essence, since the resulting control is given by V log c. the quantity of interest is the relative error of the learned eigenfunction, while existing methods are designed to minimize the absolute error.
E.5 DETAILS ON FIGURE
We give here some details the motivating plot shown in Figure I The value shown is the control L error of each the different methods in the QUADRATIC (REPULSIVE) setting. Each algorithm for 30k iterations, and the values reported the mean and 5%-95% quantiles over the last
32

---

Published as a conference paper at ICLR 2026
Deep Mid eterfandee
Vereet Peretesstionn eog eeges---eeeeeeee specec
IWW kop dignraattion
tdittve es teans exder
ROTIA Inpercidion
PRISCO
Retabv
Figure 7: Learned cigenfunctions (top row) and their logarithms (bottom row) for the RING setting with different loss functions.
Quadratic (repulsive) (d=20)
Quadratic (repuisive) (d=20)
Quadratic /repulsivel (d=20)
107
10-
poogó vlogedfi
10
10F
Average control
10
10-4
10-1
10-1
5000 10000 15000 20000 25000 Itoration
20000
40000
60000
80000
0.0
0.3
2.0
1.5
2.2
3.0
Figure 8: 2 error of V log ø for different eigenfunctions (left), average control I2 error for different methods (middle) and control L error over time (right). Legend is the same in Figure 4 and Pigure
1000 iterations. The EIGF+IDO (OURS) label refers to the combined method, where first train the eigenfunctions until convergence and then train using the relative entropy loss.
E.6 ADDITIONAL EXPERIMENT: REPULSIVE POTENTIAL
Figure shows the same plots discussed in the main text for the QUADRATIC (REPULSIVE) setting. where we obtain similar results. The reported eigenfunction error is measured in L² instead of (2)
Additionally, for this experiment setting we have tested the performance of "ergodic" control estimator. It is based only on the first eigenfunction, i.e. 8-1V log 00 used for the whole time range. The performance such an estimator is presented in Figure bby a curve labeled EIGF. It can be seen, that this approach reduces the error down with growth of time horizon T. However, using proposed approach with control as in 222 which corresponds to EIGF+IDO curve, leads to a significant improvement.
E.7 CONTROL ORJECTIVES
As mentioned before, the control objective is the final metric for evaluating the performance of the different algorithmss but due variance of the Monte Carlo estimator the difference between methods can be quite small. We report the control objectives for all experiments at convergence in Table 3 and Table 4 The value reported is the mean value of the control objective over N 65536 simulations, and the error is the standard deviation of these estimates, divided by VN.
33

---

Published as a conference paper at ICLR 2026
Centrol: Ferreve." ve.7 (ehw TIK Ienelions
Agaatinn
Impvensere 1000
Engi-800 Egg sendd
1990 HDTD001
Figure 9: In this experiment EIGF method uses an ergodic estimator based only on the first eigenfune- tion. EIFG+IDO curve corresponds the application of the proposed controller (22) with Relative Entropy loss. The figure shows L control error for different methods after 30000 iterations.
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 (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
(ours)
Log variance
32.7740 + 0.014
150.0721 - 0.12
SOCM
32.7717 + 0.014
112.3960 + 0.050
Adjoint matching
32.7725 + 0.014
114.7020 + 0.055
FBSDE
FBSDE
32.7979 = 0.014
112,6393 = 0.051
Table 4: Control objective for different methods the QUADRATIC (ANISOTROPIC) and DOUBLE
WELL settings.
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.004
SOCM
112.2728 -.12
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
(ours)
Log variance
32.5115 0.047
32.9080 + 0.0088
SOCM
31.3483 + 0.016
32.4421 + 0.0088
Adjoint matching
31.3497 0.016
32.5638 + 0.0088
FBSDE
FBSDE
31.3854 + 0.016
35.1890
= 0.011
34


Reviews:
[{"review_id": "2WXpIgvH10", "reviewer": "Reviewer_AKir", "summary": "Within the context of stochastic optimal control, this paper shows that optimal control can be obtained from the eigensystem of a Schrödinger operator S=−\\Delta + V with purely discrete spectrum. The paper also states that for a symmetric linear-quadratic regulator (LQR), the 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. For the general version, the paper presents a learning strategy of the eigensystem of the operator L via physics informed neural  networks.", "questions": "-  The presentation of the theoretical results is a bit unclear. For example, it is not clear what is the different between the presented Theorem 2 and the (Reed & Simon, 1978, Theorem XIII.67, XIII.47. Appendix B.5 announced to provide a proof of Theorem 2, but that is hardly the case, it is at best a sketch and a restatement of the results in (Reed & Simon, 1978, Theorem XIII.67, XIII.47.\n- In the numerical results, how is the check that the eigenvalues do not have a finite accumulation point? \n- The paper states that Instead of learning V_0 and \\lambda_0 jointly, it first performs training\nwith a variational loss in equation (17) to obtain an estimate for the eigenvalue \\lambda_0 and an initialization of V_0, and then ‘fine-tune’ using equation (20). Can the authors comment on the convergence rate? What affects the convergence as a function of dimensionality d?\n- Another issue is related to how effective is the proposed eigensystem learning approach when addressing the long-horizon challenge? Although the paper states under the limitations that method for determining an appropriate cutoff time exist, can computational approaches shade some light on the T_cut?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- An interesting connection between a version of stochastic optimal control and the Schrodinger operator. \n- A deep learning strategy based on PINN and several loss functions to learn the eigensystem of the operator L.", "weaknesses": "-  The presentation of the theoretical results is a bit unclear. For example, it is not clear what is the different between the presented Theorem 2 and the (Reed & Simon, 1978, Theorem XIII.67, XIII.47. Appendix B.5 announced to provide a proof of Theorem 2, but that is hardly the case, it is at best a sketch and a restatement of the results in (Reed & Simon, 1978, Theorem XIII.67, XIII.47.\n- In the numerical results, how is the check that the eigenvalues do not have a finite accumulation point? \n- The paper states that Instead of learning V_0 and \\lambda_0 jointly, it first performs training\nwith a variational loss in equation (17) to obtain an estimate for the eigenvalue \\lambda_0 and an initialization of V_0, and then ‘fine-tune’ using equation (20). Can the authors comment on the convergence rate? What affects the convergence as a function of dimensionality d?\n- Another issue is related to how effective is the proposed eigensystem learning approach when addressing the long-horizon challenge? Although the paper states under the limitations that method for determining an appropriate cutoff time exist, can computational approaches shade some light on the T_cut?\n- When discussing order-of-magnitude, what sustains the improvement in L^2 error compared to standard IDO and FBSDE methods? Some additional explanations could have help me understand also the computational complexity analysis.\n- Although this is a minor issue, when considering the cases without gradient drift (where L is non-symmetric), what are the major technical problems?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "M6ZiZTtpX2", "reviewer": "Reviewer_KCj8", "summary": "The authors present a **hybrid framework** for solving **stochastic optimal control (SOC)** problems under the **gradient-drift assumption**, combining **eigenfunctions** with **short-horizon solvers**. Using the **Cole–Hopf transform**, the **Hamilton–Jacobi–Bellman (HJB)** equation is converted into a **linear PDE**, enabling efficient solution through **eigenvalue and spectral analysis**. The framework reparameterizes the value function to learn the **top eigenfunction**, which determines the long-horizon control. Empirically, the method achieves strong performance for **long time horizons**, where existing approaches typically suffer from rapid error and computational cost growth.", "questions": "1. How is the cutoff time ($T_{\\text{cut}}$) selected or tuned in practice?\n2. Does the learned eigenfunction ($\\phi_\\theta$) correspond to the true eigenfunction of the operator (L), or is it only an approximation under the chosen loss?\n3. Is the eigenvalue ($\\lambda_0$) learned during training (fine-tuning)? If so, is its computation **differentiable** and integrated into the optimization process?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Rigorous theoretical treatment with clear assumptions, theorems, and proofs.\n2. Empirical results demonstrate improved performance over baselines as the horizon T increases.", "weaknesses": "1. The paper could better **motivate the relevance and generality** of the gradient-drift assumption, clarifying how broad and practically applicable the class of SOC problems is where the HJB can be reduced to a linear PDE.\n2. The claim that “the exponential decay of the correction term in Eq. (2) suggests that the optimal control ($u^*$) can be approximated by the gradient of the logarithm of the top eigenfunction” holds only **asymptotically**; the bounded correction term may remain **non-negligible** over ([0, T]), especially at earlier times.\n3. The modification of the loss from Eq. (15) to Eq. (20) is said to be equivalent when ($\\phi$) is unit-norm, but since the regularization coefficient ($\\alpha$) is **chosen heuristically**, there is **no guarantee** that both objectives share the same optimum or yield a normalized eigenfunction.\n4. The authors state that the loss in Eq. (21) remains sensitive when ($\\exp(-\\beta V_0)$) is small; however, this may also cause **instability**, as the loss could **diverge** when ($V_0$) becomes large.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NUfe9OL88X", "reviewer": "Reviewer_Y3as", "summary": "This paper studies the numerical resolution of continuous control PDEs (a.k.a. HJB equations) in high-dimensional state spaces using neural networks (as classical PDE solvers are prohibited by the curse of dimensionality). Contrary to most works, which focus on the computational impact of the number of samples of the trajectories, here the concern is the computational impact of the length of the trajectories. In this undiscounted, finite, but large horizon problem, the authors propose expanding the evolution operator into its eigenfunctions and keeping only the dominant eigenfunction as a numerical solution. This raises three sub-problems: 1) identifying conditions under which this approach is principled, 2) identifying or approximating the principal eigenfunction, 3) recovering an effective control from the principal eigenfunction. \n\nThe authors tackle 1) by working within a classical, moderate-complexity setting of affine control, whose (thus simplified) evolution operator can be related through a classical change of function to a linear PDE and associated Schrödinger operator. These theoretical derivations follow classical methodologies in control and PDE theory, but their highly qualitative and rigourous application yields novel results, and in particular Theorem 4. They also naturally yield an expression of the optimal control in terms of the eigenbasis (Theorem 3), which is truncated to provide a solution to sub-problem 3). This analytical method is the main body of the work, an architecture, if you will, which supports the use of neural network methods for learning eigenfunctions to solve sub-problem 2). The authors provide an extensive overview of existing methods for learning eigenfunctions, but also propose their own loss, which is expressed in units corresponding to the control problem (a reweighting which counteracts the change of functions of step 1)).", "questions": "There is one important question that, for me, casts a shadow over the phenomenon at play here, which I think ought to be answered: \n- If the horizon is becoming longer and longer, one would expect turnpike properties to kick in and the ergodic regime to become a good approximation for the behaviour of the system. This connection is suggested in Remarks 1 and 2. To what degree is this actually responsible for the improvement in computational performance? Wouldn’t a good benchmark for Figure $1$ be a solver which solves the ergodic version of the problem (though you would have to change the problem to one of the other LQRs for this to exist), which one would expect to improve with $T$ just by the law of large numbers? \n\n\nThis work has made me wonder about several research questions, which is something I like about it, but these wouldn’t change my evaluation of the work. I leave them without expecting an answer for the authors’ benefit: \n- You use only the first eigenvalue here, but you could consider higher-order approximations, which might improve performance. This performance would be inversely exponential in the spectral gap, of course, so it would be an interesting question to study the gaps in the spectra of different evolution operators to identify problems where different orders of approximation are interesting. \n- If you could study the full HJB equation (beyond control-affine assumptions, I mean) and show that the evolution operator has a spectral gap and eigenbasis, you could apply the learning components as well, no? Perhaps using hypocoercivity theory and some non-linear PDE analysis, it would be possible to obtain similar results under very general assumptions, though it would be very complicated. \n- I understand the problem with end-of-the-world effects, which cause the need to switch off the eigenfunction solver near $T$, but the performance degradation is very painful (Figure 5). Perhaps something clever can be done about it.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The paper is well-written, well-structured, and highly pedagogical, despite the difficulty of the topic and the distance between the communities of PDE theory and general ML. It is also very complete both in terms of the references in control and PDE theories and in terms of the rigour of the statements and argumentation. The theoretical contributions are impressive and support a novel multi-disciplinary approach fusing control/PDEs with ML, while being backed up by very rigourous mathematics and good experiments.", "weaknesses": "The assumptions required to obtain the theoretical results are more general than, say, linear quadratic, but they remain a far cry from many practical situations. This is highlighted in the experiments, which are restricted to toy problems. However, in my opinion, while it is a weakness, it is difficult to criticise the quality of the work on these grounds in good faith, as the step up in complexity required by weaker assumptions would be so vast.\n\n_Minor feedbacks_:\n- The later appendices (C and beyond) contain a few typos and overloaded notations; there are also some less-than-correct phrases in the main text.\n- In equation (104), shouldn’t the integral over $t$ be a classical integral and not an expectation?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)