Title: DEEP NONLINEAR STOCHASTIC OPTIMAL CONTROL FOR SYSTEMS WITH MULTIPLICATIVE UNCERTAIN-TIES

Abstract: We present a deep recurrent neural network architecture to solve a class of stochastic optimal control problems described by fully nonlinear Hamilton Jacobi Bellman partial differential equations. Such PDEs arise when one considers stochastic dynamics characterized by uncertainties that are additive, state dependent and control multiplicative. Stochastic models with the aforementioned characteristics have been used in computational neuroscience, biology, finance and aerospace systems and provide a more accurate representation of actuation than models with only additive uncertainty. Previous literature has established the inadequacy of the linear HJB theory and instead relies on a non-linear Feynman-Kac lemma resulting in a second order Forward-Backward Stochastic Differential Equations representation. However, the proposed solutions that use this representation suffer from compounding errors and computational complexity resulting in lack of scalability. In this paper, we propose a deep learning based algorithm that leverages the second order Forward-Backward SDE representation and LSTM based recurrent neural networks to not only solve such Stochastic Optimal Control problems but also overcome the problems faced by traditional approaches and shows promising scalability. The resulting control algorithm is tested on nonlinear systems from robotics and biomechanics in simulation to demonstrate feasibility and out-performance against previous methods.

1. INTRODUCTION:

Stochastic Optimal Control (SOC) is the center of decision making under uncertainty with a long history and extensive prior work both in terms of theory as well as algorithms (Stengel, 1994; Fleming & Soner, 2006) . One of the most celebrated formulations of stochastic optimal control is for linear dynamics and additive noise, known as the Linear Quadratic Gaussian (LQG) case (Stengel, 1994) . For stochastic systems that are nonlinear in the state and affine in control, the stochastic optimal control formulation results in the Hamilton-Jacobi-Bellman (HJB) equation, which is a backward nonlinear Partial Differential Equation (PDE). Solving the HJB PDE for high dimensional systems is generally a challenging task and suffers from the curse of dimensionality.

Different algorithms have been derived to address stochastic optimal control problems and solve the HJB equation. These algorithms can be mostly classified into two groups: (i) algorithms that rely on linearization and (ii) algorithms that rely on sampling. Linearization-based algorithms rely on first order Taylor's approximation of dynamics (iterative LQG or iLQG) or quadratic approximation of dynamics (Stochastic Differential Dynamic Programming) , and quadratic approximation of the cost function (Todorov & Li, 2005b; Theodorou et al., 2010b) . Application of the aforementioned algorithms is not straightforward and requires very small time discretization or special linearization schemes especially for the cases of control and/or state dependent noise. It is also worth mentioning that the convergence properties of these algorithms have not been investigated and remain open questions. Sampling-based methods include the Markov-Chain Monte Carlo (MCMC) approximation of the HJB equation (Kushner & Dupuis, 1992; Huynh et al., 2016) . MCMC-based algorithms rely on backward propagation of the value function on a pre-specified grid. Recently researchers have incorporated tensor-train decomposition techniques to scale these methods (Gorodetsky et al., 2015) . However, these techniques have been applied to special classes of systems and stochastic control problem formulations and have demonstrated limited applicability so far. Other methods include transforming the SOC problem into a deterministic one by working with the Fokker-Planck PDE representation of the system dynamics (Annunziato & Borzi, 2010; Annunziato & Borzì, 2013) . This approach requires solving the problem on a grid as well and therefore doesn't scale to high dimensional systems due to the aforementioned curse of dimensionality.

Alternative sampling-based methodologies rely on the probabilistic representation of the backward PDEs and generalization of the so-called linear Feynman-Kac lemma (Karatzas & Shreve, 1991) to its nonlinear version (Pardoux & Rascanu, 2014) . Application of the linear Feynman-Kac lemma requires the exponential transformation of the value function and certain assumptions related to the control cost matrix and the variance of the noise. Controls are then computed using forward sampling of stochastic differential equations (Kappen, 2005; Todorov, 2007; 2009; Theodorou et al., 2010a) . The nonlinear version of the Feynman-Kac lemma overcomes the aforementioned limitations. However it requires a more sophisticated numerical scheme than just forward sampling, which relies on the theory of Forward-Backward Stochastic Differential Equations (FBSDEs) and their connection to backward PDEs. The FBSDE formulation is very general and has been utilized in many problem formulations such as L 2 and L 1 stochastic control (Exarchos & Theodorou, 2018; 2016; Exarchos et al., 2018) , min-max and risk-sensitive control (Exarchos et al., 2019) and control of systems with control multiplicative noise (Bakshi et al., 2017) . The major limitation of the algorithms that rely on FBSDEs, is the compounding of errors from Least Squares approximation used at every timestep of the Backward Stochastic Differential Equation (BSDE).

Recent efforts in the area of Deep Learning for solving nonlinear PDEs have demonstrated encouraging results in terms of scalability and numerical efficiency. A Deep Learning-based algorithm was introduced by Han et al. (2018) to approximate the solution of non-linear parabolic PDEs through their connection to first order FBSDEs. However, the algorithm was tested on a simple high-dimensional linear system for which an analytical solution already exists. A more recent approach by Pereira et al. ( 2019) extended this work with a more efficient deep learning architecture and successfully applied the algorithm for control tasks on non-linear systems. However, similar to Han et al. (2016) , the approach is only applicable to stochastic systems wherein noise is either additive or state dependent.

