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 contributes a method called HinDRL, which tackles the important problem of sparse reward robotic RL tasks when supplied with demonstrations. HinDRL seeks to build off of hindsight relabeling to learn a goal-conditioned policy which is trained in a self-supervised manner. Specifically, by relabeling transitions with task-specific goal candidates from expert demonstrations, the agent will be given more task-relevant feedback. By combining the goal selection strategy with goal-conditioned rewards which are either engineered or learned through temporal time-consistency, HinDRL is demonstrated on simulated sparse-reward robotics tasks where it achieves improved final performance and demo efficiency compared to HER.