#  Learning a model is paramount for sample efficiency in reinforcement
learning control of PDEs

2nd Stefan Werner  Department of Computer Science
Paderborn University
Paderborn, Germany
<EMAIL_ADDRESS>1st Sebastian Peitz  Department of Computer
Science
Paderborn University
Paderborn, Germany
<EMAIL_ADDRESS>ORCID: 0000-0002-3389-793X  Stefan Werner, Sebastian
Peitz Both authors are with the Department of Computer Science, Paderborn
University, Paderborn, Germany. e-mail<EMAIL_ADDRESS>
###### 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.

###### Index Terms:

reinforcement learning, surrogate modeling, partial differential equations,
feedback control

##  I Introduction

Feedback control of complex physical systems is an essential building block in
almost any modern technology such as wind energy [1], combustion [2] or
nuclear fusion [3]. In all these systems, we face the task of having to take
control actions in a very short amount of time and for a system with highly
complex, distributed dynamics (typically governed by nonlinear partial
differential equations (PDEs)) that are difficult or even impossible to
observe completely. Since classical feedback control strategies have
limitations for such higly complex systems, _model predictive control_ [4] has
emerged as an important control paradigm, where a model of the system dynamics
is used to solve an optimal control problem repeatedly on a receding horizon
in order to create a feedback loop. As this calls for the solution of an
optimal control problem in a very short amount of time, surrogate models are
frequently used to accelerate the solution process, in particular when the
system dynamics are governed by PDEs. There exists a plethora of approaches,
ranging from Galerkin models based on the _Proper Orthogonal Decomposition_
[5] over the sparse identification of governing equations [6] to deep neural
networks [7]. All these modeling paradigms now face the challenge to trade
between model accuracy and computational performance, which becomes more and
more challenging for increasingly complex systems.

Due to the real-time challenges for MPC, _reinforcement learning_ [8] is
gaining more and more popularity as a very powerful and real-time capable
paradigm for feedback control. In recent years, we have seen tremendous
successes, not only in the area of games (e.g., [9]), but also in complex
technical applications such as flow control [10] or nuclear fusion [3].
However, there are two drawbacks that limit the deployment of RL agents in
real systems. The first point is that the amount of required training data
often exceeds tens to hundreds of millions of samples [11, 12] such that the
training becomes very expensive. A notable example (even though from the area
of game playing) includes the efforts of OpenAI Five to master the Dota 2
video game with an agent training on about 100,000 CPU cores for a time period
of about ten months [13]. These issues are even more severe when data
collection becomes exceedingly resource-intensive, as is the case for most
fluid or electromagnetic field simulations, see, e.g., [14]. The second
drawback is that the control performance may vary strongly, in particular for
high-dimensional state and action space dimensions, where it is challenging to
obtain the required amounts of training data. A popular approach to tackle
these issues is the usage of model-based algorithms [15], see also [16] for an
overview and a taxonomy. Therein, a surrogate model replaces the real
environment to allow for a significant increase of the training on data
created by said surrogate. Using such a model enables the agent to look ahead
at the effects of its actions on the system without ever performing them in
the actual environment. It is widely accepted that surrogate models can reduce
the amount of data an agent consumes before converging in the case that
accurate and generalizable surrogates can be learned [17, 18, 19], and in many
situations, this model can be orders of magnitude faster than evaluating the
environment model (e.g., numerically solving a PDE). Collecting data through
_model-based rollouts_ is advantageous as soon as they reduce the amount of
computation time spent on numerical simulations or the number of time-
consuming and impractical real-world trials. At the same time, models
introduce approximation errors so that we may obtain solutions inferior to
those discovered using the model-free paradigm [20]. The literature thus
studies to what extent capable function approximators can overcome modeling
bias and how to avoid for errors to propagate into control strategies [21, 22,
23, 24]. Overall, interest has shifted from stable and robust linear models
[25, 26] to shallow neural networks [27] or Gaussian processes [28, 18] as
learning machines, before advancing to the deep neural networks seen today in
state-of-the-art algorithms for model-based control [24, 29]. In the meantime,
the increased model capacity makes overfitting a substantial issue of concern
[30, 21]. In spite of that, the recent _SimPLe_ [31] and _DREAMER_ [32]
algorithms solve challenging visual control tasks, being orders of magnitude
more efficient (even though the original problem is not expensive to simulate)
than their model-free peers while rivaling their asymptotic performance.

Even though it seems clear that surrogates have the potential to improve the
learning process, the trade off in terms of sample efficiency, computational
complexity, and performance is seldom quantified, in particular in an
engineering context such as PDE control. Two recent approaches using learned
surrogates are [33] or [34]. In addition to assuming complete knowledge of the
phenomena governing the system dynamics, the first work uses a physics-
informed loss that, in essence, uses a deep neural network to compute the same
finite difference solution as the numerical simulation does, but at a
substantial slowdown in terms of runtime. The second contribution learns a
highly accurate surrogate before optimizing a control strategy on it but fails
to adjust its forecasts to the continuously changing distribution of visited
simulation states.

In this paper, we discuss and evaluate a variety of design choices associated
with model-based reinforcement learning, adopt a more sophisticated
optimization approach, and identify promising directions for future research.
We discuss the various modeling steps (Section III) that are required in order
to arrive at a model that possesses the required accuracy. We then thoroughly
study the impact on the required amount of training data (Section IV), where
an improvement by a factor of nine is observed for the case of the 1D
Kuramoto-Sivashinsky equation. In both Sections III and IV, we conduct various
ablation studies to identify the most important modeling techniques in the
offline phase, as well as online updating strategies during deployment.
Finally, we discuss the implications for the robustness of RL agents as well
as future research directions in Section VI.

##  II Reinforcement learning

Reinforcement learning (RL) aims to solve sequential decision-making problems
mathematically defined as Markov Decision Processes (MDPs) (see, e.g., [8] for
a detailed overview). It possesses a set $\mathcal{S}$ of system states and a
set $\mathcal{A}$ of actions an _agent_ may select among to exercise control.
At each discrete time step $\tau$, the agent observes state
$\boldsymbol{s}_{\tau}\in\mathcal{S}$ of the environment and responds with an
action $\boldsymbol{a}_{\tau}\in\mathcal{A}$. A stochastic transition function
$\mathcal{T}:\mathcal{S}\times\mathcal{A}\rightarrow\mathcal{P}(\mathcal{S})$
formalizes in what way the system state changes as a result thereof. In an
MDP, the state transition is independet of past states and actions, meaning
that it satisfies the Markov property. The reward signal $r_{\tau}$ then
quantifies the quality of decision $\boldsymbol{a}_{\tau}$ taken in state
$\boldsymbol{s}_{\tau}$ and is an instance of the stochastic reward function
$\mathcal{R}:\mathcal{S}\times\mathcal{A}\rightarrow\mathcal{P}(\mathbb{R})$.
Starting from an initial state
$\boldsymbol{s}_{0}\sim\mathcal{P}_{0}(\mathcal{S})$, the RL framework aims to
maximize the expected sum of discounted future rewards
$\mathbb{E}\left[\sum^{\infty}_{\tau=0}\gamma^{\tau}r_{\tau}\right]$, where
$\gamma\in[0,1]$ is the discount factor.

A policy $\pi$ is a mapping
$\pi:\mathcal{S}\rightarrow\mathcal{P}(\mathcal{A})$ modeling the probability
$\pi(\boldsymbol{a}|\boldsymbol{s})$ of taking action
$\boldsymbol{a}\in\mathcal{A}$ in state $\boldsymbol{s}\in\mathcal{S}$ of the
environment. Solving an MDP means finding an optimal policy $\pi^{*}$:

$\displaystyle\pi^{*}$
$\displaystyle=\operatorname*{arg\,max}\limits_{\pi}\mathbb{E}_{\pi}\left[\,\sum_{\tau=0}^{\infty}\gamma^{\tau}r_{\tau}\,\right].$

So-called value functions are useful concepts to formally define optimality
conditions for policies. In substance, a state value $V^{\pi}(\boldsymbol{s})$
denotes what future rewards to expect following policy $\pi$ when the
environment is in state $\boldsymbol{s}$:

$\displaystyle V^{\pi}(\boldsymbol{s})$
$\displaystyle=\mathbb{E}_{\pi}\left[\,\sum_{k=0}^{\infty}\gamma^{k}r_{\tau+k}\biggm{|}\boldsymbol{s}_{\tau}=\boldsymbol{s}\,\right].$

Therefore, a policy $\pi$ is optimal if its state value
$V^{\pi}\left(\boldsymbol{s}\right)$ at each state
$\boldsymbol{s}\in\mathcal{S}$ is at least as large as the state values of any
other policy. Similarly, state-action values
$Q^{\pi}(\boldsymbol{s},\boldsymbol{a})$ describe what rewards to expect once
action $\boldsymbol{a}$ is selected in state $\boldsymbol{s}$ and policy $\pi$
governs the behavior thereafter:

$\displaystyle Q^{\pi}\left(\boldsymbol{s},\boldsymbol{a}\right)$
$\displaystyle=\mathbb{E}_{\pi}\left[\,\sum_{k=0}^{\infty}\gamma^{k}r_{\tau+k}\biggm{|}\boldsymbol{s}_{\tau}=\boldsymbol{s},\boldsymbol{a}_{\tau}=\boldsymbol{a}\,\right]\;.$

In _Deep Reinforcement Learning_ , deep neural networks are typically used as
function approximators to learn policies and value functions. For continuous
control tasks (that is, continuous state and/or action spaces as in PDE
control), _policy gradient methods_ (e.g., the _Proximal Policy Optimization
(PPO)_ [35], the _deep deterministic policy gradient (DDPG)_ [36] or the _Soft
Actor Critic (SAC)_ [37]) are the most prominent methods.

###  II-A Model-based reinforcement learning

