[Paper Review] Learning a model is paramount for sample efficiency in reinforcement learning control of PDEs

Stefan Werner, Sebastian Peitz|arXiv (Cornell University)|Feb 14, 2023

Model Reduction and Neural Networks4 citations

TL;DR

This paper proposes a model-based reinforcement learning framework that uses a data-driven, actuated convolutional LSTM surrogate model to drastically improve sample efficiency in controlling partial differential equations (PDEs). By jointly training the RL agent and the surrogate model, and iteratively updating the model to reduce error compounding, the approach reduces data requirements by an order of magnitude compared to model-free methods, with ablation studies confirming the necessity of online model adaptation and ensemble-based rollouts.

ABSTRACT

The goal of this paper is to make a strong point for the usage of dynamical models when using reinforcement learning (RL) for feedback control of dynamical systems governed by partial differential equations (PDEs). To breach the gap between the immense promises we see in RL and the applicability in complex engineering systems, the main challenges are the massive requirements in terms of the training data, as well as the lack of performance guarantees. We present a solution for the first issue using a data-driven surrogate model in the form of a convolutional LSTM with actuation. We demonstrate that learning an actuated model in parallel to training the RL agent significantly reduces the total amount of required data sampled from the real system. Furthermore, we show that iteratively updating the model is of major importance to avoid biases in the RL training. Detailed ablation studies reveal the most important ingredients of the modeling process. We use the chaotic Kuramoto-Sivashinsky equation do demonstarte our findings.

Motivation & Objective

• To address the high sample complexity of reinforcement learning (RL) in controlling complex PDE-governed systems, which often require tens to hundreds of millions of environment interactions.
• To investigate whether learning a surrogate model in parallel with RL training can significantly reduce data requirements for sample-efficient control.
• To evaluate the impact of model error compounding on policy learning and to identify mechanisms that stabilize model-based RL in low-data regimes.
• To demonstrate the superiority of model-based RL with iterative model updates and ensemble rollouts over model-free and naive model-based baselines on chaotic PDEs.
• To provide empirical evidence and ablation studies showing that surrogate modeling is not just beneficial but paramount for sample efficiency in PDE control.

Proposed method

• A convolutional LSTM is trained as a surrogate model to predict the temporal evolution of the PDE state, with actuation integrated directly into the model architecture to improve predictive accuracy.
• The RL agent and the surrogate model are trained jointly in an end-to-end manner, enabling the agent to leverage model-based rollouts for policy improvement without relying on real environment interactions.
• Model weights are iteratively updated during training to prevent bias accumulation and reduce compounding errors, ensuring the surrogate remains aligned with the true system dynamics.
• Ensemble-based model rollouts are used to diversify exploration and stabilize learning, with a curriculum learning strategy that gradually increases rollout length.
• Ablation studies compare the full method against variants without online model updates, without ensembling, or with alternative surrogate architectures to isolate key components.
• The method is evaluated on the chaotic Kuramoto-Sivashinsky equation, a benchmark for nonlinear PDE control, using offline evaluation metrics to assess policy performance and generalization.

Figure 1: In model-based RL with learned models, collected data serves two purposes. Samples are used to learn a model of the system dynamics and to estimate update targets for the model-free agent. Using rollouts of the agent in the model, algorithms collect additional data to improve their behavio

Experimental results

Research questions

• RQ1 Can learning a surrogate model in parallel with RL training reduce the number of real environment interactions required for sample-efficient control of PDEs?
• RQ2 How does iterative model updating affect the stability and convergence of model-based RL in low-data regimes?
• RQ3 What is the impact of model error compounding on policy performance, and how can it be mitigated?
• RQ4 How do ensemble-based rollouts and curriculum learning compare to fixed-length or single-model rollouts in stabilizing the training process?
• RQ5 Does a residual-based surrogate model outperform a full-state prediction model in terms of generalization and sample efficiency?

Key findings

• The proposed model-based RL approach reduces the total number of required environment samples by approximately one order of magnitude compared to model-free baselines.
• Iterative model updating is critical to prevent bias accumulation and maintain policy performance, as static models lead to unstable learning and poor generalization.
• Ensemble-based rollouts significantly improve training stability and reduce variance, outperforming single-model rollouts that are prone to exploiting model errors.
• The surrogate ablation, which uses full-state prediction instead of residuals, fails to generalize to the true system due to persistent early-stage prediction errors, even with stable convergence.
• Random exploration alone is insufficient to cover the relevant state space, leading to divergence in policy optimization as prediction errors grow over time.
• The full method with online model adaptation, ensembling, and curriculum learning achieves the most stable and highest-performing policy, with the lowest standard deviation across trials.

Figure 2: Our surrogate model is composed of a state encoder $f_{\boldsymbol{\theta}_{\mathrm{enc}}}$ (yellow), an action encoder $f_{\boldsymbol{\psi}_{\mathrm{enc}}}$ (red) and a convolutional LSTM cell $f_{\boldsymbol{\theta}_{\mathrm{fwd}}}$ (green) modeling the transition dynamics, as well as a