Robert Dadashi
email: dadashi@google.com.

Léonard Hussenot

Matthieu Geist

Olivier Pietquin

Primal Wasserstein Imitation Learning

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.

Introduction

Reinforcement Learning (RL) has solved a number of difficult tasks whether in games [START_REF] Tesauro | Temporal difference learning and td-gammon[END_REF][START_REF] Mnih | Human-level control through deep reinforcement learning[END_REF], Silver et al., 2016] or robotics [Abbeel andNg, 2004, Andrychowicz et al., 2020]. However, RL relies on the existence of a reward function, that can be either hard to specify or too sparse to be used in practice. Imitation Learning (IL) is a paradigm that applies to these environments with hard to specify rewards: we seek to 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 [START_REF] Dean | Efficient training of artificial neural networks for autonomous navigation[END_REF][START_REF] Bagnell | Boosting structured prediction for imitation learning[END_REF][START_REF] Ross | Efficient reductions for imitation learning[END_REF] and Inverse Reinforcement Learning [START_REF] Russell | Learning agents for uncertain environments[END_REF][START_REF] Andrew | Algorithms for inverse reinforcement learning[END_REF]. In Behavioral Cloning, we seek to recover the expert's behavior by directly learning a policy that matches the expert behavior in some sense. In Inverse Reinforcement Learning (IRL), 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 an agent on it. Although IRL methods introduce an intermediary problem to solve (i.e. recovering the environment's reward) they offer the benefit of being less sensitive to distributional shift [START_REF] Dean | Efficient training of artificial neural networks for autonomous navigation[END_REF], generalizing to environments with possibly changing dynamics [START_REF] Piot | Learning from demonstrations: Is it worth estimating a reward function[END_REF], or recovering a near-optimal agent from suboptimal demonstrations [START_REF] Brown | Extrapolating beyond suboptimal demonstrations via inverse reinforcement learning from observations[END_REF][START_REF] Jacq | Learning from a learner[END_REF].

However, IRL methods are usually based on an iterative process alternating between reward estimation and RL, which might result in poor sample-efficiency. Earlier IRL methods [START_REF] Andrew | Algorithms for inverse reinforcement learning[END_REF][START_REF] Abbeel | Apprenticeship learning via inverse reinforcement learning[END_REF][START_REF] Ziebart | Maximum entropy inverse reinforcement learning[END_REF] require multiple calls to a Markov decision process solver [START_REF] Martin L Puterman | Markov decision processes: discrete stochastic dynamic programming[END_REF], whereas recent adversarial IL approaches [START_REF] Finn | Guided cost learning: Deep inverse optimal control via policy optimization[END_REF][START_REF] Ho | Generative adversarial imitation learning[END_REF][START_REF] Fu | Learning robust rewards with adversarial inverse reinforcement learning[END_REF] 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) [START_REF] Goodfellow | Generative adversarial nets[END_REF], 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 implicitely minimize an f -divergence between the state-action distribution of an expert and the state-action distribution of the learning agent [START_REF] Kamyar | A divergence minimization perspective on imitation learning methods[END_REF][START_REF] Ke | Imitation learning as f -divergence minimization[END_REF]. However the interaction between a generator (the policy) and the discriminator (the reward function) makes it a minmax optimization problem, and therefore comes with practical challenges that might include training instability, sensitivity to hyperparameters and poor sample efficiency.

In this work we use the Wasserstein distance as a measure between the state-action distributions of the expert and of the agent. Our approach is novel in the fact that we consider the problem of minimizing the Wasserstein distance through its primal formulation. We introduce an upper bound to the Wasserstein distance, which we verify empirically to be tight, and infer a conceptually simple reward function from it. The inferred reward function is non-stationary similarly to adversarial IL, but it is not re-evaluated as the agent interacts with the environment, therefore the reward function we define is computed offline. Crucially, this leads to a problem which is not a minmax optimization problem, and requires little fine tuning.

We make the following contributions. We introduce a reward function computed offline based on an upper bound of the primal form of the Wasserstein distance that does not need to be re-computed during the training procedure. We present a 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 of the task we consider (which is unknown in general). The method we present is conceptually simple and we show that it is able to recover expert behavior in MuJoCo environments with sample efficiency comparable to state-of-the-art methods both in terms of agent interactions with the environments and in terms of expert interactions with the environments.

Background and Notations

