How to Learn and Generalize From Three Minutes of Data: Physics-Constrained and Uncertainty-Aware Neural Stochastic Differential Equations
Authors: Franck Djeumou
Cyrus Neary
ufuk topcu
Venue: CoRL 2023 Oral (2023)
Keywords: Neural SDE
Physics-Informed Learning
Data-Driven Modeling
Dynamical Systems
Control
Model-Based Reinforcement Learning

Abstract: We present a framework and algorithms to learn controlled dynamics models using neural stochastic differential equations (SDEs)---SDEs whose drift and diffusion terms are both parametrized by neural networks. We construct the drift term to leverage a priori physics knowledge as inductive bias, and we design the diffusion term to represent a distance-aware estimate of the uncertainty in the learned model's predictions---it matches the system's underlying stochasticity when evaluated on states near those from the training dataset, and it predicts highly stochastic dynamics when evaluated on states beyond the training regime. The proposed neural SDEs can be evaluated quickly enough for use in model predictive control algorithms, or they can be used as simulators for model-based reinforcement learning. Furthermore, they make accurate predictions over long time horizons, even when trained on small datasets that cover limited regions of the state space. We demonstrate these capabilities through experiments on simulated robotic systems, as well as by using them to model and control a hexacopter's flight dynamics: A neural SDE trained using only three minutes of manually collected flight data results in a model-based control policy that accurately tracks aggressive trajectories that push the hexacopter's velocity and Euler angles to nearly double the maximum values observed in the training dataset.

How to Learn and Generalize From Three Minutes of Data: Physics-Constrainee and Uncertainty-Aware Neural Stochastic Differential Equations
Franck Djeumou" Cyrus Neary*
Ufuk Topeu
Center for Autonomy, The University of Texas at Austin fd feumou, eneary, (topcu)eutexase
Figure 1: With only three minutes of manually collected data, the proposed algorithms for training neural stochastic differential equations yield accurate models of hexacopter dynamics. When used for model-based control, they result in policies that track aggressive trajectories that push the hex- acopter's s velocity and Euler angles to nearly double their maximum values the training dataset. Videos of experiments are available at https://inyutl.com/22x5xvy
Abstract: We present framework and algorithms to learn controlled dynamics models using neural stochastic differential equations (SDEs)-SDD) whose drift and diffusion terms are both parametrized by neural networks. We construct the drift term to leverage a priori physics knowledge as inductive bias, and we de- sign the diffusion term to represent distance-aware estimate of the uncertainty in the learned model's predictions-it matches the system's underlying stochas- ticity when evaluated on states near those from the training dataset, and it predicts highly stochastic dynamics when evaluated on states beyond the training regime. The proposed neural SDEs be evaluated quickly enough for in model predictive control algorithms, or they can be used as simulators for model- model-based reinforcement learning. Furthermore, they make accurate predictions over long time horizons, even when trained on small datasets that cover limited regions simu-
the state space. We demonstrate these capabilities through experiments lated robotic systems, as well as by using them to model and control a hexacopter's flight dynamics: A neural SDE trained using only three minutes of manually col- lected flight data results in a model-based control policy that accurately tracks ag- gressive trajectories that push the hexacopter's 's velocity and Fuler angles nearly double the maximum values observed in the training dataset.
Keywords: Neural SDE, Physics-Informe Learning, Data-Driven Modeling, Dynamical Systems, Control, Model-Based Reinforcement Learning
Introduction
Leveraging physics-based knowledge in learning algorithms for dynamical systems can greatly im- prove the data efficiency and generalization capabilities of the resulting models, even if the dynamics are largely unknown [1]. However, before such physics-constrained learning algorithms may be used for model-based reinforcement learning (RL) or control, methods to estimate the uncertainty in the model's 's predictions required [2, 3, 5].
These authors contributed equally.
7th Conference on Robot Learning (CoRI. 2023), Atlanta, USA.

---

Physies knowledge
Me. (0. (........................ ..()))
informs structure
ge,
Known Terms
Physics-Aware Drift Term
eg. 8 ta (a u) -4)
S(...)
9og (-)
State
Neural SDE
da = [a( u) de + (1, u) x dW
SDESolve
Control Policy Input a
Learnd, Aistacce-Aware Uncertainty Estimate
Sampled state trajectories
Data coverage informs uncertainty
Time
Training
Figure 2: Overview. We propose neural stochastic differential equations (SDEs) to predict system dynamics. The SDE's drift and diffusion terms are both parametrized by neural networks. The drift term leverages priori knowledge into the structure of the corresponding neural network. The diffusion term trained to approximately capture the uncertainty the model's predictions.
To build uncertainty-awarc data-driven models that also leverage a priori physics knowledge, we present a framework and algorithms to learn controlled dynamics models using neural stochastic dif- ferential equations (SDEs)-S) SDEs whose drift and diffusion terms are both parametrized by neural networks. Given initial state and control signal as inputs, the model outputs sampled trajecto- ries of future states by numerically solving the SDE. Figure 2 illustrates an overview of the proposed approach, which comprises the following two major components.
Physics-constiin constrained drift term. We represent the drift term as a differentiable composition of known terms that capture priori system knowledge (c.g.. the rigid body dynamics for the hexacopter from Figure 1), and unknown terms that we parametrize using neural networks (c.g., the nonlincar motor-to-thrust relationship, and unmodeled aerodynamic effects). Even if the system' dynamics are largely unknown, this approach imposes useful inductive biases on the model: Unknown terms parametrized separately and their interactions defined explicitly.
Uncertainty-awarr diffusion term. Meanwhile, we construct the diffusion term to provide a distance- aware estimate of the uncertainty in the learned model's predictions-tte the neural SDE's predictions should have low stochasticity when it is evaluated on states similar to those from the training dataset, and they should be highly stochastic when it evaluated on states beyond the dataset. In this work, we use Euclidean distance to evaluate this notion of similarity. To build such a diffusion term, propose novel training objective that encourages the entries in the diffusion matrix to have local minima at each of the training data points, and to be strongly convex in a small surrounding neighborhood, while simultaneously ensuring that the neural SDE's predictions fit the training data.
By leveraging the physics- ahysics-based based knowledge encoded in the drift term's 's structure, the proposed mod- els are capable of making accurate predictions of the dynamics, even when they are trained on limited amounts of data and evaluated on inputs well beyond the regime observed during data col- lection. Furthermore, the proposed diffusion term not only allows the models to represent stochastic dynamics but also acts as conservative estimate the epistemic uncertainty in the model's pre- dictions. Finally, the proposed neural SDEs can be evaluated quickly enough for use in online model-predictee control algorithms, or they can be used as offline simulators for model-based rein- forcement learning.
We demonstrate these capabilities through experiments on simulated robotic systems, as well as by using them to model and control hexacopter's 's flight dynamics (illustrated in Figure 1). The pre- diction accuracy of the proposed neural SDEs is an order of magnitude higher than that of baseline model, and the neural SDEs also more reliably estimate the uncertainty in their own predictions. When applied to an offline model-based RL task, the control policies that result from our neural SDEs achieve similar performance to model-free RL policies, while requiring 30 less data.

---

Related work. Neural ordinary differential equations (ODEs) [6], which parametrize the right- hand side of a differential equation using a neural network, are a class of models that provide a nat- ural mechanism for incorporating existing physics and engineering knowledge into neural networks [1, []. For example, many works use the L.grangginn Hamiltonian, or Port-Hamiltonian formulation of dynamics to inform the structure of neural ODE [8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20]. persicc--nspired network architectures to learn models that are useful for control [21, 22, 23, 24, 25, 26, 27, 28, 29j. However, these models are typically deterministic and cannot estimate their own uncertainty. By contrast, designing uncertainty-aware models is a primary focus of this work. On the other hand, neural SDEs have also been recently studied in the context of generative modeling, learning stochastic dynamics, and estimating the uncertainty in neural network parameters [30, 31, 32, 33, 34, 35, 36]. However, to the best of our knowledge, propose the use of neural SDEs to capture model uncertainty while also leveraging a priori physics knowledge, and to apply them for model-based RI and control of dynamical systems.
Other works use similar
are the first to
Estimates of model uncertainty are crucial to preventing model exploitation in model-based rein- forcement learning and control, particularly in the offline setting [2, 3, 4, 5, 37]. Gaussian process regression has been successfully used to learn probabilistic models of system dynamics or of addi- tive residual terms [2, 38, 39, 40J. However, it can be challenging to scale such methods to large datasets and high-ding -igeenimensional systems, or to account for general non-additive interactions between the known and unknown portions of the dynamics. Meanwhile, many model-based RL algorithms use neural networks to parametrize the system dynamics. Such algorithms typically also rely on measures of model uncertainty, altermatively
some measure of the distance to the previously observed training data, during policy synthesis [41, 42, 3, 43, 4, 44, 45, 46, 47, 48). One popular method for uncertainty estimation is probabilistic ensembles [3], which uses ensembles of Gaussian distributions (whose mean and variance are parametrized by neural networks) to represent aleatoric and epistemic uncertainty. We similarly propose a method to learn uncertainty-awarr models of dy- namical systems. However, in contrast with existing ensemble-based methods, the proposed neural SDEs are capable of leveraging a priori physics knowledge, they only require the training and mem- ory costs associated with single model, and they provide uncertainty estimates that are explicitly related to the distance of the evaluation point to the training dataset.
2 Our Approach: Physics-Constrainet
Uncertainty-Awan Neural SDEs
Throughout the paper, we assume that a dataset D = (T..... Typ) } of system trajectories is avail- able in lieu of the system's 's model, where T = { ug) uo...... ... [Z]]] "|=)), z; E X is the state of time t4(7), uy E U is the state-dependent control input applied at time (T(), and X. U are the state and input spaces, respectively. Given any current state value at time ta. and any control signal u or sequence of control inputs, we seek to learn a model to predict trajectories aa+1:... Ih+tn, of length ny, Specifically, our objective is to train neural network models of dy- namical systems that account for stochastic dynamics, that explicitly model prediction uncertainties, and that allow for the incorporation of a priori physics- knowledge.
the system
2.1 Using Neural SDEs to Model Systems with Unknown Dynamics
We propose to neural SDEs to Iearn models of unknown dynamics, opposed to learning ensembles of probabilistic models that directly map states and control inputs to distributions over future states [49, 3]. A neural SDE is an SDE whose drift and diffusion terms are parametrized by neural networks, and its state is a stochastic process that evolves according to
dx = fo u) dt + Ez( u) dW,
(I)
where fe X x + R and Σe X' x u g are the drift and diffusson terms, respec- tively, and W is the nh -dimensional Wiener process. The notation x indicates that the SDE is either in the Itō [50] Stratonovich [51] form. The reader unfamiliar with these forms should feel free to ignore the distinction [52, 53, 30]. which becomes an arbitrary modeling choice when fo and Eg are learned. Under mild Lipschitz assumptions [54, 55] fe. Ez. we can efficiently sample a distribu- tion of predicted trajectories via numerical integration of the neural SDE [56, 57]. Furthermore, we can incorporate physics knowledge into such models, discussed further in Section 2.3.
Efficient SDE Training. Given D, we train the unknown functions fo and E. in an end-to-end fashion. For fixed 0 and o we integrate the neural SDE (1) to obtain a set of system trajectories over
3

