Stochastic Optimal Control Matching
Authors: Carles Domingo-Enrich
Jiequn Han
Brandon Amos
Joan Bruna
Ricky T. Q. Chen
Venue: NeurIPS 2024 poster (2024)
Keywords: Stochastic Optimal Control
Diffusion
Primary area: diffusion_based_models
TL;DR: We propose the first least squares loss to solve Stochastic Optimal Control problems, and show that it outperforms existing losses experimentally.

Abstract: Stochastic optimal control, which has the goal of driving the behavior of noisy systems, is broadly applicable in science, engineering and artificial intelligence. Our work introduces Stochastic Optimal Control Matching (SOCM), a novel Iterative Diffusion Optimization (IDO) technique for stochastic optimal control that stems from the same philosophy as the conditional score matching loss for diffusion models. That is, the control is learned via a least squares problem by trying to fit a matching vector field. The training loss, which is closely connected to the cross-entropy loss, is optimized with respect to both the control function and a family of reparameterization matrices which appear in the matching vector field. The optimization with respect to the reparameterization matrices aims at minimizing the variance of the matching vector field. Experimentally, our algorithm achieves lower error than all the existing IDO techniques for stochastic optimal control for three out of four control problems, in some cases by an order of magnitude. The key idea underlying SOCM is the path-wise reparameterization trick, a novel technique that may be of independent interest.