Model-based algorithms use a dynamics model to capture the changes in the
environment, i.e., they approximate the state transition map
$\mathcal{T}:\mathcal{S}\times\mathcal{A}\rightarrow\mathcal{P}(\mathcal{S})$
by a surrogate $f_{\boldsymbol{\theta}}\approx\mathcal{T}$, where
$\boldsymbol{\theta}$ are trainable parameters.

The model-based framework extends the optimization procedure of model-free
reinforcement learning with additional planning [8] and model-learning steps.
Our work focuses on models learned from data for two reasons. The phenomena
defining the physics as well as their parameter values (for example mass or
viscosity parameters) are often not known exactly. At the same time,
computational models solving the governing PDEs often do not meet real-time
constraints, which is essential for online planning.

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 behavior. In
the ideal absence of prediction errors, the model matches the environment
without its downsides of real-world trials or severe computational costs.

###  II-B Reinforcement learning for PDE control

In the literature on partial differential equations, control problems are
usually defined on a continuous time scale (in contrast to the discrete-time
decision-making underpinning reinforcement learning). The system state
$\boldsymbol{u}:\Omega\times[0,T]\rightarrow\mathbb{R}^{n}$ is a function of
time $t\in[0,T]$ and space $x\in\Omega$, and the dynamics is described by a
nonlinear partial differential operator $\mathcal{N}$, i.e.,

$\frac{\partial\boldsymbol{u}}{\partial
t}=\mathcal{N}(\boldsymbol{u},\boldsymbol{\phi}),$

with $\boldsymbol{\phi}:\Omega\times[0,T]\rightarrow\mathbb{R}^{m}$ being the
control input that may depend on both space and time. Moreover, the initial
conditions are given by
$\mathcal{I}(\boldsymbol{u},\nabla_{\boldsymbol{x}}\boldsymbol{u},\ldots)$ and
the boundary conditions by
$\mathcal{B}(\boldsymbol{u},\nabla_{\boldsymbol{x}}\boldsymbol{u},\ldots)$.

In order to draw the connection to reinforcement learning (or numerical
implementations in general), one can introduce a partial discretization in
time with a constant step size $\Delta\tau=t_{\tau+1}-t_{\tau}$,
$\tau=0,1,\ldots,p$, i.e.,

$\mathcal{T}(\boldsymbol{u}_{\tau},\boldsymbol{\phi}_{\tau})=\boldsymbol{u}_{\tau}+\int_{t_{\tau}}^{t_{\tau+1}}\mathcal{N}(\boldsymbol{u}(\cdot,t),\boldsymbol{\phi}_{\tau})\,\mathrm{d}t=\boldsymbol{u}_{\tau+1},$

where we assume that $\boldsymbol{\phi}(\cdot,t)=\boldsymbol{\phi}_{\tau}$ is
constant over the interval $[t_{\tau},t_{\tau+1})$. Note that this time step
$\Delta\tau$ is typically much larger than the time step $\Delta t$ in a
standard numerical discretization scheme. Using the above considerations, the
control task can be formalized as an optimal control problem of the following
form:

$\displaystyle\min_{\boldsymbol{\phi}}J(\boldsymbol{u},\boldsymbol{\phi})$
$\displaystyle=\min_{\boldsymbol{\phi}}\sum_{\tau=0}^{p}\ell\left(\boldsymbol{u}_{\tau},\boldsymbol{\phi}_{\tau}\right)$
(1) $\displaystyle\mbox{s.t.}\quad\boldsymbol{u}_{\tau+1}$
$\displaystyle=\mathcal{T}(\boldsymbol{u}_{\tau},\boldsymbol{\phi}_{\tau}),\qquad\tau=0,1,2,\ldots,p-1,$

where $J$ is the objective functional over the time horizon $T=p\Delta\tau$,
and $\ell$ is the _stage cost_ , e.g., a tracking term (with regularization,
including penalties on the control cost)

$\displaystyle\ell\left(\boldsymbol{u}_{\tau},\boldsymbol{\phi}_{\tau}\right)$
$\displaystyle=\left\lVert\boldsymbol{u}_{\tau}-\boldsymbol{u}_{\tau}^{\mathsf{ref}}\right\rVert_{L^{2}}^{2}+\lambda\left\lVert\boldsymbol{\phi}_{\tau}\right\rVert_{L^{2}}^{2}.$

In terms of RL, the stage cost $\ell$ can be seen as the negative reward, the
state $\boldsymbol{u}$ corresponds to $\boldsymbol{s}$ and $\boldsymbol{\phi}$
is linked to $\boldsymbol{a}$.

###  II-C The Kuramoto-Sivashinsky equation

The Kuramoto-Sivashinsky equation emerges in various reaction-diffusion
systems and other physical phenomena, but was originally developed to model
instabilities in laminar flame fronts. Since it admits chaotic behavior from a
certain domain size on (and has similarities to the Navier-Stokes equations in
fluid flow), it is one of the most frequently studied PDE systems in many
situations. The equation is defined as

$\displaystyle\frac{\partial\boldsymbol{u}}{\partial
t}=-\nabla^{2}_{\boldsymbol{x}}\boldsymbol{u}-\nabla^{4}_{\boldsymbol{x}}\boldsymbol{u}-\frac{1}{2}\boldsymbol{u}\nabla_{\boldsymbol{x}}\boldsymbol{u}+\boldsymbol{\phi}\;,$
(2)

where $\boldsymbol{u}$ is the velocity and $\boldsymbol{\phi}$ is an additive
forcing term. Similar to other works [38, 34], we study a one-dimensional
spatial domain $\Omega=[0,L]$ with periodic boundary conditions and system
size $L=22$. This spatial configuration has been shown to manifest stable
chaotic conditions, having the smallest system size to develop sustained
(weak) turbulence [39]. Again following related work [38, 40, 34], the control
consists of a superposition of several Gaussians

$\displaystyle\boldsymbol{\phi}_{x,\tau\Delta\tau}$
$\displaystyle=\sum^{4}_{i=1}\frac{\boldsymbol{a}_{\tau}[i]}{\sqrt{2\pi\sigma}}\exp\bigg{(}-\frac{(\boldsymbol{x}-\boldsymbol{x}^{(i)})^{2}}{2\sigma^{2}}\bigg{)}\;,$
(3)

located at spatial coordinates $\boldsymbol{x}^{(i)}\in\\{0,L/4,2L/4,3L/4\\}$.
Here $\boldsymbol{a}_{\tau}[i]\in[-1,1]$ is the $i$-th control output of the
agent at time step $\tau$ and $\sigma=0.4$. The applied control values
$\boldsymbol{a}_{\tau}[i]$ are changed at intervals of $\Delta\tau=0.25$ time
units. Each episode simulates $T_{\mathrm{max}}=100$ time units of the system,
that is, $400$ discrete steps, beginning with states sampled from the unforced
attractor as an initial condition. In detail, we follow the configuration
outlined in [40] and initialize the solution variable $\boldsymbol{u}$ with
uniform noise in the range $[-0.4,0.4]$. The simulation then continues
developing the state evolution for $200$ time units (about 10 Lyapunov times).
During the transient, the system is not actuated.

In accordance with [34, 40], our goal is to dampen the dissipation $D$ of the
solution variable while minimizing the amount of energy $P$ spent to power the
system as well as the actuation devices:

$\begin{split}D&=\langle\,\big{[}\nabla_{\boldsymbol{x}}^{2}\boldsymbol{u}\,\big{]}^{2}\rangle\;\;\text{
and
}\end{split}\begin{split}\;\;P&=\langle\,\big{[}\nabla_{\boldsymbol{x}}\boldsymbol{u}\,\big{]}^{2}\rangle+\langle\,\boldsymbol{u}\boldsymbol{\phi}\,\rangle\;,\end{split}$
(4)

where $\langle\,\star\,\rangle$ denotes the spatial average taken over the
physical domain $\Omega$. The reward is thus

$r_{\tau}=-\int_{\tau\Delta\tau}^{\tau\Delta\tau+\Delta\tau}D(t)+P(t)\;\,dt\;.$
(5)

For the numerical solution, we implement the finite difference simulation
outlined in [33] and use a fourth-order Runge-Kutta scheme with a fixed time
step of $\Delta t=0.001$. The second-order dissipation and fourth-order
hyperviscous damping terms are discretized by a sixth-order central difference
scheme, while a second-order upwind scheme approximates the convection term.
The spatial domain is discretized into $N=64$ equidistant collocation points
whose solution values constitute the state vector $\boldsymbol{s}$ observed
for decision-making.

##  III Learning surrogate models of forced PDEs

Before integrating our surrogate into the online learning process of model-
based RL, we evaluate its design on an offline dataset. In order to do so, we
discuss the different modeling steps in detail and assess their importance in
an ablation study on the Kuramoto-Sivashinsky system.

###  III-A Dimensionality reduction

For learning surrogates, the spatial extent of the physical domain $\Omega$
and the dimensionality of its discretization affects the amount of data
necessary in training as well as the degree to which the surrogate matches the
state evolution of the system. To achieve a dimensionality reduction, we use a
_convolutional autoencoder (CAE)_ architecture, which is particularly well
suited to cover the (on shorter time scales) local physics of PDEs. Another
merit of CAEs is the weight-sharing capability, which significantly decreases
the number of model parameters $\boldsymbol{\theta}$ and has a self-
regularizing effect. The architecture consists of an encoder and a decoder
network, as illustrated by the yellow and dark blue blocks in Figure 2. The
encoder network $f_{\boldsymbol{\theta}_{\mathrm{enc}}}$ compresses snapshots
$\boldsymbol{s}_{\tau}$ of the system to a compact latent space
$\boldsymbol{h}_{\tau}=f_{\boldsymbol{\theta}_{\mathrm{enc}}}(\boldsymbol{s_{\tau}})$,
while the decoder network $f_{\boldsymbol{\theta}_{\mathrm{dec}}}$ attempts to
recover the original input
$\hat{\boldsymbol{s}}_{\tau}=f_{\boldsymbol{\theta}_{\mathrm{dec}}}(\boldsymbol{h}_{\tau})$.
The information bottleneck between the encoder and decoder networks regulates
the degree to which the state dimensionality is reduced. Supposing that the
reconstruction loss $\mathcal{L}_{\mathrm{MSE}}$ is small, the latent variable
$\boldsymbol{h}_{\tau}$ holds similar information as the original input
$\boldsymbol{s}_{\tau}$ but in denser form. More details on the specifics of
the network architecture can be found in Appendix -A.

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 decoder network
$f_{\boldsymbol{\theta}_{\mathrm{dec}}}$ (blue) that restores the spatial
extent of its output to the physical domain. Starting off an initial condition
(left), the model predicts the state evolution using its prior output as an
input of the state variable (right). The figure does not show intermediate
scaling or normalization transformations for clarity.

