Generative PDE Control

Authors: long  wei
peiyan  hu
ruiqi  feng
yixuan  du
tao  zhang
rui  wang
yue  wang
zhi-ming  ma
tailin  wu
Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-02-08 05:53:10

Abstract: Controlling PDE is a fundamental task across science and engineering. Classical techniques for PDE control tend to be computationally demanding and recent deep learning-based approaches often struggle to optimize long-term control sequences. In this work, we introduce Diffusion generative PDE Control (DiffConPDE), a new class of method to address the PDE control problem. DiffConPDE excels by simultaneously minimizing both the learned generative energy function and the predefined control objectives across the entire trajectory and control sequence. Moreover, we enhance DiffConPDE with prior reweighting, enabling the discovery of control sequences that significantly deviate from the training distribution. We test our method in 2D jellyfish movement in a fluid environment and 1D Burgers' equation control. Our method consistently outperforms baselines. Notably, DiffConPDE unveils an intriguing fast-close-slow-open pattern observed in the jellyfish, aligning with established findings in the field of fluid dynamics.

Main paper: 
 introduction: 
 The PDE control problem injects time-variant signals to steer evolution and optimize specific objectives for a PDE. It is a fundamental task with applications, such as including controlled nuclear fusion Degrave et al. (2022) , fluid control Holl et al. (2020) , and chemical engineering Christofides & Chow (2002) . Controlling high-dimensional, complex PDE systems in an efficient way presents three significant challenges. (1) Simulation challenge. The physical system is typically high-dimensional and highly nonlinear. (2) Optimization challenge. The PDE control task requires to optimize over a highly non-convex, potentially high-dimensional control sequence on top of the physical simulation.(3) Partial observation and control. In practical applications, our ability to observe or exert control over the physical system is often constrained to a limited subset of its components.To tackle PDE control problems, various techniques have been proposed. Classical methods typically adopt numerical simulation, among which the adjoint method Lions (1971) is widely applied but computationally expensive. Regarding traditional control methods, Model Predictive Control (MPC) Schwenzer et al. (2021) is limited by high computational costs and challenges in optimizing for a globally optimal solution. Recent advances in deep learning have impressive performance. But both supervised learning (SL) Holl et al. (2020) ; Hwang et al. (2022) and reinforcement learning (RL) Farahmand et al. (2017) ; Pan et al. (2018) may fall into adversarial modes Zhao et al. (2022b) and struggle to optimize long-term control sequence. See Appendix A for more related work.In this work, we introduce Diffusion generative PDE Control (DiffConPDE), a new class of method to address the PDE control problem. We take an energy optimization perspective to implicitly capture the constraints inherent in system dynamics through the diffusion model trained using system trajectory data and control sequences. This prevents the generated system dynamics from falling out of distribution, and offers an enhanced perspective of optimization overlong-term dynamics.An essential aspect of PDE control lies in its capacity to generate near-optimal controls. 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. We then introduce a prior reweighting technique to DiffConPDE, which constitutes the our second key contribution.We demonstrate the effectiveness of DiffConPDE via extensive experiments on 2D Jellyfish flapping and 1D Burgers' equation control problems. Our method consistently outperforms widely-applied classical control methods and state-of-the-art deep learning-based methods. 