Markov decision processes We describe environments as episodic Markov Decision Processes (MDP) with finite time horizon [START_REF] Richard | Reinforcement learning: An introduction[END_REF]] (S, A, P, r, γ, ρ 0 , T ), where S is the state space, A is the action space, P is the transition kernel, r is the reward function, γ is the discount factor, ρ 0 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 a mapping from states to distributions over actions; we denote the space of all policies by Π. In RL, the goal is to learn a 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 c rather than a reward r [Puterman, 2014], which essentially moves the goal of the policy from maximizing to minimizing its return.

State action distributions Suppose a policy π visits the successive states and actions s 1 , a 1 , . . . , s T , a T during an episode, we define the empirical state-action distribution ρπ as:

ρπ = 1 T T t=1 δ st,at ,
where δ st,at is a Dirac distribution centered on (s t , a t ). Similarly, suppose we have a set of expert demonstrations D = {s e , a e } of size D, then the associated empirical expert state-action distribution ρe is defined as:

ρe = 1 D (s,a)∈D δ s,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 ν two distributions on M with finite moments, the p-th order Wasserstein distance [START_REF] Villani | Optimal transport: old and new[END_REF] is defined as:

W p p (µ, ν) = inf γ∈Γ(µ,ν) M ×M d(x, y) p dγ(x, y),
where Γ(µ, ν) is the set of all couplings between µ and ν.

In the rest of the paper we only consider distributions with finite support. A coupling between two distributions of support cardinal T and D is a doubly stochastic matrix of size T × D. We note Γ the set of all doubly stochastic matrices of size T × D, that is:

Γ = γ ∈ R T ×D + | ∀j ∈ [1 : D], T i =1 γ[i , j] = 1 D , ∀i ∈ [1 : T ], D j =1 γ[i, j ] = 1 T .
Optimizing a Wasserstein distance between distributions of state-action pairs requires the definition of a metric in the space (S, A). Defining a metric in an MDP can be highly non trivial [START_REF] Givan | Equivalence notions and model minimization in markov decision processes[END_REF][START_REF] Ferns | Metrics for finite markov decision processes[END_REF]; we extend on this subject in the appendix. From now on, we will assume the existence of a metric d : (S, A) × (S, A) → R + .

Method

In this section, we present the theoretical motivation of our approach, which is the minimization of the Wasserstein distance between the state-action distributions of the agent and the expert. We introduce a reward based on an upper bound of the Wasserstein distance inferred from a relaxation of the optimal coupling condition, and present the resulting algorithm: Primal Wasserstein Imitation Learning (PWIL).

Wasserstein distance minimization

Central to our approach is the minimization of the Wasserstein distance between the state-action distribution of the policy we seek to train ρπ and the state-action distribution of the expert ρe . In other words, we aim at optimizing the following problem:

inf π∈Π W 2 2 (ρ π , ρe ) = inf π∈Π inf γ∈Γ T i=1 D j=1 d((s π i , a π i ), (s e j , a e j )) 2 γ[i, j]. (1)
We can interpret the Wasserstein distance using the earth's movers analogy [START_REF] Villani | Optimal transport: old and new[END_REF]. Consider that the state-action pairs of the expert are D holes of mass 1 D and that the state-action pairs of the policy are piles of dirt of mass 1 T . A coupling γ is a transport strategy between the piles of dirt and the holes, where γ[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, we need knowledge of the locations of all piles of dirt. In the context of RL, this means having access to the full trajectory generated by π.

From now on, we write γ * π as the optimal coupling for the policy π, that is:

γ * π = arg min γ∈Γ T i=1 D j=1 d((s π i , a π i ), (s e j , a e j )) 2 γ[i, j].
We have:

inf π∈Π W 2 2 (ρ π , ρe ) = inf π∈Π T i=1 D j=1 d((s π i , a π i ), (s e j , a e j )) 2 γ * π [i, j] =c * i,π . (2)
In the Equation ( 2), we have introduced c * i,π , which we interpret as a cost to minimize using RL. As c * i,π depends on the optimal coupling γ * π , we can only define c * i,π 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. That is why we introduce an upper bound to the Wasserstein distance that yields a cost we can compute online, based on a suboptimal coupling strategy.

Greedy coupling

In this section we introduce the greedy coupling γ g π ∈ Γ, defined recursively for 1 ≤ i ≤ T as:

