Snowflake: Scaling GNNs to High-Dimensional Continuous Control via 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 (Wang et al., 2018; Huang et al., 2020). 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 identify the weakness with a common GNN architecture that causes this poor scaling: overfitting in the MLPs within the network that encode, decode, and propagate messages. To combat this, we introduce Snowflake, a GNN training method for high-dimensional continuous control that freezes parameters in parts of the network that suffer from overfitting. Snowflake significantly boosts the performance of GNNs for locomotion control on large agents, now matching the performance of MLPs, and with superior transfer properties.

arXiv comment20 pages, 14 figures, published at NeurIPS 2021

References in corpus (13)

• Efficient Estimation of Word Representations in Vector Space
• Neural Message Passing for Quantum Chemistry
• How transferable are features in deep neural networks?
• Trust Region Policy Optimization
• Learning Combinatorial Optimization Algorithms over Graphs
• Emergence of Locomotion Behaviours in Rich Environments
• Discriminative Embeddings of Latent Variable Models for Structured Data
• Maximum a Posteriori Policy Optimisation
• Learning to Control Self-Assembling Morphologies: A Study of Generalization via Modularity
• Deep Implicit Coordination Graphs for Multi-agent Reinforcement Learning
• Accelerated Charged Particle Tracking with Graph Neural Networks on FPGAs
• One Policy to Control Them All: Shared Modular Policies for Agent-Agnostic Control
• Can $Q$-Learning with Graph Networks Learn a Generalizable Branching Heuristic for a SAT Solver?