Reinforcement Learning for Molecular Design Guided by Quantum Mechanics

Introduction

Finding new chemical compounds with desired properties is a challenging task with important applications such as de novo drug design and materials discovery (Schneider et al. 2019). The diversity of synthetically feasible chemicals that can be considered as potential drug-like molecules was estimated to be between $10^{30}$ and $10^{60}$ (Polishchuk et al. 2013), making exhaustive search hopeless.

Recent applications of machine learning have accelerated the search for new molecules with specific desired properties. Generative models such as variational autoencoders (VAEs) (Gómez-Bombarelli et al. 2018), recurrent neural networks (RNNs) (Segler et al. 2018), and generative adversarial networks (GANs) (De Cao & Kipf 2018) have been successfully applied to propose potential drug candidates. Despite recent advances in generating valid structures, proposing truly novel molecules beyond the training data distribution remains a challenging task. This issue is exacerbated for many classes of molecules (e.g. transition metals), where such a representative dataset is not even available.

An alternative strategy is to employ RL, in which an agent builds new molecules in a step-wise fashion (e.g., Olivecrona et al. 2017, Guimaraes et al. 2018, Zhou et al. 2019, Zhavoronkov et al. 2019). Training an RL agent only requires samples from a reward function, alleviating the need for an existing dataset of molecules. However, the choice of state representation in current models still severely limits the class of molecules that can be generated. In particular, molecules are commonly described by graphs, where atoms and bonds are represented by nodes and edges, respectively. Since a graph is a simplified model of the physical representation of molecules in the real world, one is limited to the generation of single organic molecules. Other types of molecules cannot be appropriately described as this representation lacks important three-dimensional (3D) information, i.e. the relative position of atoms in space. For example, systems consisting of multiple molecules cannot be generated for this reason. Furthermore, it prohibits the use of reward functions based on fundamental physical laws; instead, one has to resort to heuristic physicochemical parameters, e.g. the Wildman-Crippen partition coefficient (Wildman & Crippen 1999). Lastly, it is not possible to impose geometric constraints on the design process, e.g. those given by the binding pocket of a protein which the generated molecule is supposed to target.

In this work, we introduce a novel RL formulation for molecular design in which an agent places atoms from a given bag of atoms onto a 3D canvas (see Fig. 1). As the reward function is based on fundamental physical properties such as energy, this formulation is not restricted to the generation of molecules of a particular type. We thus encourage the agent to implicitly learn the laws of atomic interaction from scratch to build molecules that go beyond what can be represented with graph-based RL methods. To enable progress towards designing such molecules, we introduce a new RL environment called MolGym. It comprises a suite of tasks in which both single molecules and molecule clusters need to be constructed. For all of these tasks, we provide baselines using quantum-chemical calculations. Finally, we propose a novel policy network architecture that can efficiently learn to solve these tasks by working in a translation and rotation invariant state-action space.

In summary, our contributions are as follows:

we propose a novel RL formulation for general molecular design in Cartesian coordinates (Section 2.2);

we design a reward function based on the electronic energy, which we approximate via fast quantum-chemical calculations (Section 2.3);

we present a translation and rotation invariant policy network architecture for molecular design (Section 3);

we introduce MolGym, an RL environment comprising several molecular design tasks along with baselines based on quantum-chemical calculations (Section 5.1);

we perform experiments to evaluate the performance of our proposed policy network using standard policy gradient methods (Section 5.2).

Reinforcement Learning for Molecular Design Guided by Quantum Mechanics

In this section, we provide a brief introduction to RL and present our novel RL formulation for molecular design in Cartesian coordinates.

Policy Gradient Algorithms Policy gradient methods are well-suited for RL in continuous action spaces. These methods learn a parametrized policy $\pi_{\theta}$ by performing gradient ascent in order to maximize $J(\theta)$. More recent algorithms (Schulman et al. 2015; Schulman et al. 2017) improve the stability during learning by constraining the policy updates. For example, proximal policy optimization (PPO) (Schulman et al. 2017) employs a clipped surrogate objective. Denoting the probability ratio between the updated and the old policy as $r_{t}(\theta)=\frac{\pi_{\theta}(a_{t}|s_{t})}{\pi_{\theta_{\text{old}}}(a_{t}|s_{t})}$, the clipped objective $J^{\text{CL}}$ is given by

