Learning to Control PDEs with Differentiable Physics

Philipp Holl; Vladlen Koltun; Nils Thuerey

Source-linked AI summary

Learning to Control PDEs with Differentiable Physics

Philipp Holl, Vladlen Koltun, Nils Thuerey

arXiv:2001.07457v1Jan 21, 2020cs.LGphysics.flu-dynstat.ML

TL;DR

Controlling high-dimensional nonlinear PDE systems is difficult because existing approaches are short-horizon or expensive for interactive use. This paper uses a hierarchical predictor-corrector architecture with a trajectory-planning predictor and a control network, trained end-to-end through a differentiable PDE solver. It achieves stable long-horizon control and, in a smoke-transfer task, moves 99.22% ± 0.15% of smoke while using 19.1% ± 1.0% less force than the baseline.

Problem
Existing PDE-control methods are typically limited to short time frames or few interaction parameters, while iterative optimization is costly for time-critical decisions.
Method
A hierarchical predictor-corrector scheme separates trajectory planning from control and trains both stages end-to-end with a differentiable PDE solver.
Results
The method successfully predicts and controls complex PDE systems, yielding stable control for significantly longer time spans than alternative approaches.
Takeaways & Limitations
Hierarchical inference lets the models plan ahead, while fast learned solution-manifold representations can complement classical solvers in time-critical applications.
Takeaways & Limitations
The authors do not expect learning approaches to replace iterative optimization, and the predictor can produce blurred intermediate predictions when it cannot account for prior deviations.

Abstract

from arXiv · show

Predicting outcomes and planning interactions with the physical world are long-standing goals for machine learning. A variety of such tasks involves continuous physical systems, which can be described by partial differential equations (PDEs) with many degrees of freedom. Existing methods that aim to control the dynamics of such systems are typically limited to relatively short time frames or a small number of interaction parameters. We present a novel hierarchical predictor-corrector scheme which enables neural networks to learn to understand and control complex nonlinear physical systems over long time frames. We propose to split the problem into two distinct tasks: planning and control. To this end, we introduce a predictor network that plans optimal trajectories and a control network that infers the corresponding control parameters. Both stages are trained end-to-end using a differentiable PDE solver. We demonstrate that our method successfully develops an understanding of complex physical systems and learns to control them for tasks involving PDEs such as the incompressible Navier-Stokes equations.

1 INTRODUCTION

The paper targets fast, long-horizon control of high-dimensional nonlinear PDE systems, where iterative optimization is costly and existing methods can be difficult to deploy interactively. It introduces a hierarchical predictor-corrector approach trained through differentiable physics and reports stable control over longer spans.

PDE-based physical systems require agents to perceive, predict, and direct continuous dynamics through controls.
Iterative optimization is expensive because it repeatedly starts from scratch, while interactive applications may require reactions within fractions of a second.
The hierarchical predictor-corrector scheme temporally divides control into easier subproblems and combines models specialized to different time scales.
The models represent many possible solutions, helping avoid local minima that can trap classical optimization methods.
The method is evaluated on advection-diffusion PDE control tasks, including Navier-Stokes systems, using target-state accuracy and exerted force.
The resulting control remains stable for significantly longer time spans than alternative approaches.

2 BACKGROUND

The background situates the work among data-driven PDE methods, differentiable simulators, and adjoint-based optimization. It also contributes ΦFlow, a differentiable PDE solver integrated with mainstream machine-learning frameworks.

Prior work has applied diverse learning models to PDE solution, fluid simulation, and turbulence modeling.
Differentiable solvers have supported rigid-body mechanics, manipulation planning, liquid interaction, soft-robot control, protein structures, and cloth inverse problems.
The adjoint method, also called reverse mode differentiation in machine-learning frameworks, efficiently supports gradient computation.
ΦFlow integrates differentiable PDE solving with TensorFlow and PyTorch and is publicly available.

3 PROBLEM

The problem is to reconstruct a controlled PDE trajectory from an initial observable state to a target state while minimizing applied force. The formulation accommodates partial observability and infeasible exact targets.

