Snowflake: Scaling GNNs to high-dimensional continuous control via parameter freezing
Authors: Charlie Blake
Vitaly Kurin
Maximilian Igl
Shimon Whiteson
Venue: NeurIPS 2021 Poster (2021)
Keywords: reinforcement learning
graph neural networks
locomotion control
parameter freezing

Abstract: Recent research has shown that graph neural networks (GNNs) can learn policies for locomotion control that are as effective as a typical multi-layer perceptron (MLP), with superior transfer and multi-task performance. However, results have so far been limited to training on small agents, with the performance of GNNs deteriorating rapidly as the number of sensors and actuators grows. A key motivation for the use of GNNs in the supervised learning setting is their applicability to large graphs, but this benefit has not yet been realised for locomotion control. We show that poor scaling in GNNs is a result of increasingly unstable policy updates, caused by overfitting in parts of the network during training. To combat this, we introduce Snowflake, a GNN training method for high-dimensional continuous control that freezes parameters in selected parts of the network. Snowflake significantly boosts the performance of GNNs for locomotion control on large agents, now matching the performance of MLPs while offering superior transfer properties.

Snowflake: Scaling GNNs to high-dimensional continuous control via parameter freezing
Charlie Blake*
Vitaly Kurin University of Oxford vitaly. kurinccs, ox ac.
University of Oxford thecharlleblaxee@maal.ccoo
Maximilian Igl University of Oxford maximilian. maximilian.i110@maiiii com
Shimon Whiteson
University of Oxford shimon whiteson&cs ox ac uk
Abstract
Recent research has shown that graph neural networks (GNNs) learn policies for locomotion control that are as effective as a typical multi-layer perceptron (MLP), with superior transfer and multi-task performance 55 20 However, results have so far been limited to training on small agents, with the performance of GNNs deteriorating rapidly as the number of sensors and actuators grows. key motivation for the use of GNNs in the supervised learning setting is their applicabil- ity large graphs, but this benefit has not yet been realised for locomotion control. We show that poor scaling in GNNs is result of increasingly unstable policy up- dates, caused by overfitting parts of the network during training. To combat this, we introduce SNOWFLAKE, a GNN training method for nigh-dimensional continu- network. SNOWFLAKE
ous control that freezes parameters in selected parts significantly boosts the performance of GNNs for locomotion control on large agents, now matching the performance of MLPs while offering superior transfer properties.
1 Introduction
Whereas many traditional machine learning models operate on sequential or Euclidean (grid-like) data representations, GNNs allow for gtrpt-strectured inputs. GNNs have yielded breakthroughs
in variety of complex domains, including drug discovery B3 50
fraud detection [56), computer
vision
49 4
and particle physics (18)
GNNs have also been successfully applied reinforcement learning (RI.), with promising results on locomotion control tasks with small state and action spaces. Not only are GNN policies 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, far GNNs in RL have only shown competitive performance with MLPs on dower-dimensional locomotion control tasks. For higher-dimension 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
Conresponding author. Now Graphcore, Bristol Now at Waymo, Oxford
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

