applied mathematics, computational physics, artificial intelligence

Jian-Xun Wang

# Physics-informed Dyna-Style Model-Based Deep Reinforcement Learning for
Dynamic Control

Xin-Yang Liu1 and Jian-Xun Wang1 1Department of Aerospace & Mechanical
Engineering, College of Engineering, University of Notre Dame, Notre Dame, IN,
USA<EMAIL_ADDRESS>
###### Abstract

Model-based reinforcement learning (MBRL) is believed to have much higher
sample efficiency compared to model-free algorithms by learning a predictive
model of the environment. However, the performance of MBRL highly relies on
the quality of the learned model, which is usually built in a black-box manner
and may have poor predictive accuracy outside of the data distribution. The
deficiencies of the learned model may prevent the policy from being fully
optimized. Although some uncertainty analysis-based remedies have been
proposed to alleviate this issue, model bias still poses a great challenge for
MBRL. In this work, we propose to leverage the prior knowledge of underlying
physics of the environment, where the governing laws are (partially) known. In
particular, we developed a physics-informed MBRL framework, where governing
equations and physical constraints are utilized to inform the model learning
and policy search. By incorporating the prior information of the environment,
the quality of the learned model can be notably improved, while the required
interactions with the environment are significantly reduced, leading to better
sample efficiency and learning performance. The effectiveness and merit have
been demonstrated over a handful of classic control problems, where the
environments are governed by canonical ordinary/partial differential
equations.

###### keywords:

Reinforcement learning (RL), physics-informed neural networks, flow control,
Kuramoto-Sivashinsky

##  1 Introduction

Reinforcement learning (RL) is a class of artificial intelligence (AI)
techniques that train an AI agent to learn the optimal control strategy by
interacting with the surrounding environment. Over the past few years, with
the rapid development of deep learning (DL), deep reinforcement learning (DRL)
techniques have been witnessing tremendous success in a variety of
applications. In particular, DRL has demonstrated superhuman performance at
playing Go [1, 2] and Atari games from pixels [3]. Most recently, there has
been growing interest in applying DRL for dynamic control of complex physical
systems, e.g., laminar/turbulent flows [4, 5, 6, 7, 8], active matter [9],
fish swimmers [10, 11, 12], unmanned aerial vehicles (UAV) [13, 14, 15], and
robotics [16, 17, 18].

In general, most state-of-the-art RL agents learn the desired tasks by
gathering experience directly from the environment. Namely, the optimal action
strategy is derived by interacting with the real physical system in a trial-
and-error manner. This class of RL methods is known as _model-free
reinforcement learning_ (MFRL). Due to the ease of implementation and no need
for prior knowledge of the dynamic transitions, MFRL has been widely applied
to many tasks, mainly in playing computer games [19, 3, 20]. Despite their
popularity, MFRL methods usually have _low sample efficiency_ , i.e.,
requiring a massive amount of interactions with the environment. Although the
low sample efficiency and slow convergence rate might be acceptable for
training an agent in gaming applications since the environment interactions
are nearly costless, these shortcomings will significantly limit the DRL
applications for dynamic control of physical/mechanical systems (e.g., flights
or robotics). First, real-world interactions of a mechanical system can be
very expensive and time-consuming. For instance, considering a flow control
problem with plasma actuators, it is very costly or even infeasible to train a
controller by conducting a huge amount of wind tunnel experiments with
enormous control trials, which is unlike training an AI game agent that can be
done by playing the computer games for a vast amount of times (episodes).
Second, the mechanical systems can be easily worn out from extensive action
trials, and thus the exploration of optimal control strategy in MFRL is highly
restricted to avoid possible damage to the system in real-world industry
settings. Although the off-policy MFRL algorithms with a replay buffer (e.g.,
deep Q-Networks [3] and their actor-critic extensions [21, 22, 23]) can better
utilize historical data than on-policy MFRL algorithms (e.g., trust region
policy optimization [24] and proximal policy optimization [25]), the data
efficiency is still far from sufficient.

One way to improve data efficiency is to augment the data collected from real-
world interactions with a learned transition model. This is the general idea
of the other class of DRL algorithms: _Model-based reinforcement learning_
(MBRL) [26, 27]. Using a learned model to reason about the future can avoid
the irreversible consequence of trial-and-error in the real environment and
has great potential to significantly improve data efficiency, which is thus
more appealing in applications of complex mechanical systems. In addition, the
learned transition model is independent of rewards and thus can be transferred
to other control problems in the same/similar environments. Many existing MBRL
methods rely on simple function approximators, such as Gaussian process (GP),
linear models, and Gaussian mixture models [28, 29, 30]. However, the limited
expressibility of the simple models prevents them from handling high-
dimensional problems with complex dynamic transitions. Thanks to the rapid
developments of deep learning, more and more complex high-dimensional function
approximators based on neural networks have been applied to design more
powerful MBRL algorithms. For example, Racanière et al. [31] proposed a novel
MBRL framework, Imagination-Augmented Agent (I2A), where the environment model
is constructed by a recurrent network architecture for generating imagined
trajectories to inform agent’s decisions. Kaiser et al. [32] presented a
complete MBRL method (SimPLe) using a convolutional neural network to
successfully solve Atari games with significantly fewer interactions than MFRL
methods. Hafner et al. [33] developed the Deep Planning Network (PlaNet) that
learns the latent dynamics of the environment directly from images using a
variational autoencoder and a recurrent latent network. The effectiveness of
PlaNet has been demonstrated by successfully solving a number of continuous
control tasks from pixels. Hafner et al. [34] further extended the PlaNet by
developing a novel actor-critic based MBRL method (Dreamer), which learns
long-horizon behaviors from images purely by latent imagination. Dreamer has
been evaluated on the DeepMind Control Suite and outperforms most state-of-
the-art MBRL and MFRL algorithms in every aspect.

Despite the great promise, most commonly-used MBRL approaches suffer from
model inaccuracy (i.e., model bias), preventing them from matching the success
of their model-free counterparts [26]. This is particularly true when it comes
to learning complex dynamics with high-capacity models (e.g., deep neural
networks), which are prone to overfitting in data-sparse and out-of-sample
regimes [27]. In particular, the model bias can be significantly exacerbated
for predicting long rollout horizons because of the “compound error” effect.
To mitigate this issue, rather than learning the transition deterministically,
people built the dynamic models in a probabilistic manner, where the unknown
model bias is treated as the epistemic uncertainty (i.e., model-form
uncertainty) and can be modeled in several different ways. For example,
Depeweg et al. [35] employed Bayesian neural networks (BNNs) to learn the
probabilistic dynamic transition and update the policy over an ensemble of
models sampled from the trained BNNs. Kurutach et al. [36] proposed to use an
ensemble of models to estimate the model-form uncertainty and regularize the
trust region policy optimization (TRPO). Nonetheless, the model-form
uncertainty is notoriously difficult to quantify, especially for black-box
deep learning models [37, 38]. Most recently, a more promising strategy known
as physics-informed deep learning (PIDL) has attracted increasing attention in
the scientific machine learning (SciML) community, aiming to leverage both the
advantages of deep learning and prior knowledge of underlying physics to
enable data-scarce learning. Instead of learning solely from labeled data, the
model training process is also guided by physics laws and knowledge, which
could provide rigorous constraints to the model output, alleviate overfitting
issues, and improve the robustness of the trained model in data-scarce and
out-of-sample regimes. This idea has been recently explored for solving PDEs
or modeling complex physical systems. For example, researchers have
incorporated physical constraints (e.g., realizability, symmetry, invariance)
into SciML models to develop physics-informed, data-driven turbulence models
[39, 40, 41]. People have also utilized governing equations of the physical
systems to inform or directly train deep neural networks, i.e., physics-
informed neural networks (PINNs) [42], which has been demonstrated in many
scientific and engineering applications [43, 44, 45, 46, 47].

In this work, we leverage the idea of PIDL and propose Physics-informed Model-
Based Reinforcement Learning (PiMBRL), an innovative MBRL framework for
complex dynamic control that incorporates the physical laws/constraints of the
system to alleviate the issue of model bias, reduce the real-world
interactions, and significantly improve the data efficiency. Specifically, a
novel autoencoding-based recurrent network architecture is constructed to
learn the dynamic transition in the Dyna-style MBRL framework [48], which is a
commonly-used MBRL formulation. The governing physics of the environment are
assumed to be known and are utilized to inform the model learning and RL agent
optimization. State-of-the-art Off-policy actor-critic (AC) algorithms, e.g.,
Twin Delayed Deep Deterministic Policy Gradients (TD3) [22], are used for
value/policy optimization, We have demonstrated the effectiveness and merit of
the proposed PiMBRL on a few classic dynamic control problems, where the
environments are governed by canonical ordinary/partial differential equations
(ODEs/PDEs), including cart-pole, pendulum, viscous fluid dynamics governed by
Burgers’ equation, and chaotic/turbulent dynamics governed by Kuramoto-
Sivashinsky (KS) Equation. The performance of the proposed PiMBRL algorithms
is compared with their MBRL and MFRL counterparts, and significant
improvements in terms of sample efficiency and model accuracy are observed.
The novel contributions of this work are summarized as follows: (a) we propose
a physics-informed model-based RL framework based on a novel encoder-decoder
recurrent network architecture; (b) embed the physics of the environment into
the MBRL using discretized PIDL formulation [49]; (c) demonstrate the
effectiveness of proposed methods on a variety of dynamic control problems,
particularly including nonlinear spatiotemporal chaotic systems, e.g., the KS
equation, which exhibits a wide range of dynamics from the steady to
chaotic/turbulent regimes, shedding lights on developing controllers for more
challenging fluid systems governed by Navier-Stokes equations; (d) compare the
proposed method with state-of-the-art MBRL and MFRL in terms of accuracy and
sample complexity. This work is the first attempt to use physical laws to
inform the MBRL agent optimization to the best of the authors’ knowledge.

