Learning to Control PDEs with Differentiable Physics

1 INTRODUCTION . Intelligent systems that operate in the physical world must be able to perceive , predict , and interact with physical phenomena ( Battaglia et al. , 2013 ) . In this work , we consider physical systems that can be characterized by partial differential equations ( PDEs ) . PDEs constitute the most fundamental description of evolving systems and are used to describe every physical theory , from quantum mechanics and general relativity to turbulent flows ( Courant & Hilbert , 1962 ; Smith , 1985 ) . We aim to endow artificial intelligent agents with the ability to direct the evolution of such systems via continuous controls . Such optimal control problems have typically been addressed via iterative optimization . Differentiable solvers and the adjoint method enable efficient optimization of high-dimensional systems ( Toussaint et al. , 2018 ; de Avila Belbute-Peres et al. , 2018 ; Schenck & Fox , 2018 ) . However , direct optimization through gradient descent ( single shooting ) at test time is resource-intensive and may be difficult to deploy in interactive settings . More advanced methods exist , such as multiple shooting and collocation , but they commonly rely on modeling assumptions that limit their applicability , and still require computationally intensive iterative optimization at test time . Iterative optimization methods are expensive because they have to start optimizing from scratch and typically require a large number of iterations to reach an optimum . In many real-world control problems , however , agents have to repeatedly make decisions in specialized environments , and reaction times are limited to a fraction of a second . This motivates the use of data-driven models such as deep neural networks , which combine short inference times with the capacity to build an internal representation of the environment . We present a novel deep learning approach that can learn to represent solution manifolds for a given physical environment , and is orders of magnitude faster than iterative optimization techniques . The core of our method is a hierarchical predictor-corrector scheme that temporally divides the problem into easier subproblems . This enables us to combine models specialized to different time scales in order to control long sequences of complex high-dimensional systems . We train our models using a differentiable PDE solver that can provide the agent with feedback of how interactions at any point in time affect the outcome . Our models learn to represent manifolds containing a large number of solutions , and can thereby avoid local minima that can trap classic optimization techniques . We evaluate our method on a variety of control tasks in systems governed by advection-diffusion PDEs such as the Navier-Stokes equations . We quantitatively evaluate the resulting sequences on how well they approximate the target state and how much force was exerted on the physical system . Our method yields stable control for significantly longer time spans than alternative approaches . 2 BACKGROUND . Physical problems commonly involve nonlinear PDEs , often with many degrees of freedom . In this context , several works have proposed methods for improving the solution of PDE problems ( Long et al. , 2018 ; Bar-Sinai et al. , 2019 ; Hsieh et al. , 2019 ) or used PDE formulations for unsupervised optimization ( Raissi et al. , 2018 ) . Lagrangian fluid simulation has been tackled with regression forests ( Ladicky et al. , 2015 ) , graph neural networks ( Mrowca et al. , 2018 ; Li et al. , 2019 ) , and continuous convolutions ( Ummenhofer et al. , 2020 ) . Data-driven turbulence models were trained with MLPs ( Ling et al. , 2016 ) . Fully-convolutional networks were trained for pressure inference ( Tompson et al. , 2017 ) and advection components were used in adversarial settings ( Xie et al. , 2018 ) . Temporal updates in reduced spaces were learned via the Koopman operator ( Morton et al. , 2018 ) . In a related area , deep networks have been used to predict chemical properties and the outcome of chemical reactions ( Gilmer et al. , 2017 ; Bradshaw et al. , 2019 ) . Differentiable solvers have been shown to be useful in a variety of settings . Degrave et al . ( 2019 ) and de Avila Belbute-Peres et al . ( 2018 ) developed differentiable simulators for rigid body mechanics . ( See Popovic et al . ( 2000 ) for earlier work in computer graphics . ) Toussaint et al . ( 2018 ) applied related techniques to manipulation planning . Specialized solvers were developed to infer protein structures ( Ingraham et al. , 2019 ) , interact with liquids ( Schenck & Fox , 2018 ) , control soft robots ( Hu et al. , 2019 ) , and solve inverse problems that involve cloth ( Liang et al. , 2019 ) . Like ours , these works typically leverage the automatic differentiation of deep learning pipelines ( Griewank & Walther , 2008 ; Maclaurin et al. , 2015 ; Amos & Kolter , 2017 ; Mensch & Blondel , 2018 ; van Merriënboer et al. , 2018 ; Chen et al. , 2018 ; Bradbury et al. , 2018 ; Paszke et al. , 2019 ; Tokui et al. , 2019 ) . However , while the works above target Lagrangian solvers , i.e . reference frames moving with the simulated material , we address grid-based solvers , which are particularly appropriate for dense , volumetric phenomena . The adjoint method ( Lions , 1971 ; Pironneau , 1974 ; Jameson , 1988 ; Giles & Pierce , 2000 ; Bewley , 2001 ; McNamara et al. , 2004 ) is used by most machine learning frameworks , where it is commonly known as reverse mode differentiation ( Werbos , 2006 ; Chen et al. , 2018 ) . While a variety of specialized adjoint solvers exist ( Griewank et al. , 1996 ; Fournier et al. , 2012 ; Farrell et al. , 2013 ) , these packages do not interface with production machine learning frameworks . A supporting contribution of our work is a differentiable PDE solver called ΦFlow that integrates with TensorFlow ( Abadi et al. , 2016 ) and PyTorch ( Paszke et al. , 2019 ) . It is publicly available at https : //github.com/tumpbs/PhiFlow . 3 PROBLEM . Consider a physical system u ( x , t ) whose natural evolution is described by the PDE ∂u ∂t = P ( u , ∂u ∂x , ∂2u ∂x2 , ... , y ( t ) ) , ( 1 ) where P models the physical behavior of the system and y ( t ) denotes external factors that can influence the system . We now introduce an agent that can interact with the system by controlling certain parameters of the dynamics . This could be the rotation of a motor or fine-grained control over a field . We factor out this influence into a force term F , yielding ∂u ∂t = P ( u , ∂u ∂x , ∂2u ∂x2 , ... ) + F ( t ) . ( 2 ) The agent can now be modelled as a function that computes F ( t ) . As solutions of nonlinear PDEs were shown to yield low-dimensional manifolds ( Foias et al. , 1988 ; Titi , 1990 ) , we target solution manifolds of F ( t ) for a given choice of P with suitable boundary conditions . This motivates our choice to employ deep networks for our agents . In most real-world scenarios , it is not possible to observe the full state of a physical system . When considering a cloud of smoke , for example , the smoke density may be observable while the velocity field may not be seen directly . We model the imperfect information by defining the observable state of u as o ( u ) . The observable state is problem dependent , and our agent is conditioned only on these observations , i.e . it does not have access to the full state u . Using the above notation , we define the control task as follows . An initial observable state o0 of the PDE as well as a target state o∗ are given ( Figure 1a ) . We are interested in a reconstructed trajectory u ( t ) that matches these states at t0 and t∗ , i.e . o0 = o ( u ( t0 ) ) , o∗ = o ( u ( t∗ ) ) , and minimizes the amount of force applied within the simulation domain D ( Figure 1b ) : LF [ u ( t ) ] = ∫ t∗ t0 ∫ D |Fu ( t ) |2 dx dt . ( 3 ) Taking discrete time steps ∆t , the reconstructed trajectory u is a sequence of n = ( t∗ − t0 ) /∆t states . When an observable dimension can not be controlled directly , there may not exist any trajectory u ( t ) that matches both o0 and o∗ . This can stem from either physical constraints or numerical limitations . In these cases , we settle for an approximation of o∗ . To measure the quality of the approximation of the target , we define an observation loss L∗o . The form of this loss can be chosen to fit the problem . We combine Eq . 3 and the observation loss into the objective function L [ u ( t ) ] = α · LF [ u ( t ) ] + L∗o ( u ( t∗ ) ) , ( 4 ) with α > 0 . We use square brackets to denote functionals , i.e . functions depending on fields or series rather than single values . 4 PRELIMINARIES . Differentiable solvers . Let u ( x , t ) be described by a PDE as in Eq . 1 . A regular solver can move the system forward in time via Euler steps : u ( ti+1 ) = Solver [ u ( ti ) , y ( ti ) ] = u ( ti ) + ∆t · P ( u ( ti ) , ... , y ( ti ) ) . ( 5 ) Each step moves the system forward by a time increment ∆t . Repeated execution produces a trajectory u ( t ) that approximates a solution to the PDE . This functionality for time advancement by itself is not well-suited to solve optimization problems , since gradients can only be approximated by finite differencing . For high-dimensional or continuous systems , this method becomes computationally expensive because a full trajectory needs to be computed for each optimizable parameter . Differentiable solvers resolve this issue by solving the adjoint problem ( Pontryagin , 1962 ) via analytic derivatives . The adjoint problem computes the same mathematical expressions while working with lower-dimensional vectors . A differentiable solver can efficiently compute the derivatives with respect to any of its inputs , i.e . ∂u ( ti+1 ) /∂u ( ti ) and ∂u ( ti+1 ) /∂y ( ti ) . This allows for gradientbased optimization of inputs or control parameters over an arbitrary number of time steps . Iterative trajectory optimization . Many techniques exist that try to find optimal trajectories by starting with an initial guess for F ( t ) and slightly changing it until reaching an optimum . The simplest of these is known as single shooting . In one optimization step , it simulates the full dynamics , then backpropagates the loss through the whole sequence to optimize the controls ( Kraft , 1985 ; Leineweber et al. , 2003 ) . Replacing F ( t ) with an agent F ( t|ot , o∗ ) , which can be parameterized by a deep network , yields a simple training method . For a sequence of n frames , this setup contains n linked copies of the agent and is depicted in Figure 2 . We refer to such an agent as a control force estimator ( CFE ) . Optimizing such a chain of CFEs is both computationally expensive and causes gradients to pass through a potentially long sequence of highly nonlinear simulation steps . When the reconstruction u is close to an optimal trajectory , this is not a problem because the gradients ∆u are small and the operations executed by the solver are differentiable by construction . The solver can therefore be locally approximated by a first-order polynomial and the gradients can be safely backpropagated . For large ∆u , e.g . at the beginning of an optimization , this approximation breaks down , causing the gradients to become unstable while passing through the chain . This instability in the training process can prevent single-shooting approaches from converging and deep networks from learning unless they are initialized near an optimum . Alternatives to single shooting exist , promising better and more efficient convergence . Multiple shooting ( Bock & Plitt , 1984 ) splits the trajectory into segments with additional defect constraints . Depending on the physical system , this method may have to be adjusted for specific problems ( Treuille et al. , 2003 ) . Collocation schemes ( Hargraves & Paris , 1987 ) model trajectories with splines . While this works well for particle trajectories , it is poorly suited for Eulerian solvers where the evolution of individual points does not reflect the overall motion . Model reduction can be used to reduce the dimensionality or nonlinearity of the problem , but generally requires domain-specific knowledge . When applicable , these methods can converge faster or in a more stable manner than single shooting . However , as we are focusing on a general optimization scheme in this work , we will use single shooting and its variants as baseline comparisons . Supervised and differentiable physics losses . One of the key ingredients in training a machine learning model is the choice of loss function . For many tasks , supervised losses are used , i.e . losses that directly compare the output of the model for a specific input with the desired ground truth . While supervised losses can be employed for trajectory optimization , far better loss functions are possible when a differentiable solver is available . We will refer to these as differentiable physics loss functions . In this work , we employ a combination of supervised and differentiable physics losses , as both come with advantages and disadvantages . One key limitation of supervised losses is that they can only measure the error of a single time step . Therefore , an agent can not get any measure of how its output would influence future time steps . Another problem arises from the form of supervised training data which comprises input-output pairs , which may be obtained directly from data generation or through iterative optimization . Since optimal control problems are generally not unimodal , there can exist multiple possible outputs for one input . This ambiguity in the supervised training process will lead to suboptimal predictions as the network will try to find a compromise between all possible outputs instead of picking one of them . Differentiable physics losses solve these problems by allowing the agent to be directly optimized for the desired objective ( Eq . 4 ) . Unlike supervised losses , differentiable physics losses require a differentiable solver to backpropagate the gradients through the simulation . Multiple time steps can be chained together , which is a key requirement since the objective ( Eq . 4 ) explicitly depends on all time steps through LF [ u ( t ) ] ( Eq . 3 ) . As with iterative solvers , one optimization step for a sequence of n frames then invokes the agent n times before computing the loss , each invocation followed by a solver step . The employed differentiable solver backpropagates the gradients through the whole sequence , which gives the model feedback on ( i ) how its decisions change the future trajectory and ( ii ) how to handle states as input that were reached because of its previous decisions . Since no ground truth needs to be provided , multi-modal problems naturally converge towards one solution .

Summary: In this paper, the authors outline a method for system control utilizing an "agent" formed by two neural networks and utilizing a differentiable grid-based PDE solver (assuming the PDE describing the system is known). The agent is split into a control force estimator (CFE) which applies a force to advance the state of the controlled system, and an observation predictor (OP) which predicts the trajectory needed to reach the target state. The objective is to reach the target state with minimal total amount of applied force. The order of CFE and OP calls is discussed and the importance of keeping the trajectory predictions conditioned on the actual previous state of the system, so that errors from previous steps can be taken into account.