De Novo Drug Design with Reinforcement Learning

De novo RL drug design treats molecule generation as a sequential decision-making problem. A generative model (SMILES RNN, transformer, or graph model) is used as a policy; the RL algorithm updates this policy to maximise expected reward. REINVENT (Olivecrona et al., 2017; Blaschke et al., 2020; REINVENT4, 2024) is the most widely deployed framework in industry, using a PPO/REINFORCE variant with an augmented memory and diversity filter to prevent mode collapse.

REINVENT Architecture Overview

• Prior network: A pre-trained SMILES RNN (GRU/LSTM) trained on ChEMBL. Frozen during RL. Acts as a regulariser — the agent's policy cannot drift too far from the prior (KL divergence term in the augmented likelihood objective).
• Agent network: Same architecture, initialised from the prior, updated during RL. Generates SMILES one token at a time.
• Scoring function: A weighted sum of component scores evaluated on each generated molecule. Returns a scalar reward in [0, 1].
• Augmented likelihood objective: log P_agent(SMILES) = log P_prior(SMILES) + σ · score — the agent maximises score while staying close to the prior.
• Diversity filter: Tracks scaffold counts (Murcko scaffold). Penalises oversampling of the same scaffold to prevent mode collapse.

Reward Function Design

from rdkit import Chem
from rdkit.Chem import Descriptors, QED
from rdkit.Chem import RDConfig
import os
import sys
sys.path.append(os.path.join(RDConfig.RDContribDir, "SA_Score"))
import sascorer  # RDKit SA score contribution

from dataclasses import dataclass
from typing import Callable


@dataclass
class ScoringComponent:
    """Single scoring component with a weight and transform function."""
    name: str
    weight: float
    score_fn: Callable[[Chem.Mol], float]
    min_val: float = 0.0
    max_val: float = 1.0

    def score(self, mol: Chem.Mol) -> float:
        """Return normalised score in [0, 1]."""
        raw = self.score_fn(mol)
        normalised = (raw - self.min_val) / (self.max_val - self.min_val)
        return float(max(0.0, min(1.0, normalised)))


def qed_score(mol: Chem.Mol) -> float:
    return QED.qed(mol)


def sa_score_component(mol: Chem.Mol) -> float:
    """SA score ranges 1 (easy) to 10 (hard). Invert and normalise."""
    sa = sascorer.calculateScore(mol)
    return 1.0 - (sa - 1.0) / 9.0  # -> 1 is easy, 0 is very hard


def mw_component(mol: Chem.Mol) -> float:
    """Prefer MW in 300–500 Da. Returns 1.0 in range, degrades outside."""
    mw = Descriptors.MolWt(mol)
    if 300 <= mw <= 500:
        return 1.0
    elif mw < 300:
        return mw / 300.0
    else:
        return max(0.0, 1.0 - (mw - 500) / 300.0)


def activity_model_score(mol: Chem.Mol, activity_model) -> float:
    """Predicted pIC50 normalised to [0, 1] (clipped to 4–9 range)."""
    from rdkit.Chem import AllChem
    fp = AllChem.GetMorganFingerprintAsBitVect(mol, 2, 2048)
    pIC50 = float(activity_model.predict([list(fp)])[0])
    return (pIC50 - 4.0) / 5.0  # 4 = inactive, 9 = highly active


class MultiObjectiveScoringFunction:
    """
    Weighted geometric mean of scoring components.
    Geometric mean penalises zero-scores harder than arithmetic mean,
    preventing reward hacking via single-component exploitation.
    """

    def __init__(self, components: list[ScoringComponent]):
        self.components = components
        total_weight = sum(c.weight for c in components)
        # Normalise weights to sum to 1
        for c in components:
            c.weight = c.weight / total_weight

    def __call__(self, smiles: str) -> float:
        mol = Chem.MolFromSmiles(smiles)
        if mol is None:
            return 0.0

        weighted_log_sum = 0.0
        for component in self.components:
            s = component.score(mol)
            s = max(s, 1e-6)  # Avoid log(0)
            weighted_log_sum += component.weight * float.__add__(0, __import__('math').log(s))

        import math
        return float(math.exp(weighted_log_sum))  # Geometric mean


