Stochastic Optimal Control Matching

“Could you explain why the second Ornstein-Uhlenbeck setting is called 'hard' ?”

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

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

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

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

In Table 1, we report that the variance loss takes 0.086 seconds per iteration, and that the log-variance loss takes 0.117 seconds per iteration. The computation for the two losses is very similar. We attribute the discrepancy to the speed fluctuation of the GPUs. We will run both loss training again and report an updated number.

We thank the reviewer for providing an extremely succinct description of our paper. We completely agree with the reviewer’s concerns regarding limitations, but we just want to highlight our perspective that this paper proposes a rather novel way of thinking about constructing methods for SOC problems which we believe will have a big influence on future methods that do scale well to practical applications. That is, we put together this paper as it paints a complete picture of the new proposed approach, even if it is not yet a “silver bullet” that can solve all SOC problems, and chose to further develop scalable variants in a separate follow-up work. Below, we provide detailed responses to the reviewer’s questions.

“[...] the major weakness of this paper is the lack of an additional numerical experiment, which represents a "realistic" setting [...] I am convinced that the SOCM contribution would have more impact with additional sampling numerics [...]”

We benchmark SOCM against existing methods on control problems that have been studied before in the literature (Nüsken & Richter, 2021), which allows us to clearly focus our analysis on the loss function instead of other components of the problem. These are control problems where we have access to the optimal control, which allows us to track the control $L^2$ error. Alternatively, we would only be able to compute the control objective, which as shown in Figure 3 (pg. 35), is a less sensitive metric due to numerical errors.

We also want to clarify that, as noted in the global response, the Double Well problem is far from a toy problem and is actually highly non-trivial due to its multimodal nature. We chose this problem because it is a representative challenging problem that appears in SOC interpretations of sampling problems, such as the path integral sampler. By exploiting the specific decoupling nature of the potential, we can reduce it to a 1D problem for reference solutions. We do not test SOCM on additional realistic control problems because it often has gradient variance issues due to the variance of the importance weight $\alpha$ (see Section 3 in the paper). We believe that the strength of our paper is that our framework is completely novel, and it will be the basis for the development of more algorithms that do scale well to realistic problems (ongoing work).

“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.”

In Table 1 we show the time per iteration for each loss for Quadratic OU (easy). For all losses and settings, almost all of the runtime is spent evaluating and backpropagating through the control neural network. Since the number of control computations per timestep for a given loss is the same across settings, and we are using the same neural network architecture for all settings, the time per iteration for all losses in a single setting is actually an accurate reflection of all the settings we’ve considered: the time per iteration for other settings is roughly proportional to the ratio of the number of timesteps that we are using. For SOCM, there is the additional cost of computing and backpropagating through the M function, but that also depends on the specific architecture that is used and can also be reduced by using sparse parameterizations.

“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.”

We have not tried to restrict the optimization of the reparameterization matrices to scalar or diagonal matrices, but we agree that this would allow us to trade off fast computation and low variance. We will add an experiment where we use scalar and diagonal reparameterization matrices, and we expect that it will simply be somewhere between full M and the M=I ablation result in terms of performance.

“Have you tried another parameterization of these matrices ? In particular, have you considered simple form such as $M_{\omega}(t,s) = I + \gamma(s-t) \tilde{M}_{\tilde{\omega}}(s,t)$ where $\omega = (\gamma, \tilde{\omega})$ and $\gamma(0)=0$? Otherwise, why the choice of the sigmoid ?”

We would like to clarify that in line 263, $\gamma$ is a scalar, not a one-dimensional function. Hence, arguably our form for $M_{\omega}$ is simpler than the one proposed by the reviewer. Still, the expression proposed by the reviewer makes sense and is a viable alternative.

“I find that the warm-start strategy for the optimization of the control is a very good idea [...] this strategy works in the presented settings since they are "kind of Gaussian". Do you think it may still be of interest for general SOC problems ?”

We agree with the reviewer's comment. Warm-starting the control is a strategy that makes sense when the following conditions hold simultaneously:

(i) An arbitrary initialization for control neural network causes the importance weight $\alpha$ to have high variance. If the variance of $\alpha$ is already low for an arbitrary initialization, there is no need for warm-start. This is the reason that we do not use warm-start for Quadratic OU (easy), Linear OU and Double Well. For Quadratic OU (hard), we try both no warm-start (in the main text) and warm-start (in the appendix): not warm-starting causes the learning to happen much slower for algorithms that use importance weights, although SOCM is still the best-performing loss by the end of training.

