Diversity-Aware Reinforcement Learning for de novo Drug Design

1 Introduction

The development of a novel pharmaceutical drug is a highly intricate process that can span up to a decade and incur costs exceeding US $1 billion [Wouters et al., 2020; DiMasi et al., 2016]. A key part of such an effort involves the identification of novel drug candidates that exhibit the desired molecular properties [Hughes et al., 2011]. The success in identifying drug candidates primarily depends on selecting chemical starting points that surpass a certain threshold in bioactivity toward the desired target, known as hits. High-quality hits can substantially reduce the time required to identify a viable drug candidate and be the determining factor in the success of a drug discovery campaign [Quancard et al., 2023]. Designing novel pharmaceutical molecules, or de novo drug design, is extremely challenging given the estimated number of up to 1060 possible drug-like molecules [Reymond, 2015].

Recent advances in de novo drug design utilize reinforcement learning (RL) to navigate this vast chemical space by fine-tuning a pre-trained generative model [Popova et al., 2018; Gummesson Svensson et al., 2024; Guo and Schwaller, 2024; Atance et al., 2022; Liu et al., 2021; Olivecrona et al., 2017]. Evaluations by [Gao et al., 2022] and [Thomas et al., 2022a] have demonstrated good performance when using RL to fine-tune a pre-trained recurrent neural network (RNN) to generate molecules encoded in the Simplified Molecular Input Line Entry System (SMILES) [Weininger, 1988]. Also, this approach is widely adopted in real-world applications in drug discovery [Pitt et al., 2025]. However, RL-based de novo drug design methods can easily become stuck in local optima, generating structurally similar molecules- a phenomenon known as mode collapse. This is undesirable as it prevents the agent from discovering more diverse and potentially more promising local optima. To mitigate mode collapse, [Blaschke et al., 2020b] introduced a count-based method that penalizes generated molecules based on their structure. When too many structurally similar molecules have been generated, the agent observes zero reward, instead of the actual extrinsic reward, for future generated molecules with the same structure. This is a popular way to avoid mode collapse for RL-based de novo drug design [Thomas et al., 2022b; Guo and Schwaller, 2024; Loeffler et al., 2024; Gummesson Svensson et al., 2024].

Most work mainly focuses on avoiding mode collapse to find the most optimal solution. However, the quantitative structure-activity relationship (QSAR) models utilized for in silico assessment of molecules introduce uncertainties and biases due to limited training data [Renz et al., 2019]. Thus, it is important to explore numerous modes of these models to increase the chance of identifying potential drug candidates. 1 Also, the identified (local) optimal solution might not be optimal in terms of observed safety and therapeutic effectiveness in the body. Therefore, it is meaningful to generate a diverse set of molecules. Recent work by [Renz et al., 2024] focuses on the generated molecules' diversity, finding superior performance of SMILES-based autoregressive models using RL to optimize the desired properties. They use a penalization method based on the work by [Blaschke et al., 2020b] to enable diverse molecule generation. As an alternative to penalizing the extrinsic reward, previous work in RL has shown that providing intrinsic motivation to the agent can enhance the exploration [Bellemare et al., 2016; Burda et al., 2018; Badia et al., 2020]. Recent efforts by [Park et al., 2025] and [Wang and Zhu, 2024] show the potential of memoryand prediction-based intrinsic motivation approaches in de novo drug design, demonstrating their capability to enhance the optimization of properties.

1 Extended version of our paper, including supplementary material, is available at https://doi.org/10.48550/arXiv.2410.10431

