Stochastic Optimal Control Matching

Authors: carles  domingo-enrich
jiequn  han
brandon  amos
joan  bruna
ricky t q chen
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 09:34:33

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.

Main paper: 
 introduction: 
 Stochastic optimal control aims to drive the behavior of a noisy system in order to minimize a given cost. It has myriad applications in science and engineering: examples include the simulation of rare events in molecular dynamics [37, 36, 85, 41] , finance and economics [63, 25] , stochastic filtering and data assimilation [58, 68] , nonconvex optimization [19] , sampling [9] , power systems and energy markets [8, 66] , and robotics [77, 32] . 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 [24] . Recently, it has been the basis of algorithms to sample from unnormalized densities [84, 79, 9, 71] .For continuous-time problems with low-dimensional state spaces, the standard approach to learn the optimal control is to solve the Hamilton-Jacobi-Bellman (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 a neural network and train it applying a stochastic optimization algorithm on a loss function. These methods are known as Iterative Diffusion Optimization (IDO) techniques [59] (see Subsec. 2.2).It is convenient to draw an analogy between stochastic optimal control and continuous normalizing flows (CNFs), which are a generative modeling technique where samples are 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] ), and training them is similar to solving control problems because in both cases one needs to learn high-dimensional vector fields using neural networks, in 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-CNFs) were superseded by diffusion models [75, 40, 76] and flow-matching, a.k.a. stochastic interpolant, methods [55, 1, 65, 2] , which are 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 best-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), or 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 in the bottom right plot of Figure 3 ). In light of this, a 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 we 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 a least-squares regression problem where the goal is to fit a random matching vector field which depends on a 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-entropy loss, which shows that both algorithms have the same landscape in expectation. However, SOCM has an extra flexibility in 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. 1), which is a novel technique for estimating gradients of an expectation of a functional of a 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 2 error with respect to the ground-truth control than all the existing IDO techniques, with around 10× lower error than competing methods in some instances.2 Framework 
setup and preliminaries: 
 Let (Ω, F, (F t ) t≥0 , P) be a fixed filtered probability space on which is defined a Brownian motion B = (B t ) t≥0 . We consider the control-affine problemEQUATIONs.t. dX u t = (b(X u t , t)+σ(t)u(X u t , t)) dt+ √ λσ(t)dB t , X u 0 ∼ p 0 (2) and where X u t ∈ R d is the state, u : R d × [0, T ] → R d is the feedback control and belongs to the set of admissible controls U, f : R d × [0, T ] → R is the state cost, g : R d → R is the terminal cost, b : R d × [0, T ] → R d is the base drift, and σ : [0, T ] → R d×d is the invertible diffusion coefficient and λ ∈ (0, +∞) is the noise level. In App. A we formally define the set U 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.The cost functional for the control u, point x and time t is defined as J(u; x, t) := E T t 1 2 ∥u s (X u s )∥ 2 + f s (X u s ) dt + g(X u T ) X u t = x . That is, the cost functional is the expected value of the control objective restricted to the times [t, T ] with the initial value x at time t. The value function or optimal cost-to-go at (x, t) is defined as the minimum value of the cost functional across all possible controls: V (x, t) := inf u∈U J(u; x, t).(3)Hamilton-Jacobi-Bellman equation and optimal control If we define the infinitesimal generator L := λ EQUATIONThe verification theorem [62, Sec. 2.3] states that if a function V solves the HJB equation above and has certain regularity conditions, then V is the value function (3) of the problem (1)-( 2). An implication of the verification theorem is that for every u ∈ U,EQUATIONIn particular, this implies that the unique optimal control is given in terms of the value function as u * (x, t) = -σ(t) ⊤ ∇V (x, t). Equation ( 5) can be deduced by integrating the HJB equation ( 4) over [t, T ], and taking the conditional expectation with respect to X u t = x. We include the proof of (5) in App. B for completeness.A pair of forward and backward SDEs (FBSDEs) Consider the pair of SDEsEQUATIONEQUATIONwhere Y : Ω×[0, T ] → R and Z : Ω×[0, T ] → R d are progressively measurable1 random processes.It turns out that Y t and Z t defined as Y t := V (X t , t) and Z t := σ(t) ⊤ ∇V (X t , t) = -u * (X t , t) satisfy (7) . We include the proof in App. B for completeness.An analytic expression for the value function From the forward-backward equations ( 6)- (7) , one can derive a closed-form expression for the value function V :EQUATIONwhere X t is the solution of the uncontrolled SDE (6) . This is a classical result, but we still include its proof in App. B. Given that u * (x, t) = -σ(t) ⊤ ∇V (x, t), an immediate, yet important, consequence of ( 8) is the following path-integral representation of the optimal control:EQUATIONRemark 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 score-based diffusion loss [76] can be expressed as∇ x log p t (x) = ∇ x log E Y ∼p data [ exp(-∥x-Y ∥ 2 /(2σ 2 t )) (2πσ 2 t ) d/2]. Note, however, that in (9) the gradient is taken with respect to the initial condition of the process, which requires the development of novel techniques.Conditioned diffusions Let C = C([0, T ]; R d ) be the Wiener space of continuous functions from [0, T ] to R d equipped with the supremum norm, and let P(C) be the space of Borel probability measures over C. For each control u ∈ U, the controlled process in equation ( 2) induces a probability measure in P(C), as the law of the paths X u t , which we refer to as P u . We let P be the probability measure induced by the uncontrolled process (6) , and define the work functionalEQUATIONIt turns out (Lemma 2 in App. B) that the Radon-Nikodym derivative dP u * dP satisfies dP u * dP (X) = exp λ -1 V (X 0 , 0) -W(X, 0) . Also, a straight-forward application of the Girsanov theorem for SDEs (Cor. 1) shows thatEQUATIONwhich means that the only control u ∈ U such that P u = P u * is the optimal control itself.Low-dimensional case: solving the HJB equation For low-dimensional control problems (d ≤ 3), it is 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 methods [11, 57, 4] , which approximate the derivatives and gradients of the value function using finite differences, finite element methods [47] , which involve restricting the solution to domain-dependent function spaces, and semi-Lagrangian schemes [21, 13, 12] , which trace back characteristics and have better stability than finite difference methods. See Greif [33] for an overview on these techniques, and Baňas et al. [4] for a comparison between them. Hutzenthaler et al. [44] introduced the multilevel Picard method, which leverages the Feynman-Kac and the Bismut-Elworthy-Li formulas to beat the curse of dimensionality in some settings [6, 46, 45, 43] .High dimensional methods leveraging FBSDEs The FBSDE formulation in equations ( 6)-( 7) has given rise to multiple methods to learn controls. One such approach is least-squares Monte Carlo (see Pham [63, Chapter 3] and Gobet [28] 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 E et al. [22] , Han et al. [35] . They parameterize the control using a neural network u θ , and use stochastic gradient algorithms to minimize the lossL(u θ , y 0 ) = E[(Y T (y 0 , u θ ) -g(X T )) 2 ], where Y T (y 0 , u θ ) is the process in (7) with initial condition y 0 and control u θ . This algorithm can be seen as a 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, 18, 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 (11) , the key idea is that learning the optimal control is equivalent to finding a control u such that the induced probability measure P u on paths is equal to the probability measure P u * 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 1. 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 p 0 is concentrated at a single point x init ; 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-Leibler divergence between P u and P u * :E P u [log dP u dP u * ].Upon removing constant terms and factors, this loss is equivalent to (see Lemma 3 in App. B):EQUATION)This is exactly the control objective in (1). This fact has been studied extensively [10, 31, 36, 48, 67] . Hence, the relative entropy loss is very natural and widely used; see Onken et al. [61] , Zhang and Chen [84] for examples on multiagent systems and sampling.Solving optimization problems of the form (12) has a long history that dates back to Pontryagin [64] . Note that L Adj (u) depends on u both explicitly, and implicitly through the process X u . To compute the gradient ∇ θ LAdj (u θn ) of a Monte Carlo approximation LAdj (u θn ) of L Adj (u θn ) as required by Algorithm 1, we need to backpropagate through the simulation of the trajectories, which is whyInput: State cost f (x, t), terminal cost g(x), diffusion coeff. σ(t), base drift b(x, t), noise level λ, number of iterations N , batch size m, number of time steps K, initial control parameters θ0, loss L ∈ {L Adj (12) , LCE (13) , LVar v (16), L log Varv (17) , LMom v (18)} for n ∈ {0, . . . , N -1} do Simulate m trajectories of the process X v controlled by v = u θn , e.g., using Euler-Maruyama updates if L ̸ = L Adj then 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 L(u θn ) of the loss L(u θn )Compute the gradients ∇ θ L(u θn ) of L(u θn ) w.r.t. θnObtain θn+1 with via an Adam update on θn (or another stochastic algorithm) end Output: Learned control u θ N we do not detach them from the computational graph. One can alternatively compute the gradient ∇ θ LAdj (u θn ) by explicitly solving an ODE, a technique known as the adjoint method. The adjoint method was introduced by Pontryagin [64] , popularized in deep learning by Chen et al. [20] , and further developed for SDEs in Li et al. [54] .The cross-entropy loss The cross-entropy loss is defined as the Kullback-Leibler divergence between P u * and P u , i.e., flipping the order of the two measures:E P u * [log dP u * dP u ].For an arbitrary v ∈ U, this loss is equivalent to the following one (see Prop. 3(i) in App. B):EQUATIONThe cross-entropy loss has a rich literature [38, 49, 74, 85] and has been recently used in applications such as molecular dynamics [41] . Furthermore, we note that the cross-entropy loss can be significantly simplified and written in terms of the unnormalized L 2 error of the control u with respect to the optimal control u * :EQUATION)This characterization, which is proven in Prop. 3(ii) in App. B, is relevant for us because a similar one can be written for the loss that we propose (see Prop. 2).Variance and log-variance losses For an arbitrary v ∈ U, the variance and the log-variance losses are defined as LVarv (u) = Var P v ( dP u * dP u ) and Llog Varv (u) = Var P v (log dP u * dP u ) wheneverEQUATIONThen, LVarv and Llog Varv are equivalent, respectively, to the following losses (see Lemma 4):EQUATIONEQUATIONThe 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 Varv or L log Varv in Algorithm 1, the variance is computed across the m trajectories in each batch. Moment loss For an arbitrary v ∈ U, the moment loss is defined asEQUATIONwhere Ỹ u,v T is defined in (15) . Note the similarity with the log-variance loss (17) ; the optimal value of y 0 for a fixed u is y* 0 = E[λ -1 g(X v T ) -Ỹ u,v T ],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 it 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 v = 0. 
