Year: 2021

Abstract:
A good state representation is crucial to reinforcement learning (RL) while an ideal representation is hard to learn only with signals from the RL objective. Thus, many recent works manually design auxiliary losses to improve sample efficiency and decision performance. However, handcrafted auxiliary losses rely heavily on expert knowledge, and therefore lack scalability and can be suboptimal for boosting RL performance. In this work, we introduce Automated Auxiliary loss for Reinforcement Learning (AARL), a principled approach that automatically searches the optimal auxiliary loss function for RL. Specifically, based on the collected trajectory data, we define a general auxiliary loss space of size 4.6×10 and explore the space with an efficient evolutionary search strategy. We evaluate AARL on the DeepMind Control Suite and show that the searched auxiliary losses have significantly improved RL performance in both pixel-based and state-based settings, with the largest performance gain observed in the most challenging tasks. AARL greatly outperforms state-of-the-art methods and demonstrates strong generalization ability in unseen domains and tasks. We further conduct extensive studies to shed light on the effectiveness of auxiliary losses in RL.

heading: 1 INTRODUCTION
text: 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.

heading: 2 METHODOLOGY
text: 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.

heading: 2.1 SEARCH SPACE
text: 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.

heading: 2.2 SEARCH STRATEGY
text: 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.

heading: 3 EXPERIMENTS
text: We evaluate our method and the obtained auxiliary loss functions on DMControl suite (Tassa et al., 2018). DMControl suite is powered by MuJoCo physics engine (Todorov et al., 2012) with various challenging tasks. In both pixel-based (images) and state-based (proprioceptive features) settings, we choose Soft Actor-Critic (SAC) (Haarnoja et al., 2018) as our base RL algorithm since it is the state-of-the-art model-free RL algorithm. Implementation details are given in Appendix A.

heading: 3.1 PIXEL-BASED RL
text: Experiment Settings We follow the same network architecture as CURL with a 4-layer convolutional encoder for image input. In the search phase, we compare our method AARL to SAC, SAC (no aug), and CURL with the same hyper-parameters reported in Appendix A.2.1. All methods randomly crop images from 100 × 100 to 84 × 84 as data augmentation apart from SAC (no aug). After searching, we test the generalization ability of the obtained auxiliary losses in other unseen environments in comparison with state-of-the-art model-free and model-based methods.
Search Results We apply our search algorithm on Cheetah-Run, Reacher-Easy, and Walker-Walk, which we call “training environments”. We choose these environments because they are more challenging and there might be more room for improvement with auxiliary losses. For each environment, we set the total budget of each experiment to 16k GPU hours (on NVIDIA P100) and terminate the experiment when the resource is exhausted. The evolution process is demonstrated in Figure 4, where the y-axis shows the AULC of candidates. The results show that, during the evolution process, we can easily find candidates surpassing baselines in the first stage. In the subsequent stages, the overall population score continues to grow, and most candidates can outperform baselines. This demonstrates the effectiveness of our search space and the efficiency of our search strategy.
Generalize to Other Environments To reduce the risk that auxiliary losses overfit to one particular environment and avoid selecting candidates that perform well only due to randomness, we run cross-validation on the three searched environments to decide the best candidate. The details of this process can be found in the Appendix A.4. After cross-validation, we finalize an auxiliary loss function, which we call “AARL-Pixel” (all the top candidates during evolution are reported in Appendix D), and we train agents with the obtained auxiliary loss to compare with state-of-the-art model-free and model-based methods on DMControl100k and DMControl500k. We select these environments because they are common benchmarks for state-of-the-art algorithms of pixel-based RL. (Laskin et al., 2020). Note these comparisons are additionally made on three unseen environments, that are entirely not accessible during evolution, to test the generalizability of our method and to ensure our comparisons are fair. The results are summarized in Table 3, where AARL greatly outperforms state-of-the-art methods on 11 out of 12 benchmark settings. Note that Finger-Spin, Reacher-Easy and Ball in cup-Catch are environments that are never used during the search, while AARL still performs very well on them. This result implies that AARL-Pixel is a robust and effective auxiliary loss that is potentially helpful to RL under various settings. The performance gain
is obvious in DMControl100K, where AARL-Pixel (pixel-based input) has shown great sample efficiency.
3.2 STATE-BASED RL
Experiment Settings As for state-based RL, since the state is low-dimensional, we simply use a 1-layer densely connected MLP as the state encoder as shown in Figure 5. So, for this setting, we focus on this simple encoder structure. Additional ablations on state encoder architectures are given in Section 3.4. In the search phase, we compare AARL to SAC-Identity, SAC-DenseMLP, CURL-DenseMLP. To ensure a fair comparison, all SAC related hyper-parameters are the same as those reported in the CURL paper. Details can be found in Appendix A.2.2. SAC-Identity is vanilla SAC with no state encoder, while the other three methods (AARL, SAC-DenseMLP, CURL-DenseMLP) use the same encoder architecture. Different from the pixel-based setting, there is no data augmentation in the statebased setting. Note that many environments that are challenging in pixel-based settings become easy to tackle with state-based inputs. Therefore we apply our search framework to more challenging environments for state-based RL, including Cheetah-Run, Hopper-Hop and Quadruped-Run.
Search Results Similar to pixel-based settings, we approximate AULC with the average score agents achieve at 300k, 600k, 900k, 1200k and 1500k time steps1. For each environment, we early stop the experiment when the budget of 1,500 GPU hours is exhausted. The evolution process is shown in Figure 6, where we find a large portion of candidates outperform baselines (horizontal dashed lines). The performance improvement is especially significant on Cheetah-Run, where almost all candidates in the population greatly outperform all baselines by the end of the first stage.
Generalize to Other Environments Similar to pixel-based settings, we also use cross-validation to select the best loss function, which we call “AARL-State” here (all the top candidates during evolution are reported in Appendix D), and run a thorough evaluation on all 18 environments. The results are summarized in Table 4. AARL-State again brings significant performance gain, achieving much stronger sample efficiency than SAC. It is noteworthy that AARL-State is able to outperform baseline methods in 16 out of 18 environments (with 15 unseen environments). These results show that even for tasks with lower-dimensional state space, there is still a huge potential for improvement with better auxiliary objectives. Moreover, this performance gain is especially significant
1As for Cheetah-Run, we still use average score agents achieve at 100k, 200k, 300k, 400K and 500k time steps since agents converge close to optimal score within 500k time steps.
in complex environments like Humanoid, where SAC barely learns anything at 1000K time steps, while AARL-State is able to achieve much better performance.