###  III-B Time stepping

To take full advantage of the fact that both the latent state
$\boldsymbol{h}_{\tau}$ and the auto-encoded control $\boldsymbol{v}_{\tau}$
(Appendix -A) act in the same spatial domain, albeit downsampled, we use a
convolutional LSTM model [41] to transition in latent space. Combining all
components outlined so far, we pass solution variables
$\boldsymbol{u}_{\,\tau\Delta\tau}$ and forcing terms
$\boldsymbol{\phi}_{\,\tau\Delta\tau}$ to the encoder networks
$f_{\boldsymbol{\theta}_{\mathrm{enc}}}$ and
$f_{\boldsymbol{\psi}_{\mathrm{enc}}}$ and compute
$\boldsymbol{h}_{\tau}=f_{\boldsymbol{\theta}_{\mathrm{enc}}}(\boldsymbol{u}_{\,\tau\Delta\tau})$
and
$\boldsymbol{v}_{\tau}=f_{\psi_{\mathrm{enc}}}(\boldsymbol{\phi}_{\,\tau\Delta\tau})$
before transitioning in latent space with
$f_{\boldsymbol{\theta}_{\mathrm{fwd}}}$ to predict the deterministic state
evolution of the decision-making process. Similar to other works [42, 43, 44],
the decoder network $f_{\boldsymbol{\theta}_{\mathrm{dec}}}$ is tied into a
network learning temporal residuals, i.e., it is used to predict state changes
$\Delta\boldsymbol{s}_{\tau}=\boldsymbol{s}_{\tau+1}-\boldsymbol{s}_{\tau}$ of
solution variables. State changes are scaled with the time duration
$\Delta\tau$ of control steps and modeled in latent space. Therefore, our
network architecture infers state transitions as

$\displaystyle\boldsymbol{\hat{u}}_{\,\tau\Delta\tau+\Delta\tau}$
$\displaystyle=\boldsymbol{u}_{\,\tau\Delta\tau}+\Delta\tau\cdot
f_{\boldsymbol{\theta}_{\mathrm{dec}}}(\boldsymbol{o}_{\tau})\;,$ (6)
$\displaystyle\text{with }\;\;\boldsymbol{o}_{\tau}$
$\displaystyle=f_{\boldsymbol{\theta}_{\mathrm{fwd}}}(\boldsymbol{h}_{\tau},\boldsymbol{v}_{\tau};\boldsymbol{c}_{\tau})\;,$
(7)

where targets $\Delta\boldsymbol{s}_{\tau}/\Delta\tau\approx
f_{\boldsymbol{\theta}_{\mathrm{dec}}}(\boldsymbol{o}_{\tau})$ are used to
train the transition model $f_{\boldsymbol{\theta}_{\mathrm{fwd}}}$ via end-
to-end backpropagation. The variable $\boldsymbol{c}_{\tau}$ is the cell state
of our recurrent transition model. The overall network architecture is
illustrated in Figure 2. It shows that, in each step, our model (1) transforms
control outputs $\boldsymbol{a}$ to external forcing terms $\boldsymbol{\phi}$
using a known functional, (2) encodes solution variables $\boldsymbol{u}$ and
forcing terms $\boldsymbol{\phi}$ to a latent space defined on a shared
spatial domain with encoder neural networks (yellow and red components), (3)
uses a transition model (green component) to infer encoded state changes
$\boldsymbol{o}$, and (4) decodes temporal changes $\boldsymbol{o}$ (blue
component) before adding them to the original input. In order to use our model
as a surrogate of the system, it is unrolled in time, taking the control
output of an agent as well as its prior prediction for an input.

A notable benefit of learning temporal residuals is that the network’s
predictions obey the initial conditions of the system by construction. In
fact, modeling state changes scaled with coefficient $\Delta\tau$ imposes a
temporal smoothness assumption on the predictions to thwart irregular and
erratic changes in the state evolution. In the context of model-based control,
our temporal smoothness assumption offers a significant advantage. Since the
control loop intertwines data collection, model learning, and behavior
improvement, only a limited number of system snapshots are available to the
surrogate at first. As we confirmed in our experiments, models predicting
states in place of state changes often suffer from compounding model errors in
the small data regime due to their dependence on well-functioning encoder and
decoder networks. In the small data regime, said neural networks are often
still unable to encode and decode input states to a sufficient degree of
accuracy, which can result in large model errors that can cause catastrophic
learning updates [23] for the agent. In contrast, the additive changes of our
model architecture couple its predictions to the initial condition and enforce
a close proximity to the temporal state evolution of the system for at least a
few steps of the prediction until model errors begin compounding.

##  IV Learning to control PDEs sample efficiently

As a member of the _Dyna_ family of model-based algorithms [45, 15], our
approach implements a model-free agent to learn behavior
$\pi_{\boldsymbol{\omega}}$ and improve it over time. In essence, all
descendants of Dyna follow instructions similar to those outlined in Alg. 1.
In a broad sense, the algorithm intertwines data collection, model learning,
and behavior improvement for its main steps, as shown in Alg. 1. Each
iteration begins with rolling out $\pi_{\boldsymbol{\omega}}$ to collect
additional samples that are later stored as a dataset
$\mathcal{D}_{\mathrm{env}}$. After training a deep neural network
$f_{\boldsymbol{\theta}}$ to learn the dynamics of the environment on samples
stored in $\mathcal{D}_{\mathrm{env}}$, the model serves as an approximate MDP
$\mathcal{M}^{\prime}$ of the actual system $\mathcal{M}$. Supplementary
samples from model-based rollouts of $\pi_{\boldsymbol{\omega}}$ are then
stored in a separate dataset $\mathcal{D}_{\mathrm{model}}$ as proxies for the
actual experience. Iterations of our algorithm conclude with an update of the
model-free agent. Here, data stored in $\mathcal{D}_{\mathrm{env}}$ as well as
$\mathcal{D}_{\mathrm{model}}$ is used, and the latter amount often exceeds
the other by multiple orders of magnitude.

Algorithm 1 Model-Based Reinforcement Learning

1:Initialize parameters of policy $\pi_{\boldsymbol{\omega}}$ and model
$f_{\boldsymbol{\theta}}$

2:Initialize empty datasets $\mathcal{D}_{\mathrm{env}}$ and
$\mathcal{D}_{\mathrm{model}}$

3:while not done do

4: Sample $\mathcal{M}$ using policy $\pi_{\boldsymbol{\omega}}$ $\rightarrow$
add to $\mathcal{D}_{\mathrm{env}}$

5: Train model $f_{\boldsymbol{\theta}}$ on dataset
$\mathcal{D}_{\mathrm{env}}$

6: Sample $\mathcal{M}^{\prime}$ using policy $\pi_{\boldsymbol{\omega}}$
$\rightarrow$ add to $\mathcal{D}_{\mathrm{model}}$

7: Update policy $\pi_{\boldsymbol{\omega}}$ using samples of
$\mathcal{D}_{\mathrm{model}}$ and $\mathcal{D}_{\mathrm{env}}$

8:end while

###  IV-A The model-based RL framework

A concrete implementation of the model-based principle is, therefore, all
about (1) the design of a surrogate and its training process, (2) the way in
which $\mathcal{D}_{\mathrm{model}}$ is populated with artificial samples, and
(3) the model-free agent updating $\pi_{\boldsymbol{\omega}}$ with data stored
in $\mathcal{D}_{\mathrm{env}}$ and fictitious samples in
$\mathcal{D}_{\mathrm{model}}$. In the following, we outline our approach to
each of the above steps for model-based reinforcement learning.

####  IV-A1 Model learning

Unlike related work on model-based fluid control, we do not collect snapshots
of the system in advance using random exploration. Instead, we alternate
between data collection and model learning to align the distribution of states
that PI visits with the data our surrogate is trained on and approximates
well. Our experiments (Section V-B) suggest that an online adaptation of the
model to changes in the behavior is essential for the model to maintain its
relevance. After each step of data collection, We train the model
$f_{\boldsymbol{\theta}}$ using the samples collected in
$\mathcal{D}_{\mathrm{env}}$. To mitigate overfitting in the small data
regime, we monitor the model on a validation set $\mathcal{D}_{\mathrm{val}}$
and stop training early once the validation loss $\mathcal{L}_{\mathrm{val}}$
converges. In order to encourage an adaptation of the model to recent
experience but protect $\pi_{\boldsymbol{\omega}}$ against severe changes of
inputs during policy training, we limit the minimum $P_{\mathrm{min}}$ and
maximum $P_{\mathrm{max}}$ number of gradient updates to the model. We found
that striking a balance between regular model updates but keeping shifts in
the distribution of artificial samples minor helped stabilize the
optimization. In our experiments, continuous updates of the model each
performing only a few gradient steps (small values for $P_{\mathrm{max}}$ and
$P_{\mathrm{val}}$) generally worked well and outperformed updating the model
for an unlimited number of steps. We also found that using a curriculum
(similar to [46]) for the data collection is beneficial. We use short model-
based rollouts of policy $\pi_{\boldsymbol{\omega}}$ for artificial data
collection early on, and increase their prediction horizon once more data is
available [24]. In other words, we optimize the model for short-term
performance in the small data regime and emphasize learning global dynamics
later on.

####  IV-A2 Model-Based Rollouts