method 2.1 problem setup: 
 We consider the following widely applicable PDE systems with external control signals: ∂u ∂t + F(u, ∇u, ∇ 2 u, w) = 0 (1) B(u, ∇u)| x∈∂Ω = 0(2) u| t=0 = u 0 .(3)Here, u(t, x) : [0, T ] × Ω → R du is the trajectory of the system defined on time range [0, T ] ⊂ R and spatial domain Ω ⊂ R D . Similarly, w(t, x) : [0, T ] × Ω → R dw is the external control signal. F is an operator that characterizes the dynamics of the PDE system. Eq. ( 2) is the boundary condition where B is a linear operator operating on the boundary ∂Ω. The initial condition is specified by u 0 (x). The control objective is J (u, w), whose minimization defines the PDE control task asEQUATIONwhere C(u, w) = 0 includes the PDE dynamics Eq. ( 1), the boundary condition Eq. ( 2) and the initial condition Eq. (3). It is crucial to emphasize that in numerous scenarios, explicit PDE expressions are unattainable, and we can only have access to observed control sequences and trajectory data, by which the PDE constraints C(u, w) are implicitly characterized.We model the PDE constraints as a parameterized energy-based model (EBM) E θ (u, w, c) which characterizes the distribution p(u, w|c) conditioned on conditions c of the PDE by the correspondence p(u, w|c) ∝ exp(-E θ (u, w, c)). Then the problem Eq. ( 4) can be converted to: u * , w * = arg min u,wEQUATIONwhere λ is a hyperparameter. This formulation optimizes u and w of all time steps simultaneously.Training. We train a diffusion model to estimate E θ . The denoising network (Ho et al., 2020) in the diffusion model is trained using L = E k∼U (1,K),(z,c)∼p(z,c),ϵ∼N (0,I) [∥ϵ -ϵ θ (z, c, k)∥ 2 2 ], where z denotes the concatenation of u and w and z := √ ᾱk z + √ 1 -ᾱk ϵ. The training dataset p(z, c) is either simulated using a numerical solver or collected from observation date in realistic systems.Algorithm 1 Inference for DiffConPDE 1: Require Diffusion models ϵ θ (z k , c, k) and ϵ ϕ (w k , c, k), control objective J (•), covariance matrix σ 2 k I, control conditions c, schedule ᾱk , hyperparameters λ, γ, K 2: Initialize optimization variables zK ∼ N (0, I)3: for k = K, . . . , 1 do 4: ξ1, ξ2 ∼ N 0, σ 2 k I 5: ẑk = (z k - √ 1 -ᾱk ϵ θ (z k , c, k))/ √ ᾱk 6: z k-1 = z k -η(ϵ θ (z k , c, k) + λ∇zJ (ẑ k )) + ξ1 7: w k-1 = w k-1 -η(γ -1)ϵ ϕ (w k , c, k) + ξ2 8: end for 9: u * , w * = z0 10: return u * , w *Control optimization. After ϵ θ is trained, Eq. ( 5)) can be optimized by the Langevin sampling procedure as follows. We start from an initial sample z K ∼ N (0, I), and run the processz k-1 = z k -η (∇ z (E θ (z k , c) + λJ (ẑ k )) + ξ iteratively, where ξ ∼ N 0, σ 2 k I .Here ẑk is the approximate noise-free z 0 estimated fromz k by ẑk = (z k - √ 1 -ᾱk ϵ θ (z k , c, k))/ √ ᾱk , where ᾱk := k i=1 α i .The gradient of E θ can then be replaced by our trained denoising network ϵ θ as follows:z k-1 = z k -η (ϵ θ (z k , c, k) + λ∇ z J (ẑ k )) + ξ, ξ ∼ N 0, σ 2 k I (6)where σ 2 k and η correspond to noise schedules and scaling factors used in the diffusion process, respectively. Iteration of this denoising process for k = K, K -1, ..., 1 yields a final solution z 0 = {u [1,T ],0 , w [0,T -1],0 } for the optimization problem Eq. ( 5).In PDE control, a critical challenge lies in obtaining control sequences superior to those in training datasets. To mitigate impact of the prior distribution, we propose a prior reweighting technique, which introduces an adjustable hyperparameter γ > 0, allowing for tuning the influence of this prior distribution. We denote the reweighted version of p(u, w|c) as p γ (u, w|c) := p(w|c) γ p(u|w, c)/Z = p(w|c) γ-1 p(u, w|c)/Z , where Z is a normalization constant and 0 < γ < 1. Then the energy model E (γ) (u, w, c) that learns -log (p γ (u, w|c)) can be decomposed asEQUATIONwhere E θ (u, w, c) follows Eq. ( 5), and E ϕ (w, c) is another diffusion model that learns -log p(u | c) which can be trained similarly to E θ . Then the optimization problem Eq. ( 5) can be transformed intoEQUATION)The optimization of the above problem leads to the following iteration in the diffusion model:EQUATION)where ξ 1 , ξ 2 ∼ N 0, σ 2 k I , andz k = [u k , w k ].The overview of DiffConPDE is illustrated in Figure 1 . When γ = 1, the model ϵ ϕ is not needed, and we denote this simplified version as DiffConPDE-lite. 
