How to Learn and Generalize From Three Minutes of Data: Physics-Constrained and Uncertainty-Aware Neural Stochastic Differential Equations

TL;DR

This paper introduces a physics-informed neural stochastic differential equation framework that learns accurate, uncertainty-aware dynamic models from minimal data, enabling effective long-term predictions and control in robotic systems.

Contribution

It develops a novel neural SDE approach incorporating physics-based priors and uncertainty estimation, suitable for control and simulation with limited data.

Findings

Neural SDE models accurately predict long-term dynamics.

Models trained on 3 minutes of data enable effective control of a hexacopter.

Uncertainty estimates match system stochasticity near training data.

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.