Consistent with standard practice [22, 21], we mitigate the risk of model
exploitation using an ensemble
$\\{f_{\boldsymbol{\theta}_{1}},\dots,f_{\boldsymbol{\theta}_{L_{\mathrm{ens}}}}\\}$
of dynamics models. The models deviate not only by their initial weights and
the ordering of mini-batches, but also in terms of the data used for training
and validation. Indeed, using a bootstrapping mechanism for training, [21]
find that ensembles isolate epistemic uncertainty through disagreement. Since
epistemic uncertainty is due to data scarcity, ensembling helps identifying
regions of the state space where overfitting is likely and forecast errors are
substantial. As is standard practice (see, e.g., [24]), we make use of the
elite mechanism for the ensemble. More details on our ensemble techniques can
be found in Appendix -B.

Using our ensemble of dynamics models
$\\{f_{\boldsymbol{\theta}_{1}},\dots,f_{\boldsymbol{\theta}_{L_{\mathrm{ens}}}}\\}$,
we define an approximate MDP $\mathcal{M}^{\prime}$ imitating the original
system $\mathcal{M}$. Analogous to recent works on Dyna algorithms, model-
based rollouts in our implementation branch off arbitrary system states
sampled from dataset $\mathcal{D}_{\mathrm{env}}$. Unlike rollouts beginning
at initial conditions of the decision-making process, branching rollouts off
arbitrary starting states avoids compounding model errors for states visited
during later stages of episodes. In place of sampling single starting states,
we select state-action sequences $(\boldsymbol{s}_{\tau-
K_{\mathrm{tf}}},\boldsymbol{a}_{\tau-
K_{\mathrm{tf}}},\dots,\boldsymbol{s}_{\tau},\boldsymbol{a}_{\tau})\sim\mathcal{D}_{\mathrm{env}}$
of length $K_{\mathrm{tf}}$. Sampling sequences enables us to warm-start the
memory unit of our recurrent transition model $f_{\boldsymbol{\theta}}$. Each
state-action pair of the sequence is then processed in teacher-forcing mode to
guide the transitions. We match the sequence length $K_{\mathrm{tf}}$ to the
number of teacher-forcing steps during training. In our experiments, using
$K_{\mathrm{tf}}$ state transitions to seed the cell state was essential since
we did not backpropagate gradients to its initial state at training time. Vice
versa, we do not learn the initial condition of the memory unit since model-
based rollouts start in states at different points in time for the system,
which otherwise destabilized the algorithm in our experiments. In order to
implement a similar mechanism for the first $K_{\mathrm{tf}}$ steps of an
episode, we use a simple padding method and repeat the initial state-action
tuple to extend the sequence to length $K_{\mathrm{tf}}$. A longer discussion
on our rollout approach has been shifted to Appendix -C.

####  IV-A3 Policy optimization

As our approach belongs to the Dyna family, arbitrary model-free agents can be
used to derive behavior given artificial samples. An important design choice
for our purpose is whether to optimize the behavior using on-policy or off-
policy methods. Indeed, members of the former category discard existing
experience after each update to the policy. Despite our dynamics models being
designed as lightweight surrogates of the numerical simulation, reusing
already existing experience is still more efficient. Adopting the soft-actor
critic [37] as our off-policy method (as in _MBPO_ [24]) the samples our model
synthesizes in each iteration can be kept in dataset
$\mathcal{D}_{\mathrm{model}}$ for a number of repetitions. Using artificial
samples, the model-based algorithm updates the soft-actor critic
$L_{\mathrm{updates}}$ times per iteration. Unlike the model-free framework,
parameter $L_{\mathrm{updates}}$ is set in such a way that multiple updates
(often between 20 and 40, see, e.g., [24]) are performed per environment
interaction. To ensure the stability of the agent even after updating the
behavior more than once, we collect a large number of artificial samples
$N_{\mathrm{model}}$ in each iteration (in [24], for example, 400 model
rollouts are performed for each sample of the environment). After model-based
data collection, the agent uses the samples stored in
$\mathcal{D}_{\mathrm{env}}$ and $\mathcal{D}_{\mathrm{model}}$ to update its
state-action value estimates and policy.

Our work builds on code made openly available in [47] and [48]. The soft-actor
critic agent implements double Q-learning[49] to learn estimates of state-
action values with separate networks $Q_{\boldsymbol{\theta}}$ and
$Q_{\boldsymbol{\theta}^{\prime}}$ to mitigate selection biases and uses
delayed target networks $Q_{\boldsymbol{\theta}^{-}}$ to stabilize the
optimization towards a moving target. Soft updates of
$Q_{\boldsymbol{\theta}^{-}}$ through Polyak averaging [50] keep the target
parameters consistent [11]. Similar to the TD3 algorithm [51], multiple
gradient updates of the critic are taken before altering the actor to ensure
up-to-date value estimates. The actor $\pi_{\boldsymbol{\omega}}$
parameterizes the mean and covariance of a Gaussian distribution defined over
the continuous action space. Our implementation does not use automatic entropy
tuning [37] but sets the temperature to a constant value instead. Both the
actor and critic are implemented as multilayer perceptrons of three layers
with $256$ hidden units each.

##  V Experimental Evaluation

Corresponding to Sections III and IV, we first evaluate the capabilities of
our surrogate to approximate the state evolution well when the data is
collected in advance, before considering the control aspect. More details on
the training of our model can be found in Appendix -D.

###  V-A Data-driven prediction of forced PDEs

In order to evaluate the amount of data our surrogate needs to approximate the
state evolution, we collect a dataset of $100$ simulation episodes ($40,000$
snapshots), where actions are sampled uniform at random. We split the samples
into training, validation (80% / 20%), and testing episodes. All results are
the aggregate of cross-validation scores with five folds. In general, we study
the evolution of state variable $\boldsymbol{u}$ for subsequences with
$K_{\mathrm{max}}=30$ consecutive prediction steps ($7.5$ time units) during
training and testing.

Figure 3: Prediction error of the Kuramoto-Sivashinsky equation. Each
datapoint averages the error for different testing subsequences. The models
are trained on an increasing share of the dataset composed of $40,000$
snapshots. Solid lines depict the mean of five folds, while the shaded regions
show the 95% confidence interval.

Since our later approach to model-based approach aims to use as few samples as
possible, the small data regime is critical to our work. To assess the
alignment of our model design with our overall goal, we systematically
decrease the amount of data available throughout cross-validation from $100\%$
down to $10\%$. In fact, our dataset (including validation and testing data)
equals the number of snapshots used in the work of [34] for modeling the
Kuramoto-Sivashinsky equation upfront before they derive a control strategy
with the surrogate afterward. In order to isolate the contribution of single
design decisions for our dynamics model, we introduce several ablations and
compare our results with those of a simple baseline, namely (1) the usage of
fully connected networks instead of CNNs and (2) the prediction of the state
instead of a residual update. We train each model for a maximum of $250$
epochs or terminate the optimization once the monitored loss does not decrease
for $P_{\mathrm{val}}=25$ validation epochs.

Figure 4: Comparison of the temporal state evolution of the Kuramoto-
Sivashinsky equation with forecasts of our model and the different ablations.

In Figure 3, we show the normalized mean squared state prediction error,
averaged over various initial conditions. Unsurprisingly, convolutional models
(top row) prove far more effective than fully connected networks (bottom row)
in both the small as well as large data regime, learning temporal residuals or
not. In terms of the comparison between residual and full-state prediction
(left vs. right column), the former yields a notable advantage for forecasts
over short time durations with error margins being negligible for several time
units. In the small data regime, the full-state prediction is not yet equipped
to recover state variables and, therefore, suffers from large error margins. A
notable advantage of our approach is thus the strongly reduced data
requirements supposing we limit the prediction horizon appropriately. Since
our RL approach will limit model usage after a small number of prediction
steps, the inferior accuracy after $\sim 3.25$ time units (13 discrete steps)
is of no concern to the optimization. The errors in the predicted rewards
(which have a significant impact on the learning performance) are very similar
to the state prediction. An example of forecasts for each model at an
increasing share of the $40,000$ snapshots for training is shown in Figure 4.

In general, increasing the degrees of freedom of a model enables it to learn
the state evolution of snapshots in greater detail. At the same time, models
of higher capacity are at a greater risk of overfitting to patterns in the
training set, contriving substantial errors once presented with data withheld
for testing. Considering this trade-off is essential for model-based RL since
overfitting is among its greatest challenges [21, 30, 22]. A study on this
dependency can be found in Appendix -E. For the main body of this paper, all
models have the identical structure, which we found to be a good trade off
between expressive capabilities and complexity.

###  V-B Reinforcement learning control of PDEs

At last, we evaluate the ability of our approach to model-based reinforcement
learning to stabilize systems with dynamics governed by PDEs. In this section,
we study whether (1) learning a surrogate of the global dynamics indeed
mitigates the data consumption of model-free approaches, (2) learned control
strategies can match their effectiveness, although model-based reinforcement
learning algorithms introduce approximation errors to the experience used for
behavior improvement, (3) examine the contribution of single components to our
overall approach in isolation, and (4) evaluate our algorithm for a number of
configurations.

####  V-B1 Data efficiency of learned control laws

To assess the capabilities of our control algorithm in mitigating the data
consumption of the model-free RL framework, we compare accomplished episode
returns to those of the popular PPO [35] and SAC [37] algorithms. As outlined
before, the former algorithm belongs to the family of controllers most often
encountered in the fluid mechanics literature nowadays (see, e.g., [52] for a
comprehensive survey), despite the agent discarding past experiences after
each update. Since our approach integrates SAC as a main component for policy
improvement, it is the candidate best suited to compare against. In both
cases, we use the open-source implementations made available by the stable-
baselines3 project [53] as well as their default configuration for our
experiments.

