Model-Agnostic Knowledge Guided Correction for Improved Neural Surrogate Rollout

TL;DR

This work tackles the rollout error problem in neural surrogates for transient PDE dynamics by introducing HyPER, a model- and simulator-agnostic reinforcement learning framework that learns when to invoke a physics simulator to correct surrogate trajectories. By formulating the task as an MDP with a binary action space and a cost-aware reward, HyPER achieves substantial reductions in cumulative rollout error while keeping simulator usage parsimonious. Across 2D Navier–Stokes and subsurface-flow tasks, HyPER outperforms surrogate-only baselines under in-distribution, changing-conditions, and noisy data scenarios, and demonstrates robustness to non-differentiable simulators. The approach provides a flexible, practical pathway to deploy accurate PDE surrogates in settings with limited data and varying physics, with future work aimed at more sophisticated RL policies and multi-physics extensions.

Abstract

Modeling the evolution of physical systems is critical to many applications in science and engineering. As the evolution of these systems is governed by partial differential equations (PDEs), there are a number of computational simulations which resolve these systems with high accuracy. However, as these simulations incur high computational costs, they are infeasible to be employed for large-scale analysis. A popular alternative to simulators are neural network surrogates which are trained in a data-driven manner and are much more computationally efficient. However, these surrogate models suffer from high rollout error when used autoregressively, especially when confronted with training data paucity. Existing work proposes to improve surrogate rollout error by either including physical loss terms directly in the optimization of the model or incorporating computational simulators as `differentiable layers' in the neural network. Both of these approaches have their challenges, with physical loss functions suffering from slow convergence for stiff PDEs and simulator layers requiring gradients which are not always available, especially in legacy simulators. We propose the Hybrid PDE Predictor with Reinforcement Learning (HyPER) model: a model-agnostic, RL based, cost-aware model which combines a neural surrogate, RL decision model, and a physics simulator (with or without gradients) to reduce surrogate rollout error significantly. In addition to reducing in-distribution rollout error by 47%-78%, HyPER learns an intelligent policy that is adaptable to changing physical conditions and resistant to noise corruption. Code available at https://github.com/scailab/HyPER.

Figures (8)

• Figure 1: Cumulative MSE depicting rollout error for a single trajectory of HyPER vs surrogate only methods. x's mark the timesteps during the trajectory where our RL policy calls the simulator.
• Figure 2: Overview of Hybrid PDE Predictor with RL (HyPER) with example rollout. Here $\pi_{\theta}$ is the decision model, $f_{\phi}$ is the surrogate, and $S$ is the simulator. At $t = 9.0$ in the above trajectory, the decision policy invokes the simulator, correcting the trajectory to reduce rollout error. The effect of this knowledge-guided correction can be observed by a reduction in absolute error (dotted green circle) in the figure.
• Figure 3: Predictions and absolute error snapshots of HyPER rollout for two distinct trajectories.
• Figure 4: Predictions and absolute error of HyPER vs UNet and FNO for a single trajectory and timestep. Fig. \ref{['fig:boundary-qual-error']}a shows the rollout error accumulation over the trajectory with 'x' marking the times the simulator is called. Fig. \ref{['fig:boundary-qual-error']}b shows the resolved system state for a the same trajectory at a single timestep. Fig. \ref{['fig:boundary-qual-error']}b(i) shows the ground truth field while Fig. \ref{['fig:boundary-qual-error']}b(ii)-(iv) indicate the result of UNet, FNO and HyPER rollouts respectively. Fig. \ref{['fig:boundary-qual-error']}b(v)-(vii) depict the corresponding absolute errors. We notice that only HyPER rollouts capture the correct characteristics relative to the ground truth owing to the knowledge-guided correction while SUG models are unable to faithfully resolve the trajectory under changing physical conditions.
• Figure 5: Figure \ref{['fig:time-bar-charts']}(a) shows the average time of PDE prediction of a full trajectory for each method. Figure \ref{['fig:time-bar-charts']}(b) illustrates the error per unit time (lower is better) for each method. We do not show the Sim-Only case here as we assume error is effectively zero for the simulator.