Deep Reinforcement Learning Solves Elliptic Differential Equations With Sub-Second Inference Speeds

A computational framework formulated around deep reinforcement learning has demonstrated the ability to solve classical elliptic partial differential equations with greater numerical accuracy and substantially faster inference speeds than established numerical techniques. Governing equations such as the Poisson, Helmholtz, and anisotropic diffusion equations describe steady-state physical phenomena ranging from electrostatic potential distributions to thermal conduction and linear elasticity. Conventional numerical discretization methods, including finite difference, finite element, and finite volume techniques, approximate these continuous systems on discrete spatial grids. However, fine grid resolutions cause algebraic equations to expand significantly, requiring intensive iterative convergence sweeps that hinder applications requiring rapid physical field evaluations, such as online process control and interactive design optimization.

To overcome the computational bottleneck inherent in multi-sweep iterative routines, researcher Mengyao Li developed a sequential decision-making framework using a deep Q-network architecture. Rather than relying on classical linear algebra updates or continuous automatic differentiation schemes common in physics-informed neural networks, the technique models elliptic equation solving as a discrete Markov decision process. In this formulation, an autonomous agent inspects the spatial distribution of errors across a discrete grid, selects targeted correction increments across all nodes simultaneously, and receives feedback through an engineered reward signal. This setup enables the trained neural network to bypass iterative linear system solving during deployment, outputting approximate solutions through a single feedforward evaluation.

The underlying state representation for the reinforcement learning agent is organized across uniform discrete grids where local and global field data are stacked into multi-channel arrays. For a standard grid resolution of 64 by 64 nodes, representing 4,096 discrete spatial points, each node records three distinct feature channels: the current predicted solution value, the calculated partial differential equation residual, and a binary indicator designating boundary geometry. This arrangement produces a 12,288-dimensional state tensor that feeds directly into a feature extraction network composed of three convolutional layers, batch normalization, rectified linear unit activations, two residual blocks, and global average pooling to yield a compact 512-dimensional latent feature vector.

The action space allows the agent to adjust solution magnitudes across the domain using eleven discrete increment values ranging from negative 0.10 to positive 0.10 in intervals of 0.02. Boundary nodes are constrained to an adjustment value of zero throughout the process to strictly enforce essential boundary conditions, preventing unphysical deviations at fixed domain boundaries. The agent interacts with the environment across all nodes concurrently, updating the spatial solution field in one step. The reward function driving policy learning balances three explicit physical objectives through a weighted sum: a residual penalty term weighted at 0.6, a boundary error penalty term weighted at 0.3, and a smoothing regularization penalty weighted at 0.1 designed to penalize spatial gradient oscillations.

To train the policy, the system incorporates a double deep Q-network framework paired with prioritized experience replay to mitigate overestimation bias and improve sample efficiency. The replay buffer maintains a capacity of 100,000 transitions, while the target network parameters synchronize with the active policy network every 100 training steps. The optimization protocol employs an Adam optimizer with an initial learning rate of 0.001, an initial exploration parameter of 1.0 that decays at a rate of 0.995 per step to a minimum threshold of 0.01, and a warm-up phase spanning 500 steps. The distributed computational setup executed on an Ubuntu system equipped with an Intel Xeon Gold 6248 processor and four NVIDIA Tesla V100 graphics processing units operating under PyTorch.

The training and evaluation datasets comprised 2,000 simulated configurations across rectangular, L-shaped, and circular computational geometries. The source terms were generated using Gaussian random fields with controlled length scales and analytical formulations to limit excessive high-frequency oscillations below 20 percent. A rigorous data-cleaning procedure screened the generated training profiles, removing anomalous samples exhibiting abrupt local residual spikes greater than 3.0 or boundary discrepancies exceeding 0.05. This filtration retained 1,402 clean training configurations out of the initial 1,500 candidates, establishing a structured data distribution characterized by compressed residual outliers and lower variance before model training commenced.

Rigorous numerical evaluations across three benchmark test cases demonstrated that the deep reinforcement learning solver achieved superior precision compared to standard baselines. In tests involving the Poisson equation on a rectangular domain, the solver reached a relative L2 error of 0.28 percent and an H1 semi-norm error of 0.67 percent, outperforming the finite element method at 0.32 percent and the multigrid method at 0.31 percent, while significantly surpassing physics-informed neural networks at 1.23 percent. For the Helmholtz equation with a wavenumber of 4.0, the reinforcement learning solver attained a relative L2 error of 0.39 percent compared to 0.45 percent for the finite element method and 0.44 percent for the multigrid solver. In the anisotropic diffusion case set within an L-shaped domain possessing a re-entrant corner singularity, the proposed method recorded an L2 error of 0.62 percent, maintaining higher fidelity than the finite element baseline error of 0.78 percent.

