From Uncertain to Safe: Conformal Adaptation of Diffusion Models for Safe PDE Control
Authors: Peiyan Hu
Xiaowei Qian
Wenhao Deng
Rui Wang
Haodong Feng
Ruiqi Feng
Tao Zhang
Long Wei
Yue Wang
Zhi-Ming Ma
Tailin Wu
Venue: ICML 2025 poster (2025)
Keywords: safe PDE control
PDE simulation
generative models
conformal prediction
Primary area: applications->chemistry_physics_and_earth_sciences
TL;DR: We propose Safe Diffusion Models for PDE Control, which introduces the uncertainty quantile as model uncertainty quantification to achieve optimal control under safety constraints through both post-training and inference phases.

Abstract: The application of deep learning for partial differential equation (PDE)-constrained control is gaining increasing attention. However, existing methods rarely consider safety requirements crucial in real-world applications. To address this limitation, we propose Safe Diffusion Models for PDE Control (SafeDiffCon), which introduce the uncertainty quantile as model uncertainty quantification to achieve optimal control under safety constraints through both post-training and inference phases. Firstly, our approach post-trains a pre-trained diffusion model to generate control sequences that better satisfy safety constraints while achieving improved control objectives via a reweighted diffusion loss, which incorporates the uncertainty quantile estimated using conformal prediction. Secondly, during inference, the diffusion model dynamically adjusts both its generation process and parameters through iterative guidance and fine-tuning, conditioned on control targets while simultaneously integrating the estimated uncertainty quantile. We evaluate SafeDiffCon on three control tasks: 1D Burgers' equation, 2D incompressible fluid, and controlled nuclear fusion problem. Results demonstrate that SafeDiffCon is the only method that satisfies all safety constraints, whereas other classical and deep learning baselines fail. Furthermore, while adhering to safety constraints, SafeDiffCon achieves the best control performance. The code can be found at https://github.com/AI4Science-WestlakeU/safediffcon.

From Uncertain to Safe: for Safe PDE Control
Conformal Adaptation of Diffusion Models
Peiyan Hu 12 Xiaowei Qian 13 Wenhao Deng²
Rui Wang" 14 Haodong Feng? Zhi-Ming Ma Tailin Wu
Ruiqi Feng² Tao Zhang?
Long Wei
Yue Wang
Abstract
1. Introduction
The application of deep learning for partial dif- ferential equation -onttaccnt control gaining increasing attention. However, existing methods rarely consider safety requirements cru- cial in real-world applications. To address this limitation, we propose Safe Diffusion Models for PDE Control (SafeDifiCon), which introduce the uncertainty quantile as model uncertainty quan- tification to achieve optimal control under safety constraints through both post-training and infer- ence phases. Firstly, our approach post-trains a pre-trained diffusion model to generate control sequences that better satisfy safety constraints while achieving improved control objectives via a reweighted diffusion loss, which incorporates the uncertainty quantile estimated using confor- mal prediction. Secondly, during inference, the diffusion model dynamically adjusts both its gen- eration process and parameters through iterative guidance and fine-tuning, conditioned on control targets while simultaneously integrating the esti- mated uncertainty quantile. We evaluate SafeDiff- Con on three control tasks: ID Burgers' equa- tion, 2D incompressible fluid, and controlled nuclear fusion problem. Results demonstrate that SafeDiffCon is the only method that satis- fies all safety constraints, whereas other classical and deep learning baselines fail. Furthermore, while adhering to safety constraints, SafeDiff- Con achieves the best control performance. The code can be found https: githhb.. com/ AI AI4Sciences weweeeeneeeeetttee West stttseesss/aeeiiieee ffcon.
The control of physical systems described by partial differ- ential equations (PDEs) is crucial in various scientific and engincering domains, including fluid dynamics (Hinze & Kunisch, 2001), controlled nuclear fusion (Schuster & Ari- ola, 2006), and mathematical finance (Soner, 2004). While traditional control algorithms have been extensively studied over the years (Li et al., 2006; Protas, 2008), advancements in neural networks have led the development of numerous deep learning- based approaches (Farahmand et al., 2017; Holl et al., 2020; Hwang et al., 2022). Among these, diffu- sion models have gained increasing attention to PDE control due to their ability to model high-dimensional and nonlinear data (Vahdat et al., 2022; Li et al., 2024), and have demon- strated impressive performance in various control tasks (Chi et al., 2023; Wei et al., 2024; Hu al., 2024).
However, existing deep learning- bbased based methods for PDE control generally overlook safety meaning ensuring con- trol sequences satisfy predefined constraints, mitigating risks, and preventing hazards (Dawson et al., 2022; Liu et al., 2023a). In real-world applications, controlling PDE systems particularly requires addressing safety concerns (Barros al., 1998; Argomedo et al., 2013). For instance, in fluid dynamics, minor errors can cause turbulence or structural damage, while in controlled nuclear fusion, violat- ing safety constraints may lead to catastrophic failures. This lack of safety enforcement remains a major bottleneck in ap- plying deep learning to scientific and engineering problems, posing a critical challenge for high-stakes applications.
Despite its importance, safe PDE control remains challeng- ing. First, ensuring safety requires preventing methods without formal guarantees from interacting with the envi- ronment, restricting the approach to an offline setting with pre-collected data. However, such data is often suboptimal and includes unsafe samples, creating significant gap be- tween the observed distribution and the near-optimal, safe distribution (Xu et al., 2022; Liu al., 2023a). Second, the method must resolve the inherent conflict between optimiz- ing control performance and satisfying safety constraints (e.g., aggressive control policies that reduce error may vio- late safety thresholds) (Liu et al., 2023a; Zheng et al., 2024).
"Equal contribution 'Work done as an intern at Westlake Uni- versity" "Academy of Mathematics and Systems Science, Chinese Academy of Sciences, Beijing, China 'School of Engineering, Westlake University, Hangzhou, China FFudan University, Shang- hai, China ZZhongguancun Academy, Beijing, China. Correspon- dence to: Tailin Wu <wutailin @ westlake.edu.coo edu.cn.>.
Proceedings of the n2nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

From Uncertain to Safe: Conformal Adaptation of Diffusion Models for Safe PDE Control
To address these challenges, we propose Safe Diffussion Mod- els for PDE Control (SafeDiffCon), method that adapts diffusion models to varying safety constraints by quantify- ing their uncertainty. In offline settings, training data are often sub-optimal and unsafe, which causes the distribu- tion learned by diffusion models to deviate significantly from the desired optimal and safe distribution. Inspired by conformal prediction (Vovk et al., 2005; Tibshirani et al., 2019), we propose quantifying diffusion models' safety uncertainty under distribution shifts and incorporating it through both post-training and inference phases. This un- certainty quantification, termed as uncertainty quantile, al- lows us to compute a conformal interval, within which the safety score of the diffusion model interacting with the real environment is expected to lic. By leveraging this uncer- tainty quantification for predicted safety, diffusion models can be optimized with clear safety optimization objec- tive, to satisfy various safety constraints by constraining their conformal intervals within safe boundaries. Specifi- cally, we incorporate uncertainty quantification into both the post-training and inference-time fine-tuning processes. Technically, we perform post-training with reweighted loss, where the weighting accounts for the uncertainty quantile of safety instead of the original predicted safety score. This refinement improves the model's control performance while enhancing its adherence to safety constraints, leading to the generation of more optimal and safer control sequences. To further ensure the safety of PDE control, we introduce inference-tine fine-tuning, where the diffusion model is it- eratively adjusted with guidance and fine-tuning based on specific control targets and uncertainty quantile improve the safety of the final outputs within few iterations.
(Holl al., 2020; Hwang et al., 2022), reinforcement learn- ing (RL) (Farahmand et al., 2017; Pan et al., 2018; Rabault et al., 2019), and physics-informed neural networks (PINNs) (Mowlavi & Nabi, 2023). Among these, PID is one of the earliest and most widely used method (Johnson & Moradi, 2005), known for its simplicity and effectiveness in reg- ulating physical systems; however, it faces challenges in parameter tuning and struggles with highly nonlinear or time- varying systems. Model Predictive Control (MPC) (Garcia et al., 1989; Findeisen & Allgöwer, 2002; Schwen- zer et al., 2021) is another well-known control strategy that optimizes a control sequence over a finite time horizon by solving an optimization problem at each time step, using a dynamic model of the system. As it requires solving an optimization problem in real-time each time step, may have a high computational cost and struggle to handle high- dimensional, complex systems effectivelyy Furthermore, the adjoint method (Protas, 2008) and PINNS (Mowlavi & Nabi, 2023) are also incorporated in PDE control, but they require an explicit form of the PDE. Currently, the diffusion model used in physical systems' control (Wei et al., 2024; Hu et al., 2024) integrates the learning of entire state trajec- tories and control sequences, enabling global optimization that incorporates the physical information learned by the model. However, it does not consider the important cases where safety constraints are required.
In the field of reinforcement learning (RL), several safe of- fline RL algorithms have emerged in recent years. CPQ (Xu et al., 2022) is the first practical safe offline RL method that assigns high costs to diiti--disttibution (OOD) and unsafe actions and updates the value function as well as the policy only with safe actions. COptiDICE (Lee et al., 2022) is aJCE-based method and corrects the stationary distribu- tion. CDT (Liu et al., 2023b) takes the decision transformer to solve safe offline RL problems as multi-objective opti- mization. More recent methods such as TREBI (Lin et al., 2023) and FISOR (Zheng al., 2024) address this problem through diffusion model planning, leveraging the capabil- ity of diffusion models to model high-dimensional data. However, none of these algorithms take into account the uncertainty of the predicted safety from the perspective of conformal prediction, nor do they address PDE-constrained scenarios.
Our main contributions are as follows: (1) We introduce safety constraints into deep learning- based control of PDE systems and propose our method SafeDiffCon. (2) To ad- dress significant distribution shifts, we quantify the safety uncertainty of diffusion models and design a conformal adaptation process, which incorporates post-training with reweighted loss and inference-time fine-tuning to promote safer and more optimal output distribution. (3) We design safe control tasks in the contexts of the 1D Burgers' equa- tion, 2D incompressible fluid, and the controlled nuclear fusion scenario, and evaluate various methods. Extensive experiments on these datasets demonstrate that SafeDiffCon is the only method successfully meet all safety constraints while achieving superior control objectives.
2.2. Conformal Prediction
Conformal prediction (Vovk et al., 2005) is a statistical framework that constructs prediction intervals guaranteed to contain the true label with a specified probability. Its validity could be compromised, however, by the violation of the core assumption of exchangeability due to distribution shifts in real-worte scenarios (Chernozhukov et al., 2018; Hendryeks et al., 2018). Recent studies (Maxime Cauchois & Duchi, 2024) have extended conformal prediction to accommodate
2. Related Work
2.1. Control of PDE Systems
The development of control methods physical systems critical across various scientific and engineering areas,
including PID (Li al., 2006), supervised learning (SL)
2

---