heading: 3.3 ANALYSIS OF AUXILIARY LOSS FUNCTIONS
text: We have collected a large number of data from the evolution process so far. To make the best use of these data, we now conduct a comprehensive analysis to investigate statistical relations between patterns of auxiliary loss functions and RL performance.
Typical Patterns We first analyze how RL performance is affected by typical patterns of auxiliary loss functions. Define the set of the input to the encoder that is being trained as “source”, and the set of prediction targets as “target”. We say an auxiliary loss candidate has a certain pattern if the pattern’s source is a subset of the candidate’s source, and the pattern’s target is a subset of the candidate’s target. For instance, a loss candidate of {st, at} → {st+1, st+2} has the pattern {st, at} → {st+1}, and does not have the pattern {at, st+1} → {st}. The patterns we consider include forward dynamics {st, at} → {st+1}, inverse dynamics {at, st+1} → {st}, reward prediction {st, at} → {rt}, action inference {st, st+1} → {at} and state reconstruction in the latent space {st} → {st}. For each of these patterns, we divide auxiliary loss candidates into two classes: 1) with this pattern or 2) without this pattern. We then calculate the average performance of these two classes and compare their performances to conclude whether this pattern is helpful for RL. The results are summarized in Table 5. Here a positive number indicates that this pattern is beneficial, and if the performance gain is statistically significant, the number is marked with the asterisk, indicating it is very likely to be helpful. A negative number indicates that this pattern is detrimental and correlated to worse performance. We highlight some interesting observations as follows. 1) Forward dynamics is helpful in most environments, and improves the RL performance on Reacher-Easy (Pixel) and Cheetah-Run (State) significantly (i.e., p-value<0.05). 2) State reconstruction in the latent space is able to improve RL performance in pixel-based environments while it has significant negative effects on state-based environments. We hypothesize that this is because in the pixel-based setting, data augmentation encourages the encoder to learn augmentation-invariant representations, while in the state-based setting, no augmentation is used, and thus the encoder learns no useful rep-
resentations. This also explains why CURL is underperforming in state-based experiments. Note that not all typical patterns bring performance gains. Some can even be very detrimental, and this shows that commonly-used auxiliary loss patterns do not work well in the state-based setting, further highlighting the fact that great research potential lies in this setting.
Number of Sources and Targets We then investigate whether it is more beneficial to use a small number of sources to predict a large number of targets (ntarget > nsource, e.g., use st to predict st+1, st+2, st+3), or the other way around (ntarget < nsource, e.g., use st, st+1, st+2 to predict st+3). The statistical results are shown in Table 5, where we find that auxiliary losses with more states on the target side have a significant advantage over losses with more states on the source side. This result echoes recent works (Stooke et al., 2021; Schwarzer et al., 2020) which show that predicting more states leads to strong performance gains.