γ g π [i, :] = arg min γ[i,:]∈Γi D j=1 d((s π i , a π i ), (s e j , a e j )) 2 γ[i, j] (3)
with

Γ i = γ[i, :] ∈ R D + D j =1 γ[i, j ] = 1 T , ∀j ∈ [1 : D], i i =1 γ[i , j ] ≤ 1 D
Similarly to Equation (1), Equation ( 3) can be interpreted using the earth mover's analogy. Contrary to Equation ( 1) 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 mover need 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

D j =1 γ[i, j ] = 1
T means that all the dirt that appears at the i-th timestep needs to be moved, and the

constraint ∀j ∈ [1 : D], i i =1 γ[i , j ] ≤ 1
D means that we cannot fill the holes more than their capacity 1 D . We show the difference between the greedy coupling and the optimal coupling in Figure 1, and provide the pseudo-code to derive it in Algorithm 1.

Figure 1: Illustration of the difference between the a) greedy coupling and the b) optimal coupling. We present an MDP where we drop the dependency on the action. The state space is R and the metric associated is the Euclidean distance. We note the states visited by the policy: s π 1 , s π 2 , s π 3 and the states visited by the expert: s e 1 , s e 2 , s e 3 . When the policy encounters the state s π 2 , and because we do not know s π 3 yet, the greedy coupling strategy consists in coupling it with s e 2 although the optimal coupling strategy would be to couple it with s e 3 . Note that the total cost with the greedy coupling is 7 whereas the total cost with the optimal coupling is 5. This highlights that the optimal coupling needs knowledge about the whole policy's trajectory to be derived.

An upper bound to the Wasserstein distance

We can now define an upper bound of the Wasserstein distance using the greedy coupling (since by definition it is subotpimal):

inf π∈Π W 2 2 (ρ π , ρe ) = inf π∈Π T i=1 D j=1 d((s π i , a π i ), (s e j , a e j )) 2 γ * π [i, j] ≤ inf π∈Π T i=1 D j=1 d((s π i , a π i ), (s e j , a e j )) 2 γ g π [i, j] (4)
In Section 4, we empirically validate this upper bound. We can infer a cost from Equation (1) at each timestep i:

c g i,π = D j=1 d((s π i , a π i ), (s e j , a e j )) 2 γ g π [i, j]. (5)
Note that the greedy coupling γ g π [t, .] defined in Equation ( 3) is dependent on all the previous stateactions visited by the policy π, which makes the cost c g i,π non-stationary, similarly to adversarial IL methods. Although this cost function is non-stationary, it does not have to be re-estimated as the agent interacts with the environment, and hence is said to be derived offline. We can infer a reward from the cost,

r i,π = f (c g i,π ),
where f is a monotonously decreasing function. Crucially, we 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 1). The algorithm presented is written using pseudo-code for a generic agent A which implements a policy π A , it observes tuples (s, a, r, s ) and possibly updates its components. This formulation is general enough to show that our method is independent of the RL algorithm used. The algorithm computes the greedy coupling with a complexity O((|S| + |A|)D). 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. In this section, we present the implementation of PWIL and perform an empirical evaluation that answers the following questions: 1) Can 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? We also show the dependence of PWIL on multiple of its components through an ablation study. The complete description of the architecture of the agent and the hyperparameters search is provided in appendix; we will open-source the code in the near future.

Implementation

We test our method on five environments from the OpenAI Gym MuJoCo suite [START_REF] Todorov | Mujoco: A physics engine for model-based control[END_REF][START_REF] Brockman | Openai gym[END_REF]: For each environment, multiple demonstrations are gathered using a DDPG agent with a distributional critic [START_REF] Timothy P Lillicrap | Continuous control with deep reinforcement learning[END_REF][START_REF] Barth-Maron | Distributed distributional deterministic policy gradients[END_REF] trained on the actual reward of these environments. Similarly to [START_REF] Ho | Generative adversarial imitation learning[END_REF], we subsample demonstrations by a factor of 20: we select one out of every 20 transitions with a random offset at the beginning of the episode. The subsampling of expert demonstrations makes the imitation task harder, so that a pure behavioral cloning approach does not perform to the level of the expert.

Central to our method is the existence of a metric between state-action pairs. We use the standardized Euclidean distance which is the L2 distance, weighted along each dimension by the inverse standard deviation of the expert demonstrations. From the cost c i computed using Equation ( 5), we define a reward using the following:

