JCIM 2024 | ACEGEN: 用于药物发现的生成式化学智能体强化学习
今天介绍一下发表在Journal of Chemical Information and Modeling上的一篇关于分子强化学习的论文。这项工作不仅提出了一个面向药物生成的强化学习工具包，更重要的是在方法论层面回应了当前药物设计中的几个核心难题：一是如何在复杂的化学空间中实现高效搜索，二是如何在性能优化与化学合理性之间取得平衡，三是如何通过模块化框架降低重复造轮子的问题。ACEGEN的价值不止于实现REINFORCE、REINVENT或PPO等算法，而在于它建立了一个系统化的实验平台，使不同算法和组件能够在统一环境下被对比、消融和扩展。论文强调了正则化在实际应用中的重要性，同时提醒读者，性能指标在不同任务间可能并不一致，因此更合理的基准和化学感知指标是未来发展的方向。可以说，ACEGEN不仅是一个代码库，更是推动强化学习在药物设计领域标准化、可复现和可扩展研究的里程碑。
获取详情及资源:
0 摘要
近年来，强化学习(RL)在药物设计中展现出重要价值，能够用于提出并优化具备期望性质的分子。然而，由于高级 RL 算法的复杂性以及对专用代码的高度依赖，在能力、灵活性、可靠性与效率之间取得平衡依然是巨大挑战。本研究提出了 ACEGEN，这是一个专为生成式药物设计打造的全面且精简的工具包，基于现代 RL 库 TorchRL 开发，提供经过充分测试的可复用组件。通过与其他已发表的生成建模算法进行基准对比，结果表明 ACEGEN 具备相当或更优的性能。此外，还展示了 ACEGEN 在多个药物发现案例研究中的应用。(ACEGEN 已开源，MIT 许可证下可自由使用)
1 引言
药物设计是一个复杂的过程，需要在效力、选择性、生物利用度和毒性等多种性质之间找到最佳平衡。近年来，生成建模方法被提出，作为部分自动化新分子设计的重要途径，希望在设计-合成-测试-分析循环中同时优化多个目标性质。这些模型大多依赖机器学习算法来生成候选分子，但如何高效搜索庞大的化学空间以找到最优分子，仍然是一大挑战，因为这一空间大到无法用简单枚举的方式穷尽。
强化学习(RL)近年来被认为是探索化学空间的一种潜在解决方案，尤其关注效率问题。RL 是一类利用反馈信号来指导决策过程的机器学习算法，因此可以将分子构建的决策过程调整为实现特定分子性质的目标，从而生成具有理想性质的新分子。这实际上构成了一种化学空间的搜索策略。RL 的核心在于奖励函数，它为分子分配数值，通常根据具体研究问题定制。RL 算法则通过最大化奖励来优化化学空间的探索。当化学空间无法通过打分函数直接枚举或筛选时(例如打分函数过慢，或搜索空间过大)，RL 尤其具有吸引力。
然而，RL 在药物设计中的应用与传统场景存在差异。经典 RL 通常基于博弈论，有明确环境与目标，例如“通关”或“赢得比赛”。相比之下，药物设计的目标是多维度的(既高效又低毒且可利用)，且往往依赖代理指标来代替真实目标，例如优化预测的蛋白结合亲和力，而非真正调控细胞通路。这种差异不可避免，因为实验验证成本极高。因此，为了增加发现具备所有理想性质分子的几率，必须提出多样化的解决方案，这与经典 RL“一次成功解”足够的目标截然不同。尽管如此，RL 在不同分子表示和模型架构上的应用已经展示了初步的成功。
目前已有的药物设计 RL 实现仍严重依赖定制代码，带来冗余、复杂和效率低下，限制了多样化方法的融合。而 RL 社区已有相应的基础设施来缓解这些问题，例如 TorchRL 这一现代 RL 库，提供了经过充分测试的独立模块。本文提出的 ACEGEN，便是基于 TorchRL 组件构建的药物生成工具包，能够组装出高效、可靠的药物发现智能体。这一思路已在无人机控制和组合优化等领域验证成功。此外，该框架便于研究人员在现有组件上扩展新算法，并在 PyTorch 生态中保持高质量、可维护和可拓展性。
为展示 ACEGEN 的优势，研究实现并评估了若干语言模型驱动的药物设计 RL 算法。语言模型能够学习文本中的复杂模式并生成新序列，包括分子结构。特别是，SMILES 表示因其紧凑和标准化的特点成为常用表示形式。此前多项研究已证明生成语言模型与 RL 结合的优势。因此，本文重新实现了三种基于 REINFORCE 的算法：REINFORCE、REINVENT 和 AHC，并改造了通用 RL 算法如 A2C 和 PPO，使之适应药物设计任务。进一步地，研究在 REINFORCE 中引入经验回放机制，并测试了基于 PPO 的变体 PPOD。
最后，为展示 ACEGEN 的应用潜力，研究在多个场景下对这些 RL 算法进行了评估，包括：采样效率基准测试、剖析 RL 算法组件的消融实验、更具挑战性的药物设计目标以及分子约束生成设计。ACEGEN 已在 MIT 协议下开源，代码可在 https://github.com/acellera/acegen-open 获取，以促进社区使用和发展。
图 1 | ACEGEN 实现的一般流程概览。不同的 ACEGEN 实现依赖于不同的算法，并允许用户自定义生成模型和打分函数。
2 方法
2.1 强化学习设定
强化学习(RL)任务通常形式化为马尔可夫决策过程(MDPs)，其由五元组⟨S, A, R, P, ρ₀⟩描述。其中，
这一框架可应用于分子逐步设计问题。在此情境下，参数化的强化学习策略
在这一过程中，常用的策略梯度方法包括REINFORCE、A2C(优势Actor-Critic)以及PPO(近端策略优化)。此外，还有多种技术可辅助训练。例如，奖励塑形(reward shaping)通过改造奖励函数为智能体提供更多反馈，从而生成更具信息性的信号；经验回放(experience replay)通过存储并随机采样过去的经验，提高采样效率并稳定训练过程；引入惩罚项(如KL散度损失)则能约束智能体保持接近参考策略，从而增强稳定性；另外，在每一批次中对分子进行排序，仅保留前K个最优分子，也能提升训练效率。
ACEGEN整合了上述算法与辅助技巧，形成了一个覆盖最新方法的分子生成工具集。
2.2 化学语言生成模型
分子字符串表示能够在分子图与字符串之间进行转换，因此可将分子生成任务形式化为一种自然语言处理(NLP)问题。在这一框架下，用于按序生成字符串分子的策略模型被称为化学语言模型(Chemical Language Models, CLMs)。在这种设定中，动作空间是离散的，每一个动作
ACEGEN目前提供了一个完整的语言模型实验环境，支持SMILES、DeepSMILES、SELFIES、AtomInSmiles以及SAFE语法，并配备了用户友好的词汇类工具。在附录的Section A中展示了如何轻松创建和使用词汇表与环境来进行数据生成。CLMs最初在大量无标签数据上进行无监督训练，以学习生成有效分子，通常通过 教师强制(teacher forcing) 方法来实现。经过预训练的CLM即为初始策略函数
目前，ACEGEN已提供多种预训练模型。其中包括在ChEMBL和ZINC数据集上预训练的GRU策略模型，在ChEMBL上预训练的LSTM策略模型，以及在Enamine REAL lead-like化合物数据库上预训练的GPT2策略模型。仓库中还提供了可直接使用的LSTM、GRU、GPT2、Llama2与Mamba等多种策略架构。用户也可以根据教程自由集成其他架构，而无需修改ACEGEN的核心代码。此外，ACEGEN提供了便捷的预训练脚本，既可用于官方框架，也可支持用户自定义语言模型训练。
代码设计充分考虑计算资源的灵活性，无论是单GPU还是多机多GPU分布式环境，都能高效适配，从而实现对大规模数据集的高效训练。
2.3 分子打分与评估
在药物设计任务中，评分函数必须能够反映真实的研发场景，并具备足够的灵活性，以适用于不同类型的设计挑战。然而，现有研究中常常使用过于简单的目标函数进行优化，例如带惩罚的logP，或者仅为某一特定生成模型量身定制的复杂解决方案。
ACEGEN在这一点上提供了高度可扩展性，允许用户自定义评分函数。具体而言，用户可以将评分函数定义为Python方法，输入为字符串，输出为数值，从而灵活地融入不同的优化目标。在附录的Section B中展示了如何在ACEGEN中轻松实现评分函数。此外，库中还提供了详细的教学示例，指导用户如何实现并将自定义评分函数无缝整合进完整工作流中，用于训练智能体。
在基准测试中，研究使用了MolScore库，该工具涵盖了与药物设计密切相关的多种评分函数，提供了多样性过滤器、对课程学习的支持，并内置了多种基准模式，包括MolOpt、GuacaMol等。这些功能使得分子生成与评估能够更贴近实际应用场景。
2.4 RL 智能体训练
ACEGEN支持多种强化学习智能体的训练方法，包括REINFORCE、REINVENT、AHC、A2C、PPO以及改进版本的PPOD算法。在实现上，PPOD与原始方法略有不同：省略了初始专家示范，仅重放高奖励的轨迹(而非高价值预测的轨迹)，并在每个批次中使用固定数量的回放数据。所有方法均可高度自定义，能够结合用户定义的评分函数和模型，而不仅限于库中预置的选项。
在这些方法中，REINVENT与AHC可视为REINFORCE的扩展版本。REINVENT在基础框架中引入了奖励塑形、经验回放以及对高似然序列的惩罚项；而AHC则通过对每个批次中的分子进行排序，仅保留前K个分子以优化效率。然而，研究发现奖励塑形在优势函数类方法中并不兼容，同时高似然惩罚项对性能反而有害。因此，对于A2C、PPO和PPOD，引入了一个基于 Kullback–Leibler散度(KL散度) 的损失函数项，用于约束策略偏离先验策略过远。这一机制在概念上类似于REINVENT和AHC中的奖励塑形。
引入KL约束已成为强化学习与语言模型结合的常见做法，特别是在利用人类反馈对语言模型进行对齐的研究中，具有重要的稳定性和实用价值。
2.5 从头生成、骨架修饰与片段连接
在药物研发流程中，单纯的从头分子生成(denovo generation)往往难以满足多样化需求。为此，ACEGEN提供了多种采样模式以支持不同场景。PromptSMILES是一种简便的方法，能够在仅基于完整SMILES进行教师强制训练的模型上，实现受约束的分子生成，从而保证生成结果符合特定的化学子结构。
除了从头生成，ACEGEN脚本还支持灵活配置骨架修饰(scaffold decoration)与片段连接(fragment-linking)等模式。在这些设定下，受约束采样只会改变数据收集的方式，而不影响其他任何智能体组件的运行逻辑，这一设计理念与TorchRL的思想保持一致。
在仓库中，还提供了详细的教学示例，指导如何进行受约束的分子生成，方便研究者快速上手并应用于实际的药物设计任务。
表 1 | 各算法在 MolOpt 基准任务中的表现比较，评价指标为前 10 个分子的曲线下面积（AUC）。
3 结果
为了展示ACEGEN的多种应用场景，研究者进行了化学基准测试、对REINVENT方法的消融实验、一个具有挑战性的药物设计目标测试，以及基于PromptSMILES的受约束强化学习探索。
3.1 强化学习性能基准测试
为了评估不同的RL算法，研究者在Practical Molecular Optimization (MolOpt)基准上进行了对比，该基准在MolScore中实现，涵盖了23个不同任务，每个任务在1万个分子的预算范围内优化不同目标。所有算法均使用相同的GRU策略模型架构，与原始MolOpt基准保持一致，并在经过筛选的ChEMBL子集上完成预训练。
对比的方法包括：REINFORCE，使用原始论文超参数的REINVENT，为MolOpt任务优化超参数的REINVENT-MolOpt，以及使用原始论文超参数的AHC。需要注意的是，上述算法均使用了经验回放。此外，还测试了不依赖经验回放的A2C与PPO，以及为适配经验回放而调整的PPOD。所有算法均由ACEGEN实现，相关超参数值在仓库中提供。作为验证，研究者还将自己的REINVENT实现与MolOpt论文中的实现进行了对比，结果表明简化后的版本不仅训练速度更快，而且性能更优。详细结果见附录Section C。
按照MolOpt基准的建议，表1展示了不同算法在前10个分子的AUC表现，这一指标反映了在目标导向下找到10个理想分子的采样效率，具有推动药物设计后续阶段的代表性。结果显示，PPOD在采样效率上表现最佳，其次是REINVENT-MolOpt，再次为PPO。
为了获得更全面的性能理解，表2与图2a进一步展示了一系列衡量最大性能、效率与探索度的指标(探索与利用之间存在常见的效率权衡)。结果表明：REINVENT-MolOpt在平均前10个分子指标上达到最大性能，PPOD在前10分子AUC上表现出最高效率，而AHC则在独特分子生成数量上实现最大探索度。
表 2 | 各算法在 MolOpt 基准任务上的性能比较，涵盖多项指标的分析。
3.2 用于实际药物发现的强化学习性能评估
在药物发现中，RL通过优化策略来最大化未来累计奖励，而奖励可以来自单一评分函数，或多个评分函数的加权和。然而在实践中，往往难以精确定义所需的评分函数及其权重。例如，评分函数有时并不能生成符合预期化学性质的分子，或存在明显的利用漏洞。在这种情况下，引入对先验策略的正则化是十分有益的，因为先验策略基于特定训练数据集学习了化学空间的分布，这一分布本身就代表了某种“理想化学”的示例。因此，像REINVENT这样的算法被设计为在启发式上包含正则化项，而不仅仅依赖奖励信号。
为此，研究者对性能评估指标进行了修正：对前10个分子的平均值与AUC增加了多样性约束，要求这些分子之间的ECFP4 Tanimoto相似度小于0.35。此外，还采用了**球面排斥多样性(SEDiv@1k)**来评估探索性，该指标衡量在随机抽取的1000个分子样本中，需要多少比例的分子来覆盖化学空间，被认为比单纯的“独特分子数量”更具代表性。
同时，研究者对分子引入了一系列过滤条件。基础化学过滤器(B-CF)要求：logP ≤ 4.5、可旋转键数 ≤ 7、分子量150–650 Da、原子仅限于{C, S, O, N, H, F, Cl, Br}，且不得触发已知的结构警示。此外，基于训练数据分布，还设定了目标化学过滤器(T-CF)：logP与分子量必须位于训练数据均值μ±4σ范围内，并且分子中新颖原子环境的ECFP4位点不能超过10%(如图S6所示)。两者结合称为B&T-CF。
使用这些新指标重新评估RL性能后(见表2)，发现REINVENT-MolOpt、A2C、PPO和PPOD中通过化学过滤器的分子比例更低，这与它们弱化正则化、倾向于效率优化的特性一致。尽管如此，REINVENT-MolOpt在前10个多样性分子的平均值与AUC上仍然取得最高分，其次为PPOD。在探索性方面，REINVENT的球面排斥多样性最高，其次是AHC。有趣的是，尽管REINFORCE未显式引入正则化项，却依然保持了较高比例的分子通过化学过滤器，并在绝对性能与效率上超过REINVENT。
进一步的分析(见表3与图2b)将化学过滤器与多样性过滤器显式纳入奖励函数，结果如预期一般：通过过滤器的分子比例提高，但所有算法在通过过滤后的前10个多样性分子的平均值与AUC上均有所下降。这表明，在整体优化与效率(尤其是与化学质量相关的指标)方面，基于先验策略的隐式正则化更为有效。这种现象可能源于奖励信号的具体实现方式，例如未通过过滤器的分子被直接赋予0分，而如何优化奖励函数本身就是一门“艺术”。因此，更简洁的奖励信号往往更优，而通过正则化来约束化学空间则能简化奖励函数的设计。
值得注意的是，优化算法、奖励信号与正则化的微调可能具有目标依赖性，这超出了本文的范围，但ACEGEN为这一探索提供了可能性。
图 2 | RL 算法性能的雷达图比较。 a MolOpt 基准（见表 2）；b 在奖励信号中显式加入化学约束的 MolOpt 基准（见表 3）；c 5-HT2A 案例研究（见表 5）；d REINFORCE 消融实验（见表 4）。图例位于图上方（a–c），在（d）中位于子图右侧。
3.3 REINVENT算法的消融研究
在用于CLMs微调的RL算法中，REINVENT是应用最为广泛的方法之一。从理论角度看，该算法可以视为REINFORCE、经验回放、奖励塑形以及序列似然惩罚项的组合。然而，每个机制对整体性能的具体贡献尚不明确。为此，研究者借助ACEGEN的模块化组件，对这三个关键机制进行了消融实验，以探究其在REINVENT中的作用，并使用默认超参数在MolOpt基准上进行性能比较。
从实验结果来看，表4与图2d表明，性能提升的主要来源是在REINFORCE中引入经验回放。奖励塑形能够进一步提升性能，但与经验回放结合时，并未带来额外增益。不过，若采用将智能体策略与先验策略关联的奖励塑形，则可以提高化学质量，即通过化学过滤器的有效且独特分子比例增加，但这往往伴随着整体性能的下降。这一现象再次强调了正则化在实际应用中的重要性，与前一部分的讨论相呼应。
在REINVENT的原始实现中，作者还引入了一个惩罚高似然序列的损失项。表4结果显示，该项的主要作用是增加探索性，但对其他指标均有负面影响，尤其是降低了通过化学过滤器的分子比例。更严重的是，如果在没有奖励塑形的情况下使用该损失项，会导致学习完全崩溃。
这些结果揭示了REINVENT中不同机制的作用与局限，也为如何平衡性能、化学质量与探索度提供了重要启示。
表 3 | 含或不含显式化学过滤器作为奖励信号打分函数时的算法性能比较。
3.4 案例研究：在5-HT2A上的从头分子生成
为了更贴近实际药物发现中的复杂情境，研究者将不同的RL算法应用于一个更具挑战性的目标任务，该任务此前已被提出。其目标是设计对5-HT2A受体具有选择性、且能区分与高度相关的D2受体的配体，仅依赖结构信息完成。这一挑战极具临床相关性，因为在精神病治疗中，D2受体的非靶向作用常导致不良的锥体外系副作用。
具体而言，该任务要求在32,000个分子的预算内，最小化分子对5-HT2A受体(PDB: 6A93)的对接打分，同时最大化其对D2受体(PDB: 6CM4)的对接打分。两种晶体结构均与利培酮(Risperidone)共晶，反映了两者结合口袋的高度相似性。与原始方案不同，本研究使用了rDock软件，因其许可更为宽松。需要指出的是，尽管对接打分通常不是结合亲和力的可靠预测指标，但它们在虚拟筛选中依然能够提升分子选择的富集度，并在某些情况下优于基于配体的亲和力预测方法。
表5与图2c给出了性能指标的总结。结果显示，若同时考虑优化性能与化学质量(尽管化学质量并未在目标函数中被显式约束)，AHC取得了最高的前10个分子平均值与采样效率，同时保持了较高比例的分子通过化学过滤器，并展现出较高的多样性。与偏向利用的MolOpt基准相比，该任务凸显了在实际应用中为正则化而牺牲部分性能的价值。
图3展示了一些从头生成的分子及其在5-HT2A受体上的预测对接构象。关于更多方法、分析及从化学角度解读不同RL算法差异的内容，可参见附录Section E。
表 4 | REINFORCE 算法在 MolOpt 基准上的消融研究，对比了包含或不包含经验回放、奖励塑形以及高序列似然惩罚等独立组件的表现。
表 5 | 各算法在 5-HT2A 案例研究中的性能总结。
3.5 案例研究：骨架约束生成
ACEGEN能够利用PromptSMILES实现受约束的分子生成。该方法通过在SMILES上进行迭代式CLM提示，并结合RL对预训练CLMs进行微调，从而完成特定的约束生成任务。本研究展示了两个应用场景：(1)基于已知合成反应的骨架修饰(如LibINVENT提出的方法)，以及(2)骨架修饰结合骨架约束对接，用于探索结合口袋中的生长向量。
在第一个实验中，研究者比较了不同RL算法在LibINVENT提出的两个任务上的表现。这两个任务均要求对哌嗪(piperazine)骨架进行修饰，并带有D2受体配体常见的胺连接基团。与MolOpt基准类似，预算为10,000个分子。第一个任务旨在优化D2活性预测概率(由QSAR模型评估)；第二个任务则加入了选择性反应过滤器，分别奖励哌嗪的Buchwald反应增长和胺连接基的酰胺偶联增长。表6结果显示，两类任务均实现了成功优化：AHC在第一个任务中表现最佳，而REINVENT(无论是默认参数还是MolOpt参数)在第二个任务中表现最佳。但总体而言，除A2C外，所有算法的成功率均超过90%，并且在80%至90%的分子中完全满足选择性反应目标。
图 3 | 5-HT2A 选择性任务中前 10 个分子的一些代表性示例及其在 5-HT2A 结合口袋（PDB: 6A93）中的对接构象。 共结晶配体利培酮（risperidone）作为参考。
在第二个实验中，研究者结合PromptSMILES与AHC(基于其在5-HT2A案例中的优异表现)，在阿尔茨海默病关键治疗靶点BACE1的催化位点进行骨架修饰。所用参考结构为PDB: 4B05，与典型BACE1抑制剂AZD3839共晶。该抑制剂包含一个与两个催化天冬氨酸相互作用的双环脒核心，并占据P2′、P1和P3亚口袋。研究将该双环脒核心作为PromptSMILES的提示，允许沿两个方向进行增长：一个延伸至P2′亚口袋，另一个延伸至P1和P3亚口袋。分子的奖励函数设定为：在子结构约束下最小化对接得分，同时保持合理的重原子数与可旋转键数量(见附录F)。
图4展示了在生成10,000个分子的过程中，奖励函数成功最大化(图4a,b)并对应地最小化了对接得分(图4c)。最终识别出的前10个分子的对接得分在−23.9至−15.8之间，显著优于参考分子(重对接得分为−8.44)。对接构象(图S13)进一步表明，生成分子在P1与P2亚口袋中成功恢复了芳香环的占据。
这一案例研究突出了在复杂目标下实现受约束优化的可行性与有效性。
表 6 | 各算法结合 PromptSMILES 在 LibINVENT DRD2 任务中的约束分子生成性能比较（分别在无选择性反应过滤器和有选择性反应过滤器条件下）。
图4 | 多目标奖励的优化 展示了平均奖励及其底层对接打分的优化情况。图中给出了按多目标奖励排序的前10个从头生成分子，红色部分标注了受约束的子结构，对接得分标注在下方。作为参考，共晶配体经重新对接后的得分为−8.44。其对接构象见图S13。
4 结论
本研究介绍了ACEGEN，这一全新工具集整合了强化学习(RL)在药物发现领域的最佳方法。通过借鉴通用且经过充分测试的决策库TorchRL中的RL构件，ACEGEN提供了模块化、高效且灵活的解决方案，并在其中实现了一系列基于语言模型的方法。研究展示了ACEGEN在方法基准测试、算法探索以及真实药物设计应用等多个方面的能力。
实验结果不仅加深了对REINVENT等常用算法的理解，还为在不同应用场景中选择最合适的算法提供了实践性启示，同时凸显了一个可靠而全面的工具集的重要性。总体而言，ACEGEN能够有效应对药物发现中的多样化需求，并在解决该领域复杂挑战方面展现出独特优势。这一工作也是探索RL配置空间模块化的第一步，未来将以此为基础继续拓展潜在的改进方向。
值得注意的是，要实现性能提升，必须对目标行为进行合理评估，而这依赖于更完善的基准。例如，在MolOpt基准上表现最佳的RL算法(REINVENT-MolOpt与PPOD)，在5-HT2A对接基准上的表现却并非最优，这凸显了性能评估存在不一致性。虽然本研究提出了新的化学感知指标，更贴近RL算法在实际应用中的需求，但仍需构建新的基准——既能体现如5-HT2A任务中探索的重要性，又能避免过高的计算开销。
最后，目前ACEGEN仅包含针对CLMs的RL环境与架构，未来将进一步扩展至更多类型的架构，尤其是能够直接在三维空间中进行生成的模型，这将为药物设计带来更广阔的前景。