# Generating stable molecules using imitation and reinforcement learning

Søren Ager Meldgaard  InterCat and Department of Physics and Astronomy, Aarhus
University, Denmark.  Jonas Köhler  Freie Universität Berlin, Department of
Mathematics and Computer Science, Berlin, Germany  Henrik Lund Mortensen
InterCat and Department of Physics and Astronomy, Aarhus University, Denmark.
Mads-Peter V. Christiansen  InterCat and Department of Physics and Astronomy,
Aarhus University, Denmark.  Frank Noé  Freie Universität Berlin, Department
of Mathematics and Computer Science, Berlin, Germany  Freie Universität
Berlin, Department of Physics, Berlin, Germany  Rice University, Department of
Chemistry, Houston, TX, USA  Bjørk Hammer  InterCat and Department of Physics
and Astronomy, Aarhus University, Denmark<EMAIL_ADDRESS>
###### Abstract

Chemical space is routinely explored by machine learning methods to discover
interesting molecules, before time-consuming experimental synthesizing is
attempted. However, these methods often rely on a graph representation,
ignoring 3D information necessary for determining the stability of the
molecules. We propose a reinforcement learning approach for generating
molecules in cartesian coordinates allowing for quantum chemical prediction of
the stability. To improve sample-efficiency we learn basic chemical rules from
imitation learning on the GDB-11 database to create an initial model
applicable for all stoichiometries. We then deploy multiple copies of the
model conditioned on a specific stoichiometry in a reinforcement learning
setting. The models correctly identify low energy molecules in the database
and produce novel isomers not found in the training set. Finally, we apply the
model to larger molecules to show how reinforcement learning further refines
the imitation learning model in domains far from the training data.

##  I Introduction

Discovering novel molecules or materials with desirable properties is a
challenging task because of the immense size of chemical compound space.
Further complicating the process is the costly and time-consuming process of
synthesizing and testing proposed structures. Whereas this procedure
historically was driven by a trial-and-error process, the advance of
computational quantum chemical methods allows for initial screening to select
promising molecules for experimental testing. While the computational resource
growth provided a significant speed-up in processing molecules an exhaustive
search of molecular compound space is still infeasible. Instead an automated
search for interesting candidates is desired. Examples of such search methods
include evolutionary algorithms Hartke (1993); Emilie S. Henault (2020),
basin-hopping Wales and Doye (1997) and particle swarm optimization Kennedy
and Eberhart (1995).
More recently, machine learning (ML) enhanced versions of aforementioned
methods Kolsbjerg, Peterson, and Hammer (2018); Meldgaard, Kolsbjerg, and
Hammer (2018); Winter  _et al._ (2019) and regression methods facilitating
speed-up of virtual screening James L. Melville and Hirst (2009); Carpenter
and Huang (2018) has gained considerable interest by making researchers able
to quickly identify attractive candidates for experimental testing. An added
benefit of virtual screening is the creation of numerous databases containing
structures with computed chemical and physical properties Irwin and Shoichet
(2005); Ruddigkeit  _et al._ (2012); Ramakrishnan  _et al._ (2014); Haastrup
_et al._ (2018); Winther  _et al._ (2019) leading to generative models for
discovery of novel molecules and materials Li, Zhang, and Liu (2018); Gupta
_et al._ (2018); Gómez-Bombarelli  _et al._ (2018); Zhavoronkov  _et al._
(2019); Noé  _et al._ (2019); Hoffmann and Noé (2019); Gebauer, Gastegger, and
Schütt (2019); Mansimov  _et al._ (2019); Maziarka  _et al._ (2020); Köhler,
Klein, and Noe (2020); Satorras  _et al._ (2021); Winter, Noé, and Clevert
(2021). Unlike virtual screening where candidates are selected among the
structures in the database, generative models have shown a remarkable ability
to leverage the database to produce new structures with desirable properties.
A notable limitation is a large database and no feedback-loop to improve the
generated molecules beyond what is learned from the database. To remedy this,
reinforcement learning (RL) methods have started to become a competitive
alternative Olivecrona  _et al._ (2017); Putin  _et al._ (2018); Cao and Kipf
(2018); Popova, Isayev, and Tropsha (2018); Zhou  _et al._ (2019); Schreck,
Coley, and Bishop (2019); Simm, Pinsler, and Hernandez-Lobato (2020) to
methods relying on existing databases. RL involves a model that produces
molecules and obtains properties for these molecules, by some external means
other than a database. This provides the basis for the model to learn from the
molecules it produces. Whereas generative models rely on databases for
pretraining, RL is usually done without any prior knowledge leading to initial
inefficiency as basic chemical and physical rules are learned. Instead a
reward must be defined though, so it requires a setting where a meaningful
reward function is available, such as an energy or other things to be
optimized.
In this work we built upon a previous reinforcement learning algorithm called
Atomistic Structure Learning Algorithm (ASLA) Jørgensen  _et al._ (2019) by
incorporating databases into molecular RL to improve sample efficiency while
simultaneously allowing the ML model to learn beyond the knowledge contained
in the database. First, a general purpose model is trained using a database to
create an model applicable for all stoichiometries. Then, a copy of the
general purpose model is created for each stoichiometry of interest. These
models are then further refined in a RL setting in a search for low energy
isomers for the given stoichiometry. Specifically, we utilize a very small
subset of the GDB-11 Fink, Bruggesser, and Reymond (2005) database that
consists of small organic molecules satisfying simple chemical stability rules
and contains up to 11 atoms of C, N, O and F. We demonstrate that our RL model
is able to both replicate structures in the database as well as producing
novel low energy structures. By focusing on low energy structures, we believe
we avoid the general pitfall of generative models that may produce a large
degree of unsynthesizable molecules, as e.g. shown in the work of Gao and
Connor Gao and Coley (2020). Unlike previous approaches using SMILES Weininger
(1988) or graphs, we operate directly in cartesian coordinates thus allowing
for optimization of the potential energy hence easily biasing the search
towards thermodynamically stable structures. To improve sample-efficiency we
take a model-based approach where, similar to Ref. Mortensen  _et al._ , 2020,
the potential energy surface is modeled by a neural network allowing for
approximate but cheap optimization. Finally, we introduce an architecture
based on SchNetSchütt  _et al._ (2018) with a self-attention mechanism Vaswani
_et al._ (2017) allowing the RL model to account for long-range effects. The
paper is constructed as follows: first the RL theory is outlined, followed by
a description of the neural network architecture. Then the database
pretraining and RL phase are described before demonstrating the method on a
subset of GDB-11. Finally, we apply the method to larger molecules outside the
training distribution.