r i = α exp(- βT |S| + |A| c i ). (6)
For all environments, we use α = 5 and β = 5. We pick f : x → exp(-x) as the monotonously decreasing function to get a reward from a cost because it is smooth, and has a bounded range [0, 1], which makes it well suited for a categorical distributional critic [START_REF] Marc G Bellemare | A distributional perspective on reinforcement learning[END_REF]. We test our method against a state-of-the-art imitation learning algorithm: Discriminator Actor Critic (DAC) [START_REF] Kostrikov | Discriminator-actor-critic: Addressing sample inefficiency and reward bias in adversarial imitation learning[END_REF]. DAC is based on GAIL, and introduces an off-policy algorithm (TD3 [START_REF] Fujimoto | Addressing function approximation error in actor-critic methods[END_REF]) rather than an on-policy algorithm (TRPO [START_REF] Schulman | Trust region policy optimization[END_REF]). We use the open-source code provided by the authors. For our method, we use a DDPG agent [START_REF] Timothy P Lillicrap | Continuous control with deep reinforcement learning[END_REF] with a distributional critic [START_REF] Marc G Bellemare | A distributional perspective on reinforcement learning[END_REF][START_REF] Barth-Maron | Distributed distributional deterministic policy gradients[END_REF] (see details in the appendix). We initialize the replay buffer by filling it with expert state-action pairs with maximum reward (i.e. α).

We train PWIL and DAC in the limit of 1M environment interactions (2.5M for Humanoid) on 5 seeds. We run experiments with multiple numbers of demonstrations: 1, 4, 11 (consistently with previous work [START_REF] Ho | Generative adversarial imitation learning[END_REF]Ermon, 2016, Kostrikov et al., 2019]). Every 10k environment steps, we perform 10 episodes rollouts per seed of the policy without exploration noise and report performance with respect to the environment's original reward in Figure 2. Learning curves on the inferred reward can be found in the appendix. We also provides visualisations of the final agents in https: //sites.google.com/view/wasserstein-imitation.

Results

In Figure 2, PWIL shows consistent improvement on final performance for Hopper, Walker2d and Humanoid over DAC and similar performance for HalfCheetah. For Ant, PWIL outperforms DAC for 11 demonstrations and reaches similar performance for 1 and 4 demonstrations. Overall PWIL recovers expert behavior better than DAC over the considered environments at the 1M steps mark.

In terms of sample efficiency, DAC outperforms PWIL on HalfCheetah and Ant and has similar performance for Hopper and Walker2d. This seems rather contradictory with the claim that GAIL-like methods have poor sample efficiency. However, while DAC has state-of-the-art sample efficiency on all tasks but Humanoid, it requires careful tuning of the discriminator, including e.g., a scheduled learning rate for the actor. We did not run a hyperparameter search to tune DAC on Humanoid, which might explain its poor performance on the task (results for Humanoid were not reported in [START_REF] Kostrikov | Discriminator-actor-critic: Addressing sample inefficiency and reward bias in adversarial imitation learning[END_REF]). Nevertheless, this exemplifies the brittleness we mentioned for GAIL-like approaches.

Remarkably, PWIL is able to consistently learn policies on Humanoid with an average score over 6000 even with a single demonstration, which means that the agent we train can actually walk (some approaches in IL considers that Humanoid is solved when it can stand which corresponds to a score of 5000). Interestingly, on all environments but Humanoid, both PWIL and DAC are able to learn a policy with near-optimal expert performance with a single demonstration. Wasserstein distance. In Figure 3, we show the Wasserstein distance to expert demonstrations throughout the learning procedure for PWIL and DAC. For both methods, we notice that the distance decreases while learning. PWIL leads to a smaller Wasserstein distance than DAC an all environments but HalfCheetah where it is similar.

PWIL defines a reward from an upper bound to the Wasserstein distance between the state-action distributions of the expert and the agent. We show in Figure 3 the Wasserstein distance as well as this upper bound throughout the learning procedure. Notice that our upper bound is "empirically tight", which validates the choice of the greedy coupling as an approximation to the optimal coupling.

We emphasize that this distance is useful in real settings of IL, regardless of the method used, i.e.

settings where we cannot have access to the actual reward of the task. Note that the convergence in terms of expected return (Figure 2) correlates with the convergence in terms of the Wasserstein distance (Figure 3), which opens interesting directions for e.g. early stopping of IL training procedures.

