Paper deep dive
PGFS++: Molecular Property Improvement under Synthesis and Diversity Constraints
Boqiao Zhang, Godbless James, Sai Krishna Gottipati, Andrew Fitzgibbon
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 92%
Last extracted: 8/22/2026, 1:39:28 AM
Summary
The paper introduces PGFS++, a synthesis-aware reinforcement learning framework for molecular property improvement that addresses the 'magnet output' diversity collapse problem found in its predecessor, PGFS+. While PGFS+ improved property scores by directly scoring reactant candidates, it often mapped diverse inputs to a single high-reward molecule. PGFS++ resolves this by incorporating an input-output similarity bonus into the reward function, ensuring improved properties while maintaining structural diversity and explicit synthesis routes.
Entities (8)
Relation Signals (8)
PGFS++ → extends → PGFS++
confidence 95% · PGFS++ ... extends PGFS. PGFS++ can generate molecules with desired properties and greatly improved diversity.
Input-Output Similarity Bonus → mitigates → Magnet Output
confidence 95% · This design mitigates magnet-output collapse and encourages input-specific molecular improvement
PGFS++ → suffersfrom → Magnet Output
confidence 95% · PGFS+ can introduce a new failure mode: many distinct input molecules are mapped to the same high-reward output molecule. We refer to this diversity collapse as the magnet output problem
PGFS++ → uses → Input-Output Similarity Bonus
confidence 95% · PGFS++ introduces input–output similarity shaping into the reward function
PGFS++ → optimizesfor → SEH
confidence 90% · For PGFS++, we used either r property = ∆QED or r property = ∆SEH
PGFS++ → optimizesfor → QED
confidence 90% · For PGFS++, we used either r property = ∆QED or r property = ∆SEH
PGFS++ → usesalgorithm → PPO
confidence 90% · we replace TD3... with PPO... in PGFS+ and PGFS++
PGFS++ → usesalgorithm → TD3
confidence 90% · the TD3 policy is optimized over a continuous action... in PGFS
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Improving molecular properties, such as drug-likeness or binding affinity, is a recurring task in early-stage drug discovery. However, molecules optimized in an unconstrained chemical space have limited practical value if they cannot be synthesized. Policy Gradient for Forward Synthesis (PGFS) is a synthesis-aware reinforcement learning method for molecular improvement, but its use of reactant embedding prediction makes reactant selection indirect, which, as we show, limits learning effectiveness. We first develop PGFS+, in which reaction templates and second reactants are represented by trainable embedding lookup tables. Combined with a more effective scoring function and RL algorithm, PGFS+ significantly improves the desired property. However, it exposes a reward-hacking failure mode: a powerful reactant search can map diverse input molecules to the same high-reward magnet molecule, improving the reward while collapsing the output diversity. We therefore introduce PGFS++, a synthesis-aware reinforcement learning framework for input-specific molecular improvement. Given an input molecule, PGFS++ treats it as the start of a forward-synthesis trajectory, applies learned reaction templates with compatible in-stock building blocks, and produces a molecule with improved target properties, an explicit synthesis route, and structural similarity to the input. Experiments on molecular improvement tasks show that PGFS++ improves target properties while preserving high output diversity.
Tags
Links
- Source: https://arxiv.org/abs/2608.19121v1
- Canonical: https://arxiv.org/abs/2608.19121v1
Trouble viewing inline? Open PDF directly →
Full Text
46,603 characters extracted from source content.
Expand or collapse full text
PGFS++: Molecular Property Improvement under Synthesis and Diversity Constraints Boqiao Zhang 1 , Godbless Tamaraebi James 1 , Sai Krishna Gottipati, Andrew Fitzgibbon 1,2,* 1 University of Cambridge, United Kingdom 2 Graphcore, United Kingdom * Corresponding author: awf@graphcore.ai Abstract Improving molecular properties, such as drug-likeness or binding affinity, is a recurring task in early-stage drug discovery. However, molecules optimized in an unconstrained chemical space have limited practical value if they cannot be synthesized. Policy Gradient for Forward Synthesis (PGFS) is a synthesis-aware reinforcement learning method for molecular improvement, but its use of reactant embedding prediction makes reactant selection indirect, which, as we show, limits learning effectiveness. We first develop PGFS+, in which reaction templates and second reactants are represented by trainable embedding lookup tables. Combined with a more effective scoring function and RL algorithm, PGFS+ significantly improves the desired property. However, it exposes a reward-hacking failure mode: a powerful reactant search can map diverse input molecules to the same high-reward “magnet molecule”, improving the reward while collapsing the output diversity. We therefore introduce PGFS++, a synthesis-aware reinforcement learning framework for input-specific molecular improvement. Given an input molecule, PGFS++ treats it as the start of a forward-synthesis trajectory, applies learned reaction templates with compatible in-stock building blocks, and produces a molecule with improved target properties, an explicit synthesis route, and structural similarity to the input. Experiments on molecular improvement tasks show that PGFS++ improves target properties while preserving high output diversity. Keywords: Molecular Optimization; Drug Discovery; Reinforcement Learning; Machine Learning; Forward Synthesis; Synthesis Constraints 1 Introduction Over the past decades, AI for chemistry and drug discovery has been substantially developed. Generative models can output chemically valid de novo molecules in a manner biased towards desired properties such as QED (Quantitative Estimate of Drug-likeness), but most of them do not provide explicit synthesis routes for molecules generated by the model. Reinforcement Learning (RL) is another widely used approach in drug discovery, where an agent interacts with an environment and learns to adapt its policy to maximise a reward signal, such as QED or activity against biological targets (Zhou et al., 2019; Simm et al., 2020; You et al., 2019). Still, only a few methods generate both attractive molecules and clear synthesis routes. PGFS (Policy Gradient for Forward Synthesis) (Gottipati et al., 2020) is a RL-based method that constrains the search to synthetically-accessible routes and therefore can output feasible synthesis pathways for its output molecules. RL-based methods with synthesis constraints typically require a predefined set of reaction templates from which the model learns its policy. These reactions are commonly divided into two 1 arXiv:2608.19121v1 [cs.LG] 19 Aug 2026 A. Illustrative Routes of Generating a Magnet Output MoleculeB. Model Performance (QED-Diversity Trade-off) PGFS+: Higher mean QED but collapsed diversity PGFS++: Higher mean QED and high diversity Figure 1: Reward hacking phenomenon in PGFS+. A, PGFS+ maps hundreds of input molecules to the same output molecule, causing high QED but low output diversity. This is referred to as the magnet output phenomenon. Highlighted atoms indicate the portions of the output molecule that are traced back to the corresponding input molecules. B, Trade-off between mean QED and mean diversity for PGFS, its variants, and PGFS+. Because of the magnet output, PGFS+ achieves high mean QED but has output diversity close to zero. types: uni-molecular and bi-molecular. For bi-molecular templates, one must additionally select a second reactant, denoted asR 2 , to react with the first reactantR 1 under a chosen templateT. However, the space of possibleR 2 candidates can be very large, and the action space becomes even larger when the model jointly selects bothTandR 2 . The original PGFS framework overcame this by using a continuous action space for second-reactant selection: a neuralR 2 predictor is trained to predict the embedding of the second reactant, and a nearest-neighbour search is then used to retrieve an actualR 2 from the available reactant set. However, this design has two limitations. First, the TD3 policy is optimized over a continuous action (predicted reactant embedding), whereas the actual reaction is executed using a discrete reactant retrieved from the available molecule set, creating a mismatch between the optimized action and the executedR 2 . Second, the predicted action is used directly as the query in the reactant-embedding space, which ties the policy representation to the fixed fingerprint space and may limit efficient comparison among feasible reactants. To address these limitations, we first extend PGFS to PGFS+. Instead of predicting a continuous R 2 embedding and retrieving neighbours, PGFS+ projects molecules and templates into a learned reaction space and directly scores all template-compatibleR 2 candidates (Fig. 2). This modification allows the policy to compare valid second reactants directly and removes the mismatch between continuous prediction and discrete execution. Additionally, we replace the deterministic k-N selection step with sampling at training time, which improves learning in the presence of multiple near-equivalent R 2 candidates. However, we observe that PGFS+ can introduce a new failure mode: many distinct input molecules are mapped to the same high-reward output molecule. We refer to this diversity collapse as the magnet output problem (Fig. 1). Although such outputs may achieve strong average molecular properties, they substantially reduce output diversity and weaken the ability of the model to perform input-specific molecular improvement. We therefore further develop PGFS++, which retains the direct global R 2 selection mechanism while preserving diversity. This paper makes two main contributions: 1. We present PGFS++, a synthesis-constrained reinforcement learning framework for molecular improvement that extends PGFS. PGFS++ can generate molecules with desired properties 2 Figure 2: Workflow comparison between PGFS, PGFS+, and PGFS++. At stept, all methods use anf-network conditioned on the current first reactantR ′ 1 (t) to predict the reaction templateT. PGFS uses the fingerprint representation F (R ′ 1 (t)) directly, whereas PGFS+ and PGFS++ first transform the R 1 representation through an additional MLP. For second-reactant selection, PGFS representsR 2 candidates using fixed fingerprints and retrieves candidates byk-nearest neighbours. In contrast, PGFS+ and PGFS++ represent template-compatibleR 2 candidates with a trainable embedding lookup table and score them globally using an attention-like mechanism with queryq=concat(φ 1 (R ′ 1 ),φ T (T)). OnceR ′ 2 (t) is selected, the reaction runner applies the chosen reaction and returns the product, which becomes the next state R ′ 1 (t + 1). 3 and greatly improved diversity. 2. We present PGFS+, and identified the magnet output problem, which reflects a broad fail- ure mode in synthesis-aware molecular improvement tasks. We then propose input–output similarity bonus as an effective way to resolve it in PGFS++. This design mitigates magnet- output collapse and encourages input-specific molecular improvement while maintaining output diversity. 2 Related Work 2.1 Genetic Algorithms Genetic Algorithms (GA) are a broad class of algorithms that have been used for the generation of molecules for decades (Tripp and Hernández-Lobato, 2023). They usually represent molecules as graphs (Brown et al., 2004; Jensen, 2019) or strings (Nigam et al., 2020), and recursively generate molecules that are more promising for the target molecular property. Existing GA-based de novo generation methods typically address synthetic feasibility indirectly, by incorporating the synthetic-accessibility (SA) score (Ertl and Schuffenhauer, 2009) or its variants into the optimization objective. Consequently, after a molecule with desirable predicted properties is generated, an additional retrosynthesis model or expert assessment is required to determine whether the structure can be practically synthesized. 2.2 Generative Models Another widely used class of methods for molecule generation is generative modeling. A wide range of generative approaches for drug discovery has been explored, such as variational autoencoders (VAEs) (Gómez-Bombarelli et al., 2018) and generative adversarial networks (GANs) (Cao and Kipf, 2022). Despite the advances, they exhibit similar challenges in synthetic feasibility as genetic algorithms. 2.3 Synthesis-aware molecule generation The idea of synthesis-aware molecule generation starts with Vinkers et al. (2003), who introduced SYNOPSIS. This method generates molecules from an initial dataset of available compounds by applying chemical modifications to functional groups. However, as an early rule-based approach, its optimization framework is less flexible and generally less effective than those of more recent methods. 2.4 Generative Models with Synthesis Constraints Generative models with synthesis constraints can generate molecules with desirable properties, while also providing explicit reaction routes for their synthesis. Instead of producing molecular structures first and assessing synthesizability post hoc, these methods incorporate chemical feasibility directly into the generation process. For example, Cretu et al. (2025) combines GFlowNet with a reaction-based Markov decision process to generate drug-like molecules through predefined chemical reactions and building blocks. While the method achieves competitive performance for both molecular property and output diversity, it depends heavily on reversible chemical reaction templates, which introduce additional complexity and limitations on template design. Moreover, SynFlowNet is primarily designed for de novo molecule generation, whereas PGFS++ targets input-specific molecular improvement from a given starting molecule. 4 2.5 RL-based Models with Synthesis Constraints Reinforcement learning has also been explored as a way to embed synthetic feasibility directly into molecular generation. For example, the PGFS (Policy Gradient for Forward Synthesis) framework (Gottipati et al., 2020) formulates drug discovery as a sequential forward-synthesis problem, where an agent navigates the synthetically accessible chemical space by applying valid reactions to commercially available building blocks at each step of a virtual multi-step synthesis process. Extensions of PGFS, including PGFS-TAC (Towered Actor Critic) (Gottipati et al., 2021) and PGFS-MB (Max-Bellman equation) (Gottipati et al., 2023), provide additional flexibility within this framework. However, they retain similar limitations to PGFS. 3 Methods 3.1 Overview Following PGFS, we represent molecular improvement as a sequence of unimolecular or bimolecular reactions applied to an initial molecule. Each synthetic step is decomposed into two sub-actions: selecting a reaction template and, for bi-molecular reactions, selecting a compatible second reactant from the available building blocks. The reaction template is chosen from templates available in laboratory or industrial synthesis, and the second reactant is selected from a set of available building blocks. The reinforcement learning agent first learns to select an appropriate templateTgiven the current reactantR 1 . Conditioned onR 1 andT, it then selects a second reactantR 2 to maximize the reward associated with the properties of the product molecule. This hierarchical action decomposition ensures that the generated molecules are associated with explicit synthetic routes and improves training efficiency by reducing the complexity of each decision. A key challenge is that there may be millions of possibleR 2 candidates for eachTandR 1 . PGFS addresses this by using a neural network to predict anR 2 embedding, followed byk-nearest-neighbor retrieval from the building-block set. In contrast, PGFS+ and PGFS++ directly scores feasibleR 2 candidates after applying a template-specific compatibility mask. We observed that the largeR 2 action space can be efficiently reduced to a tractable candidate set using RDKit substructure matching. We pre-compute aT-R 2 compatibility mask for each bimolecular template. For our dataset, which contains approximately 118,000 molecules and 87 bimolecular templates, pre-computing the masks for all bimolecular templates takes approximately 10 minutes on a single CPU process. This computation is performed only once and can be reused for subsequent training runs. After masking, each bimolecular template has a median of approximately 1000 compatibleR 2 candidates and a maximum of 55,000, making direct attention-like scoring feasible. As shown in Fig. 2, PGFS+ and PGFS++ concatenate learned embeddingsφ 1 (R 1 ) andφ T (T) of R 1 and T to form a query vector, q = concat(φ 1 (R 1 ),φ T (T )), and use this query to score the maskedR 2 embeddings as keys. To support this direct scoring formulation, we replace TD3, the continuous-action reinforcement learning algorithm used in PGFS, with PPO, which naturally supports a categorical policy over the indexedR 2 action space in PGFS+ and PGFS++. However, as shown in Fig. 1, when PGFS+ is trained only with a target-property reward, such as ∆QED, the policy can learn to map diverse input molecules to the same high-reward output molecule. 5 Figure 3: Performance comparison on the QED improving task among ours (PGFS++) and random search, PGFS, PGFS TAC, PGFS MB, REINVENT4 (mol2mol), SynFlowNet. The comparison with SynFlowNet used the reversible reaction templates provided in the SynFlowNet repository to avoid potential issues arising from irreversible templates. Comparison with other baselines uses a general template set derived from Button et al. (2019), containing both reversible and irreversible templates. Specifically, among 12,689 randomly selected test episodes, 12,415 (97.8%) produced the same output molecule. Empirically, we found that the collapse occurs often because the policy tends to select similar templates and second reactants across different starting molecules, producing similar or even identical high-reward outputs. In other words, the model improves different inputs by pushing them toward a small number of high-scoring magnet molecules, rather than producing input-specific improved molecules. Although we identify this behaviour in PGFS+, it reflects a broader failure mode in synthesis-aware molecular improvement tasks: when the reward depends mainly on terminal molecular properties and does not constrain input–output relatedness, the policy can be incentivised to collapse diverse inputs toward a small set of high-reward products. To mitigate this problem, PGFS++ introduces input–output similarity shaping into the reward function, encouraging the output molecule of each episode to remain structurally close to its corresponding input molecule. 3.2 Input-output similarity To discourage magnet-output collapse, PGFS++ makes similarity to the starting molecule an explicit part of the episode reward. We design the reward function incorporating input–output similarity bonus: r × episode = r property (m out ) 1 + c sim(m in ,m out ) (1) wherer property is the property-based reward, such as ∆QED and ∆SEH,c >0 control the bonus strength, andsim(m in ,m out ) denotes the Tanimoto similarity (Bajusz et al., 2015) between input 6 and output molecules per episode. When the input–output similarity falls below a tolerance thresholdτ, we cap the episode reward: r episode = ( r × episode ,sim(m in ,m out )≥ τ, minr × episode ,κ, sim(m in ,m out ) < τ, (2) where κ is the reward cap. Diversity drop as inward edit bias. We compare input and output diversity in the feature space induced by the Tanimoto kernel (Ralaivola et al., 2005; Tripp et al., 2023). Letmandndenote Morgan fingerprint vectors. For non-negative fingerprints, Tanimoto similarity is a positive-definite kernel, so there exists a feature mapψsuch that ⟨ψ(m),ψ(n)⟩ = sim(m,n). Since sim(m,m) = 1, the induced squared distance is ∥ψ(m)− ψ(n)∥ 2 = sim(m,m) + sim(n,n)− 2sim(m,n) = 2 1− sim(m,n) . (3) For an episode with input molecule m in and output molecule m out = G(m in ), define X = ψ(m in ), Y = ψ(m out ), A = Y − X, whereAis the edit displacement induced by the policy in the feature space. For an independent copy of the episode, primes denote m ′ in , m ′ out , X ′ , and Y ′ . We define the input and output diversity measures as D in =E 1− sim(m in ,m ′ in ) = 1 2 E∥X − X ′ ∥ 2 =E∥X − μ X ∥ 2 = Var(X), D out =E 1− sim(m out ,m ′ out ) = 1 2 E∥Y − Y ′ ∥ 2 =E∥Y − μ Y ∥ 2 = Var(Y ), (4) where μ · =E[·]. Since Y = X + A and μ Y = μ X + μ A , Y − μ Y = (X − μ X ) + (A− μ A ). Substituting this expression into Equation (4) gives D out = D in + V A + 2Γ, V A =E∥A− μ A ∥ 2 = Var(A)≥ 0, Γ =E[⟨X − μ X , A− μ A ⟩]. (5) Equation(5)separates the effect of editing into the non-negative edit varianceV A and the input–edit covariance 2Γ. Therefore, D out < D in ⇐⇒Γ <− V A 2 . We refer to this sufficiently negative covariance as an inward edit bias. 7 Complete collapse to a magnet molecule m ⋆ gives Y = ψ(m ⋆ ). In this case, A = ψ(m ⋆ )− X, A− μ A =−(X − μ X ), and hence V A = D in ,Γ =−D in . Substitution into Equation (5) recovers D out = 0. Bounding diversity drop with input–output similarity. Since sim(m,m) = 1, the feature-space embeddings satisfy ∥X∥ =∥Y∥ = 1. Therefore, D in = 1−∥μ X ∥ 2 , D out = 1−∥μ Y ∥ 2 . Define the average input–output similarity ̄s =E sim(m in ,m out ) =E⟨X,Y⟩. Let a =∥μ X ∥ = p 1− D in . We first bound ∥μ Y ∥ in terms of a and ̄s. Set u = μ Y /∥μ Y ∥. For any λ≥ 0, ∥μ Y ∥ =E⟨u,Y⟩ =E⟨u + λX,Y⟩− λE⟨X,Y⟩ ≤E∥u + λX∥− λ ̄s ≤ p 1 + λ 2 + 2λ⟨u,μ X ⟩− λ ̄s ≤ p 1 + λ 2 + 2λa− λ ̄s. Minimizing the right-hand side over λ≥ 0 gives ∥μ Y ∥≤ h ̄s (a), where h ̄s (a) = ( 1, ̄s≤ a, a ̄s + √ 1− a 2 √ 1− ̄s 2 , ̄s > a. Hence, the diversity drop ∆ div = D in − D out + = ∥μ Y ∥ 2 −∥μ X ∥ 2 + where [v] + = max(v, 0), satisfies ∆ div ≤ h ̄s p 1− D in 2 − (1− D in ).(6) When ̄s > √ 1− D in , this becomes ∆ div ≤ ̄s p 1− D in + p 1− ̄s 2 p D in 2 − (1− D in ).(7) When the input pool is highly diverse (D in ≈ 1), we have the approximation ∆ div ≲ 1− ̄s 2 . Therefore, the input–output similarity bonus in PGFS++ directly encourages the quantity that tightens this bound: by favouring policies with larger ̄s, it tightens the upper bound on diversity drop. 8 Figure 4: Performance comparison on the SEH improving task among ours (PGFS++) and random search, PGFS, PGFS TAC, PGFS MB, REINVENT4 (mol2mol), SynFlowNet. The comparison with SynFlowNet used the reversible reaction templates provided in the SynFlowNet repository to avoid potential issues arising from irreversible templates. Comparison with other baselines uses a general template set derived from Button et al. (2019), containing both reversible and irreversible templates. 4 Experiments We evaluated the proposed framework on a dataset containing approximately 118,000 molecules from Enamine Building Block Global Stock (Enamine, 2026) and 102 reaction templates derived from Button et al. (2019), including 15 unimolecular and 87 bimolecular templates. The test set consisted of 2,000 randomly selectedR 1 reactants from the dataset. At test time, each model was used to improve every test R 1 , and metrics were recorded for each episode. We compared PGFS++ against PGFS, PGFS-TAC, PGFS-MB, SynFlowNet, and the REIN- VENT4 mol2mol model (Loeffler et al., 2024). For REINVENT4, we fine-tuned the medium prior model using reinforcement learning for 500 steps, with 2,048 oracle-scored molecules per step. This results in approximately one million oracle calls in total, which is consistent with the number of oracle calls used in PGFS++ training. The remaining baselines were trained with their default settings. For PGFS++, we used eitherr property = ∆QEDorr property = ∆SEHas the property reward. For the hyperparameters in Equation 1 and 2, we setτ= 0.25,c= 0.5,κ QED = 0.1, andκ SEH = 0.35, reflecting the different numerical scales of the two property rewards. All other hyperparameters were kept identical to those used in PGFS. To ensure a fair comparison across synthesis-aware methods, we fixed the reaction budget to at most four reaction steps per episode. The experiments focused on two optimization objectives: QED (on a [0,1] scale) and SEH (scaled by a factor of 1 8 ). For each metric, we report a violin plot together with the median. For PGFS++, results are averaged over three random seeds. We additionally report the mean±standard deviation 9 Figure 5: Ablation of input–output similarity bonus. We compare PGFS+ without an input–output similarity bonus against PGFS++ with additive and multiplicative similarity bonuses. Bold entries indicate the best value for each metric, andunderlinedentries indicate the second-best value. Without input–output similarity bonus, PGFS+ suffers from the magnet output problem, as reflected by high mean QED, but very low output diversity, input–output similarity, and RAS. of the three seed-wise medians to quantify variability across seeds. 4.1 Metrics In addition to the two optimization objectives, we evaluated generated molecules using diversity, synthetic accessibility (SA), input-output similarity, and Route-Aware Synthesizability (RAS). Diversity measures structural variation within the output setO. It is defined as one minus the average pairwise Tanimoto similarity between Morgan fingerprints: D(O) = 1− 1 |P| X m i ,m j ∈P T (FP(m i ), FP(m j )), whereP=(m i ,m j ) :m i ,m j ∈ O,i < jis the set of unordered molecule pairs,Tdenotes Tanimoto similarity and FP denotes the Morgan fingerprint representation. Synthetic accessibility was assessed using the Ertl–Schuffenhauer synthetic accessibility (SA) score implemented in RDKit (Landrum, 2026). Although the SA score does not fully capture route-level synthetic feasibility, we used it as a simple molecule-level metric to enable comparison with non-synthesis-aware methods. For visual consistency with the other metrics, we report 1− SAin plots, so that higher values indicate better synthetic accessibility. While the SA score is widely used to estimate molecule-level synthetic accessibility, it does not fully capture the route-level information of synthesis-aware methods. RAS below addresses this limitation by incorporating information from the generation routes. Input-output similarity measures the structural similarity between the input moleculem e in and the output moleculem e out for each episodee, averaged across episodes. For each episodee, we compute sim in,out (e) = T (FP(m e in ), FP(m e out )). This metric is important for molecule improving tasks, where the optimized molecule should retain structural similarity to the starting molecule. Route-aware synthesizability (RAS) is a route-level synthetic accessibility score in [0,1], where higher values indicate more accessible synthetic routes. RAS was computed only for methods that produce explicit reaction trajectories. For an episodee, let the selected synthetic route contain reaction steps n = 1,...,N e . 10 For each step n, we obtain two raw accessibility signals and map them to normalized factors in [0, 1]: 1. Reaction feasibility (template score). Letp raw n denote the ASKCOS fast-filter plausibility (Tu et al., 2025) of the reaction at stepn, which lies in [0,1], with larger values indicating greater feasibility. We use it directly: RAS T n = p raw n . 2.Reactant accessibility (R 2 score). For bimolecular steps, lets raw n denote the BR-SAScore (Ertl and Schuffenhauer, 2009) of the selected second reactantR 2 , where smaller values indicate greater accessibility. We map this score to [0,1] by inverting and rescaling it with fixed bounds s min = 1 and s max = 10: RAS R 2 n = clip s max − s raw n s max − s min , 0, 1 . For uni-molecular steps, where no building block is added, we set RAS R 2 n = 1. Each reaction step contributes the product of its two normalized factors. RAS is then computed as the product over all steps on the selected route: RAS(e) = N e Y n=1 RAS T n × RAS R 2 n . Because every factor lies in [0,1], longer routes are penalized through multiplication by additional terms. Although we use BR-SAScore to estimate the accessibility ofR 2 , the RAS metric can incorporate other R 2 accessibility estimates, such as the price of R 2 . 4.2 Experiments for QED and SEH improvement For both the QED and SEH optimization tasks (Figs. 3 and 4), PGFS++ consistently outperformed PGFS and its variants. Compared with REINVENT4, PGFS++ achieved comparable performance. However, unlike PGFS++, REINVENT4 does not produce explicit synthesis routes for the generated molecules. We also compared PGFS++ with SynFlowNet using the reversible reaction templates provided in the SynFlowNet GitHub repository, to avoid potential issues arising from irreversible templates. Under this setting, PGFS++ also achieved competitive performance on both the QED and SEH optimization tasks. 4.3 Ablation Studies We compared the multiplicative input–output similarity bonus (Equation 1) to the additive form: r + episode = r property (m out ) + w sim(m in ,m out ), r × episode = r property (m out ) 1 + c sim(m in ,m out ) . (8) The two forms differ in how strongly they pull the output toward the input. For the additive form, ∂r + episode ∂ sim = w, 11 while for the multiplicative form, ∂r × episode ∂ sim = cr property (m out ). Thus, the multiplicative bonus becomes stronger for high-property outputs, where magnet-molecule collapse is most likely to be rewarded. The multiplicative form also provides a more stable way to control the balance between property optimization and input–output similarity. Because the similarity term scales the property reward, the bonus acts as a relative adjustment to the target-property objective. In contrast, the additive form uses an absolute bonus whose scale must be carefully tuned. In Fig. 5, we compare the two similarity-bonus forms with the hyperparameters in Equation 8 set tow=c= 0.5, while keeping other hyperparameters identical. Both additive and multiplicative similarity bonuses mitigate the magnet molecule problem observed in PGFS+ without similarity bonus. Compared with PGFS+, both variants substantially improve output diversity, input–output similarity, and RAS while maintaining competitive QED. The two bonus forms achieve similar diversity, 1− SA, input–output similarity, and RAS, but the multiplicative form better preserves QED. 5 Conclusion and Discussion In this work, we introduced PGFS++, a synthesis-constrained reinforcement learning framework for molecular improvement. PGFS++ extends PGFS by using a discrete action space, pre-computedT– R 2 masks, attention-like global scoring forR 2 selection, and input–output similarity encouragement. In experiments on QED and SEH improvement tasks, PGFS++ consistently improved over PGFS and its variants and achieved performance competitive with REINVENT4 and SynFlowNet. The framework still has a few limitations. For example, although the input-output similarity bonus can improve similarity while maintaining relatively strong performance on other metrics, the model may be sensitive to the hyperparameter values in Eq. 2. Some directions of future work are worth pursuing. First, PGFS++ currently uses Morgan fingerprints forR 1 representation, whereas alternative fingerprints, such as RLV2 and MACCS, or graph-transformer-based representations may further improve performance. Second, beyond QED and SEH, we plan to evaluate the framework on more practical objectives, such as docking scores. References Dávid Bajusz, Anita Rácz, and Károly Héberger. Why is Tanimoto index an appropriate choice for fingerprint-based similarity calculations? Journal of Cheminformatics, 7(1):20, 2015. doi: 10.1186/s13321-015-0069-3. Neil Brown, Brian McKay, Franco Gilardoni, and Johann Gasteiger. A graph-based genetic algorithm and its application to the multiobjective evolution of median molecules. Journal of Chemical Information and Computer Sciences, 44(3):1079–1087, 2004. doi: 10.1021/ci034290p. A. Button, D. Merk, J. A. Hiss, and G. Schneider. Automated de novo molecular design by hybrid machine intelligence and rule-driven chemical synthesis. Nature Machine Intelligence, 1(7):307–315, 2019. Nicola De Cao and Thomas Kipf. Molgan: An implicit generative model for small molecular graphs, 2022. URL https://arxiv.org/abs/1805.11973. 12 Miruna Cretu, Charles Harris, Ilia Igashov, Arne Schneuing, Marwin Segler, Bruno Correia, Julien Roy, Emmanuel Bengio, and Pietro Liò. Synflownet: Design of diverse and novel molecules with synthesis constraints. In International Conference on Learning Representations (ICLR), 2025. URL https://openreview.net/forum?id=uvHmnahyp1. Enamine. Enamine, 2026. URL https://enamine.net. Accessed: 2026-06-01. Peter Ertl and Ansgar Schuffenhauer. Estimation of synthetic accessibility score of drug-like molecules based on molecular complexity and fragment contributions. Journal of Cheminformatics, 1(1):8, 2009. doi: 10.1186/1758-2946-1-8. Sai Krishna Gottipati, Boris Sattarov, Sufeng Niu, Yashaswi Pathak, Haoran Wei, Shengchao Liu, Karam J. Thomas, Simon Blackburn, Connor W. Coley, Jian Tang, Sarath Chandar, and Yoshua Bengio. Learning to navigate the synthetically accessible chemical space using reinforcement learning. In Proceedings of the 37th International Conference on Machine Learning, ICML’20. JMLR.org, 2020. Sai Krishna Gottipati, Yashaswi Pathak, Boris Sattarov, Sahir, Rohan Nuttall, Mohammad Amini, Matthew E. Taylor, and Sarath Chandar. Towered actor critic for handling multiple action types in reinforcement learning for drug discovery. Proceedings of the AAAI Conference on Artificial Intelligence, 35(1):142–150, 2021. doi: 10.1609/aaai.v35i1.16087. Sai Krishna Gottipati, Yashaswi Pathak, Rohan Nuttall, Sahir, Raviteja Chunduru, Ahmed Touati, Sriram Ganapathi Subramanian, Matthew E. Taylor, and Sarath Chandar. Maximum reward formulation in reinforcement learning, 2023. URL https://arxiv.org/abs/2010.03744. Rafael Gómez-Bombarelli, Jennifer N. Wei, David Duvenaud, José Miguel Hernández-Lobato, Benjamín Sánchez-Lengeling, Dennis Sheberla, Jorge Aguilera-Iparraguirre, Timothy D. Hirzel, Ryan P. Adams, and Alán Aspuru-Guzik. Automatic chemical design using a data-driven continuous representation of molecules. ACS Central Science, 4(2):268–276, January 2018. ISSN 2374-7951. doi: 10.1021/acscentsci.7b00572. URL http://dx.doi.org/10.1021/acscentsci.7b00572. Jan H. Jensen. A graph-based genetic algorithm and generative model/monte carlo tree search for the exploration of chemical space. Chemical Science, 10(12):3567–3572, 2019. doi: 10.1039/ C8SC05372C. G. Landrum. Rdkit: Open-source cheminformatics software, 2026. URL https://github.com/rdkit/ rdkit/releases/tag/Release_2026_03_3. Hannes H. Loeffler, Jiazhen He, Angelica Tibo, Jon Paul Janet, Alexey Voronov, Lewis H. Mervin, and Ola Engkvist. Reinvent 4: Modern ai-driven generative molecule design. Journal of Chemin- formatics, 16(1):20, 2024. doi: 10.1186/s13321-024-00812-5. Akshat Kumar Nigam, Pascal Friederich, Mario Krenn, and Alán Aspuru-Guzik. Augmenting genetic algorithms with deep neural networks for exploring the chemical space. In International Conference on Learning Representations (ICLR), 2020. URL https://openreview.net/pdf?id=H1lmyRNFvr. Liva Ralaivola, Sanjay J. Swamidass, Hiroto Saigo, and Pierre Baldi. Graph kernels for chemical informatics. Neural Networks, 18(8):1093–1110, 2005. doi: 10.1016/j.neunet.2005.07.009. Gregor N. C. Simm, Robert Pinsler, and José Miguel Hernández-Lobato. Reinforcement learning for molecular design guided by quantum mechanics. In Hal Daumé I and Aarti Singh, editors, 13 Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 8959–8969. PMLR, 13–18 Jul 2020. URL https://proceedings. mlr.press/v119/simm20b.html. Austin Tripp and José Miguel Hernández-Lobato. Genetic algorithms are strong baselines for molecule generation, 2023. URL https://arxiv.org/abs/2310.09267. Austin Tripp, Sergio Bacallado, Sukriti Singh, and José Miguel Hernández-Lobato. Tanimoto random features for scalable molecular machine learning. In Advances in Neural Information Processing Systems, volume 36, 2023. Zhengkai Tu, Sourabh J. Choure, Mun Hong Fong, Jihye Roh, Itai Levin, Kevin Yu, Joonyoung F. Joung, Nathan Morgan, Shih-Cheng Li, Xiaoqi Sun, Huiqian Lin, Mark Murnin, Jordan P. Liles, Thomas J. Struble, Michael E. Fortunato, Mengjie Liu, William H. Green, Klavs F. Jensen, and Connor W. Coley. Askcos: an open source software suite for synthesis planning, 2025. URL https://arxiv.org/abs/2501.01835. H. M. Vinkers, M. R. de Jonge, F. F. Daeyaert, J. Heeres, L. M. Koymans, J. H. van Lenthe, P. J. Lewi, H. Timmerman, K. Van Aken, and P. A. Janssen. Synopsis: Synthesize and optimize system in silico. Journal of Medicinal Chemistry, 46(13):2765–2773, 2003. doi: 10.1021/jm030809x. Jiaxuan You, Bowen Liu, Rex Ying, Vijay Pande, and Jure Leskovec. Graph convolutional policy network for goal-directed molecular graph generation, 2019. URL https://arxiv.org/abs/1806. 02473. Zhenpeng Zhou, Steven Kearnes, Li Li, Richard N. Zare, and Patrick Riley. Optimization of molecules via deep reinforcement learning. Scientific Reports, 9(1):10752, July 2019. ISSN 2045-2322. doi: 10.1038/s41598-019-47148-x. URL http://dx.doi.org/10.1038/s41598-019-47148-x. 14 Supplementary Material for PGFS++: Molecular Property Improvement under Synthesis and Diversity Constraints S1 Algorithm Details Training paradigm Figure S2 compares the algorithm of PGFS (left) and PGFS++ (right). For PGFS++, at each time stept, the current moleculeR (1) t is the MDP state. The Actor encodesR (1) by a networkfapplied to its Morgan fingerprint FP(R (1) ), yielding a latent representation h = f (FP(R (1) )). A template head then samples a reaction template T ∼ Softmax(T mask ⊙ W T h), whereT mask masks out templates that are chemically infeasible forR (1) . IfT=Stop, the Actor returns immediately with no second reactant. Otherwise, a query head forms q = MLP([h,φ T (T )]) fromhand a learned template embeddingφ T (T). All building-block candidatesR (2) are represented by a table of learned embeddings as keys K = φ 2 (R (2) ). The second reactant is sampled by attention over masked keys, R (2) ∼ Softmax R (2) mask ⊙ (qK ⊤ ) , whereR (2) mask is the precomputed mask that retains only reactants compatible with templateT. The joint discrete action is therefore (T ,R (2) ), and the Actor also returnslogπ(T,R (2) | R (1) ) for PPO. The Critic is a state-value networkV(FP(R (1) )) sharing the same Morgan-fingerprint input. The environment appliesForwardReaction(R (1) ,T,R (2) ) to obtain the productR (1) ′ and assigns a step reward. Training follows on-policy PPO. The Main loop collectsn steps transitions into a rollout bufferD, sampling a fresh starting moleculeR (1) ∼ p 0 at the beginning of each episode. After each rollout, Update computes generalised advantage estimates ˆ A t = GAE(r t ,V t ) 1 Symbol Value Symbol Value α3× 10 −4 λ0.95 n steps 2048 ε0.2 B64 c e 0.05 K10 c v 0.5 γ0.99 g max 0.5 KL targ 0.02 Table S1: Additional PPO hyperparameters used in PGFS++ experiments (notations as in Equation(S1)). and returns ˆ G t = ˆ A t + V t , then maximises the clipped surrogateL clip (ρ t , ˆ A t ) with value lossL V and entropy bonusH[π], where ρ t = π(a t | R (1) t )/π old (a t | R (1) t ). Early stopping is triggered when KL(π old ,π) exceeds a target threshold KL targ . PPO hyperparameters For a fair comparison, PGFS++ strictly follows the shared experimental hyperparameters of PGFS wherever the two methods are directly comparable. The underlying learning algorithms differ, however: PGFS is trained with off-policy TD3, whereas PGFS++ is trained with on-policy PPO. Writingρ t =π(a t | R (1) t )/π old (a t | R (1) t ) for the importance ratio, the PPO objective used by PGFS++ is L(θ) = −L clip (ρ t , ˆ A t ) + c v L V − c e H[π],(S1) whereL clip is the clipped policy surrogate with clip rangeε,L V is the squared-error value loss, andH[π] is the policy entropy. Advantages ˆ A t are estimated by GAE with discountγand trace parameterλ. Optimisation uses learning rateα, rollout lengthn steps , minibatch sizeB, andK epochs per rollout, with gradient clipping atg max and early stopping whenKL(π old ,π)> KL targ . Table S1 lists the corresponding hyperparameter values. S2 Evaluation protocol When evaluating model performance, we use the following settings: • Templates. Two sets of reaction templates are used in experiments: reversible and general. The reversible template set contains only templates for which RDKit can form a backward reaction by swapping reactant and product SMARTS and recovering valid reactants. This auto-reversion is used when learning the backward policyP B in SynFlowNet. The general template set, which is derived from Button et al. (2019), includes templates that do not support a reliable auto-reversion. •Inference. At test time we use greedy (argmax) template and reactant selection, with a maximum of 4 reaction steps and an explicit Stop action. 2 MethodGeneral templates Reversible templates SynFlowNet0.7660.856 Ours 0.895 0.867 Table S2: Median QED on the test set under general vs. reversible templates. Ours denotes the matched PGFS++ QED models from the main paper. S3 Additional experimental results S3.1 SynFlowNet performance under general vs. reversible templates SynFlowNet favors reversible reaction templates to support learning of a backward policy. Table S2 compares median QED on the test set when SynFlowNet is trained and evaluated under general templates versus under reversible templates. The reversible setting yields substantially higher median QED (0.856 vs. 0.766), consistent with SynFlowNet’s dependence on template reversibility. PGFS++ does not have this restriction and attains strong QED in both settings. Accordingly, we omit SynFlowNet from the general-template comparisons in the main paper: presenting it there would disadvantage a method whose modelling assumptions are favored to reversible chemistry. For the general-template setting we instead compared to REINVENT4 as a state-of-the-art method, which has no reaction-template constraints. SynFlowNet is retained as the state-of-the-art baseline under reversible templates, and PGFS++ is trained separately in each template setting. S3.2 Seed variability of PGFS++ The main paper reports PGFS++ results aggregated over three random seeds. This subsection summarizes the corresponding seed variability. For the molecule-level results in the main paper, each metric is first averaged across the three seeds for each starting molecule, and the resulting distribution is summarized using the mean±standard deviation. Table S3 additionally reports the median obtained from each individual seed under all experimental settings. Overall, the target-property results are highly consistent across seeds. Comparing seeds 123456 and 654321 in Table S3 under general-template QED experiments, their 1−SAremains similar while RAS differs substantially, mainly due to longer routes used by seed 654321 (2.72 vs. 1.16 reactions on average) and different template /R (2) usage. This underscores the need for RAS rather than relying on SA alone. 3 Tem.Obj.Seed/stat.Prop. ↑Div. ↑1− SA ↑I/O sim. ↑RAS ↑ Gen.QED 1234560.8950.8730.7490.3890.705 6543210.8970.8670.6980.3600.267 2461350.8970.8770.7220.3710.685 Mean 0.896 0.872 0.723 0.373 0.552 s.d.0.0010.0050.0260.0150.247 Gen.SEH 1234560.8120.7720.7480.3120.545 6543210.8350.8120.6830.2980.449 2461350.8310.6630.7350.2900.563 Mean 0.826 0.749 0.722 0.300 0.519 s.d.0.0120.0770.0340.0110.061 Rev.QED 1234560.8670.8730.6770.3600.636 6543210.8670.8780.6780.3480.642 2461350.8740.8830.6960.3520.679 Mean 0.869 0.878 0.684 0.353 0.652 s.d.0.0040.0050.0110.0060.023 Rev.SEH 1234560.9130.6780.5490.2020.366 6543210.8900.6270.4940.2000.227 2461350.9020.5660.4850.1980.197 Mean 0.902 0.624 0.509 0.200 0.263 s.d.0.0120.0560.0350.0020.090 Table S3: Seed-wise median performance of PGFS++. Gen. and Rev. denote general and reversible reaction templates, respectively. The first three rows of each block report the median over generated molecules for each random seed. The final two rows report the mean and sample standard deviation, respectively, across the three seed-wise medians. S3.3 Novelty of generated molecules. We further compare PGFS++ (Ours) against REINVENT4 and SynFlowNet on novelty relative to known drug-like chemistry in MOSES, a molecular set containing approximately 1.5 million molecules. Using the same seed sweep setting as the main paper, for each generated molecule, we compute its maximum Morgan-fingerprint Tanimoto similarity to any MOSES molecule and define novelty(m) = 1−max m ′ ∈MOSES Tanimoto FP(m), FP(m ′ ) . Figure S1 shows property score (QED or SEH) against novelty for the four settings used in the main paper. For PGFS++, both novelty and property scores are averaged across seeds. Colored stars mark each method’s median (property, novelty). Table S4 reports the corresponding median statistics. Compared with REINVENT4, PGFS++ attains higher novelty at comparable target property, likely because REINVENT4 stays close to its pretraining distribution. Compared with SynFlowNet, PGFS++ remains competitive on both novelty and target property. 4 Median target propertyMedian novelty Templates Objective Seed Baseline PGFS++∆p Baseline PGFS++∆ p Gen.QED 123456 0.871 0.895 0.457 0.554 6543210.8970.613 2461350.8970.549 Mean 0.896+0.025 < 0.0010.572+0.115 0.030 Gen.SEH 123456 0.869 0.812 0.426 0.603 6543210.8350.609 2461350.8310.587 Mean0.826 −0.043 0.0260.600+0.174 0.002 Rev.QED 123456 0.856 0.867 0.564 0.567 6543210.8670.547 2461350.8740.582 Mean0.869+0.013 0.0310.565+0.001 0.879 Rev.SEH 123456 0.900 0.913 0.645 0.663 6543210.8900.673 2461350.9020.682 Mean 0.902+0.002 0.8550.673+0.028 0.036 Table S4: Target-property and novelty results on the test set. Gen. and Rev. denote general and reversible reaction templates, respectively. For each baseline, the reported value is the median over generated molecules from a single run: REINVENT4 under general templates and SynFlowNet under reversible templates. For PGFS++, the first three rows of each block report the median for each random seed, while the final row reports the mean of the three seed-wise medians. The change ∆ is calculated as the PGFS++ mean minus the corresponding baseline value, such that positive values indicate higher scores for PGFS++. Statistical significance is assessed by comparing PGFS++ with the corresponding baseline over paired starting molecules, quantified byp-values. Shaded cells indicate the PGFS++ mean values used for comparison with the baselines. Boldface indicates the better result within each setting. 5 Figure S1: Property–novelty trade-offs on the test set. Top: general templates (PGFS++ versus REIN- VENT4) for QED (left) and SEH (right). Bottom: reversible templates (PGFS++ versus SynFlowNet) for QED (left) and SEH (right). Each point is one episode’s generated molecule. Colored stars mark each method’s median (property, novelty). The black star marks an ideal high-property / high-novelty corner. Numerical values of medians are presented in Table S4. 6 PGFS 1: procedure Actor(R (1) ) 2: T ← f FP(R (1) ) ▷ f-net, Morgan FP 3: T ← GumbelSoftmax(T mask ⊙ T, τ ) 4: a← π [R (1) , T ] ▷ π-net 5: return T, a 6: 7: 8: 9: 10: 11: 12: 13: procedure Critic(R (1) , T, a) 14: return Q FP(R (1) ), T, a ▷ twin Q 1 ,Q 2 15: 16: procedure Env.Step(R (1) , T, a) 17: R (2) ← ValidReactants(T ) 18: K ← [FP(r)] r∈R (2) ▷ RLV2 / MolDSet FP 19: A← Top-k arg min r ∥a− K r ∥ 2 ▷ kNN(a,R (2) ) 20: P ←ForwardReaction(R (1) , T, r) r∈A 21: r ⋆ t , R (1) ′ ← arg max p∈P Score(R (1) →p) 22: return R (1) ′ , r ⋆ t 23: 24: 25: procedure Backward(buffer minibatch) 26: T i+1 , a i+1 ← Actor-target(R (1) i+1 ) 27: y i ← r i + γ min j=1,2 Q targ j FP(R (1) i+1 ), T i+1 , a i+1 28: minL(θ Q ) = 1 N P i y i − Q FP(R (1) i ), T i , a i 2 29: minL(θ f,π ) =− P i Q FP(R (1) i ), Actor(R (1) i ) 30: minL(θ f ) =− P i CE T i , logf (FP(R (1) i )) 31: 32: 33: 34: 35: 36: 37: 38: 39: 40: procedure Main(f, π, Q) 41: for episode = 1 to M do 42: sample R (1) 0 43: for t = 0 to N do 44:T t , a t ← Actor(R (1) t ) 45:R (1) t+1 , r t ← Env.Step(R (1) t , T t , a t ) 46:store (R (1) t , T t , a t , R (1) t+1 , r t ) in buffer 47:sample a random minibatch from buffer 48:Backward(minibatch) 49: end for 50: end for PGFS++ (Ours) 1: procedure Actor(R (1) ) 2: h← f FP(R (1) ) ▷ f-net, Morgan FP 3: T ∼ Softmax(T mask ⊙ W T h) 4: if T = Stop then 5: return (T,∅), logπ(T | R (1) ) 6: end if 7: q ← MLP [h, emb(T ) ] ▷ query head 8: K ← Embedding(R (2) )▷ learned keys 9: R (2) ∼ Softmax R (2) mask ⊙ (qK ⊤ ) ▷ attention 10: return (T, R (2) ), logπ(T, R (2) | R (1) ) 11: 12: procedure Critic(R (1) ) 13: return V FP(R (1) ) ▷ state value 14: 15: procedure Env.Step(R (1) , T, R (2) ) 16: if T = Stop then 17: return R (1) , r t 18: end if 19: R (2) ← ValidReactants(T ) ▷ masked in Actor 20: R (1) ′ ← ForwardReaction(R (1) , T, R (2) ) 21: r t ← Score(R (1) →R (1) ′ ) ▷ ∆reward + bonus 22: return R (1) ′ , r t 23: 24: 25: procedure Update(rollout D) 26: ˆ A t ← GAE(r t , V t ) 27: ˆ G t ← ˆ A t + V t 28: for k = 1 to K epoch do 29: for each minibatch B ⊂D do 30:ρ t ← π(a t | R (1) t )/π old (a t | R (1) t ) 31:L←−L clip (ρ t , ˆ A t ) + c v L V − c e H[π] 32:θ ← θ− α∇ θ L 33:if KL(π old , π) > KL targ then 34:break 35:end if 36: end for 37: end for 38: 39: 40: 41: procedure Main(π, V ) 42: for m = 1 to M do 43: D ←∅ 44: for t = 1 to n steps do 45:if new episode then 46:sample R (1) ∼ p 0 47:end if 48:(T, R (2) ), logπ ← Actor(R (1) ) 49:V ← Critic(R (1) ) 50:R (1) t+1 , r t ← Env.Step(R (1) , T, R (2) ) 51:store (R (1) , T, R (2) , logπ, V, r t ) in D 52:R (1) ← R (1) t+1 53: end for 54: Update(D) 55: end for Figure S2: Side-by-side comparison of PGFS (left) and PGFS++ (right). 7