experiments: 
 We conduct experiments on the vital and challenging 1D Burger's Equation control and 2D jellyfish movement control problems. For 1D Burgers' equation, we use baselines: (1) the adjoint method Lions (1971) with 10 and 100 time steps; (2) PID Li et al. (2006) interacting with the ground-truth solver and surrogate model; (3) an RL method named SAC Haarnoja et al. (2018) with online, offline and pseudo-online versions; and (4) the Supervised Learning method (SL) Hwang et al. (2022) . Specifically, the online and pseudo-online version of SAC interacts with the solver and the surrogate model of the solver respectively, while the offline version only uses given data. For 2D jellyfish movement control, baselines include SAC (offline), SL, and MPC Schwenzer et al. (2021) . The adjoint method and PID are inapplicable to this data-driven task. Since interaction with the solver is time-consuming, SAC (online) is not applied. Baselines are detailed in Appendix D and E. EQUATION)Here ν is viscosity, and u 0 (x) is the initial condition. The objective of control is to minimizeJ actual := Ω |u(T, x) -u d (x)| 2 dx (12)while constraining the energy cost of the control sequence[0,T ]×Ω |w(t, x)| 2 dtdx (13)subject to Eq. ( 11), where u d (x) is the given target state.We select three 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 detailed in Appendix B.2. It should be noted that online methods (PID and online SAC) have unfair advantages. Nevertheless, we included them here for the sake of clarity.In Table 1 , we report results of the optimal control error J actual of different methods. It can be observed that our DiffConPDE delivers the best results when compared to all baselines except in PO-PC where PID (solver) has an unfair edge. DiffConPDE and DiffConPDE-lite show little performance gap since the prior distribution p(w|u 0 , u T ) is conditioned on both u 0 and u T , which fully determines the optimal w. Thus, p(w|u 0 , u T ) is intrinsically the optimal distribution and there is no need to suppress it, allowing DiffConPDE-lite to deliver satisfactory results.To compare the ability of different methods to optimize J actual with constrained energy cost, we compare the Pareto frontiers of different methods. As can be observed in Figure 2 , the Pareto frontiers of DiffConPDE is consistently among the best, achieving the lowest J actual for most settings of the energy budget. Interestingly, although PID (solver) has an unfair advantage due to its interaction with the solver, it only slightly outperforms our method in the PO-PC setting (c), and requires orders of magnitude larger energy to achieve a better J actual . Although SL performs well in full observation setting (b), it encounters difficulty in partial observation scenarios (a)(c). The results demonstrate DiffConPDE's advantage to generate near-optimal control sequences. Experiment settings. The task is to control the movement of a flapping jellyfish with two wings in a 2D fluid field. The dynamic of fluid follows the 2D incompressible Navier-Stokes Equation:EQUATION)where v is the fluid's 2D velocity, and p is the pressure, constituting the PDE state u = (v, p). The initial velocity condition is v 0 (x) and the kinematic viscosity is ν. The jellyfish's boundary can be parameterized by the opening angle w t of wings due to rigidity. Consequently, the control objective is to maximize its average moving speed v, under the energy cost constraint R(w) and the periodic constraint d(w T , w 0 ) of the movement:EQUATIONsubject to Eq. ( 14) and the boundary condition that velocity vanishes near the boundary. 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 u 0 (x) and the control sequence w(t, x) are both randomly generated, and then the states 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 w, its refreshing rate is 0.1 -1 , i.e., w(t, x), t ∈ [0.1k, 0.1(k + 1)], k ∈ {0, .., 9} does not change with t. Therefore, the data size of each trajectory is [10, 128] for the state u and [11, 128] for the force f . In all settings, the initial value u(0, x) is a superposition of two Gaussian functions u(0, x) = 2 i=1 a i e - (x-b i ) 2 2σ 2 i, where a i , b i , σ i are all randomly sampled from uniform distributions:a 1 ∼ U (0, 2), a 2 ∼ U (-2, 0), b 1 ∼ (0.2, 0.4), b 2 ∼ (0.6, 0.8), σ 1 ∼ U (0.05, 0.15), σ 2 ∼ U (0.05, 0.15).Similarly, the control sequence w(x, t) is also a superposition of 8 Gaussian functionsEQUATIONwhere each parameter is independently generated as follows:b 1,i ∼ U (0, 1), b 2,i ∼ U (0, 1), σ 1,i ∼ U (0.05, 0.2), σ 2,i ∼ U (0.05, 0.2), while a 1 ∼ U (-1.5, 1.5) and for i ≥ 2, a i ∼ U (-1.5, 1.5) or 0 with equal probabilities. u(t, x), (t ̸ = 0) is then numerically simulated (using the ground-truth solver) given u(0, x) and w(t, x) based on Eq. ( 11). The setting of the dataset generation is based on a previous work (Hwang et al., 2022) .We generated 90000 trajectories for the training set and 50 for the testing set. Each trajectory takes up 32KB space and the size of the dataset sums up to 2GB.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 illustrated in Figure 5 . 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 J in different settings are not directly comparable. Following are three different settings of our experiments.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 u in this setting and measure the J actual of model control.Specifically, u(t, x), x ∈ [ 1 4 , 3 4 ] is set to zero in the dataset during training and u 0 (x), x ∈ [ 1 4 , 3 4 ] is also set to zero during testing. In this partial observation settingΩ = [1, 1 4 ] ∪ [ 3 4 , 1]. Since no information in the central 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 u at x ∈ Ω.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 ofx ∈ [ 1 4 , 3 4 ]. Ω is still [0, 1],and 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 w as before, followed by setting the central 1 2 of w to zero. To compensate for the decreased control intensity so that the magnitude of u can be roughly comparable to the full control setting, we double the magnitude of w. During the evaluation, the output control sequence is also post-processed to be zero in x ∈ [ 1 4 , 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.The final setting is the combination of the previous two settings.Only Ω = [0, 1 4 ] ∪ [ 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.In addition to the introduced explicit guidance in the main text, conditioning is also widely used to guide sampling in diffusion models (Ho & Salimans, 2021; Shu et al., 2023) . When the control objective can be naturally expressed in a conditioning form, e.g., the generated trajectory u is required to coincide with a desired target u * , we can include u = u * as a condition in c such that the sampled trajectory u from diffusion models automatically satisfying u = u * . Overall, in our proposed DiffConPDE framework, the control objective J can be optimized either using the explicit guidance ∇J or guidance conditioning depending on the specific control objectives.Since the training of models ϵ ϕ ≈ ∇ w log p(w) and ϵ θ ≈ ∇ u,w log p(u, w) are essentially the same and the latter model is exactly DiffConPDE-lite, we will introduce DiffConPDE-lite first.In general, DiffConPDE-lite follows the formulation of (Ho et al., 2020) which is also described in the main text. The data of u and 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 u and w are inconsistent, we zero-pad them into the size of 16. Then, u and w are stacked as two channels and fed into the 2D DDPM model. A 2D UNet ϵ θ is used to learn to predict ϵ. 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(w[0,T -1] , u [1,T -1] | u 0 , u T )In summary, ϵ θ takes in the current trajectory u, control w, step k, u 0 and u T as input, and predicts the noise of u and w. Note that it is not trained to predict u 0 and 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 DiffConPDE-lite.In terms of implementation, DiffConPDE is simply adding ϵ ϕ (w) to ϵ θ (u, w) during inference as shown in Section 2.3, where ϵ θ is the output of the denoising network in DiffConPDE-lite while ϵ ϕ is a new denoising network that is trained to generate w following the dataset distribution. Therefore, we only describe the model of ϵ ϕ here.ϵ ϕ takes input of w, k as in the standard DDPM and u 0 , u T as guidance conditioning. The output of ϵ ϕ (w) is of the same shape as w, so it can be treated as a network learning to sample from p(w) := p(u, w)du The output of ϵ ϕ (w) at the locations of u is thus filled with zeros.Training During training, the u 0 and 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 B.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 DiffConPDE on the dataset with control being zero in x ∈ [ 1 4 , 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.Inference During inference, u 0 and u T are set to the target u 0 and u d so that the DDPM generates samples satisfying the PDE constraint that is also conditioned on the target (u d ) or the constraint (u 0 ). In the partial observation setting, the u 0 and u T drawn from the testing set are all filled zero at the unobserved locations x ∈ [ 1 4 , 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 ϵ θ (u, w) with ϵ θ (u, w) + (γ -1)ϵ ϕ (w). It is worth noting that ϵ θ (u, w) denoises u and w simultaneously while ϵ ϕ (w) only denoises w. In our experiments, we found that adding a schedule to the output of the w network is beneficial. The results in Table 1 are generated with a reverse Sigmoid schedule followingEQUATION)where β is defined as the noise schedule in (Ho et al., 2020) . The inference of DiffConPDE-lite is simply setting γ = 1, which neglects the effect of the model ϵ ϕ (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 J = w(x, t)dxdt to produce cost-limited control sequences that are shown in Figure 2 . The gradient of the external guidance is computed and added to ϵ k in Eq. ( 17). Note that we use the predicted clean sample ŵk and ûk at the k-th step to compute ∇J since they suffer less from being noisy and leading to deviated guidance. u k and w k are computed followingx 0 ≈ xk = 1 √ ᾱk x k - √ 1-ᾱk √ᾱk ϵ k where x represents u or w as in Ho et al. (2020) . In our experiments, we use a cosine scheduling of the external guidance, and thus the final predicted noise would be ϵ k + λα k ∇J where β k is the noise schedule in Ho et al. (2020) with the cosine schedule.Evaluation After generating the trajectory u and control sequence w, we feed the control sequence w into the ground-truth solver and simulate the final state u given the generated w and the initial condition u 0 directly drawn from the testing dataset. The solver is the same as the one used in data generation in Appendix B.1. Finally, we compute J actual following Eq. ( 12). 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.We use the Lily-Pad simulator Weymouth (2015) to generate the training and testing dataset. The resolution of the 2D flow field is set to be 128 × 128. Actually, 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 w in a 2D jellyfish control experiment.Each trajectory starts from the largest opening angle and follows a cosine curve periodically with period T ′ = 200. Trajectories differ in initial angle, angle amplitude, and phase ratio τ (the ratio between the closing duration and a whole pitching duration). For each trajectory, the initial angle w 0 is generated as follows: first, sample a random angle, called mean anglew (m) ∈ [20 • , 40 • ], then sample a random angle amplitude w (a) ∈ [10 • , min(w (m) , 60 • -w (m) )].The initial w 0 is set as w 0 = w (m) + w (a) . The phase ratio τ is randomly sampled from [0.2, 0.8]. The opening angle w t of step t decreases from w (m) + w (a) to w (m) -w (a) as t grows from 0 to τ T ′ ; then w t increases from w (m) -w (a) to w (m) + w (a) as t grows from τ T ′ to T ′ . Afterwards, w t varies periodically for t > T ′ . The range of w t is [w (m) -w (a) , w (m) + w (a) ] ⊂ [10 • , 60 • ]. This setting is similar to the study of the propulsive performance of jellyfish Kang et al. (2023) . 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 ′ = 200 to 3T ′ = 600 steps with step size 10 because the simulation from t = 0 to T ′ = 200 is for initialization of the flow field. Then each trajectory is saved as a 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 6 .Besides the positions of the boundary points of wings and the opening angles 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 PDE states (fluid field) by convolution neural networks. For each trajectory, this image-like boundary representation is compatible with PDE 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 (∆x, ∆y) between the cell center to the closest point on the boundary. For each trajectory, we save PDE states, opening angles, boundary points, boundary masks and offsets, and force data. They are specified as:• PDE states u: shape [ T , 3, 64, 64]. For each step, we save the states of the fluid field consisting of velocity in x and y directions and pressure. To save space, we downsample the resolution from 128 × 128 to 64 × 64.velocity: [ T , 2, 64, 64].pressure: [ T , 1, 64, 64].• opening angels w: shape [ T ]. For each step, we save the opening angle in radians.• boundary points: shape [ 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 x and y directions are shrunk to half (64/128) of the original values.• boundary mask and offsets b: [ T , 3, 64, 64]. For each step, we save the mask of merged wings with half coordinates of boundary points and offsets in both x and y directions. The resolution is 64 × 64, compatible with that of the states.mask: [ T , 1, 64, 64].offsets: [ T , 2, 64, 64]. • force: shape [ 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 = 50 testing trajectories. Trajectories differ in the above specified parameters w (a) , w (m) and τ . Each trajectory occupies about 2MB of storage and the total dataset costs about 100GB. To generate training samples, we use sliding time windows that contain T = 20 successive time steps of states and boundaries as a sample, which corresponds to T ′ = 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.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.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 [ 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. ( 14), which makes the traditional numerical solver no longer applicable to solve this PDE 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.We use a 3D U-Net as the backbone of our diffusion model. In this paper, the architecture of the 3D U-net we employed is inspired by Ho et al. (2022) . 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 3 . The DiffConPDE-lite method learns the denoising network of the joint distribution p(u, w|c) where u is PDE states, w is the opening angle, and the conditions c consist of the initial angle w 0 , the initial state u 0 and the final angle w T = w 0 . We adopt the 3D U-Net as the backbone. To make the opening angle (of shape [T ]), align with PDE 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. in the DiffConPDE-lite method. The denoising network of p(w|c) also adopts the 3D U-Net architecture. Its input size is the same as that of p(u, w|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 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.Training. We use the MSE (mean squared error) between model prediction and the Gaussian noise as the loss function. The batch size is chosen as 16 and the training involves 200,000 iterations. The learning rate starts from 1 × 10 -3 and multiplies a factor of at the 50000th and 150000th iterations.Training details are provided in Table 4 .Inference. The pipeline of inference is shown in Figure 7 . Both diffused variables u [0,T ] and 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 7 ) to update boundary mask and offsets b [0,T ],k for each denoising k. Specifically, at each time step t ∈ [0, T ], b t,k is estimated by the initial boundary mask and offsets b 0 , and the difference of opening angle w t,k -w 0 from time step 0 to t, which is presented in the right part (after "Denoise under guidance") of Figure 7 . Details about this surrogate model are presented in F.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 w t,k in inference because the estimated w t,k does not deviate from the normalized range of noisy free w t too much.Our method introduces two kinds of inference: DiffConPDE-lite and DiffConPDE. In DiffConPDElite, we only use the denoising network of the joint distribution p(u, w|u 0 , w 0 ) for inference, while in DiffConPDE, we use the additional denoising network of the marginal distribution p(w|u 0 , w 0 ) together with that of the joint distribution for inference. These two branches are plotted in the left part of Figure 7 , where the notation [u 0 ] × T means expand initial state 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 F.2. In denoising step k, its input consists of two parts: the first one is the noise-free state û[0,T ] estimated from u [0,T ],k ; the second one is the noise-free boundary mask and offsets b[0,T ],k estimated from noise-free ŵ[0,T ] by the surrogate model to update boundaries, where ŵ[0,T ] is also estimated from w [0,T ],k . 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 J in Eq. ( 15) is computed as a summation of force and R( ŵ[0,T ] ). We fix ζ = 1000 as a default setting in Eq. ( 15) because this value can achieve a balance between scales of the average speed and the regularizer R(w). We also study the Pareto performance of varying ζ in Table 2 . Then the gradients of the objective J in terms of û[0,T ] and ŵ[0,T ] are computed and used guidance. For DiffConPDE-lite, these gradients are substracted from[u [0,T ],k , w [0,T ],k ] to generate [u [0,T ],k-1 , w [0,T ],k-1 ].For DiffConPDE, an additional term of noise (γ -1)ϵ ϕ predicted from the denoising network of the marginal distribution p(w|u 0 , w 0 ) should also be subtracted, as shown in the upper left part of Figure 7 . The effect of this term is controlled by the scale of the hyperparameter γ.Evaluation. The inference outputs opening angles 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 v, energy cost R(w), and objective J are computed as metrics. The average speed v = 1T T 0 v t dt ≈ v 0 + 1 T T -1 t=1 (T -t)F t 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(w) = T -1 t=1 (w t+1 -w t ) 2 , where the control sequence w = (w 1 , • • • , w T ) represents the predicted opening angles. The periodic term d(w 0 , w T ) = max(|w T -w 0 | -ϵ, 0) is the constraint of periodic opening angles with a small threshold ϵ = 0.01. Traditional PID controllers are limited to SISO linear time-invariant system, to address these gaps we use an NN to automatically plan PID parameters based on current errors. Proportional Integral Derivative (PID) Li et al. (2006) 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 Slama et al. (2019) ; Ding et al. (2022) 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 8 , 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 Kingma & Ba (2014) to train the model. Detailed information can be found in Table 5 . 2004) is a mathematical technique used in the control problems of PDEs. It is widely utilized in many fields such as computational fluid dynamics, structural optimization, and machine learning.The adjoint method first solves the forward problem, which is the original system of differential equations, to obtain the system's state as a function of the parameters. Then it constructs a Lagrangian function that includes the objective function, the forward problem's equations, and the associated boundary or initial conditions, all coupled together using the adjoint variables. We can derive adjoint equations by applying variational calculus, and solve the adjoint equations to get the adjoint variables. Finally, the gradients of the objective function with respect to the parameters, which are used to update the parameters, can be obtained by combining the solutions of the forward and adjoint problems using the adjoint variables.Our implementation follows the previous work Hwang et al. (2022) . While the use of the adjoint method requires the form of the equations, it can be only applied to the 1D Burgers' equation control. Also, it cannot handle cases where states of systems are partial observation. For the partial control setting, we set the output control sequences of the adjoint method to zero where the control is not allowed to be applied. We choose 10 and 100 numerical time steps for the adjoint method. Accordingly, the discrete computation of energy takes dt = 0.1 and dt = 0.01.The Soft Actor-Critic (SAC) algorithm Haarnoja et al. (2018) 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. 2018), 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 w t here is defined asr(t, u t , y T , w t ) = - Ω |u t -u d | 2 dx -α Ω |w t | 2 dx,where Ω is the space domain and α is the weight of energy. We take the Adam optimizer Kingma & Ba (2014) to train the networks and update the temperature parameter. The detailed values of hyperparameters are provided in Table 6 .The paper Hwang et al. (2022) 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 PDE constraint as Appendix F. The three CNNs respectively reconstruct u, reconstruct w and learn the transition from u t to u t+1 . More details are in Appendix F. 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 w as a learnable parameter and update it with the gradient.To enhance the accuracy, we adopt the LBFGS optimizer Liu & Nocedal (1989) , which is more accurate while slower than the Adam optimizer. We record the hyperparameters of the second stage in Table 7 .where λ 1 , λ 2 are weights of different terms. We take the Adam optimizer Kingma & Ba (2014) to update the weights of networks and the temperature parameter as in the 1D experiment. The hyperparameters are reported in Table 8 . 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. Model. The input is the boundary mask and offsets at time step 0 with shape 3 × 64 × 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 × 64 × 64. The model architecture is the U-Net Ronneberger et al. (2015) with additional scalar input, similar to the denoising network in DDPM Ho et al. (2020) , 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.Dataset. In 2D jellyfish movement control experiments, we need to train a surrogate model as a solver of the PDE 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 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 Ronneberger et al. (2015) . The model input is the states, boundary mask, and offsets at time t with shape 6 × 64 × 64 for the full observation setting and 4 × 64 × 64 for the partial observation setting. The output is the predicted states at time step t + 1, with shape 3 × 64 × 64 for the full observation setting and 1 × 64 × 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.Performance of DiffConPDE is determined by the hyperparameter γ. Since diffusion models denoise gradually, we use a varying sequence of γ = {γ k } K k=1 to substract prior in DiffConPDE. Specifically, the schedule of γ is set as γ k = 1 -ξ • β K-k , k = 1, • • • , K, where ξ is a fixed coefficient to control the scale of γ and β = {β k } K k=0 is the schedule of variances of noise in DDPM Ho et al. (2020) . In our implementation, we use the sigmoid schedule of β Jabri et al. (2022) . The total number of inference steps is K = 1000. Thus we only need to tune ξ to examine the effect of γ. When ξ < 0, DiffConPDE is prone to restrict w within its prior distribution of training dataset in inference. When ξ > 0, DiffConPDE is more likely to generate new kinds of w beyond training ones. When ξ = 0, DiffConPDE degenerates to DiffConPDE-lite. In 2D experiments, We set default ξ = 0.3 and the corresponding γ 1 = 0.7 as we empirically find this value performs well and steadily. We present the performance of DiffConPDE on the 2D jellyfish movement control task under different γ in Figure 9 and Table 10 . We can observe that the performance increases along with decreasing of γ 1 . When γ 1 < 0.6, invalid generated control sequences emerge because the prior is largely overlooked. Thus the valid interval for γ of prior reweighting on this task is [0.6, 1.0]. It is interesting to find that when γ 1 > 1, the performance decreases. This may be caused by the strict constraint of the prior distribution of p(w, c), which results in generating control sequences similar to those from training datasets and thus not good.Although DiffConPDE solves the PDE control problems with outstanding performance, there are still several limitations that provide exciting opportunity for future works. Firstly, the inference process of the diffusion model, the foundational component of DiffConPDE, may benefit from enhanced efficiency by using e.g. distillation Salimans & Ho (2022) ; Li et al. (2023) . Secondly, the training of DiffConPDE is currently conducted in an offline fashion, lacking 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 DiffConPDE 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 for the subsequent steps based on the evolving state of the environment.Moving forward, we plan to expedite the inference process by drawing insights from relevant prior works, such as Salimans & Ho (2021) ; Meng et al. (2023) ; Li et al. (2023) . Additionally, we aim to build upon our existing trajectory knowledge to generalize future trajectories. Furthermore, we intend to incorporate solver interactions into our training methodology, enhancing the model's adaptability and effectiveness. 
 
 Appendix:  


