Cross-Domain Imitation Learning via Optimal Transport

1 INTRODUCTION . Reinforcement learning ( RL ) methods have attained impressive results across a number of domains , e.g. , Berner et al . ( 2019 ) ; Kober et al . ( 2013 ) ; Levine et al . ( 2016 ) ; Vinyals et al . ( 2019 ) . However , the effectiveness of current RL method is heavily correlated to the quality of the training reward . Yet for many real-world tasks , designing dense and informative rewards require significant engineering effort . To alleviate this effort , imitation learning ( IL ) proposes to learn directly from expert demonstrations . Most current IL approaches can be applied solely to the simplest setting where the expert and the agent share the same embodiment and transition dynamics that live in the same state and action spaces . In particular , these approaches require expert demonstrations from the agent domain . Therefore , we might reconsider the utility of IL as it seems to only move the problem , from designing informative rewards to providing expert demonstrations , rather than solving it . However , if we relax the constraining setting of current IL methods , then natural imitation scenarios that genuinely alleviate engineering effort appear . Indeed , not requiring the same dynamics would enable agents to imitate humans and robots with different morphologies , hence widely enlarging the applicability of IL and alleviating the need for in-domain expert demonstrations . This relaxed setting where the expert demonstrations comes from another domain has emerged as a budding area with more realistic assumptions ( Gupta et al. , 2017 ; Liu et al. , 2019 ; Sermanet et al. , 2018 ; Kim et al. , 2020 ; Raychaudhuri et al. , 2021 ) that we will refer to as Cross-Domain Imitation Learning . A common strategy of these works is to learn a mapping between the expert and agent domains . To do so , they require access to proxy tasks where both the expert and the agent act optimally in there respective domains . Under some structural assumptions , the learned map enables to transform a trajectory in the expert domain into the agent domain while preserving the optimality . Although these methods indeed relax the typical setting of IL , requiring proxy tasks heavily restrict the applicability of Cross-Domain IL . For example , it rules out imitating an expert never seen before as well as transferring to a new robot . In this paper , we relax the assumptions of Cross-Domain IL and propose a benchmark and method that do not need access to proxy tasks . To do so , we depart from the point of view taken by previous work and formalize Cross-Domain IL as an optimal transport problem . We propose a method , that we call Gromov Wasserstein Imitation Learning ( GWIL ) , that uses the Gromov-Wasserstein distance 1Project site with videos : https : //sites.google.com/view/gwil/home to solve the benchmark . We formally characterize the scenario where GWIL preserves optimality ( theorem 1 ) , revealing the possibilities and limitations . The construction of our proxy rewards to optimize optimal transport quantities using RL generalizes previous work that assumes uniform occupancy measures ( Dadashi et al. , 2020 ; Papagiannis & Li , 2020 ) and is of independent interest . Our experiments show that GWIL learns optimal behaviors with a single demonstration from another domain without any proxy tasks in non-trivial continuous control settings . 2 RELATED WORK . Imitation learning . An early approach to IL is Behavioral Cloning ( Pomerleau , 1988 ; 1991 ) which amounts to training a classifier or regressor via supervised learning to replicate the expert ’ s demonstration . Another key approach is Inverse Reinforcement Learning ( Ng & Russell , 2000 ; Abbeel & Ng , 2004 ; Abbeel et al. , 2010 ) , which aims at learning a reward function under which the observed demonstration is optimal and can then be used to train a agent via RL . To bypass the need to learn the expert ’ s reward function , Ho & Ermon ( 2016 ) show that IRL is a dual of an occupancy measure matching problem and propose an adversarial objective whose optimization approximately recover the expert ’ s state-action occupancy measure , and a practical algorithm that uses a generative adversarial network ( Goodfellow et al. , 2014 ) . While a number of recent work aims at improving this algorithm relative to the training instability caused by the minimax optimization , Primal Wasserstein Imitation Learning ( PWIL ) ( Dadashi et al. , 2020 ) and Sinkhorn Imitation Learning ( SIL ) ( Papagiannis & Li , 2020 ) view IL as an optimal transport problem between occupancy measures to completely eliminate the minimax objective and outperforms adversarial methods in terms of sample efficiency . Heess et al . ( 2017 ) ; Peng et al . ( 2018 ) ; Zhu et al . ( 2018 ) ; Aytar et al . ( 2018 ) scale imitation learning to complex human-like locomotion and game behavior in non-trivial settings . Our work is an extension of Dadashi et al . ( 2020 ) ; Papagiannis & Li ( 2020 ) from the Wasserstein to the GromovWasserstein setting . This takes us beyond limitation that the expert and imitator are in the same domain and into the cross-domain setting between agents that live in different spaces . Transfer learning across domains and morphologies . Work transferring knowledge between different domains in RL typically learns a mapping between the state and action spaces . Ammar et al . ( 2015 ) use unsupervised manifold alignment to find a linear map between states that have similar local geometry but assume access to hand-crafted features . More recent work in transfer learning across viewpoint and embodiment mismatch learn a state mapping without handcrafted features but assume access to paired and time-aligned demonstration from both domains ( Gupta et al. , 2017 ; Liu et al. , 2018 ; Sermanet et al. , 2018 ) . Furthermore , Kim et al . ( 2020 ) ; Raychaudhuri et al . ( 2021 ) propose methods to learn a state mapping from unpaired and unaligned tasks . All these methods require proxy tasks , i.e . a set of pairs of expert demonstrations from both domains , which limit the applicability of these methods to real-world settings . Stadie et al . ( 2017 ) have proposed to combine adversarial learning and domain confusion to learn a policy in the agent ’ s domain without proxy tasks but their method only works in the case of small viewpoint mismatch . Zakka et al . ( 2021 ) take a goal-driven perspective that seeks to imitate task progress rather than match fine-grained structural details to transfer between physical robots . In contrast , our method does not rely on learning an explicit cross-domain latent space between the agents , nor does it rely on proxy tasks . The GromovWasserstein distance enables us to directly compare the different spaces without a shared space . The existing benchmark tasks we are aware of assume access to a set of demonstrations from both agents whereas the experiments in our paper only assume access to expert demonstrations . 3 PRELIMINARIES . Metric Markov Decision Process . An infinite-horizon discounted Markov decision Process ( MDP ) is a tuple ( S , A , R , P , p0 , γ ) where S and A are state and action spaces , P : S × A → ∆ ( S ) is the transition function , R : S×A→ R is the reward function , p0 ∈ ∆ ( S ) is the initial state distribution and γ is the discount factor . We equip MDPs with a distance d : S × A → R+ and call the tuple ( S , A , R , P , p0 , γ , d ) a metric MDP . Gromov-Wasserstein distance . Let ( X , dX , µX ) and ( Y , dY , µY ) be two metric measure spaces , where dX and dY are distances , and µX and µY are measures on their respective spaces2 . The Gromov-Wasserstein distance ( Mémoli , 2011 ) extends the Wasserstein distance from optimal transportation ( Villani , 2009 ) to these spaces and is defined as GW ( ( X , dX , µX ) , ( Y , dY , µY ) ) 2 = min u∈U ( µX , µY ) ∑ X 2×Y2 |dX ( x , x′ ) − dY ( y , y′ ) |2ux , yux′ , y′ , ( 1 ) where U ( µX , µY ) is the set of couplings between the atoms of the measures defined by U ( µX , µY ) = u ∈ RX×Y ∣∣∣∣∣∣ ∀x ∈ X , ∑ y∈Y ux , y = µX ( x ) , ∀y ∈ Y , ∑ x∈X ux , y = µY ( y ) . GW compares the structure of two metric measure spaces by comparing the pairwise distances within each space to find the best isometry between the spaces . Figure 1 illustrates this distance in the case of the metric measure spaces ( SE ×AE , dE , ρπE ) and ( SA ×AA , dA , ρπA ) . 4 CROSS-DOMAIN IMITATION LEARNING VIA OPTIMAL TRANSPORT . 4.1 COMPARING POLICIES FROM ARBITRARILY DIFFERENT MDPS . For a stationary policy π acting on a metric MDP ( S , A , R , P , γ , d ) , the occupancy measure is : ρπ : S ×A→ R ρ ( s , a ) = π ( a|s ) ∞∑ t=0 γtP ( st = s|π ) . We compare policies from arbitrarily different MDPs in terms of their occupancy measures . Definition 1 ( Gromov-Wasserstein distance between policies ) . Given an expert policy πE and an agent policy πA acting , respectively , on ME = ( SE , AE , RE , PE , TE , dE ) and MA = ( SA , AA , RA , PA , TA , dA ) . We define the Gromov-Wasserstein distance between πE and πA as the Gromov-Wasserstein distance between the metric measure spaces ( SE ×AE , dE , ρπE ) and ( SA ×AA , dA , ρπA ) : GW ( π , π′ ) = GW ( ( SE ×AE , dE , ρπE ) , ( SA ×AA , dA , ρπA ) ) . ( 2 ) 2We use discrete spaces for readability but show empirical results in continuous spaces . We now define an isometry between policies by comparing the distances between the state-action spaces and show that GW defines a distance up to an isometry between the policies . Figure 2 illustrates examples of simple isometric policies . Definition 2 ( Isometric policies ) . Two policies πE and πA are isometric if there exists a bijection φ : supp [ ρπE ] → supp [ ρπA ] that satisfies for all ( sE , aE ) , ( sE ′ , aE ′ ) ∈ supp [ ρπE ] 2 : dE ( ( sE , aE ) , ( sE ′ , aE ′ ) ) = dA ( φ ( sE , aE ) , φ ( sE ′ , aE ′ ) ) In other words , φ is an isometry between ( supp [ ρπE ] , dE ) and ( supp [ ρπA ] , dA ) . Proposition 1 . GW defines a metric on the collection of all isometry classes of policies . Proof . By definition 1 , GW ( πE , πA ) = 0 if and only if GW ( ( SE , dE , ρπE ) , ( SA , dA , ρπA ) ) = 0 . By Mémoli ( 2011 , Theorem 5.1 ) , this is true if and only if there is an isometry that maps supp [ ρφE ] to supp [ ρφA ] . By definition 2 , this is true if and only if πA and πE are isometric . The symmetry and triangle inequality follow from Mémoli ( 2011 , Theorem 5.1 ) . The next theorem3 gives a sufficient condition to recover , by minimizing GW , an optimal policy4 in the agent ’ s domain up to an isometry . Theorem 1 . Consider two MDPs ME = ( SE , AE , RE , PE , pE , γ ) and MA = ( SA , AA , RA , PA , pA , γ ) . Suppose that there exists four distances dSE , d A E , d S A , d A A defined on SE , AE , SA andAE respectively , and two isometries φ : ( SE , dSE ) → ( SA , dSA ) and ψ : ( AE , dSE ) → ( AS , dSA ) such that for all ( sE , aE , s ′ E ) ∈ SE ×AE × SE the three following conditions hold : R ( sE , aE ) = RA ( φ ( sE ) , ψ ( aE ) ) ( 3 ) PEsE , aE ( s ′ E ) = PAφ ( sE ) ψ ( aE ) ( φ ( s ′ E ) ) ( 4 ) pE ( sE ) = pA ( φ ( sE ) ) . ( 5 ) Consider an optimal policy π∗E in ME . Suppose that πGW minimizes GW ( π∗E , πGW ) with dE : ( sE , aE ) 7→ dSE ( sE ) + dAE ( aE ) and dA : ( sA , aA ) 7→ dSA ( sA ) + dAA ( aA ) . Then πGW is isometric to an optimal policy in MA . Proof . Consider the occupancy measure ρ∗A : SA ×AA → R given by ( sA , aA ) 7→ ρπ∗E ( φ −1 ( sA ) , ψ −1 ( aA ) ) . We first show that ρ∗A is feasible in MA , i.e . there exists a policy π ∗ A acting in MA with occupancy measure ρ∗A ( a ) . Then we show that π ∗ A is optimal in MA ( b ) and is isometric to π ∗ E ( c ) . Finally we show that πGW is isometric to π∗A , which concludes the proof ( d ) . ( a ) Consider sA ∈ SA . By definition of ρ∗A , ∑ aA∈AA ρ∗A ( sA ) = ∑ aA∈AA ρπ∗E ( φ −1 ( sA ) , ψ −1 ( aA ) ) = ∑ aE∈AE ρπ∗E ( φ −1 ( sA ) , aE ) . Since ρπ∗E is feasible in M , it follows from Puterman ( 2014 , Theorem 6.9.1 ) that∑ aE∈AE ρπ∗E ( φ −1 ( sA ) , aE ) = pE ( φ −1 ( sA ) ) + γ ∑ sE∈SE , aE∈AE PEsE , aE ( φ −1 ( sA ) ) + ρπ∗E ( sE , aE ) . By conditions 4 and 5 and by definition of ρ∗A , pE ( φ −1 ( sA ) ) + γ ∑ sE∈SE , aE∈AE PEsE , aE ( φ −1 ( sA ) ) + ρπ∗E ( sE , aE ) = pA ( sA ) + γ ∑ sE∈SE , aE∈AE PAφ ( sE ) , ψ ( aE ) ( sA ) + ρ ∗ A ( φ ( sE ) , ψ ( aE ) ) = pA ( sA ) + γ ∑ s′A∈SA , aA∈AA PAs′A , aA ( sA ) + ρ ∗ A ( s ′ A , aA ) . 3Our proof is in finite state-action spaces for readability and can be directly extended to infinite spaces . 4A policy is optimal in the MDP ( S , A , R , P , γ , d ) if it maximizes the expected return E ∑∞ t=0R ( st , at ) . It follows that∑ aA∈AA ρ∗A ( sA ) = pA ( sA ) + γ ∑ s′A∈SA , aA∈AA PAs′A , aA ( sA ) + ρ ∗ A ( s ′ A , aA ) . Therefore , by Puterman ( 2014 , Theorem 6.9.1 ) , ρ∗A is feasible in MA , i.e . there exists a policy π ∗ A acting in MA with occupancy measure ρ∗A . ( b ) By condition 5 and definition of ρ∗A , the expected return of π ∗ A in MA is then∑ sA∈SA , aA∈AA ρ∗A ( sA , aA ) RA ( sA , aA ) = ∑ sA∈SA , aA∈AA ρ∗E ( φ −1 ( sA ) , ψ −1 ( aA ) ) RE ( φ −1 ( sA ) , ψ −1 ( aA ) ) = ∑ sE∈SE , aE∈AE ρ∗E ( sE , aE ) RE ( sE , aE ) Consider any policy πA in M ′ . By condition 5 , the expected return of πA is∑ sA∈SA , aA∈AA ρπA ( sA , aA ) RA ( sA , aA ) = ∑ sE∈SE , aE∈AE ρπA ( φ ( sE ) , ψ ( aE ) ) RE ( sE , aE ) . Using the same arguments that we used to show that ρ∗A is feasible in M ′ , we can show that ( sE , aE ) 7→ ρπA ( φ ( sE ) , ψ ( aE ) ) is feasible in M . It follows by optimality of π∗E in M that∑ sE∈SE , aE∈AE ρπA ( φ ( sE ) , ψ ( aE ) ) RE ( sE , aE ) ≤ ∑ sE∈SE , aE∈AE ρπ∗E ( φ ( sE ) , ψ ( aE ) ) RE ( sE , aE ) = ∑ sA∈SA , aA∈AA ρ∗A ( sA , aA ) RA ( sA , aA ) . It follows that π∗A is optimal in M ′ . ( c ) Notice that ξ : ( sE , aE ) 7→ ( φ ( sE ) , ψ ( aE ) ) is an isometry between ( SE ×AE , dE ) and ( SA ×AA , dA ) , where dE and dA and given , resp. , by ( sE , aE ) 7→ dSE ( sE ) + dAE ( aE ) and ( sA , aA ) 7→ dSA ( sA ) + dAA ( aA ) . Therefore by definition of ρ∗A , π ∗ A is isometric to π ∗ E . ( d ) Recall from the statement of the theorem that πGW is a minimizer of GW ( π∗E , πGW ) . Since π∗A is isometric to π∗E , it follows from prop . 1 that GW ( π∗E , π∗A ) = 0 . Therefore GW ( π∗E , πGW ) must be 0 . By prop . 1 , it follows that there exists an isometry χ : ( supp [ ρ∗E ] , dE ) → ( supp [ ρπGW ] , dA ) . Notice that χ ◦ ξ−1|supp [ ρ∗A ] is an isometry from ( supp [ ρ ∗ A ] , dA ) to ( supp [ ρπGW ] , dA ) . It follows that πGW is isometric to π∗A , an optimal policy in MA , which concludes the proof . Remark 1 . Theorem 1 shows the possibilities and limitations of our method . It shows that our method can recover optimal policies even though arbitrary isometries are applied to the state and action spaces of the expert ’ s domain . Importantly , we don ’ t need to know the isometries , hence our method is applicable to a wide range of settings . We will show empirically that our method produces strong results in other settings where the environment are not isometric and don ’ t even have the same dimension . However , a limitation of our method is that it recovers optimal policy only up to isometries . We will see that in practice , running our method on different seeds enables to find an optimal policy in the agent ’ s domain . Algorithm 1 Gromov-Wasserstein imitation learning from a single expert demonstration . Inputs : expert demonstration τ , metrics on the expert ( dE ) and agent ( dA ) space Initialize the imitation agent ’ s policy πθ and value estimates Vθ while Unconverged do Collect an episode τ ′ Compute GW ( τ , τ ′ ) Set pseudo-rewards r with eq . ( 7 ) Update πθ and Vθ to optimize the pseudo-rewards end while

Summary: This paper frames cross-domain imitation learning as an optimal transport problem using the Gromov-Wasserstein distance. This problem is highly relevant to imitation learning settings where there is often substantial domain mismatch between action and state spaces, eg. a humanoid robot learning to walk from a human demonstrator. The paper introduces a reward function that can be optimised and proves that this is equivalent to minimising the Gromov-Wasserstein distance between state action occupancies of an agent and expert. Substantial discussion/ proofs are included to show that minimising the Gromov-Wasserstein distance is equivalent to recovering an optimal policy up to an isometry. This is both a blessing and a curse, as it allows for optimal policies to be recovered under extreme changes in domain or differences, but does mean that recovered policies could be entirely unsuitable due to isometry.   The paper is well written and concisely written, although does get excessively mathy at times, when a figure could be more helpful. Experimental results corroborate the proofs and propositions, and highlight the value of the proposed approach.