stochastic optimal control matching: 
 In this section we 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 (9) 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 a novel technique which introduces certain arbitrary matrix-valued functions M t , that we also optimize. Theorem 1 (SOCM loss). For each t ∈ [0, T ], let M t : [t, T ] → R d×d be an arbitrary matrixvalued differentiable function such that M t (t) = Id. Let v ∈ U be an arbitrary control. LetEQUATIONwhere X v is the process controlled by v (i.e.,EQUATIONEQUATIONL SOCM has a unique optimum (u * , M * ), where u * is the optimal control.We refer to M = (M t ) t∈[0,T ] 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 a proof sketch of Thm. 1; the full proofs for all the results in this section are in App. C.Proof sketch of Thm. 1 Let X be the uncontrolled process (6) . Consider the lossEQUATIONClearly, the only optimum of this loss is the optimal control u * . Using the analytic expression of u * in (9), the cross-term can be rewritten as (see Lemma 5 in App. C):EQUATIONIt remains to evaluate the conditional expectation ∇ x E exp -λ -1 T t f (X s , s) dsλ -1 g(X T ) X t = x , which we do by a "reparameterization trick" that shifts the dependence on the initial value x into the stochastic processes-here we introduce a free variable M t -and then applying Girsanov theorem. We coin this the path-wise reparameterization trick: Proposition 1 (Path-wise reparameterization trick for stochastic optimal control). For each t ∈ [0, T ], let M t : [t, T ] → R d×d be an arbitrary continuously differentiable function matrix-valued function such that M t (t) = Id. We have thatEQUATIONWe 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 t arise as the gradients of a perturbation to the process X t . 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,L(u) = E 1 T T 0 u(X t , t) + σ(t) T t M t (s)∇ x f (X s , s) ds + M t (T )∇g(X T ) -λ 1/2 T t (M t (s)∇ x b(X s , s) -∂ s M t (s))(σ -1 ) ⊤ (s)dB s 2 dt × exp -λ -1 T 0 f (X t , t) dt -λ -1 g(X T ) + K.If we perform a change of process from X to X v applying the Girsanov theorem (Cor. 1 in App. C), we obtain the loss L SOCM (u, M ).The following result clarifies the role of reparameterization matrices, connecting the SOCM and cross-entropy losses. Proposition 2 (Bias-variance decomposition of the SOCM loss). The SOCM loss decomposes into a bias term that only depends on u and a variance term that only depends on M :EQUATIONwhere CondVar(w;M ) = E 1 T T 0 w(t, v, X v , B, M t )- E[w(t,v,X v ,B,Mt)α(v,X v ,B)|X v t ,t] E[α(v,X v ,B)|X v t ,t] u * (X v t ,t)2 dt α(v, X v , 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 L SOCM (u, M ) with respect to u is the landscape of the cross-entropy loss L 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. 4). 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.Input: State cost f (x, t), terminal cost g(x), diffusion coeff. σ(t), base drift b(x, t), noise level λ, number of iterations N , batch size m, number of time steps K, initial control parameters θ0, initial matrix parameters ω0, loss LSOCM in (19) for n ∈ {0, . . . , N -1} do Simulate m trajectories of the process X v controlled by v = u θn , e.g., using Euler-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 LSOCM(uθ n , Mω n ) of the loss LSOCM(u θn , Mω n ) in (19) Compute the gradients∇ (θ,ω) LSOCM(uθ n , Mω n ) of LSOCM(uθ n , Mω n ) at (θn, ωn)Obtain θn+1, ωn+1 with via an Adam update on θn, ωn, resp. end Output: Learned control u θ N For good training performance, it is critical that the gradients have high signal-to-noise ratio. Looking at the SOCM loss, a good proxy for low gradient variance is to have low variance for1 T T 0 u(X v t , t) -w(t, v, X v , B, M t ) 2 dt × α(v, X v , B), and this holds when both α(v, X v , B)and w(t, v, X v , B, M t ) have low variance. Next, we present strategies to lower the variance of these two objects.Minimizing the variance of the importance weight α We want to use a vector field v such that Var[α(v, X v , 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 v . 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 set v = u * , the conditional variance Var[α(v, X v , B)|X v 0 = x init ] is zero for any x init ∈ R d .Of course, we do not have access to the optimal control u * , but it is still a good idea to set v 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 u WS (x, t) + u θ (x, t), where the warm-start u WS 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 = (M t ) t∈[0,T ] that minimizes the variance of w(t, v, X v , B, M t ) conditioned on t and X t . Note that this is exactly the term CondVar(w; M ) in the right-hand side of equation (24) . Since CondVar(w; M ) does not depend on the specific v, the optimal M does not depend on v either. And since the second term in the right-hand side of equation ( 24) does not depend on M = (M t ) t∈[0,T ] , minimizing CondVar(w; M ) is equivalent to minimizing L(u) with respect to M .Parameterizing the matrices M t vs solving for the optimal matrices In practice, we parameterize the matrices (M t ) t∈[0,T ] using a function M ω with two arguments (t, s). To enforce that M ω (t, t) = Id, we set M ω (t, s) = e -γ(s-t) Id + (1 -e -γ(s-t) ) Mω (t, s), where ω = (γ, ω), and Mω : R × R → R d×d is an unconstrained neural network. Alternatively, Thm. 4 in App. D shows that the optimal family M * = (M * t ) t∈[0,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 2 K equations and variables, K being the number of discretization time points. However, since the optimal M * does not depend on v (see Remark 1), 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. 
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 2 error incurred by the learned control u. In Figure 2 we plot the control L 2 error for each IDO algorithm described in Subsec. 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) a version of SOCM where the reparameterization matrices M t 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 https://github.com/facebookresearch/SOC-matching.At the end of training, SOCM obtains the lowest L 2 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) we show the squared norm of the gradient of each loss with respect to the parameters θ of the control: algorithms with small noise variance have low error values.In Figure 3 , we plot the control L 2 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 9 in 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 is a particularly interesting and challenging setting because its solution is highly multimodal: g has 1024 modes. Multimodality 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 g and the base drift b for QUADRATIC OU (HARD) are five times those of QUADRATIC OU (EASY). Consequently, the factor α(v, X v , B) initially has a much larger variance for the SOCM methods, and for cross-entropy. As training progresses, u θn gets closer to u * , and consequently the variance of α(v, X v , 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 App. E, we showcase a control warm-start strategy that can help and speed up convergence.We also present experimental results on two-mode Gaussian mixture sampling in increasing dimension, using the Path Integral Sampler [84] . We take Gaussians with means that are 2 units apart, and identity variance. Figure 1 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 1 shows the suboptimality gaps incurred by each algorithm. Cross-entropy, 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. If we keep increasing the dimension, SOCM also fails due to higher variance of α: 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. 
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 a least-squares problem by trying to fit a matching vector field. The training loss is optimized with respect to both the control function and a 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, a novel technique to obtain low-variance estimates of the gradient of the conditional expectation of a functional of a random process with respect to 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 to more challenging problems is that the variance of the factor α(v, X v , B) explodes when f and/or g are large, or when the dimension d is high. The large variance of α is 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.  
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: We explain the type and amount of GPUs we use in Subsec. F.3.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the 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 (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 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 is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or 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 to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to 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 as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a 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 a 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 a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: We did not use existing assets. Guidelines: 
 
 Appendix: Throughout our work, we 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 U of admissible controls is given byU = {u ∈ C 1 (R d × [0, T ]; R d ) | ∃C > 0, ∀(x, s) ∈ R d × [0, T ], b(x, s) ≤ C(1 + |x|)}.(ii) The coefficients b and σ are continuously differentiable, σ has bounded first-order spatial derivatives, and(σσ ⊤ )(x, s) is positive definite for all (x, s) ∈ R d × [0, T ]. Furthermore, there exist constants C, c 1 , c 2 > 0 such that ∥b(x, s)∥ ≤ C(1 + ∥x∥), (linear growth) c 1 ∥ξ∥ 2 ≤ ξ ⊤ (σσ ⊤ )(x, s)ξ ≤ c 2 ∥ξ∥ 2 , (ellipticity) for all (x, s) ∈ R d × [0, T ] and ξ ∈ R d .B Proofs of Sec. 2Proof of (5) By Itô's lemma, we have thatV (X u T , T ) -V (X u t , t) = T t ∂ s V (X u s , s) + ⟨b(X u s , s) + σ(X u s , s)u(X u s , s), ∇V (X u s , s)⟩ + λ 2 d i,j=1 (σσ ⊤ ) ij (X u s , s)∂ xi ∂ xj V (X u s , s) ds + S u t ,whereS u t = √ λ T t ∇V (X u s , s) ⊤ σ(X u s , s) dB s . Note that by (4), ∂ s V (X u s , s) + ⟨b(X u s , s) + σ(X u s , s)u(X u s , s), ∇V (X u s , s)⟩ + λ 2 d i,j=1 (σσ ⊤ ) ij (X u s , s)∂ xi ∂ xj V (X u s , s) = 1 2 ∥(σ ⊤ ∇V )(X u s , s)∥ 2 -f (X u s , s) + ⟨σ(X u s , s)u(X u s , s), ∇V (X u s , s)⟩ = 1 2 ∥(σ ⊤ ∇V )(X u s , s) + u(X u s , s)∥ 2 -1 2 ∥u(X u s , s)∥ 2 -f (X u s , s),and this implies that (26) and taking the conditional expectation with respect to X u t yields the final result.g(X u T ) -V (X u t , t) = T t 1 2 ∥(σ ⊤ ∇V )(X u s , s) + u(X u s , s)∥ 2 -1 2 ∥u(X u s , s)∥ 2 -f (X u s , s) ds + S u t (26) Since E[S u t | X u t = x] = 0, rearrangingProof of ( 6)-( 7) By Itô's lemma, we have thatEQUATIONNote that by (4),∂ s V (X s , s) + ⟨b(X s , s), ∇V (X s , s)⟩ + λ 2 d i,j=1 (σσ ⊤ ) ij (X s , s)∂ xi ∂ xj V (X s , s) = 1 2 ∥(σ ⊤ ∇V )(X s , s)∥ 2 -f (X s , s). Plugging this into (27) concludes the proof.Proof of (8) Since Y s = V (X s , s) and Z s = σ ⊤ (s)∇V (X s , s) = -u * (X s , s) satisfy (7) , we have thatg(X T ) = Y T = Y t - T t (f (X s , s) -1 2 ∥u * (X s , s)∥ 2 ) ds - √ λ T t ⟨u * (X s , s), dB s ⟩.Hence, recalling the definition of the work functional in (10) , we have thatEQUATIONBy Novikov's theorem (Thm. 2), we have thatE[exp(-λ -1 W(X, t))|X t ] = e -λ -1 Yt E exp λ -1/2 T t ⟨u * (X s , s), dB s ⟩ -λ -1 2 T t ∥u * (X s , s)∥ 2 ds X t = e -λ -1Yt , which concludes the proof of (8).Theorem 2 (Novikov's theorem). Let θ s be a locally-H 2 process which is adapted to the natural filtration of the Brownian motion (B t ) t≥0 . DefineEQUATION)If for each t ≥ 0,EQUATIONMoreover, the process Z(t) is a positive martingale, i.e. if (F t ) t≥0 is the filtration associated to the Brownian motion(B t ) t≥0 , then for t ≥ s, E[Z t |F s ] = Z s .Theorem 3 (Girsanov theorem). Let W = (W t ) t∈[0,T ] be a standard Wiener process, and let P be its induced probability measure over C([0, T ]; R d ), known as the Wiener measure. Let Z(t) be as defined in (29) and suppose that the assumptions of Theorem 2 hold. Let (Ω, F) be the σ-algebra associated to B T . For any F ∈ F, define the measureQ(F ) = E P [Z(T )1 F ] .Q is a probability measure because of (30) . Under the probability measure Q, the stochastic process { W (t)} 0≤t≤T defined asW (t) = W (t) - t 0 θ s ds is a standard Wiener process. That is, for any n ≥ 0 and any 0 = t 0 < t 1 < • • • < t n , the increments { W (t i+1 ) -W (t i )} n-1i=0 are independent and Q-Gaussian distributed with mean zero and covariance (t i+1 -t i )I, which means that for any α ∈ R d , the moment generating function of W (t i+1 ) -W (t i ) with respect to Q is as follows:E Q [exp(⟨α, W (t i+1 ) -W (t i )⟩)] := E P exp α, W (t i+1 ) - ti+1 0 θ s ds -W (t i ) + ti 0 θ s ds Z(T ) = exp (ti+1-ti)∥α∥ 2 2 .Corollary 1 (Girsanov theorem for SDEs). If the two SDEsdX t = b 1 (X t , t) dt + σ(X t , t) dB t , X 0 = x init dY t = (b 1 (Y t , t) + b 2 (Y t , t)) dt + σ(Y t , t) dB t , Y 0 = x initadmit unique strong solutions on [0, T ], then for any bounded continuous functional Φ on C([0, T ]), we have thatE[Φ(X)] = E Φ(Y ) exp - T 0 σ(Y t , t) -1 b 2 (Y t , t) dB t -1 2 T 0 ∥σ(Y t , t) -1 b 2 (Y t , t)∥ 2 dt = E Φ(Y ) exp - T 0 σ(Y t , t) -1 b 2 (Y t , t) d Bt + 1 2 T 0 ∥σ(Y t , t) -1 b 2 (Y t , t)∥ 2 dt , where Bt = B t + t 0 σ(Y s , s) -1 b 2 (Y s , s) ds.More generally, b 1 and b 2 can be random processes that are adapted to filtration of B. Lemma 2. For an arbitrary v ∈ U, let P v and P be respectively the laws of the SDEsdX v t = (b(X v t , t) + σ(t)v(X v t , t)) dt + √ λσ(t)dB t , X v 0 ∼ p 0 , dX t = b(X t , t) dt + √ λσ(t)dB t , X 0 ∼ p 0 .We have thatEQUATION)whereB v t := B t + λ -1/2 t 0 v(X v s , s) ds.For the optimal control u * , we have thatEQUATION)EQUATIONwhere 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) ,EQUATIONwhich implies thatdP u * dP (X) = exp λ -1/2 T 0 ⟨u * (X t , t), dB t ⟩ -λ -1 2 T 0 ∥u * (X t , t)∥ 2 dt = exp λ -1 V (X 0 , 0) -W(X, 0) .To prove (33) W(X u * , 0) = V (X u * 0 , 0) + 1 2 T t ∥u * (X u * s , s)∥ 2 ds - √ λ T t ⟨u * (X u * s , s), dB v s ⟩. Hence, dP dP u * (X u * ) = exp -λ -1/2 T 0 ⟨u * (X u * t , t), dB u * t ⟩ + λ -1 2 T 0 ∥u * (X u * t , t)∥ 2 dt = exp λ -1 -V (X u * 0 , 0) + W(X u * , 0) . Lemma 3.The following expression holds:EQUATIONProof. To prove (36) , we writelog dP u * dP u (X u ) = log dP u * dP (X u ) dP dP u (X u ) = log dP u * dP (X u ) + log dP dP u (X u ) = λ -1 V (X u 0 , 0) - T 0 f (X u t , t) dt -g(X u T ) -λ -1/2 T 0 ⟨u(X u t , t), dB t ⟩ -λ -1 2 T 0 ∥u(X u t , t)∥ 2 dt .Since E P u log dP u dP u *= -E P u log dP u * dP u , and E P u T 0 ⟨u(X u t , t), dB t ⟩] = 0, the result follows.Proposition 3. (i) The following two expressions hold for arbitrary controls u, v in the class U of admissible controls:EQUATIONEQUATION)When p 0 is concentrated at a single point x init , the terms V (x init , 0) are constant and can be removed without modifying the landscape. In other words, LCE and L CE are equal up to constant terms and constant factors.(ii) When p 0 is a generic probability measure, LCE and L CE have different landscapes, and L CE (u) =E P u * log dP u * dP u exp -λ -1 V (X u * 0 , 0). u * is still the only minimizer of the loss L CE , and for some constant K, we have thatEQUATION)Proof. We begin with the proof of (i), and prove (37) first. Note that by the Girsanov theorem (Thm. 3),EQUATIONNote that by equations ( 32) and (34) ,EQUATION)whereEQUATIONIf we plug ( 41) and ( 42) into the right-hand side of (40), we obtainE P u * log dP u * dP u (X u * ) = -E P u * (log dP u dP (X v ) + log dP dP u * (X v )) dP u * dP (X v ) dP dP u (X v ) = -E λ -1/2 T 0 ⟨u(X v t , t), dB t ⟩ + λ -1 T 0 ⟨u(X v t , t), v(X v t , t)⟩ dt -λ -1 2 T 0 ∥u(X v t , t)∥ 2 dt + λ -1 -V (X v 0 , 0) + W(X v , 0) × exp λ -1 V (X v 0 , 0) -W(X v , 0) -λ -1/2 T 0 ⟨v(X v t , t), dB t ⟩ -λ -1 2 T 0 ∥v(X v t , t)∥ 2 dt ,which concludes the proof.To show (38) , we use that by Cor. 1,dP u dP u * (X u * ) = exp -λ -1/2 T 0 ⟨u * (X u * t , t)-u(X u * t , t), dB t ⟩-λ -1 2 T 0 ∥u * (X u * t , t)-u(X u * t , t)∥ 2 dt .Hence, 37) can be removed without modifying the landscape of the loss. Hence, we are left withE P u * log dP u * dP u = -E P u * log dP u dP u * = λ -1 2 E T 0 ∥u * (X u * t , t) -u(X u * t , t)∥ 2 dt . Next, we prove (ii). The first instance of V (X v 0 , 0) in (EQUATION+ λ -1 2 T 0 ∥u(X v t , t)∥ 2 dt -λ -1 W(X v , 0) × exp λ -1 V (X v 0 , 0)-W(X v , 0) -λ -1/2 T 0 ⟨v(X v t , t), dB t ⟩-λ -1 2 T 0 ∥v(X v t , t)∥ 2 dtAnd this can be expressed asLCE (u) = E g(u; X v 0 ) exp λ -1 V (X v 0 , 0) , where g(u; x) = E -λ -1/2 T 0 ⟨u(X v t , t), dB v t ⟩ -λ -1 T 0 ⟨u(X v t , t), v(X v t , t)⟩ dt + λ -1 2 T 0 ∥u(X v t , t)∥ 2 dt -λ -1 W(X v , 0) × exp -λ -1 W(X v , 0)-λ -1/2 T 0 ⟨v(X v t , t), dB t ⟩-λ -1 2 T 0 ∥v(X v t , t)∥ 2 dt |X v 0 = x .If we consider g(u; x) as a loss function for u, note that it is equivalent to the loss LCE (u) equation in (43) for the choice p 0 = δ x , i.e., p 0 concentrated at x. Since the optimal control u * is independent of the starting distribution p 0 , we deduce that u * is the unique minimizer of g(u; x), for all x ∈ R d . In consequence, u * is the unique minimizer ofL CE (u) = E[g(u; X v 0 )].To prove (39) , note that up to a constant term, the only difference between LCE (u) and L CE (u) is the expectation is reweighted importance weight exp -λ -1 V (X v 0 , 0) . Lemma 4. (i) We can rewriteLVarv (u) = Var exp Ỹ u,v T -λ -1 g(X v T ) + λ -1 V (X v 0 , 0) , Llog Varv (u) = Var Ỹ u,v T -λ -1 g(X v T ) + λ -1 V (X v 0 , 0). When p 0 is concentrated at x init , the terms V (x init , 0) are constants and can be removed without modifying the landscape. In other words, LVarv and LlogVarv are equal to L Varv and L log Varv up to a constant term and a constant factor, respectively.(ii) When p 0 is general, LVarv and L Varv have a different landscape, and the optimum of L Varv may be different from u * . A related loss that does preserve the optimum is:LVarv (u) = E[Var P v ( dP u * dP u (X v )|X v 0 ) exp(-λ -1 V (X v 0 , 0))] = E[Var exp( Ỹ u,v T -λ -1 g(X v T ))|X v 0 ].In practice, this is implemented by sampling the m trajectories in one batch starting at the same pointX v 0 . (iii) Also, LlogVarv and L log Varv have a different landscape, and the optimum of L log Varv may be different from u * . In particular, L log Varv (u) = Var P v (logdP u * dP u (X v ) exp(-λ -1 V (X v 0 , 0))). A loss that does preserve the optimum u * is Llog Varv (u) = E[Var P v (log dP u * dP u (X v )|X v 0 ) exp(-λ -1 V (X v 0 , 0))] = E[Var Ỹ u,v T -λ -1 g(X v T )|X v 0 ].Proof. Using ( 34) and ( 31), we have thatdP u * dP (X v ) = exp λ -1 V (X v 0 , 0) -W(X v , 0) , dP dP u (X v ) = exp -λ -1/2 T 0 ⟨u(X v t , t), dB v t ⟩ + λ -1 2 T 0 ∥u(X v t , t)∥ 2 dt = exp -λ -1/2 T 0 ⟨u(X v t , t), dB t ⟩ -λ -1 T 0 ⟨u(X v t , t), v(X v t , t)⟩ dt + λ -1 2 T 0 ∥u(X v t , t)∥ 2 dt . Hence, log dP u * dP u (X v ) = log dP u * dP (X v ) + log dP dP u (X v ) = Ỹ u,v T -λ -1 g(X v T ) + λ -1 V (X v 0 , 0).Since LVarv (u) = Var P v ( dP u * dP u ) and Llog Varv (u) = Var P v (log dP u * dP u ), this concludes the proof of (i). To prove (ii), note that for general p 0 , V (X v 0 , 0) is no longer a constant, but it is if we condition on X v 0 . The proof of (iii) is analogous.C Proofs of Sec. 3C.1 Proof of Thm. 1 and Prop. 2We prove Thm. 1 and Prop. 2 at the same time. Recall that by (9) , the optimal control is of the form u * (x, t) = -σ(t) ⊤ ∇V (x, t). Consider the lossL(u) = E 1 T T 0 u(X t , t) + σ(t) ⊤ ∇V (X t , t) 2 dt exp -λ -1 T 0 f (X t , t) dt -λ -1 g(X T ) .Clearly, the unique optimum of L is -σ(t) ⊤ ∇V . We can rewrite L asL(u) = E 1 T T 0 u(X t , t) 2 + 2⟨u(X t , t), σ(t) ⊤ ∇V (X t , t)⟩ + ∥σ(t) ⊤ ∇V (X t , t) 2 dt(44) × exp -λ -1 T 0 f (X t , t) dt -λ -1 g(X T ) .Hence, we can express L as a sum of three terms: one involving ∥u(X t , t)∥ 2 , another involving ⟨u(X t , t), σ(t) ⊤ V (X t , t)⟩, and a third one, which is constant with respect to u, involving ∥∇V (X t , t)∥ 2 . The following lemma provides an alternative expression for the cross term: Lemma 5. The following equality holds:EQUATIONProof. Recall the definition of W(X, t) in (35) , which means thatEQUATIONLet {F t } t∈[0,T ] be the filtration generated by the Brownian motion B. Then, equation (9) implies thatEQUATIONWe proceed as follows:E 1 T T 0 ⟨u(X t , t), σ(t) ⊤ ∇V (X t , t)⟩ dt exp -λ -1 W(X, 0) (i) = -λE 1 T T 0 u(X t , t), σ(t) ⊤ ∇xE exp -λ -1 W(X,t) Ft E exp -λ -1 W(X,t) Ft × E exp -λ -1 W(X, t) F t exp -λ -1 t 0 f (X s , s) ds dt = -λE 1 T T 0 u(X t , t), σ(t) ⊤ ∇ x E exp -λ -1 W(X, t) F t exp -λ -1 t 0 f (X s , s) ds dt (ii) = -λE 1 T T 0 u(X t , t), σ(t) ⊤ ∇ x E exp -λ -1 W(X, t) X t = x exp -λ -1 t 0 f (X s , s) ds dt .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 an SDE.The following proposition, which we prove in Subsec. C.2, provides an alternative expression for∇ x E exp -λ -1 T t f (X s , s) ds -λ -1 g(X T ) X t = x .The technique, which is novel and we denote by Girsanov reparamaterization trick, is of independent interest and may be applied in other settings, as we discuss in Sec. 5. Proposition 1 (Path-wise reparameterization trick for stochastic optimal control). For each t ∈ [0, T ], let M t : [t, T ] → R d×d be an arbitrary continuously differentiable function matrix-valued function such that M t (t) = Id. We have thatEQUATIONPlugging ( 23) into the right-hand side of (45), we obtain thatE 1 T T 0 ⟨u(X t , t), σ(t) ⊤ ∇V (X t , t)⟩ dt exp -λ -1 T 0 f (X t , t) dt -λ -1 g(X T ) = E 1 T T 0 u(X t , t), σ(t) ⊤ T t M t (s)∇ x f (X s , s) ds + M t (T )∇g(X T ) -λ 1/2 T t (M t (s)∇ x b(X s , s) -∂ s M t (s))(σ -1 ) ⊤ (s)dB s dt × exp -λ -1 T 0 f (X t , t) dt -λ -1 g(X T) . If we plug this into the right-hand side of (44) and complete the squared norm, we get thatL(u) = E 1 T T 0 ( u(X t , t)-w(t, X, B, M t ) 2 -w(t, X, B, M t ) 2 + u * (X t , t) 2 ) dt exp -λ -1 W(X, 0)where w is defined as:w(t, X, B, M t ) = σ(t) ⊤ - T t M t (s)∇ x f (X s , s) ds -M t (T )∇g(X T ) + λ 1/2 T t (M t (s)∇ x b(X s , s) -∂ s M t (s))(σ -1 ) ⊤ (s)dB s .We also define Φ(u; X, B) asΦ(u; X, B) = 1 T T 0 ( u(X t , t) -w(t, X, B, M t )2 ) dt. Now, by the Girsanov theorem (Thm. 3), we have that for an arbitrary control v ∈ U,E[Φ(u; X, B) exp -λ -1 W(X, 0) ] = E Φ(u; X v , B v ) exp -λ -1 W(X v , 0) -λ -1/2 T 0 ⟨v(X v t , t), dB v t ⟩ + λ -1 2 T 0 ∥v(X v t , t)∥ 2 dt = E Φ(u; X v , B v ) exp -λ -1 W(X v , 0) -λ -1/2 T 0 ⟨v(X v t , t), dB t ⟩ -λ -1 2 T 0 ∥v(X v t , t)∥ 2 dt ,where B v t := B t + λ -1/2 t 0 v(X v s , s) ds. Reexpressing B v in terms of B, we can rewrite Φ(u; X v , B v ) and w(t, X v , B v , M t ) as follows:Φ(u; X v , B v ) = 1 T T 0 u(X v t , t) -w(t, X v , B v , M t ) 2 dt, w(t, X v , B v , M t ) = σ(t) ⊤ - T t M t (s)∇ x f (X v s , s) ds -M t (T )∇g(X v T ) + λ 1/2 T t (M t (s)∇ x b(X v s , s) -∂ s M t (s))(σ -1 ) ⊤ (X v s , s)dB s + T t (M t (s)∇ x b(X v s , s) -∂ s M t (s))(σ -1 ) ⊤ (X v s , s)v(X v s , s)ds .Putting everything together, we obtain thatL(u) = L SOCM (u, M ) -K,where L(u, M ) is the loss defined in (19) (note that w(t, v, X v , B, M t ) := w(t, X v , B v , M t )), andK = E 1 T T 0 ( w(t, X, B, M t ) 2 -u * (X t , t) 2 ) dt exp -λ -1 W(X, 0)To complete the proof of equation ( 24), remark that L(u) can be rewritten asL(u) = E 1 T T 0 u(X t , t) -u * (X t , t) 2 dt exp -λ -1 W(X, 0) = E 1 T T 0 u(X t , t) -u * (X t , t) 2 dt dP u * dP (X) exp(-λ -1 V (X 0 , 0)) = E 1 T T 0 u(X u * t , t) -u * (X u * t , t)2 dt exp(-λ -1 V (X u * 0 , 0)) . It only remains to reexpress K. Note that by Prop. 1, we have thatu * (X t , t) = E w(t,X,B,Mt) exp -λ -1 W(X,0) |Ft E exp -λ -1 W(X,0) |Ft = E w(t,X,B,Mt) dP u * dP (X)|Ft exp(-λ -1 V (X0,0)) E dP u * dP (X)|Ft exp(-λ -1 V (X0,0)) = E w(t,X,B,Mt) dP u * dP (X)|Ft E dP u * dP (X)|Ft = E w(t, X u * , B u * , M t )|X u * t = X tHence, using the Girsanov theorem (Thm. 3) several times, we have thatK = E 1 T T 0 w(t, X u * , B u * , M t )∥ 2 -∥E w(t, X u * , B u * , M t )|X u * t 2 dt exp(-λ -1 V (X u * 0 , 0)) (48) = E 1 T T 0 w(t, X u * , B u * , M t ) -E w(t, X u * , B u * , M t )|X u * t 2 dt exp(-λ -1 V (X u * 0 , 0)) = E 1 T T 0 w(t, X, B, M t ) -E[ w(t,X,B,Mt) exp(-λ -1 W(X,0))|Xt] E[exp(-λ -1 W(X,0))|Xt] 2 dt exp(-λ -1 W(X, 0)) = E 1 T T 0 w(t, v, X v , B, M t ) - E[w(t,v,X v ,B,Mt)α(v,X v ,B)|X v t ] E[α(v,X v ,B)|X v t ] 2 dt α(v, X v , B) ,which concludes the proof, noticing that K = Var(w; M ).Remark 1 (The optimal M t is the same for all v). Looking at equation (48) , we observe that Var(w; M ) does not depend on the base control v. Since minimizing L SOCM (u, M ) with respect to M is equivalent to minimizing Var(w; M ), we deduce that the optimal M does not depend on the vector field v.We prove a more general statement (Prop. 4), and show that Prop. 1 is a particular case of it.Proposition 4 (Path-wise reparameterization trick). Let (Ω, F, P) be a probability space, and B : Ω × [0, T ] → R d be a Brownian motion. Let X : Ω × [0, T ] → R d be the uncontrolled process given by (6), and let ψ : Ω × R d × [0, T ] → R d be an arbitrary random process such that:• For all z ∈ R d , the process ψ(•, z, •) : Ω × [0, T ] → R d is adapted to the filtration (F s ) s∈[0,T ] of the Brownian motion B.• For all ω ∈ Ω, ψ(ω, •, •) : R d × [0, T ] → R d is a twice-continuously differentiable function such that ψ(ω, z, 0) = z for all z ∈ R d , and ψ(ω, 0, s) = 0 for all s ∈ [0, T ].Let F : C([0, T ]; R d ) → R be a Fréchet-differentiable functional. We use the notation X + ψ(z, •) = (X s (ω) + ψ(ω, z, s)) s∈[0,T ] to denote the shifted process, and we will omit the dependency of ψ on ω in the proof. Then,∇ x E exp -F (X) X 0 = x (49) = E -∇ z F (X +ψ(z, •))| z=0 +λ -1/2 T 0 (∇ z ψ(0, s)∇ x b(X s , s)-∇ z ∂ s ψ(0, s))(σ -1 ) ⊤ (s)dB s × exp -F (X) X 0 = x Proof of Prop. 1.Given a family of functions (M t ) t∈[0,T ] satisfying the conditions in Prop. 1, we can define a family(ψ t ) t∈[0,T ] of functions ψ t : R d × [t, T ] → R d as ψ t (z, s) = M t (s) ⊤ z. Note that ψ t (z, t) = z for all z ∈ R d and ψ t (0, s) = 0 for all s ∈ [t, T ],and that ∇ z ψ t (z, s) = M t (s). Hence, ψ t can be seen as a random process which is constant with respect to ω ∈ Ω, and which fulfills the conditions in Prop. 4 up to a trivial time change of variable from [t, T ] to [0, T ].We also define the family(F t ) t∈[0,T ] of functionals F t : C([t, T ]; R d ) → R as F t (X) = λ -1 T t f (X s , s) ds + λ -1 g(X T ). We have that ∇ z F t (X +ψ t (z, •)) = ∇ z λ -1 T t f (X s + ψ t (z, s), s) ds + λ -1 g(X T + ψ t (z, T )) (i) = λ -1 T t ∇ z ψ t (z, s)∇f (X s +ψ t (z, s), s) ds+λ -1 ∇ z ψ t (z, T )∇g(X T +ψ t (z, T )) = λ -1 T t M t (s)∇f (X s +ψ t (z, s), s) ds+λ -1 M t (T )∇g(X T +ψ t (z, T )),where equality (i) holds by the Leibniz rule. Using that ψ t (0, s) = 0, we obtain that:∇ z F t (X + ψ t (z, •)) z=0 = λ -1 T t ∇ z ψ t (0, s)∇f (X s , s) ds + λ -1 ∇ z ψ t (T, 0)∇g(X T ),Up to a trivial time change of variable from [t, T ] to [0, T ], Prop. 1 follows from plugging these choices into equation ( 49).Remark 2. We can use matrices M t (s) that depend on the process X up to time s, since the resulting processes ψ t (•, z, •) are adapted to the filtration of the Brownian motion B. More specifically, if we let M t : R d × [t, T ] → R d×d be an arbitrary continuously differentiable function matrix-valued function such that M t (x, t) = Id for all x ∈ R d , and we define the exponential moving average of X as the process X (υ) given byX (υ) t = υ t 0 e -υ(t-s) X s ds, we have that d ds M t (X (υ) s , s) = ⟨∇M t (X (υ) s , s), dX (υ) s ds ⟩+∂ s M t (X (υ) s , s) = λ⟨∇ x M t (X (υ) s , s), X s -X (υ) s ⟩+∂ s M t (X (υ) s , s),and we can write∇ x E exp -λ -1 T t f (X s , s) ds -λ -1 g(X T ) X t = x = E -λ -1 T t M t (X (υ) s , s)∇ x f (X s , s) ds -λ -1 M t (X (υ) T , T )∇g(X T ) + λ -1/2 T t (M t (X (υ) s , s)∇ x b(X s , s) -d ds M t (X (υ) s , s))(σ -1 ) ⊤ (s)dB s × exp -λ -1 T t f (X s , s) ds -λ -1 g(X T ) X t = x .Plugging this into the proof of Thm. 1, we would obtain a variant of SOCM (Alg. 2) where the matrix-valued neural network M ω takes inputs (t, s, x) instead of (t, s). Since the optimization class is larger, from the bias-variance in Prop. 2 we 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 M ω ; one would need K(K+1)m 2 NFE per batch instead of only K(K+1) 2 , which may be too expensive if the architecture of M ω is large. A way to speed up the computation per batch is to parameterize M t using cubic splines.dX s = b(X s , s) ds + √ λσ(s) dB s , X 0 ∼ p 0 ,is the SDE for the uncontrolled process. For arbitrary x, z ∈ R d , we consider the following SDEs conditioned on the initial points:EQUATIONEQUATIONSuppose that ψ : R d × [0, T ] → R d satisfies the properties in the statement of Prop. 4. If X(x) is a solution ofd X(x) s = (b( X(x) s + ψ(z, s), s) -∂ s ψ(z, s)) ds + √ λσ(s) dB s , X(x) 0 = x, then X (x+z) = X(x) + ψ(z, •) is a solution of (50). This is becauseEQUATIONdX (x+z) s = d X(x) s + ∂ s ψ(z, s) ds = (b( X(x) s + ψ(z, s), s) -∂ s ψ(z, s)) ds + √ λσ(s) dB s + ∂ s ψ(z, s) ds = b(X (x+z) s , s) ds + √ λσ(s) dB s ,Note that we may rewrite (51) asdX (x) s = (b(X (x) s + ψ(z, s), s) -∂ s ψ(z, s)) ds + (b(X (x) s , s) -b(X (x) s + ψ(z, s), s) + ∂ s ψ(z, s)) ds + √ λσ(s) dB s , X (x) t ∼ p 0 .Hence, since ψ(z, s) is a random process adapted to the filtration of B, we can apply the Girsanov theorem for SDEs (Corollary 1) on X(x) and X (x) , and we have that for any bounded continuous functional Φ,EQUATION)We can writeEQUATIONEquality (i) holds by the definition of X (x+z) , equality (ii) holds by the factX (x+z) s = X(x)s +ψ(z, s), equality (iii) holds by equation ( 52), and equality (iv) holds by the definition of X (x) s . We conclude the proof by differentiating the right-hand side of ( 53) with respect to z. Namely,∇ x E exp -F (X) X 0 = x = ∇ z E exp -F (X) X 0 = x + z z=0 (i) = E -∇ z F (X + ψ(z, •)) + λ -1/2 T 0 (∇ z ψ(0, s)∇ x b(X s , s) -∇ z ∂ s ψ(0, s))(σ -1 ) ⊤ (s)dB s × exp -F (X) X 0 = xIn equality (i) we used (53) , and that:• by the Leibniz rule,∇ z T 0 ∥σ(s) -1 (b(X s + ψ(z, s), s) -b(X s , s) -∂ s ψ(z, s))∥ 2 ds z=0 = T 0 ∇ z ∥σ(s) -1 (b(X s + ψ(z, s), s) -b(X s , s) -∂ s ψ(z, s))∥ 2z=0 ds = 0. • and by the Leibniz rule for stochastic integrals (see [42] ),∇ z T 0 σ(s) -1 (b(X s + ψ(z, s), s) -b(X s , s) -∂ s ψ(z, s)) dB s z=0 = T 0 (∇ z ψ(0, s)∇ x b(X s , s) -∇ z ∂ s ψ(0, s))(σ -1 ) ⊤ (s) dB s .In this subsection, we provide an informal, intuitive derivation of the path-wise reparameterization trick as stated in Prop. 4. For simplicity, we particularize the functional F to F (X) = λ -1 T 0 f (X s , s) ds + λ -1 g(X T ). Consider the Euler-Maruyama discretization of the uncontrolled process X defined in (6), with K + 1 time steps (let δ = T /K be the step size). This is a family of random variables X = ( Xk ) k=0:K defined asX0 ∼ p 0 , Xk+1 = Xk + δb( Xk , kδ) + √ δλσ(kδ)ε k , ε k ∼ N (0, I).Note that we can approximateE exp -λ -1 T 0 f (X s , s) ds -λ -1 g(X T ) X 0 = x ≈ E exp -λ -1 δ K-1 k=0 f ( Xk , s) -λ -1 g( XK ) X0 = x ,and that this is an equality in the limit K → ∞, as the interpolation of the Euler-Maruyama discretization X(x) converges to the process X (x) . Now, remark that for k ∈ {0, . . . , K -1}, Xk+1 | Xk ∼ N ( Xk + δb( Xk , kδ), δλ(σσ ⊤ )(kδ)). Hence,E exp -λ -1 δ K-1 k=0 f ( Xk , s) -λ -1 g( XK ) X0 = x = C -1 (R d ) K exp -λ -1 δ K-1 k=0 f (x k , s) -λ -1 g(x K ) -1 2δλ K-1 k=1 ∥σ -1 (kδ)(x k+1 -xk -δb(x k , kδ))∥ 2 -1 2δλ ∥σ -1 (0)(x 1 -x-δb(x, 0))∥ 2 dx 1 • • • dx K ,where C = (2πδλ) K K-1 k=0 det((σσ ⊤ )(kδ)). Now, let ψ : R d × [0, T ] → R d be an arbitrary twice differentiable function such that ψ(z, 0) = z for all z ∈ R d , and ψ(0, s) = 0 for all s ∈ [0, T ]. We can writeEQUATIONIn the last equality, we used that for k ∈ {1, . . . , K}, the variables xk are integrated over R d , which means that adding an offset ψ(z, kδ) does not change the value of the integral. We also used that ψ(z, 0) = z. Now, for fixed values of x = (x 1 , . . . , xK ), and letting x0 = x, we defineG x(z) = λ -1 δ K-1 k=0 f (x k + ψ(z, kδ), s) + λ -1 g(x K + ψ(z, Kδ)) + 1 2δλ K-1 k=0 ∥σ -1 (kδ)(x k+1 +ψ(z, (k + 1)δ)-xk -ψ(z, kδ)-δb(x k +ψ(z, kδ), kδ))∥ 2 .Using that ψ(0, s) = 0 for all s ∈ [0, T ], we have that:G x(0) = λ -1 δ K-1 k=0 f (x k , s) + λ -1 g(x K ) + 1 2δλ K-1 k=0 ∥σ -1 (kδ)(x k+1 -xk -δb(x k , kδ))∥ 2 . ∇G x(z)| z=0 = λ -1 δ K-1 k=0 ∇ψ(0, kδ)∇f (x k , s) + λ -1 ∇ψ(0, Kδ)∇g(x K ) + 1 δλ K-1 k=0 (∇ z ψ(0, (k + 1)δ) -∇ z ψ(0, kδ)-δ∇ψ(0, kδ)∇b(x k , kδ)) × ((σ -1 ) ⊤ σ -1 )(kδ)(x k+1 -xk -δb(x k , kδ))And we can express the right-hand side of (54) in terms of G x(0) and ∇G x(z)| z=0 :∇ z C -1 (R d ) K exp -G x(z) dy 1 • • • dy K = -C -1 (R d ) K ∇G x(z)| z=0 exp -G x(0) dy 1 • • • dy K We define ϵ k = 1 √ δλ σ -1 (kδ)(x k+1 -xk -δb(x k , kδ)), and then, we are able to writeEQUATIONEQUATIONThen, taking the limit K → ∞ (i.e. δ → 0), we recognize (55) as Euler-Maruyama discretization of the uncontrolled process X in equation ( 6) conditioned on X 0 = x, and the last term in (56) as the Euler-Maruyama discretization of the stochastic integral λ -1/2 T 0 (∂ s ∇ z ψ(0, s) -∇ψ(0, s)∇b(X (x) s , s))(σ -1 ) ⊤ (s) dB s . Thus, lim K→∞ ∇ x E exp -λ -1 δ K-1 k=0 f ( Xk , s) -λ -1 g( XK ) = E -λ -1 T 0 ∇ψ(0, s)∇ x f (X s , s) ds -λ -1 ∇ψ(0, T )∇g(X T ) + λ -1/2 T 0 (∇ψ(0, s)∇ x b(X s , s) -∂ s ∇ψ(0, s))(σ -1 ) ⊤ (s) dB s × exp -λ -1 T 0 f (X s , s) ds -λ -1 g(X T ) X 0 = x ,L SOCM-Adj (u) := E 1 T T 0 u(X v t , t) + σ(t) ⊤ a(t, X v ) 2 dt × α(v, X v , B) ,where X v is the process controlled by v (i.e., dX t = (b(X t , t) + σ(t)v(X t , t)) dt + √ λσ(X t , t) dB t and X 0 ∼ p 0 ), α(v, X v , B) is the importance weight defined in (20) , and a(t, X v ) is the solution of the ODEda(t) dt = -∇ x b(X v t , t)a(t) -∇ x f (X v t , t), a(T ) = ∇g(X vT ), L SOCM-Adj has a unique optimum, which is the optimal control u * .Proof. The proof follows the same structure as that of Thm. 1. Instead of plugging the path-wise reparameterization trick (Prop. 1) in the right-hand side of ( 22), we make use of Lemma 6 to evaluate∇ x E exp -λ -1 T 0 f (X t , t) dt -λ -1 g(X T ) |X 0 = x .Particular cases of the result in Lemma 6 have been used in previous works such as [54, 51] . We present a more general form that covers state costs f , as well as stochastic integrals. We also present a simpler proof of the result based on Lagrange multipliers. Lemma 6 (Adjoint method for SDEs). [54, 51] Let X : Ω × [0, T ] → R d be the uncontrolled process defined in (6), with initial condition X 0 = x. We define the random process a : Ω × [0, T ] → R d such that for all ω ∈ Ω, using the short-hand a(t) := a(ω, t),da t (ω) = -∇ x b(X t (ω), t)a t (ω) -∇ x f (X t (ω), t) dt -∇ x h(X t (ω), t) dB t , a T (ω) = ∇ x g(X T (ω)), we have that ∇ x E T 0 f (X t (ω), t) dt + T 0 ⟨h(X t (ω), t), dB t ⟩ + g(X T (ω))|X 0 (ω) = x = E a 0 (ω) , ∇ x E exp - T 0 f (X t (ω), t) dt - T 0 ⟨h(X t (ω), t), dB t ⟩ -g(X T (ω)) |X 0 (ω) = x = -E a 0 (ω) exp - T 0 f (X t (ω), t) dt - T 0 ⟨h(X t (ω), t), dB t ⟩ -g(X T (ω)) |X 0 (ω) = x .Proof. We will use an approach based on Lagrange multipliers. Define a process a : Ω × [0, T ] → R d such that for any ω ∈ Ω, a(ω, •) is differentiable. For a given ω ∈ Ω, we can writeT 0 f (X t (ω), t) dt + T 0 ⟨h(X t (ω), t), dB t ⟩ + g(X T (ω)) = T 0 f (X t (ω), t) dt + T 0 ⟨h(X t (ω), t), dB t ⟩ + g(X T (ω)) - T 0 ⟨a t (ω), (dX t (ω) -b(X t (ω), t) dt -σ(t) dB t )⟩. By Lemma 7, we have that T 0 ⟨a t (ω), dX t (ω)⟩ = ⟨a T (ω), X T (ω)⟩ -⟨a 0 (ω), X 0 (ω)⟩ - T 0 ⟨X t (ω), dat dt (ω)⟩ dt. Hence, ∇ x T 0 f (X t (ω), t) dt + T 0 ⟨h(X t (ω), t), dB t ⟩ + g(X T (ω)) = ∇ x T 0 f (X t (ω), t) dt + T 0 ⟨h(X t (ω), t), dB t ⟩ + g(X T (ω)) -⟨a T (ω), X T (ω)⟩ + ⟨a 0 (ω), X 0 (ω)⟩ + T 0 ⟨a t (ω), b(X t (ω), t)⟩ + ⟨ dat dt (ω), X t (ω)⟩ dt + T 0 ⟨a t (ω), σ(t) dB t ⟩ = T 0 ∇ x X t (ω)∇ x f (X t (ω), t) dt + T 0 ∇ x X t (ω)∇ x h(X t (ω), t) dB t + ∇ x X T (ω)∇ x g(X T (ω)) -∇ x X T (ω)a T (ω) + ∇ x X 0 (ω)a 0 (ω) + T 0 ∇ x X t (ω)∇ x b(X t (ω), t)a t (ω) + ∇ x X t (ω) dat dt (ω) dt = T 0 ∇ x X t (ω) ∇ x f (X t (ω), t) + ∇ x b(X t (ω), t)a t (ω) + dat dt (ω) dt + ∇ x X T (ω) ∇ x g(X T (ω)) -a T (ω) + a 0 (ω) + T 0 ∇ x X t (ω)∇ x h(X t (ω), t) dB t .In the last line we used that ∇ x X 0 (ω) = ∇ x x = I. If choose a such thatda t (ω) = -∇ x b(X t (ω), t)a t (ω) -∇ x f (X t (ω), t) dt -∇ x h(X t (ω), t) dB t , a T (ω) = ∇ x g(X T (ω)),then we obtain that∇ x T 0 f (X t (ω), t) dt + T 0 ⟨h(X t (ω), t), dB t ⟩ + g(X T (ω)) = a 0 (ω),and by the Leibniz rule,∇ x E T 0 f (X t (ω), t) dt + T 0 ⟨h(X t (ω), t), dB t ⟩ + g(X T (ω)) = E ∇ x T 0 f (X t (ω), t) dt + T 0 ⟨h(X t (ω), t), dB t ⟩ + g(X T (ω)) = E a 0 (ω) ,and∇ x E exp - T 0 f (X t (ω), t) dt - T 0 ⟨h(X t (ω), t), dB t ⟩ -g(X T (ω)) = -E ∇ x T 0 f (X t (ω), t) dt + T 0 ⟨h(X t (ω), t), dB t ⟩ + g(X T (ω)) × exp - T 0 f (X t (ω), t) dt - T 0 ⟨h(X t (ω), t), dB t ⟩ -g(X T (ω)) = -E a 0 (ω) exp - T 0 f (X t (ω), t) dt - T 0 ⟨h(X t (ω), t), dB t ⟩ -g(X T (ω)) .Lemma 7 (Stochastic integration by parts, [60] ). LetdX t = a t dt + b t dW 1 t , dY t = f t dt + g t dW 2 t .where a t , b t , f t , g t are continuous square integrable processes adapted to a filtration (F t ) t∈[0,T ] , and W 1 , W 2 are Brownian motions adapted to the same filtration. Then,X t Y t -X 0 Y 0 = t 0 X s dY s + t 0 Y s dX s + t 0 ⟨dX s , dY s ⟩ = t 0 X s dY s + t 0 Y s dX s + t 0 ⟨b s dW 1 s , g s dW 2 s ⟩.Remark 3 (Related work to the path-wise reparameterization trick: sensitivity analysis). As shown above, the adjoint method for SDEs is an alternative to the path-wise reparameterization trick. Prior to [54] , an array of works developed methods to compute derivatives of functionals of stochastic processes with respect to generic parameters α 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 in mind (more specifically, to compute the "Greeks"). In low dimensions, dynamic programming [3] or finite differences [26, 53] work well, but they scale poorly to high dimensions. In high dimensions, several approaches have been proposed (see the section 1 of [29] for a comprehensive although dated overview):• The path-wise method (which we refer to as adjoint method) involves taking the gradient∇ α E[f (X t )] inside of the expectation as E[∇ α f (X t )] and was first described by [82] . • The likelihood method or score method [27, 69] consists in rewriting ∇ α E[f (X T )] as E[f (X T )H],where H is a random variable which is equal to ∇ α log p(α, X T ), p(α, •) being the density of the law of X T with respect to the Lebesgue measure. [82] provide explicit weights H, under the restrictions that α appears only in 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 H in the case where H also appears in the diffusion coefficient, using Malliavin calculus.The estimator of the path-wise reparameterization trick is formally similar to the likelihood method estimator, but it is different in that α is the initial condition of the process, and does not appear either in the drift nor the diffusion coefficient.Proof. Since the equality (28) holds almost surely for the pair (X, B), it must also hold almost surely for (X v , B v ), which satisfy the same SDE. That isW(X v , 0) = V (X v 0 , 0) + 1 2 T 0 ∥u * (X v s , s)∥ 2 ds - √ λ T 0 ⟨u * (X v s , s), dB v s ⟩, Thus, we obtain that α(v, X v , B) = exp -λ -1 W(X v , 0) -λ -1/2 T 0 ⟨v(X v t , t), dB v t ⟩ + λ -1 2 T 0 ∥v(X v t , t)∥ 2 dt = exp -λ -1 V (X v 0 , 0) -λ -1 2 T 0 ∥u * (X v s , s)∥ 2 ds + λ -1/2 T 0 ⟨u * (X v s , s), dB v s ⟩ -λ -1/2 T 0 ⟨v(X v t , t), dB v t ⟩ + λ -1 2 T 0 ∥v(X v t , t)∥ 2 dt, and this is equal to exp -V (X v 0 , 0) when v = u * . Since we condition on X v 0 = x init , we have obtained that the random variable takes constant value exp -V (x init , 0) almost surely, which means that its variance is zero.Theorem 4 (Optimal reparameterization matrices). Let v be an arbitrary control in U. Define the integral operatorT t : L 2 ([t, T ]; R d×d ) → L 2 ([t, T ]; R d×d ) as [T t ( Ṁt )](s) = T t Ṁt (s ′ )E χ(s ′ , X v , B)χ(s, X v , B) ⊤ × α(v, X v , B) ds ′ , where χ(t, X v , B) := T t ∇ x f (X v s , s) ds + ∇g(X v T ) + (σ -1 t ) ⊤ (t)v(X v t , t) - T t ∇ x b(X v s , s)(σ -1 s ) ⊤ (s)v(X v t , t) ds - T t ∇ x b(X v s , s)(σ -1 s ) ⊤ (s) dB s . If we define N t (s) = -E ∇g(X v T ) + T t ∇ x f (X v s ′ , s ′ ) ds ′ χ(t, X v , B) ⊤ × α(v, X v , B) , the optimal M * = (M * t ) t∈[0,T ] is of the form M * t (s) = I + s tṀ * t (s ′ ) ds ′ , where Ṁ * t is the unique solution of the following Fredholm equation of the first kind:T t ( Ṁt ) = N t .The proof of (25) shows that minimizing Var(w; M ) is equivalent to minimizingEQUATIONTo optimize with respect to M , it is convenient to reexpress it in terms of Ṁ = ( Ṁt ) t∈[0,T ] as M t (s) = I + s t Ṁt (s ′ ) ds ′ . By Fubini's theorem, we have thatT t M t (s)∇ x f (X v s , s) ds = T t I + s t Ṁt (s ′ ) ds ′ ∇ x f (X v s , s) ds = T t ∇ x f (X v s , s) ds + T t Ṁt (s) T s ∇ x f (X v s ′ , s ′ ) ds ′ ds, - T t (M t (s)∇ x b(X v s , s) -Ṁt (s))(σ -1 ) ⊤ (s)v(X v s , s) ds = T t Ṁt (s)(σ -1 ) ⊤ (s)v(X v s , s) ds - T t Ṁt (s) T s ∇ x b(X v s ′ , s ′ )(σ -1 s ′ ) ⊤ (s ′ )v(X v s , s) ds ′ ds, -λ 1/2 T t (M t (s)∇ x b(X v s , s) -Ṁt (s))(σ -1 ) ⊤ (s) dB s = λ 1/2 T t Ṁt (s)(σ -1 ) ⊤ (s)v(X v s , s) ds - T t Ṁt (s) T s ∇ x b(X v s ′ , s ′ )(σ -1 s ′ ) ⊤ (s ′) dB s ′ ds . Hence, we can rewrite (57) asG( Ṁ ) = E 1 T T 0 σ(t) ⊤ T t ∇ x f (X v s , s) ds + ∇g(X v T ) + T t Ṁt (s) T s ∇ x f (X v s ′ , s ′ ) ds ′ + ∇g(X v T ) + (σ -1 ) ⊤ (s)v(X v s , s) - T s ∇ x b(X v s ′ , s ′ )(σ -1 s ′ ) ⊤ (s ′ )v(X v s , s) ds ′ - T s ∇ x b(X v s ′ , s ′ )(σ -1 s ′ ) ⊤ (s ′ ) dB s ′ ds 2 dt × α(v, X v , B)The first variation δG δ Ṁ ( Ṁ ) of G at Ṁ is defined as the family Q = (Q t ) t∈[0,T ] of matrix-valued functions such that for any collection of matrix-valued functions P = (P t ) t∈[0,T ] ,EQUATIONT t P t (s)χ(s, X v , B)ds ds × α(v, X v , B) Now let us reexpress equation (59) as:EQUATIONHere, equality (i) holds by Lemma 8 with the choices α(t, s) The first-order necessary condition for optimality states that at the optimal Ṁ * , the first variation δG δ Ṁ ( Ṁ * ) is zero. In other words, ∂ ϵ V( Ṁ + ϵP )| ϵ=0 is zero for any P . Hence, the right-hand side of (60) must be zero for any P , which implies that almost everywhere with respect to t ∈= P t (s)χ(X v , s, B), γ(t) = σσ ⊤ (t) ∇g(X v T ) + T t ∇ x f (X v s , s) + Ṁt (s)χ(X v ,[0, T ], s ∈ [s, T ], E ∇g(X v T )+ T t (∇ x f (X v s ′ , s ′ )+ Ṁt (s ′ )χ(X v , s ′ , B)) ds ′ χ(X v , s, B) ⊤ α(v, X v , B) = 0.To derive this, we also used that σ(t) is invertible by assumption.Define the integral operatorT t : L 2 ([t, T ]; R d×d ) → L 2 ([t, T ]; R d×d ) as [T t ( Ṁt )](s) = T t Ṁt (s ′ )E χ(X v , s ′ , B)χ(X v , s, B) ⊤ × α(v, X v , B) ds ′ If we define N t (s) = -E ∇g(X v T ) + T t ∇ x f (X v s ′ , s ′ ) ds ′ χ(X v , s, B) ⊤ × α(v, X v , B), the problem that we need to solve to find the optimal Ṁt isT t ( Ṁt ) = N t .This is a Fredholm equation of the first kind. Lemma 8. If α, β : [0, T ] × [0, T ] → R d , γ : [0, T ] → R d , δ : [0, T ] → RT 0 T t α(t, s), γ(t) ds dt (i) = T 0 T -t 0 α(t, T -s), γ(t) ds dt (ii) = T 0 t 0 α(T -t, T -s), γ(T -t) ds dt (iii) = T 0 T s α(T -t, T -s), γ(T -t) dt ds (iv) = T 0 T T -s α(T -t, s), γ(T -t) dt ds (v) = T 0 s 0 α(t, s), γ(t) dt dsHere, in equalities (i), (ii), (iv) and (v) we make changes of variables of the form t → T -t, s → T -s, s ′ → T -s ′ . In equality (iii) we use Fubini's theorem.We introduce the Gaussian warm-start, a control warm-start strategy that we 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 at all times, is the foundation of our method. Proposition 6. Given Z ∼ N (0, I) define the random process Y asEQUATION)Define the control u :EQUATIONThen, if Γ 0 = σ(0), the controlled process X u defined in equation (2) has the same marginals as Y .That is, for all t ∈ [0, T ], Law(Y t ) = Law(X u t ).Proof. Following [56] , we have that∂ t X t = ∂ t µ t + ∂ t Γ(t)Z = ∂ t µ(t) + (∂ t Γ(t)) Γ(t) -1 (X t -µ(t)), ∇ log p t (x) = -Σ(t) -1 (x -µ(t)), Σ(t) = Γ(t) Γ(t) ⊤ .Now, p t satisfies the continuity equation equationEQUATIONLet D(t) = 1 2 σ(t)σ(t) ⊤ . We want to reexpress (63) as a Fokker-Planck equation of the form∂ t p t = -∇ • (v(x, t)p t )+ d i=1 d j=1 ∂ i ∂ j (D ij (t)p t ) = -∇ • (v(x, t)p t )+ d i=1 ∂ i d j=1 (D ij (t)∂ j p t ) = -∇ • (v(x, t)p t ) + ∇ • (D(t)∇p t ) = -∇ • (v(x, t)p t ) + ∇ • (D(t)∇ log p t (x)p t ) = -∇ • ((v(x, t)-D(t)∇ log p t (x))p t ).Hence, we need thatv(x, t) -D(t)∇ log p t = ∂ t µ(t) + (∂ t Γ(t)) Γ(t) -1 (x -µ(t)), =⇒ v t (x) = ∂ t µ(t) + ((∂ t Γ(t)) Γ(t) -1 (x -µ(t)) + (σσ ⊤ )(t) 2 ∇ log p t (x) = ∂ t µ(t) + (∂ t Γ(t)) Γ(t) -1 (x -µ(t)) -(σσ ⊤ )(t) 2 Σ(t) -1 (x -µ(t)). If we let Γ(t) = Γ(t) √ t, then Σ(t) = tΓ(t)Γ(t) ⊤ = tΣ(t) and ∂ t Γ(t) = ∂ t Γ(t) √ t + Γ(t) 2 √ t . That is, v(x, t) = ∂ t µ(t) + ∂ t Γ(t) √ t + Γ(t) 2 √ t Γ(t) -1 √ t (x -µ(t)) -(σσ ⊤ )(t) 2 Σ(t) -1 t (x -µ(t)) = ∂ t µ(t) + ∂ t Γ(t) Γ(t) -1 (x -µ(t)) + 1 2t (x -µ(t)) -(σσ ⊤ )(t)Σ(t) -1 2t (x -µ(t))For v to be finite at t = 0, we need that (σσ ⊤ )(0)Σ(0) -1 = I, which holds, for example, if Γ(0) = σ(0). Also, to match the form of (2), we need thatv(x, t) = b(x, t) + σ(t)u(x, t), =⇒ u(x, t) = σ(t) -1 ∂ t µ t + ∂ t Γ(t) Γ(t) -1 + I-(σσ ⊤ )(t)Σ(t) -1 2t (x -µ t ) -b(x, t) .The warm-start control is computed as the solution of a 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 µ = (µ (b) ) B b=0 , where µ (b) ∈ R d , and a linear spline Γ = (Γ (b) ) B b=0 , where Γ (b) ∈ R d×d . These linear splines take the role of µ(t) and Σ(t) in (61) . Given splines µ and Γ, we obtain the warm-start control using (62) ; for a given t ∈ [0, T ), if we let b -= ⌊Bt/T ⌋, b + = b -+ 1, ∆ = T /B, we have thatEQUATIONEQUATIONEQUATION)Algorithm 3 provides a method to learn the splines µ, Γ. It is a stochastic optimization algorithms in which the spline parameters are updated by sampling Y t in (61) at different times, computing the control cost relying on (66), and taking its gradient. K K-1 j=0 1 2 ∥û(Yij, tj)∥ 2 + f (Yij, tj) + g(YiK )Compute the gradient of LRGSOC(µn, Γn) with respect to the spline parameters (µn, Γn).Obtain µn+1, Γn+1 with via an Adam update on µn, Γn resp. (or another stochastic algorithm) end Output: Learned splines µN , ΓN , control ûN DOUBLE WELL We also use exactly the same functions as [59] , which are the following:b(x, t) = -∇Ψ(x), Ψ(x) = d i=1 κ i (x 2 i -1) 2 , f (x) = 0, g(x) = d i=1 ν i (x 2 i -1) 2 , σ 0 = I,where d = 10, and κ i = 5, ν i = 3 for i ∈ {1, 2, 3} and κ i = 1, ν i = 1 for i ∈ {4, . . . , 10}. We set T = 1, λ = 1 and x init = 0. We take K = 200 time discretization steps, and we use random seed 0.The Double Well problem is actually highly non-trivial, and is multimodal. 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 1D problems by decoupling each dimension and apply numerical methods to solve the Hamilton-Jacobi-Bellman equation for these 1D problems. It is not a problem where we actually have the ground truth control in closed form.PATH INTEGRAL SAMPLER ON MIXTURE OF GAUSSIANS We set b(x, t) = 0, f (x, t) = 0, g(x) = log(µ 0 (x)/µ(x)) = -∥x∥ 2 2 -d 2 log(2π) -log µ(x), where T = 1, and µ is the density of a mixture of two Gaussians with means ±e 1 , where e 1 = (1, 0, . . . , 0), and variance Id. Note that we take µ to be normalized, i.e. µ(x) dx = 1, or equivalently, log Z = log µ(x) dx = 0. In Figure 1 , we use the following Monte Carlo estimator of the control objective at the control u:Ŝu (X) = T 0 1 2 ∥u(X u t , t)∥ 2 + f (X u t , t) dt + g(X u T ) + ⟨u(X u t , t), dB t ⟩.Note that this estimator is unbiased because E[ ⟨u(X u t , t), dB t ⟩] = 0. This is known as the Sticking the Landing estimator, as it has zero variance when u is the optimal control [72] . The fact that E[-Ŝu (X)] ≤ log Z = 0 with equality when u = u * is stated as [84, Thm. 4 ].As a general guideline, the control function can be thought of as the analog of the score function in diffusion models; hence, a natural choice for the architecture can be U-Nets or diffusion transformers if the control task is on images, audio or video. Other domains may require different architectures. In the experiments we report, we used the architecture implemented in the class FullyConnectedUNet within the file SOC_matching/models.py. It is a simplified version of the U-Net architecture where both the down-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 is 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 (d 2 ). 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 SOC_matching/models.py, which is essentially a three layer multilayer perceptron with ReLU activations and output dimension d 2 , whose output is averaged with the identity matrix using sigmoid weights, in order to enforce that M t (t) be the identity matrix.Table 1 shows the average times per iteration for each algorithm. Each algorithm was run using a 16GB V100 GPU. 2 ).Figure 4 shows the control objective (1) for the four settings. The error bars for the control objective plots show the confidence intervals for ± one standard deviation, computed via a Monte Carlo 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 5 shows the normalized standard deviation of the importance weight for the learned control u:Var[α(u, X u , B)]/E[α(u, X u , B)]. By Lemma 1, when X u 0 = x init for an arbitrary x init (which is the case for all our experiments), this quantity is zero for the optimal control u * . Hence, the normalized standard deviation of α is an alternative metric to measure the optimality of the learned control.Figure 6 shows an exponential moving average of the norm 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 is 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 is the algorithm with smallest error as shown in Figure 3 . Understanding this phenomenon is outside of the scope of this paper.Figure 7 shows plots of the control L 2 error, 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 7 shows that SOCM is once 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 2 error is much lower than in the other figures.Figure 8 shows the value of the training loss for SOCM and its two ablations: SOCM with constant M t = I, and SOCM-Adjoint. For all such algorithms, the training loss is the sum of the L 2 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 2 error plots is the expected conditional variance of w. We observe that the expected conditional variance in the QUADRATIC OU setting is orders of magnitude smaller for SOCM than for its two ablations. For LINEAR OU, SOCM and SOCM-adjoint 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 L 2 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 × 10 -5 is smaller than the learning rates typically used for Adam, which hover from 1 × 10 -4 to 1 × 10 -3 . one standard deviation, computed via a Monte Carlo estimate using 65536 trajectories per data point. They show the standard error of the mean. We state this in Subsec. F.3.• The answer NA means that the paper does not include experiments. • The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at 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, or 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. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 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 of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.Justification: We reviewed the NeurIPS Code of Ethics and our research conforms to 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 a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).Being progressively measurable is a strictly stronger property than the notion of being a process adapted to the filtration Ft of Bt (see[50]). 


