Agentic reinforcement learning empowers next-generation chemical language models for molecular design and synthesis

TL;DR

ChemCRAFT introduces a cognitive-decoupling framework that democratizes expert-level chemistry AI by decoupling reasoning from knowledge storage through a Chemical Agent Sandbox. It builds High-Quality Reasoning Trajectories via Tool-Integrated Reasoning and Reflective Refinement, then trains compact 7B–14B models with Cold-Start SFT followed by reinforcement learning using SMILES-GRPO and a multidimensional chemical reward. The approach attains state-of-the-art performance among open-source models on ChemCoTBench tasks, surpassing cloud APIs in molecular structure understanding, optimization, and synthesis-pipeline prediction while preserving privacy and reducing inference costs. By open-sourcing the tool-trajectory dataset ChemToolDataset and the sandboxed architecture, the work demonstrates that robust scientific reasoning can be learned as a policy of tool orchestration rather than emergent from scale, enabling cost-effective, private AI-assisted chemistry at local labs.

Abstract

Language models are revolutionizing the biochemistry domain, assisting scientists in drug design and chemical synthesis with high efficiency. Yet current approaches struggle between small language models prone to hallucination and limited knowledge retention, and large cloud-based language models plagued by privacy risks and high inference costs. To bridge this gap, we introduce ChemCRAFT, a novel framework leveraging agentic reinforcement learning to decouple chemical reasoning from knowledge storage. Instead of forcing the model to memorize vast chemical data, our approach empowers the language model to interact with a sandbox for precise information retrieval. This externalization of knowledge allows a locally deployable small model to achieve superior performance with minimal inference costs. To enable small language models for agent-calling ability, we build an agentic trajectory construction pipeline and a comprehensive chemical-agent sandbox. Based on sandbox interactions, we constructed ChemToolDataset, the first large-scale chemical tool trajectory dataset. Simultaneously, we propose SMILES-GRPO to build a dense chemical reward function, promoting the model's ability to call chemical agents. Evaluations across diverse aspects of drug design show that ChemCRAFT outperforms current cloud-based LLMs in molecular structure analysis, molecular optimization, and synthesis pathway prediction, demonstrating that scientific reasoning is not solely an emergent ability of model scale, but a learnable policy of tool orchestration. This work establishes a cost-effective and privacy-preserving paradigm for AI-aided chemistry, opening new avenues for accelerating molecular discovery with locally deployable agents.

Figures (3)

• Figure 1: Overview of the data-curation pipeline and the training pipeline of our ChemCRAFT. a We build up the agentic trajectory dataset based on the pure reasoning trajectory with the interaction of our Chemical Agent Sandbox. b We implement a two-stage training scheme. The first stage involves cold-start training using token-level loss to enable the model to memorize the tool invocation format. The second stage employs GRPO to enhance the model’s understanding of tools.
• Figure 2: Experimental Analysis for Molecular Optimization and Reaction Predictions. a Distribution analysis for property improvements on LogP, QED, Solubility. b Optimized molecule visualizations for protein-activated properties, including DRD2, JNK3, GSK-3$\beta$. c Comparison between ChemCRAFT, chemical-specific models, and LLM-APIs on forward major / by product prediction and retrosynthesis prediciton. d Similarity distribution for reaction condition recommendations. e Reaction prediction visualizations on complex reaction types.
• Figure 3: a Performance on general reasoning benchmarks. Our ChemCRAFT outperforms other chemical language models on different domains. b Training curves analysis for non-tool-calling SFT, tool-augmented SFT, and Agentic-RL. c Tool-calling frequency analysis between the SFT training paradigm and agentic-RL paradigm. d ChemCRAFT has low inference cost and a more useful token length compared to other chemical multi-agent systems. e Performance and token-length correlation analysis between ChemCRAFT and other chemical language models.