##  II Theory

The objective of RL is to solve a decision problem, i.e. formulate a program
that given an input can decide on the optimal action. Specifically, we refer
to this program as an agent which is given a state, $s$, and must decide on an
action, $a$. In the general case the decision process involves multiple steps
indexed with the subscript $t$ while $T$ is used for the final step. In order
to solve the problem the agent must device a policy, $\pi$, which is a
probability distribution over the possible actions in a given state. The
optimal policy is the policy which maximizes the sum of rewards, where $r(s)$
is the reward given in state $s$. The states, actions and rewards are user-
specified and describes the problem to be solved. Following Ref. Simm,
Pinsler, and Hernandez-Lobato, 2020, the state space is defined as all
(possibly partial) molecules along with a bag of atoms not yet attached to the
molecule. For the reward we seek to minimize the potential energy, i.e. we set
intermediate rewards to zero and assign a final reward based on the potential
energy of the molecule.

$\displaystyle
r(s_{t})=\begin{cases}\max\left(\frac{E_{\text{ref}}-E(s_{T})}{\Delta
E}+1,0\right),&\mbox{if $t=T$}\\\ 0,&\mbox{else}\end{cases}$ (1)

Here $E_{\text{ref}}$ is the lowest energy observed and $\Delta E=10\text{eV}$
is the energy span from $E_{\text{ref}}$ where energy differences are
resolved. Note that $E_{\text{ref}}\leq E(s_{T})$ meaning that $0\leq
r(s_{t})\leq 1$. We scale the reward to stabilize the training as energies of
produced structures may fluctuate substantially during the RL phase. To
maximize the expected sum of rewards we define the $Q$-value,

$\displaystyle Q_{\pi}(s,a)$
$\displaystyle=\mathbb{E}_{\pi}\left[\sum_{k=t+1}^{T}r(s_{k})\middle|s_{t}=s,a_{t}=a\right]$
$\displaystyle=\mathbb{E}_{\pi}[r(s_{T})|s_{t}=s,a_{t}=a],$ (2)

i.e. the expected final reward when taking action $a$ in state $s$ and then
following policy $\pi$. The goal of the agent is to infer the optimal
$Q$-value function

$\displaystyle Q_{*}(s,a)=\max_{\pi}Q_{\pi}(s,a)\;\forall s,a,$ (3)

thereby enabling the agent to perform the best action in every state. To
improve the $Q$-value function we parameterize it by a neural network and
update it as the agent collects new experience. Specifically, we evaluate the
$Q$-values on a voxelated grid, which allows us to update the $Q$-value
towards the highest final reward observed by the agent for a specific state-
action pair (Fig. 1), which for a deterministic problem puts a lower bound on
the optimal $Q$-value Machado  _et al._ (2018). Additionally, this
discretization allows for easy inference of the optimal atom placement by
simply finding the voxel which maximizes the $Q$-value.

Figure 1: Given a state, $s_{t}$ and an action $a_{t}$, the target $Q$-value
is calculated by computing the energy of finished structures involving the
state-action pair. In the above case a hydrogen atom is placed, which has led
to four different structures. The $Q$-value is updated towards the reward for
the lowest energy structure built.

##  III Architecture

The state representation is utilized for calculating the energy of a structure
as well as a state value which is used for calculating the $Q$-value, but is
independent of the specific action. For improved data efficiency the state and
action representation must satisfy symmetries in the Hamiltonian, i.e.
translation, rotation and mirroring of the molecule as well as permutation of
identical atoms. To incorporate these properties in the state representation
we utilize pairwise distances (Fig. 2a, $\mathbf{D}$), expanded in a Gaussian
basis

$\mathbf{D}_{ijk}=e^{-\gamma(r_{ij}-\mu_{k})^{2}}f_{c}(r_{ij})$ (4)

where $\gamma=1$Å-2 and $r_{ij}$ is the distance between atom $i$ and $j$. For
$\mu_{k}$, we choose 20 values uniformly between 0 and a cutoff-radius
$r_{c}=5$Å. The cutoff function

$f_{c}(r)=\begin{cases}\frac{1}{2}\left(\cos(\frac{\pi
r}{r_{c}})+1\right),&\mbox{if $r\leq r_{c}$}\\\ 0,&\mbox{else}\end{cases}$ (5)

emphasizes the importance of the local neighborhood by decaying
$\mathbf{D}_{ijk}$ as a function of $r_{ij}$ and ensures a smooth transition
to zero at $r_{c}$. For the atoms, we utilize randomly initialized trainable
atom type embeddings (Fig. 2a, $Z_{H},Z_{C},Z_{N}$). Finally, the bag is
represented by the number of remaining atoms (Fig. 2a, $B$).
The action is represented by distances to atoms already placed (Fig. 2a,
$D_{Q}$) as well as a special query type embedding (Fig. 2a, $Z_{Q}$). When
calculating $Q$-values the query atom is placed at various voxels allowing for
inference of $Q$-values at possible positions of the next atom to be placed.
In this way, all the encoded distances are easily converted to a 3D structure.

Figure 2: (a) The molecule is represented as a distance matrix $D$ (solid
lines between atoms) in a Gaussian basis, embeddings for each atom type
$(Z_{H},Z_{C},Z_{N})$ and a bag $B$. An action constitutes an atom placed in
one of the voxels. Each action is represented as a distance to the atoms in
the state $D_{Q}$ (dashed lines, not all distances shown) and a query
embedding $Z_{Q}$. (b) To predict $Q(s,a)$ the network takes as input the
state and action representation. The green part of the network operates on the
state and the purple section includes the action. The action part is repeated
for each voxel in the state resulting in $N$ possible actions. (c) The
resulting state after the voxel with the highest $Q$-value is chosen. If the
state is terminal the structure is relaxed using the model potential.

