Initial Value Problem Enhanced Sampling for Closed-Loop Optimal Control Design with Deep Neural Networks

Optimal control aims to find a control for a dynamical system over a period of time such that a specified loss function is minimized. Generally speaking, there are two types of optimal controls: open-loop optimal control and closed-loop (feedback) optimal control. Open-loop optimal control deals with the problem with a given initial state, and its solution is a function of time for the specific initial data, independent of the other states of the system. In contrast, closed-loop optimal control aims to find the optimal control policy as a function of the state that gives us optimal control for general initial states.

By the nature of the problem, solving the open-loop control problem is relatively easy and various open-loop control solvers can handle nonlinear problems even when the state lives in high dimensions (Betts, 1998; Rao, 2009). Closed-loop control is much more powerful than open-loop control since it can cope with different initial states, and it is more robust to the disturbance of dynamics. The classical approach to obtaining a closed-loop optimal control function is by solving the associated Hamilton-Jacobi-Bellman (HJB) equation. However, traditional numerical algorithms for HJB equations such as the finite difference method or finite element method face the curse of dimensionality (Bellman, 1957) and hence can not deal with high-dimensional problems.

Since the work Han & E (2016) for stochastic optimal control problems, there have been growing interest on making use of the capacity of neural networks (NNs) in approximating high-dimensional functions to solve the closed-loop optimal control problems (Nakamura-Zimmerer et al., 2021a; b; 2020; Böttcher et al., 2022; E et al., 2022). Generally speaking, there are two categories of methods in this promising direction. One is policy search approach (Han & E, 2016; Ainsworth et al., 2021; Böttcher et al., 2022; Zhao et al., 2022), which directly parameterizes the policy function by neural networks, computes the total cost with various initial points, and minimizes the average total cost. When solving problems with a long time span and high nonlinearity, the corresponding optimization problems can be extremely hard and may get stuck in local minima (Levine & Koltun, 2014). The other category of methods is based on supervised learning (Nakamura-Zimmerer et al., 2021a; b; 2020; 2022). Combining various techniques for open-loop control, one can solve complex high-dimensional open-loop optimal control problems; see Betts (1998); Rao (2009); Kang et al. (2021) for detailed surveys. Consequently, we can collect optimal trajectories for different initial points as training data, parameterize the control function (or value function) using NNs, and train the NN models to fit the closed-loop optimal controls (or optimal values). This work focuses on the second approach and aims to improve its performance through adaptive sampling.

As demonstrated in Nakamura-Zimmerer et al. (2021b); Zang et al. (2022), NN controllers trained by the vanilla supervised-learning-based approach can perform poorly even when both the training error and test error on collected datasets are fairly small. Some existing works attribute this phenomenon to the fact that the learned controller may deteriorate badly at some difficult initial states even though the error is small in the average sense. Several adaptive sampling methods regarding the initial points are hence proposed (see Section 4 for a detailed discussion). However, these methods all focus on choosing optimal paths according to different initial points and ignore the effect of dynamics. This is an issue since the paths controlled by the NN will deviate from the optimal paths further and further over time due to the accumulation of errors. As shown in Section 6, applying adaptive sampling only on initial points is insufficient to solve challenging problems.

This work is concerned with the so-called distribution mismatch phenomenon brought by the dynamics in the supervised-learning-based approach. This phenomenon refers to the fact that the discrepancy between the state distribution of the training data and the state distribution generated by the NN controller typically increases over time and the training data fails to represent the states encountered when the trained NN controller is used. Such phenomenon has also been identified in reinforcement learning (Kakade & Langford, 2002; Long & Han, 2022) and imitation learning (Ross & Bagnell, 2010). To mitigate this phenomenon, we propose the initial value problem (IVP) enhanced sampling method to make the states in the training dataset more closely match the states that the controller reaches. In the IVP enhanced sampling method, we iteratively re-evaluate the states that the NN controller reaches by solving IVPs and recalculate new training data by solving the open-loop control problems starting at these states. Our sampling method is very versatile to be combined with other techniques like a faster open-loop control solver or better neural network structures. The resulting supervised-learning-based approach empowered by the IVP enhanced sampling can be interpreted as an instance of the exploration-labeling-training (ELT) algorithms (Zhang et al., 2018; E et al., 2021) for closed-loop optimal control problems (see Appendix A for more discussions). At a high level, the ELT algorithm proceeds iteratively with the following three steps: (1) exploring the state space and examining which states need to be labeled; (2) solving the control problem to label these states and adding them to the training data; (3) training the machine learning model.

