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 focuses on neural architecture search for constructing compact RL policies. It combines ideas from the popular ENAS and ES methods for optimisation. Recent work defined a family of compact policies by imposing a fixed Toeplitz structure. This paper introduces the so-called “chromatic network” architecture, which partitions weights of the RL network into tied sub-groups. The partitioning is searched with ENAS and shared weights are updated via ES. Experiments on continuous control benchmarks show that good performance can be obtained using a very small number of parameters. Favourable reward-compression outcomes can be achieved compared to some baseline alternatives.