The state representation (Fig. 2b, green) is updated using SchNet-blocks. In
the following we will use superscripts to index the blocks, if multiple blocks
follow each other. Fully connected layers are given as $W_{i}$ where the full
details for each fully connected layer is given in the supplementary
information. Each block operates on the distance and atom type embeddings
using continuous-filter convolutions with skip connections as introduced in
SchNet:

$\displaystyle\tilde{Z}^{l+1}_{i}=Z_{i}^{l}+\sum_{j\neq i}Z^{l}_{j}\circ
W_{d}^{l}(\mathbf{D}_{ij}),$ (6)

where $\circ$ is elementwise-multiplication and $Z_{i}\in\mathbb{R}^{d}$ are
the atom embeddings. The convolution is followed by an atomwise fully
connected layer with a skip connection:

$\displaystyle
Z^{l+1}_{i}=\tilde{Z}^{l+1}_{i}+W_{a}^{l}(\tilde{Z}^{l+1}_{i}),$ (7)

which is passed to the next SchNet block. By chaining several blocks,
information from distant atoms are passed to each atom embedding. After $L=5$
layers, the atomwise representations, $Z^{L}_{i}$, are used for predicting the
energy $(E)$ of the structure as as combination of local contributions

$\displaystyle E=\sum_{i}W_{E}(Z^{L}_{i})$ (8)

Additionally, we compute a state value ($SV$).

$\displaystyle SV=W_{SV}\left(W_{B}(B)\oplus\sum_{i}Z^{L}_{i}\right),$ (9)

where $\oplus$ is concatenation. The state value is an indication of whether
high or low $Q$-values are expected from the current state, as will be evident
in the next section. Unlike the energy, the state value must consider
remaining atoms in the bag when calculating if the current state can progress
into a valid molecule.

Having covered the part of the architecture that deals with the state, we now
turn to the action part given in Fig. 2b, purple). The query type
representation is updated using a single SchNet-block. Furthermore, a multi-
head self-attention block over all the atom representations follows. In
contrast to local models where information from distant atoms must propagate
through several blocks, self-attention allows for immediate global information
flow. In the case of predicting $Q$-values for unfinished structures this is
especially relevant, as the model must be aware of possible dangling bonds in
one end of a molecule while predicting $Q$-values in the other end. For $h=8$
heads, we calculate a query, key and value as

$\displaystyle
q^{l}_{i}=Q^{l}(Z_{i}),\;k^{l}_{i}=K^{l}(Z_{i}),\;v^{l}_{i}=V^{l}(Z_{i}),$
(10)

where the superscript index the head and $Q^{l},K^{l},V^{l}$ are linear
projections to a subspace of size $d_{h}=d/h$, where $Z_{i}\in\mathbb{R}^{d}$.
From this, a dot-product attention score is calculated

$\displaystyle\alpha^{l}_{ij}=\frac{\exp{(q^{l}_{i}\cdot
k^{l}_{j}/\sqrt{d_{h}})}}{\sum_{j}\exp{(q^{l}_{i}\cdot k_{j}/\sqrt{d_{h}})}},$
(11)

which indicates the importance of latent representation $v^{l}_{j}$ to
describe $v^{l}_{i}$. From these scores a new latent representation of $Z_{Q}$
is created as

$\displaystyle Z_{Q}^{l}=\sum_{j}\alpha^{l}_{Qj}v_{j}^{l}$ (12)

which is finally concatenated and processed by a fully connected layer with a
skip connection and layer normalization Ba, Kiros, and Hinton (2016)

$\displaystyle\tilde{Z}_{Q}=\text{Layer-
norm}[Z_{Q}+W_{c}(Z_{Q}^{1}\oplus\dots\oplus Z_{Q}^{h})]$ (13)

The representation of the query atom is then used for calculating the
$Q$-value

$\displaystyle A$ $\displaystyle=W_{A}(\tilde{Z}_{Q})$ (14) $\displaystyle Q$
$\displaystyle=\text{Softmax}(SV+A),$ (15)

where $A\in\mathbb{R}^{M+1}$ , $M$ is the number of atom types, and
$SV\in\mathbb{R}$ is added to the first $M$ entries of $A$. Entry $i$ in $Q$
specify the $Q$-value for atom type $i$ and the $M+1^{\prime}\text{th}$ entry
allows for low $Q$-values for all atom types. $SV$ enables the network to
lower the $Q$-value for all possible actions in a given state, instead of
independently learning that all actions leads to high-energy structures. By
evaluating the $Q$-value at multiple voxels, the agent follows the greedy
policy of picking the action with the highest $Q$-value, leading to the
structure in Fig. 2c.

##  IV Method

The algorithm consists of two phases. In the first phase we use a simple form
of imitation learning (IL) called behavioral cloning, where the agent learns
to build the structures in the database as well as predicting their energies
and forces. This agent serves as a starting point for all agents employed in
the next phase.
In the second phase, a RL process is started where stoichiometric-specific
agents refines the IL model by constructing new molecules and querying an
energy calculator. Based on the received energies and forces, the model
updates the $Q$-values, energies and forces before constructing a new
molecule. As the agents receive feedback from the energy calculator the
molecules continue to improve.

###  IV.1 Imitation learning phase

To create an initial model the GDB-11 database consisting of molecules
containing up to 11 heavy atoms (C,N,O,F) is utilized. Specifically, we use
all 1850 structures with 6 heavy atoms for supervised pretraining. Using RDKIT
Landrum  SMILES are transformed into 3D structures and a single point density
functional theory (DFT) calculation of the energy and forces is performed
using GPAW Mortensen, Hansen, and Jacobsen (2005); Enkovaara  _et al._ (2010)
using a localized atomic basis set Larsen  _et al._ (2009). The model is then
trained using the following loss function

$l=\rho_{E}l_{E}+\rho_{F}l_{F}+\rho_{Q}l_{Q},$ (16)

where we have