The generation of diverse sets of molecules with high (extrinsic) rewards is crucial in the drug discovery process. A diverse molecular library increases the likelihood of identifying candidates with unique and favorable pharmacological profiles, thereby enhancing the overall efficiency and success rate of drug development pipelines. While most prior research has concentrated on generating individual molecules with high (extrinsic) rewards, our work shifts the focus toward the generation of diverse molecular entities by systematically investigating various intrinsic rewards and reward penalties. This approach aims to counteract mode collapse and promote the exploration of a broader chemical space. Intrinsic rewards, inspired by human-like curiosity, encourage the RL agent to explore less familiar areas of the chemical space; while reward penalties discourage the generation of structurally similar molecules. By employing these strategies, we aim to investigate further the robustness and applicability of RL-based de novo drug design. To our knowledge, this is the first work to comprehensively study the effect such methods have on the diversity of the generated molecules. By doing so, we provide a novel framework that not only seeks optimal solutions but also ensures a wide-ranging exploration of the potential chemical space of bespoke drug candidates. This could significantly enhance the drug discovery process by providing a more diverse and promising set of molecules for further experimental validation.

2 ProblemFormulation

In this section, we introduce our framework for de novo drug design. The problem is string-based molecule generation, by fine-tuning a pre-trained policy. Following previous work, we formulate the generative process as a reinforcement learning problem where the task is to fine-tune a pre-trained generative model [Olivecrona et al., 2017]. An action corresponds to adding one token to the string representation of the molecule. A is the set of possible actions, including a start token astart and a stop token astop . The de novo drug design problem can be modeled as a Markov decision process (MDP). at ∈ A is the action taken at state st, the current state st = a0:t-1 is defined as the sequence of performed actions up to round t, the initial action a0 = astart is the start token. The transition probabilities P(st+1|st, at) = δst++at are deterministic, where P(terminal|st, astop ) = 1, ++ denotes the concatenation of two sequences and δz denotes the dirac distribution at z. If action astop is taken, the following state is terminal, stopping the current generation process and subsequently evaluating the generated molecule. The extrinsic reward is

We let T denote the round that a terminal state is visited, i.e., aT-1 = astop . The reward r(sT ) ∈ [0, 1] (only observable at a terminal state) measures the desired property, which we want to optimize, of molecule A = a1:T-2. Note that in practice, the string between the start and stop tokens encodes a molecule such that a1:T-2 is equivalent to a0:T-1 during evaluation. The objective is to fine-tune a policy πθ, parameterized by θ, to generate a structurally diverse set of molecules optimizing the property score r(·).

In practice, at each step i of the generative process, B full trajectories (until reaching a terminal state) are rolled out, to

Algorithm 1 Diversity-Aware RL framework

1: input: I,B, θprior, h 2: M← ∅ ▷ Initialize memory 3: θ ← θprior ▷ The pre-trained policy is fine-tuned 4: for i=1,. . . ,I do ▷ Generative steps 5: L(θ) ← 0 6: B ← ∅ 7: for b=1,. . . ,B do ▷ Generate batch of molecules 8: t ← 0 9: at ← a(start) ▷ Start token is always initial action 10: st+1 ← at 11: while st+1 is not terminal do 12: t ← t + 1 13: at ∼ πθ(st) 14: st+1 ← a0:t 15: end while 16: B ← B ∪ st+1 17: Observe property score r(st+1) 18: if r(st+1) ≥ h then 19: M ← M ∪ {st+1} 20: end if 21: Compute and store penalty f(st+1) 22: end for 23: for A ∈ B do 24: Compute intrinsic reward RI(A) 25: Compute diversity-aware reward ˆR(A) 26: Compute loss LA(θ) wrt ˆR(A) 27: L(θ) ← L(θ) + LA(θ) 28: end for 29: Update θ by one gradient step minimizing L(θ) 30: end for 31: output: M

obtain a batch B of generated molecules. Also, the diversityaware reward ˆR(A) (see Section 3) for each molecule A ∈ B is observed by the agent and subsequently used for finetuning. The diversity-aware reward ˆR(A) is computed using the penalty f(A) and/or intrinsic reward RI (depending on which reward function is used). Algorithm 1 illustrates our diversity-aware RL framework.

3 Diversity-Aware Reward Functions

