From Uncertain to Safe: Conformal Adaptation of Diffusion Models for Safe PDE Control

Thanks for comments. Below are our responses.

Q1. Problem formulation (Eq. 1): like offline safe RL, no boundary and time conditions.

• $C(u,w)=0$ in Eq. 1 is the PDE constraint. Following your advice, we will add boundary and time conditions along with explanations.

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.

• PDE control scenario is important and complex. Its dimensions are high, and dynamics are nonlinear.
• We evaluate several SOTA safe RL baselines like CDT and TREBI.
• 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.
• 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.
• 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

Q3. The method is some fine-tuning on the calibration set and trivial. Post-training and finetuning are similar and not novel.

• Our method goes beyond finetuning, incorporating post-training, guidance and key uncertainty quantification via conformal prediction (CP).
• 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.
• 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.

Q4. Weighted CP directly follow Eq. 7; Eq. 12 directly follows Eq. 3.

• 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.
• 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.

Q5. Exchangeability for CP breaks in sequential decision-making due to correlations between $u$, $w$, and within trajectories.

• 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.

Q6. Theories directly follow CP literature but under different settings, and are not critical to method's significance.

• Except for Theorem 4.1, theorems are largely independent of existing CP theory, with separate derivations.
• 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.

Q7. Experiment on PDE ContRol Gym for off-the-shelf baselines and fair comparison.

• 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.
• 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.
• We will cite this valuable benchmark in Related Works of the next version.

Q8. Some baselines achieve 0% unsafety (Tables 1 and 3), so strengthening constraints.

• 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.

Q9. Can this method adapt to different safety constraints without finetuning?

• 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.
• Without finetuning: We can take guidance to direct generated data to meet constraints.

Q10. References.

• We will add them in the paper.

Reference

[1] Synthetic Lagrangian turbulence by generative diffusion model.

[2] DiffusionPDE: Generative PDE-solving under partial observation.

[3] Datasets and Benchmarks for Offline Safe Reinforcement Learning.

[4] Magnetic control of tokamak plasmas through deep reinforcement learning.s

Thanks for the constructive review. Below are the responses.

Q1. Lack of evaluation of its performance in real-world scenarios.

• 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.

Q2. The performance of the proposed framework heavily relies on the collected training data.

• 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.

• 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.

Q3. The introduced framework is a little complex.

• Overall framework:
  • 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.
  • 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.
  • 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.
• 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.

Q4. Comparison of running time.

• 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.

References

[1] Denoising Diffusion Implicit Models.

Thanks for your insightful comments. Below are our responses.

Q1. Inference time cost comparison with baselines. The proposed model's inference finetune might induce significant extra time cost.

• 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].
• 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.

Q2. Why is the diffusion model chosen for the control problem, which is essentially different from the generation task?

• 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.
  • 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.
  • 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.
  • Furthermore, studies have shown that it is robust to noise which is essential in safe control problems [8].

Q3. What PDE is the Tokamak control problem driven by?

• 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].

References

[1] A generative approach to control complex physical systems.

[2] Wavelet diffusion neural operator.

[3] Diffusion policy: Visuomotor policy learning via action diffusion.

[4] Hierarchical diffusion policy for kinematics-aware multi-task robotic manipulation.

[5] Planning with diffusion for flexible behavior synthesis.

[6] Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning.

[7] Offline Reinforcement Learning via High-Fidelity Generative Behavior Modeling.

[8] Diffusion Models are Certifiably Robust Classifiers.

[9] Denoising Diffusion Implicit Models.

[10] Reconstruction of current profile parameters and plasma shapes in tokamaks.

We greatly appreciate your recognition. Below are our responses.

Q1. Ablation studies on other two tasks, 1D Burgers' Equation and tokamak control.

• 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.

Burgers':

Tokamak:

Q2. Analysis on parameters, such as coverage probability α, the split ratio of training/calibration data, and the weight of objective γ.

• 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.
• Tokamak:

• Incompressible fluid:

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?

• 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.
  • 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.
  • 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.

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?

• 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.
• 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.
• 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.

Q5. Typos.

• Thanks. We will correct them.

References

[1] Dyffusion: a dynamics-informed diffusion model for spatiotemporal forecasting.

[2] Synthetic Lagrangian turbulence by generative diffusion models.

[3] Planning with diffusion for flexible behavior synthesis.

[4] Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning.

[5] Diffusion Models are Certifiably Robust Classifiers.