Symmetry-Aware Actor-Critic for 3D Molecular Design

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 . 1 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 a plethora of machine learning approaches to accelerate this search , including generative models based on variational autoencoders ( VAEs ) ( Gómez-Bombarelli et al. , 2018 ) , recurrent neural networks ( RNNs ) ( Segler et al. , 2018 ) , and generative adversarial networks ( GANs ) ( De Cao & Kipf , 2018 ) . However , the reliance on a sufficiently large dataset for exploring unknown regions of chemical space is a severe limitation of such supervised models . Recent RL-based methods ( e.g. , Olivecrona et al . ( 2017 ) , Jørgensen et al . ( 2019 ) , Simm et al . ( 2020 ) ) mitigate the need for an existing dataset of molecules as they only require access to a reward function . 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 . It is unsuitable for encoding metals and molecular clusters as it lacks information about the relative position of atoms in 3D space . Further , geometric constraints on the design process can not be included , e.g . those given by the active site of an enzyme . A more general representation closer to the physical system is one in which a molecule is described by its atoms ’ positions in Cartesian coordinates . However , it would be very inefficient to naively learn a model based on this representation . That is because molecular properties such as 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 . In this work , we develop a 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 . 1 ( 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 auto-regressive policy that maintains the desired covariance . Building in this inductive bias enables us to generate molecular structures with more complex coordination geometry than the class of molecules that were attainable with previous approaches . Finally , we perform experiments on several 3D molecular design tasks , where we find that our approach significantly improves the generalization capabilities of the RL agent and the quality of the generated molecules . In summary , our 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 an actor-critic neural network architecture with a rotationally covariant auto-regressive policy , where the orientation of the atoms to be placed is modeled through a flexible distribution based on spherical harmonics ; • we demonstrate the benefits of our approach on several 3D molecular design tasks , including a 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 to maximize its reward . Formally , such an environment is described by a Markov decision process ( MDP ) M = ( S , A , T , µ0 , γ , r ) with states st ∈ S , actions at ∈ A , transition dynamics T : S × A 7→ S , initial state distribution µ0 , discount factor γ ∈ ( 0 , 1 ] , and reward function r : S × A 7→ R. The goal is to learn a stochastic policy π ( at|st ) that maximizes the expected discounted return J ( θ ) = Es0∼µ0 [ V π ( s0 ) ] , where the value function V π ( st ) = Eπ [ ∑T t′=t γ t′r ( st′ , at′ ) |st ] is defined as the expected discounted return when starting from state st and following policy π . Following Simm et al . ( 2020 ) , we design molecules by iteratively picking atoms from a bag and positioning them on a 3D canvas . Such a sequential decision-making problem is described by an MDP where the state st = ( Ct , Bt ) comprises both the canvas Ct and the bag Bt . The canvas Ct = C0∪ { ( ei , xi ) } t−1i=0 is a set of atoms with chemical element ei ∈ { H , C , N , O , . . . } and position xi ∈ R3 placed up to time t − 1 , where C0 can either be empty or contain a set of initially placed atoms . The number of atoms on the canvas is denoted by |Ct| . The bag Bt = { ( e , m ( e ) ) } is a multi-set of atoms yet to be placed , where m ( e ) is the multiplicity of the element e. Each action at = ( et , xt ) consists of the element et ∈ Bt and position xt ∈ R3 of the next atom to be added to the canvas . Placing an atom through action at in state st is modeled by a deterministic transition function T ( st , at ) that yields the next state st+1 = ( Ct+1 , Bt+1 ) with Bt+1 = Bt\et . The reward function r ( st , at ) = −∆E ( st , at ) is given by the negative energy difference between the resulting structure described by Ct+1 , and the sum of energies of the current structure Ct and a new atom of element et placed at the origin , i.e . ∆E ( st , at ) = E ( Ct+1 ) − [ E ( Ct ) + E ( { ( e,0 ) } ) ] . 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 , 2007 ) as implemented in SPARROW ( Husch et al. , 2018 ; Bosia et al. , 2020 ) ; see Appendix A for details . An example of a rollout is shown in Fig . 1 ( b ) . At the beginning of the episode , the agent observes the initial state ( C0 , B0 ) ∼ µ0 ( s0 ) , e.g . C0 = ∅ and B0 = SOF41 . The agent then iteratively constructs a molecule by placing atoms from the bag onto the canvas until the bag is empty.2 2.2 ROTATIONALLY COVARIANT NEURAL NETWORKS . A function f : X 7→ Y is invariant under a transformation operator Tg : X 7→ X if f ( Tg [ x ] ) = f ( x ) for all x ∈ X , g ∈ G , where G is a mathematical group . In contrast , f is covariant with respect to Tg if there exists an operator T ′g : Y 7→ Y such that f ( Tg [ x ] ) = T ′g [ f ( x ) ] . To achieve rotational covariance , it is natural to work with spherical harmonics . They are a set of complexvalued functions Y m ` : S2 7→ C with ` = 0 , 1 , 2 , . . . and m = − ` , − ` + 1 , . . . , ` − 1 , ` on the unit sphere S2 in R3 . The first few spherical harmonics are in Appendix B . They are defined by Y m ` ( ϑ , ϕ ) = ( −1 ) m √ 2 ` + 1 4π ( ` −m ) ! ( ` +m ) ! Pml ( cos ( ϑ ) ) e imϕ , ϕ ∈ [ 0 , 2π ] , ϑ ∈ [ 0 , π ] , ( 1 ) where Pm ` denotes the associated normalized Legendre polynomials of the first kind ( Bateman , 1953 ) , and each Y m ` is normalized such that ∫∫ |Y m ` ( ϑ , ϕ ) |2 sinϑdϑdϕ = 1 . Any square-integrable function f : S2 7→ C can be written as a series expansion in terms of the spherical harmonics , f ( x̃ ) = ∞∑ ` =0 ∑̀ m=− ` f̂m ` Y m ` ( x̃ ) , ( 2 ) where x̃ = ( ϑ , ϕ ) ∈ S2 . The complex-valued coefficients { f̂m ` } are the analogs of Fourier coefficients and are given by f̂m ` = ∫ f ( x̃ ) Y m∗ ` ( x̃ ) Ω ( dx̃ ) . Such a function f can be modeled by learning the coefficients { f̂m ` } 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 explicitly corresponds to a subset of atoms in the molecule . The input of CORMORANT is a spherical function f0 : S2 7→ Cd and the output is a collection of vectors f̂ = { f̂0 , f̂1 , . . . , f̂L } , where each f̂ ` ∈ τ × ( 2 ` + 1 ) is a rotationally covariant vector with τ channels . That is , if the input is rotated by R ∈ SO ( 3 ) , then each f̂ ` transforms as f̂ ` 7→ D ` ( R ) f̂ ` , where D ` ( R ) : SO ( 3 ) 7→ C ( 2 ` +1 ) × ( 2 ` +1 ) are the irreducible representations of SO ( 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 π ( a|s ) with actions a = ( e , x ) that is covariant under translation and rotation with respect to the position x , i.e. , x 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 e , i.e . the chosen element remains unchanged under such transformations ( see Fig . 1 ( 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 focal atom . If the next atom is placed relative to the focal atom , covariance under translation of x is automatically achieved and only the rotational covariance remains to be dealt with . As shown in Fig . 2 , we model the action a through a sequence of sub-actions : ( 1 ) the index f ∈ { 1 , . . . , |C| } of the focal atom around which the next atom is placed,3 ( 2 ) the element e ∈ { 1 , . . . , Ne } of the next atom from the set of available elements , ( 3 ) a distance d ∈ R+ between the focal atom and the next atom , and ( 4 ) the orientation x̃ = ( ϑ , ϕ ) ∈ S2 of the atom on a unit sphere around the focal atom . Denoting xf as the position of the focal atom , we obtain action a = ( e , x ) by mapping the local coordinates ( x̃ , d , f ) to global coordinates x = xf + d · x̃ , where 1Shorthand for { ( S , 1 ) , ( O , 1 ) , ( F , 4 ) } . 2Hereafter , we drop the time index when it is clear from the context . 3If the canvas C0 is empty , the agent selects an element e0 ∈ B0 and places it at the origin , i.e . a0 = ( e0,0 ) . x is now covariant under translation and rotation . We choose these sub-actions using the following auto-regressive policy : π ( a|s ) = π ( x̃ , d , e , f |s ) = p ( x̃|d , e , f , s ) p ( d|e , f , s ) p ( e|f , s ) p ( f |s ) . ( 3 ) A novel actor-critic neural network architecture that implements this policy is illustrated in Fig . 3 . In the following , we discuss its state embedding , actor , and critic networks in more detail .

Summary: 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.