RL-DARTS: Differentiable Architecture Search for Reinforcement Learning

1 INTRODUCTION AND MOTIVATION . Over the last decade , recent advancements in deep reinforcement learning have heavily focused on algorithmic improvements ( Hessel et al. , 2018 ; Wang et al. , 2016 ) , data augmentation ( Raileanu et al. , 2020 ; Kostrikov et al. , 2020 ) , infrastructure upgrades ( Espeholt et al. , 2018 ; Horgan et al. , 2018 ) , and even hyperparameter tuning ( Zhang et al. , 2021 ; Faust et al. , 2019 ) . The automation of such RL methods have given rise to the field of Automated Reinforcement Learning ( AutoRL ) . Surprisingly however , one relatively underdeveloped and unexplored area is automating large-scale policy architecture design . Recently , works in larger-scale RL suggest that designing a policy ’ s architecture can be just as important as the algorithm or quality of data , if not more , for various metrics such as generalization , transferrability , and efficiency . One surprising phenomenon found on Procgen ( Cobbe et al. , 2020 ) , a procedural generation benchmark for RL , was that “ IMPALA-CNN ” , a residual convolutional architecture from ( Espeholt et al. , 2018 ) , could substantially outperform “ NatureCNN ” , the standard 3-layer convolutional architecture used for Atari ( Mnih et al. , 2013 ) , in both generalization and sample complexity under limited and infinite data regimes respectively ( Cobbe et al. , 2019 ) . Furthermore , in robotics subfields such as grasping ( Kalashnikov et al. , 2018 ; Rao et al. , 2020 ) , cameras collect very detailed real-world images ( ex : 472 x 472 , 4x larger than ImageNet ( Russakovsky et al. , 2015 ) ) for observations which require deep image encoder architectures , consisting of more than 15 convolutions , raising concerns on efficiency and speed in policy training and inference . As such RL policy networks gradually become larger and more sophisticated , so does the need for understanding and automating such designs . Many efficient Neural Architecture Search ( NAS ) methods utilize the one-shot setup , where only one set of shared weight parameters is trained through an ensemble of models , otherwise known as a supernet . Differentiable Architecture Search ( Liu et al. , 2019 ) , trains the supernet by inserting architecture selection variables inside the network in a continuous relaxation approach , in order to be differentiable with respect to a loss . These architecture search parameters are then used to select a final discretized perception module to be retrained from scratch in an evaluation phase . The differentiable aspect of DARTS is highly attractive to RL , which possess two complications which SL does not : 1 . RL is an entire system . In most RL algorithms involving auto-differentiation , there are multiple different phases : a collection/exploration phase which collects environment data via forward passes/inference of a policy , in addition to a training phase which updates parameters for potentially multiple network components ( e.g . both the policy and value function in policy gradient methods ) . There can also even be a separate evaluation phase ( disjoint from evaluation in NAS ) different from the collection phase ( e.g . using argmax greedy evaluation instead of the stochastic actor , for both Q-learning and policy gradient methods ) . Such phases are also commonly distributed over different machines , adding an extra layer of complexity via inter-machine communication . 2 . RL training curves possess considerably more variance and noise than SL curves , with a much larger room for error due to sensitivity to hyperparameters ( Andrychowicz et al. , 2020 ) , which originally led to the criterion of reporting mean and standard deviation for 3 seeded runs ( Henderson et al. , 2018 ) . In fact , this noise can be nearly unbounded , as in many cases , an unlucky and catastrophic seeded training run can simply lead to a trivial reward even with an optimal architecture . Such noise may potentially be problematic for fast convergence for blackbox/evolutionary optimization methods . Furthermore , the RL loss contains important information not conveyed by just pure reward ( e.g . value function error and entropy penalty , as well as potentially exploration/novelty bonuses ) , whereas information may be lost when simply summarizing all progress to a single scalar point of data . DARTS solves ( 1 ) as it integrates naturally within the RL system with minimal modifications to the code and pipeline , by simply replacing the image encoder in standard RL algorithms . For ( 2 ) , architecture search parameters in the DARTS supernet are joint trained alongside weight parameters , and can thus also use information from multiple auxiliary losses , which can be much richer signals than potentially sparse rewards . We represent this approach in Figure 1 . However , since there has been no previous proof of DARTS ’ s feasibility in standard RL , one may hypothesize a few subtle concerns . While DARTS in SL is simpler and can focus more on generalization , a core major challenge in RLDARTS is the feasability of training the supernet itself due to a non-stationary data distribution . This is especially because using a differentiable loss ( ex : Bellman error ) as the search objective is a double-edged sword in RL , as loss is defined with respect to the replay buffer , whose data is collected by the policy itself via inference on the architecture . In SL , the data distribution is independent of the architecture , and the loss ( e.g . cross-entropy or negative log-likelihood ) is strongly correlated or even equivalent to the objective ( e.g . accuracy or density estimation ( Parmar et al. , 2018 ) ) . However in RL , a bad architecture could lead to a poor initial distribution of data , and thus a local optimum where the loss is zero while the policy is still poor . This is inherently due to the fact that the data in the replay buffer may not sufficiently cover the state and action space of the entire environment , and is common in cases such as sparse reward problems . This issue occurs even if the replay buffer data is fixed , which is the basis of offline RL ( Levine et al. , 2020 ) . The other question occurs during the discretization phase of DARTS , as any continuous relaxation method must deal with integrality gaps when rounding back to a discrete space , with this gap dependent on the optimization landscape . At the core of multiple discretization procedures ( Liu et al. , 2019 ; Chu et al. , 2020c ; Hundt et al. , 2019 ) lies the idea of op selection based on edge weights . However , RL ’ s optimization landscape is significantly different ( Ilyas et al. , 2020 ; Ahmed et al. , 2019 ) from SL , and thus even if the supernet trains in RL , it is unclear if e.g . DARTS ’ s default discretization procedure leads to a better discrete cell , which is the original goal of DARTS and NAS in general . In this paper , we carefully analyze the extent at which DARTS can optimize perception modules for RL environments . While there is a rich series of active study on the theoretical and empirical behavior of DARTS optimization in SL ( Wang et al. , 2021b ; Chen & Hsieh , 2020 ; Li et al. , 2021 ) , our work is one of the first to study its optimization in RL , a completely different paradigm . Through extensive experiments , we ultimately validate that even vanilla DARTS is capable of architecture search in RL , laying a core foundation for future work in NAS for RL . In summary , our contributions are : • We demonstrate DARTS can be integrated with existing RL algorithms in a minimally invasive and easy-to-use manner , agnostic to the specific pipeline . Remarkably , on the challenging Procgen benchmark featuring infinite data , the DARTS supernet is able to train from scratch and efficiently produce similar training curves as baselines using both on-policy ( PPO ) and off-policy ( Rainbow ) methods , both of which require inference on the network . • Due to the common issue of integrality gaps between supernet cells and discrete cells , we verify both qualitatively via visualizations and quantitatively via rewards , that discretized cells start with suboptimal architectures and gradually evolve to better architectures . However , we also demonstrate how this can fail , especially if the corresponding supernet fails to train , with further extensive ablations in the Appendix . • By evaluating the final discrete cells at the end of supernet training , we discover that many environments benefit from learned , custom-tailored architectures , suggesting that humandesigned architectures can be suboptimal . Furthermore , we find that joint-training a supernet on multiple environments can potentially produce a single transferrable cell . This opens up questions for future study , as large-scale RL can potentially benefit from more applications of NAS techniques . 2 RELATED WORKS . Most previous AutoRL methods can be thought of in terms of blackbox and evolutionary methods . For NAS on RL , these include ( Song et al. , 2021 ; Gaier & Ha , 2019 ; Stanley & Miikkulainen , 2002 ; Stanley et al. , 2009 ) which utilize CPU workers for forward pass evaluations rather than exact gradient computation on GPUs . Such methods are usually unable to train policies involving more than 10K+ parameters due to the sample complexity of zeroth order methods in high dimensional parameter space ( Agarwal et al. , 2010 ) . However , such methods have proved useful on smaller search spaces for RL , such as hyperparameter optimization ( Zhang et al. , 2021 ; Franke et al. , 2020 ; Parker-Holder et al. , 2020 ; Jaderberg et al. , 2017 ) and algorithm search/ “ learning to learn ” methods ( Faust et al. , 2019 ; Co-Reyes et al. , 2021 ) . In terms of differentiable search in RL , methods for tuning hyperparameters and environments include ( Zahavy et al. , 2020 ; Xu et al. , 2018 ; Hu et al. , 2020 ) . The only previous known application of DARTS is ( Akinola et al. , 2021 ) , which searches for the optimal way of combining the observation and action tensors together in off-policy QT-OPT ( Kalashnikov et al. , 2018 ) . However , this approach does not modify the image encoder nor uses the supernet for inference , as off-policy robotic data collected independently from the actual policy is used for training , which leaves the applicability of DARTS to inference-dependent RL as an open question addressed in our work . 3 METHODOLOGY . 3.1 VANILLA DARTS AND NAS PRELIMINARIES . We provide a brief review of DARTS which can be pictorially summarized in Figure 2 , although more details can be found in the original DARTS paper ( Liu et al. , 2019 ) . DARTS optimizes substructures called cells , where each cell contains I intermediate nodes organized in a directed acyclic graph , where each node x ( i ) , represents a feature map , and each edge ( i , j ) consists of an operation ( op ) o ( i , j ) , with later nodes x ( j ) merged ( e.g . summation ) from some previous o ( i , j ) ( x ( i ) ) . A DARTS supernet is constructed by continuously relaxing selection of ops in O , via softmax weighting , i.e . o ( i , j ) ( x ( i ) ) = ∑ o∈O p ( i , j ) o · o ( x ( i ) ) , where p ( i , j ) o = exp ( α ( i , j ) o ) ∑ o′∈O exp ( α ( i , j ) o′ ) . The cell ’ s output is by default the result of a Conv1x1 op on the depthwise concatenation of all intermediate node features , although this may be changed ( e.g . simply output last intermediate node ’ s features ) . We denote the collection of all architecture variables a ( i , j ) o as α . Since the total op space O must contain the Zero and Skip Connection ops , we mainly refer to the main op space asObase whereO = Obase ∪ { Zero , Skip } . In the standard optimization framework for SL where there is the notion of a validation set , the learning procedure consists of a bilevel optimization problem : α∗ = argmin α Lval ( θ∗ , α ) s.t . θ∗ = argmin θ Ltrain ( θ , α ) ( 1 ) where L ( θ , α ) is a loss function defined with respect to the supernet . α∗ is then discretized into a final cell , by representing each edge ( i , j ) with the highest softmax weighted op argmaxo∈O , o6=zero p ( i , j ) o and then retaining only the top K = 2 incoming edges for each intermediate node . Following common NAS practices ( Zoph et al. , 2018 ; Pham et al. , 2018 ; Zoph & Le , 2016 ) , we construct our supernet by stacking both normal ( N times ) and reduction cells ( R times ) together into blocks , which are themselves also stacked together D times ( see Figure 2 ) . Normal cells preserve the input shape on output , while reduction cells apply a stride of 2 to reduce height and width immediately on the input . Each block can have its own associated convolutional channel depth which is used throughout all cells in the block . Another common NAS practice we adopt is to train a smaller supernet ( i.e . depth 16 ) to reduce computation time , but evaluate final discretized cells on larger models ( i.e . depth 64 ) , with D = 3 layers for cheap large-scale runs and D = 5 layers for fine-grained A/B testing .

Summary: This paper studies the use of DARTS within the RL setting. Specifically, it investigates how DARTS can optimize perception modules for RL environments. RL-DARTS is optimised in a typical RL setting, using the standard RL loss function, in an end-to-end manner. The proposed approach is evaluated using the Procgen benchmark, and compared to the IMPALA-CNN architecture. Several ablation studies are performed to further examine the contributions of the components of the proposed approach.