DiffPhyCon: A Generative Approach to Control Complex Physical Systems (NeurIPS 2024)

Paper | arXiv

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:

Installation

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

Dataset and checkpoints

Dataset

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.

Training:

1D Burgers' Equation Control:

Surrogate model

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

DiffPhyCon

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.

2D Jellyfish Control:

DiffPhyCon

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

2D Smoke Control:

DiffPhyCon

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

Inference:

1D Burgers' Equation Control:

DiffPhyCon

In the diffusion_1d folder, run

bash script/eval_XOXC.sh

where XOXC stands for one of the setting names: FOPC, POFC, and POPC.

2D Jellyfish Control:

DiffPhyCon

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

2D Smoke Control:

DiffPhyCon

In the scripts/ folder, run the following script:

bash smoke_inference.sh

Citation

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}
}