1. 项目概述当强化学习在“推理时”遇见分子生成最近在药物发现和材料科学领域基于深度学习的分子生成模型已经不是什么新鲜事了。从早期的VAE、GAN到后来的流模型和扩散模型大家的目标都很明确生成具有特定性质比如高活性、低毒性、可合成性的新分子结构。但实际操作过的人都知道这事儿有个核心痛点生成容易优化难。模型可能一下子吐出几万个分子但真正符合我们苛刻的、多目标优化需求的往往凤毛麟角。这时候“推理时强化学习”这个概念开始进入我们的视野。它不像传统方法那样在模型训练阶段就把所有优化目标“硬编码”进去而是把优化这件事推迟到了模型已经训练好、准备生成样本的“推理”阶段。你可以把它想象成一个“事后诸葛亮”式的智能编辑器模型先生成一个基础分子然后一个强化学习智能体上场对这个分子进行一系列“编辑”操作比如添加、删除或改变原子/化学键每编辑一步就根据我们设定的奖励函数比如预测的活性得分、类药性分数得到一个反馈引导编辑朝着更优的方向进行。我之所以花大力气研究这个策略是因为它在实际项目中展现出了独特的灵活性。传统的分子生成优化一旦模型训练完成其优化方向就基本固定了想要临时增加一个优化目标比如突然要求分子还必须满足某种特定的ADMET性质就得重新收集数据、重新训练模型周期长、成本高。而推理时强化学习我们只需要调整推理阶段的奖励函数就能快速响应新的优化需求这在新药研发这种需求快速迭代的场景下价值巨大。本次分享我就来拆解一下我们是如何设计这套优化策略并对其效果进行系统评估的。2. 核心思路与方案选型背后的考量2.1 为什么是“推理时”强化学习在分子生成任务中我们通常面临的是一个序列决策问题从左到右依次生成表示分子的字符串如SMILES或者以图的形式逐步添加原子和键。传统的做法是在训练生成模型比如一个RNN或图神经网络时就引入强化学习用策略梯度方法直接优化期望奖励。这被称为“训练时”强化学习典型代表如REINFORCE算法在序列生成中的应用。但这种方法有几个固有的问题训练不稳定策略梯度方法的方差通常很高尤其是在生成长序列时这会导致训练过程震荡剧烈难以收敛。优化目标单一化训练时通常只能针对一个或少数几个主要奖励进行优化。当面对多目标、多约束的复杂优化时奖励函数的融合设计变得极其困难容易顾此失彼。灵活性差模型一旦训练完成其“偏好”就被固化。如果后续想微调优化方向哪怕只是调整一下不同奖励的权重都可能需要从头开始或进行复杂的微调。“推理时”强化学习的核心思想是解耦。我们将一个已经预训练好的、具备强大生成能力的模型例如一个在大规模分子库上训练好的SMILES自回归模型或图生成模型作为“基础生成器”。这个生成器的目标是学习真实的分子数据分布确保生成的都是语法正确、结构合理的分子。然后在推理阶段我们引入一个独立的强化学习智能体它的行动空间是“如何修改基础生成器产生的初始分子”。这种解耦带来了几个关键优势稳定性基础生成器是稳定预训练的强化学习智能体只负责局部、迭代式的修改动作空间相对较小训练和推理都更稳定。多目标优化的灵活性奖励函数完全在推理阶段定义和计算。我们可以轻松地组合多个预测模型QSAR模型、毒性预测器、合成可及性评分器的输出通过加权和或其他方式构建一个综合奖励。想调整优化重点直接改奖励函数的权重即可无需动基础模型。利用预训练知识我们可以直接利用在海量数据上预训练好的强大生成模型如ChemBERTa、GPT for Chemistry在其强大的分子表示和生成能力基础上进行定向优化避免了从零开始训练的数据和算力消耗。2.2 我们的核心架构设计基于以上考量我们设计了一套模块化的架构主要包括三个核心组件预训练分子生成器我们选择了一个基于Transformer架构的自回归模型在ZINC15、ChEMBL等大型分子数据库的SMILES序列上进行预训练。它的任务很简单给定一个已有的序列前缀预测下一个最可能的token原子或符号。这个模型保证了我们生成的起点至少是一个“像样”的分子。分子编辑智能体这是强化学习的核心。我们将其设计为一个基于策略网络的智能体。其状态State是当前分子图的向量表示通过一个图神经网络编码得到。其动作Action空间定义为一组可行的分子编辑操作例如AddAtom(node_type, bond_type)在指定位置添加一个特定类型的原子并以特定类型的键连接。RemoveAtom(node_id)移除一个原子及其相连的键需保证剩余子图合法。ChangeBond(node_id1, node_id2, new_bond_type)改变两个原子间化学键的类型单键、双键、三键。ChangeAtom(node_id, new_atom_type)改变一个原子的类型如C改为N。 其策略Policy是一个神经网络输入状态输出在所有可能动作上的概率分布。奖励计算引擎这是一个可插拔的模块。它接收编辑后的分子调用一系列外部或内置的评估器进行计算返回一个标量奖励。我们的奖励函数R(m)设计为多目标的加权和R(m) w1 * Activity(m) w2 * QED(m) w3 * SA_Score(m) w4 * Penalty(m)其中Activity(m)针对特定靶点的预测活性pIC50等使用预训练的QSAR模型。QED(m)类药性定量评估值在0-1之间。SA_Score(m)合成可及性评分经过归一化越高越易合成。Penalty(m)惩罚项用于约束分子量、脂水分配系数LogP、氢键供体/受体数量等理化性质在合理范围内违反约束则扣分。整个推理时优化流程是一个迭代过程生成器提供一个种子分子 - 智能体根据当前策略选择编辑动作 - 应用编辑得到新分子 - 奖励引擎计算奖励 - 智能体根据奖励更新策略或用于蒙特卡洛树搜索等规划算法- 重复此过程直至达到预设步数或收敛。3. 关键实现细节与避坑指南3.1 动作空间设计与合法性校验定义动作空间是第一步也是最容易出错的一步。一个看似简单的AddAtom操作需要考虑添加在哪里我们通常允许在所有现有原子上添加但需要避免产生过度拥挤的原子如碳原子连接超过4个键。添加什么原子常见的有机元素如C, N, O, F, Cl, S等。以什么键型连接单键、双键、芳香键等。实操心得动作空间不能定义得过大否则搜索效率极低。我们通过化学规则进行了大量剪枝。例如禁止在氢原子上继续添加原子禁止形成三元环或四元环外的超小环张力大添加双键或三键时必须检查相关原子的价态是否允许。我们实现了一个实时分子合法性校验器在任何编辑动作被最终执行前都会用RDKit等化学信息学工具快速检查结果分子是否价态合理、无奇异电荷、成键合法。如果动作会导致非法分子则在该步给予一个极大的负奖励并禁止该动作或者让智能体重新选择。这一步是保证整个流程产出有效分子的生命线。3.2 奖励函数设计的艺术与陷阱奖励函数是指引智能体前进的“指挥棒”设计不好就会南辕北辙。奖励尺度归一化不同评估指标的量纲和范围差异巨大。Activity可能从4到10QED在0到1SA_Score经过处理后可能是-10到0。直接加权求和会导致某个指标主导。我们必须对每个指标进行归一化例如使用Min-Max归一化到[0,1]区间或者使用Sigmoid函数进行变换确保它们在数值上具有可比性。稀疏奖励与稠密奖励如果我们只在生成过程的最后给出一个基于最终分子综合评分的奖励这就是稀疏奖励智能体很难学习。为此我们引入了中间奖励。例如每次编辑后如果分子的某个性质如LogP向目标范围靠近了一点就给予一个小的正向奖励。这相当于给智能体提供了更密集的反馈信号。约束条件的处理像“分子量必须在200到500之间”这种硬约束不适合直接作为负奖励的一部分因为一旦违反负奖励很大智能体可能直接“摆烂”。更好的做法是将其作为动作过滤器。在智能体选择动作时提前预判该动作是否会导致分子量超出范围如果会则直接将该动作的概率置零。对于软约束如“希望LogP在2到3之间”可以用一个基于距离的惩罚函数离目标越远惩罚越大将其融入奖励。踩坑记录初期我们曾尝试用Activity预测值直接作为主要奖励。结果智能体很快学会了“作弊”——生成一些结构极其复杂、含有特殊金属原子或奇怪环系的分子这些分子在QSAR模型上能拿到高分可能是过拟合或模型盲区但实际合成根本不可能类药性也极差。这就是典型的“奖励黑客”。后来我们引入了SA_Score和QED作为强约束并大幅提高了对合成难度和类药性的权重才将优化方向拉回到合理的化学空间。3.3 策略优化算法选择PPO与MCTS的结合在推理时这个场景下我们与环境分子编辑和评估的交互成本相对较高因为每次调用QSAR模型计算活性都需要时间。因此样本效率很重要。我们选择了近端策略优化算法作为基础。PPO通过限制每次策略更新的幅度保证了训练的稳定性非常适合这种连续动作空间实际是离散动作但概率分布连续的问题。我们会在收集一定数量的分子编辑轨迹从种子分子开始经过多步编辑得到最终分子和一系列奖励后用这些数据对策略网络进行多轮小批量更新。然而仅靠PPO在广阔分子空间中进行探索有时仍显不足。因此我们借鉴AlphaGo的思路引入了蒙特卡洛树搜索MCTS作为推理时的规划工具。具体来说在每一步编辑时智能体不仅依靠当前策略网络还会进行有限深度的MCTS模拟从当前分子状态出发模拟未来多步可能的编辑路径用快速但近似的奖励评估函数比如只计算QED和简单的惩罚项来评估叶子节点然后将搜索得到的价值信息回传辅助当前步做出更优的决策。MCTS和PPO策略网络协同工作前者负责“深谋远虑”的规划后者负责“直觉性”的动作选择并将MCTS探索到的有价值轨迹作为高质量样本用于PPO网络的后续更新。4. 完整工作流程与实操步骤4.1 阶段一环境搭建与预训练模型准备化学信息学基础环境安装RDKit和PyTorch。RDKit用于分子的解析、可视化、合法性检查和简单性质计算如LogP, TPSA。这是整个项目的基石。获取或训练基础生成模型我们使用了在ZINC数据库上预训练好的Chemformer模型作为起点。你也可以使用Hugging Face上的MolGPT或自己用SMILES数据训练一个简单的GRU/Transformer模型。关键是要确保这个生成模型能产出大量语法正确、结构多样的分子。我们准备了100万个从该模型随机生成的分子作为初始种子池。构建预测模型奖励评估器活性预测模型使用公开数据集如ChEMBL中某个靶点的数据训练一个图卷积网络GCN或消息传递神经网络MPNN作为QSAR模型。我们选择了Attentive FP它在多个基准测试上表现良好。将训练好的模型封装成一个API输入分子SMILES输出预测的pIC50。类药性与可合成性直接调用RDKit的CalcQED函数和SA_Score计算函数后者来自rdkit.Chem.rdMolDescriptors需要导入sascorer模块。定义强化学习环境使用OpenAI Gym的风格封装一个自定义环境。reset()函数从种子池随机采样一个分子作为初始状态。step(action)函数执行编辑动作调用合法性校验然后计算新状态分子和即时奖励包含中间奖励并返回done信号达到最大编辑步数或分子已足够优化。4.2 阶段二训练分子编辑智能体初始化策略网络我们采用一个GNN编码分子状态后接一个MLP输出动作概率。GNN将分子图编码为一个全局向量同时保留每个节点的特征向量以供后续选择具体原子进行编辑时使用。收集交互数据# 伪代码示意核心循环 for episode in range(total_episodes): state env.reset() # 得到一个种子分子 trajectory [] for step in range(max_steps): # 结合策略网络和MCTS选择动作 action, log_prob agent.select_action(state, use_mctsTrue) next_state, reward, done, _ env.step(action) trajectory.append((state, action, log_prob, reward, next_state, done)) state next_state if done: break # 计算优势函数等存入经验池 processed_trajectory compute_advantages(trajectory) replay_buffer.add(processed_trajectory)策略网络更新每隔一定轮次从经验池采样一批轨迹数据执行PPO的更新步骤最大化期望奖励同时约束新旧策略的KL散度不要过大。MCTS模块集成在agent.select_action函数中如果启用MCTS则从当前状态构建搜索树。每次模拟包含选择基于UCT公式、扩展添加新子节点、评估用快速奖励函数评估叶子节点、回传更新路径上节点的访问次数和价值四个步骤。最终根据根节点下各动作的访问次数分布采样或选择访问次数最多的动作作为最终输出。4.3 阶段三推理时优化与分子生成当策略网络训练到一定程度后就可以用于实际的分子优化生成了。启动优化流程用户提供或从生成器随机采样一个起点分子设定优化目标如最大化活性同时满足QED0.6, SA_Score -4.5。迭代编辑加载训练好的智能体在给定起点分子上运行优化流程。智能体会根据已学到的策略进行一系列分子编辑。结果收集与后处理记录优化路径上的所有中间分子及其评分。优化结束后通常能得到一个帕累托前沿——一组在多个目标间取得最佳权衡的分子。我们使用pymoo库进行多目标优化的结果分析与可视化。5. 效果评估体系与结果分析评估不能只看生成分子的“分数”必须多维度、多角度进行。5.1 评估指标设计我们设计了四个层次的评估指标评估维度具体指标说明与计算方法优化性能综合奖励提升值比较优化前后分子综合奖励分数的绝对提升和相对提升率。帕累托前沿质量使用超体积指标衡量优化得到的分子集在多个目标属性空间中所占的体积体积越大越好。分子质量有效性优化后分子能通过RDKit解析的比例应接近100%。独特性优化生成的分子与训练集及种子池中分子的相似度基于分子指纹越高越好但需避免过度新颖导致不合成。新颖性优化生成的分子不在训练集中的比例。化学性质主要目标属性如预测活性的实际分布、QED、SA_Score的均值与方差。次要约束符合率如分子量、LogP、氢键供体/受体数等落在预设范围内的分子比例。实用性多样性优化生成的分子集内部的结构多样性通过计算分子间Tanimoto相似度的平均值来评估越低越多样。与已知活性分子的相似性将生成的Top分子与已知的活性分子进行相似性比对评估其是否落在有希望的化学空间。5.2 对比实验设置为了证明“推理时强化学习”策略的有效性我们设置了以下基线模型进行对比Baseline 1: 随机搜索在相同的编辑动作空间内进行完全随机编辑。Baseline 2: 贪心搜索每一步只选择能带来最大即时奖励的动作。Baseline 3: 传统训练时RL使用REINFORCE算法直接训练一个从头开始的分子生成模型。Baseline 4: 基于遗传算法的优化将分子表示为SMILES字符串进行交叉、变异操作。5.3 核心结果与发现在实际针对某个激酶靶点的优化任务中我们得到了如下关键结论优化效率显著提升我们的方法推理时RL MCTS在达到相同综合奖励水平时所需的分子评估次数即调用计算昂贵的QSAR模型的次数比随机搜索少约95%比贪心搜索少约70%比传统训练时RL少约50%。这说明我们的方法能更智能地探索化学空间。成功实现多目标权衡通过调整奖励函数权重我们能够引导生成不同侧重点的分子簇。例如一组分子更偏向高活性平均预测pIC50提高1.2另一组在保持活性适度提升0.8的同时合成可及性评分SA_Score大幅改善。而传统方法很难在一次运行中得到如此清晰的帕累托前沿。生成分子质量可靠超过99.5%的优化输出分子是化学有效的。独特性高且与已知活性分子具有合理的相似性平均Tanimoto系数在0.4-0.6说明方法没有产生“天马行空”的无意义结构而是进行了合理的化学修饰。灵活性得到验证在项目中期合作方新增了一个“避免hERG通道抑制”的毒性要求。我们仅在奖励函数中加入了一个由预训练hERG毒性预测模型提供的惩罚项并调整了权重就在后续的优化轮次中成功得到了既满足原有活性、类药性要求又显著降低hERG风险的新分子。整个过程无需重新训练任何生成模型响应速度极快。6. 常见问题、排查技巧与优化方向6.1 实操中遇到的典型问题智能体陷入局部最优反复生成相似分子现象优化过程早期奖励上升很快但很快停滞生成的分子结构高度相似。排查与解决检查奖励函数是否过早地给予了某个单一指标过高的权重尝试在奖励中加入鼓励多样性的项例如对与已生成分子库相似度过高的新分子给予轻微惩罚。增加探索提高PPO算法中熵正则项的系数鼓励策略网络输出更平均的动作概率分布。或者在MCTS中调高探索常数C_puct。扩大种子池使用更多样化的种子分子作为优化起点。编辑动作经常产生无效分子现象合法性校验频繁失败导致大量负奖励学习效率低下。排查与解决强化动作空间过滤在动作选择层就加入更严格的化学规则过滤彻底排除明显非法的动作选项而不是等执行后再惩罚。简化动作空间初期可以使用更保守的动作集例如只允许“添加原子单键”和“改变原子类型”待策略稳定后再引入更复杂的动作如“添加双键”。设计更智能的合法性恢复机制当动作导致非法分子时不是简单给予负奖励并结束而是尝试一个“回退”动作或者触发一个子程序来修复该非法结构如调整氢原子数量并将此作为一个特殊的学习经验。奖励数值不稳定或出现NaN现象训练过程中回报曲线剧烈波动甚至出现数值错误。排查与解决数值裁剪对奖励值进行裁剪例如限制在[-10, 10]区间内。奖励标准化不仅对最终奖励标准化对每一步的中间奖励也进行移动平均标准化使其均值为0标准差为1。检查预测模型确保QSAR等预测模型在输入极端或训练集外分子时不会输出异常值或NaN。可以在调用前加入分子描述符的合理性检查。6.2 性能优化方向分布式并行化分子评估尤其是QSAR预测是计算瓶颈。可以将奖励计算任务分发到多个CPU核心或GPU上并行执行显著加快数据收集速度。引入课程学习让智能体从简单的优化任务开始学起。例如先只优化QED这一个简单目标然后逐步加入活性、可合成性等更复杂的目标。或者先从修改小分子开始再过渡到编辑大分子。集成更先进的规划算法可以尝试将MCTS替换为基于模型的规划算法如果能为分子编辑的动态过程学习一个快速的、近似的前向模型就能在想象中进行更高效的规划。与实验反馈闭环最终的理想状态是将虚拟生成-优化-评估的循环与真实的生化实验测试相结合。将优化出的分子进行合成与活性测试将实验数据反馈回来进一步微调奖励函数中的预测模型实现“AI设计-实验验证-模型迭代”的完整闭环。这将是下一代AI驱动药物发现的核心范式。这套基于推理时强化学习的分子生成优化策略经过我们多个项目的实战检验证明其在平衡“探索能力”与“优化效率”、兼顾“多目标柔性”与“化学合理性”方面确实提供了一条行之有效的路径。它可能不是银弹但绝对是药物化学家和计算化学家工具箱里一件越来越称手的利器。