The physical system evolves according to a PDE whose operator models physical behavior and whose external factors are separated from the control force.
The agent computes the force F(t), targeting force-solution manifolds for a fixed physical operator and suitable boundary conditions.
The agent receives only an observable state o(u), not necessarily the system’s full state.
Given initial and target observations, the desired trajectory matches both endpoints while minimizing force applied within the simulation domain.
With discrete time steps, the trajectory contains n = (t∗−t0)/∆t states.
When exact target matching is impossible because of physical or numerical constraints, an observation loss measures approximation quality.

4 PRELIMINARIES

Differentiable solvers enable gradients through complete PDE trajectories, addressing weaknesses of finite differencing and single-step supervision. The preliminaries motivate control-force estimators while identifying their computational and optimization limitations.

Differentiable solvers: Euler time stepping advances the PDE state by ∆t, but high-dimensional gradient estimation by finite differences requires expensive full-trajectory recomputation.
Differentiable solvers: Differentiable solvers compute analytic adjoint derivatives with respect to states and external inputs across arbitrary time horizons.
Iterative trajectory optimization: Single shooting simulates the full sequence, backpropagates its loss, and optimizes controls generated by linked control force estimators.
Iterative trajectory optimization: Long chains of nonlinear simulation steps make single-shooting optimization computationally expensive and can destabilize gradients far from an optimum.
Iterative trajectory optimization: Multiple shooting adds segment defect constraints, whereas collocation uses splines and is poorly suited to Eulerian solvers.
Supervised and differentiable physics losses: Supervised losses measure only single-step errors and can average over multiple valid controls, whereas differentiable physics losses backpropagate through chained future states.

5 METHOD

The method separates trajectory planning from control and organizes prediction hierarchically across time scales. Prediction refinement then updates intermediate predictions during simulation so control can remain aligned with evolving system states.

An observation predictor plans intermediate states, while a corrector estimates control forces to follow the planned trajectory.The split exposes the planned trajectory and separates long-horizon planning from action selection.
The predictor recursively estimates center states between arbitrary endpoint observations, partitioning a sequence until every time step has a prediction.This generalizes next-step prediction to states separated by arbitrary time spans.
Separate predictor networks are trained for time scales differing by factors of two, requiring O(log2 n) models for n frames.The hierarchy conditions predictions on their temporal scale and supports arbitrary time frames or resolutions when behavior can be anticipated.
Execution order: Prediction refinement combines finest-scale prediction, immediate corrector-solver execution, and prediction updates after each solver step.The selected scheme was found to produce the best results among the compared execution algorithms.
Execution order: Staggered execution lets predictions incorporate deviations once the corresponding simulated state is available, reducing divergence between actual and predicted evolution.However, some predictions remain unmodified and can become blurred because they average over deviations.
Prediction refinement: Prediction refinement re-evaluates existing predictions as simulation advances, continually conditioning final predictions on reconstructed states.It uses 3n − 2 log2(n) − 3 predictor evaluations, remains O(n), and incurs only small computational overhead.

6 RESULTS

The method is evaluated on nonlinear PDE control tasks that increase from one-dimensional Burger’s equation to two-dimensional fluid systems with indirect control. Hierarchical prediction and differentiable physics improve reconstruction quality, long-horizon stability, and control efficiency relative to supervised and iterative baselines.

Burger’s equation: 34% of the supervised model’s force is used by the differentiable objective-loss model on average when reconstructing Burger’s equation trajectories.The differentiable solver trains the network to correct temporal PDE evolution rather than only predict individual steps.
Burger’s equation: The differentiable physics and prediction-refinement variants produce trajectories close to ideal reconstructions, while simple supervised chains diverge over longer sequences.Staggered supervised execution is stable but jittering; differentiable physics improves reconstruction quality further.
Incompressible fluid flow: A multi-scale shooting baseline requires 1500 iterations to reach accuracy obtained almost instantly by the learned model.The learned solution manifold also provides useful initial guesses for iterative optimization, reducing the optimizer’s required force to 57.4% of default initialization.
Generalization: Splitting prediction from correction keeps intermediate states physically plausible and supports transitions involving multiple shapes despite training on individual shapes.This result demonstrates generalization to new shape-transition tasks.
Indirect control: 99.22% ± 0.15% of smoke reaches the target buckets with hierarchical predictions, using 19.1% ± 1.0% less force than the baseline.The baseline transfers 89% ± 2.6% of smoke, whereas direct optimization achieves only an 82% success rate and takes orders of magnitude longer than model evaluation.

