PGFS++: Molecular Property Improvement under Synthesis and Diversity Constraints

PGFS++: Molecular Property Improvement under Synthesis and Diversity Constraints

arXiv:2608.19121v1 [cs.LG] 19 Aug 2026

Boqiao Zhang1 , Godbless Tamaraebi James1 , Sai Krishna Gottipati, Andrew Fitzgibbon1,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

A. Illustrative Routes of Generating a Magnet Output Molecule

B. Model Performance (QED-Diversity Trade-off)

PGFS++: Higher mean QED and high diversity PGFS+: Higher mean QED but collapsed 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 as R2 , to react with the first reactant R1 under a chosen template T . However, the space of possible R2 candidates can be very large, and the action space becomes even larger when the model jointly selects both T and R2 . The original PGFS framework overcame this by using a continuous action space for second-reactant selection: a neural R2 predictor is trained to predict the embedding of the second reactant, and a nearest-neighbour search is then used to retrieve an actual R2 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 executed R2 . 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 R2 embedding and retrieving neighbours, PGFS+ projects molecules and templates into a learned reaction space and directly scores all template-compatible R2 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-NN selection step with sampling at training time, which improves learning in the presence of multiple near-equivalent R2 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 R2 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 step t, all methods use an f -network conditioned on the current first reactant R1′ (t) to predict the reaction template T . PGFS uses the fingerprint representation F (R1′ (t)) directly, whereas PGFS+ and PGFS++ first transform the R1 representation through an additional MLP. For second-reactant selection, PGFS represents R2 candidates using fixed fingerprints and retrieves candidates by k-nearest neighbours. In contrast, PGFS+ and PGFS++ represent template-compatible R2 candidates with a trainable embedding lookup table and score them globally using an attention-like mechanism with query q = concat(ϕ1 (R1′ ), ϕT (T )). Once R2′ (t) is selected, the reaction runner applies the chosen reaction and returns the product, which becomes the next state R1′ (t + 1).

3

and greatly improved diversity. 2. We present PGFS+, and identified the magnet output problem, which reflects a broad failure 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 magnetoutput 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 template T given the current reactant R1 . Conditioned on R1 and T , it then selects a second reactant R2 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 possible R2 candidates for each T and R1 . PGFS addresses this by using a neural network to predict an R2 embedding, followed by k-nearest-neighbor retrieval from the building-block set. In contrast, PGFS+ and PGFS++ directly scores feasible R2 candidates after applying a template-specific compatibility mask. We observed that the large R2 action space can be efficiently reduced to a tractable candidate set using RDKit substructure matching. We pre-compute a T -R2 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 compatible R2 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 (R1 ) and ϕT (T ) of R1 and T to form a query vector, q = concat(ϕ1 (R1 ), ϕT (T )), and use this query to score the masked R2 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 indexed R2 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:
× repisode = rproperty (mout ) 1 + c sim(min , mout ) (1) where rproperty is the property-based reward, such as ∆QED and ∆SEH, c > 0 control the bonus strength, and sim(min , mout ) 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: ( × repisode , sim(min , mout ) ≥ τ, repisode = (2) × min{repisode , κ}, sim(min , mout ) < τ, 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). Let m and n denote 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 min and output molecule mout = G(min ), define X = ψ(min ),

Y = ψ(mout ),

A = Y − X,

where A is 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   Din = E 1 − sim(min , m′in ) 1 = E∥X − X ′ ∥2 = E∥X − µX ∥2 = Var(X), 2  (4) Dout = E 1 − sim(mout , m′out ) 1 = E∥Y − Y ′ ∥2 = E∥Y − µY ∥2 = Var(Y ), 2 where µ· = E[·]. Since Y = X + A and µY = µX + µA , Y − µY = (X − µX ) + (A − µA ). Substituting this expression into Equation (4) gives Dout = Din + VA + 2Γ, VA = E∥A − µA ∥2 = Var(A) ≥ 0,

(5)

Γ = E[⟨X − µX , A − µA ⟩] . Equation (5) separates the effect of editing into the non-negative edit variance VA and the input–edit covariance 2Γ. Therefore, VA Dout < Din ⇐⇒ Γ < − . 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 Γ = −Din .