Ablation Study

In this section, we present the evaluation performance of PWIL over 5 seeds in the presence of ablations and report results in Table 1. We keep the hyperparameters of the original PWIL agent fixed. We provide the learning curves in the appendix.

Hopper-v2

PWIL-state

In this version of PWIL, we do not assume that we 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) [START_REF] Torabi | Behavioral cloning from observation[END_REF], Edwards et al., 2019]. The reward is defined similarly to PWIL, using a state distance rather than a 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-state recovers non-trivial behaviors on all environments. We leave the extensive study of PWIL-state for future work.

PWIL-L2

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 a 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 a drop in performance which is significant for Walker and Ant. This should not come as a surprise since a number of IL methods leverage the idea of expert transitions in the replay buffer [START_REF] Reddy | Sqil: imitation learning via regularized behavioral cloning[END_REF][START_REF] Hester | Deep q-learning from demonstrations[END_REF].

PWIL-support

In this version of PWIL, we define the reward as a function of the following cost:

∀i ∈ [1 : T ], c i,π = inf γ1,...,γ D ∈R D j=1 γj ≤ 1 T D j=1 d((s π i , a π i ), (s e j , a e j )) 2 γ j .
We can interpret this cost in the context of Section 3 as the problem of moving piles of dirt of mass 1 T into holes of infinite capacity. It is thus reminiscent of methods that consider IL as a support estimation problem [START_REF] Wang | Random expert distillation: Imitation learning via expert policy support estimation[END_REF][START_REF] Brantley | Disagreement-regularized imitation learning[END_REF]. This leads to a significant drop in performance for Ant, Walker and HalfCheetah. Perhaps suprisingly, this boosts performance on Humanoid.

Related Work

Adversarial IL Similarly to our method, adversarial IL methods aim at matching the state-action distribution of the agent with the distribution of the expert, using different measures of similarity. For instance, GAIL [START_REF] Ho | Generative adversarial imitation learning[END_REF] considers the Shannon-Jensen divergence while AIRL [START_REF] Fu | Learning robust rewards with adversarial inverse reinforcement learning[END_REF] considers the Kullback-Leibler divergence. In the context of GANs, [START_REF] Arjovsky | Wasserstein generative adversarial networks[END_REF] show that replacing the f -divergence by the Wasserstein distance through the Kantorovich-Rubinstein duality [START_REF] Villani | Optimal transport: old and new[END_REF] leads to better training stability, which a number of methods in IL have leveraged [START_REF] Li | Infogail: Interpretable imitation learning from visual demonstrations[END_REF][START_REF] Lacotte | Risk-sensitive generative adversarial imitation learning[END_REF][START_REF] Kostrikov | Discriminator-actor-critic: Addressing sample inefficiency and reward bias in adversarial imitation learning[END_REF][START_REF] Xiao | Wasserstein adversarial imitation learning[END_REF][START_REF] Liu | State alignment-based imitation learning[END_REF]. However, the Kantorovich-Rubinstein duality only holds for the W 1 distance [START_REF] Peyré | Computational optimal transport[END_REF], its implementation comes with a number of practical approximations (e.g., weight clipping to ensure Lipchitz continuity). Although these methods are based on the Wasserstein distance through its dual formulation, they drastically differ from ours as they rely on a minmax optimization problem. Recent approaches such as DAC [START_REF] Kostrikov | Discriminator-actor-critic: Addressing sample inefficiency and reward bias in adversarial imitation learning[END_REF] demonstrate that adversarial IL approaches, while based on GAIL, can be sample efficient. However, they implicitly rely on a carefully tuned discriminator (e.g. network architecture, regularization strategy, scheduled learning rates, number of updates per environment interactions) that needs to interact well with a carefully tuned RL agent. This leads to potential brittleness, which is shown on Humanoid in Section 4. In contrast, the reward function we present relies on 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 [START_REF] Piot | Boosted and reward-regularized classification for apprenticeship learning[END_REF][START_REF] Schroecker | State aware imitation learning[END_REF][START_REF] Wang | Random expert distillation: Imitation learning via expert policy support estimation[END_REF][START_REF] Brantley | Disagreement-regularized imitation learning[END_REF][START_REF] Reddy | Sqil: imitation learning via regularized behavioral cloning[END_REF]. Note that the agent might stay on the support of the expert without recovering its state-action distribution. Soft Q Imitation Learning [START_REF] Reddy | Sqil: imitation learning via regularized behavioral cloning[END_REF] assigns a reward of 1 to all expert transitions and 0 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 [START_REF] Wang | Random expert distillation: Imitation learning via expert policy support estimation[END_REF] estimates the support by using a neural network trained on expert transitions whose target is a fixed random neural network, similarly to [START_REF] Burda | Exploration by random network distillation[END_REF]. Disagreement Regularized Imitation Learning [START_REF] Brantley | Disagreement-regularized imitation learning[END_REF] 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. Our method is similar in some sense to these methods since it is 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).

