Mol-AIR
Adaptive intrinsic rewards for goal-directed molecular reinforcement learning
Mol-AIR is a reinforcement learning framework for goal-directed molecular generation, published in Journal of Chemical Information and Modeling (Park et al., 2025). I’m the first author.
Motivation
Goal-directed molecular generation searches an astronomically large chemical space (estimated at 1033–1060 drug-like compounds) for structures with target properties. RL-based generators tend to get stuck exploring only a narrow region of that space. Prior work tackled this with intrinsic (curiosity-driven) rewards, but each existing approach has a blind spot: history-based methods (e.g. counting revisited states) explore broadly but struggle to converge near a specific target structure, while learning-based methods (e.g. prediction-error rewards) converge well but taper off exploration too early in training.
Method
Mol-AIR combines the strengths of both families into a single adaptive intrinsic reward:
- History-based intrinsic reward (HIR) — counts how often a molecular state has been visited, using locality-sensitive hashing (LSH) over Morgan fingerprints, with exponential decay over visit counts.
- Learning-based intrinsic reward (LIR) — a Random Network Distillation (RND) module that rewards states a predictor network still finds hard to predict.
- Property-based extrinsic reward (PER) — the actual target chemical property score (e.g. pLogP, QED, similarity to a reference drug).
A dual-critic PPO architecture estimates separate episodic and non-episodic advantages, so credit assignment across the extrinsic and two intrinsic signals doesn’t interfere with policy optimization.
Results
Across six molecular property benchmarks (pLogP, QED, similarity to celecoxib, GSK3B, JNK3, and GSK3B+JNK3), Mol-AIR outperformed count-based, memory-based, and prediction-based intrinsic reward baselines on every task — including reaching the theoretical optimum of QED, which prior methods failed to achieve.
Proof of concept: discovering DRD2 inhibitor candidates
To test translational usefulness beyond benchmark scores, I fine-tuned Mol-AIR to generate compounds structurally related to risperidone, a known dopamine D2 receptor (DRD2) inhibitor, jointly optimizing DRD2 activity, drug-likeness (QED), and synthetic accessibility. Out of 10K generated molecules, 84 satisfied all three criteria, and molecular docking simulations (AutoDock Vina) showed the top candidates binding DRD2 with lower binding energy than risperidone itself.