Figure 5: Average evaluation episode returns for our approach compared to SAC
and PPO for up to $50,000$ (left) and $1,000,000$ (right) steps of training.
Solid lines show the mean of three trials, while shaded regions denote the
standard deviation among trials. All lines are smoothed using a Gaussian
filter with $\sigma=1$. The dotted lines on the right-hand side show the
average performance of our model-based algorithm after termination at $50,000$
time steps.

Figure 5 shows the average return of evaluation episodes with respect to the
number of samples collected for training thus far. Due to computational costs,
we terminate the execution of our model-based approach once it obtained
$50,000$ samples. Since computational workloads per interaction with the
system are lower for model-free algorithms, we execute our baselines for
$1,000,000$ samples. On the left-hand side of the figure, we show the training
curves up to the point where we terminate our algorithm, while the right-hand
side shows the asymptotic behavior of SAC and PPO. We find that our model-
based approach indeed learns suitable control strategies using substantially
fewer samples. On average, SAC takes about $430,000$ samples to break even
with the average episode returns of our approach at $50,000$ samples, while
PPO cannot match its performance even after $1,000,000$ steps. With an almost
nine-fold improvement in terms of data consumption over SAC, this demonstrates
that model-based RL is indeed a promising direction to learn control
strategies for systems governed by PDEs. Even though we stopped training early
after $50,000$ steps, the control strategy dampened the dissipation $D$ and
power consumption $P$ of the system by more than $63\%$ of the values
accomplished with random forcing. At the same time, SAC converges to a value
of about $73\%$, although using $20$ times the amount of data to do so. An
example of what actuations the controller applies to navigate the system
towards a stable state after termination is illustrated in Figure 6.

Figure 6: The spatio-temporal evolution of the Kuramoto-Sivashinsky system
(middle) and its dissipative term $D$ and power consumption $P$ (left) as a
result of the actions applied to it (right). In each row, we show an agent
attempting to stabilize the dynamics starting from an initial condition
sampled at random. The figure shows control strategies after their training
terminated after $50,000$ or $1,000,000$ steps, respectively.

####  V-B2 Ablation study

A part of the motivation for our work are the shortcomings of past studies on
model-based control for governed systems insofar as recent advances in the
reinforcement learning literature are not properly taken into account. As
outlined before, the related works [33] and [34], for example, use neither
ensembling nor rollout scheduling and in the latter case do not adapt the
model to altering state distributions, despite each feature being considered
standard practice in other fields. In order to appraise the contribution of
each component to the overall performance of our method in isolation, our work
examines the following ablations of our model-based approach.

  * •

_Offline model training ablation:_ Similar to [34], we collect 50,000
snapshots of the system in advance using random exploration. After training
our ensemble to mimic the global dynamics, the RL agent acts entirely within
the surrogate for data collection.

  * •

_Model exploitation ablation:_ Our approach establishes multiple mechanisms to
prevent compounding model errors from propagating into policy improvement
steps. Similar to [33] and [34], this ablation neither uses ensembling nor
rollout scheduling to diversify or truncate rollouts in an adaptative manner.
In contrast to our surrogate, the model is trained for a fixed number of
gradient steps in each iteration.

  * •

_Surrogate ablation:_ In place of our residual surrogate, we use the full-
state prediction model to approximate the temporal state evolution for the
algorithm as presented in Section III.

Figure 7 compares our implementation of the model-based methodology to its
ablations. It shows that an improper treatment of compounding model errors
often leads to unstable behavior optimization (purple line). Out of all
methods we have tried for learning control strategies in our experiments, we
found the _exploitation ablation_ to be the least stable, as characterized by
its wide standard deviation. Implementing model ensembling and using a
schedule to determine the length of model-based rollouts is seemingly
essential to prevent the agent from taking advantage of model errors in the
small data regime. The figure also shows that a similar issue pertains to our
_surrogate ablation_ (green line). Although convergence is more stable
overall, its approximate MDP cannot faithfully match the temporal evolution of
the PDE in the small data regime so that behavior improvements of the agent
for the surrogate do not generalize to the system. Specifically, as a result
of the _surrogate ablation_ learning solution variables in place of temporal
residuals, its forecasts sustain substantial errors at the beginning of data
collection. Our results also suggest that random exploration does not cover
all relevant regions of the state space that policies visit during their
optimization process (red line). After an initial phase of improvement, the
behavior of our _offline ablation_ diverges more and more from the state
distribution that the model used for training, increasing prediction errors
and destabilizing the optimization once the policy begins exploiting them.
Indeed, while trials are relatively stable at first, their standard deviation
dilates with an increasing number of algorithm iterations.

Figure 7: Average evaluation episode returns for our model-based approach
compared to different ablations. Solid lines show the mean of three trials,
while shaded regions denote the standard deviation among trials. All lines are
smoothed using a Gaussian filter with $\sigma=1$. The dotted lines on the
right-hand side show the average performance of the left-hand algorithms after
termination at $50,000$ time steps. Since our _offline ablation_ does not
collect additional samples, we show its performance throughout iterations of
the optimization.

##  VI Discussion and outlook

We have seen that carefully constructed surrogate models are capable of
increasing the sample efficiency by roughly one order of magnitude.
Furthermore, we found that an overall more sophisticated approach to model-
based reinforcement learning (online model adaptation, ensembling, curriculum
learning, etc.) is beneficial to stabilize the convergence, and we hope that
our findings showcase the appeal of adopting best practices. However, due to
the increased complexity, the performance of model-based RL is still inferior
in situations where the model is comparatively easy to simulate. For future
work, it will thus be highly interesting to see whether simpler yet more
efficient surrogate models can be utilized. At this point, suitable candidates
appear to be GRUs instead of LSTMs (see, e.g., [54]), sparse regreesion
techinques to identify governing equations [55, 6], and in particular the
highly popular Koopman operator [56, 57, 58], as it allows us to learn linear
models of nonlinear systems, which is very efficient both in terms of the
required training data and the run time. Finally, it might be worth looking
into recent prediction error results for these methods [59, 60, 61] and see
whether they can be transferred into guarantees for the RL process.

## Acknowledgment

S.P. acknowledges financial support by the project “SAIL: SustAInable Life-
cycle of Intelligent Socio-Technical Systems” (Grant ID NW21-059A), which is
funded by the program “Netzwerke 2021” of the Ministry of Culture and Science
of the State of Northrhine Westphalia, Germany.