heading: 3.4 ADDITIONAL ABLATION STUDY
text: Search Space Pruning As introduced in Section 2.2, we decompose the full search space into loss operator and loss inputs. Here we try to directly apply evolution strategy to the whole space without the pruning step. The comparison results are shown in Figure 7. We can see that pruning improves the evolution process, making it easier to find good candidates.
Table 6: Normalized episodic rewards of AARL (mean & standard deviation for 5 seeds of 6 environments) on statebased DMControl100K with different encoder architectures.
AARL-MLP (1-layer) AARL-MLP (4-layer) 0.919 ± 0.217 0.544 ± 0.360
AARL-DenseMLP (1-layer) AARL-DenseMLP (4-layer) 1.000 ± 0.129 0.813 ± 0.218
Encoder Architecture for State-based RL As shown in Figure 5, we choose a 1-layer densely connected MLP as the state encoder for state-based RL. Here, we conduct an ablation study on different encoder architectures in the state-based setting. The results are summarized in Table 6, where AARL with 4-layer encoders consistently perform worse than 1-layer encoders. We also note that dense connection is helpful in the state-based setting compared with naive MLP encoders.

heading: 4 RELATED WORK

heading: 4.1 REINFORCEMENT LEARNING WITH AUXILIARY LOSSES
text: Using auxiliary tasks to improve sample efficiency of RL, especially on pixel-based control tasks, has been explored in many recent works. A number of manually designed auxiliary objectives are shown to boost RL performance, including observation reconstruction (Yarats et al., 2019), reward prediction (Jaderberg et al., 2017), dynamics prediction (De Bruin et al., 2018) and contrastive learning objectives (Laskin et al., 2020; Schwarzer et al., 2020; Stooke et al., 2021). It is noteworthy that these works mainly focus on pixel-based settings, while only a limited number of works study the state-based setting (Munk et al., 2016; Ota et al., 2020). Although it might seem that the state-based setting benefits less from auxiliary tasks due to their lower-dimensional state space, we show that there is in fact a huge potential of improving state-based RL performance with auxiliary objectives.
Compared to the previous works, we point out two major advantages of our approach. 1) Instead of handcrafting an auxiliary loss with expert knowledge, AARL automatically searches for the best auxiliary loss, relieving researchers from such tedious work. 2) AARL is a principled approach that can be used in arbitrary RL settings. In both pixel-based and the rarely studied state-based settings, we discover good auxiliary losses that bring significant performance improvement.

heading: 4.2 AUTOMATED REINFORCEMENT LEARNING
text: It is well known that RL training is sensitive to hyper-parameters and environment changes (Henderson et al., 2018). Thus many works have attempted to use techniques in AutoML to alleviate human intervention. Exiting works focus on hyper-parameter optimization (Espeholt et al., 2018; Paul et al., 2019; Xu et al., 2020; Zahavy et al., 2020), reward search (Faust et al., 2019; Veeriah et al., 2019) and network architecture search (Runge et al., 2019; Franke et al., 2021). In contrast, our method aims to search for auxiliary loss functions that generalize across different environments.

heading: 4.3 AUTOMATED LOSS DESIGN
text: A few recent works in the AutoML community have been automating the design of good loss functions that outperform traditionally handcrafted ones. Specifically, AM-LFS (Li et al., 2019) defines the loss function search space as a parameterized probability distribution of softmax loss hyperparameters. AutoLoss-Zero (Li et al., 2021) proposes to search loss functions with primitive mathematical operators. These methodologies are proposed specifically for computer vision tasks.
For RL, existing works focus on searching for a better RL objective, EPG (Houthooft et al., 2018) and MetaGenRL (Kirsch et al., 2020) define the search space of loss functions as parameters of a low complexity neural network. Co-Reyes et al. (2020) defines the search space of RL loss functions as a directed acyclic graph and discovers two DQN-like regularized RL losses. Note that none of these works investigate auxiliary loss functions, which are crucial to facilitate representation learning in RL and to make RL successful in highly complex environments. To the best of our knowledge, our work is the first attempt to search for auxiliary loss functions and greatly improve RL performance.

