Original: Reinforcement Learning with Automated Auxiliary Loss Search
Revised: Reinforcement Learning with Automated Auxiliary Loss Search

Original: Abstract
Revised: Abstract

Original: A good state representation is crucial to solving complicated reinforcement learning (RL) challenges.
Revised: A good state representation is crucial to solving complicated reinforcement learning (RL) challenges.

Original: Many recent works focus on designing auxiliary losses for learning informative representations.
Revised: Many recent works focus on designing auxiliary losses for learning informative representations.

Original: Unfortunately, these handcrafted objectives rely heavily on expert knowledge and may be sub-optimal.
Revised: Unfortunately, these handcrafted objectives rely heavily on expert knowledge and may be sub-optimal.

Original: In this paper, we propose a principled and universal method for learning better representations with auxiliary loss functions, named Automated Auxiliary Loss Search (A2LS), which automatically searches for top-performing auxiliary loss functions for RL.
Revised: In this paper, we propose a principled and universal method for learning better representations with auxiliary loss functions, named Automated Auxiliary Loss Search (A2LS), which automatically searches for top-performing auxiliary loss functions for RL.

Original: Speciﬁcally, based on the collected trajectory data, we deﬁne a general auxiliary loss space of size 7 . 5 ⇥ 10 20 and explore the space with an efﬁcient evolutionary search strategy.
Revised: Speciﬁcally, based on the collected trajectory data, we deﬁne a general auxiliary loss space of size 7 . 5 ⇥ 10 20 and explore the space with an efﬁcient evolutionary search strategy.

Original: Empirical results show that the discovered auxiliary loss (namely, A2-winner ) signiﬁcantly improves the performance on both high-dimensional (image) and lowdimensional (vector) unseen tasks with much higher efﬁciency, showing promising generalization ability to different settings and even different benchmark domains.
Revised: Empirical results show that the discovered auxiliary loss (namely, A2-winner ) signiﬁcantly improves the performance on both high-dimensional (image) and lowdimensional (vector) unseen tasks with much higher efﬁciency, showing promising generalization ability to different settings and even different benchmark domains.

Original: We conduct a statistical analysis to reveal the relations between patterns of auxiliary losses and RL performance.
Revised: We conduct a statistical analysis to reveal the relations between patterns of auxiliary losses and RL performance.

Original: The codes and supplementary materials are available at https://seqml.github.io/a2ls .
Revised: The codes and supplementary materials are available at https://seqml.github.io/a2ls .

Original:  Introduction
Revised:  Introduction

Original: Reinforcement learning (RL) has achieved remarkable progress in games [31, 47, 50], ﬁnancial trading [8] and robotics [13].
Revised: Reinforcement learning (RL) has achieved remarkable progress in games [31, 47, 50], ﬁnancial trading [8] and robotics [13].

Original: However, in its core part, without designs tailored to speciﬁc tasks, general RL paradigms are still learning implicit representations from critic loss (value predictions) and actor loss (maximizing cumulative reward).
Revised: However, in its core part, without designs tailored to speciﬁc tasks, general RL paradigms are still learning implicit representations from critic loss (value predictions) and actor loss (maximizing cumulative reward).

Original: In many real-world scenarios where observations are complicated (e.g., images) or incomplete (e g., partial observable), training an agent that is able to extract informative signals from those inputs becomes incredibly sample-inefﬁcient.
Revised: In many real-world scenarios where observations are complicated (e.g., images) or incomplete (e g., partial observable), training an agent that is able to extract informative signals from those inputs becomes incredibly sample-inefﬁcient.

Original: Therefore, many recent works have been devoted to obtaining a good state representation, which is believed to be one of the key solutions to improve the efﬁcacy of RL [23, 24].
Revised: Therefore, many recent works have been devoted to obtaining a good state representation, which is believed to be one of the key solutions to improve the efﬁcacy of RL [23, 24].

Original: One of the main streams is adding auxiliary losses to update the state encoder.
Revised: One of the main streams is adding auxiliary losses to update the state encoder.

Original: Under the hood, it resorts to informative and dense learning signals in order to encode various prior knowledge and regularization [40], and obtain better latent representations.
Revised: Under the hood, it resorts to informative and dense learning signals in order to encode various prior knowledge and regularization [40], and obtain better latent representations.

Original: Over the years, a series of works have attempted to ﬁgure out the form of the most helpful auxiliary loss for RL.
Revised: Over the years, a series of works have attempted to ﬁgure out the form of the most helpful auxiliary loss for RL.

Original: Quite a few advances have been made, including observation reconstruction [51], reward prediction [20], environment dynamics prediction [40, 6, 35], etc.
Revised: Quite a few advances have been made, including observation reconstruction [51], reward prediction [20], environment dynamics prediction [40, 6, 35], etc.

Original: But we note two problems in this evolving process: (i) each of the loss designs listed above are obtained through empirical trial-and-errors based on expert designs, thus heavily relying on human labor and expertise; (ii) few works have used the ﬁnal performance of RL as an optimization objective to directly search the auxiliary loss, indicating that these designs could be sub-optimal.
Revised: But we note two problems in this evolving process: (i) each of the loss designs listed above are obtained through empirical trial-and-errors based on expert designs, thus heavily relying on human labor and expertise; (ii) few works have used the ﬁnal performance of RL as an optimization objective to directly search the auxiliary loss, indicating that these designs could be sub-optimal.

Original: [Figure]
Revised: [Figure]

Original: To resolve the issues of the existing handcrafted solution mentioned above, we decide to automate the process of designing the auxiliary loss functions of RL and propose a principled solution named
Revised: To resolve the issues of the existing handcrafted solution mentioned above, we decide to automate the process of designing the auxiliary loss functions of RL and propose a principled solution named

Original: Automated Auxiliary Loss Search (A2LS).
Revised: Automated Auxiliary Loss Search (A2LS).

Original: A2LS formulates the problem as a bi-level optimization where we try to ﬁnd the best auxiliary loss, which, to the most extent, helps train a good RL agent.
Revised: A2LS formulates the problem as a bi-level optimization where we try to ﬁnd the best auxiliary loss, which, to the most extent, helps train a good RL agent.

Original: The outer loop searches for auxiliary losses based on RL performance to ensure the searched losses align with the RL objective, while the inner loop performs RL training with the searched auxiliary loss function.
Revised: The outer loop searches for auxiliary losses based on RL performance to ensure the searched losses align with the RL objective, while the inner loop performs RL training with the searched auxiliary loss function.

Original: Speciﬁcally, A2LS utilizes an evolutionary strategy to search the conﬁguration of auxiliary losses over a novel search space of size 7 . 5 ⇥ 10 20 that covers many existing solutions.
Revised: Speciﬁcally, A2LS utilizes an evolutionary strategy to search the conﬁguration of auxiliary losses over a novel search space of size 7 . 5 ⇥ 10 20 that covers many existing solutions.

Original: By searching on a small set of simulated training environments of continuous control from Deepmind
Revised: By searching on a small set of simulated training environments of continuous control from Deepmind

Original: Control suite (DMC) [43], A2LS ﬁnalizes a loss, namely A2-winner .
Revised: Control suite (DMC) [43], A2LS ﬁnalizes a loss, namely A2-winner .

Original: To evaluate the generalizability of the discovered auxiliary loss A2-winner , we test A2-winner on a wide set of test environments , including both image-based and vector-based (with proprioceptive features like positions, velocities and accelerations as inputs) tasks.
Revised: To evaluate the generalizability of the discovered auxiliary loss A2-winner , we test A2-winner on a wide set of test environments , including both image-based and vector-based (with proprioceptive features like positions, velocities and accelerations as inputs) tasks.