$\displaystyle l_{E}$ $\displaystyle=||E-\hat{E}||^{2}$ (17) $\displaystyle
l_{F}$
$\displaystyle=\frac{1}{3N}\sum_{i=1}^{3N}\text{Huber}(F_{i},\hat{F}_{i})$
(18) $\displaystyle l_{Q}$
$\displaystyle=-\sum_{i=1}^{M+1}p_{i}\log{\hat{Q}_{i}},$ (19)

where the hat denotes values predicted by the network and $\rho_{E}=0.1$,
$\rho_{F}=0.9$ and $\rho_{Q}=1$ are empirically chosen to balance the
contributions to the total loss. $N$ is the number of atoms in the structure,
with $M$ different atom types.
The energy (eq. 17) is trained using a mean squared error (MSE) loss while
forces (eq. 18) are updated using a Huber loss given by

$\displaystyle\text{Huber}(F,\hat{F})=\begin{cases}\frac{1}{2}(F-\hat{F})^{2},&\mbox{if
$|F-\hat{F}|<1$}\\\ |F-\hat{F}|-\frac{1}{2},&\mbox{else}\end{cases}$ (20)

i.e. a squared loss for small differences and a linear loss for large errors.
The Huber loss was more stable by suppressing the effect of large force
outliers which were occasionally present in the structures built by the model.
Similarly, a cross entropy loss for the $Q$-values (eq. 19) was found to be
more stable than a MSE for the pretraining phase where only $Q$-values of 0 or
1 are required. For the cross entropy loss $\hat{Q}_{i}$ refers to the
$Q$-value for atom type $i$ and we set $p_{i}=1$ for the atom type placed in
the database and 0 for the other types. Additionally, for each action present
in the database we generate five perturbed actions by randomly moving the
selected atom and setting $p_{i}=1$ for the $M+1^{\prime}$th entry in the
$Q$-value prediction, thereby enforcing $Q$-values for all atom types to be
zero for the randomly generated actions. As this generates an imbalance in the
training data the valid actions are weighted by a factor of five.
The model is implemented in Pytorch Paszke  _et al._ (2019) and trained using
the Adam Kingma and Ba (2014) optimizer with an initial learning rate of
$5\cdot 10^{-3}$ and a batch size of 384. We use 90$\%$ of the structures for
training the remaining $10\%$ are used as a validation set for decaying the
learning rate by a factor of 2 when the validation error plateaus for 30
epochs. Once the learning rate decayed below $10^{-6}$ the validation loss
stagnated, and the final model is chosen for employment in the RL phase.

###  IV.2 Reinforcement learning phase

In the RL phase we deploy the IL model for all 135 stoichiometries present in
the training set. A cubic cell of dimension [20Å, 20Å, 20Å] and a $Q$-value
grid resolution of 0.2Å is chosen, with the initial atom restricted to the
center of the cell. We use a modified $\varepsilon$-greedy strategy, to trade-
off exploration and exploitation. To satisfy two random moves per episode in
expectation, we choose a greedy action with probability 1-2/T and a random
action with probability 2/T, where $T$ is the number of atoms in the
structure. We choose $5\%$ of random actions uniformly random, while the other
$95\%$ are uniformly sampled from the top $5\%$ of actions. Furthermore, all
hydrogen atoms are masked until all other atom types are placed which reduces
the action space without excluding any molecules. Due to the IL model an
action among the highest $Q$-values is frequently taken which ensures
exploration while simultaneously suppressing the large fraction of actions
resulting in completely invalid molecules.
To limit the action space, we only allow actions which fulfill

$\displaystyle
c_{1}[r_{\text{cov}}(a)+r_{\text{cov}}(i)]<r_{ai}<c_{2}[r_{\text{cov}}(a)+r_{\text{cov}}(i)],$
(21)

for at least one atom $i$ already present in the structure. Here $c_{1}=0.75$,
$c_{2}=1.25$ and $r_{\text{cov}}(i)$ is the covalent radius of atom $i$.
$r_{\text{cov}}(a)$ is the covalent radius of the atom placed in action $a$
and $r_{ai}$ is the distance between the new atom and a present atom $i$. That
is, the new atom must be placed within 0.75 to 1.25 times the sum of the
covalent radii of itself and an already present atom. This restricts the model
to building non-fragmented structures for a better comparison with the
database.
The agent trained as detailed in the previous section is perfectly capable of
building a variety of different molecules given a predetermined stoichiometry.
Owing to the stochastic nature of the modified $\varepsilon$-greedy policy
employed sequential builds using the same agent results in different
structures being built.
We exploit that to construct a first ensemble of structures by tasking the
agent with building molecules 200 times. It does so according to the
pretrained $Q$-values and every completed structure is subject to a relaxation
in the model energy, eq. (8), followed by snapping the molecule back on the
grid. If the relaxation results in a fragmentation of the molecule into
several pieces, the original structure is used instead. The initial
exploration phase promotes exploration of a large part of configuration space
to avoid premature convergence to a local minimum. Since the model is
pretrained the generated molecules will generally be of high quality.
Once the first 200 builds are completed, the search enters the RL mode. Here
the algorithm alternates between generating a structure and updating the model
using five mini-batches until a total of 600 additional molecules have been
generated. For the $Q$-values the loss function is changed to MSE, since the
target is now given by eq. (2), i.e. it is no longer 0 or 1, but depends on
the potential energy of the molecule. By further improving the $Q$-values and
force-predictions the search is directed towards low-energy regions of
configuration space.

###  IV.3 GDB-11 benchmark

We start by testing the performance of the current framework when it comes to
identifying molecules with stoichiometries already present in the database
used for the training. The search is conducted as a set of 64 independent
restarts following the outlined protocol: using the data-based-trained agent
for 200 builds and applying RL for 600 subsequent builds. To facilitate the
quantification of the amount of new structures found, we resort to the use of
SMILES, as an analytic tool. We stress that SMILES are not used in any way by
the agents as this would limit the scope of these. Once the search is finished
all structures are post processed by taking five gradient steps using the
final energy landscape. This is a computationally cheap step that removes
uncertainties from the original relaxation of the molecules in the less
reliable model energy landscape that had been learned up to the point of the
RL generating the particular molecule as well as unfortunate grid snapping.
These re-relaxed structures are then converted to SMILES and all unique
constitutional isomers are found. To this end, we show in Fig. 3 some of the
molecules being built when searching for C4H4O2. As illustrated, we find both
new constitutional isomers and often several stereoisomers for these. Looking
across these 64 restarts, all of the structures of this composition present in
the database (orange arrow) are found again (purple bar). However, also a
considerable number of molecules not present in the database are found (green
bar).

