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

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

heading: 1 Introduction
text: 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.

heading: 2 Background

heading: 2.1 Reinforcement Learning
text: 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].

heading: 2.2 Proximal Policy Optimisation
text: 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.

heading: 2.3 Graph Neural Networks
text: 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 .

heading: 2.4 NerveNet
text: 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.

heading: 3 Analysing GNN Scaling Challenges
text: 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.

heading: 3.1 Scaling Performance
text: To demonstrate the poor scaling of NERVENET to larger agents, we compare its performance on a selection of Centipede-n tasks to that of an MLP policy. Figure 2 shows that for the smaller Centipede-n 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 at all, whereas the MLP moves effectively.
As in previous literature [e.g., 55, 20], we are ultimately not concerned with outperforming MLPs on 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.
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 a better choice overall given the trained policy transfers better than the MLP equivalent (see Section 5 for experimental results).

heading: 3.2 Unstable policy updates
text: 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 it within the trust region). Table 1 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 a consistently higher KL-divergence than the MLP policy, indicating that PPO finds it harder to ensure stable policy updates for the GNN.
We emphasise that this discrepancy persists even with carefully-tuned hyperparameter values for limiting policy divergence. Figure 3 shows the performance of NERVENET across a range of PPO ✏-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 a key component in making GNNs scale, but we see here that PPO alone does not control the divergence sufficiently to achieve this.

heading: 3.3 Overfitting in NERVENET
text: Excessive policy divergence resulting from updates can be understood as a form of overfitting. Whereas the supervised interpretation of overfitting implies poor generalisation from training to test set, in this case we are concerned with poor generalisation across state-action distributions induced
by different iterations of the policy during training. Specifically, each update involves an optimisation step aiming to increase the expected reward over a batch of trajectories generated using the pre-update policy. The challenge for RL 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, message functions implemented as MLPs are prone to overfitting [16, p.55]. Here, we demonstrate that they also overfit (using the above interpretation) in our on-policy RL setting. Figure 4 shows the effect of applying L2 regularisation (a standard approach to reducing overfitting) to the NERVENET architecture. We regularise the parameters ✓ of NERVENET’s message function MLP M✓ , adding a ||✓||22 term to our objective function. At the optimal value of we see an improvement in performance (although still substantially inferior to using an MLP), indicating that the unregularised message-passing MLPs overfit.
We also investigate lowering the learning rate in different parts of the GNN, with the aim of identifying where overfitting is localised. If parts of the network are particularly prone to damaging overfitting, training them more slowly may reduce their contribution to policy instability across updates. Results for this 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 Fin, message function M and decoder Fout each have their learning rate set 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 a GRU (we experimented with using an MLP update function, but found that this significantly reduced performance.).