From Uncertain to Safe: Conformal Adaptation of Diffusion Models for Safe PDE Control
various distribution shifts. For example, Tibshirani et al. (2019) proposed weighted conformal prediction to handle covariate shift, where training and test data distributions differ. Podkopaev & Ramdas (2021) introduced reweighted conformal prediction and calibration techniques to address label shift using unlabeled target data. Adaptive conformal inference (Gibbs & Candes, 2021) provides valid prediction sets in online settings with unknown, time-varying distri- bution shifts without relying on exchangeability. Inspired by conformal prediction, we measure uncertainty based on the distribution of data generated by diffusion models for control problems. This is also the first time the perspective of conformal prediction has been applied to address safety control problems.
removal of ég from the prior distribution could generate a new sample that follows the data distribution p(x).
Recently, diffusion models (Wei et al., 2024; Hu et al., 2024) are applied to solve the control problem as in Eq. without the safety constraint s(u) S 80. For brevity, we only summarize the light version. These methods trans- form the physical constraint to a parameterized energy- based model (EBM) (u w) with the correspondence p(u, w) x exp( Egg(E, w)). Then the problem is con- verted to an unconstrained optimization over u and w for all physical time steps simultancously:
u w' arg min (Eo w) + I(u, w)], u.w
(2)
where is a hyperparameter. To optimize Egg a denoising
3. Preliminary
network Eg is trained to approximate Vu w) by the following loss:
3.1. Problem Setup
Ldifiusion
We consider the following safe control problem of PDE- constrained systems:
|||E - E Vva u, w] + v1 - āge, 1))1) (3)
where ãg I1 ag After Eø is trained, Eq. 2 can be optimized by sampling from an initial sample (u K w") N I). and iteratively running the following process'
w' = arg min J (u, w) s.t. C(u, w) = 0, s(u) S so
(1) where x(4,x)= [0, x Ω + Rdo is the system's state trajectory with dimension du and x(t,x) [0. 7] x Q + Rdd is the external control signal with dimension dw. They are both defined on the time range [0, C R and spatial domain Z C RDD (u, w) is the objective of the control problem, and C(u, w) = 0 is the PDE constraint.
)- - "Tx . <o([n", w*),k) (n (u= w* w=-1)== = =(n", w*) V G(a", wh) +E, E~ {~,N(0,02) of))
(4)
1. under the guidance of g = J for k = K,K Here [û w*] is the noise-free estimation of [u", w 1. [ denotes vector concatenation, and 77K and A are denoising schedules. The final sampling step yields the solution w for the optimization problem in Eq. 2.
As for the safety constraint, s(u) is the safety score and 80 is the bound of the safety score. We need to minimize the control objective while satisfying PDE constraints and constraining the safety score to stay below the bound, which requires a careful balance between safety and control per- formance. Furthermore, it is important to note that safety and control performance are not on equal footing, and the pursuit of better objective should be built upon ensuring safety (Knight, 2002; Cheng et al., 2019).
4. Method
Algorithm Algorithm of SafeDiffCon
1: Require Training set Demin. Calibration set Deat. control tar-
gets
2: 0 < Pre-train with Lattusion on Durai
3: 0 + Post-train with Gpost train D and Deal (Sec 4.2, Alg
3.2. Diffusion Models and Diffusion Control
2)
Diffusion models (Ho et al., 2020) learn data distribution from data in a generative way. They present impressive performance in a broad range of generation tasks. Diffusion models involve diffusiinn/deniisnng processes: the diffu- sion process q(x*11 xx*) (x Vagxk. (1-ax)I) corrupts the data distribution p(xo) to prior distribu- tion .V I), and the denoising process po po(x*=l(x") <x [x(* k), Iz.I makes sampling in a reverse di- rection. Here k is the diffusiin//eeoisingg step, {a j=
4: 8, w +- Fine-tune during guided inference with G and Lane-nan
on control targets and Deal (Sec 4.3, Alg 3)
5: return model parameters 0 and control sequences
In this section, we introduce our proposed method SafeDif- fCon, with its overall framework outlined in Figure 1 and Algorithm 1. Firstly, in Section 4.1, we propose the un- certainty quantile, which is the quantification of uncertain- ties. As shown in the middle of the figure, it is embedded
1K
and (ax {ox)<_, are the noise and variance schedules. In prac- tice, a denoising network Eg is trained to estimate the noise to be removed in each step. During inference, the iterative
IIn this paper, we use to represent state trajectory across time, and w to represent control signal sequence across time. u² denotes the full system trajectory across time at denoising step k.
3

---

From Uncertain to Safe: Conformal Adaptation of Diffusion Models for Safe PDE Control
Real Data
Generated Data
Feed Forward
Grad, Backward
Sample
Pre-train
Post-train
Inference-time Finetune
Training Data
Calibration Data
Line-tune
Diffusion
Diffusion
Uncertainty Quantile Q
Diffusion
Model
Model
Model
Q(1 a)
Ldiffusion
Cpost-train
Generated
Data
Control Targets
Figure 1: Overview of SafeDiffCon. First, we pre-train a diffusion model po on the training data. Then, combined with the uncertainty quantile, we post-train the model to steer distribution to safer regions with better control objectives. Finally, improve performance and safety for specific control tasks, we conduct inference-time fine-tuning, again incorporating the uncertainty quantile into the process.
throughout the algorithm, originating from the concern that the model's predictive uncertainty may cause a gap between the actual and predicted safety scores, potentially leading to unsafe events not anticipated by the model.
prediction are provided in Appendix D.
However, there is typically a distribution shift between the calibration set and the sequences generated by the diffusion model for control tasks during inference. Therefore, we account for this distribution shift, thereby introducing the shifted score set. Based on it, we get the uncertainty quan- tification metric named uncertainty quantile and finally the conformal interval. The detailed description of this process is as follows.
Secondly, in Section 4.2 and 4.3, introduce the post- training and inference-time fine-tuning based on the un- certainty quantile, respectively. In the post-training phase following pre-training, we employ a reweighted loss func- tion to guide the model's output distribution to favor regions with more optimal objectives and greater safety in general, after which the post-trained model can subsequently be used for specific control tasks. As for inference- sitrencce-tmee fine- tuning, we aim to make task-specific adjustments based on the post-trained model, enabling better control performance and safety guarantee for the specific control tasks with few iterations.
Calibration set and pre-training. To achieve the goal mentioned above, we first set aside a portion of the origi- nal training data the calibration set Deal. which will be used later to estimate the model's prediction errors. The remaining data, which will be applied to operations that can alter the model's outcomes, is referred to as the training data D After pre-training with Derin as described in Eq, 3, we get the diffusion model pe which models the joint distribution of (u. w).
4.1. Uncertainty Quantification of Diffusion Models
In offline safe control problems (Xu et al., 2022; Gu et al., 2022), the gap between pre-collected data and the target distribution exacerbates models' prediction errors, which can be critical in ensuring safety. To address this issue, employ the conformal prediction technique obtain uncertainty quantification for predicted safety, without re- quiring additional assumptions about the model and the data distribution.
Shifted score set under distribution shift. Score set, ob- tained on the calibration set, is a set of model prediction uncertainties regarding the safety score 8, which is defined as
S = ()s((u(w.)) s(;)) (u,, w.) Ufoo) (5)
Here u, and w, are the system state trajectory and control signal sequence of the 7th example, un uo(w) is the system state trajectory conditioned on control w, as predicted by the model, where 0 is the model parameter.
The intuition behind the original conformal prediction (Vovk et al., 2005) is to use the trained model on the calibration set to obtain a score set. The score set then used to endow the model's prediction with a conformal (confidence) interval. Under the assumption of exchangeability,, the actual predic- tion during inference will lie within the conformal interval with a guaranteed probability. The details of conformal
Furthermore, taking into account the distribution shift be- tween the calibration set Deal and data generated based on control targets, we apply weighting to the score set S to get the shifted score set as
(6) Intuitively, this weighting is because each sample in the calibration set has a different probability of appearing the
$ : (whrm (u;, w:) As, E SJ.
"Exchangeaiility means that exchanging examples in the cali- bration set and in inference does not alter their joint distribution.

---

From Uncertain to Safe: Conformal Adaptation of Diffusion Models for Safe PDE Control
final data distribution generated by the model. According to conformal prediction under covariate shift (Tībshirani al., 2019), here the calculation of the weights is as follows:
So we propose the reweighted post-training loss as
Cpont-mai
je --Wluw)jle (v [u, w| + vll-vvaii , (12)
puu, Ws)
w Uu W:) puu, w.)
which steers the generated data distribution towards high e-W regions. Compared to the pre-training loss in Eq. 3, the samples are reweighted, assigning higher weights to those with lower W. This loss function enables the model to adjust the generated data distribution to align with our target distribution, as described in the following theorem:
where p(u, w) is the probability density function of the calibration set and training set, and p(u, w) is the proba- bility density function of the model-genere geneeraed data during inference. Wnorm is the normalization of w defined as
w U1 Wi)
Whorn (U,, wi) E L( (uj, wj)
Theorem 4.2. Trained with Lepon train in Eq. 12, the diffusioon model's 's distribution is (u, w) x p(u, w)e" (((ww) (See Appendix A for proof.)
The specific calculation of Whorm regarding SafeDiflCon will be detailed in Section 4.2.
Specifically, in the calculation of the uncertainty quantile Q, w(u,, w.) = Cpp(u,, w,)e-W(u,,w,, /p(u,, w.). In practi- cal implementation, we let
Uncertainty quantile. Based on the shifted score set & (Eq. 6), the uncertainty quantile is Quantile((1 - (1 + 1 D: where
-[u,.w.)
is the coverage probability and [Deal the cardinality of Deal- The quantile represents the threshold below which a specific proportion of & falls. The adjustment term 1 + ccal accounts for the finite size of the calibration set, ensuring valid coverage even with limited data. Hereafter, it is abbreviated as Q(1 - or; S).
w(u; we) = Ce -W(W..)..)
(13)
and the constant C will be eliminated during normaliza- tion. Through the following theorem, we demonstrate that when the KI. divergence between p and po small, this approximation does not affect the results of the uncertainty quantile.
We can then obtain the conformal interval as:
Theorem 4.3. If DKL(Plpp)) E, then with 1 - CVe prob- ability, Quantile(1 az E Wnorm,(d.)) Quantile
Cla(1 - Deal) = [w)) - - 5), s(u)(w)) + Q(1 a; S) (9)
-11-a;;3
a; Ei= &v.), where Wnorm, and Ünom, are the nor- malized wi and W, w, = po(v,)e -Q(v)) /pls), and w e GGG().).
With least 1 a probability, the true s is covered by this interval, which can be demonstrated through the following theorem.
(See Appendix A for proof.)
The algorithm of post-training is provided in Algorithm 2 in Appendix E.2.
Theorem 4.1. Assume that samples in the calibration set Deat ~ p are independent, and the test set (u, w) p is also independent with the calibration set, Assume p is absolutely continuous with respect to P. then
4.3. Inference-time Fine-tuning
During inference, knowing the specific control tasks, we guide the model to further optimize the generation of con- trol sequences tailored to these tasks. We cyclically use the guidance generate samples and then fine-tune the model parameters with these samples, thereby iteratively improving the model's safety and performance. In this pro- cess, similar to post-training, we incorporate the uncertainty quantile Q into both the guidance and fine-tuning loss to pre- vent violations of safety constraints caused by the model's uncertainty.
P(((u)= E Cla a, Deal)) 2 a.
(10)
(See Appendix A for proof.)
4.2. Post-training with Reweighted Loss
The aim of post-training is to transform the model's distribu- tion into p" (u, w) & p(u, w) ee- -W(u.w) where W(u, w) contains both the control objective I and the safety con- straints involving the model's uncertainty. Specifically, it formulated as:
Guidance G. Guidance is the first approach we adopt to steer the model's output toward satisfying both the control objectives and safety constraints. It plays role during the sampling process of the diffusion model. The specific denoising step of implementing guidance follows Eq. 4, and the form of guidance still follows W in Eq. 11, i.e.
W(, w) max(s(u) + Q(1 - 5) 530-so.00 so, Tu,, w),
(11)
where y is the weight of J.
It means obtaining the element with (1 12--a) at((1 Paal 't
quantile from the shifted score set S.
g(u, w) W(u, w),
(14)
5

---

From Uncertain to Safe: Conformal Adaptation of Diffusion Models for Safe PDE Control
incorporating the uncertainty quantile Q to account for safety risks arising from uncertainty.
Table I: Results of ID Burgers' equation. Gray; there are unsafe trajectories. Black: all trajectories are safe. Bold: safe trajectories with lowest J.
Fine-tuning, The second approach for adjusting the out- put data distribution of the model is fine-tuning, which achieves the adjustment by optimizing the model param- eters 0. Specifically, retaining the computation graph for all denoising steps with respect f would result in an un- manageable memory overhead. Therefore, when we need to keep the computation graph during denoising for gradi- ent calculation, we only retain the computation graph of the final denoising step. To optimize the safety score and objective simultaneously, we form the fine-tune loss as:
Methods
J
Remmle V
Ryime +
Rpuin
BC
0.0001
38%
13%
1.2%
BC-Safe
0.0002
14%
3%
0.2%
PID
0.0968
0%
0%
0.0%
SL-Lag
0.0115
0%
0%
0.0%
MPC-Lag
0.0092
0%
0%
0.0%
CDT
0.0012
16%
3%
0.2%
TREBI
0.0074
0%
0%
0.0%
SafeDiffCon
0.0016
0%
0%
0.0%
Ltine-tune
Σ
W(ug, we),
(15)
(le.Wa)DDaneee
0
t 5
Vocer toed
where ug and wo are from Drampled sampled according to the guidance described above.
The entire algorithm of inference-time finetuning is pro- vided Algorithm in Appendix E.3.
5. Experiment
We design three safe PDE control scenarios, including the ID Burgers' equation, 2D incompressible fluid, and con- trolled nuclear fusion problem. These experiments aim to address the following three questions: (I) Does introducing the uncertainty quantile enable SafeDiffCon to meet safety requirements? (2) Can SafeDiffCon achieve superior control performance under safety constraints? (3) Do all compo- nents of our proposed algorithm contribute effectively to its performance?
Visualizations of the ID Burgers' equation. The top row shows the original trajectory corresponding to the control target, and the bottom row is the trajectory controlled by SafeDifiCon.
Figure 2:
external force xt.
Given target state ua(x), the control objective J is to minimize the control error between the final state uy and the target state
For comparison, we choose imitation learning method Be- havior Cloning (BC) (Pomerleau, 1988), and safe reinforce- ment learning and imitation learning methods involving BC with safe data filtering (BC-Safe), Constrained Decision Transformer (CDT) (Liu et al., 2023b) and diffusion-based method TREBI (Lin et al., 2023). Note that CDT shows the best performance in the offline Safe RL benchmark OSRL (Liu et al., 2023a). In addition, we combine the physical system control method Supervised Learning (Hwang et al., 2022) and Model Predictive Control (Schwenzer et al., 2021) with the Lagrangian approach (Chow et al., 2018) (SL-Lag, MPC-Lag) to enforce safety constraints. We also apply the classical control method PID (Li et al., 2006). We provide the code here.
J J= /= (u(c.a)))u( ud(x))
(16)
The safety score is defined as:
s(u) = sup (L.))
(u(t, x)),,
(17)
with the safety constraint 8p- If s()) > 8o, the state tra- jectory u is unsafe, otherwise, it is safe. We compute three unsafe rates to assess the safety levels of different meth- ods' control results. Rsample denotes the proportion of un- safe trajectories among total trajectories"; Rume denotes the proportion of unsafe timesteps among all timesteps; Rpoint denotes the proportion of unsafe spatial lattice points in all spatial lattice points across all time steps. More details can be found in Appendix F.I.
5.1. 1D Burgers' Equation
Experiment settings. 1D Burgers' equation is fundamen- tal equation that governs various physical systems including fluid dynamics and gas dynamics. Here we follow previous works (Hwang et al., 2022; Mowlavi & Nabi, 2023) and consider the Dirichlet boundary condition along with an
'If any point in the full trajectory is unsafe, this trajectory unsafe. So Romple is the most stringent metric.
6

