Reinforcement Learning with Chromatic Networks

1 INTRODUCTION . Consider a fixed Markov Decision Process ( MDP ) M and an agent aiming to maximize its total expected/discounted reward obtained in the environment E governed byM . An agent is looking for a sequence of actions a0 , ... , aT−1 leading to a series of steps maximizing this reward . One of the approaches is to construct a policy πθ : S → A , parameterized by vector θ , which is a mapping from states to actions . Policy πθ determines actions chosen in states visited by an agent . Such a reinforcement learning ( RL ) policy is usually encoded as a neural network , in which scenario parameters θ correspond to weights and biases of a neural network . Reinforcement learning policies πθ often consist of thousands or millions of parameters ( e.g . when they involve vision as part of the state vector ) and therefore training them becomes a challenging high-dimensional optimization problem . Deploying such high-dimensional policies on hardware raises additional concerns in resource constrained settings ( e.g . limited storage ) , emerging in particular in mobile robotics ( Gage , 2002 ) . The main question we tackle in this paper is the following : Are high dimensional architectures necessary for encoding efficient policies and if not , how compact can they be in in practice ? We show that finding such compact representations is a nontrivial optimization problem despite recent observations that some hardcoded structured families ( Choromanski et al. , 2018 ) provide certain levels of compactification and good accuracy at the same time . We model the problem of finding compact presentations by using a joint objective between the combinatorial nature of the network ’ s parameter sharing profile and the reward maximization of RL optimization . We leverage recent advances in the ENAS ( Efficient Neural Architecture Search ) literature and theory of pointer networks ( Vinyals et al. , 2015 ; Pham et al. , 2018 ; Zoph & Le , 2017 ) to optimize over the combinatorial component of this objective and state of the art evolution strategies ( ES ) methods ( Choromanski et al. , 2018 ; Salimans et al. , 2017 ; Mania et al. , 2018a ) to optimize over the RL objective . We propose to define the combinatorial search space to be the the set of different edge-partitioning ( colorings ) into same-weight classes and construct policies with learned weight-sharing mechanisms . We call networks encoding our policies : chromatic networks . We are inspired by two recent papers : ( Choromanski et al. , 2018 ) and ( Gaier & Ha , 2019 ) . In the former one , policies based on Toeplitz matrices were shown to match their unstructured counterparts accuracy-wise , while leading to the substantial reduction of the number of parameters from thousands ( Salimans et al. , 2017 ) to hundreds ( Choromanski et al. , 2018 ) . Instead of quadratic ( in sizes of hidden layers ) , those policies use only linear number of parameters . The Toeplitz structure can be thought of as a parameter sharing mechanism , where edge weights along each diagonal of the matrix are the same ( see : Fig . 1 ) . However , this is a rigid pattern that is not learned . We show in this paper that weight sharing patterns can be effectively learned , which further reduces the number of distinct parameters . For instance , using architectures of the same sizes as those in ( Choromanski et al. , 2018 ) , we can train effective policies for OpenAI Gym tasks with as few as 17 distinct weights . The latter paper proposes an extremal approach , where weights are chosen randomly instead of being learned , but the topologies of connections are trained and thus are ultimately strongly biased towards RL tasks under consideration . It was shown in ( Gaier & Ha , 2019 ) that such weight agnostic neural networks ( WANNs ) can encode effective policies for several nontrivial RL problems . WANNs replace conceptually simple feedforward networks with general graph topologies using NEAT algorithm ( Stanley & Miikkulainen , 2002 ) providing topological operators to build the network . Our approach is a middle ground , where the topology is still a feedforward neural network , but the weights are partitioned into groups that are being learned in a combinatorial fashion using reinforcement learning . While ( Chen et al. , 2015 ) shares weights randomly via hashing , we learn a good partitioning mechanisms for weight sharing . Our key observation is that ENAS and ES can naturally be combined in a highly scalable but conceptually simple way . To give context , vanilla NAS ( Zoph & Le , 2017 ) for classical supervised learning setting ( SL ) requires a large population of 450 GPU-workers ( child models ) all training one-by-one , which results in many GPU-hours of training . ENAS ( Pham et al. , 2018 ) uses weight sharing across multiple workers to reduce the time , although it can reduce computational resources at the cost of the variance of the controller ’ s gradient . Our method solves both issues ( fast training time and low controller gradient variance ) by leveraging a large population of much-cheaper CPU workers ( 300 ) increasing the effective batch-size of the controller , while also training the workers simultaneously via ES . This setup is not possible in SL , as single CPUs can not train large image-based classifiers in practice . Furthermore , this magnitude of scaling by numerous workers can be difficult with policy gradient or Q-learning methods as they can be limited by GPU overhead due to exact-gradient computation . We believe that our work is one of the first attempts to propose a flexible , rigorous approach to training compact neural network architectures for RL problems . Those may be of particular importance in mobile robotics ( Gage , 2002 ) where computational and storage resources are very limited . We also believe that this paper opens several new research directions regarding structured policies for robotics . 1.1 BACKGROUND AND RELATED WORK . Network Architecture Search : The subject of this paper can be put in the larger context of Neural Architecture Search ( NAS ) algorithms which recently became a prolific area of research with already voluminous literature ( see : ( Elsken et al. , 2019 ) for an excellent survey ) . Interest in NAS algorithms started to grow rapidly when it was shown that they can design state-of-the-art architectures for image recognition and language modeling ( Zoph & Le , 2017 ) . More recently it was shown that NAS network generators can be improved to sample more complicated connectivity patterns based on random graph theory models such as Erdos-Renyi , Barabasi-Albert or Watts-Strogatz to outperform human-designed networks , e.g . ResNet and ShuffleNet on image recognition tasks ( Xie et al. , 2019 ) . However to the best of our knowledge , applying NAS to construct compact RL policy architectures has not been explored before . On a broader scope , there has been little work in applying NAS to general RL policies , partially because unlike SL , RL policies are quite small and thus do not require a search space involving hyperparameter primitives such as number of hidden layers , convolution sizes , etc . However , we show that the search space for RL policies can still be quite combinatorial , by examining partitionings of weights in our work . Parameterizing Compact Architectures : Before NAS can be applied , a particular parameterization of a compact architecture defining combinatorial search space needs to be chosen . That leads to the vast literature on compact encodings of NN architectures . Some of the most popular and efficient techniques regard network sparsification . The sparsity can be often achieved by pruning already trained networks . These methods have been around since the 1980s , dating back to Rumelhart ( Rumelhart , 1987 ; Chauvin , 1989 ; Mozer & Smolensky , 1989 ) , followed shortly by Optimal Brain Damage ( Cun et al. , 1990 ) , which used second order gradient information to remove connections . Since then , a variety of schemes have been proposed , with regularization ( Louizos et al. , 2018 ) and magnitude-based weight pruning methods ( Han et al. , 2015 ; See et al. , 2016 ; Narang et al. , 2017 ) increasingly popular . The impact of dropout ( Srivastava et al. , 2014 ) has added an additional perspective , with new works focusing on attempts to learn sparse networks ( Gomez et al. , 2019 ) . Another recent work introduced the Lottery Ticket Hypothesis ( Frankle & Carbin , 2019 ) , which captivated the community by showing that there exist equivalently sparse networks which can be trained from scratch to achieve competitive results . Interestingly , these works consistently report similar levels of compression , often managing to match the performance of original networks with up to 90 % fewer parameters . Those methods are however designed for constructing classification networks rather than those encoding RL policies . Quantization : Weight sharing mechanism can be viewed from the quantization point of view , where pre-trained weights are quantized , thus effectively partitioned . Examples include ( Han et al. , 2016 ) , who achieve 49x compression for networks applied for vision using both pruning and weight sharing ( by quantization ) followed by Huffman coding . However such partitions are not learned which is a main topic of this paper . Even more importantly , in RL applications , where policies are often very sensitive to parameters ’ weights ( Iscen et al. , 2018 ) , centroid-based quantization is too crude to preserve accuracy . Compact RL Policies In recent times there has been increased interest in simplifying RL policies . In particular , ( Mania et al. , 2018b ) demonstrated that linear policies are often sufficient for the benchmark MuJoCo locomotion tasks , while ( Cuccu et al. , 2019 ) found smaller policies could work for vision-based tasks by separating feature extraction and control . Finally , recent work found that small , sparse sub-networks can perform better than larger over-parameterized ones ( Frankle & Carbin , 2019 ) , inspiring applications in RL ( Yu et al. , 2019 ) . Our main contributions are : 1 . We propose a highly scalable algorithm by combining ENAS and ES for learning compact representations that learns effective policies with over 92 % reduction of the number of neural network parameters ( Section 3 ) . 2 . To demonstrate the impact ( and limitations ) of pruning neural networks for RL , we adapt recent algorithms training both masks defining combinatorial structure as well as weights of a deep neural network concurrently ( see : ( Lenc et al. , 2019 ) ) . Those achieve state-ofthe-art results on various supervised feedforward and recurrent models . We confirm these findings in the RL setting by showing that good rewards can be obtained up to a high level of pruning . However , at the 80-90 % level we see a significant decrease in performance which does not occur for the proposed by us chromatic networks ( Section 4.1 , Section 4.2 ) . 3 . We demonstrate that finding efficient weight-partitioning mechanisms is a challenging problem and NAS helps to construct distributions producing good partitionings for more difficult RL environments ( Section 4.3 ) . 2 THE ARCHITECTURE OF NEURAL ARCHITECTURE SEARCH . The foundation of our algorithm for learning structured compact policies is the class of ENAS methods ( Pham et al. , 2018 ) . To present our algorithm , we thus need to first describe this class . ENAS algorithms are designed to construct neural network architectures thus they aim to solve combinatorial-flavored optimization problems with exponential-size domains . The problems are cast as MDPs , where a controller encoded by the LSTM-based policy πcont ( θ ) , typically parameterized by few hundred hidden units , is trained to propose good-quality architectures , or to be more precise : good-quality distributions D ( θ ) over architectures . In standard applications the score of the particular distribution D ( θ ) is quantified by the average performance obtained by trained models leveraging architectures A ∼ D ( θ ) on the fixed-size validation set . That score determines the reward the controller obtains by proposing D ( θ ) . LSTM-based controller constructs architectures using softmax classifiers via autoregressive strategy , where controller ’ s decision in step t is given to it as an input embedding at time t + 1 . The initial embedding that the controller starts with is an empty one . Weight Sharing Mechanism : ENAS introduces a powerful idea of a weight-sharing mechanism . The core concept is that different architectures can be embedded into combinatorial space , where they correspond to different subgraphs of the given acyclic directed base graph G ( DAG ) . Weights of the edges of G represent the shared-poolWshared from which different architectures will inherit differently by activating weights of the corresponding induced directed subgraph . At first glance such a mechanism might be conceptually problematic since , a weight of the particular edge e belonging to different architectures Ai1 , ... Aik ( see : next paragraph for details regarding weight training ) will be updated based on evaluations of all of them and different Ais can utilize e in different ways . However it turns out that this is desirable in practice . As authors of ( Pham et al. , 2018 ) explain , the approach is motivated by recent work on transfer and multitask learning that provides theoretical grounds for transferring weights across models . Another way to justify the mechanism is to observe that ENAS tries in fact to optimize a distribution over architectures rather than a particular architecture and the corresponding shared-pool of weightsWshared should be thought of as corresponding to that distribution rather than its particular realizations . Therefore the weights of that pool should be updated based on signals from all different realizations . Alternating Optimization : For a fixed parameterization θ defining policy π ( θ ) of the controller ( and thus also proposed distribution over architecturesD ( θ ) ) , the algorithm optimizes the weights of the models usingM architectures : A1 , ... , AM sampled fromD ( θ ) , where the sampling is conducted by the controller ’ s policy πcont ( θ ) . Models corresponding to A1 , ... , AM are called child models . At iteration k of the weight optimization process , a worker assigned to the architectureAi computes the gradient of the loss function LAi , Bk corresponding to the particular batch Bk of the training data with respect to the weights of the inherited edges from the base DAG . Since the loss function LBk onWshared for a fixed distribution D ( θ ) and batch Bk at iteration k is defined as an expected loss EA∼π ( θ ) [ LA , Bk ( WAshared ) ] , where WAshared stands for the projection of the set of shared-pool weights Wshared into edges defined by the induced subgraph determined by A , its gradient with respect toWshared can be estimated via Monte Carlo procedure as : ∇WsharedLBk ( Wshared ) ∼ 1 M M∑ i=1 ∇WAisharedLAi , Bk ( W Ai shared ) ( 1 ) After weight optimization is completed , ENAS updates the parameters θ of the controller responsible for shaping the distribution used to sample architectures A . As mentioned before , this is done using reinforcement learning approach , where parameters θ are optimized to maximize the expected reward EA∼π ( θ ) [ RWAshared ( A ) ] , where this time setWshared is frozen and RWAshared ( A ) is given as the accuracy obtained by the model using architectureA and weights fromWAshared on the validation set . Parameters θ are updated with the use of the REINFORCE algorithm ( Williams , 1992 ) . Controller ’ s Architecture : The controller LSTM-based architecture uses pointer networks ( Vinyals et al. , 2015 ) that utilize LSTM-encoder/decoder approach , where the decoder can look back and forth over input . Thus these models are examples of architectures enriched with attention .

Summary: This paper compresses policy networks using approaches inspired by neural architecture search. The main idea is to have a fixed size weight matrix, but learn how to share weights, so that the resulting network can be compressed by storing only unique weights values. Both the partitioning and weights are trained simultaneously, inspired by ENAS. The partitioning is modeled by an autoregressive RNN and trained via REINFORCE. The weights are modeled by a single set of weights (as opposed to a distribution) which is then updated by using a gradient approximation based on ES. The experiments carried out include comparing to existing works on policy network compression, ablating against random partitioning, as well as a few experiments meant to increase understanding of the learned partitions.