In this paper, we develop a novel Deep Neural Network (DNN) architecture for HJB PDEs that correspond to SOC problems in which noise is not only additive or state-dependent but control multiplicative as well. Such problem formulations are important in biomechanics and computational neuroscience, autonomous systems, and finance (Todorov & Li, 2005a; Mitrovic et al., 2010; Primbs, 2007; McLane, 1971; Phillis, 1985) . Prior work on stochastic optimal control of such systems considers linear dynamics and quadratic cost functions. Attempts to generalize these linear methods to the case of stochastic nonlinear dynamics with control multiplicative noise are only primitive and require special treatment in terms of methods to forward propagate and linearize the underlying stochastic dynamics (Torre & Theodorou., 2015) . Aforementioned probabilistic methods relying on the linear Feynman-Kac lemmma cannot be derived in case of control multiplicative noise.

Below we summarize the contributions of our work:

• We design a novel DNN architecture tailored to solve second-order FBSDEs (2FBSDEs).

The neural network architecture consists of Fully Connected (FC) feed-forward and Long-Short Term Memory (LSTM) recurrent layers. The resulting Deep 2FBSDE network can be used to solve fully nonlinear PDEs for nonlinear systems with control multiplicative noise. • We demonstrate the applicability and correctness of the proposed algorithm in four examples ranging from, traditionally used linear and nonlinear systems in control theory, to robotics and biomechanics. The proposed algorithm recovers analytical controls in the case of linear dynamics while it is also able to successfully control nonlinear dynamics with controlmultiplicative and additive sources of uncertainty. Our simulations show the robustness of the Deep 2FBSDE algorithm and prove the importance of considering the nature of the stochastic disturbances in the problem formulation as well as in neural network architecture.

The rest of the paper is organized as follows: in Section 2 we first introduce notation, some preliminaries and discuss the problem formulation. Next, in Section 3 we provide the 2FBSDE formulation.

The Deep 2FBSDE algorithm is introduced in Section 4. Then we demonstrate and discuss results from our simulation experiments in Section 5. Finally we conclude the paper in Section 6 with discussion and future directions.

2. STOCHASTIC OPTIMAL CONTROL:

In this section we provide notations and stochastic optimal control concepts essential for the development of our proposed algorithm and then present the problem formulation. Note that the bold faced notation will be abused for both vectors and matrices, while non-bold faced will be used for scalars.

2.1. PRELIMINARIES:

We first introduce stochastic dynamical systems which have a drift component (i.e. the non-stochastic component of the dynamics) that is a nonlinear function of the state but affine with respect to the controls. The stochastic component is comprised of nonlinear functions of the state and affine control multiplicative matrix coefficients. Let x ∈ R nx be the vector of state variables, and u ∈ R nu be the vector of control variables taking values in the set of all admissible controls U([0, T ]), for a fixed finite time horizon

T ∈ [0, ∞). Let [v(t) T w(t) T ] T t∈[0,T ] be a Brownian motion in R nw+1

, where v(t) ∈ R, w(t) ∈ R nw and the components of w(t) are mutually independent one dimensional standard Brownian motions. We now assume that functions f :

[0, T ] × R nx → R nx , G : [0, T ] × R nx → R nx×nu , Σ : [0, T ] × R nx → R nx×nw

and σ ∈ R + satisfy certain Lipschitz and growth conditions (see supplementary materials (SM) A).

Given this assumption, it is known that for every initial condition ξ ∈ R nx , there exists a unique solution x(t) t∈[0,T ] to the Forward Stochastic Differential Equation (FSDE),

EQUATION

where C(t, x(t), u(t)) = σG(t, x(t))u(t), Σ(t, x(t)) and (t) = v(t) w(t) .

2.2. PROBLEM STATEMENT AND HJB PDE:

For the controlled stochastic dynamical system above, we formulate the SOC problem as minimizing the following expected cost quadratic in control

EQUATION

where q : R nx → R + is the running state-dependent cost function, R (control cost coefficients) is a symmetric positive definite matrix of size n u × n u and φ : R nx → R + is the terminal state cost. q and φ are differentiable with continuous derivatives up to the second order. The expectation is taken with respect to the probability measure Q over the space of trajectories induced by the controlled stochastic dynamics in equation 1. We can define the value function as

EQUATION

Under the condition that the value function is once differentialble in t and twice differentiable in x with continuous derivatives, we can follow the standard stochastic optimal control derivation to obtain the HJB PDE (written without all the functional dependencies for simplicity)

EQUATION

with the optimal control of the form,

EQUATION

where, R (R + σ 2 G T V xx G). The derivation for both equations 4 and 5 can be found in SM B.

3. A FBSDE SOLUTION TO THE HJB PDE:

The theory of BSDEs establishes a connection between the solution of a parabolic PDE and a set of FBSDEs (Pardoux & Peng, 1990; El Karoui et al., 1997) . This connection has been used to solve the HJB PDE in the context of SOC problems. Exarchos & Theodorou (2018) solved the HJB PDE in the absence of control multiplicative noise dv with a set of first order FBSDEs. Bakshi et al. (2017) utilized the second order FBSDEs or 2FBSDEs to solve the fully nonlinear HJB PDE in the presence of control multiplicative noise, but did not consider any control in the FSDE. Lack of control leads to insufficient exploration and for highly nonlinear systems renders it impossible to find an optimal solution to complete the task. In light of this, we introduce a new set of 2FBSDEs

EQUATION

where u * is the optimal control given by equation 5 and the operator H is defined as

EQUATION

Note that the functional dependencies are dropped in this and the following section for simplicity.