The rest of the paper is organized as follows. The background of MFRL/MBRL and
our proposed PiMBRL algorithms are introduced in Section 2. Results of
numerical experiments on classic dynamic control problems are presented and
discussed in Section 3. The influence of model rollout length and model
accuracy threshold is further discussed in Section 4. Finally, Section 5
concludes the paper.

##  2 Methodology

###  2.1 Problem formulation and background

We consider dynamical systems equipped with localized control inputs (i.e.,
actuators),

$\frac{d{\bm{u}}}{dt}=\mathscr{F}({\bm{u}},{\bm{a}};{\bm{\mu}}),$ (1)

where ${\bm{u}}({\bm{x}},t)\in\mathbb{R}^{d_{u}}$ denotes the state variable
of the system in the spatial domain $\Omega$ and temporal domain $t\in[0,T]$,
${\bm{a}}({\bm{x}},t)\in\mathbb{R}^{d_{a}}$ represents the action variable
(i.e., control inputs), and $\mathscr{F}(\cdot)$ is a nonlinear differential
operator parameterized by ${\bm{\mu}}$. In many cases, the systems can be
assumed to possess the Markov property, referred to as Markov Decision
Processes (MDP). The discrete form can be written as,

${\bm{u}}_{t+1}=\mathcal{F}({\bm{u}}_{t},{\bm{a}}_{t};{\bm{\mu}}),$ (2)

where the state ${\bm{u}}_{t+1}$ at next time $t+1$ only depends on the state
${\bm{u}}_{t}$ and action ${\bm{a}}_{t}$ at current time step $t$, and the
time-invariant transition dynamics $\mathcal{F}$ of the environment is a
nonlinear differential functional. In the optimal control problem, the goal is
to find a series of action signals (a.k.a., policy $\pi$) that maximizes the
expected return $R(\pi)$,

$R(\pi)=\int_{0}^{T}{\mathop{\mathbb{E}}_{{\bm{u}}_{t}\sim\pi_{t}}\bigg{[}r({\bm{u}}_{t})\bigg{]}},$
(3)

where $r({\bm{u}}_{t})$ denotes the reward function of the state at time $t$,
which is a signal to assess the control agent locally. This optimal control
problem can be solved by deep reinforcement learning (DRL), either in a model-
free or model-based manner.

####  2.1.1 Value function, policy function, and Bellman equation

Before putting forth the proposed DRL algorithms, we introduce several
important concepts in DRL, including value & policy functions and Bellman
equation. Value functions are functions of a state (or a state-action pair)
that estimate the total return starting from that particular state (state-
action pair). Value function $v({\bm{u}})$ of a state ${\bm{u}}$ is known as
state-value function, while value function $q({\bm{u}},{\bm{a}})$ of a state-
action pair $({\bm{u}},{\bm{a}})$ is known as action-value function. The
state-value and action-value functions are formally defined as,

$\displaystyle v({\bm{u}})$
$\displaystyle\doteq\mathbb{E}\left[\sum_{k=1}^{\infty}{\gamma^{k}r({\bm{u}}_{t+k})}\Bigg{|}{\bm{u}}_{t}={\bm{u}}\right],$
(4a) $\displaystyle q({\bm{u}},{\bm{a}})$
$\displaystyle\doteq\mathbb{E}\left[\sum_{k=1}^{\infty}{\gamma^{k}r({\bm{u}}_{t+k})}\Bigg{|}{\bm{u}}_{t}={\bm{u}},{\bm{a}}_{t}={\bm{a}}\right],$
(4b)

where $\gamma\leq 1$ is the discount rate. A policy function $\pi$ maps states
to actions (or probabilities of actions). Namely, a policy function $\pi$ can
be defined either as a deterministic function $\pi({\bm{u}})={\bm{a}}$ or a
probability measure $\pi({\bm{a}}|{\bm{u}})$. Due to the nature of MDP, value
functions can be estimated recursively based on the Bellman equations [48],

$\displaystyle v({\bm{u}}_{t})$
$\displaystyle=\mathop{\mathbb{E}_{\pi}}_{{\bm{u}}_{t+1}\sim\mathcal{P}}\bigg{[}r({\bm{u}}_{t},{\bm{a}}_{t})+\gamma
v({\bm{u}}_{t+1})\bigg{]},$ (5a) $\displaystyle q({\bm{u}}_{t},{\bm{a}}_{t})$
$\displaystyle=\mathop{\mathbb{E}}_{{\bm{u}}_{t+1}\sim\mathcal{P}}\bigg{[}r({\bm{u}}_{t},{\bm{a}}_{t})+\gamma\mathop{\mathbb{E}}_{{\bm{a}}_{t+1}\sim\pi}q({\bm{u}}_{t+1},{\bm{a}}_{t+1})\bigg{]},$
(5b)

where $\mathcal{P}=\mathcal{P}({\bm{u}}_{t+1}|{\bm{u}}_{t},{\bm{a}}_{t})$ is
the transition probability, describing the dynamics of the environment.

####  2.1.2 Model-free and model-based reinforcement learning

As mentioned above, RL aims to find a series of actions (i.e., optimal policy)
that maximize the total return by estimating the value and/or policy function
based on the Bellman equation. Depending on whether or not learning and using
a model of the transition dynamics of the environment, RL can be classified
into two categories: model-free reinforcement learning (MFRL) and model-based
reinforcement learning (MBRL). In MFRL, the optimization process is conducted
by repeatedly interacting with the environment with a trial-and-error search,
and the model of the environment is not required (see Fig. 1a). Namely, the
state dynamics of the environment are (partially) observed as exploring
different policy strategies, and the best policy will be identified after
massive trials. MBRL, on the other hand, leverages a model
$\tilde{\mathcal{F}}$ that approximates the real environment $\mathcal{F}$ and
predicts the dynamic transition
($\tilde{\mathcal{F}}:{\bm{u}}_{t},{\bm{a}}_{t}\to{\bm{u}}_{t+1}$), which can
be learned from the interactions with the real environment. The RL agent is
then optimized by the interactions not only with the real environment but also
with the virtual environment constructed by the model (see Fig. 1b). The
learned model $\tilde{\mathcal{F}}$ can be utilized for planning with its
gradient information (e.g., SVG [50], GPS [51]) or synthesizing imagined
samples to augment real samples for better sample efficiency. The latter is
known as the Dyna-like MBRL [48, 52, 36, 32] that can directly leverage
cutting-edge MBRL algorithms.

Figure 1: Schematics of model-free reinforcement learning (MFRL) and model-
based reinforcement learning (MBRL).

The algorithms of the DRL agent optimization can be grouped into three
classes: (i) actor-only, (2) critic-only, and (3) actor-critic methods. Actor-
only (i.e., policy-gradient) methods directly optimize the policy function
$\tilde{\pi}({\bm{u}};\bm{\theta}_{\pi})$, which is often parameterized by a
deep neural network by calculating the policy gradient with respect to network
parameters $\bm{\theta}_{\pi}$. The optimization can be solved based on
stochastic gradient descent (SGD) [53] or its variants [54, 55, 56, 57].
Examples of policy-gradient methods include REINFORCE [58], TRPO [24], PPO
[25]. Critic-only (value-based) methods are a family of RL algorithms that
learn an DNN-approximated value-action function
$\tilde{q}({\bm{u}},{\bm{a}};\bm{\theta}_{q})$ based on the optimal Bellman
equation. Examples include Q-Learning [59], DQN [3], Dueling DQN [60], etc.
Since the optimization in critic-only methods is always performed off-policy,
they are more sample-efficient than the actor-only methods that are often on-
policy. However, since critic-only methods optimize the policy indirectly,
they are less stable compared to the actor-only methods. The actor-critic (AC)
methods learn a value function to support the policy gradient optimization,
and thus the AC family combines the strengths of both actor and critic
methods. As such, AC-based methods will be used in the proposed PiMBRL for
value/policy optimization.

###  2.2 Physics-informed model-based reinforcement learning

We propose a physics-informed model-based reinforcement learning (PiMBRL)
framework, where the physics knowledge (e.g., conservation laws, governing
equations, and boundary conditions) of the environment is incorporated to
inform the model learning and RL optimization. In this work, we focus on the
Dyna-style MBRL formulation with the off-policy AC-based optimization. The
proposed framework will retain the generality and optimality of model-free AC-
based DRL methods, while significantly reducing the real-world interactions by
learning a reliable environment model based on physics-informed discrete
learning.

Figure 2: Schematics of Dyna-style physics-informed model-based actor-critic
algorithm.

