This repository accompanies the paper "Solving Constrained Optimization Problems as ODE-based Models Using Reinforcement Learning", published in Transactions on Machine Learning Research (TMLR), June 2026.

Han Meng, Xinsong Feng, Yang Li, Chenan Wang, Kishansingh Rajput, Malachi Schram, Haipeng Chen

OpenReview: https://openreview.net/forum?id=QW0ZX4zRC2

CMFO (Constrained Markov Flow Optimizer) trains a flow-matching model within a reinforcement-learning paradigm as a learnable refinement mechanism for constrained optimization. A relaxComplete operator preserves the full RL reward signal while assisting constraint enforcement.

.
├── train.py                # Main training entry point for CMFO
├── load_results.py         # Aggregate experiment results into summary dicts
├── bash.sh                 # Example SLURM submission script
├── environment.yml         # Conda environment specification
├── src/
│   ├── agents/             # Flow-matching policy, models, helpers
│   ├── utils/              # Problem definitions, data utilities, default args
│   └── datasets/
│       ├── simple/         # Convex QP dataset generator
│       ├── nonconvex/      # Non-convex QPSR dataset generator
│       └── acopf/          # ACOPF-57 dataset generator + MATLAB sampling scripts

The reference environment uses Python 3.10, PyTorch 2.5.1, and CUDA 11.8.

conda env create -f environment.yml
conda activate cmfo

Key dependencies (see environment.yml for the full pinned list):

• pytorch>=2.5 (with CUDA)
• numpy, scipy, tqdm, setproctitle
• wandb (for experiment logging)
• cvxpy, osqp, qpth (optimization / differentiable QP)
• pypower (power-system utilities, for ACOPF)

Datasets are not shipped with the repository. Generate them with the scripts under src/datasets/<task>/make_dataset.py. Run each from the repository root so the relative imports resolve:

# Convex QP
python src/datasets/simple/make_dataset.py

# Non-convex QP with sine regularization
python src/datasets/nonconvex/make_dataset.py

# AC Optimal Power Flow (57-bus)
python src/datasets/acopf/make_dataset.py

Set up Weights & Biases logging:

export WANDB_API_KEY=your_real_key   # or edit the placeholder in train.py

Replace the placeholder 'yourwandbkey' on line 19 of train.py with your own W&B key, or set WANDB_MODE=disabled to skip logging entirely.

A minimal run on the convex QP benchmark:

python train.py \
  --name cmfo \
  --probType simple \
  --useCompl True \
  --penaltyWeight 1e-4 \
  --cutOff 0.5 \
  --qlWeight 0.001 \
  --bcWeight 0.0 \
  --nTimesteps 5 \
  --epochs 2000

Replace --probType simple with nonconvex or acopf57 for the other benchmarks. Per-task defaults (epochs, batch size, learning rate, constraint-tolerance settings) are defined in src/utils/default_args.py::cmfo_default_args and applied for any argument left unset on the command line.

bash.sh is the SLURM script we used on our cluster. Adjust the #SBATCH headers (partition, GPU count, time limit, conda module names) for your own environment.

sbatch bash.sh

After one or more runs finish, aggregate per-experiment statistics into summary dictionaries:

python load_results.py

This writes results/exper_status.dict and results/summary.dict.

If you find this code useful, please cite the paper:

@article{meng2026cmfo,
  title  = {Solving Constrained Optimization Problems as ODE-based Models Using Reinforcement Learning},
  author = {Meng, Han and Feng, Xinsong and Li, Yang and Wang, Chenan and Rajput, Kishansingh and Schram, Malachi and Chen, Haipeng},
  journal = {Transactions on Machine Learning Research},
  year   = {2026},
  url    = {https://openreview.net/forum?id=QW0ZX4zRC2}
}

Released under the MIT License.