Wish you were here: Hindsight Goal Selection for long-horizon dexterous manipulation

1 INTRODUCTION . Recent advances in model-free reinforcement learning ( RL ) have enabled successful applications in a variety of practical , real-world tasks ( Gu et al. , 2017 ; Levine et al. , 2018 ; Riedmiller et al. , 2018 ; OpenAI et al. , 2019 ; Sharma et al. , 2020 ; Gupta et al. , 2021 ) . Given the complexity of these tasks , additional information such as demonstrations often plays an important role ( Vecerik et al. , 2019 ; Zuo et al. , 2020 ; Davchev et al. , 2020 ; Sutanto et al. , 2020 ; Luo et al. , 2021 ) . These methods are particularly useful to robotics as they enable efficient learning with only sparse rewards , which are easier to define . However , scaling such solutions to more complex long-horizon sequential settings with limited number of demonstrations remains an open challenge ( Gupta et al. , 2020 ) . Solving long-horizon dexterous manipulation tasks is an important problem as it enables a wide range of useful robotic applications ranging from object-handling tasks , common in collaborative settings , through to contact-rich dual arm manipulations , often seen in manufacturing . Consider the bi-manual 3.5mm jack cable insertion problem as illustrated in Figure 1 . Solving such tasks with vanilla demo-driven RL often fails due to the difficulty of assigning credit over long horizons in the presence of noisy exploration and sparse rewards . Self-supervision mechanisms like Hindsight Experience Replay ( HER ) ( Andrychowicz et al. , 2017 ) offer an alternative to ease the sparse reward problem by providing a stronger learning signal through additional goal-reaching tasks which are generated from the agent ’ s trajectories . While the more frequent reward can be beneficial , the agent ’ s trajectory ∗Work done during an internship at DeepMind ; † Denotes equal contribution . distribution is often considerably more complex and non-stationary in comparison to the target-task distribution . When the final task is very complex and the main task reward is not perceived , this can lead to learning a capable , goal-reaching policy for states close to the agent ’ s initial states while unable to complete the actual task . Our contribution is based on the insight that demonstration states can be viewed as samples of the target task distribution , and can therefore be used to constrain the self-supervision process to task-relevant goals . We introduce a framework for task-constrained goal-conditioned RL that flexibly combines demonstrations with hindsight relabelling . Unlike HER , which learns a general goalconditioned agent , we train a goal-conditioned agent specialized at achieving goals which directly lead to the task solution . The approach further allows to smoothly vary the task relevance of the relabelling process . Unlike conventional goal-conditioned RL , we enable agents to solve tasks through utilising abstract goal formulations , such as inserting a 3.5mm jack , common in the context of complex sequential tasks . We achieve this through using a continuously improving target goal distribution for the online goal selection stage . We refer to our method as a Hindsight Goal Selection for Demo-Driven RL or HinDRL for short . We demonstrate that the proposed solution can solve tasks where both demonstration-driven RL and its self-supervised version with HER struggle . Specifically , we show that HinDRL can dramatically reduce the number of required demonstration by an order of magnitude on most considered tasks . 2 RELATED LITERATURE . Using demonstrations in RL is a popular tool for robot learning . Classic approaches , such as ( Atkeson & Schaal , 1997 ; Peters & Schaal , 2008 ; Tamosiunaite et al. , 2011 ) , use expert demonstrations to extract good initial policies before fine-tuning with RL . Kim et al . ( 2013 ) ; Chemali & Lazaric ( 2015 ) use demonstrations to learn an imitation loss function to bootstrap learning . However , those solutions are not applied to complex sequential tasks . Follow up work that use deep neural networks enable the application of demo-driven RL to more complex tasks such as Atari ( Aytar et al. , 2018 ; Pohlen et al. , 2018 ; Hester et al. , 2018 ) , or robotics ( Vecerik et al. , 2017 ; 2019 ; Paine et al. , 2018 ) . We build upon DPGfD ( Vecerik et al. , 2019 ) , which we explain in more details in Section 3 . The Atari literature often requires large number of environment steps ( Salimans & Chen , 2018 ; Bacon et al. , 2017 ) . This can result in hardware wear and tear if applied on a physical robotics system . Alternatively , a number of robotics solutions assume the existence of a manually defined structure , e.g . through specifying an explicit set of primitive skills ( Stulp et al. , 2012 ; Xie et al. , 2020 ) , or a manually defined curricula ( Davchev et al. , 2020 ; Luo et al. , 2021 ) . However , hand-crafting structure can be sub-optimal in practice as it varies across tasks . In contrast , restricting the search space by greedily solving for specific goals ( Foster & Dayan , 2002 ; Schaul et al. , 2015 ) can be combined with implicitly defined curricula through retroactive goal selection ( Kaelbling , 1993 ; Andrychowicz et al. , 2017 ) . Hindsight goal assignment has been successfully applied to multi-task RL ( Li et al. , 2020 ; Eysenbach et al. , 2020 ) , reset-free RL ( Sharma et al. , 2021 ) but also for batch RL ( Kalashnikov et al. , 2021 ; Chebotar et al. , 2021 ) and hierarchical RL ( Wulfmeier et al. , 2020 ; 2021 ) . However , these approaches explore based on task-agnostic goal distributions and compensate with additional policy structure such as allowing for longer training to concurrently learn a hierarchy , or using large amounts of offline data . In this work , we extend hindsight relabelling mechanisms to guide exploration along task-specific distributions implied from a limited number of demonstrations and study its performance against strong task-agnostic goal distributions . Combining hindsight relabelling with demonstrations is not a novel concept . Ding et al . ( 2019 ) relabels demonstrations in hindsight to improve generative adversarial imitation learning . The proposed solution is used to learn a goal-conditioned reward function that can be applied on demodriven RL for hindsight relabelling . Gupta et al . ( 2020 ) propose a hierarchical data relabelling mechanism that uses demonstrations to extract goal-conditioned hierarchical imitation-based policies and apply them to multi task RL . Nair et al . ( 2018a ) fuses demo-driven RL similar to us but uses HER ’ s final-goal sampling mechanism and combines it with structured resets to specific demo states to overcome the difficulties of exploration for long-horizon tasks , which would be difficult to achieve in a continuous robotics task in practice . Zuo et al . ( 2020 ) uses demonstrations to bootstrap TD3 and combines it with HER . However , in all those works , hindsight goals were always chosen directly from the agent ’ s own trajectory . Instead , we consistently select goals directly from a task distribution implied from a collection of successful demonstrations . Hindsight goal selection has previously been generalized via adopting representation learning . Florensa et al . ( 2019 ) applied HER to a pixels-to-torque reaching task where sample efficiency was not a direct constraint . Nair et al . ( 2018b ) used a β−balanced variational auto encoder ( β -VAE ) to enable learning from visual inputs with a special reward . While β -VAE is broadly a powerful tool for unsupervised representation learning , it does not take into account the temporal nature of sequential robotics tasks . This makes them sub-optimal in the context of our work ( Chen et al. , 2021 ) . In contrast , contrastive-learning based techniques such as temporal cycle consistency ( TCC ) ( Dwibedi et al. , 2019 ) can provide more informative representations as previously discussed in concurrent work ( Zakka et al. , 2022 ) . In this work we combine HinDRL , a framework for goal-conditioned demo-driven RL with both learnt and engineered representations and provide a comprehensive sensitivity analysis to the quality of the encoder and the dependency against the number of demonstrations used . 3 PRELIMINARIES . Consider a finite-horizon discounted Markov decision process ( MDP ) , M = ( S , A , P , r , γ , T ) with transition probability P : S×A×S 7→ [ 0,1 ] . Let the current state and goal s , g∈ S⊆Rns be elements in the state space S , and a ∈ A⊆ Rna denote the desired robot action . We define a sparse environmental reward r ( s ) that assigns 0/1 reward only at the final state of an episode . Let ψ ( · ) be an encoding function that embeds a given state and goal to a latent state z = ψ ( s ) and a latent goal g = ψ ( g ) . Let ζ = ( xt0 , . . . , xT ) be a trajectory with a discrete horizon T , a discount function γ ( · ) , a stateaction tuple xt = ( st , at , g ) and a trajectory return R ( ζ ) = ∑Tt=t0 γr ( st ) . In this setting , a transition ( st , at , st+1 , g , r = 0 ) can be ’ relabelled ’ as ( st , at , st+1 , ĝ ≈ st+1 , r = 1 ) and both the original and relabelled transitions can be used for training . A separate , goal conditioned reward r ( zt+1 , ẑg ) = 1 [ zt+1 = ẑg ] uses the latent state and goal and assigns reward during relabelling . Finally , we also assume access to D demonstration trajectories D = { ζ j } Dj=0 that reach the goals { ĝ j } Dj=1 retrieved from the final state of the demonstration . Then , a goal-conditioned deterministic control policy parameterised by θ , πθ ( a|z , zg ) selects an action a given a latent state z and goal zg . In this context , we define an optimal policy π∗ = argmaxπ∈π̄ J ( π ) , where J ( π ) = Es0∼p ( s0 ) , g∼π ( g ) [ R ( ζ ) ] . DPGfD ( Vecerik et al. , 2019 ) is a powerful algorithm that bootstraps learning through a BC loss applied directly to the actor . This loss is typically defined as the L2 distance between the predicted by the policy actions and the true actions of a demonstration and is decayed proportionally to the number of environmental steps . Gradual conversion from using a BC loss to using the actor-critic loss is typically done using Q filtering or gradually decaying the BC loss . We employ the latter . The algorithm uses an L2 regularised actor-critic loss with a distributional critic , introduced in ( Bellemare et al. , 2017 ) . Using BC loss with DDPG was recently applied to a goal-conditioned setting ( Ding et al. , 2019 ) but goal-conditioned policy learning was never used with demo-driven distributional RL . 4 METHODOLOGY . This section describes our method , HinDRL . First , we extend DPGfD to its goal-conditioned interpretation and introduce a mechanism that deals with intermediate goals . Next , we introduce two demonstration-driven goal sampling strategies and discuss how to use them for goal-conditioned policy learning . Finally , we discuss how to extract temporally consistent representations and define a distance-based goal-conditioned reward function . We summarise our framework in Figure 2 .

Summary: This paper studies long-horizon manipulation tasks given sparse rewards and a few demonstrations from a hindsight relabeling-based approach. Specifically, it leverages the demonstration states as relevant goals in the goal relabeling process to guide task-specific exploration.  When evaluating the policy, the online goals are sampled from a goal dataset, composed of successful states and final states from the demonstrations. When relabeling the goals for training, the goals are sampled from the set of states visited by the demonstrations. Rather than operating over the raw state observations, this method utilizes an encoder, which can be either engineered by an expert or learned through self-supervision, to extract a latent representation. Then, the rewards are relabeled in hindsight based on a thresholded distance in the latent space.