VA = Din , Substitution into Equation (5) recovers Dout = 0.

Bounding diversity drop with input–output similarity. Since sim(m, m) = 1, the feature-space embeddings satisfy ∥X∥ = ∥Y ∥ = 1. Therefore, Din = 1 − ∥µX ∥2 ,

Dout = 1 − ∥µY ∥2 .

Define the average input–output similarity   s̄ = E sim(min , mout ) = E⟨X, Y ⟩. Let a = ∥µX ∥ =

p 1 − Din .

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 ∥ ≤ hs̄ (a), where

( 1, s̄ ≤ a, √ √ hs̄ (a) = 2 2 as̄ + 1 − a 1 − s̄ , s̄ > a.

Hence, the diversity drop     ∆div = Din − Dout + = ∥µY ∥2 − ∥µX ∥2 + where [v]+ = max(v, 0), satisfies 2 p ∆div ≤ hs̄ 1 − Din − (1 − Din ) . When s̄ >

√

1 − Din , this becomes  p 2 p p ∆div ≤ s̄ 1 − Din + 1 − s̄2 Din − (1 − Din ) .

(6)

(7)

When the input pool is highly diverse (Din ≈ 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 selected R1 reactants from the dataset. At test time, each model was used to improve every test R1 , and metrics were recorded for each episode. We compared PGFS++ against PGFS, PGFS-TAC, PGFS-MB, SynFlowNet, and the REINVENT4 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 either rproperty = ∆QED or rproperty = ∆SEH as 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 18 ). 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, and underlined entries 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 set O. It is defined as one minus the average pairwise Tanimoto similarity between Morgan fingerprints: X 1 D(O) = 1 − T (FP(mi ), FP(mj )), |P | mi ,mj ∈P

where P = {(mi , mj ) : mi , mj ∈ O, i < j} is the set of unordered molecule pairs, T denotes 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 − SA in 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 molecule mein and the output molecule meout for each episode e, averaged across episodes. For each episode e, we compute simin,out (e) = T (FP(mein ), FP(meout )). 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 episode e, let the selected synthetic route contain reaction steps n = 1, . . . , Ne . 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). Let praw denote the ASKCOS fast-filter plausibility n (Tu et al., 2025) of the reaction at step n, which lies in [0, 1], with larger values indicating greater feasibility. We use it directly: RASTn = praw n . denote the BR-SAScore 2. Reactant accessibility (R2 score). For bimolecular steps, let sraw n (Ertl and Schuffenhauer, 2009) of the selected second reactant R2 , where smaller values indicate greater accessibility. We map this score to [0, 1] by inverting and rescaling it with fixed bounds smin = 1 and smax = 10:   smax − sraw n R2 , 0, 1 . RASn = clip smax − smin 2 For uni-molecular steps, where no building block is added, we set RASR 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) =

Ne Y

2 RASTn × RASR . n

n=1

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 of R2 , the RAS metric can incorporate other R2 accessibility estimates, such as the price of R2 .

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: + repisode = rproperty (mout ) + w sim(min , mout ),

× repisode = rproperty (mout ) 1 + c sim(min , mout ) .

(8)

The two forms differ in how strongly they pull the output toward the input. For the additive form, + ∂repisode

∂ sim 11

= w,

while for the multiplicative form, × ∂repisode

