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