Source-linked AI summary
Cross-Domain Imitation Learning via Optimal Transport
Arnaud Fickinger, Samuel Cohen, Stuart Russell, Brandon Amos
TL;DR
Cross-domain imitation must compare expert and agent behavior even when their dynamics and state-action spaces differ, and prior mapping approaches may require proxy tasks. GWIL uses Gromov-Wasserstein distance for this comparison and characterizes when it preserves optimality. Experiments report learning from a single cross-domain demonstration without proxy tasks in continuous-control settings.
Problem
Imitation across agents with different embodiments, dynamics, or state-action spaces is difficult, while many existing approaches assume a shared domain or require proxy tasks.
Source: §1 · p. 1Method
GWIL uses Gromov-Wasserstein distance to compare expert and agent state-action distributions and optimize imitation without proxy tasks.
Source: §1 · pp. 1–2, §4.2 · p. 6Results
GWIL’s optimality conditions are characterized, and experiments report optimal behavior from one cross-domain demonstration without proxy tasks in continuous-control settings.
Source: §1 · p. 2, §4.1 · pp. 4, 6Takeaways & Limitations
Optimal transport distances provide a useful foundation for cross-domain imitation across incomparable spaces.
Source: §6 · p. 9GWIL recovers an optimal policy only up to isometries.
Source: §4.1 · p. 6
Abstract
from arXiv · showhide
Cross-domain imitation learning studies how to leverage expert demonstrations of one agent to train an imitation agent with a different embodiment or morphology. Comparing trajectories and stationary distributions between the expert and imitation agents is challenging because they live on different systems that may not even have the same dimensionality. We propose Gromov-Wasserstein Imitation Learning (GWIL), a method for cross-domain imitation that uses the Gromov-Wasserstein distance to align and compare states between the different spaces of the agents. Our theory formally characterizes the scenarios where GWIL preserves optimality, revealing its possibilities and limitations. We demonstrate the effectiveness of GWIL in non-trivial continuous control domains ranging from simple rigid transformation of the expert domain to arbitrary transformation of the state-action space.
1 INTRODUCTION
Cross-domain imitation learning seeks to train an agent from demonstrations produced by an expert with different dynamics or state-action spaces, while many existing approaches require proxy tasks. GWIL uses Gromov-Wasserstein distance to compare agents’ stationary state-action distributions, and the paper characterizes when this approach preserves optimality and evaluates it in continuous-control settings.
1 INTRODUCTION
- Cross-domain imitation learning uses expert demonstrations across agents with different embodiments or dynamics, unlike methods limited to shared domains and state-action spaces.This setting can broaden imitation beyond in-domain expert demonstrations.Source: §1 · p. 1
- Earlier cross-domain methods learn domain mappings from proxy tasks, restricting transfer to unseen experts or new robots.Those tasks pair demonstrations from experts acting optimally in both domains.Source: §1 · p. 1
- GWIL frames cross-domain imitation as optimal transport, using Gromov-Wasserstein distance between stationary state-action distributions as a pseudo-reward.The paper also proposes a benchmark and method that do not require proxy tasks.Source: §1 · pp. 1–2, Figure 1 · p. 2
- In Euclidean spaces, rotations, translations, and reflections preserve pairwise distances, the structural property illustrated for isomorphic policies.The figure describes these as isometric transformations.Source: Figure 2 · p. 2
- The paper characterizes when GWIL preserves optimality and reports learning optimal behavior from one cross-domain demonstration without proxy tasks in continuous-control settings.The experiments span non-trivial settings.Source: §1 · p. 2
2 RELATED WORK
Prior imitation-learning work includes supervised behavioral cloning, inverse-reward learning, and occupancy-measure matching via optimal transport. Cross-domain transfer methods often depend on proxy tasks or mappings, whereas GWIL compares different spaces without requiring a shared latent space or proxy tasks.
2 RELATED WORK
- Primal Wasserstein Imitation Learning (PWIL) and Sinkhorn Imitation Learning (SIL) treat imitation as optimal transport between occupancy measures, avoiding minimax optimization.The cited work reports improved sample efficiency over adversarial methods.Source: §2 · p. 2
- The work extends Dadashi et al. (2020) and Papagiannis & Li (2020) from Wasserstein to Gromov-Wasserstein imitation, moving beyond same-domain limitations to cross-domain agents living in different spaces.Source: §2 · p. 2
- Prior transfer methods use handcrafted features, paired demonstrations, or proxy tasks to learn cross-domain mappings, limiting their applicability.Stadie et al. (2017) avoid proxy tasks but are limited to small viewpoint mismatch.Source: §2 · pp. 2–3
- Unlike mapping-based approaches, GWIL compares different spaces directly without an explicit shared latent space or proxy tasks.The paper’s experiments use expert demonstrations without requiring demonstrations from both agents.Source: §2 · p. 3
3 PRELIMINARIES
The preliminaries define metric MDPs and introduce Gromov-Wasserstein distance for comparing metric measure spaces. GW compares within-space pairwise distances to find a best alignment between spaces.
3 PRELIMINARIES
- A metric MDP is an infinite-horizon discounted MDP equipped with a distance on its state-action space.The MDP tuple includes state and action spaces, reward, transition function, initial-state distribution, and discount factor.Source: §3 · p. 3
- Gromov-Wasserstein distance compares metric measure spaces through an optimal-transport minimum over couplings between their measures.The preliminaries introduce the distance for spaces with their own distances and measures.Source: §3 · p. 3
- GW compares pairwise distances within each space to find the best isometry between the spaces.The paper applies this comparison to the expert and agent state-action spaces and their stationary distributions.Source: §3 · p. 3
4 CROSS-DOMAIN IMITATION LEARNING VIA OPTIMAL TRANSPORT
The section defines policy comparison across different MDPs through Gromov-Wasserstein distance between occupancy-based metric spaces. It gives conditions for optimality preservation and turns GW minimization into a reward that can be optimized with reinforcement learning.
4.1 COMPARING POLICIES FROM ARBITRARILY DIFFERENT MDPS
- Policy GW compares expert and agent state-action metric-measure spaces using their occupancy measures, and defines a metric on isomorphic policy classes.Source: §4.1 · p. 4
- Under theorem 1’s conditions, minimizing GW recovers an agent-domain optimal policy up to isometry.The method’s recovery guarantee is limited to policies equivalent up to isometries.Source: §4.1 · pp. 4, 6
4.2 GROMOV-WASSERSTEIN IMITATION LEARNING
- Because GW minimization requires derivatives through transition dynamics, GWIL defines a reward proxy; an agent minimizing GW is optimal for that reward.Source: §4.2 · p. 6
- GWIL approximates occupancy measures with finite trajectories and uses the resulting pseudo-rewards to train a policy with reinforcement learning.For continuous state-action spaces, the paper uses soft actor-critic; Algorithm 1 uses a single expert trajectory.Source: §4.1 · p. 6, §4.2 · p. 7
- The reward-proxy construction applies to any occupancy measure and extends prior RL methods that optimize transport quantities under uniform occupancy measures.Source: §4.2 · p. 7
5 EXPERIMENTS
The experiments test GWIL across rigid transformations and agents with different state-action spaces. GWIL recovers optimal behavior in the maze and pendulum-to-cartpole tasks; cheetah-to-walker transfer yields either forward or backward behavior in the optimal policy’s isometry class.
5 EXPERIMENTS
- The benchmark asks whether GWIL recovers optimal behavior under a rigid transformation and under different state-action spaces, using simulated continuous-control tasks.The tasks use Mujoco and the DeepMind Control Suite.Source: §5 · p. 7
- Across the tasks, GWIL recovers optimal behavior in the reflected maze and pendulum-to-cartpole transfer, while cheetah-to-walker transfer yields forward or backward isometry-class behavior.In the cheetah-to-walker task, forward motion is optimal and backward motion is suboptimal.Source: §5.1 · p. 8, Figure 3 · p. 8, §5.2 · p. 8
5.1 AGENT DOMAIN IS A RIGID TRANSFORMATION OF THE EXPERT DOMAIN
- In the PointMass Maze, reflecting the expert’s domain produces an agent domain where GWIL recovers an optimal policy, as theorem 1 predicts.Every policy in the optimal policy’s isometry class is optimal in this maze.Source: §5.1 · p. 8, Figure 3 · p. 8
5.2 AGENT AND THE EXPERT HAVE SLIGHTLY DIFFERENT STATE AND ACTION SPACES
- For pendulum-to-cartpole transfer, GWIL recovers the optimal swing-up behavior despite the non-rigid transformation and different state-action dimensions.The pendulum has 3 dimensions and the cartpole has 5.Source: §5.2 · pp. 8–9, Figure 4 · p. 8
5.3 AGENT AND THE EXPERT HAVE SIGNIFICANTLY DIFFERENT STATE AND ACTION SPACES
- For cheetah-to-walker transfer, minimizing Gromov-Wasserstein yields forward or backward walker behavior depending on the seed.The forward element is optimal; the backward element is suboptimal, though the resulting walker behaves like a cheetah.Source: §5.3 · p. 9, Figure 5 · p. 9
6 CONCLUSION
The conclusion presents optimal-transport distances as a foundational tool for cross-domain imitation across incomparable spaces. It identifies scaling, sparse-reward exploration, and temporal structure as directions for further work.
6 CONCLUSION
- Optimal-transport distances are a useful foundational tool for cross-domain imitation across incomparable spaces.Source: §6 · p. 9
- Future work includes scaling transfer to more complex environments and agents, including high-dimensional demonstrations of complex tasks.Source: §6 · p. 9
- Another proposed direction is using Gromov-Wasserstein to help agents explore sparse-reward environments when demonstrations from other agents are available.Source: §6 · p. 9
- The paper proposes comparing Gromov-Wasserstein with other optimal-transport distances that allow more flexibility in space connections and coupling invariances.Source: §6 · p. 9
- Gromov-Wasserstein ignores temporal information and trajectory ordering, motivating metrics that account for MDP temporal structure.Source: §6 · p. 9
A OPTIMIZATION OF THE PROXY REWARD
The appendix evaluates whether the proxy reward provides an easily optimized learning signal. Across three environments, SAC converges quickly and consistently to asymptotic episodic return.
A OPTIMIZATION OF THE PROXY REWARD
- Across three environments and five seeds, SAC converges quickly and consistently to asymptotic episodic return using the proxy reward.Source: A OPTIMIZATION OF THE PROXY REWARD · p. 13, Figure 6 · p. 13
- The observed convergence gives reason to expect the proxy reward to be similarly easy to optimize in other cross-domain imitation settings.Source: A OPTIMIZATION OF THE PROXY REWARD · p. 13
- The proxy reward introduced in equation 9 provides a learning signal that standard RL algorithms can easily optimize.Source: Figure 6 · p. 13, A OPTIMIZATION OF THE PROXY REWARD · p. 13
B TRANSFER TO SPARSE-REWARD ENVIRONMENTS
GWIL supports learning in sparse-reward settings from a single expert demonstration in another domain. In the sparse maze, it performs competitively with a learner using a same-domain demonstration and Wasserstein distance.
B TRANSFER TO SPARSE-REWARD ENVIRONMENTS
- The experiments give the learner a sparse task reward alongside its optimal-transport proxy reward; SAC fails to learn meaningful behavior in either task.The tasks are a point-mass maze with reward for reaching its end and cartpole with reward for staying upright for 10 consecutive steps.Source: B TRANSFER TO SPARSE-REWARD ENVIRONMENTS · p. 13
- GWIL performs similarly to a same-domain Wasserstein learner in sparse-reward environments, despite using a demonstration from another domain.The comparison includes sparse maze and cartpole tasks; the reported competitive result is specifically for the sparse maze.Source: B TRANSFER TO SPARSE-REWARD ENVIRONMENTS · p. 13, Figure 7 · p. 14
C SCALABILITY OF GWIL
GWIL's runtime is comparable to Wasserstein imitation in the sparse maze, and a GWIL walker reaches walking speed faster than a SAC walker trained to run. Its coupling computation depends on trajectory length rather than state dimension.
C SCALABILITY OF GWIL
- The optimal-coupling computation depends on the number of trajectory time steps, not the expert or agent dimension, so the method naturally scales with problem dimension.The implementation uses no entropy regularizer; entropy-regularized methods could be incorporated to further improve scalability.Source: C SCALABILITY OF GWIL · p. 14
- A GWIL walker imitating a cheetah reaches walking speed faster than a SAC walker trained to run; maze imitation takes the same order of time as Wasserstein imitation.The walking-speed comparison uses horizontal velocity 1; the SAC agent also optimizes for standing.Source: C SCALABILITY OF GWIL · p. 14, Figure 9 · p. 15, Figure 8 · p. 15