We can obtain this particular set of 2FBSDEs by substituting for the optimal control (equation 5) in the forward process (equation 1). The backward processes are obtained by applying Ito's lemma (Itô, 1944) , which is essentially the second order Taylor series expansion, to the value function V and its gradient V x . Additionally, we substitute the expression for V t from equation 4 into BSDE 1 and obtain the final form of the 2FBSDEs (detailed derivation is included in SM C) as

             dx = f dt + G(u * dt + σu * dv) + Σdw dV = -q -1 2 V T x G R-T (R + 2σ 2 G T V xx G) R-1 G T V x dt +V T x G(u * dt + σu * dv) + V T x Σdw dV x = (A + V xx f )dt + V xx G(u * dt + σu * dv) + V xx Σdw x(0) = ξ, V (T ) = φ(x(T )), V x (T ) = φ x (x(T )), (8) with A = ∂ t (V x ) + 1 2 tr(∂ xx (V x )CC T ).

Using equation 8, we can forward sample the FSDE. The second-order BSDEs (2BSDEs), on the other hand, need to satisfy a terminal condition and therefore have to be propagated backwards in time. However, since the uncertainty that enters the dynamics evolves forward in time, only the conditional expectations of the 2BSDEs can be back-propagated. For more details please refer to Shreve (2004, Chapter 2). Bakshi et al. (2017) back-propagate approximate conditional expectations, computed using regression, of the two processes. This method however, suffers from compounding errors introduced by least squares estimation at every time step. In contrast a Deep Learning (DL) based approach, first introduced by Han et al. (2018) , mitigates this problem by using the terminal condition as the prediction target for a forward propagated BSDE. This is enabled by randomly initializing the value function and its gradient at the start time and treating them as trainable parameters of a self-supervised deep learning problem. In addition, the approximation errors at each time step are compensated for by backpropagation during training of the DNN. This allowed using FBSDEs to solve the HJB PDE for high-dimensional linear systems. A more recent approach, the Deep FBSDE controller (Pereira et al., 2019) , futher extends the work successfully to nonlinear systems, with and without control constraints, in simulation that correspond to first order FBSDEs. It also introduced a LSTM-based network architecture, in contrast to using separate feed-forward neural networks at every timestep as in Han et al. (2016) , that resulted in superior performance, memory and time complexity. Extending this work, we propose a new framework for solving SOC problems of systems with control multiplicative noise, for which the value function solutions correspond to 2FBSDEs.

4. DEEP 2FBSDE CONTROLLER:

In this section, we introduce a new deep network architecture called the Deep 2FBSDE Controller and present a training algorithm to solve SOC problems with control multiplicative noise. Initialize all trainable paramters and states:

θ 0 , ψ 0 , ζ 0 , x 0 = ξ for k = 1 to K do for i = 1 to M do for t = 0 to N -1 do Compute G matrix: G i t = G x i t , t ; Network prediction: V i xx,t , A i t = f F C1 f LST M (x; θ k-1 ) , f F C2 f LST M (x; θ k-1 ) Compute optimal control: u i * t = -(R + σ 2 G iT t V i xx,t G i t ) -1 G iT t V i x,t ; Sample Brownian noise: ∆w i t ∼ N (0, I∆t); ∆v i t ∼ N (0, ∆t) Compute control: K i t = K(t, x i t ) = G i t (u i * t ∆t + σu i * t ∆v i t ) Forward propagate the SDEs: x i t+1 , V i t+1 , V i x,t+1 = f 2F BSDE (x i t , V i t , V i x,t , V i xx,t , K i t ) ← time discretized equation 8 end for end for Compute mini-batch loss: L = f Loss (V N , V * N , V x,N , V * x,N , V xx,N , V * xx,N , θ k-1 ) Gradient update: θ k , ψ k , ζ k ← Adam.step(L, θ k-1 , ψ k-1 , ζ k-1 ) end for return θ K , ψ K , ζ K

Network architecture: Inspired by the LSTM-based recurrent neural network architecture introduced by Pereira et al. (2019) for solving first order FBSDEs, we propose the network in fig. 1 adapted to the uncertainties in 2FBSDEs given by equation 8. Instead of predicting the gradient of the value function V x at every time step, the output of the LSTM is used to predict the Hessian of the value function V xx and A = ∂ t (V x ) + 1 2 tr(∂ xx (V x )CC T ) using two separate FC layers with linear activations. Notice that ∂ xx (V x ) is a rank 3 tensor and using neural networks to predict this term explicitly would render this method unscalable. We, however, bypass this problem by instead predicting the trace of the tensor product which is a vector allowing linear growth in output size with state dimensionality. Of these, V xx is used to compute the control term K = G(u * dt + σu * dv). This in turn is used to propagate the stochastic dynamics in equation 8. Both V xx and A are used to propagate V x which is then used to propagate V . This is repeated until the end of the time horizon as shown in fig. 1 . Finally, the predicted values of V , V x and V xx are compared with their targets computed using x at the end of the horizon, to compute a loss function for backpropagation.

Algorithm: Algorithm 1 details the training procedure of the Deep 2FBSDE Controller. Note that superscripts indicate batch index for variables and iteration number for trainable parameters. The value function V 0 and its gradient V x,0 (at time index t = 0), are randomly initialized and trainable.

Functions f LST M , f F C1 and f F C2 denote the forward propagation equations of standard LSTM layers (Hochreiter & Schmidhuber, 1997) and FC layers with tanh and linear activations respectively, and f 2F BSDE represents a discretized version of equation 8 using the explicit Euler-Maruyama time discretization scheme. The loss function (L) is computed using the given terminal conditions as targets, the propagated value function (V ), its gradient (V x ) and Hessian (V xx ) at the final time index as well as an L 2 regularization term as follows

L = c 1 ||V * -V N || 2 2 + c 2 ||V * x -V x,N || 2 2 + c 3 ||V * xx -V xx,N || 2 2 + c 4 ||V * || 2 2 + λ||θ|| 2 2 .

