# A Generative Approach to Control
Complex Physical Systems

Long Wei1 Peiyan Hu2∗§ Ruiqi Feng1∗ Haodong Feng1 Yixuan Du3§ Tao Zhang1
Rui Wang4§ Yue Wang5 Zhi-Ming Ma2 Tailin Wu1†
1School of Engineering, Westlake University,
2Academy of Mathematics and Systems Science, Chinese Academy of Sciences,
3Jilin University, 4Fudan University, 5Microsoft AI4Science
<EMAIL_ADDRESS>
<EMAIL_ADDRESS>Equal contribution. §Work done as an intern at
Westlake University. †Corresponding author.

###### Abstract

Controlling the evolution of complex physical systems is a fundamental task
across science and engineering. Classical techniques suffer from limited
applicability or huge computational costs. On the other hand, recent deep
learning and reinforcement learning-based approaches often struggle to
optimize long-term control sequences under the constraints of system dynamics.
In this work, we introduce Diffusion Physical systems Control (DiffPhyCon), a
new class of method to address the physical systems control problem.
DiffPhyCon excels by simultaneously minimizing both the learned generative
energy function and the predefined control objectives across the entire
trajectory and control sequence. Thus, it can explore globally and identify
near-optimal control sequences. Moreover, we enhance DiffPhyCon with prior
reweighting, enabling the discovery of control sequences that significantly
deviate from the training distribution. We test our method in 1D Burgers’
equation and 2D jellyfish movement control in a fluid environment. Our method
outperforms widely applied classical approaches and state-of-the-art deep
learning and reinforcement learning methods. Notably, DiffPhyCon unveils an
intriguing fast-close-slow-open pattern observed in the jellyfish, aligning
with established findings in the field of fluid dynamics.

##  1 Introduction

Modeling the dynamics of complex physical systems is an important class of
problems in science and engineering. Usually, we are not only interested in
predicting a physical system’s behavior but also injecting time-variant
signals to steer its evolution and optimize specific objectives. This gives
rise to the complex physical control problem, a fundamental task with numerous
applications, including controlled nuclear fusion [9], fluid control [21],
underwater devices [72] and aviation [48], among others.

Despite its importance, controlling complex physical systems efficiently
presents significant challenges. It inherits the fundamental challenge of
simulating complex physical systems, which are typically high-dimensional and
highly nonlinear, as specified in Appendix A.1. Furthermore, observed control
signals and corresponding system trajectories for optimizing a control model
are typically far from the optimal solutions of the specific control
objective. This fact poses a significant challenge of dilemma: _How to explore
long-term control sequences beyond its training distribution to seek near-
optimal solutions while making the resulting system trajectory faithful to the
dynamics of the physical system_?

To tackle physical systems control problems, various techniques have been
proposed, yet they fall short of addressing the above challenges. Regarding
traditional control methods, the Proportional-Integral-Derivative (PID)
control [33], is efficient but only suitable for a limited range of problems.
Conversely, Model Predictive Control (MPC) [59], despite a wider range of
applicability, suffers from high computational costs and challenges in global
optimization. Recent advances in supervised learning (SL) [21, 23] and
reinforcement learning (RL) [14, 47, 53], trained on system trajectories and
control signals data, have demonstrated impressive performance in solving
physical systems control problems. However, existing SL and model-based RL
methods either fall into myopic failure modes [25] that fail to achieve long-
term near-optimal solutions, or produce adversarial state trajectories [74]
that violate the physical system’s dynamics. The main reason may be that they
treat the continuous evolution of dynamics from an iterative view, both
lacking long-term vision and struggling in global optimization. See Appendix
A.2 for more related work on physical system control.

In this work, we introduce Diffusion Physical systems Control (DiffPhyCon), a
_new class_ of method to address the physical systems control problem. We take
an energy optimization perspective over system trajectory and control
sequences across the whole horizon to implicitly capture the constraints
inherent in system dynamics. We accomplish this through diffusion models,
which are trained using system trajectory data and control sequences. In the
inference stage, DiffPhyCon integrates simulation and control optimization
into a unified energy optimization process. This prevents the generated system
dynamics from falling out of distribution, and offers an enhanced perspective
over long-term dynamics, facilitating the discovery of control sequences that
optimize the objectives.

An essential aspect of physical systems control lies in its capacity to
generate near-optimal controls, even when they may deviate significantly from
the training distribution. We address this challenge with the key insight that
the learned generative energy landscape can be _decomposed_ into two
components: a prior distribution representing the control sequence and a
conditional distribution characterizing the system trajectories given the
control sequence. Based on this insight, we develop a _prior reweighting_
technique to subtract the effect of the prior distribution of control
sequences, with adjustable strength, from the overall joint generative energy
landscape during inference.

We demonstrate the effectiveness of DiffPhyCon via extensive experiments on 1D
Burgers’ equation and 2D jellyfish movement control tasks, with the additional
challenge of partial observation and partial control. In particular, the 2D
jellyfish movement control problem is an important source of inspiration for
the design of underwater and aerial devices, and its challenges come from
complex vortice behavior and fluid-solid coupling dynamics [55, 27]. We
generate such a dataset based on a CFD (computational fluid dynamic) software
to mimic the movement of a jellyfish under control signals of its flapping
behavior. To advance research in controlling complex physical systems, we
contribute our dataset as a benchmark. On both tasks, our method outperforms
widely applied classical control methods and is also competitive with recent
supervised learning and strong reinforcement learning baselines. Notably,
DiffPhyCon reveals the intriguing fast-close-slow-open pattern exhibited by
the jellyfish, aligning with findings in the field of fluid dynamics [27].

In summary, we contribute the following: (1) We develop DiffPhyCon, a novel
generative method to control complex physical systems. By optimizing
trajectory and control sequences jointly in the entire horizon by diffusion
models, DiffPhyCon facilitates global optimization of long-term dynamics and
helps to reduce myopic failure modes. (2) We introduce the prior reweighting
technique to generate control sequences that are superior to those in the
training set. (3) We demonstrate the effectiveness of our method on 1D
Burgers’ equation and 2D jellyfish movement control tasks, particularly
demonstrating advantages in scenarios with partial observations/control. (4)
We release a new benchmark of jellyfish movement control to the community of
complex physical system control.

##  2 Background

Figure 1: Overview of DiffPhyCon. The figure depicts the training (top),
inference (bottom left), and evaluation (bottom right) of DiffPhyCon. Orange
and blue colors respectively represent models learning the joint distribution
$p_{\theta}(\mathbf{u},\mathbf{w})$ and the prior distribution
$p_{\phi}(\mathbf{w})$. Through prior reweighting and guidance, DiffPhyCon is
capable of generating superior control sequences.

###  2.1 Problem Setup

We consider the following complex physical system control problem:

$\displaystyle\mathbf{w}^{*}=\operatorname*{argmin}_{\mathbf{w}}\mathcal{J}(\mathbf{u},\mathbf{w})\quad\text{s.t.}\quad\mathcal{C}(\mathbf{u},\mathbf{w})=0.$
(1)

Here
$\mathbf{u}(t,\mathbf{x}):[0,\mathcal{T}]\times\Omega\mapsto\mathbb{R}^{d_{\mathbf{u}}}$
is the system trajectory $\\{\mathbf{u}(t,\cdot),t\in[0,\mathcal{T}]\\}$
defined on time range $[0,\mathcal{T}]\subset\mathbb{R}$ and spatial domain
$\Omega\subset\mathbb{R}^{D}$, and
$\mathbf{w}(t,\mathbf{x}):[0,\mathcal{T}]\times\Omega\mapsto\mathbb{R}^{d_{\mathbf{w}}}$
is the external control signal for the physical system with dimension
${d_{\mathbf{w}}}$. $\mathcal{J}(\mathbf{u},\mathbf{w})$ denotes the control
objective. For example, $\mathcal{J}$ can be designed to measure the control
performance towards a target state $\mathbf{u}^{*}$ with cost constraints:
$\mathcal{J}\coloneqq\int\|\mathbf{u}-\mathbf{u}^{*}\|^{2}\mathrm{d}\mathbf{x}\mathrm{d}t+\int\|\mathbf{w}\|^{2}\mathrm{d}\mathbf{x}\mathrm{d}t$.
$\mathcal{C}(\mathbf{u},\mathbf{w})=0$ denotes the physical constraints.
$\mathcal{C}(\mathbf{u},\mathbf{w})=0$ can be specified either _explicitly_ by
a PDE dynamic which describes how the control signal
$\mathbf{w}(t,\mathbf{x})$ drives the trajectory $\mathbf{u}(t,\mathbf{x})$ to
evolve under boundary and initial conditions, or _implicitly_ by control
sequences and trajectory data collected from observation of the physical
system. In the latter case, we may even have access to only partial trajectory
data or engage in partial control. These situations collectively pose
significant challenges to the physical system control task.

###  2.2 Preliminary: Diffusion Models

Diffusion models [18] are a class of generative models that learn data
distribution from data. Diffusion models consist of two opposite processes:
the forward process
$q(\mathbf{x}_{k+1}|\mathbf{x}_{k})=\mathcal{N}(\mathbf{x}_{k+1};\sqrt{\alpha_{k}}\mathbf{x}_{k},(1-\alpha_{k})\mathbf{I})$
to corrupt a clean data $\mathbf{x}_{0}$ to a Gaussian noise
$\mathbf{x}_{K}\sim\mathcal{N}(\mathbf{0},\mathbf{I})$, and the reverse
parametrized process
$p_{\theta}(\mathbf{x}_{k-1}|\mathbf{x}_{k})=\mathcal{N}(\mathbf{x}_{k-1};\mu_{\theta}(\mathbf{x}_{k},k),\sigma_{k}\mathbf{I})$
to denoise from standard Gaussian
$\mathbf{x}_{K}\sim\mathcal{N}(\mathbf{0},\mathbf{I})$, where
$\\{\alpha_{k}\\}_{k=1}^{K}$ is the variance schedule. To train diffusion
models, [18] propose the DDPM method to minimize the following training loss
for the denoising network $\bm{\epsilon}_{\theta}$, a simplification of the
evidence lower bound (ELBO) for the log-likelihood of the data:

$\displaystyle\mathcal{L}=\mathbb{E}_{k\sim U(1,K),\mathbf{x}_{0}\sim
p(x),\bm{\epsilon}\sim\mathcal{N}(\mathbf{0},\mathbf{I})}[\|\bm{\epsilon}-\bm{\epsilon}_{\theta}(\sqrt{\bar{\alpha}_{k}}\mathbf{x}_{0}+\sqrt{1-\bar{\alpha}_{k}}\bm{\epsilon},k)\|_{2}^{2}],$
(2)

where $\bar{\alpha}_{k}:=\prod_{i=1}^{k}\alpha_{i}$. $\bm{\epsilon}_{\theta}$
estimates the noise to be removed to recover data $\mathbf{x}_{0}$. During
inference, iterative application of $\bm{\epsilon}_{\theta}$ from a Gaussian
noise could generate a new sample $\mathbf{x}_{0}$ that approximately follows
the data distribution $p(\mathbf{x})$. See Appendix A.3 for related work on
diffusion models.

Notation. We use $\mathbf{v}_{[n,m]}=[\mathbf{v}_{n},\cdots,\mathbf{v}_{m}]$
to denote a sequence of variables. We use $\mathbf{z}_{[0,T-1],k}$ to denote
the hidden variable of $\mathbf{z}_{[0,T-1]}$ in a diffusion step $k$. For
simplicity, we abbreviate $\mathbf{w}_{[0,T-1]}$, $\mathbf{u}_{[1,T]}$ as
$\mathbf{w}$, $\mathbf{u}$. Concatenation of two variables are denoted via
_e.g._ , $[\mathbf{u},\mathbf{w}]$.

##  3 Method

In this section, we detail our method DiffPhyCon. In Section 3.1, we introduce
our method including its training and inference. In Section 3.2, we further
propose a prior reweighting technique to improve DiffPhyCon. The overview of
DiffPhyCon is illustrated in Figure 1.

###  3.1 Generative Control by Diffusion Models

DiffPhyCon takes an energy optimization perspective to solve the problem Eq.
(1), where PDE constraints can be modeled as a parameterized energy-based
model (EBM) $E_{\theta}(\mathbf{u},\mathbf{w},\mathbf{c})$ which characterizes
the distribution $p(\mathbf{u},\mathbf{w}|\mathbf{c})$ of $\mathbf{u}$ and
$\mathbf{w}$ conditioned on conditions $\mathbf{c}$ by the correspondence
$p(\mathbf{u},\mathbf{w}|\mathbf{c})\propto\exp({-E_{\theta}(\mathbf{u},\mathbf{w},\mathbf{c})})$.
Lower $E_{\theta}(\mathbf{u},\mathbf{w},\mathbf{c})$, or equivalently higher
$p(\mathbf{u},\mathbf{w}|\mathbf{c})$, means better satisfaction of the PDE
constraints. Then the problem Eq. (1) can be converted to:

$\displaystyle\mathbf{u}^{*},\mathbf{w}^{*}=\operatorname*{argmin}_{\mathbf{u},\mathbf{w}}\left[E_{\theta}(\mathbf{u},\mathbf{w},\mathbf{c})+\lambda\cdot\mathcal{J}(\mathbf{u},\mathbf{w})\right],$
(3)

where $\lambda$ is a hyperparameter. This formulation optimizes $\mathbf{u}$
and $\mathbf{w}$ of all physical time steps simultaneously. The first term
encourages the generated $\mathbf{w}$ and $\mathbf{u}$ to satisfy the PDE
constraints $\mathcal{C}(\mathbf{u},\mathbf{w})=0$. The second term guides
optimization towards optimal objectives. The advantage of this optimization
framework is that it tasks a global optimization on $\mathbf{u}$ and
$\mathbf{w}$ of all times steps, which may obtain better solutions that are
faithful to the dynamics of the physical system. Details about the effect of
the hyperparameter $\lambda$ are provided in Appendix J.

Training. To train $E_{\theta}$, we exploit the diffusion model to estimate
the gradient of $E_{\theta}$. For convenience, we introduce a new variable
$\mathbf{z}$ to represent the concatenation of $\mathbf{u}$ and $\mathbf{w}$
as $\mathbf{z}=[\mathbf{u},\mathbf{w}]$. We use a denoising network
$\bm{\epsilon}_{\theta}$ [18], which approximates $\nabla
E_{\theta}(\mathbf{z},\mathbf{c})$ [12], to learn the noise that should be
denoised in each diffusion step $k=1,\cdots,K$. The training loss of
$\bm{\epsilon}_{\theta}$ is:

$\mathcal{L}=\mathbb{E}_{k\sim U(1,K),(\mathbf{z},\mathbf{c})\sim
p(\mathbf{z},\mathbf{c}),\bm{\epsilon}\sim\mathcal{N}(\mathbf{0},\mathbf{I})}[\|\mathbf{\bm{\epsilon}}-\bm{\epsilon}_{\theta}(\sqrt{\bar{\alpha}_{k}}\mathbf{z}+\sqrt{1-\bar{\alpha}_{k}}\bm{\epsilon},\mathbf{c},k)\|_{2}^{2}],\
\ \ $ (4)

where $\bm{\epsilon}_{\theta}$ is conditioned on $\mathbf{c}$. Regarding
training datasets, they could be generated by designing the conditions
$\mathbf{c}$ and control sequences followed by a simulation when an explicit
PDE form is available. Otherwise, they should be collected from observed pairs
of control sequences and trajectories.

Control optimization. After the denoising network is trained, the Eq. (3) can
be optimized by the Langevin sampling procedure as follows. We start from an
initial sample $\mathbf{z}_{K}\sim\mathcal{N}(\mathbf{0},\mathbf{I})$, and
iteratively run the following process