Training stability and residual reduction dynamics highlighted the critical contribution of memory replay mechanisms during the optimization process. When full experience replay was enabled, the L2 residual norm decreased rapidly from an initial value of 12.4 down to 2.3 within the first 50 steps, eventually reaching 0.16 after 20,000 training iterations. Conversely, disabling the replay buffer resulted in persistent trajectory bias, leaving the residual oscillating at 0.62 at the end of training. Systematic ablation experiments indicated that experience replay accelerated initial convergence rates by approximately 2.5 times and reduced the final residual by 72 percent, while incorporating the smoothness regularization term yielded a 12 percent error reduction.

Computational efficiency benchmarks emphasized substantial latency reductions during the forward inference stage. On a 256 by 256 grid, the reinforcement learning model executed a complete solution pass in 0.15 seconds, representing a marked acceleration relative to the multigrid solver runtime of 14.8 seconds and the finite element computation time. While physics-informed neural networks exhibited resolution-independent inference times hovering around 52 seconds, their continuous-domain formulations required prolonged prior optimization steps. The sub-second execution speed of the discrete reinforcement learning model addresses timing constraints in practical engineering tasks, such as potential flow calculations in aerodynamics and real-time thermal monitoring where loop latencies must remain beneath strict physical thresholds.

Despite these computational advantages, the study documented clear physical limitations when evaluating configurations beyond the training domain. Sensitivity tests showed that expanding the spatial correlation length scale of Gaussian source terms from 0.2 up to 0.45 raised the median relative L2 error from 0.29 percent to 0.72 percent. Furthermore, applying the trained policy to unstructured circular meshes with unseen mixed boundary conditions increased the error to 0.92 percent, while localized singularities at re-entrant corners generated error concentrations up to 1.2 percent in localized regions. The findings define operating bounds for neural PDE operators, showing that while deep reinforcement learning offers effective real-time simulation capabilities, broader deployment across irregular geometries and pure Neumann constraints will require multi-scale hierarchical schemes or rapid meta-learning adaptation.

In spatial evaluations of the Helmholtz benchmark with a wavenumber of 4.0 and a sinusoidal source distribution, predicted values closely tracked reference solutions computed on a fine grid. The reinforcement learning model produced field values ranging from negative 1.76 to positive 1.82, closely matching the reference extrema of negative 1.79 and positive 1.85. Absolute errors remained concentrated in steep gradient regions near peaks and valleys, yielding an average absolute error of 0.009 and a maximum localized discrepancy of 0.057 across the domain. From an implementation perspective, the network contains approximately 12 million parameters. Single forward inference in 32-bit floating precision requires 48 megabytes of memory for weights, with intermediate activations and the input tensor using an additional 200 megabytes.

Systematic hyperparameter sweeps revealed clear operating boundaries governing convergence and final accuracy. Across learning rates between 0.0001 and 0.01, optimal convergence occurred at 0.001, whereas shifting the rate to 0.0003 or 0.003 degraded accuracy above a 0.35 percent relative error, and a high rate of 0.01 caused training instability with errors reaching 0.89 percent. Discount factor trials showed that setting the factor to 0.99 achieved the baseline error of 0.28 percent, while reducing it to 0.90 raised error to 0.44 percent and increasing it to 0.995 produced an error of 0.31 percent. In addition, ablation tests demonstrated that adaptive learning rate scheduling with warm-up steps improved accuracy by 0.07 percentage points over a static rate, while reducing action discretization from eleven increments to five coarser choices compromised solution precision.

Beyond benchmark validation, sub-second inference enables specialized engineering tasks with strict latency requirements. In medical electrical impedance tomography, which requires solving hundreds of Poisson equations under varying electrode configurations, the model reduces single-solve times from 3.5 seconds with multigrid methods to 0.15 seconds, lowering cumulative reconstruction times from several minutes to under 30 seconds. In semiconductor thermal management, closed-loop cooling systems require solving heat conduction equations in under 50 milliseconds; the reinforcement learning model completes evaluations in 15 milliseconds on a 64 by 64 grid, satisfying embedded operational thresholds.

Subject of Research: Numerical solution of elliptic partial differential equations using deep reinforcement learning

Article Title: Solving elliptic partial differential equation based on deep reinforcement learning

Article References: Li, M. (2026). Solving elliptic partial differential equation based on deep reinforcement learning. Discover Artificial Intelligence, 6(1), Article 1731. https://doi.org/10.1007/s44163-026-02533-0

Image Credits: AI Generated

DOI: 10.1007/s44163-026-02533-0

Keywords: deep reinforcement learning, elliptic partial differential equations, Poisson equation, Helmholtz equation, anisotropic diffusion, deep Q-network, Markov decision process, numerical methods, finite element method, multigrid method, physics-informed neural networks, computational fluid dynamics, scientific computing, experience replay, discretization grids, boundary value problems