(9) A detailed justification of each loss term is included in SM D. The network is trained using any variant of Stochastic Gradient Descent (SGD) such as Adam (Kingma & Ba, 2014) until convergence. The algorithm returns a trained network capable of computing an optimal feedback control at every timestep starting from the given initial state.

5. SIMULATION RESULTS:

In this section we demonstrate the capability of the Deep 2FBSDE Controller on 4 different systems in simulation. We first compare the solution of the Deep 2FBSDE Controller to the analytical solution for a scalar linear system to validate the correctness of our proposed algorithm. We then consider a cart-pole swing-up task and a reaching task for a 12-state quadcopter. Finally, we tested on a 2-link 6-muscle (10-state) human arm model for a planar reaching task. The results were compared against the Deep FBSDE Controller in Pereira et al. (2019) , where in the effect of control multiplicative noise was ignored by only considering additive noise models resulting in first order FBSDEs, and the iLQG controller in Li & Todorov (2007) . Hereon we use 2FBSDE, FBSDE and iLQG to denote the Deep 2FBSDE, Deep FBSDE and iLQG controllers respectively. The system dynamics can be found in SM E and the simulation parameters can be found in SM F. All the comparison plots contain statistics gathered over 128 test trials. For 2FBSDE and FB-SDE, we used time discretization ∆t = 0.004 seconds for the linear system problem, ∆t = 0.02 seconds for the cart-pole and quadcopter simulations and ∆t = 0.02 seconds for the human arm simulation. For the iLQG simulations, ∆t = 0.01 seconds for cartpole and quadcopter and ∆t = 0.001 seconds for the human arm were used to avoid numerical instability. These values were hand-tuned until reasonable performance was observed from iLQG. In all plots, the solid lines represent mean trajectories, and the shaded regions represent the 68% confidence region (mean ± standard deviation).

Linear System: We consider a scalar linear time-invariant system dx(t) = Ax(t)dt + Bu(t)dt + σBu(t)dv(t) + Cdw(t),

with quadratic running state cost and terminal state cost q x(t) = 1 2 Qx(t) 2 , φ x(T ) = 1 2 Q T x(t) 2 . The dynamics and cost function parameters are set as A = 0.2, B = 1.0, C = 0.1, σ = 0.5, Q = 0, Q T = 80, R = 2, x 0 = 1.0. The task is to drive the state to 0. Note that the problem is different than the one for LQG due to the presence of control multiplicative noise (i.e. v(t) = 0). Let us assume that the value function (equation 4) has the form V t, x(t) = 1 2 P (t)x 2 (t) + S(t)x(t) + c(t), where S(T ) = 0, P (T ) = Q T , c(t) = T 0 1 2 c 2 P (s)ds. The values of P and Q are obtained by solving corresponding Riccatti equations, using ODE solvers, that can be derived in a similar manner to the LQG case in Stengel (1994, Chapter 5 ) and are used to compute the optimal control (equation 5) at every timestep. The solution obtained from 2FBSDE is compared against the analytical solution in fig. 2 . The resulting trajectories have matching mean and comparable variance, which verifies the effectiveness of the controller on linear systems.

Cartpole: We applied the controllers to cartpole dynamics for a swing-up task with a time horizon of T = 2.0 seconds. The networks were trained using a batch size of 256 each for 2000 iterations. We imposed a strict restriction on the controllers with control cost coefficient R = 0.5 and did not impose any cost or target for the cart position state. We would like to highlight in fig. 3 that both 2FBSDE and iLQG choose to pre-swing the pendulum and use the system's momentum to achieve the swing-up task. They thus respect the presence of control multiplicative noise entering the system as compared to FBSDE which tries to drive the pendulum to the desired vertical position as soon as possible by applying as much control as required. Additionally, we tested for a lower control cost R = 0.1 (see SM G for plots) and observed qualitatively different behavior from the controllers.

Quadcopter::

The controller was also tested on a 12-state quadcopter dynamics model for a task of reaching and hovering at a target position with a time horizon of 2.0 seconds. The network was trained with a batch size of 256 for 5000 iterations. Only linear and angular states are included in fig. 4 since they most directly reflect the task performance (velocity plots included in SM G). The figure demonstrates superior performance of the 2FBSDE controller over the FBSDE and iLQG controller in reaching the target state faster and maintaining smaller state variance. Moreover, these results also convey the importance of taking into account multiplicative noise in the design of optimal controllers as the state dimensionality and system complexity increases. We also tested the 2FBSDE controller for the same task (reach and hover) for a longer horizon of 3.0 seconds (see SM G for plots) to verify its stabilizing capability in presence of control multiplicative noise.

2-link 6-muscle Human Arm: This is a 10-state bio-mechanical system wherein control multiplicative noise models have been found to closely mimic empirical observations (Todorov, 2005) . The controllers were tasked with reaching a target position in 1.0 second. The networks were trained with a batch size of 256 for 2000 iterations. Additive noise in the joint-angle acceleration channels and multiplicative noise in muscle activations were considered. In fig. 5 , we show the log of terminal state cost (log φ(x N )) versus different values of additive noise standard deviations (while holding multiplicative noise standard deviation fixed at 0.1) on the left and versus different values of multiplicative noise standard deviations (keeping additive noise standard deviation fixed at 0) on the right. As seen in the fig. 5 , the performance of iLQG is very erratic compared to 2FBSDE as σ is varied. The difference in behaviors can be attributed to the fact that iLQG is only aware of the first two moments of the uncertainty entering the system while the 2FBSDE, being a sampling based controller is exposed to the true uncertainty entering the system.

