Reinforcement Learning with LLM-Guided Action Spaces for Synthesizable Lead Optimization

1. Problem Statement

Lead optimization in drug discovery involves iteratively modifying promising compounds to improve therapeutic properties (e.g., potency, selectivity, pharmacokinetics) while ensuring the modifications are synthetically feasible. Existing computational approaches face a critical trade-off:

• Property-Driven Methods: Reinforcement Learning (RL) and Genetic Algorithms often maximize property scores but frequently generate chemically invalid structures or molecules with no viable synthetic routes.
• Synthesizable Methods: Approaches that enforce synthesizability often rely on exhaustive enumeration of large reaction networks (computationally expensive) or separate projection models that lack explicit reaction pathways.
• LLM Limitations: While Large Language Models (LLMs) possess vast chemical knowledge, their direct application to molecular generation often yields chemically invalid structures due to a lack of grounded structural constraints. Furthermore, standard LLMs struggle with long-horizon sequential decision-making required for multi-step optimization.

Core Challenge: How to design a framework that simultaneously optimizes molecular properties, guarantees synthetic feasibility via explicit reaction pathways, and maintains computational efficiency during multi-step exploration.

2. Methodology: MolReAct Framework

The authors propose MolReAct, a synthesis-aware framework that formulates lead optimization as a Markov Decision Process (MDP) where the action space is dynamically constrained by validated reaction templates and guided by a tool-augmented LLM.

A. Problem Formulation (MDP)

• State ($s_t$): A chemically valid molecule represented as a SMILES string.
• Action ($a_t$): A pair of (reaction template, building block) or a "stop" action.
• Transition: Deterministic; applying a matched template with a specific building block yields a unique product.
• Reward: Only the terminal molecule ($s_T$) receives a reward from an oracle (e.g., docking score, property score). The goal is to maximize the expected terminal reward.

B. LLM Agent as a Dynamic Reaction Environment

Instead of a static action space, MolReAct uses a two-stage process to construct a compact, molecule-specific action space at each step:

1. Template Matching: The system performs substructure matching against a curated library of 115 validated reaction templates (covering uni-, bi-, and tri-molecular reactions) to identify feasible templates for the current molecule.
2. LLM-Guided Instantiation: A tool-augmented LLM agent (Llama-3.3-70B) analyzes the current molecule using specialized chemical tools (e.g., functional group detection, reactive site analysis, scaffold extraction). Based on this analysis and the target property, the LLM selects relevant templates and proposes up to 10 specific building blocks.
   • Constraint: The LLM is restricted to proposing reactions that can be executed via the matched templates, ensuring chemical validity.
   • Output: A compact set of candidate actions $A(s_t)$ passed to the policy model.

C. Trajectory Optimization via Reinforcement Learning

• Policy Model: A Qwen-3-4B language model with a linear action head. It receives the current molecule and the list of candidate reactions (template names, reactant/product SMILES) as a prompt.
• Algorithm: Group Relative Policy Optimization (GRPO).
  • For each initial molecule, the policy samples a group of $G$ independent trajectories.
  • A single terminal reward is standardized within the group to compute a group-relative advantage.
  • This advantage is assigned to every step in the trajectory, enabling trajectory-level credit assignment (crucial for multi-step optimization where immediate rewards are sparse).
• Caching Mechanism: To mitigate the high inference cost of repeated LLM calls during RL exploration, a SMILES-based caching mechanism stores the action space and transition outcomes for previously encountered molecules. This reduces end-to-end optimization time by ~43%.

3. Key Contributions

1. Synthesis-Constrained Action Space: MolReAct bridges the gap between generative freedom and synthetic reality by grounding every optimization step in validated reaction templates, ensuring all generated molecules have explicit synthetic pathways.
2. Tool-Augmented LLM for Action Generation: It leverages LLMs not as direct generators, but as dynamic "environment agents" that use chemical tools to propose chemically grounded, targeted transformations, significantly reducing the search space while maintaining diversity.
3. Trajectory-Level Policy Optimization: By employing GRPO, the framework effectively learns long-term strategies for multi-step optimization, avoiding the local optima traps common in greedy search or single-step RL.
4. Efficiency via Caching: The introduction of a SMILES-centric caching mechanism drastically reduces the computational overhead of LLM inference during RL training.

4. Experimental Results

The framework was evaluated on 14 tasks (13 from Therapeutic Data Commons and 1 structure-based docking task).

• Performance (Top-10 Score): MolReAct achieved an average Top-10 score of 0.563, outperforming the strongest synthesizable baseline (SynFormer, 0.510) by 10.4% relative improvement. It ranked first on 8 out of 14 tasks.
• Sample Efficiency: MolReAct achieved the best Area Under the Curve (AUC10) on 10 of 14 tasks, demonstrating faster convergence to high-scoring chemical regions compared to baselines.
• Ablation Studies:
  • Tool-Guided Proposals: Removing chemical tools and relying on a "Vanilla LLM" resulted in lower quality proposals, confirming the necessity of structural grounding.
  • Policy Optimization: Replacing GRPO with Greedy Search or Random Policy significantly degraded performance, proving the value of trajectory-level credit assignment.
• Synthetic Feasibility:
  • Building Block Availability: Post-hoc analysis showed that 66.1–76.2% of proposed building blocks were directly purchasable from commercial catalogs (Enamine). For non-purchasable blocks, Synthetic Accessibility (SA) scores were low (2.33–2.73), indicating they are easily synthesizable.
  • Pathway Analysis: Case studies (DRD2, GSK3β, JNK3, sEH) demonstrated that MolReAct successfully introduced specific pharmacophores (e.g., N-alkylation, aryl cyanation) that improved both oracle scores and independent docking scores (AutoDock Vina).

5. Significance and Impact

• Bridging the Gap: MolReAct effectively connects computational optimization with practical drug development by ensuring that every proposed molecule is not only property-improved but also accompanied by a template-grounded synthetic pathway.
• Scalability: By decoupling the "knowledge" (LLM for proposing actions) from the "decision" (RL policy for selecting actions), the framework scales efficiently and avoids the computational bottlenecks of enumerating massive reaction networks.
• Practical Applicability: The high rate of commercially available building blocks and the explicit synthetic routes make the output of MolReAct immediately actionable for medicinal chemists, reducing the "design-make-test" cycle time in early-stage drug discovery.

In summary, MolReAct represents a significant advancement in AI-driven drug discovery by integrating the reasoning capabilities of LLMs with the rigorous constraints of chemical synthesis and the strategic depth of reinforcement learning.