Figure 3: Subset of discovered constitutional isomers (left) and stereoisomers
(right) for C4H4O2. Further shown is the number of all unique constitutional
isomers in the database (orange arrow) as well as the molecules produced by
ASLA divided into molecules present in the database (purple bar) and novel
molecules (green bar). Figure 4: DFT energy distribution of the molecules. (a)
Molecules in the database. Blue line is the average energy and dark blue line
is the lowest energy isomer. (b) Orange: structures built in episode 400. The
red line indicates the mean energy of the structure and the dark red shows the
mean of the lowest energy structure in each run found before episode 400.
Finally, the dark red dashed line shows the lowest energy isomer among the 64
parallel runs, found before episode 400. (c) Same as (b), but now at 800
episodes. As the search progresses a shift towards lower energies is observed.

To investigate whether the produced molecules are of low energy we illustrate
the evolution of the 64 agents from episode 400 to 800 in Fig. 4. The
histograms indicate the energy distribution of the structures built in the
given episode across the 64 agents. As the search progresses the distribution
shifts towards lower energies as the model begins converging, which is further
evident by the mean of the distribution (red line) shifting down. Similarly,
when we only consider the mean of the lowest energy structure in each run
(i.e. not necessarily built in episode 400 or 800), which is depicted using
dark red we see a similar shift. In blue, the average energy of the structures
in the database, snapped to the grid ASLA operates on, is given. Initially the
database is better than the average built of the agent due to imperfect
fitting to the database but is overtaken by the agent at episode 800. If one
looks only at the best structure in the database (dark blue) the average
lowest energy approaches this limit and is marginally beaten by one of the
isomers (dark red, dashed). Fully relaxing all unique isomers in DFT indeed
reveals that the lowest energy isomer produced by ASLA coincides with the
minimum structure in the database.

Figure 5: All $Q$-values above 0.7 are shown with darker blue corresponding to
higher $Q$-values. For each state we only show the $Q$-values corresponding to
the atom placed in the given step.

To observe the decision process for the RL algorithm we plot the $Q$-values
for one of the runs in figure 5. First of all, the model has clearly learned a
great deal about chemical bonding and thus generally suppresses invalid
actions, predicting only high $Q$-values for sensible bonding sites. The agent
is fully aware of possible symmetries during the build. Starting with a single
atom in $s_{1}$, the $Q$-values have spherical symmetry. For the linear dimer
in $s_{2}$ the symmetry of the $Q$-values reduces to a ring, and for the
Y-shaped C3O backbone in $s_{5}$, two symmetric ’caps’ of high $Q$-values
develop. As evident in states $s_{4}$-$s_{8}$ where multiple sites have high
$Q$-values, the model decides between different possible isomers when building
the molecule. Bond lengths chosen early on e.g. the C-C bond lengths when
placing the C atoms from $s_{2}$ to $s_{5}$ have an impact on later
$Q$-values. The $Q$-values encircled by a black ellipse in $s_{8}$ are thus
smaller than the $Q$-value that are chosen upon building $s_{9}$. In both
cases, aldehyde (HOC-) groups are formed, but at $C$ atoms involved in double
or single C-C bonds, respectively, the latter being the energetically
preferred. Placing H atoms at both sites would have been possible causing
tautomerization involving the methyl group, but again given the initial
choices for C-C bond lengths, this is no longer preferable as reflected by the
$Q$-values in $s_{9}$. Owing to the modified $\varepsilon$-greedy policy, such
alternative molecular builds are occasionally made despite the C-C bond
lengths, and the subsequent relaxation in the on-the-fly learned model will
provide properly adjusted molecules for the RL training.

In Fig. 6 the total number of structures found after 800 episodes for the 15
stoichiometries with most isomers in the database are shown. For a figure
involving all 135 stoichiometries, see supplementary Fig. 1. Similarly to
earlier, purple bars indicate the structures present in both the search and
the database, whereas green indicate novel structures. A majority of the
structures in the database (orange line) are found within the 800 episodes
using 64 searches. In total 1769 of the 1850 molecules in the database are
found, i.e. slightly more than the 1665 structures used for the training set.
Accounting for both novel molecules and structures rediscovered from the
database, 20074 unique constitutional isomers are generated resulting in a
10-fold increase of the database. Especially for stoichiometries where
multiple bond combinations are possible, i.e. for stoichiometries containing a
lot of carbon, oxygen and nitrogen several new constitutional isomers are
discovered, whereas for hydrocarbons and CxHyFz where a lot of single bonds
are present the database covers a large fraction of the isomers.

Figure 6: Molecules discovered. Orange line: number of constitutional isomers
in the database. Purple bars: molecules found that are also present in the
database. Green bars: novel molecules found.

If the RL is neglected and the IL model is used for sampling in all 800
episodes, 1649 of the structures in the database and 15709 novel molecules are
generated for a total of 17358 unique constitutional isomers. The IL model
thus covers a large fraction of the produced structures and turning on RL
provides 15.6% more molecules.

To illustrate the difference between the IL model and the RL version, we plot
the average atomization energy per atom of the structures generated by ASLA
throughout the 800 episodes (fig. 7) and compare it to the energies of the
structures in the database when placed on the same 0.2Å resolution grid. For
each ASLA curve the average is over the 64 independents run as earlier and now
also the 135 stoichiometries. As previously observed, the model trained on the
database starts out by producing higher energy structures (red curve) than the
average structure in the database (blue curve), due to the imperfect fitting
of the database. Similarly, the average of the best structures (dark red) is
worse than the lowest energy isomer in the database (dark blue). At around 500
episodes both the average and best structures produced by ASLA marginally
outperforms the database. At 800 episodes the best structure is outperformed
by $0.04$ eV/atom on average. If one only looks at the best structure in each
of the 64 runs (dark red, dashed) the database is outperformed by $0.09$
eV/atom.