Discussion: Although the performance of iLQG for the cartpole and human arm tasks seem competitive to 2FBSDE, we would like to highlight the fact that iLQG becomes brittle as σ of the control multiplicative noise is increased (or time horizon is increased) and requires very fine time discretizations, proper regularization scheduling and fine state cost coefficient tuning in order to be able to converge. Moreover, the control cost (R) had to be tuned to a high value in order to prevent divergence at higher noise standard deviations. This is in contrast to the results in Li & Todorov (2007) , on account of two main reasons: firstly they do not consider the presence of additive noise in the angular acceleration channels in addition to multiplicative noise in the muscle activation channels and secondly the paper does not provide all details to fully reproduce the published results nor is code with muscle actuation made publicly available. Some of the crucial elements to fully implement the human arm model such as relationships between muscle lengths, muscle velocities, and the respective joint angles had to be obtained by studying the work done by Teka et al. (2017) .

6. CONCLUSIONS:

In this paper, we proposed the Deep 2FBSDE Controller to solve the Stochastic Optimal Control problems for systems with additive, state dependent and control multiplicative noise. The algorithm relies on the 2FBSDE formulation with control in the forward process for sufficient exploration. The effectiveness of the algorithm is demonstrated by comparing against analytical solution for a linear system, and against the first order FBSDE controller and the iLQG controller on systems of cartpole, quadcopter and 2-link human arm in simulation. Potential future directions of this work include application to financial models with intrinsic control multiplicative noise and theoretical analysis of error bounds on value function approximation as well as investigating newer architectures.

SUPPLEMENTARY MATERIALS A ASSUMPTIONS:

Assumption 1. There exists a constant k > 0 such that

EQUATION

)

where

C = [σGu Σ], ∀t ∈ [0, T ], ∀x, x ∈ R nx and ∀u, u ∈ R nu .

B HJB PDE DERIVATION:

Applying the dynamic programming principle (Bellman & Kalaba, 1964) to the value function we have

EQUATION

)

Then, we can approximate the running cost integral with a step function and apply Ito's lemma (Itô, 1944) to obtain

V (t, x(t)) = inf u∈U E Q q(x(t)) + 1 2 u(t) T Ru(t) dt + V (t + dt, x(t + dt)) Itô = inf u∈U E Q q(x(t)) + 1 2 u(t) T Ru(t) dt + V (t, x(t)) + Vt(t, x(t))dt + V T x (t, x(t)) f (t, x(t))dt + G(t, x(t))(udt + σu(t)dv(t)) + Σ(t, x(t))dw(t) + 1 2 tr Vxx(t, x(t))C(t, x(t), u(t))C T (t, x(t), u(t)) dt = inf u∈U q(x(t)) + 1 2 u(t) T Ru(t) dt + V (t, x(t)) + Vt(t, x(t))dt + V T x (t, x(t)) f (t, x(t))dt + G(t, x(t))u(t)dt + 1 2 tr Vxx(t, x(t))C(t, x(t), u(t))C T (t, x(t), u(t)) dt

We can cancel V (t, x(t)) on both sides and bring the terms not dependent on controls outside of the infimum to get

EQUATION

Note that the functional dependencies have been dropped for brevity. Now we can find the optimal control by taking the gradient of the terms inside the infimum with respect to u and setting it to zero. Before that we first substitute for C = σGu, Σ so that the term inside the trace operator becomes σ 2 VxxGuu T G T +VxxΣΣ T . By using the property of trace operators, we can write, tr(VxxCC T ) = tr(σ 2 VxxGuu T G T + VxxΣΣ T ) = tr(σ 2 u T G T VxxGu) + tr(VxxΣΣ T ). Note that the second term is not a function of u and therefore is zero when we take the gradient. Additionally, the first term is a scalar and taking its trace returns the same scalar. Now, taking the gradient w.r.t u, we get,

Ru * + (V T x G) T + σ 2 G T VxxGu * = 0 ⇒(R + σ 2 G T VxxG)u * = -(G T Vx) ⇒ Ru * = -(G T Vx) ⇒u * = -R-1 G T Vx.

where, R (R + σ 2 G T VxxG).

The optimal control above can be plugged back into equation 15 to obtain:

Vt + V T x f + q + 1 2 V T x G R-T R R-1 G T Vx -V T x G R-1 G T Vx + 1 2 σ 2 V T x G R-T G T VxxG R-1 G T Vx + 1 2 tr(VxxΣΣ T ) = 0 ⇒ Vt + V T x f + q + 1 2 V T x G R-T R + σ 2 G T VxxG R-1 G T Vx -V T x G R-1 G T Vx + 1 2 tr(VxxΣΣ T ) = 0 ⇒ Vt + V T x f + q + 1 2 V T x G R-T R R-1 G T Vx -V T x G R-1 G T Vx + 1 2 tr(VxxΣΣ T ) = 0 ⇒ Vt + V T x f + q - 1 2 V T x G R-1 G T Vx + 1 2 tr(VxxΣΣ T ) = 0

which gives the HJB PDE as mentioned in equation 4.

C DERIVATION OF THE 2FBSDES:

Firstly we can apply Ito's differentiation rule to V :

EQUATION

)

Since V is the value function, we can substitute in the HJB PDE for Vt and forward dynamics for dx to get:

EQUATION

For Vx, we can again apply Ito's differentiation rule:

EQUATION

)

We can again substitute in the forward dynamics for dx and get: dVx = ∂tVxdt (19)

EQUATION

We have A = ∂tVx + 1 2 tr(∂xxVxCC T ). Note that the transpose on Vxx is dropped since it is symmetric.

D LOSS FUNCTION:

The loss function used in this work builds on the loss functions used in Han et al. (2018) and Pereira et al. (2019) .

