PDE Control Gym: A Benchmark for Data-Driven Boundary Control of Partial Differential Equations

TL;DR

PDE Control Gym addresses the lack of standardized benchmarks for learning-based boundary control of PDEs by introducing a modular RL environment with three canonical PDE tasks: 1D transport, 1D reaction-diffusion, and 2D Navier-Stokes. The authors train model-free RL controllers (PPO and SAC) and compare them against classical model-based backstepping and optimization-based methods, showing that RL can stabilize these systems, albeit at higher sample costs. The benchmark decouples PDE solvers from controllers, enabling fair comparisons and rapid experimentation, with full documentation and open-source release. The work provides a foundation for data-driven PDE control and highlights directions such as reward design, network architectures, and handling time-varying dynamics.

Abstract

Over the last decade, data-driven methods have surged in popularity, emerging as valuable tools for control theory. As such, neural network approximations of control feedback laws, system dynamics, and even Lyapunov functions have attracted growing attention. With the ascent of learning based control, the need for accurate, fast, and easy-to-use benchmarks has increased. In this work, we present the first learning-based environment for boundary control of PDEs. In our benchmark, we introduce three foundational PDE problems - a 1D transport PDE, a 1D reaction-diffusion PDE, and a 2D Navier-Stokes PDE - whose solvers are bundled in an user-friendly reinforcement learning gym. With this gym, we then present the first set of model-free, reinforcement learning algorithms for solving this series of benchmark problems, achieving stability, although at a higher cost compared to model-based PDE backstepping. With the set of benchmark environments and detailed examples, this work significantly lowers the barrier to entry for learning-based PDE control - a topic largely unexplored by the data-driven control community. The entire benchmark is available on Github along with detailed documentation and the presented reinforcement learning models are open sourced.

Figures (10)

• Figure 1: Rewards for training PPO (blue) and SAC (orange) on the 1D transport PDE, 1D reaction-diffusion PDE, and 2D Navier-Stokes PDE from left to right. The solid lines represent the mean and the shaded bounds are $95\%$ confidence intervals across 5 seeds.
• Figure 2: Example of the 1D transport PDE system stabilization using backstepping, PPO, and SAC (left to right) under initial conditions $u(x, 0) = 10$. The recirculation coefficient is defined as $\beta(x)=5\cos(\gamma \cos^{-1}(x))$ with $\gamma=7.35$.
• Figure 3: Example of reaction-diffusion PDE system stabilization using backstepping, PPO, and SAC (left to right) under initial conditions $u(x, 0) = 10$. The recirculation coefficient using the Chebyshev polynomial defined as $\lambda(x)=50\cos(\gamma \cos^{-1}(x))$ with $\gamma=8$.
• Figure 4: Example of Navier-Stokes PDE tracking using optimization-based control, PPO, and SAC under initial conditions $u(x, y, 0) = 0$ at $t=0$ (top) and $t=0.2$ (bottom). Red and black arrows represent the actual and reference velocity field respectively. The background color represents the magnitude of the velocity vector.
• Figure 5: Instability of the 1D transport PDE with $\beta(x) = 5\cos(7.35\cos^{-1}(x))$ under a openloop control signal ($U(t)=0$).