Original: Extensive experiments show the searched loss function is highly effective and largely outperforms strong baseline methods.
Revised: Extensive experiments show the searched loss function is highly effective and largely outperforms strong baseline methods.

Original: More importantly, the searched auxiliary loss generalizes well to unseen settings such as (i) different robots of control; (ii) different data types of observation; (iii) partially observable settings; (iv) different network architectures; and (v) even to a totally different discrete control domain (Atarigames [1]).
Revised: More importantly, the searched auxiliary loss generalizes well to unseen settings such as (i) different robots of control; (ii) different data types of observation; (iii) partially observable settings; (iv) different network architectures; and (v) even to a totally different discrete control domain (Atarigames [1]).

Original: In the end, we make detailed statistical analyses on the relation between RL performance and patterns of auxiliary losses based on the data of whole evolutionary search process, providing useful insights on future studies of auxiliary loss designs and representation learning in RL.
Revised: In the end, we make detailed statistical analyses on the relation between RL performance and patterns of auxiliary losses based on the data of whole evolutionary search process, providing useful insights on future studies of auxiliary loss designs and representation learning in RL.

Original:  Problem Formulation and Background
Revised:  Problem Formulation and Background

Original: We consider the standard Markov Decision Process (MDP) E where the state, action and reward at time step t are denoted as ( s t , a t , r t ) .
Revised: We consider the standard Markov Decision Process (MDP) E where the state, action and reward at time step t are denoted as ( s t , a t , r t ) .

Original: The sequence of rollout data sampled by the agent in the episodic environment is ( s 0 , . . .
Revised: The sequence of rollout data sampled by the agent in the episodic environment is ( s 0 , . . .

Original: , s t , a t , r t , s t +1 , · · · , s T ) , where T represents the episode length.
Revised: , s t , a t , r t , s t +1 , · · · , s T ) , where T represents the episode length.

Original: Suppose the RL agent is parameterized by !
Revised: Suppose the RL agent is parameterized by !

Original: (either the policy ⇡ or the state-action value function Q ), with a state encoder g ✓ parameterized by ✓ ✓ !
Revised: (either the policy ⇡ or the state-action value function Q ), with a state encoder g ✓ parameterized by ✓ ✓ !

Original: which plays a key role for representation learning in
Revised: which plays a key role for representation learning in

Original: RL.
Revised: RL.

Original: The agent is required to maximize its cumulative rewards in environment E by optimizing ! , noted as R ( ! ; E ) =
Revised: The agent is required to maximize its cumulative rewards in environment E by optimizing ! , noted as R ( ! ; E ) =