(ii) The warm-started control is close enough to the optimal control, such that the importance weight $\alpha$ has low variance. As the reviewer points out, this holds when the control problem is “kind of Gaussian” (unimodal), but it does not work when multimodality is present, as in the Double Well setting.

We thank the reviewer for their comments.

“The evaluations only consider simple toy problems. Moreover, they only plot the L2 error with respect to the optimal control. However, this does not necessarily tell us about the actual task performance due to compounding errors.” “How do all the methods compare in terms of actual task performance?”

We would appreciate it if the reviewer clarified what they mean by “compounding errors”. Regarding task performance, we would also like to note that in Figure 3 in Appendix F.3, we plot the control objective (the quantity in the right-hand side of eq. 12). In Figure 3 we see that the control objective is a much less sensitive metric than the control $L^2$ error: it is harder to benchmark the losses using the control objective. From the perspective of measures over processes, the control $L^2$ error and the control objective are two sides of the same coin: the control $L^2$ error can be seen as the KL divergence between the probability measure $\mathbb{P}^{u^*}$ of the optimally controlled process and the probability measure $\mathbb{P}^{u}$ of the process controlled by $u$. And up to a constant term, the control objective is the reversed KL divergence between the same pair of measures.

We would also like to clarify that, as noted in the global response, the Double Well problem is far from a toy problem and is actually highly non-trivial due to its multimodal nature. We chose this problem because it is a representative challenging problem that appears in SOC interpretations of sampling problems, such as the path integral sampler. By exploiting the specific decoupling nature of the potential, we can reduce it to a 1D problem for reference solutions. However, our SOCM method does not utilize this strong prior knowledge and solves it as a generic multimodal high-dimensional problem.

“On the Double Well system, there is not a clear advantage compared to the variance loss and adjoint method. However, the authors do discuss how their method appears more stable than the adjoint-based ablation. [...] Why does the proposed method not work as well on the Double Well system compared to the variance baseline?”

The Double Well (Figure 2 right-side) setting is different from other ones in that the terminal cost has 1024 modes. Hence, in order to obtain a small $L^2$ error, it is necessary to learn the control well in all or almost all the modes. Each trajectory that we sample will visit at most a few modes. Assuming for simplicity that each trajectory visits a single mode, and grouping trajectories into batches, by the end of the 80000 iterations, the “effective” number of batches per mode is 80000/1024=78.125, which is too small to get errors close to zero. Yet, this setting is interesting because it shows the behavior of the adjoint loss under multimodality: its $L^2$ error has a decreasing tendency but it wavers substantially, which is undesirable. We attribute this poor behavior of adjoint to the lack of convexity of the problem and SOCM has a clear stability advantage.

Regarding the comparison to the variance loss, note that the SOCM method converges significantly faster and the we find variance objective is a poor choice most of the time with this benchmark being the only exception. We believe SOCM initially learns very fast on this problem due to the training of M (can be seen if compared to the M=I ablation) but may find sub-optimal local minima whereas the variance method has higher variance but may have a chance of finding better local minima given enough training time.

“How do these methods perform on more realistic control problems?”

We benchmark SOCM against existing methods on control problems that have been studied before in the literature (Nüsken & Richter, 2021), which allows us to clearly focus our analysis on the loss function instead of other components of the problem.

We want to clarify that, as noted in the global response, the Double Well problem is far from a toy problem and is actually highly non-trivial due to its multimodal nature. We do not test SOCM on more realistic control problems because it often has gradient variance issues due to the variance of the importance weight $\alpha$ (see Section 3 in the paper). We believe that the strength of our paper is that our framework is completely novel, and it will be the basis for the development of more algorithms that do scale well to realistic problems (ongoing work).

We thank the reviewers for their encouraging rating and for providing us the opportunity to clarify the key importance of the proposed method below.

“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.”

Firstly, SOCM is stable by construction because its loss is convex in function space, unlike the adjoint method, which is not. In particular, we see that in the Double Well setting (Figure 2 right), the adjoint-based methods are quite unstable and have large ups-and-downs in the control error during training. Secondly, we also introduce the free parameter $M$ in our novel path-wise reparameterization gradient which can significantly reduce variance, allowing us to easily outperform related importance-weighted methods such as the cross entropy method. We also ablated both of these two claims through our ablation experiments (correspondingly, these are the “SOCM-Adjoint” which uses adjoint method instead of our path-wise reparameterization gradient, and the $M=I$ ablation, which are shown in all experimental settings to lead to worse results).

