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.

An RL agent generates a molecule, which is evaluated with three reward signals — property-based extrinsic reward, and two intrinsic rewards for exploration — that feed back into training.

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.

HIR and LIR are combined into the adaptive intrinsic reward and added to PER to train the policy network end-to-end with PPO.

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.

Best property scores across six benchmarks. Mol-AIR (bottom row) outperforms all baseline intrinsic reward methods on every task.

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.

Proof-of-concept pipeline for discovering DRD2 inhibitor candidates: 10K generated molecules, filtered by DRD2 activity/QED/synthetic accessibility, then validated by docking simulation against risperidone.

References

2025

  1. Mol-AIR: Molecular Reinforcement Learning with Adaptive Intrinsic Rewards for Goal-Directed Molecular Generation
    Jinyeong Park, Jaegyoon Ahn, Jonghwan Choi, and 1 more author
    Journal of Chemical Information and Modeling. SCIE, IF 6.4 (Q1) , Feb 2025
    PMID: 39988822