cs.LGMar 2, 2026

Accelerating PDE Surrogates via RL-Guided Mesh Optimization

Authors: Yang Meng, Ruoxi Jiang, Zhuokai Zhao, Chong Liu, Rebecca Willett, Yuxin Chen

Organizations: UChicago · Fudan University · UAlbany

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Deep surrogate models for parametric partial differential equations (PDEs) can deliver high-fidelity approximations but remain prohibitively data-hungry: training often requires thousands of fine-grid simulations, each incurring substantial computational cost. To address this challenge, we introduce RLMesh, an end-to-end framework for efficient surrogate training under limited simulation budget. The key idea is to use reinforcement learning (RL) to adaptively allocate mesh grid points non-uniformly within each simulation domain, focusing numerical resolution in regions most critical for accurate PDE solutions. A lightweight proxy model further accelerates RL training by providing efficient reward estimates without full surrogate retraining. Experiments on PDE benchmarks demonstrate that RLMesh achieves competitive accuracy to baselines but with substantially fewer simulation queries. These results show that solver-level spatial adaptivity can dramatically improve the efficiency of surrogate training pipelines, enabling practical deployment of learning-based PDE surrogates across a wide range of problems.

Figures & tables

Figure 1: The RLMesh framework. An input state (e.g., PDE initial condition) is sampled from a prior distribution and passed to an RL policy that selects a small set of mesh grid points under a budget. The PDE solver is queried only at these points to produce high-fidelity but sparse observations, which are stored in an updated dataset for surrogate training. A lightweight proxy model provides fast estimates of surrogate improvement and provides reward signal for the RL agent, updating the policy and closing the loop.
Figure 2: Correlation between proxy (RBF kernel ridge, x-axis) and FNO surrogate errors (y-axis) across subset sizes. Each point corresponds to a retrained model. Spearman correlation (0.9908) indicates strong rank consistency.
Figure 3: Active learning performance across three benchmark systems. RLMesh consistently outperforms heuristic baselines under identical query budgets.
Figure 4: Burgers: effect of varying per-instance query budget BB on active learning performance.
Figure 5: Time–error analysis on Burgers. (a) Comparison across selection strategies. (b) Tradeoff across per-instance budgets using RLMesh .
Figure 6: Visualization of grid point selections by RLMesh on Burgers and Darcy. (a) Burgers: selections adapt to steep shock fronts and turning points while maintaining scaffolding in smoother regions. (b) Darcy: selections concentrate along high-contrast channels and near boundary layers, with adaptive coverage in homogeneous regions.
Figure 7: Learning progression under active acquisition. (a) Burgers: the model rapidly reduces bias and oscillatory error, with residual discrepancies confined to shock neighborhoods. (b) Darcy: coarse structure is captured early, while later acquisitions refine high-contrast channels and boundary layers.
IterationRLMeshSelf-MILCMD
10.0439 ±\pm 0.00470.0489 ±\pm 0.01080.0464 ±\pm 0.0169
60.0109 ±\pm 0.00210.0146 ±\pm 0.00140.0126 ±\pm 0.0007
120.00566 ±\pm 0.001280.00697 ±\pm 0.000840.00618 ±\pm 0.00041
180.00364 ±\pm 0.000670.00418 ±\pm 0.000380.00384 ±\pm 0.00010
Table 1: Burgers equation: RLMesh with instance-level baselines (mean ±\pm std).
Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 8: Model choice assessment for the proxy. (a) Correlation between the proxy and FNO across training subsets. (b) RMSE for candidate models across subset sizes.
Figure 9: Active learning performance on Burgers with 20 per-instance budget
Figure 10: Active learning performance on Burgers with PDE simulator. RLMesh consistently outperforms baselines under identical budgets.
Figure 11: Burgers prediction with and without mesh overlay
Figure 12: Darcy prediction with and without mesh overlay (instances 1–2).
Figure 13: Darcy predictions with and without mesh overlay (instances 3–4, continued).
Figure 14: Active learning progression: each subfigure shows (left) ground truth and model prediction, and (right) error map as iterations proceed.
Figure 15: Active learning progression: each subfigure shows (left) ground truth, (middle) model prediction, and (right) error map as the active learning iterations proceed.

Explore similar work

CardsList

  1. Learning Where to Simulate: Generative Active Sampling for Online PDE Surrogate Training

    Jun 8, 2026Pierre Cesar, Sofya Dymchenko, Abhishek Purandare +1PDE Surrogate ModelingActive Learning

    DATAMOVE · Univ. Grenoble Alpes, Inria, CNRS, Grenoble INP, LIG, Grenoble, France

  2. Compute-optimal data scaling for neural surrogates via multi-fidelity training

    Oct 12, 2026Paul Setinek, Fabian Paischer, Nils Thuerey +2Neural Surrogate ModelingPDE Surrogate Modeling

    Institute for Machine Learning, JKU Linz · Mistral AI · Technical University of Munich +1

  3. Learning to Discretize: Diffusion-Based Adaptive Mesh with Spectral Guidance

    Jul 13, 2026Zixuan Shen, Bingchuan Wang, Zhi Wang +1Neural PDE SolversPDE Surrogate Modeling

    Central South University · Central South University, Changsha, Hunan 410083, China · Nanjing University +1