Goal-Aware Cross-Entropy for Multi-Target Reinforcement Learning

1 Introduction . Reinforcement learning ( RL ) has been expanding to various fields including robotics , to solve increasingly complex problems . For instance , RL has been gradually mastering skills such as robot arm/hand manipulation on an object [ 2 , 48 , 36 , 22 ] and navigation to a target destination [ 18 , 39 ] . However , to benefit humans like the R2-D2 robot in the Star Wars , RL must extend to realistic settings that require interaction with multiple objects or destinations , which is still challenging for RL . For this reason , it is important for multi-target tasks to be considered . We use the term multi-target tasks to refer to tasks that require the agent to interact with variable goals . In a multi-target task , targets are possible goal candidates , which may be objects or key entities that play a decisive role in determining the success or failure of the task execution . The goals may be selected among the targets by the current multi-target task , specified with a cue or an instruction such as “ Bring me a { spoon , cup or specific object } ” or “ Go to the { kitchen , livingroom or specific destination } ” . The states in which the agent reaches the goal are called goal states . Reinforcement learning allows learning multi-target or instruction-based tasks in an end-to-end manner . Latest reinforcement learning studies on these tasks mainly focus on prior knowledge about targets [ 27 , 30 , 8 ] . Other studies focus on learning representation on the environment [ 46 ] or about the targets only implicitly [ 44 ] . These methods lead to insufficient understanding of the goal , and sample-inefficiency and generalization problems arise . For this matter , we propose a Goal-Aware Cross-Entropy ( GACE ) loss and Goal-Discriminative Attention Networks ( GDAN ) 2 for multi-target tasks in reinforcement learning . These methods , unlike ∗Corresponding authors . †AI Institute , Seoul National University 2Code available at https : //github.com/kibeomKim/GACE-GDAN 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . previous studies without prior knowledge , allow semantic understanding of goals including their appearances and other characteristics . The agent automatically labels and collects goal states data through trial-and-error in a self-supervised manner . Based on these self-collected data , we use GACE loss as an auxiliary loss to train a goal-discriminator that learns goal representations . Lastly , the GDAN extracts the information from the goal-discriminator as a goal-relevant query , with which an attention is performed to infer goal-directed actions . We additionally propose visual navigation and robot arm manipulation tasks as benchmarks for multi-target task experiments . These tasks involve targets which are the multiple types of objects randomly placed within the environments . In the benchmarks , we make these tasks visually complex using randomized background textures , allowing us to evaluate generalization . In summary , the contributions of this paper are as follows : • We propose a Goal-Aware Cross-Entropy loss for learning auto-labeled goal states in a selfsupervised manner , solving instruction-based multi-target tasks in reinforcement learning . • We additionally propose Goal-Discriminative Attention Networks that use the goal-relevant query from the goal-discriminator to focus exclusively on important goal-related information . • Our method achieves significantly better performances in terms of the success ratio , sampleefficiency , and generalization on visual navigation and robot arm manipulation multi-target tasks . In particular , compared with the baseline methods , our method excels in Sample Efficiency Improvement metric by 17 times in visual navigation task and by 4.8 times in manipulation task . • We present two instruction-based benchmarks for goal-based multi-target environments , to interact with multiple targets for realistic settings . These benchmarks are made publicly available , along with the implementation of our proposed method . 2 Related work . Multi-target tasks in reinforcement learning There has been a sparsity of studies that aim to solve multi-target tasks with reinforcement learning . Wu et al . [ 44 ] proposes instruction-based indoor environments , as well as attention networks for multi-target learning to respond to the given instruction . Deisenroth and Fox [ 11 ] uses a model-based policy search method for multiple targets , instead of learning goal-relevant information . The aforementioned methods are for learning dynamics model or learning targets indirectly for multiple target tasks . In our method , on the other hand , the policy directly learns to distinguish goals without the dynamics model . The terms targets and goals have been used in diverse manners in the existing literature . There are subgoal generation methods that generate intermediate goals such as imagined goal [ 31 , 14 ] or random goals [ 34 ] to help agents solve the tasks . Multi-goal RL [ 47 , 35 , 13 , 10 ] aims to deal with multiple tasks , learning to reach different goal states for each task . Scene-driven visual navigation tasks [ 49 , 12 , 30 ] specify the goal by an image . However , these tasks require the agent only to search for visually similar locations or objects , rather than gaining semantic understanding of goals and promoting generalization . Instruction-based tasks [ 1 , 40 ] focus on learning to follow detailed instructions . These tasks have a different purpose from our tasks , where we focus on having appropriate interactions with multiple targets depending on the instruction . Representation learning in reinforcement learning To improve performance in RL , several recent works have used a variety of representation learning methods . Among these methods , many make use of auxiliary tasks taught via unsupervised learning techniques such as VAE [ 24 , 19 ] for learning dynamics [ 16 , 45 , 38 ] and contrastive learning method [ 7 ] for learning representations [ 26 ] of a given environment . Nair et al . [ 32 ] uses goal states , specified as the last states of the trajectories , to learn the difference between future state and goal state for redistributing model error in modelbased reinforcement learning . The aforementioned methods are for learning representations of the environment rather than for learning those of key states for solving the task . In contrast , our method directly learns goal-relevant state representations . Jaderberg et al . [ 21 ] proposes various auxiliary tasks for representation learning in RL . One such method is reward prediction , which learns the environment by predicting the reward of the future state from three classes : { positive , 0 , negative } . This method focuses on myopic reward prediction , and not on ultimately learning the goal state to solve the task . Consequently , it is not a suitable method to apply to multi-target environments , where the goal can be selected among a diverse range of objects . Attention methods in reinforcement learning The attention method was initially introduced for natural language processing [ 3 ] but is now being actively studied in various fields such as computer vision [ 43 ] . There are also attention-based methods for reinforcement learning which use inputs from different parts of the state [ 4 , 9 ] . Other methods include the use of encoding information of sequential frames up to the last step [ 29 ] , as well as the application of self-attention [ 6 ] . Unlike these methods , our approach explicitly extracts and focuses solely on goal-related information for solving a task . 3 Method . 3.1 Preliminary . Reinforcement learning ( RL ) from Sutton and Barto [ 41 ] aims to maximize cumulative rewards by trial-and-error in a Markov Decision Process ( MDP ) . An MDP is defined by a tuple ( S , A , R , P , γ ) , where S is the set of states , A is the set of actions , R : S × A → R is the reward function , P : S × A × S → R is the transition probability distribution , and γ ∈ ( 0,1 ] is the discount factor . At each time step t , the agent observes state st ∈ S , selects an action at ∈ A according to its policy π : S → A , and receives a reward rt and next state st+1 . In finite-horizon MDPs , return Rt = Σ T−t k=0γ krt+k is accumulated discounted rewards , where T is the maximum episode length . State value function V ( s ) = E [ Rt|st = s ] is the expected return from state s. Instruction-based multi-target reinforcement learning Universal value function approximators [ 37 ] estimate state value function as V ( s , x ) , for jointly learning the task from embedded state s and goal information x . This is relevant to the notions of multi-target RL that we define as below . Multi-target environments contain N targets , or goal candidates T = { z1 , z2 , ... , zN } . These environments provide an instruction Iz that specifies which target the agent must interact with , or the goal z ∈ T . The instruction is given randomly every episode , in the form such as “ Get the Bonus ” or “ Reach the Green Box ” as in our benchmark tasks , in which cases the goals are Bonus and Green Box respectively . Hence , the state value function in such task is V ( s , Iz ) = E [ Rt|s̄t = ( s , z ) ] , where s̄t is state conditioned on the goal z . The policy π is also conditioned on the instruction Iz as π ( at|st , Iz ) where at is action at time t given state st and Iz . 3.2 Auto-labeled goal states for self-supervised learning . Prior to describing the main method in Sec . 3.3 and 3.4 , this section explains the automatic collection of goal data for self-supervised learning . Suppose that the instruction Iz specifies the target z as the goal and the episode ends when the goal is reached at time step t′ . We refer to the state st′ as a goal state , which we assume is highly correlated with the reward or the rewarding state . Throughout the training , the reached goal label z and the corresponding goal state st′ are automatically collected as a tuple ( st′ , one_hot ( z ) ) , called storage data . Rather than being manually provided the goal information , the agent actively gathers the data needed to learn the goals , relying only on the instruction Iz given by the environment . This allows the agent to learn in an end-to-end , self-supervised manner . The states for the failed episodes are also stored as a negative case with low probability ϵN . We clarify that the storage data does not serve as the prior during the learning and is instead used for training alongside the multi-target reinforcement learning task . 3.3 Goal-discriminator by goal-aware cross-entropy . Our proposed method adds an auxiliary task to the main reinforcement learning task as shown in Figure 1a . A general reinforcement learning model consists of a feature extractor σ ( · ) , which converts state st to an encoding et , and an ActorCritic f ( · ) that outputs a policy π and value V : et = σ ( st ) ( 1 ) π ( at|st , I ) , V ( st , I ) = f ( et , I ) ( 2 ) where I is the instruction . The details of the ActorCritic f ( · ) and the resulting loss LRL depend on the specific reinforcement learning algorithm that is chosen . In our visual navigation experiments , we use asynchronous advantage actor-critic ( A3C ) [ 28 ] as the main algorithm , where the loss LRL is defined as the following Lp = ∇ log π ( at|st , I ) ( Rt − V ( st , I ) ) + β∇H ( π ( at|st , I ) ) ( 3 ) Lv = ( Rt − V ( st , I ) ) 2 ( 4 ) LRL : = LA3C = Lp + 0.5 · Lv ( 5 ) where Lp and Lv respectively denote policy and value loss , Rt denotes the sum of decayed rewards from time steps t to T , and H and β denote the entropy term and its coefficient respectively . See Appendix D for algorithm details . The vanilla base algorithm is inefficient and indirect at gaining an understanding of goals . Suppose that we have a set of goal states Sz ⊂ S corresponding to each goal z . Intuitively , given an instruction Iz that specifies z , the agent must learn the value function V ( s , I ) that discriminates the goals , such that V ( szgoal , I z ) > V ( sjgoal , I z ) for szgoal ∈ Sz and s j goal ∈ Sj , ∀z ̸= j . To achieve this effect , we propose Goal-Aware Cross-Entropy ( GACE ) loss as our contribution , which trains the goal-discriminator that facilitates semantic understanding of goals alongside the policy in Figure 1a . The inference process is as follows . First , the storage data collected in Sec . 3.2 are sampled as a batch of sgoal , i ’ s ( with i as data index within the batch ) and fed into the feature extractor σ ( · ) for state encoding in Eq . 6 . Next , the encoded input data esgoal , i is fed into the goal-discriminator d ( · ) , a multi-layer perceptron ( MLP ) , in Eq . 7 . This yields a prediction ggoal , i , a vector that contains probability distribution over possible goals that sgoal , i belongs to . esgoal , i = σ ( sgoal , i ) ( 6 ) ggoal , i = d ( esgoal , i ) ( 7 ) From the output ggoal , i , we calculate the Goal-Aware Cross-Entropy loss LGACE as Eq . 8 , where M is the batch size , and zi is the automatic label corresponding to state sgoal , i . LGACE = − M−1∑ i=0 one_hot ( zi ) · log ( ggoal , i ) ( 8 ) We complete the training procedure by optimizing the overall loss Ltotal as the weighted sum of the two losses in Eq . 9 . We focus on improving the policy for performing the main task and assign weight η to LGACE for performing goal-aware representation learning for the feature extractor σ ( · ) . We note that the goal-discriminator d ( · ) is updated according to LGACE only during training , not during inference , excluding the ActorCritic . Ltotal = LRL + ηLGACE ( 9 ) The explained procedure forms a visual representation learning method , where the GACE loss makes the goal-discriminator become goal-aware without external supervision . Such goal-awareness is advantageous for sample-efficiency in multi-target environments , as well as generalization , as it makes the agent robust in a noisy environment . Such effects are demonstrated in our quantitative results and qualitative analyses .

Summary: The paper addresses problem of sample efficiency and generalization in the multi-target reinforcement learning. The author proposes goal-aware cross-entropy (GACE) loss for auto-labeling goal states. The author also develops goal-discriminative attention network (GDAN) for multi-target reinforcement learning. They evaluate the proposed method on visual navigation robot arm manipulation with multi-target environment.  The proposed method in the paper compares with A3C, SAC.