Rethinking Goal-Conditioned Supervised Learning and Its Connection to Offline RL

1 INTRODUCTION . Reinforcement learning ( RL ) enables automatic skill learning and has achieved great success in various tasks ( Mnih et al. , 2015 ; Lillicrap et al. , 2016 ; Haarnoja et al. , 2018 ; Vinyals et al. , 2019 ) . Recently , goal-conditioned RL is gaining attention from the community , as it encourages agents to reach multiple goals and learn general policies ( Schaul et al. , 2015 ) . Previous advanced works for goal-conditioned RL consist of a variety of approaches based on hindsight experience replay ( Andrychowicz et al. , 2017 ; Fang et al. , 2019 ) , exploration ( Florensa et al. , 2018 ; Ren et al. , 2019 ; Pitis et al. , 2020 ) , and imitation learning ( Sun et al. , 2019 ; Ding et al. , 2019 ; Sun et al. , 2020 ; Ghosh et al. , 2021 ) . However , these goal-conditioned methods need intense online interaction with the environment , which could be costly and dangerous for real-world applications . A direct solution to learn goal-conditioned policy from offline data is through imitation learning . For instance , goal-conditioned supervised learning ( GCSL ) ( Ghosh et al. , 2021 ) iteratively relabels collected trajectories and imitates them directly . It is substantially simple and stable , as it does not require expert demonstrations or value estimation . Theoretically , GCSL guarantees to optimize over a lower bound of the objective for goal reaching problem . Although hindsight relabeling ( Andrychowicz et al. , 2017 ) with future reached states can be optimal under certain conditions ( Eysenbach et al. , 2020 ) , it would generate non-optimal experiences in more general offline goal-conditioned RL set- ∗Corresponding Authors ting , as discussed in Appendix B.1 . As a result , GCSL suffers from the same issue as other behavior cloning methods by assigning uniform weights for all experiences and results in suboptimal policies . To this end , leveraging the simplicity and stability of GCSL , we generalize it to the offline goal-conditioned RL setting and propose an effective and theoretically grounded method , named Weighted GCSL ( WGCSL ) . We first revisit the theoretical foundation of GCSL by additionally considering discounted rewards , and this allows us to obtain a weighted supervised learning objective . To learn a better policy from the offline dataset and promote learning efficiency , we introduce a more general weighting scheme by considering the importance of different relabeling goals and the expected return estimated with a value function . Theoretically , the discounted weight for relabeling goals contributes to optimizing a tighter lower bound than GCSL . Based on the discounted weight , we show that additionally re-weighting with an exponential function over advantage value guarantees monotonic policy improvement for goal-conditioned RL . Moreover , the introduced weights build a natural connection between WGCSL and offline RL , making WGCSL available for both online and offline setting . Another major challenge in goal-conditioned RL is the multi-modality problem ( Lynch et al. , 2020 ) , i.e. , there are generally many valid trajectories from a state to a goal , which can present multiple counteracting action labels and even impede learning . To tackle the challenge , we further introduce the best-advantage weight under the general weighting scheme , which contributes to the asymptotic performance when the data is multi-modal . Although WGCSL has the cost of learning the value function , we empirically show that it is worth for its remarkable improvement . For the evaluation of offline goal-conditioned RL algorithms , we provide a public benchmark and offline datasets including a set of challenging multi-goal manipulation tasks with a robotics arm or an anthropomorphic hand . Experiments1 conducted in the introduced benchmark show that WGCSL significantly outperforms other state-of-the-art baselines in the fully offline goalconditioned settings , especially in the difficult anthropomorphic hand task and when learning from randomly collected datasets with sparse rewards . 2 PRELIMINARIES . Markov Decision Process and offline RL The RL problem can be described as a Markov Decision Process ( MDP ) , denoted by a tuple ( S , A , P , r , γ ) , where S and A are the state and action spaces ; P describes the transition probability as S ×A× S → [ 0 , 1 ] ; r : S ×A → R is the reward function and γ ∈ ( 0 , 1 ] is the discount factor ; π : S → A denotes the policy , and an optimal policy π∗ satisfies π∗ = arg maxπ Es1∼ρ ( s1 ) , at∼π ( ·|st ) , st+1∼P ( ·|st , at ) [ ∑∞ t=1 γ t−1r ( st , at ) ] , where ρ ( s1 ) is the distribution of initial states . For offline RL problems , the agent can only access a static dataset D , and is not allowed to interact with the environment during the training process . The offline data can be collected by some unknown policies . Goal-conditioned RL Goal-conditioned RL further considers a goal space G. The policy π : S × G → A and the reward function r : S × G × A → R are both conditioned on goal g ∈ G. The agent learns to maximize the expected discounted cumulative return J ( π ) = Eg∼p ( g ) , s1∼ρ ( s1 ) , at∼π ( ·|st ) , st+1∼P ( ·|st , at ) [ ∑∞ t=1 γ t−1r ( st , at , g ) ] , where p ( g ) is the distribution over goals g. We study goal-conditioned RL with sparse rewards , and the reward function is typically defined as : r ( st , at , g ) = { 1 , ||φ ( st ) − g||22 < some threshold 0 , otherwise , where φ : S → G is a mapping from states to goals . For theoretical analysis in Section 3 , we consider MDPs with discrete state/goal space , such that the reward function can be written as r ( st , at , g ) = 1 [ φ ( st ) = g ] , which provides a positive signal only when st achieves g. The value function V : S × G → R is defined as V π ( s , g ) = Eat∼π , st+1∼P ( ·|st , at ) [ ∑∞ t=1 γ t−1r ( st , at , g ) |s1 = s ] . Goal-conditioned Supervised Learning Different with goal-conditioned RL that maximizes discounted cumulative return , GCSL considers the goal-reaching problem , i.e. , maximizing the last-step reward Eg∼p ( g ) , τ∼π [ r ( sT , aT , g ) ] for trajectories τ of horizon T . GCSL iterates between two process , relabeling and imitating . Trajectories in the replay buffer are relabeled with 1Code and offline dataset are available at https : //github.com/YangRui2015/AWGCSL hindsight methods ( Kaelbling , 1993 ; Andrychowicz et al. , 2017 ) to form a relabeled dataset Drelabel = { ( st , at , g′ ) } , where g′ = φ ( si ) for i ≥ t denotes the relabeled goal . The policy is optimized through supervised learning to mimic those relabeled transitions : θ = arg maxθ E ( st , at , g′ ) ∼Drelabel log πθ ( at|st , g′ ) . 3 WEIGHTED GOAL-CONDITIONED SUPERVISED LEARNING . In this section , we will revisit GCSL and introduce the general weighting scheme which generalizes GCSL to offline goal-conditioned RL . 3.1 REVISITING GOAL-CONDITIONED SUPERVISED LEARNING . As an imitation learning method , GCSL sticks the agent ’ s policy to the relabeled data distribution , therefore it naturally alleviates the problem of out-of-distribution actions . Besides , it has the potential to reach any goal in the offline dataset with hindsight relabeling and the generalization ability of neural networks . Despite its advantages , GCSL has a major disadvantage for offline goalconditioned RL , i.e. , it only considers the last step reward r ( sT , aT , g ) and generally results in suboptimal policies . A Motivating Example We provide training results of GCSL in the PointReach task in Figure 2 . The objective of the task is to move a point from the starting position to the desired goal as quickly as possible . The offline dataset is collected using a random policy . As shown in Figure 2 , GCSL learns a suboptimal policy which detours to reach goals . This is because GCSL only considers the last step reward r ( sT , aT , g ) . As long as the trajectories reaching the goal at the end of the episode , they are all considered equally to GCSL . To improve the learned policy , a straightforward way is to evaluate the importance of samples or trajectories for policy learning using importance weights . As a comparison , Weighted GCSL ( WGCSL ) , which uses a novel weighting scheme , learns the optimal policy in the PointReach task . We will derive the formulation of WGCSL in the following analysis . Connection between Goal-Conditioned RL and SL . GCSL has been proved to optimize a lower bound on the goal-reaching objective ( Ghosh et al. , 2021 ) , which is different from the goal-conditioned RL objective that optimizes the cumulative return . In this section , we will introduce a surrogate function for goal-conditioned RL and then derive the connection between goal-conditioned RL and weighted goal-conditioned supervised learning ( WGCSL ) . For overall consistency , we provide the formulation of WGCSL as below : JWGCSL ( π ) = Eg∼p ( g ) , τ∼πb ( ·|g ) , t∼ [ 1 , T ] , i∼ [ t , T ] [ wt , i log πθ ( at|st , φ ( si ) ) ] , ( 1 ) where p ( g ) is the distribution of desired goals and πb refers to the data collecting policy . The weight wt , i is subscripted by both the time step t and the index i of the relabeled goal . When wt , i = 1 , JWGCSL reduces to GCSL ( Ghosh et al. , 2021 ) , and we denote the corresponding objective as JGCSL for comparison convenience . Note that in Eq . 1 , φ ( si ) , i ≥ t implies that any future visited states after st can be relabeled as goals . Theorem 1 . Assume a finite-horizon discrete MDP , a stochastic discrete policy π which selects actions with non-zero probability and a sparse reward function r ( st , at , g ) = 1 [ φ ( st ) = g ] , where φ is the state-to-goal mapping and 1 [ φ ( st ) = g ] is an indicator function . Given trajectories τ = ( s1 , a1 , · · · , sT , aT ) and discount factor γ ∈ ( 0 , 1 ] , let the weight wt , i = γi−t , t ∈ [ 1 , T ] , i ∈ [ t , T ] , then the following bounds hold : Jsurr ( π ) ≥ T · JWGCSL ( π ) ≥ T · JGCSL ( π ) , where Jsurr ( π ) = 1T Eg∼p ( g ) , τ∼πb ( ·|g ) [ ∑T t=1 log π ( at|st , g ) ∑T i=t γ i−1 · 1 [ φ ( si ) = g ] ] is a surro- gate function of J ( π ) . We defer the proof to Appendix B.2 , where we also show that under mild conditions , ( 1 ) Jsurr is a lower bound of log J , and ( 2 ) Jsurr shares the same gradient direction with J at πb . Theorem 1 reveals the connection between the goal-conditioned RL objective and the WGCSL/GCSL objective . Meanwhile , it suggests that GCSL with the discount weight γi−t is a tighter lower bound compared to the unweighted version . We name this weight as discounted relabeling weight ( DRW ) as it intuitively assigns smaller weights on longer trajectories reaching the same relabeled goal .

Summary: This paper proposes an extension of GCSL where the goal-conditioned BC loss is weighted by a variable that correlates with the number of steps necessary to achieve the desired goal.  The advantage of this approach is that sub-optimal trajectories to a particular goal will get downweighted to the benefit of more direct trajectories.  This effectively adds a policy improvement step to GCSL with the underlying objective function being that of arriving at goal states as fast as possible.    Results on goal-conditioned tasks show that the proposed approach (WGCSL) performs consistently better than GCSL and in some cases significantly better (HandReach-expert).