---

From Uncertain to Safe: Conformal Adaptation of Diffusion Models for Safe PDE Control
Results. In Table 1, We report the results of the control objective I and safety metrics of different methods. SafeD- iftCon can meet the safety constraint and achieve the best control objective at the same time. As shown in Figure 2, given the initial condition and the final state (control target), SafeDifiCon can control a state trajectory that satisfies the safety constraint and control target. Other methods either suffer from constraint violations or suboptimal objectives. BC and BC-Safe trained from expert trajectories fail to meet the safety constraints, showing that simple behavior cloning is not feasible under the safety constraints. SI-Lag and MPC-Lag attempt to use the Lagrangian method to balance the control objective and safety, but this coupled training program makes it difficult to find the right balance, with poor control performance. CDT uses the complex Trans- former architecture, which can achieve low control error, but it can not meet safety constraints. The diffusion- based method TREBI sacrifices too much control performance to satisfy the safety constraints, because its error bound is soft, and the safety constraints can be easily violated.
Table 2: 2D incompressible fluid control results. Gray: there are unsafe trajectories. Black: all trajectories are safe. Bold: safe trajectories with lowest J.
Methods
J
SVM 4
R 4
BC
-0.7104
0.7156
88%
BC-Safe
-0,2520
0,0330
8%
CDT
-0.7025
0,2519
30%
TREBI
-0.7019
0.0808
18%
SafeDif[Con
-0.3548
0.0000
0%
we introduce another metric called Safety Violation Mag- nitude SVM max(s so. O]. When s does not exceed the bound so, this metric is O. If exceeds 8o, the met- ric reflects the amount by which it is surpassed. From the results, it is clear that our method is the only one to success- fully satisfy the safety constraint, even surpassing methods that only consider safety like BC-Safe, which demonstrates the effectiveness of post-training and fine-tuning with the uncertainty quantile.
5.2. 2D Incompressible Fluid
5.3. Tokamak Fusion Reactor
Experiment settings. We then consider the control prob- lems of the 2D incompressible fluid, which follows the Navier-Stokes equation. Following previous works (Holl et al., 2020; Wei al., 2024; Hu et al., 2024), the control task we consider is to maximize the amount of smoke that passes through the target bucket in the fluid flow with obstacles and openings, while constraining the amount of smoke pass- ing through the dangerous region under the safety bound. Specifically, referring to Figure 3. the control objective I is defined as the negative ratio of smoke passing through the target bucket located at the center top, while the safety score corresponds the ratio of smoke entering the hazardous red region. It is important to note that there is a trade-off between controlling the flow through the hazardous region and achieving a more optimal control objective, which im- poses higher demands on the algorithm. We set the safety score bound as sp = 0.1.
Experiment settings. A tokamak is a device designed to achieve controlled nuclear fusion by confining high- temperature plasma within toroidal magnetic field, en- abling conditions suitable for fusion reactions (Federici et al., 2001). The realization of controlled fusion could provide humanity with a clean, safe, and nearly limitless energy source, free from greenhouse gases, nuclear accident risks, and fuel supply imbalances (Schuster & Ariola, 2006). Safety in a tokamak is crucial to prevent plasma instabilities and disruptions Harbeck et al., 2019), which can damage the device worth billions of dollars.
Controlling the plasma state in tokamak is critical for maintaining stability and optimizing fusion performance, which often requires precise management of parameters like beta (8p), safety factor (g9s), and internal inductance (4). Based on the previous works (Seo et al., 2021; 2022), we set the control objective to make 3, and ( reach specific target values 87 and 1o at each time step, while the safety constraint requires qao to remain above the detined safety bound at all times. In this case, the objective I and safety score 8 are formulated as follows:
Moreover, this control task particularly challenging due its specific setup: not only does it require indirect control, which means that control can only be applied the periph- eral region, but the spatial control parameters also reach as many as 1,792. As for safety, among all the training data, 53.1% % of the samples are unsafe, meaning their safety score 8 exceeds the bound sn = 0.1. The average safety score of the dataset is 0.3215. Other details can be found in Appendix G.I.
J == I
(18g(t,x)+$(())) (x))
-
(18)
Jnx J (0.))
+ (1((t x) - 17 (x))) inf 8 == (x)))((0TT)x2 t100 (t,
(19)
(20)
Results. We report results of SafeDiffCon and baselines in Table 2. Here SL-Lag and MPC-Lag fail to achieve rea- sonable control results because of the complex dynamics, and we define R as the rate of unsafe samples. Additionally,
And the safety score bound is so = - -.98.. We use the same initial state and tokamak simulation environment as in the previous work (Seo al., 2022), setting time-varying
7

---

From Uncertain to Safe: Conformal Adaptation of Diffusion Models for Safe PDE Control
t= 0
t= 7
t= 14
t = 21
t = 28
Figure 3: Visualization of the 2D incompressible fluid control by our SafeDiffCon. By controlling the fluid on the outside margin, the yellow smoke is successfully maneuvered to the center top exit while avoiding the red unsafe region.
random targets within the reasonable ranges proposed in that study. The control sequences for training are collected using the reinforcement learning model trained in the previous work (Seo al., 2022), where 71%18% of the samples in the training dataset are unsafe.
components. Furthermore, when the uncertainty quantile is removed, which means Q is set to 0, the algorithm also fails to meet safety requirements. This demonstrates that without considering the model's uncertainty, its safety predictions are inaccurate, leading to the occurrence of unsafe events.
Table 4: Results of the ablation study. We compare SafeD- iffCon with SafeDiffCon w/o post-training. inference-time fine-tuning and the uncertainty quantile Q
Table 3: Results of the tokamak fusion reactor. Gray: there unsafe trajectories. Black: all trajectories are safe. Bold: safe trajectories with lowest J.
Methods
J +
Raample &
Riime
Methods
J 4
SVM 4
R
BC
0.0610
42%
1,34%
SafeDifiCon
-0.3548
0.0000
0%
BC-Safe
0.0811
4%
0.03%
w/o post-training
-0.4911
0.0849
12%
SL-Lag
0.8812
0%
0.00%
w/o fine-tuning
-0.4692
0.0247
8%
MPC-Lag
0.8659
0%
0.00%
w/o
-0.6364
0.0623
12%
CDT
0.0071
8%
0.54%
TREBI 0.0261
0%
0.00%
6. Limitation and Future Work
SafeDiffCon
0.0121
0%
0.00%
Firstly, our method is not constrained by specific scenarios, meaning it can be applied to more real-world tasks, which is our future work. Secondly, we consider extending this method to other generative methods that require constraints, not just the diffusion model. Finally, the future, we will explore the possibility of developing stricter coverage of the safety score that sacrifices less accuracy while still ensuring safety.
Results. We provide the results of control and safety metrics in Table 3, where the definition of R.sample and Riime are the same as Section 5.1. Due to the differences in the input and output dimensions of the controller and the lack of clarity re- garding the device's specific details, PID cannot be directly applied here. Results show that SafeDiffCon successfully achieves full compliance with safety constraints. More- over, among all algorithms that satisfy safety constraints, our al algorithm achieves the lowest J. indicating that it also achieves the best control performance. These results verify our statement that our algorithm is not only safe but also effective.
7. Conclusion
In this paper, we have introduced Safe Diffusion Models for PDE Control (SafeDiffCon), novel approach that inte- grates safety constraints into deep learning-based control of PDE systems. By quantifying safety uncertainty using con- formal prediction and introducing a reweighted loss during post-training, our method effectively handles distribution shifts and ensures safety while maintaining competitive con- trol performance. Furthermore, inference-time fine-tuning allows dynamic adjustment of the model to meet specific control targets. The results of experiments on diverse control tasks, including ID Burgers' equation, 2D incompressible fluid, and controlled nuclear fusion, demonstrate that SafeD- iffCon reliably satisfies safety constraints and outperforms existing methods. This work paves the way for safer de- ployment of machine learning techniques in high-stakes real-world applications.
5.4. Ablation Study
To demonstrate the necessity and effectiveness each com- ponent of our algorithm, we conduct three ablation exper- iments on the 2D incompressible fluid control problem by removing post-training, inference-time fine-tuning, and the uncertainty quantile Q, respectively.
It can be observed that post-training and fine-tuning, as two methods that influence the model's output distribution, are both essential for meeting safety constraints. Remov- ing either one results in the algorithm failing to achieve safety, highlighting the indispensable roles of these two
8

---

From Uncertain to Safe: Conformal Adaptation of Diffusion Models for Safe PDE Control
Impact Statement
Chung, H., Kim, J., Mccann, M. T., Klasky, M. L.,, and
Ye, J. C. Diffusion posterior sampling for general noisy inverse problems. In The Eleventh International Confer- ence on Learning Representations, 2023. http:: / //openrevieg ew net/ forum? det.mmidddnnDDDnnnn00m
This paper presents work whose goal to advance the field of Machine Learning. Furthermore, our work makes Ma- chine Learning applications safer in real-world scenarios. There many potential societal consequences of our work, none of which we feel must be specifically highlighted here.
Dawson, C.,, Qin, Z., Gao, S., and Fan, C. Safe nonlinear control using robust neural lyapunov-barrier functions. In Conference on Robot Learning, pp. 1724-1735. PMLR, 2022.
Acknowledgements
We thank Yuchen Yang for the discussions and for providing feedback on our manuscript. We gratefully acknowledge the support of Westlake University Research Center for Industries of the Future and Westlake University Center for figgh-peffrmmnce Computing. The content is solely the responsibility of the authors and does not necessarily represent the official views of the funding entities.
Dhariwal, P. and Nichol, A. Diffusion Models Beat GANs on Image Synthesis, June 2021. URL http: //arxiva org abs/ htsp/bbs/21055. 05233. arXiv:2105.05233 [cs, stat].
Farahmand, A.-m., Nabi, S., and Nikovski, D. N. Deep reinforeement learning for partial differential equation control. In 2017 American Control Conference (ACC), pp. doi:
3120-3127, 2017.
References
Argomedo, F. B., Witrant, E., Prieur, C., Bremond, S., Nouailletas, R., and Artaud, J.-F. Lyapunov-based dis- tributed control of the safety-factor profile tokamak plasma. Nuclear Fusion, 53(3):033005, 2013.
Federici, G., Skinner, C. H., Brooks, N., Coad, P., Grisolia, C., Haasz, A. A., Hassanein, A., Philipps, V., Pitcher, C. S., Roth, J., et al. Plasma-material interactions in current tokamaks and their implications for next step fusion reactors. Nuclear Fusion, 41 (1 2001.
Barros, E., des Santos, and M.V.D. A safe, accurate in- travenous infusion control system. JEEE Micro, 18(5): 12-21, 1998. doi: 10. 10.1109/40.735940.
Findeisen, R. and Allgöwer, F. An introduction to nonlinear model predictive control. In 21st Benelux meeting on systems and control, volume 11, pp. 119-141. Veldhoven, 2002.
Chen, L., Lu, K., Rajeswaran, A., Lee, K., Grover, A., Laskin, M., Abbeel, P., Srinivas, A., and Mordatch, I. De- cision transformer: Reinforcement learning via sequence modeling. Advances neural information processing systems, 34:1 5084-150977 2021.
Garcia, C. E., Prett, D. M., and Morari, M. Model predictive control: Theory and practice- a survey. Automatica, 25 (3):335-348, 1989.
Cheng, R., Orosz, G., Murray, M., and Burdick, J. W. End-to-end safe reinforcement learning through barrier functions for safety-critical continuous control tasks. In Proceedings of the AAAI conference on artificial intelli- gence, volume 33, pp. 1397-3395, 2019.
Gibbs, and Candes, E. Adaptive conformal inference under distribution shift. Advances Neural Information Processing Systems, 34: 1660-166222 2021.
Gu, S., Yang. L.. Du, Y., Chen, G., Walter, F., Wang, J., and Knoll, A. A review of safe reinforcement learn- ing: Methods, theory and applications. arXiv preprint arXiv:2205. 10330, 2022.
Cheroozuukov, V., Wüthrich, K., and Yinchu, Z. Exact and robust conformal inference methods for predictive machine learning with dependent data. In Conference On learning theory, pp. 339. PMLR, 2018.
Hendrycks, D., Mazeika, M., Wilson, D., and Gimpel, K. Using trusted data to train deep networks on labels cor- rupted by severe noise. Advances in neural information processing systems, 31, 2018
Chi, C., Xu, Z., Feng, S., Cousineau, E., Du, Y., Burchfiel, B., Tedrake, R., and Song, S. Diffusion policy: Visuomo- tor policy learning via action diffusion. The International Journal of Robotics Research, pp. 027833492441233688 2023.
Hinze, M. and Kunisch, K. Second order methods for opti- mal control of time-dependent fluid flow. SIAM Journal on Control and Optimization, 40(3):925-946, 2001.
Chow, Y., Ghavamzadeh, M., Janson, L., and Pavone, M. cisk-constrained reinforcement learning with percentile risk criteria. Journal of Machine Learning Research, 18 1167):1-51, 2018.
Ho, J., Jain, A., and Abbeel, Denoising diffusion proba- bilistic models. Advances in neural information process- ing systems, 33:68400668 6851, 2020.

---

From Uncertain to Safe: Conformal Adaptation of Diffusion Models for Safe PDE Control
Ho, J., Salimans, T., Gritsenko, A., Chan, W., Norouzi, M., and Fleet, D. J. Video diffusion models. arXiv: 2022.
Lin, Q., Tang, B., Wu, Z., Yu, C., Mao, S., Xie, Q.. Wang, X., and Wang, D. Safe offline reinforcement learning with real-time budget constraints. In International Conference on Machine Learning, pp. 21 21152. PMLR, 2023.
Holl, P., Thuerey, N., and Koltun, V. Learning control pdes with differentiable physics. In International Confer- ence on Learning Representations, 2020.
Liu, Z., Guo, Z., Lin, H., Yao. Y., Zhu, J., Cen, Z., Hu, H., Yu, W., Zhang, T., Tan, J., et al. Datasets and benchmarks for offline safe reinforcement learning. arXiv preprint 2rXiv:2306.099) 2023a.
Hu, P., Wang, R., Zheng, X., Zhang, T., Feng, H., Feng, R., Wei, L.. Wang, Y., Ma, Z.-M., and Wu, T. Wavelet diffu- sion neural operator. arXiv preprint arXiv:2412.04833, 2024.
Liu, Z. Guo, Z.. Yao, Y., Cen, Z., Yu, W., Zhang, T., and Zhao, D. Constrained decision transformer for offline safe reinforcement learning. In International Conference on Machine Learning, pp. 2161 1 21630. PMLR, 2023b.
Hwang, R., Lee, J. Y.. Shin, J. Y., and Hwang, H. J. Solving pde-constrained control problems using operator learn- ing. In Proceedings of the AAAI Conference on Artificial Intelligence, pp. 4504 4512, 2022.
Maxime Cauchois, Suyash Gupta, A. A. and Duchi, J. C. Robust validation: Confident predictions even when distributions shift. Journal of the American Statis- tical Association, 0(0):1-66, 2024. doi: 10.1080/ 0162145992003.2298800 URL /ttp:::://d...ee 10 1001100//000 000.0000111245551 202.2211.9922233.....2222222 2298037.
Janner, M., Du, Y., Tenenbaum, J., and Levine, S. Plan- ning with diffusion for flexible behavior synthesis. Pro- ceedings of Machine Learning Research, 162:9902-9915, 17-23 Jul 2022.
Mowlavi, S. and Nabi, S. Optimal control of pdes using physics-informed neural networks. Journal of Computu- tional Physics, 473: 173:111733 11731, 2023.
Johnson, M. A. and Moradi, M. H. PID control. Springer, 2005.
Pan, Y., Farahmand, A.-m., White, M., Nabi, S., Grover, P., and Nikovski, D. Reinforcement learning with function- valued action spaces for partial differential equation con- trol. In Dy. J. and Krause, A. (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80, pp. 3986-399 9996-3995. PMLR, 10-15 Jul 2018.
Kates-Harbecke J., Svyatkovskiy. A., and Tang, W. Pre- dicting disruptive instabilities in controlled fusion plas- mas through deep learning. Nature, 56877753):566531, 2019.
Knight, J. C. Safety critical systems: challenges and direc- tions. In Proceedings of the 24th international conference on software engineering, pp. 547-550, 2002.
Podkopaev, A. and Ramdas, A. Distribution-free uncer- tainty quantification for classification under label shift. In de Campos, C. and Maathuis, M. H. (eds.). Proceed- ings of the Thirty-Sevent Conference on Uncertainty in Artificial Intelligence, volume 161 of Proceedings of Ma- chine Learning Research, pp. 844-853, PMLR, 27-30 Jul 2021. URL v161 w161/podkopaev2aa..tt
Lee, J., Paduraru, C., Mankowitz, D. J., Heess, N., Pre- cup, D., Kim, K.-E., and Guez, A. COptiDICE: Offline constrained reinforcement learning via stationary distri- bution correction estimation. International Confer- ence on Learning Representations, 2022. URL https: n
mll press/
//openrevies net.
6Q.
Pomerleau, D. A. Alvinn: An autonomous land vehicle in neural network. Advances in neural information processing systems, 1, 1988.
Lei, J., G'Sell, M. G.. Rinaldo, A., Tibshirani, R. J., and Wasserman, I. A. Distribution-free predictive inference for regression. Journal of the American Statistical Asso- ciation, 113:1094- - 1111, 2016. URL https:t/aaatt semanticscholest htwpiaarrrrrrraoaeeeeeeeee Cttp://orpussD--e 1174141:
Protas, B. adjoint-based optimization of pde systems with alternative gradients. Journal of Computational Physics, 227(13):6490-65) 65100 2008.
Rabault, J., Kuchta, M., Jensen, A., Réglade, U., and Cer- ardi, N. Artificial neural networks trained through deep re- inforcement learning discover control strategies for active flow control. Journal of fluid mechanics, 965:281-300 2019.
Li, T., Biferale, L.., Bonaccorso, F., Scarpolini, M. A., and Buzzicotti, M. Synthetic lagrangian turbulence by gener- ative diffusion models. Nature Machine Intelligence, pp. 1-11, 2024.
Li, Y., Ang, K. H., and Chong, G. Pid control system analysis and design. IEEE Control Systems Magazine, 26 (1):32-41, 2006. doi: 10.1 109/MCS 2000.1580155 1580152.
Schuster, E. and Ariola, M. The role of controls in nuclear fusion. In Proceedings of the 45th IEEE Conference on Decision and Control, pp. 2263-2270. IEBE, 2006.
10