where $\hat{A}_{t}$ is an estimator of the advantage function, and $\epsilon$ is a hyperparameter that controls the interval beyond which $r(\theta)$ gets clipped. To further reduce the variance of the gradient estimator, actor-critic approaches (Konda & Tsitsiklis 2000) are often employed. The idea is to use the value function (i.e. the critic) to assist learning the policy (i.e. the actor). If the actor and critic share parameters, the objective becomes

2 Setup

We design molecules by sequentially drawing atoms from a given bag and placing them onto a 3D canvas. This task can be formulated as a sequential decision-making problem in an MDP with deterministic transition dynamics, where

deterministic transition function $\mathcal{T}(s_{t},a_{t})$ places an atom through action $a_{t}$ in state $s_{t}$, returning the next state $s_{t+1}=(\mathcal{C}_{t+1},\beta_{t+1})$ with $\beta_{t+1}=\beta_{t}\backslash e_{t}$;

reward function $r(s_{t},a_{t})$ quantifies how applying action $a_{t}$ in state $s_{t}$ alters properties of the molecule, e.g. the stability of the molecule as measured in terms of its quantum-chemical energy.

3 Reward Function

where $E(e)\coloneqq E(\{(e,^{\mkern-1.5mu\mathsf{T}}\})$. Intuitively, the agent is rewarded for placing atoms so that the energy of the resulting molecules is low. Importantly, with this formulation the undiscounted return for building a molecule is independent of the order in which atoms are placed. If the reward only consisted of $E(C_{t+1})$, one would double-count interatomic interactions. As a result, the formulation in Eq. (1) prevents the agent from learning to greedily choose atoms of high atomic number first, as they have low intrinsic energy.

Quantum-chemical methods, such as the ones based on density functional theory (DFT), can be employed to compute the energy $E$ for a given $\mathcal{C}$. Since such methods are computationally demanding in general, we instead choose to evaluate the energy using the semi-empirical Parametrized Method 6 (PM6) (Stewart 2007) as implemented in the software package Sparrow (Husch et al. 2018; Bosia et al. 2019); see the Appendix for details. PM6 is significantly faster than more accurate methods based on DFT and sufficiently accurate for the scope of this study. For example, the energy $E$ of systems containing 10 atoms can be computed within hundreds of milliseconds with PM6; with DFT, this would take minutes. We note that more accurate methods can be used as well if the computational budget is available.

Policy

Building molecules in Cartesian coordinates allows to 1) extend molecular design through deep RL to a much broader class of molecules compared to graph-based approaches, and 2) employ reward functions based on fundamental physical properties such as the energy. However, working directly in Cartesian coordinates introduces several additional challenges for policy learning.

Firstly, it would be highly inefficient to naively learn to place atoms directly in Cartesian coordinates since molecular properties are invariant under symmetry operations such as translation and rotation. For instance, the energy of a molecule—and thus the reward—does not change if the molecule gets rotated, yet an agent that is not taking this into account would need to learn these solutions separately. Therefore, we require an agent that is covariant to translation and rotation, i.e., if the canvas is rotated or translated, the position $x_{t}$ of the atom to be placed should be rotated and translated as well. To achieve this, our agent first models the atom’s position in internal coordinates which are invariant under translation and rotation. Then, by mapping from internal to Cartesian coordinates, we obtain a position $x_{t}$ that features the required covariance. The agent’s internal representations for states and actions are introduced in Sections 3.1 and 3.2, respectively.

Secondly, the action space contains both discrete (i.e. element $e_{t}$) and continuous actions (i.e. position $x_{t}$). This is in contrast to most RL algorithms, which assume that the action space is either discrete or continuous. Due to the continuous actions, policy exploration becomes much more challenging compared to graph-based approaches. Further, not all discrete actions are valid in every state, e.g. the element $e_{t}$ has to be contained in the bag $\beta_{t}$. These issues are addressed in Section 3.2, where we propose a novel actor-critic neural network architecture for efficiently constructing molecules in Cartesian coordinates.

where $\text{MLP}_{\beta}$ is a multi-layer perceptron (MLP).

2 Actor

$\alpha\in[0,\pi]$ is the angle between the two lines defined by $(x,x_{f})$ and $(x,x_{n1})$, where $x_{n1}$ is the position of the atom closest to $f$; if less than two atoms are on the canvas, $\alpha$ is undefined/unused.

$\psi\in[-\pi,\pi]$ is the dihedral angle between two intersecting planes spanned by ($x$, $x_{f}$, $x_{n1}$) and ($x_{f}$, $x_{n1}$, $x_{n2}$), where $x_{n2}$ is the atom that is the second In the unlikely event that two atoms are exactly equally far from the focal atom, a random order for $x_{n1}$ and $x_{n2}$ is chosen. closest to the focal atom; if less than three atoms are on the canvas, $\psi$ is undefined/unused.