Decision: 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

Reviews: review: paper_summary: paper_summary: **Summary**

This paper introduces Stochastic Optimal Control Matching (SOCM), a novel algorithm for solving stochastic optimal control problems. Key contributions include:

1. SOCM algorithm, adapting ideas from conditional score matching in diffusion models
2. A new "path-wise reparameterization trick" for gradient estimation
3. Theoretical analysis including a bias-variance decomposition
4. Empirical evaluation showing superior performance on 3 out of 4 benchmarks

SOCM 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.
strength_weakness: strength_weakness: 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.

While 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.

The 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.

The 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.
questions: questions: 1. How might SOCM be extended to handle nonlinear or non-Gaussian systems? Are there specific challenges you foresee in this extension?

2. 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?

3. 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?

4. The path-wise reparameterization trick is an interesting contribution. Could you elaborate on potential applications of this technique outside of stochastic optimal control?

5. 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?

6. Have you explored the performance of SOCM in settings with sparse or noisy rewards, which are common challenges in reinforcement learning?
limitations: limitations: n/a
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
questions: 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: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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.
- 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.
questions: questions: - How do all the methods compare in terms of actual task performance?
- How do these methods perform on more realistic control problems?
- Why does the proposed method not work as well on the Double Well system compared to the variance baseline?
limitations: limitations: The authors discuss limitations to scaling the approach up to more challenging problems due to the variance of the importance weight.
paper_summary: paper_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.

