Learning When to Refine: Long-Horizon Reinforcement Learning for Budgeted Neural-Operator PDE Solvers

Ange Tong

cs.LGSep 19, 2026

Learning When to Refine: Long-Horizon Reinforcement Learning for Budgeted Neural-Operator PDE Solvers

Authors: Ange Tong

Organizations: National Research Tomsk State University, 36 Lenin Avenue, Tomsk 634050, Russian Federation

Abstract

Neural operators provide fast surrogates for time-dependent PDEs, but autoregressive deployment creates a refinement-allocation problem: prediction errors vary over space and time, while only a finite number of local corrections can be committed along a trajectory. We formulate this as budgeted adaptive neural-operator solving. A global Fourier neural operator advances the full field, a local operator proposes patch-wise residual corrections, and a set-aware selector chooses where to refine. A macro policy decides when and how much of the remaining refinement budget to spend. We introduce rollout-verified policy improvement (RV-PI), which evaluates feasible refinement counts through actual continuation rollouts of the learned PDE solver, converts long-horizon advantages into conservative policy targets, and accepts an update only when held-out trajectory error improves. On the shallow-water benchmark with a 32-intervention budget, RV-PI achieves a three-seed mean trajectory relative L2 error of 0.6910, improving over immediate-only policy improvement by 5.37% and RandomMacro by 2.41%. On the forcing-driven Brusselator benchmark with a 76-intervention budget, RV-PI attains 0.09954, improving over immediate-only policy improvement by 2.31% and RandomMacro by 5.32%. These results show that, under a fixed refinement budget, the value of a local correction depends on its downstream effect on the autoregressive trajectory, not only on its immediate error reduction.

Figures & tables

Fig. 1: Overview of the proposed budgeted adaptive neural-operator solver. The global operator produces a provisional transition; all local patch corrections are generated once from that state and frozen. The macro RV-PI policy decides WHEN/HOW MUCH to refine, while the conditional selector decides WHERE. Selected frozen corrections are applied simultaneously by normalized overlap blending. The trajectory budget counts committed correction interventions; ground truth is used only to construct training-time selector gains, rollout returns, and validation acceptance, and inference is ground-truth free. The initial schematic layout was prepared with assistance from OpenAI image-generation tools and subsequently edited and verified by the author against the implemented solver.

Setting	Shallow water	Brusselator V2
Native trajectory	72 frames, $256\times 256$	39 frames, $28\times 28$
Train / val / test	230 / 20 / 50	800 / 100 / 100
Learned transitions	17 (stride 4)	38
Global FNO	width 32, modes 12, depth 4	width 32, modes 10, depth 4
Local FNO	width 32, modes 12, depth 4	width 32, modes 8, depth 3
Patch partition	$8\times 8$ ; core 32; halo 8	$7\times 7$ ; core 4; halo 1

Table 1: Experimental and solver configuration.

Fig. 2: Recorded Brusselator validation behavior under one-step-oriented (V1) and deployment-aligned closed-loop (V2) operator training. Panels (a)-(b) plot the measured step-wise mean Relative $L_{2}$ and prediction-to-target RMS ratio over the 38-transition rollout on logarithmic axes; panel (c) reports the corresponding validation summary metrics. The figure is generated directly from the recorded diagnostic tables.

PDE	Method	Trajectory RelL2 mean	Trajectory std	Final RelL2 mean	Final std	Local calls
Shallow water	CoarseOnly	0.740347	0.000000	1.056560	0.000000	0
Shallow water	RandomMacro	0.708093	0.005420	1.020698	0.011187	32
Shallow water	BeamBC	0.739403	0.000000	1.080303	0.000000	32
Shallow water	ImmediateOnlyPI	0.730245	0.000571	1.082184	0.000879	32
Shallow water	SetAwareMyopicMacro	0.737156	0.000000	1.017257	0.000000	32
Shallow water	RV-PI	0.691026	0.008779	0.982092	0.035847	32

Table 2: Three-seed fixed-budget comparison. Lower trajectory and final-frame Relative

L_{2}

are better. ’Local calls’ denotes the number of committed patch corrections; deterministic rows reuse the same frozen solver/checkpoint across policy seeds, so zero standard deviation does not imply independent backbone retraining.

