Primal Wasserstein Imitation Learning
Authors: Robert Dadashi
Leonard Hussenot
Matthieu Geist
Olivier Pietquin
Venue: ICLR 2021 Poster (2021)
Keywords: Reinforcement Learning
Inverse Reinforcement Learning
Imitation Learning
Optimal Transport
Wasserstein distance

Abstract: Imitation Learning (IL) methods seek to match the behavior of an agent with that of an expert. In the present work, we propose a new IL method based on a conceptually simple algorithm: Primal Wasserstein Imitation Learning (PWIL), which ties to the primal form of the Wasserstein distance between the expert and the agent state-action distributions. We present a reward function which is derived offline, as opposed to recent adversarial IL algorithms that learn a reward function through interactions with the environment, and which requires little fine-tuning. We show that we can recover expert behavior on a variety of continuous control tasks of the MuJoCo domain in a sample efficient manner in terms of agent interactions and of expert interactions with the environment. Finally, we show that the behavior of the agent we train matches the behavior of the expert with the Wasserstein distance, rather than the commonly used proxy of performance.

Published as a conference paper at ICLR 2021
PRIMAL WASSERSTEIN IMITATION LEARNING
Robert Dadashi") Léonard Hussenot :2, Matthieu Geist', Olivier Pietquin' ! !Google Research, Brain Team
Univ. de Lille, CNRS, Inria Scool, UMR 9189 CRISIAL
ABSTRACT
Imitation Learning (II.) methods seek to match the behavior of an agent with that of an expert. In the present work, we propose a new IL method based on a conceptually simple algorithm: Primal Wasserstein Imitation Learning (PWIL), which ties to the primal form of the Wasserstein distance between the expert and the agent state-action distributions. We present reward function which derived offline, as opposed to recent adversarial IL algorithms that learn a reward function through interactions with the environment, and which requires little fine-tuning. We show that we can recover expert behavior on a variety of continuous control tasks of the MuJoCo domain in a sample efficient manner in terms of agent interactions and of expert interactions with the environment. Finally, show that the behavior of the agent we train matches the behavior of the expert with the Wasserstein distance, rather than the commonly used proxy of performance.
INTRODUCTION
Reinforcement Learning (RL) has solved number of difficult tasks whether in games Tesauro 1995 Mnih et al.) 2015 Silver et al. 2016) or robotics (Abbeel & Ng. 2004 /Andyychowicz et al. 2020). However, RL relies the existence of a reward function, that can be either hard to specify too sparse to be used in practice. Imitation Learning (IL) is paradigm that applies to these environments with hard to specify rewards: we seek solve a task by learning a policy from a fixed number of demonstrations generated by an expert.
IL methods can typically be folded into two paradigms: Behavioral Cloning, or BC (Pomerleau 1991
Bagnell et al. 2007 Ross & Bagnell 2010) and Inverse Reinforcement Learning, or IRL (Russell 1998 Ng et al. (2000.. In BC, we seek to recover the expert's behavior by directly learning a policy that matches the expert behavior in some sense. In IRI., we assume that the demonstrations come from an agent that acts optimally with respect to an unknown reward function that we seek to recover, to subsequently train agent it. Although IRL methods introduce intermediary problem (i.c. recovering the environment's reward) they are less sensitive to distributional shift (Pomerleau [199]) they generalize environments with different dynamics (Piot et al.] 2013), 3 and they can recover near-optimal agent from suboptimal demonstrations (Brown et al. 2019 J Jacq et al. 2019).
However, IRL methods are usually based on iterative process alternating between reward estimation and RI. which might result in poor sample-efficient Earlier IRL. methods (Ng et al. 2000) Abbeel &
Ng 2004 Ziebart et al.) 2008) require multiple calls a Markov decision process solver Puteeran 2014), whereas recent ad adversarial IL approaches (Finn et al. 2016 Ho & Ermon, 2016 Fu et al.
2018 interleave the learning of the reward function with the learning process of the agent. Adversarial
IL methods are based on an adversarial training paradigm similar to Generative Adversarial Networks (GANs) (Goodfellow et al., 2014). where the learned reward function can be thought of as the confusion of a discriminator that learns to differentiate expert transitions from non expert ones. These methods are well suited to the IL problem since they implicitly minimize an f-divergence between the state-action distribution of an expert and the state-action distribution of the learning agent (Ghasemipour et al., 2019; Ke et al. 2019).. However the interaction between a generator (the policy) and the discriminator (the reward function) makes it minmax optimization problem. and therefore comes with practical challenges that might include training instability, sensitivity to hyperparameters and poor sample efficiency.
"Correspondence Robert Dadashi: dadashi @goog1 1..oge.ccmm com.

---

Published as a conference paper at ICLR 2021
In this work, the expert and of the agent. Contrary to f-divergences, the Wasserstein distance is a true distance, The Wasserstein distance has gained popularity in GAN approaches (Arjovsky et al., 2017) through its dual formulation which comes with challenges (see Section [ Our approach is novel in the fact that we consider the problem of minimizing the Wasserstein distance through its primal formulation. Crucially, the primal formulation prevents the minmax optimization problem, and requires little fine tuning.
the Wasserstein distance as measure between the state-action distributions of
is smooth and it is based on the geometry of the metric space it operates
We introduce a reward function computed offline based on an upper bound of the primal form of the Wasserstein distance. As the Wasserstein distance requires a distance between state-action pairs, we show that it can be hand-defined for locomotion tasks, and that it can be learned from pixels for hand manipulation task. The inferred reward function is non-stationary, like adversarial II. methods, but it is not re-eva re-evaluated as the agent interacts with the environment, therefore the reward function we define is computed offline. We present true distance to compare the behavior of the expert and the behavior of the agent, rather than using the common proxy of performance with respect to the true return the task we consider (as it is unknown in general), Our method recovers expert behaviour comparably to existing ofte-of-the--ee methods while being based on significantly fewer hyperparameters; operates even in the extreme low data regime of demonstrations, and is the first method that makes Humanoid run with single (subsampled) demonstration.
2 BACKGROUND AND NOTATIONS
Markoy decision processes. We describe environments as episodic Markov Decision Processes (MDP) with finite time horizon (Sutton & Barto, 2018) AS,A, Y,pp, Po. T),, where S is the state space, A is the action space, P is the transition kernel, r is the reward function, 7 is the discount factor, Po is the initial state distribution and T is the time horizon. We will denote the dimensionality of S and A as |S] and |A| respectively. A policy is mapping from states to distributions over actions; we denote the space of all policies by I1. RL, the goal is learn policy π² that maximizes the expected sum of discounted rewards it encounters, that is, the expected return. Depending on the context, we might use the concept of a cost e rather than reward r (Puterman 2014), which essentially moves the goal of the policy from maximizing its return minimizing its cumlative cost.
State action distributions. Suppose a policy æ visits the successive states and actions $1. 01......... ST, or during an episode, we define the empirical state-action distribution Pa as: P. = Ei=i where 5 a. is Dirac distribution centered on (st, az). Similarly, suppose have a set of expert demonstrations D = {s", a²) of size D, then the associated empirical expert state-action distribution Pe is defined as: P.. = E(S.)ED 8s,a-
Wasserstein distance. Suppose we have the metric space (M,d) where M is a set and d is a metric on M. Suppose we have μ and y two distributions on M with finite moments, the p-th order Wasserstein distance (Villani 2008) is defined as W (14, v) = inf MMXM d(x, d0 y). where ( v) is the set of all couplings between j and v. In the remainder, we only consider distributions with finite support. A coupling between two distributions of support cardinal T and D a doubly stochastic matrix of size T x D. We note the set of all doubly stochastic matrices of size T D:
Viv [1: T7. Σ e -= =BER TXD VVj [1 DJ. Σ A ..vii
P = T
D
The Wasserstein distance between distributions of state-action pairs requires the definition of metric in the space AS. Defining metric in an MDP is non trivial (Ferns et al. 2004; Mahadevan & Maggioni 2007); we show example where the metric is learned from demonstrations in Section For now, we assume the existence of a metric d d: AS) x (S,A) + R
3 METHOD
We present the theoretical motivation of approach: the minimization of the Wasserstein distance between the state-action distributions of the agent and the expert. We introduce a reward based on upper-bound of the primal form of the Wasserstein distance inferred from relaxation of
2

---

Published as a conference paper at ICLR 2021
the optimal coupling condition, and present the resulting algorithm: Primal Wasserstein Imitation Learning (PWIL).
3.1 WASSERSTEIN DISTANCE MINIMIZATION
Central our approach is the minimization of the Wasserstein distance between the state-action distribution of the policy we seek to train P and the state-action distribution of the expert Pe. In other words, we aim optimizing the following problem:
D ΣΣ .(( (sf.
T
inf (Px. inf inf TEl
(I)
TETI BEB
im fm
In the of the paper, we only consider the 1-Wasserstein (p 1 in Equation D and leave the extensive study of the influence of the order for future work. We can interpret the Wasserstein distance using the earth's 's movers analogy (Villami 2008). Consider that the state-action pairs of the expert are D holes of mass D and that the acaton pairs of the policy are piles of dirt of mass T-1. A coupling / a transport strategy between the piles of dirt and the holes, where f(i,) j] stands for how much of the pile of dirt i should be moved towards the hole j. The optimal coupling is the one that minimizes the distance that the earth mover travels to put all piles of dirt to holes. Note that to compute the optimal coupling, context of RL, this means having access to the full trajectory generated by n. From now on, we write 87 as the optimal coupling for the policy x, that inject in Equation ①
need knowledge of the locations of all piles of dirt. In the
0% = arg min
d((s7 (69.)))))))))))) a)))
dEG
1=1
inf Wi (p: Po) wEll
inf TEll
Σ
G with G.a
Σ a((S)) a)) (3)
In Equation we have introduced ci,x. which we interpret as a cost to minimize using RL. As Gx depends on the optimal coupling 65. we can only define ci, at the very end of an episode. This can be problematic if an agent learns in an online manner or in large time-horizon tasks. Thus, we introduce an upper bound to the Wasserstein distance that yields a cost we can compute online, based on suboptimal coupling strategy.
3.2 GREEDY COUPLING
In this section we introduce the greedy coupling 0% E O, defined recursively for 1 <
i S T as:
06 arg min o j=1
Σ
d((s), a.),, (sf, of)
(3)
with θ 00/i, RF
Σ 0 - vk vke
DJ. Σ 0 k
1 D
constraint(a)
constraint(b)
Similarly to Equation Equation can be interpreted using the earth mover's 's analogy. Contrary to Equation (T) where we assume knowledge of the positions of the T piles of dirts, we now consider that they appear sequentially, and that the earth's 's mover needs to transport the new pile of dirt to holes right when it appears. To do so, we derive the distances to all holes and move the new pile of dirt to the closest remaining available holes, hence the greedy nature of it. In Equation (3), the constraint (a) means that all the dirt that appears at the i-th timestep needs to be moved, and the constraint (6) means that cannot fill the holes more than their capacity D. We show the difference between the greedy coupling and the optimal coupling in Figure] and provide the pseudo-code to derive it in Algorithm
3