heading: 5 CONCLUSION AND FUTURE WORK
text: We present AARL, a principled and universal approach for automated auxiliary loss design for RL. With this framework, we discover highly performant auxiliary loss functions that generalize beyond the training environments for both pixel-based and state-based settings. We present extensive experimental results that provide strong empirical evidence for the effectiveness of our method, and also conduct an in-depth investigation of the statistical relations between auxiliary loss patterns and RL performance. We hope our studies provide insights that will deepen the understanding of auxiliary losses in RL, and shed light on how to make RL more efficient and practical. One interesting future work direction is to further investigate the relationship between auxiliary loss patterns and RL performance, and to better understand how and why auxiliary losses can help to improve RL performance. Another direction is to study how to combine the RL objective and the auxiliary loss functions in a more effective manner.

heading: REPRODUCIBILITY STATEMENT
text: All implementation details are introduced in the main text and appendix (Appendix A). We will open source our codebase to facilitate future research.

heading: A.1 ARCHITECTURE AND SEARCH SPACE

heading: A.1.1 STATE ENCODER ARCHITECTURES
text: We demonstrate in Figure 8 the overall architecture where our auxiliary loss is used. The architecture is generally identical to those adopted in CURL (Laskin et al., 2020). “Pixel-based” and “Statebased” are the architectures we used in our experiments. “MLP” and “4-layer DenseMLP” are for ablations.

heading: A.1.2 SIAMESE NETWORK
text: For a fair comparison with baseline methods, we follow the same Siamese networks structure for representation learning as CURL (Laskin et al., 2020). As shown in Figure 1, when computing targets y for auxiliary losses, we map states to state embeddings with a target encoder. We stop gradients from target encoder θ and update θ in the exponential moving averaged (EMA) manner where θ′ = τ θ̂ + (1− τ)θ. This step, i.e., to freeze the gradients of the target encoder, is necessary when the loss is computed without negative samples. Otherwise, encoders will collapse to generating the same representation for any input. We have verified this in our early experiments.

heading: A.1.3 LOSS OPERATORS
text: Instance Discrimination Our implementation is based on InfoNCE loss (Oord et al., 2018):
L = log exp(φ(ŷ, y+)) exp(φ(ŷ, y+)) + ∑K−1 i=0 exp(φ(ŷ, yi)) (3)
The instance discrimination loss can be interpreted as the log-loss of a K-way softmax classifier whose label is y+. The difference between discrimination based loss operators lies in the discrimination objective φ used to measure agreement between (ŷ, y) pairs. Inner employs inner product φ(ŷ, y) = ŷ>y while Bilinear employs bilinear product φ(ŷ, y) = ŷWy, where W is a learnable parameter matrix. Cosine uses cosine distance φ(ŷ, y) = ŷ
>y ‖ŷ‖·‖y‖ for further matrix calculation. As
for cross entropy based loss without negative samples, we only take diagonal elements for matrix M where Mi,j = φ(ŷi, yj) for cross entropy calculation.
Mean Squared Error The implementation of MSE-based loss operators are straightforward. MSE loss operator = (ŷ − y+)2 while normalized MSE = ( ŷ‖ŷ‖ − y+ ‖y+‖ )
2. When combined with negative samples, MSE loss operator (with negative pairs) = (ŷ − y+)2 − (ŷ − y−)2 while normalized MSE (with negative pairs) = ( ŷ‖ŷ‖ − y+ ‖y+‖ ) 2 − ( ŷ‖ŷ‖ − y− ‖y−‖ ) 2.
Evaluation of Performance Expectation To reduce the large search space consists of loss input I and loss operator f , we run 15 trials for each loss operator to get an estimation of performance expectation. For each of 10 possible f in the search space (5 operators with optional negative samples), we run 5 trials on each of the 3 pixel-based environments (used in evolution) with the same loss inputs {st, at} → {st+1}, as we found that forward dynamics is a reasonable representative of our search space with highly competitive performance.

