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: This paper introduces a novel approach for improving sample efficiency in what they call "multi-target reinforcement learning". This is essentially goal-conditioned reinforcement learning where the tasks are of the form "bring me a spoon" or "go to the kitchen". Sample efficiency is improved in two ways. Firstly, they introduce a new auxiliary task of training a goal-discriminator that takes as input a state embedding and predicts which goal this state achieves. Second, they introduce a new architecture that uses attention and information from the goal discriminator to better extract task-relevant features for the policy.  Several experiments are provided to test the sample efficiency of agents trained using the method in both visual navigation tasks and robotic manipulation tasks. In both settings, their method outperforms on-policy baselines. Additional analysis is done, showing that: (1) while explicitly optimizing the state representation to improve the goal discriminator results in the fastest improvement in accuracy in the discriminator, the RL loss itself still causes the representations to contain information about the goal, just slower. This shows that the representation learning loss is helping the model learn faster what it would have to learn anyways. And finally, (2) saliency maps are provided which show that when using attention, the policy attends to more interpretable regions of the state space.