Conclusion

In this work, we present Imitation Learning as a distribution matching problem and introduce a reward function which is based on an upper bound of the Wasserstein distance between the state-action distributions of the agent and the expert. The reward function we introduce is 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 as 1 demonstration on all considered environments, even the challenging Humanoid.

A number of IL methods are developed on synthetic tasks, where the evaluation of the IL method can be done with the actual return of the task. We emphasize that in our work, we present a true measure of similarity between the expert and the agent, that we can thus use in real IL settings, that is in settings where the reward of the task cannot be specified.

We think that our work opens multiple directions of research: the state-only version of PWIL learns non-trivial behavior in all tasks we considered, and could thus be extended to changing environment dynamics settings. Another interesting direction is the setting where a metric MDP cannot be specified as naturally as in the tasks we considered, for example in visual based observations.

A Agent Architecture

The agent we use is DDPG [START_REF] Timothy P Lillicrap | Continuous control with deep reinforcement learning[END_REF] with a twin distributional critic [START_REF] Fujimoto | Addressing function approximation error in actor-critic methods[END_REF]. The actor architecture is a 4-layer neural network: the first layer has size 256 with tanh activation and layer normalization [START_REF] Lei Ba | Layer normalization[END_REF], the second layer and third layer have size 256 with elu activation [START_REF] Clevert | Fast and accurate deep network learning by exponential linear units (elus)[END_REF], the last layer is of size the dimension of the action space, with a tanh activation scaled to the action range of the environment. To enable sufficient exploration with use a Gaussian noise layer on top of the last layer with standard deviation σ = 0.2, that we clip to the action range of the environment. We evaluate the agent without exploration noise.

For the critic network we use a 4-layer neural network: the first layer has size 512 with tanh activation and layer normalization, the second layer is of size 512 with elu activation, the third layer is of size 256 with elu activation, the last layer is of dimension 101 with a softmax activation. The 101 neurons stand for the weights of the distribution supported within equal distance in the range [-150, 150], as a categorical distribution [START_REF] Marc G Bellemare | A distributional perspective on reinforcement learning[END_REF].

We use the Adam optimizer [Kingma and Ba, 2015] with = 10 -4 for both the critic and the actor. We use a batch size of 256. We clip both gradients from the critic and the actor to limit their L2 norm to 40.

We use a replay buffer of size 10 6 , a discount factor γ = 0.99 and n step returns with n = 5. We prefill the replay buffer with 50k 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 = 4 interactions with the environment. We did run a hyperparameter search on the following parameters:

Parameters

Values 10 -5 , 10 -4 , 10 -3 σ 0.1, 0.2, 0.3 k 2, 4, 8, 16

B Metrics in MDP

As PWIL is based on the Wasserstein distance, it relies on the existence of a metric in the MDP.

The problem of designing a metric in an MDP remains a challenge. Bisimulation metrics [START_REF] Ferns | Metrics for finite markov decision processes[END_REF][START_REF] Givan | Equivalence notions and model minimization in markov decision processes[END_REF] define the concept of bisimilarity in MDPs via the following recurrent definition: two states are bisimilar if they yield the same expected reward and transition to bisimulation equivalence classes with equal probability. Recent work from [START_REF] Gelada | Deepmdp: Learning continuous latent space models for representation learning[END_REF] and Castro [2020] extend bisimulation metrics to the function approximation setting. Previous work has also assumed the existence of a MDP metric for imitation learning through kernel classification [START_REF] Francisco | Learning from demonstration using mdp induced metrics[END_REF] and for policy regularization using the Wasserstein distance [START_REF] Amin | Reinforcement learning with wasserstein distance regularisation, with applications to multipolicy learning[END_REF]. We think that the problem of designing a metric and recovering the expert behavior through PWIL is an interesting direction for future work.