Fig. 3: Cross-PDE trajectory Relative $L_{2}$ under matched refinement-intervention budgets. Bars show three-seed means and markers show the recorded seed-level values. RV-PI has the lowest mean trajectory error on both benchmarks. Within each benchmark, all adaptive methods commit the same number of local corrections.

Fig. 4: Recorded WHEN/WHERE behavior on the first held-out shallow-water case (case 250, seed 42). ImmediateOnlyPI and RV-PI both commit 32 patch corrections but distribute them differently over the 17 learned transitions. Panel (a) shows temporal allocation and cumulative budget use. Panel (b) shows trajectory error, the four-patch set selected by RV-PI at transition 13, cumulative spatial selection frequency, and the most-used patches.

Reinforcement Learning

Reinforcement learning studies how agents learn sequential decisions from rewards and interaction, balancing exploration, performance, safety, and adaptation.

See all 1,717 papers

Autoregression-Free Neural Operators for Time-Dependent PDEs
May 25, 2026Jiaquan Zhang, Caiyan Qin, Haoyu Bian +7Latent Dynamics Modeling PDE Operator Learning
School of Computer Science and Engineering, University of Electronic Science and Technology of China, Chengdu, Sichuan 611731, China · School of Robotics and Advanced Manufacture, Harbin Institute of Technology, Shenzhen, Guangdong 518055, China · School of Mathematical Sciences, Capital Normal University, Beijing 100048, China +3
HERO: History-Enriched Rollout Training for Long-Horizon Autoregressive Neural Operators
Jul 31, 2026Jiaquan Zhang, Shuxu Chen, Haifan Meng +6Neural Surrogate Modeling PDE Operator Learning
School of Computer Science and Engineering, UESTC · Electronics and Information Convergence Engineering, KHU · Department of Mathematical Sciences, UOL +3
Iterative Refinement Neural Operators are Learned Fixed-Point Solvers: A Principled Approach to Spectral Bias Mitigation
May 21, 2026Xiaotian Liu, Shuyuan Shang, Xiaopeng Wang +2Neural PDE Solvers Spectral Bias
1Dartmouth College · 2CUHK Shenzhen · 3Lawrence Berkeley National Lab

May 25, 2026J/K move · Enter open · S save

Autoregression-Free Neural Operators for Time-Dependent PDEs

Jiaquan Zhang, Caiyan Qin, Haoyu Bian, Libin Cai+6

School of Computer Science and Engineering, University of Electronic Science and Technology of China, Chengdu, Sichuan 611731, China · School of Robotics and Advanced Manufacture, Harbin Institute of Technology, Shenzhen, Guangdong 518055, China · School of Mathematical Sciences, Capital Normal University, Beijing 100048, China+3

Latent Dynamics Modeling PDE Operator Learning Spatiotemporal Forecasting

Open paper AskPDF arXiv

Neural operators learn mappings from function-dependent inputs to solutions, providing an effective framework for solving partial differential equations (PDEs). For time-dependent PDEs, existing methods typically perform long-horizon prediction through autoregressive rollout directly in high-dimensional physical field spaces, where each predicted state is recursively fed back as the input for the next step. Although effective for short-term prediction, this autoregressive rollout and the lack of continuous-time modeling lead to progressive error accumulation over long-horizon rollouts. In this work, we propose Autoregression-Free Neural Operators (AFNO), which map the time evolution of PDEs into a latent space and model continuous-time vector fields within it. AFNO uses flow matching to learn the latent vector field, thereby enabling continuous evolution over extended horizons, avoiding autoregressive rollout and capturing dynamics under varying parameter configurations through explicit conditioning on physical parameters. Theoretical analysis and extensive experiments on six PDEs demonstrate that AFNO improves long-horizon prediction stability and consistently reduces rollout errors compared with the baselines.

Abstract

Reinforcement Learning

Explore similar work

Autoregression-Free Neural Operators for Time-Dependent PDEs

HERO: History-Enriched Rollout Training for Long-Horizon Autoregressive Neural Operators

Iterative Refinement Neural Operators are Learned Fixed-Point Solvers: A Principled Approach to Spectral Bias Mitigation