Original: E ⇡ [
Revised: E ⇡ [

Original: [Equation]
Revised: [Equation]

Original: In this paper, we aim to ﬁnd the optimal auxiliary loss function L Aux such that the agent can reach the best performance by optimizing !
Revised: In this paper, we aim to ﬁnd the optimal auxiliary loss function L Aux such that the agent can reach the best performance by optimizing !

Original: under a combination of an arbitrary RL loss function L RL together with an auxiliary loss L Aux .
Revised: under a combination of an arbitrary RL loss function L RL together with an auxiliary loss L Aux .

Original: Formally, our optimization goal is:
Revised: Formally, our optimization goal is:

Original: [Equation]
Revised: [Equation]

Original: where  is a hyper-parameter balancing the relative weight of the auxiliary loss.
Revised: where  is a hyper-parameter balancing the relative weight of the auxiliary loss.

Original: The left part (inner loop) of Figure 1 illustrates how data and gradients ﬂow in RL training when an auxiliary loss is enabled.
Revised: The left part (inner loop) of Figure 1 illustrates how data and gradients ﬂow in RL training when an auxiliary loss is enabled.

Original: Some instances of L RL and L Aux are given in Appendix B. Unfortunately, existing auxiliary losses L Aux are handcrafted, which heavily rely on expert knowledge, and may not generalize well
Revised: Some instances of L RL and L Aux are given in Appendix B. Unfortunately, existing auxiliary losses L Aux are handcrafted, which heavily rely on expert knowledge, and may not generalize well

Original: [Table]
Revised: [Table]

Original: in different scenarios as shown in the experiment part.
Revised: in different scenarios as shown in the experiment part.

Original: To ﬁnd better auxiliary loss functions for representation learning in RL, we introduce our principled solution in the following section.
Revised: To ﬁnd better auxiliary loss functions for representation learning in RL, we introduce our principled solution in the following section.

Original:  Automated Auxiliary Loss Search
Revised:  Automated Auxiliary Loss Search

Original: To meet our goal of ﬁnding top-performing auxiliary loss functions without expert assignment, we turn to the help of automated loss search, which has shown promising results in the automated machine learning (AutoML) community [27, 28, 48].
Revised: To meet our goal of ﬁnding top-performing auxiliary loss functions without expert assignment, we turn to the help of automated loss search, which has shown promising results in the automated machine learning (AutoML) community [27, 28, 48].

Original: Correspondingly, we propose Automated
Revised: Correspondingly, we propose Automated

Original: Auxiliary Loss Search (A2LS), a principled solution for resolving the above bi-level optimization problem in Equation 1.
Revised: Auxiliary Loss Search (A2LS), a principled solution for resolving the above bi-level optimization problem in Equation 1.

Original: A2LS resolves the inner problem as a standard RL training procedure; for the outer one, A2LS deﬁnes a ﬁnite and discrete search space (Section 3.1), and designs a novel evolution strategy to efﬁciently explore the space (Section 3.2).
Revised: A2LS resolves the inner problem as a standard RL training procedure; for the outer one, A2LS deﬁnes a ﬁnite and discrete search space (Section 3.1), and designs a novel evolution strategy to efﬁciently explore the space (Section 3.2).

Original: 3.1 Search Space Design
Revised: 3.1 Search Space Design

Original: We have argued that almost all existing auxiliary losses require expert knowledge, and we expect to search for a better one automatically.
Revised: We have argued that almost all existing auxiliary losses require expert knowledge, and we expect to search for a better one automatically.

Original: To this end, it is clear that we should design a search space that satisﬁes the following desiderata.
Revised: To this end, it is clear that we should design a search space that satisﬁes the following desiderata.

Original: • Generalization : the search space should cover most of the existing handcrafted auxiliary losses to ensure the searched results can be no worse than handcrafted losses; • Atomicity : the search space should be composed of several independent dimensions to ﬁt into any general search algorithm [30] and support an efﬁcient search scheme; • Sufﬁciency : the search space should be large enough to contain the top-performing solutions.
Revised: • Generalization : the search space should cover most of the existing handcrafted auxiliary losses to ensure the searched results can be no worse than handcrafted losses; • Atomicity : the search space should be composed of several independent dimensions to ﬁt into any general search algorithm [30] and support an efﬁcient search scheme; • Sufﬁciency : the search space should be large enough to contain the top-performing solutions.

Original: Given the criteria, we conclude and list some existing auxiliary losses in Table 1 and ﬁnd their commonalities, as well as differences.
Revised: Given the criteria, we conclude and list some existing auxiliary losses in Table 1 and ﬁnd their commonalities, as well as differences.

Original: We realize that these losses share similar components and computation ﬂow.
Revised: We realize that these losses share similar components and computation ﬂow.

Original: As shown in Figure 2, when training the RL agent, the loss ﬁrstly selects a sequence { s t , a t , r t } i + k t
Revised: As shown in Figure 2, when training the RL agent, the loss ﬁrstly selects a sequence { s t , a t , r t } i + k t

Original: = i from the replay buffer, when k is called horizon .
Revised: = i from the replay buffer, when k is called horizon .

Original: The agent then tries to predict some elements in the sequence (called target ) based on another picked set of elements from the sequence (called source ).
Revised: The agent then tries to predict some elements in the sequence (called target ) based on another picked set of elements from the sequence (called source ).

Original: Finally, the loss calculates and minimizes the prediction error (rigorously deﬁned with operator ).
Revised: Finally, the loss calculates and minimizes the prediction error (rigorously deﬁned with operator ).

Original: To be more speciﬁc, the encoder part g ✓ of the agent, ﬁrst encodes the source into latent representations, which is further fed into a predictor h to get a prediction y ; the auxiliary loss is computed by the prediction y and the target
Revised: To be more speciﬁc, the encoder part g ✓ of the agent, ﬁrst encodes the source into latent representations, which is further fed into a predictor h to get a prediction y ; the auxiliary loss is computed by the prediction y and the target

Original: ˆ y that is translated from the target by a target encoder g ˆ ✓ , using an operator f .
Revised: ˆ y that is translated from the target by a target encoder g ˆ ✓ , using an operator f .

Original: The target encoder is updated in an momentum manner as shown in Figure(details are given in Appendix C.1.2).
Revised: The target encoder is updated in an momentum manner as shown in Figure(details are given in Appendix C.1.2).

Original: Formally,
Revised: Formally,

Original: [Equation]
Revised: [Equation]

Original: where seq source , seq target ✓ { s t , a t , r t } i + k t
Revised: where seq source , seq target ✓ { s t , a t , r t } i + k t

Original: = i are both subsets of the candidate sequence.
Revised: = i are both subsets of the candidate sequence.

Original: And for simplicity, we will denote g ✓ ( s t , a t , r t , s t +1 , · · · ) as short for [ g ✓ ( s t ) , a t , r t , g ✓ ( s t +1 ) , · · · ] for the rest of this paper (the encoder g only deals with states { s i } ).
Revised: And for simplicity, we will denote g ✓ ( s t , a t , r t , s t +1 , · · · ) as short for [ g ✓ ( s t ) , a t , r t , g ✓ ( s t +1 ) , · · · ] for the rest of this paper (the encoder g only deals with states { s i } ).

Original: Thereafter, we observe that these existing auxiliary losses differ in two dimensions, i.e., input elements and operator , where input elements are further combined by horizon , source and target .
Revised: Thereafter, we observe that these existing auxiliary losses differ in two dimensions, i.e., input elements and operator , where input elements are further combined by horizon , source and target .

Original: These differences compose our search dimensions of the whole space.
Revised: These differences compose our search dimensions of the whole space.

Original: We then illustrate the search ranges of these dimensions in detail.
Revised: We then illustrate the search ranges of these dimensions in detail.

Original: Input elements.
Revised: Input elements.

Original: The input elements denote all inputs to the loss functions, which can be further
Revised: The input elements denote all inputs to the loss functions, which can be further

Original: [Figure]
Revised: [Figure]

Original: disassembled as horizon , source and target .
Revised: disassembled as horizon , source and target .

Original: Different from previous automated loss search works, the target here is not “ground-truth” because auxiliary losses in RL have no labels beforehand.
Revised: Different from previous automated loss search works, the target here is not “ground-truth” because auxiliary losses in RL have no labels beforehand.

Original: Instead, both source and target are generated via interacting with the environment in a self-supervised manner.
Revised: Instead, both source and target are generated via interacting with the environment in a self-supervised manner.

Original: Particularly, the input elements ﬁrst determine a candidate sequence { s t , a t , r t } i + k t
Revised: Particularly, the input elements ﬁrst determine a candidate sequence { s t , a t , r t } i + k t

Original: = i with horizon k .
Revised: = i with horizon k .

Original: Then, it chooses two subsets from the candidate sequence as source and target respectively.
Revised: Then, it chooses two subsets from the candidate sequence as source and target respectively.

Original: For example, the subsets can be { s t } , { s t , s t +1 } , or { s t , r t +1 , a t +2 } , { s t , s t +1 , a t +1 } , etc.
Revised: For example, the subsets can be { s t } , { s t , s t +1 } , or { s t , r t +1 , a t +2 } , { s t , s t +1 , a t +1 } , etc.

Original: Operator.
Revised: Operator.

Original: Given a prediction y and its target ˆ y , the auxiliary loss is computed by an operator f , which is often a similarity measure.
Revised: Given a prediction y and its target ˆ y , the auxiliary loss is computed by an operator f , which is often a similarity measure.

Original: In our work, we cover all different operators f used by the previous works, including inner product (Inner) [17, 42], bilinear inner product (Bilinear)
Revised: In our work, we cover all different operators f used by the previous works, including inner product (Inner) [17, 42], bilinear inner product (Bilinear)

Original: [23], cosine similarity (Cosine) [3], mean squared error (MSE) [35, 6] and normalized mean squared error (N-MSE) [39].
Revised: [23], cosine similarity (Cosine) [3], mean squared error (MSE) [35, 6] and normalized mean squared error (N-MSE) [39].

Original: Additionally, other works also utilize contrastive objectives, e.g., InfoNCE loss [33], incorporating the trick to sample un-paired predictions and targets as negative samples and maximize the distances between them.
Revised: Additionally, other works also utilize contrastive objectives, e.g., InfoNCE loss [33], incorporating the trick to sample un-paired predictions and targets as negative samples and maximize the distances between them.

Original: This technique is orthogonal to the ﬁve similarity measures mentioned above, so we make it optional and create 5 ⇥ 2 = 10 different operators in total.
Revised: This technique is orthogonal to the ﬁve similarity measures mentioned above, so we make it optional and create 5 ⇥ 2 = 10 different operators in total.

Original: Final design.
Revised: Final design.

Original: In the light of preceding discussion, with the deﬁnition of input elements and operator , we ﬁnish the design of the search space, which satisfactorily meets the desiderata mentioned above.
Revised: In the light of preceding discussion, with the deﬁnition of input elements and operator , we ﬁnish the design of the search space, which satisfactorily meets the desiderata mentioned above.

Original: Speciﬁcally, the space is generalizable to cover most of the existing handcrafted auxiliary losses; additionally, the atomicity is embodied by the compositionality that all input elements work with any operator ; most importantly, the search space is sufﬁciently large with a total size of 7 . 5 ⇥ 10(detailed calculation can be found in Appendix E) to ﬁnd better solutions.
Revised: Speciﬁcally, the space is generalizable to cover most of the existing handcrafted auxiliary losses; additionally, the atomicity is embodied by the compositionality that all input elements work with any operator ; most importantly, the search space is sufﬁciently large with a total size of 7 . 5 ⇥ 10(detailed calculation can be found in Appendix E) to ﬁnd better solutions.

Original: 3.2 Search Strategy
Revised: 3.2 Search Strategy

Original: The success of evolution strategies in exploring large, multi-dimensional search space has been proven in many works [19, 4].
Revised: The success of evolution strategies in exploring large, multi-dimensional search space has been proven in many works [19, 4].

Original: Similarly, A2LS adopts an evolutionary algorithm [37] to search for topperforming auxiliary loss functions over the designed search space.
Revised: Similarly, A2LS adopts an evolutionary algorithm [37] to search for topperforming auxiliary loss functions over the designed search space.

Original: In its essence, the evolutionary algorithm (i) keeps a population of loss function candidates; (ii) evaluates their performance; (iii) eliminates the worst and evolves into a new better population.
Revised: In its essence, the evolutionary algorithm (i) keeps a population of loss function candidates; (ii) evaluates their performance; (iii) eliminates the worst and evolves into a new better population.

Original: Note that step (ii) of “evaluating” is very costly because it needs to train the RL agents with dozens of different auxiliary loss functions.
Revised: Note that step (ii) of “evaluating” is very costly because it needs to train the RL agents with dozens of different auxiliary loss functions.

Original: Therefore, our key technical contribution contains how to further reduce the search cost (Section 3.2.1) and how to make an efﬁcient search procedure (Section 3.2.2).
Revised: Therefore, our key technical contribution contains how to further reduce the search cost (Section 3.2.1) and how to make an efﬁcient search procedure (Section 3.2.2).

Original: 3.2.1 Search Space Pruning
Revised: 3.2.1 Search Space Pruning

Original: In our preliminary experiment, we ﬁnd out the dimension of operator in the search space can be simpliﬁed.
Revised: In our preliminary experiment, we ﬁnd out the dimension of operator in the search space can be simpliﬁed.

Original: In particular, MSE outperforms all the others by signiﬁcant gaps in most cases.
Revised: In particular, MSE outperforms all the others by signiﬁcant gaps in most cases.

Original: So we effectively prune other choices of operators except MSE.
Revised: So we effectively prune other choices of operators except MSE.

Original: See Appendix D.1 for complete comparative results and an ablation study on the effectiveness of search space pruning.
Revised: See Appendix D.1 for complete comparative results and an ablation study on the effectiveness of search space pruning.

Original: [Figure]
Revised: [Figure]

Original: 3.2.2 Evolution Procedure
Revised: 3.2.2 Evolution Procedure

Original: Our evolution procedure roughly contains four important components: (i) evaluation and selection : a population of candidate auxiliary losses is evaluated through an inner loop of RL training, then we select the top candidates for the next evolution stage (i.e., generation); (ii) mutation : the selected candidates mutate to form a new population and move to the next stage; (iii) loss rejection : ﬁlter out and skip evaluating invalid auxiliary losses for the next stage; and (iv) bootstrapping initial population : assign more chance to initial auxiliary losses that may contain useful patterns by prior knowledge for higher efﬁciency.
Revised: Our evolution procedure roughly contains four important components: (i) evaluation and selection : a population of candidate auxiliary losses is evaluated through an inner loop of RL training, then we select the top candidates for the next evolution stage (i.e., generation); (ii) mutation : the selected candidates mutate to form a new population and move to the next stage; (iii) loss rejection : ﬁlter out and skip evaluating invalid auxiliary losses for the next stage; and (iv) bootstrapping initial population : assign more chance to initial auxiliary losses that may contain useful patterns by prior knowledge for higher efﬁciency.

Original: The step-by-step evolution algorithm is provided in Algorithm 1 in the appendix, and an overview of the A2LS pipeline is illustrated in Figure 1.
Revised: The step-by-step evolution algorithm is provided in Algorithm 1 in the appendix, and an overview of the A2LS pipeline is illustrated in Figure 1.

Original: We next describe them in detail.
Revised: We next describe them in detail.

Original: Evaluation and selection.
Revised: Evaluation and selection.

Original: At each evolution stage, we ﬁrst train a population of candidates with a population size P = 100 by the inner loop of RL training.
Revised: At each evolution stage, we ﬁrst train a population of candidates with a population size P = 100 by the inner loop of RL training.

Original: The candidates are then sorted by computing the approximated area under learning curve (AULC)
Revised: The candidates are then sorted by computing the approximated area under learning curve (AULC)

Original: [11, 41], which is a single metric reﬂecting both the convergence speed and the ﬁnal performance [46] with low variance of results.
Revised: [11, 41], which is a single metric reﬂecting both the convergence speed and the ﬁnal performance [46] with low variance of results.

Original: After each training stage, the top-25% candidates are selected to generate the population for the next stage.
Revised: After each training stage, the top-25% candidates are selected to generate the population for the next stage.

Original: We include an ablation study on the effectiveness of AULC in Appendix D.3.
Revised: We include an ablation study on the effectiveness of AULC in Appendix D.3.

Original: Mutation.
Revised: Mutation.

Original: To obtain a new population of auxiliary loss functions, we propose a novel mutation strategy.
Revised: To obtain a new population of auxiliary loss functions, we propose a novel mutation strategy.

Original: First, we represent both the source and the target of the input elements as a pair of binary masks, where each bit of the mask represents selected by 1 or not selected by 0.
Revised: First, we represent both the source and the target of the input elements as a pair of binary masks, where each bit of the mask represents selected by 1 or not selected by 0.

Original: For instance, given a candidate sequence { s t , a t , r t , s t +1 , a t +1 , r t +1 } , the binary mask of this subset sequence { s t , a t , r t +1 } is denoted as 110001 .
Revised: For instance, given a candidate sequence { s t , a t , r t , s t +1 , a t +1 , r t +1 } , the binary mask of this subset sequence { s t , a t , r t +1 } is denoted as 110001 .

Original: Afterward, we adopt four types of mutations, also shown in
Revised: Afterward, we adopt four types of mutations, also shown in

Original: Figure 3: (i) replacement (50% of the population): ﬂip the given binary mask with probability p = 1 2 · (3 k +3) with the horizon length k ; (ii) crossover (20%): generate a new candidate by randomly combining the mask bits of two candidates with the same horizon length in the population; (iii) horizon decrease and horizon increase (10%): append new binary masks to the tail or delete existing binary masks at the back.
Revised: Figure 3: (i) replacement (50% of the population): ﬂip the given binary mask with probability p = 1 2 · (3 k +3) with the horizon length k ; (ii) crossover (20%): generate a new candidate by randomly combining the mask bits of two candidates with the same horizon length in the population; (iii) horizon decrease and horizon increase (10%): append new binary masks to the tail or delete existing binary masks at the back.

Original: (iv) random generation (20%): every bit of the binary mask is generated from a Bernoulli distribution B (0 . 5) .
Revised: (iv) random generation (20%): every bit of the binary mask is generated from a Bernoulli distribution B (0 . 5) .

Original: Loss rejection protocol.
Revised: Loss rejection protocol.

Original: Since the auxiliary loss needs to be differentiable with respect to the parameters of the state encoder, we perform a gradient ﬂow check on randomly generated loss functions during evolution and skip evaluating invalid auxiliary losses.
Revised: Since the auxiliary loss needs to be differentiable with respect to the parameters of the state encoder, we perform a gradient ﬂow check on randomly generated loss functions during evolution and skip evaluating invalid auxiliary losses.

Original: Concretely, the following conditions must be satisﬁed to make a valid loss function: (i) having at least one state element in seq source to make sure the gradient of auxiliary loss can propagate back to the state encoder; (ii) seq target is not empty; (iii) the horizon should be within a reasonable range ( 1  k  10 in our experiments).
Revised: Concretely, the following conditions must be satisﬁed to make a valid loss function: (i) having at least one state element in seq source to make sure the gradient of auxiliary loss can propagate back to the state encoder; (ii) seq target is not empty; (iii) the horizon should be within a reasonable range ( 1  k  10 in our experiments).

Original: If a loss is rejected, we repeat the mutation to ﬁll the population.
Revised: If a loss is rejected, we repeat the mutation to ﬁll the population.

Original: Bootstrapping initial population.
Revised: Bootstrapping initial population.

Original: To improve the computational efﬁciency so that the algorithm can ﬁnd reasonable loss functions quickly, we incorporate prior knowledge into the initialization of the search.
Revised: To improve the computational efﬁciency so that the algorithm can ﬁnd reasonable loss functions quickly, we incorporate prior knowledge into the initialization of the search.

Original: Particularly, before the ﬁrst stage of evolution, we bootstrap the initial population with a prior distribution that assigns high probability to auxiliary loss functions containing useful patterns like dynamics and reward prediction.
Revised: Particularly, before the ﬁrst stage of evolution, we bootstrap the initial population with a prior distribution that assigns high probability to auxiliary loss functions containing useful patterns like dynamics and reward prediction.

Original: More implementation details are provided in Appendix C.3.
Revised: More implementation details are provided in Appendix C.3.

Original: [Figure]
Revised: [Figure]

Original: As mentioned in Section 1, we expect to ﬁnd auxiliary losses that align with the RL objective and generalize well to unseen test environments .
Revised: As mentioned in Section 1, we expect to ﬁnd auxiliary losses that align with the RL objective and generalize well to unseen test environments .

Original: To do so, we use A2LS to search over a small set of training environments , and then test the searched results on a wide range of test environments .
Revised: To do so, we use A2LS to search over a small set of training environments , and then test the searched results on a wide range of test environments .

Original: In this section, we ﬁrst introduce the evolution on training environments and search results.
Revised: In this section, we ﬁrst introduce the evolution on training environments and search results.

Original: 4.1 Evolution on Training Environments
Revised: 4.1 Evolution on Training Environments

Original: The training environments are chosen as three image-based (observations for agents are images) continuous control tasks in DMC benchmark [43], Cheetah-Run, Reacher-Easy, and Walker-Walk.
Revised: The training environments are chosen as three image-based (observations for agents are images) continuous control tasks in DMC benchmark [43], Cheetah-Run, Reacher-Easy, and Walker-Walk.

Original: For each environment, we set the total budget to 16k GPU hours (on NVIDIA P100) and terminate the search when the resource is exhausted.
Revised: For each environment, we set the total budget to 16k GPU hours (on NVIDIA P100) and terminate the search when the resource is exhausted.

Original: Due to computation complexity, we only run one seed for each inner loop RL training, but we try to prevent such randomness by cross validation (see
Revised: Due to computation complexity, we only run one seed for each inner loop RL training, but we try to prevent such randomness by cross validation (see

Original: Section 4.2).
Revised: Section 4.2).

Original: We use the same network architecture and hyperparameters conﬁg as CURL [23] (see
Revised: We use the same network architecture and hyperparameters conﬁg as CURL [23] (see

Original: Appendix C.4.1 for details) to train the RL agents.
Revised: Appendix C.4.1 for details) to train the RL agents.

Original: To evaluate the population during evolution, we measure A2LS as compared to SAC, SAC-wo-aug, and CURL, where we randomly crop images from 100 ⇥ 100 to 84 ⇥ 84 as data augmentation (the same technique used in CURL[23]) for all methods except SAC-wo-aug.
Revised: To evaluate the population during evolution, we measure A2LS as compared to SAC, SAC-wo-aug, and CURL, where we randomly crop images from 100 ⇥ 100 to 84 ⇥ 84 as data augmentation (the same technique used in CURL[23]) for all methods except SAC-wo-aug.

Original: The whole evolution process on three environments is demonstrated in Figure 4.
Revised: The whole evolution process on three environments is demonstrated in Figure 4.

Original: Even in the early stages (e.g., stage 1), some of the auxiliary loss candidates already surpass baselines, indicating the high potential of automated loss search.
Revised: Even in the early stages (e.g., stage 1), some of the auxiliary loss candidates already surpass baselines, indicating the high potential of automated loss search.

Original: The overall AULC scores of the population continue to improve when more stages come in (detailed numbers are summerized in
Revised: The overall AULC scores of the population continue to improve when more stages come in (detailed numbers are summerized in

Original: Appendix D.10).
Revised: Appendix D.10).

Original: Judging from the trend, we believe the performances could improve even more if we had further increased the budget.
Revised: Judging from the trend, we believe the performances could improve even more if we had further increased the budget.

Original: 4.2 Searched Results: A2-winner
Revised: 4.2 Searched Results: A2-winner

Original: [Figure]
Revised: [Figure]

Original: N o r m a li z e d AU
Revised: N o r m a li z e d AU

Original: L C S c o r e ( w . r .
Revised: L C S c o r e ( w . r .

Original: t l o w es t sc
Revised: t l o w es t sc

Original: o r e )
Revised: o r e )

Original: Although some candidates in the population have achieved remarkable AULC scores in the evolution (Figure 4), they were only evaluated with one random seed in one environment, making their robustness under question.
Revised: Although some candidates in the population have achieved remarkable AULC scores in the evolution (Figure 4), they were only evaluated with one random seed in one environment, making their robustness under question.

Original: To ensure that we ﬁnd a consistently-useful auxiliary loss, we conduct a cross validation.
Revised: To ensure that we ﬁnd a consistently-useful auxiliary loss, we conduct a cross validation.

Original: We ﬁrst choose the top 5 candidates of stage-5 of the evolution on Cheetah-Run (detailed top candidates during the whole evolution procedure are provided in Appendix F).
Revised: We ﬁrst choose the top 5 candidates of stage-5 of the evolution on Cheetah-Run (detailed top candidates during the whole evolution procedure are provided in Appendix F).

Original: For each of the ﬁve candidates, we repeat the RL training on all three training environments , shown in Figure 5.
Revised: For each of the ﬁve candidates, we repeat the RL training on all three training environments , shown in Figure 5.

Original: Finally, we mark the best among ﬁve (green bar in Figure 5) as our ﬁnal searched result.
Revised: Finally, we mark the best among ﬁve (green bar in Figure 5) as our ﬁnal searched result.

Original: We call it A2-winner , which has the following form:
Revised: We call it A2-winner , which has the following form:

Original: [Equation]
Revised: [Equation]

Original: and time steps).
Revised: and time steps).