Figure 7: Blue: average energy of structures in the database. Dark blue:
average energy of lowest energy molecules in the database. Red: average energy
of structures built by ASLA. Dark red: average energy of lowest energy
molecules found by ASLA. Dashed dark red: lowest energy isomer among the 64
runs.

In order to check if this energy difference amounts to new low energy
structures the specific _stereoisomers_ in the database and the structures
generated by ASLA are compared. In that case 1742 of the molecules in the
database are found. When counting _constitutional_ isomers as in Fig. 6, 1769
of the structures in the database were found, which means that in 27 cases the
stereoisomer in the database is not found, but the constitutional isomer is.
In total, 29049 stereoisomers not present in the database are found. Each
stereoisomer in both the database and generated by ASLA is then fully relaxed
using DFT, which shows that for each stoichiometry the lowest energy
stereoisomer in the database is found by ASLA for all but one stoichiometry
and in 72 of 135 stoichiometries a lower energy stereoisomer is found.

###  IV.4 Exploration beyond the database

Having shown that the ASLA framework is capable of reproducing and expanding
molecular structures of compositions already present in the database, we now
move to explore its performance when applied outside the realms of the
database.

As an example, we investigate C9H8O4, a molecule with twice as many heavy
elements as in the training set. As before, 64 independent runs are started.

Figure 8: Searching for C9H8O4 isomers. (a) Constitutional isomers found as
the search progresses. As the search progresses the RL version (red) quickly
produces more structures than the IL version. (b) Energy-density of the
generated structures. In addition to producing more structures, the RL version
(red) generates a larger fraction of low energy molecules.

In Fig. 8a the number of unique constitutional isomers generated as a function
of episode number for the IL model (blue) and the model with RL (red) is seen.
Unlike for the smaller molecules in IV. C where the IL model was sufficient to
cover a large fraction of configuration space there is now a significant
difference between the RL and IL agent. The RL model clearly outperforms the
IL model in terms of number of molecules generated as the search progresses,
showing that the on-the-fly training is able to correct inefficiencies in the
initial version. However, the ultimate goal is to generate low energy
molecules. Hence, in Fig. 8b the atomization-energy-density of the generated
molecules for both the IL model and the RL version is seen. As before all
stereoisomers are fully relaxed using DFT. The molecules generated by the RL-
enhanced model are significantly more stable than the IL model, with an
average energy difference of 0.105 eV/atom or 2.22 eV per structure. Thus, the
RL version does not only cover more of configuration space, it is also able to
focus on the low energy regions.

Figure 9: Searching for C9H8O4 isomers. 2D visualization of IL agent (squares)
and RL version (circles). All structures are colored by their energy on a
scale from grey to blue (IL) or red (RL). The 30 lowest energy structures for
both versions of the algorithms are framed in black.

In Fig. 9 we investigate a 2D visualization of the generated molecules for
both the IL model (squares colored grey to blue) and RL version (circles
colored grey to red). The structures are represented using smooth overlap of
atomic positions Bartók, Kondor, and Csányi (2013) (SOAP) in a 2D space using
t-SNE and colored by their energy using two different color schemes to help
guide the eye. The 30 lowest energy structures found for both models are
framed in black and can be seen in the supplementary. The RL version discovers
certain low energy regions such as compact aromatic molecules as shown in the
inset. The IL version, however, is not able to extrapolate into this region
based on what it has learned from the database. Interesting molecules among
the low energy isomers include the lowest energy isomer found by the search
(red star) as well as the second lowest isomer (see supplementary) that both
resemble uvitic acid. The search discovers umbellic acid (red square) and
caffeic acid (just outside top 30). Additionally, the RL model discovers a
stable isomer with a 7-membered ring (supplementary), a substructure not
present in the original database as only molecules with up to 6 heavy atoms
were used. Despite this, the RL model learns to construct such molecules which
would have been hard to discover by a pure supervised generative approach.
For the IL model we observe fewer of the common aromatic molecules as it has
no driving force towards low energy regions of configuration space, as
exemplified by the lowest energy isomer found (blue star) being 2 eV higher in
energy than the corresponding lowest energy structure in the RL version. The
missing focus on low energy parts of configuration space results in the 30
lowest energy structures being scattered across the 2D space, mostly composed
of elongated structures (such as the blue square) where subparts, unlike large
aromatic rings, are more dominant in the GDB-11 database. Utilizing the
feedback from the RL is thus crucial to correct flaws in the IL model as well
as expanding beyond knowledge in the database.

##  V Conclusion

We have presented a framework for autonomous construction of molecules. The
method relies on a preexisting database which via supervised learning provides
a base level for a model used for constructing new molecules. Further training
of the model is done in a reinforcement learning setting where feedback is
provided by single point energy calculations by a high-level quantum
mechanical total energy method. The resulting model is able to reproduce
structures in the database as well as producing novel structures. The
introduced model is able to operate directly in 3D space allowing for biasing
the search towards stable molecules. When applying the imitation learning
model in domains far from the training set, the reinforcement learning
procedure corrects initial shortcomings in the model. Further work in this
direction could investigate building molecules in specific environments, such
as organic light-emitting diode (OLED) or organic solar panels where
properties such as HOMO-LUMO gap could be included in the reward function
together with the stability.

##  VI Acknowledgments

We acknowledge support from VILLUM FONDEN (Investigator grant, Project No.
16562). This work has been supported by the Danish National Research
Foundation through the Center of Excellence “InterCat” (Grant agreement no.:
DNRF150)

##  VII References

