InVirtuoGen: Uniform-Source Discrete Flows for Drug Discovery

Repository to reproduce results from the paper.

🚀 Quick Start

Setup Environment

On ARM-based environments

mamba create -n invgen python=3.10.16 -y
mamba activate invgen
pip install -r requirements.txt
pip install -e .

We also provide a Dockerfile for aarch64-based environments - the results from the paper were produced within the corresponding Docker container and should be reproducible - although we could not check it on another cluster.

Download Checkpoints

mkdir checkpoints && cd checkpoints
mkdir checkpoints && cd checkpoints
curl -L "https://zenodo.org/records/16874868/files/invirtuo_gen.ckpt.zip?download=1" -o invirtuo_gen.ckpt.zip
curl -L "https://zenodo.org/records/16874868/files/invirtuo_gen_big.ckpt.zip?download=1" -o invirtuo_gen_big.ckpt.zip
unzip invirtuo_gen.ckpt.zip && unzip invirtuo_gen_big.ckpt.zip
cd ..

Available Models:

• invirtuo_gen.ckpt - Standard model for general use
• invirtuo_gen_big.ckpt - Larger model for optimal performance

Quick Generation:

python -m in_virtuo_gen.generate \
    --ckpt_path checkpoints/invirtuo_gen.ckpt \
    --num_samples 1000 \
    --batch_size 200

Reproduce Results in Paper

1. De Novo Molecule Generation

Reproduce Quality-Diversity Analysis:

(~6h on RTX 4090 for 3 seeds)

python -m in_virtuo_gen.evaluation.denovo \
    --checkpoint_path checkpoints/invirtuo_gen_big.ckpt \
    --dt 0.1 0.01 0.001 \
    --num_samples 1000 \
    --eta 999 \
    --num_seeds 3 \
    --device 0

2. Fragment-Constrained Generation

(<1h on RTX 4090 for 3 seeds)

Design molecules containing specific fragments while optimizing properties.

python -m in_virtuo_gen.evaluation.downstream \
--model_path checkpoints/invirtuo_gen_big.ckpt\
 --dt 0.01\
 --temperature 1\
 --noise 0\
 --eta 999\
 --device 0 --all

Supported Tasks:

• Motif Extension - Growing fixed motifs with new substituents
• Linker Design - Connecting fragments with feasible linkers
• Scaffold Decoration - Adding functional groups at defined positions
• Superstructure Generation - Assembling multiple fragments

3. Target Property Optimization

Optimize molecular properties using genetic algorithms + reinforcement learning.

Setup Vocab: (~12 hours on 16 CPU cores)

mkdir -p data
wget -O data/frags_zinc250.csv https://zenodo.org/records/16742898/files/frags_zinc250.csv
python -m in_virtuo_reinforce.preprocess.get_vocab \
    --datapath data/frags_zinc250.csv \
    --outpath in_virtuo_reinforce/vocab/zinc250k.csv

This step scores all ~250k ZINC molecules against all 23 oracles and writes in_virtuo_reinforce/vocab/zinc250k.csv (~106 MB), which the --use_prescreen path reads. The built file is not tracked in this repository, since it exceeds GitHub's 100 MB per-file limit; rebuild it with the command above. get_vocab.py checkpoints per property and supports --resume, so an interrupted build never re-scores a finished oracle.

With Prescreening ZINC250k:

State-of-the-art results on PMO benchmark

Results Table

python -m in_virtuo_reinforce.evaluation.results_table --results_root results/target_property/prescreen_3_runs --incl

Reproduce Results - Run Optimization (~36h on RTX 4090 for 3 seeds): Command:

python -m in_virtuo_reinforce.genetic_ppo --ckpt checkpoints/invirtuo_gen.ckpt --device 0 --start_t 0. --offspring_size 50 --num_timesteps 50 --num_reinforce_steps 10 --clip_eps 0.2 --dt 0.01 --experience_replay 0 --vocab_size 10 --lr 0.1 --use_prompter --use_mut --train_mut --first_pop_size 10 --aggressive_bandit --max_oracle_calls 10000 --c_neg 1 --mutation_size 20 --rl_lr 0.00001 --tot_offspring_size 100  --start_rank 1 --use_prescreen

Without Prescreening ZINC250k:

State-of-the-art results on PMO benchmark

Results Table

python -m in_virtuo_reinforce.evaluation.results_table --results_root results/target_property/no_prescreen_3_runs --include_std --exclude_prescreen

Reproduce Results - Run Optimization (~36h on RTX 4090 for 3 seeds):

python -m in_virtuo_reinforce.genetic_ppo --ckpt checkpoints/invirtuo_gen.ckpt --device 0 --start_t 0. --offspring_size 50 --num_timesteps 50 --num_reinforce_steps 5 --clip_eps 0.2 --dt 0.01 --experience_replay 300 --vocab_size 10 --lr 0.1 --use_prompter --use_mut --train_mut --first_pop_size 100 --aggressive_bandit --max_oracle_calls 10000 --c_neg 1 --mutation_size 20 --rl_lr 0.00001 --tot_offspring_size 100 --start_rank 1

