Snowflake: Scaling GNNs to high-dimensional continuous control via parameter freezing

1 Introduction . Whereas many traditional machine learning models operate on sequential or Euclidean ( grid-like ) data representations , GNNs allow for graph-structured inputs . GNNs have yielded breakthroughs in a variety of complex domains , including drug discovery [ 33 , 50 ] , fraud detection [ 56 ] , computer vision [ 49 , 43 ] , and particle physics [ 18 ] . GNNs have also been successfully applied to reinforcement learning ( RL ) , with promising results on locomotion control tasks with small state and action spaces . Not only are GNN policies as effective as MLPs on certain training tasks , but when a trained policy is transferred to another similar task , GNNs significantly outperform MLPs [ 55 , 20 ] . This is largely due to the capacity of a single GNN to operate over arbitrary graph topologies ( patterns of connectivity between nodes ) and sizes without modification . However , so far GNNs in RL have only shown competitive performance with MLPs on lower-dimensional locomotion control tasks . For higher-dimensional tasks , one must therefore choose between superior training task performance ( MLPs ) and superior transfer performance ( GNNs ) . This paper investigates the factors underlying poor GNN scaling and introduces a method to combat them . We begin with an analysis of the GNN-based NERVENET architecture [ 55 ] , which we choose for its strong zero-shot transfer performance . We show that optimisation updates for the GNN ⇤Corresponding author . Now at Graphcore , Bristol †Now at Waymo , Oxford 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . policy have a tendency to cause excessive changes in policy space , leading to performance degrading . To combat this , current state-of-the-art algorithms [ 46 , 48 , 1 ] employ trust region-like constraints , inspired by natural gradients [ 2 , 23 ] , that limit the change in policy for each update . We outline how this policy instability can be framed as a form of overfitting—a problem GNN architectures like NERVENET are known to suffer from in supervised learning , and show that parameter regularisation ( a standard remedy for overfitting ) leads to a small improvement in NERVENET performance . We then investigate which structures in the GNN contribute most to this overfitting , by applying different learning rates to different parts of the network . Surprisingly , the best performance is attained when training with a learning rate of zero in the parts of the GNN architecture that encode , decode , and propagate messages in the graph , in effect training only the part that updates node representations . We use this approach as the basis of our method , SNOWFLAKE , which freezes the parameters of particular operations within the GNN to their initialised values , keeping them fixed throughout training while updating the non-frozen parameters as before . This simple technique enables GNN policies to be trained much more effectively in high-dimensional environments . Experimentally , we show that applying SNOWFLAKE to NERVENET dramatically improves asymptotic performance and sample complexity on such tasks . We also demonstrate that a policy trained using SNOWFLAKE exhibits improved zero-shot transfer compared to regular NERVENET or MLPs on high-dimensional tasks . 2 Background . 2.1 Reinforcement Learning . We formalise an RL problem as a Markov decision process ( MDP ) . An MDP is a tuple hS , A , R , T , ⇢0i . The first two elements define the state space S and the action space A . At every time step t , the agent employs a policy ⇡ ( at|st ) to output a distribution over actions , selects action at ⇠ ⇡ ( ·|st ) , and transitions from state st 2 S to st+1 2 S , as specified by the transition function T ( st+1|st , at ) which defines a probability distribution over states . For the transition , the agent gets a reward rt = R ( st , at , st+1 ) . The last element of an MDP specifies initial distribution over states , i.e. , states an agent can be in at time step zero . Solving an MDP means finding a policy ⇡⇤ that maximises an objective , in our case the expected discounted sum of rewards J = E⇡ [ P1 t=0 trt ] , where 2 [ 0 , 1 ) is a discount factor . Policy Gradients ( PG ) [ 52 ] find an optimal policy ⇡⇤ by doing gradient ascent on the objective : ✓t+1 = ✓t + ↵r✓J |✓=✓t with ✓ parameterising the policy . Often , to reduce the variance of the gradient estimate , one learns a value function V ( s ) = E⇡ [ P1 t=0 trt | s0 = s ] , and uses it as a critic of the policy . In the resulting actor-critic method , the policy gradient takes the form : r✓J ( ✓ ) = E⇡✓ [ P t A ⇡✓ t r✓ log ⇡✓ ( at|st ) ] , where A ⇡✓ t is an estimate of the advantage function A⇡t = E⇡ [ P1 t=0 trt | at , st ] E⇡ [ P1 t=0 trt | st ] [ 47 ] . 2.2 Proximal Policy Optimisation . Proximal policy optimisation ( PPO ) [ 47 ] is an actor-critic method that has proved effective for a variety of domains including locomotion control [ 17 ] . PPO approximates the natural gradient using a first order method , which has the effect of keeping policy updates within a “ trust region ” . This is done through the introduction of a surrogate objective to be optimised : J = E⇡✓0 min ✓ ⇡✓ ( a|s ) ⇡✓0 ( a|s ) A⇡✓0 ( s , a ) , clip ⇣ ⇡✓ ( a|s ) ⇡✓0 ( a|s ) , 1 ✏ , 1 + ✏ ⌘ A⇡✓0 ( s , a ) ◆ ( 1 ) where ✏ is a clipping hyperparameter that effectively limits how much a state-action pair can cause the overall policy to change at each update . This objective is computed over a number of optimisation epochs , each of which gives an update to the new policy ⇡✓ . If during this process a state-action pair with a positive advantage A⇡✓0 ( s , a ) reaches the upper clipping boundary , the objective no longer provides an incentive for the policy to be improved with respect to that data point . This similarly applies to state-action pairs with a negative advantage if the lower clipping limit is reached . 2.3 Graph Neural Networks . GNNs are a class of neural architecture designed to operate over graph-structured data . We define a graph as a tuple G = ( V , E ) comprising a set of nodes V and edges E = { ( u , v ) | u , v 2 V } . A labelled graph has corresponding feature vectors for each node and edge that form a pair of matrices LG = ( V , E ) , where V = { vv 2 Rp | v 2 V } and E = { eu , v 2 Rq | ( u , v ) 2 E } . For GNNs we often consider directed graphs , where the order of an edge ( u , v ) defines u as the sender and v as the receiver . A GNN takes a labelled graph G and outputs a second graph G0 with new labels . Most GNN architectures retain the same topology for G0 as used in G , in which case a GNN can be viewed as a mapping from input labels LG to output labels LG0 . A common GNN framework is the message passing neural network ( MPNN ) [ 14 ] , which generates this mapping using T steps or ‘ layers ’ of computation . At each layer ⌧ 2 { 0 , . . . , T 1 } in the network , a hidden state h⌧+1v and message m⌧+1v is computed for every node v 2 V in the graph . An MPNN implementation calculates these through its choice of message functions and update functions , denoted M⌧ and U⌧ respectively . A message function computes representations from hidden states and edge features , which are then aggregated and passed into an update function to compute new hidden states : m⌧+1v = X u2N ( v ) M⌧ ( h⌧u , h ⌧ v , eu , v ) , h ⌧+1 v = U ⌧ h⌧v , m ⌧+1 v , ( 2 ) for all nodes v 2 V , where N ( v ) = { u | ( u , v ) 2 E } is the neighbourhood of all sender nodes connected to receiver v by a directed edge . The node input labels vv are used as the initial hidden states h0v . MPNN assumes only node output labels are required , using each final hidden state hTv as the output label v0v . 2.4 NerveNet . NERVENET is an MPNN designed for locomotion control , based on the gated GNN architecture [ 32 ] . NERVENET uses the morphology ( physical structure ) of the agent as the basis for the GNN ’ s input graph G , with edges representing body parts and nodes representing the joints that connect them . NERVENET assumes an MDP where the state s can be factored into input labels V , which are fed to the GNN to generate output labels : V0 = NERVENET ( G , V ) . These are then used to parameterise a normal distribution defining the stochastic policy : ⇡ ( a|s ) = N ( V0 , diag ( 2 ) ) , where the standard deviation is a separate vector of parameters learned during training . Actions a are vectors , where each element represents the force to be applied at a given joint for the subsequent timestep . The policy is trained using PPO , with parameter updates computed via the Adam optimisation algorithm [ 25 ] . Internally , NERVENET uses an encoder Fin to generate initial hidden states from input labels : h0v = Fin ( vv ) . This is followed by a message function M⌧ consisting of a single MLP for all layers ⌧ that takes as input only the state of the sender node : m⌧+1v = P u2N ( v ) MLP ( h ⌧ u ) . The update function U⌧ is a single gated recurrent unit ( GRU ) [ 9 ] that maintains an internal hidden state : h⌧+1v = GRU m⌧+1v | h⌧v . NERVENET propagates through T layers of message-passing and node-updating , before applying a decoder Fout to turn final hidden states into scalar node output labels : v0v = Fout ( hTv ) . A diagram of the NERVENET architecture can be seen in Appendix A.4 , Figure 10 . Knee 0 Root Hip 0 Hip 1 Knee 1 Body 0 Knee n-1 Body n/2-1 Hip n-1 Hip n Knee n Figure 1 : A MuJoCo rendering of Centipede-20 and its corresponding morphological graph . 3 Analysing GNN Scaling Challenges . In this section , we use NERVENET to analyse the challenges that limit GNNs ’ ability to scale . We focus on NERVENET as its architecture is more closely aligned with the GNN framework than alternative approaches to structured locomotion control ( see Section 4 ) . We use mostly the same experimental setup as Wang et al . [ 55 ] , with details of any differences and our choice of hyperparameters outlined in Appendix A.2 . We focus on environments derived from the Gym [ 8 ] suite , using the MuJoCo [ 53 ] physics engine . The main set of tasks we use to assess scaling is the selection of Centipede-n agents [ 55 ] , chosen because of their relatively complex structure and ability to be scaled up to high-dimensional inputaction spaces . The morphology of a Centipede-n agent consists of a line of n/2 body segments , each with a left and right leg attached ( see Figure 1 ) . The graph used as the basis for the GNN corresponds to the physical structure of the agent ’ s body . At each timestep in the environment , the MuJoCo engine sends a feature vector containing the positions of the agent ’ s body parts and the forces acting on them , expecting a vector to be returned specifying forces to be applied at each joint ( full details of the state representation are given in Appendix A.2 ) . The agent is rewarded for forward movement along the y-axis as well as a small ‘ survival ’ bonus for keeping its body within certain bounds , and given negative rewards proportional to the size of its actions and the magnitude of force it exerts on the ground . Existing work applying GNNs to locomotion control tasks avoid training directly on larger agents , i.e. , those with many nodes in the underlying graph representation . For example , Wang et al . [ 55 ] state that for NERVENET , “ training a CentipedeEight from scratch is already very difficult ” . Huang et al . [ 20 ] also limit training their SMP architecture to small agent types .

Summary: The authors propose a framework to address the challenge of scaling the application of graph neural network (GNN) based policies to locomotion challenges in reinforcement learning. The authors' hypothesize that the scaling challenges of GNN policies relate to overfitting of internal networks within the GNN architecture during RL training and unstable policy updates during on-policy RL training using PPO. The authors then introduce their framework, named Snowflake, to address this challenge which freezes distinct parts of the GNN during training to prevent overfitting and unstable policy updates. The authors demonstrate their framework by applying NerveNet, a GNN based policy network, to the centipede environment in the Gym + Mujoco suite. In their experiments and subsequent analysis, the authors show supporting evidence for their claim that NerveNet doesn't scale to large locomotion challenges and that the policy updates are unstable as suggested by large KL divergence between the updates when compared to regular MLP policy updates. Subsequently, the authors show how common overfitting mitigations fall short of significantly improving the performance of NerveNet in the centipede setting and provide supporting evidence that the Snowflake increases the performance of NerveNet to match regular MLP policy performance in a variety of Mujoco tasks, including various sizes in the centipede environment. The authors also claim, and support that claim, that GNN based policies, such as NerveNet, provide better generalization abilities compared to MLP based policies.