The main contributions of the paper can be summarized as follows. (1) We investigate the distribution mismatch phenomenon brought by the controlled dynamics in the supervised-learning-based approach, which explains the failure of this approach for challenging problems. We propose the IVP enhanced sampling method to update the training data, which significantly alleviates the distribution mismatch problem. (2) We show that the IVP enhanced sampling method can significantly improve the performance of the learned closed-loop controller on a uni-dimensional linear quadratic control problem (theoretically and numerically) and two high-dimensional problems (numerically), the quadrotor landing problem and the reaching problem of a 7-DoF manipulator. (3) We compare the IVP enhanced sampling method with other adaptive sampling methods and show that the IVP enhanced method gives the best performance.

Although the vanilla supervised-learning-based approach can achieve a good performance in certain problems (Nakamura-Zimmerer et al., 2021a), it is observed that its performance on complex problems is not satisfactory (see Nakamura-Zimmerer et al. (2021b); Zang et al. (2022) and examples below). One of the crucial reasons that the vanilla method fails is the distribution mismatch phenomenon. To better illustrate this phenomenon, let $\mu_{0}$ be the distribution of the initial state of interest and $\bm{u}\mathrel{\mathop{\mathchar 58\relax}}[0,T]\times\mathbb{R}^{n}\rightarrow\mathcal{U}$ be a closed-loop control function. We use $\mu_{\bm{u}}(t)$ to denote the distribution of $\bm{x}(t)$ generated by $\bm{u}$: $\dot{\bm{x}}(t)=\bm{f}(t,\bm{x}(t),\bm{u}(t,\bm{x}(t))),\bm{x}_{0}\sim\mu_{0}.$ Note that in the training process (4), the distribution of the state at time $t$ is $\mu_{\bm{u}^{*}}(t)$, the state distribution generated by the closed-loop optimal control. On the other hand, when we apply the learned NN controller in the dynamics, the distribution of the input state of $\bm{u}^{\textbf{NN}}$ at time $t$ is $\mu_{\bm{u}^{\textbf{NN}}}(t)$. The error between state $\bm{x}$ driven by $\bm{u}^{*}$ and $\bm{u}^{\textbf{NN}}$ accumulates and makes the discrepancy between $\mu_{\bm{u}^{*}}(t)$ and $\mu_{\bm{u}^{\textbf{NN}}}(t)$ increases over time. Hence, the training data fails to represent the states encountered in the controlled process, and the error between $\bm{u}^{*}$ and $\bm{u}^{\textbf{NN}}$ dramatically increases when $t$ is large. See Figures 1 (left) and 2 below for an illustration of this phenomenon.