Specifically, an AC-based RL agent will be initialized and starts interacting
with the real environment. State-action data pairs and corresponding rewards
$({\bm{u}}_{t}^{r},{\bm{a}}_{t}^{r},{\bm{u}}_{t+1}^{r},r_{t}^{r})$ are
iteratively collected from the real environment and saved into the real replay
buffer $\mathscr{D}^{r}$. These real samples are used to train the actor-
critic agent and the transition model simultaneously. The model is constructed
by an auto-encoding recurrent network, where boundary conditions (BCs) of the
system are strictly encoded, and the governing physics are imposed softly by
minimizing the violation of the conservation laws. Synthetic samples
$({\bm{u}}_{t}^{f},{\bm{a}}_{t}^{f},{\bm{u}}_{t+1}^{f},r_{t}^{f})$ are
generated by the model and are collected into the imagined replay buffer
$\mathscr{D}^{f}$, which will be leveraged to augment the real samples for the
RL update. The model training, data generation, environment interaction, and
RL agent optimization are conducted iteratively in an online manner. The
overall schematic of the proposed PiMBRL is shown in Fig. 2, and the detailed
algorithm is given by Algorithm 1. More details of AC optimization and
physics-informed model construction will be elaborated in the following
subsections.

Algorithm 1 Physics-informed Model-Based Reinforcement Learning (PiMBRL)

1:Randomly initialize policy (actor) network
$\pi({\bm{u}};\bm{\theta}_{\pi})$, value (critic) network(s)
$q({\bm{u}},{\bm{a}};\bm{\theta}_{q})$, transition model
$\tilde{\mathcal{F}}({\bm{u}},{\bm{a}};\bm{\theta}_{F})$, and replay buffers
$\mathscr{D}^{r},\mathscr{D}^{f}$ for real and fictitious environments.

2:Randomly initialize state ${\bm{u}}_{0}$ (or observed state
${\bm{u}}^{o}_{0}$), reward $r_{0}$, and done signal $d_{0}$ for real
environment.

3:for episode$=1,M$ do

4: for $i=0,T$ do

5: Execute action ${\bm{a}}_{i}=\pi({\bm{u}};\bm{\theta}_{\pi})$ in the real
environment $\mathcal{F}$;

6: Save new data pair
$({\bm{u}}^{o}_{i},{\bm{a}}_{i},{\bm{u}}^{o}_{i+1},r_{i},d_{i})$ to the real
buffer $\mathscr{D}^{r}$;

7: if episode ends then Reset the environment $\mathcal{F}$;

8: end if

9: end for

10: if sufficient ($n_{s_{M}}$) state-action pairs stored in the real buffer
$\mathscr{D}^{r}$ then

11: Sample a batch of real state-action pairs,
$\\{({\bm{u}}^{o}_{i},{\bm{a}}_{i},{\bm{u}}^{o}_{i+1})\\}_{i=1}^{n_{b_{r}}}$,
from $\mathscr{D}^{r}$

12: Update the transition model
$\tilde{\mathcal{F}}({\bm{u}},{\bm{a}};\bm{\theta}_{F})$ using the data loss
$L_{D}$ on the batch $\mathscr{D}^{r}$;

13: end if

14: if transition model meets the accuracy threshold (data loss
$L_{D}<\lambda$) then

15: for model prediction length$=1,l_{M}$ do

16: Sample a batch of states $\\{({\bm{u}}_{j+1})\\}$ from
$\\{\mathscr{D}^{r},\mathscr{D}^{f}\\}$

17: Execute actions
$\\{{\bm{a}}_{j+1}=\pi({\bm{u}}_{j+1};\bm{\theta}_{\pi})\\}^{n_{b_{f}}}_{j=1}$
in transition model $\tilde{\mathcal{F}}$;

18: Save new data pairs
$\\{({\bm{u}}^{o}_{j+1},{\bm{a}}_{j+1},{\bm{u}}^{o}_{j+2},r_{j+2},d_{j+2})\\}^{n_{b_{f}}}_{j=1}$
to buffer $\mathscr{D}^{f}$;

19: end for

20: end if

21: if enough ($n_{s_{R}}$) state-action pairs stored in
$\\{\mathscr{D}^{r},\mathscr{D}^{f}\\}$ then

22: Sample a batch of state-action pairs
$\\{({\bm{u}}^{o}_{k},{\bm{a}}_{k},{\bm{u}}^{o}_{k+1})\\}$ from the fake
buffer $\mathscr{D}^{f}$

23: Update model $\tilde{\mathcal{F}}$ according to physical loss $L_{E}$ on
sampled state-action pairs;

24: end if

25: Sample a batch of
$\\{({\bm{u}}^{o}_{l},{\bm{a}}_{l},{\bm{u}}^{o}_{l+1},r_{l},d_{l})\\}$ from
the augmented buffer $\\{\mathscr{D}^{r},\mathscr{D}^{f}\\}$

26: Update policy network $\pi({\bm{u}};\bm{\theta}_{\pi})$ and value network
$q({\bm{u}},{\bm{a}};\bm{\theta}_{q})$ on the sampled state-action pairs using
off-policy algorithms111In this paper, TD3 is used as a demonstration (details
of the TD3 is given in Algorithm 3), but other off-policy algorithms such as
Deep Deterministic Policy Gradient (DDPG) and Soft Actor-Critic (SAC) are also
applicable.(see Algorithm 3 in Appendix).

27:end for

####  2.2.1 Dyna-style model-based actor-critic optimization

We consider a generic Dyna-style model-based actor-critic (AC) optimization
algorithm, which can be easily adapted to any state-of-the-art off-policy AC
methods, e.g., DDPG, TD3, or SAC. As mentioned above, in the Dyna-style
formulation, model is used to augment real samples, and thus model-free AC-
based optimization can be directly leveraged. For a AC-based RL agent, two
neural networks $\tilde{\pi}({\bm{u}};\bm{\theta}_{\pi})$ and
$\tilde{q}({\bm{u}},{\bm{a}};\bm{\theta}_{q})$ are constructed to represent
the policy and value functions, respectively. Based on samples from both the
real environment and virtual environment simulated by the model, the policy
network $\tilde{\pi}({\bm{u}};\bm{\theta}_{\pi})$ is iteratively updated by,

$\bm{\theta}_{\pi}^{k+1}=\bm{\theta}_{\pi}^{k}+\alpha_{\pi}\nabla_{\bm{\theta}_{\pi}}J(\bm{\theta}_{\pi}^{k}),$
(6)

where $\alpha_{\pi}$ is the learning rate and
$\nabla_{\bm{\theta}_{\pi}}J(\bm{\theta}_{\pi})$ represents the policy
gradient with respect to actor network parameters $\bm{\theta}_{\pi}$, which
can be calculated based on the critic network,

$\nabla_{\bm{\theta}_{\pi}}J(\bm{\theta}_{\pi}^{k})=\mathbb{E}\left[\sum_{t=0}^{T}\nabla_{\bm{\theta}_{\pi}}\log\tilde{\pi}({\bm{u}}_{t};\bm{\theta}_{\pi}^{k})\cdot\tilde{q}({\bm{u}}_{t},{\bm{a}}_{t};\bm{\theta}_{q}^{k})\right].$
(7)

The critic network $\tilde{q}({\bm{u}},{\bm{a}};\bm{\theta}_{q})$ is optimized
by minimizing the temporal difference (TD)-based loss function,

$\bm{\theta}_{q}^{*}=\underset{\bm{\theta}_{q}}{\arg\min}\|q^{\prime}_{t}-\tilde{q}({\bm{u}}_{t},{\bm{a}}_{t};\bm{\theta}_{q})\|_{L_{2}},$
(8)

where $\|\cdot\|_{L_{2}}$ represent L2 norm and $q^{\prime}_{t}$ is estimated
based on the optimal Bellman equation,

$q^{\prime}_{t}=r_{t}+\gamma\tilde{q}\bigg{(}{\bm{u}}_{t+1},\tilde{\pi}({\bm{u}}_{t+1};\bm{\theta}_{\pi});\bm{\theta}_{q}\bigg{)}.$
(9)

####  2.2.2 Physics-informed Learning architecture for transition dynamics

We develop a physics-informed recurrent neural network to learn the dynamics
transition, aiming to map the current states and actions to the states at the
next control step
($\tilde{\mathcal{F}}:{\bm{u}}_{t},{\bm{a}}_{t}\to{\bm{u}}_{t+1}$). To better
capture the spatiotemporal dependencies, a convolutional encoder, multi-layer
perceptron (MLP) decoder, and long-short term memory (LSTM) blocks are
utilized to build the learning architecture. As shown in Fig. 3, the high-
dimensional state vector (${\bm{u}}_{t}$) at the current control step is
encoded into the latent space by a convolutional encoder. Together with the
input actions (${\bm{a}}_{t}$), the latent state vector is fed into the LSTM-
based transition network, which outputs latent intermediate transition states
between the two control steps. After a multi-layer perceptron (MLP) decoder,
the latent outputs are decoded to the full-order physical states.

Figure 3: Schematics of LSTM-based neural network architecture for transition
model

