AceGen: A TorchRL-Based Toolkit for Reinforcement Learning in Generative Chemistry

Abstract

In recent years, reinforcement learning (RL) has been increasingly used in drug design to propose molecules with specific properties under defined constraints. However, RL problems are inherently complex, featuring independent and interchangeable components with diverse method signatures and data requirements, leading existing applications to convoluted code structures. This complexity not only complicates code comprehension but also hampers modification, hindering the smooth exploration of new ideas in the field and ultimately slowing down research. In this work, we apply TorchRL - a modern general decision-making library that provides well-integrated reusable components - to make a robust toolkit tailored for generative drug design. AceGen leverages general RL solutions which enhance simplicity, making the solutions more understandable, modifiable, and reliable. We demonstrate the application of AceGen for conditioned compound library generation implementing various RL algorithms to optimize drug design targets. Furthermore, with the tools made available we propose a novel algorithm inspired by the PPOD algorithm that outperforms all baselines as benchmarked on 23 drug design relevant targets. The library is accessible at https://anonymous.4open.science/r/acegen-open-23D3.