Closed-Loop Preconditioning for Anisotropic Diffusion via Reinforcement Learning

Preconditioners are commonly held fixed throughout a linear solve, but the evolving residual provides feedback on the effectiveness of earlier corrections. We investigate whether learning to exploit this feedback can improve convergence on strongly anisotropic diffusion problems. Our Closed-Loop Preconditioner (CLP) uses reinforcement learning to incrementally update an explicit preconditioner within first-order Richardson iteration. A permutation-equivariant graph policy updates a diagonal-plus-low-rank preconditioner that is symmetric positive definite by construction, with analytic Galerkin weighting of the learned subspace. Controlled comparisons with learned static and open-loop policies support the benefit of solver feedback. CLP reduces iteration counts relative to algebraic multigrid-preconditioned Richardson and achieves shorter measured solver times than the tested classical baselines. The policy transfers from short training trajectories on small systems to substantially larger systems and supports longer solves through defect-correction cycles, without retraining.