The network is trained based on both data and physics constraints. The data
loss $L_{D}$ is defined by the mismatch between the model prediction
$\tilde{{\bm{u}}}_{t+1}$ and labeled data ${\bm{u}}_{t+1}$ obtained from
interactions with the real environment,

$L_{D}=\frac{1}{n_{b}}\sum_{j=1}^{n_{b}}\,\left\|\tilde{{\bm{u}}}_{t+1}^{(j)}-{\bm{u}}_{t+1}^{(j)}\right\|_{L_{2}},$
(10)

where $n_{b}$ is the batch size. The physical loss $L_{E}$ is constructed
based on the conservation laws of the system in their discretized form. The
residuals of the governing equations is minimized on multiple discrete
spatiotemporal snapshots. To this end, the network outputs include the states
at $N$ intermediate time steps ($\\{{\bm{u}}_{t+i/N}\\}_{i=1,N}$) as well as
the state at the next control step. The physical loss is then obtained by
taking the averaged residuals of the governing equations,

$L_{E}=\frac{1}{N}\sum_{i=1}^{N}\sum_{j=1}^{n_{b}}\,\left\|\frac{d\tilde{{\bm{u}}}_{t+i/N}^{(j)}}{dt}-\mathcal{F}(\tilde{{\bm{u}}}_{t+i/N}^{(j)},\,{\bm{a}}_{t+i/N})^{n_{b}}\right\|_{L_{2}},$
(11)

where spatial derivatives in $\mathcal{F}$ are computed using high-order
finite-difference-based spatial filtering, and temporal derivatives
$\frac{d{\bm{u}}_{t+i/N}}{dt}$ are approximated by forward Euler method.

##  3 Results

In this section, we test the proposed PiMBRL on a number of classic control
problems that are governed by ODEs or PDEs. The performance is compared
against the baseline model-free and model-based counterparts. The standard
model-free TD3 algorithm is used as the MFRL baseline (see Algorithm 2 in
Appendix), while the purely data-driven dyna-like model-based TD3 is deemed as
the MBRL baseline. The hyperparameters used in the following experiments are
summarized by Tables 1 and 2 in Appendix.

###  3.1 ODE governed environments

We first evaluate PiMBRL on two classic dynamic control benchmarks, Cart-Pole
and Pendulum, which are available in the OpenAI Gym environment. The physics
of both systems are known, which can be described by a set of ODEs. Since
there is no spatial dependence and the dimension of the system is low, we use
a low-capacity two-layer MLP with 256 neurons per layer to directly learning
the transition dynamics
${\bm{u}}_{t+1}=\tilde{\mathcal{F}}({\bm{u}}_{t},{\bm{a}}_{t})$ and model
rollout length ($l_{M}$) is set equal to the trajectory length.

####  3.1.1 Cart-Pole

We start with the Cart-Pole benchmark problem (i.e., “CartPole-v0” environment
provided in OpenAi gym), where a cart moves along a frictionless track with a
pole attached to the top of it via an unactuated joint, as shown in Fig.4 (a).
The control goal is to keep the pole from falling over by acting a horizontal
force on the cart. The reward is $+1$ for each time step as long as the pole
is upright and the cart remains in a certain region. The physics of this
system is governed by,

$\begin{cases}\displaystyle\ddot{x}=\frac{f+m_{p}\dot{\theta}^{2}\sin{\theta}-m_{p}l\ddot{\theta}}{m_{p}+m_{c}}\\\
\displaystyle\ddot{\theta}=\frac{g\sin{\theta}-\cos{\theta}\displaystyle\frac{f+m_{p}\dot{\theta}^{2}\sin{\theta}}{m_{p}+m_{c}}}{l\left(\displaystyle\frac{4}{3}-\frac{m_{p}\cos^{2}{\theta}}{m_{c}+m_{p}}\right)},\end{cases}$
(12)

where $x$ is the spatial coordinate of the cart, $\theta$ represents the angle
of the pole from vertical, and $f$ is the force that the RL agent applies on
the cart. $m_{c},m_{p}$ are the mass of the cart and pole, respectively. One
episode is considered to be ended if the pole deviates too much from vertical
position (i.e., $|\theta|>\pi/12$) or the cart leaves the designated area
(i.e., $|x|>2.4$) or one episode has more than 200 control steps. The state
observation of this environment is a four-dimensional vector,
${\bm{u}}=(x,\dot{x},\theta,\dot{\theta})$, while the action space is
discrete, consisting of two valid values $\\{-10,10\\}$. Each episode begins
at a random state
${\bm{u}}_{0}=(x_{0},\dot{x}_{0},\theta_{0},\dot{\theta}_{0})$.

(a) Cart-Pole (b) RL Performance

Figure 4: (a) Schematic diagram of the Cart-Pole environment. (b) Performance
curve of PiMBRL versus standard model-free TD3 (MFRL) and dyna-like model-
based TD3 (MBRL) in Cart-Pole environment. The solid lines indicate averaged
returns of 100 randomly selected test episodes, while the shaded area
represents the return distribution of all test samples.

Figure 4 (b) compares the performance curves of the MFRL, MBRL, and PiMBRL.
The proposed PiMBRL reaches the total return of 200 only after about 3000 time
steps in the real environment, while the vanilla dyna-like MBRL counterpart
needs much longer to achieve so and its performance is not stable as well. For
the MFRL counterpart, the total return is still below 120 even after 15000
time steps. Although both MBRL and MFRL are able to achieve the same
performance with sufficient time steps, PiMBRL only uses about 45.2% and 9.7%
time steps needed by its MBRL and MFRL counterparts, respectively. Therefore,
to achieve the same level of performance, PiMBRL can significantly reduce the
required number of interactions with the real environment, compared to the
original model-free TD3 (i.e., MFRL) and dyna-like model-based TD3 (i.e.,
MBRL).

####  3.1.2 Pendulum

The second test case is the Pendulum-v0 available in the OpenAi gym. In this
environment (see Fig. 5 (a) ), one end of the pendulum is fixed, while the
other end can swing freely. $\theta$ denotes the angle of the pendulum from
vertical position. The state contains the angle and its time derivative, i.e.,
${\bm{u}}=(\theta,\dot{\theta})$. In each episode, the pendulum starts from a
random state ${\bm{u}}_{0}\in(-1,1)\times(-1,1)$. Besides, the angular
velocity is constrained as $\dot{\theta}\in[-8,8]$, and any $\dot{\theta}$ out
of this range will be capped by the boundary value ($-8$ or $8$). The dynamics
of the pendulum system is governed by,

$\ddot{\theta}=-\frac{3g}{2l}\sin{(\theta+\pi)}+\frac{3}{ml^{2}}T,$ (13)

where $g=10$ is the acceleration of gravity, $T\in[-2,2]$ denotes the torque
that the agent applies to the pendulum, $l=1$ and $m=1$ are the length and
mass of the pendulum, respectively. The control goal here is to swing the
pendulum up and make it stays upright, meanwhile consuming as less energy as
possible. To achieve this, the reward function is defined as,

$r=-\theta^{2}-0.1\dot{\theta}^{2}-0.001T^{2}$ (14)

The performance curve is shown in Fig.5 (b). After about 3000 time steps, both
the PiMBRL and MBRL achieve averaged total return of -200 with reduced
uncertainty. In contrast, the total return of the MFRL largely fluctuates and
the average value remains less than -800. Although PiMBRL shows greater sample
efficiency over the MFRL baseline, it does not show a notable advantage over
the MBRL counterpart in the Cart-Pole case.

(a) Pendulum (b) RL performance

Figure 5: (a) Schematic diagram of pendulum environment. (b) Performance curve
of PiMBRL versus standard model-free TD3 (MFRL) and dyna-like model-based TD3
(MBRL) in the pendulum environment. The solid lines indicate averaged returns
of 100 randomly selected test episodes, while the shaded area represents the
return distribution of all test samples.

This might be due to the following two factors. First, compared to the Cart-
Pole problem, the dynamics of the pendulum environment is easier for the model
to learn, since the observation space is much smaller than that of the Cart-
Pole environment (2-dimension vs. 4-dimension). Second, the model network is
benefited less from the physical loss, since when the state $\dot{\theta}$
exceeds the limits $[-8,8]$ and capped by the boundary value, the governing
equation eq.(13) is no longer satisfied.

(a) Model trained with the physics-informed loss in PiMBRL

(b) Model trained with the data loss only in MBRL

Figure 6: Histograms of the prediction errors of the transition models for the
Pendulum environment, (a) trained with physics-informed loss in PiMBRL, and
(b) trained with the data loss only in MBRL.

Figure 6 compares the histograms of the prediction errors of the models
trained by PiMBRL and MBRL, respectively. Fig. 6 (a) shows the prediction
error of the model trained with the physics-informed loss (equation loss +
data loss), while Fig. 6 (b) shows the model prediction error in MBRL where
only the labeled data are used for training. Although the model does achieve
higher accuracy on average without using physics constraints (see two sub-
figures in the left column), in the out-of-sample regime (away from the
training set), the physics-informed model shows better performance and
robustness (see two sub-figures in the right column). Overall, the models in
both PiMBRL and MBRL are learned sufficiently well to achieve a roughly
similar RL performance.

###  3.2 PDE governed environments