Original: Note that the optimal of DMC is 1000 for all
Revised: Note that the optimal of DMC is 1000 for all

Original: The baseline methods are PlaNet [16], Dreamer [15], SAC+AE [51], SLAC [26], image-based
Revised: The baseline methods are PlaNet [16], Dreamer [15], SAC+AE [51], SLAC [26], image-based

Original: SAC [14].
Revised: SAC [14].

Original: Performance values of all baselines are referred to [23], except for Image SAC.
Revised: Performance values of all baselines are referred to [23], except for Image SAC.

Original: Learning curves of all 12 DMC environments are included in Appendix D.2.
Revised: Learning curves of all 12 DMC environments are included in Appendix D.2.

Original: [Table]
Revised: [Table]

Original:  Generalization Experiments
Revised:  Generalization Experiments

Original: To verify the effectiveness of the searched results
Revised: To verify the effectiveness of the searched results

Original: [Table]
Revised: [Table]

Original: we conduct various generalization experiments on a wide range of test environments in depth.
Revised: we conduct various generalization experiments on a wide range of test environments in depth.

Original: Implementation details and more ablation studies are given in Appendix C and Appendix D.
Revised: Implementation details and more ablation studies are given in Appendix C and Appendix D.

Original: Generalize to unseen image-based tasks.
Revised: Generalize to unseen image-based tasks.

