Title: Stochastic Optimal Control Matching
Venue: NeurIPS 2024
Accepted: yes
Keywords: Stochastic Optimal Control

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

Stochastic Optimal Control Matching

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.

1 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

2.1 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 problem

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.

Cost functional and value function 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 ) ) d t + 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:

Hamilton-Jacobi-Bellman equation and optimal control If we define the infinitesimal generator L := λ 2 ∑ d i,j =1 ( σσ ⊤ ) ij ( t ) ∂ x i ∂ x j + ∑ d i =1 b i ( x, t ) ∂ x i , the value function solves the following Hamilton-Jacobi-Bellman (HJB) partial differential equation:

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

In 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 SDEs

where Y : Ω × [0 , T ] → R and Z : Ω × [0 , T ] → R d are progressively measurable 1 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 :

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

Remark this equation involves the gradient of logarithm of a conditional expectation, which is reminiscent of the vector fields that are learned when training diffusion models. For example, the target vector field for variance-exploding 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 functional

It turns out (Lemma 2 in App. B) that the Radon-Nikodym derivative d P u ∗ d P satisfies d P u ∗ d P ( X ) = exp ( λ -1 ( V ( X 0 , 0) -W ( X, 0) )) . Also, a straight-forward application of the Girsanov theorem for SDEs (Cor. 1) shows that

which means that the only control u ∈ U such that P u = P u ∗ is the optimal control itself.

2.2 Existing approaches and related work

Low-dimensional case: solving the HJB equation For low-dimensional control problems ( d ≤ 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ˇ nas 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 loss L ( 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 d P u d P u ∗ ] . Upon removing constant terms and factors, this loss is equivalent to (see Lemma 3 in App. B):

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 ∇ θ ˆ L Adj ( u θ n ) of a Monte Carlo approximation ˆ L Adj ( u θ n ) of L Adj ( u θ n ) as required by Algorithm 1, we need to backpropagate through the simulation of the trajectories, which is why

Algorithm 1 Iterative Diffusion Optimization (IDO) algorithms for stochastic optimal control

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 , loss L ∈ {L Adj (12) , L CE (13) , L Var v (16) , L log Var v (17) , L Mom v (18) }

1 for n ∈ { 0 , . . . , N -1 } do

2

Simulate m trajectories of the process X v controlled by v = u θ n , e.g., using Euler-Maruyama updates

3 if L ̸ = L Adj then detach the m trajectories from the computational graph, so that gradients do not backpropagate;

4 Using the m trajectories, compute an m -sample Monte Carlo approximation ˆ L ( u θ n ) of the loss L ( u θ n )

5 Compute the gradients ∇ θ ˆ L ( u θ n ) of ˆ L ( u θ n ) w.r.t. θ n

6 Obtain θ n +1 with via an Adam update on θ n (or another stochastic algorithm)

7

end

Output:

Learned control u θ N

we do not detach them from the computational graph. One can alternatively compute the gradient ∇ θ ˆ L Adj ( 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 d P u ∗ d P u ] . For an arbitrary v ∈ U , this loss is equivalent to the following one (see Prop. 3(i) in App. B):

The 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 ∗ :

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 ˜ L Var v ( u ) = Var P v ( d P u ∗ d P u ) and ˜ L log Var v ( u ) = Var P v (log d P u ∗ d P u ) whenever E P v | d P u ∗ d P u | < + ∞ and E P v | log d P u ∗ d P u | < + ∞ , respectively. Define

Then, ˜ L Var v and ˜ L log Var v are equivalent, respectively, to the following losses (see Lemma 4):

The variance and log-variance losses were introduced by Nüsken and Richter [59]. Unlike for the cross-entropy loss, the choice of the control v does lead to different losses. When using L Var v or L log Var v 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 as

where ˜ Y 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 ) -˜ Y 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 .

3 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. Let L SOCM : L 2 ( R d × [0 , T ]; R d ) × L 2 ([0 , T ] 2 ; R d × d ) → R be the loss function defined as

where X v is the process controlled by v (i.e., dX v t = ( b ( X v t , t ) + σ ( t ) v ( X v t , t )) d t + √ λσ ( t ) d B t and X v 0 ∼ p 0 ), and

L 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 loss

Clearly, 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):