heading: A.2 TRAINING DETAILS

heading: A.2.1 HYPER-PARAMETERS IN THE PIXEL-BASED SETTING
text: We use the same hyper-parameters for AARL, SAC (no aug), SAC and CURL during the search phase, to ensure a perfectly fair comparison. When evaluating the searched auxiliary loss, we use a slightly larger setting (e.g., larger batch size) to train RL agents sufficiently. A full list is shown in Table 7.

heading: A.2.2 HYPER-PARAMETERS IN THE STATE-BASED SETTING
text: We use the same hyper-parameters for AARL, SAC-Identity, SAC-DenseMLP and CURLDenseMLP, shown in Table 8. As training in state-based environments is substantially faster than pixel-based environments, there is no need to balance training cost and agent performance. We use this setting for both the search phase and the final evaluation phase.

heading: A.3 EVOLUTION STRATEGY
text: Horizon-changing Mutations There are two kinds of mutations that are able to change horizon length. One is to decrease horizon length. Specifically, we remove the last time step, i.e.,
(at+k, rt+k, st+k+1) if the target horizon length is k. The other is to increase horizon length, in which we append three randomly generated bits to the given masks at the end. We do not shorten the horizon when it becomes too small (less than 1), or lengthen the horizon when it is too long (exceeding 10).
Mutating Source and Target Masks When mutating a candidate, the mutation on the source and the target masks are independent to each other except for horizon change mutation where two masks could either both increase horizon or decrease horizon.
Loss-rejection Protocol The loss-rejection protocol makes sure that invalid loss functions do not go into the expensive computation of RL training. Concretely, the following conditions must be satisfied to make a valid loss function: 1) having at least one state embedding in m̂ to make sure the gradient of auxiliary loss backward propagates to the state encoder, and 2) target m is not empty. If a loss is rejected, we repeat the mutation to fill up the population.
Initialization At each initialization, we randomly generate 100 auxiliary loss functions (every bit of masks are generated from Bernoulli(p) where p = 0.5.) and generate 25 auxiliary loss functions with prior probability, which makes the auxiliary loss have some features like forward dynamics prediction or reward prediction. The prior probability for generating forward dynamics pattern is: 1) every bit of states from target is generated from Bernoulli(p) where p = 0.2; 2) every bit of actions from source is generated from Bernoulli(p) where p = 0.8; 3) every bit of states from target is generated by flipping the states of source; 4) The other bits are generated from Bernoulli(p) where p = 0.5. The prior probability for generating reward prediction pattern is: 1) every bit of rewards from target is generated from Bernoulli(p) where p = 0.8; 2) Every bit of states and actions from target are 0; 3) The other bits are generated from Bernoulli(p) where p = 0.5.

heading: A.4 CROSS VALIDATION
text: After the evolutionary search is done, we choose top-5 candidates for cross-validation across environments. This step is to avoid overfitting to one single environment. For each top-5 loss candidate, we run 3 trials on each of these 3 training environments (45 runs in total). Then we sort these top-5 candidates by cross-validation overall performance (mean AULC score on three environments) and retrieve the top-1 candidate as our final searched loss. We do this in pixel-based environments and state-based environments separately. For pixel-based RL, the top-5 candidates are the top candidates
from stage 5 of Cheetah-Run (Pixel). For state-based RL, the top-5 are the top candidates from stage 4 and 5 of Cheetah-Run (State).

heading: A.5 BASELINES IMPLEMENTATION
text: Pixel-based Setting CURL (Laskin et al., 2020) is the main baseline to compare with in the pixel-based setting, which is considered to be the state-of-the-art pixel-based RL algorithm. CURL learns state representations with a contrastive auxiliary loss. PlaNet (Hafner et al., 2019b) and Dreamer (Hafner et al., 2019a) are model-based methods that generate synthetic rollouts with a learned world model. SAC+AE (Yarats et al., 2019) uses a reconstruction auxiliary loss of images to boost RL training. SLAC (Lee et al., 2019) leverages forward dynamics to construct a latent space for RL agents. Note that there are two versions of SLAC with different gradient Updates per agent step: SLACv1 (1:1) and SLACv2(3:1). We adopt SLACv1 for comparison since all methods only make one gradient update per agent step. Pixel SAC are just vanilla SAC (Haarnoja et al., 2018) agents with images a as inputs respectively.
State-based Setting As for the state-based setting, we compare AARL-State with SAC-Identity, SAC and CURL. SAC-Identity is the vanilla state-based SAC where states are directly fed to actor/critic networks. SAC and CURL use the same architecture of 1-layer densely connected MLP as a state encoder. Note that both AARL and baseline methods use the same hyper-parameter reported in Table 8 without additional hyper-parameter tuning.