## References

  * Hartke (1993) B. Hartke, The Journal of Physical Chemistry 97, 9973 (1993).
  * Emilie S. Henault (2020) J. H. J. Emilie S. Henault, Maria H. Rasmussen, PeerJ Physical Chemistry 2 (2020), 10.7717/peerj-pchem.11.
  * Wales and Doye (1997) D. J. Wales and J. P. K. Doye, The Journal of Physical Chemistry A 101, 5111 (1997).
  * Kennedy and Eberhart (1995) J. Kennedy and R. Eberhart, in  _Proceedings of ICNN’95 - International Conference on Neural Networks_, Vol. 4 (1995) pp. 1942–1948 vol.4.
  * Kolsbjerg, Peterson, and Hammer (2018) E. L. Kolsbjerg, A. A. Peterson, and B. Hammer, Phys. Rev. B 97, 195424 (2018).
  * Meldgaard, Kolsbjerg, and Hammer (2018) S. A. Meldgaard, E. L. Kolsbjerg, and B. Hammer, The Journal of Chemical Physics 149, 134104 (2018).
  * Winter  _et al._ (2019) R. Winter, F. Montanari, A. Steffen, H. Briem, F. Noé, and D.-A. Clevert, Chem. Sci. 10, 8016 (2019).
  * James L. Melville and Hirst (2009) E. K. B. James L. Melville and J. D. Hirst, Combinatorial Chemistry and High Throughput Screening 12, 332 (2009).
  * Carpenter and Huang (2018) K. A. Carpenter and X. Huang, Current Pharmaceutical Design 24, 3347 (2018).
  * Irwin and Shoichet (2005) J. J. Irwin and B. K. Shoichet, Journal of Chemical Information and Modeling 45 (2005), 10.1021/ci049714+.
  * Ruddigkeit  _et al._ (2012) L. Ruddigkeit, R. van Deursen, L. C. Blum, and J.-L. Reymond, Journal of Chemical Information and Modeling 52 (2012), 10.1021/ci300415d.
  * Ramakrishnan  _et al._ (2014) R. Ramakrishnan, P. O. Dral, M. Rupp, and O. A. von Lilienfeld, Scientific Data 1 (2014), 10.1038/sdata.2014.22.
  * Haastrup  _et al._ (2018) S. Haastrup, M. Strange, M. Pandey, T. Deilmann, P. S. Schmidt, N. F. Hinsche, M. N. Gjerding, D. Torelli, P. M. Larsen, A. C. Riis-Jensen, J. Gath, K. W. Jacobsen, J. J. Mortensen, T. Olsen, and K. S. Thygesen, 2D Materials 5, 042002 (2018).
  * Winther  _et al._ (2019) K. T. Winther, M. J. Hoffmann, J. R. Boes, O. Mamun, M. Bajdich, and T. Bligaard, Scientific Data 6, 75 (2019).
  * Li, Zhang, and Liu (2018) Y. Li, L. Zhang, and Z. Liu, “Multi-objective de novo drug design with conditional graph generative model,” (2018), arXiv:1801.07299 .
  * Gupta  _et al._ (2018) A. Gupta, A. T. Müller, B. J. H. Huisman, J. A. Fuchs, P. Schneider, and G. Schneider, Molecular Informatics 37, 1700111 (2018).
  * Gómez-Bombarelli  _et al._ (2018) R. Gómez-Bombarelli, J. N. Wei, D. Duvenaud, J. M. Hernández-Lobato, B. Sánchez-Lengeling, D. Sheberla, J. Aguilera-Iparraguirre, T. D. Hirzel, R. P. Adams, and A. Aspuru-Guzik, ACS Central Science 4, 268 (2018).
  * Zhavoronkov  _et al._ (2019) A. Zhavoronkov, Y. A. Ivanenkov, A. Aliper, M. S. Veselov, V. A. Aladinskiy, A. V. Aladinskaya, V. A. Terentiev, D. A. Polykovskiy, M. D. Kuznetsov, A. Asadulaev, Y. Volkov, A. Zholus, R. R. Shayakhmetov, A. Zhebrak, L. I. Minaeva, B. A. Zagribelnyy, L. H. Lee, R. Soll, D. Madge, L. Xing, T. Guo, and A. Aspuru-Guzik, Nature Biotechnology 37, 1038 (2019).
  * Noé  _et al._ (2019) F. Noé, S. Olsson, J. Köhler, and H. Wu, Science 365, eaaw1147 (2019).
  * Hoffmann and Noé (2019) M. Hoffmann and F. Noé, “Generating valid euclidean distance matrices,” (2019), arXiv:1910.03131 .
  * Gebauer, Gastegger, and Schütt (2019) N. W. A. Gebauer, M. Gastegger, and K. T. Schütt, “Symmetry-adapted generation of 3d point sets for the targeted discovery of molecules,” (2019), arXiv:1906.00957 .
  * Mansimov  _et al._ (2019) E. Mansimov, O. Mahmood, S. Kang, and K. Cho, Scientific Reports 9, 20381 (2019).
  * Maziarka  _et al._ (2020) Ł. Maziarka, A. Pocha, J. Kaczmarczyk, K. Rataj, T. Danel, and M. Warchoł, Journal of Cheminformatics 12, 2 (2020).
  * Köhler, Klein, and Noe (2020) J. Köhler, L. Klein, and F. Noe, in  _Proceedings of the 37th International Conference on Machine Learning_, Proceedings of Machine Learning Research, Vol. 119, edited by H. D. III and A. Singh (PMLR, 2020) pp. 5361–5370.
  * Satorras  _et al._ (2021) V. G. Satorras, E. Hoogeboom, F. B. Fuchs, I. Posner, and M. Welling, “E(n) equivariant normalizing flows for molecule generation in 3d,” (2021), arXiv:2105.09016 .
  * Winter, Noé, and Clevert (2021) R. Winter, F. Noé, and D.-A. Clevert, “Permutation-invariant variational autoencoder for graph-level representation learning,” (2021), arXiv:2104.09856 .
  * Olivecrona  _et al._ (2017) M. Olivecrona, T. Blaschke, O. Engkvist, and H. Chen, Journal of Cheminformatics 9, 48 (2017).
  * Putin  _et al._ (2018) E. Putin, A. Asadulaev, Y. Ivanenkov, V. Aladinskiy, B. Sanchez-Lengeling, A. Aspuru-Guzik, and A. Zhavoronkov, Journal of Chemical Information and Modeling 58, 1194 (2018).
  * Cao and Kipf (2018) N. D. Cao and T. Kipf, “Molgan: An implicit generative model for small molecular graphs,” (2018), arXiv:1805.11973 .
  * Popova, Isayev, and Tropsha (2018) M. Popova, O. Isayev, and A. Tropsha, Science Advances 4, eaap7885 (2018).
  * Zhou  _et al._ (2019) Z. Zhou, S. Kearnes, L. Li, R. N. Zare, and P. Riley, Scientific Reports 9, 10752 (2019).
  * Schreck, Coley, and Bishop (2019) J. S. Schreck, C. W. Coley, and K. J. M. Bishop, ACS Central Science 5, 970 (2019).
  * Simm, Pinsler, and Hernandez-Lobato (2020) G. Simm, R. Pinsler, and J. M. Hernandez-Lobato, in  _Proceedings of the 37th International Conference on Machine Learning_, Proceedings of Machine Learning Research, Vol. 119, edited by H. D. III and A. Singh (PMLR, 2020) pp. 8959–8969.
  * Jørgensen  _et al._ (2019) M. S. Jørgensen, H. L. Mortensen, S. A. Meldgaard, E. L. Kolsbjerg, T. L. Jacobsen, K. H. Sørensen, and B. Hammer, The Journal of Chemical Physics 151, 054111 (2019).
  * Fink, Bruggesser, and Reymond (2005) T. Fink, H. Bruggesser, and J.-L. Reymond, Angewandte Chemie International Edition 44, 1504 (2005).
  * Gao and Coley (2020) W. Gao and C. W. Coley, Journal of Chemical Information and Modeling 60, 5714 (2020), https://doi.org/10.1021/acs.jcim.0c00174 .
  * Weininger (1988) D. Weininger, Journal of Chemical Information and Computer Sciences 28, 31 (1988).
  * Mortensen  _et al._ (2020) H. L. Mortensen, S. A. Meldgaard, M. K. Bisbo, M.-P. V. Christiansen, and B. Hammer, Phys. Rev. B 102, 075427 (2020).
  * Schütt  _et al._ (2018) K. T. Schütt, H. E. Sauceda, P. Kindermans, A. Tkatchenko, and K. Müller, The Journal of Chemical Physics 148, 241722 (2018).
  * Vaswani  _et al._ (2017) A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. u. Kaiser, and I. Polosukhin, in  _Advances in Neural Information Processing Systems_, Vol. 30, edited by I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (Curran Associates, Inc., 2017).
  * Machado  _et al._ (2018) M. Machado, M. Bellemare, E. Talvitie, J. Veness, M. Hausknecht, and M. Bowling, Journal of Artificial Intelligence Research , 5573 (2018).
  * Ba, Kiros, and Hinton (2016) J. L. Ba, J. R. Kiros, and G. E. Hinton, “Layer normalization,” (2016), arXiv:1607.06450 .
  * (43) G. Landrum, “Rdkit: Open-source cheminformatics; http://www.rdkit.org,” .
  * Mortensen, Hansen, and Jacobsen (2005) J. J. Mortensen, L. B. Hansen, and K. W. Jacobsen, Phys. Rev. B 71, 035109 (2005).
  * Enkovaara  _et al._ (2010) J. Enkovaara, C. Rostgaard, J. J. Mortensen, J. Chen, M. Dułak, L. Ferrighi, J. Gavnholt, C. Glinsvad, V. Haikola, H. A. Hansen, H. H. Kristoffersen, M. Kuisma, A. H. Larsen, L. Lehtovaara, M. Ljungberg, O. Lopez-Acevedo, P. G. Moses, J. Ojanen, T. Olsen, V. Petzold, N. A. Romero, J. Stausholm-Møller, M. Strange, G. A. Tritsaris, M. Vanin, M. Walter, B. Hammer, H. Häkkinen, G. K. H. Madsen, R. M. Nieminen, J. K. Nørskov, M. Puska, T. T. Rantala, J. Schiøtz, K. S. Thygesen, and K. W. Jacobsen, Journal of Physics: Condensed Matter 22, 253202 (2010).
  * Larsen  _et al._ (2009) A. H. Larsen, M. Vanin, J. J. Mortensen, K. S. Thygesen, and K. W. Jacobsen, Phys. Rev. B 80, 195112 (2009).
  * Paszke  _et al._ (2019) A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Köpf, E. Yang, Z. DeVito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala, “Pytorch: An imperative style, high-performance deep learning library,” (2019), arXiv:1912.01703 .
  * Kingma and Ba (2014) D. P. Kingma and J. Ba, “Adam: A method for stochastic optimization,” (2014), 1412.6980 .
  * Bartók, Kondor, and Csányi (2013) A. P. Bartók, R. Kondor, and G. Csányi, Phys. Rev. B 87, 184115 (2013).

