Weakly-Supervised Trajectory Segmentation for Learning Reusable Skills

Learning useful and reusable skill , or sub-task primitives , is a long-standing problem in sensorimotor control . This is challenging because it ’ s hard to define what constitutes a useful skill . Instead of direct manual supervision which is tedious and prone to bias , in this work , our goal is to extract reusable skills from a collection of human demonstrations collected directly for several end-tasks . We propose a weakly-supervised approach for trajectory segmentation following the classic work on multiple instance learning . Our approach is end-to-end trainable , works directly from high-dimensional input ( e.g. , images ) and only requires the knowledge of what skill primitives are present at training , without any need of segmentation or ordering of primitives . We evaluate our approach via rigorous experimentation across four environments ranging from simulation to real world robots , procedurally generated to human collected demonstrations and discrete to continuous action space . Finally , we leverage the generated skill segmentation to demonstrate preliminary evidence of zero-shot transfer to new combinations of skills . Result videos at https : //sites.google.com/view/trajectory-segmentation/ . 1 INTRODUCTION . Humans have an uncanny ability to generalize from one task to another using either few , or at times , no new examples . This wouldn ’ t be possible if they were to learn each new task from scratch . Humans rather extract reusable skills from already learned tasks and compose them to generalize to new tasks seamlessly . However , learning such repeatable skills has been a long standing challenge in sensorimotor control , partly because it is hard to define what constitutes a useful skill in itself . One way to layout the scope of a skill is either by designing a corresponding reward function , or collecting expert demonstrations . For instance , consider a skill of reaching for an object . One can easily learn a policy for this skill by either using reinforcement learning with l2 distance as reward ( Sutton & Barto , 1998 ) , or by imitation learning from kinesthetic demonstrations ( Argall et al. , 2009 ; Hussein et al. , 2017 ) . However , neither of these approaches provide a natural form of supervision because the way this skill is performed in isolation can be drastically different from the way it could be used as part of some end task . For instance , reaching for a cup for pushing is very different from the way one would reach for a cup to pick it up for pouring . Therefore , defining a skill directly in a supervised manner could easily lead to biased set of examples . A promising alternative is to learn skills that are already embedded in some useful end tasks . Previous works have explored this in the context of an agent ’ s own exploration ( Eysenbach et al. , 2018 ; Nair et al. , 2018 ; Pathak et al. , 2018 ) , where , the agent learns goal conditioned skill policies using data collected during its exploration phase . These skills are then used to plan for novel tasks at inference . However , exploration in itself is an open research problem , and hence , such approaches have difficulty in scaling to complex skills . In this work , we follow an alternative paradigm where we extract skills from a collection of human demonstrations gathered to perform different end tasks . A straightforward way to extract reusable skills would be to get an expert to label each time-step of demonstrations with the corresponding skill label . However , this per-step segmentation supervision is tedious for the expert and too expensive to scale . Moreover , such a labeling would be biased towards what expert thinks is the right segmentation than towards a segmentation which helps learn the skills better . This leads to the question : is it possible to use the expert knowledge to only know the types of skills present in demonstration and figure out the segmentation from the data itself ? Inspired by the classic work in multiple instance learning ( MIL ) ( Andrews et al. , 2002 ) , we propose a weakly-supervised approach for segmentation of human demonstration trajectories into primitive repeatable skills . Our approach assumes the access to only trajectory-level labels , i.e. , what primitive skills , aka sub-tasks , are present in the demonstration . The key insight is to learn a primitive skills classifier model conditioned on the input sensory data ( e.g . demonstration images ) , and incorporate a per-time step reasoning structure in this classifier . An overview of our approach is shown in Figure 1 . Our model generates per time-step primitive skill label prediction estimates which are then accumulated via differentiable function to generate trajectory-level predictions . In contrast to classic MIL , where only the most confident prediction across time-steps is trained , our full model is trained end-to-end using trajectory-level multi-class loss directly from raw sensory images . However , we are training with only trajectory-level supervision , then why should our per-step predictions of segmentation model converge to meaningful skill primitives ? Data comes to rescue ! Since our model trains across a variety of demonstrations , and hence , it would have seen plenty of demonstration trajectories that contain a certain skill primitive ( positives ) and plenty that do not ( negatives ) . The classification loss would force the segmentation model to focus on discriminative cues that are common across all positives , and absent from negatives . These discriminative cues corresponding to each skill would encourage the per time-step predictions to gradually correspond to the correct ground truth time-step labels . We evaluate our approach in four different environments : ( a ) As proof of concept , we start with a simple 2D navigation task in grid-world setup where the demonstrations are programatically generated . ( b ) We then discuss result in robotics setup with continuous control actions , in particular , Jaco robotic arm performing button pushing demonstrations in a touch pad with procedurally generated demonstrations . ( c ) We test our approach in a robotic setup with actual human demonstrations collected on the RoboSuite benchmark . ( d ) Finally , we evaluate on a real robot dataset with actual human demonstrations collected kinesthetically . Across all these environments , our approach outperforms the other variants of MIL and achieves decent segmentation accuracy . We then show zero-shot generalization to tasks containing novel permutations of skills in Jaco environment . 2 METHOD : SEGMENTING DEMONSTRATIONS INTO SKILL PRIMITIVES . Given a collection of human demonstration trajectories , our goal is to learn a labeling for skill primitives at each time-step of the sequence , i.e. , per time-step skill segmentation . Let X be a human demonstration trajectory denoted by X = { x1 , x2 , x3 . . . xT } , where T is the length of the demonstration and xt denotes a tuple of observation ( which is raw sensory image in our case ) at time t and action taken from it . Note that the action data is optional for the purpose of skill segmentation , but can be useful post segmentation for learning skill policies via imitation . Let Y = { y1 , y2 , y3 . . . yT } be the latent ground truth labeling of skill primitives in the sequence . Each label yt belongs to one of the k labels from a set of all skill classes C = { 1 , . . . , k } , i.e. , yt ∈ C. These per time-step labels are not only tedious for expert to annotate , but also difficult to scale . In this work , we do not assume access to yt during training , and learn the per time-step segmentation in a weakly-supervised fashion by only using trajectory-level 1-bit label during training , i.e. , whether a skill class is present in the trajectory or not . After training , our model is able to directly segment demonstrations at inference , without requiring any labels of any kind . An overview of our method is shown in Figure 1 . The marginal probability of a skill primitive at each time-step of demonstration can be written as P ( yt|θ , { xt , xt−1 . . . , x1 } ) where θ is the parameter vector of the segmentation model represented by a neural network in our case . If we had access to the true sequence labels Y , the network parameters θ can be easily learned via maximum log-likelihood by representing the probability as follows : P ( yt|θ , { xt , xt−1 . . . , x1 } ) = 1 Zt exp ( f ( yt ; θ , { xt , xt−1 . . . , x1 } ) ) ( 1 ) where Zt is the partition function at t , defined as Zt = ∑ k∈C exp ( ft ( k ; θ , { xt , xt−1 . . . , x1 } ) ) . The output of the function ft corresponds to the logit score value generated by the neural network . In order to model temporal dependency on across observation time-steps xt , we represent f ( . ) via a recurrent neural network , in particular , LSTM ( Hochreiter & Schmidhuber , 1997 ) . 2.1 WEAKLY-SUPERVISED TRAJECTORY SEGMENTATION . We are given a dataset of demonstration trajectories during training , D = { X1 , . . . , Xn } , where n is total number of demonstrations available for training . Each demonstration trajectory is weakly labelled with what skill primitives are contained within the trajectory . Neither do we have access to which time-step densely correspond to which skill primitive , nor to the permutation in which the skills are executed in . Instead , we are only given a set of skill primitive labels CX ∈ C present in the demonstration trajectory X . Although our supervision is only at trajectory-level , we do not directly predict output labeling CX from input demonstration X . Instead , we instill the structure of per-step prediction in our weakly supervised segmentation model by first computing the per-step classification score f ( yt ; θ , { xt , xt−1 . . . , x1 } ) and then accumulate it across all time-steps to compute the probability of a class in the whole trajectory . This weakly-supervised setup is captured by classical paradigm of multiple instance learning ( MIL ) ( Andrews et al. , 2002 ) . At inference , we use this per time-step score to compute the probability of skill primitives at each time t as described in Equation ( 1 ) . There are multiple ways one could accumulate these per time-step scores discussed as follows . 2.2 ACCUMULATION OF TIME-STEP PREDICTIONS . Intuitively , we would like an estimator that could generate an aggregated score for the class depending on how much each time-step votes for that class . We would ideally like the highest score value across time-steps to contribute most to the decision whether a class label ŶX is present in the trajectory X or not . One simple way to achieve that is to employ element-wise max operator , which is also the de facto approach in MIL to accumulate element-level scores . However , this would amount to passing gradients only to the most confident time-step and will completely eradicate the role of other time-steps . This is especially problematic in case of sequential trajectories because no skill primitive will be of only 1 time-step long . Hence , instead of max , we use a soft approximation to it which can take into account the contribution of all time-steps . In particular , we use log-sum-exp operator . Given the the logit score f ( yt ; θ , { xt , xt−1 . . . , x1 } ) at each time-step , the trajectory-level logit score g for class c ∈ C is computed as follows : g ( c ; θ , X ) = log ( T∑ t=1 exp ( f ( yt = c ; θ , { xt , xt−1 . . . , x1 } ) ) ) ( 2 ) We perform this operation for all c ∈ C and use softmax over g ( c ; θ , X ) to compute trajectory-level probability distribution Q ( c|X , θ ) . Finally , the parameters θ are optimized to maximize Q ( c|X , θ ) for each class c with respect to the ground truth trajectory level tags CX . Note that this optimization is fully differentiable through Equation 2 , and hence can be optimized via stochastic gradient descent in the end-to-end fashion . Pinheiro & Collobert ( 2015 ) has also showed the effectiveness of a temperature-based variant of log-sum-exp operation for semantic segmentation in images . However , these per-step scores f would be almost uniformly random in the beginning of the training process due to the absence of per-step supervision . Since we are training with only trajectorylevel supervision , why should these per-step predictions should ever converge to meaningful skill segmentation ? It turns out to be the case because we are learning across large variety of demonstration examples . Hence , for each skill primitive we would have seen plenty of positive as well as negative trajectories . The loss suppresses the negative classes and encourages the positive ones , hence our segmentation model would be forced to focus on discriminative cues that are exclusively common among the trajectories containing the skill primitives and not common to the cues that help distinguish other skills . Since our trajectory-level segmentation is based on a deterministic transformation of per-step predictions , each per-step score will then be forced to focus on those discriminative cues . The discriminative nature encourages the per time-step predictions to slowly drift towards the true latent ground truth segmentation which are not available directly .

Summary: This paper tackles the problem of learning to label individual timesteps of sequential data, when given labels only for the sequence as a whole. The authors take an approach derived from the multiple-instance learning (MIL) literature that involves pooling the per-timestep predictions into a sequence-level prediction and to learn the per-timestep predictions without having explicit labels. They evaluate several pooling techniques and conclude that the log-sum-exp pooling approach is superior. The learned segmentations are used to train policies for multiple control skills, and these are used to solve hierarchical control tasks where the correct skill sequence is known.