C PWIL Learning Curves

1 DT

 1 Algorithm 1 PWIL: Primal Wasserstein Imitation Learning Input: Expert demonstrations D = {s e j , a e j } j∈[1:D] , Agent A, number of episodes N for k = 1 to N do Copy expert demonstrations with weight: D := {s e j , a e j , w e j } j∈[1:D] , with w e j = Reset environment, initial state s for i = 1 to T do Take action a := π A (s), observe next state s Initialize weight w π := 1 Initialise cost c := 0 while w π > 0 do Compute s e , a e , w e := arg min (s e ,a e ,w e )∈D d((s, a), (s e , a e )) if w π ≥ w e then c := c + w e d((s, a), (s e , a e )) 2 w π := w π -w e D .pop(s e , a e , w e ) else c := c + w π d((s, a), (s e , a e )) 2 w e := w e -w π w π := 0 r := f (c) A observes (s, a, r, s ) A updates itself Update state s := s 4 Experiments

Figure 2 :

 2 Figure 2: Mean and standard deviation return of the evaluation policy over 10 rollouts and 5 seeds, reported every 10k environment steps. The return here is in term of the environment's original reward. Top row: 1 demonstration, bottom row: 11 demonstrations.

Figure 3 :

 3 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 5 seeds, reported every 10k environment steps. For PWIL, we include the upper bound on the Wasserstein distance based on the greedy coupling defined in Equation (4). Top row: 1 demonstration, bottom row: 11 demonstrations.

Figure 4 :

 4 Figure 4: Mean and standard deviation of the return of the DDPG agent while learning (over the 5 seeds). The return here is in term of the reward defined with Equation (6). Top row: 1 demonstration, middle row: 4 demonstrations, bottom row 11 demonstrations.

Figure 5 :

 5 Figure 5: Mean and standard deviation of the original environment returns of the evaluation policy over 10 rollouts and 5 seeds, reported every 10k environment steps. Top row: 1 demonstation, middle row: 4 demonstrations, bottom row: 11 demonstrations.

Figure 6 :

 6 Figure 6: 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 5 seeds, reported every 10k environment steps. For PWIL, we include the upper bound on the Wasserstein distance based on the greedy coupling defined in Equation (4). Top row: 1 demonstration, middle row: 4 demonstrations, bottom row: 11 demonstrations.

Figure 7 :

 7 Figure 7: Mean of the evaluation performance of different variants of PWIL over 10 rollouts and 5 seeds, reported every 10k environment steps. The return here is in term of the environment's original reward. Top row: 1 demonstation, middle row: 4 demonstrations, bottom row: 11 demonstrations.

Table 1 :

 1 Ablation study of PWIL. Evaluation performance of variants of PWIL trained for 1M steps (2.5M for Humanoid) on 11 demonstrations. The numbers are the averages and standard deviations of the returns for 50 rollouts (10 rollouts per seed).

		Walker2d-v2	Ant-v2	Humanoid-v2	HalfCheetah-v2
	Expert	3548.0 ± 28.8 6131.7 ± 801.1	4177.0 ± 71.4	8966.3 ± 82.1	8879.1 ± 75.4
	PWIL	3578.7 ± 14.3 5959.5 ± 142.5	4016.6 ± 69.3	6318.5 ± 2247.7 8975.4 ± 294.1
	PWIL-state	3542.6 ± 55.7 3038.1 ± 1760.8 3277.2 ± 662.4	6640.8 ± 1960.2 9057.1 ± 727.8
	PWIL-L2	3464.7 ± 46.2 4688.0 ± 2407.2 3212.2 ± 887.7	254.8 ± 95.3	2580.6 ± 2477.9
	PWIL-nofill	3535.6 ± 70.7 5251.7 ± 677.7	3624.1 ± 151.9	5597.8 ± 2357.9 8828.8 ± 137.2
	PWIL-support	3548.8 ± 40.6 3349.4 ± 1999.4 2626.0 ± 1175.9 7137.8 ± 1599.5 2285.2 ± 1260.1

Table 2 :

 2 DDPG hyperparameters search.For the reward function, we did run a hyperparameter search on α and β with the following values:

	Parameters	Values
	α	1, 5, 10
	β	1, 5, 10

Table 3 :

 3 Reward function hyperparameters search.