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: This paper introduces a method for searching for the best auxiliary loss function from a huge search space automatically, while the best auxiliary loss function is defined as the one that encourages the agent to get a higher return. The searching space of auxiliary loss functions is finite and discrete. It is a combination of (1) manually designed similarity measures, which is used for encouraging the prediction and target to be similar in the auxiliary task, and (2) the binary mask, which is used for selecting auxiliary loss inputs. In this work, the size of the search space is around $4.6 \times 10^{19}$. The space is pruned firstly by random sampling from the similarity measure space and choosing the one showing the highest averaged performance, then the evolutionary algorithm is applied to select top candidates for loss inputs, which is a method in the literature. The chosen auxiliary losses are empirically shown to be helpful with increasing the learning efficiency and have generalization ability to new tasks.