3.b Ablations:

Results Table

python  -m in_virtuo_reinforce.evaluation.results_table --ablation_mode \
 --results_path   results/target_property/prescreen_with_exp  results/target_property/prescreen_eta999_start_t_0.2 results/target_property/prescreen_eta999 results/target_property/prescreen_no_bandit results/target_property/prescreen_no_mutation   results/target_property/prescreen_no_ppo results/target_property/no_prescreen_no_experience_replay results/target_property/no_prescreen_no_prompter_no_mut \
  --model_names  "With Experience Replay" "Sampling with eq.~\ref{eq:our_update}, $t_{start}=0.2$" "Sampling with eq.~\ref{eq:our_update}" "No Bandit" "No Mutation" "No PPO" "No Prescreen, No Experience Replay" "No Prompter, No Mutation, No Prescreen"

Commands: With Experience Replay

-m in_virtuo_reinforce.genetic_ppo --ckpt checkpoints/invirtuo_gen.ckpt --device 0 --start_t 0. --offspring_size 50 --num_timesteps 50 --num_reinforce_steps 10 --clip_eps 0.2 --dt 0.01 --experience_replay 300 --vocab_size 10 --lr 0.1 --first_pop_size 10 --aggressive_bandit --max_oracle_calls 10000 --c_neg 1 --mutation_size 20 --rl_lr 0.00001 --tot_offspring_size 100 --start_seed 1 --start_rank 1 --use_prescreen --use_prompter --use_mutation --train_mutation

Commands: No Bandit:

python -m in_virtuo_reinforce.genetic_ppo --ckpt checkpoints/invirtuo_gen.ckpt --device 0 --start_t 0. --offspring_size 50 --num_timesteps 50 --num_reinforce_steps 10 --clip_eps 0.2 --dt 0.01 --experience_replay 0 --vocab_size 10 --lr 0.1 --first_pop_size 10 --aggressive_bandit --max_oracle_calls 10000 --c_neg 1 --mutation_size 20 --rl_lr 0.00001 --tot_offspring_size 100 --start_rank 1 --use_prescreen --use_prompter --use_mutation --train_mutation --no_bandit

Commands: No Mutation:

python -m in_virtuo_reinforce.genetic_ppo --ckpt checkpoints/invirtuo_gen.ckpt --device 0 --start_t 0. --offspring_size 20 --num_timesteps 50 --num_reinforce_steps 10 --clip_eps 0.2 --dt 0.01 --use_prompter --experience_replay 300 --lr 0.1 --vocab_size 10 --first_pop_size 10 --aggressive_bandit --max_oracle_calls 10000 --c_neg 1 --mutation_size 20 --rl_lr 0.00005 --tot_offspring_size 100 --use_prescreen --num_seeds 1 --start_rank 1

Commands: no PPO:

python -m in_virtuo_reinforce.genetic_ppo --ckpt checkpoints/invirtuo_gen.ckpt --device 0 --start_t 0. --offspring_size 10 --num_timesteps 0 --num_reinforce_steps 0 --clip_eps 0.2 --dt 0.01 --use_prompter --experience_replay 300 --lr 0.1 --vocab_size 10 --first_pop_size 100 --aggressive_bandit --max_oracle_calls 10000 --c_neg 1 --mutation_size 20 --rl_lr 0.00005 --tot_offspring_size 100 --use_prescreen --num_seeds 1 --start_rank 1

Commands: No Experience Replay:

-m in_virtuo_reinforce.genetic_ppo --ckpt checkpoints/invirtuo_gen.ckpt --device 0 --start_t 0. --offspring_size 50 --num_timesteps 50 --num_reinforce_steps 5 --clip_eps 0.2 --dt 0.01 --experience_replay 0 --vocab_size 10 --lr 0.1 --first_pop_size 100 --aggressive_bandit --max_oracle_calls 10000 --c_neg 1 --mutation_size 20 --rl_lr 0.00001 --tot_offspring_size 100 --start_rank 1 --use_prompter --use_mutation --train_mutation

Commands: No Prompter, No Mutation:

python -m in_virtuo_reinforce.genetic_ppo --ckpt checkpoints/invirtuo_gen.ckpt --device 0 --start_t 0. --offspring_size 100 --num_timesteps 50 --num_reinforce_steps 10 --clip_eps 0.2 --dt 0.01 --experience_replay 300 --lr 0.1 --first_pop_size 100 --aggressive_bandit --max_oracle_calls 10000 --c_neg 1 --mutation_size 20 --rl_lr 0.00005 --tot_offspring_size 100 --num_seeds 1 --start_rank 1

Commands: Our Sampling:

