Flow Matching Reinforcement Learning Trainer

This repository contains the reinforcement learning (RL) fine-tuning framework for reproducing the results of the paper "Controllable Molecular Generation with Fine-tuned Flow-matching Model".

The trainer fine-tunes a pretrained SemlaFlow model to generate molecules with desired properties using reward-guided optimization. During training, molecules are sampled from the flow-matching model, evaluated by external reward functions, and used to optimize the model through a reward-weighted objective.

The framework supports both:

• Unconditional molecular generation
• Protein pocket-conditioned molecular generation

---

Software Requirements

Supported platforms

The software has been developed and tested on:

• Linux (recommended)
• Python 3.10+
• NVIDIA GPUs with CUDA 11.8 or newer

Required dependencies

Install the required packages using Conda and pip:

conda create --name rl python=3.11

pip install -r requirements.txt

---

Running the code

Training is launched using a YAML configuration file:

python rl_fastrl_trainer.py --config rl_configs/rl_fastrl.yml

Example with command-line overrides:

python rl_fastrl_trainer.py \
    --config rl_configs/rl_fastrl.yml \
    --seed 42 \
    --regular_weight 0.1

Training statistics, including reward values and molecule distributions, are logged to Weights & Biases.

---

Configuration

All experiment settings are specified in the YAML configuration file (rl_configs/rl_fastrl.yml).

Section	Field	Description
train	pretrained_ckpt_path	Path to the pretrained SemlaFlow checkpoint.
rl	reward	Reward function (psa, logp, pac, pac_strain, pac_logp, psa_strain_logp, complex_eng_strain).
rl	use_protein	Enables protein pocket conditioning.
rl	protein, native_ligand	Protein pocket and reference ligand (required when protein conditioning is enabled).
rl	n_min_atoms, n_max_atoms	Minimum and maximum number of atoms to sample.
rl	batch_size, n_epochs, regular_weight	Main reinforcement learning hyperparameters.
rl	ckpt_base_path, save_every	Checkpoint output directory and saving frequency.

---

Citation

If you use this software in your research, please cite:

@article{wang2025controllable,
  title={Controllable Molecular Generation with Fine-tuned Flow-matching Model},
  author={Wang, Kunyu and Janet, Jon Paul and Tibo, Alessandro},
  year={2025},
  publisher={ChemRxiv}
}