Because the Deep 2FBSDE Controller propagates 2 BSDEs, in addition to using the propagated value function (Vt), the propagated gradient of the value function (Vx,t) can also be used in the loss function to enforce that the network meets both the terminal constraints i.e. φ(xN ) and φx(xN ) respectively. Moreover, since the terminal cost function is known, its Hessian can be computed and used to enforce that the output of the F C1 layer at the terminal time index (Vxx,N ) be equal to the target Hessian of the terminal cost function φxx(xN ). Although this is enforced only at the terminal time index (Vxx,N ), because the weights of a recurrent neural network are shared across time, in order to be able to predict Vxx,N accurately all of the prior predictions Vxx,t will have to be adjusted accordingly, thereby representing some form of propagation of the Hessian of the value function.

Additionally, applying the optimal control, which uses the network prediction, to the forward process introduces an additional gradient path through the system dynamics at every time step. Although this makes training difficult (gradient vanishing problem) it allows the weights to now influence what the next state (i.e. at the next time index) will be. As a result, the weights can control the state at the end time index and hence the target V * (xN ) for the neural network prediction itself. This can be added to the loss function which can accelerate the minimization of the terminal cost to achieve the task objectives. The final form of our loss function is

L = c1 ||V * -VN || 2 2 + c2 ||V * x -Vx,N || 2 2 + c3 ||V * xx -Vxx,N || 2 2 + c4 ||V * || 2 2 + λ||θ|| 2 2 ,

where

V * = φ(xN ), V * x = φx(xN ) and V * xx = φxx(xN ).

E NON-LINEAR SYSTEM DYNAMICS AND SYSTEM PARAMETERS:

For all simulation experiments, we used the same quadratic state cost of the form q(x) = 1 2 (x -xtarget) T Q(xxtarget) for both the running and terminal state costs.

E.1 CARTPOLE:

The cartpole dynamics is given by d

   x θ ẋ θ   =      ẋ θ mp sin θ(l θ2 +g cos θ) mc+mp sin 2 θ -mpl θ2 cos θ sin θ-(mc+mp)g sin θ l(mc+mp sin 2 θ)      dt +     0 0 1 mc+mp sin 2 θ -1 l(mc+mp sin 2 θ)     (udt + σudv) + 02×2 02×2 02×2 I2×2

dw.

The model and dynamics parameters are set as mp = 0.01 kg, mc = 1 kg, l = 0.5 m, σ = 0.125. The initial pole and cart position are 0 rad and 0 m with no velocity, and the target state is a pole angle of π rad and zeros for all other states. The control costs R are 0.5 and 0.1 for experiments in fig. 3 and fig. 6 respectively. The state cost matrix is the same for both experiments at Q = diag [0.0, 6.0, 0.3, 0.3].

E.2 QUADCOPTER:

The quadcopter dynamics used can be found in ElKholy (2014) . Its dynamics is given by

d                    x y z φ θ ψ ẋ ẏ ż φ θ ψ                   =                     ẋ ẏ ż φ θ ψ 0 0 g Iyy Ixx ψ θ -Izz Ixx θ ψ Izz Iyy φ ψ -Ixx Iyy ψ φ Ixx Izz θ φ - Iyy Izz φ θ                     dt +                    0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 -1 m (sin φ sin ψ + cos φ cos ψ sin θ) 0 0 0 -1 m (cos φ sin ψ sin θ -cos ψ sin φ) 0 0 0 -1 m cos φ cos θ 0 0 0 0 l Ixx 0 0 0 0 l Iyy 0 0 0 0 1 Izz                       u1 u2 u3 u4    dt + σ    u1 u2 u3 u4    dv) + 06×6 06×6 06×6 I6×6 dw,

where

   u1 u2 u3 u4    =    1 1 1 1 0 -1 0 1 1 0 -1 0 d -d d -d       τ1 τ2 τ3 τ4    .

The states are positions, velocities, angles and angular rates, and the controls are the four rotor torques (τ ). 

E.3 2-LINK 6-MUSCLE HUMAN ARM:

The forward rigid-body dynamics of the 2-link human arm as stated in Li & Todorov (2004) is given by:

θ = M(θ) -1 (τ -C(θ, θ) -B θ)

where θ ∈ R 2 represents the joint angle vector consisting of θ1 (shoulder joint angle) and θ2 (elbow joint angle). M(θ) ∈ R 2×2 is the inertia matrix which is positive definite and symmetric. C(θ, θ) ∈ R 2 is the vector centripetal and Coriolis forces. B ∈ R 2×2 is the joint friction matrix. τ ∈ R 2 is the joint torque applied on the system. The torque is generated by activation of 6 muscle groups. Following are equations for components of each of the above matrices and vectors:

M = ã1 + 2ã2cosθ2 ã3 + ã2cosθ2 ã3 + ã2cosθ2 ã3 C = -θ2(2 θ1 + θ2) θ1 a2sinθ2 B = b11 b12 b21 b22 ã1 = I1 + I2 + M2l 2 1 ã2 = m2L1s2 ã3 = I2

Where b11 = b22 = 0.05, b12 = b21 = 0.025, mi is the mass (1.4kg,1kg), li is the length of link i (30cm, 33cm), si is the distance between joint center and mass of link i (11cm, 16cm), Ii represents for the moment of inertia (0.025kgm 2 ,0.045kgm 2 ). The activation dynamics for each muscle is given by: ȧi

= (1 + σdv)ui -ai t , ∀ i ∈ [1, 6] t = t deact