Decision: The reviewers both clearly indicate that this paper is a clear accept. It proposes to use diffusion models for PDE Control and tests both 1D and 2D examples. The topic is very relevant to the workshop and I vote for acceptnace.

Reviews: review: title: title: Good paper, clear accept, low confidence
main_review: main_review: Summary: This paper introduces an approach to PDE control problems; the goal is to learn to control the PDE behavior. The DiffConPDE model learns under the condition that the explicit form of the PDE may not be known, rather what is known is observed control sequences and trajectory data. A generative denoising diffusion model is used to generate the trajectory, and an inference algorithm estimates the optimal control signal.

My research area is in solving PDEs with ML, not PDE control. I also am not familiar with the mathematics of generative diffusion models (sorry!), so I do not understand the meaning of equations (6) (9) and (10). For those reasons, my review comes with relatively low confidence.

This paper is well-written, well-motivated, clearly of interest to the workshop, and seems to achieve good results (both in the main text and the appendix). I am quite happy with this paper.

It is possible that there are flaws with this paper that I am not aware of (better baselines, prior work with similar ideas, etc). However, I don't have any reason to suspect these flaws exist.

I believe this paper is clearly in the top 50% of accepted papers. I am not familiar enough with this research area to decide that this paper falls into the top 15% of accepted papers, but I would not disagree if someone more familiar with the research area felt that it deserved such a rating.
title: title: Strong paper with strong empirical results, possible spotlight
main_review: main_review: # Summary
This work explores the use of diffusion models in the setting of controlled partial differential equation (PDE). In this setting, the goal is to minimize a control objective $\mathcal J (\mathbf u, \mathbf w)$ that is a function of both the trajectory of the system $\mathbf u$ and the external control signal $\mathbf w$. The dynamics of the trajectory $\mathbf u$ are governed by a PDE system that depends on the control signal $\mathbf w$ and initial and boundary conditions $\mathcal c$.