= c rproperty (mout ). ∂ sim 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 to w = 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-computed T – R2 masks, attention-like global scoring for R2 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 for R1 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 Cheminformatics, 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é III 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 (1) step t, the current molecule Rt is the MDP state. The Actor encodes R(1) by a network f applied 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(Tmask ⊙ WT h), where Tmask masks out templates that are chemically infeasible for R(1) . If T = Stop, the Actor returns immediately with no second reactant. Otherwise, a query head forms q = MLP([h, ϕT (T )]) from h and a learned template embedding ϕT (T ). All building-block candidates R(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,
(2) R(2) ∼ Softmax Rmask ⊙ (qK⊤ ) , (2)

where Rmask is the precomputed mask that retains only reactants compatible with template T . The joint discrete action is therefore (T , R(2) ), and the Actor also returns log π(T, R(2) | R(1) ) for PPO. The Critic is a state-value network V (FP(R(1) )) sharing the same Morgan-fingerprint input. The ′ environment applies ForwardReaction(R(1) , T, R(2) ) to obtain the product R(1) and assigns a step reward. Training follows on-policy PPO. The Main loop collects nsteps transitions into a rollout buffer D, sampling a fresh starting molecule R(1) ∼ p0 at the beginning of each episode. After each rollout, Update computes generalised advantage estimates Ât = GAE({rt , Vt }) 1

Symbol

Value

Symbol

Value

α nsteps B K γ

3 × 10−4 2048 64 10 0.99

λ ϵ ce cv gmax KLtarg

0.95 0.2 0.05 0.5 0.5 0.02

Table S1: Additional PPO hyperparameters used in PGFS++ experiments (notations as in Equation (S1)).

and returns Ĝt = Ât + Vt , then maximises the clipped surrogate Lclip (ρt , Ât ) with value loss LV and entropy bonus H[π], where (1)

(1)

ρt = π(at | Rt )/πold (at | Rt ). Early stopping is triggered when KL(πold , π) exceeds a target threshold KLtarg . 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. (1) (1) Writing ρt = π(at | Rt )/πold (at | Rt ) for the importance ratio, the PPO objective used by PGFS++ is L(θ) = − Lclip (ρt , Ât ) + cv LV − ce H[π], (S1) where Lclip is the clipped policy surrogate with clip range ϵ, LV is the squared-error value loss, and H[π] is the policy entropy. Advantages Ât are estimated by GAE with discount γ and trace parameter λ. Optimisation uses learning rate α, rollout length nsteps , minibatch size B, and K epochs per rollout, with gradient clipping at gmax and early stopping when KL(πold , π) > KLtarg . 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 policy PB 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

Method SynFlowNet Ours

General templates

Reversible templates

0.766 0.895

0.856 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−SA remains 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.

Gen.

Gen.

Rev.

Rev.

Obj.

Seed/stat.

Prop. ↑

Div. ↑

1 − SA ↑

I/O sim. ↑

RAS ↑

QED

123456 654321 246135 Mean s.d.

0.895 0.897 0.897 0.896 0.001

0.873 0.867 0.877 0.872 0.005

0.749 0.698 0.722 0.723 0.026

0.389 0.360 0.371 0.373 0.015

0.705 0.267 0.685 0.552 0.247

SEH

123456 654321 246135 Mean s.d.

0.812 0.835 0.831 0.826 0.012

0.772 0.812 0.663 0.749 0.077

0.748 0.683 0.735 0.722 0.034

0.312 0.298 0.290 0.300 0.011

0.545 0.449 0.563 0.519 0.061

QED

123456 654321 246135 Mean s.d.

0.867 0.867 0.874 0.869 0.004

0.873 0.878 0.883 0.878 0.005

0.677 0.678 0.696 0.684 0.011

0.360 0.348 0.352 0.353 0.006

0.636 0.642 0.679 0.652 0.023

SEH

123456 654321 246135 Mean s.d.

0.913 0.890 0.902 0.902 0.012

0.678 0.627 0.566 0.624 0.056

0.549 0.494 0.485 0.509 0.035

0.202 0.200 0.198 0.200 0.002

0.366 0.227 0.197 0.263 0.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 Tanimoto FP(m), FP(m′ ) . m ∈MOSES

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 property Templates Objective Seed

Gen.

Gen.

Rev.

Rev.

QED

SEH

QED

SEH

123456 654321 246135 Mean 123456 654321 246135 Mean 123456 654321 246135 Mean 123456 654321 246135 Mean

Baseline PGFS++

0.871

0.869

0.856

0.900

0.895 0.897 0.897 0.896 0.812 0.835 0.831 0.826 0.867 0.867 0.874 0.869 0.913 0.890 0.902 0.902

∆

Median novelty p

Baseline PGFS++

∆

p

+0.025 < 0.001

0.554 0.613 0.549 0.572

+0.115 0.030

0.026

0.603 0.609 0.587 0.600

+0.174 0.002

0.031

0.567 0.547 0.582 0.565

+0.001 0.879

0.855

0.663 0.673 0.682 0.673

+0.028 0.036

0.457

0.426 −0.043

0.564 +0.013

0.645 +0.002

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 by p-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 REINVENT4) 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