As shown in Fig. 3 (right), these internal coordinates can then be mapped back to Cartesian coordinates $x$.

Model This action representation suggests a natural generative process: first choose next to which focal atom the new atom is placed, then select its element, and finally decide where to place the atom relative to the focal atom. Therefore, we assume that the policy factorizes as

Maintaining Valid Actions As the agent places atoms onto the canvas during a rollout, the number of possible focal atoms $f$ increases and the number of elements $e$ to choose from decreases. To guarantee that the agent only chooses valid actions, i.e. $f\in\{1,\dots,n(\mathcal{C})\}$ and $e\in\beta$, we mask out invalid focal atoms and elements by setting their probabilities to zero and re-normalizing the categorical distributions. Neither the agent nor the environment makes use of ad-hoc concepts like valence or bond connectivity—any atom on the canvas can potentially be chosen.

3 Critic

where $\text{MLP}_{\phi}$ is an MLP that computes value $V$ (see Fig. 4).

4 Optimization

We employ PPO (Schulman et al. 2017) to learn the parameters $(\theta,\phi)$ of the actor $\pi_{\theta}$ and critic $V_{\phi}$, respectively. While most RL algorithms can only deal with either continuous or discrete action spaces and thus require additional modifications to handle both (Masson et al. 2016; Wei et al. 2018; Xiong et al. 2018), PPO can be applied directly as is. To help maintain sufficient exploration throughout learning, we include an entropy regularization term over the policy. However, note that the entropies of the continuous and categorical distributions often have different magnitudes; further, in this setting the entropies over the categorical distributions vary significantly throughout a rollout: as the agent places more atoms, the support of the distribution over valid focal atoms $f$ increases and the support of the distribution over valid elements $e$ decreases. To mitigate this issue, we only apply entropy regularization to the categorical distributions, which we find to be sufficient in practice.

Related Work

Deep Generative Models A prevalent strategy for molecular design based on machine learning is to employ deep generative models. These approaches first learn a latent representation of the molecules and then perform a search in latent space (e.g., through gradient descent) to discover new molecules with sought chemical properties. For example, Gómez-Bombarelli et al. 2018; Kusner et al. 2017; Blaschke et al. 2018; Lim et al. 2018; Dai et al. 2018 utilized VAEs to perform search or optimization in a latent space to find new molecules. Segler et al. 2018 used RNNs to design molecular libraries. The aforementioned approaches generate SMILES strings, a linear string notation, to describe molecules (Weininger 1988). Further, there exist a plethora of generative models that work with graph representations of molecules (e.g., Jin et al. 2017; Bradshaw et al. 2019a; Li et al. 2018a; Li et al. 2018b; Liu et al. 2018; De Cao & Kipf 2018; Bradshaw et al. 2019b). In these methods, atoms and bonds are represented by nodes and edges, respectively. Brown et al. 2019 developed a benchmark suite for graph-based generative models, showing that generative models outperform classical approaches for molecular design. While the generated molecules are shown to be valid (De Cao & Kipf 2018; Liu et al. 2018) and synthesizable (Bradshaw et al. 2019b), the generative model is restricted to a (small) region of chemical space for which the graph representation is valid, e.g. single organic molecules.

3D Point Cloud Generation Another downside of string- and graph-based approaches is their neglect of information encoded in the interatomic distances. To this end, Gebauer et al. 2018; Gebauer et al. 2019 proposed a generative neural network for sequentially placing atoms in Cartesian coordinates. While their model respects local symmetries by construction, atoms are placed on a 3D grid. Further, similar to aforementioned approaches, this model depends on a dataset to exist that covers the particular class of molecules for which one seeks to generate new molecules.

Reinforcement Learning Olivecrona et al. 2017, Guimaraes et al. 2018, Putin et al. 2018, Neil et al. 2018 and Popova et al. 2018 presented RL approaches based on string representations of molecules. They successfully generated molecules with given desirable properties but, similar to other generative models using SMILES strings, struggled with chemical validity. You et al. 2018 proposed a graph convolutional policy network based on graph representations of molecules, where the reward function is based on empirical properties such as the drug-likeliness. While this approach was able to consistently produce valid molecules, its performance still depends on a dataset required for pre-training. Considering the large diversity of chemical structures, the generation of a dataset that covers the whole chemical space is hopeless. To address this limitation, Zhou et al. 2019 proposed an agent that learned to generate molecules from scratch using a Deep Q-Network (DQN) (Mnih et al. 2015). However, such graph-based RL approaches are still restricted to the generation of single organic molecules for which this representation was originally designed. Further, graph representations prohibit the use of reward functions based on fundamental physical laws, and one has to resort to heuristics instead. Finally, geometric constraints cannot be imposed on the design process. Jørgensen et al. 2019 introduced an atomistic structure learning algorithm, called ALSA, that utilizes a convolutional neural network to build 2D structures and planar compounds atom by atom.

