Symmetry-Aware Actor-Critic for 3D Molecular Design
Authors: Gregor N. C. Simm
Robert Pinsler
Gábor Csányi
José Miguel Hernández-Lobato
Venue: ICLR 2021 Poster (2021)
Keywords: deep reinforcement learning
molecular design
covariant neural networks

Abstract: Automating molecular design using deep reinforcement learning (RL) has the potential to greatly accelerate the search for novel materials. Despite recent progress on leveraging graph representations to design molecules, such methods are fundamentally limited by the lack of three-dimensional (3D) information. In light of this, we propose a novel actor-critic architecture for 3D molecular design that can generate molecular structures unattainable with previous approaches. This is achieved by exploiting the symmetries of the design process through a rotationally covariant state-action representation based on a spherical harmonics series expansion. We demonstrate the benefits of our approach on several 3D molecular design tasks, where we find that building in such symmetries significantly improves generalization and the quality of generated molecules.

Published as a conference paper at ICLR 2021
SYMMETRY-AWARE ACTOR-CRITIC FOR 3D MOLECULAR DESIGN
Gregor N. C. Simm,' Robert Pinsler,' Gábor Csányi,' & José Miguel Hernández-Lobato
DDparrmmet of Engineering, University of Cambridge, Cambridge, UK 2 AAlan Turing Institute, London, UK (gncs2, rp586, gc121, jmh2333 ) cam. ac uk
ABSTRACT
Automating molecular design using deep reinforcement learning (RL) has the po- tential to greatly accelerate the search for novel materials. Despite recent progress on leveraging graph representations to design molecules, such methods are fun- damentally limited by the lack of diree-dimensionll (3D) information. In light of this, we propose novel actor-critic architecture for 3D molecular design that can generate molecular structures unattainable with previous approaches. This is achieved by exploiting the symmetries of the design process through rotationally covariant state-action representation based on a spherical harmonics series expan- sion. We demonstrate the benefits of our approach on several 3D molecular design tasks, where we find that building in such symmetries significantly improves gen- eralization and the quality of generated molecules.
INTRODUCTION
The search for molecular structures with desirable properties is a challenging task with important
applications in de novo drug design and materials discovery (Schneider et al. 2019). There exist plethora of machine learning approaches to accelerate this search, including generative models based et al.. 2018), recurrent neural networks (RNNs) (Segler et al. 2018), and generative adversarial networks (GANs) (De Cao & Kipi 2018). However, the reliance on sufficiently large dataset for exploring unknown regions of chemical space is severe limitation of such supervised models. Recent RL.-based methods (e.g.. Oliveerona cct al. 20177. ,, Jørgensen et al. (2019), Simm et al. (2020)) mitigate the need for an existing dataset of molecules they only require access to reward function.
on variational autoencoders (VAEs) (tómez-Bombanee
Most approaches rely on graph representations of molecules, where atoms and bonds are represented by nodes and edges, respectively. This is a strongly simplified model designed for the description of single organic molecules. is unsuitable for encoding metals and molccular clusters as it lacks information about the relative position of atoms in 3D space. Further, geometric constraints design process cannot be included, e.g. those given by the active site of an enzyme. more general representation closer the physical system one which molecule is described by its atoms' positions in Cartesian coordinates. However, it would be very inefficient to naively learn model based on this representation. That is because molecular properties such the energy are invariant (i.e. unchanged) under symmetry operations like translation or rotation of all atomic positions. A model without the right inductive bias would thus have to learn those symmetries from scratch.
the
In this work, we develop novel RL approach for designing molecules in Cartesian coordinates that explicitly encodes these symmetry operations. The agent builds molecules by consecutively placing atoms such that if the generated structure is rotated or translated, the agent's action is rotated and translated accordingly; this way, the reward remains the same (see Fig. [ (a)). We achieve this through a rotationally covariant state representation based on spherical harmonics, which we integrate into a novel actor-critic network architecture with an rego-segressive policy that maintains the desired covariance. Building in this inductive bias enables to generate molecular structures with more complex coordination geometry than the class of molecules that were attainable with previous approaches. Finally, we perform experiments several 3D molecular design tasks, where

---

Published as a conference paper at ICLR 2021
find that approach significantly improves the generalization capabilities of the RL agent and the quality of the generated molecules.
(a)
(b)
St
ag
81+1
(
x
80
81
82
83
87
ag
02
43
OF
R
R
2
V
Y
Os
S
TO
1
T3
T6
C
C
0
0
(o, Rr)
(a) Illustration of a rotation-covaricn state-action representation. the structure is rotated by R. the position of the action transforms accordingly. (b) Rollout with bag Bn SOF4, The agent builds a molecule by repeatedly taking atoms from the bag and placing them onto the 3D canvas. Bonds connecting atoms are only for illustration and not part of the MDP.
Figure 1:
In summary, contributions are as follows:
we propose the first approach for 3D molecular design that exploits symmetries of the design process by leveraging a rotationally covariant state representation;
we integrate this state representation into actor-critic neural network architecture with a rotationally covariant auto-regressive policy, where the orientation of the atoms to be placed is modeled through flexible distribution based on spherical harmonics;
we demonstrate the benefits of our approach on several 3D molecular design tasks, includ- ing newly proposed task that showcases the generalization capabilities of our agent.
2 BACKGROUND
2.1 REINFORCEMENT LEARNING FOR MOLECULAR DESIGN
In the standard RL setting (Sutton& Barto 2018). an agent interacts with the environment max- imize its reward. Formally, such an environment is described by a Markov decision process (MDP) T.
M AS, A. T, Mu. Y, r) with states s, E S, actions a₂ A, transition dynamics T S A S,
initial state distribution (lu. discount factor 7 E (0, 1], and reward function r : S x A + R. The goal to learn a stochastic policy x (a, 8t) that maximizes the expected discounted return
(0) = Eeo-po (V (so)], where the value function (*(ss) = E₂ DD' 7" rssr, (rsr),ar)la [se] is defined the expected discounted return when starting from state s, and following policy म.
Following Simm et al. (200)), we design molecules by iteratively picking atoms from a bag and positioning them on 3D canvas. Such sequential decision-making problem is described by an MDP where the state = (Ce, B,) comprises both the canvas C, and the bag B. The canvas C, ((e, 1) is a set of atoms with chemical element e; E (H, C, N, O,... and position rr R placed up time / 1, where Cu can either be empty or contain set of initially placed atoms. The number of atoms on the canvas is denoted by [c,. The bag B, {(e. m(e))} is multi-set of atoms yet be placed, where m(e) is the multiplicity of the element e. Each action a, (Et, 2.) consists of the element e E B, and position Ia E R3 of the next atom to be added to the canvas. Placing an atom through action a, in state 8, modeled by a deterministic transition (C.+1,B.4) with B.+1 B, (e
function ((tt, ae) that yields the next state S141
The reward function r(st, a.) E(ss a;) is given by the negative energy difference between the resulting structure described by C,-1, and the sum of energies of the current structure C, and new atom of element é placed at the origin, i.e. AE(s, a,) = E(C(+1) [E(C,) + (e( O}}}] Intuitively, the reward encourages the agent to build stable, low-energy structures. We evaluate the energy using the fast semi-empirical Parametrized Method 6 (PM6) (Stewart) (007) as implemented
in SPARROW (Husch al. 2018 Bosia al. 2020); Appendix for details.
2

---

Published as a conference paper at ICLR 2021
An example of rollout is shown in Fig. (b). At the beginning of the episode, the agent observes the initial state (Co, Bb) ~ [s())), e.g. Co 0 and B, SOF I The agent then iteratively constructs a molecule by placing atoms from the bag onto the canvas until the bag is empty |
2.2 ROTATIONALLY COVARIANT NEURAL NETWORKS
A function f x +> y is invariant under a transformation operator T. X + X' if f(Tg fx) for all x g g G, where G is mathematical group. In contrast, is covariant with respect to Tg if there exists an operator T : y + y such that f(T, [x]) = =]][(x)]. To achieve rotational covariance, it natural to work with spherical harmonics. They are a set of complex- valued functions Y S + C with i = 0, 1, ..... and m -L, -f 1.....& 1, l on the unit sphere S in 3 The first few spherical harmonics are in Appendix B They are defined by
Yfth (0.+)) 2l + (e m)!
4m (e+ m)m Ptth (cos(
yp E [0, 2T),
[0,T],
where Pmn denotes the associated normalized Legendre polynomials of the first kind (Bateman, 1953), and each Y is normalized such that [X]" (1) s2] sin dd0dpp = 1. Any square-iicegralc function f S: + C can be written as a series expansion in terms of the spherical harmonics,
f(t) = Σ Σ )'Y,"
t=
where a - (0, E S2. The complex-valued coefficients { } are the analogs of Fourier coeffi- cients and are given by fT f( (d)). Such a function f can be modeled by learning the coefficients { using CORMORANT (Anderson et al. 2019}. a neural network architecture for predicting properties of chemical systems that works entirely in Fourier space. A key feature is that each neuron is covariant to rotation but invariant to translation; further, each neuron explic- itly corresponds to a subset of atoms in the molecule. The input of CORMORANT is a spherical function f S + Cd and the output a collection of vectors f (So, f where each fe E r x (2f + 1) is a rotationally covariant vector with r channels. That is, if the input is rotated by R E SO(3), then each fe transforms as fe + (R( where D²(R): SO(3) + ((28+1)x(26+11 are the irreducible representations of S0(3), also called the Wigner D-matrices.
3 COVARIANT POLICY FOR MOLECULAR DESIGN
An efficient RL agent needs to exploit the symmetries of the molecular design process. Therefore, we require a policy r[ass) with actions a (c, .)) that is covariant under translation and rotation with respect to the position x, i.e., should rotate (or translate) accordingly if the atoms on the canvas C are rotated (or translated). In contrast, the policy needs to be invariant to the element i.e. the chosen element remains unchanged under such transformations (see Fig. [ (a)). Since learning such a policy is difficult when working directly in global Cartesian coordinates, we instead follow Simm et al. (2020) and use an action representation that is local with respect to an already placed local atom. If the next atom is placed relative to the focal atom, covariance under translation of is automatically achieved and only the rotational covariance remains to be dealt with.
As shown in Fig. 2] model the action through sequence of sub-actions: (1) the index ff E [l..... [C of the focal atom around which the next atom is placed (2) the element (1, (1,..., No) of the next atom from the set of available elements, (3) distance d RR4 be- tween the focal atom and the next atom, and (4) the orientation ã = (0, sp) E $ of the atom on a unit sphere around the focal atom. Denoting zy as the position of the focal atom, we obtain action a (e,z) x) by mapping the local coordinates (á, d, f) to global coordinates I = If + d I, where
Shorthand for (S. 1), (0, 1), (F,4))
Hereafterr we drop the time index when it is clear from the context.
If the canvas Co is empty, the agent selects an element eo 15, and places it the origin, i.e. ag (eo, 0).
3

---

Published as a conference paper at ICLR 2021
is now covariant under translation and rotation. We choose these sub-actions using the following auto-regressive policy:
d, = p( 1o
$3
a3
54
C3
C4
F
S
B3
B4
o
Action representation of the auto auto-regressive policy. The agent chooses focal atom f. element e, distance d, and orientation Σ. We then map back to global coordinates æ to obtain action (e, v). Bonds between atoms are only for illustration.
Figure 2:
novel actor-critic neural network architecture that implements this policy illustrated in Fig. In the following, we discuss its state embedding, actor, and critic networks in more detail.
3.1 STATE EMBEDDING
The state embedding network transforms canvas and bag B to obtain a rotationally covariant and translationally invariant representation. For that, concatenate vectorized representation of the bag with each atom on the canvas and feed it into CORMORANT, i.c 5ot + CORMORANT(C, B),, Lens where goo ( I St E ()C)xTx(2++11 and is the number of channels. For the sake of exposition, we assume a single channel for each element in the bag, i.e. T = No (cf. Fig 3): in practice, we use up to four channels per element.
However, not every sub-action in Eq. (3p should transform equally under rotation and translation. While the orientation 7 needs to be covariant under rotation, the choice of focal atom f. element e, and distance d have to be invariant to rotation and translation. For these sub-actions, we additionally require an invariant state representation. To obtain such a representation inv e RIC|x we employ a combination of transformations from (nndernnsoon et al. (2019) as listed in Appendix (e.g. for f = 0, one can simply select the og component), which we collectively denote as Tinv
3.2 ACTOR
The distribution s(f]s) over the focal atom is modeled cate- gorical, f ~ Cat(f;hy), where h are the logits for each atom in C predicted by a multi-layer perceptron (MLP). Likewise, the distribution over the element is given by p(e(f,)) Cat(e; he) with fe MLP, (s)"") where siy' is the invariant representation for the focal atom. Since the num- ber of possible focal atoms f increases and the set of available elements e decreases during rollout, mask invalid focal atoms & (l,..., [L,...,/C.]]] [Ce]] and elements e B, by setting their probabilities to zero and re- re-normaliiing the categorical distributions. The agent does not make use of chemical concepts like bond connectivity to aid the choice of the focal atom.
Focal Atom and Element
Distance We select the channel Te corresponding to element from s to obtain 8 = {st)y Imm..... and sTe (5) Then model the distribution over the distance d be- tween the focal atom and the next atom to be placed as a mixture of M Gaussians, p(dle, f. = Em=1 Tm Nim, o where Fm is the mixing coefficient of the m-th Gaussian N(um, o² The mixing coefficients and the means are predicted by a mixture density network (MDN) (Bishop) 1994., i.e. {Tm, pm m==
M MDN(s) (Te)) The standard deviations {am { are global parameters. We guarantee that the sampled distances are positive by clipping values below zero
Combining Invariant and Covariant Features The choice of distance d can significantly affect the orientation ह of the atom. For example, if d has the length of a triple bond, then & will be very different from if was a single bond. Thus, we condition sow on distance d through non-linear and learnable transformation that preserves rotational covariance. We then this representation
f,e
4

---

Published as a conference paper at ICLR 2021
State Embedding
T8 TO
TF
2
1
2
C
CORMORANT
TrN
t= 0
3
8
B
inv
8 2
cov
Actor
h
slov
h
TF
Tinv
inv
MLP
MLP.
Cat(f:hy)
e ~ Cat(ee; he)
inv f
8 f.e
Cov
inv f.c
MDN
~ d
Teov
{Y;")
Critic
kuv
MLP
Σ
MLP,
V
Illustration of the state embedding, actor, and critic networks. Both canvas and bag B are fed to the state embedding network CORMORANT to obtain rotation-covariant (s*or) and -invariant (s" state representations. The actor network then samples the different sub-actions highlighted in bold. The critic takes the invariant representation giv compute a value V.
Figure 3:
to model a spherical distribution over T. Kondor Trivedi (0188) showed that a linear transfor- mation with learnable parameters is only covariant if the operation combines fragments with the same &. Further, the Clebsch- Gordan (CG) non-linearity allows one to combine two covariant fea- tures such that the result is still covariant. Thus, we obtain rotationally covariant representation ;= (f.......... + Te (d, ( conditioned on all previous sub-actions as follows:
fz
con 8
@ d St,
cav
@ (d 5f. Bog d. sTe)) We Ae,
(4)
where @ denotes the appropriate concatenation of matrices, and We is learnalee complex-valued matrix. As in Anderson al. (2019), perform the CG product Øcy only channel-wise to reduce computational complexity.
Next, utilize f obtain rotationally covariant spherical distribution for the orientation based on the series expansion in Eq. (2). Taking inspiration from commonly used distributions (Jammalamadaka & Terdik 2019), we propose to use the following expression:
Orientation
Lyeas
(( Z exp -B
(T))
PIY (3)
(5)
t=U
m=-
where 3 E a scaling parameter, and the term 1/VT with & Σtoo DE - regularizes
the distribution that it does not approach delta function. The normalization constant Z is esti- mated via Lebedev quadrature (Lebedev) 1975 1977). We sample from the distribution in Eq. (5)
using rejection sampling (Bishop) 2009) with uniform proposal distribution g (4T) Note
that contrast to more commonly used parametric distributions (e.g. von Mises-Fisher), this for- mulation allows to model coltiiiodalltiess. We discuss alternatives to Eq. (5) in Appendix
5

---

Published as a conference paper at ICLR 2021
3.3 CRITIC
The critic needs to compute a value V for the state that is invariant under translation and rotation. Given s!" we apply permutation-invariiin set encoding (Zaheer al., 2017) of the atoms, i.c.
ICI Σ
V(s) = MLP,
MLP (O.P)
(6)
Finally, we use PPO (Schulman et al. 2017) to learn the parameters of the actor-critic architecture. To encourage sufficient exploration, we add an entropy regularization term over the categorical sub- action distributions of the auto-regressive policy. For offline evaluation, we evaluate the policy without any exploration by choosing the most probable action. While the mode of the distributions for f and e is available in closed form, we approximate the global mode of the distributions over d and ō by evaluating the density S samples and picking the one with the highest density.
4 RELATED WORK
Reinforcement Learning for Molecular Design There exists large variety of RL-based ap- proaches for molccular design using either string- or graph-based representations of molecules Oliveerona et al. 2017 Guimaraes et al. 2018 Putin et al. [018 Neil et al., 2018 Popova et al.) Zhou et 2019.. However, the choice of representation limits the molecules that can be generated to a (small) region of chemical space for which the representation is applicable, i.e. single organic molecules. Such representations also prohibit the use of reward functions based on quantum-mechaniee mechanical propertiex; instead, heuristics are often used. Lastly, geometric constraints on the design process cannot be imposed as the representation does not include any 3D information.
2018| You al. 2
Molecular Design in Cartesian Coordinates Another down- side of string- and graph-based approaches is their neglect of information encoded in the interatomic distances. In light of this, Gebauer et al.) (2018) 2019) proposed a supervised gener- ative neural network for sequentially placing atoms in Cartesian coordinates. While the model respects local symmetries by con- struction, atoms are placed on 3D grid. Similar to other su- pervised approaches, one further requires a dataset that covers the particular class of molecules to be generated. Hammer and coworkers (Jorgensen et al., | Meldgaard et al.] (2020) em- ployed a Deep Q-Network (Mnih et al.) 2015) to build planar compounds and crystalline surfaces by placing atoms on grid. Recently, (Sim et al. 2020)) presented an RL formulation for molecular design in continuous 3D space. The agent models the position of the next atom to be placed in internal coordinates- i.e. the distance, angle, and dihedral angle with respect already existing atoms-whicl are invariant under translation and rota- tion. By mapping from internal to Cartesian coordinates, they then obtain a policy that is covariant under these symmetry oper- ations. However, as shown in Fig 4. the angle and dihedral angle are only defined with respect to two reference points, which are chosen be the two closest points to a focal atom. In highly symmetric states, e.g. as commonly encountered in materials, this representation fails to distinguish different configurations as one can- not uniquely select the two closest atoms as reference points anymore. In contrast, we do not rely on such reference points as the agent directly samples the orientation from spherical distribution.
(a)
(b)
OF
Figure 4: Example of two figurations (a) and (b) that the agent by Simm al. (020) can- not distinguish. While the values for distance d, angle a and dihe- dral angle & are the same, choos- ing different reference points (in red) leads to different action. This is particularly problematic in symmetric states, where one cannot uniquely determine these reference points.
Covariant Neural Networks in Chemical Science Prior work employed rotationally covariant neural networks to predict translation- and rotation-invariant physical properties (Thomas et al.
2018 Kondor et al., 2018) Weiler et al., 2018: Anderson et al. 2019 Miller et al. 2020 Finzi al. 2020 Fuchs al. 2020), e.g. scalars such the electronic energy. In contrast,
propose a translation inttrrannnvvaiant and rotation-covariant neural network architecture for generating molecules. For more general treatment of covariance (or equivariance) in RI., see van der Pol et al. (2020).

---

Published as a conference paper at ICLR 2021
5 EXPERIMENTS
We perform experiments to answer the following questions: (1) is the agent able to learn how to build highly symmetric molecules in Cartesian coordinates from scratch, (2) can we increase the validity, diversity, and stability of generated molecules, and (3) does our approach lead to improved generalization? We address (1) and (2) by evaluating the agent on a diverse range of tasks from the MoLGYM benchmark suite (Simm et al. 2020), and (3) on newly proposed stochastic-bag task (see Section [5.]] where bags are sampled from distribution over bags. In Appendix G
we show
with an additional experiment that the agent can learn to place water molecules around given solute to form solvation shell
We compare our approach (COVARIANT) against the RL agent proposed by [Sim et al. 2020) which iteratively builds molecules on a 3D canvas by working in internal coordinates (INTERNAL). As additional baseline, we consider a classical, optimization-based agent (OPT) with access to a black-box function that yields the energy E(C) and the atomic forces F(C) for a given canvas The agent constructs molecules by alternating between randomly placing atom and optimizing the structure. Moreover, the agent applies several heuristics inspired by fundamental chemical con- cepts guide the placement of atoms. To make the comparisons fair, we grant OPT comparable computational budget in terms of the total number of energy computations. Finally, for some exper- iments, the best possible performance based on quantum-chemical calculations can be reported. See Appendices D and F for more details on the baselines and experimental settings, and Appendix H for additional runtime comparison between the agents.
5.1 STOCHASTIC-BR BAG TASK
In Simm et al. (2020), set of molecular design tasks was introduced: the single-bag task assesses an agent's 's ability to build single stable molecules, whereas the multi-bag task focuses on building several molecules of different composition and size at the same time. A limitation of these tasks is that the initial bags were selected such that they correspond to known formulas, which in practice might not be known priori. In the stochastic-bag task, we relax this assumption by sampling from a more general distribution over bags. Before each episode, we construct a bag B = {(e, m(e)) by sampling counts (m(e1), .., m (emaa)) Mult(Ç, (C),,, where the bag size is sampled uniformly from the interval [Sminn ). Here, we obtain an empirical distribution pe from the multiplicities m(e) of given bag B*, For example, with B {H, 2), (O, 1)} we obtain PH = 2 and po 1. Since sampled bags might no longer correspond to valid molecules when placed completely, we discard bags where the sum of valence electrons over all atoms contained in the bag is odd. This ensures that the agent can build closed-shell system.
5.2 RESULTS
Building Highly Symmetric Molecules First, we evaluate the ability to build stable molecules featuring high symmetry and coordination numbers (e.g. trigonal bipyramidal, square pyramidal, and octahedral) on the single-bag task with bags SOF IF s. SOF s. and 5F As shown in Fig. (la),, COVARIANT can solve the task for SOF and IF; within 30 000 to 40 000 steps, whereas INTERNAL fails to build low-energy configurations as it cannot distinguish highly symmetric intermediates (cf. Fig. Further results in Fig. 5 (b) for SOF, and SF, show that COVARIANT is capable of building such structures. Likewise, OPT found the optimal structures for all four bags. While the constructed molecules are small in size, they would be unattainable with graph- or string-basee methods as such representations lack important 3D information. For example, RDKIT (Landrum, 2019, a state-of-the-arr library for 3D structure generation of organic molecules, failed at this task.
Validity, Diversity, and Stability of Generated Molecules Since string and graph representa- tions are not well-suited for designing molecules with complex 3D structure, it is difficult to directly compare to most prior work. To still enable comparisons, we follow the GuacuMol benchmark (Brown et al. 2019) and report the chemical validity, diversity, and stability of the molecules gen- erated by the agents for different experiments. A generated structure is considered valid if it can be successfully converted to a molecular graph by the tool X XYZ2MOL (Tensen) 2019 Kim & Kim
*Source code of the agent and environment is available at htttps::: / / thub. com/ / gncs/ mocs/mo'> gym For the calculation of F(C) and F(C) we employ PM6; the same method in the reward function.
7

---

Published as a conference paper at ICLR 2021
0.8
Covariant
0.6
Internal
0.8
0.4
Average Return
0.2
0.0
0.2
-02
00.4
-0.4
-0.8
-0.6
10
20
30
40
10
20
30
40
Steps 1000
Steps 1000
Figure 5: (a) Average offline performance on the single-bag task with bags SOF, (left) and IF's (right) across 10 seeds. In the lower right, molecular structures generated by the agents are shown. Dashed lines denote the optimal return for each experiment. Error bars show two standard devia- tions. (b) Further molecular structures generated by COVARIANT, namely SOF, and SF
Table 1: Validity, diversity, and stability (RMSD in Å) of generated structures.
Task
Experiment
Validity (f is better)
Diversity (?)
RMSD (.)
OFr
INTERNAL COVARIANT
INTERNAL
COVARIANT
INTERNAL
COVARIANT
CaHNNO;
0.06
0.70
0.90
19
35
65
0.32
0.30
C.H.N
0.80
0.70
0.29
Single-bag
CjHμO
0.06
0.90
0.80
2
0.42
0.22
0.05
0.50
0.80
85
0.80
0.76
C;HgN:O;
0.03
0.60
0.70
5
58
118
0.57
0.52
Multi-bag
0.54
0.78
0.89
22
19
42
0.04
0.04
Stochastic-bag
C:HnOO C,H,N,O,
0.05
0.40
0.60
10
26
59
0.65
0.71
0.03
0.10
0.80
Stochastic-bag (gen.)
C:Inn
0.00
0.00
0.13
0
38
68
n/a
1.02
C:H,N:O;
0.05
2015). The validity reported in Table li the ratio of valid molecules generated during offline evalu- ation at the end of training over 10 seeds. Two molecules are considered identical if their molecular graphs yield the same SMILES strings under RDKIT. The diversity shown in Table [ is the total number of unique and valid structures generated during offline evaluation during training over 10 seeds In the two stochastic- experiments, the agents are trained on bags of sizes from the interval [16, 22] sampled with B = C2HgN202 and C2H respectively. Finally, to assess the stability of the generated molecules, valid structures generated in the last iteration underwent struc- ture optimization using the PM6 method (see Appendix for details). Then, the root-mean-squal deviation of atomic positions (RMSD, Å) between the original and the optimized structure was computed. In Table [ the median RMSD is given per experiment.
Results are listed in Table We observe that COVARIANT significantly outperforms the other agents most experiments both in terms of validity and diversity. The difference is particularly large for the more challenging stochastic-bag tasks, where COVARIANT does similarly well as on the single- bag experiments. This finding is confirmed in Fig. (a), showing that the exact stoichiometry does not need to be known a priori for the agent to build valid molecules. Moreover, the structures generated by COVARIANT are overall slightly more stable compared INTERNAL. In contrast, OPT often fails to build valid structures. Inspection of the generated structures reveals that for larger bags the agent tends build multi-molecular clusters. which are considered invalid in this experiment. TTe stability for OPT is omitted as all of its valid structures are stable by definition.
Compared to graph-based approaches (e.g., Jin (201bb a Liu et al. (2018); De Cao & Kipf (011); Bradshaw et al. (2019b)), the average validity
(017): Bradshaw
(2019a) Li et al.
Foor a fairer comparison in the stochastic-bag task, we use the structures generated during offline evaluation, instead of those generated during training in Simm al. 202))
8

---

Published as a conference paper at ICLR 2021
and diversity achieved by COVARIANT still relatively low. This can partly be explained by the fact that state-of-the-art graph-based approaches have the strict rules of chemical bonding in organic molecules encoded into their models. But result, they are limited to generating single organic molecules and cannot build molecules for which these rules do not apply (e.g., hypervalent iodine compounds such IFs). In terms of stability, the supervised generative model by [Gebaurr et al. (2019) reported average RMSD of approximately 0.25 Å. While their approach and the considered molecules are significantly different from ours, this suggests that the generated structures more stable compared to COVARIANT. Nonetheless, the RL approach presented in this work remains particularly attractive if no dataset exists on which such supervised model can be trained.
Generalization To evaluate the generalization capabilities of all agents unseen bags, train on distribution over bags with B = C2H10 and O, and test on sets of larger, out- of-distribution bags {CoH1400g, O, O, C2H CgHisO),, and (CaH) CgH112N203, 2,H14 Oz2 respectively. As shown in Fig. o (b), COVARIANT obtains higher av- erage returns with lower variance compared to INTERNAL on C2H10O2, while performing only slightly worse than the agent trained directly the test bags (purple). Results for C2H&N,O2 are in Appendix G Although the difference performance seems chemical validity is often determined by the last 10% of the returns. Indeed, Table | and Fig. 9 in Appendix G highlight the higher quality of the structures generated by COVARIANT, indicating better generalization to unseen bags of larger size. OPT fails at this task.
marginal, we stress that
(a)
(b)
3
Average Return
3
2-
2
Covariant Covariant (Stoch.) Interual (Stoch.)
Covariani Covariant (Stoch.) Internal (Stoch.)
25
50
75
100
125
150
175
50
100
150
200
Steps x 1000
Steps 1000
Figure Average offline performance on the stochastic-bag task with 5 C2H10 evaluated on (a) C2H10O2 and (b) larger, unseen bags { (CaHl140s C7H1e0, C7H1g O) over 10 seeds. For comparison, we show an agent trained only on the test bags (purple). Error bars are two standard deviations. Molecular structures generated by COVARIANT (Stochastic) shown.
6 CONCLUSION
We proposed novel covariant actor-critic architecture based on spherical harmonics for design- ing highly symmetric molecules in 3D. We showed empirically that exploiting symmetries of the molccular design process improves the quality of the generated molecules and leads to better gener- alization. In future work, we aim employ more accurate quantum-chemice methods (e.g., density functional theory) required for building transition metal complexes or structures in which weak inter- molecular interactions are important. For that, however, the sample-efficiency of our agent needs to be improved. Finally, we aim to explore reward functions specifically tailored towards drug design.
ACKNOWLEDGEMENTS
We thank A. J. Tripp and K. Jensen for useful discussions and feedback. RP receives funding from iCASE grant #1950384 with support from Nokia. JMHL acknowledges support from Turing AI Fellowship under grant EP/V023756/1. This work has been performed using resources operated by the University of Cambridge Research Computing Service (funded by grant EP/P020259/1).

---

Published as a conference paper at ICLR 2021
REFERENCES
Brandon Anderson, Truong Son Hy, and Risi Kondor. Cormorant: Covariant Molecular Neural Networks. Advances in Neural Information Processing Systems 32, pp. 14537-14546. Curran Associates, Inc., 2019.
Harry Bateman. Higher Transcendental Functions, volume I-III. McGraw-Hill Book Company, New York, 1953.
Christopher M Bishop. Mixture density networks. Technical report, Aston University, 1994.
Christopher M. Bishop. Rejection sampling. In Pattern Recognition and Machine Learning, Infor- mation Science and Statistics, pp. 552--330. Springer, New York, 2009. ISBN 978N9-88-387-31077
Francesco Bosia, Tamara Husch, Alain C. Vaucher, and Markus Reiher. QCScine Sparrow: 2020.
Release
2.0.1. https://doi.org/
John Bradshaw, Matt J. Kusner, Brooks Paige, Marwin H. S. Segler, and José Miguel Hernández- Lobato. A generative model for electron paths. In International Conference Learning Repre- sentations, 2019a.
John Bradshaw, Brooks Paige, Matt J Kusner, Marwin Segler, and José Miguel Hernández-Lobato. A Model Search for Synthesizable Molecules. In Advances in Neural Information Processing Systems, pp. 7935-7947, 2019b.
Nathan Brown, Marco Fiscato, Marwin H.S. Segler, and Alain C. Vaucher. GuacaMol: Benchmark- ing Models for de Novo Molecular Design. J. Chem. Inf. Model., 1993)1009--1008 2019.
Nicola De Cao and Thomas Kipf. MolGAN: An implicit generative model for small molecular graphs. arXiv preprint arXiv:1805.11977 2018.
Mare Finzi. Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson. Generalizing Convo- lutional Neural Networks for Equivariance Lie Groups on Arbitrary Continuous Data. arXiv preprint arXiv:2002. /2880, 2020.
Fabian В Fuchs, Daniel E Worrall, Volker Fischer, and Max Welling. SE (3)-transformers: 3D roto-translation equivariant attention networks. arXiv preprini arXiv:2006. 10503, 2020.
Niklas W. A. Gebauer, Michael Gastegger, and Kristof T. Schütt,, Generating equilibrium molecules with deep neural networks. arXiv preprint arXiv:
2018.
Niklas W. A. Gebauer, Michael Gastegger, and Kristof T. Schütt. Symmetry-adapted generation of 3D point sets for the targeted discovery of molecules. In Advances Neural Information Processing Systems, pp. 7576, 2019.
Rafael GömezzBom Jennifer N. Wei, David Duvenaud, José Miguel Hernández-Lobate Benjamín Sánchez-L Sänchez-Lengelingg Dennis Sheherla, Jorge Aguilera-Iparraguire, Timothy D. Hirzel, Ryan P. Adams, and Alán Aspuru-Guzit Automatic Chemical Design Using a Data-Driven Con- tinuous Representation of Molecules. ACS Cent. Sci., 4(2):268-276, 2018.
Gabriel Lima Guimaraes, Benjamin Sanchez-Lengeling Carlos Outeiral, Pedro Luis Cunha Farias, and Alán Aspiru-Guzik. Objective-RReifforee Generative Adversarial Networks (ORGAN) for Sequence Generation Models. arXiv preprint arXiv: 1705. 10843, 2018.
J. D. Hunter. Matplotlib: A 2D graphics environment. Comput. Sci. Eng. 9(3):90-95, 2007.
Tamara Husch and Markus Reiher. Comprehensive Analysis of the Neglect of Diatomic Differential Overlap Approximation. Chem. Theory Comput., 14(10)5519-5119, 2018.
Tamara Husch, Alain C. Vaucher, and Markus Reiher. Semiempirical molccular orbital models based the neglect of diatomic differential overlap approximation. Int. J. Quantum Chem., 118 (24): 2018
S. Rao Jammalamadaka and György H. Terdik. Harmonic analysis and distribution-free inference for spherical distributions. J. Multivar. Anal, 171:436-451, 2019.
10

---

Published as a conference paper at ICLR 2021
Jan Jensen. XYZ2Mol. http://ggiibbboom///yssseggaaggp 2019.
Wengong Jin, Connor Coley, Regina Barzilay, and Tommi Jaakkola. Predicting Organic Reaction Outcomes with WeisfeilerLLemman Network. In Advances in Neural Information Processing Sys- tems, pp. 2607-2616, 2017.
Mathias S. Jørgensen, Henrik L. Mortensen, Søren A. Meldgaard, Esben L. Kolsbjerg, Thomas L. Jacobsen, Knud H. Sørensen, and Bjørk Hammer. Atomistic structure learning. J. Chem. Phys., 151(5):054113 2019.
Yeonjoon Kim and Woo Youn Kim. Universal Structure Conversion Method for Organic Molecules: From Atomic Connectivity Three-Dimensional Geometry. Bull. Korean Chem. Soc., 36(7): 1769-1777, 2015.
Risi Kondor and Shubhendu Trivedi. On the Generalization of Equivariance and Convolution in Neural Networks to the Action of Compact Groups. In International Conference on Machine Learning, pp. 2747-2755. PMLR, 2018.
Risi Kondor, Zhen Lin, and Shubhendu Trivedi. Clebsch- Gordan Nets: A Fully Fourier Space Spherical Convolutional Neural Network. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-BBanchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems 31, pp. 10117-10126. Curran Associates, Inc., 2018.
Gregory Landrum. RDKit 2019.09.3. http://www.edkit.rrg.. 2019. (Accessed: 22. January 2019).
V. I. Lebedev. Values of the nodes and weights of ninth to seventeenth order gauss-marks markov quadra- ture formulae invariant under the octahedron group with inversion. Zh. Vychisl. Mat. Mat. Fiz, 15(1):44-51, 1975.
I. Lebedev. Spherical quadrature formulas exact to orders 25-29. Sibirsk. Mat. Zh., 18(1):99-107, 1977.
Yibo Li, Liangren Zhang, and Zhenming Liu. Multi-objective de novo drug design with conditional graph generative model. J. Cheminf., 10(1)333, 2018a.
Yujia Li, Oriol Vinyals, Chris Dyer, Razvan Pascanu, and Peter Battaglia. Learning Deep Generative Models of Graphs. arXiv preprint /803.v:/003..3333 2018b.
Qi Liu, Miltiadis Allamanis, Marc Brockschmidt, and Alexander Gaunt. Constrained Graph Vari- ational Autoencoders for Molecule Design. In Advances Neural Information Processing Sys- tems, pp. 7795-7804, 2018.
Wes McKinney. Data Structures for Statistical Computing in Python. Stéfan van der Walt and Jarrod Millman (eds.), Proceedings of the 9th Python in Science Conference, pp. 51-56, 2010.
Søren Meldgaard, Henrik Mortensen, Mathias Jørgensen, and Bjørk Hammer. Structure prediction of surface reconstructions by deep reinforcement learning. J. Phys.: Condens. Matter, 32(40):404005, 2020.
Benjamin Kurt Miller, Mario Geiger, Tess E. Smidt, and Frank Noé. Relevance of Rotationally Equivariant Convolutions for Predicting Molecular Properties. arXiv preprint arXiv:2008, 08461, 2020.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Marc G. Belle- mare, Alex Graves, Martin Riedmiller, Andreas K. Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wier- stra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforeement learning. Nature, 518(7540):599-533 2015.
Daniel Neil, Marwin Segler, Laura Guasch, Mohamed Ahmed. Dean Plumbley, Matthew Sellwood, and Nathan Brown. Exploring Deep Recurrent Models with Reinforcement Leaminn for Molecule Design. OpenReview, 2018. URL htttp:://openreee ew wttpeewweeweeeteitnimmmoommm
11

---

Published as a conference paper at ICLR 2021
Mareus Olivecrona, Thomas Blaschke, Ola Engkvist, and Hongming Chen. Molecular de-novo design through deep reinforcement learning. J. Cheminf., 91)):48, 2017.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Vetoto Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch: An Imperative Style, HiighPeeformance Deep Learning Library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' A'ché-é-Boor Buch E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems 32, pp. 8026-8037. Curran Associates, Inc., 2019.
Mariya Popova, Olexandr Isayev, and Alexander Tropsha. Deep reinforcement learning for de novo drug design. Sci. Adv. 4(7)xeaap7885, 2018.
Evgeny Putin, Arip Asadulaev, Yan Ivanenkov, Vladimir Aladinskiy, Benjamin Sanchez-Lengelingg Alán Aspuru-Guzik, and Alex Zhavoronkov. Reinforced Adversarial Neural Computer for de Novo Molecular Design. J. Chem. Inf. Model., 5866:1194-1204, 2018.
Petra Schneider, W. Patrick Walters, Alleyn T. Plowright, Norman Sieroka, Jennifer Listgarten, Robert A. Goodnow, Jasmin Fisher, Johanna M. Jansen, José S. Duca, Thomas S. Rush, Matthias Zentgraf, John Edward Hill, Elizabeth Krutoholow, Matthias Kohler, Jeff Blaney, Kimito Funatsu, Chris Luebkemann, and Gisbert Schneider. Rethinking drug design in the artificial intelligence era. Nat. Rev. Drug Discovery, pp. 1-12, 2019.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 2017.
Marwin H. S. Segler, Thierry Kogej, Christian Tyrchan, and Mark P. Waller. Generating Focused Molecule Libraries for Drug Discovery with Recurrent Neural Networks. ACS Cent. Sci., 4(1): 120-131, 2018
Gregor NC Simm, Robert Pinsler, and José Miguel Hermández-Lobato. Reinforcement learning for molecular design guided by quantum mechanics. International Conference on Machine Learning, 2020.
abs hbtp://..s//oo...oo/ 07717
James P. Stewart. Optimization of parameters for semiempirical methods V: Modification of NDDO approximations and application to 70 elements. J. Mol. Model., 13(12)1 2007.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation- and translation 2018.
translationeqquivrraan equivariant neural networks for 3D point
clouds. arXiv prepint
Elise van der Pol, Daniel E Worrall, Herke van Hoof, Frans A Oliehoek, and Max Welling. MDP homomorphic networks: Group symmetries in reinforcement learning. arXiv preprint 10X1:22006.16988, 2020.
Maurice Weiler, Mario Geiger, Max Welling, Wouter Boomsma, and Taco S Cohen. 3D Steerable CNNs: Learning Rotationally Equivariant Features in Volumetric Data. In S. Bengio, H. Wal- lach, H. Larochelle, Grauman, N. Cesa-Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems 31, pp. 10381-10392. Curran Associates, Inc., 2018.
Jiaxuan You, Bowen Liu, Zhitao Ying, Vijay Pande, and Jure Leskovec. Graph Convolutional Policy Network for Goal-Directed Molecular Graph Generation. In Advances in Neural Information Processing Systems, pp. 6410-6421, 2018.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander J Smola. Deep sets. In Advances in Neural Information Processing Systems, pp- 3391-3401 2017.
Zhenpeng Zhou, Steven Kearnes, Li Li, Richard N. Zare, and Patrick Riley. Optimization of Molecules via Deep Reinforcement Learning, Sci. Rep., 91)::1110, 2019.
12

---

Published as a conference paper at ICLR 2021
A REWARD CALCULATION
In the reward function, the energy E has to be computed using quantum-chemical methods. For that, we use the fast semi-empirical Parametrized Method 6 (PM6) (Stewart 2007). In particular, we use the implementation in the software package SPARROW (Husch al. 2018 Bosta et al. (020). For each calculation, a molecular charge of zero and the lowest possible spin multiplicity are chosen. All calculations spin-unrestricted.
Limitations of semi-empirical methods are highlighted in, for example, recent work by Husch & Reiher 2018). More accurate methods such approximate density functionals need to be employed especially for systems containing transition metals.
Further, we enforce that atoms are not placed too close (< 0.6 Å) nor too far away from each other (> 2.0 A). the agent places an atom outside these boundaries, the minimum reward of -0.6 is awarded and the episode terminates. Further, the environment encourages the agents to build single molecular structures by terminating the episode and return a reward -0.6 elements forming stable bimolecular compounds (e.g. Hz) placed too far away from other atoms the canvas.
B SPHERICAL HARMONICS
The spherical harmonics form orthonormal basis of the Hilbert space of square-integrable func- tions L.((5)). The first few spherical harmonics are given by:
1 (0,+)) 2
/ V 8T
3 4
3 8
(0, =
sin de +p (1) V
cos 0, () -
sin de (8)
15
15
Y2 (0,,0) V 32
sin De -120 Y2 (0. .p)
V 8e
cos sin de -ip
(1,)))
(3 cos2 0 - 1)
(9)
16T
15 cos O sin de#f,
(1. V
(2) = 32
sin 02.
8
The spherical harmonics are normalized such that:
L Jo |Y (0, s2 sin 0dodp =
Vl,m.
(10)
C
CALCULATION OF INVARIANT FEATURES
One can obtain scalar invariats from the covariant features f Anderson et al. 2019
Take the component 0: E fe
Calculate the scalar product with itself: (fe) Re(Ea(f)) + Im(E2(f)), where &
Emm- fm f-m.
Calculate the iov3)-invrrian norm: Ex(fe) Ein S (fre) where denotes the complex conjugate.
The invariant components are then concatenated fim +- Tin = Ei( @ (
E2())) os(f)))
13

---

Published as a conference paper at ICLR 2021
D PROBABILITY DISTRIBUTION FOR ORIENTATION
In the main paper, propose the expression in Eq. alternative, equally valid expression
5 for the distribution p( An
e, =
Σ 2=
YE (T)
(11)
where the term 1/vk with k = Et= / 12 imentally that agent using this expression performs worse when generating molecular structures featuring complex geometries. We hypothesize that this is because the distribution cannot get peaked enough for Lmax S 5. As larger Lmox would result significant increase in computational com- plexity, we chose the expression in Eq. (5) over that in Eq.
normalizes the distribution. We found exper-
The normalization constant Z in Eq. (5) is estimated via Lebedev quadrature with 1730 angular grid points (Lebedev 1975 1977). We sample from the distribution in Eq. (5) using rejection sampling
Bishop 2009) with unitorm proposal distribution g(Ø) 1. In rejection sampling, one first draws a sample from Pp from (.F) Then, one generates random number ug from the uniform distribution over [0, Mf)aa) where M is such that Mq(#) 2 p(ö|d, f, s). We determine M by evaluating p(ī|d, s) ug p d, e, s) then the sample is accepted.
4x
uniform grid on 8 employing Fibonacci 'sunflower' grid. Finally,
We ran experiment to compare the COVARIANT agent described in the main paper with an agent employing Eq. for the distribution p(old, e, s). In Fig. 7 the hypothesis that the distri- bution in Eq. 11 cannot get narrow enough is confirmed. After 40 000 steps, the online performance of the alternative agent COVARIANT (ALT.) converges to around 0.5, which is significantly lower compared COVARIANT. Note that the difference is smaller when considering the offline return because the estimated global mode for each distribution could still be similar.
1.0
Average Online Return
-0.2
Covariant Covarianl (Alt.)
-0.4
Steps x 1000
Figure 7: Comparison of COVARIANT agent (red) with an agent employing alternative distribution for the orientation (green). The average online performance the single-bag task with B SF across seeds is shown. Dashed lines denote the optimal return for the experiment. Error bars indicate two standard deviations.
E BASELINES
E. OPT AGENT
Below, we detail the algorithm of the OPT agent. At the beginning of each experiment, the agent given a canvas Co. bag Bg, and a black-box function that can compute the energy E(C) and the atomic forees F(C) for given canvas. We assume total charge of zero and low-spin configura- tion. At the end of each experiment, we compute the total reward obtained for the final structure on canvas Cr and report the total number of energy and gradient computations.
14

---

Published as a conference paper at ICLR 2021
1. If the canvas Ce is not empty, randomly choose focal atom f from the list of available atoms on the canvas. An atom is considered available if its number of neighbors is less than predefined number that depends on its element (e.g., one for hydrogen and four for carbon). Two atoms on the canvas neighbors if their Euclidean distance is below 1.5 A. If there are no available atoms on the canvas, a focal atom is randomly chosen from the list of atoms the canvas.
2. Randomly choose element e, from the bag B,-
3. Randomly place the atom ag = 2e) on a sphere with radial distance d = 1.1 Å around zy obtain C++1, If the canvas empty, place the atom the origin.
4. Optimize only the position of a, using F to obtain 1,
5. Compute the energy difference (t ((C++1,op)) [E(C) Elfe, 0})l.
6. If AE()) > return e, the bag and go back step 1.
7. Optimize canvas C1+1.op1 using F to obtain Cz+1-
8. Increment by 1.
9. If the bag is not empty, go back to step 1.
In the experiments, the different agents need to be given a comparable computational budget to ensure meaningful comparison of their performance. This is difficult as they use different compu- tational resources: OPT runs CPU whereas INTERNAL and COVARIANT perform many of their computations on a GPU. However, we found experimentally that the quantum-chemical calculations are the most computationally expensive ones. These calculations are performed in the same way for all approaches. Therefore, we believe that by granting each approach the same number of PM6 calculations achieve fair comparison.
E.2 OPTIMAL RETURN
The optimal return for the single-bag tasks derived in the following way. First, obtained molecular structures for the complexes SOF,, IFs, SF,, and SOF,- Subsequently, we performed a structure optimization using the PM6 method. Since the undiscounted return is path-independent, we determined the return R(s) by computing the total interaction energy in the canvas C, i.e.
Σ Elles 0}) }
[C|
} -B(C). E(C).
R(s))
(12)
F EXPERIMENTAL DETAILS
F.1 COMPUTING INFRASTRUCTURE
Experiments were run on an Intel Xeon E5-2650 v4 2.2GHz 12-core processor (96GiB RAM) and an Nvidia P100 GPU (16GiB). Our agent is implemented in the deep learning framework PyTorch Paszke et al. 2019). Data analysis was performed with the Python libraries matplotlic otb (Hunter 2007) and pandas (McKinney 2010).
F.2 IMPLEMENTATION DETAILS
where the dimensions of sinv and iny f.c are possible, we made similar architectural choices as Simm et al. (2020), e.g. regarding the number of hidden units/layers, activation functions, and initialization schemes. We initialize the biases of each network with 0 and each weight matrix as a (semi- Jothhgoolll matrix. After each hidden layer, a ReLU non-linearity is employed. As explained in the main text, both MLP, and MLP, function to guarantee that only valid actions are chosen. To model the distance d, we employ a Gaus- sian mixture model consisting of M Gaussians. As treat the standard deviations (am m==1
The model architecture is summarized in Table
in = (Lmax - + 2) 2 and diny = (Lmax + 2) To 2, respectively.
masked softmax activation
15

---

Published as a conference paper at ICLR 2021
Table 2:
Model architecture for actor and critic networks.
Network Dimensions per layer
Output activation
MLP, dtt, 128, dive 128, Emax
masked softmax
MLP.
masked softmax
MDN
d 128,
linear (mm) tanh (fmm)
MLP &
dinv 128, 128 128, 128,
linear linear
MLP,
as global parameters, the MDN has 6 outputs. Further, we rescale the means Mmm E [-1,1] to Mm E [dminn dmax]. If the sampled distance is negative, clip the value 0.001.
Hyperparameters for CORMORANT are listed in Table 3 our experiments, we found it important to use multiple filters per element (e.g. 4) and to set Lmax 4. This gives the model enough flexibility to represent complex spherical distributions while remaining computationally tractable. For more details on CORMORANT, the original work (Anderson al. 2019). Further hyperpa- rameters used our experiments are Table 4 PPO is known to be relatively robust with respect to the choice of hyperparameters, and found the default values be sufficient in most cases. Within the actor, the scaling parameter 3 is important to avoid that the spherical distribution ap- proaches delta distribution. Note that values of / can vary significantly across experiments and might require some tuning. Lastly, the right number of samples for the global mode estimation of the spherical distribution generally depends on the shape of the distribution. In particular, we would expect that more samples are required as the distribution becomes more peaked. Since we avoid pathological behaviors by scaling the distribution with 8, we found 1024 to be sufficient for all our experiments.
Table 3: Hyperparameters for CORMORANT (Anderson et al. 2019) used in all experiments.
Hyperparameter
Value
Number of Clebsch-Gordal layers
3
Lmnx in spherical harmonics series expansion
4
Number of filters element
4
Number of filters
Te N
Table 4: Hyperparameters for the single-bag, multi-hag and stochastic-bag tasks. Values in paren- theses were only used for the single-bag task. For further details how the PPO hyperparameters are defined, please refer to Schulman et al. (2017).
Hyperparameter
Search Set
Value
Range
[dmin dmux] (A)
[0.95, 1.80]
Number of workers
10
PPO clipping
0.2
PPO gradient clipping
0.5
PPO GAE parameter
0.95
PPO value function coefficient C1
PPO entropy coefficient C2
(0.01, 0.05)
0.01
Number of optimization epochs
Adam stepsize
3 10-4
Discount factor y
0.99
Time horizon T
20 - [8]
Distance clipping
0.001
Scaling factor /
[- -100,,,,,,,,,11000 100
100 (-10)
Number of samples S for mode estimation
1024
16

---

Published as a conference paper at ICLR 2021
G
ADDITIONAL RESULTS
(a)
(b)
Average Return
3
3
2
Covariant Covariant (Stoch.)
Covariant
Covariant (Stoch.)
Internal (Stoch.)
Internal (Stoch.)
Steps x 1000
Steps x 1000
Figure 8: Average offline performance on the stochastic-bag task with 5 = C7HgN2Op evaluated on (a) C2H2N2O2 and (b) larger, unseen bags (43 O, C7 CgH1sO) across 10 seeds. For comparison, we show an agent that is trained only on the test bags (purple). Error bars indicate two standard deviations.
In Fig. 9 a selection of molecular structures generated by the three agents during offline evaluation shown. The agents are trained distribution over bags with B' 17 and tested on sets larger, out-of-distribution bags [C3H 16 17H CsHisO) (cf. Fig. 6 in the main text). From visual inspection of the structures, can be seen why OPT fails at this task: it tends to generate molecular clusters, often containing Hg. Similarly, INTERNAL commonly builds H2O molecules instead of constructing single organic molecule out of the atoms in the bag. By contrast, COVARIANT often builds valid molecules. Further, its generated structures are more often "branched" than those of INTERNAL, indicating higher degree of complexity.
(a)
(b)
(c)
g
8
Figure 9: Selection of molecular structures generated by (a) OPT, (b) INTERNAL, and (c) COVARI- ANT during the last offline evaluation. The agents are trained distribution over bags with B* 1o O2 and tested on the out-of-distribution bags (CoHi4Og, C,H1, C2H C.H
Next, we assess the ability of COVARIANT to generate solvation clusters type molecular struc- ture that cannot be built with graph-based approaches. Following [Sim et al. (2020), we task the agent to place water molecules around a formaldehyde molecule that is already on the canvas at the beginning of each episode. In addition, the reward function is augmented with penalty term for placing atoms far away from the center, i.e. r(st, a4) = - - - ||| where p is a parameerr that is set to 0.01 (see Simm et al. (2020) for details). Therefore, the agent needs to place the water molecules such that hydrogen bonds can be formed between water molecules and between water molecules and the solute.
17

---

Published as a conference paper at ICLR 2021
2.0
Average Return
S
0.5
-0.5
Covariant Invariant
I
50
75
100
125
150
Steps 1000
Average offline performance on solvation task with 5 H2O molecules and formalde- hyde the solute across 10 seeds. Error bars show two standard errors. The dashed line denotes the optimal return. A selection of molccular clusters generated by the COVARIANT agent is shown.
Figure 10:
From Fig. 10 it can be seen that COVARIANT can solve this task by constructing stable H₂O molecules and placing them in the vicinity of the solute. From visual inspection of the generated structures, can be observed that in many cases COVARIANT arranges the molecules such that in- termolecular bonds can be formed. However, should be noted that the quantum-chemical method used in the reward function is not very well suited for modeling these interactions. Finally, Fig, I0 shows that while INTERNAL learns faster at the beginning of training, COVARIANT slightly out- performing INTERNAL towards the end.
H RUNTIME EVALUATION
We compared the runtimes between OPT. INTERNAL, and COVARIANT. For instance, for the single- bag task with the bag C,HsNO,, T 240 steps of the last rollout took COVARIANT and INTERNAL on average 12 and 11 seconds (s), respectively. The final offline evaluation took the agents on average 4 and 1s, respectively. This speed difference mainly due to the relatively slow rejection sampling procedure in COVARIANT. Each iteration, policy optimization took on average 2 and 6s for the agents COVARIANT and INTERNAL, respectively. In this case, INTERNAL slower than COVARIANT as it performed around twice many epochs during optimization due early stopping. Since there is no training for OPT, this agent was overall faster than the others. Further, we note that the largest fraction of time spent the quantum- chenttmmm calculations which are the same for all agents. The time the quantum-chemical calculation takes to converge depends not only on the size but also the geometry of the input structure. The entire experiment took COVARIANT approximately 4 hours, INTERNAL 5 hours, and OPT 3 hours.
18


Reviews:
[{"review_id": "vePHhu3GSg", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper proposes an actor-critic neural network architecture for autoregressive generation of 3D molecular structures with reinforcement learning (RL). It builds upon the RL approach by Simm et al. (2020) which makes use of internal coordinates in order to deal with the symmetries that occur when placing atoms in the molecular design process.\nThe main contribution is the introduction of a covariant state-action representation extracted with Cormorant (Anderson et al., 2019) and a covariant, spherical distribution for the orientation of newly placed atoms (w.r.t. a focal atom) replacing the internal coordinates of the former approach. It is shown that these architectural and conceptual changes allow to build highly symmetric structures (which were problematic with the former approach) and improve the generalization as well as the quality of generated molecules in several tasks from the RL environment MolGym.\n\nThe paper has a clear structure, is well-written and includes several experiments to verify the main claims. To the best of my knowledge, the proposed architecture is the first model for generation of 3d molecules that uses covariant features to leverage the rotational symmetries of molecules. This seems to improve the model by Simm et al. (2020), especially when it comes to symmetric structures such as $SF_6$. The proposed covariant formulation could be beneficial for other existing or future generative models for 3d molecules, as well.\n\nThe experimental section could be improved by additional experiments to allow for comparison between the proposed architecture and other approaches for generation of molecular graphs and 3d molecules (Liu et al , Gebauer et al. 2019).  In many real-world molecular design tasks, a variety of (possibly diverse) molecules that exhibit certain target properties are desirable. It would be interesting to see whether the proposed model can also be applied for these purposes instead of the rather limited task to generate stable molecule from a predefined bag.\n\nPros\n----\n- well motivated neural network design incorporating the necessary symmetries\n- extended experiments show generalization performance\n\nCons\n----\n- It does not become quite clear, but as far as I understand only the most probable action is chosen during offline evaluation (Sec 3.3). An additional experiment on how the validity and diversity depend on the degree of exploration would be useful.\n- The ms could be improved by additional experiments that allow for comparison with previous molecule generation models on different tasks (see Liu et al 2018, Mansimov et al 2019, Gebauer et al 2019).\n\nUpdate:\nI raised my score to reflect the added metric and experiments that improved the paper.", "extra_scores": {}, "extra_text": {}}, {"review_id": "_2GJs9eji1t", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This work presents an approach for 3D molecular design using reinforcement learning that exploits rotational symmetries in molecular conformations. The formulation involves an MDP that selects atoms from a “bag” and positions them in 3D space. The reward function is based on PM6 energies to encourage the generation of low-energy structures. The manuscript is well-written and provides ample context and appropriate references.\n\nThe generative formulation directly builds on that of Simm et al. 2020 to choose a focal atom, choose an element to add, choose a distance, then choose an orientation. The consideration of rotational covariance during new atom placement and the use of Cormorant (Anderson et al. 2019) for embedding seem to be the primary novelties of this work. Thus, Simm et al.’s approach relying on an internal coordinate representation serves as the primary baseline in the empirical evaluations here; “Opt” is meant to serve as an essentially-random baseline and performs predictably poorly. One could even envision using a 2D generative model with a heuristic conformer generation scheme followed by geometry optimization using PM6 forces as a stronger baseline approach.\n\nThe benefit of Covariant over Internal in terms of actual optimization performance (Figure 5, Figure 6) appears marginal. Did the authors perform any experiments to optimize more than just PM6 energy (e.g., a computed property) to compare the two methods? It’s not clear how significant the observed differences are and to what extent achieving rotational covariance helps in this task. An ablation study that decouples the effects of the rotationally-invariant critic and the rotationally-covariant actor could be informative. How sensitive are either of these models to hyperparameter selection?\n\nCovariant does appear to generate a greater number of valid structures (Table 1), although this evaluation maps the 3D coordinates back to a SMILES string to check for validity and uniqueness, which seems not to be the intended application of this work. For that evaluation, one could consider using purely graph- or SMILES-based approaches and more traditional generative methods. \n\nA primary advantage of this formulation seems to be the ability to generate molecular clusters (in contrast to 2D approaches) and periodic materials (in contrast to Simm et al.). The current set of experiments do not demonstrate this ability in my opinion. There also is not evidence that “building in this inductive bias enables [one] to generate geometrically complex molecular structures that were unattainable with previous approaches.”\n\nThe integration of rotational covariance with 3D molecular generation is elegant and, to my knowledge, novel. I don’t disagree with the premise that one *should* strive to have these invariances, but I don’t believe that the current set of experiments actually demonstrate the superiority. \n", "extra_scores": {}, "extra_text": {}}, {"review_id": "QLINRb5DBPX", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "### Summary of the paper\nThe paper utilizes a actor-critic framework for 3D molecular design. The central part of the approach is a rotation equivariant network (Comorant). For each atom, Comorant learns a state representation so that it is equivariant under rotation. The method is evaluated on 9 different molecules and it outperforms Simm et al.'s method in terms of validity and diversity.\n\n### Strength and weakness\n1. The method adopts a better model architecture (Comorant) to learn the state representation. However, Comorant is developed by Anderson et al. Therefore it cannot be counted as the contribution of this paper. The actor-critic formulation is also standard in RL. The sequential decision process is also similar to Simm et al. The technical innovation is not original enough.\n\n2. The evaluation protocol is problematic. In table 1, authors only report validity and diversity of generated compounds. The validity is defined as \"successfully parsed by RDKit\". To my knowledge, RDKit validity checking is based on 2D constraints such as valence, aromaticity and kekulization. It does not capture 3D information at all. Are the generated compound stable? What is the RMSD of generated compounds? Simm et al. 2020 reported RMSD to measure the structural stability. Why RMSD metric is missing? This is important because 2D graph generation models already satisfy validity quite well (100% validity for even large molecules). \n\n3. Motivation is not clear. In the paper, authors state that the choice of focal atom, element and distance have to be *invariant* to rotation. It seems like invariant representation is sufficient. Moreover, the advantage over Simm et al. is not clear in section 4. Why highly symmetric states are problematic for prior work? How does this method solves this issue? I am afraid that the experiment section does not address this, since most test cases are not \"highly symmetric\" to my knowledge.\n\n### Overall evaluation\nI vote for rejection. My major reason is the problematic evaluation protocol. RMSD must be added to evaluate the stability of compounds. Moreover, I would like to see evaluation on \"highly symmetric\" cases. I think it's important since authors state that this is the major limitation of prior work.\n\n### Question\n1. Is the method scalable to large molecules? How is the runtime of your model compared to Simm et al.?\n\n### Post rebuttal feedback\nIt's good to see that RMSD experiments are added and the results are better than Simm et al. Therefore, I am raising my score to 6. I also realized that the validity calculation is different from standard graph generation methods. The validity results now look reasonable to me.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper provides a new covariant approach to 3D molecular generation motivated by the desire handle compounds with symmetries. To this end, the method uses equivariant state representations for autoregressive generation, built largely from recently proposed covariant molecular networks (comorant), and integrating such representations within an existing actor-critic RL generation framework (Simm et al). The selection of focal atom, element to add, and the distance are realized in an equivariant manner while the compound valuation remains invariant to rotation. The approach is clean and well-executed. The authors added additional experiments (e.g., RMSD demonstrating stability of generated compounds) to further reinforce the case for the method. 

Decision: Accept (Poster)