policy have tendency to cause excessive changes in policy space, leading to performance degrading. To combat this, current state-of-the art algorithms [46 48 II employ trust region-like constraints, inspired by natural gradients [2] 23 that limit the change in policy for each update. We outline how 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.
this policy instability can be framed as form of overfitting
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 attained when training with 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 their initialised values, keeping them fixed throughout training while updating the non-frozen parameters as before. This simple technique enables GNN policies be trained much more effectively in nigh-dimensional environments.
Experimentally, we show that applying SNOWFLKKE to NERVENET dramatically improves asymp- totic performance and sample complexity on such tasks. We also demonstrate that 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 (S, A. R, T, Pu)- The first two elements define the state space S and the action space A. At every time step I, the agent employs a policy ({)]))) to output a distribution over actions, selects action d₂ (-1ss), and transitions from state St S to 8++1 S, as specified by the transition function (8 18,, 4x) which defines probability distribution over states. For the transition, the agent gets reward r2 Rsx, (t, 8t+1). The last element of an MDP specifies initial distribution over states, i.c., states an agent can be in at time step zero.
Solving an MDP means finding a policy x² that maximises an objective, in case the expected discounted sum of rewards J = Ex IEt= 1===0"7r r're]. where 7 E [o, 1) is a discount factor. Policy Gradients (PG) [52 find optimal policy x* by doing gradient ascent the objective: 6s+1 8, + aV 6==0, with 9 parameterising the policy.
Often, to reduce the variance of the gradient estimate, one learns value function V(s) E,, 1Σt= 7'r, 8p = s). and uses it as a critic of the policy. In the resulting actor-critic method, the
policy gradient takes the form: [oJ( = Erg 1E: A7 Vo log (ar(88))) where A70 is estimate
of the advantage function AT = E, E r're ae. st] E, (E T'r St 47
2.2 Proximal Policy Optimisation
Proximal policy optimisation (PPO) 47 is an actor-critic method that has proved effective for variety of domains including locomotion control 177 PPO approximates the natural gradient using first order method, which has the effect of keeping policy updates within a "trust region". This done through the introduction of a surrogate objective to be optimised:
Te(als) AF (s, (s,a), a), clip Tor (als)
(a(a)s) To (als))
e, 1 + (T))
E Ery
min
where e is a clipping hyperparameter that effectively limits how much state-action pair can cause the overall policy to change at each update. This objective is computed over number of optimisation epochs, each of which gives an update to the new policy my. If during this process state-action pair with positive advantage T (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 negative advantage if the lower clipping limit is reached.

---

2.3 Graph Neural Networks
GNNs are class of neural architecture designed to operate over graph-structured data. We define graph as tuple G = (V, E) comprising set of nodes V and edges E = {(, v) u, E V). A labelled graph has corresponding feature vectors for each node and edge that form a pair of matrices Lg (V, E), where V = (ve E IP e ]] and E (eu,o E R4 (u, v) E E). For GNNs we often consider directed graphs, where the order of an edge (u, defines u as the sender and as the receiver.
GNN takes labelled graph 9 and outputs a second graph g' with new labels. Most GNN architectures retain the same topology for G as used in G, in which case GNN can be viewed as mapping from input labels Lg output labels Cg.
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 {0, (O.....I 1] in the network, a hidden state h. and message m is computed for every node V in the graph.
An MPNN implementation calculates these through its choice of message functions and update functions, denoted M and U respectively. 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:
++1 Σ M (h h. eu,,), m
h U m,t))
wEN(2)
for all nodes v E V, where N(v) = {u (u, v) E F3 is the neighbourhood of sender nodes connected to receiver v by directed edge. The node input labels v. used as the initial hidden states h. MPNN assumes only node output labels are required, using each final hidden state h as the output label v
2.4 NerveNet
NERVENET 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 MDP where the state can be factored into input labels V, which are fed to the GNN generate output labels: V NERVENET(G, V). These are then used to parameterise normal distribution defining the stochastic policy: <[ass) AV(V', diag(o")). where the standard deviation is separate vector of parameters learned during training. Actions a are vectors, where each element represents the force to be applied at given joint for the subsequent timestep. The policy is trained using PPO, with parameter updates computed via the Adam optimisation algorithm 255
Internally, NERVENET uses an encoder Fin to generate initial hidden states from input labels: h. Fin (v.). This is followed by message function M consisting of single MLP for all layers 7 that takes as input only the state of the sender node: m = EueN(e) MLP (h.). The update function U is single gated recurrent unit (GRU) [ that maintains an internal hidden state: h = GRU [m [ NERVENET propagates through layers of message-passing and node-updating, before applying decoder Fou to turn final hidden states into scalar node output
labels: v/. Four(h (h.)). diagram of the NERVENET architecture can seen in Appendix A.4
Figure
Knee
Root
Body 0.
Body n/2-f
Hip 1
Hip
Knee
Knee
Figure 1: MuloCo rendering of Centipede-22 and its corresponding morphological graph.
3

---

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 D. 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 [ suite, using the MuJoCo 53 physics engine.
The main set of tasks we use to assess scaling the selection of Cent Centipede-ie agents [55 chosen because of their relatively complex structure and ability to be scaled up to high-dimensional input- action spaces.
The morphology of a Centipedes agent consists of a line of n/2 body segments, each with left and right leg attached (see Figure D The graph used as the basis for the GNN corresponds the physical structure of the agent's s body. At each timestep in the environment, the MuloCo engine sends 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 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 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 CentipedeEight from scratch is already very difficult". Huang also limit training their SMP architecture to small agent types.
et al.
20
3.1 Scaling Performance
To demonstrate the poor scaling of NERVENET to larger agents, compare its performance a selection of Centipede- tasks to that of an MLP policy. Figure 2 shows that for the smaller Centipede- agents both policies are similarly effective, but as the size of the agent increases, the performance of NERVENET drops relative to the MLP. A visual inspection of the behaviour of these agents shows that for Centipede-20, NERVENET barely makes forward progress all, whereas the MLP moves effectively.
As in previous literature [e.g.. the specific training task, but rather matching their training task performance so that the additional benefits of GNNs can be realised. In our setting we particularly wish to leverage the strong transfer benefits of GNNs-as demonstrated by Wang et al. (55) -resulting from their capacity to process inputs of arbitrary size and structure.
55
20
we are ultimately not concerned with outperforming MLPs on
NerveNet
MLP
Centipede-6
Centipede-8
Centipede- 12
Centipede-20
6000
Average Reward
6000
6000
3000
4000
4000
4000
2000
2000
1000
0
0
Timesteps (107)
Figure 2: Comparison of the scaling of NERVENET relative to MLP-based policy. Performance is
similar for the smaller agent sizes, but NERVENET scales poorly to the larger agents.

---

Table 1: KL-divergence from the policy before each update to the policy after, calculated over each batch. train on 10? timesteps, recording the table the mean taken over last 10% steps.
Policy KL-divergence
Policy type Centipede-6
Centipede-8 Centipede-12 Centipede-22
MLP
0.021
0.024
0.031
0.044
NERVENET 0.115
0.137
0.118
0.123
In other words, the focus of this paper is on deriving a method that can close the gap in Figure 2 as doing so makes GNNs better choice overall given the trained policy transfers better than the MLP equivalent (see Section for experimental results).
3.2 Unstable policy updates
As outlined in Section 2.2 one of the key challenges for on-policy RL is preventing individual updates from causing excessive changes in policy space (i.e., keeping within the trust region). Table shows the extent to which this problem contributes to NERVENET's poor scaling, calculating the average KL-divergence from the pre-update policy to the post-update policy for both policy types. NERVENET has consistently higher Kl-divergence than the MLP policy, indicating that PPO finds harder to ensure stable policy updates for the GNN.
We emphasise that this discrepancy persists even with carefully-tuni hyperparameter values for limiting policy divergence. Figure shows the performance of NERVENET across range of PPO e-clipping values (see Section 2.2 and in all cases NERVENET is still substantially inferior to an MLP (note that our experiments on NERVENET always use the best value of 0.1 found here). As we demonstrate later (in Figure 8 controlling policy divergence effectively is key component in making GNNs scale, but we see here that PPO alone does not control the divergence sufficiently to achieve this.
3.3 Overfitting in NERVENET
Excessive policy divergence resulting from updates can be understood form of overfitting. Whereas the supervised interpretation of overfitting implies poor generalisation from training to test set, in this case concerned with poor generalisation across state-action distributions induced
Edge Norm
Average Reward
Policy Update Divergence
1600
0.5
Final Average Reward
1400
0.4
1500
1200
-.3
1000
1000
0.2
500
0.1
0.05
0.10
0.15
0.20
2
3
PPO E clipping paramneeer
Timesteps (10')
Figure 3: Final performance of NER- VENET on Centipede-20 after mil- lion timesteps, across range of clip- ping hyperparameter values. As e in- creases (i.e.. clipping is reduced) the KL divergence from the old to new policy (blue) increases. This improves perfor- mance (orange) up to point, after which it begins to deteriorate.
L2 Penalty:
1e-2
1e-4
1e-6
1e-3
1e-5
Figure 4: L2 regularisation for NERVENET'S mes- sage function across range of values for the L2 penalty X, trained on Centipede -eiideeee...... 20. Increasing this penalty reduces the I.2 norm of the weights learned (left). Improved performance for higher values of À (right) indicates the presence of over- fitting for the message function.
5

---

by different iterations of the policy during training. Specifically, each update involves optimisation step aiming to increase the expected reward over batch of trajectories generated using the pre-upal pre-update policy. The challenge for RI. algorithms is that the agent is then evaluated and trained on trajectories generated using the post-update policy, i.e., a different distribution to the one optimised on.
For MPNN architectures like NERVENET, it is a known deficiency that in the supervised setting, mes- sage functions implemented as MLPs are prone to overfitting [T6. p.55]. Here, we demonstrate that they also overfit (using the above interpretation) in our on-policy RL setting. Figure / shows the effect of applying L2 regularisation (a standard approach to reducing overfitting) to the NERVENET architecture. We regularise the parameters O of NERVENET'S mes- sage function MLP Ma. adding X 0| term to our objective function. At the optimal value of À we an improvement in performance (although still sub- stantially inferior to using MLP), indicating that the unregularised message- pasasingg MLPs overfit.
0.0001
2400
Final Avarage Reward
Learning Rate
1e-05 3e-05
2000
1800
1400
0.0
1200
update
messago encoder decoder
GNN Function
We also investigate lowering the learning rate in dif- ferent parts of the GNN, with the aim of identifying where overfitting is localised. If parts of the network are particularly prone damaging overfitting, train- ing them more slowly may reduce their contribution to policy instability across updates. Results for this
Figure 5: Colour-coded final NERVENET performance after 5M training steps on Centipede -0 when changing learning rates for individual GNN components, compared to the base learning rate of 10-4.
experiment can be seen in Figure 5
Not only does lowering the learning rate in parts of the model improve performance, but surprisingly the best performance is obtained when the encoder Fi. message function M and decoder Four each have their learning rate to zero. The encoder and decoder play a similar role to the message function, all of which are implemented as MLPs, whereas the update function U is GRU (we experimented with using an MI.P update function, but found that this significantly reduced performance.).
3.4 Snowflake
Training with learning rate of zero is equivalent to parameter freezing (e.g., Brock al. (7), where parameters are fixed to their initialised values throughout training. NERVENET can learn a policy with some of its functions frozen, learning still takes place in the un-frozen functions. For instance, we consider freezing the encoder, this results in an arbitrary mapping of input features to the initial long key
hidden states. As still train the update function that processes this representation, information from the input features is not lost via the arbitrary encoding, the update function can still learn useful representations. The same logic applies using a frozen decoder or message function.
Based on the effectiveness of parameter freezing within parts the network, we propose a simple technique for improving the training of GNNs via brased optimisation, which name frozen graph structure). SNOWFLAKE assumes a GNN architec-
SNOWFLAKE (a naturally--ocurriii ture made up internally of functions Fg.... ....... . F where 0 denotes the parameters of given function. } of these functions. Their parameters
Prior to training we select a fixed subset 2 C (F)..... are then placed in SNOWELAKE'S s frozen set (= (0 Fe E Z}. During training, SNOWFLAKE excludes parameters in ( from being updated by the optimiser, instead fixing them to whatever values the GNN architecture uses as an initialisation. Gradients still flow through these operations during backpropagation, but their parameters are not updated. In practice, we found optimal performance for Ç {Fm. Four, Mr }, i.e. when freezing the encoder, decoder and message function of the GNN. not stated otherwise, this is the architecture we refer to SNOWFLAKE in subsequent sections. A visual representation of SLOKE applied to the NERVENET model can be seen in Figure Appendix A.4
FT)
For our experiments, we initialise the values in the GNN using the orthogonal initialisation We found this to be slightly more effective for frozen and unfrozen training than uniform and Xavier

---

initialisations [15. For our message function, which has input and output dimensions of the same size, we find that performance with the frozen orthogonal initialisation is similar to that of simply using the identity function instead of an MLP. However, in the general case where the input and output dimensions of functions in the network differ (such as in the encoder and decoder, or in GNN architectures where layers use representations of different dimensionality), this simplification is not possible and freezing is required.
4 Related Work
Structured Locomotion Control Several different graph neural network-like architectures 15 have been proposed to learn policies for locomotion control. Wang et al. [55] introduce NERVENET, which trains GNN based on the agent's morphology, along with a selection of scalable benchmarks. NERVENET achieves multi-task and transfer learning across morphologies, even in the zero-shot setting (i.e., without further training). which standard MLP-based policies fail to achieve. Sanchez- Gonzalez et al. [2] a GNN-based architecture for learning a model of the environment, which is then used for model-predictive control.
Huang et al. 20 propose Shared Modular Policies (SMP), which focuses on multi-task training and shows strong generalisation to out-o--distributtin agent morphologies using a single policy. The architecture of SMP has similarities with GNN, but requires a tree-based description of the agent's s morphology, and replaces size- and permutation-invane
aggregation with a fixed-cardinality MLP. Pathak et al. [3]] propose dynamic graph networks (DGN), where a GNN is used to learn a policy enabling multiple small agents to cooperate by combining their physical structures.
Amorpheus 2 uses an architecture based on transformers [54] to represent locomotion policies. Transformers can be seen as GNNs using attention for edge-to-vertex aggregation and operating on a fully connected graph, meaning computational complexity scales quadratically with the graph size.
For all of these existing approaches GNN-based locomotion control, training is restricted small agents. In the case of NERVENET and DGN, emphasis is placed on the ability to perform zero-shot transfer larger agents, but this still incurs a significant drop performance.
Graph-Bane Based Reinforcement Learning GNNs have recently gained traction in RL due to their support for variable sized inputs and outputs, enabling new RI applications and enhancing the capabilities of agents on existing benchmarks.
Khalil al. 24 apply DQN 38 to combinatorial optimisation problems using Structure2 [0 for function approximation. Lederman et al. 300 use policy gradient methods to learn heuristics of a quantified Boolean formulae solver, while Kurin et al.
28 use DQN
38 with graph networks [5]
to learn the branching heuristic of a Boolean SAT solver. Klissarov and Precup 2 use a GNN to
represent an MDP, which is then used to learn a form of reward shaping. Deac et al. [ similarly use a GNN MPD representation to generalise Value Iteration Nets [5I] to a broader class of MDPs.
Other approaches involve the construction of graphs based on factorisation of the environmental multi-agent RI,, researchers have used a similar approach to model the relationship between agents, well as between environmental objects [60 20 BT] In this setting, increasing the number of agents can result additional challenges, such combinatorial explosion of the action space. Our approach can be potentially useful to the above work, in improving scaling properties across a variety of domains.
state into objects with associated attributes 4 34
Random Embeddings and Parameter Freezing Embeddings represent feature vectors that have been projected into a new, typically sower-dimensional space that is easier for models to process. Bingham and Mannila [ show that multiplying features by randomly generated matrix (c.g.. with entries sampled from a Gaussian distribution) preserves similarity well and empirically attains comparable performance to PCA. Wang et al. [57 use this approach to apply Bayesian Optimisation to high dimensional datasets by randomly projecting them into a smaller subspace. For natural language applications, commonly used pre-trained embeddings (c.g., word2vec [40], Ve [ have been shown to offer only a small benefit over random embeddings on benchmark datasets p 22 and may offer no benefit industry-scale data 3
More generally, random embeddings can be induced by freezing typically-learned parameters within model to fixed values throughout training. This approach has been explored for transformer
7

---

architectures, where fixed attention weights (either show no significant drop performance, and even freezing intermediate feedforward layers still enables surprisingly effective learning 3 A similar technique can also be found in common fine-tuning methods, where parameters are pre-trained on another, possibly unsupervised objective, but frozen during training except for the final layer [e.g., 58, 19
[59] or hand-crafted (1
Experiments
We present experiments evaluating the performance of SNOWFLAKE when applied to NERVENET, and compare against regular NERVENET and MLP policies. We evaluate each model on selection of MuloCo tasks, including three standard tasks from the Gym suite []] and the Centipede agents from Wang et al. (55). Note that we do not train on even larger Centipede-nn agents due to wall-clock simulation time becoming prohibitively large.
All training statisties are calculated as the mean across six independent runs (unless specified otherwise), with the standard error across runs indicated by the shaded areas on each graph. The average reward typically has high variance, so smooth our results plot the mean taken sliding window of 30 data points. Further experimental details are outlined in Appendix
Scaling to High-Dimensional Tasks Figure 6 compares the scaling properties of the regular NERVENET model with SNOWFLAKE. As the size of the agent increases, SNOWFLAKE significantly outperforms NERVENET with comparable asymptotic performance the MLP. This indicates that SNOWFLAKE is successful in addressing the deficiencies of regular NERVENET training, and that freezing overfitting parameters is effective training strategy this setting. This holds true across locomotive agents with substantially different morphologies.
Zero-shot transfer An important motivation for improving GNN scaling is to harness their transfer capabilities on large tasks. Regular NERVENET is limited by the fact that can only effectively train on and transfer between small agent sizes M We show in Figure that SNOWFLAKE attains exceptional zero-shot transfer performance across centipede sizes, surpassing alternative methods. SNOWFLAKE is the only method that train single policy that is effective Centipede-22 through to 12.
SNOWFLAKE therefore achieves initial objective: combining the strong training task perfor- mance of an MLP and the strong transfer performance of regular NERVENET. As a consequence, SNOWFLAKE-trainet GNNs offer the most promising policy representation for locomotion control tasks where transfer is a desirable property.
A We found that regular NERVENET policy trained to achieve high training task performance on smaller agent (c.g., Centipede-6) does not transfer effectively to larger ones, as reflected in 55 Figure 4].
Snowflake
NerveNet
MLP
Centipede-6
Centipede-8
Centipede-12
Centipede-20
5000
5000
5000
2000
2500
2500
Average Reward
0.0
0.5
1.0
0.0
0.5
1.0
Hopper
Walker
Humanoid
2000
4000
4000
1000
2000
2000
0
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
Timesteps (10?)
Figure 6: Comparison of the performance of SNOWFLAKE training, regular NERVENET and the MLP- based policy. SNOWFLAKE enables effective scaling to the larger agents, significantly outperforming regular NERVENET and comparable to using an MLPP- policy.

---

Palicy KL Divergence
% of Batch Clipped
18 20 (train)
100%
Test Centipede-n Size
3444
2108
3594
3000
Awerage Reward
0.06
75%
3369
1926
609
0.04
50%
-2000
0.02
12
2609
1420
0.00
0%
1000
Timesteps (10%)
140
137
44
Snowflake
NerveNet
Snowflake NerveNet
MLPAA
Figure 8: The effect of SNOWFLKKE policy divergence and PPO clipping on Centipede- 20. By freezing parts of the network that overfit, SNOWFLAKF. reduces the policy KI. divergence lead- ing to less clipping during training.
Figure 7: Zero-shot transfer performance for SNOWFLAKE, NERVENET, and MLP models trained on Centipede-2 20, evalu- ated across a range of sizes.
Policy Stability and Sample Efficiency By reducing overfitting parts of the GNN, SNOWFLAKE mitigates the effect of harmful policy updates seen with regular NERVENET. As a consequence, the policy can train effectively on smaller batch sizes. This is demonstrated in Figure y which shows the performance of NERVENET trained regularly versus using SNOWFLAKE the batch size decreases.
potential benefit of training with smaller batch sizes is improved sample efficiency, as fewer timesteps are taken in the environment update. However, smaller batch sizes also lead to increased policy divergence due increased noise in the gradient estimate. When the policy divergence is too great, performance begins to decrease, limiting how small the batch can be. However due to a reduction policy divergence as a result of SNOWFLAKE, we can afford to use smaller batch sizes while still keeping the policy under control. This provides a wider motivation for the use of SNOWFLLK than just scaling to larger agents: also improves sample efficiency across agents regardless of size.
The success SNOWFAKKE in scaling to larger agents can also be understood in this context. Without FLAKE, for NERVENET to attain strong performance on large agents an infeasibly large batch size would be required, leading to poor sample efficiency. The more stable policy updates enabled by SNOWFLKK make solving these large tasks tractable.
PPO Clipping SNOWFLAKE'S improved policy stability also reduces the amount of clipping performed by PPO across each training batch. Figure X shows the percentage of state-action pairs that are clipped for regular NERVENET versus SNOWFLAKE on the Centipede- 20 agent, as a result of reduced KL divergenec
When NERVENET is trained without using SNOWFLAKE larger percentage of state-action pairs are clipped during PPO updates- a consequence of the greater policy divergence caused by overfitting, For PPO if too many data points reach the clipping limit during optimisation, the algorithm is only able to learn on small fraction of the experience collected, reducing the effectiveness of training. strengths that because reduces policy divergence requires less severe restrictions to keep the policy within the trust region. The combination of this effect and the ability to train well on smaller batch sizes enables SNOWFLAKE'S strong performance on the largest agents.
One of SNOWFLLKE
6 Conclusion
We proposed SNOWFLAKE, a method that enables GNN-based policies to be trained effectively on much larger locomotive agents than previously possible. We no longer observe substantial difference in performance between using GNNs to represent a locomotion policy and the standard approach of using MLPs, even on the most challenging morphologies. As consequence, GNNs
'It may seem counter-intuitive that the KL divergence increases over time. This is due to the standard deviation of the policy (a learned parameter) reducing as training progresses, with the agent trading exploration for exploitation.

---

Batch Size:
2048
1024
512
256
Centipede-6
Hopper
Walker
Centipede-6
Hopper
Walker
Average Reward
4K
4K
2K
2K
2K
1K
2K
2
3
Timesteps (10%)
Timesteps (10%)
(a) NERVENET
(b) SNOWELAKE
Figure 9: Effectiveness of SNOWFLAKE across smaller batch sizes relative to standard NERVENET training. SNOWFLAKE is able to use smaller batch sizes, leading to improved sample efficiency. This due to SNOWFLAKE reducing policy divergence across updates. Corresponding policy divergence plots can be found in the appendix.
now offer alternative MLPs for more than just simple tasks, and if the additional features of GNNs such as strong transfer are requirement, then they are likely to be more effective choice. We have also provided insight into why poor scaling oceurs for certain GNN architectures, and why parameter freezing is effective in addressing the overfitting problem we identify.
Limitations of our work include the upper-limit the size of agents were able simulate, the use of single algorithm and architecture, and focus only on locomotion control tasks. Future work may include applying alternative RL algorithms and GNN architectures, schemes for automating the selection of frozen parts of the network, and applying SNOWFLAKE- methods to wider range of learning problems.
Acknowledgments and Disclosure of Funding
VK is doctoral student at the University of Oxford funded by Samsung R&D Institute UK through the AIMS program. SW has received funding from the European Research Council under European Union's Horizon 2020 research and innovation programme (grant agreement number 637713). The experiments were made possible by a generous equipment grant from NVIDIA.
10


Reviews:
[{"review_id": "qkosu-Iks9c", "reviewer": "Reviewer_PQSJ", "summary": "The paper considers training RL agents by using graph neural networks (GNNs) as the function approximator, therefore inducing a useful inductive bias. Training GNNs is not without its hurdles and thus far previous work has focused on relatively small agents, putting into question scalability. This work attempts to address this issue by propposing parameter freezing in parts of the GNN. The hypothesis for this is that policy updates become unstable with GNNs, to which authors provide empirical evidence. The experiments also show better performance in more complex agents and zero-shot transfer than the NerveNet baseline.", "questions": "", "limitations": "", "rating": 7, "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": "The intersection of GNNs and RL is an interesting and promising field of study with many applications where a graph structure can be leveraged. One such way is to consider the graph formed by the agent's anatomy, which can then be used for better transfer as shown by the authors. As GNNs become hard to train with scaling agents, the paper proposes a perhaps surprising way to fix this issue: to freeze some parts of the GNN. \n\nThe paper reads nicely and the motivation for freezing the weights is well defined, at least for the case when PPO is used. It is however not clear if this method would work or is relevant for the case then another algorithm is used, such as A2C or DDPG. What should we expect for such cases? \n\n\nAnother potential downside is that the transfer experiments focused exclusively on the centipede-n domain. Although it is a challenging one, it is also a limited setting. However, throughout the paper the authors refer to the \"transfer properties\" of GNNs, which is quite a general claim (for example in L147, where there is also no citation backing up the claim) . Such a claim would need to be backed up by experimenting on different domains, or by being more specific as to the scope of the claim.\n\nAlso, do the authors have any hypothesis as to why the KL divergence (L156) is higher when using GNNs? The idea of overfitting is an interesting one, but then again why would GNNs overfit? Perhaps L183 should be expanded a bit to give this important insight. \n\nFigure 5 gives us a plot of the performance with respect to changes in individual parts of the GNNs (i.e. freeze one part at a time). However Snowflake freezes multiple parts of the GNNs at once. Is there any empirical evidence as to why this is done?\n\nWhen freezing the message function, the authors acknowledge that this is very similar (at least in terms of performance) to using an identity function. The structure of the GNN then becomes more close to Graph Convolutional Networks [1]. Have the authors tried using GCN as a baseline?  In the same direction, it would be very informative to have an ablation or a baseline where instead of freezing the encoder and/or decoder, an identity function is used instead of a random network.\n\nConcerning related work, the graph-based RL section could benefit from two relevant works [2,3] where GNNs are used in two quite different manners. In [2] the authors use GNNs for planning by doing more of a bread-first search approximation to the graph, where as in [3] the authors use GNNs for reward shaping by doing more of a depth-first search approximation to the graph. Adding these two references would highlight the flexibility with which GNNs can be combined with RL.\n\nMore of a minor comment, but in the background section of RL, I believe some standard references could be added as well as improving the notation (for example usually the expectation is usually explicit with respect to the policy $\\pi$, as well as the value function and advantage function).\n\nFinally, the number of runs is 6 form what I gathered in the appendix. I believe that for MuJoCo this could be easily pushed to 10 in order to provide more credibility, which is important in the current state of DRL.\n\n[1] Kipf and Welling, Semi-Supervised Classification with Graph Convolutional Networks\n[2] Deac et al., XLVIN: eXecuted Latent Value Iteration Nets\n[3] Klissarov et al., Reward Propagation using Graph Convolutional Networks\n\n\n\n\n\n=================Post-rebuttal===================\n\nI appreciate the authors rebuttal. I think this paper is a solid contribution, and the authors rebuttal has added to this feeling. For this reason I will keep my score to accept.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The societal impact is discussed in the appendix, but in this context is not particularly significant.\n\nFor the limitations, the authors mention that in section 5 it is discussed how their approach is hard to scale beyond centipede-20. Is it possible to provide the wall clock performance across different configurations and compare it to an MLP agent?", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "eKhcN6Ogijs", "reviewer": "Reviewer_soZG", "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.\nThe 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.", "questions": "", "limitations": "", "rating": 7, "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": "Nits:\n* the checklist should start on a separate page\n\nOverall, I thought the paper outlined their problem description, hypothesis and supporting experimental evidence pretty well. The paper provided data for their initial analysis of why GNN policies have trouble scaling to larger locomotion tasks, which is improvement based on the ICML reviews, and then show that their method improves the performance of NerveNet on various Mujoco tasks.\n\n**Originality:** The paper proposes a novel method to address an important challenge in applying GNN based architectures to RL settings. The method can be generally applied to all GNN based architectures that are used in RL settings and potentially in other supervised learning settings.\n\n**Quality:** The paper outlines their hypothesis and the underlying well and provides supporting evidence for their analysis of the problem, such as KL policy divergence and experiments on common regularization techniques. The experiments and results when using the method generally the method generally support the authors' claims. To further improve the paper, it would be good to clarity on a couple of items:\n* How does the Snowflake choose which parts of the GNN to freeze and which parts to train? Does the method enable different learning rates across the network for the unfrozen layers (I think the authors claim it does)? If so, how are those chosen? Does applying Snowflake basically increase the space of hyperparameters in the algorithm or is there a way to systemically determine the additional choices?\n* Do you expect the method to perform similarly on off-policy RL methods, such as TD3 and SAC? Why or why not?\n* A discussion on the limitations of the method\n\n**Clarity:** The paper was generally well written and provides helpful presentation of the data. Further improvements in clarity can be made by having a diagram (or algorithm block) of how the Snowflake method chooses what to freeze as also mentioned above.\n\n**Significance:** The results in the paper indicate that Snowflake can help the adoption of GNN policies to RL problems, which would enable solving more and more complex RL challenges with greater generalization properties. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The discussion on potential societal impacts, data generation and computation infrastructure in the appendix is pretty thorough. More details on limitations and potential future work of the method would make the paper stronger. ", "review_text": "", "extra_scores": {"ethical_concerns": "N/A", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}, {"review_id": "QViUEI5Va6o", "reviewer": "Reviewer_X8TF", "summary": "* The authors study a previously published method for learning locomotion in 3D agents, called NerveNet, which models the agent as a graph of joints linked by limbs and uses a Graph Neural Network (GNN) mechanism to compute joint forces. Previous work has shown that this method transfers better than MLPs across morphologies, but also scales poorly to larger agents.  \n\n* The authors identify a problem that they suggest causes this poor scaling - namely, instability of policy updates, even at performance-optimal choices of hyperparameters. They choose to call this \"overfitting\" - a contentious choice no doubt, but why not. They show that a L2 regularization penalty on weights alleviates the problem a bit, but not enough. \n\n* Crucially, they also discover that reducing learning rates to zero for the encoding and decoding components of the model (i.e. using random frozen embeddings) considerably improves performance. \n\n* Somehow they decide to give a name to the fact of using frozen random embeddings, calling it \"Snowflake\". \n\n* They compare this method to standard NerveNet and MLPs on various locomotion tasks. They show that \"Snowflake\", unlike vanilla NerveNet, performs about as well as MLPs even on large morphologies (after enough training). Importantly, they also show that \"Snowflake\" trained on a given morphology transfers even better to a smaller morphology than NerveNet (and that MLPs simply can't do that at all).\n\n* In short, the contribution of this paper (as I understand it) is an identification of a source for poor NerveNet performance, a solution to this problem (just use random frozen embeddings), and a measurement of the improved performance (matches MLPs in performance, considerably beats them in transfer).", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "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": "The paper is interesting and convincing. My main question is about significance: does the contribution (improving a previously known technique) warrant a NeurIPS presentation?\n\nI suppose it could, given the strength of the performance improvement in comparison to MLP (largely closing the gaps), the remarkable transfer results, and the possibility that this finding might apply to other GNN methods (which are evidently a hot topic right now). \n\nConversely, I'm a bit puzzled by the decision to coin a new name for simply using frozen random embeddings. Is it really warranted? \n\nIn the end, although I like the paper, I could see the decision go either way.\n\nSeparate from this overall uncertainty about significance, I have various clarity concerns that should be easily fixable.\n\n1- In the (nice and useful) description of the GNN method, the notation makes it look like the inputs and outputs v_v and v'_v have the same dimensionality, and that the v' replace the v for the next step (since that's what happens between theta' and theta in the description of PPO that precedes it) ?? Surely that is not the case - it seems to me that the \"input labels\" v_v are fixed, unchanging descriptors of the vertices, while the \"output labels\" are basically just arbitrary outputs with little connection to the labels? If so this should be clarified (and perhaps the notation could be tweaked a bit).\n\n2- Importantly, unless I missed it, the exact inputs are not clearly described.  In p.4, l. 129, we are told that they consist of \"positional information\" about joints and body parts. What does this information look like? The precise coding of inputs might have a great impact on performance and especially on generalization across body sizes. For example, is the \"position\" along the body axis fed as a simple natural integer from 0 to n/2, or is it a fractional number between 0 and 1, normalized to the total length n/2 of the body? Clearly the latter might help any architecture generalize to various lengths. Therefore, I believe a clear description of input format and information should be in the main text. \n\n3- Relatedly, the same line mentions body parts (the edges of the network), but in section 2.4 it seems that only joint (vertices) descriptors are used? Could this be clarified?\n\n4- Supposedly, in the experimental comparisons, MLPs have a roughly similar number of trainable parameters to NervNet/Snowflake? Can this be confirmed? In addition, it would be nice to mention how much computational cost the method incurs in comparison to MLPs (see below).\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Unless I missed it (very possible), it would be useful to have a brief mention of computational costs. Is NerveNet/Snowflake slower or faster than MLPs (with similar number of trainable parameters) *per timestep*, for the particular settings used in Section 5? Even a single sentence or two would be enough.\n\nSocietal impacts do not seem applicable here.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "0yetMovishw", "reviewer": "Reviewer_Qzdd", "summary": "This paper studies some of the optimization challenges with training large graph neural networks for reinforcement learning tasks, in particular locomotion control tasks. The paper proposes that one of the larger challenges in optimizing graph neural networks is the significant variance between updates across the network graph. The method proposes to randomly freeze nodes and components of nodes in the graph network. The freezing of these different blocks inside of the graph Network helps reduce the effects of large changes in the output of the network making the overall optimization less noisy and more stable resulting in better performance.\n\nPros\n- The analysis in the paper does appear to be rather thorough. Especially going into the specifics about the difficulties in training neural networks for this particular type of problem. Also at the same time showing some of the limitations of training graph neural networks for example nerve net and its difficulty to perform well on larger agent morphologies.\n- The paper then proposes a random perimeter freezing algorithm that increases the overall learning performance and final policy quality of the nerve net like design.\n- Along with the performance gain the method is still able to show significant transfer capabilities across some different morphologies.\n\nCons\n- One of the more subtle challenges in the contribution of this paper is that the method overall is somewhat simplistic. It could be thought of as performing a type of block dropout while training the network in order to reduce the amount of parameter change when making policy updates.\n- One of the other more nuanced difficulties in the contribution of the paper is that it focuses on training graph neural networks for locomotion tasks on reinforcement learning. This could be seen as a somewhat specific area of research that may not have a more broad impact on the community.\n- Related to the last point it would be helpful to compare this method to the amorphous work. That work calls out the challenges with training and the optimization of graph neural networks that may limit their overall application. In this new work proposed by the authors shows they've now been able to overcome some of the difficulties in training such that it may perform better than the amorphous algorithm as well.\n", "questions": "", "limitations": "", "rating": 5, "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": "Additional comments on the work:\n- In figure 5 it shows that setting the learning rate to practically zero for the message part of the CNN leads to the highest final average award. This also seems to be corroborated in the recent paper Amorpeuos [1], which uses Transformers to train morphological type models. Does this indicate that the training of GNNs especially the message passing system makes the overall optimization complex and could indicate that GNNs are also generally limited by the difficulty of their optimization?\n- [1] This reference should be fixed in the paper to reference the proper version of the paper that has been published at ICLR 2021.\n- The paper mentions the Amorpeuos prior work that uses a transformer to represent locomotion policies yet it does not compared to in the paper. In order to get a broader picture of how this method compares to the most recent state-of-the-art, it would be important to compare to this structure. Especially since one of the main points in the Amorpeuos paper was that training GNNs, especially the message passing system, is too complex in that the benefits of GNNs will be hindered by the difficulties of this optimization. However, the method proposed in this paper should be addressing these optimization limitations and therefore should perform better than the Amorpheus work.\n- The analysis in figure eight with respect to how the policy KL Divergence increases over training is somewhat counterintuitive. PPO on its own does learn to adjust the policy distribution over time and as the policy should converge to at least a local optimum one would think the KL Divergence would also shrink yet in these graphs we don't see a downward trend later in learning. Does this imply that the policy is not finished learning and there is still room to learn a more optimal policy?\n\n\n\nRefernces:\n[1] Kurin, V., Igl, M., Rocktaschel, T., Boehmer, W., & Whiteson, S. (2021, April). My body is a cage: the role of morphology in graph− based incompatible control. In Proceedings of the International Conference on Learning Representations. OpenReview.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "To increase the contrubution of the paper it will be important to include comparisons to other prior works, like Amorpeuos. ", "review_text": "", "extra_scores": {"ethical_concerns": "N/A", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2.5"}, "extra_text": {}}]

Author rebuttal:

Meta-review: There was quite an extensive discussion between the reviewers on this paper after the author response. There is no doubt that the the paper proposes a simple (in a good way) fix to train GNNs with PPO. So it makes a valuable contribution to the community. There was also agreement that the paper could potentially be made a lot stronger with additional work, e.g., the paper is mainly based on empirical evidence for the proposed method so evaluating freezing part by part (reviewer PQSJ) would be very interesting, more comparisons to other work, etc. The main difference was in the weighting of these two aspects.
Decision: Accept (Poster)