Original: We ﬁrst investigate the generalizability of A2-winner to unseen image-based tasks by training agents with A2-winner on common DMC tasks and compare with model-based and model-free baselines that use different auxiliary loss functions (see Appendix C.5 for details about baseline methods).
Revised: We ﬁrst investigate the generalizability of A2-winner to unseen image-based tasks by training agents with A2-winner on common DMC tasks and compare with model-based and model-free baselines that use different auxiliary loss functions (see Appendix C.5 for details about baseline methods).

Original: The results are summarized in Tablewhere A2-winner greatly outperforms other baseline methods on most tasks, including unseen test environments .
Revised: The results are summarized in Tablewhere A2-winner greatly outperforms other baseline methods on most tasks, including unseen test environments .

Original: This implies that A2-winner is a robust and effective auxiliary loss for image-based continuous control tasks to improve both the efﬁciency and ﬁnal performance.
Revised: This implies that A2-winner is a robust and effective auxiliary loss for image-based continuous control tasks to improve both the efﬁciency and ﬁnal performance.

Original: [Table]
Revised: [Table]

Original: Generalize to totally different benchmark domains.
Revised: Generalize to totally different benchmark domains.

Original: To further verify the generalizability of A2-winner on totally different benchmark domains other than DMC tasks, we conduct experiments on the Atari 2600 Games [1], where we take
Revised: To further verify the generalizability of A2-winner on totally different benchmark domains other than DMC tasks, we conduct experiments on the Atari 2600 Games [1], where we take

