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 proposes to combine differentiable neural architecture search (DARTS) with standard reinforcement learning (RL) frameworks by searching the model structures (convolutional cells) for the policy and value functions. The authors have applied the method to infinitely procedurally generated Procgen benchmark and demonstrated the benefits of DARTS in RL on search efficiency in terms of time and compute. In addition, the proposed method can be easily integrated with existing RL pipelines by simply replacing the image encoder with a DARTS supernet, compatible with both off-policy and on-policy RL algorithms. The authors further show that the supernet gradually learns better cells with more training iterations, leading to alternative architectures that can be highly competitive against manually designed policies and verify previous design choices for RL policies.