Watch and Match: Supercharging Imitation with Regularized Optimal Transport
Abstract:
Imitation learning holds tremendous promise in learning policies efﬁciently for complex decision making problems.
Current state-of-the-art algorithms often use inverse reinforcement learning (IRL), where given a set of expert demonstrations, an agent alternatively infers a reward function and the associated optimal policy.
However, such IRL approaches often require substantial online interactions for complex control problems.
In this work, we present Regularized Optimal Transport (ROT), a new imitation learning algorithm that builds on recent advances in optimal transport based trajectory-matching.
Our key technical insight is that adaptively combining trajectory-matching rewards with behavior cloning can signiﬁcantly accelerate imitation even with only a few demonstrations.
Our experiments onvisual control tasks across the DeepMind Control Suite, the OpenAI Robotics Suite, and the Meta-World Benchmark demonstrate an average of 7 . 8 × faster imitation to reach 90% of expert performance compared to prior state-of-the-art methods.
On real-world robotic manipulation, with just one demonstration and an hour of online training, ROT achieves an average success rate of 90.1% across 14 tasks.
Introduction
Imitation Learning (IL)
[1, 2, 3] has a rich history that can be categorized across two broad paradigms,
Behavior Cloning (BC) [1] and Inverse Reinforcement Learning (IRL) [4].
BC uses supervised learning to obtain a policy that maximizes the likelihood of taking the demonstrated action given an observation in the demonstration.
While this allows for training without online interactions, it suffers from distributional mismatch during online rollouts [5].
IRL, on the other hand, infers the underlying reward function from the demonstrated trajectories before employing RL to optimize a policy through online environment rollouts.
This results in a policy that can robustly solve demonstrated tasks even in the absence of task-speciﬁc rewards [6, 7].
Although powerful, IRL methods suffer from a signiﬁcant drawback – they require numerous expensive online interactions with the environment.
There are three reasons for this: (a) the inferred reward function is often highly non-stationary, which compromises the learning of the associated behavior policy [7]; (b) even when the rewards are stationary, policy learning still requires effective exploration to maximize rewards [8]; and (c) when strong priors such as pretraining with BC are applied to accelerate policy learning, ensuing updates to the policy cause a distribution shift that destabilizes training [9, 10].
Combined, these issues manifest themselves on empirical benchmarks, where IRL methods have poor efﬁciency compared to vanilla RL methods on hard control tasks [11].
In this work, we present Regularized Optimal Transport (ROT) for imitation learning, a new method that is conceptually simple, compatible with high-dimensional observations, and requires minimal additional hyperparameters compared to standard IRL approaches.
To demonstrate the effectiveness of ROT, we run extensive experiments on 20 simulated tasks across DM Control [12], OpenAI
Robotics [13], and Meta-world [14], and 14 robotic manipulation tasks on an xArm (see Fig. 3).
Our main ﬁndings are summarized below.
ROT outperforms prior state-of-the-art imitation methods, reaching 90% of expert performance 7 . 8 × faster than our strongest baselines on simulated visual control benchmarks.
On real-world tasks, with a single human demonstration and an hour of training, ROT achieves an average success rate of 90.1% with randomized robot initialization and image observations.
This is signiﬁcantly higher than behavior cloning (36.1%) and adversarial IRL (14.6%).
ROT exceeds the performance of state-of-the-art RL trained with rewards, while coming close to methods that augment RL with demonstrations (Appendix I.3).
Unlike standard RL methods, ROT does not require hand-speciﬁcation of the reward function.
Ablation studies demonstrate the importance of every component in ROT, particularly the role that soft Q-ﬁltering plays in stabilizing training and the need for OT-based rewards during online learning (Appendix I.2 & Appendix I.4).
Open-source code and demonstration data will be publicly released on our project website.
Videos of our trained policies can be seen here: rot-robot.github.io .
Regularized Optimal Transport
A fundamental challenge in imitation learning is to balance the ability to mimic demonstrated actions along with the ability to recover from states outside the distribution of demonstrated states.
Behavior Cloning (BC) specializes in mimicking demonstrated actions through supervised learning, while Inverse Reinforcement Learning (IRL) specializes in obtaining policies that can recover from arbitrary states.
Regularized Optimal Transport (ROT) combines the best of both worlds by adaptively combining the two objectives.
This is done in two phases.
In the ﬁrst phase, a randomly initialized policy is trained using the BC objective on expert demonstrated data.
This ‘BC-pretrained’ policy then serves as an initialization for the second phase.
In the second phase, the policy is allowed access to the environment where it can train using an IRL objective.
To accelerate the IRL training, the
BC loss is added to the objective with an adaptive weight.
Details of each component are described below, with additional algorithmic details in Appendix D.
2.1 Phase 1: BC Pretraining
BC corresponds to solving the maximum likelihood problem shown in Eq.
Here T e refers to expert demonstrations.
When parameterized by a normal distribution with ﬁxed variance, the objective can be framed as a regression problem where, given inputs s e , π BC needs to output a e .
[Equation]
After training, it enables π BC to mimic the actions corresponding to the observations seen in the demonstrations.
However, during rollouts in an environment, small errors in action prediction can lead to the agent visiting states not seen in the demonstrations [5].
This distributional mismatch often causes π BC to fail on empirical benchmarks [15, 11] (see Fig.
4 (a) in Appendix C).
2.2 Phase 2: Online Finetuning with IRL
Given a pretrained π BC model, we now begin online ‘ﬁnetuning’ of the policy π b ≡ π ROT in the environment.
Since we are operating without explicit task rewards, we use rewards obtained through OT-based trajectory matching, which is described in Appendix A. These OT-based rewards r OT enable the use of standard RL optimizers to maximize cumulative reward from π b ≡ π ROT .
In this work we use n-step DDPG [16], a deterministic actor-critic based method that provides high-performance in continuous control [8].
Finetuning with Regularization π BC is susceptible to distribution shift due to accumulation of errors during online rollouts [5] and directly ﬁnetuning π BC also leads to subpar performance (refer to Fig.
4 in Appendix C).
To address this, we build upon prior work in guided RL [15] and ofﬂine
RL [9], and regularize the training of π ROT by combining it with a BC loss as seen in Eq.
[Equation]
[Figure]
[Figure]
ode
[Figure]
a r d
[Figure]
ss
[Figure]
e
[Figure]
and_place
[Figure]
ss
[Figure]
metaworld
[Figure]
door
[Figure]
open
[Figure]
Here, Q ( s, a ) represents the Q-value from the critic which is optimized using OT-based rewards during the actor-critic policy optimization.
α is a ﬁxed weight, while λ ( π ) is a policy-dependent adaptive weight that controls the contributions of the two loss terms.
D β refers to the replay buffer for online rollouts.
Adaptive Regularization with Soft Q-ﬁltering While prior work [15, 17] use hand-tuned schedules for λ ( π ) , we propose a new adaptive scheme that removes the need for tuning.
This is done by comparing the performance of the current policy π ROT and the pretrained policy π BC on a batch of data sampled from the replay buffer for online rollouts D β .
More precisely, given a behavior policy π BC ( s ) , the current policy π ROT ( s ) , the Q-function Q ( s, a ) and the replay buffer D e , we set λ as:
[Equation]
The strength of the BC regularization hence depends on the performance of the current policy with respect to the behavior policy.
This ﬁltering strategy is inspired by Nair et al.
[18], where instead of a binary hard assignment we use a soft continuous weight.
Experimental comparisons with hand-tuned decay strategies are presented in Appendix I.2.
Considerations for image-based observations Since we are interested in using ROT with highdimensional visual observations, additional machinery is required to ensure compatibility.
Following prior work in image-based RL and imitation [8, 11], we perform data augmentations on visual observations and then feed it into a CNN encoder.
Similar to Cohen et al.
[11], we use a target encoder with Polyak averaging to obtain representations for OT reward computation.
This is necessary to reduce the non-stationarity caused by learning the encoder alongside the ROT imitation process.
Further implementation details and the training procedure can be found in Appendix D.
Experiments
Our experiments are designed to answer the following questions: (a) How efﬁcient is ROT for imitation learning?
(b) How does ROT perform on real-world tasks?
Additional results and details about simulated and robot tasks have been provided in Appendix I.
3.1 How efﬁcient is ROT for imitation learning?
Performance of ROT for image-based imitation is depicted on select environments in Fig.
To reach 90% of expert performance, ROT is on average 8 .
7 × faster on DeepMind Control tasks, 2 . 1 × faster on OpenAI Robotics tasks, and 8 . 9 × faster on Meta-world tasks.
We also ﬁnd that the improvements of ROT are most apparent on the harder tasks, which are shown in Fig.
Appendix I.1 shows results on all 20 simulated tasks, along with experiments that exhibit similar improvements in state-based settings.
[Figure]
3.2 How does ROT perform on real-world tasks?
We devise a set of 14 manipulation tasks on our xArm robot to compare the performance of ROT with BC and our strongest baseline RDAC, an adversarial IRL method that combines DAC [7] with our pretraining and regularization scheme.
The BC policy is trained using supervised learning on a single expert demonstration collected by a human operator.
ROT and RDAC ﬁnetune the pretrained
BC policy through 1 hour of online training, which amounts to ∼ 6k environment steps.
Since there is just one demonstration, our tasks are designed to have random initializations but ﬁxed goals.
Note that a single demonstration only demonstrates solving the tasks from one initial condition.
Evaluation results across 20 different initial conditions can be seen in Fig.
We observe that ROT has an average success rate of 90.1% over 20 evaluation trajectories across all tasks as compared to 36.1% for BC and 14.6% for RDAC.
The poor performance of BC can be attributed to distributional mismatch due to accumulation of error in online rollouts and different initial conditions.
The poor performance of RDAC can be attributed to slow learning during the initial phase of training.
More detailed evaluations of RDAC on simulated environments is present in Appendix I.4.
Conclusion and Limitations
In this work, we have proposed a new imitation learning algorithm, ROT, that demonstrates improved performance compared to prior state-of-the-art work on a variety of simulated and robotic domains.
However, we recognize a few limitations in this work: (a) Since our OT-based approach aligns agents with demonstrations without task-speciﬁc rewards, it relies on the demonstrator being an ‘expert’.
Extending ROT to suboptimal, noisy and multimodal demonstrations would be an exciting problem to tackle.
(b) Performing BC pretraining and BC-based regularization requires access to expert actions, which may not be present in some real-world scenarios particularly when learning from humans.
Recent work on using inverse models to infer actions given observational data could alleviate this challenge [19].
(c) On robotic tasks such as Peg in box (hard) and Pressing a switch from Fig. 2, we ﬁnd that ROT’s performance drops substantially compared to other tasks.
This might be due to the lack of visual features corresponding to the task success.
For example, in the ‘Peg’ task, it is visually difﬁcult to discriminate if the peg is in the box or behind the box.
Similarly for the ‘Switch’ task, it is difﬁcult to discern if the button was pressed or not.
This limitation can be addressed by integrating more sensory modalities such as additional cameras, and tactile sensors in the observation space. Acknowledgments
We thank Ben Evans, Anthony Chen, Ulyana Piterbarg and Abitha Thankaraj for valuable feedback and discussions.
This work was supported by grants from Honda, Amazon and ONR awards
N00014-21-1-2404 and N00014-21-1-2758.
A Background
Reinforcement Learning (RL)
We study RL as a discounted inﬁnite-horizon Markov Decision
Process (MDP) [20, 21].
For pixel observations, the agent’s state is approximated as a stack of consecutive RGB frames [22].
The MDP is of the form ( O , A , P, R, γ, d 0 ) where O is the observation space, A is the action space, P : O×A → ∆( O ) is the transition function that deﬁnes the probability distribution over the next state given the current state and action, R :
O × A → R is the reward function, γ is the discount factor and d 0 is the initial state distribution.
The goal is to ﬁnd a policy π : O → ∆( A ) that maximizes the expected discount sum of rewards E π [Σ ∞ t =
γ t R ( o t , a t )] , where o 0 ∼ d 0 , a t ∼ π ( o t ) and o t +1 ∼ P ( .
| o t , a t ) .
Imitation Learning (IL)
The goal of imitation learning is to learn a behavior policy π b given access to either the expert policy π e or trajectories derived from the expert policy T e .
While there are a multitude of settings with differing levels of access to the expert [23], this work operates in the setting where the agent only has access to observation-based trajectories, i.e. T e ≡ { ( o t , a t )
Tt =0 } Nn =0 .
Here N and T denotes the number of trajectory rollouts and episode timesteps respectively.
We choose this speciﬁc setting since obtaining observations and actions from expert or near-expert demonstrators is feasible in real-world settings [24, 25] and falls in line with recent work in this area [26, 6, 7].
Inverse Reinforcement Learning (IRL) IRL [4, 27] tackles the IL problem by inferring the reward function r e based on expert trajectories T e .
Then given the inferred reward r e , policy optimization is used to derive the behavior policy π b .
Prominent algorithms in IRL [7, 6] requires alternating the inference of reward and optimization of policy in an iterative manner, which is practical for restricted model classes [27].
For compatibility with more expressive deep networks, techniques such as adversarial learning [6, 7] or optimal-transport [28, 26, 11] are needed.
Adversarial learning based approaches tackle this problem by learning a discriminator that models the gap between the expert trajectories T e and behavior trajectories T b .
The behavior policy π b is then optimized to minimize this gap through gap-minimizing rewards r e .
Such a training procedure is prone to instabilities since r e is updated at every iteration and is hence non-stationary for the optimization of π b .
Optimal Transport for Imitation Learning (OT)
To alleviate the non-stationary reward problem with adversarial IRL frameworks, a new line of OT-based approaches have been recently proposed [28, 26, 11].
Intuitively, the closeness between expert trajectories T e and behavior trajectories T b can be computed by measuring the optimal transport of probability mass from T b → T e .
During policy learning, the policy π φ encompasses a feature preprocessor f φ which transforms observations into informative state representations.
Some examples of a preprocessor function f φ are an identity function, a mean-variance scaling function and a parametric neural network.
In this work, we use a parametric neural network as f φ .
Given a cost function c : O × O → R deﬁned in the preprocessor’s output space and an OT objective g , the optimal alignment between an expert trajectory o e and a behavior trajectory o b can be computed as
[Equation]
[Equation] where M = { µ ∈ R T × T : µ 1 = µ T 1 = 1 T 1 } is the set of coupling matrices and the cost c can be the Euclidean or Cosine distance.
In this work, inspired by [11], we use the entropic Wasserstein distance with cosine cost as our OT metric, which is given by the equation
[Equation]
(5) where the cost matrix C t,t (cid:48)
=
c ( f φ ( o b ) , f φ ( o e )) .
Using Eq.
5 and the optimal alignment µ ∗ obtained by optimizing Eq. 4, a reward signal can be computed for each observation using the equation
[Equation]
Intuitively, maximizing this reward encourages the imitating agent to produce trajectories that closely match demonstrated trajectories.
Since solving Eq.
4 is computationally expensive, approximate solutions such as the Sinkhorn algorithm [29, 28] are used instead.
Actor-Critic based reward maximization Given rewards obtained through OT computation, efﬁcient maximization of the reward can be achieved through off-policy learning [7].
In this work, we use
Deep Deterministic Policy Gradient (DDPG) [16] as our base RL optimizer which is an actor-critic algorithm that concurrently learns a deterministic policy π φ and a Q-function Q θ .
However, instead of minimizing a one step Bellman residual in vanilla DDPG, we use the recent n-step version of
DDPG from Yarats et al.
[8] that achieves high performance on visual control problems.
B Related Work
Imitation Learning (IL) IL [30] refers to the setting where agents learn from demonstrations without access to environment rewards.
IL can be broadly categorized into Behavior Cloning (BC) [1, 23] and Inverse Reinforcement Learning (IRL) [4, 27].
BC solely learns from ofﬂine demonstrations but suffers on out-of-distributions samples [5] whereas IRL focuses on learning a robust reward function through online interactions but suffers from sample inefﬁciency [7].
Deep IRL methods can be further divided into two categories: (1) adversarial learning [31] based methods, and (2) state-matching [32, 33] based methods.
GAIL [6] is an adversarial learning based formulation inspired by maximum entropy IRL [34] and GANs [31].
There has been a signiﬁcant body of work built up on GAIL proposing alternative losses [35, 36, 37], and enhancing its sample efﬁciency by porting it to an off-policy setting [7].
There have also been visual extensions of these adversarial learning approaches [38, 39, 40, 11].
However, although adversarial methods produce competent policies, they are inefﬁcient due to the non-stationarity associated with iterative reward inference [11].
Optimal Transport (OT) OT [32, 33] is a tool for comparing probability measures while including the geometry of the space.
In the context of IL, OT computes an alignment between a set of agent
[Figure]
and expert observations using distance metrics such as Sinkhorn [41], Gromov-Wasserstein [42],
GDTW [43], CO-OT [44] and Soft-DTW [45].
For many of these distance measures, there is an associated IL algorithm, with SIL [28] using Sinkhorn, PWIL [26] using greedy Wasserstein,
GDTW-IL [43] using GDTW, and GWIL [46] using Gromov-Wasserstein.
Recent work from Cohen et al.
[11] demonstrates that the Sinkhorn distance [28] produces the most efﬁcient learning among the discussed metrics.
They further show that SIL is compatible with high-dimensional visual observations and encoded representations.
Inspired by this, ROT adopts the Sinkhorn metric for its
OT reward computation, and improves upon SIL through adaptive behavior regularization.
Behavior Regularized Control Behavior regularization is a widely used technique in ofﬂine
RL [47] where explicit constraints are added to the policy improvement update to avoid bootstrapping on out-of-distribution actions [48, 49, 50, 51, 52, 53].
In an online setting with access to environment rewards, prior work [15, 10] has shown that behavior regularization can be used to boost sample efﬁciency by ﬁnetuning a pretrained policy via online interactions.
For instance, Jena et al.
[17] demonstrates the effectiveness of behavior regularization to enhance sample efﬁciency in the context of adversarial IL.
ROT builds upon this idea by extending to visual observations, OT-based IL, and adaptive regularization, which leads to improved performance (see Appendix I.4).
We also note that the idea of using adaptive regularization has been previously explored in RL [18].
However, ROT uses a soft, continuous adaptive scheme, which on initial experiments provided signiﬁcantly faster learning compared to hard assignments.
C Challenges in Online Finetuning from a Pretrained Policy
In this section, we study the challenges with ﬁnetuning a pretrained policy with online interactions in the environment.
Fig.
4 illustrates a task where an agent is supposed to navigate the environment from the top left to the bottom right, while dodging obstacles in between.
The agent has access to a single expert demonstration, which is used to learn a BC policy for the task.
Fig.
4 (a) shows that this
BC policy, though close to the expert demonstration, performs suboptimally due to accumulating errors on out-of-distribution states during online rollouts [5].
Further, Fig. 4 (b) uses this BC policy as an initialization and naively ﬁnetunes it with OT rewards (described in Section ?? ).
Such naive ﬁnetuning of a pretrained policy (or actor) with an untrained critic in an actor-critic framework exhibits a forgetting behavior in the actor, resulting in performance degradation as compared to the pretrained policy.
This phenomenon has also been reported by Nair et al.
[9] and we provide a detailed discussion in Appendix C.1.
In this paper, we propose ROT which addresses this issue by adaptively keeping the policy close to the behavior data during the initial phase of ﬁnetuning and reduces this dependence over time.
Fig.
4 (c) demonstrates the performance of our approach on such ﬁnetuning.
It can be clearly seen that even though the BC policy is suboptimal, our proposed adaptive regularization scheme quickly improves and solves the task by driving it closer to the expert demonstration.
In Fig.
4 (d), we demonstrate that even if the agent was initialized at points outside the expert trajectory, the agent is still able to learn quickly and complete the task.
This generalization to starting states would not be possible with regular BC.
C.1 Issue with Fine-tuning Actor-Critic Frameworks
In this paper, we use n -step DDPG proposed by Yarats et al.
[8] as our RL optimizer for actorcritic based reward maximization.
DDPG [16] concurrently learns a deterministic policy π φ using deterministic policy gradients (DPG) [54] and a Q-function Q θ by minimizing a n-step Bellman residual (for n-step DDPG).
For a parameterized actor network π φ ( s ) and a critic function Q θ ( s, a ) , the deterministic policy gradients (DPG) for updating the actor weights is given by [Equation]
[Equation]
Here, ρ β refers to the state visitation distribution of the data present in the replay buffer at time t .
From Eq.
7, it is clear that the policy gradients in this framework depend on the gradients with respect to the critic value.
Hence, as mentioned in [9, 10], naively initializing the actor with a pretrained policy while using a randomly initialized critic results in the untrained critic providing an exceedingly poor signal to the actor network during training.
As a result, the actor performance drops immediately and the good behavior of the informed initialization of the policy gets forgotten.
In this paper, we propose an adaptive regularization scheme that permits ﬁnetuning a pretrained actor policy in an actor-critic framework.
As opposed to Rajeswaran et al.
[15], Jena et al.
[17] which employ on-policy learning, our method is off-policy and aims to leverage the sample efﬁcient characteristic of off-policy learning as compared to on-policy learning [7].
D Algorithmic Details
Algorithm 1 ROT:
Regularized Optimal Transport
Require:
Expert [Equation] , a t )
Tt =0 } Nn =
Pretrained policy π BC
Replay buffer D , Training steps T , Episode Length L
Task environment env
Parametric networks for RL backbone (e.g., the encoder, policy and critic function for DrQ-v2)
A discriminator D for adversarial baselines
Algorithm: π ROT ← π BC (cid:46) Initialize with pretrained policy for each timestep t
= 1 ...
T do if done then r 1: L
= rewarder OT (episode) (cid:46) OT-based reward computation
Update episode with r 1: L and add ( o t , a t , o t +1 , r t ) to D o t = env.reset () , done =
False , episode =
[ ] end if a t = π ROT ( o t )
o t +1 , done = env.step ( a t ) episode.append ([ o t , a t , o t +1 ])
Update backbone-speciﬁc networks and reward-speciﬁc networks using D end for
D.1 Implementation
Algorithm 1 describes our proposed algorithm, Regularized Optimal Transport (ROT), for sample efﬁcient imitation learning for continuous control tasks.
Further implementation details are as follows:
Algorithm and training procedure Our model consists of 3 primary neural networks - the encoder, the actor and the critic.
During the BC pretraining phase, the encoder and the actor are trained using a mean squared error (MSE) on the expert demonstrations.
Next, for ﬁnetuning, weights of the pretrained encoder and actor are loaded from memory and the critic is initialized randomly.
We observed that the performance of the algorithm is not very sensitive to the value of α and we set it to 0 .
03 for all experiments in this paper.
A copy of the pretrained encoder and actor are stored with ﬁxed weights to be used for computing λ ( π ) for soft Q-ﬁltering.
Actor-critic based reward maximization We use a recent n-step DDPG proposed by Yarats et al. [8] as our RL backbone.
The deterministic actor is trained using deterministic policy gradients (DPG) [54] given by Eq. 7.
The critic is trained using clipped double Q-learning similar to Yarats et al.
[8] in order to reduce the overestimation bias in the target value.
This is done using two
Q-functions, Q θ 1 and Q θ 2 .
The critic loss for each critic is given by the equation
[Equation]
(8) where D β is the replay buffer for online rollouts and y is the target value for n-step DDPG given by
[Equation]
Here, γ is the discount factor, r is the reward obtained using OT-based reward computation and ¯ θ 1 , ¯ θ 2 are the slow moving weights of target Q-networks.
Target feature processor to stabilize OT rewards The OT rewards are computed on the output of the feature processor f φ which is initialized with a parametric neural network.
Hence, as the weights of f φ change during training, the rewards become non-stationary resulting in unstable training.
In order to increase the stability of training, the OT rewards are computed using a target feature processor f φ (cid:48) [11] which is updated with the weights of f φ every T update environment steps.
For state-based observations, f φ corresponds to a ’trunk’ network which is a single layer neural network.
For pixel-based observations, f φ includes DrQ-v2’s encoder followed by the ’trunk’ network.
D.2 Hyperparameters
The complete list of hyperparameters is provided in Table 1.
Similar to Yarats et al.
[8], there is a slight deviation from the given setting for the Walker Stand/Walk/Run task from the DeepMind
Control suite where we use a mini-batch size of 512 and a n -step return of 1.
E Environments
Table 2 lists the different tasks that we experiment with from the DeepMind Control suite [12, 55],
OpenAI Robotics suite [56] and the Meta-world suite [57] along with the number of training steps and the number of demonstrations used.
For the tasks in the OpenAI Robotics suite, we ﬁx the goal while keeping the initial state randomized.
No modiﬁcations are made in case of the DeepMind
Control suite and the Meta-world suite.
The episode length for all tasks in DeepMind Control issteps, for OpenAI Robotics is 50 steps and Meta-world is 125 steps (except bin picking which runs for 175 steps).
F Demonstrations
For DeepMind Control tasks, we train expert policies using pixel-based DrQ-v2 [8] and collectdemonstrations for each task using this expert policy.
The expert policy is trained using a stack of 3 consecutive RGB frames of size 84 × 84 with random crop augmentation.
Each action in the environment is repeated 2 times.
For OpenAI Robotics tasks, we train a state-based DrQ-v2 with hindsight experience replay [58] and collect 50 demonstrations for each task.
The state representation comprises the observation from the environment appended with the desired goal location.
For this, we did not do frame stacking and action repeat was set to 2.
For Meta-World tasks, we use a single expert demonstration obtained using the task-speciﬁc hard-coded policies provided in their open-source implementation [57].
[Table]
G Robot Tasks
In this section [Table] we describe the suite of manipulation experiments carried out on a xArm robot in this paper.
(a) Door Close: Here, the robot arm is supposed to close an open door by pushing it to the target.
(b) Hang Hanger: While holding a hanger between the grippers, the robot arm is initialized at a random position and is tasked with putting the hanger at a goal region on a closet rod.
(c) Erase Board: While holding a board duster between the grippers, the robot arm is tasked with erasing markings drawn on the board while being initialized at a random position.
(d) Reach: The robot arm is required to reach a speciﬁc goal after being initialized at a random position.
(e) Hang Mug: While holding a mug between the grippers, the robot arm is initialized at a random position and is tasked with hanging the mug on a speciﬁc hook.
Suite
Tasks
Allowed Steps # Demonstrations
DeepMind Control
Acrobot Swingup
[Table]
Cartpole Swingup
Cheetah Run
Finger Spin
Hopper Stand
Hopper Hop
Quadruped Run
Walker Stand
Walker Walk Walker Run
OpenAI Robotics
Fetch Reach 1 .
5 ×
Fetch Push
Fetch Pick and Place
Meta-World
Hammer 1 ×
Drawer Close
Door Open
Bin Picking
Button Press Topdown
Door Unlock.
xArm Robot
Close Door 6 ×
Hang Hanger
Erase Board
Reach
Hang Mug
Hang Bag
Turn Knob
Stack Cups
Press Switch
[Table]
(f) Hang Bag: While holding a tote between the grippers, the robot arm is initialized at a random position and is tasked with hanging the tote bag on a speciﬁc hook.
(g) Turn Knob: The robot arm is tasked with rotating a knob placed on the table by a certain angle after being initialized at a random position.
We consider a 90 degree rotation as success.
C l o s e
D oo r
H a ng
H a ng e r
E r a s e
B o a r d
P ou r
O p e n
B ox
P e g ( H a r d )
P e g ( M e d )
P e g ( E a s y )
P r e ss
S w it
c
h
S t ac k
C up s
T u r n
K nob
H a ng
B a g
H a ng
M ug
R eac h
T u r n K nob
H a ng B a g
P e g ( H a r d )
E r a s e B o a r d
O p e n B ox
(h) Stack Cups: While holding a cup between the gripper, the robot arm is required with stacking it on another cup placed on the table.
(i) Press Switch: With the gripper kept closed, the robot arm is required to press a switch (with an LED light) placed on the table.
(j) Peg (Easy, Medium,
Hard): The robot arm is tasked with inserting a peg, hanging by a wire, into a bucket placed on the table.
This task has 3 variants - Easy, Medium, Hard - with the size of the bucket decreasing from Easy to Hard.
(k) Box Open:
In this task, the robot arm is supposed to open the lid of a box placed on the table by lifting a handle provided in the front of the box.
(l) Pour: While holding a cup containing some item (in our case, almonds), the robot arm is supposed to move towards another cup placed on the table and pour the item into this cup.
Evaluation procedure For each task, we obtained a set of 20 random initializations and evaluate all of the methods (BC, RDAC and ROT) over 20 trajectories from the same set of initializations.
These initializations are different for each task based on the limits of the observation space for the task.
H Baselines
Throughout the paper, we compare ROT with several prominent imitation learning and reinforcement learning methods.
Here, we give a brief description of each of the baseline models that have been used.
(a) Expert:
For each task, the expert refers to the expert policy used to generate the demonstrations for the task (described in Appendix F).
(b) Behavior Cloning (BC): This refers to the behavior cloned policy trained on expert demonstrations.
(c) Adversarial IRL (DAC): Discriminator Actor Critic [7] is a state-of-the-art adversarial imitation learning method [6, 37, 7].
Since DAC outperforms prior work such as GAIL[6] and AIRL[35], it serves as our primary adversarial imitation baseline.
(d) State-matching IRL (OT): Sinkhorn Imitation Learning [28, 26] is a state-of-the-art statematching imitation learning method [59] that approximates OT matching through the Sinkhorn Knopp algorithm.
Since ROT is derived from similar OT-based foundations, we use SIL as our primary state-matching imitation baseline.
(e) RDAC:
This is the same as ROT, but instead of using state-matching IRL (OT), adversarial IRL (DAC) is used.
(f) Finetune with ﬁxed weight:
This is similar to ROT where instead of using a time-varying adaptive weight
λ
( i ) , only the ﬁxed weight λ 0 is used.
λ 0 is set to a ﬁxed value of 0 . 03 .
(g) Finetune with ﬁxed schedule:
This is similar to ROT that uses both the ﬁxed weight λ 0 and the time-varying adaptive weight λ 1 ( i ) .
However, instead of using Soft Q-ﬁltering to compute λ 1 ( i ) , a hand-coded linear decay schedule is used. (h) DrQ-v2 (RL): DrQ-v2 [8] is a state-of-the-art algorithm for pixel-based RL.
DrQ-v2 is assumed to have access to environment rewards as opposed to ROT which computes the reward using OT-based techniques. (i) Demo-DrQ-v2: This refers to DrQ-v2 but with access to both environment rewards and expert demonstrations.
The model is initialized with a pretrained BC policy followed by RL ﬁnetuning with an adaptive regularization scheme like ROT.
During RL ﬁnetuning, this baseline has access to environment rewards. (j) BC+OT: This is the same as the OT baseline but the policy is initialized with a pretrained BC policy.
No adaptive regularization scheme is used while ﬁnetuning the pretrained policy. (k) OT+BC Reg.
: This is the same as the OT baseline with randomly initialized networks but during training, the adaptive regularization scheme is added to the objective function.
I Experimental Results
Our experiments are designed to answer the following questions: (a) How efﬁcient is ROT for imitation learning?
(b) How does ROT perform on real-world tasks?
(c) How important is the choice of IRL method in ROT?
(d) Does soft Q-ﬁltering improve imitation? (e) How does ROT compare to standard reward-based RL?
Simulated tasks We experiment with 10 tasks from the DeepMind Control suite [12, 55], 3 tasks from the OpenAI Robotics suite [56], and 7 tasks from the Meta-world suite [57].
For DeepMind
Control tasks, we train expert policies using DrQ-v2 [8] and collect 10 demonstrations for each task using this policy.
For OpenAI Robotics tasks, we train a state-based DrQ-v2 with hindsight experience replay [58] and collect 50 demonstrations for each task.
For Meta-world tasks, we use a single hard-coded expert demonstration from their open-source implementation [57].
Full environment details can be found in Appendix E and details about the variations in demonstrations and initialization conditions can be found in Appendix F.
Robot tasks Our real world setup for each of the 14 manipulation tasks can be seen in Fig.
We use an Ufactory xArm 7 robot with a xArm Gripper as the robot platform for our real world experiments.
However, our method is agnostic to the speciﬁc robot hardware.
The observations are RGB images from a ﬁxed camera.
In this setup, we only use a single expert demonstration collected by a human operator with a joystick and limit the online training to a ﬁxed period of 1 hour.
Descriptions of each task and the evaluation procedure is in Appendix G.
Primary baselines We compare ROT with baselines against several prominent imitation learning methods.
While a full description of our baselines are in Appendix H, a brief description of the two strongest ones are as follows:
Adversarial IRL (DAC): Discriminator Actor Critic [7] is a state-of-the-art adversarial imitation learning method [6, 37, 7].
DAC outperforms prior work such as GAIL [6] and AIRL [35], and thus it serves as our primary adversarial imitation baseline.
Trajectory-matching IRL (OT): Sinkhorn Imitation Learning [28, 26] is a state-of-the-art trajectory-matching imitation learning method [59] that approximates OT matching through the Sinkhorn Knopp algorithm
Since ROT is derived from similar OT-based foundations, we use SIL as our primary state-matching imitation baseline.
I.1 How efﬁcient is ROT for imitation learning?
In addition to the results provided in Sec. 3.1, Fig. 7 and Fig. 8 shows the performance of ROT for pixel-based imitation on 10 tasks from the DeepMind Control suite, 3 tasks from the OpenAI
Robotics suite and 7 tasks from the Meta-world suite.
On all but one task, ROT is signiﬁcantly more sample efﬁcient than prior work.
Finally, the improvements from ROT hold on state-based
[Figure]
p i s o d
[Figure]
_ r e
[Figure]
a r d
[Figure]
e p
i s o d e _
r e w a
[Figure]
d
[Figure]
i
[Figure]
o d e _ r e w a r d
[Figure]
BC
OT
DAC
ROT (Ours) observations as well(see Fig. 9).
Table 3 provides a comparison between the factor of speedup of
ROT to reach 90% of expert performance compared to prior state-of-the-art [7, 11] methods.
I.2 Does soft Q-ﬁltering improve imitation?
To understand the importance of soft Q-ﬁltering, we compare ROT against two variants of our proposed regularization scheme: (a) A tuned ﬁxed BC regularization weight (ignoring λ ( π ) in Eq. 2); (b) A carefully designed linear-decay schedule for λ ( π ) , where it varies from 1 . 0 to 0 .
in the ﬁrst 20k environment steps [15].
As demonstrated in Fig.
10, ROT is on par and in some cases exceeds the efﬁciency of a hand-tuned decay schedule, while not having to hand-tune its regularization weights.
We hypothesize this improvement is primarily due to the better stability of adaptive weighing as seen in the signiﬁcantly smaller standard deviation on the Meta-world tasks.
[Figure]
I.3 How does ROT compare to standard reward-based RL?
We compare the performance of ROT against DrQ-v2 [8], a state-of-the-art algorithm for image-based
RL.
As opposed to the reward-free setting ROT operates in, DrQ-v2 has access to environments rewards.
The results in Fig.
11 show that ROT handily outperforms DrQ-v2.
This clearly demonstrates the usefulness of imitation learning in domains where expert demonstrations are available over reward-based RL.
We also compare against a demo-assisted variant of DrQ-v2 agent using the same pretraining and regularization scheme as ROT.
Interestingly, we ﬁnd that our soft Q-ﬁltering based regularization can accelerate learning of RL with task rewards, which can be seen in the high performance of the demo-assisted variant of DrQ-v2.
I.4 How important are the design choices in ROT?
Importance of pretraining and regularizing the IRL policy Fig.
12 compares the following variants of ROT on set of pixel-based tasks: (a) Training the IRL policy from scratch (OT); (b)
Finetuning a pretrained BC policy without BC regularization (BC+OT); (c) Training the IRL policy
[Figure]
i
[Figure]
o d
[Figure]
_ r
[Figure]
w a r d dmc_hopper_
[Figure]
i
[Figure]
o d e _ r e w
[Figure]
r d
[Figure]
u cc
[Figure]
_ p e r c e n t a g e
[Figure]
BC
OT
DAC
ROT [Equation] [Figure] from scratch with BC regularization (OT+BC Reg.).
We observe that pretraining the IRL policy (BC+OT) does not provide a signiﬁcant difference without regularization.
This can be attributed to the ‘forgetting behavior’ of pre-trained policies, studied in Nair et al.
Interestingly, we see that even without BC pretraining, keeping the policy close to a behavior distribution (OT+BC Reg.) can yield improvements in efﬁciency over vanilla training from scratch.
Our key takeaway from these experiments is that both pretraining and BC regularization are required to obtain sample-efﬁcient imitation learning.
[Table]
>
[Table]
[Table] [Table]
>
[Table]
OpenAI
[Table]
and Place
[Table]
> [Table] [Table] > 1M
[Table]
DAC)
[Table]
Close
[Table]
> 1M
[Table]
OT)
[Table]
> 1M > 1M (OT)
[Table]
>
[Table]
[Table]
>
[Table] [Table] [Table] > 1M
[Table]
OT) [Table] > 1M (OT)
Choice of IRL method In ROT, we build on OT-based IRL instead of adversarial IRL.
This is because adversarial IRL methods require iterative reward learning, which produces a highly nonstationary reward function for policy optimization.
In Fig.
13, we compare ROT with adversarial
IRL methods that use our pretraining and adaptive BC regularization technique (RDAC).
We ﬁnd that our soft Q-ﬁltering method does improve prior state-of-the-art adversarial IRL (RDAC vs. DAC in Fig. 13).
However, our OT-based approach (ROT) is more stable and on average leads to more efﬁcient learning.
Choice of Q-ﬁltering method In ROT, we adopt a soft Q-ﬁltering method as opposed to the hard assignment strategy proposed by Nair et al.
Fig.
14 shows a comparison between the performance of soft Q-ﬁltering and hard Q-ﬁltering.
We observe that though the two strategies have comparable performance in most cases, soft Q-ﬁltering exhibits better sample efﬁciency and more stable training in some tasks.
This justiﬁes our choice of opting for soft Q-ﬁltering as opposed to hard assignment.
[Figure]
ss
[Figure]
en
[Figure]
age
[Figure]
age
[Figure]
Expert
BC
Finetune with ﬁxed weight
Finetune with ﬁxed
[Figure]
ROT (
[Figure] [Figure]
[Figure] [Figure]
[Equation]
[Figure] [Figure]
_door
[Figure] [Figure]