Original: Efﬁcient Rainbow [44] as the base RL algorithm and add A2-winner to obtain a better state representation.
Revised: Efﬁcient Rainbow [44] as the base RL algorithm and add A2-winner to obtain a better state representation.

Original: Results are shown in Table 3 where A2-winner outperforms all baselines, showing strong evidence of the generalization and potential usages of A2-winner .
Revised: Results are shown in Table 3 where A2-winner outperforms all baselines, showing strong evidence of the generalization and potential usages of A2-winner .

Original: Note that the base RL algorithm used in Atari is a value-based method, indicating that A2-winner generalizes well to both value-based and policy-based RL algorithms.
Revised: Note that the base RL algorithm used in Atari is a value-based method, indicating that A2-winner generalizes well to both value-based and policy-based RL algorithms.

Original: Generalize to different observation types.
Revised: Generalize to different observation types.

Original: To see whether A2-winner (searched in image-based environments) is able to generalize to the environments with different observation types, we test
Revised: To see whether A2-winner (searched in image-based environments) is able to generalize to the environments with different observation types, we test

Original: A2-winner on vector-based (inputs for RL agents are proprioceptive features such as positions, velocities and accelerations) tasks of DMC and list the results in Table 4.
Revised: A2-winner on vector-based (inputs for RL agents are proprioceptive features such as positions, velocities and accelerations) tasks of DMC and list the results in Table 4.

Original: Concretely, we compare
Revised: Concretely, we compare

Original: A2-winner with SAC-Identity, SAC and CURL, where SAC-Identity does not have state encoder while the others share the same state encoder architecture (See Appendix C.1.1 and Appendix D.6 for detailed implementations).
Revised: A2-winner with SAC-Identity, SAC and CURL, where SAC-Identity does not have state encoder while the others share the same state encoder architecture (See Appendix C.1.1 and Appendix D.6 for detailed implementations).

Original: To our delight, A2-winner still outperforms all baselines in 12 out ofenvironments, showing A2-winner can also beneﬁt RL performance in vector-based observations.
Revised: To our delight, A2-winner still outperforms all baselines in 12 out ofenvironments, showing A2-winner can also beneﬁt RL performance in vector-based observations.

Original: Moreover, the performance gain is particularly signiﬁcant in more complex environments like
Revised: Moreover, the performance gain is particularly signiﬁcant in more complex environments like

Original: Humanoid, where SAC barely learns anything at 1000K time steps.
Revised: Humanoid, where SAC barely learns anything at 1000K time steps.

Original: In order to get a deeper understanding of this phenomenon, we additionally visualize the Q loss landscape for both methods in Appendix D.7.
Revised: In order to get a deeper understanding of this phenomenon, we additionally visualize the Q loss landscape for both methods in Appendix D.7.

Original: [Table]
Revised: [Table]

Original: ± 391 363 ±
Revised: ± 391 363 ±

Original: [Figure]
Revised: [Figure]

Original: [Figure]
Revised: [Figure]

Original: 286 ±9 ±7 ±7 ±
Revised: 286 ±9 ±7 ±7 ±

Original: [Figure]
Revised: [Figure]

Original: 311 ±299 ±16 ±2 ±2 ±
Revised: 311 ±299 ±16 ±2 ±2 ±

Original: [Figure]
Revised: [Figure]

Original: Run
Revised: Run

Original: [Figure]
Revised: [Figure]

Original: 75 ±88 ±1 ±1 ±1 ±
Revised: 75 ±88 ±1 ±1 ±1 ±

Original: [Figure]
Revised: [Figure]

Original: Training
Revised: Training

Original: [Figure]
Revised: [Figure]

Original: Unseen
Revised: Unseen

Original: [Figure]
Revised: [Figure]

Original: Generalize to different hypothesis spaces.
Revised: Generalize to different hypothesis spaces.

Original: The architecture of a neural network deﬁnes a hypothesis space of functions to be optimized.
Revised: The architecture of a neural network deﬁnes a hypothesis space of functions to be optimized.

Original: During the evolutionary search in Section 4.1, the encoder architecture has been kept static as a 4-layer convolutional neural network.
Revised: During the evolutionary search in Section 4.1, the encoder architecture has been kept static as a 4-layer convolutional neural network.

Original: Since encoder architecture may have a large impact on the RL training process [34, 2], we test A2-winner with three encoders with different depth of neural networks.
Revised: Since encoder architecture may have a large impact on the RL training process [34, 2], we test A2-winner with three encoders with different depth of neural networks.

Original: The result is shown in
Revised: The result is shown in

Original: Figure 6.
Revised: Figure 6.

Original: Note that even though the auxiliary loss is searched with a 4-layer encoder, the 6-layer convolutional encoder is able to perform better in both two environments.
Revised: Note that even though the auxiliary loss is searched with a 4-layer encoder, the 6-layer convolutional encoder is able to perform better in both two environments.

Original: This proves that the auxiliary loss function of A2-winner is able to improve
Revised: This proves that the auxiliary loss function of A2-winner is able to improve

Original: RL performance with a deeper and more expressive image encoder.
Revised: RL performance with a deeper and more expressive image encoder.

Original: Moreover, the ranking of RL performance (6-layer > 4-layer > 2-layer) is consistent across the two environments.
Revised: Moreover, the ranking of RL performance (6-layer > 4-layer > 2-layer) is consistent across the two environments.

Original: This shows that the auxiliary loss function of A2-winner does not overﬁt one speciﬁc architecture of the encoder.
Revised: This shows that the auxiliary loss function of A2-winner does not overﬁt one speciﬁc architecture of the encoder.

Original: [Figure]
Revised: [Figure]

Original: Generalize to partially observable scenarios.
Revised: Generalize to partially observable scenarios.

Original: Claiming the generality of a method based on conclusions drawn just on fully observable environments like DMC is very dangerous.
Revised: Claiming the generality of a method based on conclusions drawn just on fully observable environments like DMC is very dangerous.

Original: Therefore, we conduct an ablation study on the Partially Observable Markov
Revised: Therefore, we conduct an ablation study on the Partially Observable Markov

Original: Decision Process (POMDP) setting to see whether
Revised: Decision Process (POMDP) setting to see whether

Original: A2-winner is able to perform well in POMDP.
Revised: A2-winner is able to perform well in POMDP.

Original: We random mask 20% of the state dimensions (e.g.,dimensions -> 12 dimensions) to form a POMDP environment in DMC.
Revised: We random mask 20% of the state dimensions (e.g.,dimensions -> 12 dimensions) to form a POMDP environment in DMC.

