[Paper Review] Reinforcement Learning for Molecular Design Guided by Quantum Mechanics

Gregor N. C. Simm, Robert Pinsler+1 more|arXiv (Cornell University)|Feb 18, 2020

Machine Learning in Materials ScienceMaterials Science37 references21 citations

TL;DR

This paper proposes a reinforcement learning framework for molecular design that operates directly in 3D Cartesian coordinates, enabling the generation of diverse molecular structures—including clusters and non-organic molecules—by placing atoms sequentially on a canvas. The method uses quantum-chemical energy calculations as a physics-based reward, and demonstrates successful learning via a rotation- and translation-invariant policy network in a new environment called MolGym, achieving high validity and structural stability in generated molecules.

ABSTRACT

Automating molecular design using deep reinforcement learning (RL) holds the promise of accelerating the discovery of new chemical compounds. Existing approaches work with molecular graphs and thus ignore the location of atoms in space, which restricts them to 1) generating single organic molecules and 2) heuristic reward functions. To address this, we present a novel RL formulation for molecular design in Cartesian coordinates, thereby extending the class of molecules that can be built. Our reward function is directly based on fundamental physical properties such as the energy, which we approximate via fast quantum-chemical methods. To enable progress towards de-novo molecular design, we introduce MolGym, an RL environment comprising several challenging molecular design tasks along with baselines. In our experiments, we show that our agent can efficiently learn to solve these tasks from scratch by working in a translation and rotation invariant state-action space.

Motivation & Objective

• To overcome limitations of graph-based molecular generation in reinforcement learning, which restricts models to single organic molecules and heuristic reward functions.
• To enable the design of complex molecular systems, including clusters and inorganic compounds, by operating in 3D Cartesian space rather than graph representations.
• To develop a reward function grounded in fundamental physical properties, specifically electronic energy approximated via fast quantum-chemical methods.
• To introduce MolGym, a benchmark environment with diverse molecular design tasks and quantum-chemistry-based baselines for evaluating RL agents.
• To design a policy network architecture invariant to translation and rotation, enabling stable and generalizable learning in 3D molecular space.

Proposed method

• The agent operates in a Markov decision process where states represent atomic coordinates and actions correspond to placing atoms at specific 3D positions on a canvas.
• The state representation is made invariant to translation and rotation by using relative coordinates and inter-atomic distances as input features.
• The reward function is derived from quantum-mechanical energy, computed via the PM6 semi-empirical method, to guide the agent toward low-energy, stable molecular configurations.
• A novel actor-critic neural network architecture is proposed, using message-passing layers to encode atomic interactions while preserving invariance to global translations and rotations.
• The training uses standard policy gradient methods (e.g., PPO) to optimize the policy end-to-end, with exploration enabled through stochastic action sampling.
• The MolGym environment includes tasks such as single-molecule formation, multi-bag synthesis, and solvation modeling, each with predefined target stoichiometries and baseline performance metrics.

Experimental results

Research questions

• RQ1 Can reinforcement learning in Cartesian coordinates generate stable, valid molecular structures beyond the scope of graph-based methods?
• RQ2 Can a physics-based reward function based on quantum-chemical energy improve the quality and stability of generated molecules compared to heuristic descriptors?
• RQ3 Can a translation- and rotation-invariant policy network effectively learn molecular design tasks from scratch without prior structural inductive biases?
• RQ4 To what extent does the proposed method scale to larger or more complex molecular systems, such as clusters or molecules with multiple components?
• RQ5 How does the performance of the agent compare to existing benchmarks in terms of validity, diversity, and structural stability?

Key findings

• The agent achieved 90% validity on the C2H2O2 task and 70–90% validity across various single-bag tasks, with median RMSD values below 0.5 Å, indicating high structural stability.
• For the multi-bag task, the agent generated 22 unique, valid molecules with a median RMSD of 0.05 Å, demonstrating effective learning on complex, multi-component systems.
• In solvation tasks, the agent achieved 80–90% validity and median RMSD values of 0.92–1.06 Å, showing capability in modeling solvent interactions.
• The method generated up to 61 unique molecules for the C7H8N2O2 task, indicating strong diversity despite the complexity of the target formula.
• The translation- and rotation-invariant policy network enabled stable training and generalization across different molecular tasks without requiring explicit geometric supervision.
• The use of quantum-chemical energy as a reward function led to the generation of low-energy, stable molecular configurations, as confirmed by post-optimization energy minimization.