$\begin{gathered}\mathbf{z}_{k-1}=\mathbf{z}_{k}-\eta\left(\nabla_{\mathbf{z}}(E_{\theta}(\mathbf{z}_{k},\mathbf{c})+\lambda\mathcal{J}(\hat{\mathbf{z}}_{k})\right)+\mathbf{\xi},\quad\mathbf{\xi}\sim\mathcal{N}\bigl{(}\mathbf{0},\sigma^{2}_{k}\mathbf{I}\bigl{)},\end{gathered}$
(5)

where $\sigma^{2}_{k}$ and $\eta$ correspond to noise schedules and scaling
factors used in the diffusion process, respectively. Here
$\hat{\mathbf{z}}_{k}$ is the approximate noise-free $\mathbf{z}_{0}$
estimated from $\mathbf{z}_{k}$ by:

$\hat{\mathbf{z}}_{k}=(\mathbf{z}_{k}-\sqrt{1-\bar{\alpha}_{k}}\bm{\epsilon}_{\theta}(\mathbf{z}_{k},\mathbf{c},k))/\sqrt{\bar{\alpha}_{k}},$
(6)

We calculate $\mathcal{J}$ in (5) based on $\hat{\mathbf{z}}_{k}$ instead of
directly using $\mathbf{z}_{k}$ because otherwise noise in $\mathbf{z}_{k}$
could bring errors to $\mathcal{J}$. Then $\nabla E_{\theta}$ can be replaced
by our trained denoising network $\bm{\epsilon}_{\theta}$ as follows:

$\mathbf{z}_{k-1}=\mathbf{z}_{k}-\eta\left(\bm{\epsilon}_{\theta}(\mathbf{z}_{k},\mathbf{c},k)+\lambda\nabla_{\mathbf{z}}\mathcal{J}(\hat{\mathbf{z}}_{k})\right)+\xi,\quad\mathbf{\xi}\sim\mathcal{N}\bigl{(}\mathbf{0},\sigma^{2}_{k}\mathbf{I}\bigl{)}$
(7)

Iteration of this denoising process for $k=K,K-1,...,1$ yields a final
solution $\mathbf{z}_{0}=\\{\mathbf{u}_{[1,T],0},\mathbf{w}_{[0,T-1],0}\\}$
for the optimization problem Eq. (3).

Figure 2: Intuition of Prior Reweighting. The top surface illustrates the
landscape of $\mathcal{J}(\mathbf{u},\mathbf{w})$, where the high-dimensional
variables $\mathbf{u}$ and $\mathbf{w}$ are represented using one dimension.
The middle and lower planes depict probability heatmaps for the reweighted
distribution $p^{\gamma}(\mathbf{w})p(\mathbf{u}|\mathbf{w})/Z$. Adjusting
$\gamma$ from $\gamma=1$ (middle plane) to $0<\gamma<1$ (lower plane), a
better minimal of $\mathcal{J}$ (red dot in the lower plane) gains the chance
to be sampled. This contrasts with the suboptimal red point in the middle
plane highly influenced by the prior $p(\mathbf{w})$.

Guidance conditioning. In addition to the above introduced explicit guidance,
conditioning is also widely used to guide sampling in diffusion models [19,
60]. When the control objective can be naturally expressed in a conditioning
form, e.g., the generated trajectory $\mathbf{u}$ is required to coincide with
a desired target $\mathbf{u}^{*}$, we can include $\mathbf{u}=\mathbf{u}^{*}$
as a condition in $\mathbf{c}$ such that the sampled trajectory $\mathbf{u}$
from diffusion models automatically satisfying $\mathbf{u}=\mathbf{u}^{*}$.
Overall, in our proposed DiffPhyCon framework, the control objective
$\mathcal{J}$ can be optimized either using the explicit guidance
$\nabla\mathcal{J}$ or guidance conditioning depending on the specific control
objectives.

###  3.2 Prior Reweighting

Motivation. In physical systems control, a critical challenge lies in
obtaining control sequences superior to those in training datasets, which
often deviate significantly from achieving the optimal control objective.
Although guidance of the control objective is incorporated in our diffusion
model, generated control sequences are still highly influenced by the prior
distribution of control sequences in training datasets. This inspires us to
explore strategies to mitigate the effect of this prior, aiming to generate
near-optimal control sequences.

We address this challenge with the key insight that the energy-based model
$E(\mathbf{u},\mathbf{w},\mathbf{c})$ can be decomposed into two components:
one is $E^{(p)}(\mathbf{w},\mathbf{c})$ derived from the prior distribution
$p(\mathbf{w}|\mathbf{c})$ of control sequences, and the other
$E^{(c)}(\mathbf{u},\mathbf{w},\mathbf{c})$ representing the conditional
probability distribution $p(\mathbf{u}|\mathbf{w},\mathbf{c})$ of trajectories
with respect to given control sequences. This decomposition has the following
form

$E(\mathbf{u},\mathbf{w},\mathbf{c})=E^{(p)}(\mathbf{u},\mathbf{c})+E^{(c)}(\mathbf{u},\mathbf{w},\mathbf{c}),$
(8)

by the corresponding decomposition of distribution
$p(\mathbf{u},\mathbf{w}|\mathbf{c})=p(\mathbf{w}|\mathbf{c})p(\mathbf{u}|\mathbf{w},\mathbf{c})$.
We propose a prior reweighting technique, which introduces an adjustable
hyperparameter $\gamma>0$ as an exponential of $p(\mathbf{w}|\mathbf{c})$,
allowing for the tuning of the influence of this prior distribution. Then we
have a reweighted version of $p(\mathbf{u},\mathbf{w}|\mathbf{c})$ as
$p_{\gamma}(\mathbf{u},\mathbf{w}|\mathbf{c})=p(\mathbf{w}|\mathbf{c})^{\gamma}p(\mathbf{u}|\mathbf{w},\mathbf{c})/Z$,
which is also a probability distribution and can be further transformed to

$p_{\gamma}(\mathbf{u},\mathbf{w}|\mathbf{c})=p(\mathbf{w}|\mathbf{c})^{\gamma-1}p(\mathbf{u},\mathbf{w}|\mathbf{c})/Z,$
(9)

where $Z$ is a normalization constant. In particular, when $0<\gamma<1$, this
approach is advantageous as it flattens the distribution
$p(\mathbf{u},\mathbf{w}|\mathbf{c})$, thereby increasing the likelihood of
sampling from low probability region of $p(\mathbf{u},\mathbf{w}|\mathbf{c})$,
where the optimal solutions of the problem Eq. (3) probably lie.

Integrating Eq. (9) into Eq. (8), we have

$E^{(\gamma)}(\mathbf{u},\mathbf{w},\mathbf{c})=(\gamma-1)E^{(p)}(\mathbf{w},\mathbf{c})+E_{\theta}(\mathbf{u},\mathbf{w},\mathbf{c})-\log{Z},$
(10)

where
$E^{(\gamma)}(\mathbf{u},\mathbf{w},\mathbf{c})=-\log{(p_{\gamma}(\mathbf{u},\mathbf{w}|\mathbf{c}))}+{const}$
is the reweighted energy-based model associated with
$E_{\theta}(\mathbf{u},\mathbf{w},\mathbf{c})$ in Eq. (3), relying on the
hyperparameter $\gamma$. Then the optimization problem Eq. (3) can be
transformed to

$\displaystyle\mathbf{u}^{*},\mathbf{w}^{*}=\operatorname*{argmin}_{\mathbf{u},\mathbf{w}}\left[E^{(\gamma)}(\mathbf{u},\mathbf{w},\mathbf{c})+\lambda\cdot\mathcal{J}(\mathbf{u},\mathbf{w})\right].$
(11)

Optimization of this problem encourages sampling from the low likelihood
region of $p(\mathbf{w}|\mathbf{c})$ while minimizing the control objective,
which possesses the capability to generate control sequences that are more
likely to be near-optimal than its degenerate version $\gamma=1$ in the
original optimization problem Eq. (3). The intuition of prior reweighting is
illustrated in Figure 2.

Training. To learn the reweighted energy
$E^{(\gamma)}(\mathbf{u},\mathbf{w},\mathbf{c})$, we parameterize its gradient
as a summation of two parts by taking the gradient of both sides of Eq. (10):

$\nabla
E_{\phi,\theta}^{(\gamma)}(\mathbf{u},\mathbf{w},\mathbf{c})=(\gamma-1)\nabla
E^{(p)}_{\phi}(\mathbf{w},\mathbf{c})+\nabla
E_{\theta}(\mathbf{u},\mathbf{w},\mathbf{c}),$ (12)

where $E^{(p)}_{\phi}(\mathbf{w},\mathbf{c})$ parameterizes the energy based
model $E^{(p)}(\mathbf{w},\mathbf{c})$ corresponding to
$p(\mathbf{w}|\mathbf{c})$. Note that $\nabla\log{Z}$ vanishes here because it
is a constant. Notice that $\nabla
E_{\theta}(\mathbf{u},\mathbf{w},\mathbf{c})$ has already been trained by Eq.
(4). $\nabla E^{(p)}_{\phi}(\mathbf{w},\mathbf{c})$ can be trained similarly
by the following loss function

$\mathcal{L}=\mathbb{E}_{k\sim U(1,K),(\mathbf{w},\mathbf{c})\sim
p(\mathbf{w},\mathbf{c}),\bm{\epsilon}\sim\mathcal{N}(\mathbf{0},\mathbf{I})}[\|\mathbf{\bm{\epsilon}}-\bm{\epsilon}_{\phi}(\sqrt{\bar{\alpha}_{t}}\mathbf{w}+\sqrt{1-\bar{\alpha}_{k}}\bm{\epsilon},\mathbf{c},k)\|_{2}^{2}],\
\ \ $ (13)

where $\bm{\epsilon}_{\phi}$ is the conditional denoising network that
approximates $\nabla E^{(p)}_{\phi}(\mathbf{w},\mathbf{c})$.

Control optimization. With both $\bm{\epsilon}_{\theta}$ and
$\bm{\epsilon}_{\phi}$ trained, Eq. (11) can be optimized by running:

$\displaystyle\mathbf{z}_{k-1}$
$\displaystyle=\mathbf{z}_{k}-\eta(\bm{\epsilon}_{\theta}(\mathbf{z}_{k},\mathbf{c},k)+\lambda\nabla_{\mathbf{z}}\mathcal{J}(\hat{\mathbf{z}}_{k}))+\xi_{1},\quad\xi_{1}\sim\mathcal{N}\bigl{(}0,\sigma^{2}_{k}\mathbf{I}\bigl{)}$
(14) $\displaystyle\mathbf{w}_{k-1}$
$\displaystyle=\mathbf{w}_{k-1}-\eta(\gamma-1)\bm{\epsilon}_{\phi}(\mathbf{w}_{k},\mathbf{c},k)+\xi_{2},\quad\xi_{2}\sim\mathcal{N}\bigl{(}0,\sigma^{2}_{k}\mathbf{I}\bigl{)},$
(15)

iteratively, where $\mathbf{z}_{k}=[\mathbf{u}_{k},\mathbf{w}_{k}]$. The
difference between this iteration scheme and Eq. (7) is that it uses an
additional step to update $\mathbf{w}_{k}$ based on the predicted noise of
$\bm{\epsilon}_{\phi}$. This guides
$\mathbf{z}_{k}=[\mathbf{u}_{k},\mathbf{w}_{k}]$ to move towards the
reweighted distribution $p_{\gamma}(\mathbf{u},\mathbf{w}|\mathbf{c})$ while
aligning with the direction to decrease the objective by its guidance in the
iteration of $\mathbf{z}_{k}$. The complete algorithm is present in Algorithm
1. Detailed discussion and results about how to set the hyperparameter
$\gamma$ are provided in Appendix J.

DiffPhyCon-lite. The introduction of the prior reweighting technique in
DiffPhyCon involves training and evaluation of two models, thus bringing in
additional computational cost. It is gratifying to note that we can balance
the control performance and computational overhead of DiffPhyCon by adjusting
the parameter $\gamma$. When $\gamma=1$, the model $\bm{\epsilon}_{\phi}$ is
not needed, and we denote this simplified version of DiffPhyCon as DiffPhyCon-
lite.

Algorithm 1 Inference for DiffPhyCon

1: Require Diffusion models
$\bm{\epsilon}_{\theta}(\mathbf{z}_{k},\mathbf{c},k)$ and
$\bm{\epsilon}_{\phi}(\mathbf{w}_{k},\mathbf{c},k)$, control objective
$\mathcal{J}(\cdot)$, covariance matrix $\sigma_{k}^{2}I$, control conditions
$\mathbf{c}$, schedule $\bar{\alpha}_{k}$, hyperparameters $\lambda,\gamma,K$

2: Initialize optimization variables
$\mathbf{z}_{K}=[\mathbf{u}_{K},\mathbf{w}_{K}]\sim\mathcal{N}(\bm{0},\mathbf{I})$

3: for $k=K,\ldots,1$ do

4:
$\hat{\mathbf{z}}_{k}=(\mathbf{z}_{k}-\sqrt{1-\bar{\alpha}_{k}}\bm{\epsilon}_{\theta}(\mathbf{z}_{k},\mathbf{c},k))/{\sqrt{\bar{\alpha}_{k}}}$

5:
$\mathbf{z}_{k-1}=\mathbf{z}_{k}-\eta(\bm{\epsilon}_{\theta}(\mathbf{z}_{k},\mathbf{c},k)+\lambda\nabla_{\mathbf{z}}\mathcal{J}(\hat{\mathbf{z}}_{k}))+\xi_{1},\xi_{1}\sim\mathcal{N}\bigl{(}0,\sigma^{2}_{k}\mathbf{I}\bigl{)}$
// transition to next diffusion step

6:
$\mathbf{w}_{k-1}=\mathbf{w}_{k-1}-\eta(\gamma-1)\bm{\epsilon}_{\phi}(\mathbf{w}_{k},\mathbf{c},k)+\xi_{2},\xi_{2}\sim\mathcal{N}\bigl{(}0,\sigma^{2}_{k}\mathbf{I}\bigl{)}$
// prior reweighting

7: end for

8: return $\mathbf{u}^{*},\mathbf{w}^{*}=\mathbf{z}_{0}$

##  4 Experiments

In this section, we aim to answer the following questions: (1) Can DiffPhyCon
present superiority over traditional, supervised learning, and reinforcement
learning methods for physical systems control? (2) Does the proposed prior
reweighting technique help achieve better control objectives? (3) Could
answers to (1) and (2) be generalized to more challenging partial observation
or partial control scenarios? To answer these questions, we conduct
experiments on the 1D Burgers’ equation and 2D jellyfish movement control
problem, both of which are vital and challenging.

The following state-of-the-art control methods are selected as baselines. For
1D Burgers’ equation, we use (1) the classical and widely used control
algorithm Proportional-Integral-Derivative (PID) [33] interacting with our
trained surrogate model of the solver; (2) Supervised Learning method (SL)
[23]; RL methods including (3) Soft Actor-Critic (SAC) [15] with offline and
surrogate-solver versions; (4) Behaviour Cloning (BC) [50]; and (5) Behavior
Proximal Policy Optimization (BPPO) [75]. Specifically, the surrogate-solver
version of SAC interacts with our trained surrogate model of the solver, while
the offline version only uses given data. BC and BPPO are also in offline
version. For 2D jellyfish movement control, baselines include SL, SAC
(offline), SAC (surrogate-solver), BC, BPPO, and an additional classical
multi-input multi-output algorithm Model Predictive Control (MPC) [59]. PID is
inapplicable to this data-driven task. Detailed descriptions of baselines are
provided in Appendix F and Appendix G. We provide the code to reproduce the
result for our method and baselines here.

###  4.1 1D Burgers’ Equation Control

Experiment settings. The Burgers’ equation is a governing law occurring in
various physical systems. We consider the 1D Burgers’ equation with the
Dirichlet boundary condition and external force $w(t,x)$, which is also
studied in [23, 43].

$\displaystyle\begin{cases}\frac{\partial u}{\partial t}=-u\cdot\frac{\partial
u}{\partial x}+\nu\frac{\partial^{2}u}{\partial x^{2}}+w(t,x)&\text{in
}[0,T]\times\Omega\\\
u(t,x)=0\quad\quad\quad\quad\quad\quad\quad\quad&\text{on
}[0,T]\times\partial\Omega\\\
u(0,x)=u_{0}(x)\quad\quad\quad\quad\quad\quad&\text{in
}\\{t=0\\}\times\Omega.\end{cases}$ (16)

Here $\nu$ is the viscosity parameter, and $u_{0}(\mathbf{x})$ is the initial
condition. Subject to Eq. (16), given a target state $u_{d}(x)$, the objective
of control is to minimize the control error $\mathcal{J}_{\text{actual}}$
between $u_{T}$ and $u_{d}$, while constraining the energy cost
$\mathcal{J}_{\text{energy}}$ of the control sequence $w(t,x)$:

$\mathcal{J}_{\text{actual}}\coloneqq\int_{\Omega}|u(T,x)-u_{d}(x)|^{2}\mathrm{d}x,\
\
\mathcal{J}_{\text{energy}}\coloneqq\int_{[0,T]\times\Omega}|w(t,x)|^{2}\mathrm{d}t\mathrm{d}x$
(17)

To make the evaluation challenging, we select three different experiment
settings that correspond to different real-life scenarios: partial
observation, full control (PO-FC), full observation, partial control (FO-PC),
and partial observation, partial control (PO-PC), which are elaborated on in
Appendix C.2 and illustrated in Appendix B. These settings are challenging for
classical control methods such as PID since they require capturing the long-
range dependencies in the system dynamics. Note that the reported metrics in
different settings are not directly comparable.

In this experiment, DiffPhyCon uses the guidance conditioning (Section 2) to
optimize $\mathcal{J}_{\text{actual}}$ and explicit guidance to optimize the
energy cost, which is elaborated in Appendix C.

Table 1: Best $\mathcal{J}_{\text{actual}}$ achieved in 1D Burgers’s equation control. Bold font denotes the best model, and underline denotes the second best model.  | PO-FC | FO-PC | PO-PC
---|---|---|---
PID (surrogate-solver) | - | 0.09115 | 0.09631
SL | 0.09752 | 0.00078 | 0.02328
SAC (surrogate-solver) | 0.01577 | 0.03426 | 0.02149
SAC (offline) | 0.03201 | 0.04333 | 0.03328
BC | 0.02836 | 0.00856 | 0.00952
BPPO | 0.02771 | 0.00852 | 0.00891
DiffPhyCon-lite (ours) | 0.01139 | 0.00037 | 0.00494
DiffPhyCon (ours) | 0.01103 | 0.00037 | 0.00494
Figure 3: Pareto frontier of $\mathcal{J}_{\text{energy}}$ vs.
$\mathcal{J}_{\text{actual}}$ of different methods for 1D Burgers’ equation.

Results. In Table 1, we report results of the control error
$\mathcal{J}_{\text{actual}}$ of different methods. It can be observed that
DiffPhyCon delivers the best results compared to all baselines. Specifically,
DiffPhyCon decreases $\mathcal{J}_{\text{actual}}$ of the best baseline by
30.1%, 52.6%, and 44.6% in the PO-FC, FO-PC, and PO-PC settings respectively.
From Table 1, DiffPhyCon and DiffPhyCon-lite show little performance gap. This
is because, the prior distribution of $w$ that our DiffPhyCon-lite learned is
conditioned on both $u_{0}$ and $u_{T}$, which fully determines the optimal
$w$. Therefore, $p(w|u_{0},u_{T})$ is intrinsically the optimal distribution
and thus DiffPhyCon-lite can already deliver satisfactory performance.

To compare the ability of different methods to optimize
$\mathcal{J}_{\text{actual}}$ with constrained energy cost
$\mathcal{J}_{\text{energy}}$, we compare the Pareto frontiers of different
methods in Figure 3. We vary the hyperparameter $\lambda$ to control the
tradeoff between $\mathcal{J}_{\text{actual}}$ and the energy cost since most
baselines have this hyperparameter. As can be observed in Figure 3, the Pareto
frontiers of DiffPhyCon are consistently among the best, achieving the
_lowest_ $\mathcal{J}_{\text{actual}}$ for most settings of the energy budget.
Although SL performs well in full observation setting (b) where the system
dynamics can be more easily predicted, it encounters difficulty in partial
observation scenarios (a)(c). The results demonstrate DiffPhyCon’s ability to
generate near-optimal control sequences compared to baselines. More
visualization results are provided in Appendix B. For efficiency evaluation,
please refer to Table 19 in Appendix I.

###  4.2 2D Jellyfish Movement Control

Experiment settings. This task is to control the movement of a flapping
jellyfish with two wings in a 2D fluid field where fluid flows at a constant
speed. The jellyfish is propelled by the fluid when its wings flap. Its moving
speed and efficiency are determined by the mode of flapping. This task
represents a common phenomenon in aquatic locomotion, which involves complex
fluid-solid interactions and is an important source of inspiration for the
design of underwater and aerial devices [27]. For this task, fluid dynamics
follows the 2D incompressible Navier-Stokes Equation:

$\displaystyle\begin{cases}\frac{\partial\mathbf{v}}{\partial
t}+\mathbf{v}\cdot\nabla\mathbf{v}-\nu\nabla^{2}\mathbf{v}+\nabla p=0\\\
\nabla\cdot\mathbf{v}=0\\\
\mathbf{v}(0,\mathbf{x})=\mathbf{v}_{0}(\mathbf{x}),\end{cases}$ (18)

where $\mathbf{v}$ represents the 2D velocity of the fluid, and $p$ represents
the pressure, constituting the PDE state $\mathbf{u}=(\mathbf{v},p)$. The
initial velocity condition is $\mathbf{v}_{0}(\mathbf{x})$ and the kinematic
viscosity is $\nu$. We assume that each wing is rigid, so the jellyfish’s
boundary can be parameterized by the opening angle $\mathbf{w}_{t}$ of wings.
Long-term movement of jellyfish usually presents a periodic flapping mode.
Consequently, the control objective is to maximize its average moving speed
$\bar{v}$ determined by the pressure of the fluid, under the energy cost
constraint $R(\mathbf{w})$ and the periodic constraint
$d(\mathbf{w}_{T},\mathbf{w}_{0})$ of the movement:

$\mathcal{J}=-\bar{v}+\zeta\cdot
R(\mathbf{w})+d(\mathbf{w}_{T},\mathbf{w}_{0}),$ (19)

subject to Eq. (18) and the boundary condition that the velocity of fluid
vanishes near the boundary. The hyperparameter $\zeta$ is set to be 1000. We
evaluate in two settings: full observation, where the full state
$\mathbf{u}=(\mathbf{v},p)$ is observed; and partial observation, where only
pressure is observed. This task is very challenging due to the complicated
dynamics of fluid-solid interactions and vortices behaviour [55, 27],
especially in the scenario of partial observation where the missing of
$\mathbf{v}$ in Eq. (18) restricts the information available to generate well-
informed control signals. Details of the experiment are provided in Appendix
E.

Open-source Dataset Description. We use the Lily-Pad simulator [68] to
generate a dataset describing the movement of jellyfish under the control of
its flapping behavior. Statistics about the dataset are listed in Table 2. The
feature of this dataset is that it contains a rich variation of vortices
behavior and fluid-solid interaction dynamics, determined by violent changes
in opening angles of wings and open-close phase ratio. It would serve as an
important benchmark for studying complex physical system control problems.
This dataset is released via an anonymous link. Details about the dataset are
provided in Appendix D.

Table 2: Jellyfish movement dataset outline. training trajectories | test trajectories | resolution | #fluid features | trajectory length
---|---|---|---|---
30000 | 1000 | 128$\times$128 | 3 | 40
Table 3: 2D jellyfish movement control results. Bold font denotes the best model, and underline denotes the second best model. | Full observation | Partial observation
---|---|---
|  $\bar{v}$ $\uparrow$ |  $R(\mathbf{w})$ $\downarrow$ |  $\mathcal{J}$ $\downarrow$ |  $\bar{v}$ $\uparrow$ |  $R(\mathbf{w})$ $\downarrow$ |  $\mathcal{J}$ $\downarrow$
MPC | 25.72 | 0.0112 | 109.17 | -150.51 | 0.1791 | 329.59
SL | -76.94 | 0.1286 | 205.57 | -102.98 | 0.1188 | 221.79
SAC (surrogate-solver) | -166.96 | 0.0069 | 18.14 | -153.09 | 0.0057 | 158.82
SAC (offline) | -158.66 | 0.0069 | 165.58 | -206.21 | 0.0058 | 211.96
BC | 30.48 | 0.0629 | 32.44 | 20.08 | 0.0556 | 35.48
BPPO | 107.67 | 0.0867 | -20.93 | 54.83 | 0.0518 | -3.02
DiffPhyCon-lite (ours) | 95.04 | 0.0746 | -20.47 | 2.92 | 0.0779 | 74.97
DiffPhyCon (ours) | 279.87 | 0.2058 | -74.11 | 150.21 | 0.1269 | -23.32

Results. Evaluation results are presented in Table 3. We see that our method
outperforms the baselines by a large margin in optimizing the control
objective $\mathcal{J}$. At a cost of slightly increasing the control cost
$R(\mathbf{w})$, control sequences generated by our method achieve a much
faster average speed than baselines. In the full observation setting, the
control objective achieved by DiffPhyCon is -53.18 lower than the best
baseline BPPO, while the average speed exhibits an increment of 172.2 over it.
This demonstrates that diffusion models are effective for this challenging
control task by performing global optimization of trajectory and control
sequences in a generative approach. Comparison between DiffPhyCon-lite and
DiffPhyCon reveals that by flattening the prior distribution of control
sequences, another significant improvement is further achieved. Configuration
of the hyperparameter $\gamma$ in DiffPhyCon and performance with respect to
varying $\gamma$ is presented in Figure 14 in Appendix J. Even in the more
challenging partial observation setting, DiffPhyCon still exhibits substantial
advantages over existing methods. This reflects our method has a strong
control capability under inadequate information.

Figure 4:  Comparison of generated control curves of three test jellyfish. The
resulting control objective $\mathcal{J}$ for each curve is presented. Figure
5: Movement and fluid field visualization on the jellyfish controlled by
DiffPhyCon as in the middle subfigure of Figure 4.

Figure 4 visualizes generated opening angle curves of different methods on
three test jellyfish. Opening angle curves of DiffPhyCon-lite show an obvious
fast-close-slow-open shape, which is proven to produce high speed in jellyfish
movement [27]. The reason is that the fast closing of wings leads to a thrust
from fluid in the early state, resulting in a long-term high speed, and
followed by a slow opening to reduce resistance. While this mode of movement
appears rarely in the training dataset, DiffPhyCon-lite could generate such
control sequences for most test samples. This reflects that diffusion models
under guidance are effective in optimizing the control objective. Furthermore,
DiffPhyCon makes this mode of movement more aggressive, with sharper change of
the opening angle in the beginning and end stage within a period. This
provides strong evidence that reweighting the prior distribution of control
sequences allows for flexible sampling over a flattened distribution, thus
control sequences with low prior but good objectives are more likely to be
sampled. The movement and the resulting fluid field of the jellyfish
corresponding to the middle subfigure of Figure 4 controlled by DiffPhyCon is
illustrated in Figure 5 and more examples are provided in Figure 13 in
Appendix H.6. Conversely, opening angles obtained by baselines are inferior.
The reason may be that they predict opening angles sequentially and hard to
globally optimize the objective with three conflicting terms: average speed,
$R(\mathbf{w})$, and $d(\mathbf{w}_{T},\mathbf{w}_{0})$. We further study such
myopic failure mode of SAC in Appendix H.7. More comparisons about the
variation of the weight $\zeta$ are listed in Table 16 and Table 17 in
Appendix H.6. For efficiency evaluation, please refer to Table 20 in Appendix
I.

We also extend the experiment setting to a high-dimensional control signal
setting, where the wings of the jellyfish are assumed to be soft. We find that
our method is still competitive with baselines. This demonstrates the good
scalability of our method when facing higher-dimensional physical systems.
Details about this evaluation are provided in Appendix K.

##  5 Limitation

There are still several limitations of our DiffPhyCon that provide exciting
opportunity for future works. Firstly, our approach is data-driven, so it is
not guaranteed to achieve optimal solutions and also lacks estimation of the
error gap between the generated control sequence and the optimal ones.
Secondly, the training of DiffPhyCon is currently conducted in an offline
fashion, without interaction with a ground-truth solver. Incorporating solvers
into the training framework could facilitate real-time feedback, enabling the
model to adapt dynamically to the environment and discover novel strategies
and solutions. Furthermore, our proposed DiffPhyCon presently operates in an
open-loop manner, as it does not consider real-time feedback from solvers.
Integrating such feedback would empower the algorithm to adjust its control
decisions based on the evolving state of the environment.

##  6 Conclusion

In this work, we have introduced DiffPhyCon, a novel methodology for
controlling complex physical systems. It generates control sequences and state
trajectories by jointly optimizing the generative energy and control
objective. We further introduced prior reweighting to enable the discovery of
control sequences that diverge significantly from training. Through
comprehensive experiments, we demonstrated our method’s superior performance
compared to classical, deep learning, and reinforcement learning baselines in
challenging physical systems control tasks. We discuss future work in Appendix
L and state social impact in M.

## References

  * [1] Anurag Ajay, Yilun Du, Abhi Gupta, Joshua Tenenbaum, Tommi Jaakkola, and Pulkit Agrawal.  Is conditional generative modeling all you need for decision-making?  arXiv preprint arXiv:2211.15657, 2022.
  * [2] Fan Bao, Min Zhao, Zhongkai Hao, Peiyao Li, Chongxuan Li, and Jun Zhu.  Equivariant energy-guided SDE for inverse molecular design.  In The Eleventh International Conference on Learning Representations, 2023.
  * [3] Gerben Beintema, Alessandro Corbetta, Luca Biferale, and Federico Toschi.  Controlling rayleigh–bénard convection via reinforcement learning.  Journal of Turbulence, 21(9-10):585–605, 2020.
  * [4] Johannes Brandstetter, Rianne van den Berg, Max Welling, and Jayesh K Gupta.  Clifford neural layers for PDE modeling.  In The Eleventh International Conference on Learning Representations, 2023.
  * [5] Johannes Brandstetter, Daniel E. Worrall, and Max Welling.  Message passing neural PDE solvers.  In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022.
  * [6] Salva Rühling Cachay, Bo Zhao, Hailey James, and Rose Yu.  Dyffusion: A dynamics-informed diffusion model for spatiotemporal forecasting.  arXiv preprint arXiv:2306.01984, 2023.
  * [7] Shengze Cai, Zhiping Mao, Zhicheng Wang, Minglang Yin, and George Em Karniadakis.  Physics-informed neural networks (pinns) for fluid mechanics: A review.  CoRR, abs/2105.09506, 2021.
  * [8] Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shuran Song.  Diffusion policy: Visuomotor policy learning via action diffusion.  arXiv preprint arXiv:2303.04137, 2023.
  * [9] Jonas Degrave, Federico Felici, Jonas Buchli, Michael Neunert, Brendan Tracey, Francesco Carpanese, Timo Ewalds, Roland Hafner, Abbas Abdolmaleki, Diego de Las Casas, et al.  Magnetic control of tokamak plasmas through deep reinforcement learning.  Nature, 602(7897):414–419, 2022.
  * [10] Prafulla Dhariwal and Alexander Nichol.  Diffusion models beat gans on image synthesis.  Advances in neural information processing systems, 34:8780–8794, 2021.
  * [11] Jie Ding, Min Wu, and Min Xiao.  Nonlinear decoupling control with pi λ d μ neural network for mimo systems.  IEEE Transactions on Neural Networks and Learning Systems, pages 1–8, 2022.
  * [12] Yilun Du, Conor Durkan, Robin Strudel, Joshua B Tenenbaum, Sander Dieleman, Rob Fergus, Jascha Sohl-Dickstein, Arnaud Doucet, and Will Sussman Grathwohl.  Reduce, reuse, recycle: Compositional generation with energy-based diffusion models and mcmc.  In International Conference on Machine Learning, pages 8489–8510. PMLR, 2023.
  * [13] MA Elhawary.  Deep reinforcement learning for active flow control around a circular cylinder using unsteady-mode plasma actuators.  arXiv preprint arXiv:2012.10165, 2020.
  * [14] Amir-massoud Farahmand, Saleh Nabi, and Daniel N. Nikovski.  Deep reinforcement learning for partial differential equation control.  In 2017 American Control Conference (ACC), pages 3120–3127, 2017.
  * [15] 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.
  * [16] Elie Hachem, Hassan Ghraieb, Jonathan Viquerat, Aurélien Larcher, and P Meliga.  Deep reinforcement learning for the control of conjugate heat transfer.  Journal of Computational Physics, 436:110317, 2021.
  * [17] Haoran He, Chenjia Bai, Kang Xu, Zhuoran Yang, Weinan Zhang, Dong Wang, Bin Zhao, and Xuelong Li.  Diffusion model is an effective planner and data synthesizer for multi-task reinforcement learning.  Advances in neural information processing systems, 36, 2024.
  * [18] Jonathan Ho, Ajay Jain, and Pieter Abbeel.  Denoising diffusion probabilistic models.  Advances in neural information processing systems, 33:6840–6851, 2020.
  * [19] Jonathan Ho and Tim Salimans.  Classifier-free diffusion guidance.  In NeurIPS 2021 Workshop on Deep Generative Models and Downstream Applications, 2021.
  * [20] Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David J Fleet.  Video diffusion models.  arXiv:2204.03458, 2022.
  * [21] Philipp Holl, Nils Thuerey, and Vladlen Koltun.  Learning to control pdes with differentiable physics.  In International Conference on Learning Representations, 2020.
  * [22] Benjamin Holzschuh, Simona Vegetti, and Nils Thuerey.  Solving inverse physics problems with score matching.  In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
  * [23] Rakhoon Hwang, Jae Yong Lee, Jin Young Shin, and Hyung Ju Hwang.  Solving pde-constrained control problems using operator learning.  In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 4504–4512, 2022.
  * [24] Allan Jabri, David Fleet, and Ting Chen.  Scalable adaptive computation for iterative generation.  arXiv preprint arXiv:2212.11972, 2022.
  * [25] Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine.  Planning with diffusion for flexible behavior synthesis.  In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162, pages 9902–9915. PMLR, 17–23 Jul 2022.
  * [26] Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine.  Planning with diffusion for flexible behavior synthesis.  Proceedings of Machine Learning Research, 162:9902–9915, 17–23 Jul 2022.
  * [27] Linlin Kang, An-Kang Gao, Fei Han, Weicheng Cui, and Xi-Yun Lu.  Propulsive performance and vortex dynamics of jellyfish-like propulsion with burst-and-coast strategy.  Physics of Fluids, 35(9), 2023.
  * [28] Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  arXiv preprint arXiv:1412.6980, 2014.
  * [29] Remi Lam, Alvaro Sanchez-Gonzalez, Matthew Willson, Peter Wirnsberger, Meire Fortunato, Ferran Alet, Suman Ravuri, Timo Ewalds, Zach Eaton-Rosen, Weihua Hu, et al.  Learning skillful medium-range global weather forecasting.  Science, 382(6677):1416–1421, 2023.
  * [30] Leon Lapidus and George F Pinder.  Numerical solution of partial differential equations in science and engineering.  John Wiley & Sons, 1999.
  * [31] A Larcher and E Hachem.  A review on deep reinforcement learning for fluid mechanics: An update.  Physics of Fluids, 34(11), 2022.
  * [32] Yanyu Li, Huan Wang, Qing Jin, Ju Hu, Pavlo Chemerys, Yun Fu, Yanzhi Wang, Sergey Tulyakov, and Jian Ren.  Snapfusion: Text-to-image diffusion model on mobile devices within two seconds.  arXiv preprint arXiv:2306.00980, 2023.
  * [33] Yun Li, Kiam Heong Ang, and G.C.Y. Chong.  Pid control system analysis and design.  IEEE Control Systems Magazine, 26(1):32–41, 2006.
  * [34] Zongyi Li, Nikola Borislavov Kovachki, Kamyar Azizzadenesheli, Burigede liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar.  Fourier neural operator for parametric partial differential equations.  In International Conference on Learning Representations, 2021.
  * [35] 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.
  * [36] Jacques Louis Lions.  Optimal control of systems governed by partial differential equations, volume 170.  Springer, 1971.
  * [37] Dong C Liu and Jorge Nocedal.  On the limited memory bfgs method for large scale optimization.  Mathematical programming, 45(1-3):503–528, 1989.
  * [38] Nan Liu, Shuang Li, Yilun Du, Antonio Torralba, and Joshua B Tenenbaum.  Compositional visual generation with composable diffusion models.  In European Conference on Computer Vision, pages 423–439. Springer, 2022.
  * [39] Antoine McNamara, Adrien Treuille, Zoran Popović, and Jos Stam.  Fluid control using the adjoint method.  ACM Transactions On Graphics (TOG), 23(3):449–456, 2004.
  * [40] Chenlin Meng, Robin Rombach, Ruiqi Gao, Diederik Kingma, Stefano Ermon, Jonathan Ho, and Tim Salimans.  On distillation of guided diffusion models.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 14297–14306, 2023.
  * [41] Rajat Mittal and Gianluca Iaccarino.  Immersed boundary methods.  Annu. Rev. Fluid Mech., 37:239–261, 2005.
  * [42] Keith W Morton and David Francis Mayers.  Numerical solution of partial differential equations: an introduction.  Cambridge university press, 2005.
  * [43] Saviz Mowlavi and Saleh Nabi.  Optimal control of pdes using physics-informed neural networks.  Journal of Computational Physics, 473:111731, 2023.
  * [44] Alex Nichol and Prafulla Dhariwal.  Improved Denoising Diffusion Probabilistic Models.  ArXiv, February 2021.
  * [45] Alexander Quinn Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Pamela Mishkin, Bob Mcgrew, Ilya Sutskever, and Mark Chen.  GLIDE: Towards photorealistic image generation and editing with text-guided diffusion models.  Proceedings of Machine Learning Research, 162:16784–16804, 17–23 Jul 2022.
  * [46] Guido Novati, Siddhartha Verma, Dmitry Alexeev, Diego Rossinelli, Wim M Van Rees, and Petros Koumoutsakos.  Synchronisation through learning for two self-propelled swimmers.  Bioinspiration & biomimetics, 12(3):036001, 2017.
  * [47] Yangchen Pan, Amir-massoud Farahmand, Martha White, Saleh Nabi, Piyush Grover, and Daniel Nikovski.  Reinforcement learning with function-valued action spaces for partial differential equation control.  In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, volume 80, pages 3986–3995. PMLR, 10–15 Jul 2018.
  * [48] Aditya A Paranjape, Jinyu Guan, Soon-Jo Chung, and Miroslav Krstic.  Pde boundary control for flexible articulated wings on a robotic aircraft.  IEEE Transactions on Robotics, 29(3):625–640, 2013.
  * [49] Tobias Pfaff, Meire Fortunato, Alvaro Sanchez-Gonzalez, and Peter W Battaglia.  Learning mesh-based simulation with graph networks.  arXiv preprint arXiv:2010.03409, 2020.
  * [50] Dean A Pomerleau.  Alvinn: An autonomous land vehicle in a neural network.  Advances in neural information processing systems, 1, 1988.
  * [51] Ilan Price, Alvaro Sanchez-Gonzalez, Ferran Alet, Timo Ewalds, Andrew El-Kadi, Jacklynn Stott, Shakir Mohamed, Peter Battaglia, Remi Lam, and Matthew Willson.  Gencast: Diffusion-based ensemble forecasting for medium-range weather.  arXiv preprint arXiv:2312.15796, 2023.
  * [52] Bartosz Protas.  Adjoint-based optimization of pde systems with alternative gradients.  Journal of Computational Physics, 227(13):6490–6510, 2008.
  * [53] 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.
  * [54] 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.  J. Comput. Phys., 378:686–707, 2019.
  * [55] Leif Ristroph and Stephen Childress.  Stable hovering of a jellyfish-like flying machine.  Journal of The Royal Society Interface, 11(92):20130992, 2014.
  * [56] Olaf Ronneberger, Philipp Fischer, and Thomas Brox.  U-net: Convolutional networks for biomedical image segmentation.  In Medical Image Computing and Computer-Assisted Intervention–MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18, pages 234–241. Springer, 2015.
  * [57] Tim Salimans and Jonathan Ho.  Progressive distillation for fast sampling of diffusion models.  In International Conference on Learning Representations, 2021.
  * [58] Alvaro Sanchez-Gonzalez, Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovec, and Peter Battaglia.  Learning to simulate complex physics with graph networks.  In International conference on machine learning, pages 8459–8468. PMLR, 2020.
  * [59] Max Schwenzer, Muzaffer Ay, Thomas Bergs, and Dirk Abel.  Review on model predictive control: An engineering perspective.  The International Journal of Advanced Manufacturing Technology, 117(5-6):1327–1349, 2021.
  * [60] Dule Shu, Zijie Li, and Amir Barati Farimani.  A physics-informed diffusion model for high-fidelity flow field reconstruction.  Journal of Computational Physics, 478:111972, 2023.
  * [61] Sabrine Slama, Ayachi Errachdi, and Mohamed Benrejeb.  Neural adaptive pid and neural indirect adaptive control switch controller for nonlinear mimo systems.  Mathematical Problems in Engineering, 2019, 2019.
  * [62] Jiaming Song, Chenlin Meng, and Stefano Ermon.  Denoising diffusion implicit models.  arXiv preprint arXiv:2010.02502, 2020.
  * [63] Kiwon Um, Robert Brand, Yun Raymond Fei, Philipp Holl, and Nils Thuerey.  Solver-in-the-loop: Learning from differentiable physics to interact with iterative pde-solvers.  Advances in Neural Information Processing Systems, 33:6111–6122, 2020.
  * [64] 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.
  * [65] Pantelis R Vlachas, Georgios Arampatzis, Caroline Uhler, and Petros Koumoutsakos.  Multiscale simulations of complex systems by learning their effective dynamics.  Nature Machine Intelligence, 4(4):359–366, 2022.
  * [66] Marin Vlastelica, Tatiana Lopez-Guevara, Kelsey R Allen, Peter Battaglia, Arnaud Doucet, and Kim Stachenfeld.  Diffusion generative inverse design.  2023\.
  * [67] Sifan Wang, Hanwen Wang, and Paris Perdikaris.  Learning the solution operator of parametric partial differential equations with physics-informed deeponets.  Science advances, 7(40):eabi8605, 2021.
  * [68] Gabriel D Weymouth.  Lily pad: Towards real-time interactive computational fluid dynamics.  arXiv preprint arXiv:1510.06886, 2015.
  * [69] Tailin Wu, Takashi Maruyama, and Jure Leskovec.  Learning to accelerate partial differential equations via latent global evolution.  In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, LA, USA, November 28 - December 9, 2022, 2022.
  * [70] Tailin Wu, Takashi Maruyama, Long Wei, Tao Zhang, Yilun Du, Gianluca Iaccarino, and Jure Leskovec.  Compositional generative inverse design.  arXiv preprint arXiv:2401.13171, 2024.
  * [71] Yanjie Ze, Gu Zhang, Kangning Zhang, Chenyuan Hu, Muhan Wang, and Huazhe Xu.  3d diffusion policy.  arXiv preprint arXiv:2403.03954, 2024.
  * [72] Shuang Zhang, Xinyu Qian, Zhijie Liu, Qing Li, and Guang Li.  Pde modeling and tracking control for the flexible tail of an autonomous robotic fish.  IEEE Transactions on Systems, Man, and Cybernetics: Systems, 52(12):7618–7627, 2022.
  * [73] Min Zhao, Fan Bao, Chongxuan Li, and Jun Zhu.  Egsde: Unpaired image-to-image translation via energy-guided stochastic differential equations.  Advances in Neural Information Processing Systems, 35:3609–3623, 2022.
  * [74] Qingqing Zhao, David B Lindell, and Gordon Wetzstein.  Learning to solve PDE-constrained inverse problems with graph networks.  Proceedings of Machine Learning Research, 162:26895–26910, 17–23 Jul 2022.
  * [75] Zifeng Zhuang, Kun Lei, Jinxin Liu, Donglin Wang, and Yilang Guo.  Behavior proximal policy optimization.  arXiv preprint arXiv:2302.11312, 2023.

##  Appendix A Additional Related Work

###  A.1 Physical Systems Simulation

Complex physical systems simulation forms the foundation of systems control.
While classical numerical techniques for simulating physical systems are
renowned for their accuracy, they are often associated with significant
computational expenses [42, 30]. Recently, neural network-based solvers show a
significant advantage over classical solvers in accelerating simulations. They
could be roughly divided into three primary classes: data-driven methods [34,
58, 49, 5, 69, 4, 29], Physics-Informed Neural Networks (PINNs) [54, 7, 67],
and solver-in-the-loop methods [63, 65]. Most of them use an iterative
horizontal prediction framework. Instead, we treat the system trajectory as a
whole variable and use diffusion models to learn an explicit simulator
conditioned on control sequences. A notable work is by [6], which introduces
diffusion models for temporal forecasting. While both our work and [6]’s
employ diffusion models, we tackle a different task of physical system
control. Furthermore, we incorporate the control objective into the inference
and introduce prior reweighting to tune the influence of the prior.

###  A.2 Physical Systems Control

For physical systems whose dynamics are described by PDEs, the adjoint methods
[36, 39, 52] have been the most widely used approach for system control in the
last decades. It is accurate but computationally expensive. Deep learning-
based methods have emerged as a powerful tool for modeling physical systems’
dynamics. Supervised learning (SL) [21, 23] trains parameterized models to
directly optimize control using backpropagation through time over the entire
trajectory. For example, [21] proposes a hierarchical predictor-corrector
scheme to control complex nonlinear physical systems over long time frames. A
more recent work proposed by [23] designs two stages which respectively learn
the solution operator and search for optimal control. Different from these
methods, we do not use the surrogate model, and learn both state trajectories
and control sequences in an integrated way. Reinforcement learning (RL) [14,
47, 53] treats control signals as actions and learns policies to make
sequential decisions. Particularly in the field of fluid dynamics [31],
reinforcement learning has been applied to a multitude of specific problems
including drag reduction [53, 13], conjugate heat transfer [3, 16] and
swimming [46, 64]. But they implicitly consider physics information and
sequentially make decisions. In contrast, we generalize the entire
trajectories, which results in a global optimization with consideration of
physical information learned by models. Recently, PINNs are also incorporated
in PDE control [43], but they require an explicit form of PDE dynamics, while
our method is data-driven and can deal with a broader range of complex
physical system control problems without explicit PDE dynamics.

###  A.3 Diffusion Models

Diffusion models [18] have significantly advanced in applications such as
image and text generation [10, 45], inverse design [70, 66], inverse problem
[22], physical simulation [6, 51], and decision-making [26, 1, 17]. In
particular, recent progress in robot control shows that diffusion models have
significant advantages over existing reinforcement learning methods for action
planning [8, 71]. Generating diverse yet consistent samples poses a challenge.
For diversity, methods [38, 2, 73, 12] that integrate score estimates from
various models have been effective. For consistency, guidance diffusion
techniques [10, 19] have been utilized to generate condition-specific samples.
Our approach differs by flattening the joint distribution to achieve better
control by slightly expanding beyond the prior distribution range.

##  Appendix B 1D Burgers’ Equation Visualization

We present more visualization results of our method and baselines under three
settings: FOPC, POFC, and POPC in Figure 6, Figure 7 and Figure 8,
respectively. Under each setting, we present the results of five randomly
selected samples from the test dataset. The goal of control is to make the
final state $\mathbf{u}_{T}$ ($T=10$) close to the target state.

Figure 6: Visualizations results of 1D Burgers’ equation control under the FO-
PC (full observation, partiall control) setting. The curve for the system
state $\mathbf{u}_{t}$ of each time step $t=0,\cdots,10$ under control is
plotted for our method (DiffPhyCon) and baselines. The $x$-axis is the spatial
coordinate and the $y$-axis is the value of the system state.

From these visualization results in in Figure 6, Figure 7, and Figure 8, it
can be observed that under the control of our DiffPhyCon, the system state
could converge smoothly to the target state given different initial states,
and the final state $\mathbf{u}_{T}$ always coincides with the target state.
Furthermore, this observation is consistent under all three settings: FOPC,
POFC, and POPC, implying that our DiffPhyCon is effective in addressing the
partial observation and partial control challenges. In contrast, the baselines
showed inferior results. Even the best baseline, BPPO, presents obvious
mismatching with the target state on some samples.

Figure 7: Visualizations results of 1D Burgers’ equation control under the PO-
FC (partiall observation, full control) setting. The curve for the system
state $\mathbf{u}_{t}$ of each time step $t=0,\cdots,10$ under control is
plotted for our method (DiffPhyCon) and baselines. The $x$-axis is the spatial
coordinate and the $y$-axis is the value of the system state.

Figure 8:  Visualizations results of 1D Burgers’ equation control under the
PO-PC (partially observation, partially control) setting. The curve for the
system state $\mathbf{u}_{t}$ of each time step $t=0,\cdots,10$ under control
is plotted for our method (DiffPhyCon) and baselines. The $x$-axis is the
spatial coordinate and the $y$-axis is the value of the system state.

##  Appendix C Additional Details for 1D Burgers’ Equation Control

###  C.1 Data Generation

We use the finite difference method (called solver or ground-truth solver in
the following) to generate the training data for the 1D Burgers’ equation.
Specifically, the initial value $\mathbf{u}_{0}(x)$ and the control sequence
$\mathbf{w}(t,x)$ are both randomly generated, and then the states
$\mathbf{u}(t,x)$ are numerically computed using the solver.

In the numerical simulation (using the ground-truth solver), a domain of
$x=[0,1],~{}t=[0,1]$ is simulated. The space is discretized into $128$ grids
and time into $10000$ steps. However, in the dataset, only 10 time stamps are
stored. For the control sequence $\mathbf{w}$, its refreshing rate is
$0.1^{-1}$, i.e., $\mathbf{w}(t,x),~{}t\in[0.1k,0.1(k+1)],~{}k\in\\{0,..,9\\}$
does not change with $t$. Therefore, the data size of each trajectory is
$[11,128]$ for the state $\mathbf{u}$ and $[10,128]$ for the control
$\mathbf{w}$.

In all settings, the initial value $\mathbf{u}(0,x)$ is a superposition of two
Gaussian functions
$\mathbf{u}(0,x)=\sum_{i=1}^{2}a_{i}e^{-\frac{(x-b_{i})^{2}}{2\sigma_{i}^{2}}}$,
where $a_{i},b_{i},\sigma_{i}$ are all randomly sampled from uniform
distributions: $a_{1}\sim U(0,2),~{}a_{2}\sim
U(-2,0),~{}b_{1}\sim(0.2,0.4),~{}b_{2}\sim(0.6,0.8),~{}\sigma_{1}\sim
U(0.05,0.15),~{}\sigma_{2}\sim U(0.05,0.15)$. Similarly, the control sequence
$\mathbf{w}(x,t)$ is also a superposition of 8 Gaussian functions

$\mathbf{w}(t,x)=\sum_{i=1}^{8}a_{i}e^{-\frac{(\mathbf{x}-b_{1,i})^{2}}{2\sigma_{1,i}^{2}}}e^{-\frac{(t-b_{2,i})^{2}}{2\sigma_{2,i}^{2}}},$
(20)

where each parameter is independently generated as follows: $b_{1,i}\sim
U(0,1),~{}b_{2,i}\sim U(0,1),~{}\sigma_{1,i}\sim
U(0.05,0.2),~{}\sigma_{2,i}\sim U(0.05,0.2)$, while $a_{1}\sim U(-1.5,1.5)$
and for $i\geq 2$, $a_{i}\sim U(-1.5,1.5)$ or $0$ with equal probabilities.
$\mathbf{u}(t,x),~{}(t\neq 0)$ is then numerically simulated (using the
ground-truth solver) given $\mathbf{u}(0,x)$ and $\mathbf{w}(t,x)$ based on
Eq. (16). The setting of the dataset generation is based on a previous work
[23].

We generated $90000$ trajectories for the training set and $50$ for the
testing set. Each trajectory takes up $32$KB space and the size of the dataset
sums up to $2$GB.

###  C.2 Experimental Setting

During inference, alongside the control sequence $\mathbf{w}(t,x)$, our
diffusion model generates states $\mu(t,x)$, and some models produce surrogate
states $\mu(t,x)$ when feeding the control $\mathbf{w}(t,x)$ into the
corresponding surrogate model. However, our reported evaluation metric
$\mathcal{J}_{\text{actual}}$ is always computed by feeding the control
$\mathbf{w}(t,x)$ into the ground truth numerical solver to get
$\mathbf{u}_{\text{g.t.}}(t,x)$ and computed following Eq. (17).

Following are three different settings of our experiments.

####  C.2.1 Partial Observation, Full Control

In realistic scenarios, the system is often unable to be observed completely.
Generally speaking, it is impractical to place sensors everywhere in a system,
so the ability of the model to learn from incomplete data is imperative. To
evaluate this, we hide some parts of $\mathbf{u}$ in this setting and measure
the $\mathcal{J}_{\text{actual}}$ of model control.

Specifically, $\mathbf{u}(t,x),~{}x\in[\frac{1}{4},\frac{3}{4}]$ is set to
zero in the dataset during training and
$\mathbf{u}_{0}(x),~{}x\in[\frac{1}{4},\frac{3}{4}]$ is also set to zero
during testing. In this partial observation setting
$\Omega=[1,\frac{1}{4}]\cup[\frac{3}{4},1]$. Since no information in the
central $\frac{1}{2}$ space is ever known, the model does not know what will
influence the control outcome of the unobserved states. Therefore, controlling
the unobserved states is not a reasonable task and they are excluded from the
evaluation metric.

This setting is particularly challenging not only because of the uncertainty
introduced by the unobserved states but also the generation of the control in
the central locations that implicitly affect the controlled $\mathbf{u}$ at
$x\in\Omega$.

####  C.2.2 Full Observation, Partial Control

This is another setting of practical relevance, where only a fraction of the
system can be controlled. The control sequence is enforced to be zero in the
central locations of $x\in[\frac{1}{4},\frac{3}{4}]$. $\Omega$ is still
$[0,1]$, and $\mathcal{J}$ is evaluated on all of the observed states, though.

Some modifications to the dataset should be mentioned. The generation of the
data involves first generating $\mathbf{w}$ as before, followed by setting the
central $\frac{1}{2}$ of $\mathbf{w}$ to zero. To compensate for the decreased
control intensity so that the magnitude of $\mathbf{u}$ can be roughly
comparable to the full control setting, we double the magnitude of
$\mathbf{w}$. During the evaluation, the output control sequence is also post-
processed to be zero in $x\in[\frac{1}{4},\frac{3}{4}]$.

It is worth noting that in this setting, even when the control energy is not
limited at all, it is still challenging to find a perfect control since the
model has to learn how to indirectly impose control on the central locations.

####  C.2.3 Partial Observation, Partial Control

The final setting is the combination of the previous two settings. Only
$\Omega=[0,\frac{1}{4}]\cup[\frac{3}{4},1]$ is observed, controlled and
evaluated.

It is worth noting that some models require accessing the current state to
produce output. If the model interacts with the ground truth solver instead of
a surrogate model, then the result would be unfairly good since the
information of the unobserved states is leaked through the interaction.

###  C.3 Model

Since the training of models
$\bm{\epsilon}_{\phi}\approx\nabla_{\mathbf{w}}\log p(\mathbf{w})$ and
$\bm{\epsilon}_{\theta}\approx\nabla_{\mathbf{u},\mathbf{w}}\log
p(\mathbf{u},\mathbf{w})$ are essentially the same and the latter model is
exactly DiffPhyCon-lite, we will introduce DiffPhyCon-lite first.

####  C.3.1 DiffPhyCon-lite

In general, DiffPhyCon-lite follows the formulation of [18] which is also
described in the main text. The data of $\mathbf{u}$ and $\mathbf{w}$ is fed
in as images of size $(N_{t},N_{x})$ where $N_{t}$ is the number of time steps
(11 and 10 respectively) and $N_{x}$ is the spatial grids (128). Since the two
$N_{t}$s for $\mathbf{u}$ and $\mathbf{w}$ are inconsistent, we zero-pad them
into the size of $16$. Then, $\mathbf{u}$ and $\mathbf{w}$ are stacked as two
channels and fed into the 2D DDPM model.

A 2D UNet $\bm{\epsilon}_{\theta}$ is used to learn to predict
$\bm{\epsilon}$. It is structured into three main components: the downsampling
encoder, the central module, and the upsampling decoder. The downsampling
encoder is made up of four layers, each layer consisting of two ResNet blocks,
one linear Attention block, and one downsampling convolution block. The
central module also consists of two ResNet blocks and one linear Attention
block. Each upsampling layer is the same as the downsampling layer except the
downsampling block is replaced by the upsampling convolution block.

In our experiments, we found that the control result is best when learning the
conditional probability distribution of
$p(\mathbf{w}_{[0,T-1]},\mathbf{u}_{[1,T-1]}\mid\mathbf{u}_{0},\mathbf{u}_{T})$
In summary, $\bm{\epsilon}_{\theta}$ takes in the current trajectory
$\mathbf{u}$, control $\mathbf{w}$, step $k$, $\mathbf{u}_{0}$ and
$\mathbf{u}_{T}$ as input, and predicts the noise of $\mathbf{u}$ and
$\mathbf{w}$. Note that it is not trained to predict $\mathbf{u}_{0}$ and
$\mathbf{u}_{T}$ which are used as a condition, but there are still model
outputs at the corresponding locations for the data shape consistency across
different design choices of DiffPhyCon-lite. The hyperparameters in different
settings are listed in Table 4.

####  C.3.2 DiffPhyCon

In terms of implementation, DiffPhyCon is simply adding
$\bm{\epsilon}_{\phi}(\mathbf{w})$ to
$\bm{\epsilon}_{\theta}(\mathbf{u},\mathbf{w})$ during inference as shown in
Section 3.2, where $\bm{\epsilon}_{\theta}$ is the output of the denoising
network in DiffPhyCon-lite while $\bm{\epsilon}_{\phi}$ is a new denoising
network that is trained to generate $\mathbf{w}$ following the dataset
distribution. Therefore, we only describe the model of $\bm{\epsilon}_{\phi}$
here.

$\bm{\epsilon}_{\phi}$ takes input of $\mathbf{w},k$ as in the standard DDPM
and $\mathbf{u}_{0},\mathbf{u}_{T}$ as guidance conditioning. The output of
$\bm{\epsilon}_{\phi}(\mathbf{w})$ is of the same shape as $\mathbf{w}$, so it
can be treated as a network learning to sample from
$p(\mathbf{w})\coloneqq\int p(\mathbf{u},\mathbf{w})\mathrm{d}\mathbf{u}$ The
output of $\bm{\epsilon}_{\phi}(\mathbf{w})$ at the locations of $\mathbf{u}$
is thus filled with zeros. The model hyperparameters are also listed in Table
4.

Table 4: Hyperparameters of the UNet architecture and training for the results
of 1D Burgers’ equation in Table 1.

Hyperparameter name | Full observation | Partial observation | Partial observation
---|---|---|---
Partial Control | Full Control | Partial Control
UNet $\bm{\epsilon}_{\phi}(\mathbf{w})$
Initial dimension | 32 | 32 | 32
Downsampling/Upsampling layers | 4 | 4 | 4
Convolution kernel size | 3 | 3 | 3
Dimension multiplier | $[1,2,4,8]$ | $[1,2,4,8]$ | $[1,2,4,8]$
Resnet block groups | 8 | 8 | 8
Attention hidden dimension | 32 | 32 | 32
Attention heads | 4 | 4 | 4
UNet $\bm{\epsilon}_{\theta}(\mathbf{u},\mathbf{w})$
Initial dimension | 128 | 128 | 64
Downsampling/Upsampling layers | 4 | 4 | 4
Convolution kernel size | 3 | 3 | 3
Dimension multiplier | $[1,2,4]$ | $[1,2,4,8]$ | $[1,2,4,8]$
Resnet block groups | 8 | 8 | 8
Attention hidden dimension | 32 | 32 | 32
Attention heads | 4 | 4 | 4
Training
Training batch size | 16 | 16 | 16
Optimizer | Adam | Adam | Adam
Learning rate | 1e-4 | 1e-4 | 1e-4
Training steps | 190000 | 170000 | 190000
Learning rate scheduler | cosine annealing | cosine annealing | cosine annealing
Inference
Sampling iterations | 1000 | 1000 | 1000
Intensity of energy guidance $\mathcal{J}=\int\|\mathbf{w}\|^{2}\mathrm{d}x\mathrm{d}t$ | 0 | 0 | 0
Scheduler of energy guidance $\mathcal{J}=\int\|\mathbf{w}\|^{2}\mathrm{d}x\mathrm{d}t$ | cosine | cosine | cosine

###  C.4 Training and Evaluation

##### Training

During training, the $\mathbf{u}_{0}$ and $\mathbf{u}_{d}$ without noise are
fed into the model and the model outputs at the corresponding locations are
excluded from the loss. In the partial observation settings, the unobserved
data is invisible to the model during both training and testing as introduced
in Appendix C.2. We simply pad zero in the corresponding locations of the
model input and also exclude these locations in the training loss. Therefore,
the model only learns the correlation between the observed states and control
sequences. In the partial control setting, we train our DiffPhyCon on the
dataset with control being zero in $x\in[\frac{1}{4},\frac{3}{4}]$. In this
way, the model naturally learns to output zero at the “non-controllable”
locations.

We use the MSE loss to train the denoising UNets and other training
hyperparameters are listed in Table 4.

##### Inference

During inference, $\mathbf{u}_{0}$ and $\mathbf{u}_{T}$ are set to the target
$\mathbf{u}_{0}$ and $\mathbf{u}_{d}$ so that the DDPM generates samples
satisfying the physical constraint that is also conditioned on the target
($\mathbf{u}_{d}$) or the constraint ($\mathbf{u}_{0}$). In the partial
observation setting, the $\mathbf{u}_{0}$ and $\mathbf{u}_{T}$ drawn from the
testing set are all filled zero at the unobserved locations
$x\in[\frac{1}{4},\frac{3}{4}]$, which is the same as the data used to train
the UNets. In the partial control scenarios,

During inference, we replace the denoising network’s output
$\bm{\epsilon}_{\theta}(\mathbf{u},\mathbf{w})$ with
$\bm{\epsilon}_{\theta}(\mathbf{u},\mathbf{w})+(\gamma-1)\bm{\epsilon}_{\phi}(\mathbf{w})$.
It is worth noting that $\bm{\epsilon}_{\theta}(\mathbf{u},\mathbf{w})$
denoises $u$ and $\mathbf{w}$ simultaneously while
$\bm{\epsilon}_{\phi}(\mathbf{w})$ only denoises $\mathbf{w}$. In our
experiments, we found that adding a schedule to the output of the $\mathbf{w}$
network is beneficial. The results in Table 1 are generated with a reverse
Sigmoid schedule following

$\bm{\epsilon}_{k}=\bm{\epsilon}_{\theta}(\mathbf{u}_{k},\mathbf{w}_{k},k,\mathbf{u}_{0},\mathbf{u}_{T})+\frac{}{}(\gamma-1)\beta_{K-k}\bm{\epsilon}_{\phi}(\mathbf{w}_{k},k,\mathbf{u}_{0},\mathbf{u}_{T}),$
(21)

where $\beta$ is defined as the noise schedule in [18]. The inference of
DiffPhyCon-lite is simply setting $\gamma=1$, which neglects the effect of the
model $\bm{\epsilon}_{\phi}(\mathbf{w})$.

When trying to regulate the control energy, however, it is not as natural to
learn a conditional model. Therefore, we use the external guidance
$\mathcal{J}=\int\mathbf{w}(x,t)\mathrm{d}x\mathrm{d}t$ to produce cost-
limited control sequences that are shown in Figure 3. The gradient of the
external guidance is computed and added to $\bm{\epsilon}_{k}$ in Eq. (21).
Note that we use the predicted clean sample $\hat{\mathbf{w}}_{k}$ and
$\hat{\mathbf{u}}_{k}$ at the $k$-th step to compute $\nabla\mathcal{J}$ since
they suffer less from being noisy and leading to deviated guidance.
$\mathbf{u}_{k}$ and $\mathbf{w}_{k}$ are computed following
$\mathbf{x}_{0}\approx\hat{\mathbf{x}}_{k}=\frac{1}{\sqrt{\bar{\alpha}_{k}}}\mathbf{x}_{k}-\frac{\sqrt{1-\bar{\alpha}_{k}}}{\sqrt{\bar{\alpha}_{k}}}\bm{\epsilon}_{k}$
where $\mathbf{x}$ represents $\mathbf{u}$ or $\mathbf{w}$ as in [18]. In our
experiments, we use a cosine scheduling of the external guidance, and thus the
final predicted noise would be
$\bm{\epsilon}_{k}+\lambda\alpha_{k}\nabla\mathcal{J}$ where $\beta_{k}$ is
the noise schedule in [18] with the cosine schedule.

##### Evaluation

After generating the trajectory $\mathbf{u}$ and control sequence
$\mathbf{w}$, we feed the control sequence $\mathbf{w}$ into the ground-truth
solver and simulate the final state $\mathbf{u}$ given the generated
$\mathbf{w}$ and the initial condition $\mathbf{u}_{0}$ directly drawn from
the testing dataset. The solver is the same as the one used in data generation
in Appendix C.1. Finally, we compute $\mathcal{J}_{\text{actual}}$ following
Eq. (17). In the partial observation setting, the MSE is computed only on the
observed region, and in the control setting, the generated control will first
be set to zero in the uncontrolled region before being fed into the solver.

##  Appendix D Jellyfish Movement Dataset

We use the Lily-Pad simulator [68] to generate the Jellyfish Movement dataset,
which serves as a benchmark for physical system control research and also the
dataset for our 2D evaluation task. Lily-Pad adopts the Immersed Boundary
Method (IBM) [41] to simulate fluid-solid dynamics. The resolution of the 2D
flow field is set to be $128\times 128$. The flow field is assumed to be
boundless in Lily-Pad. The head of the jellyfish is fixed at $(25.6,64)$. Its
two wings are represented by two identical ellipses, where the ratio between
the shorter axis and the longer axis is 0.15. At each moment, the two wings
are symmetric about the central horizontal line $y=64$. For each wing, we
sample $M=20$ points along the wing to represent the boundary of the wing. The
opening angle of the wings is defined as the angle between the longer axis of
the upper wing and the horizontal line. It acts as the control sequence
$\mathbf{w}$ in a 2D jellyfish control experiment.

Each trajectory starts from the largest opening angle and follows a cosine
curve periodically with period $T^{{}^{\prime}}=200$. Trajectories differ in
initial angle, angle amplitude, and phase ratio $\tau$ (the ratio between the
closing duration and a whole pitching duration). For each trajectory, the
initial angle $\mathbf{w}_{0}$ is generated as follows: first, sample a random
angle, called mean angle $\mathbf{w}^{(m)}\in[20^{\circ},40^{\circ}]$, then
sample a random angle amplitude
$\mathbf{w}^{(a)}\in[10^{\circ},\min(\mathbf{w}^{(m)},60^{\circ}-\mathbf{w}^{(m)})]$.
The initial $\mathbf{w}_{0}$ is set as
$\mathbf{w}_{0}=\mathbf{w}^{(m)}+\mathbf{w}^{(a)}$. The phase ratio $\tau$ is
randomly sampled from $[0.2,0.8]$. The opening angle $\mathbf{w}_{t}$ of step
$t$ decreases from $\mathbf{w}^{(m)}+\mathbf{w}^{(a)}$ to
$\mathbf{w}^{(m)}-\mathbf{w}^{(a)}$ as $t$ grows from 0 to $\tau
T^{{}^{\prime}}$; then $\mathbf{w}_{t}$ increases from
$\mathbf{w}^{(m)}-\mathbf{w}^{(a)}$ to $\mathbf{w}^{(m)}+\mathbf{w}^{(a)}$ as
$t$ grows from $\tau T^{{}^{\prime}}$ to $T^{{}^{\prime}}$. Afterwards,
$\mathbf{w}_{t}$ varies periodically for $t>T^{{}^{\prime}}$. The range of
$\mathbf{w}_{t}$ is
$[\mathbf{w}^{(m)}-\mathbf{w}^{(a)},\mathbf{w}^{(m)}+\mathbf{w}^{(a)}]\subset[10^{\circ},60^{\circ}]$.
For each trajectory, we simulate for $600$ simulation steps, i.e., 3 periods.
To save space, we only save the piece of trajectory from $T^{{}^{\prime}}=200$
to $3T^{{}^{\prime}}=600$ steps with step size $10$ because the simulation
from $t=0$ to $T^{{}^{\prime}}=200$ is for initialization of the flow field.
Then each trajectory is saved as a $\tilde{T}=(600-200)/10=40$ steps long
sequence. An example of the simulated fluid field and the corresponding curve
of opening angles are shown in Figure 9.

Besides the positions of the boundary points of wings and the opening angles
$\mathbf{w}$, we also use another kind of image-like representation of the
boundaries of wings as this representation contains spatial information that
can be more effectively learned along with physical states (fluid field) by
convolution neural networks. For each trajectory, this image-like boundary
representation is compatible with physical states in shape. At each time step,
boundaries of two wings are merged and then represented as a tensor of shape
[3, 64, 64], where it has three features for each grid cell: a binary mask
indicating whether the cell is inside a boundary (denoted by 1) or in the
fluid (denoted by 0), and a relative position $(\Delta\mathbf{x},\Delta y)$
between the cell center to the closest point on the boundary. For each
trajectory, we save system states, opening angles, boundary points, boundary
masks and offsets, and force data. They are specified as:

  * •

system states $\mathbf{u}$: shape $[\tilde{T},3,64,64]$. For each step, we
save the states of the fluid field consisting of velocity in $\mathbf{x}$ and
$y$ directions and pressure. To save space, we downsample the resolution from
$128\times 128$ to $64\times 64$.

    * –

velocity: $[\tilde{T},2,64,64]$.

    * –

pressure: $[\tilde{T},1,64,64]$.

  * •

opening angels $\mathbf{w}$: shape $[\tilde{T}]$. For each step, we save the
opening angle in radians.

  * •

boundary points: shape $[\tilde{T},2,M,2]$. For each step, we save the
boundary points on the upper and lower wings. Each wing consists of $M=20$
points and each point consists of 2 coordinates. To make boundary points
compatible with the downsampling of states, coordinates of $\mathbf{x}$ and
$y$ directions are shrunk to half ($64/128$) of the original values.

  * •

boundary mask and offsets $b$: $[\tilde{T},3,64,64]$. For each step, we save
the mask of merged wings with half coordinates of boundary points and offsets
in both $\mathbf{x}$ and $y$ directions. The resolution is $64\times 64$,
compatible with that of the states.

    * –

mask: $[\tilde{T},1,64,64]$.

    * –

offsets: $[\tilde{T},2,64,64]$.

  * •

force: shape $[\tilde{T},2]$. For each step, the simulator outputs the
horizontal and vertical force from the fluid to the jellyfish. The horizontal
force is regarded as a thrust to jellyfish if positive and a drag otherwise.

We generate $n=30,000$ training trajectories and $n=200$ testing trajectories.
Trajectories differ in the above specified parameters
$\mathbf{w}^{(a)},\mathbf{w}^{(m)}$ and $\tau$. Each trajectory occupies about
2MB of storage and the total dataset costs about 100GB.

##  Appendix E Additional Details for 2D Jellyfish Movement Control

###  E.1 Dataset Preparation

Based on our generated dataset in Appendix D, we prepare training samples for
the 2D jellyfish movement control task as follows. We use sliding time windows
that contain $T=20$ successive time steps of states and boundaries as a
sample, which corresponds to $T^{{}^{\prime}}=200$ original simulation steps
and constitutes exactly a period of wing movement. In this way, each
trajectory can produce 20 samples. Therefore, we get 6 million training
samples in total. In each training sample, the initial and the final time
steps share the same opening angle due to periodicity, which serves as the
conditions for control. For each test trajectory, we select the opening angle
of the jellyfish in the initial time and the initial states as the control
condition for both the initial and final time and state initial condition.

Figure 9: Example of a flapping jellyfish in Lily-Pad simulator (a) and the
corresponding curve of opening angle (b). The control signal is the opening
angels $\mathbf{w}$ of the wings of a jellyfish.

###  E.2 Experimental Setting

####  E.2.1 Full Observation

In this setting, we assume all the states of the fluid field are observable.
That is, both the velocity of $x$ and $y$ directions and pressure are
available in all the time steps of the training dataset and the initial time
of the testing dataset.

####  E.2.2 Partial Observation

In this setting, we assume only partial states are observed. A typical
scenario in fluid simulation and control is that we can only observe pressure
data while the velocity data is not easy to access. That is, only pressure is
available in all the time steps of the training samples and the initial time
of the testing samples, hence the state tensor is of shape
$[\tilde{T},1,64,64]$. Notice that even if only pressure is available, we can
still compute the force of fluid on the jellyfish and consequently the control
objective because force is fully determined by the shape of the jellyfish and
pressure. The challenge of this partial observation setting is that the
velocity variable $v$ is missing in Eq. (18), which makes the traditional
numerical solver no longer applicable to solve this physical system control
problem. However, this challenge could be well addressed by our method since
it could learn the relationship between control and pressure despite missing
of the velocity data, and use the accessible control objective as guidance for
flapping control.

###  E.3 Model

####  E.3.1 Architecture

We use a 3D U-Net as the backbone of our diffusion model, in both DiffPhyCon-
lite and DiffPhyCon methods (detailed in the following subsection). In this
paper, the architecture of the 3D U-net we employed is inspired by [20]. To
better capture temporal conditional dependencies, we modify the previous
space-only 3D convolution into space-time 3D convolut ion. Notably, we did not
perform any scaling on the temporal dimension during downsampling or
upsampling. Specifically, our U-net consists of three main modules: the
downsampling encoder, the middle module, and the upsampling decoder. The
downsampling encoder is composed of three layers, each incorporating two
residual modules, one spatial attention module, one temporal attention module,
and one downsampling module. The middle module consists of two residual
modules, one spatial attention module, and one temporal attention module.
Meanwhile, the upsampling decoder consists of four layers, each containing two
residual modules, one spatial attention module, one temporal attention module,
and one upsampling module. The input shape of our U-net is [batch size,
frames, channels, height, width]. During convolution, the operation is
performed on the [frames, height, width] dimensions. The output shape follows
the same structure. Further details are provided in Table 5.

Table 5: Hyperparameters of 3D-Unet architecture. Hyperparameter name | Value
---|---
Kernel size of conv3d | (3, 3, 3)
Padding of conv3d | (1,1,1)
Stride of conv3d | (1,1,1)
Kernel size of downsampling | (1, 4, 4)
Padding of downsampling | (1, 2, 2)
Stride of downsampling | (0, 1, 1)
Kernel size of upsampling | (1, 4, 4)
Padding of upsampling | (1, 2, 2)
Stride of upsampling | (0, 1, 1)
attention heads | 4

####  E.3.2 DiffPhyCon-lite

The DiffPhyCon-lite method learns the denoising network of the joint
distribution $p(\mathbf{u},\mathbf{w}|\mathbf{c})$ where $\mathbf{u}$ is
physical states, $\mathbf{w}$ is the opening angle, and the conditions
$\mathbf{c}$ consist of the initial angle $\mathbf{w}_{0}$, the initial state
$\mathbf{u}_{0}$ and the final angle $\mathbf{w}_{T}=\mathbf{w}_{0}$. We adopt
the 3D U-Net as the backbone. To make the opening angle (of shape $[T]$),
align with physical states (of shape $[T,3,64,64]$ in full observation setting
and $[T,1,64,64]$ in partial observation setting) in shape, we expand the
opening angle to shape $[T,1,64,64]$ along spatial dimension by value copy.
Besides, we also adopt the boundary mask and offsets representation, whose
shape is $[T,3,64,64]$, determined by the opening angles as an auxiliary model
input because they contain explicit spatial features, which makes model
learning more effective. Then states, boundary mask and offsets, and expanded
opening angle are stacked along the channel dimension and we get a tensor of
shape $[T,7,64,64]$ in full observation setting or $[T,5,64,64]$ in partial
observation setting as the model input. The model output contains predicted
noise of states and open angles. Thus its shape is $[T,4,64,64]$ in the full
observation setting or $[T,2,64,64]$ in the partial observation setting, where
the last channel corresponds to the predicted noise of opening angles and
other channels correspond to predicted noise of states.

Figure 10: Inference of our DiffPhyCon-lite and DiffPhyCon in the 2D experiment. Table 6: Hyperparameters of network architecture and training for the 2D experiment. Hyperparameter name | full observation | partial observation
---|---|---
Batch size | 16 | 16
Optimizer | Adam | Adam
Initial learning rate | 0.001 | 0.001
Loss function | MSE | MSE

####  E.3.3 DiffPhyCon

DiffPhyCon learns the denoising network of the joint distribution
$p(\mathbf{u},\mathbf{w}|\mathbf{c})$ and the marginal distribution
$p(\mathbf{w}|\mathbf{c})$. The denoising network of
$p(\mathbf{u},\mathbf{w}|\mathbf{c})$ is exactly the same as the one
introduced in the DiffPhyCon-lite method. The denoising network of
$p(\mathbf{w}|\mathbf{c})$ also adopts the 3D U-Net architecture. Its input
size is the same as that of $p(\mathbf{u},\mathbf{w}|\mathbf{c})$ in both full
and partial observation settings. The difference is that the input states
feature is replaced by the expansion of the initial state $\mathbf{u}_{0}$
along the time dimension by value copy. The output is the predicted noise of
opening angles, whose shape is $[T,1,64,64]$, no matter the full observation
setting or partial observation setting.

###  E.4 Training, Inference, and Evaluation

Training. We use the MSE (mean squared error) between model prediction and the
Gaussian noise as the loss function Eq. (4). The batch size is chosen as 16
and the training involves 200,000 iterations. The learning rate starts from
$1\times 10^{-3}$ and multiplies a factor of 0.1 at the 50000th and 150000th
iterations. Training details are provided in Table 6. The training is
performed on two NVIDIA Tesla A100 GPUs with 80 GB memory for about 3 days.

Inference. The pipeline of inference is shown in Figure 10. Both diffused
variables $\mathbf{u}_{[0,T]}$ and $\mathbf{w}_{[0,T]}$ are initialized from
Gaussian prior and gradually denoised from denoising step $k=1000$ to $k=0$
based on denoising networks and guidance. Because we introduce the boundary
mask and offsets as auxiliary inputs, the model input and output are not
consistent in shape. Thus we introduce a surrogate model (shown as "Boundary
updater" block in Figure 10) to update boundary mask and offsets $b_{[0,T],k}$
for each denoising $k$. Specifically, at each time step $t\in[0,T]$, $b_{t,k}$
is estimated by the initial boundary mask and offsets $b_{0}$, and the
difference of opening angle $\mathbf{w}_{t,k}-\mathbf{w}_{0}$ from time step
$0$ to $t$, which is presented in the right part (after "Denoise under
guidance") of Figure 10. Details about this surrogate model are presented in
H.3. Notice that although this surrogate model is trained on noise-free data,
we do not worry too much about its generalization to the noisy scalar
$\mathbf{w}_{t,k}$ in inference because the estimated $\mathbf{w}_{t,k}$ does
not deviate from the normalized range of noisy free $\mathbf{w}_{t}$ too much.

Our method introduces two kinds of inference: DiffPhyCon-lite and DiffPhyCon.
In DiffPhyCon-lite, we only use the denoising network of the joint
distribution $p(\mathbf{u},\mathbf{w}|\mathbf{u}_{0},\mathbf{w}_{0})$ for
inference, while in DiffPhyCon, we use the additional denoising network of the
marginal distribution $p(\mathbf{w}|\mathbf{u}_{0},\mathbf{w}_{0})$ together
with that of the joint distribution for inference. These two branches are
plotted in the left part of Figure 10, where the notation
$[\mathbf{u}_{0}]\times T$ means expand initial state $\mathbf{u}_{0}$ (of
shape $[3,64,64]$) along time dimension by value copy to form a tensor of
shape $[T,3,64,64]$.

As for guidance, we use a surrogate force model to approximate the force of
fluid on jellyfish. This model is detailed in Subsection H.2. In denoising
step $k$, its input consists of two parts: the first one is the noise-free
state $\hat{\mathbf{u}}_{[0,T]}$ estimated from $\mathbf{u}_{[0,T],k}$ by Eq.
(6); the second one is the noise-free boundary mask and offsets
$\hat{b}_{[0,T],k}$ estimated from noise-free $\hat{\mathbf{w}}_{[0,T]}$ by
the surrogate model to update boundaries, where $\hat{\mathbf{w}}_{[0,T]}$ is
also estimated from $\mathbf{w}_{[0,T],k}$ as in Eq. (6). The model output is
force. Here we only use the horizontal force. Notice that the force could be
computed via the surrogate force model no matter whether states are fully or
partial observation in that force is irrelevant to the velocity of the fluid.
The control objective $\mathcal{J}$ in Eq. (19) is computed as a summation of
force and $R(\hat{\mathbf{w}}_{[0,T]})$. We fix $\zeta=1000$ as a default
setting in Eq. (19) because this value can achieve a balance between scales of
the average speed and the regularizer $R(\mathbf{w})$. We also study the
Pareto performance of varying $\zeta$ in Table 3. Then the gradients of the
objective $\mathcal{J}$ in terms of $\hat{\mathbf{u}}_{[0,T]}$ and
$\hat{\mathbf{w}}_{[0,T]}$ are computed and used in guidance. For DiffPhyCon-
lite, these gradients are substracted from
$[\mathbf{u}_{[0,T],k},\mathbf{w}_{[0,T],k}]$ to generate
$[\mathbf{u}_{[0,T],k-1},\mathbf{w}_{[0,T],k-1}]$. For DiffPhyCon, an
additional term of noise $(\gamma-1)\bm{\epsilon}_{\phi}$ predicted from the
denoising network of the marginal distribution
$p(\mathbf{w}|\mathbf{u}_{0},\mathbf{w}_{0})$ should also be subtracted, as
shown in the upper left part of Figure 10. The effect of this term is
controlled by the scale of the hyperparameter $\gamma$, as studied in Appendix
J. The inference is performed on one NVIDIA Tesla-V100 GPU with 32 GB memory
for about 3 hours for 50 testing samples.

Evaluation. The inference outputs opening angles $\mathbf{w}_{[0,T]}$ of
$T=20$ steps for 50 testing samples. In simulation, for each testing sample,
the ground-truth first $T=20$ steps of the opening angles (which corresponds
to 200 simulation steps) are directly input to the Lily-Pad simulator for the
reason of generating initial states $u_{0}$ of fluid, which is followed by the
predicted control sequences of opening angles (interpolated to 200 steps of
opening angles). The simulator outputs the horizontal force of fluid on the
jellyfish for each simulation step. Finally, average speed $\bar{v}$, energy
cost $R(\mathbf{w})$, and objective $\mathcal{J}$ are computed as metrics. The
average speed $\bar{v}=\frac{1}{T}\int_{0}^{T}v_{t}\mathrm{d}t\approx
v_{0}+\frac{1}{T}\sum_{t=1}^{T-1}(T-t)F_{t}$ $\bar{v}$, where $v_{0}$ is the
initial speed and $F_{t}$ is the horizontal thrust from the fluid. The mass of
the jellyfish is assumed to be 1. The energy cost term
$R(\mathbf{w})=\sum_{t=1}^{T-1}(\mathbf{w}_{t+1}-\mathbf{w}_{t})^{2}$, where
the control sequence $\mathbf{w}=(\mathbf{w}_{1},\cdots,\mathbf{w}_{T})$
represents the predicted opening angles. The periodic term
$d(\mathbf{w}_{0},\mathbf{w}_{T})=\max(|\mathbf{w}_{T}-\mathbf{w}_{0}|-\bm{\epsilon},0)$
is the constraint of periodic opening angles with a small threshold
$\epsilon=0.01$.

##  Appendix F 1D Baselines

###  F.1 PID

Figure 11: The architecture of ANN-PID. To use MIMO PID controller to control
$U_{t}$ to $U_{d}$, we train a neural-network-based PID parameter planner to
output MIMO PID parameters based on $Err_{t}$, then use the PID controller to
output the control sequence $f_{t}$.

Proportional Integral Derivative (PID) [33] control is a versatile and
effective control method widely used in various real-world control scenarios.
It operates by utilizing the difference (error) between the desired target and
the current state of a system. PID control is often considered the go-to
option for many control problems due to its simplicity and usefulness.
However, despite its popularity, PID control does encounter certain
challenges, such as parameter adaptation and limitations when applied to
Single Input Single Output (SISO) systems. In our specific context, the 1D
Burgers’ Equation Control problem presents a Multiple Input Multiple Output
(MIMO) control scenario, which makes it infeasible to directly employ PID
control to regulate the Burgers’ equation. Inspired by the early works [61,
11] using a neural network as a PID parameter adapter, we have integrated deep
learning with PID control to tackle the MIMO control problem. As shown in
Figure 11, ANN(artificial neural network) PID uses a neural network as a PID
parameter adapter to output multiple sets of PID parameters and do multiple
sets of SISO PID control.

The neural network to output PID parameters comprises two 1D convolutional
layers, 2 fully connected layers, and 4 corresponding activation layers. We
use the $L1$ loss of the current state and target state as training loss and
the Adam optimizer [28] to train the model. Detailed information can be found
in Table 7.

Table 7: Hyperparameters of network architecture and training for ANN-PID. Hyperparameter name | Full observation | Partial observation
---|---|---
Kernel size of conv1d | 3 | 3
Padding of conv1d | 1 | 1
Stride of conv1d | 1 | 1
Activation function | Softsign | Softsign
Batch size | 16 | 16
Optimizer | Adam | Adam
Learning rate | 0.0001 | 0.0001
Loss function | MAE | MAE

As PID itself is a SISO control method, ANN-PID uses a neural network to get
multiple sets of PID parameters to do multiple SISO PID controls for MIMO
control in the Burgers’ equation. But here, ANN-PID requires the dimensions of
inputs and outputs to be the same, so it can only cope with full observation,
full control control problems, and partial observation, partial control
control problems. Besides, the ANN-PID controller has 2 training setups,
including directly interacting with the solver and interacting with the 1D
surrogate model in Appendix H.

###  F.2 SAC

Table 8: Hyperparameters of 1D SAC. The full observation partial control, partial observation full control, and partial observation partial control settings share the same hyperparameters. Hyperparameter name | Value
---|---
Hyperparameters for 1D Burgers’ equation control: |
Discount factor for reward | 0.5
Target smoothing coefficient | 0.05
Learning rate of critic loss | 0.0003
Learning rate of entropy loss | 0.003
Learning rate of policy loss | 0.003
Training batch size | 8192
Number of episodes | 1500
Number of model updates per simulator step | 50
Value target updates per step | 15
Size of replay buffer | 1000000
Number of trajectories interacted with the environment per step | 1
Number of layers of critic networks | 3
Number of hidden dimensions of critic networks | 4096
Number of layers of the policy network | 5
Number of hidden dimensions of the policy network | 4096
Activation function | ReLU
Clipping’s range of policy network’s standard deviation output | $\left[e^{-20},e^{2}\right]$

The Soft Actor-Critic (SAC) algorithm [15] is a cutting-edge reinforcement
learning method. Conceptualized as an improvement over traditional Actor-
Critic methods, SAC distinguishes itself by introducing an entropy
regularization term into the loss function, which encourages the policy to
explore more efficiently by maximizing both the expected cumulative reward and
the entropy of the policy itself.

Compared with Deep Deterministic Policy Gradient (DDPG) algorithm [35, 47],
SAC’s entropy regularization encourages more effective exploration and
prevents early convergence to suboptimal policies, a limitation often seen
with DDPG’s deterministic approach. Additionally, SAC’s twin Q-networks
mitigate the overestimation bias that can affect DDPG’s value updates, leading
to more stable learning. The automatic tuning of the temperature parameter in
SAC further simplifies the delicate balance between exploration and
exploitation, reducing the need for meticulous hyperparameter adjustments.
Consequently, these features render SAC generally more sample-efficient and
robust, particularly in complex and continuous action spaces.

During training, experience for training is stored in a replay buffer and
sampled randomly to update the networks. All data in the training set are in
the replay buffer at the beginning. For offline SAC, the replay buffer is
unchanged, while online SAC alternates between collecting experience by
interacting with the environment and updating the networks with the replay
buffer. And offline SAC only uses the surrogate model trained with the
training set instead of the real environment to collect experience. The policy
network is updated to maximize the expected return, considering both the
Q-value and the entropy term. The critic networks are updated to minimize the
distance between their Q-value predictions and the target Q-values. SAC also
employs a target critic network for the critic networks, which are slowly
updated with the weights of the main critic network to stabilize training. For
the inference, SAC uses the policy network to determine the action by
selecting the action with the highest probability.

In practice, to help the system approximate the target state accurately and
quickly, we need to include the distance between the states of every time step
and the target state in the reward. So the reward function of time step $t$,
state $u_{t}$, target state $u_{T}$ and action $\mathbf{w}_{t}$ here is
defined as

$\displaystyle
r(t,\mathbf{u}_{t},y_{T},\mathbf{w}_{t})=-\int_{\Omega}|\mathbf{u}_{t}-\mathbf{u}_{d}|^{2}\mathrm{d}\mathbf{x}-{\alpha}\int_{\Omega}|\mathbf{w}_{t}|^{2}\mathrm{d}\mathbf{x},$

where $\Omega$ is the space domain and $\alpha$ is the weight of energy. We
take the Adam optimizer [28] to train the networks and update the temperature
parameter. The detailed values of hyperparameters are provided in Table 8.

###  F.3 Supervised Learning

The paper [23] proposes a supervised-learning-based control algorithm that
takes a neural operator as a surrogate model to solve control problems. It
contains two stages. In the first stage, we take a neural operator to learn
the physical constraint as Appendix H. The three CNNs respectively reconstruct
$u$, reconstruct $\mathbf{w}$ and learn the transition from $u_{t}$ to
$u_{t+1}$. More details are in Appendix H. In the second stage, these three
neural networks are used as surrogate models to calculate the gradient of the
objective function with respect to the control input. We consider the control
$\mathbf{w}$ as a learnable parameter and update it with the gradient.

To enhance the accuracy, we adopt the LBFGS optimizer [37], which is more
accurate while slower than the Adam optimizer. We record the hyperparameters
of the second stage in Table 9.

Table 9: Hyperparameters of the second stage of the 1D supervised learning
method.

Hyperparameter name | Value
---|---
Hyperparameters for 1D Burgers’ equation control: (full observation partial control) |
Learning rate of $\mathbf{w}$ updating | 0.1
Number of epochs | 300
Weight of objective function loss | 500
Weight of reconstruction loss | 0.03
Termination tolerance on first order optimality of LBFGS optimizer | $4\times 10^{-7}$
Termination tolerance on parameter changes LBFGS optimizer | $4\times 10^{-7}$
Hyperparameter name | Value
Hyperparameters for 1D Burgers’ equation control: (partial observation partial/full control) |
Learning rate of $\mathbf{w}$ updating | 0.1
Number of epochs | 300
Weight of objective function loss | 50000
Weight of reconstruction loss | 3
Termination tolerance on first order optimality of LBFGS optimizer | $4\times 10^{-7}$
Termination tolerance on parameter changes LBFGS optimizer | $4\times 10^{-7}$

###  F.4 BC

The Behavior Cloning (BC) algorithm [50] is a fundamental imitation learning
method. BC aims to learn policies directly from expert demonstrations,
leveraging supervised learning to map states to actions. This approach
bypasses the need for exploration typically required in reinforcement learning
by directly mimicking the behavior observed in the provided demonstrations. BC
does not require interaction with the environment during training, which
simplifies the learning process and reduces the computational resources
needed.

In Behavior Cloning, the policy network is trained using supervised learning
techniques, where the objective is to minimize the difference between the
predicted actions and the actions taken by the expert in the training data.
The loss function is typically the mean squared error between the predicted
and expert actions. The training data, consisting of state-action pairs, is
collected from expert demonstrations and stored in a dataset. In the
inference, we evaluated the same objective function as the SAC, which is the
MSE between the final state after $T-1$ steps’ control and the target. For the
partially observed or partially controlled settings, we concatenate zeros in
the masked dimensions to make the same dimension as the input of the policy
network, but the zeros are not considered during the calculation of the
metrics. The detailed values of hyperparameters are provided in Table 10.

Table 10: Hyperparameters of 1D BC. The full observation partial control, partial observation full control, and partial observation partial control settings share the same hyperparameters. Hyperparameter name | Value
---|---
Hyperparameters for 1D Burgers’ equation control: |
Learning rate | $1\times 10^{-4}$
Training batch size | 512
Number of episodes | $5\times 10^{5}$
Size of replay buffer | $2\times 10^{6}$
Number of layers of policy networks | 2
Number of hidden dimensions of policy networks | 1024
Activation function | ReLU

###  F.5 BPPO

The Behavior Proximal Policy Optimization (BPPO) algorithm [75] is an advanced
reinforcement learning method that builds upon the strengths of Proximal
Policy Optimization (PPO) while incorporating elements from behavior cloning.
BPPO is an offline algorithm that monotonically improves behavior policy in
the manner of PPO. Owing to the inherent conservatism of PPO, BPPO restricts
the ratio of learned policy and behavior policy within a certain range,
similar to the offline RL methods which make the learned policy close to the
behavior policy. By leveraging the inherent conservatism of online on-policy
algorithms, BPPO addresses the overestimation issue commonly encountered in
offline RL settings. The algorithm starts by estimating a behavior policy
using behavior cloning and then iteratively improves a target policy using the
Proximal Policy Optimization (PPO) objective with a behavior constraint.
Through this process of policy improvement, advantage estimation, and policy
update, BPPO aims to refine the target policy while ensuring it remains close
to the behavior policy. By combining the strengths of online on-policy methods
with tailored offline RL techniques, BPPO demonstrates promising results on
the D4RL benchmark, outperforming state-of-the-art offline RL algorithms.

During the training, BPPO first initializes the behavior policy $\pi_{\beta}$
and target policy $\pi_{\theta}$. Then, it estimates the behavior policy
$\pi_{\beta}$ using behavior cloning to mimic the behavior demonstrated in the
offline dataset. Next, the target policy $\pi_{\theta}$ is optimized by the
PPO objective with a behavior constraint, ensuring that the target policy
remains close to the behavior policy. After that, it estimates the advantage
function $A^{\pi_{\beta}}$ using the behavior policy $\pi_{\beta}$ to evaluate
the quality of actions taken by the target policy. Finally, updating the
target policy by maximizing the PPO objective with that estimated advantage
function, and adjusting the policy parameters to improve performance. In the
implementation, a state value network and Q value network are pre-trained
using the state, action, and reward from the offline dataset.

In practice, to help the system approximate the target state accurately and
quickly, we need to include the distance between the states of every time step
and the target state in the reward. So the reward function of time step $t$,
state $u_{t}$, target state $u_{d}$ and action $\mathbf{w}_{t}$ here is
defined as

$\displaystyle
r(t,\mathbf{u}_{t},\mathbf{u}_{d},\mathbf{w}_{t})=-\int_{\Omega}|\mathbf{u}_{t}-\mathbf{u}_{d}|^{2}\mathrm{d}\mathbf{x}-{\alpha}\int_{\Omega}|\mathbf{w}_{t}|^{2}\mathrm{d}\mathbf{x},$

where $\Omega$ is the space domain and $\alpha$ is the weight of energy. We
take the Adam optimizer [28] to train the networks and update the temperature
parameter. The detailed values of hyperparameters are provided in Table 11.

Table 11: Hyperparameters of 1D BPPO. The full observation partial control, partial observation full control, and partial observation partial control settings share the same hyperparameters. Hyperparameter name | Value
---|---
Hyperparameters for 1D Burgers’ equation control: |
State value network: |
Learning rate of value network | $1\times 10^{-4}$
Steps of value network | $2\times 10^{6}$
Number of layers of value network | 3
Batch size of value network | 512
Number of hidden dimensions of value network | 512
Q value network: |
Learning rate of Q network | $1\times 10^{-4}$
Steps of Q network | $2\times 10^{6}$
Number of layers of Q network | 2
Batch size of Q network | 512
Number of hidden dimensions of Q network | 1024
Target Q network updates per step | 2
Soft update factor | 0.005
Discount factor for reward | 0.99
Behavior cloning: |
Learning rate of BC | $1\times 10^{-4}$
Training batch size of BC | 512
Number of episodes of BC | $5\times 10^{5}$
BPPO: |
Number of episodes of BPPO | $1\times 10^{2}$
Number of layers of policy networks | 2
Number of hidden dimensions of policy networks | 1024
Learning rate of BPPO | $1\times 10^{-5}$
Training batch size of BPPO | 512
Clip ratio of BPPO | 0.25
Weight decay factor | 0.96
Weight of advantage function | 0.9
Size of replay buffer | $2\times 10^{6}$
Activation function | ReLU

##  Appendix G 2D Baselines

###  G.1 MPC and SL

Model Predictive Control (MPC) [59] is a control strategy that solves an
optimization problem repeatedly to determine the optimal control inputs for a
dynamic system. It operates over a finite prediction horizon, optimizing a
cost function and applying only the first control action. In the 2D jellyfish
movement control problem, MPC uses the control sequences $\mathbf{w}$ and the
fluid states as internal state variables. Without the need to train a control
agent model, MPC relies on 2D surrogate models mentioned in Appendix H to
estimate future states based on the current state and control. We use
backpropagation to compute the gradient, update the control action sequences,
and optimize the control objective $\mathcal{J}$ in Eq. (19). For every time
step, we get the optimized sequences from this time step forward in this way,
and only the first control sequence of the optimized control sequences is
applied. Compared with MPC, the Supervised learning (SL) method [23] only
optimizes the entire control sequences and employs the entire sequences.

MPC is an optimization technique that aims to optimize the control of complex
dynamic systems by considering future predictions. This approach can optimize
performance measures over a future time horizon, handle systems with multiple
variables and constraints, adapt to changes in the system behavior, and offer
good performance even in the presence of nonlinearity. Nevertheless, using MPC
comes with a high cost, both in terms of computational resources and time.
Additionally, adapting the optimization hyperparameters for MPC can be a
challenging task. In our experiment, both MPC and SL face difficulties when
trying to generate smooth opening angle control curves, even when constraints
$R(\hat{\mathbf{w}})$ are included in the optimization objective
$\mathcal{J}$. In the case of multi-objective optimization problems, it
becomes even more challenging for them to simultaneously achieve both the
higher speed (bigger $\bar{v}$) and the control curves smoothness (smaller
$R(\hat{\mathbf{w}})$).

###  G.2 SAC

For the 2D case, the algorithm and basic architecture of SAC are the same as
the 1D case in Appendix F.2. When designing the reward function for the 2D
jellyfish movement control, we find the periodic condition of the opening
angle curves is hard to constrain. So to satisfy the periodic condition
better, we include the distance between $\mathbf{w}_{t}$ of every time step
$t$ and $\mathbf{w}_{0}$. Also, we both consider the squared and absolute
error of $(\mathbf{w}_{t}-\mathbf{w}_{0})$ since they respectively constrain
the periodic condition when $(\mathbf{w}_{t}-\mathbf{w}_{0})$ is large and
small. As a result, the reward function of time step $t$, force $F_{t}$,
opening angle $(\mathbf{w}_{t-1},\mathbf{w}_{t})$ and condition angle
$\mathbf{w}_{0}$ is defined as

$\displaystyle
r(t,\mathbf{w}_{t-1},\mathbf{w}_{t},\mathbf{w}_{0})=(T-t)*F_{t}-\lambda_{1}(\mathbf{w}_{t}-\mathbf{w}_{t-1})^{2}-\lambda_{2}((\mathbf{w}_{t}-\mathbf{w}_{0})^{2}+|\mathbf{w}_{t}-\mathbf{w}_{0}|),$

where $\lambda_{1}$, $\lambda_{2}$ are weights of different terms. We take the
Adam optimizer [28] to update the weights of networks and the temperature
parameter as in the 1D experiment. The hyperparameters are reported in Table
12. In particular, we take the best checkpoint to evaluate the final
performance, thus the actual number of training episodes for each setting
ranges from about 100 to about 300.

Table 12: Hyperparameters of 2D SAC. The full observation and partial observation settings share the same hyperparameters. Hyperparameter name | Value
---|---
Hyperparameters for 2D Jellyfish movement control: |
Weight of the constraint of periodic condition $\beta$ | 0.001
Discount factor for reward | 0.5
Target smoothing coefficient | 0.05
Learning rate of critic loss | 0.0003
Learning rate of entropy loss | 0.0003
Learning rate of policy loss | 0.0003
Training batch size | 2048
Number of episodes | 350
Number of model updates per simulator step | 20
Value target updates per step | 15
Size of replay buffer | 11400001
Number of trajectories interacted with the environment per step | 5
Activation function | ELU
Clipping’s range of policy network’s standard deviation output | $\left[e^{-5},e^{-2}\right]$

###  G.3 BC

For the 2D task, the algorithm and basic architecture of BC are the same as
the 1D case in Appendix F.4. In Behavior Cloning, the policy network is
trained using supervised learning techniques, where the objective is to
minimize the difference between the predicted actions and the actions taken by
the non-expert in the training data. The loss function is typically the mean
squared error between the predicted actions and actions in training data. The
training data, consisting of state-action pairs, is collected from non-expert
random demonstrations and stored in a dataset. In the inference, we evaluated
the same objective function as the SAC. The detailed values of hyperparameters
are provided in Table 13.

Table 13: Hyperparameters of 2D BC. The full observation and partial observation settings share the same hyperparameters. Hyperparameter name | Value
---|---
Hyperparameters for 2D Jellyfish movement control: |
Learning rate | $1\times 10^{-4}$
Training batch size | 512
Number of episodes | $5\times 10^{3}$
Size of replay buffer | $2\times 10^{6}$
Activation function | ELU
Clipping’s range of policy network’s standard deviation output | $[-5,-2]$

###  G.4 BPPO

For the 2D task, the algorithm and basic architecture of BPPO are the same as
the 1D case in Appendix F.5. When designing the reward function for the 2D
jellyfish movement control, we find the periodic condition of the opening
angle curves is hard to constrain. So to satisfy the periodic condition
better, we include the distance between $\mathbf{w}_{t}$ of every time step
$t$ and $\mathbf{w}_{0}$. Also, we both consider the squared and absolute
error of $(\mathbf{w}_{t}-\mathbf{w}_{0})$ since they respectively constrain
the periodic condition when $(\mathbf{w}_{t}-\mathbf{w}_{0})$ is large and
small. As a result, the reward function of time step $t$, force $F_{t}$,
opening angle $(\mathbf{w}_{t-1},\mathbf{w}_{t})$ and condition angle
$\mathbf{w}_{0}$ is defined as

$\displaystyle
r(t,\mathbf{w}_{t-1},\mathbf{w}_{t},\mathbf{w}_{0})=(T-t)*F_{t}-\lambda_{1}(\mathbf{w}_{t}-\mathbf{w}_{t-1})^{2}-\lambda_{2}((\mathbf{w}_{t}-\mathbf{w}_{0})^{2}+|\mathbf{w}_{t}-\mathbf{w}_{0}|),$

where $\lambda_{1}$, $\lambda_{2}$ are weights of different terms. The
hyperparameters are reported in Table 14.

Table 14: Hyperparameters of 2D BPPO. The full observation and partial observation settings share the same hyperparameters. Hyperparameter name | Value
---|---
Hyperparameters for 2D jellyfish movement control: |
State value network: |
Learning rate of value network | $1\times 10^{-4}$
Steps of value network | $2\times 10^{3}$
Number of layers of value network | 3
Batch size of value network | 512
Q value network: |
Learning rate of Q network | $1\times 10^{-4}$
Steps of Q network | $2\times 10^{3}$
Number of layers of Q network | 3
Batch size of Q network | 512
Target Q network updates per step | 2
Soft update factor | 0.005
Discount factor for reward | 0.99
Behavior cloning: |
Learning rate of BC | $1\times 10^{-4}$
Training batch size of BC | 512
Number of episodes of BC | $5\times 10^{3}$
BPPO: |
Number of episodes of BPPO | $1\times 10^{2}$
Number of layers of policy networks | 3
Learning rate of BPPO | $1\times 10^{-5}$
Training batch size of BPPO | 512
Clip ratio of BPPO | 0.25
Weight decay factor | 0.96
Weight of advantage function | 0.9
Size of replay buffer | $2\times 10^{6}$
Activation function | ELU

##  Appendix H Surrogate Models

###  H.1 1D Surrogate Model

For the control problem of 1D Burgers’ equation, our 1D surrogate model is
based on the previous paper [23], which uses 2 autoencoders to model dynamics
in the latent space. The neural simulator architecture and training details
are shown in Table 15.

Table 15: Hyperparameters of 1D surrogate model. Hyperparameter name | Full observation, | Partial observation, | Partial observation,
---|---|---|---
partial control | full control | partial control
Autoencoder of state
Convolution kernel size | 5 | 5 | 5
Convolution padding | 2 | 2 | 2
Activation function | ELU | ELU | ELU
Latent vector size | 256 | 128 | 128
Autoencoder of force
Convolution kernel size | 5 | 5 | 5
Convolution padding | 2 | 2 | 2
Activation function | ELU | ELU | ELU
Latent vector size | 256 | 256 | 256
Training
Training batch size | 5100 | 5100 | 5100
Optimizer | Adam | Adam | Adam
Learning rate | 1e-3 | 1e-3 | 1e-3
Training epochs | 500 | 500 | 500
Learning rate scheduler | cosine annealing | cosine annealing | cosine annealing

###  H.2 2D Force Models

Dataset. In 2D jellyfish movement control experiments, we train a force
surrogate to approximate the computation of the average speed of the jellyfish
for the guidance of inference, which is implemented by a neural network and is
thus differentiable. The training data consists of pressure, boundary, and
force data in the training trajectories. Each training trajectory amounts to
$\tilde{T}=40$ training samples. Therefore, we have 1.2 million training
samples and 4 thousand testing samples in total.

Model. The model’s input contains pressure, boundary mask, and offsets with
shape $4\times 64\times 64$ at a certain time step, and the output is the
corresponding forces of $x$ and $y$ directions. The model architecture is the
down-sampling part of a U-Net [56] that embeds the input features into a
512-dimensional hidden representation; then we use a linear function with
output dimension two to output forces.

Training. We use MSE (mean squared error) loss between the ground truth and
predicted forces to train the force surrogate model. The optimizer is Adam
[28]. The batch size is 64. The model is trained for 10 epochs. The learning
rate starts from $1\times 10^{-4}$ and multiplies a factor of 0.1 every three
epochs. After training, the relative $l_{2}$ test error is 0.4%.

###  H.3 2D Boundary Mask and Offsets Updater

Dataset. In 2D jellyfish movement control experiments, we train a boundary
mask and offsets updater surrogate to approximate the transition of boundary
mask and offsets from time step $0$ to $t$. Thus each training trajectory
amounts to $\tilde{T}-1=39$ training samples.

Model. The input is the boundary mask and offsets at time step $0$ with shape
$3\times 64\times 64$, and the difference of the opening angle from $0$ to
$t$. The output is the boundary mask and offsets at time step $t$ with shape
$3\times 64\times 64$. The model architecture is the U-Net [56] with
additional scalar input, similar to the denoising network in DDPM [18], where
the input scalar diffusion step is replaced by the angle difference in our
model.

Training. We use MSE (mean squared error) loss between the ground truth and
predicted boundary mask and offsets to train this surrogate model.
Hyperparameters of training are the same as those of the force surrogate
model.

###  H.4 2D Simulator

Dataset. In 2D jellyfish movement control experiments, we need to train a
surrogate model as a solver of the physical dynamics for the baseline methods
like SAC (online) and MPC, because of their iterative nature. Conversely, our
diffusion method does not need this surrogate model. This model approximates
the transition of states under the boundary condition from time step $t$ to
$t+1$. Thus each training trajectory amounts to $\tilde{T}-1=39$ training
samples. We train two versions of this model for full/partial observation
settings.

Model. This surrogate model is also implemented by the U-Net [56]. The model
input is the states, boundary mask, and offsets at time $t$ with shape
$6\times 64\times 64$ for the full observation setting and $4\times 64\times
64$ for the partial observation setting. The output is the predicted states at
time step $t+1$, with shape $3\times 64\times 64$ for the full observation
setting and $1\times 64\times 64$ for the partial observation setting.

Training. We use MSE loss between the ground truth and predicted states to
train this surrogate model. Hyperparameters of training are the same as those
of the force surrogate model.

###  H.5 DiffPhyCon for 1D Burgers’ Equation

As DiffPhyCon can jointly generate the trajectory $\mathbf{u}$ and control
sequence $\mathbf{w}$, we demonstrate the MSE of the diffused trajectory
$\mathbf{u}$ from the trajectory computed by the ground-truth solver
$\mathbf{u}_{\text{ground truth}}$ given the diffused control sequence
$\mathbf{w}$. This should be differentiated from $\mathcal{J}$ reported in
Table 1 which is the MSE of the trajectory computed by the ground-truth solver
given $\mathbf{w}$ from the target trajectory $\mathbf{u}^{*}$. The former MSE
is denoted as $\mathcal{E}\coloneqq\|\mathbf{u}-\mathbf{u}_{\text{ground
truth}}\|$ in Figure 12.

As can be observed from Figure 12, there is a clear correlation between
$\mathcal{J}$ and $\mathcal{E}$. This is because, the higher the prediction
error is, the less accurate the solution to the inverse control problem would
be. The correlation proves that our DiffPhyCon learns the dynamics of the
system (i.e., the dependency of $\mathbf{u}$ on $\mathbf{w}$), based on which
the control sequence is generated. In partial observation settings, the
randomness is larger than in the full observation setting, which demonstrates
the challenge of controlling partially observed systems. The reference line
denotes $\mathcal{J}=\mathcal{E}$ where the control objective completely
depends on the prediction accuracy. In Figure 12 (b), samples cluster tightly
around the line, which implies DiffPhyCon can find near-optimal control
sequence based on the predicted trajectory in the corresponding level of
prediction accuracy. This ability is remarkable in that DiffPhyCon does not
rely on explicit gradient descent like those in the adjoint method and SL.

Figure 12: $\mathcal{J}$ and state MSE $\mathcal{E}$ of different samples in
different settings.

###  H.6 2D Jellyfish Movement

In this subsection, we record extensive results of 2D jellyfish movement
control. To further compare baselines with DiffPhyCon in more settings, we
change $\zeta$ in $\mathcal{J}$ and results of $\zeta=500$, $\zeta=1000$ and
$\zeta=2000$ are listed in Table 16 and Table 17. From the results, it is easy
to observe that DiffPhyCon outperforms other baselines in diverse settings
with different objectives.

We also present more simulation results of our method in Figure 13. Each line
represents an example from the test dataset. We plot five snapshots of
boundary and fluid field for each example.

Table 16: Full observation setting of 2D jellyfish movement control. Bold font
denotes the best model, and underline denotes the second best model.

| $\zeta=500$ | $\zeta=1000$ | $\zeta=2000$
---|---|---|---
|  $\bar{v}$ $\uparrow$ |  $R(\mathbf{w})$ $\downarrow$ |  $\mathcal{J}$ $\downarrow$ |  $\bar{v}$ $\uparrow$ |  $R(\mathbf{w})$ $\downarrow$ |  $\mathcal{J}$ $\downarrow$ |  $\bar{v}$ $\uparrow$ |  $R(\mathbf{w})$ $\downarrow$ |  $\mathcal{J}$ $\downarrow$
MPC | -38.21 | 0.1743 | 212.49 | 25.72 | 0.0112 | 109.17 | 44.24 | 0.0755 | 31.29
SL | 90.09 | 0.2570 | 166.89 | -76.94 | 0.1286 | 205.57 | -62.96 | 0.0648 | 127.8
SAC (pseudo-online) | -134.13 | 0.0204 | 154.49 | -166.96 | 0.0069 | 178.14 | -309.18 | 0.0040 | 313.19
SAC (offline) | -159.83 | 0.0119 | 171.73 | -158.66 | 0.0069 | 165.58 | -278.52 | 0.0051 | 283.58
BC | 15.61 | 0.0589 | 43.33 | 30.48 | 0.0629 | 32.44 | 33.5 | 0.0677 | 34.23
BPPO | 88.64 | 0.0776 | -11.07 | 107.67 | 0.0867 | -20.93 | -20.12 | 0.0621 | 82.27
DiffPhyCon-lite | 151.57 | 0.0939 | -57.69 | 95.04 | 0.0746 | -20.47 | 150.53 | 0.0923 | -58.23
DiffPhyCon | 310.02 | 0.2344 | -75.64 | 279.87 | 0.2058 | -74.11 | 270.56 | 0.1755 | -95.08

Table 17: Partial observation setting of 2D jellyfish movement control. Bold
font denotes the best model, and underline denotes the second best model.

| $\zeta=500$ | $\zeta=1000$ | $\zeta=2000$
---|---|---|---
|  $\bar{v}$ $\uparrow$ |  $R(\mathbf{w})$ $\downarrow$ |  $\mathcal{J}$ $\downarrow$ |  $\bar{v}$ $\uparrow$ |  $R(\mathbf{w})$ $\downarrow$ |  $\mathcal{J}$ $\downarrow$ |  $\bar{v}$ $\uparrow$ |  $R(\mathbf{w})$ $\downarrow$ |  $\mathcal{J}$ $\downarrow$
MPC | -179.85 | 0.2534 | 433.3 | -150.51 | 0.1791 | 329.59 | -116.75 | 0.1397 | 256.46
SL | 44.04 | 0.2895 | 245.43 | -102.98 | 0.1188 | 221.79 | -152.45 | 0.0661 | 218.53
SAC (pseudo-online) | 21.87 | 0.0798 | 57.96 | -153.09 | 0.0057 | 158.82 | -206.43 | 0.0048 | 211.2
SAC (offline) | -149.32 | 0.0156 | 164.93 | -206.21 | 0.0058 | 211.96 | -258.96 | 0.0029 | 261.87