Original: As demonstrated in Figure 7,
Revised: As demonstrated in Figure 7,

Original: A2-winner consistently outperforms CURL and SAC-DenseMLP in the POMDP setting in HopperHop and Cheetah-Run, showing that A2-winner is not only effective in fully observable environments but also partially observable environments.
Revised: A2-winner consistently outperforms CURL and SAC-DenseMLP in the POMDP setting in HopperHop and Cheetah-Run, showing that A2-winner is not only effective in fully observable environments but also partially observable environments.

Original: To search or not?
Revised: To search or not?

Original: As shown above, the searched result A2-winner can generalize well to all kinds of different settings.
Revised: As shown above, the searched result A2-winner can generalize well to all kinds of different settings.

Original: A natural question here is, however, for a new type of domain, why not perform a new evolution search, instead of simply using the previously searched result?
Revised: A natural question here is, however, for a new type of domain, why not perform a new evolution search, instead of simply using the previously searched result?

Original: To compare these two solutions, we conduct another evolutionary search similar to Section 4.1 but replaced the three image-based tasks with three vector-based ones (marked by † in Table 4) from scratch.
Revised: To compare these two solutions, we conduct another evolutionary search similar to Section 4.1 but replaced the three image-based tasks with three vector-based ones (marked by † in Table 4) from scratch.

Original: More details are summarized in Appendix D.5.
Revised: More details are summarized in Appendix D.5.

Original: We name the searched result as “ A2-winner-v ”.
Revised: We name the searched result as “ A2-winner-v ”.

Original: As shown in
Revised: As shown in

Original: Table 4, A2-winner-v is a very strong-performing loss for vector-based tasks, even stronger than
Revised: Table 4, A2-winner-v is a very strong-performing loss for vector-based tasks, even stronger than

Original: A2-winner .
Revised: A2-winner .

Original: Actually, A2-winner-v is able to outperform baselines in 16 out of 18 environments (with 15 unseen test environments ), while A2-winner only outperforms baselines in 12 out ofenvironments.
Revised: Actually, A2-winner-v is able to outperform baselines in 16 out of 18 environments (with 15 unseen test environments ), while A2-winner only outperforms baselines in 12 out ofenvironments.

Original: However, please note that it costs another 5k GPU hours (on NVIDIA P100) to
Revised: However, please note that it costs another 5k GPU hours (on NVIDIA P100) to

Original: [Table]
Revised: [Table]

Original: search for A2-winner-v while there is no additional cost to directly use A2-winner .
Revised: search for A2-winner-v while there is no additional cost to directly use A2-winner .

Original: It is a trade-off between lower computational cost and better performance.
Revised: It is a trade-off between lower computational cost and better performance.

Original:  Analysis of Auxiliary Loss Functions
Revised:  Analysis of Auxiliary Loss Functions

Original: In this section, we analyze all the loss functions we have evaluated during the evolution procedure as a whole dataset in order to gain some insights into the role of auxiliary loss in RL performance.
Revised: In this section, we analyze all the loss functions we have evaluated during the evolution procedure as a whole dataset in order to gain some insights into the role of auxiliary loss in RL performance.

Original: By doing so, we hope to shed light on future auxiliary loss designs.
Revised: By doing so, we hope to shed light on future auxiliary loss designs.

Original: We will also release this “dataset” publicly to facilitate future research.
Revised: We will also release this “dataset” publicly to facilitate future research.

Original: Typical patterns.
Revised: Typical patterns.

Original: We say that 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 .
Revised: We say that 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 .

Original: For instance, a loss candidate of { s t , a t } !
Revised: For instance, a loss candidate of { s t , a t } !

Original: { s t +1 , s t +2 } has the pattern { s t , a t } !
Revised: { s t +1 , s t +2 } has the pattern { s t , a t } !

Original: { s t +1 } , and does not have the pattern { a t , s t +1 } ! {
Revised: { s t +1 } , and does not have the pattern { a t , s t +1 } ! {

Original: s t } .
Revised: s t } .

Original: We then try to analyze whether a certain pattern is helpful to representation learning in RL in expectation.
Revised: We then try to analyze whether a certain pattern is helpful to representation learning in RL in expectation.

Original: Speciﬁcally, we analyze the following patterns: (i) forward dynamics { s t , a t } !
Revised: Speciﬁcally, we analyze the following patterns: (i) forward dynamics { s t , a t } !

Original: { s t +1 } ; (ii) inverse dynamics { a t , s t +1 } !
Revised: { s t +1 } ; (ii) inverse dynamics { a t , s t +1 } !

Original: { s t } ; (iii) reward prediction { s t , a t } !
Revised: { s t } ; (iii) reward prediction { s t , a t } !

Original: { r t } ; (iv) action inference { s t , s t +1 } !
Revised: { r t } ; (iv) action inference { s t , s t +1 } !

Original: { a t } and (v) state reconstruction in the latent space { s t } ! {
Revised: { a t } and (v) state reconstruction in the latent space { s t } ! {

Original: s t } .
Revised: s t } .

Original: For each of these patterns, we categorize all the loss functions we have evaluated into (i) with or (ii) without this pattern.
Revised: For each of these patterns, we categorize all the loss functions we have evaluated into (i) with or (ii) without this pattern.

Original: We then calculate the average RL performances of these two categories, summarized in Table 5.
Revised: We then calculate the average RL performances of these two categories, summarized in Table 5.

Original: Some interesting observations are as follows.
Revised: Some interesting observations are as follows.

Original: (i) Forward dynamics is helpful in most tasks and improves RL performance on Reacher-Easy (image) and Cheetah-Run (vector) signiﬁcantly (p-value < 0.05).
Revised: (i) Forward dynamics is helpful in most tasks and improves RL performance on Reacher-Easy (image) and Cheetah-Run (vector) signiﬁcantly (p-value < 0.05).

Original: (ii) State reconstruction in the latent space improves RL performance in image-based tasks but undermines vector-based tasks.
Revised: (ii) State reconstruction in the latent space improves RL performance in image-based tasks but undermines vector-based tasks.

Original: The improvements in image-based tasks could be attributed to the combination of augmentation techniques, which, combined with reconstruction loss, enforces the extraction of meaningful features.
Revised: The improvements in image-based tasks could be attributed to the combination of augmentation techniques, which, combined with reconstruction loss, enforces the extraction of meaningful features.

Original: In contrast, no augmentation is used in the vector-based setting, and thus the encoder learns no useful representations.
Revised: In contrast, no augmentation is used in the vector-based setting, and thus the encoder learns no useful representations.

Original: This also explains why CURL performs poorly in vector-based experiments.
Revised: This also explains why CURL performs poorly in vector-based experiments.

Original: (iii)
Revised: (iii)

Original: In the vector-based setting, some typical human-designed patterns (e
Revised: In the vector-based setting, some typical human-designed patterns (e

Original: g., reward prediction, inverse dynamics, and action inference) can be very detrimental to RL performance, implying that some renowned techniques in loss designs might not work well under atypical settings.
Revised: g., reward prediction, inverse dynamics, and action inference) can be very detrimental to RL performance, implying that some renowned techniques in loss designs might not work well under atypical settings.

Original: Number of Sources and Targets.
Revised: Number of Sources and Targets.

Original: We further investigate whether it is more beneﬁcial to use a small number of sources to predict a large number of targets ( n target > n source , e.g., using s t to predict s t +1 , s t +2 , s t +3 ), or the other way around ( n target < n source , e.g., using s t , s t +1 , s t +2 to predict s t +3 ).
Revised: We further investigate whether it is more beneﬁcial to use a small number of sources to predict a large number of targets ( n target > n source , e.g., using s t to predict s t +1 , s t +2 , s t +3 ), or the other way around ( n target < n source , e.g., using s t , s t +1 , s t +2 to predict s t +3 ).