In this section, we define the diversity-aware reward functions examined in this study. We investigate two approaches to encourage diversity among generated molecules by modifying the extrinsic reward: (1) penalize the extrinsic reward, and (2) provide intrinsic reward (intrinsic motivation). Moreover, we also investigate the combination of these approaches by integrating two intrinsic reward approaches with a penalty function on the extrinsic reward. Given an extrinsic reward R(A), the agent will receive a reward signal at the end of the generation sequence in the form of

for every generated molecule A. Hence, we impose reward shaping [Laud, 2004] where the reward observed by the agent is a linear function of the non-linear extrinsic reward, as illustrated in Figure 1. The penalty defines the importance of the extrinsic reward for each molecule to determine the exploitation rate adaptively. Sufficient exploitation is necessary to find high-quality solutions. In contrast, the intrinsic reward provides an additive bonus to encourage the agent to continually explore (independent of a molecule's extrinsic reward). We suggest several novel domain-specific penalties and intrinsic rewards, and, to the best of our knowledge, this domain-specific combination is novel. Such a linear combination avoids adding unnecessary complexity to the reward objective, but is necessarily not optimal. We want to optimize a complex extrinsic reward function while enforcing continuous exploration of a large solution space. The aim is to find a diverse set of high-quality solutions.

3.1 Extrinsic Reward Penalty

We propose and examine five different functions to penalize the extrinsic reward by discretely or continuously decreasing it based on the number of previously generated structurally similar molecules. These are based on binary, error, linear, sigmoid, or hyperbolic tan functions. To the best of our knowledge, utilizing error and hyperbolic tangent functions is novel for our application. Below, we define all penalty functions for clarity.

Non-binary functions will provide an incremental change of the (extrinsic) reward signal over time. Therefore, it is potentially more informative and can incrementally incentivize the agent to find new local optima, while still exploiting the current local optima. On the other hand, a binary function implies a sharp change of the reward function, where the agent quickly needs to find new optima.

Identical Molecular Scaffold Penalty (IMS). The IMS penalty was first introduced by [Blaschke et al., 2020b] and has thereafter been used in several works, e.g., [Loeffler et al., 2024; Guo and Schwaller, 2024]. It is based on molecular scaffolds, which is one of the most important and commonly used concepts in medicinal chemistry. The IMS penalty uses the molecular scaffold defined by [Bemis and Murcko, 1996], which is obtained by removing all side chains (or R groups). In this work, we also study the Topological scaffold, which is obtained from the molecular scaffold by converting all atom types into carbon atoms and all bonds into single bonds. Note that in this work we use the molecular scaffold since it is less general and has demonstrated good performance in earlier works [Blaschke et al., 2020b; Guo and Schwaller, 2024; Thomas et al., 2022b]. The Topological scaffold is therefore exclusively applied to assess the molecules' diversity and is not incorporated into any penalty or intrinsic reward method defined hereafter.

Let us define the reward function ˆRIMS(A) for the IMS penalty. For each generated molecule A with a reward of at least h, its molecular scaffold SA is computed and put in memory. A molecule fulfilling the (extrinsic) reward threshold h is commonly known as a predicted active molecule, denoted simply as active. If mmolecules with the same scaffold have been generated, future molecules of the same scaffold are given a reward of 0 to avoid this scaffold. Given a generated molecule A with an extrinsic reward of at least h and its corresponding molecular scaffold SA, the reward function of the IMS penalty method is then defined by

where SA is the molecular scaffold of molecule A and N[S] is the number of molecules with molecular scaffold S in memory, i.e., with an extrinsic reward of at least h. If a molecule A corresponds to an extrinsic reward smaller than h, the extrinsic reward is provided to the agent without any modification. Hence, only predicted active molecules are penalized.

Error Function Identical Molecular Scaffold Penalty (ErfIMS). The Error Function Identical molecular scaffold Penalty is a soft (non-binary) version of the IMS penalty method. It uses the error function to incrementally decrease extrinsic rewards based on the number of molecules in the molecular scaffold

where erf(x) = 2 √π ∫ x 0 e-t2 dt is the error function. Given the threshold h for the extrinsic reward R(·), the reward function for a molecule A is defined by

Linear Identical molecular scaffold Penalty (LinIMS). The linear identical molecular scaffold penalty linearly reduces the extrinsic score based on the number of generated molecules in memory with the same molecular scaffold. We define the linear penalty function by

The reward function of a molecule A is defined by, given the threshold h for the extrinsic reward R(·),

where [·]+ denotes the positive part of a function. This is equivalent to the linear penalty proposed by [Blaschke et al., 2020b].

Sigmoid Identical Molecular Scaffold Penalty (SigIMS). The sigmoid identical molecular scaffold penalty uses a sigmoid function to gradually reduce the extrinsic reward based on the number of molecules in memory with the same scaffold. A molecule is put in memory if it has an extrinsic reward of at least h. Given a molecule A, the sigmoid function in this work is defined as

This is equivalent to the sigmoid penalty function proposed by [Blaschke et al., 2020b]; therefore, we use the same parameters. Given a molecule A, we define the reward function as follows

Tanh Identical Molecular Scaffold Penalty (TanhIMS). The tanh identical molecular scaffold penalty utilizes the hyperbolic tangent function to incrementally decrease the extrinsic reward based on the number of molecules generated with the same molecular scaffold, up to and including the current step (i.e., those stored in memory). For a molecule A, the following hyperbolic tangent function is used to incrementally penalize the extrinsic reward

In the following experiments we use ctanh = 3 since this factor implies ftanh ≈ 0 around N[SA] = 25 for a bucket size m = 25. For a molecule A, we define the reward function for the TanhIMS penalty as follows

3.2 Intrinsic Reward

We examine eight methods to provide intrinsic reward to the agent, namely diverse actives (DA), minimum distance (MinDis), mean distance (MeanDis), minimum distance to random coreset (MinDisR), mean distance to random coreset (MeanDisR), KL-UCB, random network distillation (RND) and information (Inf). To the best of our knowledge, all methods, except for RND, are novel in the context of de novo drug design.

Diverse Actives (DA). We define the diverse actives intrinsic reward based on the diverse hits metric by [Renz et al., 2024], which is based on the #Circles metric proposed by [Xie et al., 2023]. Given a set of possible centers, the #Circles metric counts the number of non-overlapping circles with equivalent radius in the distance metric space. An active molecule is defined as a molecule with a reward of at least h. Following the work by [Renz et al., 2024] but using the terminology of predicted active molecules rather than hit molecules, we define the number of diverse actives for distance threshold D by

where H is a set of predicted active molecules, P is the power set, d(x, y) is the distance between molecules x and y. Note that there is a substantial difference between a set of actives and a set of diverse actives. Determining the number of diverse actives is analogous to determining the packing number of the set H in the distance metric space [Renz et al., 2024]. Let ∆µ be the difference in the number of diverse actives, between two sets H and ˜H of active molecules, defined by

Moreover, let Hi be the batch of generated actives in the current generative step i, Ci-1 the set of previously generated diverse actives and Ci = Ci-1 ∪ Hi. We define the reward function using diverse actives as an intrinsic reward by

Note that the intrinsic reward ∆µ (Ci-1 ∪Hi, Ci-1;D) is sparse since a new batch does not necessarily increase the number of non-overlapping circles. On the other hand, the intrinsic reward can be substantially larger than the extrinsic reward R(A) ∈ [0, 1], providing strong intrinsic motivation towards a specific area.

Minimum Distance (MinDis). Minimum distance is a distance-based intrinsic reward. A bonus reward is given based on the minimum distance to previously generated diverse actives (see definition of diverse actives above). Let Hi be a batch of generated actives in the current generative step i and Ci-1 be a set of previously generated diverse actives. Then the reward function of MinDis for a molecule A and reward threshold h (of predicted active molecules) is defined as follows

where ˜ Ci-1 := Ci-1 ∪ (Hi \ {A}), and d(x, y) is a distance metric between molecules x and y. In this work, we use the distance metric based on the Jaccard index [Jaccard, 1912], also known as the Tanimoto distance, widely used to measure chemical (dis-)similarity.

Mean Distance (MeanDis). Mean distance is also a distance-based intrinsic reward, but where the intrinsic reward is defined as the mean dissimilarity (distance) to previously generated diverse actives and the current batch of actives. Let Hi be a batch of generated actives in the current generative step i and Ci-1 be a set of previously generated diverse actives (see definition above of diverse actives). We then define the reward function of MeanDis of molecule A by

Minimum Distance to Random Coreset (MinDisR). Minimum distance to random coreset is a distance-based intrinsic reward similar to MinDis. The difference is that MinDisR is based on the distance between the actives from the current generative step and a random set of previously generated actives. Given the set Hi of actives generated in the current generative step i, a molecule A ∈ Hi generated in the current step, and a (uniform) random set X of previously generated actives, the reward function is defined by

where ˜ X := X ∪ (Hi \ {A}) and d(x, y) is the distances between molecules x and y. In this work, X consists of 5000 randomly sampled actives, uniformly sampled without replacement from the set of previously generated actives Hi-1. If 5000 actives have not been generated at generative step i, all previously generated actives are used.

Mean Distance to Random Coreset (MeanDisR). Mean distance to random coreset is an intrinsic reward given by the mean distance to a random coreset of actives. Given a set Hi of actives generated in the current generative step i, a molecule A generated in the current generative step and a random set X of previously generated actives, we define the reward function of MeanDisR as

KL-UCB. The KL-UCB intrinsic reward is based on the KL-UCB algorithm by [Garivier and Capp´e, 2011] for the multi-armed bandit problem. It defines an improved upper confidence bound to handle the trade-off between exploration and exploitation in the multi-armed bandit problem. In our study, this trade-off is crucial as the agent must determine the optimal balance between exploiting and exploring various molecular structures. We compute the KL-UCB intrinsic reward for a molecule A by

where n is the total number of generated molecules up to and including the current generative step i, KL(p, q) = p log pq + (1 - p) log 1-p 1-q is the Bernoulli Kullback-Leibler divergence and c = 0 is used for optimal performance in practice [Garivier and Capp´e, 2011]. Moreover, SA is the scaffold of molecule A, Σ[S] is the sum of rewards of actives with scaffold S in memory and N[S] is the total number of actives with scaffold S in memory. Given a molecule A, extrinsic reward R(A) and reward threshold h, we define the reward function of KL-UCB by

This implies that actives are given a reward corresponding to the upper confidence bound of the mean extrinsic reward of the actives with the same scaffold.

Random Network Distillation (RND). Random network distillation [Burda et al., 2018] is an exploration technique in reinforcement learning that provides an intrinsic reward based on the prediction error of a neural network. Specifically, it employs a fixed, randomly initialized neural network f and a predictive neural network ˆ fϕ trained to mimic the outputs of the fixed network. The intrinsic reward is derived from the prediction error between these two networks. This error serves as a measure of novelty, incentivizing the RL agent to explore less familiar regions of the parameter space, and potentially enhancing the exploration of less familiar regions of the chemical space. We adapt RND as an intrinsic reward for generated active molecules, i.e., molecules with an extrinsic reward of at least h. In this work, f and ˆ fϕ have identical architecture as the pre-trained policy (see Section 4). We let the predictive network ˆ fϕ be initialized to this pre-trained policy.

For a molecule A, we define

where πf(at|st) and πϕ(at|st) are the policies induced by the fixed and predictive network, respectively. Moreover, let ∆ ˆ f (A; ϕ) be the squared norm of the difference between these networks for a molecule A, defined by

where ϕ is the weights of the prediction network ˆ f of the current generative step.

Let ϕ be the weights of the predictive network ˆ f up to the current generative step, we define the reward function of a molecule A by

Since R(A) ∈ [0, 1], we rescale the prediction error over the batch of active molecules generated in the current generative step i by

Information (Inf). We define an information-inspired intrinsic reward function based on the number of actives in each scaffold and scaffolds generated up to and including the current generative step i. Let A be a molecule, SA its scaffold, N[S] the number of active molecules with scaffold S in memory, and S the set of unique molecular scaffolds in memory up to (including) current generative step i. We define the the scaffold (pseudo-)probability of molecule A by

We use this scaffold probability to define the scaffold information by

Given a set of active molecules Hi generated at the current generative step i, where A ∈ Hi, the normalized scaffold information is defined by

In practice, we only normalize if |Hi| > 2. Using the scaffold information to define the information-based intrinsic reward, we define the reward function by

3.3 Combining Penalty and Intrinsic Reward

We investigate two combinations of intrinsic reward and extrinsic reward penalty. To the best of our knowledge, combining these approaches is novel.

Tanh RandomNetwork Distillation (TanhRND). We define a soft version of random network distillation by combining RND and TanhIMS. The extrinsic reward is penalized as defined by TanhIMS and a (non-penalized) intrinsic reward based on RND is provided to the agent. We define the TanhRND reward function by

Tanh Information (TanhInf). We also propose and examine a reward function combining (extrinsic) reward penalty and information-based intrinsic reward. We use TanhIMS to penalize the extrinsic reward and Inf to provide a (nonpenalized) intrinsic reward. For a molecule A, we define the TanhInf reward function by

where h is the reward threshold, R(·) is the extrinsic reward, Hi is the set of active molecules generated in the current generative step.

4 Experimental Evaluation

We now describe experiments designed to examine the efficacy of our diversity-aware reward functions.

4.1 Experimental Setup

We run experiments on two extrinsic reward functions, namely the c-Jun N-terminal Kinases-3 (JNK3) and Glycogen Synthase Kinase 3 Beta (GSK3β) oracles provided by Therapeutics Data Commons [Huang et al., 2021; Huang et al., 2022; Velez-Arce et al., 2024]. These are well-established molecule binary bioactivity label optimization tasks. To compute an extrinsic reward in [0, 1], each oracle utilizes a random forest classifier trained on data from the ExCAPE-DB dataset [Sun et al., 2017] using extended-connectivity fingerprints with radius 3 [Rogers and Hahn, 2010]. These oracles only provide rewards to valid molecules and, therefore, we assign invalid molecules an extrinsic reward of -1 to distinguish them from the penalized molecules. Additionally, previously generated (predicted) active molecules are each assigned zero reward.

For distance-based intrinsic rewards, the Jaccard distance is computed based on Morgan fingerprints [Rogers and Hahn, 2010] computed by RDKit [Landrum, 2006], with a radius of 2 and a size of 2048 bits. The distance threshold for the diverse actives-based approaches is fixed to D = 0.7, as suggested by [Renz et al., 2024] since there is a significant decrease in the probability of similar bioactives beyond this threshold [Jasial et al., 2016]. Moreover, scaffoldbased diversity-aware reward functions (see Section 3) utilize molecular scaffolds and a bucket size of m = 25. The reward that the agents see is only modified for molecules with an extrinsic reward of at least 0.5, i.e., we define predicted active molecules as molecules reaching an (extrinsic) reward of at least h = 0.5.

The molecular generative model builds directly on REINVENT [Segler et al., 2018; Olivecrona et al., 2017; Blaschke et al., 2020a; Loeffler et al., 2024] and consists of a long short-term memory (LSTM) network [Hochreiter, 1997] using SMILES to represent molecules as text strings. REINVENT utilizes an on-policy RL algorithm optimizing the policy πθ to generate molecules with higher reward. The algorithm is based on the augmented log-likelihood defined by

where A = a1:T-2 is a generated molecule, σ is a scalar value, πθprior is the (fixed) prior policy. We use the pre-trained policy by [Blaschke et al., 2020a] as the prior policy. It is pre-trained on the ChEMBL database [Gaulton et al., 2017] to generate drug-like bioactive molecules. The action space A consists of 34 tokens including start and stop tokens, i.e., |A| = 34. The policy πθ is optimized by minimizing the squared difference between the augmented log-likelihood and policy likelihood given a sampled batch B of SMILES

Previous work has shown that minimizing this loss function is equivalent to maximizing the expected return, as for policy gradient algorithms [Guo and Schwaller, 2024]. Evaluations by both [Gao et al., 2022] and [Thomas et al., 2022a] have concluded good performance compared to both RL-based and non-RL-based approaches for de novo drug design. The generative process has a budget of I = 2000 generative steps, where a batch of |B| = 128 molecules is generated in each step. Each experiment is evaluated by 20 independent runs of the RL fine-tuning process.

4.2 Comparison of Diversity-Aware Reward Functions

For both the GSK3β- and JNK3-based extrinsic reward, we evaluate the quality by extrinsic reward per generative step, and diversity by the total number of molecular scaffolds, topological scaffolds and diverse actives after I = 2000 generative steps.

Glycogen Synthase Kinase 3 Beta (GSK3β). For the GSK3β oracle, Figure 2a compares the average extrinsic reward (moving average using a window size of 101), over 20 independent runs, per generative step. We observe that the extrinsic rewards converge to comparable values across the diversity-aware reward functions. To evaluate the diversity, Figures 3a to 3c display boxplots comparing the number

(b) JNK3

Figure 2: Average extrinsic reward (moving average using a window size of 101), over 20 independent runs, per generative step for both the GSK3β and the JNK3 oracle.

of molecular scaffolds, topological scaffolds and diverse actives, respectively, over 20 independent runs with a budget of 2000 generative steps. Only active molecules with an extrinsic reward of at least h = 0.5 are displayed, where each orange line and green triangle display the median and mean, respectively, across 20 reruns. Inf and TanhInf generate substantially more molecular and topological scaffolds, whereas TanhRND is also among the top methods. Concerning molecular scaffolds, we observe that the 7 most promising methods have a lower variability,i.e., interquartile range, among the 20 reruns, while the other methods often show a higher variability. Hence, the most promising methods can consistently generate a large number of molecular scaffolds. In terms of topological scaffolds, TanhRND exhibits both low variability and the largest number of topological scaffolds, together with TanhInf and Inf.

MinDis and TanhRND generate more diverse actives than the other methods we have investigated, but MinDis shows a higher variability among the reruns. Overall, we observe a higher variability for the number of topological scaffolds and diverse actives among the reruns. This can partially be explained by the fact that the scaffolds-based methods only consider molecular scaffolds, but more molecular scaffolds

(c) Diverse actives do not necessarily lead to more topological scaffolds and diverse actives. However, as we observe here, if utilized appropriately, it can still enhance the overall diversity.

Figure 3: Evaluation of diversity on the GSK3β oracle. The boxplots compare the total number of molecular scaffolds, topological scaffolds and diverse actives after 2000 generative steps. Each boxplot displays the evaluation across 20 independent runs of each method, where the orange line and green triangle display the median and mean, respectively. Only active molecules are considered when computing the diversity metrics.

In general, applying RND alone does not yield among the highest diversities, but the combination of TanhIMS and RND can improve the diversity of TanhIMS, especially for topological and diverse actives. Hence, we observe that proper structure-based scaling of the extrinsic reward can help to improve diversity, while adding a prediction-based intrinsic reward, i.e., RND, can further enhance and stabilize diversity. Combining Inf, a structure-based intrinsic reward, with TanhIMS does not seem to enhance the diversity further, but it also does not reduce the diversity. We also observe enhanced performance in scaffold diversity for the non-binary penalty functions alone. Still, this structural information does not improve the similarity-based diversity, i.e., diverse actives.

c-Jun N-terminal Kinases-3 (JNK3). Figure 2b displays the average extrinsic reward (moving average using a window size of 101), over 20 independent runs, per generative step for the JNK3 oracle. We observe a substantially lower extrinsic reward when not using any intrinsic motivation. The extrinsic rewards are significantly lower compared to the GSK3β experiments, highlighting that this is a more difficult optimization problem where the agent is more prone to get stuck in a local optimum if not sufficient exploration is carried out. Note that the average extrinsic rewards of the penalty-based methods do not reach the threshold of h = 0.5, where extrinsic rewards are penalized. However, by combining the penalty with an intrinsic reward, it is possible to escape lowrewarding local optima.

We will now also see how the diversity of the generation process has been affected by the different diversity-aware reward functions. Figures 4a to 4c show boxplots displaying the total number of molecular scaffolds, topological scaffolds, and diverse actives, respectively, after 2000 generative steps of 20 reruns. Only active molecules are considered, where each orange line and green triangle displays the median and mean, respectively, across 20 reruns. Generally, substantially fewer molecular scaffolds, topological scaffolds, and diverse actives are found compared to the experiments on the GSK3β extrinsic reward function (see Figures 3a to 3c). This is because the JNK3-based extrinsic reward is more difficult to optimize (see Figures 2a and 2b) and, therefore, the number of actives is lower. In terms of molecular and topological scaffolds (see Figures 4a and 4b), it is evident that the methods that only penalize the extrinsic reward generate less diverse molecules. We observe that TanhInf, TanhRND and Inf generate substantially more molecular and topological scaffolds while their boxes (in the boxplots) do not intersect with the other methods. However, for TanhInf and Inf, we observe a higher variability in terms of topological scaffolds, while TanhRND demonstrates a more consistent number of topological scaffolds among the 20 reruns.

The differences between the methods are not as apparent in terms of diverse actives, where most methods generate around 20 diverse actives. Only TanhRND shows the capability to generate more than 30 diverse actives. We notice that TanhRND has a higher variability in terms of diverse actives,

(c) Diverse actives but its box does not intersect with the other methods. The higher variability of TanhRND is possibly because it is difficult to improve the number of diverse actives and, therefore, there is a natural variability in how much it can be improved in each rerun.

Figure 4: Evaluation of diversity on the JNK3 oracle. The boxplots compare the total number of molecular scaffolds, topological scaffolds and diverse actives after 2000 generative steps. Each boxplot displays the evaluation across 20 independent runs of each method, where the orange line and green triangle display the median and mean, respectively. Only active molecules are considered when computing the diversity metrics.

Overall, TanhRND shows a high diversity, with only TanhInf and Inf showing significantly more molecular scaffolds. On the other hand, for RND and TanhIMS on their own, we observe no substantial improvement in diversity compared to the other methods. This is likely because JNK3 is more difficult to optimize (see Figure 2b) due to having fewer optima compared to GSK3β. Hence, we observe that the combination of RND and TanhIMS considerably improves the overall diversity for a more sparse optimization task.

5 Conclusion

Our comprehensive study proposes and evaluates several novel intrinsic rewards and reward penalties to enhance the diversity of de novo drug design using reinforcement learning (RL). Our approach balances exploration and exploitation to promote a more diverse generation of molecules. Our results consistently show that methods incorporating both intrinsic reward and reward penalty generate significantly more diverse actives, molecular scaffolds, and topological scaffolds. In particular, using structure-based scaling of the extrinsic reward and a prediction-based method to encourage exploration in the agent's state space, we observe improved diversity in terms of both structure and similarity. Yet, a single type of information does not fully enhance diversity. Our work opens several future directions for studying how to incorporate domain and agent information into the reward signal efficiently.

Acknowledgments

This work was partially supported by the Wallenberg AI, Autonomous Systems and Software Program (WASP) funded by the Knut and Alice Wallenberg Foundation. The experimental evaluation was enabled by resources provided by the National Academic Infrastructure for Supercomputing in Sweden (NAISS), partially funded by the Swedish Research Council through grant agreement no. 2022-06725.