Unlike the environments that can be described by ODEs, the systems governed by
PDEs are much more complicated in terms of the dimension of spatiotemporal
solution space, dramatically increasing the level of difficulty in learning
the dynamic model as well as the policy and value functions. In this section,
we evaluate the proposed PiMBRL structure on two continuous control problems
in the environments governed by Burgers’ equation and Kuromoto-Sivashinsky
(KS) equation, respectively. The learned transition model is only used to
predict a limited rollout length ($l_{M}$) in one trajectory to control the
model error accumulation, making it remain at a relatively low level.

####  3.2.1 Burgers’ equation

For the first PDE-based control problem, we consider a 1-D Burgers’ equation
with periodic boundary condition,

$\frac{\partial u}{\partial t}+\frac{1}{2}u\frac{\partial u}{\partial
x}=\nu\frac{\partial^{2}u}{\partial x^{2}}+f(x,t),\quad x\in[0,l],\
t\in[0,2\pi],$ (15)

where $x$ is the spatial coordinate, $\nu=0.01$ is the kinematic viscosity,
and $f(x,t)$ denotes the source term, defined as,

$f(x,t)=a_{1}(t)\exp\left[{\left(-15(\frac{x}{l}-0.25)\right)}^{2}\right]+a_{2}(t)\exp\left[{\left(-15(\frac{x}{l}-0.75)\right)}^{2}\right]$
(16)

with the control parameters ${\bm{a}}=(a_{1},a_{2})\in{[-0.025,0.075]}^{2}$.

The control problem is defined as training the RL agent to match a reference
trajectory. Namely, the RL agent is trained to find the optimal strategy of
controlling the source term with two control parameters $a_{1},a_{2}$, in
order to match a predefined reference trajectory profile ${\bm{u}}_{re}$,

$u_{re}(x,t)=0.05\sin{t}+0.5,\quad t\in[0,2\pi].$ (17)

Each episode starts from a randomly generated initial condition,

$u(x,0)=0.2\,c\exp\left[{(-5(\frac{x}{l}-0.5))}^{2}\right]+0.2\,(1-c)\left(0.5\sin{4\pi\frac{x}{l}}+0.5\right),$
(18)

where $c$ is randomly sampled from a uniform distribution on $[\,0,1)$. That
is, the trained RL is expected to finally match the reference trajectory,
starting from any randomly generated initial state by eq.(17). The observation
is set as the discrepancy between the PDE state and reference state at the
same control step ${\bm{u}}^{o}={\bm{u}}-{\bm{u}}_{re}$. The environment is
simulated numerically based on the finite difference methods, where the
convection term and diffusion term are discretized by the second-order upwind
scheme and fourth-order central difference scheme, respectively. Euler method
is used for the time integration. The simulated environment is defined on a
spatial mesh of 150 grid points and the numerical time step is set as 0.01.
The control signal is applied every 500 numerical steps and one episode
contains 60 control steps. The reward function is defined as
$-10\left\|{\bm{u}}^{o}\right\|_{L_{2}}$.

(a) Uncontrolled (b) Uncontrolled actions & rewards

(c) Controlled (d) Controlled actions & rewards

Figure 7: Results of one test episode (a) with random control signals and (c)
with trained RL controller. The corresponding actions and reward curves of the
(b) uncontrolled episode and (d) RL controlled episode.

Without the RL training, random control signals are applied to the system, and
the corresponding spatiotemporal state surface ${\bm{u}}$ is shown in Fig. 7a.
We can see that the state surface is unsmooth, and a large discrepancy remains
between the uncontrolled state with the reference state. Figure 7c shows one
of the test episodes controlled by the trained PiMBRL agent. The corresponding
actions & rewards are given in Fig. 7d. Although the initial state is far from
the reference, the controlled surface gradually approaches and finally matches
the reference state after $t=40$. The corresponding action curves are smooth,
suggesting that the RL agent successfully learns an effective control
strategy.

Figure 8: Performance curve of PiMBRL versus standard MFRL and MBRL in the
Burgers’ equation environment. The solid lines indicate averaged returns of
100 test episodes, while the shaded area represents the return distribution of
all test samples.

Figure 8 shows the performance curves of the PiMBRL, MBRL, MFRL tested on 100
randomly selected initial conditions. The PiMBRL reaches the total return of
0.1 only after 800 time steps, while it takes the MBRL about 1300 time steps
to achieve a similar level of performance. The MFRL counterpart can not reach
the same level of performance within 1400 time steps. Again, our PiMBRL shows
significant advantages in terms of sample efficiency, since it only uses about
65% of time steps required by MBRL and 46.7% time steps required by MFRL to
achieve the control goal.

####  3.2.2 Kuramoto-Sivashinsky (KS) Equation

In the last case, we evaluate the proposed PiMBRL on the control of a
nonlinear, chaotic dynamic system governed by the 1-D Kuramoto-Sivashinsky
(KS) equation, which is more challenging. The system governed by the KS
equation often exhibits spatiotemprally chaotic or weakly turbulent behavior,
and thus the KS equation is widely used as a model system for turbulence study
[61]. In this case, the KS environment is controlled by four actuators
distributed equally in space to minimize the energy dissipation and total
input power. The physics of this system is governed by the KS equation,

$\frac{\partial u}{\partial t}+\frac{\partial^{2}u}{\partial
x^{2}}+\frac{\partial^{4}u}{\partial x^{4}}+\frac{1}{2}u\frac{\partial
u}{\partial x}=f(x,t),\ \ x\in[0,l],\ t\in[0,\infty),$ (19)

where $u$ is the state variable, and $f$ represents the source term (i.e.,
actuator) defined by.

$f(x,t)=\sum_{i=1}^{4}\frac{a_{i}(t)e^{-(x-x_{i})^{2}/2}}{\sqrt{2\pi}},$ (20)

where $x_{i}\in\\{0,\,l/4,\,l/2,\,3l/4\\}$ is the spatial locations of the
actuator, and ${\bm{a}}=\\{a_{i}(t)\\}_{i=1,2,3,4}\in[-0.5,0.5]^{4}$ defines
the control parameters. To achieve the control goal, the reward function is
defined as follows,

$r=-\frac{1}{Tl}\int_{t_{0}}^{t_{0}+T}\int_{0}^{l}\left((\frac{\partial^{2}u}{\partial
x^{2}})^{2}+(\frac{\partial u}{\partial x})^{2}+uf\right)\,dx\,dt$ (21)

where $T$ is the time length of one control step. The environment is simulated
numerically based on the finite difference method, where the convection term
is discretized by the second-order upwind scheme, the second and fourth
derivatives are discretized by the 6th-order central difference scheme, and
the 4th-order Runge-Kutta scheme is used for time integration with time
stepping size of 0.001 on the 1D domain $l=8\pi$ discretized by a mesh of 64
grid points.

Figure 9: Contours of spatiotemporal state $u$ of four randomly-selected
episodes. The top one is an uncontrolled episode, while the other three are
controlled by the trained RL agent.

Each control step contains 250 numerical steps, and one episode consists of
400 control steps. Each episode starts with a random initial condition sampled
from the attractor of the unforced KS equation. Figure 9 shows the
spatiotemporal states of four test episodes with randomly sampled initial
states. The top one is an uncontrolled episode, where nonlinear chaotic
behavior is developed along the time axis. In contrast, the “turbulence” in
the other three episodes controlled by the agent can be quickly stabilized
after 200 time steps, showing the effectiveness of the PiMBRL controller.

Figure 10 shows the performance curve of PiMBRL versus that of the MFRL. It is
clear that the PiMBRL agent reaches higher averaged total returns with fewer
time steps than the MFRL counterpart does. The PiMBRL performance curve is
consistently above that of the MFRL approach, and meanwhile, less uncertainty
is observed. In PiMBRL, the agent only uses about $30.2\%$ of the time steps
required by its MFRL counterpart to reach the averaged total return of $-55$.
In this case, a model-free fine-tuning approach [62] is applied when the
averaged return is above -55 to further improve the PiMBRL performance. This
is because when the RL agent is trained to achieve a high accuracy level, the
model-based exploration does not help too much while the model bias becomes
the bottleneck. At this point, the RL agent can be fine-tuned by interacting
with the real environment without using the model [62]. As such, the PiMBRL
and MFRL performance curves will eventually converge to each other with a
large number of training time steps.

Figure 10: Performance curve of PiMBRL versus MFRL and MBRL in KS equation
environment. Solid lines shows the average return of all the test episodes
while the shaded area represents the distribution range of the 200 test
episodes

The MBRL result is not plotted here because the accuracy ($L_{D}$) of the
learned model purely based on labeled data does not meet the accuracy
threshold $\lambda=0.01$, and thus the MBRL performance curve is identical to
that of the MFRL counterpart. In contrast to the ODE-based pendulum system
shown above, the spatiotemporal dynamics of the chaotic system is much more
challenging for the model to learn purely based on the limited amount of
labeled data. Incorporating physics constraints can significantly help this
scenario and improve the model learning performance. This can be clearly seen
in Fig. 11, where the comparison of model prediction errors of PiMBRL (left)
and MBRL (right) are shown. The mean squared error distribution of the model
trained based on the physics-informed loss is much less than that of the model
trained purely based on the labeled data.