---

From Uncertain to Safe: Conformal Adaptation of Diffusion Models for Safe PDE Control
Schwenzer, M., Ay, M., Bergs, T., and Abel, D. Review model predictive control: An engineering perspective. The International Journal of Advanced Manufacturing Technology, 117(5):1327-1399 2021.
Seo, J., Na, Y-S.,, Kim, B., Lee, C., Park, M., Park, S., and Lee, Y. Feedforward beta control in the kstar tokamak by deep reinforcement learning. Nuclear Fusion, 61(10): 106010, 2021.
Sco, J., Na, Y-S.,, Kim, B., Lee, C., Park, M., Park, S., and Lee, Y. Development of operation trajectory design algorithm for control of multiple Od parameters using deep reinforcement learning in kstar. Nuclear Fusion, 62 (8):086049, 2022.
Soner, H. M. Stochastic optimal control in finance. Scuola normale superiore, 2004.
Song, J., Meng, C., and Ermon, S. Denoising diffusion implicit models. In International Conference on Learning Representations, 2021. URL https: /tttpp://eeeeee pppnrevien.eeeeeee net/ www.ooooiddisiii.oee (tt.:/uumiidsst1iarrccile
Tibshirani, R. J., Barber, R. F., Candès, E. J., and Ramdas, Conformal prediction under covariate shift. In Neural Information Processing Systems, 2019. URL http://c...ssmttcccccc..ll.l.ll http//oopuSD::111100000
Vahdat, A., Williams, F., Gojeic, Z., Litany, O., Fidler, S., Kreis, K., al. Lion: Latent point diffusion models for 3d shape generation. Advances in Neural Information Processing Systems, 35:10021-10039, 2022.
Vovk, V., Gammerman, A., and Shafer, G. Algorithmic learning in random world, volume 29, Springer, 2005.
Wei, L., Hu, P., Feng, R., Feng, H., Du, Y., Zhang, T., Wang, R., Wang. Y., Ma, Z.,M., and Wu, A generative approach to control complex physical systems. Advances in Neural Information Processing Systems, 2024.
Xu, H., Zhan, X., and Zhu, X. Constraints penalized q- learning for safe offline reinforcement learning. Pro- ceedings of the AAAI Conference on Artificial Intelli- gence, volume 36, pp. 8753-8760, 2022.
Zheng, Y., Li, J., Yu, D., Yang, Y., Li, S. E., Zhan, X., and Liu, Safe offline reinforcement learning with feasibility-guiied diffusion model. arXiv preprint 10000:2401.10000 2024.

---