## References

  * [1] S. Aubrun, A. Leroy, and P. Devinant, “A review of wind turbine-oriented active flow control strategies,” _Experiments in Fluids_ , vol. 58, no. 10, pp. 1–21, 2017.
  * [2] Z. Wu, D. Fan, Y. Zhou, R. Li, and B. R. Noack, “Jet mixing optimization using machine learning control,” _Experiments in Fluids_ , vol. 59, no. 8, pp. 1–17, 2018.
  * [3] J. Degrave, F. Felici, J. Buchli, M. Neunert, B. Tracey, F. Carpanese, T. Ewalds, R. Hafner, A. Abdolmaleki, D. de Las Casas _et al._ , “Magnetic control of tokamak plasmas through deep reinforcement learning,” _Nature_ , vol. 602, no. 7897, pp. 414–419, 2022.
  * [4] L. Grüne and J. Pannek, _Nonlinear model predictive control_. Springer, 2017.
  * [5] K. Kunisch and S. Volkwein, “Galerkin proper orthogonal decomposition methods for a general equation in fluid dynamics,” _SIAM Journal on Numerical Analysis_ , vol. 40, pp. 492–515, 2002.
  * [6] E. Kaiser, J. N. Kutz, and S. L. Brunton, “Sparse identification of nonlinear dynamics for model predictive control in the low-data limit,” _Proceedings of the Royal Society A_ , vol. 474, no. 2219, p. 20180335, 2018\.
  * [7] K. Bieker, S. Peitz, S. L. Brunton, J. N. Kutz, and M. Dellnitz, “Deep model predictive flow control with limited sensor data and online learning,” _Theoretical and computational fluid dynamics_ , vol. 34, no. 4, pp. 577–591, 2020.
  * [8] R. S. Sutton and A. G. Barto, _Reinforcement learning: An introduction_. MIT press, 2018.
  * [9] D. Silver, A. Huang, C. J. Maddison, A. Guez, L. Sifre, G. Van Den Driessche, J. Schrittwieser, I. Antonoglou, V. Panneershelvam, M. Lanctot _et al._ , “Mastering the game of go with deep neural networks and tree search,” _nature_ , vol. 529, no. 7587, pp. 484–489, 2016.
  * [10] J. Rabault, M. Kuchta, A. Jensen, U. Réglade, and N. Cerardi, “Artificial neural networks trained through deep reinforcement learning discover control strategies for active flow control,” _Journal of fluid mechanics_ , vol. 865, pp. 281–302, 2019.
  * [11] V. Mnih, K. Kavukcuoglu, D. Silver, A. A. Rusu, J. Veness, M. G. Bellemare, A. Graves, M. Riedmiller, A. K. Fidjeland, G. Ostrovski _et al._ , “Human-level control through deep reinforcement learning,” _nature_ , vol. 518, no. 7540, pp. 529–533, 2015.
  * [12] B. Zoph and Q. V. Le, “Neural architecture search with reinforcement learning,” _arXiv:1611.01578_ , 2016.
  * [13] C. Berner, G. Brockman, B. Chan, V. Cheung, P. Debiak, C. Dennison, D. Farhi, Q. Fischer, S. Hashme, C. Hesse _et al._ , “Dota 2 with large scale deep reinforcement learning,” _arXiv:1912.06680_ , 2019.
  * [14] P. Ma, Y. Tian, Z. Pan, B. Ren, and D. Manocha, “Fluid directed rigid body control using deep reinforcement learning,” _ACM Transactions on Graphics (TOG)_ , vol. 37, no. 4, pp. 1–11, 2018.
  * [15] R. S. Sutton, “Dyna, an integrated architecture for learning, planning, and reacting,” _ACM Sigart Bulletin_ , vol. 2, no. 4, pp. 160–163, 1991.
  * [16] T. Wang, X. Bao, I. Clavera, J. Hoang, Y. Wen, E. Langlois, S. Zhang, G. Zhang, P. Abbeel, and J. Ba, “Benchmarking model-based reinforcement learning,” _arXiv:1907.02057_ , 2019.
  * [17] C. G. Atkeson and J. C. Santamaria, “A comparison of direct and model-based reinforcement learning,” in _Proceedings of international conference on robotics and automation_ , vol. 4. IEEE, 1997, pp. 3557–3564.
  * [18] M. Deisenroth and C. E. Rasmussen, “Pilco: A model-based and data-efficient approach to policy search,” in _Proceedings of the 28th International Conference on machine learning (ICML-11)_. Citeseer, 2011, pp. 465–472.
  * [19] T. M. Moerland, J. Broekens, and C. M. Jonker, “Model-based reinforcement learning: A survey,” _arXiv:2006.16712_ , 2020.
  * [20] H. W. Andersen and M. Kümmel, “Evaluating estimation of gain directionality: Part 1: Methodology,” _Journal of Process Control_ , vol. 2, no. 2, pp. 59–66, 1992.
  * [21] K. Chua, R. Calandra, R. McAllister, and S. Levine, “Deep reinforcement learning in a handful of trials using probabilistic dynamics models,” _Advances in neural information processing systems_ , vol. 31, 2018.
  * [22] T. Kurutach, I. Clavera, Y. Duan, A. Tamar, and P. Abbeel, “Model-ensemble trust-region policy optimization,” _arXiv:1802.10592_ , 2018.
  * [23] H. P. Van Hasselt, M. Hessel, and J. Aslanides, “When to use parametric models in reinforcement learning?” _Advances in Neural Information Processing Systems_ , vol. 32, 2019.
  * [24] M. Janner, J. Fu, M. Zhang, and S. Levine, “When to trust your model: Model-based policy optimization,” _Advances in Neural Information Processing Systems_ , vol. 32, 2019.
  * [25] S. Levine and P. Abbeel, “Learning neural network policies with guided policy search under unknown dynamics,” _Advances in neural information processing systems_ , vol. 27, 2014.
  * [26] R. Lioutikov, A. Paraschos, J. Peters, and G. Neumann, “Sample-based informationl-theoretic stochastic optimal control,” in _2014 IEEE International Conference on Robotics and Automation (ICRA)_. IEEE, 2014, pp. 3896–3902.
  * [27] X. Xiong, F. Wörgötter, and P. Manoonpong, “Neuromechanical control for hexapedal robot walking on challenging surfaces and surface classification,” _Robotics and Autonomous Systems_ , vol. 62, no. 12, pp. 1777–1789, 2014.
  * [28] M. Seeger, “Gaussian processes for machine learning,” _International journal of neural systems_ , vol. 14, no. 02, pp. 69–106, 2004.
  * [29] D. Hafner, T. Lillicrap, I. Fischer, R. Villegas, D. Ha, H. Lee, and J. Davidson, “Learning latent dynamics for planning from pixels,” in _International conference on machine learning_. PMLR, 2019, pp. 2555–2565.
  * [30] A. Nagabandi, G. Kahn, R. S. Fearing, and S. Levine, “Neural network dynamics for model-based deep reinforcement learning with model-free fine-tuning,” in _2018 IEEE International Conference on Robotics and Automation (ICRA)_. IEEE, 2018, pp. 7559–7566.
  * [31] L. Kaiser, M. Babaeizadeh, P. Milos, B. Osinski, R. H. Campbell, K. Czechowski, D. Erhan, C. Finn, P. Kozakowski, S. Levine _et al._ , “Model-based reinforcement learning for atari,” _arXiv:1903.00374_ , 2019.
  * [32] D. Hafner, T. Lillicrap, J. Ba, and M. Norouzi, “Dream to control: Learning behaviors by latent imagination,” _arXiv:1912.01603_ , 2019.
  * [33] X.-Y. Liu and J.-X. Wang, “Physics-informed dyna-style model-based deep reinforcement learning for dynamic control,” _Proceedings of the Royal Society A_ , vol. 477, no. 2255, p. 20210618, 2021.
  * [34] K. Zeng, A. J. Linot, and M. D. Graham, “Data-driven control of spatiotemporal chaos with reduced-order neural ode-based models and reinforcement learning,” _arXiv:2205.00579_ , 2022.
  * [35] J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov, “Proximal policy optimization algorithms,” _arXiv:1707.06347_ , 2017.
  * [36] T. P. Lillicrap, J. J. Hunt, A. Pritzel, N. Heess, T. Erez, Y. Tassa, D. Silver, and D. Wierstra, “Continuous control with deep reinforcement learning,” _arXiv:1509.02971_ , 2015.
  * [37] T. Haarnoja, A. Zhou, P. Abbeel, and S. Levine, “Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor,” in _International conference on machine learning_. PMLR, 2018, pp. 1861–1870.
  * [38] M. A. Bucci, O. Semeraro, A. Allauzen, G. Wisniewski, L. Cordier, and L. Mathelin, “Control of chaotic systems by deep reinforcement learning,” _Proceedings of the Royal Society A_ , vol. 475, no. 2231, p. 20190351, 2019\.
  * [39] P. Cvitanović, R. L. Davidchack, and E. Siminos, “On the state space geometry of the kuramoto–sivashinsky flow in a periodic domain,” _SIAM Journal on Applied Dynamical Systems_ , vol. 9, no. 1, pp. 1–33, 2010.
  * [40] K. Zeng and M. D. Graham, “Symmetry reduction for deep reinforcement learning active control of chaotic spatiotemporal dynamics,” _Physical Review E_ , vol. 104, no. 1, p. 014210, 2021.
  * [41] X. Shi, Z. Chen, H. Wang, D.-Y. Yeung, W.-K. Wong, and W.-c. Woo, “Convolutional lstm network: A machine learning approach for precipitation nowcasting,” _Advances in neural information processing systems_ , vol. 28, 2015.
  * [42] S. Mo, N. Zabaras, X. Shi, and J. Wu, “Deep autoregressive neural networks for high-dimensional inverse problems in groundwater contaminant source identification,” _Water Resources Research_ , vol. 55, no. 5, pp. 3856–3881, 2019.
  * [43] N. Geneva and N. Zabaras, “Modeling the dynamics of pde systems with physics-constrained deep auto-regressive networks,” _Journal of Computational Physics_ , vol. 403, p. 109056, 2020.
  * [44] P. Ren, C. Rao, Y. Liu, J.-X. Wang, and H. Sun, “Phycrnet: Physics-informed convolutional-recurrent network for solving spatiotemporal pdes,” _Computer Methods in Applied Mechanics and Engineering_ , vol. 389, p. 114399, 2022.
  * [45] R. S. Sutton, “Integrated architectures for learning, planning, and reacting based on approximating dynamic programming,” in _Machine learning proceedings 1990_. Elsevier, 1990, pp. 216–224.
  * [46] S. Chiappa, S. Racaniere, D. Wierstra, and S. Mohamed, “Recurrent environment simulators,” _arXiv:1704.02254_ , 2017.
  * [47] P. Tandon, “pytorch_sac,” https://github.com/pranz24/pytorch-soft-actor-critic, 2018.
  * [48] D. Yarats, A. Zhang, I. Kostrikov, B. Amos, J. Pineau, and R. Fergus, “pytorch_sac_ae,” https://github.com/denisyarats/pytorch_sac_ae, 2021\.
  * [49] H. Van Hasselt, A. Guez, and D. Silver, “Deep reinforcement learning with double q-learning,” in _Proceedings of the AAAI conference on artificial intelligence_ , vol. 30, no. 1, 2016.
  * [50] B. T. Polyak and A. B. Juditsky, “Acceleration of stochastic approximation by averaging,” _SIAM journal on control and optimization_ , vol. 30, no. 4, pp. 838–855, 1992.
  * [51] S. Fujimoto, H. Hoof, and D. Meger, “Addressing function approximation error in actor-critic methods,” in _International conference on machine learning_. PMLR, 2018, pp. 1587–1596.
  * [52] J. Viquerat, P. Meliga, and E. Hachem, “A review on deep reinforcement learning for fluid mechanics: an update,” _arXiv:2107.12206_ , 2021.
  * [53] A. Raffin, A. Hill, A. Gleave, A. Kanervisto, M. Ernestus, and N. Dormann, “Stable-baselines3: Reliable reinforcement learning implementations,” _Journal of Machine Learning Research_ , vol. 22, no. 268, pp. 1–8, 2021\. [Online]. Available: http://jmlr.org/papers/v22/20-1364.html
  * [54] S. Pandey, P. Teutsch, P. Mäder, and J. Schumacher, “Direct data-driven forecast of local turbulent heat flux in rayleigh–bénard convection,” _Physics of Fluids_ , vol. 34, no. 4, p. 045106, 2022.
  * [55] S. L. Brunton, J. L. Proctor, and J. N. Kutz, “Discovering governing equations from data by sparse identification of nonlinear dynamical systems,” _Proceedings of the national academy of sciences_ , vol. 113, no. 15, pp. 3932–3937, 2016.
  * [56] J. L. Proctor, S. L. Brunton, and J. N. Kutz, “Dynamic mode decomposition with control,” _SIAM Journal on Applied Dynamical Systems_ , vol. 15, no. 1, pp. 142–161, 2016.
  * [57] M. Korda and I. Mezić, “Linear predictors for nonlinear dynamical systems: Koopman operator meets model predictive control,” _Automatica_ , vol. 93, pp. 149–160, 2018.
  * [58] S. Peitz and S. Klus, “Koopman operator-based model reduction for switched-system control of pdes,” _Automatica_ , vol. 106, pp. 184–191, 2019\.
  * [59] F. Nüske, S. Peitz, F. Philipp, M. Schaller, and K. Worthmann, “Finite-data error bounds for Koopman-based prediction and control,” _Journal of Nonlinear Science_ , vol. 33, p. 14, 2023.
  * [60] S. Peitz and K. Bieker, “On the Universal Transformation of Data-Driven Models to Control Systems,” _Automatica_ , vol. 149, p. 110840, 2023.
  * [61] C. Zhang and E. Zuazua, “A quantitative analysis of Koopman operator methods for system identification and predictions,” _Comptes Rendus. Mécanique_ , 2023.
  * [62] A. Vahdat and J. Kautz, “Nvae: A deep hierarchical variational autoencoder,” _Advances in Neural Information Processing Systems_ , vol. 33, pp. 19 667–19 679, 2020.
  * [63] S. Elfwing, E. Uchibe, and K. Doya, “Sigmoid-weighted linear units for neural network function approximation in reinforcement learning,” _Neural Networks_ , vol. 107, pp. 3–11, 2018.
  * [64] K. He, X. Zhang, S. Ren, and J. Sun, “Deep residual learning for image recognition,” in _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 2016, pp. 770–778.
  * [65] T. Salimans and D. P. Kingma, “Weight normalization: A simple reparameterization to accelerate training of deep neural networks,” _Advances in neural information processing systems_ , vol. 29, 2016.
  * [66] J. L. Ba, J. R. Kiros, and G. E. Hinton, “Layer normalization,” _arXiv:1607.06450_ , 2016.
  * [67] S. Ioffe and C. Szegedy, “Batch normalization: Accelerating deep network training by reducing internal covariate shift,” in _International conference on machine learning_. PMLR, 2015, pp. 448–456.
  * [68] M. D. Zeiler, D. Krishnan, G. W. Taylor, and R. Fergus, “Deconvolutional networks,” in _2010 IEEE Computer Society Conference on computer vision and pattern recognition_. IEEE, 2010, pp. 2528–2535.
  * [69] H. Noh, S. Hong, and B. Han, “Learning deconvolution network for semantic segmentation,” in _Proceedings of the IEEE international conference on computer vision_ , 2015, pp. 1520–1528.
  * [70] T. Wang and J. Ba, “Exploring model-based planning with policy networks,” _arXiv:1906.08649_ , 2019.
  * [71] G. Z. Holland, E. J. Talvitie, and M. Bowling, “The effect of planning shape on dyna-style planning in high-dimensional state spaces,” _arXiv:1806.01825_ , 2018.
  * [72] R. J. Williams and D. Zipser, “A learning algorithm for continually running fully recurrent neural networks,” _Neural computation_ , vol. 1, no. 2, pp. 270–280, 1989.
  * [73] Y. Bengio, J. Louradour, R. Collobert, and J. Weston, “Curriculum learning,” in _Proceedings of the 26th annual international conference on machine learning_ , 2009, pp. 41–48.
  * [74] D. P. Kingma and J. Ba, “Adam: A method for stochastic optimization,” _arXiv:1412.6980_ , 2014.
  * [75] H. Jaeger, “Tutorial on training recurrent neural networks, covering bppt, rtrl, ekf and the” echo state network” approach,” 2002.
  * [76] M. Tan and Q. Le, “Efficientnet: Rethinking model scaling for convolutional neural networks,” in _International conference on machine learning_. PMLR, 2019, pp. 6105–6114.