heading: 3.4 Snowflake
text: Training with a learning rate of zero is equivalent to parameter freezing (e.g., Brock et al. [7]), where parameters are fixed to their initialised values throughout training. NERVENET can learn a policy with some of its functions frozen, as learning still takes place in the un-frozen functions. For instance, if we consider freezing the encoder, this results in an arbitrary mapping of input features to the initial hidden states. As we still train the update function that processes this representation, so long as key information from the input features is not lost via the arbitrary encoding, the update function can still learn useful representations. The same logic applies to using a frozen decoder or message function.
Based on the effectiveness of parameter freezing within parts of the network, we propose a simple technique for improving the training of GNNs via gradient-based optimisation, which we name SNOWFLAKE (a naturally-occurring frozen graph structure). SNOWFLAKE assumes a GNN architecture made up internally of functions F 1✓ , . . . , F n ✓ , where ✓ denotes the parameters of a given function. Prior to training we select a fixed subset Z ✓ {F 1✓ , . . . , Fn✓ } of these functions. Their parameters are then placed in SNOWFLAKE’s frozen set ⇣ = {✓ | F✓ 2 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 ⇣ = {Fin, Fout,M⌧}, i.e. when freezing the encoder, decoder and message function of the GNN. If not stated otherwise, this is the architecture we refer to as SNOWFLAKE in subsequent sections. A visual representation of SNOWFLAKE applied to the NERVENET model can be seen in Figure 11, Appendix A.4.
For our experiments, we initialise the values in the GNN using the orthogonal initialisation [44]. 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.

heading: 4 Related Work
text: Structured Locomotion Control Several different graph neural network-like architectures [45, 5] have been proposed to learn policies for locomotion control. Wang et al. [55] introduce NERVENET, which trains a 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. SanchezGonzalez et al. [42] use 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-of-distribution agent morphologies using a single policy. The architecture of SMP has similarities with a GNN, but requires a tree-based description of the agent’s morphology, and replaces size- and permutation-invariant aggregation with a fixed-cardinality MLP. Pathak et al. [39] 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 [29] 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 to GNN-based locomotion control, training is restricted to small agents. In the case of NERVENET and DGN, emphasis is placed on the ability to perform zero-shot transfer to larger agents, but this still incurs a significant drop in performance.
Graph-Based Reinforcement Learning GNNs have recently gained traction in RL due to their support for variable sized inputs and outputs, enabling new RL applications and enhancing the capabilities of agents on existing benchmarks.
Khalil et al. [24] apply DQN [38] to combinatorial optimisation problems using Structure2Vec [10] for function approximation. Lederman et al. [30] 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 [26] use a GNN to represent an MDP, which is then used to learn a form of reward shaping. Deac et al. [11] similarly use a GNN MPD representation to generalise Value Iteration Nets [51] to a broader class of MDPs.
Other approaches involve the construction of graphs based on factorisation of the environmental state into objects with associated attributes [4, 34]. In multi-agent RL, researchers have used a similar approach to model the relationship between agents, as well as between environmental objects [60, 21, 31]. In this setting, increasing the number of agents can result in additional challenges, such as 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.
Random Embeddings and Parameter Freezing Embeddings represent feature vectors that have been projected into a new, typically lower-dimensional space that is easier for models to process. Bingham and Mannila [6] show that multiplying features by a randomly generated matrix (e.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 (e.g., word2vec [40], GloVe [37]) have been shown to offer only a small benefit over random embeddings on benchmark datasets [27, 12] and may offer no benefit on industry-scale data [3].
More generally, random embeddings can be induced by freezing typically-learned parameters within a model to fixed values throughout training. This approach has been explored for transformer
architectures, where fixed attention weights (either Gaussian-distributed [59] or hand-crafted [41]) show no significant drop in performance, and even freezing intermediate feedforward layers still enables surprisingly effective learning [35]. 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].

heading: 5 Experiments
text: 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 a selection of MuJoCo tasks, including three standard tasks from the Gym suite [8] and the Centipede-n agents from Wang et al. [55]. Note that we do not train on even larger Centipede-n agents due to wall-clock simulation time becoming prohibitively large.
All training statistics 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 to smooth our results we plot the mean taken over a sliding window of 30 data points. Further experimental details are outlined in Appendix A.2.
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 to the MLP. This indicates that SNOWFLAKE is successful in addressing the deficiencies of regular NERVENET training, and that freezing overfitting parameters is an effective training strategy in 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 it can only effectively train on and transfer between small agent sizes.3 We show in Figure 7 that SNOWFLAKE attains exceptional zero-shot transfer performance across centipede sizes, surpassing alternative methods. SNOWFLAKE is the only method that can train a single policy that is effective on Centipede-20 through to 12.
SNOWFLAKE therefore achieves our initial objective: combining the strong training task performance of an MLP and the strong transfer performance of regular NERVENET. As a consequence, SNOWFLAKE-trained GNNs offer the most promising policy representation for locomotion control tasks where transfer is a desirable property.
3We found that a regular NERVENET policy trained to achieve high training task performance on a smaller agent (e.g., Centipede-6) does not transfer effectively to larger ones, as reflected in [55, Figure 4].
Policy Stability and Sample Efficiency By reducing overfitting in 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 9, which shows the performance of NERVENET trained regularly versus using SNOWFLAKE as the batch size decreases.
A potential benefit of training with smaller batch sizes is improved sample efficiency, as fewer timesteps are taken in the environment per update. However, smaller batch sizes also lead to increased policy divergence due to 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 in 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 SNOWFLAKE than just scaling to larger agents: it also improves sample efficiency across agents regardless of size.
The success of SNOWFLAKE in scaling to larger agents can also be understood in this context. Without SNOWFLAKE, 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 SNOWFLAKE 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 8 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 divergence4.
When NERVENET is trained without using SNOWFLAKE a 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 a small fraction of the experience collected, reducing the effectiveness of training. One of SNOWFLAKE’s strengths is that because it reduces policy divergence it 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.

heading: 6 Conclusion
text: We proposed SNOWFLAKE, a method that enables GNN-based policies to be trained effectively on much larger locomotive agents than was previously possible. We no longer observe a 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 a consequence, GNNs
4It 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.
now offer an alternative to MLPs for more than just simple tasks, and if the additional features of GNNs such as strong transfer are a requirement, then they are likely to be a more effective choice. We have also provided insight into why poor scaling occurs 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 on the size of agents we were able to simulate, the use of a single algorithm and architecture, and a 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-like methods to a wider range of learning problems.
Acknowledgments and Disclosure of Funding
VK is a 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 the 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.

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

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.  Pros - 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. - 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. - Along with the performance gain the method is still able to show significant transfer capabilities across some different morphologies.  Cons - 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. - 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. - 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. 

* 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.    * 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.   * 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.   * Somehow they decide to give a name to the fact of using frozen random embeddings, calling it "Snowflake".   * 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).  * 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).

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.