Figure 11: Histograms of the prediction errors (MSE) of the transition models
for the KS system, (left) trained with physics-informed loss in PiMBRL, and
(right) trained with the data loss only in MBRL. Both are trained on the same
real buffer with $10^{5}$ time steps stored.

##  4 Discussion

###  4.1 Influence of model rollout length

The rollout length of the model is important to the RL performance. On the one
hand, a long-term model prediction could help the agent see “deeper and
further”, improving the exploration rate and increasing the sample efficiency
in the Dyna-like MBRL framework. On the other hand, accurately predicting a
long trajectory is always challenging due to the error accumulation effect,
and inaccurate prediction data could be harmful to the RL training, which is a
trade-off.

Figure 12: Model predicted snapshots of the KS environment with different
rollout lengths compared with the true environment at three randomly selected
initials.

Figure 12 shows the model prediction performance for the KS environment with
four different rollout lengths ($l_{M}=3,8,40,120$). For all three randomly
selected initials, we can see that the states predicted by the model agree
well with the ground truth when the rollout length is small ($l_{M}<=8$).
However, large discrepancies can be observed when $l_{M}>=40$ and the model
predictions significantly deviate from the ground truth when the rollout
length is 120.

We are more interested in how the model rollout length affects the PiMBRL
performance, so we investigated the influence of four different rollout
lengths on the reward curves. The case setting is the same as that used in
section 33.23.2.2 except for the rollout length, and the RL performance is
evaluated on 200 randomly generated episodes.

Figure 13: Performance curves of PiMBRL with four different model rollout
lengths compared with MFRL base line. (Left bottom) Zoom-in view of the dashed
box. (right bottom) Zoom-in view of the range between $1\times 10^{5}$ to
$4\times 10^{5}$ time steps.

Figure 13 shows the performance curves of PiMBRL with different model
prediction lengths. As expected, the RL performance slightly deteriorates if
the rollout length of the model is either too short or too long. For the four
rollout lengths ($l_{M}=3,8,40,120$), the RL agent achieves the best
performance with $l_{M}=8$ at almost any stage of the entire training process.
When $l_{M}<=8$ or $l_{M}>=8$, the RL convergence speed is relatively slow
before entering the fine-tuning stage due to the trade-off between exploration
and model bias. However, even after the model-free fine tuning, the RL agent
with a longer rollout length ($l_{M}>=40$) still suffers from the low-quality
model prediction data and is very difficult to be further improved (see the
comparison of the green and blue curves in Fig. 13).

###  4.2 Influence of model accuracy threshold.

In model-based RL, the transition model is trained along with the RL agent
(i.e., value and policy networks) from scratch, and the model-generated data
usually can be effectively used once the model is trained to reach a certain
level of accuracy. As mentioned above, we use the model accuracy threshold
parameter ($\lambda$) to determine when the model-predicted data should be
utilized for the RL training. Here, we would like to study how this parameter
of $\lambda$ affects the PiMBRL performance.

Using the KS environment as an example, Fig. 14 shows the performance curves
of PiMBRL with three different threshold values, i.e.,
$\lambda=0.02,0.01,0.005$. Overall, the influence of $\lambda$ values on the
final performance of the RL agent is negligible since all cases converge to
the same level of total return. This is because the model is trained together
with the RL agent and can be improved over the entire RL training process.
Actually, with the same amount of training time steps in the real environment,
the models can reach a similar accuracy level regardless of threshold values
$\lambda$. However, the RL agent with a relatively large threshold value
($\lambda=0.02$) performs slightly better and has a relatively faster
convergence rate at the early stage of the training (time step 40000 to
60000). A higher threshold allows the agent to access more data generated by
the model earlier, which leads to a better exploration rate and thus slightly
higher sample efficiency (see Fig. 14b).

(a) Full view (b) Zoom-in view

Figure 14: Performance curves of PiMBRL with three different model accuracy
threshold values $\lambda=0.02,0.01,0.005$, where (a) is the full view of the
entire training history and (b) shows the zoom-in view of the dashed box area
in (a).

##  5 Conclusion

In this work, we presented an innovative model-based reinforcement learning
framework (PiMBRL) for dynamic control, which leverages the physical laws and
constraints of the environment to alleviate the model-bias issue and
significantly improve the sample efficiency. In particular, an autoencoding
based recurrent network structure is devised to learn the spatiotemporal
dynamics of the environments. The merit and effectiveness of the proposed
PiMBRL framework have been demonstrated over a set of classic dynamic control
problems, where the environments are governed by canonical ODEs or PDEs,
including viscous Burgers’ equations and KS equations with chaotic behaviors.
Compared to the model-free or purely data-driven model-based reinforcement
learning counterparts, our PiMBRL shows a significant improvement in model
predictive accuracy and RL sample efficiency. Moreover, the effects of
different hyper-parameters used in PiMBRL (e.g., model rollout length $l_{M}$
and model accuracy threshold $\lambda$) are studied, and how these parameters
affect the RL performance is discussed.

The data that support the findings of this study will be openly available in
GitHub at https://github.com/Jianxun-Wang/PIMBRL upon publication.

X.-Y. Liu: participated in the design of the study, implemented the entire
framework, carried out all numerical experiments, and drafted the manuscript.
J.-X. Wang: conceived of the study, drafted and revised the manuscript, and
supervised the project.

We declare we have no competing interests.

This work is funded by the National Science Foundation under award numbers
CMMI-1934300 and OAC-2047127 and startup funds from the College of Engineering
at University of Notre Dame.

## Appendix

###  5.1 Twin-delayed deep deterministic policy gradient (TD3)

Algorithm 2 Model-free Twin-delayed deep deterministic policy gradient (TD3)

1:Initialize policy (actor) network $\pi({\bm{u}};\bm{\theta}_{\pi})$, value
(critic) networks $q_{1}({\bm{u}},{\bm{a}};\bm{\theta}_{q_{1}})$,
$q_{2}({\bm{u}},{\bm{a}};\bm{\theta}_{q_{2}})$, empty the replay buffer
$\mathscr{D}^{r}$ and reset the environment $\mathcal{F}$.

2:Make a copy of policy and value networks as target networks
$\pi_{targ}\leftarrow\pi$, $q_{targ,1}\leftarrow q_{1}$, $q_{targ,2}\leftarrow
q_{2}$

3:for time steps $t=1,N$ do

4: Execute action ${\bm{a}}_{i}=\pi({\bm{u}};\bm{\theta}_{\pi})$ in the
environment $\mathcal{F}$;

5: Save new data pair
$({\bm{u}}^{o}_{i},{\bm{a}}_{i},{\bm{u}}^{o}_{i+1},r_{i},d_{i})$ to buffer
$\mathscr{D}^{r}$;

6: if episode ends then Reset the environment $\mathcal{F}$;

7: end if

8: if $t\mod\texttt{update\\_every}==0$ then

9: if enough state-action pairs stored in
$\\{\mathscr{D}^{r},\mathscr{D}^{f}\\}$ then

10: for $k=1,I_{RL}$ do

11: Sample $J$ state-action pairs
$\\{({\bm{u}}^{o}_{j},{\bm{a}}_{j},{\bm{u}}^{o}_{j+1},r_{j},d_{j})\\}$ from
buffer $\mathscr{D}^{r}$

12: Compute target
$Q_{j}=r_{j}+\gamma(1-d_{j})\min_{i=1,2}\left\\{q_{targ,i}({\bm{u}}^{o}_{j+1},\pi_{targ}({\bm{u}}^{o}_{j+1}))\right\\}$

13: Update value networks via gradient descent,
$\nabla_{\bm{\theta}_{q_{i}}}\frac{1}{J}\sum_{j=1}^{J}\left[q_{i}({\bm{u}}^{o}_{j},{\bm{a}}_{j})-Q_{j}\right]^{2},\qquad
i=1,2$

14: if $k\mod 2==0$ then

15: Update policy network via gradient descent:
$\nabla_{\bm{\theta}_{\pi}}\frac{1}{J}\sum_{j=1}^{J}q_{1}({\bm{u}}^{o}_{j},\pi({\bm{u}}^{o}_{j}))$

16: Update target networks:
$\begin{split}\bm{\theta}_{q_{targ,i}}\leftarrow\rho\bm{\theta}_{q_{targ,i}}+(1-\rho)\bm{\theta}_{q_{i}},\qquad
i=1,2\\\
\bm{\theta}_{\pi_{targ}}\leftarrow\rho\bm{\theta}{\pi_{targ}}+(1-\rho)\bm{\theta}{\pi}\qquad\qquad\qquad\end{split}$

17: end if

18: end for

19: end if

20: end if

21:end for

TD3 hyper-parameters used in this work is shown in Table 1,

Parameter | discount rate $\gamma$ | $P_{D}$ | $I_{RL}$ |  | Policy / Value
---
network optimizer
| Policy / Value
---
network learning rate
Value |  | 0.99
---
(0.977 in KS)
2 | 50 | Adam | 1e-3
Table 1: TD3 hyper-parameters used in this work Algorithm 3 TD3-based policy
and value network update algorithm used in MBRL/PiMBRL

