Tai Dang, Hieu Tran, Long-Hung Pham, Sang Truong, Edward A Pham, Jeffrey S Glenn, Thang Luong
Expand abstractCollapse
The latest advances of diffusion models in biomolecular structure prediction are still based on distance-based supervised training without adequately accounting for physically plausible interactions needed for real-world drug discovery. We introduce $\mathit{\boldsymbol{\pi}\text{-Fold}}$, a novel physics-informed policy-based folding approach that augments reinforcement learning with physics-aware scoring functions for training biomolecular diffusion models. Specifically, we fine-tune AlphaFold3-based models using Group Relative Policy Optimization and explicit physics-based rewards such as Rosetta all-atom energy, AutoDock Vina binding affinity, and ligand strain energy. Crucially, we demonstrate that these physics-based priors translate to enhanced performance on critical downstream prediction tasks. $\mathit{\boldsymbol{\pi}\text{-Fold}}$ achieves state-of-the-art results on SKEMPI for protein--protein mutation effects, CASP16 and OpenFE for protein--ligand binding affinities, Run-N-Poses for ligand geometric validity and interface quality across diverse structural fidelity metrics. On FoldBench, $\mathit{\boldsymbol{\pi}\text{-Fold}}$ achieves the best overall MolProbity score of 1.30, halving steric clashes and nearly eliminating rotamer outliers. Human evaluation further confirms that $\mathit{\boldsymbol{\pi}\text{-Fold}}$ generates structures that are meaningful for downstream drug discovery applications. Our code is available at \href{https://anonymous.4open.science/r/pifold-45B8/README.md}{github.com/anonymous/pifold}.