Code base for Variational conditional normalizing flows for computing second-order mean field control problems
- Run the following script to prepare the environment
git clone git@github.com:jiaxi98/cnf_ot.git
python -m venv venv/cnf_ot
source venv/cnf_ot/bin/activate
cd cnf_ot
pip install .- Modify config/mfc.yaml for different settings (see following instructions ) and solve the problem via:
python cnf_ot/mfc/solvers.pyOur codebase supports solving three kinds of problems: optimal transport, regularized Wasserstein proximal operator, and Fokker-Planck equation. Check our paper for mathematical formulations of these problems and the parameters.
optimal transport: switch type to ``ot'' to solve
type: ot # ot, rwpo, fp
ot:
subtype: free # free, obstacle
subtype "free" corresponds to the original optimal transport problem while subtype "obstacle" corresponds to optimal transport with a soft obstacle between source and target, passing through which has an additional cost.
regularized Wasserstein proximal operator: switch type to ``rwpo'' to solve
type: rwpo # ot, rwpo, fp
rwpo:
T: 1
beta: 1
a: 1
pot_type: quadratic # quadratic, double_well
Fokker-Planck equation: switch type to ``fp'' to solve
type: fp # ot, rwpo, fp
fp:
T: 1
a: 1 # drift coeff
sigma: .5
velocity_field_type: gradient # gradient, nongradient
If you find this codebase useful for your research, please consider citing:
@article{zhao2025variational,
title={Variational conditional normalizing flows for computing second-order mean field control problems},
author={Zhao, Jiaxi and Zhou, Mo and Zuo, Xinzhe and Li, Wuchen},
journal={arXiv preprint arXiv:2503.19580},
year={2025}
}