To overcome this problem, we propose the following IVP enhanced sampling method. The key idea is to improve the quality of the NN controller iteratively by enlarging the training dataset with the states seen by the NN controller at previous times. Given predesigned (not necessarily even-spaced) temporal grid points $0=t_{0}<t_{1}<\dots<t_{K}=T$, we first generate a training dataset $S_{0}$ by solving open-loop optimal control problems on the time interval $[0,T]$ starting from points in $X_{0}$, a set of initial points sampled from $\mu_{0}$, and train the initial model $\hat{\bm{u}}_{0}$. Under the control of $\hat{\bm{u}}_{0}$, the generated trajectory deviates more and more from the optimal trajectory. So we stop at time $t_{1}$, i.e., compute the IVPs using $\hat{\bm{u}}_{0}$ as the closed-loop control and points in $X_{0}$ as the initial points on the time interval $[0,t_{1}]$, and then on the interval $[t_{1},T]$ solve new optimal paths that start at the end-points of the previous IVPs. The new training dataset $S_{1}$ is then composed of new data (between $t_{1}$ and T) and the data before time $t_{1}$ in the dataset $S_{0}$, and we train a new model $\hat{\bm{u}}_{1}$ using $S_{1}$. We repeat this process to the predesigned temporal grid points $t_{2},t_{3},\cdots$ until end up with $T$. In other words, in each iteration, the adaptively sampled data replaces the corresponding data (defined on the same time interval) in the training dataset (the size of training data remains the same). The whole process can be formulated as Algorithm 1, and we refer to Figure 1 for an illustration of the algorithm’s mechanism. We call this method IVP enhanced sampling method because the initial points of the open-loop optimal control problems are sampled by solving the IVP with the up-to-date NN controller. It is worth noting that the later iterations require less effort in labeling data as the trajectories are shorter and thus easier to solve.

It is worthwhile mentioning that the IVP enhanced sampling method is versatile enough to combine other improvements for closed-loop optimal control problems, such as efficient open-loop control problem solvers (Kang et al., 2021; Zang et al., 2022) or specialized neural network structures (Nakamura-Zimmerer et al., 2020; 2021b; 2022). One design choice regarding the network structure in the IVP enhanced sampling method is whether to share the same network among different time intervals. We choose to use the same network for all the time intervals in the following numerical examples, but the opposite choice is also feasible.

In this section, we review existing literature on adaptive sampling methods for the closed-loop optimal control problem.

We start with the methods in imitation learning (Hussein et al., 2017), which aim to learn the expert’s control function. Our task can be viewed as an imitation problem if we take the optimal control function as the expert’s control. With the same argument, we know the distribution mismatch phenomenon also exists therein. However, there is a key difference regarding the mechanism of data generation between the two settings: in imitation learning, it is often assumed that one can easily access the expert’s behavior at every time-state pair while in the optimal control problem, it is much more computationally expensive to access since one must solve an open-loop optimal control problem. This difference affects algorithm design fundamentally. Take the forward training algorithm (Ross & Bagnell, 2010), a popular method in imitation learning for mitigating distribution mismatch, as an example. To apply it to the closed-loop optimal control problem, we first need to consider a discrete-time version of the problem with a sufficiently fine time grid: $0=t_{0}<t_{1}<\dots<t_{K^{\prime}}=T$. At each time step $t_{i}$, we learn a policy function $\bar{\bm{u}}^{i}\mathrel{\mathop{\mathchar 58\relax}}\mathbb{R}^{n}\rightarrow\mathcal{U}$ where the state $x$ in the training data are generated by sequentially applying $\bar{\bm{u}}^{0},\dots,\bar{\bm{u}}^{i-1}$ and the labels are generated by solving the open-loop optimal solutions with $(t_{i},x)$ as the initial time-state pair. Hence, the open-loop control solver is called with numbers proportionally to the discretized time steps $K^{\prime}$, and only the first value on each optimal control path is used for learning. In contrast, in Algorithm 1, we can use much more values over the optimal control paths in learning, which allows its temporal grid for sampling to be much coarser than the grid in the forward training algorithm, and the total cost of solving open-loop optimal control problems is much lower.

