MolMem: Memory-Augmented Agentic Reinforcement Learning for Sample-Efficient Molecular Optimization

Idea

MolMem 在多轮 molecular optimization 中同时利用静态分子 exemplars 和训练过程积累的 editing skills，并让 retrieval 与 agentic reinforcement learning 共同工作。模型遇到属性改进停滞时，可以借鉴邻近分子或过去成功的结构编辑，再继续提出满足 lead similarity 约束的新分子。

Research Problem

LLM 分子优化在有限 oracle budget 下容易反复尝试无效编辑，单次生成又难以利用先前候选的反馈。仅检索已知高分分子可能变成复制 exemplar，纯 RL 则需要反复支付探索成本，且缺少把有效局部变换迁移到其他结构的显式机制。

Insight

具体分子和抽象编辑规律提供不同粒度的先验：前者给出可行邻域，后者说明在何处做何种改变及预期作用。将检索限定在 plateau 时触发，既能针对探索困难补充信息，也减少常态上下文开销；把高回报编辑转成可检索技能，使每轮训练的经验成为后续 rollout 的输入。

Method

1. 将问题写成多轮 MDP：状态包括优化目标、lead molecule、历史 SMILES 和奖励，动作是新候选 SMILES，oracle 返回属性与相似性相关反馈。先用满足相似性约束的离线编辑对进行 SFT warm start。
2. Static Exemplar Memory 从 ChEMBL 建立带预计算 QED、LogP 的分子库，用 ECFP4 Morgan fingerprints 与 FAISS 检索；停滞时先找当前分子的近邻，再按与原始 lead 的相似度筛选、按目标属性排序，并惩罚直接复制检索分子。
3. Evolving Skill Memory 从高奖励增量的编辑对出发，利用 maximum common substructure 提取增删片段、scaffold、functional groups 与 descriptor 变化，概括成 Action–What–Where–Effect 技能卡；以编辑前 fingerprint 和 functional-group tags 建立任务相关索引。
4. 停滞时按 fingerprint Tanimoto 和 functional-group Jaccard 检索、排序技能；两类记忆同时触发时随机选择一种控制上下文预算。PPO 配合逐步奖励和 GAE 优化多轮策略，成功轨迹继续更新技能库。