The authors propose that there is a generative model of $p(u, w \vert c)$. The negative log-density of this distribution is known as a parameterized energy-based model. The controlled PDE optimization problem can be recast as minimizing the negative log-density of $p(u, w \vert c)$ with $\mathcal J(u, w)$ as an added regularization term (Equation 5).

To solve the new optimization problem, the authors first train a diffusion model to learn $p(u, w \vert c)$. The key observation is that the original diffusion model presented in [1] yields an estimate of the gradient of $E_{\theta}$ at each step (resembling Langevin dynamics). With this observation, the authors modify the standard sampling step of a diffusion model to incorporate the gradient of the control objective $\mathcal J$. This leads to an optimization routine that targets the desired objective.

The authors also propose a scheme for minimizing the impact of choice of prior on the control sequences $p(w \vert c)$. This works by tempering the density of $p(w \vert c)^{\gamma}$ (I believe there was a typo in the text that I have corrected here). Similar to the $\beta$ parameter in the original diffusion paper, $\gamma$ also changes with each step of the diffusion. Specifically, $\gamma_{1} = 0.7$ and $\gamma_{K} = 1$, so the "prior component" of the target distribution becomes less diffuse in later steps of the diffusion.

The authors present results on a 2-d jellyfish control task in the main body. Their proposed method outperforms all other competing methods. The optimal strategy their method found was corroborated in other works as being an effective strategy for jellyfish movement.