Another popular method in imitation learning is DAGGER (Dataset Aggregation) (Ross et al., 2011), which can also be applied to help sampling in the closed-loop optimal control problem. In DAGGER, in order to improve the current closed-loop controller $\hat{\bm{u}}$, one solves IVPs using $\hat{\bm{u}}$ over $[0,T]$ starting from various initial states and collect the states on a time grid $0<t_{1}<\dots<t_{K-1}<T$. The open-loop control problems are then solved with all the collected time-state pairs as the initial time-state pair, and all the corresponding optimal solutions are used to construct a dataset for learning a new controller. The process can be repeated until a good controller is obtained. The time-state selection in DAGGER is also related to the distribution mismatch phenomenon, but somehow different from the IVP enhanced sampling. Take the data collection using the controller $\hat{\bm{u}}_{1}$ in the first iterative step for example. The IVP enhanced sampling focuses on the states at the time grid $t_{1}$ while DAGGER collects states at all the time grids. If $\hat{\bm{u}}_{1}$ is still far from optimal, the data collected at later time grids may be irrelevant to or even mislead training due to error accumulation in states. In Appendix G, we reports more theoretical and numerical comparison between DAGGER and IVP enhanced sampling, which indicates DAGGER performs less satisfactorily.

Except for the forward training algorithm and DAGGER, there are other adaptive sampling methods for the closed-loop optimal control problems. Nakamura-Zimmerer et al. (2021a) propose an adaptive sampling method that prefers to choose the initial points with large gradients of the value function as the value function tends to be steep and hard to learn around these points. Landry et al. (2021) propose to sample the initial points on which the errors between predicted values from the NN and optimal values are large. These two adaptive sampling methods both focus on finding points that are not learned well but ignore the influence of the accumulation of the distribution mismatch over time brought by controlled dynamics. We will show in Section 6 that the IVP enhanced sampling method can outperform such sampling methods.

In this section, we analyze the superiority of the IVP sampling method by considering the following uni-dimensional linear quadratic regulator (LQR) problem:

where $T$ is a positive integer, $t_{0}\in[0,T]$ and $x_{0}\in\mathbb{R}$. Classical theory on linear quadratic control (see, e.g. Sontag (2013)) gives the following explicit linear form of the optimal controls:

We consider the following two models to approximate the closed-loop optimal control function with parameter $\theta$:

Since there will be no error in learning a linear model when the data is exact, to mimic the errors encountered when learning neural networks, throughout this section, we assume the data has certain noise. To be precise, for any $t_{0}\in[0,T]$ and $x_{0}\in\mathbb{R}$, the open-loop optimal control solver gives the following approximated optimal path:

where $\epsilon>0$ is a small positive number to indicate the scale of the error and $Z$ is a normal random variable whose mean is $m$ and variance is $\sigma^{2}$. In other words, the obtained open-loop control is still constant in each path, just like the optimal open-loop control, but perturbed by a random constant. The random variables in different approximated optimal paths starting from different $t_{0}$ or $x_{0}$ are assumed to be independent.

We compare the vanilla supervised-learning-based method and IVP enhanced sampling method theoretically for the first model and numerically for the second model (in Appendix B). In the vanilla method, we randomly sample $NT$ initial points from a standard normal distribution and use corresponding optimal paths to learn the controller. In the IVP enhanced sampling method, we randomly sample $N$ initial points from a standard normal distribution, set the temporal grid points for sampling as $0<1<\dots<T-1<T$, and perform Algorithm 1. In both methods, the open-loop optimal control solver is called $NT$ times in total.

Theorem 1 compares the performance of the vanilla method and the IVP enhanced sampling method under Model 1 (5) . The more detailed statement and proof can be found in Appendix B. This theorem shows that both the distribution difference and performance difference with respect to the optimal solution for the vanilla method will increase when $T$ increases, while they are always constantly bounded for the IVP enhanced sampling method. Therefore, compared to the vanilla method, the IVP enhanced sampling method mitigates the distribution mismatch phenomenon and significantly improves the performance when $T$ is large.

In this section, we test the IVP enhanced sampling method on the optimal landing problem of a quadrotor. We consider the full quadrotor dynamic model with 12-dimensional state variable and 4-dimensional control variable (Bouabdallah et al., 2004; Madani & Benallegue, 2006; Mahony et al., 2012). We aim to find optimal landing paths from some initial states $\bm{x}_{0}$ to a target state $\bm{x}_{T}=0$ with minimum control efforts during a fixed time duration $T=16$. The open-loop optimal solutions are obtained by solving the corresponding two-point boundary value problems with the space-marching technique (Zang et al., 2022). See Appendix C, D and E for more details.

