XIRL:
Cross-embodiment Inverse Reinforcement Learning
Anonymous Author(s)
AffiliationAddress email
Abstract: We investigate the visual cross-embodiment imitation setting, in whichagents learn policies from videos of other agents (such as humans) demonstrating thesame task, but with stark differences in their embodiments – shape, actions, end-effectordynamics, etc.
In this work, we demonstrate that it is possible to automatically discoverand learn vision-based reward functions from cross-embodiment demonstrationvideos that are robust to these differences.
Specifically, we present a self-supervisedmethod for Cross-embodiment Inverse Reinforcement Learning (XIRL) that leveragestemporal cycle-consistency constraints to learn deep visual embeddings that capturetask progression from offline videos of demonstrations across multiple expert agents,each performing the same task differently due to embodiment differences.
Prior to ourwork, producing rewards from self-supervised embeddings typically required alignmentwith a reference trajectory, which may be difficult to acquire under stark embodimentdifferences.
We show empirically that if the embeddings are aware of task-progress,simply taking the negative distance between the current state and goal state in thelearned embedding space is useful as a reward for training policies with reinforcementlearning.
We find our learned reward function not only works for embodiments seenduring training, but also generalizes to entirely new embodiments.
Additionally, whentransferring real-world human demonstrations to a simulated robot, we find that XIRLis more sample efficient than current best methods.
Introduction
The ability to learn new tasks from third-person demonstrations holds the potential to enable robots toleverage the vast quantities of tutorial videos that can be gleaned from the world-wide web (e g., YouTubevideos).
However, distilling diverse and unstructured videos into motor skills with vision-based policiescan be daunting – as the videos themselves are often not only captured from different camera viewpointsin different environments, but also with different experts that may use different tools, objects or strategiesto perform the same task.
Perhaps most critically, there often exists a clear embodiment gap between thehuman expert demonstrator, and the robot hardware that executes the learned policies.
One approach toclose this gap is to learn a mapping between the human and robot embodiment [1], which is a non-trivialintermediate problem in itself.
Despite considerable progress in learning policies with paired observations and actions (e.g., collected viateleoperation) [2, 3], much less work has been done in getting robots to learn policies from tasks definedonly by third-person observations of demonstrations [4, 5, 6].
This is a surprisingly challenging problem,as different embodiments are likely to use unique strategies that suit them and allow them to make progresson a given task.
For example, if asked to “place five pens into a cup”, a human hand is likely to scoopup all pens before slipping them into the cup, whereas a two-fingered gripper might instead need to pickand place each pen individually.
Both strategies complete the task, but generate different state-actiontrajectories.
In this setting, it may be difficult to acquire labeled frame-to-frame correspondences betweenexpert demonstration videos and learned embodiments [6], particularly across a multitude of embodimentsor experts.
The approach we investigate instead is whether we can successfully learn tasks through alearned notion of task progress that is invariant to the embodiment performing the task.
In this work, we propose to enable agents to imitate video demonstrations of experts – including ones withdifferent embodiments – by using a task-specific, embodiment-invariant reward formulation trained via
XIRL
Policy Learning on Novel Embodiment Use Temporal Cycle-consistencyTo Learn an Encoder Generate Reward Function Train Agent with using RL and Learned Reward Distance to Goal R e w a r d Sof--Actor Critic
Reward Learning from Demonstrations
[Figure]
E n
[Figure]
o d e r
Goal
[Figure]
E n c o d e r E n c o d e r E n c o d e r
[Figure]
temporal cycle-consistency (TCC) [7].
We demonstrate how we can leverage an encoder trained with TCCto define dense rewards for downstream reinforcement learning (RL) policies via simple distances in thelearned embedding space.
Simulated experiments across four different embodiments show that these learnedrewards are capable of generalizing to new embodiments, enabling unseen agents to learn the task via reinforcement, and surprisingly in some cases, exceeding the sample efficiency of the same agent learned withground truth sparse rewards.
We also demonstrate the effectiveness of our approach for learning robot policies using human demonstration on the puck-pushing environment from [6], where our reward is first learnedon real-world human demonstrations, then used to teach a Sawyer arm how to perform the task in simulation.
Our contributions are as follows: ( i)
We introduce Cross-embodiment Inverse Reinforcement Learning
(XIRL), an effective, label-free framework for tackling cross-embodiment visual imitation learning.
Our core contribution is to use self-supervised learning on third-person demonstration videos to define dense reward functions amenable for downstream reinforcement learning policies, (ii) Along with XIRL, we release a cross-embodiment imitation learning benchmark, X-MAGICAL which features multiple simulated agents with different embodiments performing the same manipulation task, including one thousand expert demonstrations for each agent, (iii) We show that XIRL significantly outperforms alternative methods on both the X-MAGICAL benchmark and the human-to-robot transfer benchmark from [6], and discuss our observations, which point to interesting areas for future research, (iv) Finally, we introduce a real-world dataset, X-REAL (Cross-embodiment Real demonstrations), of a manipulation task performed with nine different embodiments, which can be used to evaluate cross-embodiment reward learning methods 1 .
Code and datasets will be made available
Related Work
Traditional formulations of imitation learning [8, 1, 9] assume access to a corpus of expert demonstrationdata which includes both state and action trajectories of the expert policy.
In the context of third-personimitation learning, including when learning from expert agents with different embodiments, obtainingaccess to ground-truth actions is difficult or impossible 2 .
Inferring expert actions.
To address this issue, several approaches either try to infer expert actions [10, 11,12] – for example by training an inverse dynamics model on agent interaction data [10] – or employ forwardprediction on the next state to imitate the expert without direct action supervision [13].
While these methodssuccessfully address learning from observation-only demonstrations, they either do not support skill transferto different policy embodiments at all, or they cannot take advantage of multiple embodiments in order toimprove generalization to unseen policy configurations.
We explicitly address these problems in this work.
Imitation via learned reward functions.
In contrast to imitation via supervised methods, such as
BCO [10], a recent body of work [4, 5, 14, 6, 15] has focused on learning reward functions from expertvideo data and then training RL policies to maximize this reward.
In [4], the authors combine ImageNetpre-trained visual features with an L2-norm distance reward to match policy and expert observationsin a latent feature space.
In their follow-up work [5], the reward is computed in a viewpoint-invariant representation that is self-supervised on video data.
While both these methods are compelling in theiruse of cheap unlabeled data to learn invariant rewards, the use of a time index as a heuristic for definingweak correspondence is a constraining limitation for tasks that need to be executed at different speeds,or are not strictly monotonic (e.g., have ambiguous sub-task ordering).
In [14] a dense reward is learnedvia unsupervised learning on YouTube data and the authors make no assumption about time alignment.
However, in their work, the expert and learned policy are executed in the same domain and embodiment,an assumption we relax in our work.
Framed in a multi-task learning setting, [16] propose training policieswith morphologically different embodiments first on a similar set of proxy tasks, in order to learn a latentspace to map between domains, and then sharing skills on a held-out task from one policy to another.
Atime-index heuristic is used to define a metric reward when performing RL training of the new task.
Inour work, the learned embedding finds correspondences in a fully-unsupervised fashion, without the needfor such strict time alignment.
In [17], a small sub-set of states is human labeled for goal success and aconvolutional network is then trained to detect successful task completion from image observations, whereon-policy samples are used as negatives for the classifier.
By contrast, our learned embedding encodestask progress in its latent representation without the use of expensive human labels.
Imitation via domain adaptation.
An additional category of approaches to third-person imitation learningare those that perform domain adaptation of expert observations [18, 19, 20, 21].
For instance, in [18]a CycleGAN [22] architecture is used to perform pixel-level image-to-image translation between policydomains, which is then used to construct a reward function for a model-based RL algorithm.
A similarmodel-free approach is proposed in [19].
In [20], a generative model is used to predict robot high-levelsub-goals conditioned on a third-person demonstration video, and a lower-level control policy is trained ina task-agnostic manner.
Similarly, [21] uses high level task conditioning from zero-shot transfer of expertdemonstrations, but they use KL matching to perform both high and low-level imitation.
In contrast to thesemethods, the unsupervised TCC alignment in this work avoids performing explicit domain adaptation orpixel-level image translation by instead learning a robust and invariant feature space in a fully offline fashion.
Reinforcement learning with demonstrations.
Recent work in offline-reinforcement learning [6]explicitly tackles the problem of policy embodiment and domain shift.
Their method, Reinforcement
Learning from Videos (RLV), uses a labelled collection of expert-policy state pairs in conjunction withadversarial training to learn an inverse dynamics model jointly optimized with the policy.
In contrast, weavoid the limitation of collecting human-labeled dense state correspondences by using a self-supervisedalgorithm (i.e., TCC [7]) which uses cycle-consistency to automatically learn the correspondence betweenstates of two domains.
We also show that this formulation improves generalization to unseen embodiments.
Since the problem setup is similar to ours, we also compare to their method as a baseline.
Approach
Our overall XIRL framework (Figure 1) addresses the cross-embodiment visual imitation problem(Section 3.1).
The framework consists of first using TCC to self-supervise embodiment-invariant visualrepresentations (Section 3.2), then using a novel embodiment-invariant visual reward function (Section 3.3)to perform cross-embodiment visual imitation via reinforcement learning (Section 3.4).
3.1 Problem Formulation
Our objective is to extract an agent-invariant definition of a task, via a learned reward function, from adataset of videos of separate agents performing the same task.
In particular, we are interested in agentsthat may solve the task in entirely different ways due to differences in their end-effector, shapes, dynamics,etc., which we refer to as embodiment differences.
For example, consider how differently a vacuum gripperand a parallel-jaw gripper will grasp an object as a result of their respective end-effectors.
Such a setupis quite common in robotic imitation learning, where we might have access to observation data of humansdemonstrating a task, but want to teach a robot to perform it.
It is very likely that the way the humanexecutes the task will diverge from how the robot would execute it.
We define a dataset of multiple agents performing the same task T as D =
(cid:83) ni
=1 D i , where D i is anagent-specific dataset containing observations of only agent i performing task T .
Each agent’s dataset
D i is a collection of videos defined as D i
= { v 1 i ,v 2 i ...v Ki } , where v ji represents the video of the j thdemonstration of agent i successfully performing the task T .
We would like to highlight that D onlycontains observation data, i.e., it does not store the actions taken by the respective agents.
We useself-supervised representation learning techniques to learn task-specific representations from this dataset.
3.2 Representation Learning
In this work, we use TCC [7] to learn task-specific representations in a self-supervised way.
The methodhas been shown to learn useful representations on videos of the same action for temporally fine-grained [Figure] [Figure] [Figure]
[Figure]
[Figure] [Figure] downstream tasks. In their paper, the authors show that TCC representations can predict frame-level taskprogress, such as predicting how much water is in a cup during pouring, without requiring any humanannotations.
Task progress can provide dense signals for learning a new task and we would like to bakethis property into our learned reward.
Another advantage of TCC is that it does not require supervisionfor frame-level alignment (i.e., which frames in two videos correspond to each other).
Such frame-to-framecorrespondences are required by a prior method, RLV [6], to achieve successful reinforcement learningon the considered tasks, but we would like to avoid this type of manual supervision.
We train an image encoder φ , that ingests an image I and returns an embedding vector φ ( I ) , using TCC.
TCC assumes that there exist semantic temporal correspondences between two video sequences, eventhough they are not labeled and the actions may be executed at different speeds.
Note that in our case, theassumption holds since all sequences in our dataset originate from the same task.
Furthermore, even if twodifferent agents accomplish the task very differently, there will be a common set of states or frames thatwill temporally correspond.
When we apply the TCC loss, we compare frames of one agent performingthe task with frames from another, and search for temporal similarities in how both execute the task.
By performing multiple such comparisons and cycling-back (described in the next paragraph), TCCencodes task progress in its latent representation, a property that is useful for learning task-specific, yetembodiment-invariant reward functions.
For completeness, we describe the training technique below.
We define a dataset D all = { v 1 ,v 2 ...v N } that contains the videos of all agents executing the same task
T .
We are able to merge the datasets of different agents since TCC does not require embodimentlabels (i.e., IDs) during training.
We first sample a random batch of videos and embed all their framesusing the aforementioned image encoder φ .
For each video v i , this results in a sequence of embeddings
V i = { φ ( v 1 i ) ,φ ( v 2 i ) ...,φ ( v L i i ) } , where L i is the length of the i th video.
From this mini-batch of sequences ofvideo frame embeddings, we choose a pair of sequences V i and V j and compute their TCC loss.
In particular,we randomly sample a frame embedding from sequence V i – say V ti corresponding to the t
th frame ofvideo v i – and compute the soft nearest-neighbor of V ti in sequence V j in the embedding space as follows:
[Equation]
, where
[Equation]
We then cycle-back [7] to the first sequence V i by computing the soft-nearest neighbor of (cid:102) V tij with allthe frames in V i .
The probability of cycling-back to the k th frame in V i can be computed as:
[Equation]
The expected frame index we cycle-back to is then µ tij = (cid:80) L i k β kijt k .
Since we know the index of theframe that started the cycle, in this case t , we can minimize the mean-squared (MSE) error loss betweent and the closest index retrieved via soft-nearest neighbor, i.e., (cid:102) V tij .
The loss for a single frame is thus:
L t ij =( µ t ij − t ) 2 .
Finally, we minimize the average loss L over all frames in video v i with all other videosin the dataset v j , defined as L = (cid:80) ijt L tij .
3.3 Reward Function
Once the encoder φ has been trained on demonstrations of different agents performing the same task T ,we want to use it to transfer information about the task from one agent to another.
We do this by leveraging [Table] φ to generate rewards via distances to goal observations in the learned embedding space.
Specifically, wedefine the goal embedding g as the mean embedding of the last frame of all the demonstration videos inour offline dataset D all .
Concretely, g = (cid:80) Ni =1 φ ( v Lii ) / N , where L i is the length of video v i .
Our reward rthen is the scaled negative distance of the current state embedding to the goal embedding g i.e., r ( s )
=
− 1 / κ ·(cid:107) φ ( s ) − g
(cid:107) 22 , where φ ( s ) is the state embedding at the current timestep and κ is a scale parameterthat ensures the distances are in a range amenable for training reinforcement learning algorithms [23].
Wefound it effective to set κ to be the average distance of the first frame’s embedding to the goal embeddingfor all the demonstrations in the dataset.
Defining r in this manner gives us several advantages: (a) it isdense, encoding both task completion and task progress, (b) it does not require any correspondence witha reference trajectory [5, 14] and (c) it sidesteps the need for a finite library of reference trajectories, unlikeprior work [5, 16] that define time-indexed rewards relative to some reference trajectory.
Thus, agents withtrajectories of varying lengths (due to embodiment-specific constraints) can efficiently leverage this rewardbecause the learned encoder can map different strategies to a common notion of task progress.
3.4 Reinforcement Learning
Using the pre-trained frozen encoder φ , we define a Markov Decision Process (MDP) for any agent asthe tuple (cid:104)S , A , P, r (cid:105) where S is the set of possible states, A is the set of possible actions, P is the statetransition probability matrix encoding the dynamics of the environment (including the agent) and r isthe learned reward function (defined in Section 3.3).
Notice how we are able to use the same rewardfunction r for any agent – even ones that the encoder may not have seen during training.
Furthermore,this reward function solely depends on the learned encoder.
Hence, the task represented by the MDP nowdepends solely on how well the encoder has learned task-specific representations since we do not use theenvironment reward to either define the task or learn the policy.
This is an important distinction because weare expecting the encoder to generalize to new states that an agent might encounter during training, as theexpert demonstrations in our dataset only contain successful trajectories.
It is also possible augment sparserewards (which only define task success or failure) with our learned dense reward while training policies.
Experimental Setup
4.1 X-MAGICAL Benchmark
We introduce a cross-embodiment imitation learning benchmark, X-MAGICAL, which is based on theimitation learning benchmark MAGICAL [24], implemented on top of the physics engine PyMunk [25].
In this work, we consider a simplified 2D equivalent of a common household robotic sweeping task,wherein an agent has to push three objects into a predefined zone in the environment (colored in pink).
Wechoose this task specifically because its long-horizon nature highlights how different agent embodimentscan generate entirely different trajectories.
The reward in this environment is defined as the fraction ofdebris swept into the zone at the end of the episode.
Multiple Embodiments in X-MAGICAL.
We create multiple embodiments by designing agents withdifferent shapes and end-effectors that induce variations in how each agent solves a task.
In Figure 1, weshow three of these embodiments and some sample trajectories that solve the sweeping task.
Please see
Appendix B for a detailed description of the benchmark.
Three agents are shaped like a stick and they differonly in length.
We call them short-stick , medium-stick and long-stick based on the length of their body.
Theagent in the last row is called gripper : it is circular in shape and has two arms that can actuate.
All agentsare capable of two actions - a rotation around their axis and a translation in a forward/backward directionalong this axis (similar to the agent in the MAGICAL benchmark).
All agents have a two-dimensionalaction space and use force and torque control to change their position and orientation respectively.
Thegripper agent has an additional degree of freedom for opening or closing its fingers.
The default stateof the gripper’s fingers is open.
For all agents, the state representation is a 16-dimensional vector withthe following information: ( x, y ) position of the agent, (cos θ, sin θ ) where θ is the agent’s 2D orientation,and for each of the three debris: its ( x, y ) position, its distance to the agent and its distance to the goal zone.
We frame-stack [26] three consecutive state vectors to encode temporal and velocity information,resulting in a final state dimension of 48 .
Demonstrations and Different Embodiment Strategies in X-MAGICAL.
To learn task-specific representations for this task, we collect 1000 demonstrations per agent, where each demonstration consists insweeping all three debris, initialized with random positions, into the target zone.
This is the dataset D all (described in Section 3.1) containing observation-only agent-specific demonstrations.
In Figure 2, we highlightthe differences that exist between the trajectories taken by these agents.
Figure 2a shows a heat map of thefrequency of visits ( state visitation count ) of each agent at every 2D position in the grid, across all demonstrations in the dataset.
We plot the 2D projection of the state visitation count onto the XY plane with a binwidth of 0 . 1 .
Yellow encodes higher state visitation whereas blue encodes lower state visitation.
We observethat agent long-stick has less coverage of the environment as opposed to agent gripper , which has significantly more coverage.
Similarly, we show the distribution of debris locations for all agents across all demonstrations in Figure 2b.
In Figure 2c, we plot a randomly sampled trajectory from each agent’s demonstrationpool.
We use transparency to encode the start (lighter) and end (darker) of the trajectory.
It is clear from thisfigure that each agent solves the task in a different manner.
Additionally, there is a significant difference inthe time taken by each agent to execute this task, as shown in Table 1.
Agent long-stick is able to finish thetask the quickest because of its long shape that can sweep all the debris at once, while agents short-stick andgripper take longer because they have to frequently push or grasp one debris at a time.
These differences arethe types of challenges that a representation must overcome to successfully generalize across embodiments.
As such, X-MAGICAL serves to create a highly simplified version of a real-world scenario where we mightwant to learn new tasks from an observation dataset of humans performing these tasks in highly diverse ways.
4.2 Baselines
Here, we describe the alternative reward functions we baseline our method against, color coded to matchtheir appearance in Figs.
3, 4, and 5. 1) ImageNet : We use an ImageNet pre-trained ResNet-18 encoder withno additional self-supervised training, i.e., we load the pre-trained weights, discard the classification head,and use the 512-dimensional embedding space from the previous layer.
2) Goal classifier :
We follow [27]and train a goal frame classifier on a binary classification task where the last frame of all the demonstrationsis considered positive and all the others are considered negatives.
We use the output probabilities of theclassifier as the reward function.
3) LIFS :
We implement the method from [16] which learns a feature spacethat is invariant to different embodiments using a contrastive loss function paired with an autoencodingloss. 4)
TCN : single-view Time-Contrastive Network (TCN) [5] with positive and negative frame windowsof 1 and 4 respectively.
For more details regarding baseline implementations, see Appendix D.1.1.
4.3 Implementation Details
Our encoder for all experiments and methods is a ResNet-18 [28] initialized with ImageNet pre-trainedweights.
We replace the classification head with an embedding layer outputting a 32-dimensional vector.
The encoder is trained on images of resolution 224 × 224 with ADAM [29] and a learning rate of 10 − 5which is decreased by a factor of 2 when the validation loss plateaus.
Note that our learned reward isagnostic to the RL algorithm used – in this work, we opt for Soft-Actor Critic (SAC) [30], which is areinforcement learning algorithm that has been successfully used to train policies for continuous controltasks [31].
Once the TCC encoder is trained, we use it to embed the observation frames as the agentinteracts with the environment.
Experiments
We execute a series of experiments to evaluate whether the learned reward functions are effective at visualimitation.
Specifically, our experiments seek to answer the following questions: first (Section 5.1), inthe same -embodiment case, where the demonstration dataset D contains the embodiment of the learningagent, does our method enable successful reinforcement learning for that agent?
Next (Section 5.2), weinvestigate our primary interest, the cross -embodiment case, where the demonstration dataset D does notcontain the embodiment of the learning agent.
To additionally test our approach using real-world data(Section 5.3), we use the dataset from [6] to leverage real-world human demonstrations to learn policiesin simulation.
Note that each embodiment’s performance is evaluated over 50 episodes and all figuresplot the mean performance over 5 random seeds, with a standard deviation shading of ± 0 .
Videos ofour results are in the supplementary video.
5.1 Results on Learning from Same-Embodiment Demonstrations
In this experiment, we want to validate whether our approach of using a reward function trained with
TCC is good enough to train agents to perform the task defined in a dataset of expert demonstrations.
Note that the learned reward function has to be robust enough such that it can provide a useful signal for
[Figure] [Figure]
[Table]
new states an agent might encounter while learning a policy and interacting with the environment.
In
Figure 3, we compare our method XIRL with baselines described in Section 4.2.
We find XIRL is moresample-efficient than the other learned reward baselines.
We attribute this sample efficiency to the fact thatthe TCC embeddings encode task progress which helps the agent learn to reach for objects and goal zoneswhile interacting with the environment, rather than exploring in a purely random manner.
This experimentprovides evidence that XIRL’s reward function is suitable for downstream reinforcement learning.
5.2 Results on Learning from Cross-Embodiment Demonstrations
After verifying that XIRL works on the embodiments it was trained with, we move to the experiments thatanswer the core question in our work: can XIRL generalize to unseen embodiments ?
In this section weconduct experiments where the reward function is learned using embodiments that are different from theones on which the policy is trained.
As noted in Table 1, the timescales with which the agents execute thetask can vary significantly.
We conduct four experiments, each one corresponding to holding-out one agentfrom the expert demonstration set.
In each such experiment, we train an encoder on demonstrations fromthe remaining three agents.
We compare with reward functions learned using TCN, LIFS, and goal frameclassifiers.
While both TCC and TCN are contrastive losses, the former makes explicit comparisons acrossdifferent embodiments, whereas the latter implicitly relies on an encoder shared across embodiments tolearn the cross-embodiment representation.
In Figure 4, we show that XIRL generalizes to new agentssignificantly better than the other learned reward baselines.
5.3 Results on Learning from Real-World Cross-Embodiment Demonstrations
In this experiment, we test how well we can learn rewards from more challenging real-world humandemonstration videos.
To do so, we use the dataset introduced in [6].
We train only on real-world humandemonstrations and use the learned reward to train a Sawyer arm in sim to push a puck to a designatedtarget zone.
We compare our results to training the policy on the sparse reward from the environment and
RLV method presented in [6].
Unlike RLV, we (i) do not use human labeling of paired frames, and (ii)do not use any data from the robot domain in learning the reward function.
Like RLV we also use thesparse reward from the environment.
As demonstrated in Figure 5, we find our approach can improve thesample-efficiency of learning in the environment, compared to using RLV or solely the environment reward.
5.4 Qualitative comparison between learned reward functions vs. handcrafted rewards
In Figures 6 and 7, we visualize and compare the XIRL reward function with the environment’s reward(ground truth) for the Sweeping task from X-MAGICAL (see Sec.
4.1 for details) and the Puck Pushingand Drawer Opening tasks from [6] 3 .
We find that the learned reward is highly correlated with theground truth reward from the environment for both successful demos (first column in Figures 6 and 7)and unsuccessful demos (second column in Figures 6 and 7).
It is especially encouraging to see that ina sparser reward environment like the Drawer Opening task, XIRL provides a dense signal that shouldallow the agent to learn the task more efficiently.
Additionally, we can see that in the example of the failed
[Figure]
collision trajectory (i
[Figure]
e., second row third column in Fig. 7), where the arm collides with the drawer ratherthan opening it, XIRL is able to provide it with a partial reward (i
[Figure]
e., for correctly moving towards thedrawer) as opposed to the environment reward which remains zero.
[Figure]
task. The learned reward was trained on the three other agents.
[Figure]
Conclusion
This paper presents XIRL, a framework for learning vision-based reward functions from videos of expertdemonstrators exhibiting different embodiments.
XIRL uses TCC to self-supervise a deep visual encoderfrom videos, and uses this encoder to generate rewards via simple distances to goal observations in the embedding space.
XIRL enables unseen agents with new embodiments to learn the demonstrated tasksvia IRL.
Reward functions from XIRL are fully self-supervised from videos, and we can successfullylearn tasks without requiring manually paired video frames [6] between the demonstrator and learner.
Inthis sense, our method presents favorable scalability to an arbitrary number of embodiments or expertswith varying skill levels.
Experiments show that policies learned via XIRL are more sample efficientthan multiple baseline alternatives, including TCN [5], LIFS [16], and RLV [6].
While our experimentsdemonstrate promising results for learning policies in simulated environments using rewards learned fromboth simulated and real-world videos, we have yet to show policy learning on a real robot, which we lookforward to trying post-COVID.
Appendix A Expanded Related Work
Traditional formulations of imitation learning [8, 1, 9] assume access to a corpus of expert demonstrationdata which includes both state and action trajectories of the expert policy.
In the context of third-personimitation learning, including when learning from expert agents with different embodiments, obtainingaccess to ground-truth actions is difficult or impossible.
Inferring expert actions.
To address this issue, several approaches either try to infer expert actions[10, 11, 12, 32] – for example by training an inverse dynamics model on agent interaction data [10] – oremploy forward prediction on the next state to imitate the expert without direct action supervision [13].
Inthe case of [33], a video-based action classifier trained on a large-scale human activity dataset is leveragedto provide rewards for single-task RL policies, which are then used to provide expert state-action pairsfor multi-task behavior cloning.
While these methods successfully address learning from observation-onlydemonstrations, they either do not support skill transfer to different policy embodiments at all, or theycannot take advantage of multiple embodiments in order to improve generalization to unseen policyconfigurations.
We explicitly address these problems in this work.
Imitation via learned reward functions.
In contrast to imitation via supervised methods, such as
BCO [10], a recent body of work [4, 5, 14, 6, 15, 34] has focused on learning reward functions from expertvideo data and then training RL policies to maximize this reward.
In [4], the authors combine ImageNetpre-trained visual features with an L2-norm distance reward to match policy and expert observationsin a latent feature space.
In their follow-up work [5], the reward is computed in a viewpoint-invariantrepresentation that is self-supervised on video data.
While both these methods are compelling in theiruse of cheap unlabeled data to learn invariant rewards, the use of a time index as a heuristic for definingweak correspondence is a constraining limitation for tasks that need to be executed at different speeds,or are not strictly monotonic (e.g., have ambiguous sub-task ordering).
In [14] a dense reward is learnedvia unsupervised learning on YouTube data and the authors make no assumption about time alignment.
However, in their work, the expert and learned policy are executed in the same domain and embodiment,an assumption we relax in our work.
Framed in a multi-task learning setting, [16] propose training policieswith morphologically different embodiments first on a similar set of proxy tasks, in order to learn a latentspace to map between domains, and then sharing skills on a held-out task from one policy to another.
Atime-index heuristic is used to define a metric reward when performing RL training of the new task.
Inour work, the learned embedding finds correspondences in a fully-unsupervised fashion, without the needfor such strict time alignment.
In [17], a small sub-set of states is human labeled for goal success and aconvolutional network is then trained to detect successful task completion from image observations, whereon-policy samples are used as negatives for the classifier.
By contrast, our learned embedding encodestask progress in its latent representation without the use of expensive human labels.
Imitation via domain adaptation.
An additional category of approaches to third-person imitation learningare those that perform domain adaptation of expert observations [18, 19, 20, 21].
For instance, in [18]a CycleGAN [22] architecture is used to perform pixel-level image-to-image translation between policydomains, which is then used to construct a reward function for a model-based RL algorithm.
A similarmodel-free approach is proposed in [19].
In [20], a generative model is used to predict robot high-levelsub-goals conditioned on a third-person demonstration video, and a lower-level control policy is trained ina task-agnostic manner.
Similarly, [21] uses high level task conditioning from zero-shot transfer of expertdemonstrations, but they use KL matching to perform both high and low-level imitation.
In contrast to thesemethods, the unsupervised TCC alignment in this work avoids performing explicit domain adaptation orpixel-level image translation by instead learning a robust and invariant feature space in a fully offline fashion.
Inspired by maximum entropy inverse RL [35, 36] and generative adversarial networks [37], the seminal
GAIL algorithm [38] performs distribution matching between the expert and policy’s state-action occupancyvia an adversarial formulation; a discriminator is trained with on-policy samples, which is then used as areward in an RL framework.
Many recent works [39, 40, 41, 42, 43] build upon GAIL in order to performobservation-only imitation learning using state-only occupancy matching [39], domain adaptation viadomain confusion [40, 41], and state-alignment using a variational autoencoder next-state predictor [42].
Likewise, the algorithm proposed in [15] combines a metric learning loss that uses temporal video coherence to learn a robust skill representation with an entropy-regularized adversarial skill transfer loss.
Finally, the authors of [44] propose an adversarial formulation for learning across domains with dynamics,embodiment and viewpoint mismatch.
In contrast to these methods, our unsupervised reward is robust todomain shift without requiring online fine-tuning or the additional complexity of dynamic reward learning.
Reinforcement learning with demonstrations.
Recent work in offline-reinforcement learning [6]explicitly tackles the problem of policy embodiment and domain shift.
Their method, Reinforcement
Learning from Videos (RLV), uses a labelled collection of expert-policy state pairs in conjunction with adversarial training to learn an inverse dynamics model jointly optimized with the policy.
In contrast, weavoid the limitation of collecting human-labeled dense state correspondences by using a self-supervisedalgorithm (i.e., TCC [7]) which uses cycle-consistency to automatically learn the correspondence betweenstates of two domains.
We also show that this formulation improves generalization to unseen embodiments.
Since the problem setup is similar to ours, we also compare to their method as a baseline.
Appendix B X-MAGICAL Benchmark Details
In this section, we provide more information about our X-MAGICAL benchmark, including a descriptionof the task, an overview of the different embodiments and details regarding horizons, the success metricand the environment reward.
We encourage the reader to read [24] for an in-depth description of the base
MAGICAL benchmark.
B.1 Action and Observation Space
[Figure]
We use a continuous action space for our Sweeping task.
All embodiments have a 2D action space with theexception of the gripper agent which has an additional degree of freedom to open and close its arms.
Thefirst degree of freedom is for longitudinal movement (forward/backward), the second degree of freedomis for angular movement (left and right rotation) and the third degree of freedom, if applicable, is a grippingaction (push fingers closed/allow fingers to open).x-MAGICAL provides both state and image-based observations.
The state observations are used as inputto the RL policies whereas the pixel observations are used by the pretrained encoder to generate rewards.
The state vector contains the ( x,y ) position of the agent, (cos θ, sin θ ) where θ is the agent’s 2D orientation,and for each of the three debris: its ( x,y ) position, its distance to the agent and its distance to the goalzone.
For the RGB image, we employ an allocentric, top-down perspective (Figure 8) with full view ofthe workspace.
Similar to MAGICAL, we use an 8 Hz control rate, thus with a frame stacking value of3 , this corresponds to roughly 0 .
3 seconds of interaction.
B.2 Detailed Task and Embodiment Description
In the Sweeping task, the robotic agent must push 3 debris blocks to the pink zone at the top of the workspace.
The agent’s position is constrained to always spawn below the debris.
Both the agent’s position and thedebris positions are randomized at every environment reset.
Specifically, we sample the same y-coordinatefor all three debris, then randomly space them out from each other (different x-coordinate).
The horizon forthe longstick agent is H =50 time steps since it can solve the task much faster than the other embodimentsthanks to its morphology.
The horizon for all other embodiments is H =100 .
The ground-truth environmentreward is defined as 13 · (cid:80)
i
=
{ d i ∈ G } , i.e., the fraction of total debris present inside the goal zone G .
In Figure 9, we provide a film strip demonstration of each embodiment solving the Sweeping task witha plot of the environment reward as a function of time.
For this visualization specifically, we manuallyteleoperate each agent and disable the environment horizon limit.
B.3 Demonstration Data
To collect demonstration data for each embodiment, we train an oracle policy with SAC using theground-truth environment reward.
We then rollout the policy in the environment, discarding any potentiallyunsuccessful demonstrations, until we are left with 1000 demonstrations per embodiment.
A comprehensiveoverview of the hyperparameters used for reinforcement learning are detailed in Table 7 in Appendix G.2.
[Figure]
d
Appendix C X-REAL: A Real-World Cross-Embodiment Dataset
[Table]
To test reward learning in the real world on more challengingmanipulation tasks, we collect a real-world dataset named
[Table]
X-REAL (Cross-embodiment Real-world demonstrations),which contains 93 demonstration videos of different embod-
[Table]
iments (manifested as different manipulator end-effectors)
[Table]
solving the same manipulation task in the real-world: transferring five pens to two cups consecutively .
This is a multi-stepmanipulation task where the pens on the table need to be liftedto one cup and then moved again to a separate cup.
The different end-effectors consist in a human hand as well as sixtools purchased from Amazon and displayed in Figure 10.
In contrast to the dataset used in Section 5.3, there is visualdiversity among the different end-effectors in X-REAL, andthere is also significant variation in how the task is solved and how long it takes to solve the task.
Someend-effectors (e.g., tweezers) can only move one pen at a time, while others (e.g., human hand with fivefingers) can move all pens at once.
Additionally, the demonstrations are not collected in a constrainedfashion that tries to mimic the robot.
We report the mean and standard deviation of demonstration lengthsfor each embodiment in Table 2.
Note the variation in demo lengths across different embodiments.
[Figure]
C.1 Data Collection & Hardware Setup
The hardware and data collection setup is shown in Figure 11.
We use a GoPro Hero8 mounted on a tripodto record the demonstrations and use voice commands to efficiently start and stop the video recordings.
The Hero8 records RGB images with a resolution of 1920 × 1080 at 30 frames per second.
C.2 Reward Learning from Real-world Multi-Cross-Embodiment Demonstrations
a
[Figure]
c
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
In this section, we learn reward functions from videos in X-REAL,and demonstrate that our method is capable of handling the visualcomplexities of the real world without requiring annotations ofend-effectors, objects, or their states.
We train the encoder onall embodiments in the training set and present examples of thelearned XIRL rewards on video demonstrations from the validationset in Figure 12.
Specifically, we visualize two embodiments: the
RMS Grabber Reacher (top row) and the human 1 Hand 5 Fingers(bottom row) and for each embodiment, we show both a successfuland unsuccessful trajectory.
In the top row, both the successful and unsuccessful demonstrationsfollow a similar trajectory at the start of the task execution.
Thesuccessful one nets a high reward for placing the pens consecutivelyinto the mug then into the glass cup, while the unsuccessful oneobtains a low reward because it drops the pens outside the glass cuptowards the end of the execution.
In the bottom row, for the 1 Hand5 Fingers embodiment, we observe that not completing the taskand more specifically, leaving the pens in the first cup, generatesa reward that is roughly half (image row 2, plot orange curve) theone achieved by a successful execution (image row 1, plot bluecurve).
These results are encouraging – they show that our learned encoder can represent fine-grainedvisual differences relevant to the task.
Additionally, the training process for this visual reward did notrequire any additional environment instrumentation (apart from a camera), a desirable property for scalingto more complex, multi-step manipulation tasks.
[Figure]
Appendix D Additional Experimental Details
Our codebase is implemented in PyTorch [45].
Experiments were performed on a desktop machine with an
AMD Ryzen 7 2700X CPU (8 Cores/16 Threads, 3.7GHz base clock), 32GB RAM, and a single NVIDIA
GeForce RTX 2080 Ti GPU.
D.1 Representation Learning
[Table]
Each representation learning run – specifically training and evaluating a representation and computing thefinal goal embedding vector – took an average of 25 - 30 minutes of wall clock time.
D.1.1 Additional Baseline Details
In this section, we provide a comprehensive overview of the baselines briefly described in Section 4.2.
ImageNet : We use an ImageNet pre-trained ResNet-18 with no additional training, i.e., we load thepre-trained weights, discard the classification head, and use the 512 -dimensional embedding space fromthe penultimate layer.
[Figure]
Goal Classifier : A common strategy for using a learned visual
[Table]
reward function is to train a goal frame classifier [27] ona binary classification task where the last frame of all thedemonstrations is considered positive and all the others areconsidered negatives.
Since demonstration sequences do notnecessarily end exactly when the task is solved, we randomlysample the negatives from frames that are at least 5 indicesaway from the final frame.
This was chosen heuristically byexamining the dataset.
Note that unlike [27], we do not haveaccess to any unsuccessful trajectories in the dataset.
[Figure]
LIFS : We re-implement the approach from [16], which learnsa feature space that is invariant to different embodiments usinga loss function that encourages corresponding pairs of embodiment states across demonstrations to be close in the embedding
[Table]
space.
We use the time-based alignment method described intheir paper to find these corresponding pairs which assumes
[Table]
each embodiment performs the task at the same rate.
To prevent the embeddings from collapsing to a constant value, theyuse an additional reconstruction loss to encourage the networkdecoders to preserve as much domain-invariant information as possible.
We found early stopping to becrucial in preventing LIFS from collapsing to trivial embeddings.
For this baseline, we also randomly
[Table]
sample N evenly-spaced frames from a video to construct each mini-batch.
[Table]
TCN : We re-implement the single-view variant of Time-
[Table]
Contrastive Network (TCN) [5] with positive and negative
[Table]
frame windows of 1 and 4 respectively.
We use the n-pairloss [46] to implement TCN as done in [47].
Different from[5], we do not use a time-indexed reward which is not applicable to agents with different embodiments.
Like XIRL, we
[Table]
use the negative distance to the average goal embedding as reward.
For this baseline, we randomly sample N evenly-spacedframes from a video to construct each mini-batch.
[Table]
D.1.2 Data Augmentation & Preprocessing
We apply data augmentation during training using the Albumentations library [48].
Concretely, thisinvolves the following transformations:
• RandomResizedCrop : This transformation takes a random crop from the original image, thenresizes it to a final output height and width.
The lower and upper bounds for the random crop
area are set to [0 .
, 1 . 0] , the lower and upper bounds for the random aspect ratio of the crop areset to [0 . 75 , 1 . 33] and the final output size is set to 224 × 224 .
We apply this transformation witha probability of 1 . 0 and denote it by C .
• ColorJitter :
This transformation varies the brightness, contrast, saturation and hue of an inputimage.
Our parameters for these respective changes are 0 . 4 , 0 . 4 , 0 . 1 and 0 .
We apply it with aprobability of 0 . 8 and denote it by J .
• Normalize : Lastly, we divide the pixel values by 255 to scale their range to [0 , 1] .
We apply itwith a probability of 1 . 0 and denote it by N .
We apply the same randomly sampled transformations to all the sampled frames from the same video.
However, independently sampled transformations are applied for each such frame stack in the mini-batch.
Note that the order of transformation matters, i.e., we apply the following composed transform: N ◦ J ◦ C .
D.1.3 Training and Evaluation
All our representations are trained using an ADAM optimizer with β 1 =0 .
99 , β 2 =0 .
999 and weight decayof 10 − 5 .
We use a plateau learning rate scheduler: an initial learning rate of 10 − 5 is decayed by a factor of2 if the validation loss fails to decrease for 10 consecutive epochs.
While the representations are evaluated on the downstream policy learning performance, we also computethe following quantitative metrics and qualitative results on the train and validation sets to diagnose ourrepresentations:
• Kendall’s Tau :
A metric ranging from [ − 1 , 1] that measures how well-aligned two sequences
are in time.
We refer the reader to [7] for a more in-depth explanation.
• Nearest-Neighbor Alignment Video : We randomly select one demonstration as a referencevideo.
We use nearest-neighbor in the embedding space to align a test video with the referencevideo.
See Appendix F and the supplemental video for example visualizations for both same andcross -embodiment settings.
These videos highlight how well the embedding space encodes the
task progress across different embodiments.
D.1.4 Other Hyperparameters
For a comprehensive list of hyperparameters used for representation learning on the X-MAGICALenvironment, the Puck Pushing environment, and X-REAL (Appendix C), see Appendix G.1.
D.2 Policy Learning
The SAC [30] implementation we use is based off of [49].
Each run on the X-MAGICAL benchmark,i.e., the training and evaluation of a specific reward learning method on a specific embodiment with asingle seed, took an average of 00 h 27 m, 01 h 42 m, 03 h 56 m and 03 h 56 m wall clock times for long-stick ,medium-stick , short-stick and gripper , respectively.
Note the difference in run times is due to the fact thateach embodiment is trained for a different number of total training steps since each converges at differentrates: 75 k, 225 k, 500 k and 500 k for long-stick , medium-stick , short-stick and gripper respectively.
For the
Puck Pushing experiments in Section 5.3, each run took an average wall clock time of 01h25m for 200 ktimesteps.
Note that run times for both the above environments are recorded while performing up to 5 seedruns in parallel.
D.2.1 Soft-Actor Critic Architecture
We use clipped double Q-learning [50, 51] for the critic, where each Q -function is parameterized bya 3-layer multi-layer perceptron (MLP) with ReLU activations.
The actor is implemented as a tanhdiagonal-Gaussian, and is also parameterized by a 3-layer MLP which outputs mean and covariance.
Both actor and critic MLPs have a hidden size of 1024 – the weights are initialized with orthogonal [52]initialization, while the biases are initialized to zero.
D.3 Policy Input
As mentioned in Section 4.1 and Appendix B.1, we construct our observational inputs by stacking 3consecutive state vectors.
The input to the policy is thus a flattened vector in R 48 .
For the Puck Pushingenvironment described in Section 5.3, we construct the observational input by stacking 3 consecutive statevectors containing the 3D Cartesian coordinates of the robot end-effector and the planar 2D coordinates of
D.3.1 Training and Evaluation Setup
We first collect 5000 seed observations with a uniform random policy, after which we sample actions usingthe SAC policy.
We then perform one gradient update every time we receive a new environment observation.
When evaluating our agent every 5000 steps, we take the mean policy output (i.e., no sampling) and averagethe final success rate over 50 evaluation episodes.
D.4 Other Hyperparameters
For a comprehensive list of hyperparameters used for policy learning on the X-MAGICAL and Puck
Pushing environments, see Appendix G.2.
Appendix E Additional ExperimentsE.1 XIRL vs. Ground-truth Environment Reward
In this section, we compare the performance of our method XIRL against the ground-truth environmentreward on the same-embodiment experiment from Section 5.1 and the cross-embodiment experiment from
Section 5.2.
We observe that in both same-embodiment (Figure 13, top ) and cross-embodiment (Figure 13,bottom ) settings across multiple embodiments, XIRL is either more or just as sample-efficient as theenvironment reward.
This highlights XIRL’s ability to provide denser reward information via encoding taskprogress, as opposed to the sparser ground-truth environment reward, which was shown in Figure 6(a)
[Figure] [Figure]
In this section
[Figure]
we compare the performance of XIRL against SimCLR [53]
[Figure]
a constrative pretrainingtechnique that has exhibited SOTA self-supervised performance on ImageNet.
We implement two SimCLR baselines: (a) SimCLR is a ResNet18 trained on the x-MAGICAL demonstration dataset with theconstrastive pretraining pipeline described in [53]
[Figure]
and (b) SimCLR ImageNet is a ResNet18 pretrained on
ImageNet with SimCLR, with no further pretraining on x-MAGICAL.
Below, we present results for thelongstick and mediumstick embodiments of the x-MAGICAL benchmark:
[Figure]
We find that the SimCLR objective performs poorly when trained solely on the x-MAGICAL dataset,whereas the SimCLR baseline pertained on ImageNet (without finetuning on x-MAGICAL) does muchbetter on the longstick embodiment (which is easier to solve).
For the mediumstick embodiment, bothperform poorly.
XIRL, shown in blue, performs significantly better.
This highlights that overall, both visualpretraining on cross-embodiment demonstrations and the inductive biases offered by the TCC loss arerequired to obtain good performance on downstream RL tasks.
[Figure]
Appendix F Qualitative ResultsF.1 t-SNE Visualizations
To better understand and compare our learned representations, we visualize the t-SNE projection of thelearned XIRL and Goal Classifier embedding spaces for 4 video demonstrations of the shortstick agent in
Figure 16.
We observe that:
• Trajectories for different demonstrations overlap and are well-aligned in the XIRL embeddingspace.
In contrast, there is significantly less structure in the Goal Classifier space.
[Figure]
• Distances to the goal (top left corner in the top figure) correlate well to task progress.
[Figure]
We also provide video visualizations of the t-SNE embeddings in the supplementary video, which highlight
[Figure]
some more properties.
[Figure]
F.2 Nearest-Neighbor Retrieval
[Figure]
We provide nearest-neighbor alignment videos in the supplementary video.
[Figure]
F.3 Other Visualizations
[Figure]
Please see the supplemental video for videos showing trained policy rollouts on the Sweeping and Puck
Pushing tasks for the cross-embodiment setting, as well as interactive visualizations of the learned rewardfor the above environments and X-REAL.
Appendix G Hyperparameters
In this section, we give a comprehensive overview of the hyperparameters used for representation learningand policy learning on the Sweeping task from X-MAGICAL, the Puck Pushing task from [6], and
X-REAL.
G.1 Representation Learning
We used mostly the same hyperparameters to train the XIRL encoder across all environments.
The mainparameters that vary are the embedding dimension and the number of sampled frames.
[Table]
G.2 Policy Learning
Most hyperparameters used for downstream reinforcement learning are identical across X-MAGICALand Puck Pushing .
What changes is the total number of training steps for each embodiment, since someconverge much faster than others.
[Table]
−
[Table]
Entropy
[Table] [Table]