##  VIII Supplementary

###  VIII.1 Network details

In table. 1 a detailed overview of the blocks in the network is given. The
activation function is a shifted softplus Schütt  _et al._ (2018) (ssp) given
by

$\text{ssp}(x)=\ln(0.5e^{x}+0.5)$ (22)

Bag representation | $W_{B}$ | (5, 16, spp, 32)
---|---|---
Atomwise layers | $W_{a}$ | (64, 64, ssp)
Distances representation | $W_{d}$ | (20, 64, ssp)
State Value | $W_{SA}$ | (96, 32, ssp, 1)
Advantage | $W_{A}$ | (96, 32, ssp, 6)
Energy | $W_{E}$ | (64, 32, ssp, 16, ssp, 8, ssp, 1)
Concatenation layer | $W_{c}$ | (64, 64)
Table 1: Dimension of data as it passes through the fully connected layers.
The activation function is is the shifted softplus indicated by ssp.

Figure 1: Overview of all constitutional isomers built. Orange line indicates
number of isomers in the database. Purple are isomers built by ASLA which are
also part of the database. Green structures are new isomers.

Figure 2: 30 lowest energy structures for the RL run. Figure 3: 30 lowest
energy structures for the IL run. Energies are relative to the lowest energy
isomer found in the RL version.