We sample $N=500$ initial points for generating training data and use a fully-connected neural network to approximate the optimal control. The temporal grid points on which we do IVP enhanced sampling is $0<10<14<16$. After learning, we use learned models to run the initial value problem at 500 training initial points and show the similarity between paths controlled by the NN controller and their corresponding training data. In Figure 2, the left sub-figure shows the average pointwise distance between data reached by the NN controller and corresponding training data at different times. The right sub-figure shows the maximum mean discrepancy (Borgwardt et al., 2006) between these two datasets using Gaussian kernel $k(x,y)=\exp({-\frac{\|x-y\|^{2}}{2}})$. In both figures, there are jumps at $t=10$ and $14$ since the NN-controlled path is continuous across time while training data is discontinuous at locations where we do IVP enhanced sampling. It can be seen that without adaptive sampling (after iteration 0), the discrepancy between the states reached by the NN controller and training data is large. With our method, they get closer to each other as the iteration goes.

Next, we check the performance of the learned NN controller. We repeat experiments five times using different random seeds. Figure 3 (left) shows the cumulative distribution function of the ratio between NN-controlled cost and optimal cost on 200 test initial points. The light solid line represents the mean of the cumulative distribution for five experiments, while the shaded area represents the mean cost ratio $\pm$ the standard deviation. We also test the performance of the ensembles, which use the average output of five networks. The ensemble results from three iterations are represented by the three dark curves in the figure. Without any adaptive sampling, the performance of model $\hat{\bm{u}}_{0}$ is poor. In contrast, the NN model is very close to the optimal policy after two rounds of adaptive sampling. The ensemble of networks after iteration 2 performs the best. More statistics of the results are provided in Appendix E.

Then we compare our method with four other methods. We also repeat every experiment five times with different random seeds. As data generation is the most time-consuming part, for the sake of fairness, we keep the number of solving open-loop problems the same (1500) among all the methods (except the last method). The first method is training a model on directly sampled 1500 optimal paths (called vanilla sampling). The second method is the adaptive sampling (AS) method proposed by Nakamura-Zimmerer et al. (2021a) that chooses initial points with large gradient norms. This is equivalent to choosing initial points whose optimal control has large norms, and we refer to this method as AS w. large u. With a little modification, the third method, AS w. large v, is to choose initial points whose total costs are large under the latest NN controller. The last adaptive sampling method, AS w. bad v, is a variant of the SEAGuL algorithm (Sample Efficient Adversarially Guided Learning) (Landry et al., 2021). The original SEAGuL algorithm proposes to use a few gradient ascent updates to find initial points with large gaps between the learned values and optimal values. Here we give this method more computational budget to solve more open-loop optimal control problems to find such initial points. The cumulative distribution functions of cost ratios of the above methods are shown in Figure 3 (right), which clearly demonstrate the superiority of the IVP enhanced sampling method. More details about the implementation and results of these methods are provided in Appendix E.

We also test the NN controllers obtained from different sampling methods in the presence of observation noises. The detailed results are provided in Appendix E. It is observed that when measurement errors exist, closed-loop controllers are more reliable than the open-loop controller and the one trained by the IVP enhanced sampling method performs best among all the considered methods. Finally, we test four different choices of temporal grid points in Algorithm 1. The results listed in Appendix E show that our algorithm is robust to the choice of temporal grid points.

In this section, we consider the optimal reaching problem on a 7-DoF torque-controlled manipulator, the KUKA LWR iiwa R820 14 (Kuka, ; Bischoff et al., 2010). See the figure in appendix F for an illustration of this task. Let $\bm{x}=(\bm{q},\bm{v})=(\bm{q},\dot{\bm{q}})$ be the state of the system where $\bm{q},\bm{v}\in\mathbb{R}^{7}$ are joint angles and velocities of the manipulator, respectively. Our goal is to find the optimal torque $\bm{u}\in\mathcal{U}\subset\mathbb{R}^{7}$ that drives the manipulator from $\bm{x}_{0}=(\bm{q}_{0},0)$ to $\bm{x}_{1}=(\bm{q}_{1},0)$ in $T=0.8$ seconds and minimizes a quadratic type cost. See Appendix F for details of the problem and the experiment configurations. To obtain training data, we use differential dynamic programming (Jacobson & Mayne, 1970) implemented in the Crocoddy library (Mastalli et al., 2020).