7 CONCLUSIONS

The paper concludes that differentiable physics and hierarchical predictor-corrector inference can control complex PDE systems over long sequences. The learned solution-manifold representation supports fast inference while remaining complementary to iterative optimization.

Conclusion: Deep networks combined with a differentiable physics solver successfully predict and control complex physical systems.The conclusion attributes long-sequence reconstruction to hierarchical treatment of physical behavior at different time scales.
Conclusion: Differentiable solvers improve solution quality by letting networks learn how decisions affect future states, while hierarchical inference supports planning ahead.The paper reports that hierarchical inference outperforms traditional sequential agents in its experiments.
Scope and deployment: The deployment agent uses restricted observations, although the PDE solver still requires full state information during simulation.The information restriction applies to the agent at deployment, not to the solver used for training and simulation.
Scope and deployment: Learning approaches are presented as complementary to iterative optimization, providing fast solution-manifold inference and useful initial guesses for classical solvers.The paper does not claim that learned methods replace iterative optimization.

A IMPLEMENTATION DETAILS OF THE DIFFERENTIABLE PDE SOLVER

The differentiable PDE solver builds fluid simulations from differentiable grid operations and adjoint pressure solves, while hierarchical execution schemes reduce long-horizon control into recursively structured predictions.

The solver is publicly available, MIT-licensed, and implemented with machine-learning frameworks for automatic differentiation and neural-network integration.
Differential operators such as gradients, divergences, curls, and Laplacians are implemented with tensor operations that support framework-provided gradients.
Differentiable PDE building blocks: Semi-Lagrangian advection back-traces along an interpolated vector field and linearly interpolates neighboring values to transport fields differentiably.
Fluid simulation: Incompressible Navier-Stokes stepping applies diffusion, buoyancy, and pressure projection in sequence, with pressure obtained from a sparse Poisson system.
Differentiable pressure solve: The pressure gradient is computed by solving the adjoint pressure problem, avoiding unrolling hundreds of conjugate-gradient iterations while preserving automatic differentiation.
Execution schemes: Staggered execution recursively halves sequence length while increasing prediction count, whereas prediction refinement performs additional evaluations through recursive reconstruction.

C NETWORK ARCHITECTURES AND TRAINING

The networks use modified residual U-nets tailored to observation prediction and indirect control, with supervised pretraining followed by differentiable-physics training through long simulation chains.

Network architecture: All neural networks use modified U-net architectures with residual blocks replacing regular convolutions at each level.
Observation predictor: The fluid observation predictor receives current and target states, progressively downsamples through five residual blocks, and decodes spatial content through upsampling and skip concatenation.
Indirect control network: The indirect-control network outputs two velocity feature maps and uses a dense layer at 4x4 resolution to expand its receptive field.
Training: Supervised training uses ADAM, learning rates decreasing from 10^-3 to 10^-5, and typically converges within a few epochs.
Training: Differentiable-solver training starts at 10^-4 and is more expensive because each optimization step includes forward and backward simulation through the full chain.

D.1 BURGER’S EQUATION

The Burger’s-equation experiment evaluates hierarchical control on shock-wave dynamics, showing that differentiable-physics training produces efficient, stable trajectories compared with supervised control chains.

Experimental setup: Burger’s equation is simulated on a one-dimensional grid with 32 samples over 32 time steps, where opposing shock waves weaken until the stronger wave survives.
Experimental setup: The experiment compares staggered execution and prediction refinement, alongside supervised and differentiable-physics training conditions.
Results: The CFE chains fail to converge to the target, while differentiable-physics control produces a state resembling the target and supervised control deviates from an optimal trajectory.
Results: Hierarchical models converge toward the target with much less force than CFE chains, whereas supervised training induces oscillatory control and later counter-actions.
Results: Differentiable-physics methods significantly outperform supervised counterparts in required force and can learn trajectories close to ground-truth efficiency.
Results: Up to 13% less force than ground truth reaches the target on the sequence shown in Figure 4 through deliberate overshoot and later correction.
Force analysis: Large inferred forces are exponentially rare under the L2 regularizer, while supervised methods can produce a second high-force peak near the final step.
Comparison: Single-shooting reaches near-optimal solutions but requires around 300 iterations, compared with around 60 iterations to match the staggered prediction scheme.