# Review

Overall, with an interesting method and strong empirical results, this is a nice paper that should be a strong acceptance to the workshop. In fact, I believe this would be a good candidate for a spotlight paper, since it demonstrates a nice, effective use of deep learning to solve real challenges in numerical methods.

## Originality
This paper builds upon well-established diffusion models to facilitate controlled PDE optimization. [1] commented that the diffusion sampling process resembles Langevin dynamics and follows the gradient of the log probability density. This work exploits that fact by  modifying the diffusion sampling algorithm to incorporate the gradient of the control objective $\mathcal J$. To my knowledge, this is a new idea, though I am not deeply familiar with the diffusion model literature.

## Quality
The overall quality of the work is high for a workshop paper. The method is well-explained, and the empirical evaluation is comprehensive for a workshop paper. The empirical evaluation clearly demonstrates the effectiveness of the proposed method at an assigned task.

While the quality is high for a workshop paper, I have some remarks below that should be answered in a full paper.

1. The probability distribution of $u, w$ needs to be more carefully defined and characterized. The set $\{u: \mathcal C(u, w) = 0 \}$ will have Lebesgue measure 0 on the space where $u$ is defined ($[0, \mathcal{T}] \times \Omega$). As a result, the density of $p(u \vert w, c)$ is a point mass on this space. It does not seem that this has caused any practical issues in the current empirical setting. However, samples from the trained diffusion model almost certainly do not satisfy $\mathcal C(u, w)$, and this could pose problems in situations where the ground truth simulator is not available.

