Official repo for the paper DiffPhyCon: A Generative Approach to Control Complex Physical Systems.
Long Wei*, Peiyan Hu*, Ruiqi Feng*, Haodong Feng, Yixuan Du, Tao Zhang, Rui Wang, Yue Wang, Zhi-Ming Ma, Tailin Wu
NeurIPS 2024.
We introduce a novel method, DiffPhyCon, for controlling complex physical systems using generative models, by minimizing the learned generative energy function and specified control objectives. Further, we enhance DiffPhyCon with prior reweighting, enabling the discovery of control sequences that significantly deviate from the training distribution.
Framework of DiffPhyCon:
Movement and fluid field visualization on the jellyfish controlled by DiffPhyCon:
Density of smoke (left) and control signals (middle and right) controlled by DiffPhyCon:
Run the following commonds to install dependencies. In particular, when run the smoke control task, the python version must be 3.8 due to the requirement of the Phiflow software.
conda env create -f environment.yml
conda activate DiffPhyCon
The checkpoints of both our DiffPhyCon and baselines on the three tasks (1D Burgers', 2D jellyfish, 2D smoke), and our released Jellyfish datasets can be downloaded in link. To run the following training and inference scripts locally, replace the path names in filepath.py by your local paths.
python /model/pde_1d_surrogate_model/burgers_operator.py --date_time "2024-01-09_1d_surrogate_partial_ob_partial_ctr_autoregress-3" --epochs 500 --gpu 0 --train_batch_size 151 --autoregress_steps 3 --dataset_path "/dataset_control_burgers/free_u_f_1e5_front_rear_quarter" --is_partially_observable 1 --is_partially_controllable 1 --lr 0.001
In the diffusion_1d folder, run
bash script/train_pw_XOXC.sh
bash script/train_XOXC.sh
where XOXC stands for one of the setting names: FOPC, POFC, and POPC.
In the scrpits/ folder, for the full observation scenario, run the following script to train a diffusion model for joint distribution of the state trajectory and control signals:
bash jellyfish_train_joint_full.sh
run the following script to train an additional diffusion model for the control signels prior:
bash jellyfish_train_w_full.sh
Similarly, for the partial observation scenario, run the following two scripts, respectively:
bash jellyfish_train_joint_partial.sh
bash jellyfish_train_w_partial.sh
In the scrpits/ folder, run the following script to train a diffusion model for joint distribution of the state trajectory and control signals:
bash smoke_train_joint.sh
run the following script to train an additional diffusion model for the control signels prior:
bash smoke_train_w.sh
In the diffusion_1d folder, run
bash script/eval_XOXC.sh
where XOXC stands for one of the setting names: FOPC, POFC, and POPC.
In the scripts/ folder, for the full/partial observation scenarios, run the following two scripts, repectively:
bash jellyfish_inference_full.sh
bash jellyfish_inference_partiall.sh
In the scripts/ folder, run the following script:
bash smoke_inference.sh
-
CinDM (ICLR 2024 spotlight): We introduce a method that uses compositional generative models to design boundaries and initial states significantly more complex than the ones seen in training for physical simulations.
-
LAMP (ICLR 2023 spotlight): First fully DL-based surrogate model that jointly optimizes spatial resolutions to reduce computational cost and learns the evolution model, learned via reinforcement learning.
If you find our work and/or our code useful, please cite us via:
@inproceedings{
wei2024diffphycon,
title={DiffPhyCon: A Generative Approach to Control Complex Physical Systems},
author={Wei, Long and Hu, Peiyan and Feng, Ruiqi and Feng, Haodong and Du, Yixuan and Zhang, Tao and Wang, Rui and Wang, Yue and Ma, Zhi-Ming and Wu, Tailin},
booktitle={The Thirty-eighth Annual Conference on Neural Information Processing Systems},
year={2024},
url={https://openreview.net/forum?id=MbZuh8L0Xg}
}