Original: Statistical results are shown in Table 5, where we ﬁnd that auxiliary losses with more states on the target side have a signiﬁcant advantage over losses with more states on the source side.
Revised: Statistical results are shown in Table 5, where we ﬁnd that auxiliary losses with more states on the target side have a signiﬁcant advantage over losses with more states on the source side.

Original: This result echoes recent works [42, 39]: predicting more states leads to strong performance gains.
Revised: This result echoes recent works [42, 39]: predicting more states leads to strong performance gains.

Original:  Related Work
Revised:  Related Work

Original: Reinforcement Learning with Auxiliary Losses.
Revised: Reinforcement Learning with Auxiliary Losses.

Original: Usage of auxiliary tasks for learning better state representations and improving the sample efﬁciency of RL agents, especially on image-based tasks, has been explored in many recent works.
Revised: Usage of auxiliary tasks for learning better state representations and improving the sample efﬁciency of RL agents, especially on image-based tasks, has been explored in many recent works.

Original: A number of manually designed auxiliary objectives are shown to boost RL performance, including observation reconstruction [51], reward prediction [20], dynamics prediction [6] and contrastive learning objectives [23, 39, 42].
Revised: A number of manually designed auxiliary objectives are shown to boost RL performance, including observation reconstruction [51], reward prediction [20], dynamics prediction [6] and contrastive learning objectives [23, 39, 42].

Original: It is worth noting that most of these works focus on image-based settings, and only a limited number of works study the vector-based setting [32, 35].
Revised: It is worth noting that most of these works focus on image-based settings, and only a limited number of works study the vector-based setting [32, 35].

Original: Although people may think that vector-based settings can beneﬁt less from auxiliary tasks due to their lower-dimensional state space, we show in our paper that there is still much potential for improving their performance with better learned representations.
Revised: Although people may think that vector-based settings can beneﬁt less from auxiliary tasks due to their lower-dimensional state space, we show in our paper that there is still much potential for improving their performance with better learned representations.

Original: Compared to the previous works, we point out two major advantages of our approach.
Revised: Compared to the previous works, we point out two major advantages of our approach.

Original: (i) Instead of handcrafting an auxiliary loss with expert knowledge, A2LS automatically searches for the best auxiliary loss, relieving researchers from such tedious work.
Revised: (i) Instead of handcrafting an auxiliary loss with expert knowledge, A2LS automatically searches for the best auxiliary loss, relieving researchers from such tedious work.

Original: (ii) A2LS is a principled approach that can be used in arbitrary RL settings.
Revised: (ii) A2LS is a principled approach that can be used in arbitrary RL settings.

Original: We discover great auxiliary losses that bring signiﬁcant performance improvement in image-based and the rarely studied vector-based settings.
Revised: We discover great auxiliary losses that bring signiﬁcant performance improvement in image-based and the rarely studied vector-based settings.

Original: Automated Reinforcement Learning.
Revised: Automated Reinforcement Learning.

Original: RL training is notoriously sensitive to hyper-parameters and environment changes [18].
Revised: RL training is notoriously sensitive to hyper-parameters and environment changes [18].

Original: Recently, many works attempted to take techniques in AutoML to alleviate human intervention, for example, hyper-parameter optimization [7, 36, 49, 53], reward search [9, 45] and network architecture search [38, 10].
Revised: Recently, many works attempted to take techniques in AutoML to alleviate human intervention, for example, hyper-parameter optimization [7, 36, 49, 53], reward search [9, 45] and network architecture search [38, 10].

Original: In contrast to these methods which optimize a new conﬁguration for each environment, we search for auxiliary loss functions that generalize across different settings such as (i) different robots of control; (ii) different data types of observation; (iii) partially observable settings; (iv) different network architectures; (v) different benchmark domains.
Revised: In contrast to these methods which optimize a new conﬁguration for each environment, we search for auxiliary loss functions that generalize across different settings such as (i) different robots of control; (ii) different data types of observation; (iii) partially observable settings; (iv) different network architectures; (v) different benchmark domains.

Original: Automated Loss Design.
Revised: Automated Loss Design.

Original: In the AutoML community, it has become a trend to design good loss functions that can outperform traditional and handcrafted ones.
Revised: In the AutoML community, it has become a trend to design good loss functions that can outperform traditional and handcrafted ones.

Original: To be speciﬁc, to resolve computer vision tasks, AM-LFS [27] deﬁnes the loss function search space as a parameterized probability distribution of the hyper-parameters of softmax loss.
Revised: To be speciﬁc, to resolve computer vision tasks, AM-LFS [27] deﬁnes the loss function search space as a parameterized probability distribution of the hyper-parameters of softmax loss.

Original: A recent work, AutoLoss-Zero [28], proposes to search loss functions with primitive mathematical operators.
Revised: A recent work, AutoLoss-Zero [28], proposes to search loss functions with primitive mathematical operators.

Original: For RL, existing works focus on searching for a better RL objective, EPG [19] and MetaGenRL [21] deﬁne the search space of loss functions as parameters of a low complexity neural network.
Revised: For RL, existing works focus on searching for a better RL objective, EPG [19] and MetaGenRL [21] deﬁne the search space of loss functions as parameters of a low complexity neural network.

Original: Recently, [4] deﬁnes the search space of RL loss functions as a directed acyclic graph and discovers two
Revised: Recently, [4] deﬁnes the search space of RL loss functions as a directed acyclic graph and discovers two

Original: DQN-like regularized RL losses.
Revised: DQN-like regularized RL losses.

Original: Note that none of these works investigates auxiliary loss functions, which are crucial to facilitate representation learning in RL and to make RL successful in highly complex environments.
Revised: Note that none of these works investigates auxiliary loss functions, which are crucial to facilitate representation learning in RL and to make RL successful in highly complex environments.

Original: To the best of our knowledge, our work is the ﬁrst attempt to search for auxiliary loss functions that can signiﬁcantly improve RL performance.
Revised: To the best of our knowledge, our work is the ﬁrst attempt to search for auxiliary loss functions that can signiﬁcantly improve RL performance.

Original:  Conclusion and Future Work
Revised:  Conclusion and Future Work

Original: We present A2LS, a principled and universal framework for automated auxiliary loss design for RL.
Revised: We present A2LS, a principled and universal framework for automated auxiliary loss design for RL.

Original: By searching on training environments with this framework, we discover a top-performing auxiliary loss function A2-winner that generalizes well to a diverse set of test environments .
Revised: By searching on training environments with this framework, we discover a top-performing auxiliary loss function A2-winner that generalizes well to a diverse set of test environments .

Original: Furthermore, we present an in-depth investigation of the statistical relations between auxiliary loss patterns and RL performance.
Revised: Furthermore, we present an in-depth investigation of the statistical relations between auxiliary loss patterns and RL performance.

Original: 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 efﬁcient and practical.
Revised: 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 efﬁcient and practical.

Original: Limitations of our current work lie in that searching requires an expensive computational cost.
Revised: Limitations of our current work lie in that searching requires an expensive computational cost.

Original: In the future, we plan to incorporate more delicate information such as higher-order informationsuch as [12] of the inner-loop RL training procedure to derive more efﬁcient auxiliary loss search methods.
Revised: In the future, we plan to incorporate more delicate information such as higher-order information [12] of the inner-loop RL training procedure to derive more efﬁcient auxiliary loss search methods.