---

Published as a conference paper at ICLR 2021
S
S
81
c
$3
82
$3
81
5g
$2
82
$3
a)
b)
Figure I: Illustration of the difference between the a) greedy coupling and the b) optimal coupling. We present MDP where drop the dependency metric associated is the Euclidean distance. We note the states visited by the policy: sf s3 sg and the states visited by the expert: si, sq. sg. When the policy encounters the state s². and because we do not know sg yet, the greedy coupling strategy consists in coupling it with sz although the optimal coupling strategy would be to couple it with s.. Note that the total cost (omitting the constant coupling multiplication factor) with the greedy coupling is 7 whereas the total cost with the optimal coupling This highlights that the optimal coupling needs knowledge about the whole policy's trajectory to be derived.
the action. The state space is IR and the
We can now define an upper bound of the Wasserstein distance using the greedy coupling (since by definition is suboptimal):
T
inf Wi Wi inf sell
Σ
d((s) a (5),,
xell
i=1 ;=
T
S inf ΣΣ #Ell ;=1j=1
d((<)]
(5))))))))))
(4)
In Section 4 we empirically validate this bound. We infer a cost from Equation ①1 at each timestep i:
C/ = Dd((s a.), a(), (3))))
(5)
Note that the greedy coupling [t, defined in Equation (3) depends on all the previous state-actions visited by the policy π, which makes the cost cf.- non-stationary, similarly to adversarial IL methods. We can infer a reward from the cost, T;, = f(c],)) where f a monotonically decreasing function. This reward function is an episodic history dependent reward function r(sg. ag...., St, ax) where r is uniquely defined from expert demonstrations (hence is said to be derived offline). Crucially, have defined a reward function that we seek to maximize, without introducing an inner minimization problem (which is the case with adversarial IL approaches).
We can now derive an algorithm building on this reward function (Algorithm) D The algorithm presented is written using pseudo-code for generic agent A which implements policy aA. observes tuples (s, a, &') and possibly updates its components. The runtime of the algorithm for single reward step computation has complexity O([[S) (Al)D D²) (see Appendix E.2 for details), and therefore has complexity ((([S) (A) D2) for computing rewards across the entire episode. Note that in the case where T = D, computing the greedy coupling is simply a lookup of the expert state-action pair that minimizes the distance with the agent state-action pair, followed by a pop-out of this minimizing expert state-action pair from the set of expert demonstrations.
Algorithm I inntroduces an agent with the capability to observe and update itself, which is directly in- spired from ACME's 's formalism of agents (Hoffman et al. (2020). and is general enough to characterize a wide range of agents. We give an example rundown the algorithm in Section
4 EXPERIMENTS
In this section, we present the implementation of PWIL and perform an empirical evaluation, based on the ACME framework (Hoffman et al. 2020). We test PWIL on MuJoCo locomotion tasks and compare it to the state-of-the-an Gill-like algorithm DAC (Kostrikov et al.) 2019) and against the common BC baseline. As DAC is based on TD3 (Fujimoto et al.) 2018) which is a variant of DDPG-based agent for fair comparison: D4PG (Barth
DDPG (Lillicrap et al.) 2016),, we Maron et al., 2018). We also provide a proof-of-concee experiment on a door opening task where the demonstrations do not include actions and are pixel-based and therefore the metric of the
4

---