where t deact = 66msec, a ∈ R 6 is a vector of muscle activations, u ∈ R 6 is a vector of instantaneous neural inputs. As a result of these dynamics, six new state variables (a) will be incorporated into the dynamical system. With the muscle activation dynamics proposed above, the joint torque vector can be computed as follows:

τ = M (θ)T (a, l(θ), v(θ, θ))

wherein T (a, l, v) ∈ R is the tensile force generated by each muscle, whose expression is given by:

T (a, l, v) = A(a, l)(F l (l)FV (l, v) + FP (l)) A(a, l) = 1 -exp - a 0.56N f (l) N f (l) N f (l) = 2.11 + 4.16 1 l -1 FL(l) = exp - l 1.93 -1 1.03 1.87 FV (l, v) = -5.72-v -5.72+(1.38+2.09l)v , if v ≤ 0 0.62-(-3.12+4.21l-2.67l 2 )v 0.62+v

, otherwise

Fp(l) = -0.02 exp(13.8 -18.7l)

The Where the last column consists of moment arms which are approximated according to (Li & Todorov, 2004, figure B) . l ∈ R 6×1 is the length of each muscle, which is calculated by approximating the range of each muscle group in Li & Todorov (2004) using the data of joint angle ranges and linear relationships with corresponding joint angles as in Teka et al. (2017) . Each of the fitted linear functions for muscle lengths and cosine functions for moment arms are available in the provided MATLAB code.

The above system is a deterministic model of the 2-link 6-muscle human arm. However, in the main paper we consider experiments with both additive and control multiplicative stochasticities. A state-space SDE model for the human arm is given as follows,

d                                  θ1 θ2 θ1 θ2 a1 a2 a3 a4 a5 a6                                  =                                  θ1 θ2 θ1 θ2 -a 1 t deact -a 2 t deact -a 3 t deact -a 4 t deact -a 5 t deact -a 6 t deact                                  dt +    04×2 1 t deact I6×6    (udt + σudv) +        02×2 Σ2×2 06×2        dw.

where, θ = [ θ1, θ2] T = M(θ) -1 (τ -C(θ, θ) -B θ) and u = [u1, u2, u3, u4, u5, u6] T

In this system, the control cost coefficient matrix is R = 2I6×6. The state cost matrix is Q = diag[10.0, 10.0, 0.1, 0.1, 0.001, 0.001, 0.001, 0.001, 0.001, 0.001]. where iterations is the number of epochs. c1 to c4 represent of the weight of the components in the loss function in SM D. λ is the weights for L2 regularization of neural network. dt is the time discretization. T is the time horizon for the task. h is the output size of each LSTM layer (cell state dimension is the same).

F SIMULATION PARAMETERS:



G ADDITIONAL PLOTS G.1 CARTPOLE SWING-UP WITH LOWER CONTROL COST:

The plots in fig. 6 show the comparison of the 2FBSDE, FBSDE and iLQG controllers when the restriction on the controls is lowered as compared to that in the main paper. Here, the control cost is set to R = 0.1 as compared to R = 0.5 in the main paper. Clearly, the behavior is very different. Now, all three controllers try to reach the targets as soon as possible by applying the required control effort thus ignoring the presence of control multiplicative noise in the system. 

:



Figures and Tables:

Figure 1: Deep 2FBSDE neural network architecture. The blocks F C 1,2 are fully connected layers with linear activations while the LST M block represents recurrent layers of stacked LSTM cells with standard nonlinear activations. These layers are parameterized by θ and shared temporally. Additionally, V (x 0 , t 0 ) and V x (x, t 0 ) are parameterized by ψ and ζ respectively. The self-supervised targets V * N , V * x,N are the terminal conditions from equation 8 and V * xx,N = φ xx (x N ). Time discretization: In order to approximate numerical solutions of the Stochastic Differential Equations (SDEs) we choose the explicit Euler-Maruyama time-discretization scheme (Kloeden & Platen, 2013, Chapter 9) similar to(Pereira et al., 2019;Han et al., 2018). Here we overload t as both the continuous-time variable and discrete time index and discretize the task horizon 0 < t < T as t = {0, 1, • • • , N }, where T = N ∆t. This is also used to discretize all variables as step functions if their discrete time index t lies in the interval [t∆t, (t + 1)∆t). We use subscript t to denote the discretized variables in our proposed algorithm Alg. 1. Algorithm 1: Finite Time Horizon Deep 2FBSDE Controller Given: ξ, f , G, Σ, σ: Initial state, drift, actuation dynamics, state dependent noise matrix and control multiplicative noise std. deviation; φ, q, R: Cost function parameters; N : Task horizon, K: Number of iterations, M : Batch size; ∆t: Time discretization; λ: regularization parameter; Parameters: V 0 = V (x 0 ; ψ): Value function at t = 0 parameterized by trainable weight ψ; V x,0 = V x (x; ζ): Gradient of value function at t = 0 parameterized by trainable weights ζ; θ: Weights and biases of all fully-connected and LSTM layers;Initialize all trainable paramters and states:θ 0 , ψ 0 , ζ 0 , x 0 = ξ for k = 1 to K do for i = 1 to M do for t = 0 to N -1 do Compute G matrix: G i t = G x i t , t ; Network prediction: V i xx,t , A i t = f F C1 f LST M (x; θ k-1 ) , f F C2 f LST M (x; θ k-1 ) Compute optimal control: u i * t = -(R + σ 2 G iT t V i xx,t G i t ) -1 G iT t V i x,t ; Sample Brownian noise: ∆w i t ∼ N (0, I∆t); ∆v i t ∼ N (0, ∆t) Compute control: K i t = K(t, x i t ) = G i t (u i * t ∆t + σu i * t ∆v i t ) Forward propagate the SDEs: x i t+1 , V i t+1 , V i x,t+1 = f 2F BSDE (x i t , V i t , V i x,t , V i xx,t , K i t )← time discretized equation 8 end for end for Compute mini-batch loss: L = f Loss (V N , V * N , V x,N , V * x,N , V xx,N , V * xx,N , θ k-1 ) Gradient update: θ k , ψ k , ζ k ← Adam.step(L, θ k-1 , ψ k-1 , ζ k-1 ) end for return θ K , ψ K , ζ K