D.2 INCOMPRESSIBLE FLUID FLOW

The incompressible-fluid experiments apply hierarchical prediction and correction to high-dimensional flow reconstruction and shape transitions, with differentiable training sharpening predictions and reducing required force.

Experimental setting: The fluid setting models complex, potentially chaotic incompressible Navier-Stokes dynamics with low viscosity and minimal diffusion to preserve a challenging control environment.
Experimental setting: Velocity is hidden from observation while smoke density is observed, and the system state combines velocity and density fields.
Control representation: The CFE uses a stream-function output whose curl becomes the velocity update, numerically simplifying enforcement of incompressibility.
Datasets: The datasets use 128 × 128 resolution and more than 16,000 effective continuous control parameters per time step.
Datasets: Flow reconstruction targets smoke density after 64 simulated time steps, while shape transitions use geometric shapes and evaluate sequences of length n = 16.
Training: CFE pretraining uses supervised velocity loss, while OPs are pretrained independently and then jointly optimized end-to-end with differentiable physics.
Flow reconstruction: Differentiable-physics prediction largely removes blurry long-term predictions, and refinement conditions predictions on prior reconstructed states for closer alignment.
Shape transition: Most shape-transition targets are closely matched, and prediction-correction generalizes to multiple weakly interacting shapes through intermediate predicted states.

D.3 INCOMPRESSIBLE FLUID WITH INDIRECT CONTROL

The indirect-control setup restricts forcing to peripheral regions, requiring the model to coordinate many parameters through global fluid dynamics. The learned method places substantially more smoke in the target bucket than the reference while using less force.

Setup: The network controls only peripheral regions, with more than 5000 control parameters per step, while obstacles and solid boundaries impose additional constraints.Smoke remains outside the controlled area, so incompressibility is the only mechanism influencing its velocity there.
Results: 99.22% ± 0.15% of smoke reached the target buckets while requiring 19.1% ± 1.0% less force than the straight-line reference.The comparison averages results over 100 test examples.
Results: The straight-line reference placed 89%±2.6% of smoke into the target bucket, compared with 99.22% ± 0.15% for the full algorithm.Performance was evaluated using both target-bucket smoke density and total applied force.
Comparison: Iterative optimization reached 82.1 ± 7.3 of the smoke in the correct bucket after roughly 660 optimization steps.The target began being reached in some examples after around 60 steps, followed by 600 further iterations of convergence.

D.4 COMPARISON TO SHOOTING METHODS

The paper compares its neural reconstruction with shooting-based optimization, whose long-horizon gradients can become unreliable in nonlinear Navier–Stokes dynamics. The neural method produces near-optimal trajectories much faster and can also initialize iterative refinement effectively.

Single-shooting: Single-shooting produces strong reconstruction artifacts because recurrent backpropagation through nonlinear Navier–Stokes dynamics makes gradients noisy and unreliable.The comparison uses differentiable-physics optimization of control velocities for a single input.
Multi-scale shooting: Multi-scale shooting improves reliability by first solving a coarsely discretized problem and then iteratively refining the discretization.The initial resolution uses 1/16 of the original width and height, reducing control parameters and nonlinear effects.
Efficiency: The neural network’s trajectories correspond to about 1500 multi-scale optimization steps, but inference takes 0.5 seconds instead of 131 seconds for one 16-frame sequence.The optimization ran on a GTX 1080 Ti, and the time gap grows for longer sequences because network inference scales favorably.
Trajectory quality: The neural method can represent a solution manifold, whereas multi-scale optimization may find unintuitive trajectories with noticeable detours despite similar observation loss.The paper illustrates this behavior for density-matching examples.
Qualitative comparison: Figure 17 contrasts direct shooting, multi-scale shooting, and the refined neural method on one shape-transition reconstruction.The neural model infers its shown solution in a single pass and is described as generalizing to a large class of inputs.
Hybrid refinement: Initializing multi-scale optimization with neural-network velocities makes both observation and force losses decrease from the beginning.The neural reconstruction is close enough to the optimum that the described multi-resolution approach is unnecessary.