2. I would like to see more theoretical results about the convergence of the control optimization to a valid solution. From my comment above (1), the solution yielded by algorithm (1) will almost certainly not satisfy $C(u, w) = 0$. It would be beneficial to know the conditions under which the solution is arbitrarily close to satisfying $C(u, w) = 0$, even if these conditions are not met in practice. Knowing the failure modes of this methodology is important for knowing the limitations and addressing them in future work.

## Clarity
The writing is clear throughout. The visuals of the experiment are quite nice. A few concepts, such as the prior re-weighting scheme, required me to reference the appendix to get a full understanding. Specifically, in the main text, it is not made clear that $\gamma$ depends on the index $k$, similar to the parameter $\beta_{k}$. However, this is likely a consequence of the 4-page limit.

## Significance
This work provides a blueprint on how deep learning can solve real problems better than classical methods, which embodies the spirit of this workshop. I also believe that this approach could be useful in a broad range of applications, even beyond controlled PDES.

## General Feedback
Currently, you train the diffusion model to estimate $E_{\theta}$ by having it target the distribution with density proportional to $\exp(-E_{\theta}(u, w, c))$. What if, instead, you trained the diffusion model to target the distribution with density proportional to $\exp(-E_{\theta}(u, w, c) - \lambda \mathcal J(u, w))$? This will lead to more weight on samples with a lower $\mathcal J$.