heading: B FURTHER EXPERIMENT RESULTS

heading: B.1 LEARNING CURVES FOR AARL ON PIXEL-BASED DMCONTROL
text: We benchmark the performance of AARL to the best-performing pixel-based baseline (CURL). As shown in Figure 9, the sample efficiency of AARL outperforms CURL on 10 out of 12 environments. Note that the learning curves of CURL may not match to the data in Table 3, this is because we use the data reported in CURL paper for tabular while we rerun CURL for learning curves plotting, where we find the performance of our rerunning CURL is sightly below the CURL paper.
Ball in cup-Catch Cartpole-Swingup Cartpole-Swingup sparse Cheetah-Run

heading: Pendulum-Swingup Reacher-Easy Walker-Stand Walker-Walk

heading: B.2 EFFECTIVENESS OF AULC SCORES
text: To illustrate intuitively why we use the area under learning curve instead of other metrics, we select top-10 candidates with different evolution metrics. Figure 10 demonstrates the usage of AULC score could well balance both sample efficiency and final performance. The learning curves of the top-10 candidates selected by AULC score look better than the other two metrics (that select top candidates simply with 100k step score or 500k step score).

heading: B.3 HISTOGRAM OF AUXILIARY LOSS ANALYSIS
text: The histogram of each pattern analysis is shown in Figure 11.

heading: C SEARCH SPACE COMPLEXITY ANALYSIS
text: The search space size is the size of loss input space multiplied by the size of loss operator space.
For loss input, we calculate separately for each possible horizon length k. When length is k, the interaction sequence length (st, at, rt, · · · , st+k) has length (3k+1). For binary mask m̂, there are 23k+1 different options. There are also 23k+1 distinct binary mask m to select targets. Therefore, there are 26k+2 combinations when horizon length is fixed to k. As our maximum horizon is 10, we enumerate k from 1 to 10, resulting in ∑10 i=1 2 6i+2.
For loss operator, we can learn intuitively from Table 2 that there are 5 different similarity measures with or without negative samples, resulting in 5× 2 = 10 different loss operators. In total, the size of the entire space is
10× 10∑ i=1 26i+2 ≈ 4.6× 1019.

heading: D TOP-PERFORMING AUXILIARY LOSSES

heading: D.1 AARL-PIXEL AND AARL-STATE
text: The auxiliary loss of AARL-Pixel is:
{st+1, at+1, at+2, at+3} → {rt, rt+1, st+2, st+3} (4)
which is the third-best candidate of stage 4 in Cheetah-Run (Pixel).
The auxiliary loss of AARL-State is:
{st, at, at+1, st+2, at+2, at+3, rt+3, at+4, rt+4, at+5, at+7, st+8, at+8, rt+8} → {st+1, st+3, at+4, st+6, st+9}
(5)
which is the fourth-best candidate of stage 4 in Cheetah-Run (State).
These two losses are chosen because they are the best-performing loss functions during crossvalidation (Appendix A.4).

heading: D.2 DURING EVOLUTION
text: We report all the top-5 auxiliary loss candidates during evolution in this section.

Review summaries:
In this paper, the authors intend to automatically search for the optimal auxiliary loss. To achieve this goal, the authors propose Automated Auxiliary loss for Reinforcement Learning (AARL), which uses an evolutionary search strategy to explore the very large loss space. The paper conducts extensive experiments on the DeepMind Control Suite and shows that the searched auxiliary losses have significantly improved RL performance in both pixel-based and state-based settings. 

Reinforcement learning agents can benefit from auxiliary tasks. This paper introduces an evolutionary algorithm to automate the design of auxiliary loss functions, based on prior approaches like forward dynamics, inverse dynamics, contrastive state representation learning, etc. The evolution solves a bilevel optimization problem in which the inner loop is regular RL training while the outer loop evolves the loss function. Experiments are conducted on the Deepmind Control Suite for both pixel and state-based observations. 

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. 

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.