stochastic filtering Stochastic Optimal Control Matching
Carles Domingo-Enrich NYU & FAIR, Meta
Jiequn Han Flatiron Institute
cd2754@nyu.edu
Brandon Amos FAIR, Meta
Joan Bruna
Ricky T. Q. Chen FAIR, Meta rtqichen@mett com
NYU & Flatiron Institute bruna@cims ndu
bda@meta. com
Abstract
Stochastic optimal control, which has the goal of driving the behavior of noisy systems, is broadly applicable in science, engineering and artificial intelligence. Our work introduces Stochastic Optimal Control Matching (SOCM), novel Iterative Diffusion Optimization (IDO) technique for stochastic optimal control that stems from the same philosophy as the conditional score matching loss for diffusion models. That is, the control is learned via a least squares problem by trying fit matching vector field. The training loss, which is closely connected to the cross-entropy loss, is optimized with respect to both the control function and family of reparameterization matrices which appear in the matching vector field. The optimization with respect to the reparameterization matrices aims at minimizing the variance of the matching vector field. Experimentally, our algorithm achieves lower eгror than all the existing IDO techniques for stochastic optimal control for three out of four control problems, in some cases by an order of magnitude. The key idea underlying SOCM is the path-wise reparameterization trick, novel technique that may of independent interest.
Introduction
Stochastic optimal control aims to drive the behavior of a noisy system in order to minimize given cost. has myriad applications in science and engineering: examples include the simulation of rare events in molecular dynamics 337 36
85 41, finance and economics (63 25
and data ussimilation [58 68). nonconvex optimization [19], sampling 9 power systems and energy markets [8] [66), and robotics [77
Stochastic optimal has also been impactful in fields such as mean- field games [17). optimal transport [80] [81). backward stochastic differential equations (BSDEs) 14] and large deviations [2 Recently, it has been the basis of algorithms to sample from unnormalized densities 84
32].
For continuous-time problems with low-dimensional state spaces, the standard approach to learn the optimal control is to solve the
(HJB) partial differential equation (PDE) by gridding the space and using classical numerical methods. For high-dimensional problems, a large number of works parameterize the control using neural network and train applying stochastic optimization algorithm on loss function. These methods are known as Iterative Diffusion Optimization (IDO) techniques [59 (see Subsec. 2.2].
is convenient to draw an analogy between stochastic optimal control and continuous normalizing flows (CNFs), which are generative modeling technique where samples generated by solving an ordinary differential equation (ODE) for which the vector field has been learned, initialized at a Gaussian sample. CNFs were introduced by [20 (building on top of Rezende and Mohamed (70),
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

and training them is similar to solving control problems because in both cases one needs to learn high-dimensional vector fields using neural networks, continuous time.
The first algorithm developed to train normalizing flows was based on maximizing the likelihood of the generated samples (20) Sec. 4]. Obtaining the gradient of the maximum likelihood loss with respect to the vector field parameters requires backpropagating through the computation of the ODE trajectory, or equivalently, solving the adjoint ODE in parallel to the original ODE. Maximum likelihood CNFs (ML-CNRs) were superseded by diffusion models 75 40 (66 and flow-matching. a.k.a. stochastic interpolant, methods [55] [ 63 2), which currently the preferred algorithms to train CNFs. Aside from architectural improvements such as the UNet [73], a potential reason for the success of diffusion and flow matching models is that their functional landscape is convex, unlike for ML-CNFs. Namely, vector fields are learned by solving least squares regression problems where the goal is to fit a random matching vector field. Convex functional landscapes in combination with overparameterized models and moderate gradient variance can yield very stable training dynamics and help achieve low error.
Returning to stochastic optimal control, one of the hest-performing IDO techniques amounts to choosing the control objective (equation [1] as the training loss (see (12) ). As in ML-CNFs, computing the gradient of this loss requires backpropagating through the computation of the trajectories of the SDE 2 equivalently, using an adjoint method. The functional landscape of the loss is highly non-convex, and the method is prone to unstable training (see green curve the bottom right plot of Figure 3). In light of this, natural idea is to develop the analog of diffusion model losses for the stochastic optimal control problem, to obtain more stable training and lower error, and this is what set out to do in our work. Our contributions are as follows:
We introduce Stochastic Optimal Control Matching (SOCM), a novel IDO algorithm in which the control is learned by solving least-squares regression problem where the goal is to fit random matching vector field which depends on family of reparameterization matrices that are also optimized.
We derive a bias-variance decomposition of the SOCM loss (Prop. 2). The bias term is equal to an existing IDO loss: the cross- cross-entropy loss, which shows that both algorithms have the same land- scape in expectation. However, SOCM has an extra flexibility the choice of reparameterization matrices, which affect only the variance. Hence, we propose optimizing the reparameterization matrices to reduce the variance of the SOCM objective.
The key idea that underlies the SOCM algorithm is the path-wise reparameterization trick (Prop. l}. which is novel technique for estimating gradients of an expectation of a functional of random process with respect to its initial value. It is of independent interest and may be more generally applicable outside of the settings considered in this paper.
We perform experiments on four different settings where we have access to the ground-truth control. For three of these, SOCM obtains a lower L error with respect to the ground control than all the existing IDO techniques, with around 10x lower error than competing methods in some instances.
2 Framework
2.1 Setup and Preliminaries
Let (FF)) P) be a fixed filtered probability space on which is defined a Brownian motion
B = (Br)220- We consider the control-affine problem
min E [[ (1) ( 1)1++ d(+g(
(1)
uEll
s.t. dx.'' (b(X)' t) +o (t)
dt Voo ~ Po
(2)
and where XI" E Rd is the state, u Rd [0, - Rd is the feedback control and belongs to the set of admissible controls U. f: Rd x [0, T] + 12 is the state cost, g I2 + IR is the terminal cost, b : Rdd x [0, + Rd is the base drift, and a: [o. 7] Rdxd is the invertible diffusion coefficient and X E (0, +o0) is the noise level. In AnpApp A we formally define the set 24 of admissible controls and describe the regularity assumptions needed on the control functions. In the remainder of the section we introduce relevant concepts in stochastic optimal control; we provide the most relevant proofs in App. B and refer the reader to Oksendal [60] Chap. 11] and Nüsken and Richter [59] Sec. 2] for a similar, more extensive treatment.

---

Cost functional and value function The cost functional for the control u, point r and time t is defined J (u; x, == L[ (x),, (x(")) dt o(X)) = =]. That is, the cost functional is the expected value of the control objective restricted to the times [t. T] with the initial value z at time t. The value function or optimal cost-to-go at (a, t) is defined as the minimum value of the cost functional across all possible controls:
Vix, t) == inf J (u: x, x,t).
Hamiloon-JacoiiBellnee equation and optimal control If we define the infinitesimal gener- ator == 2/2 1 (ao (ttt)),,,an Ozy + +E (z t)0x,, , the value function solves the following manilton-Jacobi-elllan (HJB) partial differential equation:
(a. + [,V(,,)) 1) 211((r T (V),,,))l' (xx, =0,
V(x, T) g(x))
(4)
The verification theorem [62] Sec. 2.3] states that if function V solves the HJB equation above and has certain regularity conditions, then V the value function (3) of the problem implication of the verification theorem is that for every a E U,
An
V(a, t) + + J VV + +XI?(XX 8) ds X = = = Ju, x, x,t). t).
(5)
In particular, this implies that the unique optimal control is given in terms of the value function as u² (zz t) = (t) VV t). Equation (5) can be deduced by integrating the HJB equation (4) over
[T, and taking the conditional expectation with respect to X;' x. We include the proof of (5) in
App. B for completeness.
pair of forward and backward SDEs (FBSDEs)
Consider the pair of SDEs
dx, = b(x, f) dt d.
Xo ~ Pu.
(6)
dy; = (- () 2 1Z.11) (Z,
Y- g(x))
where Y: 12 [0, + R and Z f2x [0, T + Rd are progressively measurable measurable-rrnddm prosspss turns out that Y, and Z, defined as Y, := V(X,, and Z, == a(1) VV((..)) -u" (X,,) 1) satisfy 77. We include the proof in App. B for completeness.
processes.
An analytic expression for the value function From the can derive a closed-form expression for the value function V:
equations
6 7 one
V(z, = = El cxp (-x-1 S:" f(X,, 8) ds x x).
(8)
where X, the solution of the uncontrolled SDE 6 This a classical result, but we still include its proof in App. B Given that u² (z, t) (t( V V (z, t), immediate, yet important, consequence of (8) is the following path-integral representation of the optimal control:
(9) Remark this equation involves the gradient of logarithm of a conditional expectation, which is reminiscent of the vector fields that are learned when training diffusion models. For example, the target vector field for variance-exploding based diffusion loss (76) can be expressed as
"² (o) TV, log [[eep eepppxx-> (-x-1 S f(X,, ds-dd---- d =x]
Va log pr (I) = V. log Ey~paas! (irp-|lr-y||///2a0
Note, however, that in ( the gradient is taken with respect the initial condition of the process, which requires the development of novel techniques.
(2x07)a/2
Conditioned diffusions Let C C((o, R)) be the Wiener space of continuous functions from [0, to Rd equipped with the supremum norm, and let P(C) be the space of Borel probability measures over C. For each control u E l4, the controlled process in equation (2) induces a probability measure in P(C), the law of the paths X²" which we refer to as 1p" We let IP be the probability measure induced by the uncontrolled process (6). and define the work functional
W(X,t) = S," f(X,, ds g(Xr)-
(10)
Being progressively measurable is strictly stronger property than the notion of being a process adapted to the filtration F; of B, (see (50]).
3

---

turns out (emma 2 in App. B that the Radon-Nikodym derivative dF" dF satisfies dP (X) exp (A (V(Xo, 0) (XX, 0))) Also, a straight-forward application of the Girsanov theorem for SDEs Cor. D shows that
aF
drth d..
(X") -xxx) (-2-1/2 So (u (XV tt)- (B), So (XF .t) ((p dt),
(11)
which means that the only control u E U such that 1 Ip" is the optimal control itself.
2.2 Existing approaches and related work
Low-dimensional case: solving the HJB equation For low-dimensional control problems (d S. 3). it possible to grid the domain and use a numerical PDE solver to find a solution to the HJB equation 4 The main approaches include finite difference methody [1] 57 4. which approximate the derivatives and gradients of the value function using finite differences, finite
finite element methods (47), which involve restricting the solution to domain-dependent function spaces, and semi-Lagrangian schemes [2 12). which trace back characteristics and have better stability than finite difference methods. See Greif (33] for overview these techniques, and Bañas et al. [4] for comparison between them. Hutzenthaler et al. [44] introduced the multilevel Picard method, which leverages the Feynman-Kar Kac and the Bismut-Elworthy- formulas to beat the curse of dimensionality in some
settings 6 46 45 43
High dimensional methods leveraging FBSDEs The FBSDE formulation in equations has given rise to multiple methods to learn controls. One such approach is least- Monte Carlo (see Pham [3] Chapter 3] and Gobet (8 for an introduction, and Gobet et al. [30], Zhang et al. [83) for an extensive analysis), where trajectories from the forward process (6) are sampled, and then regression problems are solved backwards in time to estimate the expected future cost in the spirit of dynamic programming. A second method that exploits FBSDEs was proposed by et al. [22]. Han et al. [35]. They parameterize the control using a neural network ug. and use stochastic gradient algorithms
minimize the loss L(ue, (o) = (E2 (yo. ug) (Xx where Yr (yo: ug) is the process in ( with initial condition you and control ug. This algorithm can be seen as shooting method, where the initial condition and the control are learned to match the terminal condition. Multiple recent works have combined neural networks with FBSDE Monte Carlo methods
for parabolic and elliptic PDEs (5 I8 86). control 7 39
multi-agent games 34 15 16
see 23
for a more comprehensive review.
Many of the methods referenced above and some additional ones can be seen from a common perspective using controlled diffusions. As observed in equation (TT). the key idea is that learning the optimal control is equivalent to finding control u such that the induced probability measure Pd on paths is equal to the probability measure F" for the optimal control. In the paragraphs below we cover several loss that fall into this framework. All the losses below can be optimized using a common algorithmic framework, which we describe in Algorithm [ For more details, we refer the reader to Nüsken and Richter [59], which introduced this perspective and named such methods Iterative Diffusion Optimization (IDO) techniques. For simplicity, we introduce the losses for the setting in which the initial distribution pu is concentrated at single point Finit; we cover the general setting in App. B
The relative entropy loss and the adjoint method The relative entropy loss is defined as the Kullback-Leiblrr divergence between Fa and Ip"" and factors, this loss is equivalent to (see Lemma 3 in App. B
Eg [log
d3" Upon removing constant terms
(u) = E So (3 [u(XX, t)l +(X," t)) dt +x(+)]
(12)
This is exactly the control objective in (T. This fact has been studied extensively 1T 31 36
48
Hence, the relative entropy loss is very natural and widely used; see Onken et al. 661. Zhang and Chen (84) for examples on multiagent systems and sampling.
Solving optimization problems of the form [2 has long history that dates back to Pontryagin 64 Note that Lad5(u) (u) depends u both explicilly, and implicitly through the process X". To compute the gradient VoLis (uo..) of a Monte Carlo approximation (xaj(ue.)) of L (uo.) ) as required by Algorithm we need to backpropagate through the simulation of the trajectories, which is why

---

Algorithm Iterative Diffusion Optimization (IDO) algorithms for stochastic optimal control
Input: State cost f(z, t), terminal cost g(x). diffusion coeff. a (t), base drift b(z,t), t),, noise level A, num- ber of iterations N. batch m. number time steps K, initial control parameters Pg. loss
C Addj ( Lce [3]. LVar. (66) Var 17 Mom,, (8)
for E , NN- do
Simulate trajectories of the process X" controlled by uo,, e.g. using Euler-Maruy updates if * Ladj then detach the trajectories from the computational graph, that gradients do backpropagate:
Using the m trajectories. compute an m--ammle Monte Carlo approximation ((uo,) of the loss [(ue,)
Compute gradients VoL(uo,) of [(uo..) w.r.t. fr Obtain 4 with via an Adam update on 0, (or another stochastic algorithm)
end
Output: Learned control uaN
we do not detach them from the computational graph. One can alternatively compute the gradient (us.) by explicitly solving an ODE, technique known as the adjoint method. The adjoint method was introduced by Pontryagin [64, popularized in deep learning by Chen al. 20) and
further developed for SDEs in Li et al. [54]
The cross-entropy loss The cross-entropy loss defined as the Kullback-Leibh divergence
between P"" and P" i... flipping the order of the two measures: Ep. llog v E Ut this loss is equivalent to the following one (see Prop. 3(1) in App. B
d2d dit
For arbitrary
(CE(w): (- 1
1/2 5o (u(X), t), d - -1 S (u(XF, 1), u(X , dt S |I((X; t)l 1/2
x exp X X 1-12 (v(X), t), dB,)
A fo ||((X 1) d))))d)))
(13)
The cross-entropy loss has rich literature such as molecular dynamics [41). Furthermore, we note that the cross-entropy loss can be significantly simplified and written in terms of the unnormalized 1!2 error of the control u with respect to the optimal control u":
38 49 74 85
and has been recently used in applications
CCE(u) = - [Xx"(X"" ,t) u(Xp 1)1 dt x exp (-X-IV(Xy .))]
(14)
This characterization, which is proven in Prop. ii) in App. is relevant for because similar one can be written for the loss that we propose (see Prop. 2).
Variance and lor-variance losses
For an arbitrary e U4 the variance and the log- losses dd" dp"
defined as Lvar, (u) Varg=
<d=" dP
and Var, (u) Var (log dp= dp"
whenever Ep.
+oo and EE.- log
d0"" +00, respectively. Define
dpw
yr. - 10 (x(X),),,((X,,)) x-1 10" ()x X- fo (u(X),t),dB)
(15)
Then, Evar, and Glog Var, are equivalent, respectively, to the following losses (see Lemma 4
Lvar, (u) = Var( exp ( x 9(x)))
(16)
C (u) := - xxx))
(17)
The variance and log-variance losses were introduced by Nüsken and Richter (59). Unlike for the cross-entropy loss, the choice of the control v does lead to different losses. When using L or the variance is computed across the trajectories in each batch.
clug in Algorithm
Moment loss
For an arbitrary v E U4 the moment loss defined as
Lmmm. (u, yo) B(( + yo A -- (+
(18)
5

---

where YT. is defined in [5]. Note the similarity with the oog-variance loss 17 the optimal value
of yo for fixed is yo | g(X)) [[(=='((X))
")"",, and plugging this into 18 yields exactly the log-variance loss. The moment loss was introduced by Hartmann et al. [39 Section III.B], and is a
generalization of the FBSDE method pioneered by E et al. [22], Han et al. [35 and referenced earlier
in this subsection, which corresponds to setting e = 0.
Stochastic Optimal Control Matching
In this section present our loss, Stochastic Optimal Control Matching (SOCM). The corresponding method, which we describe in Algorithm 2 falls into the class of IDO techniques described in Subsec. 2.2] The general idea is to leverage the analytic expression of u" in ( to write a least squares loss for u, and the main challenge is to reexpress the gradient of a conditional expectation with respect to the initial condition of the process. We do that using novel technique which introduces certain arbitrary mutrix-valuee functions Mc. that we also optimize.
Theorem 1 (SOCM loss). For each
[0, T), ler M, [t + Baddd be an arbitrary matrix- valued differentiable fanction such that Mc(l) = Id. Let v E u be an arbitrary control. Let CSOCM: (Pd x T) R2) x (10, Rdxd) + R be the loss function defined as
[so M) == E lo- ||((F, t) w v, X", B, M,)l dt x a(v, B)], (19)
where X" is the process controlled by u (i.e., dx;" = (6(x), 6) o0)) vTo(t) dBg and Xğ po). and
w(l, X", B, -
+ 5 M)
+\ S:t
a(t,X",B) = exp -- S00 f(X;, ds x
(20)
--12 le
A= So 0)l at).
LSOCM has a unique optimum (u'M'" where w' is the optimal control.
We refer to M (M.)) as the family of reparametrization matrices, to the random vector field w as the matching vector field, and to as the importance weight. We present proof sketch of Thm. I the full proofs for all the results in this section are in App. C
Proof sketch Th.. 1
Let X be the uncontrolled process (6). Consider the loss
[(u) [ Jo [lu(X,,)) n² t)l 1) dt exp ( (- - 50d f(Xe, t) dt X --- x(x)(]
[ 1o ([lu(X,,))) - ( 2(wu(((,..) t), u² (X,,t)) + (X,, t) t) dt
(21)
x exp -- 5 s(x,, 1) dt - 9 (X)))]
Clearly, the only optimum of this loss is the optimal control u". Using the analytic expression of u² in the cross-term can be rewritten as (see Lemma 5 in App.
E E So (u(Xe, w² (X,, 1)) di exp (- (-x-1 X-1 fo -xx,)))(--''((xx))
A I (u(Xe,t), T V.E exp (-x-1 fe² -x =x))(22)
x exp (- 1o f(x,, dt].
It remains to evaluate the conditional expectation V_E exp - x-1 f(X,, ds X x 'x(()))) on the initial value x into the stochastic processes- we introduce free variable M,-and then applying Girsanov theorem. We coin this the path-wise reparameterization trick:
x|. which we do by "reparametteizatton trick" that shifts the dependence
Proposition (Path-wise reparameterizattoon trick for stochastic optimal control). For each t E [0 let M, [t, 17 + dd be arbitrary continuously differentiable function matrix-vlued funciiioo

---

such that Mt(t) = Id. We have that
[. (- x- J(X (1Xx,3))d =
B[( - 1:¹
(23)
8, + -- 5 - (.)BB.)
exp (-X 1/ 8) - X =
We prove a more general form of this result (Prop. 4) in Subsec. C.2 and also provide an intuitive derivation in Subsec. C.3 In the proof of Prop. 4 the reparameterization matrices M, arise as the gradients of a perturbation to the process X,. Similar ideas can potentially be applied to derive losses for generative modeling. If we plug (23) into the right-hand side of (22). and then this back into (21 and we complete the square, we obtain that for some constant K independent of u,
E(u) = 7 Lo" [lu(Xe,t) + +o( Me(s)V s) ds + M,(T)VV((T))
- -2 -M (s))BB.)] dt
exp (-X-1 S0" d) + K.
If we perform a change of process from X to X" applying the Girsanov theorem Cor. in App. C we obtain the loss LsOCM (u, (u,M) M).
The following result clarifies the role of reparameterization matrices, connecting the SOCM and cross-entropy losses.
Proposition (Bias- decomposition of the SOCM loss). The SOCM loss decomposes into bias term that only depends on u and a variance term that only depends on M:
LSOCM .)
CondVar (w; M) + [ 1 [xxx"'
1)] (X,"
dt d .0]
(24)
Unnormalized expected conditional variance w
Unnormalized bios of
where
CandWar(w; =M=RG/O [wt(t, x", XB, M.)
E[w[[t,.X" B.M, B. BIX "a)x,
dt a(v, X", B)]
(25)
Remark that the bias term in equation 24 is equal to the characterization of the cross-entropy loss in (14). In other words, the landscape of (s, M) with respect to is the landscape of the cross-entropy loss (CE(u). Thus, the SOCM loss can be seen as some form of variance reduction method for the cross-entropy loss, and performs substantially better experimentally (Sec. 4p. Yet, the expressions of the SOCM loss and the cross-entropy loss are very different; the former is a least squares loss and is expressed in terms of the gradients of the costs.
Algorithm 2 Stochastic Optimal Control Matching (SOCM)
Input: State cost f(x, 1), terminal cost g(z), diffusion coeff. (t). base drift b(r, noise level X, number of iterations N, batch size m, number of time steps K, initial control parameters fg. initial matrix parameters loss LSOCM (9)
for (O...., N
Simulate m trajectories the process X" controlled by u ue, e.g., using Euler-Mane Maruyama updates Detach the m trajectories from the computational graph, so that gradients do not backpropagate
Using the m trajectories, compute an m-sample Monte-Carlo approximation ÊsoCM (Ne.,, Ma,,) of the loss LSocM Ma..) in (19)
Compute the gradients ((s...) (sw.)soo Mwn) of Lsocu Mu.)) at (fn,wa) Obtain 8,4 Whi1 with via an Adam update on On, Wh, resp.
end
Output: Learned control nox
For good training performance, it is critical that the gradients have high signal-to-noise ratio. Looking at the SOCM loss, good proxy for low gradient variance is to have low variance for w(t, u. X", B, M,) x a(v, X", B), and this holds when both (vv, B)
T J [uXy,,, and wit, v, X", , B, M,) have low variance. Next, we present strategies to lower the variance of these two objects.

---

Minimizing the variance of the importance weight a We want to use a vector field v such that Varla(u, B)] is as low as possible, As shown by the following lemma, which is well-known in the literature, setting v to be the optimal control u² actually achieves variance zero when we condition on the starting point of the controlled process X The proof of this result can be found in Hartmann et al. (38), but we include it in [Subsec. C.5 for completeness.
Lemma 1. When we
u", the conditional variance Varla(e, X", B)( -mmit] is zero for
any Tinit E IR4.
Of course, we do not have access to the optimal control u², but it is still a good idea to set as the closest vector field to u" that we have access to, which is typically the currently learned control. In some instances, one may benefit from using a warm-started control parameterized as (ws(2, t) + up t). where the warm-start uws is a reasonably good control obtained via a different strategy (see App. E
Minimizing the variance of the matching vector field w We are interested in finding the family M = 20, that minimizes the variance of we(t, v, X B, M,) conditioned on and X Note that this is exactly the term CondVar(w: M) in the right-hand side of equation (24). Since Cond CondVar Var(w; M) does not depend on the specific v, the optimal M does not depend on either. And
since the second term in the right-hand side of equation (24) does not depend on M = (M.)
minimizing Cond Var(w; M) is equivalent minimizing C(u) with respect to M.
Parameterizing the matrices M, vs solving for the optimal matrices In practice, we parameterize the matrices (Me) using a function M.. with two arguments (4, s). To enforce that (l,
Id, we set (t, s) e (- ll-ca e ())()))) Ma s), where
(x, W). and Ma X - Raxd is an unconstrained neural network. Alternatively. [Thm. in App. D shows that the optimal family M = (M;) [ee[o.T] can be characterized as the solution of a linear equation in infinite dimensions (a Fredholm equation of the first kind). The discretized linear system has d' K equations and variables, K being the number of discretization time points. However, since the optimal M does not depend on (sce Remark I), this is a computation that must be done only once and that may be affordable in some settings. We did not test this approach experimentally.
(Thm.
4 Experiments
We consider four experimental settings that we adapt from Nüsken and Richter (59] QUADRATIC ORNSTEIN UHLENBECK (EASY). QUADRATIC ORNSTEIN UHLENBECK (HARD), LINEAR ORN- STEIN UHLENBECK and DOUBLE WELL. We describe them in detail in App. F For all of them, we have access to the ground-truth optimal control, which means that we are able to estimate the L error incurred by the learned control u. In Figure 2 we plot the control L error for each IDO algorithm described in Subseec. 2.2] and for the SOCM algorithm (Algorithm 2). for the QUADRATIC OU (EASY) and (HARD) settings. We also include two ablations of SOCM: (i) version of SOCM where the reparameterization matrices Mt are set fixed to the identity I, (ii) SOCM-Adjoint, where we estimate the conditional expectation in equation (23) using the adjoint method for SDEs instead of the path-wise reparameterization trick (see Subsec. C.4). Code can be found at htttps://gitte http://tthtb.ccccaccccbbeeeeeeee httm///aaccbooeeeeeere///CC
At the end of training. SOCM obtains the lowest L error, improving over all existing methods by a factor of around ten. The two SOCM ablations come in second and third by a substantial difference, which underlines the importance of the path-wise reparameterization trick. The best among existing methods is the adjoint method (the relative entropy loss). In Figure 2 (bottom) norm of the gradient of each loss with respect to the parameters 0 of the control: algorithms with small noise variance have low error values.
show the squared
In Figure 3 we plot the control L error for LINEAR ORNSTEIN UHLENBECK and DOUBLE WELL. For LINEAR OU, the error is around five times smaller for SOCM than for any existing method. For DOUBLE WELL, the SOCM algorithm achieves the third smallest error, slightly behind the variance loss and the adjoint method, but the latter shows instabilities. As we show in Figure 9nn App. F these instabilities are inherent to the adjoint method and they do not disappear for small learning rates. Both in Figure 2 and Figure 3 we observe that learning the reparameterization matrices is critical to obtain gradient estimates with high signal-to-noise ratio. DOUBLE WELL a particularly interesting and challenging setting because its solution is highly multimodal: g has 1024 modes. Multimodality
8

---

is a feature observed in realistic settings, and is hard to handle because it involves learning the control correctly in each mode.
The costs f and and the base drift for QUADRATIC OU (HArd) are five times those of QUADRATIC OU (EASY). Consequently, the factor a(v, Xx, B) initially has much larger variance for the SOCM methods, and for cross-entropy. As training progresses, ue, gets closer to u², and consequently the variance of a(v, X", B) decreases, which in turn makes learning easier. This explains the initial slow decrease in the control error, followed by a fast drop that places SOCM well below existing algorithms. In In|Ap. E. showcase a control warm-start strategy that can help and speed up convergence.
We also present experimental results on two-mode Gaussian mixture sampling increasing dimen- sion, using the Path Integral Sampler [84]. We take Gaussians with means that are 2 units apart, and identity variance. Figure shows control objective estimates obtained after running the Adjoint, SOCM, and Cross-entropy algorithms for 40000 iterations, at dimensions d = 2, 8, 16, 32, 64, and error bars show standard errors. By Theorem 4 of [84), we know that the optimal value of the control objective is zero; Figure shows the suboptimality gaps incurred by each algorithm. Cross-entropy.
Control Objective estimate Dimension
Figure 1: This plot shows the control objective values for different algorithms (Adjoint, SOCM, and Cross-entropy) across multiple dimensions, with error bars indicating the standard deviations. The y-axis is restricted to [0, 0.1] for better visibiliity of the lower range values; cross-entropy takes value 2.915 + 0.008 at d = 64.
0.04
0.02
0.00
which uses the same importance weight as SOCM, performs worse than the other two losses for all dimensions, and its results are particularly poor for dimension 64, because the variance of is too large for learning to happen. In this case, we see that SOCM has better variance reduction than cross-entropy, despite both using importance weighted objectives for training. We observe that the values for SOCM are slightly below that of Adjoint for most dimensions, which confirms that our method is better for this range of dimensions. we keep increasing the dimension, SOCM also fails due to higher variance of a: for n 128, the control objective estimates for the Adjoint, SOCM, and Cross-Entropy losses are 0.146 0.001, 7.49 + 0.01, and 12.61 + 0.02, respectively.
5
Conclusion
Our work introduces Stochastic Optimal Control Matching, a novel Iterative Diffusion Optimization technique for stochastic optimal control that stems from the same philosophy as the conditional score matching loss for diffusion models. That is, the control is learned via least-squares problem by trying to fit a matching vector field. The training loss is optimized with respect to both the control function and family of reparameterization matrices which appear in the matching vector field. Optimizing the reparameterization matrices reduces the variance of the matching vector field. Experimentally, our algorithm achieves lower error than all existing IDO techniques in four settings.
One of the key ideas for deriving the SOCM algorithm is the path-wise reparameterization trick, novel technique to obtain low-variance estimates of the gradient of the conditional expectation of a functional of random process with respect its initial value. An interesting future direction is to use the path-wise reparameterization trick to decrease the variance of the matching vector field for diffusion models. The main roadblock when we try to apply SOCM more challenging problems is that the variance of the factor a(v, X" B) explodes when f and/or g are large, or when the dimension d is high. The large variance of due to the mismatch between the probability measures induced by the learned control and the optimal control, and it decreases as the learned control approaches the optimal control.
The research presented is foundational, but it may serve as the basis of algorithms that improve the quality of generative models.

---

Quadratic Ornstein Uhlenbeck, easy
Quadratic Uhlenbeck. easy
(ounta)
(ours))
522C latton)
550133 Hy (abracion)
SOCM-Ad (abiation))
12F
Adgoint (aclation)
10r
AAdjoitt
Entrusy
Cradlent sppt. (TMAA c.C11
Croea Entrrpy
LLng
-24
LLog
Normntt
Moment
20"7g
2DP
10a4
20000
30000
10000
Quedratic Ornstein Utlenbeck. herd, no warm start (d = 20)
Quadretic Orsstein Unlenbeck. hard. no warm start (df 201
53CM (oun)
5012 (sur))
5DCH (abaern)))
SOCM (atrarion)
SOCH-Ad (ablettont
10a
50CM-Agpoie (actation)
107
ee--Adjjint
Entmay
Crosa Entropy
Variance
Noment
Crudent eOm (T744)
Homent
20*
Contrel
TOP
12
10-3
127
10000 20000 30000 40000
50000
60000
70200
80000
20000 20000
30003
4000C
50002
50300
70000
53000
Num. bereians
Narrs Deriations
Figure 2: Plots of the 1.2 error incurred by the learned control (left), and the norm squared of the gradient with respect to the parameters 0 of the control (right), for the QUADRATIC ORNSTEIN UHLENBECK (EASY) (top) and (HARD) (bottom) settings and for each IDO loss. Both plots show exponential moving averages.
Linear Grastein Uhlenbeck id 10)
Double
10
SDCH (ours)
50CM lours)
-SCC As (ablar)
---SCCH Ni (abracion)
S-C.-OOMMMMA Arr Cablation
..... -0C35 Adgoirt (ablation)
10r
AAgeent
---Adjolrt Antjoint
CCrss Entrosy
---Crose Entropy
-Log Varionce
LLogog Varfance
Morrant
Var 1LM
20*
10000
20020
30000
4000C
600CB
13020
20000
50000
60000
BO000
Borations
Nam.
Figure 3: Plots of the 2 error of the learned control for the LINEAR ORNSTEIN UHLENBECK and DOUBLE WELL settings.
Funding disclosure Funded by respective author affiliations.
References
[1] Michael S. Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants, 2022.
[2] Michael S Albergo, Nicholas M Boffi, and Eric Sanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08797, 2023.
[3] Jonathan Baxter and Peter Bartlett. infinite-horizon policy-gradient estimation. Journal of Artificial Intelligence Research, 15:319-350, 2001.
10

---

[4] L'ubomíí Bañas, Herbert Dawid, Tsiry Avisoa Randrianasolo, Johannes Storn, and Xin Wen. Nu- merical approximation system of hamilton-jacobb aucobi-bellman equations arising in innovation dynamics. Journal of Scientific Computing, 92, 2022.
[5] Christian Beck, Sebastian Becker, Philipp Grohs, Nor Jaafari, and Arnulf Jentzen. Solv- ing stochastic differential equations and Kolmogorov equations by means of deep learning. arXiv:1806.00422. 2018.
[6] Christian Beck, Fabian Hornung, Martin Hutzenthaler, Arnulf Jentzen, and Thomas Kruse. Overcoming the curse of dimensionality in the numerical approximation of Allen-Cahn partial differential equations via truncated full-history recursive multilevel Picard approximations. arXiv: 2019.
[7] Sebastian Becker, Patrick Cheridito, and Amulf Jentzen. Deep optimal stopping. Journal of Machine Learning Research, 20, 2019.
[8] Andrea Belloni, Luigi Piroddi, and Maria Prandini. A stochastic optimal control solution to the energy management of a microgrid with storage and renewables. In 2016 American Control Conference (ACC), pages 2340-2345, 2016.
[9] Julius Berner, Lorenz Richter, and Karen Ullrich. An optimal control perspective on diffusion- based generative modeling, 2023.
[10] Joris Bierkens and Hilbert Kappen. Explicit solution of relative entropy weighted control. Systems & Control Letters, 72:36 2014.
[11] J. Bonnans, Elisabeth Ottenwaelter, and Hasnaa Zidani. A fast algorithm for the two dimensional hjb equation of stochastic control. M2AN. Mathematical Modelling and Numerical Analysis. ESAIM, European Series in Applied and Industrial Mathematics, 38, 07 2004.
[12] Elisa Calzola, Elisabetta Carlini, Xavier Dupuis, and Francisco Silva. A semi-Lagrangian scheme for Hamilton-Jacoa Numerische Mathematik, page 153, 2022.
equations with oblique derivatives boundary conditions.
[13] E. Carlini, A. Festa, and N. Forcadel. A semi-Lagrangian scheme for Hamilton-Jacoii-Belln equations on networks. SIAM J. Numer. Anal., 58(6):3165-3196, 2020.
[14] René Carmona. Lectures on BSDEs, stochastic control, and stochastic differential games with financial applications, volume SIAM, 2016.
[15] René Carmona and Mathicu Laurière. Convergence analysis of machine learning algorithms for the numerical solution of mean field control and games i: The ergodic case. SIAM Journal on Numerical Analysis, 59(3): 19(3):1455-1453 1455-1485, 2021.
[16] René Carmona and Mathieu Laurière. Convergence analysis of machine learning algorithms for the numerical solution of mean field control and games: li finite horizon case. The Annals of Applied Probability, 32(6):4065-4105 2022.
[17] René Carmona, François Delarue, et al. Probabilistic Theory of Mean Field Games with Applications 1-Il. Springer, 2018.
[18] Quentin Chan- Wai-Nam, Joseph Mikael, and Xavier Warin. Machine learning for semilinear PDEs. Journal of Scientific Computing, 79(3):1667-1712, 2019.
[19] Pratik Chaudhari, Adam Oberman, Stanley Osher, Stefano Soatto, and Guillaume Carlier. Deep relaxation: partial differential equations for optimizing deep neural networks. Research in the Mathematical Sciences, 5(3):30, 2018.
[20] Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. In Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc 2018.
[21] Kristian Debrabant and Espen R. Jakobsen. Semi-lagrangian schemes for linear and fully non-linear diffusion equations. Mathematics of Computation, 82(283):1433-144 2013.

---

[22] W. E, Jiequn Han, and Amullf Jentzen. Deep learning- based numerical methods for high- dimensional parabolic partial differential equations and backward stochastic differential equa- tions. Communications in Mathematics and Statisticx, 5(4))349-33) 2017.
[23] Weinan E, Jiequn Han, and Arnulf Jentzen. Algorithms for solving high dimensional pdes: from nonlinear monte carlo to machine learning. Nonlinearity, 35(1):278, 2021.
[24] Jin Feng and Thomas G Kurtz. Large deviations for stochastic processes. Number 131. American Mathematical Soc., 2006.
[25] Wendell H Fleming and Jerome L Stein. Stochastic optimal control, international finance and debt. Journal of Banking & Finance, 28(5):979-99) 2004.
[26] Paul Glasserman and David D. Yao. Some guidelines and guarantees for common random numbers. Manage. Sci., 38(6):884 908, jun 1992.
[27] Peter W. Glynn. Stochastic approximation for monte carlo optimization. In Proceedings of the 18th Conference on Winter Simulation, page 356 365..65.4 Association for Computing Machinery, 1986.
[28) Emmanuel Gobet. Monte-Carlo methods and stochastic processes: from linear to non-linear. CRC Press, 2016.
[29] Emmanuel Gobet and Rémi Munos. Sensitivity analysis using Itô Malliavin calculus and martingales, and application to stochastic optimal control. SIAM Journal on control and optimization, 43(5):1676 1713, 2005.
[30] Emmanuel Gobet, Jean- Phillpppp Lemor, Xavier Warin, et al. regression- -basdd Monte Carlo method to solve backward stochastic differential equations. The Annals of Applied Probability, 133222212-2 2005.
[31] Vicenę Gómez, Hilbert Kappen, Jan Peters, and Gerhard Neumann. Policy search for path integral control. In Joint European Conference on Machine Learning and Knowledge Discovery in Dutabases, pages 482-497. Springer, 2014.
[32] Alex Gorodetsky, Sertac Karaman, and Youssef Marzouk. High-dimensional stochastic optimal control using continuous tensor decompositions. International Journal of Robotics Research, 37(2-3), 2018.
[33] Constantin Greif. Numerical methods for hamitoon-jacobi-bellnna equations. 2017.
[34] Jiequn Han and Ruimeng Hu. Deep fictitious play for finding markovian nash equilibrium in multi-agent games. In Mathematical and scientific machine learning, pages 221-245. PMLR, 2020.
[35] Jiequn Han, Arnulf Jentzen, and W. E. Solving high-dimensional partial differential equations using deep learning. Proceedings of the National Academy of Sciences, 115(34):8505-831 -5510, 2018.
[36] Carsten Hartmann and Christof Schütte. Efficient rare event simulation by optimal nonequi- librium forcing. Journal of Statistical Mechanics: Theory and Experiment, 2012(11):PP 1004, 2012.
[37] Carsten Hartmann, Ralf Banisch, Marco Sarich, Tomasz Badowski, and Christof Schüitee. Characterization of rare events in molecular dynamics. Entropy, 16(1):350-376, 2014.
[38] Carsten Hartmann, Lorenz Richter, Christof Schütte,, and Wei Zhang. Variational characteriza- tion of free energy: Theory and algorithms. Entropy, 19(11), 2017.
[39] Carsten Hartmann, Omar Kebiri, Lara Neureither, and Lorenz Richter. Variational approach to rare event simulation using least-squares regression. Chaos: An Interdisciplinary Journal of Nonlinear Science, 29())063107, 2019.
[40] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, volume 33. Curran Associates, Inc., 2020.
12

---

[41] Lars Holdijk, Yuanqi Du, Ferry Hooft, Priyank Jaini, Bernd Ensing, and Max Welling. Stochastic optimal control for collective variable free sampling of molecular transition paths, 2023.
[42] James E. Hutton and Paul Nelson. Interchanging the order differentiation and stochastic integration. Stochastic Processes and their Applications, 18(2):371-377 1984.
[43] Martin Hutzenthaler and Thomas Kruse. Multilevel picard approximations of nigh-dimensional semilinear parabolic differential equations with gradient-dependent nonlinearities. SIAM Journal on Numerical Analysis, 58(2):929 9010999611110 2020.
[44] Martin Hutzenthaler, Arnulf lentzen, Thomas Kruse, et al. Multilevel picard iterations for solving smooth semilinear parabolic heat equations. arXiv preprint arXiv: /607.02 03295, 2016.
[45] Martin Hutzenthaler, Arnulf Jentzen, Thomas Kruse, Tuan Anh Nguyen, and Philippe von Wurstemberger. Overcoming the curse of dimensionality in the numerical approximation of semilinear parabolic paralooi partial differential equations. arXiv:/807.01212, 2018.
[46] Martin Hutzenthaler, Arnulf Jentzen, and Thomas Kruse. Overcoming the curse of dimension- ality in the numerical approximation of parabolic partial differential equations with gradient- dependent nonlinearities. arXiv: 2019.
[47] Max Jensen and lain Smears. On the convergence of finite element methods for hamilton- jacobi-bellman equations. SIAM Journal on Numerical Analysis, 51(1)1337116,, 2013.
[48] Hilbert J Kappen, Vicenç Gómez, and Manfred Opper. Optimal control as a graphical model inference problem. Machine learning, 87(2):159-183 110:1199--122000 2012.
[49] Hilbert Johan Kappen and Hans Christian Ruiz. Adaptive importance sampling for control and inference. Journal of Statistical Physics, 162(5):1244-1266) 2016.
[50] I. Karatzas and S. Shreve. Brownian Motion and Stochastic Calculus. Graduate Texts in Mathematics (113) (Book 113). Springer New York, 1991.
[51] Patrick Kidger, James Foster, Xuechen Li, Harald Oberhauser, and Terry Lyons. Neural sdes as infinite-dimensional gans. In International Conference on Machine Learning, 2021.
[52] Harold Kushner and Paul G Dupuis. Numerical methods for stochastic control problems in continuous time, volume 24. Springer Science & Business Media, 2013.
[53] Pierre L'Ecuyer and Gaétan Perron. On the convergence rates of ipa and fdc derivative estimators. Operations Research, 42(4))643-656, 1994.
[54] Xuechen Li, Ting-Kam Leonard Wong, Ricky TQ Chen, and David Duvenaud. Scalable gradi- ents for stochastic differential equations. In International Conference on Artificial Intelligence and Statistics, pages 3870-3882. PMLR, 2020.
[55] Yaron Lipman, Ricky T. Q Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling, 2022.
[56] Guan-Horng Liu, Yaron Lipman, Maximilian Nickel, Brian Karrer, Evangelos A. Theodorou, and Ricky Q Chen. Generalized schrödinger bridge matching, 2023.
[57] Jingtang Ma and Jianjun Ma. Finite difference methods for the hamilton-jacobi-bellne equations arising in regime switching utility maximization. Sei. Comput., 85(3):55, 2020.
[58] Sanjoy K Mitter. Filtering and stochastic control: A historical perspective. IEEE Control Systems Magazine, 16(3):67-76. 1996.
[59] Nikolas Nüsken and Lorenz Richter. Solving high-dimensional Hamilton-Jacobi Jaw.Jaab----llaannn pdes using neural networks: perspectives from the theory of controlled diffusions and measures on path space. Partial differential equations and applications, 2:1-48, 2021.
[60] Bernt Oksendal. Stochastic differential equations: an introduction with applications. Springer Science & Business Media, 2013.
13

---

[61] Derek Onken, Levon Nurbekyan, Xingjian Li, Samy Wu Fung, Stanley Osher, and Lars Ruthotto. neural network approach for high- sigh-dimensional optimal control applied multiagent path finding. IEFE Transactions on Control Systems Technology, 31(1):235-255) j 2023.
[62] Grigorios A Pavliotis. Stochastic processes and applications: diffusion processes. the Fokker- Planck and Langevin equations, volume 60. Springer, 2014.
[63] Huyên Pham. Continuous- time stochastic control and optimization with financial applications, volume 61. Springer Science & Business Media, 2009.
[64] L.S. Pontryagin. The Mathematical Theory of Optimal Processes. Interscience Publishers, 1962.
[65]
Aram-Alexandre Pooladian, Heli Ben-Hamu, Carles Domingo-Enrich, Brandon Amos, Yaron Lipman, and Ricky T. Q. Chen. Multisample flow matching with optimal transport couplings. International Conference on Machine Learning, 2023.
[66] Warren B. Powell and Stephan Meisel. Tutorial on stochastic optimization in energy- part i: Modeling and policies. IEEE Transactions on Power Systems, 31(2):1459-1467, 2016.
[67] Konrad Rawlik, Marc Toussaint, and Sethu Vijayakumar. On stochastic optimal control and re- inforcement learning by approximate inference. In Tweny-Third International Joint Conference on Artificial Intelligence, 2013.
[68] Sebastian Reich. Data assimilation: The Schrödinger perspective. Acta Numerica, 28:635-711 1, 2019.
[69] Martin I. Reiman and Alan Weiss. Sensitivity analysis for simulations via likelihood ratios. Oper: Res., 37:830-844, 1989.
[70] Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows In Proceedings of the 32nd International Conference on Machine Learning. 2015.
[71] Lorenz Richter and Julius Berner. Improved sampling via learned diffusions. In The Twelfth International Conference on Learning Representations, 2024.
[72] Geoffrey Roeder, Yuhuai Wu, and David K Duvenaud. Sticking the landing: Simple, lower- variance gradient estimators for variational inference. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
[73] Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assie Assister-Assisted Intervention-MICCAA 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18, pages 234-241. Springer, 2015.
[74] Reuven Rubinstein and Dirk Kroese. The cross-entropy method: a unified approach to combinatorial optimization, Monte-Carke simulation and machine learning. Springer Science & Business Media, 2013.
[75] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. arXiv preprint 1eXiiv:907..5000 2019.
[76] Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations (ICLR 2021), 2021.
[77] Evangelos Theodorou, Freek Stulp, Jonas Buchli, and Stefan Schaal. An iterative path integral stochastic optimal control approach for learning robotic tasks. IFAC Proceedings Volumes, 44 (1):11594-11601, 2011. 18th IFAC World Congress.
[78] Ramon Van Handel. Stochastic calculus, filtering, and stochastic control. Course notes, URL http:///ww, prince- ton. edud/van/acmm2277/cMMI1 2007.
[79] Francisco Vargas, Will Sussman Grathwohl, and Arnaud Doucet. Denoising diffusion samplers.
In The Eleventh International Conference Learning Representations, 2023,
14

---

[80] C. Villani. Topics in Optimal Transportation. Graduate studies in mathematics. American Mathematical Society, 2003.
[81] C. Villani. Optimal Transport: Old and New. Grundlehren der mathematischen Wissenschaften. Springer Berlin Heidelberg, 2008.
[82] Jichuan Yang and Harold J. Kushner. A monte carlo method for sensitivity analysis and paramet- ric optimization of nonlinear stochastic systems. SIAM Journal on Control and Optimization, 29(5):1216- 12t112216-119,,,,, 1991.
[83] Jianfeng Zhang et al. A numerical scheme for BSDEs. The annals of applied probability, 14(1): 459-488, 2004.
[84] Qinsheng Zhang and Yongxin Chen. Path integral sampler: A stochastic control approach for sampling. International Conference on Learning Representations, 2022.
[85] Wei Zhang, Han Wang, Carsten Hartmann, Marcus Weber, and Christof Schüitee Applications of the cross-entropy method to importance sampling and optimal control of diffusions. SIAM Journal on Scientific Computing. (166))2265-A26722 2014.
[86] Mo Zhou, Jiequn Han, and Jianfeng Lu. Actor-critic method for high dimensional static Hamilton-lacoiiBellnn partial differential equations based neural networks. SIAM Journal on Scientific Computing, 44043:4600 2021.
15

---

Contents
Introduction
2
Frameworki
2.1
Setup and Preliminaries
2.2
Existing approaches and related work
Stochastic Optimal Control Matching
Experiments
5
Conclusion
Technical assumptions
16
Proofs Sfe
16
Proofs Sfe
21
C.1
Proof of Thm. and Prop. 2
21
C.2
Proof of the path-wise reparameterization trick (Prop.
24
Informal derivation of the path-wise reparameterization trick
26
C.4
SOCM-Adjoint: replacing the path wise reparameterizzttin trick with the adjoint
method
28
C,5
Proof of Lemma
30
D Optimal reparameterization matrices
30
Control warm-starting
32
Experimental details and additional plots
34
F.1
Experimental details
34
Model architectures
35
Additional tables and plots
35
A Technical assumptions
Throughout our work, make the same assumptions as (59), which are needed for all the objects considered to be well-defined. Namely, we assume that:
(i) The set 24 of admissible controls is given by
U {u E C² (IR" x (d.T):Rd) EC O, V(z,s) Hd x [0, b(z, C(1 + [z])
(ii) The coefficients and a are continuously differentiable, a has bounded first-order spatial derivatives, and (ao (z, s) is positive definite for all (x, s) Rd [0, T). Furthermore, there exist constants C, C1, C2 > 0 such that
|b(, s)1l < C(1 + |
(linear growth) (ellipticity)
ea 1 << [zz s){ < 1211
for all (x E Rd x [) and g E Rd.
B
Proofs of Sec. 2
Proof of 5 By Itô's lemma, we have that
V v(X#, - - S:" (a.V(X), + (6XX s)(3)) VV(X;,
+ Ei (10 v(Xy, ds + Sp,
16

---

where Sp -K= -V VV s)r oxx dB, Note that by 4
( 8) + (6(X", + + s) ( s), VV(X,, 8))
+ \/5 0,,
1 ((x))))))00
s)(() (V ())
_1( +)
[1 3)|| - ,)
and this implies that
- - -
S ( VV) + s)1| 11lu(XX s)|| - f(X , 8)) ds + i
(26)
Since E[S Xi = x| = 0, rearranging 26 and taking the conditional expectation with respect to X yields the final result.
Proof of (6--7) By Itô's lemma, we have that
dV (X., s) (X,V 9) + (b(X., (V(),,
2/ (oo= s)0,, Oz, V(X,, ds VAVV(X), so oxx,, dB,,
(27)
Note that by ④
0, V, + VV + E Jy(X,, s)0,, 0z, ((...)
-V) |) f(x,, s).
Plugging this into 27 concludes the proof.
Proof of 8 Since Y, = V(X,, 5) and Z, = =s=(s)VV(X,,,) -u" satisfy we have
that
g(Tr) = Yr = Y (f(X,.s) ( b s)| ds - VA S1 (w s),
Hence, recalling the definition of the work functional in TOP, we have that
W(X,1) = Y, + 2 1e" ||u² (X., ds vx s),
(28)
By Novikov's theorem (Thm. 2p. have that
([exp(>>->W(x ())
=e- Y exp (X-1/2 1/2 S. (X,, s), xx Se ||u" ds)
X,] X.
=
which concludes the proof of (8)
Theorem (Novikov's theorem). Let 02 a locally-Hi H2 process which is adapted to the natural
filtration of the Brownian motion (Bi).2o Define
Z(t) exp (So 0, dB. 56 |1.,1 ds).
(29)
for each 2 0,
E exp (Jo ||0,11 [[s]] < +00,
then for each 0,
E[Z(t)] =
(30)
Moreover, the process Z(t) is a positive martingale, i.e. if (Fi) is the filtration associated to the Brownian motion (B:) then for t 2 s, [.] = Z..
17

---

Theorem 3 (Girsanov theorem). Let W = (WD)ejo.T] be standard Wiener process, and let IP be its induced probability measure over TI; R2 known as the Wiener measure. Let Z(t) be as defined in (29 and suppose that the assumptions of Theorem 2 hold. Let (D, be the a-algebra associated to BT. For any E F, define the measure
((F)
(T)
Q is probability measure because of (30 Under the probability measure Q, the stochastic process
[W (1)) defined as
((I) ((1) So 0, ds
is a standard Wiener process. That is, for any 2 0 and any 0 = < 41 <... < ty the increments (W+++1) - (t:)) are independent and Q Ga-Ssisn distributed with mean zero and covariance (te+1 t,)L, which means thatfor any Rd, the moment generating function of W (t (t.))
with respect to Q is as follows:
Eq (exp((a, ((++1)) ((t))))
:= Ep exp (aa, W (t++1)) - fo fli 0g ds - ((t.) + So² 0g ds)) ((T)) = exp
(
2
Corollary (Girsanov theorem for SDEs). the two SDEs
dX, b1 t) + t)
Xo = Kinit
dY, = (Y,,), + b2(Y,, 1)) dt + oY,, , 1) d),
Yo= =Initt
admit unique strong solutions on [0, T). then for any bounded continuous we have that
foniiiiuuuuuuuoooiiioo functional & C([0, T)),
E[((()) = EE (Y exp (- So (2 3.3 100Y,,0)%%,,,,,,111
(Y) exp S0 (y, + 1/2 So tb t) dt)].
where B³ = Bi + 1n¹ oY.. S) b2 s) ds. More generally, b, and b2 can be random processes that are adapted to filtration of B.
Lemma 2. For arbitrary
u, let Ip and P be respectively the laws of the SDEs
axς (b(Xg, () + (X)(X)) 1)) di VTo(I)dBr, XØ po: dx;- = b(X,, t) dt + (io(t)dB. Xo
We have that
dP diee
(X") = exp
- X- 1/2 1o- (v(XF, t),")) + x-1 fod [I((X, - -1/2 1or (v(XF, t), dB,) x-1 --o' | t)I
(31)
= exp
d de (X) exp () 1/2 f0" (u(Xe, e, t), dB.) =2
x-1 So" [lo(X,,))) dr).
(32)
where :² B, - 1/2 J 0 v(Xg 5) ds. For the optimal control u", have that
dP upp (x cxp (> V(X;" 0) + W(X" ())),
(33)
dp=
dP (X) exp
(34)
where the functional W is defined in 10
Proof. The proof of 31 32 follows directly from Cor. 1
To prove 34 we use that by 28).
W(X,0) V (Xo. 0) + 1 5 I|u s) ds v (u (X..
(35)
which implies that
(X) = exp (X- 10 (x,, (), - - So (x,,)))))) cxx (X- (V 0) - WWX,0))
de
18

---

To prove 33 we use that since 'xp"" b(X;" dt + (\o(t)dB)" equation (35 holds if we
replace X and B by X" and B" which reads
W(X" .0 = V(Xg .0 + S (x, ()) ds - vAJi Se (w" ,s dB.)
Hence,
= exp ( - - 50 (xF (,,B; + + .013 (x"))=exp
= exp (A-1 V 0) + W(X" (0))).)))
Lemma 3. The following expression holds;
Eps log
= - "[J'"" () + t(),,)) dt + 9 0)] (36)
d
Proof. To prove 36 we write
log dp=
(X") log
AP" dF
(X (x dP (X*)) = log dPn 2
dF (X") + log 12 (X")
= (V(X8,0) - So ((xr,))d (00-g(()))
1/2 So (u(XF, 1), dBB) - - So [|((X;,))] ()
Since Epp [log
- Ep-
[lo
de and Ep. [ (u(XP, dBM 0. the result follows.
Proposition 3. (i) The following two expressions hold for arbitrary controls u, in the class U of admissible controls:
(w) Eg log dP"'r = x-112 50 (), - - (u(X), t), o( f)) dt
(37)
+ - Jo |uu( [x;
0) - w(x",0))
x cxp () (V 0) - W(X", 00)
-1-1/2 50 (w( (X; 1), d - 232 So O)
Éce(u) = A 11 (XF', ,0) .0) u(xp' ,1)l" de].
(38)
When po is concentrated at a single point Finit. the terms V (Iinit, 0) are constant and can be removed without modifying the landscape. other words, LCe and LCE are equal up to constant terms and constant factors.
(ii) When po is a generic probability measure, LCF and LCv, have different landscupes, and Lcr(u) (p O]]]] u" is still the only minimizer of the loss LCE, and for some
ap" Epu [log new exp( constant K, we have that
LCr(u, 0) L ||ue (XP", .1) u(Xy" 1)11 exp (- ----V .0)] + K. (39)
Proof. We begin with the proof of (i), and prove (37 first. Note that by the Girsanov theorem TTh. 3).
Eps-
dp" [log de (X=))
= Epu-
llog
dp" (X")) )]
- Epe-
dp" [log
(X" log dP (X"")
de
d= (40)
=-Epo (log dr" (Xx) + log pe-
dP
(xx)) dp=
(X") (X")
Note that by equations
(32 and 33
log
dP"" dF (X*)= = 112 0ed (u(XF, (u,X) dB.)) 2/2 of [1(
= -12 1o_ (u(Xy, t), dB,) + - fo (u(X), t), o(Xy, dt 1 fo [x(Xy t)ll
1/2
at,
AP log JP = - (- V(gg, 0)
W(X)0)).
(41)
19

---

where B = B₁ + 112 So v(Xg, 8) ds. Also,
1/2
62 de (Xx) exp (A (V 0) W(X", ())),
(42)
dP (X*) = exp - - 1/2 So¹ (o(X), (), dB.)) - - So |e((X; (o(x),))1 dt).
doe
we plug (11 and 42 into the right- handt-handdee side of (40), we obtain
Epp
log dp"
Bpp- [1lg 4f (X*) log Tith (X (x²) (X))
=- 1/2 1oe (), + - lo (u(X), t) dt
- Sod
-- W(X",
x exp (X (V 0) w(x", 0)) -- 1/ (u (), dB.) - x-1
2 So (0x
which concludes the proof.
To show (38. we use that by Cor.
exp (-x 1/2 Lo (XF .t) uxp - So [xx (()-((x'
(t)x ,)]l
Hence,
Ep-- [log
dr" dF"
Ep..
[log
drt d= 2
x-1 [ |x'(X] .) - u(X²", ,)]l dt].
Next, we prove (ii). The first instance of V(X. 0) in 377 can be removed without modifying the landscape of the loss. Hence, we are left with
LCE(U)= = E - X
1/2 11/2 So (u(X), (u(X).),dde. dB;) - x-1 1o' (u(X;, (), o(X, ()) dt
(43)
+ + oo x exp (A-1
- 1/2 fo (), - 1-1 S0 1xx (x((2
(1)]
And this can be expressed as
ECE(U) = =[g((; X8) exp V 0)),
where
g(u:z) = [ So" (u(X) dB,) - - fo (u(XP, t), u(X), t)) dt
+ A 5fT d 0))
X exp ----' ((xx)))--000 00 Se (u t), dB,) 1/T J0 lI((X +)ll x2]
A-1
If we consider g(u; x²) as a loss function for u. note that it equivalent the loss LCg(u) equation in 43 for the choice po 5,, i.e., Po concentrated at Since the optimal control u² is independent of the starting distribution Po. we deduce that u* is the unique minimizer of g(u; x:), for all TR... In consequence, u" is the unique minimizer of LCE(u) ([g(w; Xi)).
To prove 39.. note that up constant term, the only difference between (Ce(u) and (CE(u) is the expectation is reweighted importance weight exp A - ((x 0)).
Lemma 4. (i) We can rewrite
Evar. (u) = Var( exp (F - x x1 ((FF) + - -
Var. (u) Var Var(Yz (Yr," +X
0)).
When po concentrated at Dinit, the terms VI 0) are constants and can be removed without modifying the landscape. In other words, Lvvar, and fly are equal to Loaar and flog up constant term and constant factor, respectively.
20

---

(i) When is general, Evar, and Vaa, have different landscape, and the optinum of Lvar, may be different from u². related loss that does preserve the optimum is:
Evar. (u) E, Vare-
dP" dF" (xx)) = E. Var( exp(Y1 X ((xx)))x)
0))
In practice, this is implemented by sampling the m trajectories in one batch starting at the same point X@-
Elog and clog have different landscape, and the optimum of C may be different (ii) Also, Var. from u In particular, L (u) Vary= (log
dr""
A loss that does
dpn
preserve the optimum u is
Blog
(u)
E Varp (log dr
de" (xx))
0))
E[Var (y x 1 ((x))(x)))
Proof. Using 34 and 31
we have that
dP" (X²)= = exp (A- (V W(X m()),
dpn (Xx) = exp ( (-2-1/2 1/2 So (u(X;,)) dB,)) + x Jn [lu(X;,)))
exp
( x 1/2 -1/2 So (m(X;, (), -X (u(Xi, oXF, di
+ x So [u a(l
Hence,
dpn log dpu
(X") = log dp=
(X") + log dpu (X*) = y
dP
X - + + 0).
dp
Since Larr, (u) Varp-
d dim
and alog (u) Var (log dr" Var,
dF this concludes the proof of (i).
To prove (ii), note that for general pu. V(X;, 0) is no longer a constant, but it is if we condition on Xj. The proof of (iii) is analogous.
C Proofs of Sec. 3
C.1 Proof of hm. and Prop. 2
We prove Thm. 1 and Prop. at the same time. Recall that by ()) the optimal control is of the form (z. (t) (x. t). Consider the loss
É(u) [E lor |lu(X,, +o( (x) exp (-x-1 la f(X,,t) X A-'g(XT))]]
Clearly, the unique optimum of 2 (t) VV. We can rewrite Ż as
Ž(w) E[+ 5od ([lu(X,, 1) + 2m X (1) VV(X,, 1) + x exp (- 5-' tt -t
VVX( 011) +lo() al(44)
Hence, we express E as of three terms: involving [|m(X,, t) another involv- ing (u(X,, t), (t) t)), and third one, which is constant with respect u, involving (X, 2)... The following lemma provides an alternative expression for the cross term:
Lemma 5. The following equality holds:
E 17 (u, T (V)(X,,)) dt exp (- x-1 fo- dt V.E = - So (), t), (V. exp (- 1ed f(X, xx'--x)9XX)) Xt = =2(445)
exp (- -1 5/ f(x ds) dt].
21

---

Proof. Recall the definition of W(X, t) in 35 which means that
W(X, 0) W(X, t) + So f(X,, 8) ds.
(46)
Let (Fr) be the filtration generated by the Brownian motion B. Then, equation ( implies that
att) T tV,, = Aa(()
V exp
(-x !W(X.))) Fc
(47)
2 exp
(-X (>x-1W(X,0)
5i
We proceed follows:
E + 1o (u(X,, t). a(t) VV
(V) t)) dt exp --x X-1
AE T 10 (u(X,,t), t); o1)) V.B
V.E exp (- (-x- |x
exp (-X-1W(X.+))
x (= Fi] exp ( - x-1 So s(Xs; 8) ds) dt]
= - 00' (u(x).)),o exp ( - 6)) [i]) exp -- J0 f(x,: 8) ds) dt]
n [E Oo (1) exp ( (-X-
x =]}> exp (-X-1 So f(x 8) ds) at).
Here, (i) holds by equation 47 the law of total expectation and equation (46). and (ii) holds by the Markov property of the solution of SDE.
The following proposition, which we prove in Subsec. C.2 provides an alternative expression for
V.E[ exp (-x-+ S." ((X,,,) ds - (Xr)) X, fx=xi x]- The technique, which is novel and we denote by Girsanov reparamaterizatioo trick, of independent interest and may be applied in other settings, as we discuss in Sec. 5)
Proposition (Path-wise reparameterization trick for stochastic optimal control). For each E T7. let Me : [t. 27 + Rdxd be an arbitrary continuously differentiable function matrix-valued functioooo such that (t) = Id. We have that
V E exp
[exp (- (-x-1Sh Se" x
E x-1 1:" (())
(23)
+ 11/ 5:
exp (-x-1 1/" -x-''((XT))) Xr =x].
Plugging 23 into the right-hand side of 45 we obtain that
E [ (m(X),,),,(1)) VV ()) di exp (-x-- So di - 9Xx))
= = for ((1) (J. ((..)V..(X..
3) ds g(( (s)dB.) dt x exp (-x-1 So t(),,t)dt x ''(()))))
1/2 (M
If we plug this into the right-hand side of (44 and complete the squared norm, we get that
[()) E fo (IIu( X,, 1) w(t, X,, B, M.) || - (t, X, B, MI]l
+ +||u² (Xe, (X dx (-x --xx-W((X0))
where wi is defined as:
Xt (tt -)
+ (. "
We also define b(u; X, B) as
( X, B) = (1 @lt x, X, B. 1)
22

---

Now, by the Girsanov theorem Thm. 3), we have that for an arbitrary control U,
B X, B) Bx) exp
-- 0))1
[I; X", Bd) exp - --X"" :0) - X-1/2 So (v(X), + + Lo |e((x; (1l a()
exp - -->---(x- 0) - X-1 Jo (u(X), 1), dB.) - 2 So |xx 2) ]t)],
where By = Bg + X-1/2 J0 xxx;; , ds. ds.
Reexpressing B in terms of B, we can rewrite
Buu; X" )d) and (t, X . _,, .M, as follows:
(uu X", B") = 1 10- [ X", M.) dt,
,(, B", M,) (1)) ( (- --i M())Va(XX) M(()VV)) + +112
s)dB.
+ (M -
Putting everything together, we obtain that
Ī(u) LSOCM (u, M) - K,
where C(u, M) is the loss defined in (note that wtt, U, X", B, M,) == Xt, ,B, M,), and
K = E= So (t, x, B, M) - ||, ()] dt exp (-x-Iw(X, 0))
To complete the proof of equation
24 remark that E(u) can be rewritten as
Ž(v) = [[+ 1or [u(Xe (u(),() w² (X,, 1)|l di exp ( x--x L + Sor [u(Xe,t) n² (X,, }) dt dP- exp (x) 0))]
= fo" |u(XF", t) - (X,' |) dt
O))].
only remains to reexpress K. Note that by Prop. I we have that
"² (Xe,t) E ([1(t..,BBMM) exp (-x.)
exp (-X 1W(X.0) F.
e ((t, 2OM
(X)
exp(-A VV(X..0) R( (x)
2F (X) exp(-> (V(X..0)
(X)
Xt, B"
Mi X X,]
Hence, using the Girsanov theorem Thm. 3) several times, we have that
K E + So [ittt B= .M) [tt B= x.) 1| dt
0))
(48)
E 7 5o |[ttt BB"
B" M.) Et X" B"" M) dt exp(-> 'Xx' .0)]
[E So [[tt,X, B, cxpl W(X X || dt exp(- - - 0))
-lexpi-x-WWC o)
=Ei fo [,,(,,,, v, X", B, E| X B, B.M, (",,x,,,,x"")XX BIX 2[a(1,X (B1x;
(l,a X)],
which concludes the proof, noticing that = Var(w; M).
Remark 1 (The optimal M, is the same for all v). Looking at equation 448 we observe that
Var(w; M) does not depend on the buse control v. Since minimizing LsoCM M) with respect to M is equivalent minimizing Var(w; M), we deduce that the optimal M does not depend on the vector field v.
23

---

C.2 Proof of the path-wise reparameterization trick (Prop.
We prove a more general statement (Prop. 4), and show that Prop. is a particular case of
Proposition 4 (Path-wise reparameterization trick), Let (R, IP) be a probability space, and B : Q x 0. [| + Rd be Brownian motion. Let :2 x [0, + R be the uncontrolled process given by 16. and let vb Q R2 x 00 T] Rd be an arbitrary random process such that:
For all & E RR, the process 10., 2. : 9 x [0,T] + Rd is adapted to the filtration (Fs)se(o,T) of the Brownian motion B.
For all w E 2,, (b(w,,,,' Rd x [0, 7] - IRd is twice-continuously difserentiable function such that (W 0) for all IRd, and (iw,, 0, for all E [0, T).
Let F: C((0,T):*)) I be Fréchet-differentiable functional. We use the notation X + (Zz.)) (X, (w) + (ww, 2, 8) s)),E( to denote the shifted process, and we will omit the dependency of ih on in the proof. Then,
V exp (- F(X)) Xo x
(49)
=B((-V (FX=V=((X+(+ === +A-1/2 (V=W(0.))V.MMX...> ( "
exp (- ((X)) xi
Proof of Prop. 1
Given family of functions (M))E(( satisfying the conditions in Prop. can define family (thx) (e) of functions 1/2 Rd [t. T) - I&d as (2, 8) M2(s) Note that (2, t) for all E Rd and (0.
for all E [t. 7J. and that (2, s) = (s). Hence, dh can be seen as random process which is constant with respect to w E Ωt and which fulfills the conditions in Prop. 4 up to a trivial time change of variable from [4, T] to [0,
We also define the family (Fi)rEjo,T) of functionals F. C(t,, ([ T): + R as F((X) = X-1 Se f(X,, s)
do + (X We have that
V. V.F((Xtu(Zz))
(>(x-' ! 1: (ss +hx @ X- //t.:///TT...........o/////////.
5: -
where equality (i) holds by the Leibniz rule. Using that (0, 0. we obtain that:
V. F.Fi(X + (2(z, 1>0 --1 S v v. (0, (0,
ds +
Up to trivial time change of variable from [t. T to [0. T). Prop. 1 follows from plugging these choices into equation (49).
We can use matrices M((s) that depend on the process X up to time s, since the resulting processes the ( 2, are adapted to the filtationn of the Brownian motion B. More specifically, if let M, IRd x [t, T + Raxd be arbitrary continuously differentiable function matrix-valued function such that M((x,t) t) = Id for all E Idd, and we define the exponential moving average of X as the process (() given by
Remark 2.
x( So -v(t-s) X, de,
we have that
d " ()) (VM s), dxlx) (0, -)) (.) s), x)).x.x(!')) (xx
da
24

---

and we can write
F exp - x- Se s) - X - = =
- X-1 1e M v 8) - ( (XT)
+ + 1: (M.(X.) " 8)) 1
x exp S-x-1 J(X 8) ds g Xx =
Plugging this into the proof of Thm. 1 we would obtain a variant of SOCM (Alg. 2 where the matrix-valued neural network Ma takes inputs 8,2) instead of (t, s)- Since the optimization class is larger, from the bias-variance in Prop. 2we deduce that this variant would yield a lower variance of the vector field w, and likely an algorithm with lower error. This is at the expense of an increased number of function evaluations (NFE) of Mai one would need K(K + 1mm NFE per batch instead of only K(K+1), which may be too expensive if the architecture of Ms is large. A way to speed up the computation per batch is to parameterize Me using cubic splines.
Proof of Prop. Recall that
dX, = b( VAa(s) dB,, Xo ~ po:
the SDE for the uncontrolled process. For arbitrary x, 2 E Rd, we consider the following SDEs conditioned on the initial points:
dx((+=) b(x! (1))) ds + VAo dB,
x x(*+2) = x xx
(50)
dx! =b(X!") 5) ds + (s) dB,,
x( I.
(51)
Suppose that 3 Rd x [0, + Rd satisfies the properties in the statement of Prop. 4 If (() is solution of
dx! = (b(x{ (x) + (2,,5),s) - (s) + (io(s)
x!( 2,
then X (x++2) = {1)) + (2, -) a solution of 50). This is because x x((+z) x((²) + (2,20) =
(a) + x + z, and
dx (x{=+)) = =xx(²) + 02 8) ds
(M(X!") + (z2 8), 8) 8)) (T) + 0.4 ds
b(x((t2) s) ds + (X)
Note that we may rewrite (1 as
dx((") (b(x!")) (2, s). 0,0 8)) ds
+ (b(x!!)) -bb(( + (22 s), s) + 0,ú(2, ds + Vio(s) dB,,
x(() ~ po
Hence, since (22 s) is random process adapted to the filtration of B, apply the Girsanov theorem for SDEs (Corollary) on (x) and (x), and we have that for any bounded continuous functional &
E ( (X{()))
6(x!*) () 0,b((z, 8)) dB, E (x(((((x)) (=)) exp (Lo )-122(()) + +(,(2,),,
1/2 T |1-1 (b(x!') + 6(x!=) s) 0. s)) ds)]
(52)
25

---

We can write
E exp ( -
+ = --ee((-F(xxett)) -(X(+t))) ( E exp ( F(X(()) + (2(
( E[ exp ( F(X + ((2)))
x exp (So' X (1)Fo(s))((((3) (0000 0,
- -5 _-Th ==,,,)) 0,
(v) exp (-F(X+ 3,,
dB de) Xo = x
- ( -
(53)
Equality (i) holds by the definition of x(*+=), equality (ii) holds by the fact x! X (2, 8).
equality (iii) holds by equation (52). and equality (iv) holds by the definition of x! We conclude the proof by differentiating the right-hand side of (53) with respect to z. Namely,
Vr V..[expp (-F(X)) Xa- =] V= ( (- F()(()) Xo + =] =]11=0
@ BI( - V.F(X + (z, + 1/2 So
v=0,v(0, s))oo 3)
(s)dB.) (T))dBB)
x exp ( - FXx) =
In equality (i) we used 53 and that:
by the Leibniz rule,
V lo [lo(s) (b(ss + (22 s), s) b(X,, - s)), ==0
= S V=llo(s) + (2,,),, s) - b(X,, - 0, (22 s)][|3 == ds = 0.
and by the Leibniz rule for stochastic integrals (see 42
V= (bx) + (12 s) - b(x,, 8) - 0, 8)) dB.)) ==0
= S (V.b(0.sVV,,b(X, - (.,0v((0))))0 1T dB,.
C.3 Informal derivation of the path-wise reparameterization trick
In this subsection, we provide an informal, intuitive derivation of the path-wise reparameteriza- tion trick as stated in Prop. 4 For simplicity, we particularize the functional F F(X) S-"S f(X,, s) ds + > (XT).. Consider the Muler-Maruyama discretization of the uncontrolled process X defined in (6). with K + 1 time steps (let = T/K K be the step size). This is a family of random variables X = (XK))= k= defined as
Xo ~ po: X = X2 + 36(8,,k0) ko) + Vo)o(kō)En
E ~ I).
Note that we can approximate
E [[exp( ( (- f0-
ds sls
=x]
E exp E[exp
1- 0 - x]
and that this is equality in the limit K + x0, as the interpolation of the Fuler-Maruyama discretization (x) converges the process (()), Now, remark that for k {O...., K 1}.
X & ~ N(Xx + 5b( Xx, k6), 8N (A((oT)((k)))) Hence,
E B[ex exp (- 1----87 1 -O f(Ex, 8) - X g((())))) ()) = x2]
= C-1 Massyy exp ( - x Ek= f 8) - x (((x))
1 Σ
K- (
- - 556( (xx, 12
20
k=
10 (00) - - 5b(2, 12) d.². dē
28A
26

---

(2m8AA) ( dot(( (kk))) ((((())))) Now, let vh R x [0. 7] + I be an arbitrary
where C = -=
twice differentiable function such that (²2, 0) = 2 for all E R2, and (0, s) = 0 for all 8 E [0. T]. We can write
VoLl
(x 1x =] X g (xx)) Xo = x x +
exp (- (-X-1 s(X,,)) X 9(X)))))))))))
= V. exp - x X
Ek= K-
(- exp ( x 2k=0 [ g
1
0 16 ko]]ll
2d1
2 -000)))00-))))-000000 10000 d 1==0
...
V. exp
g(0x + ((22 ",
28A 1221=1110 -
+ ob -
kō), I?
26A 10 (0000)10((22)000=00(000000)00-0)
dă ... ddx) ==0-
(54)
In the last equality, we used that for l E (1..... }). the variables St. are integrated over Rd, which means that adding an offset (1z k5) does not change the value of the integral. We also used that (1, 0) = 2. Now, for fixed values of a = (f1,..... ãx), and letting dn = x, we define
Gx(2)= = K-
f(ex + (1, k5), 8) + x g (),, (l- - 1b 15) 5b( A.))
+ 28x
K= (A5)
K-
Using that (0 = 0 for all s E [0, T) we have that:
G2 = X M f(tx, 5) + A g((x) + 28X ( -2x k6)] k2)
= = K-
k=0 (0, (0. (4
+ +
+ V10-00000-1100---700000000000100 ko))
x ((o σ
ko))
And we can express the right-hand side of (54) in terms of (0) and VG2 (=))
V=(C-1 1 [1 K exp G=(=)) dy ... dyk)
-- (R( ==) exp ( (- (0) dy dy.
We define EK = 10 0 110001((00 (k00) 5x - 86(; ko)). and then, we are able to write V8h
+ .x + Ölb(EA kd) + (kDoek
t0 =
(55)
(0) - Ek= [ 8)
+ D 12,
VGa(=))= 1 Σ == V., - Vv Ko) (
(56)
+V K (0, -00) Vb) (o)) ((:))ex
Then, taking the limit K + D (i.e. - 0), we recognize (55) as Mareraaa discretiza- tion of the uncontrolled process X in equation ( conditioned on X, x, and the last term in (56) as the Euler- Maruyama discretization of the stochastic integral X- 1/2 fo (a, (0 s)
(0. (0) (s))) dB,. Thus,
lim +0 V exp ( --- - f().s)) (9((x)
= EE - (- l T () (s)
1/2 1-' (V ().00ssVVVbbX,.. V.V
exp (-x-1 S/ 8) - x
dx Xo
Xo =x],
which concludes the derivation.
27

---

C.4 SOCM-Adjoint: replacing the path-wise reparameterization trick with the adjoint method
Proposition 5. Lel LSOCM-Adj Adj L2 x [0, + R be the loss function defined as
x(, B)] where X" is the process controlled by (i.e., dx, = (b(X,, + o (X, ()) (X d)B
(soCM-Ad(uu)- == E So [|x(Xy, (
and Xo ~ Po). a(v,X", B) is the importance weight defined in (20) and att, X") is the solution of the ODE
da(t)
dt = (xb((t,t)(() - (x)
a(T) Vg((z),
Ado has a unique optimum, which is the optimal control u".
Proof. The proof follows the same structure as that of Thm.. Instead of plugging the path-wise reparameterization trick (Prop. 1) in the right-hand side of (22h. we make use of Lemma 6jto evaluate V exp (-x-1 So" f(X,, t) dt 3-'g(Xn Xo- =x]. Particular cases of the result in Lemma have been used in previous works such as [54 [T. We present more general form that covers state costs f. as well as stochastic integrals. We also present simpler proof of the result based on Lagrange multipliers.
Lemma 6 (Adjoint method for SDEs). 154 57 Let X : 12x (0, + Rd be the uncontrolled process defined in (6), with initial condition Xo I. We define the random process a 9 x [0, + Rd such that for all w E 52, using the short-hand a(t) == a(w,t),
dor(w)
- Vab(X,(w), (a)((w)))))))))))))))))))))))( - (x)(((()),)) at V.N((x(w),)) ar (r(W) Veg(XT(w))
have that
Vo f(Xi(w),) t) dt + fo (h())(w))t), dB,) + (XX = 2] = E [[[o[w]]] V exp (- [(Xe(w),))d -5o (I(X)(w), (), dB.) - ((Xx(())) =1| -B[lo((w exp ( S0" f(Xe(w), dt So" (h(X)(w), t), dB,) ((XT(w))) (o(w) x]]
Proof. We will use an approach based on Lagrange multipliers. Define a process a Ω x (0, [) - Rd such that for any w E 9t a(w,,) is differentiable. For a given w E Z we can write
Sd' ((X)(w)) () dt + 1ed (M(X,(W)) dB.) + ((X((w)) = So ()) dt fo (M() f), +B))) -o ((x)))) - () - (() (.).))))))))
By Lemma 7, we have that
1o (ar(w), = (aT(W) X)(w)) - (an(w), Xo(w)) 10'( (X,(w), des (w)) dt.
Hence,
Va (55" ((Xi(w),tdttt (6 (M(, t), dB,) +
= =v2(So (() ( (M f), dB,) +
Xr())) (ao (o) for
+ (1) de
+ (((
= 0 (. V.
+ ( V (w)) dt
- XKK.6///55////000000444444440000 + d (w)
+ (xXr((\V ar(w)) + ao (w) + fo- (xX
28

---

In the last line we used that V voNo(w) = Vxz = If choose such that
dar (w) Vab(Xe Vab(X,(()))a((()) (x))( (V),((((w) d)d ar (w) V V=g((x((w))
then obtain that and by the Leibniz rule,
Ve f(X,(w), ff (M(,()))) dB,) + +(Xx(w))) - ao(w),
V Lo f(X.(w), ( (M(, f), dB)) + [((((w))
= V t( de (M(X)W)) dB.) + [[(x(w)))) Bloo(w))
and
V V (- (-S" f(x,(w), - (h(X,(W)) ( (), dB.) - ((XT(w))))
= - f(X,(w),t) dt (b(X)(w).)) dBi) + (X(r(wW))
exp (- S0¹ ((x((w),
(h(X)(())) t), dB,) ((XT(w)))
- exp S0 f(X,(w), dt f (h(x t), aBi) - s(Xr(w))))
Lemma 7 (Stochastic integration by parts, (60)). Let
dX, a, dt + b, dW?
dY, f. dt + gr dW?.
where au by fo 9t are continuous square integrable processes adapted to filtration (F)
and W1, W are Brownian motions adapted to the same filtration.. Then,
X,Y: -oo%o = 10' K,dY, + Y dX, + (dx,, dY.) - x. + So dx, (b, dw? dw?).
Remark 3 (Related work to the path-wise reparameterization trick: sensitivity analysis). As shown above, the adjoint method for SDEs an alternative to the path-wise reparameterization trick. Prior to 134y, an array of works developed methods to compute derivatives of functionals of stochustic pro- cesses with respect to generic parameters a that appear either in the drift or diffusion coefficients (52). This area is known as sensitivity analysis, and has been developed largely with financial applications compute the "Greeks"), In low dimensions, dynamic programming [3] or finite differences [26 531 work well, but they scale poorly to high dimensions. In high dimensions, several approaches have been proposed (see the section of [29] for comprehensive although dated overview):
in mind (more specifically,
The path-wise method (which we refer to as adjoint method) involves taking the gradient [xE()(X)) inside of the expectation as EEV a. (Xx)] and was first described by (82).
The likelihood method or score method 127 699 consists in rewriting V Vo [f as (ff(Xr) where H. is a random variable which equal V. lo p(a, XT). p(a,-) be- ing the density of the law of Xx. with respect to the Lebesgue measure. (82) provide explicis weights H, under the restrictions that a appears only the drift of the SDE (and not in the diffusion coefficient) and that the diffusion coefficient is elliptic, using the Girsanov theorem. (29] provide an expression for in the case where H also appears in the diffusiio coefficient, using Malliavin calculus.
The estimator of the path-wise reparameterization trick is formully similar to the likelihood method estimator, but is different that a is the initial condition of the process, and does not appear either in the drift nor the diffusion coefficient.
29

---

C,5 Proof of Lemma
Proof. Since the equality (28 holds almost surely for the pair (X, B), it must also hold almost surely for ,X",B), B"), which satisfy the same SDE. That is
W(X", -(X*,0) = V 0) + 2 S0 | (X";
8) ds
vXJo (u* o),
Thus, obtain that
alv, X", , B) = exp ( 1-----(- .0) 1--1/ Sor t), + - So |o( t)l dt)
exp - 0) - So I ||| ds x~ 1/2 Lo- (u' (X,, dB;)
- 1/2 10 (u(X;' 1), AB + ス Lo o 1) dt).
and this is equal to exp V(Xx: 0)) when u²" Since we condition on X Fiit. we have obtained that the random variable takes constant value exp ( (( 0)) almost surely, which means that its variance is zero.
D Optimal reparameterization matrices
Theorem 4 (Optimal reparameteriization matrices). Let u be an arbitrary control in U. Define the + P[; as
integral operator Te
[T((M))(s)) 1." (x X", B),,,, BYT xa(e, X", B))B)]ds
where
x(t, B) = St Vef(X), 5) ds + Vg(x)) + (T S:" ()(, dB,.
-.t s(oo"))))
If we define Ns(s) = -B([Vg(H)) + J s')s B) alv, X", B)}, the
optimal M = (etor) is of the form Ms(s) = + S M, (s') ds', where Mi is the unique solution of the following Fredholm equation of the first kind:
((() = N,.
The proof of 25 shows that minimizing Var(w; M) is equivalent to minimizing
[ J0 [[,[t,, | dt ac B)]
(57)
To optimize with respect to M, convenient reexpress it in terms of M (M.)) as M₂(s) = 1 Je¹ Mc(s') ds'. By Fubini's theorem, we have that
S M,(s)V ds = Stt (1+ Je M,(s') ds') V ds Mc(s) J Vaf(Xy, ds, ds, = = x)
-((M) -M)))))))MM(()))())))))))) (o (s)u
- M(()) s)ds --ddss-- 1:T M((s) 1-T ((0 8) ds' ds,
-X1/2 S (M.()) 8) - (s)
=X'/2 MM(s)(o (s) 1h" Me(s) l Vab(xy, ((
Hence, we can rewrite 57 as
G(M)=E( fo ( (S. v.f(Xy ds Vg((H)
+ 1 (.) Vef(Xy, s') ds' Vo (s)))sss)
- -f" V'V.b (') (x))(xXx)
5) ds' - 1º" V.bKy, (T) (5') ())) ds) 2 dt
x a(v X",B))
30

---

The first variation 69 (M) of G at M is defined as the family Q (Qu)e( of valuee functions such that for any collection of matrix-valued functions P = (P.)
V.V(M + l V(M+ (P.Q) = S0 S ((R(()))(()))) ds ds
where M + +eP = (M, + e ePi)eEto,T Now, note that
V.V(M + (P)] 0. 1o' [I] s)
+ 5 (M)(s) + <P. V.
Vo((H)
+ (o
-1 (17 alce -
Vab(X),'))0 (s)) ( ds) ||
= lo (o(t))(t) (Se Vx
+ S ( V.f(X 1O
[ "1 ds),
1. (.
1s V,bx - s' dB;) ds) dt
x xa(e,X"BB)
(58)
If we define
X(s, x", B) = -- V (X + Vg(X)) + (o-)))))))))))))))())
Vxx,, (
can rewrite 58
0, ((
odd MM)xss) ds).
(59)
S.' P((sxx(s X", B)ds) ds x alv, X", B)]
Now let us reexpress equation (9) as:
E L (ao (V + L (V.J(X;,)) + M((s)x(s, XV, B))ds) ds),
STT P,(s)x(s, X", Xx,))ds dt a(u, X",
E T J J (P)()) B),
oo (I((VVV(()) + 1fd () X", B)) di x alv, X", B)]
@ D [ H (oo + (V.](X))))) (V, Mc(5'x(s x², B)) ds') ds') B)T,
P((s)) dt x a(v, B]]
= So (ooo (V) (V (M( B) alv, X", X", B))
P((s)) ds.
(60)
Here, equality (i) holds by Lemma with the choices a(t,s) P((s)x(X", ss B) >(t)
oo T (t)((V + 1 (V=f(Xy, + (X)s)) s, B)) ds. Equality (ii) follows from the fact that for any matrix A and vectors b, (Ab, c) = c where ( denotes the Frobenius inner product.
Ab = Tr(cT Ab) = Tr(Abe = (B, cb )F
The first-order necessary condition for optimality states that at the optimal Mr, the first variation &q (M²) is zero. In other words, (M eP) (P) is zero for any P. Hence, the right-hand side 8M
31

---

of 60 must be zero for any P, which implies that almost everywhere with respect to f E [O,T],
S. T).
E { (Va M,('))) ss', xs') B) X", B)] B))e,, To derive this, we also used that (1) is invertible by assumption.
Define the integral operator T: 1. ([t, T): Rdxd) (1t, T];
[T(())))(s) 1/" M 8' B)X 8, B) a(v,X", X", B) ds' we define Ns(s) = (B(( + 1 problem that we need to solve to find the optimal M. is
B) x a(u, X B)]. the
Te(M) N₁.
This is Fredholm equation of the first kind.
Lemma 8. a, 3 : [0, 7] [], Rd, y: [0, 7] Rd. 8 [0, Iddd are arbitrary
integrable functions, we have that
fo (Se
(()) dt for So s).t)s).7(t) dt
Proof. We have that:
So 5-" (a(t, s), 7()) ds @ So (a(t,T s), (()) ds dt
@ So SJ (a(T T,T (T, 1)) dt (m fo S- (alT 1,T s), ()) ds
@ 0h (a - (()), 1)) di _ 5 Jo (((1,,),((1)) dt ds
Here, in equalities (i), (ii), (iv) and (v) we make changes of variables of the form t + T - t.
T s, + T s' equality (iii) we use Fubini's 's theorem.
E
Control warm-starting
We introduce the Gaussian warm-start, control warm-start strategy that adapt from [56]. and that we use in our experiments in Figure 7 Their work tackles generalized Schrödinger bridge problems, which are different from the control setting in that the final distribution is known and there is no terminal cost. The following proposition, that provides an analytic expression of the control needed for the density of the process to be Gaussian all times, is the foundation of our method.
Proposition 6. Given Z N(O, 1) define the random process Y
Y, =(() + (())Z,
where ull) E IR,, î(1) = viT(I) E Rdxd
(61)
Define the control u 12d [0, + Rd as
u(x, = (an)) + ((OTT(1))T())- + I-(00 XO (x ((t)) M(, 1))
2
Then, iT (0), the controlled process X defined in equation That is, for all t E [), Law(Y,) = Law(X.))
has the same marginals Y.
Proof. Following
56 we have that
0,x, 0,44 +
= =.,(() + (O)(()))()))X -xtt)(t)),
V (o) (x
= (()(()))
Now, Pt satisfies the continuity equation equation
d.P. -V ((0,p(() (e -a p(()))p))))
(63)
32

---

Let D(t) = (t) T. We want to reexpress 63 as a Fokker-Plande equation of the form
O.p: -- (u(z, t)pp + Σj (D,((1, - (u(z,
a. X 20,,
= - (o(x,00p)) (D()))p)) = -V - (v(2,))))))) V (V.((()Vogg)))
= - V-VVV.((le
Hence, we need that
v(z 1) V logpt 0,1(1) + (0.(()))(()) u(t)),
vx(x)= = 0,u(t) + ((0I(()))(()) (--m(t))) + (a 1 1(1) log
2
= =a,u(t) + (0((()))(0)) ((= - ((0) E(t) (x p(tt).
If we let [tt) = ((tt) Vt, then Σ(1) =)(()(()) =t)( and 0,s(1) = 0,l + r(a) r(0
That is,
2v
o(x, 1) 0.p(()) + (a.r() 2/7 Vt (x m(()) (0o (()) =(0) (x pl())
= =0, + (orr(t))((()>>> - (x(((t) + # (x -p()) (oo (()
(x p(t))
2
For e be finite at
we need that (ar T (0)) 1, which holds, for example,
((0) = (0)) Also, to match the form of 2 we need that
u(z, t) = b(x, t) + + (t)u(x,,
= u(z, t) = (t) 11 (B + ((Or())))(()) + I-(aa (1E01 (x -p) - b(x,t)).
The warm-start control is computed the solution of Restricted Gaussian Stochastic Optimal Control problem, where we constrain the space of controls to those that induce Gaussian paths as described in Prop. 6 In practice, we learn a linear spline = (n(²)) where μم E Rd, and a linear spline T= ()) >0=p- where (() E RRdd These linear splines take the role of (u(t) and E() in 66). Given splines fu and I, obtain the warm-start control using 62 for a given i 00, T),
we let 6_ = [,t/T], b4 = b_ 1L.+.. A = T/B, we have that
(tt) (t- (t-b-A))))) ا) + (b4 t)n
Tru((t (0 (0-
(64)
P(I)= (1--A))( + AOD---
0T((1) r((+) -p(h_
(65)
ú(2, t) = (o.j(()) + (OT(() I- (C)) (0)
((()) b(x, (66)
2t
Algorithm provides a method to learn the splines μr, r. It is a stochastic optimization algorithms in which the spline parameters updated by sampling Y, in (61) different times, computing the control cost relying on (66). and taking its gradient.
Algorithm 3 Restricted Gaussian Stochastic Optimal Control
Input: State cost [æx f), terminal cost g(x). diffusion coeff. ot)), base drift b(x, t). noise level A, number of iterations N. batch size number of time steps K, number of spline knots B, initial mean spline knots
(n4p (b) Jmm initial noise spline knots To (r()) y
for = 0: (N 1)} do
Sample m i.i.d. variahles (7.)== NO, T) and m times (tt) Unif([0, T]).
for j = 0 do
Set 1, jT/K, K, and compute (t,,,, 0,14, (ty), (t),t), (t,) according (64
65
using Po..
1'm
for i = 1: m do compute Y4 ((t,)) + V (t,)) and ân (Yij, f,) using 66)
end
Compute LRGSOC Tw) Σr Σバー |YY t f(y,,,,) &)) + (YYx))
Compute the gradient of ÉRCsOc(pm. rw) with respect to the spline parameters (no, Tw).
Obtain pnt1, 1't11 with via an Adam update on pm.. 1 n resp. (or another stochastic algorithm)
10 end
Output: Learned splines μu. 1'N. control dn
33

---

Once we have access to the restricted control i, we can warm-start the control in Algorithms | and by introducing u as an offset. That is, we parameterize the control as ug uN + Ug.
F
Experimental details and additional plots
F.1
Experimental details
The control L error curves show the following quantity:
F [|lu (XF' u(X;", =t)lfe->~-viss / (v-x- .0]
That is, we sample trajectories using the optimal control, and compute the error using a Monte Carlo estimate. In all experiments, the distribution X6 delta, which means that do need to compute (XX6' 0). We keep an exponential moving average (EMA) estimate the control L error, which we show in the plots. To compute it, we sample ten batches of optimally controlled trajectories every 10 training iterations, and we update the quantity with the average of the ten batches, using EMM coefficient 0.02. All other quantities shown in the plots are also smoothed out using EMA with coefficient 0.01, except for control objective values, which are computed the average of 65536 samples, every 5000 training steps.
For all losses and all settings, train the control using Adam with learning rate 1 x 10 For SOCM, we train the reparametrization matrices using Adam with learning rate x 10-2. We use batch size 128 unless otherwise specified. When used, with B = 20 knots, K 200 time steps, and batch size m 512, and we use Adam with learning rate 3 10-1 for N 60000 iterations.
the warm-start algorithm (Algorithm 3
QUADRATIC ORNSTEIN-HLLENBEC The choices for the functions of the control problem are:
b(x,1)= Ar, ((x,()=xx1 Px, g(x) = x Tux o0l) ou-
where Q is a positive definite matrix. Control problems of this form are better known linear quadratic regulator (LQR) and they admit a closed form solution [78 Thm. 6.5.1]. The optimal control given by:
1 -2, Fex,
where F, is the solution of the Ricatti equation
dF.
A F,A 2F.0000
dt
with the final condition Pr Q.
Within the QUADRATIC OU class, we consider two settings:
Easy: We set 20, A 0.21, P 0.2I, 0.1I, du I,
T
Pinit = (0, I). We do not use warm-start for any algorithm. We take K = 50 time random seed O.
discretization steps, and
Hard: We set d = 20, A = I, We use the Gaussian warm-starl (App. ED. We take batch size m = 64 and K = 150 time discretization steps, and we use random seed O.
I, 0.5I,
=, = 1, Tinit = 00. 7).
LINEAR ORNSTEIN- CHISTEIN-ULLENRCCE
The functions of the control problem are chosen follows:
b(x, t) = Ax, f(x, t) = O. g(x) = (7,
(t) = dp-
The optimal control for this class of problems is given by 59 Sec. A.4]:
in (x) - of A (AT(TT-)) (T-1)
We exactly the same functions 59 sample (Eij)gesss once the beginning of the simulation, and set:
10,
A - (5 1suss = 11 on
o= = I+ I
T = 1,
X = 1, Dinit 0.5, 7).
We take K 100 time discretization steps. and random seed 0.
34

---

DOUBLE WELL We also use exactly the same functions as [59]. which are the following:
b(x; 1) = (x),,, ((x)= = Σ (1 1))
f(x) 0, g(x) Σ Us (1( - 1) On I,
for {4,.... 10). We set T= 1, À = and Tinit 0. We take = 200 time discretization steps, and we use random seed O. The Double Well problem is actually highly non-trivial, and is multimodal. The only reason produce a "ground truth" control to compare to in this setting is that we use significant knowledge of the problem; we analytically reduce ID problems by decoupling each dimension and apply numerical methods to solve the Hamilton-Jacobb Jacbbi-Bllaan equation for these ID problems. is not problem where actually have the ground truth control in closed form.
where d = 10, and H; 5, 3 for i (1,2,3) and K;
PATH INTEGRAL SAMPLER ON MIXTURE OF GAUSSIANS We set
b(x,1) () 0, (xx = O,
g(x) = log((" = = ||=||* 2 log(2n) - log (1)))
where T = 1, and je is the density of a mixture of two Gaussians with means Hej, where e1 = (1,0,... ..,,....),, 0). and variance 1d. Note that take μ be normalized, i.e. f f(xx) dæ 1, or equivalently, log log p(xx) dz) = O. In Figure 1 of the control objective at the control u:
we use the following Monte Carlo estimator
(x = ( (allu(XF, f)I + f(X), t)) dt + g(x;) + (u((X))) dB:).
Note that this estimator is unbiased because E[ (X," f), dB,)] = 0. This is known as the Sticking the Landing estimator, as has zero variance when
the optimal control 772. The fact that E| [I-S((X) & log = with equality when u = u² is stated as [84 Thm. 4].
F.2 Model architectures
As general guideline, the control function can be thought of as the analog of the score function in diffusion models; hence, natural choice for the architecture can be U-Nets diffusion transformers the control task on images, audio or video. Other domains may require different architectures. In the experiments report, used the architecture implemented in the class Ful FullyConnectedUhe within the file SOC_ma meatching/modee py. is a simplified version of the U-Net architecture where both the nown-sampling and up-sampling layers are fully connected with ReLU activations, and the horizontal layers are linear transformations, We use three down-sampling and up-sampling steps, with widths 256, 128 and 64 (hence, the first down-sampling step actually an up-sampling, because the data dimensions in our experiments range from 10 to 20).
The reparameterization matrices have an unusual trait, which is that their input dimension is small (two) while their output dimension is large (dP). Hence, the kind of functions that they need to learn are low dimensional and hence easy. In our case, we used the architecture implemented in the class SigmoidMLP within the file S0C_mat matchiin//oooeeeeee multilayer perceptron with ReLU activations and output dimension df, whose output is averaged with the identity matrix using sigmoid weights, in order to enforce that M,(t) be the identity matrix.
.yy which essentially three layer
F.3 Additional tables and plots
Table shows the average times per iteration for each algorithm. Each algorithm was run using a 16GB V V100 GPU.
SOCM SOCM M. = SOCM adj.
Adj.
Cross entropy
Log-variance
Moment
Variance
0.222
0.090
0.099
0.169
0.086
0.117
0.087
0,086
Table I: Time per iteration (exponential moving average) various algorithms seconds per iteration, for the QUADRATIC OU (EASY) experiments (Figure 2h.
Figure 4 shows the control objective L for the four settings. The error bars for the control objective plots show the confidence intervals for + one standard deviation, computed via Monte Carlo
35

---

estimate using 65536 trajectories per data point. They show the standard error of the mean. As expected, SOCM also obtains the lowest values for the control objective, up to the estimation error.
Figure shows the normalized standard deviation of the importance weight for the learned control u: v [/ar X [ X", B)]. By Lemma 1 when Xü = Zinit for an arbitrary Zinit (which is the for all our experiments), this quantity is zero for the optimal control u². Hence, the normalized standard deviation of a is an alternative metric to measure the optimality of the learned control.
Figure 6 shows exponential moving average of the squared of the gradient for LINEAR OU and DOUBLE WELL. For LINEAR OU, the minimum gradient norm is achieved by the adjoint method, while for DOUBLE WELL it achieved by the cross entropy loss. The training instabilities of the adjoint method become apparent as well. Interestingly, in both settings the algorithms with smallest gradients are not SOCM, which the algorithm with smallest error as shown in Figure Understanding this phenomenon is outside of the scope of this paper.
Figure 7
shows plots of the control L eπror, the norm squared of the gradient, and the control objective for the QUADRATIC OU (HARD) setting, using a warm-start strategy detailed in App. E Figure shows that SOCM is again the algorithm that achieves the lowest error and the smallest gradients. Remark that the warm-start control is a reasonable approximation of the optimal control,
as the initial control L error is much lower than in the other figures.
Figure 8 shows the value of the training loss for SOCM and its ablations: SOCM with constant Mr = I, and SOCM-Adjoint For all such algorithms,, the training loss is the sum of the L error of the learned control u. and the expected conditional variance of the matching vector field w. Thus, the difference between the training loss plots and the L error plots is the expected conditional variance of w. We observe that the expected conditional variance the QUADRATIC OU setting is orders of magnitude smaller for SOCM than for its two ablations. For LINEAR OU, SOCM and SOCM-adjoin have similar expected conditional variance, and a possible explanation is that the LINEAR OU setting is very simple. In the DOUBLE WELL setting, the SOCM-adjoint training loss curve has spikes that are probably caused by instabilities of the adjoint method. These spikes can be attributed mostly to the expected conditional variance term, since the corresponding I error curve in Figure 3 does not present them.
Figure 9 shows that the instabilities of the adjoint method are inherent to the loss, because they also appear at small learning rates: 3 x 10-5 is smaller than the learning rates typically used for Adam,
which hover from 1 x 10-4 to 1 x 10-3
36

---

Quedratic Uhlenbeck, essy
Quadratic Uhlenbeck, hard. warm
--50CM M. lablation
(nurs)
160
-S.OSM Ry (ablatioon
50CM (ours)
5.40
SOCN -OCN-Adgern at
....
-0CCM-Mjjlit cab arion)
"Ad
AAluin
5.41
Crous Entropy
Legesbetence
120-
-ogg/orincce
-Cres Entropy
Momont
dbjective
5.00
80
5.08
5.85
40
40000
10000
20000
40000
50000
Num,
Linea Ornstein Uhlenbeck (d 10)
Double Well (d= 10)
-2.02
50CM
-OCM (attatoon)
3.0
50CM Laclation
-2.55
Entrayy
Entmay
Control otteciie
2.8
2.7
-2.05
-2,87
-7.88
2.4
2.3
10000
20000
30000
40000
50000
60000
10000
20000
20000
40000
50000
600C0
70000
50000
Num, Baaation
Num. Banations
Figure 4: Plots of the control objective for the four settings.
Quadratic Ornstein Uhlenbeck, easy - 201
Quadratic Ornstein Uhlenbeck. hard, warm start (d 20)
KG (ablution)
Intematied)))
Adjolnt
Importance 11 Incmmalocal
SDCM-Adjjint oab etion)
Adjoint
-Lag Entropy
-Cress Entropy
Momment
Mo ment
Variance
Var ance
101
10000
20000
30000
40000
50000
60003
10000
20000
30000
40000
50000
G000D
70000 50000
Num, Renations
Num. hresttaog ont Double 10)
Linear Ornstein Ulienbeck
101
SOCM tounsl
5DCM tound)
SOCM lateation)
-.CC Bt = (actstionn
Importarce weight 26d dev. Insmmalond)
S0CM Adointt loblation
Impertercee de dee (ncomstioe))
.... SOCW Adjoinc ation)
AAlomn
--Aajoint
Crrss Entropy
Cross Entrapy
-Log-lariance
Lep--errannee
Mornett
Momentt
Variance
Variance
x 10"
20000
50000
40200 50000
50000 10000
80000
Rorotiers
Figure 5: Plots of the normalized standard deviation of the importance weights: Varlaau, X", X",
37

---

Linear Omstain Uhienbeck
Double Well
50CM (ours)
SOCM lactabion)
10"
SDOC M. (eclation)
SOCW-Md (sblation)
Gridert romm sed (EMA 0,013
10
Adjoint
Cross Entropy
Coxteet moom (FFMA c..11
100
Cross Encrosy
"Lep-aarance
"Lag-Hrriance
10"
Vir Veriance
Moreent
Variance
Momment
ZD
2D7
10
103
10"3
101
Figure 6: Plots of the norm squared the gradient for the LINEAR ORNSTEIN UHLENBECK and DOUBLE WELL settings.
20)
SSCM lourst
30
-DCM ui (stration)
4'
SCCM adjoitt (ostation)
30'
Connall ener (ete (PNA 0.1)
Log
Entropy
Gondent wptt (PMA 0.11
307
50OM (ours)
-00M Gatiittion
Adjinn (oblation)
10
-Lpg
Kommet
Varfance
LOC00
20000
1000C
40000
50000
Berations
Quacratic Ornstein Uhlenbeck, hard, warm start id 20)
Quadratic Ornstein Uhlenbeck. hard, warm start (d 20)
inportence wagght (pemaitre))
28.76
Entrepy
Fnimay
-Log
-Lop
Control oe
28.72
x1-
20.88
10000
20000
30003
10000
50000
60000
10000
20000
30000
10000
50000
60COC
Num.
Num. iteratiooo
Figure Plots the control 2 error, the norm squared of the gradient, and the control objective for the QUADRATIC ORNSTEIN-UHLENBEER (HARD) setting, without using warm-start.
38

---

Quadratic Ornstein Uhlenbeck, easy
Quadratic Ornstein Unlenbeck, hard,
107
-OCC
(ears)
-SOCM
[aataatio)
-SOCM H. (ackation
10
-0CW Adomm i2blation
-DCMMsddpre (40 aeion)
Tainiig lns (ENA 0.333
307
Tatring (7M
301
10
10-1
10-3
Sc000
Linear Ornstein Utienbeck id 101
Double Well id 10)
10
5OCM
[bblatioo)
SOCM
Taaiing (os, (EMA 0.111
Thinkng leee (PMA 0.11
30-1
SOCM toues)
10
-00C At (obtation)
30-7
50CWAdem leblation)
10'
10a4
10"
20
10000
20000
30000
40000
50000
60000
10000
20000
30000
40000
50000
60000
70000
50000
Num. Feeatlon
Num. Deeatt
Figure 8: Plots of the training loss for SOCM and its ablations: SOCM with constant M, I, and SOCM-Adjoint.
Adont UR- 1e Addant 18- 30-
AdjontLn-Iee
Adjortt 18mm
Control unce trnA na11 nal)
1.20
Centurt somm upt 1780
10000
20000 25000
35000
20000
15000
30000
40000
Figure 9: Plots of the control 1 error and the norm squared of the gradient for the adjoint method on DOUBLE WELL, for two different values of the Adam learning rate. The instabilities the adjoint method persist for small learning rates, signaling inherent issue with the loss.
39

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: In the abstract and introduction state that we introduce SOCM, novel algorithm to solve stochastic optimal control problems, and we claim that it outperforms all existing algorithms in three out of four settings. All of these claims supported by paper.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Yes, in the second-to- -o-sadt paragraph of See. we discuss the main limitation of SOCM (our algorithm), which is that the variance of the importance weight is too large in certain settings. Regarding the computational efficiency of our algorithm, Table shows a comparison between our algorithm and existing ones.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model vell-specifiiattio, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
40

---

3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: All the theoretical results stated in the paper are numbered and have a corre- sponding detailed proof in the appendices. In the main text, we indicate the precise location the proofs. We provide a proof sketch for Theorem in the main paper, and a full proof in the appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide description of experimental details in Subsec. F.T. and of model architectures in Subsee. F.2 We also provide a link to the GitHub repository that contains the code for this paper.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
41

---

(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We also provide link to the GitHub repository that contains the code for this paper.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https SonddeSubmisssooPPoon
/ wttpps.cc/public////
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide description of experimental details in Subsec. F.I and of model architectures in Subsee. F2
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: Most of the plots in the paper do not contain error bars, because the lines show exponential moving averages, which means that the error is already smoothed out and very small. The error bars for the control objective plots do show the confidence intervals for +
42

---

one standard deviation, computed via a Monte Carlo estimate using 65536 trajectories per data point. They show the standard eтror of the mean. We state this Subsec. F.3
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We explain the type and amount of GPUs we use Subsec. F3
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (c.g., preliminary or failed experiments that didn't 't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: btetps://neuripe: httppps.cccppppppeeecccceieiiiia
Answer: [Yes]
Justification: We reviewed the NeurIPS Code of Ethics and research conforms it. We made sure to preserve our anonymity.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
43

---

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [No]
Justification: The research presented is foundational, and the code released does not have direct societal impact. Yet, it may serve as the basis to develop algorithms that improve the quality of generative models. We state this in Sec. 5
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g.. disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied partieular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [No]
Justification:
The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (c.g., code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
We did use existing assets.
Guidelines:
44

---

The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g.. website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paaassttssppersseaae paperswi thcode. com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and the documentation provided alongside the assets?
Answer: [Yes]
Justification: We present code implementing our algorithm. The code folder contains a README which explains how reproduce the experiments. We describe algorithm throughout the paper, and provide details in App. F
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and sereenshots, if applicable, as well as details about compensation (if any)?
Answer: [No]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
45

---

Answer:
[No]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
46


Reviews:
[{"review_id": "PQTkGwx9k6", "reviewer": "Reviewer_Su1g", "summary": "**Summary**\n\nThis paper introduces Stochastic Optimal Control Matching (SOCM), a novel algorithm for solving stochastic optimal control problems. Key contributions include:\n\n1. SOCM algorithm, adapting ideas from conditional score matching in diffusion models\n2. A new \"path-wise reparameterization trick\" for gradient estimation\n3. Theoretical analysis including a bias-variance decomposition\n4. Empirical evaluation showing superior performance on 3 out of 4 benchmarks\n\nSOCM learns a control function by fitting a matching vector field via least squares, while optimizing reparameterization matrices to reduce variance. The method is currently limited to linear Gaussian models and requires knowledge of certain parameters. Experiments demonstrate SOCM's effectiveness on theoretical benchmarks, outperforming existing methods in most cases. The paper provides a solid theoretical foundation but lacks exploration of real-world applications or non-linear systems.", "questions": "1. How might SOCM be extended to handle nonlinear or non-Gaussian systems? Are there specific challenges you foresee in this extension?\n\n2. The paper focuses on theoretical benchmarks. Have you considered applying SOCM to any real-world stochastic optimal control problems? If so, what challenges did you encounter or do you anticipate?\n\n3. How does the computational complexity of SOCM scale with the dimensionality of the problem? Could you provide a more detailed comparison of computational requirements with existing methods?\n\n4. The path-wise reparameterization trick is an interesting contribution. Could you elaborate on potential applications of this technique outside of stochastic optimal control?\n\n5. The paper mentions that SOCM requires knowledge of certain model parameters. In practical scenarios where these parameters might not be known precisely, how sensitive is SOCM to parameter misspecification?\n\n6. Have you explored the performance of SOCM in settings with sparse or noisy rewards, which are common challenges in reinforcement learning?", "limitations": "n/a", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper introduces Stochastic Optimal Control Matching (SOCM), a novel algorithm for solving stochastic optimal control problems. Its originality lies in adapting ideas from conditional score matching in diffusion models to the domain of optimal control. This creative combination represents an interesting cross-pollination between two active areas of research.\n\nThe quality of the theoretical work is notable. The authors provide a comprehensive mathematical foundation for their method, including detailed proofs and a novel \"path-wise reparameterization trick\". This theoretical rigor is a significant strength of the paper.\n\nIn terms of clarity, the paper is well-structured and clearly written. The authors effectively guide the reader from the problem formulation through the theoretical development to the empirical results. The use of illustrative examples and detailed appendices aids in understanding the complex mathematical concepts presented.\n\nThe significance of this work lies in its potential to improve the efficiency of solving stochastic optimal control problems. The empirical results, showing improved performance over existing methods on multiple benchmarks, underscore the practical impact of this approach. However, the significance is somewhat limited by the current restrictions to linear Gaussian models.", "weaknesses": "The primary weakness of this paper is its limited scope and applicability. The method is currently restricted to linear Gaussian models and requires knowledge of certain model parameters. This significantly narrows its potential impact on the broader field of stochastic optimal control. The authors should discuss potential approaches to extend SOCM to more general settings, such as nonlinear or non-Gaussian systems.\n\nWhile the empirical results are promising, they are limited to theoretical benchmarks. The paper would be strengthened by including experiments on real-world problems or more complex simulated environments. This would help demonstrate the method's practical utility and potential for broader impact.\n\nThe scalability of the method is not thoroughly addressed. As the dimensionality of the problem increases, how does the computational complexity of SOCM compare to existing methods? A more detailed analysis of computational requirements and scaling properties would be valuable.\n\nThe comparison with existing methods, while showing SOCM's superior performance, could be more comprehensive. Including comparisons with the most recent state-of-the-art methods would provide a clearer picture of SOCM's relative performance in the current landscape of stochastic optimal control algorithms.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for an accurate list of our contributions; however, we think there may be a misunderstanding regarding the scope of applications. We detail our responses to the reviewer’s concerns and questions below:\n\n**“The method is currently restricted to linear Gaussian models and requires knowledge of certain model parameters. [...] The paper mentions that SOCM requires knowledge of certain model parameters. In practical scenarios where these parameters might not be known precisely, how sensitive is SOCM to parameter misspecification?”**\n\nWe are unsure about this question and we think there may be a misunderstanding. We would appreciate it if the reviewer could point us to the comment that they refer to.\n\nSOCM does not actually require any explicit assumptions on the model parameters or a linear-Gaussian model. To clarify our interpretation of this concern, the linear-Gaussian model assumption usually refers to state transitions $p(x_{t+h} | x_t, u_t)$ being Gaussian distributed with linear dependence on the state $x_t$ and independent control variables $u_t$. In the continuous-time formulation, the updates under a linear-Gaussian model would be of the form $dX_t = Ax_t + Bu_t + \\sigma(t) dW_t$. However, in our formulation we use neural networks to model a control velocity field, resulting in non-linear updates. That is, the updates we have are of the form $dX_t = u_t(x_t) + \\sigma(t) dW_t$, which includes the linear-Gaussian model but also generalizes to nonlinear dependencies. \nThe only assumption SOCM actually makes is that the objective functional depends on the control velocity field only through a quadratic form (equation 1). This is required for the path-integral representation in Theorem 1 but this is a typical assumption made in nearly all stochastic optimal control methods. SOCM requires the same amount of knowledge as the existing methods: we just need to know the base SDE, the running cost and the terminal cost.\n\n**“How might SOCM be extended to handle nonlinear or non-Gaussian systems? Are there specific challenges you foresee in this extension?”**\n\nAs we detail in the proof sketch of Thm. 1, SOCM relies on the path-integral representation of the optimal control (eq. 8). As far as we know, such path-integral representations hold for control systems with arbitrary base drift and linear dependency of the drift on the control, and it is a typical assumption made by works in the literature (e.g. Nüsken & Richter, 2021). If path-integral representations exist or are developed for more general control problems, we believe that our technique may be adapted to handle those.\n\n**“The paper focuses on theoretical benchmarks. Have you considered applying SOCM to any real-world stochastic optimal control problems? If so, what challenges did you encounter or do you anticipate?”**\n\nSOCM performs well when the importance weight $\\alpha$ defined in eq. 20 has low variance, which holds when its exponent has low variance. In our experimental section, we show that when $\\alpha$ has low variance, SOCM outperforms existing methods. However, when $\\alpha$ has high variance, the signal-to-noise ratio of the stochastic gradient is low, and the performance of the algorithm is compromised (much like it happens for the existing cross-entropy method, which also contains the factor $\\alpha$). We do not regard SOCM as the solution to solve all stochastic optimal control problems, but rather as a new perspective which can lead to the development of more algorithms (ongoing work).\n \n**“How does the computational complexity of SOCM scale with the dimensionality of the problem? Could you provide a more detailed comparison of computational requirements with existing methods?”**\n\nTable 1 in Section F.3 shows the time per iteration for each of the loss functions that we consider: SOCM takes 0.22 seconds, while the adjoint loss takes 0.169 seconds, and all other methods are around 0.1 seconds. The reason that SOCM takes longer is that there is an additional neural network that is trained: the M function. The second paragraph in Section F.2 describes the neural network we use for the M function. In our setup, the total cost of M evaluations per iteration is $O(d^2 K^2)$, where $d$ is the dimension of the control system and $K$ is the number of discretization steps. The dependency $O(d^2)$ stems from the fact that we parameterize the whole matrix $M$; an alternative would be to parameterize $M$ as a diagonal matrix, which would result in a cost $O(d)$ at the expense of higher gradient variance. To further reduce the variance while using more computational resources, it is also possible to choose a function $M$ that depends on the iterate $X_t$ (see Remark 2). We will provide more clarity regarding the computation cost in the paper.\n\n**“The path-wise reparameterization trick is an interesting contribution. Could you elaborate on potential applications of this technique outside of stochastic optimal control?”**\n\nThe path-wise reparameterization trick can be used as a drop-in replacement of the adjoint method when computing gradients of functionals on ODE/SDE trajectories, either with respect to the starting point or with respect to neural network parameters. The path-wise reparameterization trick provides a natural way to enable variance reduction, thanks to the arbitrary function M. One can potentially apply this trick to any setting where neural ODEs/SDEs are used, and also to estimate scores of SDEs with arbitrary drifts.", "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": "rknhJZVyH4", "reviewer": "Reviewer_gdJ8", "summary": "This paper presents stochastic optimal control matching (SOCM), which is an iterative diffusion optimization for optimal control aiming to fit a matching vector field. The authors introduce a new loss function and address the analysis and design of a learning-based control method.", "questions": "As stated in Introduction, the work is motivated by stabilizing the unstable training of conventional IDO, which comes from the non-convexity of the loss.  Could the authors comment and/or perform some motivating experiments to show the stability of the training by SOCM?  They can emphasize the contribution of this paper.", "limitations": "No limitations in this work.  This work is devoted to the theoretical analysis of control system design, and it does not directly bring a negative social impact.", "rating": 7, "confidence": 1, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "The work is nicely motivated in Introduction, showing the drawbacks of traditional works.  The proposed control method is supported by the uniqueness analysis of the control logic (Theorem 1) and the sophisticated design methods (Propositions 1 and 2). In the reviewer's understanding, they are technically correct.", "weaknesses": "As stated in Algorithm 2 below, reducing noise in the gradient is crucial for the presented algorithm.  This weakness is addressed by Lemma 1 and extensions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewers for their encouraging rating and for providing us the opportunity to clarify the key importance of the proposed method below.\n\n**“Could the authors comment and/or perform some motivating experiments to show the stability of the training by SOCM? They can emphasize the contribution of this paper.”**\n\nFirstly, SOCM is stable by construction because its loss is convex in function space, unlike the adjoint method, which is not. In particular, we see that in the Double Well setting (Figure 2 right), the adjoint-based methods are quite unstable and have large ups-and-downs in the control error during training. \nSecondly, we also introduce the free parameter $M$ in our novel path-wise reparameterization gradient which can significantly reduce variance, allowing us to easily outperform related importance-weighted methods such as the cross entropy method. \nWe also ablated both of these two claims through our ablation experiments (correspondingly, these are the “SOCM-Adjoint” which uses adjoint method instead of our path-wise reparameterization gradient, and the $M=I$ ablation, which are shown in all experimental settings to lead to worse results).", "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": "CJBNWrtFax", "reviewer": "Reviewer_cXHg", "summary": "This paper proposes a novel algorithm for approximating the solution to the Hamilton-Jacobi-Bellman (HJB) equation with a neural network control policy. Rather than backpropagating through rollouts of the dynamics, the authors develop a least-squares objective which resembles the score-matching loss used in diffusion models. However, this requires computing gradients of a diffusion process with respect to its initial condition. To address this, the authors develop a novel path-wise reparameterization trick which relies on a family of reparameterization matrices. They show how to optimize these matrices to reduce the variance of the objective estimate. They demonstrate that their method obtains a lower error with respect to the ground-truth control on toy problems, sometimes by an order of magnitude.", "questions": "- How do all the methods compare in terms of actual task performance?\n- How do these methods perform on more realistic control problems?\n- Why does the proposed method not work as well on the Double Well system compared to the variance baseline?", "limitations": "The authors discuss limitations to scaling the approach up to more challenging problems due to the variance of the importance weight.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The proposed objective function acts as a form of variance reduction for the cross-entropy loss when solving stochastic optimal control problems.\n- The novel reparameterization trick for estimating gradients of diffusion processes with respect to its initial condition may be more broadly applicable.\n- On toy problems, their method appears to generally outperform other approaches in solving the HJB equation for the optimal controls.\n- The paper is well organized and overall written well. It provides a thorough related work section and does a good job explaining the novelty and results.", "weaknesses": "- The evaluations only consider simple toy problems. Moreover, they only plot the L2 error with respect to the optimal control. However, this does not necessarily tell us about the actual task performance due to compounding errors.\n- On the Double Well system, there is not a clear advantage compared to the variance loss and adjoint method. However, the authors do discuss how their method appears more stable than the adjoint-based ablation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their comments.\n\n**“The evaluations only consider simple toy problems. Moreover, they only plot the L2 error with respect to the optimal control. However, this does not necessarily tell us about the actual task performance due to compounding errors.” “How do all the methods compare in terms of actual task performance?”**\n\nWe would appreciate it if the reviewer clarified what they mean by “compounding errors”. Regarding task performance, we would also like to note that in Figure 3 in Appendix F.3, we plot the control objective (the quantity in the right-hand side of eq. 12). In Figure 3 we see that the control objective is a much less sensitive metric than the control $L^2$ error: it is harder to benchmark the losses using the control objective. From the perspective of measures over processes, the control $L^2$ error and the control objective are two sides of the same coin: the control $L^2$ error can be seen as the KL divergence between the probability measure $\\mathbb{P}^{u^*}$ of the optimally controlled process and the probability measure $\\mathbb{P}^{u}$ of the process controlled by $u$. And up to a constant term, the control objective is the reversed KL divergence between the same pair of measures.\n\nWe would also like to clarify that, as noted in the global response, the Double Well problem is far from a toy problem and is actually highly non-trivial due to its multimodal nature. We chose this problem because it is a representative challenging problem that appears in SOC interpretations of sampling problems, such as the path integral sampler. By exploiting the specific decoupling nature of the potential, we can reduce it to a 1D problem for reference solutions. However, our SOCM method does not utilize this strong prior knowledge and solves it as a generic multimodal high-dimensional problem.\n\n**“On the Double Well system, there is not a clear advantage compared to the variance loss and adjoint method. However, the authors do discuss how their method appears more stable than the adjoint-based ablation. [...] Why does the proposed method not work as well on the Double Well system compared to the variance baseline?”**\n\nThe Double Well (Figure 2 right-side) setting is different from other ones in that the terminal cost has 1024 modes. Hence, in order to obtain a small $L^2$ error, it is necessary to learn the control well in all or almost all the modes. Each trajectory that we sample will visit at most a few modes. Assuming for simplicity that each trajectory visits a single mode, and grouping trajectories into batches, by the end of the 80000 iterations, the “effective” number of batches per mode is 80000/1024=78.125, which is too small to get errors close to zero. Yet, this setting is interesting because it shows the behavior of the adjoint loss under multimodality: its $L^2$ error has a decreasing tendency but it wavers substantially, which is undesirable. We attribute this poor behavior of adjoint to the lack of convexity of the problem and SOCM has a clear stability advantage.\n\nRegarding the comparison to the variance loss, note that the SOCM method converges significantly faster and the we find variance objective is a poor choice most of the time with this benchmark being the only exception. We believe SOCM initially learns very fast on this problem due to the training of M (can be seen if compared to the M=I ablation) but may find sub-optimal local minima whereas the variance method has higher variance but may have a chance of finding better local minima given enough training time.\n\n**“How do these methods perform on more realistic control problems?”**\n\nWe benchmark SOCM against existing methods on control problems that have been studied before in the literature (Nüsken & Richter, 2021), which allows us to clearly focus our analysis on the loss function instead of other components of the problem.\n\nWe want to clarify that, as noted in the global response, the Double Well problem is far from a toy problem and is actually highly non-trivial due to its multimodal nature. We do not test SOCM on more realistic control problems because it often has gradient variance issues due to the variance of the importance weight $\\alpha$ (see Section 3 in the paper). We believe that the strength of our paper is that our framework is completely novel, and it will be the basis for the development of more algorithms that do scale well to realistic problems (ongoing work).", "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": "kPzoZv19ow", "reviewer": "Reviewer_Uos4", "summary": "In this paper, the authors propose a novel learning algorithm, Stochastic Optimal Control Matching (SOCM), to numerically solve general formulations of Stochastic Optimal Control (SOC) problems involving affined-controlled diffusion processes. They build upon the Iterative Diffusion Optimization (IDO) [1] framework, which consists in iteratively refining a parametric controlled diffusion process by minimizing at each iteration (with stochastic gradient descent) a specific objective function with respect to the parametric control. Previous works had for instance considered the relative entropy loss, the cross-entropy loss, the log-variance loss or the moment loss as the objective function. In SOCM, it is a least-squares regression loss (with multiplicative importance weight) which aims at fitting the parametric control to a vector field which depends on a family of *reparameterization matrices* (also optimized). The design of this objective function relies on standard tools from SOC theory, as well as an original contribution, the *path-wise reparameterization trick*, to compute gradients of conditional expectations of a functional applied on a random process. The authors show that the SOCM loss can be decomposed as the sum of a bias term, that is linked to the cross-entropy loss, and variance term, which is only affected by the *reparameterization matrices*. Hence, these extra parameters can be seen as a way to reduce the variance of the SOCM loss, which motivates their introduction. Moreover, this loss has the computational advantage to avoid computing gradients of the control along the path measure (which is the main drawback of the relative-entropy loss). Finally, the authors conduct numerical experiments to compare their approach to existing designs of IDO losses. They consider four different settings ($d\\in \\{10,20\\}$) with access to the ground-truth control, which allows them to compute the $L^2$ error on the control. Using this metric, their results indicate better performance on most of the settings while maintaining a certain training stability.\n\n[1] Solving high-dimensional Hamilton–Jacobi–Bellman pdes using neural networks: perspectives from the theory of controlled diffusions and measures on path space. Nüsken et al. 2021", "questions": "- Have you tried to restrict the optimization of the reparameterization matrices to scalar matrices ? I have the feeling that this choice may align a low computational budget with a good expressivity.\n- Have you tried another parameterization of these matrices ? In particular, have you considered simple form such as $M_{w}(t,s)=I_d + \\gamma(s-t)\\tilde{M}_{\\tilde{w}}(s,t)$ where $w=(\\gamma, \\tilde{w})$ and $\\gamma(0)=0$ ? Otherwise, why the choice of the sigmoid ?\n- I find that the warm-start strategy for the optimization of the control is a very good idea, as it benefits from the tractability of the stochastic interpolants with the lightness of the spline formulation. However, I have the feeling that this strategy works in the presented settings since they are \"kind of Gaussian\". Do you think it may still be of interest for general SOC problems ?\n- Could you explain why the second Ornstein-Uhlenbeck setting is called 'hard' ?\n- Could you provide the results on $L^2$ error of the control without EMA, as it is presented in [1] ?\n- I am quite surprised of the relatively computational budget induced by the use of the log-variance loss, could you comment on this ?\n\n[1] Solving high-dimensional Hamilton–Jacobi–Bellman pdes using neural networks: perspectives from the theory of controlled diffusions and measures on path space. Nüsken et al. 2021", "limitations": "The main limitation of the approach is discussed in Section 5: it is the variance of the importance weight in the SOCM loss, which may blow up.", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "- The paper is very well-written, it is a pleasure to read it. In particular, the authors pay attention to define their notation, introduce with clarity the SOC framework, state clear mathematical statements, recall (and prove) standard results from SOC theory, provide intuition on theoretical results (with details on the proof and meaningful comments). This is really good work.\n- The relation to prior work is clearly well established: in particular, the comparison between the SOCM loss and the previous IDO losses (presented in Section 2.2) is well highlighted. \n- This papers introduces an interesting contribution that may be applied beyond this framework (in particular, in the generative community where terms involving gradients of expectations often appear) : this is the path-wise reparameterization trick, which is proved to be decisive in the numerics.", "weaknesses": "- In my opinion, the major weakness of this paper is the lack of an additional numerical experiment, which represents a \"realistic\" setting (for instance, where the expression of the control is not known). For instance (as mentioned by the authors), a significant line of recent research has considered the sampling problem via a SOC perspective, see for example [1,2,3]. I am convinced that the SOCM contribution would have more impact with additional sampling numerics comparing SOCM, relative entropy [1,2] and log-variance [3] losses, for challenging distributions (namely, multi-modal distributions in relatively high dimension). In this case, the quality of the methods would be assessed with sampling metrics. This weakness explains my current score (although I really appreciate the paper).\n- I find that the complexity/accuracy tradeoff of the IDO losses (including SOCM) is not well highlighted to me. The table provided by the authors only considers one setting. To have a full picture, it should be given for all settings.\n\n[1] Path Integral Sampler. Zhang et al. 2022.\n\n[2] Denoising Diffusion Sampler. Vargas et al. 2023\n\n[3] Improving sampling via learned diffusions. Richter et al. 2023", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for providing an extremely succinct description of our paper. We completely agree with the reviewer’s concerns regarding limitations, but we just want to highlight our perspective that this paper proposes a rather novel way of thinking about constructing methods for SOC problems which we believe will have a big influence on future methods that do scale well to practical applications. That is, we put together this paper as it paints a complete picture of the new proposed approach, even if it is not yet a “silver bullet” that can solve all SOC problems, and chose to further develop scalable variants in a separate follow-up work. Below, we provide detailed responses to the reviewer’s questions.\n\n**“[...] the major weakness of this paper is the lack of an additional numerical experiment, which represents a \"realistic\" setting [...] I am convinced that the SOCM contribution would have more impact with additional sampling numerics [...]”**\n \nWe benchmark SOCM against existing methods on control problems that have been studied before in the literature (Nüsken & Richter, 2021), which allows us to clearly focus our analysis on the loss function instead of other components of the problem. These are control problems where we have access to the optimal control, which allows us to track the control $L^2$ error. Alternatively, we would only be able to compute the control objective, which as shown in Figure 3 (pg. 35), is a less sensitive metric due to numerical errors.\n\nWe also want to clarify that, as noted in the global response, the Double Well problem is far from a toy problem and is actually highly non-trivial due to its multimodal nature. We chose this problem because it is a representative challenging problem that appears in SOC interpretations of sampling problems, such as the path integral sampler. By exploiting the specific decoupling nature of the potential, we can reduce it to a 1D problem for reference solutions. We do not test SOCM on additional realistic control problems because it often has gradient variance issues due to the variance of the importance weight $\\alpha$ (see Section 3 in the paper). We believe that the strength of our paper is that our framework is completely novel, and it will be the basis for the development of more algorithms that do scale well to realistic problems (ongoing work).\n\n**“I find that the complexity/accuracy tradeoff of the IDO losses (including SOCM) is not well highlighted to me. The table provided by the authors only considers one setting. To have a full picture, it should be given for all settings.”**\n\nIn Table 1 we show the time per iteration for each loss for Quadratic OU (easy). For all losses and settings, almost all of the runtime is spent evaluating and backpropagating through the control neural network. Since the number of control computations per timestep for a given loss is the same across settings, and we are using the same neural network architecture for all settings, the time per iteration for all losses in a single setting is actually an accurate reflection of all the settings we’ve considered: the time per iteration for other settings is roughly proportional to the ratio of the number of timesteps that we are using. For SOCM, there is the additional cost of computing and backpropagating through the M function, but that also depends on the specific architecture that is used and can also be reduced by using sparse parameterizations.\n\n**“Have you tried to restrict the optimization of the reparameterization matrices to scalar matrices ? I have the feeling that this choice may align a low computational budget with a good expressivity.”**\n\nWe have not tried to restrict the optimization of the reparameterization matrices to scalar or diagonal matrices, but we agree that this would allow us to trade off fast computation and low variance. We will add an experiment where we use scalar and diagonal reparameterization matrices, and we expect that it will simply be somewhere between full M and the M=I ablation result in terms of performance.\n\n**“Have you tried another parameterization of these matrices ? In particular, have you considered simple form such as $M_{\\omega}(t,s) = I + \\gamma(s-t) \\tilde{M}_{\\tilde{\\omega}}(s,t)$ where $\\omega = (\\gamma, \\tilde{\\omega})$ and $\\gamma(0)=0$? Otherwise, why the choice of the sigmoid ?”**\n\nWe would like to clarify that in line 263, $\\gamma$ is a scalar, not a one-dimensional function. Hence, arguably our form for $M_{\\omega}$ is simpler than the one proposed by the reviewer. Still, the expression proposed by the reviewer makes sense and is a viable alternative.\n\n**“I find that the warm-start strategy for the optimization of the control is a very good idea [...] this strategy works in the presented settings since they are \"kind of Gaussian\". Do you think it may still be of interest for general SOC problems ?”**\n\nWe agree with the reviewer's comment. Warm-starting the control is a strategy that makes sense when the following conditions hold simultaneously:\n\n(i) An arbitrary initialization for control neural network causes the importance weight $\\alpha$ to have high variance. If the variance of $\\alpha$ is already low for an arbitrary initialization, there is no need for warm-start. This is the reason that we do not use warm-start for Quadratic OU (easy), Linear OU and Double Well. For Quadratic OU (hard), we try both no warm-start (in the main text) and warm-start (in the appendix): not warm-starting causes the learning to happen much slower for algorithms that use importance weights, although SOCM is still the best-performing loss by the end of training.\n\n(ii) The warm-started control is close enough to the optimal control, such that the importance weight $\\alpha$ has low variance. As the reviewer points out, this holds when the control problem is “kind of Gaussian” (unimodal), but it does not work when multimodality is present, as in the Double Well setting.", "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:
We thank the reviewers for their helpful comments. We would like to clarify and reemphasize our contributions, and to provide a global response to some issues that have been raised by multiple reviewers.
 
We acknowledge that our method has limitations due to the use of importance weighting and is not yet a "silver bullet" that can handle non-trivial SOC problems. However, it provides an unconventional view of SOC methods and provides a new direction to explore SOC methods as least squares objectives. That is, we believe this paper paints a complete picture in deriving a new framework for least squares objective, and scalable variants are better covered in a separate follow-up.

Our paper makes two key contributions:

(i) The formulation of least squares objectives that directly regress onto the optimal control, leading to the proposed SOCM objective. We have empirically found that SOCM easily outperforms existing importance-weighted objectives such as the popular cross entropy method (either in faster training, better final result, or both). Compared to adjoint-based methods, we find that SOCM exhibits more stable training (as can be seen on the Double Well) and often much lower control errors (as in OU Quadratic and OU Linear).

(ii) Our proposed path-wise reparameterization gradient is orthogonal to the SOCM objective and is a general method for computing gradients of cost functionals. In particular, the path-wise reparameterization gradient has a built-in variance reduction option in the form of the matrix $M$, which we see significantly improves convergence speed and performance. Both of these claims are closely ablated in our experiments (as “SOCM-Adjoint” and “M=I” in our experimental results).

A common point among reviewers is that the paper only considers toy control problems. The main reasons we chose such problems are that (i) they were used by (Nüsken & Richter, 2021) as a benchmarking suite and (ii) because we are able to compute (or closely approximate) ground truth solutions for the control function, and thus we can assess the $L^2$ error incurred when learning the control. The alternative metric is the control objective functional itself, which is much less informative due to numerical errors (see Fig. 3 in App. F.3).

Furthermore, there is (understandably) some misunderstanding about the difficulty of some of the benchmark problems as it is not emphasized in the paper. In particular, the Double Well problem is actually highly non-trivial, is multimodal, and is also closely related to the SOC interpretations of sampling problems such as path integral sampler (Ref. 81: Zhang & Chen, 2022). The only reason we can produce a "ground truth" control to compare to in this setting is that we use significant knowledge of the problem; we analytically reduce it to a 1D problem and apply numerical methods to solve this 1D problem. It is not a problem where we actually have the ground truth control in closed form.

We hope that these points help clarify the main concerns raised by reviewers. We also provide detailed responses to each reviewer separately.

Meta-review: All reviewers agree that this paper presents two key contributions that are of significant value: (1) the formulation of least-squares objectives to perform regression onto the optimal control of a stochastic optimal control problem (inspired in conditional score matching form diffusion models) and (2) a novel reparameterization trick for gradient estimation, which has a built-in variance reduction component that can be more broadly applicable beyond SOC. The clarity, presentation and quality of this work is excellent, and as pointed out by several reviewers.

During the rebuttal, some initial clarification was needed about which specific class of SOC problems is considered in the work. This was clarified quickly: the work considers a special case of path-integral control problems (non-linear dynamics with additive cost) where state and controls have the same dimensionality.

The only weakness shared by reviewers is that experimental evaluation did not include a real-world realistic scenario. All reviewers agree that this is not a not a reason for rejecting the paper and that overall the strengths outweigh this weakness. I fully agree and thus recommend acceptance.

I do suggest authors to consider the following points to improve the manuscript for its final revised version:

- Emphasize the complexity of the double well task. This was done during the rebuttal and should be incorporated in the revised paper.
- Incorporate as much as possible the additional GMM experiment discussed with reviewer Uos4, including certain details to ensure fair comparison such as adding results using log-variance and another sampling metric [2]. Even if it is not a realistic real-world application, it definitely helps to better understand the variance reduction behavior of the proposed algorithm in relation to other methods.
- Although the paper provides a thorough related work section, the paper would improve if the authors relate their method with other very relevant papers such as [1][3] and [4]. In the latter, a method that bridges cross-entropy and the relative entropy (KL-control cost) is proposed to reduce the gradient variance in path-integral control problems. 

[1] Improved sampling via learned diffusions. Richter et al. ICLR 2024

[2] Beyond ELBOs: A Large-Scale Evaluation of Variational Methods for Sampling. Blessing et al. 2024

[3] Denoising diffusion samplers. Vargas et al. ICLR 2023

[4] Adaptive Smoothing for Path-Integral Control. Thalmeier et al. JMLR 2020
Decision: Accept (poster)