Bou, Albert
Test-Time Training Scaling for Chemical Exploration in Drug Design
Thomas, Morgan, Bou, Albert, De Fabritiis, Gianni
Chemical language models for molecular design have the potential to find solutions to multi-parameter optimization problems in drug discovery via reinforcement learning (RL). A key requirement to achieve this is the capacity to "search" chemical space to identify all molecules of interest. Here, we propose a challenging new benchmark to discover dissimilar molecules that possess similar bioactivity, a common scenario in drug discovery, but a hard problem to optimize. We show that a population of RL agents can solve the benchmark, while a single agent cannot. We also find that cooperative strategies are not significantly better than independent agents. Moreover, the performance on the benchmark scales log-linearly with the number of independent agents, showing a test-time training scaling law for chemical language models.
REINFORCE-ING Chemical Language Models in Drug Design
Thomas, Morgan, Bou, Albert, De Fabritiis, Gianni
Chemical language models, combined with reinforcement learning, have shown significant promise to efficiently traverse large chemical spaces in drug design. However, the performance of various RL algorithms and their best practices for practical drug design are still unclear. Here, starting from the principles of the REINFORCE algorithm, we investigate the effect of different components from RL theory including experience replay, hill-climbing, baselines to reduce variance, and alternative reward shaping. Additionally we demonstrate how RL hyperparameters can be fine-tuned for effectiveness, efficiency, or chemical regularization as demonstrated using the MolOpt benchmark.
Aviary: training language agents on challenging scientific tasks
Narayanan, Siddharth, Braza, James D., Griffiths, Ryan-Rhys, Ponnapati, Manu, Bou, Albert, Laurent, Jon, Kabeli, Ori, Wellawatte, Geemi, Cox, Sam, Rodriques, Samuel G., White, Andrew D.
Language agents [1-4] are AI agents [5] that integrate LLMs [6-8] as core components. LLMs excel at zero-shot generalization [9, 10], providing a notable advantage over traditional AI agents, such as those based on handcrafted rules or reinforcement learning, which often struggle to generalize to new environments [11]. While LLMs can exhibit flawed reasoning and logic when used in isolation [12-14], constructing a language agent by grounding LLMs in an environment with observational feedback can mitigate these issues. Early work on language agents used LLMs to directly output actions in the external environment [15-17], while more recently, language agents have been augmented with internal reasoning [18, 19] and planning [20, 21] procedures, as well as long-term memory storage [22, 23]. An emergent research challenge is to pose a theoretical description of the learning problem solved by language agents [4, 24] and to develop efficient methods to optimize the components of a language agent [24-26]. Here, we define common language agent tasks as language decision processes (LDPs) and frame language agents as stochastic computation graphs [27] that may be trained to solve LDPs. We show that pre-existing agents [18, 19, 21] can be implemented within our stochastic computation graph framework and introduce a simple and extensible software package named LDP that enables modular interchange of environments, agents, and optimizers, simplifying experimentation across a variety of settings. These authors jointly supervise technical work at FutureHouse.
ACEGEN: Reinforcement learning of generative chemical agents for drug discovery
Bou, Albert, Thomas, Morgan, Dittert, Sebastian, Ramírez, Carles Navarro, Majewski, Maciej, Wang, Ye, Patel, Shivam, Tresadern, Gary, Ahmad, Mazen, Moens, Vincent, Sherman, Woody, Sciabola, Simone, De Fabritiis, Gianni
In recent years, reinforcement learning (RL) has emerged as a valuable tool in drug design, offering the potential to propose and optimize molecules with desired properties. However, striking a balance between capabilities, flexibility, reliability, and efficiency remains challenging due to the complexity of advanced RL algorithms and the significant reliance on specialized code. In this work, we introduce ACEGEN, a comprehensive and streamlined toolkit tailored for generative drug design, built using TorchRL, a modern RL library that offers thoroughly tested reusable components. We validate ACEGEN by benchmarking against other published generative modeling algorithms and show comparable or improved performance. We also show examples of ACEGEN applied in multiple drug discovery case studies. ACEGEN is accessible at \url{https://github.com/acellera/acegen-open} and available for use under the MIT license.
TorchRL: A data-driven decision-making library for PyTorch
Bou, Albert, Bettini, Matteo, Dittert, Sebastian, Kumar, Vikash, Sodhani, Shagun, Yang, Xiaomeng, De Fabritiis, Gianni, Moens, Vincent
PyTorch has ascended as a premier machine learning framework, yet it lacks a native and comprehensive library for decision and control tasks suitable for large development teams dealing with complex real-world data and environments. To address this issue, we propose TorchRL, a generalistic control library for PyTorch that provides well-integrated, yet standalone components. We introduce a new and flexible PyTorch primitive, the TensorDict, which facilitates streamlined algorithm development across the many branches of Reinforcement Learning (RL) and control. We provide a detailed description of the building blocks and an extensive overview of the library across domains and tasks. Finally, we experimentally demonstrate its reliability and flexibility and show comparative benchmarks to demonstrate its computational efficiency. TorchRL fosters long-term support and is publicly available on GitHub for greater reproducibility and collaboration within the research community. The code is open-sourced on GitHub.
NAPPO: Modular and scalable reinforcement learning in pytorch
Bou, Albert, De Fabritiis, Gianni
Reinforcement learning (RL) has been very successful in recent years but, limited by its sample inefficiency, often requires large computational resources. While new methods are being investigated to increase the efficiency of RL algorithms it is critical to enable training at scale, yet using a code-base flexible enough to allow for method experimentation. Here, we present NAPPO, a pytorch-based library for RL which provides scalable proximal policy optimization (PPO) implementations in a simple, modular package. We validate it by replicating previous results on Mujoco and Atari environments. Furthermore, we provide insights on how a variety of distributed training schemes with synchronous and asynchronous communication patterns perform. Finally we showcase NAPPO by obtaining the highest to-date test performance on the Obstacle Tower Unity3D challenge environment. The full source code is available.