Experiments

We perform experiments to evaluate the performance of the policy introduced in Section 3. While prior work has focused on building molecules using molecular graph representations, we are interested in designing molecules in Cartesian coordinates. To this end, we introduce a new RL environment called MolGym in Section 5.1. It comprises a set of molecular design tasks, for which we provide baselines using quantum-chemical calculations. See the Appendix for details on how the baselines are determined. Source code of the agent and environment is available at https://github.com/gncs/molgym.

We use MolGym to answer the following questions: 1) can our agent learn to construct single molecules in Cartesian coordinates from scratch, 2) does our approach allow building molecules across multiple bags simultaneously, 3) are we able to scale to larger molecules, and 4) can our agent construct systems comprising multiple molecules?

We propose three different tasks for molecular design in Cartesian coordinates, which are instances of the MDP formulation introduced in Section 2.2: single-bag, multi-bag, and solvation. More formally, the tasks are as follows:

Single-bag Given a bag, learn to design stable molecules. This task assesses an agent’s ability to build single stable molecules. The reward function is given by $r(s_{t},a_{t})=-\Delta_{E}(s_{t},a_{t})$, see Eq. (1). If the reward is below a threshold of $-0.6$, the molecule is deemed invalid and the episode terminates prematurely with the reward clipped at $-0.6$. $\Delta_{E}$ is on the order of magnitude of $-0.1$ Hartree, resulting in a reward of around $0.25$ for a well placed atom.

Multi-bag Given multiple bags with one of them being randomly selected before each episode, learn to design stable molecules. This task focuses on the agent’s capabilities to learn to build different molecules of different composition and size at the same time. The same reward function as in the single-bag task is used. Offline performance is evaluated in terms of the average return across bags. Similarly, the baseline is given by the average optimal return over all bags.

2 Results

In this section, we use the tasks specified in Section 5.1 to evaluate our proposed policy. We further assess the chemical validity, diversity and stability of the generated structures. Experiments were run on a 16-core Intel Xeon Skylake 6142 CPU with 2.6GHz and 96GB RAM. Details on the model architecture and hyperparameters are in the Appendix.

Learning across Multiple Bags We train the agent on the multi-bag task using all formulas contained in the QM9 dataset (Ruddigkeit et al. 2012; Ramakrishnan et al. 2014) with up to $5$ atoms, resulting in $11$ bags (see Table 1). Despite their small size, the molecules feature a diverse set of bonds (single, double, and triple) and geometries (linear, trigonal planar, and tetrahedral). From the performance and from visual inspection of the generated molecular structures shown in Fig. 6, it can be seen that a single policy is able to build different molecular structures across multiple bags. For example, it learned that a carbon atom can have varying number and type of neighboring atoms leading to specific bond distance, angles, and dihedral angles.

Quality Assessment of Generated Molecules In the spirit of the GuacaMol benchmark (Brown et al. 2019), we assess the molecular structures generated by the agent with respect to chemical validity, diversity and structural stability for each experiment. To enable a comparison with existing approaches, we additionally ran experiments with the bag $\text{C}{\vphantom{\text{X}}}_{\smash[t]{\text{7}}}\text{H}{\vphantom{\text{X}}}_{\smash[t]{\text{8}}}\text{N}{\vphantom{\text{X}}}_{\smash[t]{\text{2}}}\text{O}{\vphantom{\text{X}}}_{\smash[t]{\text{2}}}$, the stoichiometry of which is taken from the GuacaMol benchmark (Brown et al. 2019).

The results are shown in Table 2. To determine the validity and stability of the generated structures, we first took the terminal states of the last iteration for a particular experiment. Structures are considered valid if they can be successfully parsed by RDKit (Landrum 2019). However, those consisting of multiple molecules were not considered valid (except in the solvation task). The validity reported in Table 2 is the ratio of valid molecules over 10 seeds.