python -m in_virtuo_reinforce.genetic_ppo --ckpt checkpoints/invirtuo_gen.ckpt --device 0 --start_t 0. --offspring_size 50 --num_timesteps 50 --num_reinforce_steps 10 --clip_eps 0.2 --dt 0.01 --experience_replay 0 --vocab_size 10 --lr 0.1 --first_pop_size 10 --aggressive_bandit --max_oracle_calls 10000 --c_neg 1 --mutation_size 20 --rl_lr 0.00001 --tot_offspring_size 100 --use_prompter --use_mutation --train_mutation --use_prescreen --eta 999 --num_seeds 1

Commands: Our Sampling, different start t:

python -m in_virtuo_reinforce.genetic_ppo --ckpt checkpoints/invirtuo_gen.ckpt --device 0 --start_t 0.2 --offspring_size 50 --num_timesteps 50 --num_reinforce_steps 10 --clip_eps 0.2 --dt 0.01 --experience_replay 0 --vocab_size 10 --lr 0.1 --first_pop_size 10 --aggressive_bandit --max_oracle_calls 10000 --c_neg 1 --mutation_size 20 --rl_lr 0.00001 --tot_offspring_size 100 --identifier 898948 --start_seed 0 --start_rank 1 --use_prompter --use_mutation --train_mutation --eta 999 --use_prescreen --num_seeds 1

3.c Comparison in the Unlimited-Oracle-Call Regime (Graph-GRPO)

Some methods, notably Graph-GRPO, report PMO scores without counting the oracle calls spent on training the generative model, and only count the calls of the final optimization run. We consider this at odds with the purpose of the benchmark, which is sample efficiency with respect to an expensive oracle, but report results under that convention for completeness (see the appendix of the paper).

Those experiments use a two-phase protocol (a finetuning phase of at most 100k oracle calls per task, followed by the standard 10k-call PMO run from the finetuned policy) that is not part of the code on main. The code, submission scripts, per-oracle result CSVs and table-building recipe for them live on the grpo branch:

git checkout grpo
cat EVAL_RUNS.md   # which run produced which table, and the exact commands

4. Lead Optimization:

(We were unable to install the required dependencies for these experiments in the Dockerfile on an aarch64 environment. On ARM systems, the setup works by simply running pip install -r requirements.txt). Also we can not guarantee exact reproducability here, due to the stochasticity of the docking with openVINA.

Commands

python -m in_virtuo_reinforce.ppo_docking --device 1 --start_t 0. --offspring_size 50 --seed 0 --max_oracle_calls 1000 --num_reinforce_steps 50 --clip_eps 0.5 --start_task 0 --experience_replay_size 100 --dt 0.01 --c_neg 1 --dt 0.01 --tot_offspring 100 --vocab_size 10 --num_timesteps 50 --start_task 0

reproduce table

 python -m in_virtuo_reinforce.evaluation.results_table_lead --sim_results_dir results/lead_optimization/sim_constraint

4.b Without Similarity Constraints:

 python -m in_virtuo_reinforce.ppo_docking --device 1 --start_t 0. --offspring_size 50 --seed 0 --max_oracle_calls 1000 --num_reinforce_steps 50 --clip_eps 0.5 --start_task 0 --experience_replay_size 100 --dt 0.01 --c_neg 1 --dt 0.01 --tot_offspring 100 --vocab_size 10 --num_timesteps 50 --no_sim_constraint

reproduce table

 python -m in_virtuo_reinforce.evaluation.results_table_lead --mode nosim --no_sim results/lead_optimization/no_sim_constraint

Hugging Face:

Paper + Model also available at: https://huggingface.co/papers/2509.26405

📄 Citation

If you use this work, please cite:

@misc{kaech2025refinedrugsdontcomplete,
      title={Refine Drugs, Don't Complete Them: Uniform-Source Discrete Flows for Fragment-Based Drug Discovery},
      author={Benno Kaech and Luis Wyss and Karsten Borgwardt and Gianvito Grasso},
      year={2025},
      eprint={2509.26405},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2509.26405 },
}

📧 Contact

📬 [email protected]

⚖️ License and Terms of Use

The InVirtuoGEN source code is licensed under the Creative Commons Attribution-Non-Commercial ShareAlike International License, Version 4.0 (CC-BY-NC-SA 4.0) (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc-sa/4.0/

The InVirtuoGEN model parameters are made available under the InVirtuoGEN Model Parameters Terms of Use (the "Terms"); you may not use these except in compliance with the Terms. You may obtain a copy of the Terms at WEIGHTS_TERMS_OF_USE.md.

Unless required by applicable law, InVirtuoGEN and its output are distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. You are solely responsible for determining the appropriateness of using InVirtuoGEN, or using or distributing its source code or output, and assume any and all risks associated with such use or distribution and your exercise of rights and obligations under the relevant terms. Output are predictions with varying levels of confidence and should be interpreted carefully. Use discretion before relying on, publishing, downloading or otherwise using the InVirtuoGEN Assets.