1:With buffer $\mathscr{D}^{r}$, and $\mathscr{D}^{f}$, policy network
$\pi({\bm{u}};\bm{\theta}_{\pi})$, value networks
$q_{1}({\bm{u}},{\bm{a}};\bm{\theta}_{q_{1}})$,
$q_{2}({\bm{u}},{\bm{a}};\bm{\theta}_{q_{2}})$ and their corresponding target
networks $\pi_{targ},q_{targ,1},q_{targ,2}$

2:for $k=1,I_{RL}$ do

3: Sample $J$ state-action pairs
$\\{({\bm{u}}^{o}_{j},{\bm{a}}_{j},{\bm{u}}^{o}_{j+1},r_{j},d_{j})\\}$ from
buffer $\\{\mathscr{D}^{r},\mathscr{D}^{f}\\}$

4: Compute target
$Q_{j}=r_{j}+\gamma(1-d_{j})\min_{i=1,2}\left\\{q_{targ,i}({\bm{u}}^{o}_{j+1},\pi_{targ}({\bm{u}}^{o}_{j+1}))\right\\}$

5: Update value networks by gradient descent:
$\nabla_{\bm{\theta}_{q_{i}}}\frac{1}{J}\sum_{j=1}^{J}\left[q_{i}({\bm{u}}^{o}_{j},{\bm{a}}_{j})-Q_{j}\right]^{2},\qquad
i=1,2$

6: if $k$ mod 2 == 0 then

7: Update policy network by gradient descent:
$\nabla_{\bm{\theta}_{\pi}}\frac{1}{J}\sum_{j=1}^{J}q_{1}({\bm{u}}^{o}_{j},\pi({\bm{u}}^{o}_{j}))$

8: Update target networks:
$\begin{split}\bm{\theta}_{q_{targ,i}}\leftarrow\rho\bm{\theta}_{q_{targ,i}}+(1-\rho)\bm{\theta}_{q_{i}},\qquad
i=1,2\\\
\bm{\theta}_{\pi_{targ}}\leftarrow\rho\bm{\theta}{\pi_{targ}}+(1-\rho)\bm{\theta}{\pi}\qquad\qquad\qquad\end{split}$

9: end if

10:end for

Hyper-parameters of PiMBRL used in section 3 is summarized in Table 2.

Environment |  | Model accuracy
---
threshold $L_{D}$
| Model rollout
---
length $l_{M}$
| Maximum
---
episode length
$n_{s_{M}}$ | $n_{s_{R}}$
Cart-Pole |  $1$e$-4$ |  $200$ (Full length) | $200$ | 800 | 1000
Pendulum |  $1$e$-2$ |  $200$ (Full length) | 200 | 6000 | 12000
Burgers’ |  $1$e$-2$ | $1$ | 60 | 120 | 120
KS |  $1$e$-2$ | $3$ | 400 | 6000 | 12000
Table 2: Hyper-parameters of PiMBRL used in this work (section 3)