It remains to evaluate the conditional expectation ∇ x E [ exp ( -λ -1 ∫ T t f ( X s , s ) d s -λ -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 that

We prove a more general form of this result (Prop. 4) in Subsec. C.2 and also provide an intuitive derivation in Subsec. C.3. In the proof of Prop. 4, the reparameterization matrices M 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 ,

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 :

where

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

Algorithm 2 Stochastic Optimal Control Matching (SOCM)

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 L SOCM in (19)

1 for n ∈ { 0 , . . . , N -1 } do

2 Simulate m trajectories of the process X v controlled by v = u θ n , e.g., using Euler-Maruyama updates

3 Detach the m trajectories from the computational graph, so that gradients do not backpropagate

4 Using the m trajectories, compute an m -sample Monte-Carlo approximation ˆ L SOCM ( u θ n , M ω n ) of the loss L SOCM ( u θ n , M ω n ) in (19)

5 Compute the gradients ∇ ( θ,ω ) ˆ L SOCM ( u θ n , M ω n ) of ˆ L SOCM ( u θ n , M ω n ) at ( θ n , ω n )

6 Obtain θ n +1 , ω n +1 with via an Adam update on θ n , ω n , resp.

7

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 for 1 T ∫ T 0 ∥ ∥ u ( X v t , t ) -w ( t, v, X v , B, M t ) ∥ ∥ 2 d t × α ( 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.

4 Experiments

We consider four experimental settings that we adapt from Nüsken and Richter [59]: QUADRATIC ORNSTEIN UHLENBECK (EASY), QUADRATIC ORNSTEIN UHLENBECK (HARD), LINEAR ORNSTEIN 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,

ControlObjectiveestimatevsDimension

0.10

Adjoint

SOCM

Cross-entropy

0.08

estimate

0.06

objective

Control

0.04

0.02

0.00

0

10

20

30

40

50

60

Dimension

Figure 1: This plot shows the control objective values for different algorithms (Adjoint, SOCM, and Cross-entropy) across multiple dimensions, with error bars indicating the standard deviations. The y-axis is restricted to [0, 0.1] for better visibility of the lower range values; cross-entropy takes value 2 . 915 ± 0 . 008 at d = 64 .

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.

5 Conclusion

Our work introduces Stochastic Optimal Control Matching, a novel Iterative Diffusion Optimization technique for stochastic optimal control that stems from the same philosophy as the conditional score matching loss for diffusion models. That is, the control is learned via 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.

QuadraticOrnsteinUhlenbeck,easy(d=20)

101

SOCM(ours)

SOCM(ours)

-SOCMMt=/(ablation)

-SOCMMt=/(ablation)

..SOCM-Adjoint(ablation)

108

SOCM-Adjoint (ablation)

10°

--Adjoint

--CrossEntropy

Gradient norm sqd.(EMA 0.01)

--CrossEntropy

-Adjoint

error (EMA)

.-Log-Variance

106

--Log-Variance

Moment

Moment

10-1

104

102

10-2

100

10-3

0

10000

20000

30000

40000

50000

60000

0

10000

20000

30000

40000

50000

60000

Num. iterations

Num. iterations

10′

SOCM(ours)

SOCM(ours)

-SOCM Mt=/(ablation)

-SOCM Mt=/(ablation)

:SoCM-Adjoint(ablation)

106

SOCM-Adjoint(ablation)

101

--Adjoint

-CrossEntropy

Gradient norm sqd.(EMA 0.01)

-Adjoint

-CrossEntropy

Control L2 error (EMA)

-Log-Variance

105

..-Log-Variance

-Moment

..- Moment

104

100

103

102

10-1

101

100

0

10000

20000

30000

40000

50000

60000

70000

80000

0

10000

20000

30000

40000

50000

60000

70000

80000

Num.iterations

Num. iterations

Linear Ornstein Uhlenbeck(d=10)

Double Well (d=10)

101

SOCM(ours)

SOCM(ours)

-SOCMM=/(ablation)

-SOCMM=/(ablation)

..SOCM-Adjoint(ablation)

....SOCM-Adjoint(ablation)

100

=--Adjoint

---Adjoint

=--CrossEntropy

---Cross Entropy

error(EMA)

.-Log-Variance

Control L2error (EMA)

..-Log-Variance

..Moment

....Moment

10-1

ControlL²e

100

10-2

10-3

10-4

0

10000

20000

30000

40000

50000

60000

0

10000

20000

30000

40000

50000

60000

70000

80000

Num. iterations

Num.iterations

Funding disclosure Funded by respective author affiliations.

References

Michael S. Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants, 2022.

Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08797 , 2023.

Jonathan Baxter and Peter L Bartlett. Infinite-horizon policy-gradient estimation. Journal of Artificial Intelligence Research , 15:319-350, 2001.

L'ubomír Baˇ nas, Herbert Dawid, Tsiry Avisoa Randrianasolo, Johannes Storn, and Xin Wen. Numerical approximation of a system of hamilton-jacobi-bellman equations arising in innovation dynamics. Journal of Scientific Computing , 92, 2022.

Christian Beck, Sebastian Becker, Philipp Grohs, Nor Jaafari, and Arnulf Jentzen. Solving stochastic differential equations and Kolmogorov equations by means of deep learning. arXiv:1806.00421 , 2018.

Christian Beck, Fabian Hornung, Martin Hutzenthaler, Arnulf Jentzen, and Thomas Kruse. Overcoming the curse of dimensionality in the numerical approximation of Allen-Cahn partial differential equations via truncated full-history recursive multilevel Picard approximations. arXiv:1907.06729 , 2019.

Sebastian Becker, Patrick Cheridito, and Arnulf Jentzen. Deep optimal stopping. Journal of Machine Learning Research , 20, 2019.

Andrea Belloni, Luigi Piroddi, and Maria Prandini. A stochastic optimal control solution to the energy management of a microgrid with storage and renewables. In 2016 American Control Conference (ACC) , pages 2340-2345, 2016.

Julius Berner, Lorenz Richter, and Karen Ullrich. An optimal control perspective on diffusionbased generative modeling, 2023.

Joris Bierkens and Hilbert J Kappen. Explicit solution of relative entropy weighted control. Systems & Control Letters , 72:36-43, 2014.

J. Bonnans, Elisabeth Ottenwaelter, and Hasnaa Zidani. A fast algorithm for the two dimensional hjb equation of stochastic control. M2AN. Mathematical Modelling and Numerical Analysis. ESAIM, European Series in Applied and Industrial Mathematics , 38, 07 2004.

Elisa Calzola, Elisabetta Carlini, Xavier Dupuis, and Francisco Silva. A semi-Lagrangian scheme for Hamilton-Jacobi-Bellman equations with oblique derivatives boundary conditions. Numerische Mathematik , page 153, 2022.

E. Carlini, A. Festa, and N. Forcadel. A semi-Lagrangian scheme for Hamilton-Jacobi-Bellman equations on networks. SIAM J. Numer. Anal. , 58(6):3165-3196, 2020.

René Carmona. Lectures on BSDEs, stochastic control, and stochastic differential games with financial applications , volume 1. SIAM, 2016.

René Carmona and Mathieu Laurière. Convergence analysis of machine learning algorithms for the numerical solution of mean field control and games i: The ergodic case. SIAM Journal on Numerical Analysis , 59(3):1455-1485, 2021.

René Carmona and Mathieu Laurière. Convergence analysis of machine learning algorithms for the numerical solution of mean field control and games: Ii-the finite horizon case. The Annals of Applied Probability , 32(6):4065-4105, 2022.

René Carmona, François Delarue, et al. Probabilistic Theory of Mean Field Games with Applications I-II . Springer, 2018.

Quentin Chan-Wai-Nam, Joseph Mikael, and Xavier Warin. Machine learning for semilinear PDEs. Journal of Scientific Computing , 79(3):1667-1712, 2019.

Pratik Chaudhari, Adam Oberman, Stanley Osher, Stefano Soatto, and Guillaume Carlier. Deep relaxation: partial differential equations for optimizing deep neural networks. Research in the Mathematical Sciences , 5(3):30, 2018.

Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. In Advances in Neural Information Processing Systems , volume 31. Curran Associates, Inc., 2018.

Kristian Debrabant and Espen R. Jakobsen. Semi-lagrangian schemes for linear and fully non-linear diffusion equations. Mathematics of Computation , 82(283):1433-1462, 2013.

W. E, Jiequn Han, and Arnulf Jentzen. Deep learning-based numerical methods for highdimensional parabolic partial differential equations and backward stochastic differential equations. Communications in Mathematics and Statistics , 5(4):349-380, 2017.

Weinan E, Jiequn Han, and Arnulf Jentzen. Algorithms for solving high dimensional pdes: from nonlinear monte carlo to machine learning. Nonlinearity , 35(1):278, 2021.

Jin Feng and Thomas G Kurtz. Large deviations for stochastic processes . Number 131. American Mathematical Soc., 2006.

Wendell H Fleming and Jerome L Stein. Stochastic optimal control, international finance and debt. Journal of Banking & Finance , 28(5):979-996, 2004.

Paul Glasserman and David D. Yao. Some guidelines and guarantees for common random numbers. Manage. Sci. , 38(6):884-908, jun 1992.

Peter W. Glynn. Stochastic approximation for monte carlo optimization. In Proceedings of the 18th Conference on Winter Simulation , page 356-365. Association for Computing Machinery, 1986.

Emmanuel Gobet. Monte-Carlo methods and stochastic processes: from linear to non-linear . CRC Press, 2016.

Emmanuel Gobet and Rémi Munos. Sensitivity analysis using Itô-Malliavin calculus and martingales, and application to stochastic optimal control. SIAM Journal on control and optimization , 43(5):1676-1713, 2005.

Emmanuel Gobet, Jean-Philippe Lemor, Xavier Warin, et al. A regression-based Monte Carlo method to solve backward stochastic differential equations. The Annals of Applied Probability , 15(3):2172-2202, 2005.

Vicenç Gómez, Hilbert J Kappen, Jan Peters, and Gerhard Neumann. Policy search for path integral control. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases , pages 482-497. Springer, 2014.

Alex Gorodetsky, Sertac Karaman, and Youssef Marzouk. High-dimensional stochastic optimal control using continuous tensor decompositions. International Journal of Robotics Research , 37(2-3), 3 2018.

Constantin Greif. Numerical methods for hamilton-jacobi-bellman equations. 2017.

Jiequn Han and Ruimeng Hu. Deep fictitious play for finding markovian nash equilibrium in multi-agent games. In Mathematical and scientific machine learning , pages 221-245. PMLR, 2020.

Jiequn Han, Arnulf Jentzen, and W. E. Solving high-dimensional partial differential equations using deep learning. Proceedings of the National Academy of Sciences , 115(34):8505-8510, 2018.

Carsten Hartmann and Christof Schütte. Efficient rare event simulation by optimal nonequilibrium forcing. Journal of Statistical Mechanics: Theory and Experiment , 2012(11):P11004, 2012.

Carsten Hartmann, Ralf Banisch, Marco Sarich, Tomasz Badowski, and Christof Schütte. Characterization of rare events in molecular dynamics. Entropy , 16(1):350-376, 2014.

Carsten Hartmann, Lorenz Richter, Christof Schütte, and Wei Zhang. Variational characterization of free energy: Theory and algorithms. Entropy , 19(11), 2017.

Carsten Hartmann, Omar Kebiri, Lara Neureither, and Lorenz Richter. Variational approach to rare event simulation using least-squares regression. Chaos: An Interdisciplinary Journal of Nonlinear Science , 29(6):063107, 2019.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems , volume 33. Curran Associates, Inc., 2020.

Lars Holdijk, Yuanqi Du, Ferry Hooft, Priyank Jaini, Bernd Ensing, and Max Welling. Stochastic optimal control for collective variable free sampling of molecular transition paths, 2023.

James E. Hutton and Paul I. Nelson. Interchanging the order of differentiation and stochastic integration. Stochastic Processes and their Applications , 18(2):371-377, 1984.

Martin Hutzenthaler and Thomas Kruse. Multilevel picard approximations of high-dimensional semilinear parabolic differential equations with gradient-dependent nonlinearities. SIAM Journal on Numerical Analysis , 58(2):929-961, 2020.

Martin Hutzenthaler, Arnulf Jentzen, Thomas Kruse, et al. Multilevel picard iterations for solving smooth semilinear parabolic heat equations. arXiv preprint arXiv:1607.03295 , 2016.

Martin Hutzenthaler, Arnulf Jentzen, Thomas Kruse, Tuan Anh Nguyen, and Philippe von Wurstemberger. Overcoming the curse of dimensionality in the numerical approximation of semilinear parabolic partial differential equations. arXiv:1807.01212 , 2018.

Martin Hutzenthaler, Arnulf Jentzen, and Thomas Kruse. Overcoming the curse of dimensionality in the numerical approximation of parabolic partial differential equations with gradientdependent nonlinearities. arXiv:1912.02571 , 2019.

Max Jensen and Iain Smears. On the convergence of finite element methods for hamiltonjacobi-bellman equations. SIAM Journal on Numerical Analysis , 51(1):137-162, 2013.

Hilbert J Kappen, Vicenç Gómez, and Manfred Opper. Optimal control as a graphical model inference problem. Machine learning , 87(2):159-182, 2012.

Hilbert Johan Kappen and Hans Christian Ruiz. Adaptive importance sampling for control and inference. Journal of Statistical Physics , 162(5):1244-1266, 2016.

I. Karatzas and S. Shreve. Brownian Motion and Stochastic Calculus . Graduate Texts in Mathematics (113) (Book 113). Springer New York, 1991.

Patrick Kidger, James Foster, Xuechen Li, Harald Oberhauser, and Terry Lyons. Neural sdes as infinite-dimensional gans. In International Conference on Machine Learning , 2021.

Harold Kushner and Paul G Dupuis. Numerical methods for stochastic control problems in continuous time , volume 24. Springer Science & Business Media, 2013.

Pierre L'Ecuyer and Gaétan Perron. On the convergence rates of ipa and fdc derivative estimators. Operations Research , 42(4):643-656, 1994.

Xuechen Li, Ting-Kam Leonard Wong, Ricky TQ Chen, and David Duvenaud. Scalable gradients for stochastic differential equations. In International Conference on Artificial Intelligence and Statistics , pages 3870-3882. PMLR, 2020.

Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling, 2022.

Guan-Horng Liu, Yaron Lipman, Maximilian Nickel, Brian Karrer, Evangelos A. Theodorou, and Ricky T. Q. Chen. Generalized schrödinger bridge matching, 2023.

Jingtang Ma and Jianjun Ma. Finite difference methods for the hamilton-jacobi-bellman equations arising in regime switching utility maximization. J. Sci. Comput. , 85(3):55, 2020.

Sanjoy K Mitter. Filtering and stochastic control: A historical perspective. IEEE Control Systems Magazine , 16(3):67-76, 1996.

Nikolas Nüsken and Lorenz Richter. Solving high-dimensional Hamilton-Jacobi-Bellman pdes using neural networks: perspectives from the theory of controlled diffusions and measures on path space. Partial differential equations and applications , 2:1-48, 2021.

Bernt Oksendal. Stochastic differential equations: an introduction with applications . Springer Science & Business Media, 2013.

Derek Onken, Levon Nurbekyan, Xingjian Li, Samy Wu Fung, Stanley Osher, and Lars Ruthotto. A neural network approach for high-dimensional optimal control applied to multiagent path finding. IEEE Transactions on Control Systems Technology , 31(1):235-251, jan 2023.

Grigorios A Pavliotis. Stochastic processes and applications: diffusion processes, the FokkerPlanck and Langevin equations , volume 60. Springer, 2014.

Huyên Pham. Continuous-time stochastic control and optimization with financial applications , volume 61. Springer Science & Business Media, 2009.

L.S. Pontryagin. The Mathematical Theory of Optimal Processes . Interscience Publishers, 1962.

Aram-Alexandre Pooladian, Heli Ben-Hamu, Carles Domingo-Enrich, Brandon Amos, Yaron Lipman, and Ricky T. Q. Chen. Multisample flow matching with optimal transport couplings. In International Conference on Machine Learning , 2023.

Warren B. Powell and Stephan Meisel. Tutorial on stochastic optimization in energy-part i: Modeling and policies. IEEE Transactions on Power Systems , 31(2):1459-1467, 2016.

Konrad Rawlik, Marc Toussaint, and Sethu Vijayakumar. On stochastic optimal control and reinforcement learning by approximate inference. In Twenty-Third International Joint Conference on Artificial Intelligence , 2013.

Sebastian Reich. Data assimilation: The Schrödinger perspective. Acta Numerica , 28:635-711, 2019.

Martin I. Reiman and Alan Weiss. Sensitivity analysis for simulations via likelihood ratios. Oper. Res. , 37:830-844, 1989.

Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In Proceedings of the 32nd International Conference on Machine Learning , 2015.

Lorenz Richter and Julius Berner. Improved sampling via learned diffusions. In The Twelfth International Conference on Learning Representations , 2024.

Geoffrey Roeder, Yuhuai Wu, and David K Duvenaud. Sticking the landing: Simple, lowervariance gradient estimators for variational inference. In Advances in Neural Information Processing Systems , volume 30. Curran Associates, Inc., 2017.

Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assisted Intervention-MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18 , pages 234-241. Springer, 2015.

Reuven Y Rubinstein and Dirk P Kroese. The cross-entropy method: a unified approach to combinatorial optimization, Monte-Carlo simulation and machine learning . Springer Science &Business Media, 2013.

Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. arXiv preprint arXiv:1907.05600 , 2019.

Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations (ICLR 2021) , 2021.

Evangelos Theodorou, Freek Stulp, Jonas Buchli, and Stefan Schaal. An iterative path integral stochastic optimal control approach for learning robotic tasks. IFAC Proceedings Volumes , 44 (1):11594-11601, 2011. 18th IFAC World Congress.

Ramon Van Handel. Stochastic calculus, filtering, and stochastic control. Course notes, URL http://www. prince- ton. edu/rvan/acm217/ACM217 , 2007.

Francisco Vargas, Will Sussman Grathwohl, and Arnaud Doucet. Denoising diffusion samplers. In The Eleventh International Conference on Learning Representations , 2023.

C. Villani. Topics in Optimal Transportation . Graduate studies in mathematics. American Mathematical Society, 2003.

C. Villani. Optimal Transport: Old and New . Grundlehren der mathematischen Wissenschaften. Springer Berlin Heidelberg, 2008.

Jichuan Yang and Harold J. Kushner. A monte carlo method for sensitivity analysis and parametric optimization of nonlinear stochastic systems. SIAM Journal on Control and Optimization , 29(5):1216-1249, 1991.

Jianfeng Zhang et al. A numerical scheme for BSDEs. The annals of applied probability , 14(1): 459-488, 2004.

Qinsheng Zhang and Yongxin Chen. Path integral sampler: A stochastic control approach for sampling. In International Conference on Learning Representations , 2022.

Wei Zhang, Han Wang, Carsten Hartmann, Marcus Weber, and Christof Schütte. Applications of the cross-entropy method to importance sampling and optimal control of diffusions. SIAM Journal on Scientific Computing , 36(6):A2654-A2672, 2014.

Mo Zhou, Jiequn Han, and Jianfeng Lu. Actor-critic method for high dimensional static Hamilton-Jacobi-Bellman partial differential equations based on neural networks. SIAM Journal on Scientific Computing , 43(6):A4043-A4066, 2021.

Contents

A Technical assumptions

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 by

(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 Proofs of Sec. 2

Proof of (5) By Itô's lemma, we have that

where S u t = √ λ ∫ T t ∇ V ( X u s , s ) ⊤ σ ( X u s , s ) d B s . Note that by (4),

and this implies that

Since E [ S u t | X u t = x ] = 0 , rearranging (26) and taking the conditional expectation with respect to X u t yields the final result.

Proof of (6) -(7) By Itô's lemma, we have that

Note that by (4),

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 that

Hence, recalling the definition of the work functional in (10), we have that

By Novikov's theorem (Thm. 2), we have that

which concludes the proof of (8).

Theorem 2 (Novikov's theorem) . Let θ s be a locallyH 2 process which is adapted to the natural filtration of the Brownian motion ( B t ) t ≥ 0 . Define

Moreover, 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 .

If for each t ≥ 0 , then for each t ≥ 0 ,

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 measure

Q is a probability measure because of (30) . Under the probability measure Q , the stochastic process { ˜ W ( t ) } 0 ≤ t ≤ T defined as

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 -1 i =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:

Corollary 1 (Girsanov theorem for SDEs) . If the two SDEs

admit unique strong solutions on [0 , T ] , then for any bounded continuous functional Φ on C ([0 , T ]) , we have that

where ˜ B t = B t + ∫ t 0 σ ( Y s , s ) -1 b 2 ( Y s , s ) d s . 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 SDEs

We have that

where B v t := B t + λ -1 / 2 ∫ t 0 v ( X v s , s ) d s . For the optimal control u ∗ , we have that

where the functional W is defined in (10) .

Proof. The proof of (31)-(32) follows directly from Cor. 1. To prove (34), we use that by (28),

which implies that

To prove (33), we use that since d X u ∗ t = b ( X u ∗ t , t ) d t + √ λσ ( t )d B u ∗ t , equation (35) holds if we replace X and B by X u ∗ and B u ∗ , which reads

W ( X u ∗ , 0) = V ( X u ∗ 0 , 0) + 1 2 ∫ T t ∥ u ∗ ( X u ∗ s , s ) ∥ 2 d s -√ λ ∫ T t ⟨ u ∗ ( X u ∗ s , s ) , d B v s ⟩ . Hence,

Lemma 3. The following expression holds:

Proof. To prove (36), we write

Since E P u [ log d P u d P u ∗ ] = -E P u [ log d P u ∗ d P u ] , and E P u [ ∫ T 0 ⟨ u ( X u t , t ) , d B 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:

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, ˜ L CE and L CE are equal up to constant terms and constant factors.

(ii) When p 0 is a generic probability measure, ˜ L CE and L CE have different landscapes, and L CE ( u ) = E P u ∗ [ log d P u ∗ d P 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 that

Proof. We begin with the proof of (i), and prove (37) first. Note that by the Girsanov theorem (Thm. 3),

Note that by equations (32) and (34),

where B v t := B t + λ -1 / 2 ∫ t v ( X v s , s ) d s . Also,

If we plug (41) and (42) into the right-hand side of (40), we obtain

which concludes the proof.

To show (38), we use that by Cor. 1,

Hence,

Next, we prove (ii). The first instance of V ( X v 0 , 0) in (37) can be removed without modifying the landscape of the loss. Hence, we are left with

And this can be expressed as

where

If we consider g ( u ; x ) as a loss function for u , note that it is equivalent to the loss ¯ L CE ( 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 of L CE ( u ) = E [ g ( u ; X v 0 )] .

To prove (39), note that up to a constant term, the only difference between ¯ L CE ( u ) and L CE ( u ) is the expectation is reweighted importance weight exp ( -λ -1 V ( X v 0 , 0) ) .

Lemma 4. (i) We can rewrite

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, ˜ L Var v and ˜ L log Var v are equal to L Var v and L log Var v up to a constant term and a constant factor, respectively.

(ii) When p 0 is general, ˜ L Var v and L Var v have a different landscape, and the optimum of L Var v may be different from u ∗ . A related loss that does preserve the optimum is:

In practice, this is implemented by sampling the m trajectories in one batch starting at the same point X v 0 .

(iii) Also, ˜ L log Var v and L log Var v have a different landscape, and the optimum of L log Var v may be different from u ∗ . In particular, L log Var v ( u ) = Var P v (log d P u ∗ d P u ( X v ) exp( -λ -1 V ( X v 0 , 0))) . A loss that does preserve the optimum u ∗ is

Proof. Using (34) and (31), we have that

Hence,

Since ˜ L Var v ( u ) = Var P v ( d P u ∗ d P u ) and ˜ L log Var v ( u ) = Var P v (log d P u ∗ d P 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. 3

C.1 Proof of Thm. 1 and Prop. 2

We 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 loss

Clearly, the unique optimum of ˜ L is -σ ( t ) ⊤ ∇ V . We can rewrite ˜ L as

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:

Proof. Recall the definition of W ( X,t ) in (35), which means that

Let {F t } t ∈ [0 ,T ] be the filtration generated by the Brownian motion B . Then, equation (9) implies that

We proceed as follows:

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 ) d s -λ -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 that

Plugging (23) into the right-hand side of (45), we obtain that

If we plug this into the right-hand side of (44) and complete the squared norm, we get that

where ˜ w is defined as:

We also define Φ( u ; X,B ) as

Now, by the Girsanov theorem (Thm. 3), we have that for an arbitrary control v ∈ U ,

where B v t := B t + λ -1 / 2 ∫ t 0 v ( X v s , s ) d s . 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:

Putting everything together, we obtain that

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 ) ), and

To complete the proof of equation (24), remark that ˜ L ( u ) can be rewritten as

It only remains to reexpress K . Note that by Prop. 1, we have that

Hence, using the Girsanov theorem (Thm. 3) several times, we have that

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 .

C.2 Proof of the path-wise reparameterization trick (Prop. 1)

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,

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 ) d s + λ -1 g ( X T ) . We have that

where equality (i) holds by the Leibniz rule. Using that ψ t (0 , s ) = 0 , we obtain that:

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 by

we have that

and we can write

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.

Proof of Prop. 4. Recall that

is the SDE for the uncontrolled process. For arbitrary x, z ∈ R d , we consider the following SDEs conditioned on the initial points:

Suppose that ψ : R d × [0 , T ] → R d satisfies the properties in the statement of Prop. 4. If ˜ X ( x ) is a solution of

then X ( x + z ) = ˜ X ( x ) + ψ ( z, · ) is a solution of (50). This is because X ( x + z ) 0 = ˜ X ( x ) 0 + ψ ( z, 0) = ˜ X ( x ) 0 + z = x + z , and

Note that we may rewrite (51) as

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 Φ ,

We can write

Equality (i) holds by the definition of X ( x + z ) , equality (ii) holds by the fact X ( 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,

In equality (i) we used (53), and that:

by the Leibniz rule,

and by the Leibniz rule for stochastic integrals (see [42]),

C.3 Informal derivation of the path-wise reparameterization trick

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 ) d s + λ -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 = ( ˆ X k ) k =0: K defined as

Note that we can approximate

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 } , ˆ X k +1 | ˆ X k ∼ N ( ˆ X k + δb ( ˆ X k , kδ ) , δλ ( σσ ⊤ )( kδ )) . Hence,

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 write

In the last equality, we used that for k ∈ { 1 , . . . , K } , the variables ˆ x k 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 , . . . , ˆ x K ) , and letting ˆ x 0 = x , we define

Using that ψ (0 , s ) = 0 for all s ∈ [0 , T ] , we have that:

And we can express the right-hand side of (54) in terms of G ˆ x (0) and ∇ G ˆ x ( z ) | z =0 :

We define ϵ k = 1 √ δλ σ -1 ( kδ )(ˆ x k +1 -ˆ x k -δb (ˆ x k , kδ )) , and then, we are able to write

Then, 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 ) d B s . Thus,

which concludes the derivation.

C.4 SOCM-Adjoint: replacing the path-wise reparameterization trick with the adjoint method

Proposition 5. Let L SOCM -Adj : L 2 ( R d × [0 , T ]; R d ) → R be the loss function defined as

where X v is the process controlled by v (i.e., dX t = ( b ( X t , t ) + σ ( t ) v ( X t , t )) d t + λσ ( X t , t ) d B 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 ODE

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 ) ,

we have that

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 write

By Lemma 7, we have that

In the last line we used that ∇ x X 0 ( ω ) = ∇ x x = I . If choose a such that

then we obtain that

and by the Leibniz rule,

and

Lemma 7 (Stochastic integration by parts, [60]) . Let

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,

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.

C.5 Proof of Lemma 1

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 is

Thus, we obtain that

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.

D Optimal reparameterization matrices

Theorem 4 (Optimal reparameterization matrices) . Let v be an arbitrary control in U . Define the integral operator T t : L 2 ([ t, T ]; R d × d ) → L 2 ([ t, T ]; R d × d ) as

where

If we define N t ( s ) = -E [( ∇ g ( X v T ) + ∫ T t ∇ x f ( X v s ′ , s ′ ) d s ′ ) χ ( 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 ˙ M ∗ t ( s ′ ) d s ′ , where ˙ M ∗ t is the unique solution of the following Fredholm equation of the first kind:

The proof of (25) shows that minimizing Var( w ; M ) is equivalent to minimizing

To optimize with respect to M , it is convenient to reexpress it in terms of ˙ M = ( ˙ M t ) t ∈ [0 ,T ] as M t ( s ) = I + ∫ s t ˙ M t ( s ′ ) d s ′ . By Fubini's theorem, we have that

Hence, we can rewrite (57) as

The first variation δ G δ ˙ M ( ˙ M ) of G at ˙ M 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 ] ,

If we define

we can rewrite (58) as

Now let us reexpress equation (59) as:

Here, equality (i) holds by Lemma 8 with the choices α ( t, s ) = P t ( s ) χ ( X v , s, B ) , γ ( t ) = σσ ⊤ ( t ) ( ∇ g ( X v T ) + ∫ T t ( ∇ x f ( X v s , s ) + ˙ M t ( s ) χ ( X v , s, B ) ) d s . Equality (ii) follows from the fact that for any matrix A and vectors b, c , ⟨ Ab, c ⟩ = c ⊤ Ab = Tr( c ⊤ Ab ) = Tr( Abc ⊤ ) = ⟨ B,cb ⊤ ⟩ F , where ⟨· , ·⟩ F denotes the Frobenius inner product.

The first-order necessary condition for optimality states that at the optimal ˙ M ∗ , the first variation δ G δ ˙ M ( ˙ M ∗ ) is zero. In other words, ∂ ϵ V ( ˙ M + ϵ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 ∈ [0 , T ] , s ∈ [ s, T ] ,

To derive this, we also used that σ ( t ) is invertible by assumption.

Define the integral operator T t : L 2 ([ t, T ]; R d × d ) → L 2 ([ t, T ]; R d × d ) as

If we define N t ( s ) = -E [( ∇ g ( X v T ) + ∫ T t ∇ x f ( X v s ′ , s ′ ) d s ′ ) χ ( X v , s, B ) ⊤ × α ( v, X v , B ) ] , the problem that we need to solve to find the optimal ˙ M t is

This is a Fredholm equation of the first kind.

Lemma 8. If α, β : [0 , T ] × [0 , T ] → R d , γ : [0 , T ] → R d , δ : [0 , T ] → R d × d are arbitrary integrable functions, we have that

Proof. We have that:

Here, 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.

E Control warm-starting

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 as

Define the control u : R d × [0 , T ] → R d as

Then, 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

Now, p t satisfies the continuity equation equation

10

Let D ( t ) = 1 2 σ ( t ) σ ( t ) ⊤ . We want to reexpress (63) as a Fokker-Planck equation of the form

Hence, we need that

If we let ˜ Γ( t ) = Γ( t ) √ t , then ˜ Σ( t ) = t Γ( t )Γ( t ) ⊤ = t Σ( t ) and ∂ t ˜ Γ( t ) = ∂ t Γ( t ) √ t + Γ( t ) 2 √ t . That is,

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 that

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 -= ⌊B t/T ⌋ , b + = b -+1 , ∆ = T/ B , we have that

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.

Algorithm 3 Restricted Gaussian Stochastic Optimal Control

Output: Learned splines µ N , Γ N , control ˆ u N

Once we have access to the restricted control ˆ u N , we can warm-start the control in Algorithms 1 and 2 by introducing ˆ u N as an offset. That is, we parameterize the control as u θ = ˆ u N + ˜ u θ .

F Experimental details and additional plots

F.1 Experimental details

The control L 2 error curves show the following quantity:

That is, we sample trajectories using the optimal control, and compute the error using a Monte Carlo estimate. In all our experiments, the distribution X u ∗ 0 is a delta, which means that we do not need to compute V ( X u ∗ 0 , 0) . We keep an exponential moving average (EMA) estimate of the control L 2 error, which we show in the plots. To compute it, we sample ten batches of optimally controlled trajectories every 10 training iterations, and we update the quantity with the average of the ten batches, using EMA coefficient 0.02. All other quantities shown in the plots are also smoothed out using EMA with coefficient 0.01, except for control objective values, which are computed as the average of 65536 samples, every 5000 training steps.

For all losses and all settings, we train the control using Adam with learning rate 1 × 10 -4 . For SOCM, we train the reparametrization matrices using Adam with learning rate 1 × 10 -2 . We use batch size m = 128 unless otherwise specified. When used, we run the warm-start algorithm (Algorithm 3) with B = 20 knots, K = 200 time steps, and batch size m = 512 , and we use Adam with learning rate 3 × 10 -4 for N = 60000 iterations.

QUADRATIC ORNSTEIN-UHLENBECK The choices for the functions of the control problem are:

where Q is a positive definite matrix. Control problems of this form are better known as linear quadratic regulator (LQR) and they admit a closed form solution [78, Thm. 6.5.1]. The optimal control is given by:

where F t is the solution of the Ricatti equation

with the final condition F T = Q . Within the QUADRATIC OU class, we consider two settings:

Easy: We set d = 20 , A = 0 . 2 I , P = 0 . 2 I , Q = 0 . 1 I , σ 0 = I , λ = 1 , T = 1 , x init = 0 . 5 N (0 , I ) . We do not use warm-start for any algorithm. We take K = 50 time discretization steps, and we use random seed 0.

Hard: We set d = 20 , A = I , P = I , Q = 0 . 5 I , σ 0 = I , λ = 1 , T = 1 , x init = 0 . 5 N (0 , I ) . We use the Gaussian warm-start (App. E). We take batch size m = 64 and K = 150 time discretization steps, and we use random seed 0.

LINEAR ORNSTEIN-UHLENBECK The functions of the control problem are chosen as follows:

The optimal control for this class of problems is given by [59, Sec. A.4]:

We use exactly the same functions as [59]: we sample ( ξ ij ) 1 ≤ i,j ≤ d once at the beginning of the simulation, and set:

We take K = 100 time discretization steps, and we use random seed 0.

DOUBLE WELL We also use exactly the same functions as [59], which are the following:

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

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 ) d x = 1 , or equivalently, log Z = log ( ∫ µ ( x ) d x ) = 0 . In Figure 1, we use the following Monte Carlo estimator of the control objective at the control u :

Note that this estimator is unbiased because E [ ∫ ⟨ u ( X u t , t ) , d B 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 [ -ˆ S u ( X )] ≤ log Z = 0 with equality when u = u ∗ is stated as [84, Thm. 4].

F.2 Model architectures

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.

F.3 Additional tables and plots

Table 1 shows the average times per iteration for each algorithm. Each algorithm was run using a 16GB V100 GPU.

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 .

QuadraticOrnsteinUhlenbeck,easy (d=20)

SOCM(ours)

160

SOCM(ours)

5.96

5.94

...SOCM-Adjoint(ablation)

---Adjoint

--Cross Entropy

-SOCM Mt=/(ablation)

140

---Cross Entropy

--Adjoint

..SOCM-Adjoint(ablation)

-SOCMMt=/(ablation)

Control objective

.-Log-Variance

-Moment

Control objective

120

.--Log-Variance

·Moment

5.92

100

5.90

80

5.88

60

5.86

40-

5.84

10000

20000

30000

40000

50000

60000

10000

20000

30000

40000

50000

60000

70000

80000

Num.iterations

Num.iterations

Linear Ornstein Uhlenbeck(d=10)

Double Well (d =10)

2.82

SOCM(ours)

SOCM(ours)

-SOCM Mt=/(ablation)

3.0-

-SOCM M=/(ablation)

-2.83

.. SOCM-Adjoint(ablation)

...SOCM-Adjoint(ablation)

2.84

---Cross Entropy

--Adjoint

2.9

---Cross Entropy

---Adjoint

.-Log-Variance

-.-Log-Variance

Control objective

2.85

Moment

Control objective

2.8

... Moment

2.7

-2.86

2.6

-2.87

2.5

-2.88

2.4

2.89

2.3

-2.90

10000

20000

30000

40000

50000

60000

10000

20000

30000

40000

50000

60000

70000

80000

Num. iterations

Num. iterations

QuadraticOrnsteinUhlenbeck,easy(d=20)

Importance weight std.dev.(normalized)

. SOCM-Adjoint(ablation)

SOCM(ours)

-SOCM Mt=/(ablation)

Importance weight std.dev.(normalized)

101

.. SOCM-Adjoint(ablation)

-SOCM Mt=/(ablation)

SOCM(ours)

---Cross Entropy

--Adjoint

-Adjoint

--CrossEntropy

-Log-Variance

-Log-Variance

Moment

-Moment

-Variance

100

100

0

10000

20000

30000

40000

50000

60000

0

10000

20000

30000

40000

50000

60000

70000

80000

Num.iterations

Num.iterations

Linear Ornstein Uhlenbeck(d=10)

Double Well (d =10)

SOCM(ours)

SOCM(ours)

SOCM Mt=/(ablation)

3 × 100

-SOCM M=/(ablation)

Importance weight std dev. (normalized)

SOCM-Adjoint (ablation)

Importance weight std.dev.(normalized)

.. SOCM-Adjoint(ablation)

--Adjoint

---Adjoint

100

--CrossEntropy

--CrossEntropy

-Log-Variance

2 ×100

--Log-Variance

Moment

Moment

10-1

100

10-2

6×10-1

0

10000

20000

30000

40000

50000

60000

0

0000000009000000000000000000700000

80000

Num. iterations

Num. iterations

Linear Ornstein Uhlenbeck(d=10)

Double Well (d=10)

107

SOCM(ours)

SOCM(ours)

-SOCMMt=/(ablation)

107

-SOCMM=/(ablation)

...SOCM-Adjoint(ablation)

SOCM-Adjoint (ablation)

Gradient norm sqd. (EMA 0.01)

105

---Adjoint

---Cross Entropy

Gradient norm sqd. (EMA 0.01)

106

---Adjoint

---Cross Entropy

.-Log-Variance

..-Log-Variance

...Moment

105

Moment

103

104

101

103

10-1

102

中

10-3

101

HHH

0

10000

20000

30000

40000

50000

60000

0

10000

20000

30000

40000

50000

60000

70000

80000

Num. iterations

Num.iterations

Quadratic OrnsteinUhlenbeck,hard,warm start(d=20)

Quadratic OrnsteinUhlenbeck,hard,warm start(d=20)

SOCM(ours)

104

-SOCMMt=/(ablation)

4 ×10-2

...SOCM-Adjoint(ablation)

103

error (EMA 0.01)

--Adjoint

-CrossEntropy

Gradient norm sqd. (EMA 0.01)

3×10-2

.-Log-Variance

102

Moment

-Variance

101

2×10-2

100

Control L²

SOCM(ours)

SOCMMt=/(ablation)

10-1

SOCM-Adjoint(ablation)

-Adjoint

10-2

-Cross Entropy

10-2

-Log-Variance

10-3

Moment

0

10000

20000

30000

40000

50000

60000

0

10000

20000

30000

40000

50000

60000

Num.iterations

Num.iterations

28.78

SOCM(ours)

SOCM(ours)

-SOCM Mt=/(ablation)

-SOCMM=/(ablation)

...SOCM-Adjoint(ablation)

Importance weight std.dev.(normalized)

6.6×10-1

.. SOCM-Adjoint(ablation)

28.76 -

---Adjoint

--Adjoint

--CrossEntropy

--CrossEntropy

28.74

-.-Log-Variance

6.4 ×10-1

..-Log-Variance

Control objective

.....Moment

... Moment

28.72

6.2 ×10-1

28.70

28.68

6×10-1

28.66

5.8×10-1

1111

28.64

10000

20000

30000

40000

50000

60000

0

10000

20000

30000

40000

50000

60000

Num. iterations

Num. iterations

QuadraticOrnsteinUhlenbeck,easy(d=20)

107

SOCM(ours)

SOCM(ours)

-SOCMMt=/(ablation)

-SOCM M=/(ablation)

101

SOCM-Adjoint(ablation)

105

·SOCM-Adjoint (ablation)

Training loss (EMA 0.01)

Training loss (EMA 0.01)

100

103

101

10-1

10-1

10-2

10-3

10-3

0

10000

20000

30000

40000

50000

60000

0

10000

20000300004000050000

60000

70000

80000

Num.iterations

Num. iterations

Linear OrnsteinUhlenbeck(d=10)

Double Well (d =10)

1010

SOCM(ours)

101

-SOCM M=/(ablation)

SOCM-Adjoint (ablation)

100

108

Training loss (EMA 0.01)

Training loss (EMA 0.01)

··.....................

10-1

SOCM(ours)

106.

10-2

-SOCM Mt=/(ablation)

SOCM-Adjoint (ablation)

10-3

104

10-4

102

10-5

0

10000

20000

30000

40000

50000

60000

0

10000

20000

30000

40000

50000

60000

70000

80000

Num.iterations

Num. iterations

AdjointLR=1e-4

AdjointLR=1e-4

.-AdjointLR=3e-5

108

--Adjoint LR=3e-5

107

Gradient norm sqd. (EMA 0.01)

Control L2 error (EMA 0.01)

106

105

100

104

103

102

101

0

5000

10000

150002000025000

30000

35000

40000

0

5000

10000

150002000025000

30000

35000

40000

Num. iterations

Num. iterations

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: In the abstract and introduction we state that we introduce SOCM, a novel algorithm to solve stochastic optimal control problems, and we claim that it outperforms all existing algorithms in three out of four settings. All of these claims are supported by our paper.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: Yes, in the second-to-last paragraph of Sec. 5 we discuss the main limitation of SOCM (our algorithm), which is that the variance of the importance weight α is too large in certain settings. Regarding the computational efficiency of our algorithm, Table 1 shows a comparison between our algorithm and existing ones.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: All the theoretical results stated in the paper are numbered and have a corresponding detailed proof in the appendices. In the main text, we indicate the precise location of the proofs. We provide a proof sketch for Theorem 1 in the main paper, and a full proof in the appendix.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We provide a description of experimental details in Subsec. F.1, and of model architectures in Subsec. F.2. We also provide a link to the GitHub repository that contains the code for this paper.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We also provide a link to the GitHub repository that contains the code for this paper.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We provide a description of experimental details in Subsec. F.1, and of model architectures in Subsec. F.2.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: Most of the plots in the paper do not contain error bars, because the lines show exponential moving averages, which means that the error is already smoothed out and very small. The error bars for the control objective plots do show the confidence intervals for ±

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.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, 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.

8. 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).

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

Answer: [Yes]

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

10. 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).

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

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: We present code implementing our algorithm. The code folder contains a README which explains how to reproduce the experiments. We describe our algorithm throughout the paper, and provide details in App. F.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [No]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [No]

Justification: The paper does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, the authors propose a novel learning algorithm, Stochastic Optimal Control Matching (SOCM), to numerically solve general formulations of Stochastic Optimal Control (SOC) problems involving affined-controlled diffusion processes. They build upon the Iterative Diffusion Optimization (IDO) [1] framework, which consists in iteratively refining a parametric controlled diffusion process by minimizing at each iteration (with stochastic gradient descent) a specific objective function with respect to the parametric control. Previous works had for instance considered the relative entropy loss, the cross-entropy loss, the log-variance loss or the moment loss as the objective function. In SOCM, it is a least-squares regression loss (with multiplicative importance weight) which aims at fitting the parametric control to a vector field which depends on a family of *reparameterization matrices* (also optimized). The design of this objective function relies on standard tools from SOC theory, as well as an original contribution, the *path-wise reparameterization trick*, to compute gradients of conditional expectations of a functional applied on a random process. The authors show that the SOCM loss can be decomposed as the sum of a bias term, that is linked to the cross-entropy loss, and variance term, which is only affected by the *reparameterization matrices*. Hence, these extra parameters can be seen as a way to reduce the variance of the SOCM loss, which motivates their introduction. Moreover, this loss has the computational advantage to avoid computing gradients of the control along the path measure (which is the main drawback of the relative-entropy loss). Finally, the authors conduct numerical experiments to compare their approach to existing designs of IDO losses. They consider four different settings ($d\\in \\{10,20\\}$) with access to the ground-truth control, which allows them to compute the $L^2$ error on the control. Using this metric, their results indicate better performance on most of the settings while maintaining a certain training stability.\n\n[1] Solving high-dimensional Hamilton–Jacobi–Bellman pdes using neural networks: perspectives from the theory of controlled diffusions and measures on path space. Nüsken et al. 2021\n\nSTRENGTHS:\n- The paper is very well-written, it is a pleasure to read it. In particular, the authors pay attention to define their notation, introduce with clarity the SOC framework, state clear mathematical statements, recall (and prove) standard results from SOC theory, provide intuition on theoretical results (with details on the proof and meaningful comments). This is really good work.\n- The relation to prior work is clearly well established: in particular, the comparison between the SOCM loss and the previous IDO losses (presented in Section 2.2) is well highlighted. \n- This papers introduces an interesting contribution that may be applied beyond this framework (in particular, in the generative community where terms involving gradients of expectations often appear) : this is the path-wise reparameterization trick, which is proved to be decisive in the numerics.\n\nWEAKNESSES:\n- In my opinion, the major weakness of this paper is the lack of an additional numerical experiment, which represents a \"realistic\" setting (for instance, where the expression of the control is not known). For instance (as mentioned by the authors), a significant line of recent research has considered the sampling problem via a SOC perspective, see for example [1,2,3]. I am convinced that the SOCM contribution would have more impact with additional sampling numerics comparing SOCM, relative entropy [1,2] and log-variance [3] losses, for challenging distributions (namely, multi-modal distributions in relatively high dimension). In this case, the quality of the methods would be assessed with sampling metrics. This weakness explains my current score (although I really appreciate the paper).\n- I find that the complexity/accuracy tradeoff of the IDO losses (including SOCM) is not well highlighted to me. The table provided by the authors only considers one setting. To have a full picture, it should be given for all settings.\n\n[1] Path Integral Sampler. Zhang et al. 2022.\n\n[2] Denoising Diffusion Sampler. Vargas et al. 2023\n\n[3] Improving sampling via learned diffusions. Richter et al. 2023\n\nQUESTIONS:\n- Have you tried to restrict the optimization of the reparameterization matrices to scalar matrices ? I have the feeling that this choice may align a low computational budget with a good expressivity.\n- Have you tried another parameterization of these matrices ? In particular, have you considered simple form such as $M_{w}(t,s)=I_d + \\gamma(s-t)\\tilde{M}_{\\tilde{w}}(s,t)$ where $w=(\\gamma, \\tilde{w})$ and $\\gamma(0)=0$ ? Otherwise, why the choice of the sigmoid ?\n- I find that the warm-start strategy for the optimization of the control is a very good idea, as it benefits from the tractability of the stochastic interpolants with the lightness of the spline formulation. However, I have the feeling that this strategy works in the presented settings since they are \"kind of Gaussian\". Do you think it may still be of interest for general SOC problems ?\n- Could you explain why the second Ornstein-Uhlenbeck setting is called 'hard' ?\n- Could you provide the results on $L^2$ error of the control without EMA, as it is presented in [1] ?\n- I am quite surprised of the relatively computational budget induced by the use of the log-variance loss, could you comment on this ?\n\n[1] Solving high-dimensional Hamilton–Jacobi–Bellman pdes using neural networks: perspectives from the theory of controlled diffusions and measures on path space. Nüsken et al. 2021\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- The proposed objective function acts as a form of variance reduction for the cross-entropy loss when solving stochastic optimal control problems.\n- The novel reparameterization trick for estimating gradients of diffusion processes with respect to its initial condition may be more broadly applicable.\n- On toy problems, their method appears to generally outperform other approaches in solving the HJB equation for the optimal controls.\n- The paper is well organized and overall written well. It provides a thorough related work section and does a good job explaining the novelty and results.\n\nWEAKNESSES:\n- The evaluations only consider simple toy problems. Moreover, they only plot the L2 error with respect to the optimal control. However, this does not necessarily tell us about the actual task performance due to compounding errors.\n- On the Double Well system, there is not a clear advantage compared to the variance loss and adjoint method. However, the authors do discuss how their method appears more stable than the adjoint-based ablation.\n\nQUESTIONS:\n- How do all the methods compare in terms of actual task performance?\n- How do these methods perform on more realistic control problems?\n- Why does the proposed method not work as well on the Double Well system compared to the variance baseline?\n\nLIMITATIONS:\nThe authors discuss limitations to scaling the approach up to more challenging problems due to the variance of the importance weight.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-11"}

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

STRENGTHS:
The work is nicely motivated in Introduction, showing the drawbacks of traditional works.  The proposed control method is supported by the uniqueness analysis of the control logic (Theorem 1) and the sophisticated design methods (Propositions 1 and 2). In the reviewer's understanding, they are technically correct.

WEAKNESSES:
As stated in Algorithm 2 below, reducing noise in the gradient is crucial for the presented algorithm.  This weakness is addressed by Lemma 1 and extensions.

QUESTIONS:
As stated in Introduction, the work is motivated by stabilizing the unstable training of conventional IDO, which comes from the non-convexity of the loss.  Could the authors comment and/or perform some motivating experiments to show the stability of the training by SOCM?  They can emphasize the contribution of this paper.

LIMITATIONS:
No limitations in this work.  This work is devoted to the theoretical analysis of control system design, and it does not directly bring a negative social impact.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n**Summary**\n\nThis paper introduces Stochastic Optimal Control Matching (SOCM), a novel algorithm for solving stochastic optimal control problems. Key contributions include:\n\n1. SOCM algorithm, adapting ideas from conditional score matching in diffusion models\n2. A new \"path-wise reparameterization trick\" for gradient estimation\n3. Theoretical analysis including a bias-variance decomposition\n4. Empirical evaluation showing superior performance on 3 out of 4 benchmarks\n\nSOCM learns a control function by fitting a matching vector field via least squares, while optimizing reparameterization matrices to reduce variance. The method is currently limited to linear Gaussian models and requires knowledge of certain parameters. Experiments demonstrate SOCM's effectiveness on theoretical benchmarks, outperforming existing methods in most cases. The paper provides a solid theoretical foundation but lacks exploration of real-world applications or non-linear systems.\n\nSTRENGTHS:\nThe paper introduces Stochastic Optimal Control Matching (SOCM), a novel algorithm for solving stochastic optimal control problems. Its originality lies in adapting ideas from conditional score matching in diffusion models to the domain of optimal control. This creative combination represents an interesting cross-pollination between two active areas of research.\n\nThe quality of the theoretical work is notable. The authors provide a comprehensive mathematical foundation for their method, including detailed proofs and a novel \"path-wise reparameterization trick\". This theoretical rigor is a significant strength of the paper.\n\nIn terms of clarity, the paper is well-structured and clearly written. The authors effectively guide the reader from the problem formulation through the theoretical development to the empirical results. The use of illustrative examples and detailed appendices aids in understanding the complex mathematical concepts presented.\n\nThe significance of this work lies in its potential to improve the efficiency of solving stochastic optimal control problems. The empirical results, showing improved performance over existing methods on multiple benchmarks, underscore the practical impact of this approach. However, the significance is somewhat limited by the current restrictions to linear Gaussian models.\n\nWEAKNESSES:\nThe primary weakness of this paper is its limited scope and applicability. The method is currently restricted to linear Gaussian models and requires knowledge of certain model parameters. This significantly narrows its potential impact on the broader field of stochastic optimal control. The authors should discuss potential approaches to extend SOCM to more general settings, such as nonlinear or non-Gaussian systems.\n\nWhile the empirical results are promising, they are limited to theoretical benchmarks. The paper would be strengthened by including experiments on real-world problems or more complex simulated environments. This would help demonstrate the method's practical utility and potential for broader impact.\n\nThe scalability of the method is not thoroughly addressed. As the dimensionality of the problem increases, how does the computational complexity of SOCM compare to existing methods? A more detailed analysis of computational requirements and scaling properties would be valuable.\n\nThe comparison with existing methods, while showing SOCM's superior performance, could be more comprehensive. Including comparisons with the most recent state-of-the-art methods would provide a clearer picture of SOCM's relative performance in the current landscape of stochastic optimal control algorithms.\n\nQUESTIONS:\n1. How might SOCM be extended to handle nonlinear or non-Gaussian systems? Are there specific challenges you foresee in this extension?\n\n2. The paper focuses on theoretical benchmarks. Have you considered applying SOCM to any real-world stochastic optimal control problems? If so, what challenges did you encounter or do you anticipate?\n\n3. How does the computational complexity of SOCM scale with the dimensionality of the problem? Could you provide a more detailed comparison of computational requirements with existing methods?\n\n4. The path-wise reparameterization trick is an interesting contribution. Could you elaborate on potential applications of this technique outside of stochastic optimal control?\n\n5. The paper mentions that SOCM requires knowledge of certain model parameters. In practical scenarios where these parameters might not be known precisely, how sensitive is SOCM to parameter misspecification?\n\n6. Have you explored the performance of SOCM in settings with sparse or noisy rewards, which are common challenges in reinforcement learning?\n\nLIMITATIONS:\nn/a", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}