All valid generated structures underwent a structure optimization using the PM6 method (see Appendix for more details). Then, the RMSD (in Å) between the original and the optimized structure were computed. In Table 2, the median RMSD over all generated structures is given per experiment. In the approach by Gebauer et al. 2019, an average RMSD of $\approx$ 0.25 Å is reported. Due to significant differences in approach, application, and training procedure we forego a direct comparison of the methods.

Further, two molecules are considered identical if the SMILES strings generated by RDKit are the same. The diversity reported in Table 2 is the total number of unique and valid structures generated through training over 10 seeds.

Discussion

This work is a first step towards general molecular design through RL in Cartesian coordinates. One limitation of the current formulation is that we need to provide bags for which we know good solutions exist when placed completely. While being able to provide such prior knowledge can be beneficial, we are currently restricted to designing molecules of known formulas. A possible solution is to provide bags that are larger than necessary, e.g. generated randomly or according to some fixed budget for each element, and enable the agent to stop before the bag is empty.

Compared to graph-based approaches, constructing molecules by sequentially placing atoms in Cartesian coordinates greatly increases the flexibility in terms of the type of molecular structures that can be built. However, it also makes the exploration problem more challenging: whereas in graph-based approaches a molecule can be expanded by adding a node and an edge, here, the agent has to learn to precisely position an atom in Cartesian coordinates from scratch. As a result, the molecules we generate are still considerably smaller. Several approaches exist to mitigate the exploration problem and improve scalability, including: 1) hierarchical RL, where molecular fragments or entire molecules are used as high-level actions; 2) imitation learning, in which known molecules are converted into expert trajectories; and 3) curriculum learning, where the complexity of the molecules to be built increases over time.

Conclusion

We have presented a novel RL formulation for molecular design in Cartesian coordinates, in which the reward function is based on quantum-mechanical properties such as the energy. We further proposed an actor-critic neural network architecture based on a translation and rotation invariant state-action representation. Finally, we demonstrated that our model can efficiently solve a range of molecular design tasks from our MolGym RL environment from scratch.

In future work, we plan to increase the scalability of our approach and enable the agent to stop before a given bag is empty. Moreover, we are interested in combining the reward with other properties such as drug-likeliness and applying our approach to other classes of molecules, e.g. transition-metal catalysts.

Acknowledgements

We would like to thank the anonymous reviewers for their valuable feedback. We further thank Austin Tripp and Vincent Stimper for useful discussions and feedback. GNCS acknowledges funding through an Early Postdoc.Mobility fellowship by the Swiss National Science Foundation (P2EZP2_181616). RP receives funding from iCASE grant #1950384 with support from Nokia.

References

Appendix A Quantum-Chemical Calculations

For the calculation of the energy $E$ we use the fast semi-empirical Parametrized Method 6 (PM6) (Stewart 2007). In particular, we use the implementation from the software package Sparrow (Husch et al. 2018; Bosia et al. 2019). For each calculation, a molecular charge of zero and the lowest possible spin multiplicity are chosen. All calculations are spin-unrestricted.

Limitations of semi-empirical methods are highlighted in, for example, recent work by Husch & Reiher 2018. More accurate methods such as approximate density functionals need to be employed especially for systems containing transition metals.

For the quantum-chemical calculations to converge reliably, we ensured that atoms are not placed too close ($<$ 0.6 Å) nor too far away from each other ($>$ 2.0 Å). If the agent places an atom outside these boundaries, the minimum reward of $-0.6$ is awarded and the episode terminates.

Appendix B Learning the Dihedral Angle

We experimentally validate the benefits of learning $|\psi|\in[0,\pi]$ and $\kappa\in\{-1,1\}$ instead of $\psi\in[-\pi,\pi]$ by comparing the two models on the single-bag task with bag $\text{CH}{\vphantom{\text{X}}}_{\smash[t]{\text{4}}}$ (methane). Methane is one of the simplest molecules that requires the model to learn a dihedral angle. As shown in Fig. 9, learning the sign of the dihedral angle separately (with $\kappa$) speeds up learning significantly. In fact, the ablated model (without $\kappa$) fails to converge to the optimal return even after $100\,000$ steps (not shown).

Appendix C Experimental Details