[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
strength_weakness: strength_weakness: - 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).
- 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.

[1] Path Integral Sampler. Zhang et al. 2022.

[2] Denoising Diffusion Sampler. Vargas et al. 2023

[3] Improving sampling via learned diffusions. Richter et al. 2023
questions: 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.
- 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 ?
- 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 ?
- Could you explain why the second Ornstein-Uhlenbeck setting is called 'hard' ?
- Could you provide the results on $L^2$ error of the control without EMA, as it is presented in [1] ?
- I am quite surprised of the relatively computational budget induced by the use of the log-variance loss, could you comment on this ?

[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: 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.

Comments: title: Answer to the rebuttal
Thanks and raise score
Rebuttal (2/2)
Rebuttal (2/2)
comment: We would like to thank the reviewer for their helpful comments. We have taken their suggestion into consideration, and present experimental results on two-mode Gaussian mixture sampling in increasing dimension, using the Path Integral Sampler [81]. Namely, we set $p_0 = \delta_{x_0}$, $b(x,t)=0$, $f(x,t)=0$, $T=1$, and $g(x)=\log (\mu^0(x)/\mu(x)) = - \|x\|^2/2 - d/2 \log (2\pi) - \log \mu(x)$, where $\mu$ is the density of a mixture of two Gaussians with means $\pm 1$ and variance $1$. 
Note that we take $\mu$ to be normalized, i.e. $\int \mu(x) dx = 1$, or equivalently, $\log Z := \log (\int \mu(x) dx) = 0$. 

For context, if we let $\hat{S}^u(X) = \int_0^T \frac{1}{2}\|u(X_t,t)\|^2 \, dt + \int_0^T \langle u(X_t,t), dB_t  \rangle + \log (\mu^0(x)/\mu(x))$, Theorem 4 of the path integral sampler paper states that $- \mathbb{E}[\hat{S}^u(X^u)] \leq \log Z = 0$, for any control $u$, and that equality holds when $u = u^*$. Hence, the quantity $\mathbb{E}[\hat{S}^u(X^u)]$, which we report in the “- ELBO” column, allows us to benchmark different SOC algorithms: the smaller the better. A perfect SOC algorithm would yield zero in this setting. We also track the regular Control Objective, which is equal in expectation to the negative ELBO, because $\mathbb{E}[ \int_0^T \langle u(X_t,t), dB_t  \rangle] = 0$ by the martingale property of stochastic integrals. 
We show the standard error for each quantity. We ran the experiments using the architectures described in the paper, for a total of 40000 iterations.

| Algorithm | Dimension | Control Objective | - ELBO |
|----------|----------|----------|----------|
| Adjoint             | 2    | 0.00282 +/- 0.00321 | 0.00515 +/- 0.00041 | 
| SOCM             | 2    | 0.00317 +/- 0.00320 | 0.00541 +/- 0.00042 |
| Cross-entropy | 2    | 0.00450 +/- 0.00320 | 0.00677 +/- 0.00046 |
| Adjoint             | 8    | 0.01560 +/- 0.00435 | 0.01157 +/- 0.00058 |
| SOCM             | 8    | 0.01495 +/- 0.00434 | 0.01104 +/- 0.00057 |
| Cross-entropy | 8    | 0.01817 +/- 0.00433 | 0.01400 +/- 0.00064 |
| Adjoint             | 16  | 0.02356 +/- 0.00548 | 0.01909 +/- 0.00075 | 
| SOCM             | 16  | 0.02242 +/- 0.00548 | 0.01802 +/- 0.00073 |
| Cross-entropy | 16  | 0.03288 +/- 0.00544 | 0.02803 +/- 0.00091 |
| Adjoint             | 32  | 0.04271 +/- 0.00726 | 0.03544 +/- 0.00102 | 
| SOCM             | 32  | 0.04013 +/- 0.00726 | 0.03287 +/- 0.00098 |
| Cross-entropy | 32  | 0.07167 +/- 0.00718 | 0.06445 +/- 0.00138 |
| Adjoint             | 64  | 0.07669 +/- 0.00991 | 0.06576 +/- 0.00141 | 
| SOCM             | 64  | 0.07143 +/- 0.00992 | 0.06150 +/- 0.00136 |
| Cross-entropy | 64  | 2.90879 +/- 0.00816 | 2.91517 +/- 0.00846 |

Cross-entropy, which uses the same importance weight $\alpha$ as SOCM, performs worse than the other two losses for all dimensions, and its results are particularly poor for dimension 64. This is because the variance of $\alpha$ 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. Note that $\alpha = \exp(...)$ where the exponent scales linearly with dimension (can be seen from Eq 20).

We observe that the -ELBO for SOCM is slightly below that of Adjoint for most dimensions, which confirms that our method is better for this range of dimensions, but if we were to keep increasing the dimension, SOCM should eventually also fail due to higher variance of $\alpha$. We are still running the higher dimensions, and will also include a case where Adjoint overtakes SOCM.
Thank you for your response.
First, I would like to thank the authors for providing precise answers to my questions and my comments, this is much appreciated. Second, I would like to re-emphasize that I acknowledge the high quality of their work, even if it is not a 'silver bullet' as they call it; I am still convinced that it can be the basis of future impactful works.

Although the authors have addressed all the points I have raised, I would like to re-insist on the first one: the extension of the SOC formulation to sampling tasks (following the PIS/DDS framework), which has become very important in the sampling community. The authors explain that they **did not test SOCM on additional realistic control problems because it often has gradient variance issues due to the variance of the importance weight**. Without being as ambitious as considering "realistic" settings, have you tried to apply your method to sample from a Gaussian mixture with two modes in increasing dimension ? Does the "variance" issue appear in this synthetic setting ?
Thank you for your response to my comments. I have read your rebuttal and am happy to raise my score.
**“Could you explain why the second Ornstein-Uhlenbeck setting is called 'hard' ?”**

We refer to the second OU setting as hard because for losses that rely on importance weights (SOCM, SOCM-Adjoint, Cross-Entropy), the importance weight $\alpha$ has a high variance due to the magnitude of the running and terminal costs being larger. When the importance weight $\alpha$ has a high variance, the signal-to-noise ratio for the gradients becomes small, making learning hard (at least initially). For the OU hard setting, SOCM still manages to outperform the other methods at advanced stages of training because as the learned control approaches the optimal control, the variance of $\alpha$ drops. We will show plots for an OU setting with even larger matrices for the costs of the problem (see answer to Reviewer gdJ8), and we will see that learning becomes impossible for SOCM.

**“Could you provide the results on  error of the control without EMA, as it is presented in [1] ?”**

We believe that showing EMA plots is more informative because their variance is lower. At later stages of training, the EMA value depends mostly on the previous 100 values (the EMA coefficient is 0.02), which means that the EMA value is very close to the actual value.

**“I am quite surprised of the relatively computational budget induced by the use of the log-variance loss, could you comment on this ?”**

In Table 1, we report that the variance loss takes 0.086 seconds per iteration, and that the log-variance loss takes 0.117 seconds per iteration. The computation for the two losses is very similar. We attribute the discrepancy to the speed fluctuation of the GPUs. We will run both loss training again and report an updated number.
**“Have you explored the performance of SOCM in settings with sparse or noisy rewards, which are common challenges in reinforcement learning?”**

In our benchmark problem Linear Ornstein Uhlenbeck, we only assume there is a terminal cost/reward and no intermediate state cost/rewards. This is a type of sparse reward (only one per episode) setting that reinforcement learning typically studies. However, we think the main difference to typical reinforcement learning applications is that control problems often assume the state costs are differentiable, which allows the use of gradient-based / adjoint methods for solving control problems, which can significantly reduce the effect of the sparse or noisy reward problem.