We use the QRNet (Nakamura-Zimmerer et al., 2020; 2021b) as the backbone network in this example (see Appendix F for details) and evaluate networks trained in six different ways: four of them are trained using Algorithm 1 with different choices of temporal grid points for adaptive sampling and two of them are trained by vanilla sampling method with $300$ (Vanilla300) and $900$ (Vanilla900) trajectories separately. All four networks (AS1–AS4) trained by the IVP enhanced sampling have initial training data of $100$ trajectories and three iterations ($K=3$), i.e., each of them requires solving the open-loop solution $300$ times in total for generating training data. The difference of the four networks lies in that they use different temporal grids ($t_{1}$ and $t_{2}$) for enhanced sampling. Each experiment has been independently run five times. We again plot the cumulative distribution functions of cost ratios (clipped at 2.0) between the NN-controlled cost and optimal cost in Figure LABEL:fig:manipulator (left: mean $\pm$ standard deviation and ensemble of AS4 in different iterations; middle: average of final networks trained by different schemes). More details and results are provided in Appendix F. We find that adding more data in the vanilla sampling method has very limited effects on improvement while the IVP enhanced sampling greatly improves the performance. Furthermore, such improvement is again robust to different choices of temporal grid points (AS1–AS4).

We additionally test the performance of the network trained by our adaptive sampling method in the presence of measurement errors. At each simulation timestep, the network’s input states are corrupted by a noise uniformly sampled from $[-\sigma,\sigma]^{14}$ for $\sigma=10^{-5},10^{-4},10^{-3}$. See Figure LABEL:fig:manipulator (right) for the results on the best model trained from AS1–AS4. The NN controller performs well at $\sigma=10^{-4}$, and there are more than $60\%$ of cases in which our controller achieves a ratio less than $2.0$ at $\sigma=10^{-3}$.

In this work, we propose the IVP enhanced sampling method to overcome the distribution mismatch problem in the supervised-learning-based approaches for the closed-loop optimal control problem. Both theoretical and numerical results show that the IVP enhanced sampling method significantly improves the performance of the learned NN controller and outperforms other adaptive sampling methods. There are a few directions worth exploring in future work. In the IVP enhanced sampling method, one choice we need to make is the temporal grid points for adaptive sampling. We recommend that at each iteration, one can compute the distance between the training data and data reached by the NN controller at different times (see Figure 2 for an example) and choose the time at which the distance starts to increase quickly as the temporal grid for adaptive sampling. We observe that the IVP enhanced sampling method performs well using this strategy. It will be ideal to make this process more systematic. Another direction is to design more effective approaches utilizing the training data. In Algorithm 1 (lines 8–9), at each iteration, we replace parts of the training data with the newly collected data, and hence some optimal labels are thrown away, which are costly to obtain. An alternative choice is to augment data directly, i.e., setting $S_{i}=\hat{S}_{i}\bigcup S_{i-1}$ in line 9. Numerically, we observe that this choice gives similar performance to the version used in Algorithm 1, which suggests that so far the dropped data provides little value for training; see Appendix  E and  F for details. But it is still possible to find smarter ways to utilize them to improve performance. We also need to evaluate the IVP enhanced sampling method for problems with other features like state/control constraints. Furthermore, the IVP enhanced sampling method can be straightforwardly applied to learning general dynamics from multiple trajectories as the controlled system under the optimal policy can be viewed as a special dynamical system. It is of interest to investigate its performance in such general settings. Finally, theoretical analysis beyond the LQR setting is also an interesting and important problem.