###  -A Neural network details

Our encoder network consists of residual cells similar to those proposed in
[62] for image generation, as they mitigate vanishing and exploding gradients
at deeper layers of the architecture with so-called skip connections. The
cells used in our work stack two $3\times 3$ convolutional layers with _SiLU_
activation functions [63] along the residual branch, see Figure 8 for an
illustration. Similar to [64], we use a $1\times 1$ strided convolution unit
along the shortcut connection to downsample the input. Strides are also
applied in the first $3\times 3$ convolution of the residual branch to reduce
the spatial dimensionality and dictate the degree to which the spatial extent
is reduced. We typically half the spatial dimensionality with each stacked
residual cell in the encoder network and use convolution instead of pooling
layers for downsampling. Lastly, we use _weight normalization_ [65] and _layer
normalization_ [66] over _batch normalization_ [67] layers due to the well-
known shortcomings of the latter during inference in recurrent structures. For
our decoder network $f_{\boldsymbol{\theta}_{\mathrm{dec}}}$, we use
deconvolutions [68, 69] without unpooling layers and with SiLU activation
functions. After the spatial extent of the physical domain has been restored,
we introduce further convolutional layers to the architecture, which enhanced
the performance in our experiments.

For the control term, we experimented with sharing the encoder network
$f_{\boldsymbol{\theta}_{\mathrm{enc}}}$ between encoding states
$\boldsymbol{u}_{\tau\Delta\tau}$ and forcing terms
$\boldsymbol{\phi}_{\tau\Delta\tau}$. However, since we only consider an
additive control term in our application, we ultimately found training a
separate encoder network
$\boldsymbol{v}_{\tau}=f_{\psi_{\mathrm{enc}}}(\boldsymbol{\phi}_{\tau\Delta\tau})$
to be more effective. In general, the action encoder $f_{\psi_{\mathrm{enc}}}$
matches the model architecture of $f_{\boldsymbol{\theta}_{\mathrm{enc}}}$ but
implements fewer convolutional filters due to the simple Gaussian shape of the
exerted force in the case of the Kuramoto-Sivashinsky control problem
considered here.

Besides parameter sharing, another benefit of the convolution operation is its
ability to obey the boundary conditions. In the case of periodic boundaries,
as studied in our work, a circular padding mechanism is applied at each border
of the spatial domain. Similar to other inductive biases like spatial
locality, encoding such prior system knowledge often improves data usage and
generalization of forecasts. Using circular padding for all convolution
operators in the encoder, decoder, and transition model helped significantly
when learning the state evolution in our experiments.

Figure 8: Residual block consisting of a skip connection (lower branch)
downsampling the inputs with $1\times 1$ convolutions and a main branch using
$3\times 3$ convolutions. Layer normalization [66] is used to normalize the
output activations for single input instances, which accelerates the training.
The illustration is adopted from [62].

###  -B Implementation details on ensembles

Other works [22, 70, 24] alternate the dynamics model
$f_{\boldsymbol{\theta}_{\ell}}$ used to predict state transition
$\boldsymbol{\hat{s}}_{\tau+1}=f_{\theta_{\ell}}(\boldsymbol{\hat{s}}_{\tau},\boldsymbol{a}_{\tau};\,\boldsymbol{c}_{\tau,\ell})$
at each step. Owing to the fact that the memory unit of our model maintains a
state of its past inputs, switching transition models after each step is only
feasible if we invoke all models of the ensemble for each step. Otherwise, the
memory kept as cell state $\boldsymbol{c}_{\tau,\ell}$ would be outdated.
Saving computation, our work explores an alternative design choice and selects
single models of ensemble
$\\{f_{\boldsymbol{\theta}_{1}},\dots,f_{\boldsymbol{\theta}_{L_{\mathrm{ens}}}}\\}$
and uses them throughout the entire model-based rollout, similar to what is
considered in [21] for model predictive control. Our hope is that the use of
an intermediate replay buffer $\mathcal{D}_{\mathrm{model}}$ and training our
model-free agent off-policy will have a similar decorrelation effect, since
agents still update their parameters with artificial experience generated
through several models at once. Indeed, our experiments show that there is no
significant difference between both implementations for the decision-making
problem studied in our work. The compute savings for model-based rollouts
amount to a factor equal to the number of elites in the ensemble.

###  -C Model-based rollouts

In order for model-based algorithms to succeed, planning must be advantageous
to updating the behavior with existing samples or collecting additional data.
On this matter, a recent study [71] formulates an essential condition for
models to benefit the optimization. Specifically, unlike existing experience,
a model has to generalize and has to be queried at unseen regions of the state
space. A working assumption is thus that the model has to generalize quicker
towards the global dynamics than the policy $\pi_{\boldsymbol{\omega}}$ can
learn and generalize behavior. A prolonged rollout of
$\pi_{\boldsymbol{\omega}}$ in the surrogate is thus inevitable for planning
to generate samples unfamiliar to the agent [71]. Vice versa, with an
increasing rollout length, model errors compound and could cause catastrophic
learning updates to the agent [23]. As proposed in [24], we adopt a scheduler
for the rollout length $K_{\mathrm{rll}}$ in order to trade off compounding
model errors against the novelty of samples. A rollout length scheduler deters
the agent from exploiting model errors, specifically in the small data regime,
where overfitting is often an issue of high-capacity function approximators.
In general, we use our previous results for the valid prediction time of
Section IV for guidance to limit the number of data collection steps. Our
implementation uses a linear scheduler that increments the length
$K_{\mathrm{rll}}$ up to a maximum number of steps. In Figure 9, we illustrate
the competing effects of sample novelty and compounding model errors in
detail. The figure shows an ideal situation where the model quickly
generalizes to other regions of the state space once more and more data
becomes available (contour lines widening on the right-hand side), while the
ability of the agent to generalize its behavior lags behind (marginal widening
of contour lines on the left-hand side). Despite model-based rollouts
increasing in depth, model-based rollouts continue to be accurate due to the
generalizing model (bottom right) and collect data on state space regions
unfamiliar to the agent (bottom left).

Figure 9: Intuition of rollout length scheduling. In the small data regime
(top), the model does not generalize to unfamiliar regions of the state space
(top right) so that rollouts must be truncated after few steps. In the large
data regime (bottom), forecast can branch off collected data for several steps
before compounding model errors become substantial (bottom right). In an ideal
scenario, the model can generate unfamiliar experience for the agent (left-
hand side), since it generalizes more efficiently.

