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

arXiv: https://arxiv.org/abs/2601.04521 · HTML: https://mlsanity.com/paper/2601.04521/

## Key facts

- Cited by 1 papers in the map (data through 7 Oct 2026).
- Builds on 22 of its 53 references that are in the map.
- First posted on arXiv 8 Jan 2026.

## Builds on (most cited first, 22 of 22)

- [Improved Training of Wasserstein GANs](https://mlsanity.com/paper/1704.00028/) (2017)
- [Automatic chemical design using a data-driven continuous representation of molecules](https://mlsanity.com/paper/1610.02415/) (2016)
- [ChemBERTa: Large-Scale Self-Supervised Pretraining for Molecular Property Prediction](https://mlsanity.com/paper/2010.09885/) (2020)
- [DiffDock: Diffusion Steps, Twists, and Turns for Molecular Docking](https://mlsanity.com/paper/2210.01776/) (2022)
- [GuacaMol: Benchmarking Models for De Novo Molecular Design](https://mlsanity.com/paper/1811.09621/) (2018)
- [Molecular Sets (MOSES): A Benchmarking Platform for Molecular Generation Models](https://mlsanity.com/paper/1811.12823/) (2018)
- [Molecular De Novo Design through Deep Reinforcement Learning](https://mlsanity.com/paper/1704.07555/) (2017)
- [Deep Reinforcement Learning for De-Novo Drug Design](https://mlsanity.com/paper/1711.10907/) (2017)
- [MolecularRNN: Generating realistic molecular graphs with optimized properties](https://mlsanity.com/paper/1905.13372/) (2019)
- [Mapping the Space of Chemical Reactions Using Attention-Based Neural Networks](https://mlsanity.com/paper/2012.06051/) (2020)
- [Keeping it Simple: Language Models can learn Complex Molecular Distributions](https://mlsanity.com/paper/2112.03041/) (2021)
- [Augmenting Genetic Algorithms with Deep Neural Networks for Exploring the Chemical Space](https://mlsanity.com/paper/1909.11655/) (2019)
- [Gotta be SAFE: A New Framework for Molecular Design](https://mlsanity.com/paper/2310.10773/) (2023)
- [Structure-based drug design with geometric deep learning](https://mlsanity.com/paper/2210.11250/) (2022)
- [Application of generative autoencoder in de novo molecular design](https://mlsanity.com/paper/1711.07839/) (2017)
- [Group SELFIES: A Robust Fragment-Based Molecular String Representation](https://mlsanity.com/paper/2211.13322/) (2022)
- [JANUS: Parallel Tempered Genetic Algorithm Guided by Deep Neural Networks for Inverse Molecular Design](https://mlsanity.com/paper/2106.04011/) (2021)
- [Molecular Property Prediction Based on Graph Structure Learning](https://mlsanity.com/paper/2312.16855/) (2023)
- [FastFlows: Flow-Based Models for Molecular Graph Generation](https://mlsanity.com/paper/2201.12419/) (2022)
- [Recent advances in the Self-Referencing Embedding Strings (SELFIES) library](https://mlsanity.com/paper/2302.03620/) (2023)
- [Target Specific De Novo Design of Drug Candidate Molecules with Graph Transformer-based Generative Adversarial Networks](https://mlsanity.com/paper/2302.07868/) (2023)
- [REINFORCE-ING Chemical Language Models for Drug Discovery](https://mlsanity.com/paper/2501.15971/) (2025)

## Built on by (most cited first, 1 of 1)

- [Molecular Identifier Visual Prompt and Verifiable Reinforcement Learning for Chemical Reaction Diagram Parsing](https://mlsanity.com/paper/2603.15011/) (2026)

## Abstract (arXiv)

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.