## References

  * Silver et al. [2017] David Silver, Julian Schrittwieser, Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, et al.  Mastering the game of go without human knowledge.  _nature_ , 550(7676):354–359, 2017.
  * Silver et al. [2018] David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Marc Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, et al.  A general reinforcement learning algorithm that masters chess, shogi, and Go through self-play.  _Science_ , 362(6419):1140–1144, 2018.
  * Mnih et al. [2013] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller.  Playing atari with deep reinforcement learning.  _arXiv preprint arXiv:1312.5602_ , 2013.
  * Rabault et al. [2019] Jean Rabault, Miroslav Kuchta, Atle Jensen, Ulysse Réglade, and Nicolas Cerardi.  Artificial neural networks trained through deep reinforcement learning discover control strategies for active flow control.  _Journal of fluid mechanics_ , 865:281–302, 2019.
  * Ghraieb et al. [2021] H Ghraieb, J Viquerat, A Larcher, P Meliga, and E Hachem.  Single-step deep reinforcement learning for open-loop control of laminar and turbulent flows.  _Physical Review Fluids_ , 6(5):053902, 2021.
  * Ren et al. [2021] Feng Ren, Jean Rabault, and Hui Tang.  Applying deep reinforcement learning to active flow control in weakly turbulent conditions.  _Physics of Fluids_ , 33(3):037121, 2021.
  * Fan et al. [2020] Dixia Fan, Liu Yang, Zhicheng Wang, Michael S Triantafyllou, and George Em Karniadakis.  Reinforcement learning for bluff body active flow control in experiments and simulations.  _Proceedings of the National Academy of Sciences_ , 117(42):26091–26098, 2020.
  * Bucci et al. [2019] Michele Alessandro Bucci, Onofrio Semeraro, Alexandre Allauzen, Guillaume Wisniewski, Laurent Cordier, and Lionel Mathelin.  Control of chaotic systems by deep reinforcement learning.  _Proceedings of the Royal Society A_ , 475(2231):20190351, 2019.
  * Falk et al. [2021] Martin J Falk, Vahid Alizadehyazdi, Heinrich Jaeger, and Arvind Murugan.  Learning to control active matter.  _arXiv preprint arXiv:2105.04641_ , 2021.
  * Gustavsson et al. [2017] Kristian Gustavsson, Luca Biferale, Antonio Celani, and Simona Colabrese.  Finding efficient swimming strategies in a three-dimensional chaotic flow by reinforcement learning.  _The European Physical Journal E_ , 40(12):1–6, 2017.
  * Verma et al. [2018] Siddhartha Verma, Guido Novati, and Petros Koumoutsakos.  Efficient collective swimming by harnessing vortices through deep reinforcement learning.  _Proceedings of the National Academy of Sciences_ , 115(23):5849–5854, 2018.
  * Zhu et al. [2021] Yi Zhu, Fang-Bao Tian, John Young, James C Liao, and Joseph CS Lai.  A numerical study of fish adaption behaviors in complex environments with a deep reinforcement learning and immersed boundary–lattice Boltzmann method.  _Scientific Reports_ , 11(1):1–20, 2021.
  * Hwangbo et al. [2017] Jemin Hwangbo, Inkyu Sa, Roland Siegwart, and Marco Hutter.  Control of a quadrotor with reinforcement learning.  _IEEE Robotics and Automation Letters_ , 2(4):2096–2103, 2017.
  * Wada et al. [2021] Daichi Wada, Sergio A Araujo-Estrada, and Shane Windsor.  Unmanned aerial vehicle pitch control using deep reinforcement learning with discrete actions in wind tunnel test.  _Aerospace_ , 8(1):18, 2021.
  * Deng et al. [2021] Yingjie Deng, Tao Liu, and Dingxuan Zhao.  Event-triggered output-feedback adaptive tracking control of autonomous underwater vehicles using reinforcement learning.  _Applied Ocean Research_ , 113:102676, 2021.
  * Bhagat et al. [2019] Sarthak Bhagat, Hritwick Banerjee, Zion Tsz Ho Tse, and Hongliang Ren.  Deep reinforcement learning for soft, flexible robots: brief review with impending challenges.  _Robotics_ , 8(1):4, 2019.
  * Nagabandi et al. [2020] Anusha Nagabandi, Kurt Konolige, Sergey Levine, and Vikash Kumar.  Deep dynamics models for learning dexterous manipulation.  In _Conference on Robot Learning_ , pages 1101–1112. PMLR, 2020.
  * Li et al. [2020] Tianyu Li, Nathan Lambert, Roberto Calandra, Franziska Meier, and Akshara Rai.  Learning generalizable locomotion skills with hierarchical reinforcement learning.  In _2020 IEEE International Conference on Robotics and Automation (ICRA)_ , pages 413–419. IEEE, 2020.
  * Botvinick et al. [2019] Matthew Botvinick, Sam Ritter, Jane X Wang, Zeb Kurth-Nelson, Charles Blundell, and Demis Hassabis.  Reinforcement learning, fast and slow.  _Trends in cognitive sciences_ , 23(5):408–422, 2019.
  * Hessel et al. [2018] Matteo Hessel, Joseph Modayil, Hado Van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Azar, and David Silver.  Rainbow: Combining improvements in deep reinforcement learning.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 32, 2018.
  * Lillicrap et al. [2015] Timothy P Lillicrap, Jonathan J Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra.  Continuous control with deep reinforcement learning.  _arXiv preprint arXiv:1509.02971_ , 2015.
  * Fujimoto et al. [2018] Scott Fujimoto, Herke Hoof, and David Meger.  Addressing function approximation error in actor-critic methods.  In _International Conference on Machine Learning_ , pages 1587–1596. PMLR, 2018.
  * Haarnoja et al. [2018] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine.  Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor.  In _International Conference on Machine Learning_ , pages 1861–1870. PMLR, 2018.
  * Schulman et al. [2015] John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz.  Trust region policy optimization.  In _International conference on machine learning_ , pages 1889–1897. PMLR, 2015.
  * Schulman et al. [2017] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov.  Proximal policy optimization algorithms.  _arXiv preprint arXiv:1707.06347_ , 2017.
  * Moerland et al. [2020] Thomas M Moerland, Joost Broekens, and Catholijn M Jonker.  Model-based reinforcement learning: A survey.  _arXiv preprint arXiv:2006.16712_ , 2020.
  * Plaat et al. [2020] Aske Plaat, Walter Kosters, and Mike Preuss.  Model-based deep reinforcement learning for high-dimensional problems, a survey.  _arXiv preprint arXiv:2008.05598_ , 2020.
  * Deisenroth and Rasmussen [2011] Marc Deisenroth and Carl 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)_ , pages 465–472. Citeseer, 2011.
  * Tassa et al. [2012] Yuval Tassa, Tom Erez, and Emanuel Todorov.  Synthesis and stabilization of complex behaviors through online trajectory optimization.  In _2012 IEEE/RSJ International Conference on Intelligent Robots and Systems_ , pages 4906–4913. IEEE, 2012.
  * Levine and Koltun [2013] Sergey Levine and Vladlen Koltun.  Guided policy search.  In _International conference on machine learning_ , pages 1–9. PMLR, 2013.
  * Racanière et al. [2017] Sébastien Racanière, Théophane Weber, David P Reichert, Lars Buesing, Arthur Guez, Danilo Rezende, Adria Puigdomenech Badia, Oriol Vinyals, Nicolas Heess, Yujia Li, et al.  Imagination-augmented agents for deep reinforcement learning.  In _Proceedings of the 31st International Conference on Neural Information Processing Systems_ , pages 5694–5705, 2017.
  * Łukasz Kaiser et al. [2020] Łukasz Kaiser, Mohammad Babaeizadeh, Piotr Miłos, Błażej Osiński, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, Afroz Mohiuddin, Ryan Sepassi, George Tucker, and Henryk Michalewski.  Model based reinforcement learning for Atari.  In _International Conference on Learning Representations_ , 2020.  URL https://openreview.net/forum?id=S1xCPJHtDB.
  * Hafner et al. [2019] Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson.  Learning latent dynamics for planning from pixels.  In _International Conference on Machine Learning_ , pages 2555–2565. PMLR, 2019.
  * Hafner et al. [2020] Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi.  Dream to control: Learning behaviors by latent imagination.  In _International Conference on Learning Representations_ , 2020.  URL https://openreview.net/forum?id=S1lOTC4tDS.
  * Depeweg et al. [2016] Stefan Depeweg, José Miguel Hernández-Lobato, Finale Doshi-Velez, and Steffen Udluft.  Learning and policy search in stochastic dynamical systems with bayesian neural networks.  _arXiv preprint arXiv:1605.07127_ , 2016.
  * Kurutach et al. [2018] Thanard Kurutach, Ignasi Clavera, Yan Duan, Aviv Tamar, and Pieter Abbeel.  Model-ensemble trust-region policy optimization.  In _International Conference on Learning Representations_ , 2018.  URL https://openreview.net/forum?id=SJJinbWRZ.
  * Abdar et al. [2021] Moloud Abdar, Farhad Pourpanah, Sadiq Hussain, Dana Rezazadegan, Li Liu, Mohammad Ghavamzadeh, Paul Fieguth, Xiaochun Cao, Abbas Khosravi, U Rajendra Acharya, et al.  A review of uncertainty quantification in deep learning: Techniques, applications and challenges.  _Information Fusion_ , 2021.
  * Wang et al. [2018] Jian-Xun Wang, Christopher J Roy, and Heng Xiao.  Propagation of input uncertainty in presence of model-form uncertainty: a multifidelity approach for computational fluid dynamics applications.  _ASCE-ASME J Risk and Uncert in Engrg Sys Part B Mech Engrg_ , 4(1), 2018.
  * Wang et al. [2017] Jian-Xun Wang, Jin-Long Wu, and Heng Xiao.  Physics-informed machine learning approach for reconstructing reynolds stress modeling discrepancies based on dns data.  _Physical Review Fluids_ , 2(3):034603, 2017.
  * Ling et al. [2016] Julia Ling, Andrew Kurzawski, and Jeremy Templeton.  Reynolds averaged turbulence modelling using deep neural networks with embedded invariance.  _Journal of Fluid Mechanics_ , 807:155–166, 2016.
  * Duraisamy et al. [2019] Karthik Duraisamy, Gianluca Iaccarino, and Heng Xiao.  Turbulence modeling in the age of data.  _Annual Review of Fluid Mechanics_ , 51:357–377, 2019.
  * Raissi et al. [2019] Maziar Raissi, Paris Perdikaris, and George E Karniadakis.  Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations.  _Journal of Computational Physics_ , 378:686–707, 2019\.
  * Sun et al. [2020] Luning Sun, Han Gao, Shaowu Pan, and Jian-Xun Wang.  Surrogate modeling for fluid flows based on physics-constrained deep learning without simulation data.  _Computer Methods in Applied Mechanics and Engineering_ , 361:112732, 2020.
  * Gao et al. [2021a] Han Gao, Luning Sun, and Jian-Xun Wang.  Super-resolution and denoising of fluid flow using physics-informed convolutional neural networks without high-resolution labels.  _Physics of Fluids_ , 33(7):073603, 2021a.
  * Arzani et al. [2021] Amirhossein Arzani, Jian-Xun Wang, and Roshan M. D’Souza.  Uncovering near-wall blood flow from sparse data with physics-informed neural networks.  _Physics of Fluids_ , 33(7):071905, 2021.  10.1063/5.0055600.  URL https://doi.org/10.1063/5.0055600.
  * Sahli Costabal et al. [2020] Francisco Sahli Costabal, Yibo Yang, Paris Perdikaris, Daniel E Hurtado, and Ellen Kuhl.  Physics-informed neural networks for cardiac activation mapping.  _Frontiers in Physics_ , 8:42, 2020.
  * Rao et al. [2021] Chengping Rao, Hao Sun, and Yang Liu.  Physics-informed deep learning for computational elastodynamics without labeled data.  _Journal of Engineering Mechanics_ , 147(8):04021043, 2021.
  * Sutton and Barto [2018] Richard S Sutton and Andrew G Barto.  _Reinforcement learning: An introduction_.  MIT press, 2018.
  * Gao et al. [2021b] Han Gao, Luning Sun, and Jian-Xun Wang.  PhyGeoNet: physics-informed geometry-adaptive convolutional neural networks for solving parameterized steady-state PDEs on irregular domain.  _Journal of Computational Physics_ , 428:110079, 2021b.
  * Heess et al. [2015] Nicolas Heess, Gregory Wayne, David Silver, Timothy Lillicrap, Tom Erez, and Yuval Tassa.  Learning continuous control policies by stochastic value gradients.  In C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Garnett, editors, _Advances in Neural Information Processing Systems_ , volume 28. Curran Associates, Inc., 2015.  URL https://proceedings.neurips.cc/paper/2015/file/148510031349642de5ca0c544f31b2ef-Paper.pdf.
  * Levine and Abbeel [2014] Sergey Levine and Pieter Abbeel.  Learning neural network policies with guided policy search under unknown dynamics.  In _NIPS_ , volume 27, pages 1071–1079. Citeseer, 2014.
  * Luo et al. [2018] Yuping Luo, Huazhe Xu, Yuanzhi Li, Yuandong Tian, Trevor Darrell, and Tengyu Ma.  Algorithmic framework for model-based deep reinforcement learning with theoretical guarantees.  _arXiv preprint arXiv:1807.03858_ , 2018.
  * Amari [1967] Shunichi Amari.  A theory of adaptive pattern classifiers.  _IEEE Transactions on Electronic Computers_ , (3):299–307, 1967.
  * Konečnỳ and Richtárik [2013] Jakub Konečnỳ and Peter Richtárik.  Semi-stochastic gradient descent methods.  _arXiv preprint arXiv:1312.1666_ , 2013.
  * Johnson and Zhang [2013] Rie Johnson and Tong Zhang.  Accelerating stochastic gradient descent using predictive variance reduction.  _Advances in neural information processing systems_ , 26:315–323, 2013.
  * Defazio et al. [2014] Aaron Defazio, Francis Bach, and Simon Lacoste-Julien.  Saga: A fast incremental gradient method with support for non-strongly convex composite objectives.  In _Advances in neural information processing systems_ , pages 1646–1654, 2014.
  * Schmidt et al. [2017] Mark Schmidt, Nicolas Le Roux, and Francis Bach.  Minimizing finite sums with the stochastic average gradient.  _Mathematical Programming_ , 162(1-2):83–112, 2017.
  * Williams [1992] Ronald J Williams.  Simple statistical gradient-following algorithms for connectionist reinforcement learning.  _Machine learning_ , 8(3):229–256, 1992.
  * Watkins and Dayan [1992] Christopher JCH Watkins and Peter Dayan.  Q-learning.  _Machine learning_ , 8(3-4):279–292, 1992.
  * Wang et al. [2016] Ziyu Wang, Tom Schaul, Matteo Hessel, Hado Hasselt, Marc Lanctot, and Nando Freitas.  Dueling network architectures for deep reinforcement learning.  In _International conference on machine learning_ , pages 1995–2003. PMLR, 2016.
  * Cvitanović et al. [2010] Predrag Cvitanović, Ruslan L Davidchack, and Evangelos Siminos.  On the state space geometry of the kuramoto–sivashinsky flow in a periodic domain.  _SIAM Journal on Applied Dynamical Systems_ , 9(1):1–33, 2010.
  * Nagabandi et al. [2018] Anusha Nagabandi, Gregory Kahn, Ronald S Fearing, and Sergey 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)_ , pages 7559–7566. IEEE, 2018.