AARL: Automated Auxiliary Loss for Reinforcement Learning

1 INTRODUCTION . Reinforcement learning ( RL ) has been a hot research topic and has shown significant progress in many fields ( Mnih et al. , 2013 ; Silver et al. , 2016 ; Gu et al. , 2017 ; Vinyals et al. , 2019 ) . Recent RL research focuses on obtaining a good representation of states as it is shown to be the key to improve sample efficiency of RL ( Laskin et al. , 2020 ) . This is because in high-dimensional environments , applying RL directly on the complex state inputs is incredibly sample-inefficient ( Lake et al. , 2017 ; Kaiser et al. , 2019 ) , making RL hard to scale to real-world tasks where interacting with the environment is costly ( Dulac-Arnold et al. , 2019 ) . Standard RL paradigm learns the representation from critic loss ( value prediction ) and / or actor loss ( maximizing cumulative reward ) , which hardly extracts informative representations in challenging environments like pixel-based RL and complex robotics systems ( Lake et al. , 2017 ; Kober et al. , 2013 ) . This motivates adding auxiliary losses in support of learning better latent representations of states . The usage of auxiliary loss encodes prior knowledge of RL and its environments , and puts regularization on training ( Shelhamer et al. , 2016 ) . Typical auxiliary losses are manually designed by human experts , including observation reconstruction ( Yarats et al. , 2019 ) , reward prediction ( Jaderberg et al. , 2017 ) and environment dynamics prediction ( Shelhamer et al. , 2016 ; De Bruin et al. , 2018 ; Ota et al. , 2020 ) . However , such auxiliary loss designs and choices rely heavily on human knowledge of what might be helpful for RL training , which requires extensive human efforts to scale to new tasks and can be suboptimal for improving RL performance ( Yang & Nachum , 2021 ) . In this paper , we rigorously treat auxiliary loss functions for RL as a first-class problem to explicitly address the question : what is the universal approach to find a good auxiliary loss function for RL ? Considering that the automated machine learning ( AutoML ) community has shown promising results with automated loss search in computer vision tasks ( Li et al. , 2019 ) , we propose to automate the process of designing auxiliary loss functions of RL . Specifically , we formulate the task as a bi-level optimization , where we try to find the best auxiliary loss function , which , to the most extent , helps train a good RL agent . The inner-level problem is a typical RL training problem , while the outer-level can be seen as a search problem . To tackle this , we first design a novel search space of auxiliary loss functions , which is a combination of two components : loss input and loss operator . As shown in Table 1 , the search space covers many existing handcrafted losses , and it is also substantially large ( as large as 4.6 × 1019 ) . We then propose an efficient search strategy to explore the space . This is done in a two-step manner . We first finalize the loss operator specification and then use an evolutionary strategy that performs mutations on configurations of loss input , to identify the top-performing loss inputs quickly . We evaluate our search framework on both pixel-based and state-based environments in the DeepMind Control suite ( Tassa et al. , 2018 ) . Extensive experiments show that the searched auxiliary loss functions greatly outperform state-of-the-art methods and can easily transfer to unseen environments ( that are never used during the search ) , showing that our proposed method is robust and effective . We highlight the main contributions of this paper below : • We introduce a principled and universal approach for auxiliary loss design in RL . To the best of our knowledge , we are the first to derive the optimal auxiliary loss function with an automatic process . Our framework can be easily applied to arbitrary RL tasks to search for the best auxiliary loss function . • We demonstrate that AARL significantly outperforms state-of-the-art methods in both pixel-based and state-based environments . The searched auxiliary loss functions show strong generalization ability to transfer to unseen environments . • We analyze the derived auxiliary loss functions and deliver some insightful discussions that we hope will deepen the understanding of auxiliary losses in RL . We will also open source our codebase to facilitate future research . 2 METHODOLOGY . We consider the standard Markov Decision Process ( MDP ) setting where we specify the state , action and reward at time step t as ( st , at , rt ) . The sequence of interaction data is denoted as ( st , at , rt , · · · , st+k ) , where k represents the horizon of this sequence . Suppose we have an RL agent Mω parameterized by ω and R ( Mω ; E ) is the agent performance ( i.e. , cumulative discounted reward ) in environment E . The goal of AARL is to find the optimal auxiliary loss function L , such that , when ω is optimized under an arbitrary RL objective function LRL ( e.g. , actor-critic loss ) together with L , the agent achieves the best performance . Formally , we have max L R ( Mω∗ ( L ) ; E ) , s.t . ω∗ ( L ) = argmin ω ( LRL ( ω ; E ) + λL ( ω ; E ) ) , ( 1 ) where λ is a hyper-parameter that controls the relative weight of the auxiliary task . We solve this bi-level optimization problem with AutoML techniques . The inner optimization is a standard RL training procedure . For the outer one , we define a finite and discrete search space ( Section 2.1 ) , and use a variation of evolution strategy to explore the space ( Section 2.2 ) . We will explain the details in the rest of this section . 2.1 SEARCH SPACE . An auxiliary loss function L can be viewed as a combination of two components : 1 ) loss input I and 2 ) loss operator f . We define a search space as shown in Figure 1 , where I is a pair of binary masks ( m̂ , m ) that selects from interaction data to compute ŷ and y , and f is an operator that aggregates the selected inputs into an expression of the loss function with scalar output . Loss Input Unlike supervised machine learning tasks with an explicit prediction and ground truth for the loss function , auxiliary losses in RL have no ground truths beforehand . Instead , they are generated accordingly upon interaction with the environment itself in a self-supervised manner . As shown in Figure 1 , we take some tokens from the sequence ( st , at , rt , · · · , st+k ) with binary mask vector m̂ and feed them into the encoder θ̂ which maps states to latent representations . The predictor module then tries to predict targets ( encoded with a momentum encoder θ , as is typically done in recent works ) , which are another few tokens we select from the sequence with another binary mask vectorm . Details about momentum encoder are given in Appendix A.1.2 . In other words , the search space for loss input is a pair of binary masks ( m̂ , m ) , each of which is up to length ( 3k + 1 ) if the length of an interaction data sequence , i.e. , horizon , is limited to k steps . In our case , we set the maximum horizon length kmax = 10 . Loss Operator We optimize ŷ to be as similar as possible to y . Therefore , we make loss operator f cover commonly-used similarity measures , including inner product ( Inner ) ( He et al. , 2020 ; Stooke et al. , 2021 ) , bilinear inner product ( Bilinear ) ( Laskin et al. , 2020 ) , cosine similarity ( Cosine ) ( Chen et al. , 2020 ) , mean squared error ( MSE ) ( Ota et al. , 2020 ; De Bruin et al. , 2018 ) and normalized mean squared error ( N-MSE ) ( Schwarzer et al. , 2020 ) . Some existing works , e.g. , contrastive objectives like InfoNCE loss ( Oord et al. , 2018 ) , also incorporate the trick to sample un-paired predictions and targets as negative samples and maximize the distances between them . We find this technique applicable to all the loss operators mentioned above and thus incorporate the discriminative version of these operators in our search space . Search Space Complexity As shown in Table 1 , many existing manually designed auxiliary losses can naturally fit into our loss space thus are our special cases , which proves that this loss space is reasonable , flexible , and general . However , this is also at the cost that the space is substantially large . In total , the size of the entire space is 10× ∑10 i=1 2 6i+2 ≈ 4.6× 1019 . The detailed derivation can be found in Appendix C . 2.2 SEARCH STRATEGY . Search Space Pruning Considering that the loss space is extremely large , an effective optimization strategy is inevitably required . Directly grid-searching over the whole space is infeasible because of unacceptable computational cost . Thus some advanced techniques such as space pruning and an elaborate search strategy are necessary . Our search space can be seen as a combination of the space for the input I and the space for the operator f . Inspired by AutoML works ( Dai et al. , 2020 ; Ying et al. , 2019 ) that search for hyper-parameters first and then neural architectures , we approximate the joint search of input and operator in Equation ( 1 ) in a two-step manner . The optimal auxiliary loss { I∗ , f∗ } can be optimized as : max L R ( Mω∗ ( L ) ; E ) = maxI , f R ( Mω∗ ( I , f ) ; E ) ≈ maxI R ( Mω∗ ( I , f∗ ) ; E ) where f∗ ≈ argmax f EI [ R ( Mω∗ ( I , f ) ; E ) ] ( 2 ) To decide the best loss operator , for every f in the loss operator space , we estimate EI [ R ( Mω∗ ( I , f ) ; E ) ] with a random sampling strategy . We run 15 trials for each loss operator and calculate the average of their normalized return on three environments . More details are available in Appendix A.1.3 . Surprisingly , as summarized in Table 2 , the simplest MSE without negative samples outperforms all other loss operators with complex designs . Therefore , this loss operator is chosen for the rest of this paper . Evolution Even with the search space pruning , the rest of the space still remains large . To efficiently identify top candidates from the search space , we adopt an evolutionary algorithm ( Real et al. , 2019 ) . The pipeline of the evolution process is illustrated in Figure 2 . In each stage , we first train and evaluate a population of candidates with size P = 100 , where each candidate corresponds to a pair of binary masks ( m̂ , m ) . The candidates are then sorted by the approximated area under learning curve ( AULC ) ( Ghiassian et al. , 2020 ; Stadie et al. , 2015 ) , which well captures both convergence speed and final performance ( Viering & Loog , 2021 ) and reduces the variance of RL evaluation . The top-25 % candidates are selected after each training stage , and we perform four types of mutations on the selected candidates to form a new population for the next stage : replacement ( 50 % of the population ) ; crossover ( 20 % ) ; horizon decrease and horizon increase ( 10 % ) , as shown in Figure 3 . Moreover , the last 20 % of the new population is from random generation . As for each bit of masks , replacement operation flips the given bit with probability p = 12· ( 3k+1 ) , where k is the horizon length . Crossover generates a new candidate by randomly combining the mask bits of two candidates with the same horizon length in the population . We also incorporate the knowledge from the existing auxiliary loss designs by bootstrapping the initial population with a prior distribution ( Co-Reyes et al. , 2020 ) , so that we can quickly converge to a reasonable loss function . More implementation details are available in Appendix A.3 .

Summary: The paper introduces an approach to automatically discover optimal auxiliary loss for RL, the method is called (AARL). The authors claim that AARL outperform baselines on both pixel-based and state-based task on the DeepMind Control Suite. Also using their method the author analyse different auxiliary losses to identify common patterns.