Antibiotic Discovery with Deep RL + Graph Neural Networks
Multi-task GATv2 regressor and three-phase PPO agent that together generate novel antibiotic candidates against S. aureus and E. coli. Final project for the Advanced Machine Learning course, JHU MS in AI program (May 2026).
Code, README, and full reproducibility instructions: github.com/jsf3467v/antibiotic-discovery
Files in this repository
gnn_best.pt— trained multi-task GATv2 regressor. Test AUROC 0.83 (S. aureus) / 0.87 (E. coli) on a scaffold-split test set of 7,824 compounds. The E. coli result is at the empirical replicate-noise ceiling.policy_final.pt— final RL policy from phase 3, used to generate the 20,031-molecule candidate pool.policy_prior.pt— behavior-cloned policy used as the KL anchor during PPO training.surrogate.pt— fingerprint MLP surrogate, used for fast inner-loop reward calls during rollouts.paper.pdf— full paper with tables, figures, methodology, results, and discussion.
Headline results
The RL agent produces 20,031 unique valid molecules and beats random, hill-climbing, and SMILES-RNN baselines with Bonferroni-corrected p < 10⁻¹⁶. The agent's pool achieves the closest distributional match to known active antibiotics (FCD 26.1 vs ≥43.8) of any method tested. Approximately 95% of generated molecules trigger at least one Brenk structural alert and would require medicinal-chemistry refinement before any synthesis. Full breakdown including caveats and limitations is in the paper.
Loading checkpoints
import torch
# Example: load the trained GNN regressor
state_dict = torch.load("gnn_best.pt", weights_only=True)
# Then instantiate MultiTaskGNN from src/gnn.py in the GitHub repo
# and call .load_state_dict(state_dict)
See src/evaluate.py in the GitHub repository for the full loading and evaluation pattern.
Citation
@misc{keith2026antibiotic,
author = {Keith, Arlene},
title = {Deep Reinforcement Learning with Graph Neural Networks for Antibiotic Discovery},
year = {2026},
howpublished = {Final project, Advanced Machine Learning course, JHU MS in AI program},
url = {https://github.com/jsf3467v/antibiotic-discovery}
}
License
MIT. See LICENSE in the GitHub repository.