# Example setup
components = [
    ScoringComponent("QED",       weight=1.0, score_fn=qed_score,       min_val=0.0, max_val=1.0),
    ScoringComponent("SA_Score",  weight=1.0, score_fn=sa_score_component, min_val=0.0, max_val=1.0),
    ScoringComponent("MW",        weight=0.5, score_fn=mw_component,    min_val=0.0, max_val=1.0),
]
scoring_fn = MultiObjectiveScoringFunction(components)

Diversity Filter and Scaffold Tracking

from rdkit.Chem.Scaffolds import MurckoScaffold
from collections import defaultdict

class DiversityFilter:
    """
    Penalise oversampling of the same Murcko scaffold.
    Molecules from a scaffold sampled more than `minscore_scaffold_count`
    times receive a diversity penalty.
    """

    def __init__(self, minscore: float = 0.4, scaffold_memory_size: int = 1000):
        self.minscore = minscore
        self.scaffold_memory_size = scaffold_memory_size
        self.scaffold_memory: dict[str, int] = defaultdict(int)

    def penalty(self, smiles: str, raw_score: float) -> float:
        if raw_score < self.minscore:
            return raw_score  # Don't penalise already-low scorers

        mol = Chem.MolFromSmiles(smiles)
        if mol is None:
            return 0.0

        scaffold_smiles = MurckoScaffold.MurckoScaffoldSmilesFromSmiles(smiles)
        self.scaffold_memory[scaffold_smiles] += 1

        count = self.scaffold_memory[scaffold_smiles]
        # Penalise proportionally to excess count above threshold (3)
        if count > 3:
            penalty = min(1.0, (count - 3) * 0.1)
            return raw_score * (1.0 - penalty)

        return raw_score

    def top_scaffolds(self, n: int = 10) -> list[tuple[str, int]]:
        return sorted(self.scaffold_memory.items(), key=lambda x: x[1], reverse=True)[:n]

GuacaMol and MOSES Benchmarks

Avoiding Reward Hacking

• Use geometric mean instead of arithmetic sum. This ensures all components must be reasonably good — a single high-scoring component cannot compensate for a near-zero component.
• Include synthetic accessibility (SA score, SYBA, or SCScore) as a mandatory component. Without it, RL readily generates high-scoring but unsynthesisable structures.
• Filter known problematic substructures: PAINS, reactive groups, Pan-Assay Interference Compounds. Apply as binary filters (score = 0) rather than penalties.
• Monitor tanimoto similarity to training set. Generated molecules collapsing to near-duplicates of known actives is a sign of mode collapse, not de novo design.
• Periodically validate top hits with physics-based docking. RL reward functions are surrogates — cross-validate with AutoDock Vina or Glide on top-scoring molecules.

Key Takeaways

• REINVENT's augmented likelihood framework is the de facto standard for RL-based SMILES generation in industry. REINVENT4 supports SMILES, SMILES fragments, and Mol2Mol scaffold-hopping modes.
• Reward function design is the most critical — and most domain-knowledge-intensive — part of the RL pipeline. Involve medicinal chemists early to set reasonable score thresholds.
• Geometric mean reward aggregation and diversity filters are non-negotiable for realistic campaigns. Arithmetic mean leads to mono-scaffold collapsing within ~500 RL steps.
• GuacaMol goal-directed benchmarks (MPO tasks) are the standard for comparing RL generative models. REINVENT4 and graph-based RL agents (MCTS-based) are state-of-the-art as of 2024.

Try it live

Explore a related model hands-on: open AI Drug Design: Molecular Generation, Docking & ADMET Pipeline and change its parameters while it runs. Everything stays in your browser.