PGFS++ (Ours)

(1) 1: procedure Actor(R )
2: T ← f FP(R(1) ) ▷ f -net, Morgan FP 3: T ← GumbelSoftmax(T ⊙ T, τ) mask
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 Q1 , Q2 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 minr ∥a − Kr ∥2 ▷ kNN(a, R(2) ) (1) 20: P ← {ForwardReaction(R , T, r)}r∈A ′ 21: rt⋆ , R(1) ← arg maxp∈P Score(R(1) → p) ′ 22: return R(1) , rt⋆ 23: 24: 25: procedure Backward(buffer minibatch) (1) 26: Ti+1 , ai+1 ← Actor-target(Ri+1 )
(1) targ 27: yi ← ri + γ minj=1,2 Qj FP(Ri+1 ), Ti+1 , ai+1
2 P (1) 28: min L(θQ ) = N1 i yi − Q FP(Ri ), Ti , ai P (1) (1)
29: min L(θf,π ) = − i Q FP(Ri ), Actor(Ri )
P (1) 30: min L(θf ) = − i CE Ti , log f (FP(Ri )) 31: 32: 33: 34: 35: 36: 37: 38: 39: 40: procedure Main(f, π, Q) 41: for episode = 1 to M do (1) 42: sample R0 43: for t = 0 to N do (1) 44: Tt , at ← Actor(Rt ) (1) (1) 45: Rt+1 , rt ← Env.Step(Rt , Tt , at ) (1) (1) 46: store (Rt , Tt , at , Rt+1 , rt ) in buffer 47: sample a random minibatch from buffer 48: Backward(minibatch) 49: end for 50: end for

(1) 1: procedure Actor(R )
(1) 2: h ← f FP(R ) ▷ f -net, Morgan FP 3: T ∼ Softmax(Tmask ⊙ WT 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
(2) 9: R(2) ∼ Softmax Rmask ⊙ (qK ⊤ ) ▷ attention 10: return (T, R(2) ), log π(T, R(2) | R(1) ) 11: (1) 12: procedure Critic(R
) 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) , rt 18: end if 19: R(2) ← ValidReactants(T ) ▷ masked in Actor ′ 20: R(1) ← ForwardReaction(R(1) , T, R(2) ) ′ 21: rt ← Score(R(1) → R(1) ) ▷ ∆reward + bonus ′ 22: return R(1) , rt 23: 24: 25: procedure Update(rollout D) 26: Ât ← GAE({rt , Vt }) 27: Ĝt ← Ât + Vt 28: for k = 1 to Kepoch do 29: for each minibatch B ⊂ D do (1) (1) 30: ρt ← π(at | Rt ) / πold (at | Rt ) clip 31: L ← −L (ρt , Ât ) + cv LV − ce H[π] 32: θ ← θ − α∇θ L 33: if KL(πold , π) > KLtarg 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 nsteps do 45: if new episode then 46: sample R(1) ∼ p0 47: end if 48: (T, R(2) ), log π ← Actor(R(1) ) 49: V ← Critic(R(1) ) (1) 50: Rt+1 , rt ← Env.Step(R(1) , T, R(2) ) 51: store (R(1) , T, R(2) , log π, V, rt ) in D (1) 52: R(1) ← Rt+1 53: end for 54: Update(D) 55: end for

Figure S2: Side-by-side comparison of PGFS (left) and PGFS++ (right).

7

Related documents

• EvoDiffMol: evolutionary diffusion framework for 3D molecular design with optimized properties.#375259
• GROMACS 4:  Algorithms for Highly Efficient, Load-Balanced, and Scalable Molecular Simulation#6939
• AI-Based Retrosynthesis of Known and Unexplored Covalent Organic Framework Monomers.#469097
• RetroMPA: A Molecular Property-Aware Auxiliary Framework for Enhancing Retrosynthesis Prediction#620525
• RetroMPA: A Molecular Property-Aware Auxiliary Framework for Enhancing Retrosynthesis Prediction#618846