[Paper Review] ChemGymRL: An Interactive Framework for Reinforcement Learning for Digital Chemistry

Chris Beeler, Sriram Ganapathi Subramanian+11 more|arXiv (Cornell University)|May 23, 2023

Innovative Microfluidic and Catalytic Techniques Innovation4 citations

TL;DR

This paper introduces ChemGymRL, an open-source, customizable reinforcement learning (RL) framework modeled after real-world chemistry labs, enabling training of RL agents on virtual chemical benches such as reaction, extraction, and distillation. It demonstrates that PPO and PPO-XL agents achieve optimal performance across multiple tasks, outperforming A2C, SAC, TD3, and DQN, especially in complex synthesis and purification workflows.

ABSTRACT

This paper provides a simulated laboratory for making use of Reinforcement Learning (RL) for chemical discovery. Since RL is fairly data intensive, training agents `on-the-fly' by taking actions in the real world is infeasible and possibly dangerous. Moreover, chemical processing and discovery involves challenges which are not commonly found in RL benchmarks and therefore offer a rich space to work in. We introduce a set of highly customizable and open-source RL environments, ChemGymRL, based on the standard Open AI Gym template. ChemGymRL supports a series of interconnected virtual chemical benches where RL agents can operate and train. The paper introduces and details each of these benches using well-known chemical reactions as illustrative examples, and trains a set of standard RL algorithms in each of these benches. Finally, discussion and comparison of the performances of several standard RL methods are provided in addition to a list of directions for future work as a vision for the further development and usage of ChemGymRL.

Motivation & Objective

• Address the data scarcity and safety challenges in training RL agents for real-world chemical synthesis by providing a simulated laboratory environment.
• Bridge the gap between autonomous laboratories and digital chemistry by enabling rapid, safe, and scalable RL experimentation.
• Create a modular, extensible testbed supporting diverse RL paradigms such as hierarchical, model-based, and constrained RL.
• Enable researchers to train and benchmark RL agents on realistic chemical processes with structured rewards and state transitions.
• Facilitate future development of lab manager agents that coordinate multiple bench-specific agents toward complex synthesis goals.

Proposed method

• Design ChemGymRL as a set of interconnected OpenAI Gym-compatible environments simulating real chemistry lab benches: reaction (RxN), extraction (ExT), distillation (DiT), and characterization.
• Implement a vessel-based state representation where materials and chemical states are tracked across benches, enabling sequential processing.
• Define environment dynamics using simplified differential equations for reaction kinetics and solvent evaporation, with configurable parameters.
• Assign sparse, task-specific rewards based on molar yield of target compounds or successful completion of purification steps.
• Support training of standard RL algorithms (PPO, A2C, DQN, SAC, TD3) using a modular, extensible codebase with full documentation and tutorials.
• Enable agent transfer between benches by modeling inter-bench material flow, mimicking real lab workflows.

Figure 1: (a) The ChemGymRL environment. Individual agents operate at each bench, working towards their own goals. The benches pictured are extraction ( ExT ), distillation ( DiT ), and reaction ( RxN ). The user determines which materials the bench agents have access to and what vessels they start

Experimental results

Research questions

• RQ1 Can standard RL algorithms learn optimal policies for complex chemical synthesis and purification tasks in a simulated lab environment?
• RQ2 How do different RL algorithms (e.g., PPO, A2C, DQN) compare in performance across distinct chemical processes like reaction, extraction, and distillation?
• RQ3 To what extent can RL agents generalize across variations in initial conditions, such as presence of impurities (e.g., salt) in distillation tasks?
• RQ4 Can RL agents learn behavior that matches or exceeds heuristic strategies in terms of efficiency and outcome quality?
• RQ5 How extensible and adaptable is the ChemGymRL framework to support advanced RL paradigms like curriculum learning, hierarchical RL, or inverse RL?

Key findings

• PPO and PPO-XL agents achieved optimal returns in both Wurtz reaction experiments, outperforming A2C, SAC, TD3, and DQN, especially in complex tasks with multiple targets.
• In the Wurtz distillation task, all algorithms outperformed random actions, with PPO achieving the highest returns, demonstrating effective learning of solvent evaporation control.
• In the Wurtz extraction task, PPO surpassed heuristic baselines, while A2C and DQN failed to improve over inaction, indicating algorithmic sensitivity to task structure.
• When salt was present in the distillation vessel, DQN agents persisted in learning policies as if salt were absent, revealing a failure to adapt to new conditions.
• A2C and DQN agents in the distillation task with salt did not modify behavior despite the added complexity, unlike PPO and PPO-XL, which adapted similarly to the heuristic policy.
• In the absence of salt, all top-performing agents (PPO, PPO-XL, A2C, DQN) learned policies nearly identical to the heuristic, including timely termination after solvent evaporation.

Figure 2: A visualization of the reaction bench ( RxN ) observation and action space. (a) An example of a UV-Vis spectra that would been seen in an observation and (b) The icons representing each action in RxN .