Published as a conference paper at ICLR 2021
Algorithm 1 PWIL: Primal Wasserstein Imitation Learning
Input:
Expert demonstrations D (5}, aj [e(::))) Agent A, number of episodes N for = to N do
Copy expert demonstrations with weight: D' ag wz JE(I:D) with wef 1/
Reset environment, initial state 8
1 to I' do
Take action := (s),, observe next state s' Initialize weight u= : 1.
T. cost c==00:0
while nT > 0 do
Compute s a", w == arg (ss,,,,,))) d(((s,)) (89,
if tx 2 w" then
c: c w wdd(d((s a), (sf,a))) a²))
wº = wx we
DD' D'ppops a², ,²)
else
c: w= w= (8 (s))
w" = w w
w =
r= [cc)
A observes (s, a, r, 8'), updates itself
MDP has to be learned. In this setup. we use SAC the direct RL method. We answer the following questions: 1) Does PWIL recover expert behavior? 2) How sample efficient is PWIL? 3) Does PWIL actually minimize the Wasserstein distance between the distributions of the agent and the expert? 4) Does PWIL extend to visual based observations where the definition of an MDP metric is not straightforward? We also show the dependence of PWIL on multiple of its components through an ablation study. The complete description of the experiments is given in Appendix A We provide experimental code and videos of the trained agents here: https //sites.gcoqle...
7/sites wwo. com
wwtp:wwwwccccceceeeooe.. wetm:/view/nnne www.ee/wwwssestteeee www.immitaioon
4.1 IMPLEMENTATION
We test PWIL following the same evaluation protocol as GAIL and DAC, on six environments from the OpenAl Gym MuloCo suite (Todorov et al.) 2012: Brockman et al., 2016): Reacher-v2, Ant-v2, Humanoid-v2, Walker2d-v2, HalfCheetah-v2 and Hopper-v2. For each environment, multiple demonstrations gathered using D4PG agent trained on the actual reward of these environments (which was chosen as it is a base agent in the ACME framework). We subsample demonstrations by factor of 20: select
out of every 20 transitions with random offset at the beginning of the episode (expect for Reacher for which we do not subsample since the episode length 50). The point of subsampling is simulate low data regime, that a pure behavioral cloning approach suffers the behavioral drift and hence does not perform to the level of the expert. We run experiments with multiple number of demonstrations: {1, 4, 11} (consistently with the evaluation protocol of DAC).
Central to our method is metric between state-action pairs. We use the standardized Euclidean distance which is the L2 distance on the concatenation of the observation and the action, weighted along each dimension by the inverse standard deviation of the expert demonstrations. From the cost e; defined in Equation (5). we define the following reward:
3T
r; = a exp( VISI+
(6)
acts as normalizer For all environments, we use = 5 and B = 5. The scaling factor VISI+AA
T
on the dimensionality of the state and action spaces and the time horizon of the task. We pick f: x + cxp(-2) as the monotonically decreasing function.
We test our method against state-of-the-art imitation learning algorithm: DAC and BC. DAC is based on GAIL: it discriminates between expert and non-expert states and uses an off-policy algorithm TD3 which is variant of DDPG instead of the on-policy algorithm TRPO (Schulman
5

---

Published as a conference paper at ICLR 2021
et al. 2015). We used the soun-s code provided by the authors. The Humanoid environment is not studied in the original DAC paper. The hyperparameters the authors provided having poor performance, we led an exhaustive search detailed in Section A.2 and present results for the best hyperparameters found.
For method, with different direct RL agents in Appendix We initialize the replay buffer by filling with expert state-action pairs with maximum reward a (however, makes it an approximation to replaying the actual trajectories). We train PWIL and DAC in the limit of 1M environment interactions (2.5M for Humanoid) on 10 seeds. Every 10k environment steps, we perform 10-episode rollouts per seed of the policy without exploration noise and report performance with respect to the environment's original reward in Figure
D4PG agent (see details in Appendix A and provide additional experiments
that we used subsampled trajectories which
4.2 RESULTS
In Figure PWIL shows improvement on final performance for Hopper, Ant, HalfCheetah and Humanoid over DAC, similar performance for Walker2d, and is worse on Reacher. On Humanoid, although we ran an exhaustive hyperparameter search (see Section (.2) DAC has poor performance, thus exemplifying the brittleness of GAIL-like approaches. On the contrary, PWIL has near-optimal performance, even with a single demonstration.
In terms of sample efficiency, DAC outperforms PWIL on HalfCheetah, Ant and Reacher and has similar performance for Hopper and Walker2d. Note that ablated versions of PWIL are more sample- efficient than DAC HalfCheetah and Ant (Section B). This seems rather contradictory with the claim that GAll--lke methods have poor sample efficiency. However the reward function defined in DAC requires careful tuning (e.g. architecture, optimizer, regularizers, learning schedule) which demands experimental cycles and therefore comes with a large number of environment interactions.
Remarkably, PWIL is the first method able consistently learn policies Humanoid with an average score over 7000 even with a single demonstration, which means that the agent can actually run, (other works consider Humanoid is "solved" for score of 5000 wich corresponds to standing).
Walkerdd
HolCChetthh
Humanois
Reacher
2000
3xc-
5000
6200
3000
X00
4000
1000
2002
-00
1003
0.2h
5.5
0.75
3.10
0.00
023
330
0.75
1:30
000
9.2
030
312
6.5
3.2
532
0.7
1.3
200
023
0.90
1/00
Hooper
Walker2d
Ant
HolCChettah
Humanoid
Reachar
1002
5002
2500
4000
-20
2500
4000
-25
1002
1130
2300
2000
-25
Aapert
middon
165
165
165
166
1e5
1e5
Figure 2: Mean and standard deviation return of the evaluation policy over 10 rollouts and 10 seeds, reported every 10k environment steps. The return is in term of the environment's original reward. Top row: demonstration, bottom row: 11 demonstrations.
Wasserstein distance. In Figure 3 we show the Wasserstein distance to expert demonstrations throughout the learning procedure for PWIL DAC. The Wasserstein distance is evaluated the end of an episode using transportation simplex (Bonneel et al., 2011): we used the implementation from Flamary & Courty 2017). For both methods, we notice that the distance decreases while learning. PWIL leads to a smaller Wasserstein distance than DAC on all environments but HalfCheetah where similar and Reacher where it larger. This should not come as a surprise since PWIL's objective is to minimize the Wasserstein distance. PWIL defines reward from upper bound to the Wasserstein distance between the state-action distributions of the expert and the agent, that we include as well in Figure 3 Notice that upper bound is "empirically tight", which validates the choice of the greedy

---

Published as a conference paper at ICLR 2021
Hopper
Walker2d
Art
Ha/CChesthh
Humanood
Reacher
300
DAC
PMTL WJ Blund)
1500
73
390 633
0.50
0.75
1.30
0.30
925
333
975
0.00
0.25
0.50
0.75
1.33
0.00
0.25
0.30
0.75
1.30
1e5
0.00
0.5
0.0
0.55
170
Figure 3: Mean of the Wasserstein distance between the state-action distribution of the evaluation policy and the state-action distribution of the expert over 10 rollouts and 10 seeds. Agents were trained with 11 demonstrations. For PWIL, we include the upper bound on the Wasserstein distance based on the greedy coupling defined in Equation
coupling as approximation to the optimal coupling. We emphasize that this distance is useful in real settings of IL, regardless of the method used, where we cannot have access to the actual reward of the task.
4.3 ABLATION STUDY
In this section, we present the evaluation performance of PWIL in the presence of ablations and report final performances in Figure 4 The learning curves for all ablations can be found in Appendix We keep the hyperparameters of the original PWIL agent fixed.
Hoppwr
Walcer2d
Art
HalfCheetahs
Humanod
Beacher
OXT
6X1
$300
4300
2000
280
33
32
31
11
-03
11
Domonstations
Figure 4: Mean and standard deviation of the evaluation performance of PWIL and variants at the IM environment interactions mark (2.5M for Humanoid). Results are computed over 10 seeds and 10 episodes for each seed.
PWIL-state. In this version of PWIL, we do not assume that have access to the actions taken by the expert. Therefore, we try to match the state distribution of the agent with the state distribution of the expert (instead of the state-action distribution). The setup is referred to as Learning from Observation (LfO) (Toabi et al., 2018 Edwards et al.) 2019). The reward is defined similarly to PWIL, using a state distance rather than state-action distance. Note that in this version, we cannot pre-fill the replay buffer with expert state-action pairs since we do not have access to actions. Remarkably, PWIL--tate recovers non-trivial behaviors on all environments.
PII In this version of PWIL, we use the Euclidean distance between state-action pairs, rather than the standardized Euclidean distance. In other words, we do not weight the state-action distance by the inverse standard deviation of the expert demonstrations along each dimension. There is significant drop in performance for all environments but Hopper-v2. This shows that the performance of PWIL is sensitive to the quality of the MDP metric.
PWIL-nofill. In this version, we do not prefill the replay buffer with expert transitions. This leads to drop in performance which is significant for Walker and Ant. This should not come surprise since number of IL methods leverage the idea of expert transitions in the replay buffer (Reddy et al. 2020 Hester et al.) 2018).
PWIL-support. In this version of PWIL, we define the reward as a function of the following cost:
Vi E T), Ct
inf 8 ER Σ 02 2 J=1
Σ
(()
We can interpret this cost in the context of Section as the problem of moving piles of dirt of mass + into holes of infinite capacity. It is thus reminiscent of methods that consider II. support
7

---

Published as a conference paper at ICLR 2021
estimation problem (Wang et al., 2019 Brantley et al., 2020 This leads significant drop in performance for Ant, Walker and HaliChectah.
4.4 DOOR OPENING TASK: REWARD FROM PIXEL- BASED OBSERVATIONS
In this section, we show that PWIL extends to risual-based setups, where the MDP metric has to be learned. We use the door opening task from Rajeswaran et al. (2018), with human demonstrations generated with Haptix (Kumar Todorov 2015). The goal of the task is for the controlled hand to open the door (Figure ()) whose location changes at each episode. We add an early termination constraint when the door opened (hence suppressing the incentive for survival), We assume that we can only access the high-dimensional visual rendering of the demonstrations rather than the internal state (and hence have no access to actions, that is LfO setting).
The PWIL reward is therefore defined using distance between the visual rendering of the current state and the visual rendering of the demonstrations. This distance is learned offline through self-
Figure Visual rendering
of the door opening task.
supervised learning: we construct an embedding of the frames using Temporal Sycle-Consisteneyy Learning (TCC) (Dwibedi et al. 2019); TCC builds an embedding by aligning frames of videos of task coming from multiple examples (in this case the demonstrations). The distance learned is thus the L2 distance between embeddings. The distance learned on the 25 human demonstrations provided by Rajeswaran et al. (2018) which are then re-used to define the PWIL reward. From the cost C; defined in Equation (5). we define the following reward: r; a exp(-33Te) with a 5, B We used SAC (Haarnoja et al., (0188) as the forward RL algorithm, which was more robust than D4PG on the task. We train the agent on 1.5M environment steps and show that we indeed recover expert behaviour and consistently solve the task in Figure
Wasserstein Distance
Imitation Returns
40
Task success rate
1.0
100
30
0.5
50
20
human
10
Steps
Figure 6: We present the average and standard deviation of evaluation performance of PWIL the door environment, Metrics are computed every 10k environment steps and aggregated over 10 seeds and 10 rollouts. Left: the imitation returns of the agent. Center: Wasserstein distance in the embedding space to expert demonstrations. Right: Success rate of the agent, i.e. ratio of episodes where the agent manages to open the door.
5 RELATED WORK
Adversarial IL. Similarly our method, adversarial IL methods aim matching the state-action distribution of the agent with the distribution of the expert, using different measures of similarity. For instance, GAIL (Ho Ermon 2016) considers the Shannon-Jenser divergence while AIRL (Fu et al. 2018) considers the Bullback-Leible divergence. In the context of GANs, Arjovsky et al. (2017) show that replacing the fddivergence by the Wasserstein distance through the Kantorovich-Rubinser duality [Villani] 2008) leads to better training stability, which number of methods in II. have leveraged (Li et al.| 2017 Lacotte et al., 2019; Kostrikov et al. 2019 Xiao et al., 2019 Liu et al.) 2020 Pacchiano et al. 2020). However, the Kantorovich-Rubiie
duality only holds for the Wi distance (Peyré et al. 2019). its implementation comes with a number of practical approximations (e.g., weight clipping ensure Lipchitz continuity). Although these methods based the Wasserstein distance through its dual formulation, they drastically differ from ours as they rely on minmax optimization problem. DAC (Kostrikov et al. 2019) demonstrates that adversarial II.
8

---

Published as a conference paper at ICLR 2021
approaches, while based GAIL, can be sample efficient. However, they rely carefully tuned discriminator (e.g. network architecture, regularization strategy, scheduled learning rates) which requires to interact with the environment. In contrast, the reward function we present relies only two hyperparameters.
Expert support estimation, Another line of research in IL consists in estimating the state-action support of the expert and define a reward that encourages the agent to stay on the support of the expert (Piot al. 2014 [Schroecker & Isbell| 2017: Wang al. 2019 Brantley et al.) 2020 Reddy et al. 2020). Note that the agent might stay on the support of the expert without recovering its state-action distribution. Soft Q Imitation Learning (Reddy et al. 2020) assigns a reward of 1 to all expert transitions and O to all transitions the expert encounters; the method learns to recover expert behavior by balancing out the expert transitions with the agent transitions in the replay buffer of a value-based agent. Random Expert Distillation (Wang et al. 2019) estimates the support by using a neural network trained on expert transitions whose target a fixed random neural network. Disagreement Regularized Imitation Learning (Brantley et al. 2020 estimates the expert support through the variance of an ensemble of BC agents, and use a reward based on the distance to the expert support and the KL divergence with the BC policies. PWIL differs from these methods since it based on the distance to the support of the expert, with a support that shrinks through "pop-outs" to enforce the agent to visit the whole support. We showed in the ablation study that "pop-outs" are sine qua non for recovering expert behaviour (see PWIL-support).
Trajectory-Based II. Similarly to the door opening experiment, Aytar al. (2018) learn em- bedding offline and reinforces a reward based on the distance in the embedding space to a single demonstration. However, it is well suited for deterministic environments since they use a single trajectory that the agent aims at reproducing while the door experiment shows that PWIL is able to generalize the door and handle locations are picked at random. Peng al. (2018) reinforces a tem- poral reward function (the reward at time t is based on the distance of the state of the agent to the state of the expert at time 2). Contrary to PWII., the imitation reward function is added to a task- reward function; the distance to expert demonstrations is defined along handpicked dimensions of the observation space; the temporal constraint makes it hard to generalize to environment with stochastic initial states.
Offline reward estimation. Existing approaches derive a reward function without interaction with the environment (Boularias et al. 2011 Klein al. 2013 2012 Piot et al., 2016). They typically
require strong assumptions on the structure of the reward (e.g. linearity in some predefined features). PWIL also derives reward offline, without structural assumptions, however it is non-stationary since it depends on the state-action visited in the episode.
MMD-IL. Generative Moment Matching Imitation Learning (Kim & Park 2018) assumes metric on the MDP and aims at minimizing the Maximum Mean Discrepancy (MMD) between the agent and the expert. However, the MMD (like the Wasserstein distance) requires complete rollouts of the agent, and does not present a natural relaxation like the greedy coupling. GMMIL is based on an on-policy agent which makes significantly less sample-efficien than PWIL. We were not able to implement a version of MMD based on an off- oollyo agent.
6 CONCLUSION
In this work, we present Imitation Learning as distribution matching problem and introduce reward function which based on an upper bound of the Wasserstein distance between the state-action distributions of the agent and the expert. A number of II methods are developed on synthetic tasks, where the evaluation of the method can be done with the actual return of the task. We emphasize that in our work, present direct measure of similarity between the expert and the agent, that can thus use in real IL settings, that is in settings where the reward of the task cannot be specified.
The reward function we introduce is learned offline, in the sense that it does not need to be updated with interactions with the environment. It requires little tuning (2 hyperparameters) and it can recover near expert performance with as little 1 demonstration on all considered environments, even the challenging Humanoid. Finally our method extends to the harder visual-bace setting, based on a metric learned offline from expert demonstrations using self-supervised learning.

---

Published as a conference paper at ICLR 2021
ACKNOWLEDGMENTS
We thank Zafarali Ahmed, Adrien Ali Taiga, Gabriel Dulac-Arold, Johan Ferret, Alexis Jaeq and Saurabh Kumar for their feedback on an earlier version of the manuscript. We thank Debidatta Dwibedi and Ilya Kostrikov for helpful discussions.
REFERENCES
Pieter Abbeel and Andrew Ng. Apprenticeship learning via inverse reinforcement learning. In International Conference Machine Learning, 2004.
OpenAl: Marcin Andrychowicz, Bowen Baker, Maciek Chociej, Rafal Jozefowicz, Bob McGrew,
Jakub Pachocki, Arthur Petron, Matthias Plappert, Glenn Powell, Alex Ray, al. Learning dexterous in-hand manipulation. The International Journal of Robotics Research, 2020.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International Conference on Machine Learning, 2017.
Yusuf Aytar, Tobias Pfaff, David Budden, Thomas Paine, Ziyu Wang, and Nando de Freitas. Playing hard exploration games by watching youtube. In Advances in Neural Information Processing Systems, 2018.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey Hinton. Layer normalization. arXiv preprint orX1v:1607.06450. 2016.
JA Bagnell, Joel Chestnutt, David Bradley, and Nathan D Ratliff. Boosting structured prediction for imitation learning. In Advances in Neural Information Processing Systems, 2007.
Gabriel Marth-Maron, Matthew W Hoffman, David Budden, Will Dabney, Dan Horgan, Dhruva Tb, Alistair Muldal, Nicolas Heess, and Timothy Lillicrap. Distributed distributional deterministic policy gradients. International Conference on Learning Representations, 2018.
Marc G Bellemare, Will Dabney, and Rémi Munos. A distributional perspective reinforcement learning. International Conference on Machine Learning. 2017.
Nicolas Bonneel, Michiel Van De Panne, Sylvain Paris, and Wolfgang Heidrich. Displacement interpolation using lagrangian mass transport. In Proceedings of the 2011 SIGGRAPH Asia Conference, 2011.
Abdeslam Boularias, Jens Kober, and Jan Peters. Relative entropy inverse reinforcement learning. In International Conference Artificial Intelligence and Statistics, 2011.
Kianté Brantley, Wen Sun, and Mikael Henaff. Disagreement-rguularizeed imitation learning. In International Conference on Learning Representations, 2020.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym, 2016.
Daniel Brown, Wonjoon Goo, Prabhat Nagarajan, and Scott Niekum. Extrapolating beyond sub- optimal demonstrations via inverse reinforcement learning from observations. In International Conference on Machine Learning, 2019.
Djork-Arné Clevert, Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network learning by exponential linear units (elus). International Conference on Learning Representations, 2016.
Debidatta Dwibedi, Yusuf Aytar, Jonathan Tompson, Pierre Sermanet, and Andrew Zisserman. Temporal cycle-consistency learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), June 2019.
Ashley Edwards, Himanshu Sahni, Yannick Schroecker, and Charles Isbell. Imitating latent policies from observation. In International Conference on Machine Learning, 2019.
10

---

Published as a conference paper at ICLR 2021
Norm Ferns, Prakash Panangaden, and Doina Precup. Metrics for finite markov decision processes. In Conference on Uncertainty in Artificial Intelligence, 2004.
Chelsea Finn, Sergey Levine, and Pieter Abbeel. Guided cost learning: Deep inverse optimal control via policy optimization. International Conference on Machine Learning, 2016.
Rémi Flamary and Nicolas Courty. Pot python optimal transport library. Gitllub: hutp:://itthu. com/rflamary/POT. 2017.
Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adversarial inverse reinforce-
ment learning. International Conference on Learning Representations, 2018.
Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor- critic methods. In International Conference Machine Learning, 2018.
Seyed Kamyar Seyed Ghasemipour, Richard Zemel, and Shixiang Gu. A divergence minimization perspective imitation learning methods. Conference Robot Learning, 2019.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley. Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Advances in Neural Information Processing Systems, 2014.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. International Conference on Machine Learning, 20184 8a.
Tuomas Haarnoja, Aurick Zhou, Kristian Hartikainen, George Tucker, Sehoon Ha, Jie Tan, Vikash Kumar, Henry Zhu, Abhishek Gupta, Pieter Abbeel, et al. Soft actor-critic algorithms and applications. arXiv preprint 1812.00 2018b.
Todd Hester, Matej Vecerik, Olivier Pietquin, Marc Lanctot, Tom Schaul, Bilal Piot, Dan Horgan, John Quan, Andrew Sendonaris, Ian Osband, al. Deep q-learning from demonstrations, In AAAI Conference on Artificial Intelligence, 2018.
Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. In Advances in Neural
Information Processing Systems, 2016.
Matt Hoffman, Bobak Shahriari, John Aslanides, Gabriel Barth-Maron, Feryal Behbahani, Tamara Norman, Abbas Abdolmaleki, Albin Cassirer, Fan Yang, Kate Baumli, et al. Acme: research framework for distributed reinforcement learning. arXiv preprint 20Xi::2006.00999, 2020.
Alexis Jacq, Matthieu Geist, Ana Paiva, and Olivier Pietquin. Learning from a learner. In International Conference on Machine Learning, 2019.
Liyiming Ke, Matt Barnes, Wen Sun, Gilwoo Lee, Sanjiban Choudhury, and Siddhartha Srinivasa. Imitation learning as fa f-divergence minimization. arXiv preprins arXiv: /905. 12888, 2019.
Kee-Eung Kim and Hyun Soo Park. Imitation learning via kernel mean embedding. In AAAI
Conference on Artificial Intelligence, 2018.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. International
Conference on Learning Representations, 2015.
Edouard Klein, Matthieu Geist, Bilal Piot, and Olivier Pietquin. Inverse reinforcement learning through structured classification. In Advances in Neural Information Processing Systems, 2012.
Edouard Klein, Bilal Piot, Mathieu Geist, and Olivier Pietquin. A cascaded supervised learning approach to inverse reinforcement learning. In Joint European conference on machine learning and knowledge discovery in databases, 2013.
Ilya Kostrikov, Kumar Krishna Agrawal, Debidatta Dwibedi, Sergey Levine, and Jonathan Tompson. ietiimiinttor-cto--criiiii Addressing sample inefficiency and reward bias in adversarial imitation learning. International Conference on Learning Representations, 2019.
11

---

Published as a conference paper at ICLR 2021
V. Kumar and E. Todorov. Mujoco haptix: A virtual reality system for hand manipulation. In 2015 TEEE-RAS 15th International Conference on Humanoid Robots (Humanoids), 2015.
Jonathan Lacotte, Mohammad Ghavamzadeh, Yinlam Chow, and Marco Pavone. Risk-sensitive
generative adversarial imitation learning. In International Conference on Artificial Intelligence and Statistics, 2019.
Yunzhu Li, Jiaming Song, and Stefano Ermon. Infogail: Interpretable imitation learning from visual demonstrations. In Advances in Neural Information Processing Systems, 2017.
Timothy P Lillicrap, Jonathan J Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. International Conference on Learning Representations, 2016.
Fangchen Liu, Zhan Ling, Tongzhou Mu, and Hao Su. State alignment-based imitation learning. International Conference on Learning Representations, 2020.
Sridhar Mahadevan and Mauro Maggioni. Proto-value functions: A laplacian framework for learning representation and control in markov decision processes. Journal of Machine Learning Research, 2007.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, al. Human-level control through deep reinforcement learning. Nature, 2015.
Andrew Ng. Stuart Russell, et al. Algorithms for inverse reinforcement learning. In International Conference Machine Learning, 2000.
Aldo Pacchiano, Jack Parker-Holder Yunhao Tang, Krzysztof Choromanski, Anna Choromanska, and Michael Jordan. Learning to score behaviors for guided policy optimization. In International Conference on Machine Learning, 2020.
Xue Bin Peng, Pieter Abbeel, Sergey Levine, and Michiel van de Panne. Deepmimic: Example-guie guiiee deep reinforcement learning of physics-based character skills. ACM Transactions on Graphics (TOG), 2018.
Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport. Foundations and Trends in Machine Learning, 2019.
Bilal Piot, Matthieu Geist, and Olivier Pietquin. Learning from demonstrations: Is worth estimating reward function? In Joint European Conference on Machine Learning and Knowledge Discovery in Databases, 2013.
Bilal Piot, Matthieu Geist, and Olivier Pietquin. Boosted and regaaa-regularizeed classification for apprenticeship learning. In International conference on autonomous agents and multi-agent systems, 2014.
Bilal Piot, Matthieu Geist, and Olivier Pietquin. Bridging the gap between imitation learning and inverse reinforcement learning. IEEE transactions on neural networks and learning systems, 2016.
Dean A Pomerleau. Efficient training of artificial neural networks for autonomous navigation. Neural computation, 1991.
Martin L Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 2014.
Aravind Rajeswaran, Vikash Kumar, Abhishek Gupta, Giulia Vezzani, John Schulman, Emanuel Todorov, and Sergey Levine. Learning complex dexterous manipulation with deep reinforcement learning and demonstrations. In Proceedings of Robotics: Science and Systems (RSS), 2018.
Siddharth Reddy, Anca D Dragan, and Sergey Levine. Sqil: imitation learning via regularized behavioral cloning. International Conference on Learning Representations, 2020.
Stéphane Ross and Drew Bagnell. Efficient reductions for imitation learning. In International Conference Artificial Intelligence and Statisties, 2010.
12

---

Published as a conference paper at ICLR 2021
Stuart Russell. Learning agents for uncertain environments. In Conference Computational learning theory, 1998.
Yannick Schroecker and Charles L Isbell. State aware imitation learning. In Advances in Neural Information Processing Systems, 2017.
John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In International Conference on Machine Learning, 2015.
David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Marc Lanctot, et al. Mastering the game of go with deep neural networks and tree search. Nature, 2016.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. 2018.
Gerald Tesauro. Temporal difference learning and td id-gammon. Communications of the ACM, 1995.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In International Conference on Intelligent Robots and Systems, 2012.
Faraz Torabi, Garrett Warnell, and Peter Stone. Behavioral cloning from observation. In International
Joint Conference on Artificial Intelligence, 2018.
Cédric Villani. Optimal transport: old and new. 2008.
Ruohan Wang, Carlo Ciliberto, Pierluigi Vito Amadori, and Yiannis Demiris. Random expert distillation: Imitation learning via expert policy support estimation. International Conference Machine Learning, 2019.
Huang Xiao, Michael Herman, Joerg Wagner, Sebastian Ziesche, Jalal Etesami, and Thai Hong Linh. Wasserstein adversarial imitation learning. arXiv preprint arXiv:1906.088 13, 2019.
Brian D Ziebart, Andrew Maas, Andrew Bagnell, and Anind K Dey. Maximum entropy inverse reinforcement learning. In AAAI Conference Artificial Intelligence, 2008.
13

---

Published as a conference paper at ICLR 2021
A LOCOMOTION EXPERIMENTS
A.1 PWIL IMPLEMENTATION
The agent we use is D4PG (Barth-Maron Maron et al. 2018). We use the default architecture from the ACME framework (Hoffman al. 2020).
The actor architecture is 4-yyyer neural network: the first layer has size 256 with tanh activation and layer normalization (Ba al. 2016), the second layer and third layer have size 256 with elu activation (Clevert et al. (016), the last layer is of size the dimension of the action space, with tanh activation scaled the action range of the environment. To enable sufficient exploration with use Gaussian noise layer top of the last layer with standard deviation action range the environment. evaluate the agent without exploration noise.
0.2, that we clip to the
For the critic network we use 4-layer neural network: the first layer has size 512 with tanh activation and layer normalization, the second layer of size 512 with elu activation, the third layer is of size 256 with elu activation, the last layer is of dimension 201 with softmax activation. The 201 neurons stand for the weights of the distribution supported within equal distance in the range [-150, 150], as categorical distribution (Bellemare al. 2017).
We use the Adam optimizer Kingma Ba 2015) with Aa = 5 10-6 for the actor and A = 7 10-5 for the critic. We use batch size of 256. We clip both gradients from the critic and the actor limit their L2 norm to 40.
We replay buffer of size 105; discount factor 0 0.99 and step returns with prefill the replay buffer with 50000 state-action pairs from the set of demonstrations (which means that we put multiple times the same expert transitions in the buffer). We perform updates on the actor and the critic every k interactions with the environment. The hyperparameters search is detailed Table
We
Parameters
Values
A
10 2 2,5x
5, 10 1,10
10
A
5 1,- 10 5,10
g
0.1,0.2,0.3
k
2,4,8,16
Table DDPG hyperparameters search.
For the reward function, we did run hyperparameter search on a and / with the following values for a and /3 in (1,5, 10}.
A.2 DAC IMPLEMENTATION
We used the open-source implementation of DAC provided by Kostrikov et al. (2019). For Humanoid (which was not reported in the paper), we did run large hyperparameter search described in Table Out of the 729 experiments, 7 led to an average performance above 1000 at the 2.5M training environment steps mark.
Actor learning rare (Ir) 0.001, 0.005, 0.0001
IF 0.001, 0L0MMS, 0.0001
Discrimintto
regularizer
Explocaiio noise
Decay
0.001, 0.0005, 0.0001
5. 10,
0.3
0.3, 0.S. 0.8
Table 2: Hyperparameter search on Humanoid for 11 demonstrations. (tested/best)
A.3 BC IMPLEMENTATION
We used 3-layer neural network, the first layer has size 128 with relu activation, the second layer has size 64 and the last layer has size of the action space with a tanh activation scaled to the action range of the environment. We found out that normalizing the observations with the average and standard deviation of the expert demonstrations' observations helped performance. We trained the network using the mean squared error as a loss, with an Adam optimizer. We ran an hyperpameter search on the learning rate in (10-5, 10-1 10-3) and the batch size (128, 256).
14

---

Published as a conference paper at ICLR 2021
A.4 PWIL LEARNING CURVES
We present the learning curves of PWIL, in terms of the reward defined in Equation well as the original reward of the task.
Hopper
Walker2d
Ant
HallChestahh
Humanod
Reacher
3.0
-06
100
154
100
300
200
202
16
60-
100
20
160
160
166
160
1eb
100
Hopper
Waller2d
HattChoctah
Humanold
600
340
600
22
429
209
2.30
0.25
3.75
1.30
000
0.35
658
6.15
3.23
0.00
0.25
0.50
0.75
1.0
020
236
5.22
8.16
1.30
3.30
0.56
3.90
0.75
Hopper
Walker2d
Ant
HattChoctah
Humanold
Reachor
1200
1030
1re0-
325
000
300
465
42
166
160
166
166
165
166
Figure 7: Mean and standard deviation of the imitation return of the agent. We perform 10 rollouts over 10 seeds every 10k environment steps over 1M steps. The return here is in term of the reward defined with Equation (6) Top row: demonstration, middle row: 4 demonstrations, bottom row demonstrations.
Hooper
Walker2d
Ant
HalfCheetah
Humanoid
Reachar
2500
5004
-20
2000
4005
1500
3002
2500
1000
1002
2007
1130
2200
3000
-40
random
14b
14b
166
146
1eb
185
Walke/2d
HollChhctth
Humanold
Reacher
3000
4299
1200
-10
1000
4000
-25
1530
2300
500
-40
rozent
0.2
615
1.00
035
132
0.00
0.5
0.5
312
0.5
325
0.5
133
7.00
0.30 030
Hooper
Walker2d
Ant
HalfChwoial:
Hamanoid
Reac her
3500
3000
8000
4002
3500
6000
2000
500
2030
4300
PHLL
1000
X00
-20
1402
5000
Girdkr
3.00
165
165
165
305
165
Figure 8: Mean and standard deviation of the original environment returns of the evaluation policy over 10 rollouts and 10 seeds, reported every 10k environment steps over 1M steps. Top row: demonstation, middle row: 4 demonstrations, bottom row: 11 demonstrations.
15

---

Published as a conference paper at ICLR 2021
Hopper
Walker2d
Art
HarCChesah
Humanood
Reacher
3.3
DXC
PMTL, WJ
4.
3.5
2.5
330
023
0.50
0.75
1.30
0.30
323
6.0
0.7
1.9
0.00
0.25
0.30
0.75
1.33
0.00
0.25
0.30
0.75
1.30
1e5
0.00
0.35
0.00
0.75
100
Hopper
Welker2d
Art
HellC heet et
Hummenckd
Reest
DAC
5.3
PArt Wi Baure
20
Watvedden
15
3.5
14
1.00
225
0.50
235
1.00
0.30
325
6.5
0.75
1.00
0.0
0.25
0.50
0.75
1.39
0.00
0.25
0.30
0.75
1.30
1e5
0.00
0.25
0.00
0.75
1.00
Hoppar
Walker2d
Art
HafCChestah
Humanokd
Reacher
20.0
14
DAC
PWYL WL Dound
4.5
25
393 0.5
0.75
1.30
0.30
925
3:33
0.00 0.25
0.50
0.75
1.39
0.00
0.5
0.30
0.5
1.30
1e5
0.0 035
050
1.00
Figure 9: Mean and standard deviation of the Wasserstein distance between the state-action distri- bution of the evaluation policy and the state-action distribution of the expert over 10 rollouts and 10 seeds, reported every 10k environment steps. We include the upper bound the Wasserstein distance based on the greedy coupling defined in Equation Top row: demos, middle row: demos, bottom row: 11 demos.
B ABLATION STUDY
We present the learning curves of PWIL in the presence of ablations.
Walkerd
Humannid
300
4E00
CO00
5000
7ML
3030
0200
3000
1500
cvril
6200
1000
Pit Potk-capport
1230
2200
-00
0.33
0.25
1.54
0.5
1.6
0.00
635
3.3
0.75
133
0.00
035
0.00
0.35
L.00
644
3.25
0.0
0.55
1.00
330
0.5
830
0.75
Hoppper
Walee2d
Ant
HoltChetthh
Humanold
Reacher
2040
1003
5002
2500
3000
400
-20
300
2259
4000
1240
MolLosppent
1003
2002
3000
leb
160
1e6
1e6
1e5
Hozper
Walkerdd
HaltChectah
Humanoid
Reachar
2000
4030
6000
2500
PMLL
6000
2500
2050
4200
Pott-capper
P00147
1030
2200
3000
0/5
6J5
435
6.75
1:30
6.5
2.25
076
1mb
100
10c
24b
180
Figure 10: Evaluation performance of different variants of PWIL over 10 rollouts and 10 seeds, re- ported every 10k environment steps over 1M steps. The return is in term of the environment's original reward. Top row: demonstration, middle row: demonstrations, bottom row: 11 demonstrations,
16

---

Published as a conference paper at ICLR 2021
C INFLUENCE OF THE DIRECT RL ALGORITHM
In this section, we study the influence of the direct RL method on the performance on PWIL on locomotion tasks. We compare the performance PWIL with D4PG with the performance of PWIL with TD3 and SAC. We use a custom implementation of both SAC and TD3 which is reproduction of the authors' implementation. For SAC we used the same default hyperparameters: for TD3 we removed the delayed policy update which led to better performance. We evaluate both methods on same reward function D4PG. We remove the pre-warming of the replay buffer for SAC. For TD3, we found out that we need different werrannn for each environment to get good performance (contrary to D4PG): we used 1000 state-action pairs for Reacher, Hopper and Walker2d and no pre-warming for Humanoid, Ant and HallCheetah.
Hoppor
Walker2d
Reacher
330
480
Figure 11: Mean and standard deviation of the evaluation performance of PWIL-D4PG, PWIL-SAC and PWIL-TD3 at the 1M environment interactions mark (2.5M for Humanoid). Results are computed over 10 seeds and 10 episodes for each seed.
We show in Figure II that both PWIL-SAC and PWIL-TD3 recover expert-like performance on Hopper, Ant and HalfCheetah and Reacher. For Walker2d, there large standard deviation in performance, since some seeds lead to near-optimal expert performance, and some seeds lead to poor performance. Remark that we led minimal hyperparameter tuning on the default versions of SAC and TD3, which might explain the differences in performance.
17

---

Published as a conference paper at ICLR 2021
D DOOR OPENING EXPERIMENTS
In this section we present the details of the experiment presented in Section 4.4
D.1 ENVIRONMENT
The default door opening environment has a fixed horizon of 200 timesteps. We add an early termination condition to the environment when the door is opened to demonstrate that PWIL extends to tasks where there is no incentive for survival. Note that similarly to AIRL (Fu et al., 2018), early termination states are considered to be absorbing. Therefore if the expert reaches early termination state early, its state distribution as a large weight on these states, which leads the PWIL agent to similarly reach the early termination state. 25 demonstrations were collected from humans using virtual reality system (Kumar & Todorov (2015)).
D.2 EMBEDDING
As we assume that we only have access to the visual rendering of the demonstrations (with resolution 84 x 84), we build lower dimensional latent space. We use a self-supervised learning method: TCC, learn the latent space. We used the classification version of the algorithm detailed in Dwibedi et al. (2019) (Section 3.2).
The encoding network is the following: convolutional layer with 128 filters of size and stride 4, relu activation, convolutional layer with 64 filters of size 4 and stride 2, relu activation, convolutional layer with 64 filters size 3 and stride 1, a relu activation, a linear layer of size 128, a relu activation and output layer of size 32.
We separated the demonstrations into a train set of 20 demonstrations and a validation set of 5 demonstrations. Similarly to Dwibedi et al. (2019), we select the encoding network by taking the that maximizes the alignment according to Kendall's tau score on the validation set.
0.9
Kendail's
0.6
We ran hyperparameter search over the batch size {1, 4, 16, 3, the subsampling ratio (1, 4, 8) and the learn- ing rate (10-2, 10-3, 10 10-4) (using the Adam optimizer). The best performing encoder had an average Kendall's 's tau score of 0.98 on the training set and 0.91 on the validation set.
-0.2
10
Iteration
Figure 12: Kendall's tau score of the encoding network on the train and vali- dation sets. An iteration corresponds to 1000 gradient updates.
D.3 AGENT
In this setup, we found that SAC leads to better results than D4PG. We detail the implementation of SAC below.
The actor is a 3-layer network that outputs the parameters of normal distribution with diagonal covariance matrix. We enforce the terms of the covariance matrix to be greater than 10-3, The first layer has size 256 with relu activation, the second layer has size 256 with relu activation and the last layer has size 48 (the action space has size 24).
The critic is 3-layer network. We pass the concatenation of the state and the action of the environment input. There two hidden layers of size 256 with relu activation. We do twin critic.
We use an adaptive temperature a as described by Haarnoja et al. (0018)).
For the three losses we use the Adam optimizer, and use Polyak averaging for the critic networks with a rate т. We ran an hyperparameter search over the learning rates Xay Ae, Aa in {10-2, 3 10-1), over 7 (0.01, 0.005), over the batch size {64, 128, 256} and over the number of interactions with the environment between each actor, critic and temperature update {1, 4, 8..
18

---

Published as a conference paper at ICLR 2021
E
ALGORITHM DETAILS
E RUNDOWN
We present an example rundown PWIL in Figure 13
Start episode.
W ==
4'4
61 :=
Initialization
si 82
83 1
1 1.
1
wº = [0.
1 1'4
=
4
x
d(sT. +f)
$9 82
83 84
Reward for first state:
r1 = exp (-c1)
1
1
wº = 0.00
4
C1
4
x
d(ss
d(s), $5)
si
82
8/ 84
1
w" = 10,0,
02 ==
1
Initialization
$2
83 81
1
w. 10.0.0.
02
x
d(sg, 5g)
Reward for second state: 1'2 exp (-C2)
82
8/ 2
We = 10,0,0,0 62 = dl 11 (9)
d( d(x2,
End episode.
Figure 13: PWIL example rundown. We drop the dependency on the action. The example task has an horizon T 2, and the demonstration set has D = 4 elements. We take f : + exp (-c).
E.2 RUNTIME
We detail the runtime of PWIL in this section. A single step reward computation consists in computing the distance from the current observation (a vector with dimension [S] |A| with D vectors (the expert observations), which has complexity O(([S) + D). The next step to get the minimum of these distances (complexity OD))) for [1 times (the number times the algorithm executes the while loop). Therefore the complexity of single step reward computation is O((SS) + [A)) 2²) and the complexity of the reward computation of an episode (([S) + [A)) D2).
19


Reviews:
[{"review_id": "xzD-F9c79O", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: \nThis paper proposes to use Wasserstein distance in the primal form for imitation learning. Compared with its dual form and f-divergence minimization variants, it avoids the unstable minimax optimization. In order to compute the Wasserstein distance in primal form, they also propose a greedy approximation. Their experiments demonstrate that this method has a better performance compared with baseline methods.\n\nPros:\n+ The greedy approximation of the primal Wasserstein distance is a clean solution, and it works well for MuJoCo tasks, even if in the LfO setting\n+ The paper presents a series of experiments and ablation studies, which showed quite strong performance. In the ablation, the experiments about PWIL-support is quite convincing. The agent may stay on the supports for other the density/support estimation-based methods, while PWIL doesn't.\n+ The paper is well-written and easy to follow. It provides enough implementation details and codes, which is reproducible\n\nSome questions & concerns:\n+ I'm wondering if the uniform distribution assumption can be improved, by incorporating certain density estimation methods. \n+ Why the complexity of the algorithm is O((|S| + |A|)D)? It should also be linear in T\n+ One limitation is that you need carefully selected metrics, e.g. L1, standardized L2. And it's generally hard to compute Wasserstein distance in the image domain. \n+ For the visual imitation experiments, the figures only contain PWIL without the results of other baselines. They can still be trained on the feature from TCC. \n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "n-UN6LMCvPR", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper develops an imitation-learning (IL) method starting from the primal form of the Wasserstein distance, creating an upper-bound (by replacing optimal coupling with a greedy coupling), and converting that into a practical, scalable algorithm (PWIL). Experiments on the standard MuJoCo tasks and a pixel-based door opening task show that the method is competitive to the recent IL approaches in terms of performance and sample-efficiency.\n\nI am currently on the fence about this paper. I would like the authors to address the following concerns:\n\n1.\tThe experiments in Figure 2 compare PWIL (Distributional DDPG) with DAC (TD3). The difference in the underlying off-policy RL algorithm somewhat muddles the comparative evaluation of the reward-deriving strategy – greedy coupling in PWIL vs. discriminator training in DAC. From what I understand, PWIL should work seamlessly with TD3 as well. Is there a particular reason why the authors went with distributional-RL? Providing PWIL+TD3 results would be an equitable comparison to DAC.\n\n2.\tDiscussion of potential limitations of the current approach -- the paper does strong relaxations when moving from theory to practice. I don’t have issues about this, but it would certainly be useful to include comments about the potential pitfalls of the approach. For instance, are there MDPs where the following the greedy coupling strategy provides deceptive rewards leading to sub-optimal imitation? Another example is the handling of the reward bias. Equation (6) is explicitly biased to provide a positive reward, and therefore, cannot handle cases where the episodes should be terminated at the goal (like the popular LunarLander Gym environment). Section 4.4 adds a manually designed constraint to get around this, but unlike DAC, I do not see a principled way to manage this with PWIL.\n\n3.\tComparing rewards to those from adversarial IL – In the case of adversarial IL, the rewards are non-stationary as the discriminator parameters change throughout the training. However, for discriminator parameters at any particular iteration, the rewards are Markovian (depend only on current state and action). The rewards for PWIL seem to be non-Markovian (Equation 5), yet they have been integrated with the RL machinery (policy-gradients like DDPG) designed for Markovian rewards. Is this observation correct, and if yes, is there an intuition for why this should work?\n\n4.\tI assume the blue curve in Figure 3 is the actual Wasserstein distance. It would be helpful to add a line or two regarding the method used to obtain the optimal coupling to plot this.\n\n\n===== POST REBUTTAL UPDATE =====\n\n\nI have updated my score based on the response. I hope the authors can add PWIL+TD3 results (point 1. in review) in the next revision.", "extra_scores": {}, "extra_text": {}}, {"review_id": "f3gIxvEBnzd", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "### Summary \nThe authors proposed an imitation learning algorithm that utilizes the primal form of Wasserstein distance to match agent’s and expert’s state-action visitation distributions. They considered the upper bound of the primal form and devise the optimization method based on greedy coupling which makes learning suitable for sequential problems. With standardized Euclidean distance and exponential smoothing, the proposed method PWIL is shown to perform well for both MuJoCo and Door Opening Tasks and highly outperforms the baseline (DAC) for Humanoid. \n\n### Quality \nTheoretical derivations are sound, and experiments are well-tuned and designed and highly support the authors’ claims.\n\n### Clarity \nThe submission is clearly written overall, but I added minor comments on clarity in `Detailed Comments`\n\n### Originality \nThe idea of submission is novel from the use of the primal form of Wasserstein distance in concatenation with greedy coupling. \n\n### Significance \nI think the contribution of this work is significant in the sense that it considers the primal form of Wasserstein distance and eventually show the decrease of Wasserstein distance through experiments, whereas the existing works haven’t explicitly shown the performance in terms of probabilistic metrics. \n\n### Detailed comments \n(p. 1, Abstract) We present a reward function which is derived offline, as opposed to recent adversarial IL algorithms that learn a reward function through interactions with the environment, and which requires little fine-tuning.\n- I agree that PWIL doesn’t require neural networks and relevant gradient updates for reward learning but wasn’t sure about the term “offline” is appropriate to describe this. PWIL involves the internal loop for reward updates during environment simulation. This is a bit confusing at the first glance since I believe “offline” is used when the environment interaction is not allowed, e.g., offline reinforcement learning. (My first bias was that the reward is learned before the interaction, fixed, and then used during training without any updates.) I guess there’ll be a better expression for it (like a non-parameter reward?) or simply ignoring the term “offline” is much better. \n- A sentence in the introduction “The inferred reward function is non-stationary, like adversarial IL methods, but it is not re-evaluated as the agent interacts with the environment, therefore the reward function we define is computed offline.” tries to clarify this term, but I don’t understand why it is said that the reward is *not re-evaluated* since the most inner loop of **Algorithm 1** is related to the reevaluation of rewards. \n\n(p. 1, Introduction) Our method recovers expert behaviour better than existing state-of-the-art methods while being based on significantly fewer hyperparameters\n- In Figure 1, PWIL seems to perform comparably to DAC, not better than DAC, except for Humanoid. \n\n(p. 2, Background and Notations) $\\delta$ is the discount factor\n- It seems like $\\delta$ is used instead of $\\gamma$ since $\\gamma$ is used to indicate a stochastic matrix. However, since most potential readers may be familiar with the terms in RL and imitation learning, I’d rather use $\\gamma$ for the discount factor and a different letter for coupling. \n- Also, it collides with Drac distribution in the following paragraph. \n\n(p. 2, Background and Notations) from maximizing to minimizing its return\n- $\\rightarrow$ from maximizing its return to minimizing cumulative costs?\n\n(p. 2, Background and Notations) requires the definition of a metric in the space\n- $\\rightarrow$ requires the definition of a metric $d$ in the space (to easily link with one in the definition of Wasserstein distance)\n\n(p. 2, Method) We introduce a reward based on an upper-bound of the Wasserstein distance\n- $\\rightarrow$ We introduce a reward based on an upper-bound of the Primal form of  Wasserstein distance\n\n(p. 3, Wasserstein distance minimization) This can be problematic if an agent learns in an online manner or in large time-horizon tasks. Thus, we introduce an upper bound to the Wasserstein distance that yields a cost we can compute online, based on a suboptimal coupling strategy.\n- I agree that this is a key idea of the submission, but I was wondering if it isn’t available to calculate the cost even when off-policy RL is applied as was done in the experiments. \n\n(p. 3, Greedy Coupling) Eq (4)\n- It seems to me removing $\\inf_{\\pi\\in\\Pi}$ is possible and fits the whole equation inside the format. \n\n(p. 4, Figure 1) Note that the total cost with the greedy coupling is 7 whereas the total cost with the optimal coupling is 5.\n- With the definition in Eq (5), I think the multiplication of $\\gamma_\\pi^q[i, j]$ (which is the uniform distribution) is missing. \n- Also, it would be better to add detailed calculations in the Appendix for curious readers. \n\n(p.4, Greedy Coupling)  The algorithm computes the greedy coupling with a complexity $\\mathcal{O}((|\\mathcal{S}|+|\\mathcal{A}|)D)$\n- How’s the complexity calculated in detail?\n\n(P.4, Experiments) As DAC is based on TD3 which is a variant of DDPG, we use a DDPG-based agent for fair comparison.\n- This makes me a bit confused since it was non-trivial for me to link off-policy RL with Algorithm 1 directly. One thing I’m really curious about is how the reward comes from the replay buffer used in a DDPG-based agent. \n\n(P.4, Experiments) Figure 2\n- How’s the computational cost for overall evaluation? Specifically, I think there’s no problem with the evaluation of Wasserstein distance when a small number of the expert demonstration is given, but the calculation time will hugely increase for a large number of expert data (which I believe can be in practical scenarios).\n\n\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "8mKnvizVnLC", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors present a well written paper with many detailed experiments.  I am happy to revisit my evaluation after a fruitful discussion :)\n\nPros\n\nWesserstein distance between demonstration and policy is setup\nUsing the earth mover’s analogy with the knowledge of T piles of dirt, distances to the holes are computed. Then the piles of dirt are moved to the closest holes in a greedy fashion.\nAn offline learning policy is presented\nSample complexity for learning is suggested to be low \nThe model seems to be able to learn from a single example to run\n\nCons\n\nThe Sample complexity for training is still high. While I do agree that the number of demonstrations required is lower.\n\nThe win is not obvious or strong except for the humanoid condition. This is my biggest concern. Can we explore more complex environments? Or the effects of more demonstrations? It would be nice to understand where this model truly succeeds or fails.\n\nIn the door opening experiment it would be great to know how the model performs while varying the number of expert demonstrations (say). \n\nMight be tangential, but this paper seems relevant at least in introducing the idea of inverse reinforcement learning and W-distance\nhttps://arxiv.org/pdf/1906.08113.pdf\n\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: It is common in imitation learning to measure and minimize the differences between the agent’s and expert’s visitation distributions. This paper proposes using Wasserstein distance for this, named PWIL, by considering the upper bound of its primal form and taking it as the optimization objective. The effectiveness of the approach is demonstrated by an extensive set of experiments. 

Overall, reviewers reached general agreement that this paper makes a good contribution to the conference, and given the overall positive reviews, I also recommend accepting the paper.

Decision: Accept (Poster)