Overall, model-based rollouts proceed in the following way. After a sampled
sequence of existing experience $(\boldsymbol{s}_{\tau-
K_{\mathrm{tf}}},\boldsymbol{a}_{\tau-
K_{\mathrm{tf}}},\dots,\boldsymbol{s}_{\tau},\boldsymbol{a}_{\tau})\sim\mathcal{D}_{\mathrm{env}}$
is processed in teacher-forcing mode to obtain $\hat{\boldsymbol{s}}_{\tau}$,
the policy $\pi_{\boldsymbol{\omega}}$ selects an action
$\boldsymbol{a}_{\tau}$ as its control output. A dynamics model
$f_{\boldsymbol{\theta}_{\ell}}$ of the ensemble then predicts the state
transition
$\hat{\boldsymbol{s}}_{\tau+1}=f_{\boldsymbol{\theta}_{\ell}}(\hat{\boldsymbol{s}}_{\tau},\boldsymbol{a}_{\tau};\,\boldsymbol{c}_{\ell,\tau})$.
In practice, dynamics model $f_{\boldsymbol{\theta}_{\ell}}$ compresses the
state-action pair to a latent space using encoder networks
$f_{\boldsymbol{\theta}_{\ell,\mathrm{enc}}}$ and
$f_{\boldsymbol{\psi}_{\ell,\mathrm{enc}}}$, performs a time step using the
convolutional LSTM $f_{\boldsymbol{\theta}_{\ell,\mathrm{fwd}}}$ and restores
a state space representation with decoder network
$f_{\boldsymbol{\theta}_{\ell,\mathrm{dec}}}$. Using the known reward function
$\mathcal{R}$, we estimate the reward of the action as
$\hat{r}_{\tau}=\mathcal{R}\left(\hat{\boldsymbol{u}}_{\tau\Delta\tau},\boldsymbol{\phi}_{\tau\Delta\tau}\right)$
and insert the transition as an additional sample into dataset
$\mathcal{D}_{\mathrm{model}}$. Rollouts are stopped manually after
$K_{\mathrm{rll}}$ steps or conclude once a terminal state of the system is
reached. The studied decision-making problem in our work does not define a
terminal condition except for the fixed time limit of $T_{\mathrm{max}}$.

###  -D Model training details

Given sequences of system states
$\\{\boldsymbol{s}_{\tau}\\}_{\tau=0}^{K_{\mathrm{seq}}}$ and control outputs
$\\{\boldsymbol{a}_{\tau}\\}_{\tau=0}^{K_{\mathrm{seq}}-1}$ for training, we
first use our surrogate model $f_{\boldsymbol{\theta}}$ to predict
$\\{\boldsymbol{\hat{s}}_{\tau}\\}_{\tau=1}^{K_{\mathrm{seq}}}$. Afterwards,
we train the model end-to-end while minimizing the mean squared error
$\mathcal{L}_{\mathrm{MSE}}\propto\sum_{i}\sum_{\tau=1}^{K_{\mathrm{seq}}}\left\lVert\boldsymbol{\hat{s}}_{\tau}^{(i)}-\boldsymbol{s}_{\tau}^{(i)}\right\rVert^{2}_{2}$
averaged over batches and sequence lengths. To address the slow convergence of
the optimization and instabilities during training, we adopt teacher forcing
[72] for guidance. In doing so, we split state sequence
$\\{\boldsymbol{s}_{\tau}\\}_{\tau=0}^{K_{\mathrm{seq}}}$ into separate
portions $\\{\boldsymbol{s}_{\tau}\\}_{\tau=0}^{K_{\mathrm{tf}}}$ and
$\\{\boldsymbol{s}_{\tau}\\}_{\tau=K_{\mathrm{tf}}+1}^{K_{\mathrm{seq}}}$ and
replace model predictions $\boldsymbol{\hat{s}}_{\tau+l}$ with states
$\boldsymbol{s}_{\tau+l}$ in the first $K_{\mathrm{tf}}$ steps of the
recurrent feedback loop to mitigate compounding errors propagating to later
prediction steps at the beginning of our training process. After
$K_{\mathrm{tf}}$ steps with teacher feedback, the model predicts sequence
$\\{\boldsymbol{\hat{s}}_{\tau}\\}_{\tau=K_{\mathrm{tf}}+1}^{K_{\mathrm{seq}}}$
using its prior output $\boldsymbol{\hat{s}}_{\tau+l}$ as inputs for the
succeeding prediction step $\boldsymbol{\hat{s}}_{\tau+l+1}$. In an attempt to
reconcile the benefits of teacher forcing in terms of convergence speed with
the necessity of consecutive prediction steps, we use a curriculum for
training [73]. At first, we keep the number of consecutive state transitions
${K_{\mathrm{seq}}}-K_{\mathrm{tf}}$ small to learn the encoding and decoding
of state variables quickly. After that, we shift the focus towards learning
temporal transitions and increase the sequence length ${K_{\mathrm{seq}}}$
towards the target sequence length $K_{\mathrm{max}}+K_{\mathrm{tf}}$, that
is, predict $K_{\mathrm{max}}$ steps in feedback mode.

Given a training set of $N_{\mathrm{eps}}$ episodes composed of
$T_{\mathrm{max}}/\Delta\tau$ state transitions each, we subsample sequences
$\\{\boldsymbol{s}_{\tau},\boldsymbol{a}_{\tau},\dots,\boldsymbol{s}_{\tau+K_{\mathrm{seq}}},\boldsymbol{a}_{\tau+K_{\mathrm{seq}}}\\}$
starting at arbitrary points within episodes. In each training and validation
epoch, we subsample $K_{\mathrm{subseq}}$ sequences equal to the number of
consecutive segments in the dataset which are non-overlapping and of length
$K_{\mathrm{seq}}$. While we subsample our training sequences anew after each
epoch, the validation set remains fixed throughout the optimization procedure.

We use the Adam optimizer [74] with an adaptive learning rate (initial
learning rate of $0.001$, momentum of $0.9$) to minimize the training loss
$\mathcal{L}_{\mathrm{MSE}}$. Standard techniques, including gradient clipping
(threshold of $0.5$) and regularization, are used. Using truncated
backpropagation through time [75], we limit the maximum number of steps
$K_{\mathrm{TBTT}}$ the gradient signal is backpropagated. All inputs as well
as prediction targets are scaled to the unit interval or normalized based on
statistics derived from the training set. In our experiments, learning
normalized target values of the temporal residuals
$\Delta\boldsymbol{s}_{\tau}/\Delta\tau=(\boldsymbol{s}_{\tau+1}-\boldsymbol{s}_{\tau})/\Delta\tau$
at each step worked significantly better than backpropagating through
reconstructed state variables
$\boldsymbol{\hat{s}}_{\tau+1}=\boldsymbol{\hat{s}}_{\tau}+\Delta\boldsymbol{\hat{s}}_{\tau}$.

Table I shows the network architecture for our dynamics model. In essence, the
state and action encoder networks compress the number of spatial dimensions of
the state variable four-fold.

TABLE I: Network architecture of our dynamics model.

Network | Layer | Filters | Filter | Stride | Parameters | Output
---|---|---|---|---|---|---
|  |  | Size |  |  |
State $\boldsymbol{s}$ | — | — | — | — | — | $1\times 64$
State Encoder | Res. Block | 4 | 3 | 1 | 64 | $4\times 64$
Res. Block | 8 | 3 | 2 | 320 | $8\times 32$
Res. Block | 16 | 3 | 2 | 1280 | $16\times 16$
Forcing Term $\boldsymbol{\phi}$ | — | — | — | — | — | $1\times 64$
Action Encoder | Res. Block | 2 | 3 | 1 | 20 | $2\times 64$
Res. Block | 4 | 3 | 2 | 80 | $4\times 32$
Res. Block | 6 | 3 | 2 | 204 | $6\times 16$
Transition Model | Conv. LSTM | 16 | 3 | 1 | 4800 | $16\times 16$
State Decoder | DeConv. | 16 | 3 | 2 | 784 | $16\times 32$
DeConv. | 8 | 3 | 2 | 392 | $8\times 64$
Res. Block | 4 | 5 | 1 | 176 | $4\times 64$
Conv. | 1 | 7 | 1 | 13 | $1\times 64$

###  -E Study on the neural network sizes

We here study several variants of our dynamics model with an increasing number
of parameters. In doing so, we evaluate surrogates similar to the encoder,
decoder and transition model outlined in Table I, but scale the width $\beta$
of each network, which defines the exponential number of channels $\beta^{N}$
of its layers in terms of their depth $N$ [76]. We evaluate the predictive
capabilties of surrogates with network widths and overall number of parameters
outlined in Table II. In this setting, the network in our previous study
(which we are also going to use in the next section) corresponds to
$\beta\approx 3$.

TABLE II: Number of learnable parameters at increasing model widths. Width $\beta$ | 1.5 | 2.0 | 2.5 | 3.0 | 3.5 | 4.0
---|---|---|---|---|---|---
Model Parameters | 2612 | 4137 | 7636 | 17542 | 36409 | 77443
Figure 10: Valid prediction time of models at increasing model widths and
several thresholds. Solid lines depict the mean of five folds, while the
shaded regions show the 95% confidence interval.

In order to determine the time span throughout which forecasts match the
spatio-temporal evolution of state variables and attributed rewards, we
compute the valid prediction time as the farthest point in time
$T_{\mathrm{vpt}}$ with normalized mean squared errors of state and reward
being consistently below a fixed threshold $\epsilon$. Therefore,
$T_{\mathrm{vpt}}$ gives a coarse estimate on the number of steps we expect to
be safe in the context of model-based reinforcement learning. In Figure 10, we
illustrate the valid prediction time of models at increasing network widths
for a number of thresholds and dataset sizes given for training. Surprising to
us, models of larger network widths are not susceptible to overfitting if only
few snapshots are available, with our model at width $\beta=4.0$ outperforming
the network at width $\beta=1.5$ consistently. (In contrast, we observed an
inverse trend on fully connected architectures.) As a matter of fact, the
valid prediction time consistently increases with growing network width but
plateaus marginally after $\beta=3.0$.

Since computation time increases at an exponential rate with the width of a
neural network [76], we fix $\beta\approx 3.0$ for subsequent experiments.
Intermediate experiments in our work suggested that increasing model capacity
also improves data efficiency and overall performance of learned control
strategies. Still, a comprehensive evaluation of our algorithm for model-based
control turned out to be too resource-intensive and time-consuming for larger
dynamics models.