We thank the reviewer for an accurate list of our contributions; however, we think there may be a misunderstanding regarding the scope of applications. We detail our responses to the reviewer’s concerns and questions below:

“The method is currently restricted to linear Gaussian models and requires knowledge of certain model parameters. [...] The paper mentions that SOCM requires knowledge of certain model parameters. In practical scenarios where these parameters might not be known precisely, how sensitive is SOCM to parameter misspecification?”

We are unsure about this question and we think there may be a misunderstanding. We would appreciate it if the reviewer could point us to the comment that they refer to.

SOCM does not actually require any explicit assumptions on the model parameters or a linear-Gaussian model. To clarify our interpretation of this concern, the linear-Gaussian model assumption usually refers to state transitions $p(x_{t+h} | x_t, u_t)$ being Gaussian distributed with linear dependence on the state $x_t$ and independent control variables $u_t$. In the continuous-time formulation, the updates under a linear-Gaussian model would be of the form $dX_t = Ax_t + Bu_t + \sigma(t) dW_t$. However, in our formulation we use neural networks to model a control velocity field, resulting in non-linear updates. That is, the updates we have are of the form $dX_t = u_t(x_t) + \sigma(t) dW_t$, which includes the linear-Gaussian model but also generalizes to nonlinear dependencies. The only assumption SOCM actually makes is that the objective functional depends on the control velocity field only through a quadratic form (equation 1). This is required for the path-integral representation in Theorem 1 but this is a typical assumption made in nearly all stochastic optimal control methods. SOCM requires the same amount of knowledge as the existing methods: we just need to know the base SDE, the running cost and the terminal cost.

“How might SOCM be extended to handle nonlinear or non-Gaussian systems? Are there specific challenges you foresee in this extension?”

As we detail in the proof sketch of Thm. 1, SOCM relies on the path-integral representation of the optimal control (eq. 8). As far as we know, such path-integral representations hold for control systems with arbitrary base drift and linear dependency of the drift on the control, and it is a typical assumption made by works in the literature (e.g. Nüsken & Richter, 2021). If path-integral representations exist or are developed for more general control problems, we believe that our technique may be adapted to handle those.

“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?”

SOCM performs well when the importance weight $\alpha$ defined in eq. 20 has low variance, which holds when its exponent has low variance. In our experimental section, we show that when $\alpha$ has low variance, SOCM outperforms existing methods. However, when $\alpha$ has high variance, the signal-to-noise ratio of the stochastic gradient is low, and the performance of the algorithm is compromised (much like it happens for the existing cross-entropy method, which also contains the factor $\alpha$). We do not regard SOCM as the solution to solve all stochastic optimal control problems, but rather as a new perspective which can lead to the development of more algorithms (ongoing work).

“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?”

Table 1 in Section F.3 shows the time per iteration for each of the loss functions that we consider: SOCM takes 0.22 seconds, while the adjoint loss takes 0.169 seconds, and all other methods are around 0.1 seconds. The reason that SOCM takes longer is that there is an additional neural network that is trained: the M function. The second paragraph in Section F.2 describes the neural network we use for the M function. In our setup, the total cost of M evaluations per iteration is $O(d^2 K^2)$, where $d$ is the dimension of the control system and $K$ is the number of discretization steps. The dependency $O(d^2)$ stems from the fact that we parameterize the whole matrix $M$; an alternative would be to parameterize $M$ as a diagonal matrix, which would result in a cost $O(d)$ at the expense of higher gradient variance. To further reduce the variance while using more computational resources, it is also possible to choose a function $M$ that depends on the iterate $X_t$ (see Remark 2). We will provide more clarity regarding the computation cost in the paper.

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

The path-wise reparameterization trick can be used as a drop-in replacement of the adjoint method when computing gradients of functionals on ODE/SDE trajectories, either with respect to the starting point or with respect to neural network parameters. The path-wise reparameterization trick provides a natural way to enable variance reduction, thanks to the arbitrary function M. One can potentially apply this trick to any setting where neural ODEs/SDEs are used, and also to estimate scores of SDEs with arbitrary drifts.

“Have you explored the performance of SOCM in settings with sparse or noisy rewards, which are common challenges in reinforcement learning?”

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