TSSR: Two-Stage Swap-Reward-Driven Reinforcement Learning for Character-Level SMILES Generation

TL;DR

This work addresses the fragility of SMILES-based molecular generation by introducing TSSR, a two-stage reinforcement learning framework that first repairs syntactic errors through local token swaps and then reduces chemical inconsistencies using RDKit diagnostics. Implemented with a GRU-based generator and PPO, TSSR can train from scratch (P-RL) or fine-tune pretrained models (F-RL), improving syntactic validity, chemical validity, and novelty on the MOSES benchmark. The approach provides interpretable, dense reward signals without requiring hand-crafted grammars, enabling more diverse and chemically plausible molecule generation while maintaining computational efficiency. The results demonstrate that Stage One and Stage Two synergy yields higher validity and novelty, with potential for integration into broader de novo design pipelines and extensions to other representations and objectives.

Abstract

The design of reliable, valid, and diverse molecules is fundamental to modern drug discovery, as improved molecular generation supports efficient exploration of the chemical space for potential drug candidates and reduces the cost of early design efforts. Despite these needs, current chemical language models that generate molecules as SMILES strings are vulnerable to compounding token errors: many samples are unparseable or chemically implausible, and hard constraints meant to prevent failure can restrict exploration. To address this gap, we introduce TSSR, a Two-Stage, Swap-Reward-driven reinforcement learning (RL) framework for character-level SMILES generation. Stage one rewards local token swaps that repair syntax, promoting transitions from invalid to parseable strings. Stage two provides chemistry-aware feedback from RDKit diagnostics, rewarding reductions in valence, aromaticity, and connectivity issues. The reward decomposes into interpretable terms (swap efficiency, error reduction, distance to validity), is model agnostic, and requires no task-specific labels or hand-crafted grammars. We evaluated TSSR on the MOSES benchmark using a GRU policy trained with PPO in both pure RL (P-RL) from random initialization and fine-tuning RL (F-RL) starting from a pretrained chemical language model, assessing 10,000 generated SMILES per run. In P-RL, TSSR significantly improves syntactic validity, chemical validity, and novelty. In F-RL, TSSR preserves drug-likeness and synthesizability while increasing validity and novelty. Token-level analysis shows that syntax edits and chemistry fixes act jointly to reduce RDKit detected errors. TSSR converts a sparse terminal objective into a denser and more interpretable reward, improving both syntactic and chemical quality without reducing diversity. TSSR is dataset-agnostic and can be adapted to various reinforcement learning approaches.

Figures (3)

• Figure 1: Example a Two-Stage, Swap-Reward-driven (TSSR) reinforcement learning (RL) framework for character-level SMILES generation.
• Figure 2: Examples of TSSR Stage Two fixes: Invalid SMILES to Chemically valid with 2 Fixes Each
• Figure 3: TSSR workflow: Generation, Pre-Training, Reinforcement Learning