The model architecture is summarized in Table 3. We initialize the biases of each MLP with $0$ and each weight matrix as a (semi-)orthogonal matrix. After each hidden layer, a ReLU non-linearity is used. The output activations are shown in Table 3. As explained in the main text, both $\text{MLP}_{f}$ and $\text{MLP}_{e}$ use a masked softmax activation function to guarantee that only valid actions are chosen. Further, we rescale the continuous actions $(\mu_{d},\mu_{\alpha},\mu_{\psi})\in^{3}$ predicted by $\text{MLP}_{\text{cont}}$ to ensure that $\mu_{d}\in[d_{\text{min}},d_{\text{max}}]$, $\mu_{\alpha}\in[0,\pi]$ and $\mu_{\psi}\in[0,\pi]$. For more details on the SchNet, see the original work (Schütt et al. 2018b).

C.2 Hyperparameters

We manually performed an initial hyperparameter search on a single holdout validation seed. The considered hyperparameters and the selected values are listed in Table 4 (single-bag), Table 5 (multi-bag) and Table 6 (solvation). The hyperparameters used for SchNet are shown in Table 7.

Appendix D Baselines

Below, we report how the baselines for the single-bag and multi-bag tasks were derived. First, we took all molecular structures for a given chemical formula (i.e. bag) from the QM9 dataset (Ruddigkeit et al. 2012; Ramakrishnan et al. 2014). Subsequently, we performed a structure optimization using the PM6 method (as described in Section A) on the structures. This was necessary as the structures in this dataset were optimized with a different quantum-chemical method. Then, the most stable structure was selected and considered optimal for this chemical formula; the remaining structures were discarded. Since the undiscounted return is path independent, we determined the return $R(s)$ by computing the total interaction energy in the canvas $\mathcal{C}$, i.e.

where $N$ is the number of atoms placed on the canvas.

The baseline for the solvation task was determined in the following way. 12 molecular clusters were generated by randomly placing $n$ $\text{H}{\vphantom{\text{X}}}_{\smash[t]{\text{2}}}\text{O}$ molecules around the solute molecule (in the main text $n=5$). Subsequently, the structure of these clusters was optimized with the PM6 method (as described in Section A). Similar to Eq. (9), the undiscounted return of each cluster can be computed:

where the distance penalty $\rho=0.01$. Finally, the maximum return over the optimized clusters was determined.

Appendix E Additional Results

In Fig. 10, we show a selection of molecular structures generated by trained models for the bags $\text{C}{\vphantom{\text{X}}}_{\smash[t]{\text{4}}}\text{H}{\vphantom{\text{X}}}_{\smash[t]{\text{7}}}\text{N}$ and $\text{C}{\vphantom{\text{X}}}_{\smash[t]{\text{3}}}\text{H}{\vphantom{\text{X}}}_{\smash[t]{\text{8}}}\text{O}$. Further, since the agent is agnostic to the concept of molecular bonds, it is able to build multiple molecules if it results in a higher return. An example of a bimolecular structure generated by a trained model for the bag $\text{C}{\vphantom{\text{X}}}_{\smash[t]{\text{3}}}\text{H}{\vphantom{\text{X}}}_{\smash[t]{\text{8}}}\text{O}$ is shown in Fig. 11. Finally, in Fig. 12, we showcase a set of generated molecular structures that are not chemically valid.

E.2 Solvation Task

In Fig. 13, we report the average offline performances of agents placing 5 $\text{H}{\vphantom{\text{X}}}_{\smash[t]{\text{2}}}\text{O}$ molecules around the solutes (i.e, $\mathcal{C}_{0}$) acetonitrile and ethanol. As can be seen, the agents are able to accurately place water molecules such that they interact with the solute. However, we stress that more accurate quantum-chemical methods for computing the reward are required to describe hydrogen bonds to chemical accuracy.

In Fig. 14, we compare the average offline performance of two agents placing in total 10 H and 5 O atoms around a formaldehyde molecule. One agent is given $5$ $\text{H}{\vphantom{\text{X}}}_{\smash[t]{\text{2}}}\text{O}$ bags consecutively following the protocol of the solvation task as described in the main text, another is given a single $\text{H}{\vphantom{\text{X}}}_{\smash[t]{\text{10}}}\text{O}{\vphantom{\text{X}}}_{\smash[t]{\text{5}}}$ bag. Their average offline performances are shown in Fig. 14 in blue and red, respectively. It can be seen that giving the agent $5$ $\text{H}{\vphantom{\text{X}}}_{\smash[t]{\text{2}}}\text{O}$ bags one at a time instead of a single $\text{H}{\vphantom{\text{X}}}_{\smash[t]{\text{10}}}\text{O}{\vphantom{\text{X}}}_{\smash[t]{\text{5}}}$ bag improves performance.

E.3 Generalization and Transfer Learning