---

the time horizon [T)(T), (T)). We then define the training objective
Jatta
Σ
- L
OGg- xz)
Pp
(2)
"pp j=i-1p=
(ii JED
"p p=1 SDESolve(z;; M4,...,, Ui+n,-1; J0, Ex);
where (Ii, 04y ritne) E 'D denotes any sequence of transitions of length n, from the dataset, the predicted trajectory x : ,itn. is the p-th path sampled by any differentiable numerical integrator SDFSolve, and "p is the total number of sampled paths. We optimize the parameters and & by minimizing Taata using either automatic differentiation through SDESolve, or via the adjoint sensitivity method for SDEs [33].
For given time index j, the summand in (2) is proportional the negative log-likelihood of the observed state I with respect to fixed-variance Gaussian distribution with covariance matrix S centered around the predicted state zzf. Minimizing Jaam is thus equivalent to minimizing the likeli- hood of the observed trajectories under the neural SDE, given Gaussian state observation model. This formulation has been a common choice in the literature [33, 30), and is suitable for many robotics and engineering settings. However, we expect that more nuanced forms of aleatoric uncer- tainty can be captured by adapting Jdata. e.g., using kernel density estimation to approximate the SDE-generated distribution over future states when defining the data likelihood.
2.2 Using Neural SDEs to Represent Distance-Aware Estimates of Uncertainty
We emphasize that direct optimization of the model as presented in Section 2.1 would result in the diffusion approaching zero, if the drift term alone could be used to fit the training data well [31]. Instead, we propose to design, constrain, and train the diffusion term to represent modeling uncertainties. In particular,
propose approach to learn distance-aware diffusion term 2ddad that may capture the true stochasticity of the dynamics when it is evaluated near points in the training dataset (in terms Euclidean distance), while constrained to output highly stochastic predictions for points that are far from the dataset. A feature of our approach is that Dgad is neural network trained by locally sampling points near the training dataset, and yet its predictions respect useful global properties. Specifically, we translate this notion of distance-awarenesss into several mathematical properties that Idlad must satisfy and propose suitable loss function for its training. For notational simplicity, we use [xu]] to denote the vector that results from concatenating z; and uz.
Bounded diffusion far from training dataset. On evaluation points sufficiently far from the train-
ing dataset, the entries in the diffusion term should saturate at maximum values. We accordingly
propose to write >d ([xu) == m ([xu)) @ (n)(((())) where nh : X' x U + [0, 1] parametrized scalar-valuee function that encodes the notion of distance awareness, he [0, 1]
[0, 1.11"xx,,, XRw is a parameterized element-wise monotonic function, Σ X x 2 + R", x"n is
matrix specifying element wise maximum values, and @ denotes the element- wise product.
Small diffusion values near the training dataset. On points included in the training dataset, Daad should not introduce additional stochasticity related to epistemic uncertainty. We ensure that the model's 's outputs are less stochastic when evaluated on points near the training dataset than on those far from it, by requiring that each entry of the diffusion term has a local minima at each pair (z,, u;) in D. We impose this property on ddad by training The to have zero-gradient values at every point in the training dataset. This yields the loss function Tgrad
Σ [x
[ [xui ([xx])]
Increasing diffusion values along paths that move away from the training dataset. As the query point moves away from the points in the training dataset, the magnitude of the diffusion term should monotonically increase. Let I' be any path along which the distance from the current point to the nearest training datapoint always increases. Then, along I' the entries of Iddd should monotonically increase. We enforce this property via local strong convexity constraints near the training dataset. Specifically, for every datapoint (I;, u.) E DD and a fixed radius > O, we design The to be strongly convex within ball B, (x, us) == {(x,
[[[xu] [xu]a]] r) with strong convexity constant > 0. That is, for any (z. u), (a', o ud) E (x,,, me), we desire that the constraint ( u, u, u²) holds, where SC,, is defined as:
(I, u, ud) = ([xul)) - ny ([xu))- V(xuj(s< T ([xu]' - [xu]) - u4 [|xu]x - [xu]]]²

---

Now, to eliminate the strong convexity constants as tuneable hyperparameters, we parametrize
function Mc : X' u + R, using a neural network. Note that while Pę is a continuous function over X x U, we only ever evaluate it at points in the training dataset. This is effectively equivalent to learning separate values of μ4 for every (ri, (4) E We then define the loss functions
f0, if (x, u, u') (SC u, u, w))?, u,), otherwise
Js Σ
E
and Ju
Σ (BL.)))D
Me (2 u.)
(BLUDED
(a',),
((((a)))))))))
where ((1,; Uz); r) Gaussian distribution with mean (2i, u4) and standard deviation r. Here, Ju is a regularization loss term that encourages high values of pag This ensures that The reaches its maximum value close to the boundaries of the training dataset.
Modeling details. In practice, has sigmoid at its output, no(x, u) = sigmoid (I, u)). b), where W. b E RR and we restrict each entry of W greater than Thus, ho (2; u)) sigmoid(NNx((x uW + b), which allows it to learn heterogeneous distributions of noise from the output of 7/-
Furthermore, we define (z) == sigmoid W sigmoid
Empirical illustration of 17-- To demonstrate the effectiveness of the designed constraints and loss functions, we train 7/ to optimize Arrad + >seTse + AI.Jn on two illustrative datasets of 2D trajectories. We parametrize N and le as neural networks with swish activa- tion functions, 2 hidden layers of size 32 each for The and size each for juç. We pick the val- ues Agrad = Ase = A = 1, and the ball radius r = 0.05. Figure illustrates the learned 77 evaluated the grid [-0.2, 0.2] x [-0.2, 0.2].
Datasel
0.4
Figure 3: Output values of Te after training on illustrative sample datasets in 2D domain.
We observe that Th outputs low values near the training dataset, that these values increase as we move away from the dataset, and that the values are close to 1 when evaluated far from the dataset.
2.3 Leveraging A Priori Physics Knowledge in Neural SDEs
We represent the drift term fe (I) as the composition of a known function derived from priori knowledge and a collection of unknown functions that must be learned from data. That is, we write fo(x, u) F(x, u, 96, (:)..... go, (.),, where F is a known differentiable function and go, (:).....oo, are unknown terms within the underlying model. The inputs to these functions could themselves be arbitrary functions of the states and control inputs, or the outputs of the other unknown terms. Furthermore, by designing the element-wise maximum values D (xr, u) for the diffusion term, we can encode prior notions of uncertainty into the neural SDE by specifying the maximum amount of stochasticity that the model will predict when evaluated on any given region of the state space. We provide several examples of such physics-basee knowledge Section 3. The complete neural SDE learning problem can now be formulated
minimize 0...0..
AdataJdata Agrad + AsoJse + AynJ
Experimental Results
We evaluate approach by comparing its performance for modeling and model-based control tasks against state-of-the-arn techniques such as probabilistic ensembles [49, 3, 58), system identification- based algorithms [59, 60, 61], and neural ODEs [6, 1]. Experimental details and additional results
are available in the appendix. Code is available at https//gihub.coomwwwwsst wuwushrek/sde4mmb.t
3.1 Spring-Mass-DDant Daaser: Generalization and Uncertainty Estimation with Limited Data
We start with the problem of modeling the dynamics of an uncontrolled dass-e system given 5 noisy trajectories of seconds each. The equations of motion are given by the state æ = =[g. where g is the position of the mass and q is its velocity.
5

---

SDE model. We train the informed model dæ = 14. 96, (z)]dt - n @ (1()((2) dW, where the vector max specifies the maximum level of stochasticity outside of the training dataset. We note that while this SDE model assumes is unknown, it knows that the rate of change of is
Neural SDE improves uncertainty estimates and accuracy baselines for comparison, we train neural ODE (which has the same architecture as the SDE model but excludes the diffusion term), and an ensemble of probabilistic (Gaussian) models [58].
probabilistic ensembles. As
Figure 4 shows the prediction accu- racy and uncertainty estimates of the learned models on grid discretiza- tion of the state space. We use each point of the grid as an initial state for the models to predict 100 trajec- tories over time horizon of 0.2 sec- onds. The first row shows that neural SDE and ODE models have high pre- diction accuracy. The neural SDE's accuracy decreases only slightly out- side of the training dataset. On the other hand, for this low data regime, the probabilistic ensemble model at least one order of magnitude less accurate than the neural SDE at points near the training dataset. The second row illustrates the variance the model predictions, which we interpret as a measure of uncertainty. While the neural ODE has notion of prediction uncertainty, the proposed neural SDE yields uncertainty estimates that agree with the availability of the training data. There is no clear relationship between the training dataset and the probabilistic ensemble's 's uncertainty predictions.
Neural ODE
Prob. Ensemble
Our Model
0.2
Datases
Prod
002
0.2
1.0
Unnotainn
0.8
0.5
-02
0.0
0.2
-0.2
0.0
0.2
-0.2
0.0
0.2
Figure 4: Model prediction errors and uncertainty esti- mates, when trained on the illustrated dataset.
Neural SDE (Ours)
Neural ODE
Ground Truth
Neural SDE generalizes beyond the training dataset. Figure 5 shows the state trajectory predicted by the neu- ral SDE and the neural ODE, from a given initial state. We observe that both models generalize well be- yond the training dataset and are suit- able for long-term prediction. How- ever, the neural SDE's predictions are slightly more accurate than those of the neural ODE.
-10
0.2
2
Pred. Error
1.5
0
q
-0.2
0.5
4
2
6
Time
Time
Figure 5: Neural SDE improves accuracy over neural ODE.
3.2 Cartpole Swingup: Offline Model-Based Reinforcement Learning with Neural SDEs
We now consider the cartpole swingup problem. The system state defined by x = IP Pz. 0 and the scalar control input u. Here, is the position of the cart, is the angle of the pendulum, and u the force applied the cart.
Data collection. In contrast to the Hass-Spring-Damaae experiments, we now utilize a moderate amount of data ensure that the probabilistic ensemble methods are able to learn reasonably curate models of the dynamics. Specifically, we build two training datasets, each of which consists of 100 trajectories of 200 system interactions. The first dataset, referred to random dataset, is collected by applying random control inputs. The second dataset, referred to as on-policy dataset, is collected by applying inputs generated from policy trained through Proximal Policy Optimization (PPO) [62] and 600, 000 interactions with the ground truth dynamics model.
Experimental setup. We consider two physics-constraings neural SDE models with differ- ent prior knowledge. The first model, simply referred to as neural SDE, is given by dæ = [e,, 90, (x, u), 0, 9o, (x, w)]dt + max @ (n (x)) x dW while the second model, referred to as neural SDE with side info, exploits the known control-affine structure of the dynamics and is given by: dæ
[p=, 9o, (xr) go, 0)u, 0, gog (x) go, (0, 0)zjdl max (1 (æ)) dW.
We train both neural SDE models, neural ODE models with access to the same physics information, and ensemble of probabilistic (Gaussian) models [58], using both the random and the on on-policy

---

nandon dataset
on policy dataset
Episodic Reward
513
303
89
Neural
Dynamics
(Ous)
kissemble
wi Side Info.
w/ Side Info.
(Ous)
Elsemble
(Ours
Policy evaluated using dynamics
Policy evalusted using learned dynamics
Figure 6: Reward RL policies trained using learned dynamics.
datasets. For each learned model, we train control policy via the PPO algorithm while using the model as environment simulator. We compare the performance of the learned policies with the baseline policy used to generate the on-policy dataset. We note that the learned policies are trained on a dataset that is 30x smaller than that required to train this baseline model-free RL policy. Figure 6 illustrates the mean episodic reward values achieved when evaluating the learned policies using the true dynamics (hatch pattern) vs using the learned dynamics models (grid pattern). The RL policies trained using the proposed neural SDE models consistently achieve higher rewards than policies trained using neural ODEs or probabilistic ensemble models of the dynamics.
Offline model-based RL using neural SDEs is as performant as model-free RL, while requiring 30x fewer interactions with the environment. For the random dataset, the model- based policy trained using the neural SDE achieves comparable performance to the model-free baseline policy. For the on-policy datset (which is more limited in its coverage of the Mate-action space), the neural SDE results in a model-based based policy that less performant than the baseline policy. However, this gap is closed when considering the neural SDE that incorporates more physics knowledge.
Neural SDEs robust against model exploitation. For the on-policy dataset, the probabilistic ensemble results in model-based policy that performs much worse when operating in the true envi- romment than it does when using the learned model simulator. However, the neural SDE appears to be much more robust to this symptom of model exploitation. We note that many model-based RL algorithms are designed explicitly prevent the policy optimization procedure from exploit- ing model inaccuracies. We did not employ such algorithms in this experiment but instead used an implementation of the PPO algorithm with default parameters [63].
3.3 Hexacopter System: Real-Time, Learning- Based Control with 3 Minutes of Data
We now demonstrate that by using only basic knowledge of rigid- dynamics and 3 minutes of flight data, approach bridges the sim-to-real gap and enables real-time control of hexacopter.
Experimental setup. Our hexacopter uses a CubePilot Orange flight controller running PX4 [64], and a Beelink MINIS 12 as the main computational unit. The CubePilot fuses measurements from the onboard IMU and motion capture system estimate the state x == [pw, lw; qw B: WB] frequency of 100 Hz. Here pw and vw the position and velocity in the world frame, qw B is the quaternion representing the body orientation, and wi is the angular rate in the body frame. The Beelink unit compiles our SDE models, solves model predictive control (MPC) problem at each new state estimate, and sends back the resulting motor and desired angular rate commands.
Data collection: 3 minutes worth of data. We collect 3 system trajectories by manually flying the hexacopter via radio-based remote controller. We store the estimated states æ and input motor commands [o, 1]6 at frequency of 100 Hz. We obtain total of 203 seconds worth of flight data. analysis shows that 95% of the data corresponds to the hexacopter operating below the speed of 1.71m/s. absolute roll of 18 and absolute pitch of 13 while the maximum absolute speed, roll, and pitch attained are respectively 2.7 m/s, 23", and 23
SDE model. The physics-constraing model leverages the structure of 6-dof rigid body dynamics while parametrizing the aerodynamics forces and moments, the motor command thrust function, and (geometric) parameters of the system such as the mass and the inertia matrix:
[w WB
(gwB (Ta(u) + f1ee fo (zx @w + gw X JowB
=
mo
di + of max o holn x dW,
Jo (u) + (x
7

---

where feat [tw. wp). x denotes the cross product, qw B is the conjugate of qw B, the vector o the maximum diffusion calibrated by overestimating sensor noise, the variables mg and Jo djag(Jō, , . J8) represent the system mass and inertia matrix, the neural network functions f0 and Mg represent the residual forces and moments due to unmodelled and higher order aerodynamic effects, the parametrized functions T3 and Mg estimate the map between motor commands, thrusts, and moment values. We use gw to denote the gravity vector. Our model assumes the elementary knowledge that dpw = owdt and dgw 0.5qw (0, wnldf, which completes the full SDE model.
Neural SDE improves prediction accu- racy over system identification while also providing uncertainty estimates. As baseline model, we remove the diffu- sion term in our SDE model and use non- linear system identification (SysID) [61, 60] algorithm
Neural SDE (Ours)
SysID
Ground Truth
2/
Wy
learn the parametrized terms mg, Ja, fø M2 Te and Mp. The nonlinear unknown functions are parametrized using polynomial basis. Figure shows the model predictions along trajectory that not included in the training dataset. We split the trajectory into segments of I second and perform open-loop predic- tion of the state evolution given the segment's ' initial state and the motor commands. The proposed neural SDE demonstrates improved accuracy and stability in comparison with the SyslI) baseline. We include additional comparisons to neural ODE and SysID baselines for the prediction and control of a high-fidelity quadcopter simulation in the appendix.
-5
47
50
55
60
47
50
55
60
Time
Time
Figure 7: Our neural SDE achieves accurate and stable long-horizon predictions.
Neural SDE results in highly performant MPC even when operating beyond the training regime. Figure shows the result of using the learned neural SDE in an MPC algorithm to track a lemniscate trajectory. The trajectory is obtained through generic minimum snap trajectory gen- eration [65] without any knowledge of the hexacopter dynamics. Figure 8 demonstrates the ability of the neural SDE to generalize far beyond the limited training dataset, and to achieve high per- formance in terms of tracking accuracy when used for MPC. For example, for the first 10 seconds, the hexacopter reaches speed up 1.7 m/s, yet the tracking error remains only em. We provide additional results in the supplementary material for tracking an aggressive circle trajectory.
1
15
RMSE [em]
Py
3
10
>
0
MPC+SDE
2
Reference
1
-1
-2
10
20 Time
30
40
20 Time
40
Pr
Figure 8: Tracking performance on the lemniscate trajectory. We obtain a root mean squared error (RMSE) of 15 em despite operating at speeds up 3.6 m/s and roll, pitch angles up to 32", nearly twice as high as the maximum values from the training dataset (shown by the dashed red line).
Discussion and Limitations
We present a framework and algorithms for learning uncertainty-aware models of controlled dy- namical systems, that leverage a priori physics knowledge. Through experiments on a hcxacopter and on other simulated robotic systems, demonstrate that the proposed approach yiclds data- efficient models that generalize beyond the training dataset, and that these learned models result in performant model-based reinforcement learning and model predictive control policies.
While our formulation is general enough to extend to partial state observations easily, our current experiments use noisy observations of the system state. Future work will extend the proposed ap- proach to learn dynamics from more varied observations of the state, e.g., from image observations. Furthermore, future work will explore methods to improve the neural SDE's ability to capture nu- anced forms of aleatoric uncertainty, e-g., using kernel density estimation to approximate the SDE's state distribution when defining the loss function.

---

Acknowledgments
This work was supported in part by AFOSR FA9550-19-110005 and NSF 2214939.
References
[1] Djeumou, C. Neary, Goubault, S. Putot, and U. Topeu. Neural networks with physics- informed architectures and constraints for dynamical systems modeling. In Proceedings of The 4th Annual Learning for Dynamics and Control Conference, volume 168 of Proceedings of Machine Learning Research, 2022.
[2] M. Deisenroth and C. E. Rasmussen, Pilco: A model-based and data-efficient approach to pol- icy search. In Proceedings of the 28th International Conference machine learning (ICML- I1), pages 465-472, 2011.
[3] K. Chua, R. Calandra, R. McAllister, and Levine. Deep reinforcement learning in a handful of trials using probabilistic dynamics models, In Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
[4] T. Kurutach, I. Clavera, Y. Duan, A. Tamar, and P. Abbeel. Model-ensemble trust-region policy optimization. In International Conference on Learning Representations, 2018.
[5] S. Levine, A. Kumar, G. Tucker, and Fu. Offline reinforcement learningg Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.0164. 2020.
[6] R. T. Q. Chen, Y. Rubanova, J. Bettencourt, and D. K. Duvenaud. Neural ordinary differ- ential equations. In Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
[7] C. Rackauckas,, Y. Ma, J. Martensen, C. Warner, K. Zubov, R. Supekar, D. Skinner, A. Ramad- han, and A. Edelman. Universal differential equations for scientific machine learning. arXiv preprint arXiv:2001.04385, 2020.
[8] S. A. Desai, M. Mattheakis, D. Sondak, P. Protopapas, and S. J. Roberts. Port-hamiltonian neural networks for learning explicit time-dependent dynamical systems. Physical Review E, 104(3):034312, 2021.
[9] Eidnes, A. Stasik, C. Sterud, E. Bøhn, and Riemer-Sorensen. Eseudo-hamiltoniar neural networks with state-dependent external forces. Physica D: Nonlinear Phenomena, 446: 133673, 2023.
[10] Z. Chen, J. Zhang, M. Arjovsky, and L. Bottou. Symplectic recurrent neural networks. In International Conference on Learning Representations, 2020.
[11] A. Zhu, P. Jin, and Y. Tang. Deep hamiltonian networks based on symplectic integrators. arXiv preprint 20X4::20004138300 2020.
[12] C. Bllnc-Blanchete,, S. Veer, Majumdar, and N. E. Leonard. Lagnetvip: lagrangian neural network for video prediction. arXiv preprint arXiv:2010. /2932, 2020.
[13] Y. D. Zhong, B. Dey, and A. Chakraborty. Extending lagrangian and hamiltonian neural net- works with differentiable contact models. In Advances in Neural Information Processing Sys- tems, volume 34. Curran Associates, Inc., 2021.
[14] Y. D. Zhong, Dey, and A. Chakraborty. Benchmarking energy-conserving neural networks for learning dynamics from data. In Proceedings of the 3rd Conference on Learning for Dy- namics and Control, volume 144. PMLR, 2021.
[15] Y. D. Zhong, B. Dey, and A. Chakraborty. Dissipative symoden: Encoding hamiltonian dy- namies with dissipation and control into deep learning. In ICLR 2020 Workshop on Integration of Deep Neural Models and Differential Equations, 2020.

---

[16] M. A. Rochrl, T. A. Runkler, V. Brandtstetter, M. Tokic, and S. Obermayer. Modeling sys- tem dynamics with ahysics-informed neural networks based on lagrangian mechanics. IFAC- PnpersOnl.ine, 53(2):9195-9200 2020.
[17] C. Neary and U. Topeu. Compositional learning of dynamical system models using port- hamiltonian neural networks. In 5th Annual Learning for Dynamics & Control Conference, 2023.
[18] M. Finzi, K. A. Wang, and A. G. Wilson. Simplifying hamiltonian and lagrangian neural networks via explicit constraints. In Advances Neural Information Processing Systems, volume 33. Curran Associates, Inc., 2020.
[19] S. Greydanus, M. Dzamba, and J. Yosinski. Hamiltonian neural networks. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[20] M. Cranmer, S. Greydanus, S. Hoyer. P. Battaglia, D. Spergel, and S. Ho. Lagrangian neural networks. arXiv preprint arXiv:2003.0400 2020.
[21] X.-Y. Liu J.- Wang. Physics-informee dyna-style model-based deep reinforcement learn- ing for dynamic control. Proceedings of the Royal Society A: Mathematical, Physical and Engineering Sciences, 477(2255).
[22] G. Shi, X. Shi, M. O'Connell, R. Yu, K. Azizzadenesheli, A. Anandkumar, Y. Yue, and S.-J. Chung. Neural lander: Stable drone landing control using learned dynamics. In International Conference Robotics and Automation, pages 9784 9790, 2019.
[23] S. S.-E. Plaza, R. Reyes-Báez, and B. Jayawardhana. Total energy shaping with neural in- terconnection and damping assignment passivity based control. In Proceedings of The 4th Annual Learning for Dynamics and Control Conference, volume 168. PMLR, 2022.
[24] L. Furieri, C. L. Galimberti, M. Zakwan, and G. Ferrari-Trecate Distributed neural network control with dependability guarantees: a compositional port-hamiltonian approach. In Pro- ceedings of The 4th Annual Learning for Dynamics and Control Conference, volume 168. PMLR, 2022.
[25] Wong, V. Makoviychuk, A. Anandkumar, and Zhu. Oscar: Data-driven operational space control for adaptive and robust robot manipulation. In 2022 International Conference on Roboties and Automation, 2022.
[26] K. Menda, J. K. Gupta, Z. Manchester, and M. J. Kochenderfer. Structured mechanical models for efficient reinforcement learning. In Workshop on Structure and Priors Reinforcement Learning. International Conference on Learning Representations, 2019.
[27] K. Gupta, K. Menda, Z. Manchester, and M. Kochenderfer. Structured mechanical models for robot learning and control. In Proceedings of the 2nd Conference on Learning for Dynamics and Control, volume 120. PMLR, 2020.
[28] T. Duong and N. Atanasov. Hamiltonian based neural ode networks on the se (3) manifold for dynamics learning and control. In Roboties: Science and Systems (RSS), 2021.
[29] M. Lutter, C. Ritter, and J. Peters. Deep lagrangian networks: Using physics as model prior for deep learning. International Conference on Learning Representations, 2019.
[30] P. Kidger. On neural differential equations. PhD thesis, University of Oxford, 2021.
[31] W. Xu, Q. Chen, X. Li, and D. Duvenaud. Infinitely deep bayesian neural networks with stochastic differential equations. In Proceedings of The 25th International Conference on Artificial Intelligence and Statisties, volume 151, PMLR, 2022.
[32] L. Kong, J. Sun, and C. Zhang. SDE-net: Equipping deep neural networks with uncertainty estimates. In Proceedings of the 37th International Conference on Machine Learning, volume 119, PMLR, 2020.
10

---

[33] X. Li, T.-K. L. Wong, R. T. Chen, and D. K. Duvenaud. Scalable gradients and variational inference for stochastic differential equations. In Symposium on Advances Approximate Bayesian Inference. PMLR, 2020.
[34] J. Jia and A. R. Benson. Neural jump stochastic differential equations. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[35] L. Yang, T. Gao, Y. Lu, J. Duan, and T. Liu, Neural network stochastic differential equation models with applications to financial data forecasting. Applied Mathematical Modelling, 115: 279-299, 2023.
[36] P. Kidger, J. Foster, X. Li, and T. J. Lyons. Neural sdes as infinite-dimensional gans. In Proceedings of the 38th International Conference Machine Learning, volume 139 of Pro- ceedings of Machine Learning Research. PMLR, 2021.
[37] M. Vlastelica, S. Blaes, C. Pinneri, and G. Martius. Risk-averse zero-order trajectory optimiza- tion. In Proceedings of the 5th Conference on Robot Learning, volume 164. PMLR, 2022.
[38] Kabzan, L. Hewing, A. Liniger, and M. N. Zeilinger. Learning basedd model predictive control for autonomous racing. IEEE Robotics and Automation Letters, 4(4):3363-33700 2019. doi: 100..010000.RA2.1..22922220
[39] A. Liu, G. Shi, S.-J Chung, A. Anandkumar, and Yue. Robust regression for safe exploration in control. In Learning for Dynamics and Control, pages 608-619. PMLR, 2020.
[40] M. Saveriano, Y. Yin, P. Falco, and D. Lee. Data-efficient control policy search using residual dynamics learning. In 2017 IEEE/RS.J International Conference on Intelligent Robots and Systems (IROS), pages 4709-4715, 2017. 100:10.110//R00.200...855 8.806633..33.
[41] G. An, S. Moon, J.-H. Kim, and H. O. Song. Uncertainty-basee offline reinforcement learning with diversified q-ensemble. In Advances in Neural Information Processing Systems, vol- ume 34. Curran Associates, Inc., 2021.
[42] T. Yu, A. Kumar, R. Rafailov, A. Rajeswaran, S. Levine, and C. Finn. Combo: Conserva- tive offline mosel-based policy optimization. In Advances in Neural Information Processing Systems, volume 34. Curran Associates, Inc., 2021.
[43] M. Janner, J. Fu, M. Zhang, and Levine. When trust your model: Model-based policy optimization. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[44] Depeweg, J. M. Hernández-Lobato. Doshi Velez, and S. Udluft. Learning and policy search in stochastic dynamical systems with bayesian neural networks. In International Con- ference on Learning Representations, 2017.
[45] T. Yu, G. Thomas, L. Yu, S. Ermon, J. Y. Zou, Levine, C. Finn, and Ma. Mopo: Model- based offline policy optimization. In Advances in Neural Information Processing Systems, volume 33. Curran Associates, Inc., 2020.
[46] R. Kidambi, A. Rajeswaran, P. Netrapalli, and T. Joachims. Morel: Model-based offline re- inforcement learning. In Advances in Neural Information Processing Systems, volume 33. Curran Associates, Inc., 2020.
[47] J. Jeong, X. Wang, M. Gimelfarb, H. Kim, B. abdulhai, and S. Sanner. Conservative bayesian model-based value expansion for offline policy optimization. In The 11th International Con- ference on Learning Representations, 2023.
[48] A. Kumar, A. Zhou, G. Tucker, and S. Levine. Conservative q-learning for offline reinforce- ment learning. In Advances in Neural Information Processing Systems, volume 33. Curran Associates, Inc., 2020.
[49] B. Lakshninarayanan, A. Pritzel, and C. Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.

---

[50] K. Itó. On stochastic differential equations. Number 4. American Mathematical Soc., 1951.
[51] R. L. Stratonovich. A new representation for stochastic integrals and equations. Siam Journal on Control, 4:362-371, 1966.
[52] N. G. Van Kampen. Itō versus stratonovich. Journal of Statistical Physics, 24:175-187, 1981.
[53] S. Massaroli, M. Poli, S. Peluchetti, J. Park, A. Yamashita, and H. Asama. Learning stochastic optimal policies via gradient descent. IEEE Control Systems Letlers, 6:1094-1099, 2021.
[54] B. Øksendal. Stochastic differential equations. Springer Berlin Heidelberg, 2003.
[55] H. Kunita. Stochastic flows and stochastic differential equations, volume 24, Cambridge university press, 1997.
[56] G. N. Milstein. Numerical integration of stochastic differential equations, volume 313. Springer Science & Business Media, 1994.
[57] E. Kloeden, E. Platen, and H. Schurz. Numerical solution of SDE through computer experi- ments. Springer Science & Business Media, 2002.
[58] L. Pineda, B. Amos, A. Zhang, N. O. Lambert, and R. Calandra. Mbrl-lib: A modular library for model-based reinforcement learning, arXiv preprint arXiv:2 1015599 2021.
[59] S. L. Brunton, J. L. Proctor, and J. N. Kutz. Discovering governing equations from data by sparse identification of nonlinear dynamical systems. Proceedings of the national academy of sciences, 1 113(15)3333--3337 2016.
[60] E. Kaiser, J. N. Kutz, and S. L. Brunton. Sparse identification of nonlinear dynamics for model predictive control in the low-data limit. Proceedings. Mathematical,, Physical, and Engineering Sciences, 474, 2017.
[61] B. M. de Silva, K. Champion, M. Quade, J.-C. Loiseau, N. Kutz, and S. L.. Brunton. Pysindy: a python package for the sparse identification of nonlinear dynamics from data. arXiv preprint arXiv:2004.0m 08424, 2020.
[62] Schulman, Wolski, P. Dhariwal, A. Radford, and O. Klimov, Proximal policy optimization algorithms. arXiv preprint arXiv: 1707,06777 2017.
[63] A. Serrano-Muñoz, N. Arana-Arexolaleibb D. Chrysostomou, and S. Bøgh. skrl: Modular and flexible library for reinforcement learning. arXiv preprint arXiv:2202.0382. 2022.
[64] L. Meier, D. Honegger, and M. Pollefeys. Px4: A mode-based multithreaded open source robotics framework for deeply embedded platforms. In 2015 IEEE international conference on robotics and automation (ICRA), pages 6235-6244 6240. IEEE, 2015.
[65] D. Mellinger and V. R. Kumar. Minimum snap trajectory generation and control for quadrotors. 2011 IEEE International Conference on Robotics and Automation, pages 2520-2525, 2011.
[66] Bradbury, Frostig, P. Hawkins, M. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke, J. VanderPlas, S. Wanderman-Milne, and Q. Zhang. JAX: composable transfor- mations of Python+NumPP programs, 2018 URL http://gthhuh.com/gooog/jix
[67] D. P. Kingma and Ba. Adam: method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
[68] N. Koenig and A. Howard. Design and use paradigms for gazebo, open-source multi- robot simulator. 2004 IEEE/RSJ international conference on intelligent robots and systems (IROS)IEEE Cat. No. 04CH37566), volume 3, pages 2149-2154, IEEE, 2004.
12

---

How to Learn and Generalize From Three Minutes of Data: Physics-Constrainet and Uncertainty-Aware Neural Stochastic Differential Equations
Appendix
A Implementation and Modeling Details
We implement all the numerical experiments using the python library JAX [66], in order to take advantage of its automatic differentiation and just-in-time time compilation features. We use Python 3.8.5 for the experiments and train all our models on laptop computer with an Intel i9-9900 3.1 GHz CPU with GB of RAM and GeForce RTX 2060, TU106. We provide the code for all experiments at https://gituub.com wtt.:/cmm//w//aaee wwwp///wwwwwseee////e/eebbaaaaae
We use the Adam optimizer [67] for all optimization prob- lems. This includes when training the neural ODE model, the probabilistic ensemble model, and the system identification-based model. We use the default hyperparameters for the optimizer, except for the learning rate, which we linearly decay from 0.01 to 0.001 over the first 10000 gradient steps. We use early stopping criteria for all our experiments. We use a batch size of 512 for the neural ODE, SDE models, and the system identification-base model. Instead, we use batch size of 32 for the ensembles of probabilistic models.
Training optimizer hyperparameters.
Model design. We use the same batch size and learning rate scheduler for the optimizer, the same neural network architecture for Ty, hgs and jıc, and the same parameters upp Adaa Agrad Age- Only Ap, and max vary across experiments. Specifically, network with swish activation functions and 2 hidden layers of size 32 each while the parameters W. of he are matrices of size corresponding to the dimension of max We parametrize Pę as a feedforward neural network with tanh activation functions and hidden layers of size 8 each. For the loss function penalty terms, zero-gradient loss, and Aac = 0.01 for the strong convexity loss. We use "p = for training the neural SDE model. For the ensembles of probabilistic models,
parametrize Th as a feedforward neural
Adata 1.0 for the data loss, Agrad 0.01 for the
always use 5 models in the ensemble, where each model is Gaussian with mean and variance parametrized by feedforward neural networks with 2 hidden layers of size 64 each and silu actívation functions.
Training times of neural ODEs and SDEs. For the Cartpole swing-up example, training a neural ODE model took 2 minutes and 4 seconds on average, with standard deviation of 23 seconds over 5 training runs. In contrast, training our neural SDE model took 4 min and 30 seconds on average, with standard deviation of 42 seconds over training runs. We used the Euler integration scheme for the ODE and the Euler-Maruy Maruyama integration scheme for the SDE. While training neural SDEs, we use single particle for all the examples the paper as it reduces training times without affecting model performance.
Inference times of neural ODEs and SDEs. For the Cartpole swing-up example, we compare inference times of neural ODE and neural SDE for prediction horizon n, = 100 and by repeating the simulation over 10 different initialization. Inference times for neural ODE was 0.16 0.05ms.
Using 50 particles, the inference time of neural SDE is 1.11 - 0.1ms. For the hexacopter system, we also compare inference times of neural ODE and neural SDE for a prediction horizon n,, = 30 and by repeating the simulation over 10 different initialization. Inference times for neural ODE was 0.07 + 0.05ms. In contrast, using a single partiele (np 1), the inference time for neural SDE 0.13 + 0.06ms. Using 50 particles, the inference time of neural SDE is 0.7 + 0.15mss
13

---

B
Supplementary Mas-Spring-Danper Details and Results
(Ours) Neural SDE
Neural ODE
Ground Truth
Training Datapoint
0.2
0.2
0
-0.2
-0.2
0
2
6
2
4
8
Time (s)
Time (8)
-10
-2
Cum. Avg. Pred. Eлror
2
0.1
1.5
0
i
-0.1
0.5
0
-0.1
0
0.1
2
Time (s)
Figure 9: Prediction the neural SDE and neural ODE models over time horizon of 8 seconds for an initial condition Tinit [0.15, -0.15] outside of the training dataset. The neural SDE generalizes well beyond the training dataset while providing accurate coverage of the groundtruth trajectory and slightly improving accuracy over the neural ODE model.
The equations of motion are given by the state x = [21, x2]] [g. g] with á = d2 and má2 = -bxr kz,, where q the position of the mass, m is the mass, b is the damping coefficient, and the spring constant. We consider the case where m = 1, = 0.5, and = 1; all the quantities being the international system of units.
Data collection: Noisy and extremely scarce amount of data. We collect two dataset D1 and D2 of 5 trajectories each. The trajectories are obtained from the known dynamics with the initial positions randomly sampled in the top right quadrant [0.1, 0.1] x [0.05, 0.15] for Dj and in the more broad region [-0.1, 0.1] [-0.1, 0.1] for D22 Each trajectory has length of 5 seconds and is integrated through the Euler method with a discrete step size 0.01 second. Besides, we add a zero-mean Gaussian noise with standard deviation of [0.005, 0.01] to each state measurement.
Specifically, we use the first dataset to show that our neural SDE framework provides interpretable uncertainty estimates and better prediction accuracy than Gaussian ensemble models in the ex- tremely low and non-diverse data regime. The second dataset will be used to show that our neural SDE model improves prediction accuracy over neural ODE when the dataset is sufficiently diverse, even low data regime.
Benchmark models. We assume that the dynamics of the mass-spring are unknown, and use the proposed neural SDE framework train predictive models from the training dataset D) and Dz. Our neural SDE model has the following structure:
da [22, go (x)]dl + a max @ he (1t (z:) * dW,
where go is feedforward neural network with tanh activation functions and 2 hidden layers of size 4 and 16, respectively. The vector o max = [o - diffusion values outside of the training dataset formally defined in Section 2.2.
= [0.001, 0.02) provides the desired
We compare the neural SDE models with neural ODE and ensemble of probabilistic (Gaussian) models. The neural ODE models are trained with the same architecture as the neural SDE models
14

---

without the diffusion term. We use Euler- Maruyama and Euler methods as the SDESolve algo- rithms, respectively. with step size of 0.01 and horizon n, 50 during training. Further, use Ay = 1 for encouraging large strong convexity coefficients and ball radius = 0.05 to enforce the strong convexity property locally via sampling the neighborhood of the dataset.
Neural SDE generalizes beyond the training dataset. Figure 5 shows the state evolution of the neural SDE and neural ODE models trained on the more diverse dataset D2 and evaluated for an initial state outside the training dataset. We can observe that the SDE generalizes well beyond the training dataset and is suitable for long-term prediction. On the plot representing the evolution of q function of time,
observe that the noise at the beginning is high outside of the training dataset. We emphasize that we do not show the trained Gaussian ensemble model on this plot as it diverges over 0.3 seconds of integration time.
start from a point
15

---

C
Supplementary Experimental Results for the Cartpole Swingup Task
Benchmark models. We use feedforward neural networks with tanh activation functions and hidden layers of size 8 and 24 respectively for ga, and ge, in the non-control-affine case. In the control-affine setting, we use feedforward neural networks with tanh activation functions and hidden layers of size 8 and 24 respectively for gø, and geg. while Iø, and 9d, smaller networks with tanh activation functions and hidden layers size and respectively. The vector o maa [0.005, 0.05, 0.004, 0.01] provides the desired diffusion values beyond the dataset coverage.
We compare the neural SDE models with neural ODE and ensemble of probabilistic (Gaussian) models. The neural ODE models are trained with the same architecture as the neural SDE models without the diffusion term. We use Euler-Manun Marer-Maruyama and Euler methods as the SDESolve algo- rithms, respectively, with step size of 0.02 and a horizon
30 during training. Further, we use A = 10 for encouraging large strong convexity coefficients and ball radius r 0.2 to enforce the strong convexity property locally.
Our model
Gaussian ensemble
Dataset
Random datasst
5
0.8
0.6
15
15
0.4
10
On-policy dataset
5
0
0.2
-5
-10
0.0
-15
2
2
0
Figure 10: Uncertainty estimates of the neural SDE and Gaussian ensemble models on the on-policy and random datasets.
Neural SDE provides distance-aware uncertainty estimates. Figure 10 illustrates the uncer- tainty estimates of the neural SDE and probabilistic ensemble both the on-policy and random datasets. The neural SDE model provides low uncertainty estimates around the region of the state space covered by the training dataset and high uncertainty estimates outside of the training dataset. Instead, although the probabilistic ensemble seems to provide high uncertainty estimates far away from the dataset, it often provides high uncertainty estimates around the dataset.
Neural SDE improves prediction accuracy over the probabilistic ensemble. Figure 11 illus- trates the prediction accuracy of the neural SDE and probabilistic ensemble on trajectories of second. We observe that, even with moderate amount of data, the neural SDE provides more accurate trajectory predictions than probabilistic ensembles. This observation matches our results in the Mass- Mass-SSping-- Damper experiments, where the neural SDE also provides more accurate pre-
16

---

(Ours) Neural SDE
Probabilistic Ensemble
Ground Truth
2
0
0.5
Pr
P2
sin
0
0.5
-4
Cum. Pred. Error
cos(0) 00)
5
2
0
-0.5
0
1
5
0
0.2 0.4 0.6 0.8
0.2 0.4 0.6 0.8
0
0.2
0.4
0.6
0.8
Time (s)
Time (s)
Time
Figure 11: Comparison of the prediction accuracy of the neural SDE and probabilistic ensemble on trajectories of second.
dictions than the probabilistic ensemble in the low data regime, while the probabilistic ensemble's predictions explode even for short horizon trajectories.
Orund Ru51
Ncural sDe (Durt
Noural ODE
0.3
-10
33
0.6
-04
0.4
2 2-1
0.2
-35
0.0
-20
-10
-0.2
4.5
+04
0.4
0.7
0.0
-0.2
0.4
0.3
-3
0.5
0.2
2,4
3.0
0.0
0.2
0.4
0.0
0.5
0.5
0.2
0.4
0.6
0.8
Figure 12: Model accuracy of black-box neural SDEs and neural ODEs on the Cartpole system when predicting trajectories in the training dataset.
Neural SDEs yield models with improved robustness to data noise and generalization beyond the training dataset compared neural ODEs. Figure 12 and Figure illustrate the prediction accuracy of the neural SDE and ODE models learned on the on-policy dataset when evaluated on randomly chosen initial states and control inputs inside and outside of the dataset, respectively. By inside of the training dataset, we mean that the initial state and the control actions are picked from the dataset. Instead, by outside of the training dataset, we mean trajectories generated by picking random actions in the environment training dataset. The learned models are evaluated at initial states randomly picked either inside or outside the training dataset, and by applying the corresponding control inputs in open loop manner. The accuracy plot (bottom and far-right figure) shows the prediction error of the learned models when evaluated on 25 initial state predictions. When evaluated near the training dataset, the
opposed to the suboptimal policy employed to generate the
17

---

-0.1
-0.3
0.0
-0.5
-0.0
1.2
-2.00
-0925
-6450
-0435
11000
00
Time 151
Tome 55
Time
Figure 13: Model accuracy of black-box neural SDEs and neural ODEs on the Cartpole system when predicting trajectories outside of the training dataset.
neural SDE-based model slightly improves prediction accuracy over neural ODEs, while providing uncertainty estimates. The performance gap becomes larger when evaluating outside the training dataset. Figure 13 shows that the neural SDE model improves accuracy over the neural ODE model when evaluated on unseen data. This figure demonstrates the generalization and robustness to noise of neural SDEs compared to neural ODEs. This gap in the prediction accuracy could explain why the reinforcement learning policy trained using our neural SDE outperforms the neural ODE- Ebased policy Figure 6.
Groutd Truth
Naouall 5DC (Ount an Into.
Noural ODC Ar SDde B91.
1.0
-0.3
-044
1440
0.0
-0.6
-1.0
0.500
0.525
0.1
-0:000
-0.325
Cur
-1.000
0.
0.1
Figure 14: Model accuracy of neural SDEs and neural ODEs models with the control-affine side information when predicting trajectories outside of the training dataset.
Uncertainty estimates matter for model-based control: With side information, our neural SDE-based policies outperform neural ODE-based policies despite the two models showing similar prediction accuracy. With the control-affine dynamics side information, there is no gap prediction accuracy between the neural SDE and ODE model when evaluating on the training dataset and outside of the training dataset. This contrasts the black-box neural SDE and ODE mod- els in Figure 12 and 13, where the neural ODE model has lower prediction accuracy. However, the policy trained using our neural SDE outperforms the neural ODE-based policy as illustrated in
18

---

Figure 6. This result emphasizes the need for uncertainty estimates in offline model based reinforce- ment learning, as the neural SDE and ODE models used for training exhibit comparable predictive capabilities yet perform differently when doing control.
19

---

D Supplementary Hexacopter Details and Results
This numerical experiment shows that with basic knowledge of rigid-body dynamics as prior physics knowledge for our neural SDE, we can learn accurate and uncertainty-awat predictive models for multirotor systems from only 3 minutes of manual flight. Then, using our learned model in a non- linear model predictive control (NMPC) framework, we show incredible tracking performance on aggressive trajectories, despite how the reference trajectories push the system to operate far beyond what was seen during training. Besides, we provide additional results in Gazebo [68]. a high-fideligg platform for cettare-in-the-lore simulation of multirotors, as safe approach to compare the con- trollers based on the learned neural SDE, ODE, and system identification- models.
9w
Figure 15: Our experimental Hexacopter and the frame used for modeling its dynamics.
Experimental Setup. In hardware experiments,
custom-built hexacopter with CubePilot Orange as the flight controller running PX4 [64]. The hexacopter is equipped with 920KV brushless motors, 10 inch and two- wo-bbldded propellers, and it features the DJI F550 frame, which has a 550mm diagonal motor to motor distance. In contrast, in our the Iris quadcopter with PX4 firmware providing the simulation tools. We emphasize how our method works and integrates well with both hexacopter and quadcopter systems.
Gazebo simu-
lation,
The full state of the system, whether it is hexacopter
quadcopter, is given by [Pu, Py: Ps, Uny Uni Ua qwo qss Qy qos Way wys We], where pw = [Px, Py, Pz] is the position the world frame, fw = [Ua Uyy V-] is the velocity in the world frame, Qw = (9 Qz By, Qz] is the unit quater- nion representing the body orientation, wi [W2, Wy. W,] is the angular rate in the body frame, and the world and body frames follow respectively the traditional East-North-U Up and Forward-Left Up, shown in Figure 15. The state is estimated at a frequency of 100 Hz using the PX4 implementation of Extended Kalman Filter that fuses the measurements from the onboard IMU and Vicon motion capture system. Besides the flight controller that handles the state estimation and motor control, the main computational unit the hexacopter to compile our neural SDE models implemented in JAX, receive state estimates from the flight con- troller, solve the stochastic NMPC, and send back the motor commands and desired angular the flight controller for low-level motor control. In software in the loop simulation, we could track desired trajectories via sending motor commands directly output from our MPC. However, due to the latency during hardware experiments, we instead send desired angular rates from the NMPC to the flight controller, which then uses a PI controller track the desired angular rates.
Beelink MINIS 12. Its primary task is
to
3
4
2
2
Angle -10
0
-2
-5
Roll
Pitch
Yaw
0
-8
2
3
3.7
3
3.7
Time -100
Time -100
Figure 16: Hexacopter system: The velocity magnitude and Euler angles attained during data collection. The hexacopter mostly operates in the low speed and low Euler angles regime.
20

---

Data Collection. seek a predictive model for the hexacopter and quadcopter dynamics that can be used to autonomously track aggressive trajectories. To this end. we collect system trajectories by manually flying the hexacopter via radio-based remote controller. For the quadcopter in Gazebo, we use a Joystick controller to collect three trajectories of the system. During data collection, we store the estimated state at a frequency of 100 Hz, as well as the desired motor input commands u = [nn, u2, u3, U4, U5 us] for the hexacopter and u = [ut, U2 13, ug] for the quadcopter. Figure 16 shows the velocity magnitude and Euler angles from the collected dataset in the case of the hexacopter hardware experiments. We obtain a total of 203 seconds worth of flight data, with the 3 trajectories being respectively 73, 66, and 64 seconds long. Besides, Figure 16 shows that 95% of the collected data corresponds to the hexacopter operating below the speed of 1.7 1.71m/s, absolute roll of 18", absolute pitch of 13 and absolute yaw of 24". Instead, the maximum absolute speed, roll, pitch, and yaw attained are respectively 2.7 m/s, 23", 23", and 80".
Benchmark Models. We use the proposed neural SDE framework to train a model of the hexa- copter and quadcopter dynamics with the limited dataset described above. Our neural SDE model takes advantage of the general structure of 6-dof rigid body dynamics while having as unknown terms: The aerodynamics forces and moments, the motor command to thrust function, and (geo- metric) parameters of the system such as the mass and the inertia matrix. Specifically, our physics- informed neural SDE model is given by:
UW 1 (qwB B (To(u) + f000 (zfena)) (WB)= + 9w mg
Pw ow
d
dt + o max ② (I * dW,
QWB WB
gg BQWBWB BWB
J (Mo(u) + Mg (eata)) wB Jown
(4)
where x. feat = [ww, wn]. x denotes the cross product, @w 8 is the conjugate of gwn, the product qu between quaternion q and vector u is define the quaternion product between q and the 4-D vector [0; u). the vector o max = [1. 1, 1, 10, 10, 10. 1, 1, 1, 50, 50, 50) 10-3 is the maximum diffusion term, the variables me and Jo diag(Jğ) JJ. JJ) represent the system mass and inertia matrix, the neural network functions fø²= and M, represent the residual forces and moments due unmodelled and higher order aerodynamic effects, the parametrized functions Ta and Ma provide estimate of the motor command to thrust and moment values, and fw = [0, 0, -9.81] T is the gravity vector. Specifically, we parametrize ft Mp as feedforward neural networks with tanh activation functions and 2 hidden layers of size 8 and 16, respectively. The motor thrust forces and moments are learned via [TT] MMT"
[0,0, Tz. My My M3] Ag [T (u)).... Thmot (us) where matrix of learnable parameters constrained by the geometry of the hexacopter, and T.mor is a parametrized function that maps the motor commands to the thrust forces. We use polynomial functions for Tgmot and empirically found that a degree of 1 as Tymot (z) = agz + So works particularly well for control purpose compared to higher order polynomials. We emphasize the diffusion terms on pw and qW are low as their dynamics are known, and the noise in the estimation will come from integrating the noisy velocity and angular rate components. For training the neural SDE model, we use the derivative-free Milstein method as the SDESolve algorithm with Ase = for encouraging large strong
Ag
step size of 0.05 and a horizon n, = 20. Further, convexity coefficients and a vector of ball radius r = 0.1 to enforce the strong convexity property locally. The quadcopter modeling is the same as for the hexacopter, with the thrust forces and moments being a function of the 4 control inputs instead of 6 for the hexacopter.
To illustrate the prediction accuracy of our model, we compare it with a system identification--art approach that uses the same formulation as our SDE model (4) but without the diffusion term and the residual neural network terms fi and Mgss. Precisely, with system identification, we seek to identify all the parameters mg. Jo. Agrix, Cgg, and By by estimating ă from the dataset using finite difference and then solving least square problem to fit the system's 's differential equation to the data. This is a simple approach for modeling multirotors that has been widely used in the literature. Additionally, as another baseline, we train a neural ODE model with the same architecture as the neural SDE model but without the diffusion term.
Nonlinear Model Predictive Control. We seek to use our learned SDE for autonomous control of the hexacopter. To this end, we employ a receding horizon model predictive control approach.
21

---

Such an approach uses the learned SDE model to predict future system trajectories over a fixed time horizon, and then optimize the control imputs to minimize cost function that penalizes the control effort and deviation from reference trajectory. For numerical optimization, we discretize the state and control inputs into do equal time intervals over the horizon H second, yielding constrained optimization problem of the following form at each state measurement Tinit Zz.
minimize Mi
IE
Σ (ar
T ref
(Q( xx + ur Rus
(5)
k=1
subject to { ........
xh fp=1 SDESolve(RO, u; (4)),
(6)
50 Finit: U1,.... Un, [0, 1],
where oxt is the reference state at time t = t + kH//n,, we use qra to denote the term (q(g"=)))))) 2zy,, the positive definite matrices Q and R penalize the deviation from the reference trajectory and the control effort, respectively, and пp is the number of particles used for the SDE solver. In all our experiments, we used np = 1. Q diiag(100, 100, 200, 5, 10, 1, 1, 100, 1, 1, 1), 1). Besides, we solve the above optimization problem using an adaptive
and R diag(1, 1, learning rate, Nesterov acceleration-based projected gradient descent; all implemented in JAX.
(Ours) Neural SDE
Nonlincar System Identification
Ground Truth
Px
Py
Pz
M
2
2
0.5
0p
2/y
0
0
-2
-5
R
P
Y
100
50
100
50
-100
50
0
0.5
2
-2
-1
-0.5
4
-2
50
55
60
50
55 Time
60
50
35
60
Time
Time
Figure 17:
Hexacopter system: Time evolution of the hexacopter state predicted by the learned SDE model and the vanilla system identification-baser model. The quantities R, P, Y denote the roll, pitch, and yaw angles, respectively, with units in degrees.
Hexacopter modeling: Neural SDE improves prediction accuracy over system identification while providing uncertainty estimates. Figure 17 a complete version of Figure 7, where the prediction of the remaining states of the system is also provided. The neural SDE model provides
22

---

uncertainty estimates while also providing better long- long-horizon prediction capabilities compared to the system identification-but -aasd model.
RMSE -10
Reference NMPC+SDE
1.0
0.5
0
0.0
1
0.5
3
1.0
1.5
Time
10
1.300
10
T 10
1.275
-2
-2
1.250
VV
W4
-4
20
20
Time
Time
Time
Time
Figure 18: Hexacopter system: Tracking performance achieved on the circle trajectory. The MPC based on our learned SDE achieves an RMSE of 20 em for 35 seconds trajectory while the hexa- copter reaches speed up to 3.6 m/s, roll and pitch angles up 32", and yaw angle up 120".
Hexacopter control: Neural SDE yields high-performance MPC algorithms even when oper- ating far beyond what it has been trained on. Figure 18 and Figure 19 show the results of our experiments on a circle trajectory and lemniscate trajectory, respectively. These reference trajecto- ries are obtained by minimum snap trajectory generation [65] without any prior knowledge of the hexacopter dynamics. For both trajectories, we show the time evolution of the velocity, roll, pitch, yaw, and tracking accuracy during autonomous control. These plots demonstrate the ability of our learned SDE to generalize far beyond what has been trained on. In fact, we can observe that the hexacopter must reach velocity up to 3.6 m/s, roll and pitch angles up to 32", and yaw angle up to 120", order to track the reference trajectories. of the training data regime, learned SDE is able to generalize to these extreme conditions and achieves a high tracking accuracy of 20 and 15 em for the circle and lemniscate trajectories, respectively. We also note that the tracking accuracy is not uniform across the trajectory, and it is lower when the hexacopter is moving faster. This is expected as the hexacopter is sides, the performance of our control approach is further displayed on the plot showing the evolution of the altitude, where altitude error of less than 5 cm is achieved consistently.
emphasize how these values are outside shown and detailed in the data collection section. Despite this, the
difficult to control when it is moving faster. Be-
Quadcopter Gazebo simulation: Neural SDEs improve model accuracy and model predictive control performance compared to neural ODEs and system identification baselines. Figure 20 shows the predictive capabilities of the neural SDE and baseline models on the trajectories col- lected from the Gazebo simulation of the Quadcopter. The plot shows better approximation of system identification-basel models due to less noisy dataset than the real-world hexacopter collected dataset. We observe that the neural ODE improves prediction accuracy over the system identification-base techniques but has slightly lower accuracy than the neural SDE model, specifi- cally for extended integration steps.
23

---

RMSE -10 -2
15
1.0
Reference NMPC+$DE
0.5
-1.0
30
Time
1.650
1.625
10
1-110
-2
-4
20
40
20
40
Time
Time
Time
Time
Figure 19: Hexacopter system: Tracking performance achieved on the lemniscate trajectory. The MPC based on our learned SDE achieves an RMSE of 15 em for 40 seconds trajectory while the hexacopter reaches speed up to 3.4 m/s, roll angle up to 32 pitch up to 19", and yaw up to 13".
Ground Tuth
Neurs! SDE
Sys
Neural ODE
1.4
E
1.0
0.8
0.5
3.0
0.5
(
(matt))
0.0
-0.5
20
331
40
[(s]
Figure 20: Quadrotor Gazebo simulation: Comparison of the prediction accuracy of neural SDE,
ODE, and system identification- -based model.
Moreover, Figure 21 and Figure show the results using the different learned models for predictive control in Gazebo simulation. We observe that the neural SDE-based controller more accurately tracks the specified trajectories than the SysID-based controller. The gap between their tracking accuracies appears to be largest when the trajectories require the highest velocities and the largest

---

angular rates. This observation indicates that the neural SDE has relatively high predictive accuracy, particularly in regions of the state space that are beyond the training dataset. Besides, observe slight improvement in control performance between the neural SDE and neural ODE model.
Sys
Noural
COE
2.0
0.20
2.3
15
RMGSE (mmi2)
HMSE
0.2
0.5
0.05
20
20
-1.0
1.0
X[m]
0.6
2.0
1.32
0.4
0.4
1.28
(nell
0.2
0.0
1.26
-0.2
0.5
1.2
-0.4
20
20
20
[14]]
[(G]]
Figure 21: Quadrotor Gazebo simulation: Tracking performance on the Circle trajectory.
Releewnce
Neurial SDD
Neural DOF
0.30
1.0
0.20
RVSSE
0.5
0.10
0.05
100
20
30
20
0.0
3.5
-0.5
2.0
1.5
-1.0
0.5
10
30
x[mi]
1.32
M
1.2
0.0
174
0.2
-0.2
-0.4
1.20
0..5
10
20
115]
tEsl
Figure 22: Quadrotor Gazebo simulation: Tracking performance the lemniscate trajectory.
25


Reviews:
[{"review_id": "qvHmtg2C-k", "reviewer": "Reviewer_5AP6", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "This paper proposes to use Neural Stochastic Differential Equations  (NSDEs), i.e. a differentiable SDE integrator, to learn dynamics models from data. A main motivation for this is the ability of NSDEs to be distance-aware. Specifically, the model is uncertain about states that are deemed far from the data distribution. The drift and diffusion terms of the NSDE can parameterised as trainable models. ", "summary_of_recommendation": "Please see the review above.", "technical_quality": "Good: Generally convincing, but some additional experiments or analyses needed.", "clarity_of_presentation": "Very Good:  Well organized and clear.", "potential_impact": "", "robotics_focus": "Sufficient demonstration on hardware", "quality_of_limitations": "Additional details required", "questions_for_rebuttal": "- Justification of NSDEs over NODEs, and some clarity over the computation (and also training stability) trade-offs between using an NSDE over a NODE. Please list out inference and training times wrt to NSDEs and NODEs, and ideally discuss an example where NSDEs show marked improvement over NODEs. NODEs have not been compared on the use-case problems, so it is unclear to the reader the practical benefits of NSDE over NODE.\n\n- Comparisons and/or discussions against previous Gaussian Process-based approaches to similar problems. ", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Overall, I find the main idea of the paper to be clear -- SDEs have been shown to be capable of modeling epistemic uncertainty (Kong et al, ICML 2020), and so it is natural to attempt to learn unknown dynamics, with additional requirements of uncertainty modeling. Neural ODEs (NODEs) and SDEs also allow for easy incorporation of prior knowledge, by having a mixture of trainable apriori specified terms to represent terms in the ODE or SDE. My questions and concerns are as follows:\n\n- The empirical evaluation is critically underwhelming. It is unclear whether the benefits of the uncertainty estimates provided by the NSDE are of any benefit, over a NODE with its dynamics parameterised in the same way as the drift term of the SDE. NSDEs are much more unstable to train and slower to roll out than NODEs -- *is the heavy machinery worth it, wrt to NODes?* Would love to see examples where the NODE fails and some discussion of this.\n\n- Also critically lacking is the comparison against Gaussian Process-based approaches. The uncertainty estimates for NSDEs were visually compared against ensemble-based approaches to extract uncertainty. Gaussian processes have been well studied and are the gold standard for epistemic uncertainty quantification. Additionally, GPs have been used within an MPC setting for race cars (Hewing, Kabzan, and Zeilinger) and drones (Liu, Shi, Chung, Anandkumar, Yue, L4DC 2020), to learn unknown parameters (with additional uncertainty estimates available) within a controller. It would be good to explore the trade-offs between GPs and NSDEs -- why NSDEs over GPs?\n\nOverall the paper is fairly well-written, and the method appears sounds -- I am happy to increase my score after these concerns have been made.", "extra_scores": {"impact": "3: The work is incremental and unlikely to have much impact, even though it may be technically correct and well executed."}, "extra_text": {}}, {"review_id": "fGldqfrkQA", "reviewer": "Reviewer_DHcC", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "The paper presents a method for learning dynamics models using physics constrained neural stochastic differential equations. The approach demonstrates that one can learn meaningful uncertainty models in low-data regimes while accurately predicting drone dynamics with just three minutes of data and physics priors. The physical robot results are convincing and demonstrate the effectiveness of the proposed approach.", "summary_of_recommendation": "The work demonstrates a novel approach to model-based control using a neural stochastic differential equation. The results are impressive and demonstrate clear applicability on a physical system. The comments mentioned earlier should be addressed to highlight the limitations of the proposed work, but the approach appears novel and exciting. Based on the rebuttal and responses, an accept is recommended.", "technical_quality": "Excellent: Very convincing and thorough.", "clarity_of_presentation": "Very Good:  Well organized and clear.", "potential_impact": "", "robotics_focus": "Sufficient demonstration on hardware", "quality_of_limitations": "Additional details required", "questions_for_rebuttal": "Is the proposed approach amenable to other metric choices?\n\nAt what limit does the approach fail to capture the true dynamics (is this a function of data or metric choices)?", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper presents a novel approach to model-based learning and control using a neural stochastic differential equation. The key insight of the work is that effective modeling of low-data regimes is key to have conservative behavior when unsure and confident, aggressive maneuvers high-data regimes. \n\nA strength of the work is the empirical evidence tied to the theoretical justification for the proposed work. The sufficient ablation studies compare with existing methods for model learning and demonstrate clear improvement. \n\nOne weakness that the paper addresses is the single-use case with the drone. Though it is possible to extend this relatively quickly. In addition, the paper does assume full observability which is highlighted in the limitations. \n\nAnother, perhaps more critical weakness of the work is the choice of metric for the uncertainty. The paper assumes a Euclidean metric is sufficient to model distance in the model variance. However, Euler angles are not Euclidean and do not exhibit the properties of a vector space. This can lead to issues with large angles and misrepresented uncertainty values. It is critical that the paper address this assumption in the limitations as this could change the problem structure all together.", "extra_scores": {"impact": "4: The work contains interesting new ideas that will potentially have major impact in robotics or machine learning."}, "extra_text": {}}, {"review_id": "cEZmgrGHUh", "reviewer": "Reviewer_S722", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "The paper proposes to use neural SDEs to model robot dynamics that consist of an informed drift term (governing equations, rigid body dynamics and learned dynamics parameters) and  diffusion term that captures the uncertainty of the dynamics. The main idea behind the diffusion term is that uncertainty is greater when prediction come farther from the data. The term is learned by enforcing local minima in regions of the data and strong convexity of the within a ball around the data points.", "summary_of_recommendation": "The main thing that is bothering me (why it is a weak reject) is the way that the predictive trajectory distribution is treated (see the rebuttal questions) + secondary things.", "technical_quality": "Good: Generally convincing, but some additional experiments or analyses needed.", "clarity_of_presentation": "Good: Minor revision needed to improve organization or clarity.", "potential_impact": "", "robotics_focus": "Sufficient demonstration on hardware", "quality_of_limitations": "Limitations are addressed clearly", "questions_for_rebuttal": "you should cite work from [1],  Indeed, your reference to PETS considers only sampling, but doesn't separate uncertainties and doesn't take them into account, whereas Risk Averse Trajectory Optimization that builds on top of it does.\n\neq. 2 - the likelihood of the trajectory is indeed not Gaussian, this is not the correct likelihood given the Neural SDE, therefore it's not an NLL loss given oyour probabilistic model. Likelihood computation can be done by probability flow ODE + instant change of variables. Or am I missing something here? This is very critical.\n\nline 138 - what do you mean by elementwise maximum values for the $\\Sigma_{max}$? Is this just the maximum over the concatenated vector [x,u]?\n\nparagraph line 139 - this indeed enforces local minima, but it doesn't enforce that no additional uncertainty is introduced when you are on-data?\n\nfigure 4 - nice figure\n\nfigure 5 - why is there no error band for the prediction error here? Also, how is the prediction error computed for the Neural SDE? Since it is a sampling method, do you make an MC estimate of the expected error? Over how many samples? The evaluation metrics are not very clear to me.\n\nfigure 6, PPO experiment - why didn't you report results for the neural ODE case? Except being curious about them, I believe that you need them for completeness.\n\n[1] Vlastelica, Marin, Sebastian Blaes, Cristina Pinneri, and Georg Martius. \"Risk-averse zero-order trajectory optimization.\" In 5th Annual Conference on Robot Learning. 2021.", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "## clarity\n\nThe paper is mostly clear, up to a few things that I think need to be elaborated (see the questions). There is a good amount of experiments that support the claim that the paper is making, that this Neural SDE formulation is good for accounting for uncertainty and for being used as a model.\n\n## strengths\n\nThe work is novel to the extent of my knowledge, providing an inductive bias to the drift term in order to improve dynamics modelling is well-motivated, also the theoretical framework for dealing with the problem of uncertainty estimation is correct.\n\nExperimental evaluation validates the claims of the paper.\n\nThere are results on real hardware.\n\n## weaknesses\n\nclarity can be improved regarding the computation of the diffusion term and the experiments.\n\nthe diffusion coefficient seems overly engineered\n\nsome claims in the paper might be wrong (see comments)", "extra_scores": {"impact": "3: The work is incremental and unlikely to have much impact, even though it may be technically correct and well executed."}, "extra_text": {}}, {"review_id": "G8biWIAkaTZ", "reviewer": "Reviewer_BxRN", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "In this paper the authors propose a structured dynamics learning methodology that fits neural stochastic differential equations (SDEs) to observed trajectory data. The SDE consists of a drift term and diffusion term capturing the deterministic and stochastic components of the dynamics respectively; the authors propose that the design of the drift term can be physics-based with unmodeled/error terms parameterized as neural networks, and that the diffusion term may capture both the aleatoric and epistemic uncertainty (the latter estimated using distance to training data). The resulting learned dynamics may be employed in model-based RL or model predictive control algorithms to provide safe, uncertainty-aware policies even in the absence of full data coverage. Experiments on simple simulated systems as well as in hardware on a hexacopter system tracking aggressive trajectories demonstrate the improved performance over baselines and data efficiency of the approach.", "summary_of_recommendation": "This paper achieves its stated goal of demonstrating the capability of neural SDEs to capture model uncertainty while incorporating strong inductive bias; these aims are of broad relevance to the robot learning research community and therefore I recommend this paper's acceptance at CORL 2023. The experiments provide convincing evidence for this method's potential, though I believe that from a practitioner's perspective additional evaluation may be needed to show that this approach should be selected over other state-of-the-art model-based methods.", "technical_quality": "Very Good:  Convincing.", "clarity_of_presentation": "Very Good:  Well organized and clear.", "potential_impact": "", "robotics_focus": "Sufficient demonstration on hardware", "quality_of_limitations": "Limitations are addressed clearly", "questions_for_rebuttal": "From the comments/questions above, please prioritize:\n- justification of the loss formulation Eq. (2), and\n- describe the performance (at least in simulation) of the SysID baseline when deployed in closed-loop MPC.", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The experimental results are undeniably impressive and the paper is well-written with clearly described insights derived from each of the experiments. I am pleased to recommend its acceptance; please find below a list of comments/questions that if addressed may help readers better appreciate the impact of this work.\n- *SDE training loss:* The loss (2) sums up the log-likelihoods of the data with respect to each fixed-variance Gaussian centered around integrated SDE paths. Though it's noted that this is a \"common choice in the literature [33, 30]\" it seems that this treatment of likelihood may misrepresent multimodal data distributions (i.e., in such a case I believe the SDE would learn to predict the mean -- indeed, I'm not sure how (2) encourages the model to capture any aleatoric uncertainty), though notably none of the systems considered in this work exhibit multimodal uncertainty. That is, I had initially expected that (2) would essentially be using a set of particles sampled from the SDE dynamics (i.e., $\\text{SDESolve}$) to produce a kernel density estimate of the SDE's trajectory distribution pdf which would then be used to evaluate the likelihood of the observed data (when implemented this would be not unlike (2), just with a logsumexp). In any case some additional explanation/justification would be beneficial for readers to understand why this loss formulation is chosen.\n- *Drift term architecture:*\n  - The interpretation of $\\eta_\\psi$ and $h_\\phi$ is not completely clear to me -- if $\\eta_\\psi$ is intended to encode distance to the training dataset why does it need to saturate at 1 (especially when $h_\\phi$, which also saturates at 1, immediately applies a $\\text{sigmoid}^{-1}$ to the output of $\\eta_\\psi$)?\n  - Is there an advantage to representing aleatoric and epistemic uncertainty jointly? That is, one might imagine decomposing the total uncertainty into a sum of two terms: (i) a non-negative aleatoric uncertainty term and (ii) a distance-based epistemic uncertainty term that is encouraged to take zero values on points in the training dataset (compared to zero-gradient, as for $\\eta_\\psi$, which requires the strong convexity loss to avoid saddle points).\n- *Hexacopter MPC comparison:*\n  - In SysID what degree polynomial basis is used to fit the nonlinear unknown functions (mentioned on line 322)? I believe that with appropriate constants for $m$ and $J$, linear gains defining $T_\\theta$ and $M_\\theta$, and assuming $f^{res} = M^{res} = 0$ roboticists have historically been able to implement multicopter tracking MPC controllers. Notably only open loop predictions are compared in Figure 7 (where perhaps the erratic SysID predictions may be attributed to Runge's phenomenon with high-order polynomial fitting). Is 3 minutes also enough for a traditional SysID + MPC approach to track trajectories (though perhaps with greater error)?\n  - Though neural SDEs are topical today, it is not clear that this method provides benefit over other state-of-the-art approaches, e.g., [R1] which uses Gaussian processes that incorporate a priori physics knowledge and provide estimated uncertainty, similar to the motivations of this work.\n- *Stochastic NMPC implementation:* Though more detail is provided in the Appendix, I believe that some concept of how the stochastic NMPC is implemented should be stated in the main text (i.e., that a control sequence is chosen to minimize expected cost over SDE predicted trajectories, estimated empirically using samples without explicitly taking any uncertainty estimates (e.g., diffusion term magnitude) into account).\n\n[R1] G. Torrente, E. Kaufmann, P. Fohn, and D. Scaramuzza, \"Data-Driven MPC for Quadrotors,\" *IEEE Robotics and Automation Letters*, 2021. Available at https://arxiv.org/abs/2102.05773.pdf.", "extra_scores": {"impact": "4: The work contains interesting new ideas that will potentially have major impact in robotics or machine learning."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper presents a technique to learn the system dynamics and associated uncertainty using neural stochastic differential equations where one term models the drift and the other captures the diffusion. This makes the model able to capture the epistemic uncertainty when inferring points far from the data distribution. Experiments were performed on simple simulated systems as well as on a real hexacopter.

Reviewers were positive about the novelty of the work but were concerned about the experimental section and lack of comparisons to equivalent formulations such as GP-based dynamical systems. Authors have mostly addressed these concerns and have elaborated on the limitations of the work in the revised version.
Decision: Accept (Oral)