Figure 2: The 2FBSDE solution is very similar to the analytical solution for a scalar linear system.

Figure3: Cartpole Swing-up Task: 2FBSDE is most aware of the presence of control multiplicative noise and uses least control effort to perform the task. This is evident from the cart velocity plot, a directly actuated state, which tries to stay as close to zero as possible.

Figure 4: Quadcopter Reach and Hover task: 2FBSDE clearly out-performs both FBSDE and iLQG, as indicated by the smaller variance in all states.

Figure 5: Comparison of 2FBSDE and iLQG for Human Arm system: Plots show terminal cost versus additive (Σ) and control multiplicative noise standard deviations (σ). The FBSDE results are omitted as failure (divergence) occurred in all test trials. When additive noise is varied (left), 2FBSDE outperforms iLQG at lower Σ and performance deteriorates much slower than iLQG as Σ increases. When multiplicative noise is varied (right), iLQG exhibits very erratic behaviour while the terminal cost gradually increases with σ for 2FBSDE.

The model and dynamics parameters are set as m = 0.47 kg, Ixx = Iyy = 4.86 × 10 -3 kg m 2 , Izz = 8.8 × 10 -3 kgm 2 , l = 0.225 m, d = 0.05. The initial state conditions used are all zeros, and the target state is 1 meter each in the north-east-down directions and all other states zero. The control cost matrix is R = 2I4×4. The state cost matrix is Q = diag[20.0, 20.0, 50.0, 5.0, 5.0, 5.0, 1.25, 1.25, 5.0, 0.25, 0.25, 0.25].

Fig.7includes a single control trajectory as well as the mean and variance for the human arm simulation to demonstrate the feasibility of control trajectories from 2FBSDE.

Figure 6: Comparison of 2FBSDE, FBSDE and iLQG on Cartpole with control cost of 0.1.

Figure 7: Single control trajectory as well as mean and variance of 2FBSDE on Human Arm. The single trajectory demonstrates the feasibility of the resulting controls.

Figure 8: Comparison of 2FBSDE, FBSDE and iLQG on the quadcopter velocities and angular rates. Again the variance of 2FBSDE trajectories are smaller than FBSDE and iLQG.

Figure 9: Phase plot at control multiplicative standard deviation σ = 0.1 and additive noise standard deviation Σ = 0.1I. The yellow circle denotes the starting point, and yellow star denotes the target.

Figure 10: Quadcopter state trajectories of 2FBSDE for 3 seconds. The plots demonstrate the controller's capability to reach and hover. G.4 QUADCOPTER LONGER TIME HORIZON

The precise training/network parameters are shown in the following table

Table content: <table><tr><td>system</td><td colspan="6">batch size iterations c1 c2 c3 c4</td><td>λ</td><td colspan="2">dt (sec) T (sec)</td><td>h</td></tr><tr><td>CartPole</td><td>256</td><td>2000</td><td>1</td><td>1</td><td>1</td><td>1</td><td>0.0005</td><td>0.02</td><td>2</td><td>[8, 8]</td></tr><tr><td>QuadCopter</td><td>256</td><td>5000</td><td>1</td><td>1</td><td>1</td><td>1</td><td>0.0005</td><td>0.02</td><td>2</td><td>[8, 8]</td></tr><tr><td>HumanHand</td><td>256</td><td>2000</td><td>1</td><td>0</td><td>1</td><td>1</td><td>0.0005</td><td>0.02</td><td>2</td><td>[8, 8]</td></tr></table>

Peer review (iclr2020):
**Summary:** The paper contains a method tailored to a certain kind of SOC problems involving multiplicative noise. The central idea is to use a recurrent network to transform the observations into a representation that can be used with solvers specifically tailored towards that class of problems. **Decision:** I recommend to accept the paper for publication. **Arguments for decision:** The paper clearly adresses an important problem and poroposes a method capable of solving it. The method appears to be theoretically founded and the experimental validation seems solid. The relevance of the method is there as the problem class is prevalent in practical applications. The venue is a good fit as well, as the focus is the representation of a control problem in a way that allows more efficient solutions. **Feedback for improvement:** - The type setting could be improved at times. E.g. below equation (1). - I feel that the term *exploration* is overloaded. While it serves as an explicit mean to reduced the sample complexity of methods in RL, it appears to be about avoiding premature convergence in this work. I am too unfamiliar with the relevant SOC literature to judge how well the term fits, but coming from a ML background I stumbled over this expression. - Some of the experimental details, e.g. the exact choice of time discretisation, don*t appear motivated well. - The paper needs to respect [1, 2] in the related work and show relations. From the perspective of learning state representations for optimal control, both works are relevant. - Is it necessary to start the discussion from the continuous case? While I appreciate the elegance of starting out with a continuous problem and then discretising at the last step, it felt like a barrier to understanding in my case, as my understanding of continuous optimal control is limited–and I feel the audience of ICLR might have the same problem. **References:** [1] Watter, Manuel, et al. *Embed to control: A locally linear latent dynamics model for control from raw images.* *Advances in neural information processing systems*. 2015. [2] Banijamali, Ershad, et al. *Robust locally-linear controllable embedding.* *arXiv preprint arXiv:1710.05373*(2017).