From Uncertain to Safe: Conformal Adaptation of Diffusion Models for Safe PDE Control
A. Theory
In this section, we provide the proof of the theorems mentioned in the main text.
Theorem A.I. Assume samples the calibration set Deat ~ P are independent, and the test set (u, w) p is also independent with the calibration set. Assume p is absolutely continuous with respect to p, then
P(s(u) E Clo(1 ae,aa) 2 1 a.
(21)
Proof. According to Lemma in the previous work (Tibshiraní et al., 2019), the calibration set and test set are weighted exchangeable (defined in Definition in the previous work (Tibshirani et al., 2019)), with the weight function w.
According to Lemma in the previous work (Tibshirani al., 2019), let the (u, w) pairs be the random variables, can obtain that
P(s(u)= S s(uo(w)) + Q(1 o; o); 2
(22)
Similarly, it can be derived that
Q a;; S s(u)) 2 1
(23)
Thus
P(s(u) E Cla(1 a, Deal)) 2 a.
(24)
Theorem A.2. Trained with Cpoct-train
Cpost-train = B.t.:rr...A.laww////t.w........ -ww.w//
W(u.w) (W 1-- lelvoolva/u, w] + vI āxe, k)]l
(25)
the diffusion model's distribution is p" (u, w) x p(u, w ( W(u.w)
Proof. Trained with the diffusion loss
Laiftusion
e0(Vāz w] + VI ō.E, k))))
(26)
the diffusion model's distribution is thus to learn the distribution p", the loss should be
C Bww.e....../wuw//........... co(van(u, w| + vI age, k)
(27)
-EL /
cle-colva VI ãxe, k)ll p" (u, w)dudw]
(28)
x xEL.oq.k). 1 le--elvāalu,, + VI ã.c, 1) 3 p(u, w/e W(n.) - dudw) dudw]
(29)
Fu.ciu.i..au////umoooooooo -W(u.w) colvvā w] + vi -1-aa -,,, k)1
(30)
Theorem A.3. the KL. divergence of the real distribution p and the distribution Pe learned by the model less than e,
then with 1- probabiliy the quantile Quantile(1 a;; Σ Wnorm,gd., Quantile(1 a > i==l 8 where
Wnorm, and i Wnomm, are the normalized wi and w. w, Pole, le W((,)
and w; -(v.))
p(v:)
12

---

From Uncertain to Safe: Conformal Adaptation of Diffusion Models for Safe PDE Control
Proof.
Since (pl(p) f p(x) log pla) dæ e, we can get pa(x)
P(A {Ip(2) - po(x)) > 8)) SA (p(pa) pa po(
(31)
So (p(x) pro
(32)
6
27V(p,pe)
(33)
5
2
(xl(pllpe)
(34)
8 2
V2e
(35)
5
where (p, po) is the total variation distance between p and po- Eq. 34 is obtained using Pinsker's inequality. This
22e inequality means that with 1 _ probability, |Ap(o)] = (p) p(v) < 5.
Let 5 be small quantity and p- min .N (p(1:)), since
(1 + Ap(rz)]| 2 Σ
Ap(U))
Ap(uz)]
E
2 Σ
(36)
p
Pj
p
and
(1 + Aplied )
W
Wnerm,
Σe W/
=
Σ En (1 + Sp(r)) )ūs
(37)
m(n))
thus
+ Ap(#:)
1 + Ap(es)
plue)
mnomm S Whorm, 5
p(e:)
bomm
(38)
Ap(r:)
Ap(r))
By rearranging terms, we obtain
Ap(v:)/p(e - - Ap(e:)/p(v.) + [Ap(u:)1/B
Whorm norm,
max
(39)
1+ AAp(o:)\/p
1 Ap(2;)1/p-
11/p(u;) + 11/p-1
Ap((;))
(40)
1 1 8/p-
ClAp(1)
(41)
Clo).
(42)
Because can merge the or values that the same, assume t1 Quantile(a; E 50..) UR. which means
1/2
UN- Suppose
Whorm 2 1 a
Wnomm <1-a a.
(43)
With Eq. 42, we can get that
Wnorm,i 2 1 Cr Ckò,
Wnorm,i <1- a Ckô.
(44)
We can let small enough to make l satisfy
k-
Wnorm, 2 1 - a Σ
Ünorm, < 1 - a,
(45)
which means Quantile(o; E Wnorm, )og) Quantile(re (o: Σ Wnorm, )o
13

---

From Uncertain to Safe: Conformal Adaptation of Diffusion Models for Safe PDE Control
B. Additional Experiment Results
B.1. Inference Time
We have accelerated the efficiency of inference. On one hand, by introducing post-training, the data distribution is already close to the desired distribution before inference. On the other hand, we have significantly accelerated the diffusion model's sampling time using DDIM (Song al., 2021). The Table 5 below presents comparison of the inference time between method and the baseline the Burgers' equation A800 with 8 CPUs. It shows that method achieves moderate speed and is much faster compared to TREBI, which is also based on the diffusion model.
Table Inference time on the Burgers' equation on A800 with 8 CPUs.
Methods
BC
PID
SL-Lag
MPC-Lag
CDT
TREBI
SafeDiffCon
Inference Time (min)
0.1351
0.1091
0.8842
26,2905
0.0890
13.5525
2.3575
B.2. Ablation Study
In this subsection, we provide results of ablation studies the 11) Burgers' equation and the tokamak fusion reactor. We also compare SafeDiffCon and SafeDifiCon without post-training, inference-time fine-tuning, and the uncertainty quantile Q. The results in Table 6 still show that the absence of any module affects the model, causing it to fail to meet the safety constraints. Therefore, each module is effective.
Table 6: Results of the ablation study on the tokamak fusion reactor: We compare SafeDiffCon with SafeDiffCon w/o post-training, inference-time fine-tuning and the uncertainty quantile Q.
Methods
J 4
Reample
Riime
SafeDiffCon
0.0094
0%
0%
w/o post-training
0.0153
10%
0.52%
w/o fine-tuning
0.0210
50%
7.74%
w/o
0.0269
28%
0.65%
The only exception the result of SafeDiffCon without post-training in Table 7. Due to the high precision requirements of this control scenario, large-scale post-training may easily disrupt the dynamics originally learned by the model during training. Therefore, directly applying fine-tuning allows the model adapt specific control tasks without significant loss of accuracy, enabling both safe and precise control.
Table 7: Results of the ablation study on the Burgers' equation. We compare SafeDifiCon with SafeDifTCon w/o post-training, inference-time fine-tuning and the uncertainty quantile Q.
Methods
J 4
R Raamplo
Raie &
R point
SafeDiffCon
0.0016
0%
0%
0%
w/o post-training
0.0016
0%
0%
0%
w/o fine-tuning
0.0007
40%
20%
3%
w/o Q
0.0006
30%
10%
1%
B.3. Analysis on Parameters
We conduct experimental analysis on coverage probability a, the split ratio of trainingccalibration data, and the weight of objective 7. The Table and 10 show that they not impact much on performance, indicating that the model is robust.
14

---

From Uncertain to Safe: Conformal Adaptation of Diffusion Models for Safe PDE Control
Table 8: Impact of the coverage probability a on the tokamak fusion reactor.
Split Ratio
J +
Reample +
Riime
0.8
0.01017
0%
0%
0.85
0.0091
0%
0%
0.9
0.0094
0%
0%
0.95
0.0095
0%
0%
Table Impact of the split ratio of training/calibration data on the tokamak fusion reactor.
Qt
J
Rammle +
Reime
0.005
0.0124
4%
0.03%
0.01
0.0093
0%
0%
0.02
0.0094
0%
0%
C. Visualization of Experiment Results
C.1. ID Burgers' Equation
In this section, we provide additional visualizations of the control results for the ID Burgers' equation, shown in Figure In these figures, the top row represents the original trajectories corresponding to the control targets, while the bottom row displays the trajectories controlled by SafeDiffCon. It can be observed that SafeDiffCon successfully controls the trajectories, preventing boundary violations and guiding them to the desired final state.
15

---

From Uncertain to Safe: Conformal Adaptation of Diffusion Models for Safe PDE Control
Table 10: Impact of the split ratio of the weight of objective 7 on the 2D incompressible fluid.
J +
SVM +
R
0.01
0.3548
0%
0.1
0.4953
0.004
2%
0.3
0.498
0.003
2%
t=
10 Lower Beand
Upea Bound
10
Upeel Nose
Figure 4: Visualization of the ID Burgers' equation.
C.2. 2D Incompressible Fluid
Here, we provide additional visualizations of the control problems 2D incompressible fluid. From the figures, we can observe that SafeDiffCon can successfully control the smoke to avoid the red hazardous region and reach the target bucket as well.
16

---

From Uncertain to Safe: Conformal Adaptation of Diffusion Models for Safe PDE Control
t=0
t=7
t=14
t=21
t=28
Figure 5: Visualizations of the 2D incompressible fluid control problem..
D. Additional Details for Conformal Prediction
Conformal prediction is a flexible framework that provides prediction intervals with guaranteed coverage probabilities for new, unseen data points, under the assumption that the data are exchangeable.
Theoretical Foundations The exchangeability assumption is cornerstone of conformal prediction. requires that the order of the data points does not affect their joint distribution, meaning that any permutation of the indices yields an identical distribution. In particular, exchangeability holds for independent and identically distributed (i.i.d.) samples, common assumption in machine learning tasks. Ensuring exchangeability guarantees the validity of the prediction intervals constructed using conformal prediction.
Implementation Details To implement conformal prediction, the dataset first split into two subsets: proper training set (Drin) and calibration set (Dea). A predictive model μ* trained on the training set using specified learning algorithm
17

---

