ACEGEN: A Reinforcement Learning Toolkit for Generative Drug Design
Published on May 15, 2024

The advent of advanced machine learning techniques has opened new frontiers. Among these, reinforcement learning (RL) stands out for its potential to optimize and propose novel molecules with desired properties. Enter ACEGEN, an open source toolkit designed to use RL for generative chemical design. Developed by our team at Acellera Therapeutics and collaborators, ACEGEN offers a modular and configurable toolkit tailored specifically for drug discovery.
The Need for ACEGEN
Drug design is a complex and multifaceted process involving the identification of biomolecules that exhibit optimal properties such as potency, selectivity, bioavailability, and toxicity. Traditional methods often fall short in efficiently navigating the vast chemical space, which contains an astronomical number of potential compounds. Reinforcement learning, with its trial-and-error approach and adaptive learning capabilities, presents a promising solution to this challenge.
What Makes ACEGEN Unique?
ACEGEN leverages TorchRL, a robust decision-making library within the PyTorch ecosystem. This integration ensures ACEGEN benefits from well-tested, modular components that enhance efficiency and reliability. The toolkit provides pre-trained models for various architectures, including GRU, LSTM, GPT-2, Llama2, and Mamba, which can be seamlessly integrated with custom architectures.
Key Features of ACEGEN
1. Modular and Flexible Design
ACEGEN is flexible, allowing researchers to incorporate custom scoring functions and models. The toolkit supports various molecular grammars, such as SMILES, DeepSMILES, SELFIES, and AtomInSmiles, providing a common interface for different molecular representations.
2. Comprehensive RL Algorithms
The toolkit includes implementations of several RL algorithms, such as REINFORCE, REINVENT, AHC, A2C, PPO, and PPOD. These algorithms are fully configurable, enabling customization to meet specific research needs. Notably, their implementations and evaluation are described in the linked publication.
3. Efficient Training and Benchmarking
ACEGEN includes training and benchmarking workflows. The toolkit also offers a tutorial for integrating custom scoring functions and models.
Benchmarking Performance
To assess the performance of different RL algorithms, ACEGEN was benchmarked using the Practical Molecular Optimization (MolOpt) framework. The results for algorithms implemented in ACEGEN, including PPOD and REINVENT-MolOpt, are reported in the linked publication with the benchmark definition and evaluation measures.
Looking Ahead
ACEGEN provides reusable components for generative molecular design. It is available under the MIT licence so the scientific community can inspect, use and contribute to its development.
For more information and to access the ACEGEN toolkit, visit the ACEGEN GitHub repository.
Check the publication here.
Call for collaborations
We want to hear from you if you’re a pharma or biotech company looking to leverage generative AI on your chemical libraries or synthesis rules. We’re looking to partner with small and large companies to make ACEGEN an even more powerful solution for drug discovery.
If you’re interested, get in touch.
References:
- Bou, A., Thomas, M., Dittert, S., et al. "ACEGEN: Reinforcement Learning of Generative Chemical Agents for Drug Discovery." GEM workshop, ICLR 2024.
- Blaschke, T., Arús-Pous, J., Chen, H., et al. "REINVENT 2.0: An AI Tool for De Novo Drug Design." Journal of Chemical Information and Modeling, 60(12):5918-5922, 2020.
- Gao, W., Fu, T., Sun, J., et al. "Sample Efficiency Matters: A Benchmark for Practical Molecular Optimization." Advances in Neural Information Processing Systems, 35:21342-21357, 2022.