Tutorial

De Novo Drug Design with Reinforcement Learning and REINVENT

Difficulty: Advanced Time: 25 min read

Introduction

Traditional high-throughput screening involves physically testing millions of compounds to find a viable drug candidate. De novo drug design flips this paradigm: instead of searching a library, we use Generative AI to "invent" entirely new molecules that are optimized for specific biological targets. In this tutorial, we will use REINVENT (an open-source AI tool originally developed by AstraZeneca) to generate novel chemical structures using a Recurrent Neural Network (RNN) optimized via Reinforcement Learning (RL).

Architecture Overview


graph TD
    Prior[Pre-trained Prior RNN] -. Initializes .-> Agent[Actor / Agent RNN]
    
    Agent -->|Generates SMILES| Env[Environment / Scoring Function]
    Env -->|Calculates Reward| Loss[Loss Function]
    
    Prior -->|Calculates NLL| Loss
    Loss -->|Policy Gradient Update| Agent
    
    subgraph Environment Components
        Env --> QED[Drug-Likeness QED Score]
        Env --> Tox[Toxicity Filters]
        Env --> Docking[Docking / Predictive Model Score]
    end

Prerequisites

  • Python 3.8+
  • RDKit (for cheminformatics)
  • A GPU (highly recommended for RL training)
  • Basic understanding of SMILES strings

Step 1: Install Dependencies

Because cheminformatics environments can be finicky, it is highly recommended to use Conda.

# Clone the repository
git clone https://github.com/MolecularAI/Reinvent.git
cd Reinvent

# Create and activate the conda environment
conda env create -f reinvent.yml
conda activate reinvent.v3.2

Step 2: Understanding the Architecture

REINVENT uses an Actor-Critic Reinforcement Learning framework.

  • The Prior (Agent): A pre-trained RNN that knows how to generate valid SMILES strings (representing molecules) based on the ChEMBL database.
  • The Scoring Function (Environment): A custom function that evaluates the generated SMILES string. It rewards the agent if the molecule has desired properties (e.g., high predicted binding affinity, low toxicity, synthesizability) and penalizes it otherwise.
Over time, the Agent learns to shift its generative distribution to produce molecules that score highly in the Environment, while a KL-divergence penalty ensures it doesn't deviate too far from the Prior (to ensure it still produces chemically valid structures).

Step 3: Defining the Configuration JSON

REINVENT is primarily driven by JSON configuration files rather than writing raw Python scripts. You define the paths to your prior models and setup your scoring function parameters.

{
  "logging": {
    "sender": "http://127.0.0.1",
    "recipient": "local",
    "logging_path": "./logs/rl_experiment",
    "job_name": "De_Novo_Design_Tutorial",
    "job_id": "001"
  },
  "parameters": {
    "actor": "./data/random.prior.new",
    "critic": "./data/random.prior.new",
    "learning_rate": 0.0001,
    "batch_size": 128,
    "n_steps": 500
  },
  "scoring_function": {
    "name": "custom_product",
    "parallel": false,
    "parameters": [
      {
        "component_type": "qed",
        "name": "QED Score",
        "weight": 1
      },
      {
        "component_type": "custom_alerts",
        "name": "Toxicity Filters",
        "weight": 1,
        "specific_parameters": {
          "smiles": [
            "[*;R1]", "[*;R2]"
          ]
        }
      }
    ]
  }
}

In this simple example, our scoring function rewards molecules with high QED (Quantitative Estimate of Drug-likeness) and penalizes molecules containing specific toxic substructures (defined as SMARTS patterns).

Step 4: Running the RL Loop

Once your configuration is set, executing the Reinforcement Learning loop is straightforward.

python input.py rl_config.json

As the script runs, you will see the average score of the generated batch of molecules increase. The RNN is learning to bias its generation towards the rules you defined in the JSON.

Troubleshooting: Mode Collapse

A common failure mode in generative drug design is Mode Collapse. This happens when the Agent discovers a single, highly-rewarded molecular scaffold and repeatedly generates identical or nearly-identical molecules, ignoring chemical diversity.

  • Increase the Diversity Filter: REINVENT includes a Diversity Filter parameter. Set it to enforce structural variation across epochs (e.g., penalizing repeated Murcko scaffolds).
  • Tune the Learning Rate: If the learning rate is too high, the Agent rushes to the nearest local maximum and gets stuck. Try reducing it to 0.00005.
  • Balance the Scoring Function: If one component of your scoring function (e.g., Molecular Weight) is overwhelmingly easy to optimize, the model will exploit it. Adjust the weight parameters to force the model to balance multiple complex properties.

Conclusion

REINVENT is an incredibly powerful tool for generative chemistry. In a real-world scenario, you would replace the simple QED score with a machine learning model predicting binding affinity against a specific protein target (like EGFR or KRAS). By letting the RL agent explore the vast chemical space (estimated at 10^60 possible molecules), you can discover novel, patentable scaffolds that human chemists might never consider.