From Uncertain to Safe: Conformal Adaptation of Diffusion Models for Safe PDE Control
A. Once trained, the model generates predictions for the calibration These predictions are used compute "conformity scores', which measure the model's accuracy for each calibration point. Specifically, for each instance in the calibration set, the conformity score S; is defined as:
S = (mo - Yil, i E Deal-
Additionally, worst-case score of is included to account for extreme scenarios.
( (S) of the set of conformity scores is calculated, where represents the desired significance level Then quantile şi - (e.g,, 0.05 for 95% confidence interval).
Given a new data point x,1 the prediction interval for its corresponding output is calculated as:
Ca(Xn+1) = [pe qi-a(S), ue
Ho (X=+1) + q1-o
This interval provides an estimate for the range within which the true value Y is expected to lie, with a coverage probability of at least - a. Thus, conformal prediction offers a flexible and robust method for constructing prediction intervals that account for both the model's accuracy and the variability in the data.
Theoretical Guarantees Conformal prediction provides theoretical guarantees for finite samples (Vovk et al., 2005; Lei al., 2016). Specifically, for any new data point, the prediction interval satisfies the following probabilistic bound:
P(Y,+1 E (a((x+1)) 2
This ensures that the true label Yn+ will fall within the predicted interval at least - a percent of the time. This framework, based on the assumption exchangeability, provides a robust method for generating reliable prediction intervals, even in settings with limited sample sizes.
E. Details of the Method Implementation
E.1. Conditionally Sample (w(w)
In our proposed algorithm, we need to sample from the conditional distribution p(u[w) with the model that learns the joint distribution p(u, w). To achieve this, at each denoising step of the sampling process, we replace the noisy w in the input of the denoising network with the actual clean w that serves as the condition (Chung et al., 2023). In fact, this situation represents a special case of the data distribution that the denoising network encounters during training, where the part noisy, while the part remains noise-free.
E.2. Algorithm of Post-training
Here we provide the specific algorithm of post-training in Algorithm 2.
E.3. Algorithm of Inference-time Fine-tuning
In this subsection, we provide the entire algorithm of inference-time fine-tuning in Algorithm
Algorithm 2 Post- Post-training of SafeDiffCon
1: Require Training set Duain. Calibration set Deat coverage probability a, number of epochs N, number of updating steps per epoch M
2: for 1,...,N N do
Compute the shifted score set S with D Eq.
4: Get the uncertainty quantile Q(1 ax; Š)
5: for 1,..., M do
Take gradient descent step on Cposstmtinn with Diain Eq.
7: end for
8: end for
9:
return 0
18

---

From Uncertain to Safe: Conformal Adaptation of Diffusion Models for Safe PDE Control
Algorithm 3 Inference of SafeDifiCon
1: Require Calibration set Deal, coverage probability number of iterations N
2: for n ,,..., , do
3: Compute the shifted score set with Daa // Eq.
Get the uncertainty quantile Q(1 oç
5: Sample the control sequence w with guidance // Eq. 14
6: Take gradient descent step on V Ciine-tuue // Eq. 15
7: end for
8: Sample the control sequence with guidance // Eq. 14
return
F. Additional Details for 1D Burgers' Experiment
F.1. Experiment Setting
Following the previous works (Holl et al., 2020; Wei al., 2024), generate the ID Burgers' equation dataset. This equation is formulated as follows:
@u(t.z) -u(1,
Duft.z)
((t.))) in [0, Ω on [0, an
w((,x)
U2
Ox
u(t, x) =0
(46)
u(0, x) = uo(x)
{t = 03 x Z
where denotes the viscosity parameter, while un signifies the initial condition. We
0.01, T and Ω [0, 1].
During inference, alongside the control sequence w(t, diffusion model generates states u(t, Our reported evaluation metric J is always computed by feeding the control w(t, x) into the ground truth numerical solver to get (t, 2) and computed following Eq. (16). Following Eq. (5.1), we consider the safety constraint and define the safety score as s. In our experiment, the bound of safety score Su is set to 0.64, 89.7% samples are unsafe among the training set, 90% of samples are unsafe among the calibration set and all of the samples in the test set are unsafe. The details the ID Burgers' equation dataset for safe PDE control problem listed in Table 11.
Table 11: Details of ID Burgers' equation dataset.
Training Set
Calibration Set
Test Set
Unsafe Trajectories 34,985
900
50
Safe Trajectories 4,015
100
F.2. Model
The model architecture in this experiment follows the Denoising Diffusion Probabilistic Model (DDPM) (Ho et al., 2020).
For control tasks, we condition on ug, ur and apply guidance to generate the full trajectories of U|r,7] flo.r and the safety score The hyperparameters for the 2D-Unet architecture recorded in Table 12.
G. Additional Details for 2D Fluid Experiment
G.1. Experiment Setting
The 2D environment is modeled by the Navier-Stokes equation:
at + v. Vv VV Wr
V. 0,
(47)
where v is the velocity, p is the pressure, f the external force and is the viscosity coefficient.
19

---

From Uncertain to Safe: Conformal Adaptation of Diffusion Models for Safe PDE Control
Table 12: Hyperparameters of 2D-Unet architecture in ID experiment.
Hyperparameter Name
Value
Initial dimension
128
Convolution kernel size
3
Dimension multiplier
[1,2,4.8]
Resnet block groups
Attention hidden dimension
32
Attention heads
4
Number of training steps
200000
DDIM sampling iterations
100
of DDIM sampling
Following works from Holl al. (2020); Wei et al. (2024); Hu et al. (2024), 2D incompressible fluid dataset. The control objective and data generation is the same as before (Wei et al., 2024). The main difference between our data and previous ones is that we consider the safety constraint here. We define the safety score as the percentage of smoke passing through a specific region. This reflects the need to limit the amount of pollutants passing through certain areas real-world scenarios, such as in watershed.
the package PhiFlow to generate the
We simulate the fluid on x28x 128 grid. The selected hazardous region is [44, 36]
[40,64]. Since the optimal path for smoke, starting from left-biased position, is likely to pass through this hazardous region, this poses greater challenge for the algorithm: how to balance safety and achieving optimal objective, making this more difficult problem.
G.2. Model
In this paper, the design of the three-dimensional U-net we is based on the previous work (Ho al., 2022). In experiment, utilize spatio-temporal 3D convolutions. The U-net consists of three key components: downsampling encoder, central module, and an upsampling decoder.
The diffusion model conditions on the initial density and uses guidance as previous mentioned to to generate the full trajectories of density, velocity, control, the objective J and the safety score. As J and s are scalers, we repeat them to match other channels. The hyperparameters for the 3D U-net architecture are listed Table 13.
Table 13: Hyperparameters of 3D-Unet architecture in 2D experiments.
Hyperparameter Name
Value Number of attention heads
Kernel size of conv3d Padding of conv3d
3,33
(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 Padding of upsampling
(1,4.4)
(1,2,2)
Stride of upsampling
(0, 1,
Number of training steps
200000
DDIM sampling iterations
100
1) of DDIM Sampling
Intensity of guidance in control
100
Weight of safety term in guidance
10000
20

---

From Uncertain to Safe: Conformal Adaptation of Diffusion Models for Safe PDE Control
H. Additional Details for Tokamak Fusion Reactor
H.1. Experiment Setting
Following the previous work (Sco et al., 2022), the environment used in this experiment is a data-driven simulator designed to replicate the plasma behavior in the KSTAR tokamak. It is constructed using long short-term memory (LSTM) neural networks to capture the time-dependent dynamies of plasma, such as transport processes and flux diffusion. The LSTM-based model is trained five years of experimental data from KSTAR, incorporating various tokamak control variables inputs, including plasma current, toroidal magnetic field, line-averaged density, neutral beam injection and electron cyclotron heating powers, and boundary shape parameters.
The simulator predicts the evolution of key zero-dimensional (OD) plasma parameters (1 - Op. 4, and (9s) on a 100 ms timescale in an autoregressive manner, using the current state and control inputs. This enables fast and experimentally relevant predictions of plasma responses. Using this model avoids the high costs and risks associated with real tokamak experiments while ensuring sufficient accuracy and efficiency for training.
I. Baselines
L.I. CDT
Constraints Decision Transformer (CDT) (Liu et al., 2023b) models control as a multi-task regression problem, extending the Decision Transformer (DT) (Chen et al., 2021). It sequentially predicts returns-to-go, costs-to-go, observations, and actions, making actions dependent on previous returns and costs. The authors propose two techniques to adapt the model for safety-constrained scenarios:
1. Stochastic Policy with Entropy Regularization: This technique aims to reduce the risk of constraint violations due to out-of-distiibttion actions. In deterministic policy, the model selects single action based on its learned policy, which may result in unsafe actions when faced with states not well represented in the training data. By using stochastic policy, the model samples actions from distribution, encouraging the exploration of wider action space. Entropy regularization further enforces diversity in the sampled actions, making the model more robust in uncertain or underrepresented situations. This approach reduces the likelihood of selecting unsafe actions when faced with states outside the distribution of the training set.
2. Pareto-Frontier-Basee Data Augmentation: and adhering to safety constraints by leveraging a Pareto-frontier of the training data. The Pareto-fronties consists of trajectories that provide the highest possible return under specific safety constraints. Fitting polynomial the helps identify confficting high-return and safety constraint pairs, which are then used for augmentation. The augmentation generates synthetic trajectories by relabeling safe trajectories from the -roettffrrotier with higher returns and assigning higher or equal safety constraints. This encourages the model to imitate the most rewarding, safe trajectories when the desired return given the safety constraint is infeasible.
The technique tries to resolve the conflict between maximizing returns
In the 2D incompressible fluid experiment, the model fails to extrapolate safe trajectories with higher returns due to the training data covering a broad range of costs, while the desired cost lies in a narrow range. The augmentation treats all safe trajectories on the Pareto-frontier equally, without emphasizing the region of interest. This complexity. however, could potentially be addressed by SafeDiffCon, which is indicated in Table We use the official CDT implementation and follow DT guidelines to sweep desired returns and cost constraints in testing time. In the ID Burgers experiment, modify the control objective from the original mean squared error J between the prediction and target, to an exponential form exp (-J). This new objective is bounded within [0, 1J. which better aligns with the reward wwward-mximmzzing setup in reinforcement learning used by CDT. For the 2D incompressible fluid setup, the state, action, and cost prediction heads each consist of 3-layer MLPs with the transformer's s hidden dimension as the inner size.
1.2. BC-AII
The Behavior Cloning (BC) algorithm, introduced by (Pomerleau, 1988), is a foundational technique in imitation learning. BC is designed to derive policies directly from expert demonstrations, utilizing supervised learning to associate states with corresponding actions. This method eliminates the necessity for exploratory steps commonly required in reinforcement
21

---

From Uncertain to Safe: Conformal Adaptation of Diffusion Models for Safe PDE Control
Table 14: Hyperparameters of 1D CDT.
ID Burgers'
All data
Safe filtered
State Dimension
256
256
Action Dimension
128
128
Hidden Dimension
1024
1024
Number of Transformer Blocks
2
Number of Attention Heads
8
Horizon (Sequence Length)
5
5
Learning Rate
1e-4
le-4
Batch Size
64
64
Weight Decay
le-5
le-5
Learning Steps
1,000,000
1,000,000
Learning Rate Warmup Steps
500
500
Pareto-Frontier Fitted Polynomial Degree
Augmentation Data Percentage
0.3
0.3
Max Augment Reward
10.0
10.0
Min Augment Reward
1.0
1.0
Target Entropy
-128
-128
Testing Time Sweep Returns
9.0,9.9
9.0, 9.9
Testing Time Sweep Costs
0.0, 1.0, 2.0, 3.0
0.0, 1.0, 2.0, 3.0
learning by replicating the actions observed in expert demonstrations. One of the significant advantages of BC is that it does not involve interacting with the environment during the training phase, which streamlines the learning process and diminishes the demand for computational resources.
In this approach, a policy network trained using standard supervised learning strategics aimed at reducing the discrepancy between the actions predicted by the model and those performed by the expert in the dataset. The commonly used loss function for this purpose the mean squared error between the predicted actions and expert actions. The dataset for training comprises state-action pairs harvested from these expert demonstrations. In the work, employ the implementation (Liu et al., 2023a).
I.3. BC-Safe
Following the baseline in (Liu et al., 2023a), the BC-Safe fed with only safe trajectories filtered from the training dataset, satisfies most safety requirements, although with conservative performance and lower rewards. Others are same BC-All, except for the safe trajectories.
1.4. SL-Lag
Hwang et al. (2022) introduces supervised learning (SL) based method to control PDE systems. It first trains neural surrogate model to capture the PDE dynamics, which includes a VAE to compress PDE states and controls into the latent space and another model to learn the PDE's time evolution in the latent space. To obtain the optimal control sequence, SL can compute the gradient J. where is the control objective and f is the input control sequence. Then iterative gradient optimization can be exccuted to improve the control sequence.
To ensure that the optimal control is compatible with the hard constraint in our experiments, we follow Chow et al. (2018) apply the Lagrange optimization method to the constrained optimization. Specifically, we iteratively solve the optimization problem below:
max min J(w) + ((s((w)) c). A20
(48)
We denote the modified SI. method SI.-Lag.
22

---

From Uncertain to Safe: Conformal Adaptation of Diffusion Models for Safe PDE Control
Table 15: Hyperparameters of 2D CDT.
2D incompressible fluid
All data
Safe filtered
State Dimension
3x64x64
3x64x64
Action Dimension
2x64x64
2 2x64x61 66
Hidden Dimension
512
512
Number of Transformer Blocks
3
Number of Attention Heads
8
Horizon (Sequence Length)
10
10
Learning Rate
le-4
1e-4
Batch Size
8
Weight Decay
le-5
le-5
Learning Steps
1,000,000
1,000,000
Learning Rate Warmup Steps
500
500
Pareto-Fronttier Fitted Polynomial Degree
4
Augmentation Data Percentage
0.3
0.3
Max Augment Reward
32.0
32.0
Min Augment Reward
1.0
1.0
Target Entropy
-2x66xx64)
-(2x64x64)
Testing Time Sweep Returns
18.0, 32.0
18.0, 32.0
Testing Time Sweep Costs
0.0, 0.1, 0.2
0.0, 0.1, 0.2
Table 16: Hyperparameters of network architecture and training for SL-Lag in 1D Burgers' experiment.
Hyperparameter name
Value
Initialization value w Optimizer w
0.001
LBFGS
Learning rate of w
0.1
Initialization value of the Lagrange multiplier X
Optimizer of
SGD
Learning rate of X Iteration of A Loss function
10
2
MSE
1.S. MPC-Lag
We adopt the MPC with stochastic gradient descent and choose the planning step K = 1. The simulation model is the same as the one in SL-Lag. Also, the way we combine MPC with the Lagrangian method is also the same as SL-Lag. The hyperparameters are reported in Table 17.
L6. TREBI
In Lin et al. (2023), the diffusion model is adopted for the planning task under safety budgets. It generates trajectory under safety constraints using classifier- classifier-uiidance Dhariwal & Nichol (2021) by adding a safety loss to the reward guidance following Diffuser Janner al. (2022).
However, the original setting is different from our experiments. In our ID Burgers' equation control, our objective is that certain state equals the target state, which in in-painting diffusion condition Janner et al. (2022) is appropriate for. Furthermore, as in our method and in (Wei et al., 2024), a conditional diffusion model can be learned to tackle the objective more directly. In addition, TREBI follows the setting in Diffuser where the interaction with the environment is allowed which in our experiments becomes an MPC method. Note that the reported results of our method do not involve interaction with the surrogate model (though our method can casily adapted to be an MPC method). Thus, the results of TREBI in Table and 2 have an unfair advantage.
23

---

From Uncertain to Safe: Conformal Adaptation of Diffusion Models for Safe PDE Control
Table 17: Hyperparameters of for MPC-Lag in ID Burgers' experiment.
Hyperparameter name
Value
Initialization value of w Optimizer of w Learning rate of w
0.01
LBFGS
0.01
Initialization value of the Lagrange multiplier X
Optimizer of Learning rate of X Iteration of À Loss function
SGD
10
2
MSE
Therefore, for ID Burgers' experiment, we conducted different experiments on TREBI including (1) planning multiple times with interaction with the surrogate model or (2) planning only once, and with target state conditioning target state guidance. The target state guidance + planning multiple times turned out the best and is reported in Table 1. For 2D smoke control, the target is not state constraint but reward, and the planning multiple-step setting is too computationally expensive. To this end, we use reward guidance with planning one single time, which is identical to the ablation study of our method in Table 4. The hyperparameters of the ID experiment are reported in Table 18, and those of 2D are in Table 19.
Table 18: Hyperparameters of network architecture and training for TREBI in 1D Burgers' experiment.
Hyperparameter name
Value
Reward guidance intensity
50
Safety guidance intensity
Cost budget
0.64
Number of guidance steps
10
Denoising steps
200
Sampling algorithm
DDPM
U-Net dimension
64
U-Net dimension mitiplications
1, 4.
Planning horizon
8 steps
Optimizer
Adam
Learning rate
0.0002
Batch size
Loss function
MSE
I.7. PID
Proportional Integral Derivative (PID) control (Li et al., 2006) a versatile and effective method widely employed in numerous control scenarios. For ID control task, we mainly implement the PID bascline adapted from (Wei et al., 2024). More detailed configurations can be found in Table 20
24

---

From Uncertain to Safe: Conformal Adaptation of Diffusion Models for Safe PDE Control
Table 19: Hyperparameters of network architecture and training for TREBI in 2D Incompressible Fluid Control experiment.
Hyperparameter name
Value
Reward guidance intensity
5000
Safety guidance intensity
0.01
Cost budget
0.1
Number of guidance steps
1
Denoising steps
20
Sampling algorithm
DDPM
3D U-Net dimension
8
3D U-Net dimension mitiplications
2.
Planning horizon
8 steps
Optimizer
Adam
Learning rate
0.0002
Batch size
16
Loss function
MSE
Table 20: Hyperparameters of network architecture and training for ANN PID.
Hyperparameter name
Value
Kernel size of convld Padding of conv ld
Stride of conv ld
Activation function
Softsign
Batch size
16
Optimizer
Adam
Learning rate
0.0001
Loss function
MAE
25


Reviews:
[{"review_id": "Fhob7uvXpK", "reviewer": "Reviewer_54oM", "summary": "The paper introduces SafeDiffCon, a method integrating safety constraints into deep learning-based control of PDE systems through diffusion models. Addressing the gap in existing methods that neglect safety, SafeDiffCon employs conformal prediction to estimate uncertainty quantiles, which guide both post-training and inference phases. Evaluated on 1D Burgers equation, 2D incompressible fluid flow, and a nuclear fusion control problem, SafeDiffCon uniquely satisfies all safety constraints across tasks, outperforming classical and deep learning baselines in control performance. Key contributions include the integration of uncertainty quantification via conformal prediction, safety-constrained diffusion training, and adaptive inference mechanisms, demonstrating robust and safe PDE control in complex scenarios.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We greatly appreciate your recognition. Below are our responses.\n>Q1. Ablation studies on other two tasks, 1D Burgers' Equation and tokamak control.\n- Thanks for the suggestion. We conduct ablation studies on these two tasks. Results still show that the absence of any module affects the model, causing it to fail to meet the safety constraints. Therefore, each module is effective. These will be added to the manuscript.\n\nBurgers':\n|Methods|J|$R_{sample}$|$R_{time}$|$R_{point}$|\n|-|-|-|-|-|\n|SafeDiffCon|0.0011|0%|0%|0%|\n|w/o post-training|0.0014|4%|50%|0.01%|\n|w/o fine-tuning|0.0007|40%|20%|3%|\n|w/o Q|0.0006|30%|10%|1%|\n\nTokamak:\n|Methods|J|$R_{sample}$|$R_{time}$|\n|-|-|-|-|\n|SafeDiffCon|0.0094|0%|0%|\n|w/o post-training|0.0153|10%|0.52%|\n|w/o fine-tuning|0.0210|50%|7.74%|\n|w/o Q|0.0269|28%|0.65%|\n\n>Q2. Analysis on parameters, such as coverage probability α, the split ratio of training/calibration data, and the weight of objective γ.\n- Great suggestion! We conduct experimental analysis on these parameters. The tables show that they do not impact much on performance, indicating that the model is robust. This will be added to the next version.\n- Tokamak:\n\n|α|J|$R_{sample}$|$R_{time}$|\n|-|-|-|-|\n|0.8|0.01017|0%|0%|\n|0.85|0.0091|0%|0%|\n|0.9|0.0094|0%|0%|\n|0.95|0.0095|0%|0%|\n\n|Split ratio|J|$R_{sample}$|$R_{time}$|\n|-|-|-|-|\n|0.005|0.0124|4%|0.03%|\n|0.01|0.0093|0%|0%|\n|0.02|0.0094|0%|0%|\n- Incompressible fluid:\n\n|γ|J|SVM|$R_{sample}$|\n|-|-|-|-|\n|0.01|0.3548|0|0%|\n|0.1|0.4953|0.004|2%|\n|0.3|0.498|0.003|2%|\n\n>Q3. How reliable are the conformal intervals in SafeDiffCon for safety scores on novel control tasks, given the distribution shift between calibration data and the optimal control distribution during inference?\n- Good questions! Novel control tasks include **three cases**: new equation forms (i.e., dynamics), (1) new safety constraints and bounds, and (2) a shift in the optimal control distribution and (3) the distribution of the calibration set.\n  - We primarily discuss **the third case** in the paper. Both theoretical and experimental results demonstrate that our method can adapt to this distribution shift, ensuring that the actual score is covered. However, the theoretical analysis assumes that the diffusion model's approximation of the training set distribution during pretraining is not too poor.\n  - We lack discussion of **the first and second cases** in the paper, which will be added in the next version. Since Eq. 5 and Eq. 11 are tied to specific dynamics and safety constraints, the conformal interval doesn't generalize to new ones. However, we can **skip the post-training and use inference-time finetuning** to compute conformal prediction intervals based on new tasks, ensuring reliable coverage guarantees.\n\n>Q4. Why choose diffusion models? What advantages do they offer over other generative or deep learning methods? Where specifically does the synergy between your uncertainty quantification framework and the diffusion model architecture manifest? Could similar conformal techniques applied to other models achieve comparable safety and performance?\n- Since our approach to uncertainty quantification is **probabilistic**, we require a model with an explicit probabilistic framework. Diffusion models meet this requirement, as they allow us to **analyze the relationship between the generated distribution and training set distribution**. This enables us to incorporate probability-related theory into its training and sampling algorithms, such as our consideration of distribution shift and the design of reweighted loss.\n- With multi-step denoising, diffusion models have **strong modeling capabilities**, allowing it to handle high-dimensional, long-term problems. It has been used in weather and PDE system modeling tasks [1, 2]. Additionally, because it denoises from the Gaussian distribution and **models the entire trajectory rather than transition pairs**, it can perform **global optimization**, and has been applied in many decision-making scenarios [3, 4]. Furthermore, studies have shown that it is **robust**, which means it can withstand noise attacks [5] and is critical in safe control. **Experimental results** also demonstrate its superiority in control problems.\n- **Alternative model architectures** can also draw from ideas of conformal prediction. However, due to the strong capabilities of diffusion models and their ability to explicitly model the relationship between training and generated distributions, they integrate with conformal prediction more naturally and effectively.\n\n>Q5. Typos.\n- Thanks. We will correct them.\n\n**References** \n\n[1] Dyffusion: a dynamics-informed diffusion model for spatiotemporal forecasting.\n\n[2] Synthetic Lagrangian turbulence by generative diffusion models.\n\n[3] Planning with diffusion for flexible behavior synthesis.\n\n[4] Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning.\n\n[5] Diffusion Models are Certifiably Robust Classifiers.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "Yes. I have checked the proofs.", "experimental_designs_or_analyses": "Following previous works, the designs of experiments are clear and well-organized. The analysis is also convincing.\n\nTo improve the quality of the paper, I suggest the authors provide more ablation analysis on the other two tasks, 1D Burgers  Equation and nuclear fusion control, to better demonstrate the effectiveness of the proposed method.", "relation_to_broader_scientific_literature": "The work builds upon existing diffusion models for physical control (like Wei et al., 2024; Hu et al., 2024) but extends them with conformal prediction techniques (Vovk et al., 2005; Tibshirani et al., 2019) to handle distribution shifts between training data and desired safe controls. The approach combines post-training with reweighted loss functions and inference-time fine-tuning, enabling models to satisfy safety constraints while optimizing control objectives. This relates to safe offline reinforcement learning methods like CPQ (Xu et al., 2022), COptiDICE (Lee et al., 2022), and TREBI (Lin et al., 2023), but differs by specifically addressing PDE-constrained systems and using uncertainty quantiles based on conformal prediction to provide safety guarantees. The work demonstrates applications in three different physical domains (1D Burgers' equation, 2D incompressible fluid dynamics, and controlled nuclear fusion), showing greater effectiveness than traditional control methods (PID, MPC-Lag) and deep learning alternatives.", "essential_references_not_discussed": "To the best of my knowledge, the paper has discussed all essential references.", "other_strengths_and_weaknesses": "Strengths:\n\n- The paper addresses a critical gap in existing methods by integrating safety constraints into deep learning-based control of PDE systems.\n\n- This paper is well-written and clearly organized, with a strong motivation and clear contributions.\n\n- The experimental results are well-presented and provide a comprehensive evaluation of the proposed method in various physical domains, highlighting its effectiveness.\n\nWeakness:\n\n- The effectiveness of the uncertainty quantile seems to be related to the choice of various parameters, such as coverage probability $\\alpha$, the split ratio of training/calibration data, and the weight of objective $\\gamma$. It is suggested to provide more ablation analysis on these parameters to better understand their impact on the performance of the proposed method.\n\n- As mentioned above, more ablation analysis on the other two tasks, 1D Burgers Equation and nuclear fusion control, would further demonstrate the effectiveness of the proposed method.", "other_comments_or_suggestions": "There are a few typos in the paper that need to be corrected:\n\n- On page 6, in Figure 2's caption,  the authors refer to their method as \"SafeConPhy\" instead of \"SafeDiffCon\".\n\n- Same on page 19, in appendix H.1, the authors refer to their method as \"SafeConPhy\" instead of \"SafeDiffCon\".\n\n- On page 22, in the \"H.7. PID\" section, there's a typo in the first sentence where \"Propercentageal\" is written instead of \"Proportional\" when describing the PID control method.", "questions_for_authors": "- To what extent can the conformal prediction intervals in SafeDiffCon provide reliable coverage guarantees for safety scores when applied to novel control tasks during inference, especially considering the distribution shift between calibration data and the optimal control distribution?\n\n- While the paper presents an interesting approach to safe PDE control, I'm curious about the specific choice of diffusion models as the foundation. Could you elaborate on the intrinsic advantages that diffusion models provide for your method compared to other generative or deep learning approaches? Where specifically does the synergy between your uncertainty quantification framework and the diffusion model architecture manifest? Would applying similar conformal adaptation techniques to alternative model architectures yield comparable safety guarantees and performance improvements? These are my main concerns of the paper.\n\n## Post Rebuttal\nMost of my concerns are appropriately addressed. I choose to increase the score.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HdEsh8mJmT", "reviewer": "Reviewer_9N9K", "summary": "This paper introduces SafeDiffCon, a method that integrates safety constraints into diffusion models for PDE control tasks. By leveraging conformal prediction to quantify model uncertainty, the approach employs post-training with a reweighted loss and inference-time fine-tuning to align generated control sequences with safety requirements. Experiments on 1D Burgers’ equation, 2D fluid dynamics, and nuclear fusion control demonstrate that SafeDiffCon uniquely satisfies safety constraints while achieving superior control performance compared to baselines like BC, MPC-Lag, and TREBI. The key innovation lies in combining uncertainty-aware conformal adaptation with diffusion models to address distribution shifts in offline settings.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your insightful comments. Below are our responses.\n>Q1. Inference time cost comparison with baselines. The proposed model's inference finetune might induce significant extra time cost.\n- Great suggestion. We would like to point out that we **have accelerated the efficiency of inference**. On one hand, by introducing post-training, the data distribution is already close to the desired distribution before inference. On the other hand, we have significantly accelerated the diffusion model's sampling time using DDIM [9].\n- The table below presents a **comparison of the inference time between our method and the baseline** on the Burgers' equation on A800 with 8 CPUs. It shows that our method achieves a moderate speed and is much faster compared to TREBI, which is also based on the diffusion model.\n\n|Methods|BC|PID|SL-Lag|MPC-Lag|CDT|TREBI|Ours|\n|-|-|-|-|-|-|-|-|\n|Inference Time (min)|0.1351|0.1091|0.8842|26.2905|0.0890|13.5525|2.3575|\n\n>Q2. Why is the diffusion model chosen for the control problem, which is essentially different from the generation task?\n- In fact, diffusion models have already been **widely used** in decision-making tasks, including scenarios such as PDE systems [1, 2], robotics [3, 4], and traditional reinforcement learning [5, 6, 7]. This is because the task can naturally be modeled as a probabilistic distribution over state and control sequences, and the diffusion model has many unique advantages in handling such tasks. \n  - It has superior modeling capabilities due to its ability to capture complex distributions and generate high-fidelity outputs by progressively refining predictions over multiple denoising steps.\n  - By denoising from a Gaussian distribution, it models the entire trajectory, meaning it learns to generate states at different time steps simultaneously, aiding in capturing long-range dependencies. So it can perform **global optimization**.\n  - Furthermore, studies have shown that it is **robust** to noise which is essential in safe control problems [8].\n\n>Q3. What PDE is the Tokamak control problem driven by?\n- Tokamak control involves the coupling of multiple PDEs. The primary equation is the Grad–Shafranov equation, along with others such as the Heat Transport Equation and the Skin Effect Equation [10].\n\n**References**\n\n[1] A generative approach to control complex physical systems.\n\n[2] Wavelet diffusion neural operator.\n\n[3] Diffusion policy: Visuomotor policy learning via action diffusion.\n\n[4] Hierarchical diffusion policy for kinematics-aware multi-task robotic manipulation.\n\n[5] Planning with diffusion for flexible behavior synthesis.\n\n[6] Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning.\n\n[7] Offline Reinforcement Learning via High-Fidelity Generative Behavior Modeling.\n\n[8] Diffusion Models are Certifiably Robust Classifiers.\n\n[9] Denoising Diffusion Implicit Models.\n\n[10] Reconstruction of current profile parameters and plasma shapes in tokamaks.", "claims_and_evidence": "The paper claims that SafeDiffCon achieves optimal control under safety constraints, and the experiment results support the claim.", "methods_and_evaluation_criteria": "The post-training and fine-tuning methods are suitable for constrained problems. The benchmark datasets Burgers, NS, and Tokamak make sense for the control applications.", "theoretical_claims": "The theoretical claims seem to be correct.", "experimental_designs_or_analyses": "The inference efficiency, e.g., the inference time cost is not compared. The proposed model's inference finetune might induce significant extra time cost.", "relation_to_broader_scientific_literature": "Related to Control Theory, Numerical Methods.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "No", "other_comments_or_suggestions": "No", "questions_for_authors": "1. Can you provide inference time cost comparison with baselines?\n2. Why is the diffusion model chosen for the control problem, which is essentially different from the generation task?\n3. What PDE is the Tokamak control problem driven by?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tWX9Xxfe43", "reviewer": "Reviewer_hkaj", "summary": "This paper introduces an approach that maintains safety constraints in PDE-constrained control problems. It employs uncertainty estimation with conformal prediction to optimize control while preserving safety. It fine-tunes a diffusion model using conformal prediction to produce safe control sequences. The experimental results in tasks like 1D Burger equation show the effectiveness of the proposed method.\n\n\n## update after rebuttal\nMy major concerns, such as running time, have been resolved. Therefore, I maintain my score as weak accept.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the constructive review. Below are the responses.\n>Q1. Lack of evaluation of its performance in real-world scenarios.\n- In fact, the tokamak control in the paper is a near real-world experiment for controlled nuclear fusion which is a highly-nonlinear and highly coupled system. This environment is trained using real data collected from the KSTAR tokamak device (https://github.com/jaem-seo/KSTAR_tokamak_simulator). Additionally, its simulation has been tested with many real discharges, showing reasonable predictions and acceptable prediction accuracy.\n\n>Q2. The performance of the proposed framework heavily relies on the collected training data.\n- Compared to baselines, our method **does not rely more heavily on the collected training data**. The table below shows the average unsafety rate of the training set and the generated data. It can be seen that the distribution of data generated by our method differs from that of the training data. Moreover, only our method achieves full safety.\n\n| Datasets|Training data|Generated data|\n| -| -| -|\n| Burgers' equation|89.7%|0%|\n| Incompressible fluid|53.1%|0%|\n| Tokamak fusion reactor|71.2%|0%|\n- As mentioned in the third and fourth paragraphs of the Introduction, interacting with the environment in the safe PDE control problem is hazardous. Therefore, we choose a more appropriate **offline setting**, where we can only rely on the data that has already been collected.\n\n>Q3. The introduced framework is a little complex.\n- Overall framework:\n  - The method begins with the introduction of **uncertainty quantification** through the concept of uncertainty quantiles. This is integrated throughout the algorithm to address potential gaps between the actual and predicted safety scores, which could lead to unsafe events not anticipated by the model. By introducing them, we address the critical issue of predictive uncertainty and its impact on safety scores, enabling more robust and safe control actions.\n  - After pre-training, we employ a reweighted loss function to **post-train** the model’s output distribution, guiding it toward regions with better safety and more optimal objectives.\n  - During the **inference** phase, task-specific fine-tuning is performed on the post-trained model, allowing it to achieve better control performance and stronger safety guarantees for specific tasks, with minimal adjustments needed.\n- To help understand, we have present our proposed method with the framework outlined in **Figure 1, Algorithm 1 and the first paragraph of Method**. We are happy to **provide a clearer explanation** of the algorithm at the beginning of the Method section to enhance readability.\n\n>Q4. Comparison of running time.\n- Thanks for the good suggestion. The table below presents the **comparison of the inference time between our method and baselines** on the Burgers' equation on A800 with 8 CPUs. We would like to highlight that we have enhanced the inference efficiency using several techniques. Firstly, by introducing post-training, the data distribution is already closely aligned with the desired distribution prior to inference. Secondly, we significantly speed up the sampling process of the diffusion model using DDIM [1]. As a result, our method demonstrates a reasonable inference speed and is considerably faster than TREBI, which is also based on the diffusion model.\n\n|Methods|BC|PID|SL-Lag|MPC-Lag|CDT|TREBI|Ours|\n|-|-|-|-|-|-|-|-|\n|Inference Time (min)|0.1351|0.1091|0.8842|26.2905|0.0890|13.5525|2.3575|\n\n**References** \n\n[1] Denoising Diffusion Implicit Models.", "claims_and_evidence": "Yes. I think the claims are correct and clear.", "methods_and_evaluation_criteria": "Yes. The evaluation criteria follow previous works and effectively assess performance. And the method is proposed to solve the safety issue,", "theoretical_claims": "Yes. I think it is majorly correct.", "experimental_designs_or_analyses": "The paper designs good experimental designs that assesses SafeDiffCon performance across three PDE control tasks: 1D Burgers’ equation, 2D incompressible fluid, and controlled nuclear fusion. The chosen baseliens are also relatively new. Therefore, the experimental design is valid.", "relation_to_broader_scientific_literature": "This research makes a good contribution by combining diffusion models, conformal uncertainty quantification, and safe control frameworks within PDE-constrained control systems which fills a research gap.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "**Strengths**\n- This paper is well-written. The proof is well performed.\n- The evaluation is comprehensive, and the experimental results show the effectiveness of the proposed method.\n\n**Weakness**\n- Lack of evaluation of its performance in real-world scenarios.\n- The performance of the proposed framework heavily relies on the collected training data.\n- The introduced framework is a little complex, and the time needed is unclear.", "other_comments_or_suggestions": "1. It would be better for the authors to provide a table about the running time of the proposed methods and the baselines.", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "q2qh0U2HX9", "reviewer": "Reviewer_yisF", "summary": "This paper proposes a safe PDE control method based on the diffusion model inspired by conformal prediction. Specifically, the authors propose two new phases of post-training and inference-time fine-tuning to accommodate the quantified uncertainty score from the calibration set for the safety constraint and objective. Experiments show that the proposed method can achieve better overall results and meet safety constraints compared to baselines.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for comments. Below are our responses.\n>Q1. Problem formulation (Eq. 1): like offline safe RL, no boundary and time conditions.\n- \\(C(u,w)=0\\) in Eq. 1 is the PDE constraint. Following your advice, we will add boundary and time conditions along with explanations.\n\n>Q2. PDE-safe control setting may not be ideal, as safe RL can formulate the same. The method is decoupled from PDE and applicable to safe RL.\n- PDE control scenario is **important and complex**. Its dimensions are high, and dynamics are nonlinear.\n- We evaluate several **SOTA safe RL baselines** like CDT and TREBI.\n- Our paper's motivation lie in **highlighting importance of safe PDE control and introducing it into deep learning-based control**. To address this, we design methods and develop **experimental scenarios, datasets, and baselines across domains**.\n- Given **PDEs' high-dimensional and complex dynamics**, we choose diffusion models for their strong modeling power [1, 2]. To satisfy **initial conditions**, we apply conditional generation to enforce strict adherence.\n- To address the concern, we provide results on **a safe RL benchmark** [3]. https://anonymous.4open.science/r/Safe-Diffusion-Models-for-PDE-Control-213C/Rebuttal.md\n\n>Q3. The method is some fine-tuning on the calibration set and trivial. Post-training and finetuning are similar and not novel.\n- Our method goes **beyond finetuning**, incorporating post-training, guidance and key uncertainty quantification via conformal prediction (CP).\n- It adjusts **not only on the calibration set*. Post-training uses both training and calibration sets, while finetuning and guidance are applied to control targets and the calibration set.\n- Post-training and inference-time finetuning adjust outputs **completely differently**: Post-training uses a reweighted loss (Eq. 12), while finetuning uses a sampling-based loss (Eq. 15) with additional guidance.\n\n>Q4. Weighted CP directly follow Eq. 7; Eq. 12 directly follows Eq. 3.\n- We do **not directly use Eq. 7**. As noted in Section 4.2, applying it requires estimating the intractable ratio between generated and dataset distributions. To address this, we design a loss (Eq. 12) and prove Theorem 4.2 and 4.3 to justify using Eq. 13 for valid coverage.\n- Eq. 12 **does not directly follow Eq. 3**. Loss in Eq. 12 is designed to both enable ratio estimation and promote safer, more optimal distributions. Its form is derived in Appendix A.\n\n>Q5. Exchangeability for CP breaks in sequential decision-making due to correlations between \\(u\\), \\(w\\), and within trajectories.\n- As noted after Eq. 5, the distribution is over full state and control trajectory pairs, so dependencies between \\(u\\) and \\(w\\), and within trajectories don't affect exchangeability.\n\n>Q6. Theories directly follow CP literature but under different settings, and are not critical to method's significance.\n- Except for Theorem 4.1, theorems are largely **independent of existing CP theory**, with separate derivations.\n- Theorems are **essential**: Theorem 4.2 aligns post-training distribution with the desired form, and Theorems 4.1–4.3 jointly ensure valid conformal intervals for the post-trained model.\n\n>Q7. Experiment on PDE ContRol Gym for off-the-shelf baselines and fair comparison. \n- This benchmark is an **online** setting **not for safe PDE control**. As noted in Introduction, interaction with environments in safe control is risky, so we focus on offline control. In the absence of safe PDE benchmarks, we **contribute safe datasets and environments**. Our experiments follow **widely used prior works** (PhiFlow: 1.6k stars) and use **the same equations** as PDE ContRol Gym, including Burgers' and NS.\n- **This benchmark's baselines** are **not for safe control**, while we compare **SOTA safe RL methods**. Our comparison is **fair**, with open-sourced code in the paper and tuned hyperparameters.\n- We will **cite** this valuable benchmark in Related Works of the next version.\n\n>Q8. Some baselines achieve 0% unsafety (Tables 1 and 3), so strengthening constraints.\n- There are two criteria for safe control methods: **safety and control accuracy under safety**. If only one method is safe, we can't compare the second criterion, limiting comprehensive analysis.\n\n>Q9. Can this method adapt to different safety constraints without finetuning?\n- **Different constraints**: Our method remains applicable by computing each safety score minus its bound, combining them via smooth max [4], and constraining the result to be below 0, as in the current method.\n- **Without finetuning**: We can take guidance to direct generated data to meet constraints.\n\n>Q10. References.\n- We will add them in the paper.\n\n**Reference** \n\n[1] Synthetic Lagrangian turbulence by generative diffusion model.\n\n[2] DiffusionPDE: Generative PDE-solving under partial observation.\n\n[3] Datasets and Benchmarks for Offline Safe Reinforcement Learning.\n\n[4] Magnetic control of tokamak plasmas through deep reinforcement learning.s", "claims_and_evidence": "From the perspective of PDE control theory, this paper solves a similar problem of in-domain PDE control, but the problem formulation in Eq (1) is more like offline safe RL because it misses the boundary and time conditions for the PDE dynamical system. I don't think the PDE safe control setting is the best manner to showcase the idea of a diffusion model with a conformal prediction for the safe sequential decision-making task. That is to say, as shown in (Liu et al. 2023a), there are tons of offline safe RL datasets and benchmarks to formulate the constrained optimization problem in Eq (1)  and validate the methodology. The method part is disconnected from the motivation of the safe PDE control problem and can work on other general safe RL settings if it is a genuinely effective method.", "methods_and_evaluation_criteria": "There are some concerns and comments for method part.\n\n- Overall, the method is basically some fine-tuning over the calibration set so that the performance gets better, which is trivial and lacks in-depth insight. Post-train and inference-time fine-tuning are similar and not novel and impressive to the audience.\n\n- The distribution shift in Eq (7) is directly handled through previous weighted conformal prediction (Tibshirani et al. 2019), and post-training of Eq (12) directly follows the preliminary part of Eq (3), which greatly weakens the contribution and novelty.\n\n- Regarding conformal prediction, I agree that exchangeability is the key assumption but it may not hold in sequential decision-making settings. However, the reasons are not only the distribution shift but also the correlation of sequential data. It seems to be problematic to use the sample $(u,w)$ as the probability density function in Eq (7) because of such correlation between $u,w$ and within each trajectory as well.", "theoretical_claims": "The theories and proofs directly follow the conformal prediction literature but under different settings. The theorems are not essential to the significance of the proposed method since it is very mature in conformal prediction.", "experimental_designs_or_analyses": "I appreciate the different PDE settings and multiple baselines in the experiment part. However, it is expected to experiment on the PDE control benchmark https://github.com/lukebhan/PDEControlGym for off-the-shelf baselines and for fair comparison. Besides, in Table 1 and Table 3, some baselines can achieve 0% regarding safety constraints so it is expected to make the safety constraint harder to be satisfied to highlight the performance of the proposed method. I also wonder if the proposed method can do inference-time scaling, e.g. to adapt to different safety constraints without fine-tuning.", "relation_to_broader_scientific_literature": "It is benificial to scientific ML and AI4Science.", "essential_references_not_discussed": "Yes, some essential references are not discussed. PDE control problem has been well studied in the control and system community [1,2,3,4]. Another recent paper studies safe PDE boundary control [5], which should be discussed and compared. Also, regarding safe RL and safe control, control barrier function (CBF) based methods are missing [6,7].\n\n\n\n\n---\n[1] Krstic et al.  Boundary control of PDEs: A course on backstepping designs, 2008\n\n[2] Smyshlyaev et al. Adaptive control of parabolic PDEs, 2010\n\n[3] Zhang et al.  Controlgym: Large-scale control environments for benchmarking reinforcement learning algorithms, 2024\n\n[4] Bhan et al. Pde control gym: A benchmark for data-driven boundary control of partial differential equations, 2024\n\n[5] Hu et al. On the Boundary Feasibility for PDE Control with Neural Operators, 2024\n\n[6] Dawson et al. Safe control with learned certificates: A survey of neural lyapunov, barrier, and contraction methods for robotics and control. 2023\n\n[7] Zhao et al. Model-free safe control for zero-violation reinforcement learning, 2021", "other_strengths_and_weaknesses": "See above.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "See above.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper introduces SafeDiffCon, a method integrating safety constraints into deep learning-based control of PDE systems through diffusion models.  In relation to existing work, this work builds upon existing diffusion models for physical control but extends them with conformal prediction techniques to handle distribution shifts between training data and desired safe controls.  There was limited discussion on this paper: reviewer 54oM requested an ablation analysis, which the authors provided, and which lead to an increased score.  Reviewer yisF raised two key concerns, firstly that the method is "essentially finetuning" and second that they would like to see evaluation on the PDE control suite.  These concerns were adequately addressed in rebuttal, again leading to an increased score.  After discussion there is consensus that this paper be accepted, and I concur.
Decision: Accept (poster)