Latent Wasserstein Adversarial Imitation Learning
Authors: Siqi Yang
Kai Yan
Alex Schwing
Yu-Xiong Wang
Venue: ICLR 2026 Poster (2026)
Keywords: adversarial imitation learning
wasserstein distance
latent state space
Primary area: reinforcement learning
TL;DR: We propose a Wasserstein adversarial imitation learning method with ICVF-learned metric for imitation learning from observation.

Abstract: Imitation Learning (IL) enables agents to mimic expert behavior by learning from demonstrations. However, traditional IL methods require large amounts of medium-to-high-quality demonstrations as well as actions of expert demonstrations, both of which are often unavailable. To reduce this need, we propose Latent Wasserstein Adversarial Imitation Learning (LWAIL), a novel adversarial imitation learning framework that focuses on state-only distribution matching. It benefits from the Wasserstein distance computed in a dynamics-aware latent space. This dynamics-aware latent space differs from prior work and is obtained via a pre-training stage, where we train the Intention Conditioned Value Function (ICVF) to capture a dynamics-aware structure of the state space using a small set of randomly generated state-only data. We show that this enhances the policy's understanding of state transitions, enabling the learning process to use only one or a few state-only expert episodes to achieve expert-level performance. Through experiments on multiple MuJoCo environments, we demonstrate that our method outperforms prior Wasserstein-based IL methods and prior adversarial IL methods, achieving better results across various tasks.

Published as a conference paper at ICLR 2026
LATENT WASSERSTEIN ADVERSARIAL LEARNING
IMITATION
Kai Yan University of Illinois Urbana-Champaign
Siqi Yang
Alexander G. Schwing
Yu-Xiong Wang
deiglyang, kaiyan3, aschwing, yxw] www.iioniiei...o https:/ // /ttp:///ggtttt.btteegg hat.:b...oocccctyyaggg5/////a
ABSTRACT
Imitation Learning (IL.) enables agents to mimic expert behavior by learning from demonstrations. However, traditional methods require large amounts of medium-to-hin -uuiiih--ault demonstrations well as actions of expert demon- strations, both of which are often unavailable. To reduce this need, we propose Latent Wasserstein Adversarial Imitation Learning (LWAIL), novel adversar- ial imitation learning framework that focuses on state-only distribution matching. It benefits from the Wasserstein distance computed in dynamics-awe -ware latent space. This dynamics-aware latent space differs from prior work and is obtained via pre-training stage, where we train the Intention Conditioned Value Function (ICVF) to capture a dynamics-awand structure of the state space using a small set of randomly generated state-only data. We show that this enhances the policy's s understanding of state transitions, enabling the learning process to use only one or few state-only expert episodes to achieve expert-level performance. Through experiments on multiple MuJoCo environments, we demonstrate that our method outperforms prior Wasserstein- based methods and prior adversarial IL meth- ods, achieving better results across various tasks.
INTRODUCTION
As powerful tool for solving sequential decision-making problems, Reinforcement Learning (RL. has achieved remarkable success in recent years across various fields, such as gaming (Silver al., 2016) and training large language models (Ramamurthy et al., 2023; Guo et al., 2025). However, RL relies heavily on well-defined reward signals (Jain et al., 2021; Li et al., 2021), which can be difficult to obtain in real-world settings (e.g., robot control (Ibarz et al., 2021) with varied target tasks) or may require careful, enviromment-speciff considerations (Yu et al., 2020).
Imitation Learning (IL) offers compelling alternative by learning directly from expert demonstra- tions, thus bypassing the need for explicit reward design. Within IL, a major challenge are the frequently missing expert actions, encouraging adoption of Imitation Learning from Observations (LfO),, which uses only sequences of expert states. However, even these state-only demonstrations are often expensive amount of expert data. Adversarial Imitation Learning (AIL) methods (Torabi et al., 2018b; 8b; Zhu al., 2020), which learn by matching the distribution of agent states that of the expert, are a pop- ular approach to LfO. Many such AIL methods measure the distribution difference by f-divergence, such as KL. 2- JS-divergence. However, f f-divergences require "distribution coverage", i.e., the distributions in f-divergence must be on the same support set to avoid numerical error. This could lead to additional theoretical constraints, especially when an additional dataset of non-expert data is involved; for example, offline LfO methods such SMODICE (Ma al., 2022) require the non-expert state distribution to fully cover the expert state distribution, which does not necessarily hold in practice especially when the non-expert data quality is low (e.g., random).
acquire, making crucial to develop methods that can learn from a minimal
To address these limitations, the Wasserstein distance has gained popularity (Arjovsky et al., 2017; Zhang et al., 2020). However, many prior Wasserstein IL works that employ the Kantorovich- Rubinstein (KR) dual (Zhang et al., 2020; Garg et al., 2021; Sun et al., 2021) overlook an important issue: the distance metric between individual states is rather simplistic. Indeed, to compute the

---

Published as a conference paper at ICLR 2026
State A
Random (s.z)
Expert (2,4)
Agent
Env
(b()
+
(8))
r(s,
(5,55)
(5)
State B
Expert State
f(v)
(())))
(s)
Random (s.5)
State B State State
Expert (s.
f()) 1 fG) f(-) if
(a) Better distance metric
(b) Pre-training stage
(c) Imitation stage
motivation for better distance metric and outline of the algorithm. Panel (a) illustrates a case where the Euclidean distance between states not a good metrie: state is closer to expert state C, but it is apparently less desirable than more distant state A, as it cannot reach C. To fix, we use random data and ICVF to find more meaningful embedding space, as pre-training
Figure I: Illustrating
shown in the lower half. Panel (b), together with the lower half of panel (a), shows stage: we first train ICVF to obtain Φ(s), which serves as a reward for our agent in the online stage shown in panel (c). Fire indicates trainable modules, and snowflakes indicate frozen modules.
Wasserstein distance between distributions of states, a "ground metric" cost for measuring the distance between any two individual states is first defined. For this, the Euclidean distance is com- mon (Zhang et al., 2020; Garg et al., 2021). However, it fails to capture the environment's dynamics (see Fig. ((a for an illustration). This deficiency can severely mislead the learning process. Exist- ing solutions bypass the issue by using the primal Wasserstein distance, which however introduces other complexities and requires surrogates (Yan et al., 2024a; Dadashi et al., 2021). This raises a critical question: Can learn a distance metric that encodes the environment's 's dynamics from minimum amount of state-only low-quality data, allowing an agent to learn efficiently from just few, state-only expert trajectories?
To answer this question, we introduce a novel Wasserstein AlL algorithm that directly addresses this foundational issue. We propose two-stage process. First, the pre-training stage, we leverage small (1% of online rollouts) number of unstructured, low-quality (e.g., random) state-only data to train an CeenttionCCnnitioonne Value Function (ICVF) (Ghosh et al., 2023). The resulting em- bedding captures a rich, dynamics-awarc notion of reachability between states. Second, during the imitation stage, we freeze this ICVF embedding and use the Euclidean distance in this new latent space as the cost function within a standard Wasserstein AIL framework. By modifying the agent and discriminator to operate in this learned space, we fix the core issue of prior methods.
We find that this simple approach allows agent to effectively match the expert's state distribution, enabling highly efficient imitation from minimal expert data. We validate our approach on Umaze and challenging locomotion tasks in the MuJoCo environment from the D4RL benchmark (Fu al., 2020), and show that the latent space grasps the transition dynamics much better than the vanilla Euclidean distance. In terms of imitation learning, we achieve strong results using only a single trajectory of state-based expert data.
We summarize contributions follows: 1) dynamics- -wwree ground metric even if only small amount of (possibly low-quality) state-only data is available; this is first direct remedy for the known geometric limitations of the Euclidean distance in prior Wasserstein occupancy matching methods. 2) We propose a simple but effective improvement to current Wasserstein AlL methods achieving expert-level performance with a single state-only expert trajectory; 3) We empirically show that our method outperforms various baselines on multiple testbeds, proving that a better distance metric can greatly benefit AIL.
We show that the ICVF latent space captures
2 PRELIMINARIES
Markoy Decision Process (MDP). A Markov Decision Process (MDP) is a mathematical frame- work which describes the interactions of agent with environment at discrete time steps. It is
2

---

Published as a conference paper at ICLR 2026
defined by the tuple A. where S represents the state space, and A denotes the action space. The state transition probability function P(s"|s, a) defines the likelihood of transitioning to state s' E S after taking an action a A in the current state S. The reward function R(s.a) a) E IR specifies the immediate reward received after taking action a in state 8. 7 E [0, 1) is the discount factor, determining the importance of future rewards relative immediate ones.
At each time step t, the agent observes the current state 8, E S, selects an action ag E A, receives a reward r; st; Ør), and transitions to the next state s according to the transition function P. A complete running process is called a Irajectory. The goal of the agent is to learn a policy π S + A that maximizes the expected cumulative discounted reward G. = Ek= 7 7TT4+ In this paper, we focus on the state and state occupancy, which are the visitation frequency of states and state- pairs. Given policy æ, the state occupancy is defined as dg (s) = (1 - Ei= o 'Pr(ss = 8) and the state-pair occupancy is given by d% (s, s') = (1 (1-7) Ei=o 7 Pr(st
5t-1 = s').
Wasserstein Distance. Wasserstein distance, also known as Earth Mover's Distance (EMD) (Kan- torovich, 1939), is widely used to measure the distance between two probability distributions. For the metric space (M, c) where M is set and c : M x M + IR is metric, the I-Wasserstein distance' between two distributions p(z) and (2z) on the metric space (M, c) is defined as:
Wi inf / c( v) dII v).
(1(p4) MXM
Intuitively, this equation quantifies the optimal way to "move" mass from p while minimiz- ing the total movement, as described by the joint distributions [p,, q) with marginals and g. A more popular form adopted by the machine learning community is the Kantorovich-Rubinstee dual (Kantorovich & Rubinstein, 1958) of the I-Wasserstein distance, which reads as follows:
(KR)
(i q) = sup (E - E [f())))
11
Here, ll S 1 restricts function f to be 1-Lipschitz, i.c., for any x, x', [(=))((')
< 1. As the most prominent way to compel Lipschitzness is regularization of the gradient (Gulrajani et al., 2017: Stanczuk et al., 2021) (i.e., ff(x)) ---=zol limits the distance metric to be Euclidean (Stanczuk et al., 202 1), which is often undesirable (Yan
o(2,z")
[(x)=/(20) for local z), the 1-Lipschitz constraint inherently
et al., 2024a). In this paper, we fix this issue by introducing an ICVF-learned distance metric.
Intention Conditioned Value Function (ICVF). ICVF (Ghosh et al., 2023) is an embedding algo- rithm which learns an augmented value function (ss, 84, 2). It is defined as the unnormalized state occupancy of the future state 54 S starting from state E S with policy nu that is optimal for reaching some intention (i.e., goal state) z E S. Formally,
20
o <'((s: = =$4 so = 8
V(s, S4, 2) ETz
Σ
s0==s++o $ = 1'pss==s+ke
= -E
Σ
(3)
where (1)-))))) is the transition probability from 8; to 81+1 when acting according to T., which can be seen pseudo-policy that takes the next state "pseudo-action". I(s; = =8+) serves set of pseudo-reward labels: its value is 1 if the condition is true, and 0 otherwise. Thus, V with all different Te, Ps and $ can be jointly learned from a small, random state-only dataset with offline RL, such as IQL (Kostrikov et al., 2022).
With such value function learned, the next step is extract the dynamic-aware embedding. To easily do so, the value function is structured as follows:
(4) Here, @o(s) E Rd is the state representation that maps a state into a latent space, To(z) E Rdxd the matrix of intention, and (s4) E Ed is the outcome representation (see Appendix for details). As we later prove in Sec. 3.2, under certain assumptions, the state-pair occupancy of the expert policy is approximately a linear combination of do (s), which justifies our choice of using the ICVF embedding. This design is also empirically validated in Sec. 3.2 and Sec. 4.
Va 84, (o(s))TT TT (2))bo(4))
Unless otherwise specified, we will discuss W-Wasserstein distance in this paper.
3

---

Published as a conference paper at ICLR 2026
3 LATENT WASSERSTEIN ADVERSARIAL IMITATION LEARNING (LWAIL)
This section is organized as follows: In Sec. 3.1 we first define our goal and frame it using a Wasser- stein adversarial state occupancy matching objective with KR duality. We then point out its inherent shortcomings and propose the ICVF-trained latent space metric as a solution in Sec. 3.2. Finally, we introduce our algorithm in Sec. 3.3. See Fig. for an overview of our work.
3.1 WASSERSTEIN ADVERSARIAL. STATE OCCUPANCY MATCHING
Our goal is to learn a policy x using three sources of information: few-shot, state-only expert dataset E, another small dataset with state-only random transitions (s,s') (either given collected with an untrained policy), and online interactions. Inspired by recent state occupancy matching works (Kostrikov et al., 2020; Garg et al., 2021; Ma et al., 2022; Kim et al., 2022a), here, we minimize the I- I-Wasserstein distance between state-pair occupancy distributions of the policy т,, i.e., dge(s, s"), and of the empirical policy of the expert, i.e., d (s, s'). Formally, we address min Wi(d (8, 5',, (8,
(5) where s and s' are adjacent states in trajectory. Note, while many occupancy matching works, such as SMODICE (Ma et al., 2022) and LobsDICE (Kim et al., 2022a), use f-divergences, we opt to use the I-Wasserstein distance because it provides a smoother measure and leverages the underlying geometric property of the state space, unlike f-divergences.
However, the Wasserstein distance itself is hard to compute as it is inherently constrained linear programming problem (see Eq. (I)), which is difficult to solve via gradient descent. While there exist workarounds such as convex regularizers (Yan et al., 2024a), surrogates (Dadashi al., 2021), I,, and direct matching of trajectories (Luo et al., 2023; Bobrin et al., 2024), here, we choose the widely adopted KR dual (Kantorovich & Rubinstein, 1958) as our objective. Combined with policy optimization, this results in the final objective
min max (B(s,"))dyy, (fs)s')] (e (f(s, I 1251
(6)
where the I-Lipschitz constraint can be encouraged by prominent methods such as gradient regular- ization (Gulrajani et al., 2017). With constraint 'addressed," Eq. (6) is a bi-level optimization and can be optimized iteratively by any RL algorithm. Specifically, since d is independent of n, the objective for finding policy π is max, )~ddy, [-f(s,
I-f f=s((sss)) s')]. This can be optimized with any RL
algorithm using reward r(s, a) = (ss)) s')]. From an adversarial imitation learning perspective, I [s, 8') can be interpreted as a discriminator that outputs a high score f(s,, 8') for expert state pairs and low score 8s,,'') for non-expert ones.
3.2 DYNAMICS-AWARE DISTANCE METRIC WITH ICVF EMBEDDING SPACE
While the objective in Eq. (6) already provides a viable solution, it has a subtle limitation: As mentioned in Sec. 2, the metric c(s, 5') is limited be Euclidean in practice due to the Lipschitz f(>'))
constraint c(s.5) < 1. However, as illustrated in Fig. 1, a Euclidean distance in the raw state space often fails to capture the true relation between states. This subtle reliance on the Fuclidean distance is often ignored by prior KR duality-based imitation learning methods, c.g., IQlearn (Garg et al., 2021) and WDAIL (Zhang et al., 2020), but crucial for performance (Yan et al., 2024b).
To address this issue, we aim to learn a latent state representation, in which the Buclidean distance serves more effective metric that permits capturing the environment's dynamics and relationship between states from a small amount of randomiy collected, unlabeled data, even without access to ground-truth actions and rewards. To achieve this goal, we benefit from the ICVF (Ghosh et al., 2023) framework. We found that using as the cost c(s, s') = ())) dhe (s') |l the root mean squared difference between state representations in latent space, i.e., (s)s), ), is not only empirically effective as validated in Sec. 4, but it also theoretically benefits learning. More specifically, we have:
Theorem 3.1. In determmiemiiiit MDP with 7 1, For any converged policy n. learned in ICVF with goal z, there exists a vector n such that, for any adjacent (s,s') E I, d (s, s') se n² (s))), ie., the state-pair occupancy is approximately linear combination of
4

---

Published as a conference paper at ICLR 2026
State Space
Latent Space
State Space
Latent Space
10
Dimmasstor
Dimunston
Dimonntoo
-30 <5
10
23
30
25
Drmersion
Gimensiio
Dimension
Dimmession
(a) HalfCheetah
(b) Walker2d
Figure 2: t-SNE visualization of the same trajectory in the original state space and the embedding latent space. The color of the points represents the ground-truth reward of the state (greener is higher). States connected by lines are adjacent in the trajectory. We observe that an ICVF-trained embedding provides a much more dynamics-aware metric than the Euclidean distance. For Hopper and Ant environment embeddings, see Appendix F.4.2.
We defer the proof to Appendix Intuitively, the idea is that the value function V(s, s', 2), our optimization target d (s, s"), and the expert state-pair occupancy d (s, s'), all model the proba- bility of visiting s' after 8 with policy иz. As they share the same underlying linear structure of Φo. which captures the environment's dynamics, learning in the space of dhe reduce the learning difficulty by aligning via ICVF the computation structure of Wasserstein optimization and transition dynamics. The benefits of an alignment have also been observed other contexts (Xu al., 2020).
To further show the effectiveness of the distance metric c(s, s') || (s) dgg (s') ||2, we provide a t-SNE (Van der Maaten & Hinton, 2008) visualization of the same trajectory both the raw state space and the latent space in Fig. 2 using environments we evaluated in Sec. 4.2. The result shows that the latent space better captures the dynamic relationship between states. This finding highlights that the Euclidean distance in the latent space is more suitable metric for Wasserstein distance- based state matching. We also provide ablations in Sec. 4.4 and Appendix
3.3 LWAIL
We introduce proposed method, LWAIL, which consists of two stages: pre-training and imitation. We will also provide pseudo-code to summarize our approach.
Pre-training. The pre-training stage consists of three steps. First, collect random transition data from the environment by using a randomly initialized policy. This can be skipped if a random dataset is available, e.g., in setting following Yue et al. (2024). Second, train the value function via the loss given in Eq. (11) using IQL, and retrieve the projection function é from the value function formulated in Eq. (4). Third, train f(-,-) using the following objective with frozen latent variable mapping & and frozen, untrained, random policy क.
(E(,,)))dz (f((s)) ((((()()(())))))))) ((s')) - (e.e
11
Here, f serves the discriminator (from an adversarial learning perspective). the reward func- tion for the policy in later imitation (from an RL perspective), and the KR dual function (from Wasserstein perspective). To encourage the Lipschitz constraint, f is trained with gradient penalty, following WGAN-GP (Gulrajani et al., 2017).
Imitation. the online imitation learning stage, we again freeze the learned embedding & and replace and s' with their latent space representations, @08) and d(s'). Then the imitation learning problem in Eq. (6) can be addressed via
min max (IE (E(,,)))<d [f(((s), o(s))) -oaa ((6()) I1
(8)
Following the standard off-policy approach, the agent interacts with the environment to gather data and iteratively updates the value function and policy. Once the policy has collected a batch of tra- jectories, update the discriminator network f based on Eq. (8). Following Sec. 3.1, using adversarial learning framework, we then use ff to generate rewards for the downstream reinforce- ment learning algorithm, for which we employ TD3 (Fujimoto et al., 2018), robust, off-policy
5

---

Published as a conference paper at ICLR 2026
Algorithm LWAIL
Require:
State-only expert dataset E, state-action random dataset I (optional), initial policy ,, discriminator /. replay buffer B, update frequency m
Pretrain:
1: Collect transitions into buffer B with random actions or use random dataset provided
2: Use ICVF to pre-train the representation network ø (Eq. (11))
3: Pre-train f with initial policy π (inner level of Eq. (8))
4: Imitation:
5: while t S do
6: Collect transitions (s, a, s', done) using
7: Get pseudo-reward rp o(-((é(s)) ((s'))
8: Add (s, s', r.p. done) replay buffer B
9: if t mod m = 0 then
Update (inner level of Eq. (8)) 11: end if
Sample mini-batch of N transitions from B to perform TD3 update end while
13:
reinforcement learning method selected due its stability and effectiveness. Slightly different from Sec. 3.1 however, the reward for the TD3 policy is defined as r(s, s') = o (m(s), o(s"). is the sigmoid function that normalizes the reward to the range [0, 1], stabilizing the downstream RL algorithm. ())(s), o(s') is 1-sample estimation of (g/"P()")s..) |- ((((s) ((s')]] for transi- tion (s, s¹), following Kim et al. (Kim et al., 2022a). f and π are then iteratively updated until the policy æ is properly trained. The entire procedure of our method is summarized in Alg.
4 EXPERIMENTS
In this section, we assess the efficacy of LWAIL across multiple benchmark tasks. Specifically, we study the following questions: 1) How is our assigned reward f(((s), o(s')) different from the ground-truth reward? 2) Can our algorithm work well on complicated continuous control environ- ments? 3) How much the ICVF embedding LWAIL contribute its performance, and is LWAIL robust environment noise? We will answer 1) in Sec. 4.1, 2) in Sec. 4.2, and 3) in Sec. 4.4.
4.1 SIMPLE ENVIRONMENT ON MAZE2D
We first evaluate LWAIL in the Umaze environment, which offers intuitive visualizations of how I.A.I.. effectivlly learns reward representations for downstream tasks.
Experimental and Dataset Setup. We the maze2d-umaze-vo environment from D4RL (Fu et al., 2020), where 2D point mass moves from a random start to specific target at coordinates (1,1). Observations consist of (zt, y) positions and velocities; actions are linear forces in æ, y. The sparse reward is 1 when the goal is reached (distance < 0.5m). The ICVF model is trained on the D4RL random dataset, followed by Wasserstein learning with Euclidean and latent-space distances.
Results. After convergence, the reward map is compared in Fig. (b) and (c). The state used to calculate rewards is the position with zero velocity. The environment walls are shown orange. These results indicate that the ICVF-lcarned metric captures trajectory dynamics, enhancing reward feedback during online inverse RL exploration. The reward curve in Fig. (d) further shows that our method converges effectively on Maze2d, surpassing TD3 with ground-trui sparse rewards.
4.2 MUJoCo ENVIRONMENT
Baselines. We test variety of baselines in this section, which can be categorized into four types: 1) classic imitation methods, including GAIL (Ho & Ermon, 2016). AIRL (Fu et al., 2018) and the plain Behavior Cloning (BC); 2) Wasserstein-ban based imitation methods, including PWIL (Dadashi et al., 2021), WDAIL (Zhang et al., 2020) and IQlearn (Garg et al., 2021); 3) LfO methods, including BCO (Torabi al., 8a), GAIfO (Torabi et al., 2018b), DIFO (Huang al., 2024), LS-IQ (Al-

---

Published as a conference paper at ICLR 2026
Normal Space
Lasent Space
Mazee Pert momaane
45
(a) Umaze
(b) Normal Reward
(c) Learned Reward
(d) Maze Performance
Figure 3: Results Maze2D, Panel (a) shows the environment setup with the start (green) and the goal (red). Panel (b) and (c) illustrate the reward distribution without and with ICVF embedding
(orange represents the wall). Panel (d) shows the learning curves: LWAIL TD3 with ground- truth rewards. We note that the reward is more dynamics-aware with LWAIL.
Hopper
Halfcheeteh
160
100
50
60
40
32
24
05
20
1.0
0.2
0.4
0.6
0.8
1.0
0.0
0.
24
0.6
0.1
1.5
6.
02
0.4
Staps
PWTL
IQteam
GAIO
LINIIL (ours)
WDAL
BC
OPCLO
Figure 4: Performance on the MuJoCo environments. Overall, our method consistently delivers strong performance across all tasks. For LWAIL, the number of gradient steps matches the number of online interaction samples. For clarity, we present training curves for main baselines here; please refer to Fig. for the complete version.
Hafez et al., 2023), DiffAIL (Wang et al., 2024), OPOLO (Zhu et al., 2020) and DACFO (LfO variant of DAC (Kostrikov al., 2019); serves as baseline for OPOLO (Zhu al., 2020)); offline to online imitation learning, which includes OLLIE (Yue et al., 2024). Some methods, such GAIL. and WDAIL require expert action. For these methods, we report the results with extra access to expert actions. We report the mean and standard deviation from 5 independent runs with different seeds. The performance is measured by the D4RL-defined normalized reward (higher better).
Experimental and Dataset Setup. We evaluate method on four standard MuJoCo (Todorov et al., 2012) environments: Hopper, HalfCheetah, Walker2D, and Ant. The expert data for both method and the baselines consists of single trajectory from the D4RL expert dataset. The corresponding normalized expert scores are 113.23, 88.42, 106.84, and 1 116.97, respectively. For ICVF pre-training, we use 10K state pairs collected from random rollouts, which merely 1% of our online data. All baselines, including ours, are trained with a single state-only trajectory as expert data. We report the normalized average reward over 10 evaluation trajectories, providing both the mean and standard deviation (higher values indicate better performance). Each method is trained with 1M online samples. Additional hyperparameters are provided in Appendix E.3.
Results. Fig. 4 shows the reward curves of each method on MuJoCo environments, while Tab, summarizes the final results. Both the figure and the table show that our method achieves compelling results and convergence compared to the baselines across tasks. BCO, DACFO, OPOLO, and PWIL. perform well in some environments, while other methods struggle.
4.3 NAVIGATION ENVIRONMENT
To further evaluate method in setting similar that discussed in Fig. 1, we test the mazeld- medium, maze2d-large, and antmaze-umaze-N navigation tasks. For maze2d-medium and maze2d- large, increase the difficulty and investigate the agent's generalization capabilities, inject Gaussian noise into the agent's 's initial state, with the noise level being the standard deviation of
We are unable to run the github version of OLLIE due to trivirivial typos in their code, and we are not able to reproduce I.SIQQ results, likely due to version mismatch for unspecified packages. For both methods, directly report the final numbers with expert trajectory from their paper instead.

---

Published as a conference paper at ICLR 2026
Table Performance comparison on MuJoCo environments. The first part of the table lists methods that require extra access to expert actions; the second part is for LfO methods. represents results reported from the original paper, as we are unable to reproduce the results on our machine. Results are averaged over 50 trajectories. It is apparent that our method outperforms most baselines, even those with access to expert actions. DIFO is the most competitive expert state-only baseline.
Hopper
HalfCheetah
Walker2D
Ant
Average
BC
1.51 + 0.61
1.92 + 0.80
2.06 = 0.99
12.74 + 2.34
4.56
GAIL
7.78 + 2.13
-0.33 + 0.60
2.14 + 1.03
-1.98 + 4.41
1.90
AIRL
1.16 116+00.4 0.43
6.02 6.02 + 3.59
0.83 + 0.95
3.30 11.39
2.83
WDAIL
107.72
38.30 1.09
126.07 + 19.36
87.59 + 13.12
89.92
DiffAIL
95.52
33.36 +: 28.14
133,90 + 1.41
80.56 +: 13.09
85.84
OLLIE
71.10 3.5*
35.50 4.0*
59.80 1
57.10 7.0*
55,87
PWIL
78.93 + 39.00
20.81 + 33.21
84.01 = 26.53
105.62 = 2.36
72.34
IQleam
86.24 21.92
2.51 + 1.05
7.07 1 7.23
7.39 1 0.13
25.80
BCO
21.31 4.06
4.08 + 1.72
0.88
25.33 1
12.90
DACIO
109.46 0.39
61.52 1: 0.76
0.76
45.28 37.26
113.40 1 10.20
82.41
GAIRO
58.74 9.07
29.79 2.12
32.73
12.99 2.77
38,56
OPOLO
99.24 5.49
58.98 =
37.07 = 12.67
129.6633.6
81.19
DIFO
98.50 12.60
78.62 18.68
99.43 11.72
93.47 8
92.51
LS-IQ
76.84 4
41.64 + 3.64
102.02 + 2.54
132.28 3.29
88.20*
IWAIL (ours)
108.84 + 0.79
90.40 3.71
106.51 + 1.14
90.53 + 3.26
99.07
Gaussian distribution. This perturbation forces the agent to start from unseen observations that deviate from the expert demonstrations.
Maze2d. The quantitative results summarized Tab. demonstrate that LWAIL significantly out- performs the baseline with diverse initial state. As the noise level increases, LWAIL without ICVF suffers a catastrophic performance drop, failing to navigate successfully.
contrast, LWAIL with ICVF embedding maintains consistent high performance across all noise levels. This suggests that the ICVF embeddings in LWAIL successfully capture the environment's ' dynamics, allowing the agent handle unseen observations and recover from unfamiliar states.
Table Performance on mazz2ddmedium and maze2d-large large with different levels of perturbed initial states. We report the normalized returns across different noise levels.
Maze Type
Noise Level
LWAIL
No Embedding
Maze2d-medium
0.0
144.59 + 0.87
117.91 1 0.30
0.2
143.42 + 0.91
-17.00 + 0.01
0.5
135.12 : 3.12
-16.98 + 0.02
Maze2d-large
0.0
156.68 + 0.60
148.25 + 6.41
0.2
157.04 = 0.71
-11.11 0.04
0.5
157.06 0.81
-11.13 0.02
Antmaze. Tab. 3 shows the performance antmaze, for which use WDAIL, IQ-Leamn and SMODICE (Ma et al., 2022) as the basclines. The result shows that LWAIL with ICVF works comparably well with the baselines, and again illustrates the importance of ICVF embedding. Note, different from SMODICE, which uses different hyperparameters such as divergences (KL vs. x²). our method uses the same hyperparameter for all environments.
Table 3: Performance on antmaze-umaze-e
LWAIL
No Embedding
WDAIL
IQ-Leam
SMODICE
34.76 + 7.21
0.03 + 0.01
31.03 + 5.62
0.00 + 0.00 44.44
8

---

Published as a conference paper at ICLR 2026
Table 4: Ablation of different embedding methods with LWAIL. The results show that ICVF em- beddings outperform other contrastive learning- based embeddings.
Hopper
HaltCheetah
Walker2D
Ant
Average
LWAIL 108.84 + 0.79
90.40 3.71
106.51 + 1.14
90.53 + 3.26
99.07
PW-DICE
110.60 = 0.77
46.07 = 27.95
106.63 1.03
85.36 = 8.12
87.16
CURL
105.70 1.22
87.62 + 5.10
102.97 4.19
52.03 + 8.33
87.08
No Embedding
108.34 3.42
85.98 + 3.42
62.39 i 20.43
40.72 = 18.95
74.36
Table 5: LWAIL with and without environment noise on MuloCo environments. The result shows that LWAIL is robust transition noise in the environment.
Hopper
HalfCheetah
Walker2D
Ant
Average
LWAIL.
110.52 + 1.06
86.71 = 5.67
105.30 + 2.33
80.56 + 13.09
95.77
LWAIL with noise 0.1
110.25 + 1.78
82.04 6.68
104.98 + 1.44
79.95 + 12.08
94.31
LWAIL with noise 0.2
108.27 + 2.38
79.29 5.21
104.61
2.34
77.10 12.08
92.32
LWAIL with noise 0.5
107.93
2.38
48.39 3.95
41.28 29.66
-12.94 = 9.12
46.17
4.4 ABLATION STUDY
Due to the page limit, we only present part of the ablations. See Appendix F for more ablations on the dataset, algorithm components, baselines, visualizations, and results on more environments.
How well does our embedding work? To further verify the effectiveness of the ICVF embedding beyond theoretical insights, we compare our method with ICVF embeddings to use of other embed- dings. We identify two contrastive learning-basee baselines that are most suitable for our scenario: CURL (Laskin et al., 2020) and PW-DICE (Yan al., 2024a). Both methods use InfoNCE (Oord et al., 2018) their contrastive loss for better state embeddings. Their difference: 1) CURL updates embeddings with an auxiliary loss during online training, while PW-DICE updates embeddings be- fore all other training; 2) CURL compares the current state with different noises added as positive contrast examples, while PW-DICE uses the next states as positive contrast samples. Tab. summa- rizes the results. The result shows that 1) state emheddings generally aid learning, and thus good state embedding is necessary; and 2) our proposed method works best.
How robust is LWAIL with respect to environment noise? One limitation of our Thm. 1.1 is that it assumes the environment to be near-deterministic. To empirically validate the robustness of our method with noisy transitions, we inject Gaussian noise with varying standard deviations into the actions a E 1J" performed in the MuJoCo environments. This noise introduces controlled ran- domness, emulating real- world uncertainty and variability in system responses. Other experimental settings follow the main experiments. The results are listed in Tab. 5. The results clearly show that our method is robust to stochastic environments.
Applying ICVF embedding on baselines. also show that our proposed solution outperforms ex- isting methods with ICVF embedding, both Wasserstein-basee (IQlearn, WDAIL) and f-divergence based. The results summarized in Tab. 6. We find that 1) our method outperforms prior methods with ICVF embedding, and 2) ICVF best suits our AIL framework with TD3 as the downstream RL algorithm (see Tab. 13 for our method with other downstream RL algorithms).
Table 6: ICVF with other methods. Our method far outperforms baselines with ICVF embeddings.
Hopper
HalfCheetah
Walker2D
Ant
Average
LWAIL
108.84 1 0.79
90.40 + 3.71
106.51 1 1.14
90.53 1 3.26
99.07
ICVI
110.02 1 0.53
30.07 + 2.32
68.68 1: 9.16
3.42 1 1.01
53.04
ICIE GAIL+IVVF
29.80 + 10.12
3.82 0.98
6.54 + 1.23
8.91 - 0.45
12.27
8.96 1 2.09
0.12 0.40
3.98 1.41
-3.09 0.85
2.49

---

Published as a conference paper at ICLR 2026
5 RELATED WORK
Due to the space limit in the main paper, we only discuss three areas of the literature that are most related to our work. See Appendix B for an extended discussion of related work.
Wasserstein-based imitation learning. The Wasserstein distance (Kantorovich, 1939) widely adopted in [L/R (Xiao et al., 2019; Agarwal et al., 2021; Fickinger et al., 2022). It provides a geometry-aware measure between policy occupancies with more informative learning signals (Ar- jovsky et al., 2017). Among different forms of the Wasserstein distance, the primal form (Dadashi et al., 2021; Luo et al., 2023; Yan et al., 2024a; Bobrin et al., 2024) and the dual (Kantorovich & Rubinstein, 1958; Zhang et al., 2020; Garg al., 2021; Sun et al., 2021) are most prominent. The former allows for larger freedom in its underlying metric, but requires a regularizer (Yan al., 2024a), surrogates (Dadashi et al., 2021), tories (Luo et al., 2023; Bobrin et al., 2024). The latter is easier to optimize with gradient descent, but its metric is limited to Euclidean, which is often suboptimal (Stanczuk et al., 2021; Yan al., 2024a). Our work chooses the latter but overcomes its shortcomings, Among all these works, IQ- learn (Garg et al., 2021) is most similar to ours. Our objective in Eq. (6) is a special case of IQ-learn with Wasserstein distance. However, three key differences exist: 1) IQ-learn uses SAC (Haarnoja et al., 2018) instead of TD3; 2) IQ-learn focuses on x2 diverrgeee
direct match between trajec-
the online setting, which was found to be less effective in several prior works (Ma et al., 2022; Yan et al., 2024a); 3) We point out and overcome the metric limitation by adopting ICVF, which is not considered in IQ-learn.
Offinne-toonline II.. While offline II. (Zolna et al., 2020; Ma et al., 2022; Kim et al., 2022a) and online IL (Ho & Ermon, 2016; Fu et al., 2018) are both well-studied areas, offlne-to-online IL is relatively under-explored. especially when compared with offline-to-online RL (Schmitt et al., 2018; Kostrikov et al., 2022: Zhang et al., 2023a) which combines the best of offline RL (high data efficiency) and online RL (active data collection). While there some works using offline data to aid online imitation (Watson et al., 2024) by building dynamic models (Chang et al., 2021; Yue al., 2023) aligning discriminator and policy (Yue et al., 2024), they differ in two aspects from our proposed LWAIL: 1) their solution requires medium- medium-to-high quality offline data and does not work well with random offline data, with which our ICVF-lcarned metric works well; 2) they require state-action pairs for expert demonstrations, while our method only requires expert states.
State embedding. Many works have explored the possibility of learning a good state space em- bedding that better captures the dynamics of the environment and boosts RL performance (Zhang al., 2021: Ghosh et al., 2023; Modi et al., 2024). These works can be roughly categorized into two groups: 1) the 'theoretical group', which focuses on state equivalence (also known "bisim- ulation")) (Zhang et al., 2021; Kemertas & Aumentado-Armstroogg 2021; Le Lan et al., 2021) and the low-rank property of the MDP (Agarwal et al., 2020; Uchara et al., 2022; Modi et al., 2024); and 2) the 'empirical group' often tested on visual RL with high-dimensional input (Anand et al., 2019; Laskin et al., 2020; Yarats al., 2022; Giammarino
al., 2023)), which focuses on repre- sentation learning (Ha & Schmidhuber, 2018; Hafner et al., 2023; Bruce et al., 2024), autoencoder methods (Senthilnath
al., 2024), and contrastive learning (Sermanet al., 2017; Anand et al., 2019; Laskin et al., 2020). The recently proposed ICVF (Ghosh et al., 2023) studies an empirical, intention-based method for state embedding computation. It was shown to be effective in down- stream tasks (Ghosh et al., 2023; Bobrin et al., 2024). Our work is the first to leverage ICVF state embeddings to overcome the metric limitation of the KR duality of the Wasserstein distance.
CONCLUSION
We propose a novel adversarial imitation learning approach for state distribution matching with Wasserstein distance. Unlike prior methods that rely on Euclidean distance metrics, we optimize this distance metric by leveraging an embedding learned by the Intention Conditioned Value Function (ICVF), which captures environmental dynamics. This allows us to better align the expert-agent state distributions, even with sparse state-only demonstrations. Through theoretical analysis and multiple experiments, we demonstrate that the ICVF-leamed distance metric outperforms baselines, enabling more efficient and accurate imitation from limited expert data with only one expert trajectory. We believe our work provides a new direction for improving state-only imitation learning by using the Wasserstein distance while addressing the limitations of traditional distance metrics.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
This work was supported in part by NSF under Grants 2008387, 2045586, 2106825, and 2519216, the DARPA Young Faculty Award, the ONR Grant N00014-66-1-2999, the NIFA Award 2020- 67021-32799, the Amazon- Illinois Center on AI for Interactive Conversational Experiences, the Toyota Research Institute, and the IBM-Illinois Discovery Accelerator Institute. This work used computational resources, including the NCSA Delta and DeltaAl supercomputers through alloca- tions C1S230012, CIS230013, and CIS240419 from the Advanced Cyberinfrastructure Coordination Ecosystem: Services & Support (ACCESS) program, well the TACC Frontera supercomputer and Amazon Web Services (AWS) through the National Artificial Intelligence Research Resource (NAIRR) Pilot.
ETHICS STATEMENT
The proposed imitation learning method enables more efficient policy learning from minimal expert demonstrations (even state-only) and random data, potentially democratizing access to RL applica- tions in robotics and assistive technologies while reducing reliance on costly expert annotations. is important, however, to utilize our work adequately in real-life applications avoid potential misuse (e.g. military purposes) and job loss.
REPRODUCIBILITY STATEMENT
We include the procedure of our algorithm in Alg. 1. For environments used in experiments, list their details in Sec. 4.1 (for Maze2d) and Appendix E. I (for MuJoCo); for datasets in our experi- ment, list their statistics in Appendix E.2; for the hyperparameters of our method, list them in Tab. in Appendix E.3; for implementation of the baselines, their related repositories and licenses, we summarize them Appendix E.4. Finally, we state our computational resource consumption in Appendix H. We published our code at https: : //gtttb....//ggg/// com/ /aaa/aaaaccayaaaaaaaaaaaaaa LLMAII.
REFERENCES
Alekh Agarwal, Sham Kakade, Akshay Krishnamurthy, and Wen Sun. Flambe: Structural complex- ity and representation learning of low rank mdps. In NeurIPS, 2020.
Rishabh Agarwal, Marlos C.
Machado, Pablo Samuel Castro, and Mare G. Bellemare. Contrastive behavioral similarity embeddings for generalization in reinforcement learning. In ICLR, 2021.
Firas Al-Hafez, Davide Tateo, Oleg Arenz, Guoping Zhao, and Jan Peters. Ls-iq: Implicit reward regularization for inverse reinforcement learning. In /CLR, 2023.
Ankesh Anand, Evan Racah, Sherjil Ozair, Yoshua Bengio, Marc-Alexandre Côté, and R Devon Hjelm. Unsupervised state representation learning atari. In NeurlPS, 2019.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In ICML, 2017.
Kai Arulkumaran and Dan Ogawa Lillrank. A pragmatic look at deep imitation learning, In ACML, 2023.
Nir Baram, Oron Anschel, and Shie Mannor. Model-based adversarial imitation learning. In NIPS, 2016.
Maksim Bobrin, Nazar Buzun, Dmitrii Krylov, and Dmitry V Dylov. Align your intents: Offline imitation learning via optimal transport. arXiv preprint arXiv:2402.13037 2024.
Jake Bruce, Michael D Dennis, Ashley Edwards, Jack Parker-Holder Yuge Shi, Edward Hughes, Matthew Lai. Aditi Mavalankar, Richie Steigerwald, Chris Apps, et al. Genie: Generative inter- active environments. In ICML., 2024.
11

---

Published as a conference paper at ICLR 2026
Pablo Samuel Castro. Scalable methods for computing state similarity in deterministic markov decision processes. In AAAI, 2020.
Jonathan Chang, Masatoshi Uchara, Dhruv Sreenivas, Rahul Kidambi, and Wen Sun. Mitigating covariate shift in imitation learning via offline data with partial coverage. In NeurIPS, 2021.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. NeurlPS, 2018.
Xinshi Chen, Shuang Li, Hui Li, Shaohua Jiang, Yuan Qi, and Le Song. Generative adversarial user model for reinforcement learning based recommendation system. In ICML, 2019.
Robert Dadashi, Léonard Hussenot, Matthieu Geist, and Olivier Pietquin. Primal wasserstein imita- tion learning. In ICLR, 2021.
Peter Englert, Alexandros Paraschos, Mare Peter Deisenroth, and Jan Peters. Probabilistic model- based imitation learning, Adaptive Behavior, 2013.
Arnaud Fickinger, Samuel Cohen, Stuart Russell, and Brandon Amos. Cross-domain imitation
learning via optimal transport. In ICLR, 2022.
Farama Foundation. Pendulum, 2025. envi www.ooentt/caasiesneenooeot
URL https: /atp:///yymnaiiii um. farama org/ wwt.oo//eenuuummmmmmm
Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adversarial inverse rein- forcement learning. In ICLR, 2018.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprini arXiv:2004.07219 2020.
Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor- critic methods. In ICML, 2018
Divyansh Garg, Shuvam Chakraborty, Chris Cundy, Jiaming Song, and Stefano Ermon. Iq-learm: Inverse soft-q learning for imitation. In NeurlPS, 2021.
Dibya Ghosh, Chethan Anand Bhateja, and Sergey Levine. Reinforcement learning from passive data via latent intentions. ICML, 2023.
Vittorio Giammarino, James Queeney, and Ioannis Ch Paschalidis. Adversarial imitation learning from visual observations using latent information. arXiv preprini urXiv:2309. 17371, 2023.
Adam Gleave, Mohammad Taufeeque, Juan Rocamonde, Erik Jenner, Steven H. Wang, Sam Toyer, Maximilian Ernestus, Nora Belrose, Scott Emmons, and Stuart Russell. imitation: Clean imitation learning implementations. arXiv preprint urXiv:2211.11972. 2022.
Ishaan Gulrajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron Courville. Im- proved training of wasserstein gans. In NeurlPS, 2017.
Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. Deepseek-rl: Incentivizing reasoning capability in Ilms via reinforcement learning, arXiv preprint arXiv:2501. 2025.
David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv: 1803. 10/22,
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In ICML, 2018.
Danijar Hafner, Timothy Lillicrap. Mohammad Norouzi, and Jimmy Ba. Mastering atari with dis- crete world models. ICLR, 2021.
Danijar Hafner. Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301.04104 2023.
Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. In NIPS, 2016.
12

---

Published as a conference paper at ICLR 2026
Bo-Ruei Huang, Chi Yang. Chun-Mao Lai, Dai-Jie Wu, and Shao-Hua Sun. Diffusion imita- tion from observation. In NeurIPS, 2024.
Julian Ibarz, Jie Tan, Chelsea Finn, Mrinal Kalakrishnan, Peter Pastor, and Sergey Levine. How to train your robot with deep reinforeement learning: lessons we have learned. The International Journal of Robotics Research, 2021.
Unnat Jain, Iou-Jen Liu, Svetlana Lazebnik, Aniruddha Kembhavi, Luca Weihs, and Alexander Schwing. GridToPix: Training Embodied Agents With Minimal Supervision. In ICCV, 2021.
Shengyi Jiang, Jingcheng Pang, and Yang Yu. Offline imitation learning with misspecified simu- lator. In NeurIPS, 2020.
Lconid Kantorovich and Gennady S. Rubinstein. On a space of totally additive functions. Vestnik Leningradskogo Universiteta, 1958.
Leonid V Kantorovich. Mathematical methods of organizing and planning production. Management science, 1939.
Mete Kemertas and Tristan Aumentado-Armsttrme Towards robust bisimulation metric learning. In NeurlPS, 2021.
Rahul Kidambi, Jonathan Chang, and Wen Sun. Mobile: Model-based imitation learning from observation alone. In NeurlPS, 2021.
Geon- Geon-Hyeong Kim, Jongmin Lee, Youngsoo Jang, Hongseok Yang, and Kee-Eung Kim. Lobsdice: Offline learning from observation via stationary distribution correction estimation. In NeurIPS, 2022a.
Geon-Hyeong Kim, Seokin Jongmin Lee, Wonseok Jeon, HyeongJoo Hwang, Hongseok Yang, and Kee-Eung Kim. Demodice: Offine imitation learning with supplementary imperfect demon- strations. In ICLR, 2022b.
Victor Kolev, Rafacl Rafailov, Kyle Hatch, Jiajun Wu, and Chelsea Finn. Efficient imitation learning with conservative world models. In LADC, 2024.
Iyy Kostrikov, Kumar Krishna Agrawal, Debidatta Dwibedi, Sergey Levine, and Jonathan Tomp- son. Discriminatoracctorrriiiic Addressing sample inefficiency and reward bias in adversarial imitation learning. In ICLR, 2019.
Ilya Kostrikov, Ofir Nachum, and Jonathan Tompson. Imitation learning via off- distribution matching. In ICLR, 2020.
Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q- learning. In ICLR, 2022.
Ashish Kumar, Saurabh Gupta, and Jitendra Malik. Learning navigation subroutines from egocentric videos. In CoRL, 2019.
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised representa- tions for reinforcement learning, In ICML, 2020.
Charline Le Lan, Mare G Bellemare, and Pablo Samuel Castro. Metrics and continuity in reinforce- ment learning. In AAAI, 2021.
HyeAnn Lee and Donghwan Lee. Suppressing overestimation in q-learning through adversarial behaviors. arXiv preprint 1rxiv:2310.06286, 2023.
Jongmin Lee, Wonseok Jeon, Byung-Jun Lee, Joelle Pineau, and Kee-Eung Kim. Optidice: Offline policy optimization via stationary distribution correction estimation. In ICML., 2021.
Kevin Li, Abhishek Gupta, Ashwin Reddy, Vitchyr H. Pong, Aurick Zhou, Justin Yu, and Sergey Levine. Mural: Meta-learning uncertainty-awant rewards for outcome outcome-driven reinforcement learning. In ICML, 2021.
13

---

Published as a conference paper at ICLR 2026
Shangzhe Li, Zhiao Huang, and Hao Su. Reward-free world models for online imitation learning. In ICML,, 2025.
Timothy P. Lillicrap. Jonathan Hunt, Alexander Pritzel, Nicolas Manfred Otto Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learn- ing. In ICLR, 2016.
HaoChih Lin, Baopu Li, Xin Zhou, Jiankun Wang, and Max Q-H Meng. No need for interactions: Robust model-based imitation learning using neural ode. In ICRA, 2021.
Yicheng Luo, Zhengyao Jiang. Samuel Cohen, Edward Grefenstette, and Marc Peter Deisenroth. Optimal transport for offline imitation learning. In ICLR, 2023.
Yecheng Ma, Andrew Shen, Dinesh Jayaraman, and Osbert Bastani. Versatile offline imitation from observations and examples via regularized occupancyyy matching. In ICML, 2022.
Aditya Modi, Jinglin Chen, Akshay Krishnamurthy, Nan Jiang, and Alekh Agarwal. Model-free representation learning and exploration in low-rank mdps. JMLR, 2024.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predic- tive coding. arXiv preprint arXiv:1807.03745. 2018.
Jyothish Pari, Nur Muhammad (Mahi) Shafiullah, Sridhar Pandian Arunachalam, and Lerrel Pinto. The surprising effectiveness of representation learning for visual imitation. In RSS, 2022.
Rafael Rafailov, Tianhe Yu, Aravind Rajeswaran, and Chelsea Finn. Visual adversarial imitation learning using variational models. In NeurIPS, 2021.
Rajkumar Ramamurthy, Prithviraj Ammanabrolu, Kianté Brantley, and Jack Hessel. reinforce- ment learning (not) for natural language processing: Benchmarks, baselines, and building blocks for natural language policy optimization. In ICLR, 2023.
Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. reduction of imitation learning and struc- tured prediction no-regret online learning. In AISTAYS, 2011.
Simon Schmitt, Jonathan J Hudson, Augustin Zidek, Simon Osindero, Carl Doersch, Wojeiech M Czarnecki, Joel Z. Leibo, Heinrich Kuttler, Andrew Zisserman, Karen Simonyan, et al. Kickstart- ing deep reinforcement learning. arXiv preprint arXiv: (exXXiv::000.000 /803.08835, 2018
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 2017.
J Senthilnath, Bangjian Zhou, Zhen Wei Ng, Decksha Aggarwal, Rajdeep Dutta, Ji Wei Yoon, Aye Phyu Phyu Aung, Keyu Wu, Min Wu, and Xiaoli Li. Self-evolving autoencoder embedded q- network. arXiv preprint arXiv:2402. 1604, 2024.
Pierre Sermanet, Corey Lynch, Jasmine Hsu, and Sergey Levine. Time-contrastive networks: Self- supervised learning from multi-view observation. In CVPRW, 2017.
David Silver. Aja Huang. Chris Maddison, al. Mastering the game of go with deep neural networks and tree search. Nature, 2016.
Jan Stanczuk, Christian Etmann, Lisa Maria Kreusser, and Carola-Bibiana Schönlieb. Wasser- stein gans work because they fail (to approximate the wasserstein distance). arXiv preprint arXiv:2103.01678, 2021.
Mingfei Sun, Anuj Mahajan, Katja Hofmann, and Shimon Whiteson. SoftDICE for imitation learn- ing: Rethinking off-policy distribution matching. arXiv preprint arXiv:2106.03155, 2021.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Bud- den, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, al. Deepmind control suite. arXiv preprint ar / 000690, 2018.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In IROS, 2012.
14

---

Published as a conference paper at ICLR 2026
Faraz Torabi, Garrett Warnell, and Peter Stone. Behavioral cloning from observation. In IJCAI, 2018a.
Faraz Torabi, Garrett Warnell, and Peter Stone. Generative adversarial imitation from observation. In ICML Workshop on Imitation, Intent, and Interaction, 2018b.
Masatoshi Uehara, Xuezhou Zhang, and Wen Sun. Representation learning for online and offline rl in low-rank mdps. ICLR, 2022.
Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. JMLR, 2008.
Bingzheng Wang, Guoqiang Wu, Teng Pang, Yan Zhang, and Yilong Yin. Diffail: Diffusion adver- sarial imitation learning. In AAAI, 2024.
Joe Watson, Sandy Huang, and Nicolas Heess. Coherent soft imitation learning. In NeurlPS, 2024.
Alan Wu, AJ Piergiovanni, and Michael S. Ryoo. Model-based behavioral cloning with future image similarity learning. CoRL,, 2019.
Alan Wu, AJ Piergiovanni, and Michael S Ryoo. Model-based behavioral cloning with future image
similarity learning. In CoRL, 2020.
Huang Xiao, Michael Herman, Joerg Wagner, Sebastian Ziesche, Jalal Etesami, and Thai Hong
Linh. Wasserstein adversarial imitation learning, arXiv preprint 1906. 081 13, 2019.
Danfei Xu and Misha Denil. Positive-unlabeled reward learning. In CoRL, 2019.
Keyulu Xu, Jingling Li, Mozhi Zhang, Simon S Du, Ken-ichi Kawarabayashi, and Stefanic Jegelka. What neural networks reason about? In ICLR, 2020.
Kai Yan, Alex Schwing, and Yu-Xiong Wang. Offline imitation from observation via primal wasser- stein state occupancy matching. In ICML, 2024a.
Kai Yan, Alex Schwing, and Yu-Xiong Wang. A simple solution for offline imitation from observa- tions and examples with possibly incomplete trajectories. In NeurIPS, 2024b.
Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous con- trol: Improved luga-neeeented reinforcement learning. In ICLR, 2022.
Zhao-Heng Yin, Weirui Ye, Qifeng Chen, and Yang Gao. Planning for sample efficient imitation learning. NeurlPS, 2022.
Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-worid: A benchmark and evaluation for multi-task and meta reinforcement learning. In CoRL, 2020.
Sheng Yue, Guanbo Wang, Wei Shao, Zhaofeng Zhang, Sen Lin, Ju Ren, and Junshan Zhang, Clare: Conservative model-based reward learning for offline inverse reinforcement learning. In ICLR, 2023.
Sheng Yue, Xingyuan Hua, Ren, Sen Lin, Junshan Zhang, and Yaoxue Zhang. Ollie: Imitation learning from offline pretraining to online finetuning. In ICML, 2024.
Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invari- ant representations for reinforcement learning without reconstruction. In ICLR, 2021.
Haichao Zhang, We Xu, and Haonan Yu. Policy expansion for bridging offline-to-online reinforce- ment learning. In ICLR, 2023a.
Ming Zhang, Yawei Wang, Xiaoteng Ma. Li Xia, Jun Yang. Zhiheng Li, and Xiu Li. Wasserstein distance guided adversarial imitation learning with reward shape exploration. In DDCLS, 2020.
Wenjia Zhang, Haoran Xu. Haoyi Niu, Peng Cheng, Ming Li, Heming Zhang, Guyue Zhou, and Xianyuan Zhan. Discriminator-guian model-base offline imitation learning. In CoRL, 2023b.
15

---

Published as a conference paper at ICLR 2026
Zhuangdi Zhu, Kaixiang Lin, Bo Dai, and Jiayu Zhou. Off off-policy imitation learning from observa- tions. In NeurlPS, 2020.
Konrad Zoina, Alexander Novikov, Ksenia Konyushkova, Caglar Gulcehre, Ziyun Wang, Yusuf Aytar, Misha Denil, Nando de Freitas, and Scott E. Reed. Offline learning from demonstrations and unlabeled experience. In Offline Reinforcement Learning Workshop at NewrIPS, 2020.
16

---

Published as a conference paper at ICLR 2026
APPENDIX: LATENT WASSERSTEIN ADVERSARIAL IMITATION LEARNING
Our appendix is organized as follows. We first discuss limitations of our work in Sec. A, and more areas of related work in Sec. B. In Sec. C, we discuss extended preliminaries of our work, including TD3 and a more detailed explanation of Eq. (11) and ICVE. In Sec. D, we provide a mathemati- cal proof for Thm. 3.1. In Sec. E, we provide the details of the environment in experiments (Sec. E.I). the dataset used in our experiments (Sec. E.2), the hyperparameters we used for our method (Sec. E.3), and the details for our baselines (Sec. E.4). We provide variety of ablations in Sec. F. In Sec. G, we summarize the notation used in our paper. Finally, in Sec. H, we state the computational resources used for running our experiments.
In particular, we provide a detailed analysis of our method in Sec. F, which includes the following content:
Dataset ablations: How will the performance change if size (Sec. F2.1) and quality (Sec. F.2.2) of the offline state-only dataset I change? How will the performance change if the size of the expert dataset E changes (Sec. F.2.3), and how does LWAIL work with incomplete expert trajectories (Sec. F.2.4)?
Algorithm component ablations: What happens if we use other RL algorithms (e.g., PPO (Schulman et al., 2017) or DDPG (Lillicrap et al., 2016)) in the imitation stage (Sec. F.3.1). and the Sigmoid reward mapping that we use in LWAIL always beneficial for RL on MuJoCo environments (Sec. F.3.2)?
Detailed performance analysis: How robust is our method with respect to different hyper- parameters (Sec. 1.4.))? What the t-SNE visualizations of the trajectories in the embed- ding space for the Ant and Hopper environment (Sec. F.4.2)? How does our pseudo-rewan look like during training (Sec. F.4.3), and how does it perform compared to ground-truth reward signal (Sec. F.4.4)?
Baseline analysis: How do the baselines perform when they are trained for a longer time (Sec. 1.51
More environments: How does LWAIL perform beyond MuJoCo environments
(Sec. 16.1), and can LWAIL learn from expert trajectories with mismatched dynamics (Sec. R6.2)?
A LIMITATIONS
Similar to other prior AIL methods such as WDAIL (Zhang et al., 2020), our pipeline requires an iterative update of the actor-critic agent and the discriminator during online training. The update frequency needs to be balanced during training. Also, testing our method on more complicated environments, such as image-based ones, is interesting future research.
B EXTENDED RELATED WORK
Imitation by occupancy matching adversarial training. GAIL (Ho & Ermon, 2016) is one of the first works to study adversarial imitation learning. A bi-level optimization of "RL over in- verse RL" is considered, which corresponds to state-action occupancy matching between an expert and the learner's policy. Follow-up works (Fu al., 2018; Kostrikov al., 2019; Torabi et al., 2018b; Zhu et al., 2020) have further explored the adversarial training paradigm, jointly training 1) a discriminator distinguish policy occupancy from expert occupancy; and 2) the policy fitting expert demonstrations. While some occupancy matching works like DIstribution Corrected Estima- tion (DICE) (Ma al., 2022; Kim al., 2022a; Yan al., 2024a) are trained using a single-level optimization instead of an adversarial setup, they essentially derive a closed-form solution for the policy given the discriminator. Most works in this field, however, focus on f-divergence (especially KL (Zhu et al., 2020) or x" (Ma et al., 2022)) minimization, making them not only unaware of the underlying geometry of the states which leads to less informative learning signals, but also requires them to consider the same support set for the expert and learner occupancy distributions, which could lead to additional theoretical constraints and numerical issues (Ma et al., 2022). In contrast, work studies the Wasserstein distance and overcomes initial limitations.
17

---

Published as a conference paper at ICLR 2026
Remark B.1. Note, the key difference between DICE methods and our Wasserstein AlL framework is that DICE methods are generally not compatible with Wasserstein distance. DICE methods usu- ally use f-divergences, especially KL and x divergggess because the Bellman constraint is linear one; to turn the optimization into an unconstrained one, the primal objective must be non-linear. Un- fortunately, Wasserstein distance is linear objective, hence incompatible with DICE. While there are some exceptions, they must add an additional non-linear regularizer to bypass this issue, such as state-action KL divergence between learner and non-expert data (Yan al., 2024a) an entropy regularizer (Sun et al., 2021). In contrast, our current Wasserstein AlL framework does not exhibit this issue.
Remark B.2. Theoretically, with infinite data, perfect convergence and while ignoring regulariz- ers (e.g., offline II. methods like LobsDICE (Kim al., 2022a) use a pessimistic KI--divegeence term with suboptimal dataset occupancy distribution, which deviates from the optimal solution), all cunsition-distributte
matching methods optimizing s"), (s, a) or (s, a, can encode optimal paths regardless of the underlying distance metric. However, this idcalization does not remove the practical need for geometry-ave aware embedding in real policy learning. Two reasons are central:
Neural networks struggle when semantically different states are numerically close: If two dynamically distinct states are nearly identical in raw Euclidean space, the value function cannot reliably separate them. A good embedding makes these differences linearly separable and eases optimization substantially.
Wasserstein AIL fundamentally depends on the quality of the underlying metric: Un- like f-divergences, Wasserstein distances defined directly through the metric structure. A poor metric yields a discontinuous, highly non-smooth loss landscape. Enforcing Lip- schitzness with annEuulideen metries is difficult in practice; ambedding-based Euclidean metrics provide a practical and effective solution.
Thus, even though minimizing transition distributions is theoretically sufficient asymptotically, in practice, a good state embedding is crucial for learnability, stability, and sample efficiency under Wasserstein AIL.
Imitation from observation. Imitation (Learning) from Observation (LfO) aims retrieve an ex- pert policy without labeled actions. This is particularly interesting in robotics, where the expert action can be either inapplicable during cross-embodiment imitation (Sermanet et al., 2017) or un- available when imitating from videos (Pari et al., 2022). The three primary strategies of LfO can be categorized as follows: 1) minimizing an occupancy divergence through DICE methods (Zhu et al., 2020; Lee et al., 2021; Ma et al., 2022; Kim et al., 2022a;b; Yan et al., 2024a) or iterative inverse-RL updates (Torabi al., 2018b; Xu & Denil, 2019; ZoIna et al., 2020); 2) predicting miss- ing actions through inverse dynamics modeling (Torabi et al., 201 2018a; Kumar et al., 2019); and 3) similarity- based reward assignment (Sermanct et al., 2017: Chen et al., 2019; Wu et al., 2019). Our work belongs to the first category, and adopts the Wasserstein distance as the measure between occupancies, which improves results over prior works.
Model-based Imitation Learning. While ICVF implicitly encodes environment dynamics into its embedding space, some imitation learning work takes a more explicit, model-based way to build dynamic model world model (Yue al., 2023). For example, RMBIL (Lin al., 2021) uses a Neural ODE (Chen et al., 2018) to model state transition, while CMIL (Kolev et al., 2024) uses DreamerV2 (Hafner al., 2021); many other studies (Chang et al., 2021: Jiang al., 2020) simply utilize normal neural networks as inverse dynamic models. With such dynamic models, a wide range of options have been explored, including recovering expert action (Torabi et al., 201 aa). alignment to new dynamics (Jiang et al., 2020). planning in advance (Yin et al., 2022; Li et al., 2025; Englert et al., 2013; Wu et al., 2020), generating data from simulated rollouts (Chang et al., 2021: Zhang et al., 2023b; Kidambi et al., 2021; Rafailov et al., 2021), and making the environment differen- tiable (Baram et al., 2016). Among these more explicit model- methods, ICVF embedding is special as it does not require any action; while many aforementioned methods do not need expert action (Torabi et al., 2018a; Kidambi et al., 2021), they still need non-expert actions to build the dynamic model.
18

---

Published as a conference paper at ICLR 2026
C EXTENDED PRELIMINARIES
TD3. Twin Delayed Deep Deterministic Policy Gradient (TD3) (Fujimoto et al., 2018) extends the Deep Deterministic Policy Gradient (DDPG) (Lillicrap et al., 2016) algorithm, designed to mitigate the overestimation bias commonly found in Q-learning. TD3 introduces three key modifications: 1) Clipped Double Q-learning: TD3 maintains two Q-networks, Qo and Që₃ parameterized by 0₃ and 0g respectively. and uses the smaller of the two as the critic loss to reduce overestimation (Lee & Lee, 2023). The clipped value stabilizes training and results. 2) Delayed Policy Updates: To further stabilize learning, the policy is updated less frequently than the critic, reducing the chance of policy updates based on inaccurate Q-values. 3) Target Policy Smoothing: To address overfitting to deterministic policies, when calculating the target y for the critic loss, a Gaussian noise with variance o² > 0 is clipped with a threshold oo > 0 before being added to the target action a'. This regularizes the policy, making it more robust to small state changes. TD3 improves upon DDPG and works well, particularly in high-diment sighddimessionll continuous action spaces. In this paper, we adopt TD3 for the downstream RL component of our method.
Extended Introduction of Intention Conditioned Value Function (ICVF). Intuitively, V(s, 24 is designed evaluate the likelihood of the following question: How likely I to see 84 in the future trajectory if act with the intention of reaching = from state 8? The learn- ing of ICVF is similar other value-learning algorithms. For policy (--1,,__,, ICVF satisfies the following Bellman equation:
V (= S4, 2) Egw 1) [I(s 8+) +. (VT s+)))
Here, (s, is a transition and Pa (=(4tals)) is the transition probability from 8; to 8141 when acting according to intent 2 E S.3 Here, serves as a pseudo-action label and I(s = 8+) serves as a pseudo-reward label. Thus, the original reward or action are not needed ICVF training.
The original paper adopts implicit Q-learning (IQL) for ICVF learning. In one update batch, we sample transition 88,8)), potential future outcome 84, and intent z. Similar to the original IQL (Kostrikov et al., 2022), we update the critic with asymmetric critic losses to avoid out-of- distribution overestimation. To do this, we adjust the weights of the critic loss based on the positivity of the advantage. Note, as care about whether the transition (s, s') corresponds acting with intention z, our goal 8+ is equal to z." The advantage A is defined as:
A I(s s_) + Vo 2) Vo(s,2,2).
(10)
Following that, the critic loss is defined as:
((V)) = ((s.)))...)) [|a - I(A < O) - I(s = s+) - ~VVol s4,,)
(11)
where 2 E S and 84 E S are sampled from random states the dataset, from random future states in the trajectory, as the current state with different probability, respectively. Note, as we want to penalize overestimation, the coefficient for overestimation Ja - 0] a should be greater than underestination Ja 1| 1 a, and thus > 0.5. In LWAIL, we original ICVF paper.
= 0.9 following the
Remark C.I. The principle of ICVF embedding is related to bisimulation metrics (Castro, 2020). The bisimulation metric trains an encoder such that in its embedding space, the distance between the embedding approximately equals the bisimulation metric, which is "reward difference + Wasser- stein distance of transition distribution in the embedding space." " For example, the "distance between the embedding" can be approximated by another neural network, and the Wasserstein distance is by- passed by the deterministic property of MDP (Castro, 2020). Alternatively, the "distance between the embedding" can be Manhattan distance with 2- Wassersseeinn distance for a convenient closed form (Zhang et al., 2021). Mcanwhile, ICVF embedding trains a value function V(s, s_, 2) con- ditioned on intention (goal state) a and reward function being the visit probability of the state s4. (s,, $+, 2) is designed to be a product of the embedding and other components. They are similar as they:
2 Thhoretically, can also be task-speeific labels, e.g., indices for different tasks, but ICVF mainly moti- vated by goal-based RI. Thus, in both the main part of ICVF and in our work, we have 2 E 'Intuitivell, use "the best actions with the highest estimated value" of the next step in traditional Bellman operators; here, apparently the intention of reaching 2 the best choice if the target future state is
19

---

Published as a conference paper at ICLR 2026
Both use Wasserstein distance (ICVF itself does not use Wasserstein distance, but LWAIL does);
Both address the issue described in Fig. ((a), which describes the similarity of two states regardless of their numerical values. Bisimulation does so by considering embedded state similarity within one step reach (propagated to trajectories during learning), while ICVF does so by discrimination in different values for the value function with the same intention z and the same reward based 8+.
Both produce representations that can generalize to unseen reward functions. ICVF does so by carefully designing the value function architecture, such that all downstream value functions with state-based rewards are linear representations of the embedding: bisimula- tion does so by encoding all the causal ancestors of the reward under certain conditions (see Thm. Zhang et al. (2021) for details).
The key difference is: bisimulation is action- focuuedf while ICVF is state-focused. This difference leads to different advantages of the two embeddings: ICVF work with only the state available, but the notion of "intention" assumes goal-based downstream rewards for generalization guarantee; bisimulation metric requires actions for training, and is designed for denser reward functions (since the metric directly compares reward). The ability of ICVF utilizing low-quality state-only data which is otherwise unusable aligns with our motivation.
D MATHEMATICAL PROOFS
Here we provide the proof of Thm. 3.1, which we restate for readability:
Theorem D.1. In a near-deterministic MDP with 7 < 1, For any converged policy az learned in ICVF with goal there exists vector such that, for any adjacent 8') E I, (8, as n @o(s), i.e, the state-pair occupancy is approximately a linear combination of Φ.
Proof. We first consider the trivial case: z is unreachable from In this case, we have (s, 54,2) O for any 84, which is apparently a linear combination of e with coefficient
Next, consider the definition of the ICVF value function V (s, 2) E R. For goal current state s and some possible future state 8+,
8
(" 84,
Σ 7 p(s; = = = s);
(12)
where denotes the future steps. Meanwhile, the state and state pair occupancy d' (s), d= 84) are defined as
d5(s) (1
p'p(s: = =),
(13)
(s. (1 11-7)
7 1pp(( St s).
Therefore for any adjacent (8,84) E
de (s) 4" s, 84, 2)
8
=(1 7)
7' fp(s =))
Σ
xt 1- =4450
-(1 7))
Σ ppp(st =s) 8s
[( = =
=4)+==))+++ + Σ
7- =5 = s)) (14)
i=2
80
F (s, (8, 84) + (1 -7)
Σ
p'p(s: = 8) Σ
7 1-1
p( = 8 = 8)
1=0
1=2
20

---

Published as a conference paper at ICLR 2026
S4 Case
in Case
Figure 5: An illustration of "case 1" (green states) and "case 2" (orange states) in our proof, and the 2 is the intention (i.e. goal). The transparent states are longer paths, which are unlikely be visited with converged policy and in a near-deterministic MDP.
We now consider two different cases of 84. They correspond to the two cases illustrated in Fig. 5. We will prove that either the second term is negligible (case I), or it is multiple of (s, (case 2).
Case 1: 84 is on the path towards reaching 2.
Since is steadily reachable from in near-deterministic environment, visiting s' directly from s instead of later returning from other states s' shortens the trajectory towards z. It should thus be preferred by any converged policy n. with y < 1. Thus, we have
p(5++1 = 8418; 8) >> p(8++ 8418; = s) k: 2 2 and (s, 84) E T,
(15)
8 8+10= = s)] term is negligible. In this case, dh, (8. (s) 8( (s) d (z) and we have 1) =
which means the (1 - 7) EE 8) 2/
d (s) (=)))oo(s4).
Case 2: 8.4 is on the path after reaching 2.
In this case, π. should "loop back" to as soon as the MDP allows. As the MDP is near- deterministie, suppose the smallest state cycle looping back towards is T. Without loss of gener- ality, we assume = s-; for other states in the loop, factor of yt applies if the agent is still steps away from 84 in the loop.
Then, we have
0, (p(10-1= st 1 s i2 2, 1) mod T 0.
i2 (i 1) mod T + 0,
p(s+ti= = /8 =
(16)
Thus, we have
d V (=
Σ Epp(s p(ss = 8) p
1-1 p(s s4 s)
=(1 7)
- s4|s, =) 0
Σ
(17)
2=2
1
d' 88,
and thus similar to case 1, (1 (1- 7d==")d== (6) To())boo((()))
Remark D.2. An example of case in the proof is a navigation task; in a navigation task, the agent towards its goal without hovering over any intermediate state s'. An example of case is keeping a pendulum inverted after reaching its balance (Foundation, 2025), which requires the agent stay at the current state. Note, all MuJoCo environments tested in Sec. 4 satisfy the assumption in Thm. 3. as the noise on the initial states is mild while the transition functions are deterministic. We also find LWAII to empirically robust when the assumption does not hold, e.g., with noise in transition (Sec. 4.4).
21

---

Published as a conference paper at ICLR 2026
E EXPERIMENTAL DETAILS
E.1 ENVIRONMENTS
We use five MuJoCo (Todorov al., 2012) and D4RL (Fu al., 2020) environments: Maze2d, Hopper, HalfCheetah, Walker2D and Ant. The environment specifications for maze2d are provided Sec. 4.1, In this section, we will briefly introduce the other MuloCo environments. Fig. provides an illustration of those environments.
(a) Hopper
(b) Halfcheetah
(c) Walker2D
Tif
(d) Ant
(c) Maze2d
Antmaze
Figure 6: Illustration of the MuJoCo (Todorov et al., 2012) environments we test in Sec. 4.2 and Sec. 4.3.
1. Hopper. The hopper environment (as well the other three environments) is locomotion task. In Hopper, the agent needs to control a single-legged robot leaping forward 2D space with x- and z-axis. The 11-dimensional state space encompasses joint angles and velocities of the robot, while the 3-dimensional action space corresponds to torques applied on each joint.
2. Halfcheetah. In the Halfcheetah environment, the agent needs to control a cheetah-shaper robot to sprint forward. also operates in a 2D space with z- and æ-axis, but has a 17- dimensional state representing joint positions and velocities, and 6-dimensional action space that modulates joint torques.
Walker2D. As implied by its name, in Walker2D, the agent needs to control a 8-DoF bipedal robot to walk in the two dimensional space. It has a 27-dimensional state space and an 8-dimensional action space.
4. Ant. Different from the previous three environments, the Ant environment is 3D set- ting where the agent navigates four-legged robotic ant moving towards particular direc- tion. The is represented by 111 dimensions, including joint coordinates and velocities, while the action space has 8 dimensions.
5. Maze2d. In maze2d, the agent controls a 2-DoF, force-actuated ball a maze to reach the target goal. The action space is 2-dimensional, which represents the force applied on each
22

---

Published as a conference paper at ICLR 2026
dimension; the observation space is 4-dimensional, describing the velocity and location of the agent.
6. Antmaze. Antmaze is more difficult version of maze2d. where the agent needs nav- igate through the maze, not with a ball, but with a robotic ant similar to that in the ant environment. The state space is 29-dimensional and the action space is 8-dimensional.
E.2 DATASETS
For expert datasets of the MuJoCo and navigation environments, we use 1 trajectory from the D4RL expert dataset (with Apache-2.0 license) as E,, which has 1000 steps, and randomly sample trajec- tories with total of 10K transitions from the D4RL random dataset I. Some baselines, such as PWIL (Dadashi et al., 2021) employ subsampling hyperparameter, which creates a low-data training task by taking only one state/state- action pair from every 20 steps of the expert demonstra- tion. For fairness, we set all baselines' subsampling factors to be i.e., no subsampling except for Sec. F.2.4.
Table 7: The basic statistics of the random datasets from D4RL (Fu al., 2020) applied in our MuloCo experiments (dataset for navigation tasks are random explorations without reward). We use the first 10K transitions in each dataset for ICVF training.
It is apparent that all these data are of very low quality compared to an expert, yet our ICVF-leamed metric still works well.
Dataset
Normalized Reward (Expert is 100)
Hopper-random--22
1.19 + 1.16
HalfCheetabh-random-y
0.07 + 2.90
2aikkor--random
0.01 = 0.09
Ant-random-v2
6.36 = 10.07
E.3 HYPERPARAMETERS
Tab. summarizes the hyperparameters for our method. We use the same settings for all environ- ments, and keep hyperparameters identical to TD3 (Fujimoto et al., 2018) and ICVF (Ghosh et al., 2023) whenever possible.
E.4 BASELINES
We use several different Github repositories for baselines. We default settings of those repos, except for the number of expert trajectories (which set to 1) and the subsampling factor (see Appendix E.2). We always follow the original hyperparameters in their paper if possible; otherwise, we use hyperparameters from similar experiments. Below are the repos we used our experiments for each baseline:
BC (Ross al., 2011), GAIL (Ho & Ermon, 2016), AIRL (Fu et al., 2018); We use the imitation (Gleave et al., 2022) library, which provides clean implementations of several imitation learning algorithms and has a MIT license.
OPOLO (Zhu el al., 2020), DACJO (Kostrikov et al., 2019), BCO (Torabi et al., 2018a), GAIJO (Torabi et al., 2018b): We use OPOLO's official code (https: //tp::///ttttt com/ illidanlab. aalllllbbppllooooooooo opolo code), where DACfO, BCO and GAlfO are integrated base- lines, which does not have license.
OLLIE (Yue al., 2024); We tried to use the official code but it can't be executed due to non-trivial typos. Thus we use their reported numbers on D4RL random dataset instead.
PWIL: We another widely adopted imitation learning repository (Arulku- maran & Ogawa Lillrank, 2023) fattpa://qitee wwm//ao//////eeeeeee.. /at/iiiiini// Emitation-learni which has an MIT license.
23

---

Published as a conference paper at ICLR 2026
Table 8: Summary of the hyperparameters of LWAIL.
Type
Hyperparameter
Value
Note
ICVE
Network Size of
[256, 256]
same as original paper
Disc.
Network Size
[64, 64]
Activation Function
ReLU
Learning Rate
0.001
Update Epoch
40 steps
Update interval
4000
Batch Size
4000
Optimizer
Adam
Gradient Penalty coefficient
10
Actor
Network Size
[256. 256]
Activation Function
ReLU
Learning Rate
0.0003
Training length
IM steps
Batch Size
256
Optimizer
Adam
Critic
Network Size
[256, 256]
Activation Function
ReLU
Learning Rate
0.001
Training Length
1M steps
Batch Size
256
Optimizer
Adam
0.99
discount factor
WDAIL: We their official code (nttps://github. wwt.://i.ooooo//ggggcggca httpss//gghhab..oomm//gggggggg Adversarial aawaariialmmmtatiiiiaea hww.tattooo--eeerrnnn/eeeer//eee which does not have a li- cense.
IQleam:
We use their official code (tttsss//gttaabb.oo/////////11/3
tree/main) with research-only license.
DIFO:
their official code (https://githen com/ hwtrobbttieennnggaba//rooo
We
mminn) with a research-only license.
DiffAIL: We their official code ftttps://githut wwm/commm////// Grou-- SDU/ DiffAIL) which does not have license. As the training process of DiffAIL is highly unstable (especially on the walker environment), select the runs that reach 1M step possible.
LS-IQ: We tried to use their official code (https: https:///gttubb is-iq) with MIT license, but the result cannot be reproduced their dataset machine, likely due to an environment version mismatch. Thus, we use their reported numbers with expert trajectory.
SMODICE: We use their official code SMODICE) which does not have license.
F MORE ABLATIONS
In this section, we provide additional ablation results our method. We report normalized reward (higher is better) for all results.
F.1 COMPLETE TRAINING CURVES
Tab. presents the full training corresponding to the main results.
24

---

Published as a conference paper at ICLR 2026
Hopper
Halfcheetah
Wel k
Ant
150
150
60-
100
125
Norma lovel Newwedd
100
40-
0
-50
5.
Stadd Staps
PWIL
AIRL
IQlearn
BCD
GAMO
OPOLD
DITAIL
UNAIL (ours)
DACIO
OLUE
Figure 7: Performance on the MuJoCo environments. Full training curves for all tasks. Note that DiffAIL occasionally collapses on Walker2D; we therefore report the average over successful runs only.
Table 9: Ablation the size of the random dataset I for ICVF pretraining. Our method works well with a small amount of offline data for ICVE.
Size
Hopper
HalfCheetah
Walker2D
Ant
Average
10K (original)
108.84 1 0.79
90,40 : 3.71
106.51 +: 1.14
90.53 = 3.26
99.07
50K
92.35 24.77
89.88 - 4.59
104.54 1.66
89.54 = 6.74
94.08
100K
109.74 + 1.29
91.97 = 1.51
100.54 9.42
73.48 26.17
93.93
IM
110.52 + 1.06
86.71 = 5.67
105.30 2.33
80.56 = 13.09
95.77
F.2
DATASET ABLATIONS
F.2.1 ABLATIONS ON THE SIZE OF THE RANDOM DATASET I
To investigate the effect of the size of the auxiliary dataset I used for ICVF pretraining, we vary the number of random transitions from 10K up to 1M. Tab. 9 reports the results. We find that our method remains robust even with only 10K random transitions. While performance on some envi- ronments fluctuates when the dataset is very large or moderately sized, the overall results indicate that only small amount of random data is sufficient to achieve strong performance. This confirms the efficiency of our approach in terms of auxiliary data requirements.
F.2.2 ABLATIONS ON THE QUALITY THE RANDOM DATASET
We further examine whether the auxiliary dataset must be purely random. To this end, we replace the random dataset with the D4RL medi um and expert datasets, and report the results in Tab. 10. We observe that our method performs consistently well regardless of whether is random, medium, expert. This demonstrates that approach is fully compatible with existing offline datasets and does not rely on the randomness assumption. Notably, the performance is comparable across different choices of further validating the robustness and flexibility of our framework.
Table 10: Ablation on the quality of dataset I for ICVF pretraining. Our method works equally well with random non-random data sources.
Hopper
HalfCheetah
Walker2D
Ant
Average
Random dataset
108.84 0.79
90.40 3.71
106.51 + 1.14
90.53 + 3.26
99.07
Medium dataset
106.24 + 2.67
90.12 + 2.89
103.85 : 1.96
85.04 + 3.71
96.31
Expert dataset
108.11 + 1.84
89.82 + 3.24
106.42 1 1.63
84.57 + 2.98
97.23
F.2.3 ABLATIONS ON THE SIZE OF THE EXPERT DATASET E
To demonstrate the robustness of our method even the expert data scarce, we test our method with 5 expert trajectories and the whole expert dataset (1M transitions). Tab. 11 summarizes the sults. We observe consistent compelling performance regardless of the number of expert trajectories.
25

---

Published as a conference paper at ICLR 2026
Table 11: Ablation on using multiple trajectories as expert demonstrations. Our method shows consistent expert-level performance regardless of the number of expert demonstrations.
Hopper
HalfCheetah
Walker2D
Average
1 trajectory
108.84 0.79
90.40 + 3.71
106.51 + 1.14
90.53 + 3.26
99.07
5 trajectories
107.42 6.83
91.96 2.34
107.05 1.42
88.37 : 9.87
98.70
10 trajectories
108.11 5.21
92.15 2.67
106.64 1.63
89,42 = 8.11
99.08
All expert dataset
109.21 3.64
94.07 2.98
104.59 2.05
90.72 + 9.14
99.65
F.2.4 EXPERT DATASET E WITH SUBSAMPLED TRAJECTORIES
To validate the robustness of our policy, we provide results with subsampled expert trajectories, videly-adopted scenario in many prior works such PWIL and IQ-learn. Only small portion of the complete expert trajectories are present. Our subsample ratio is 10, ie., we take expert state pair of adjacent 10 pairs. Tab. 12 summarizes the results, which show that 1) our method with subsampled trajectories outperforms Wasserstein-basee baselines such as WDAIL (Zhang et al., 2020) and IQlearn (Garg et al., 2021), and 2) the performance of our method is not affected by incomplete expert trajectories.
Table 12: Ablation on subsampled expert trajectories. The result shows that LWAIL is robust to subsampled expert demonstrations and outperforms other baselines with subsampled expert demon- strations.
Hopper
HalfCheetah
Walker2D
Ant
Average
LWAIL
108.84 1 0.79
90.40 3.71
106,51 :: 1.14
90.53 1 3.26
99.07
LWAIL w./ subsample
109.00 + 0.46
86.73 +: 7.02
106.13 2.47
83.21 +: 8.80
96.27
WDAII w./ subsample
108.21 + 4.90
35.41 2.07
114.32 2.07
83.87
10.92
85.45
IQlcarn subsample
60,26 14.21
4.12 = 1.03
8.31 = 1.48
5.32 + 3.87
19.50
F.3 ALGORITHM COMPONENT ABLATIONS
F.3. DOWNSTREAM RL ALGORITHM
We used TD3 as our downstream RL algorithm rather than other algorithms such PPO and DDPG with an entropy regularizer. Our choice is motivated by better efficiency and stability, especially because TD3 is an off off-policy algorithm which more robust to the shift of the reward function and adversarial training pipeline. We ablate this choice of the downstream RL algorithm and show that TD3 outperforms PPO and DDPG in our framework. Tab. 13 summarizes the results.
Table 13: Ablation on downstream RL algorithms. The result shows that TD3 works much better than PPO and DDPG.
Hopper
HalfCheetah
Walker2D
Ant
Average
LWAIL+TD3 (original)
108.8 84 1 0.79
90.40 = 3.71
106.51 +: 1.14
90.53 + 3.26
99.07
LWAII +PPO
65.21 + 4.81
1.02 + 0.21
24.13 + 2.14
9.12 + 0.85
24.87
LWAIL+DDPG
72.32 + 21.76
79.52 = 5.21
5.86 + 3.13
-1.67 = 9.37
29.01
F.3.2 SIGMOID REWARD MAPPING
We adopt the sigmoid function to regulate the output of our neural networks for better stability (similar to WDAIL (Zhang et al., 2020)). However, one cannot naively apply the sigmoid to the reward function for better performance. To show this, compare to TD3 with sigmoid function applied to the ground- -ruth reward. The result is illustrated in Tab. 14. The result shows that a naive sigmoid mapping of the reward does not improve TD3 results.
26

---

Published as a conference paper at ICLR 2026
Table 14: Results of TD3 with and without sigmoid applied on the ground reward. The results show that applying the sigmoid function does not yield better performance.
Environment
Hopper
HalfChectah
Walker2D
Ant
Maze2D
Average
TD3
105. 54 + 1.32
76.13 3.85
89.68 = 0.84
89.21 # 1.0
120.14 14 0.32
96.14 62.27
SigmSid reward
84.23 + 4.86
30.76 14.33
42.55 9.07
34.79 20.01
119.03 0.18
F.4 DETAILED PERFORMANCE ANALYSIS
F.4.1 HYPERPARAMETER SENSITIVITY ANALYSIS
We present additional ablation studies hyperparameters in Tab. 15. Please refer to Tab. 8 for detailed definitions and default values. While Adversarial Imitation Learning methods are typically sensitive to hyperparameter configurations, performance across a reasonable range variations.
results demonstrate that LWAIL maintains robust
Notably, we did not perform an extensive grid search to maximize performance for the main results; instead, we adopted standard settings from prior work to ensure a fair comparison. As shown in the table, specific parameter tuning (e.g., poddatejintervn can yield results even superior to the reported defaults. Regarding the update interval, we observe that while it has marginal impact on asymptotic performance, it significantly influences training stability and the smoothness of learning curves.
Table 15:
Ablation study on hyperparameters. We compare different Learning Rates (LR), discrim- inator epochs, and discriminator update intervals.
Hopper
HalfCheetah
Walker2D
Ant
Average
LWAIL (Default)
110.12 + 1.05
87.20 + 5.60
104.88 2.40
81.05 12.90
95.81
Critic LR=3e-3
109.15 + 1.47
66.58 + 8.31
99.25 + 2.44
74,50 + 14.92
87.37
Critic LR=3e-4
108.96 + 1.25
85.91 + 5.42
103.87 + 2.73
83.84 + 13,90
95.65
Actor LR=1e-3
110.20 + 1.42
85.34 6.89
104.55 + 2.67
76.26 +: 14.11
94.09
Actor LR=1e-4
111.52 + 1.31
83.76 + 6.75
101.85 2.59
84.35 + 12.66
95.37
Disc. I.R33e-3
107.47 + 1.38
88.21 + 6.12
103.62 2.70
82.72 + 14.03
95.51
Disc. LR=le-4
109.25 + 1.28
84.66 + 5.54
101.12 2.51
83.48 12.91
94.63
Update Epoch=10
108.15 + 1.17
81.12 5.44
98.00 + 2.87
80.59
14.33
91.97
Update Epoch=50
105.82 + 1.33
86.85 + 6.22
99.26 + 2.46
74.97 13.77
91.73
Update Interval=1k
109.12 + 1.43
88.54 1 6.59
102.46 + 2.40
85,95 + 13.66
96.52
Update Interval=2k
106.29 + 1.19
89.37 - 5.77
103.19 + 2.28
81.22 + 12.84
95.02
Update Interval=8k
103.12 + 1.51
86.79 + 6.44
101.41 + 2.65
88.67 + 13.92
95.00
F.4.2 MORE EMBEDDING VISUALIZATIONS
We only visualize embedding for HalfCheetah and Walker2d in the main paper due to space lim- its. In Fig. 8, we visualize results for the embedding learned on the Hopper and Ant environments.
F.4.3 PSEUDO-REWARD METRIC CURVE
To validate the effect using sigmoid and ICVF embedding for our pseudo-reward generated by f. conduct two experiments:
1) Run a standard setting of LWAIL, and compare pscudo-rewards generated by f with the sigmoid function, and pseudo-rewards without the sigmoid function for the MuJoCo environments. This is illustrated in Fig.
2) Run standard LWAIL and LWAIL without ICVF embedding, and compare pseudo-rewards (with the sigmoid function) for the MuJoCo environments. This is illustrated in Fig. 10.
The value is the total reward accumulated over evaluation trajectory. The result clearly shows that both ICVF-embedding and the sigmoid function are very important for pseudo-reward stability and positive correlation with ground-truth reward.
27

---

Published as a conference paper at ICLR 2026
State Space
Latent Space
State Space
Latent Space
15.
10
10.
10
Dimension 2
Dimension 2
Dimendion 2
Dimension
-2
-10-
-15
-16
20
Dimension
Dimansion
Dimension 1
Dimansion
(a) Hopper
(b) Ant
Figure 8: Visualization of the same trajectory in the original state space and the embedding (latent) space. The color of the points represents the ground-trufe reward of the state. We observe that an ICVF-trained embedding provides much more dynamics-awan metric than the vanilla Euclidean distance. See Fig in the main paper for HalfCheetah and Walker2d environment visualizations.
5300
I0000
SCCO
2030
3300
1939
-1000
12
Figure 9: The reward curves of pseudo- and ground-truth reward in single training session, where pseudo-reward is generated by f following Alg. and serves as the reward signal for our downstream TD3. We note that the pseudo-reward is much more stable and positively correlated with the ground- truth reward when using sigmoid function.
Pepper
Har Etfeesoh
VSt A1-24
620
420-
500
N
n2o-
5c0-
200
220
47
2.5
pirt and Inthh agmatt, with smboddrig)
www www.---------- d Fw aigoord Deiddingg
Figure 10: The pseudo-reward curves with and without ICVF embedding in a single training session. We note that without ICVE, the pseudo-reward is generally less stable (c.g., fluctuation halfcheetah and sudden drop in walker2d and ant) and sometimes less correlated with ground fround-truth reward (e.g., environment).
28

---

Published as a conference paper at ICLR 2026
F.4.4 COMPARISON TO
REWARD SIGNAL
To show the validity our learned reward signal, compare our method with ground-truth guided TD3. Mean and standard deviation are obtained from three independent runs with different seeds. The result is illustrated in Fig. The result shows that LWAIL with our learned reward can achieve comparable or better performance than a human -deaand ground-trus reward.
Happer
Halfcheetah
130
105
200
072
000
1.0
600
172
5/8
0.2
10.5
16
57
02
5-4
035
0.
Figure 11: Ablations of the pseudo reward learned by the discriminator and the ground-truid truth reward. Our learned reward (red) performs equal to or better than TD3 (blue) with ground-truth reward.
F.5 BASELINE ANALYSIS
F.5.1 OTHER METHODS WITH LONGER TRAINING
As extra (albeit very small) random offline dataset I, raises potential concer:: is the comparison of our method unfair to baselines with purely online imitation learning? To address such concern, we more steps for several baselines and compare the performance of the baselines and our method. This is illustrated in Tab. 16. We observe that our method still works better than baselines which more offline (for OLLIE) or online (for online imitation learning methods) data.
Environment
Hopper
HalfCheetah
Walker2D
Ant
Average
OLLIE ( (IM offiine IM online)
71.10 35
35.50 +: 4.0*
59.80 18.5
57.10 7.0"
$5,87
OPOLO (2M online)
95.05 + 6.50
53.15 + 8.90
43.25 = 15.10
119.20 + 4.65
77.66
DACRO online)
30
65.05
40.70 42.00
DIPO (2M online)
95.14 14.30
73.87 + 17.86
98.07 = 13.88
87.62 10.75
88.68
DiTAIL (2M oaline)
108.56 + 0.07
24.38 + 21.45
73.06 19:40
114.62 + 4.73
80.16
LWAIL (0,0IM ollline 1M online)
108 84 1
90.40
3.71
106.51 1: 1.14
90.53
3.26
99.07
I.WAI. (IM offline online)
110.52 + 1.06
86.71 + 5.67
105.30 4: 2.33
80.56 1: 13.09
95.77
Table 16: Ablations of LWAIL and baselines trained with more steps (LWAIL with 1M offline uses IM offline random data, and 0.0IM 0.01M offline data). Results with * are reported from the main paper. The result clearly shows that our method still works better than baselines trained to see an equal number of even more data. Note, as DiffAIL consistently collapses before 2M steps, we use the average of the longest runs.
F.6 MORE ENVIRONMENTS
F. BALL-IN-CUP ENVIRONMENT
To test method more diverse environments, report the performance of our method on the ball-in-cup environment the DeepMind control task (Tassa et al., 201 8). This is an 8-dim state and 2-dim action task where the agent needs to put ball into connected cup, illustrated in Fig. 12.
We report the performance of our method with and without ICVF in Tab. 17. The result shows that our method works well this task, and ICVF plays vital role.
F.6.2 MISMATCHED DYNAMICS
It is worth noting that the very motivation of LWAIL is to find latent space that aligns well with the environment's 's true dynamics. Despite this, we agree that there might cases where the latent space employed in LWAIL does not align with the true dynamics due inaccurate data, e.g., mismatched
29

---

Published as a conference paper at ICLR 2026
U
Figure 12: An illustration of the ball-in-cup environment.
LWAIL
LWAIL w./o. ICVF
84.85 + 3.55
14.80 + 4.52
Table 17: Normalized reward (0-100 as in the main paper, higher is better) LWAIL with and without ICVF on the ball-in-cup environment. ICVF plays vital role on this environment.
dynamics between expert demonstrations and the actual environment. To test such cases, we use the halfcheetah mismatched experts scenario analyzed in SMODICE (Ma et al., 2022): for expert demonstration, the torso of the cheetah agent is halved in length, thus causing inaccurate alignment. We compared method with the results reported in the SMODICE paper. Tab. 18 summarizes the final average normalized reward (higher is better). Results show that 1) our method works better than several baselines including SMODICE; and 2) our method is robust to mismatched dynamics.
Table 18: Performance on the Halfcheetah environment with mismatched dynamics. Our method outperforms baselines.
Normalized Reward
LWAIL
24.31 + 4.51
SMODICE
23.2 7.43
SAIL
0 + 0
ORIL
2.47 + 0.32
G LIST OF NOTATIONS
Tab. 19 summarizes the symbols that appear in our paper.
H COMPUTATIONAL RESOURCES
All our experiments are performed with an Ubuntu 20.04 server, which has 128 AMD EPYC 7543 32-Core Processor and single NVIDIA RTX A6000 GPU. The overall training time of our method is comparable to, or even faster than, that of existing methods. Specifically, ICVF training takes approximately 80-90 minutes, while online training requires about 65-75 minutes for MuJoCo en- vironments. In contrast, PWIL requires 8-9 hours to complete, and DACfO takes 5-6 hours.
30

---

Published as a conference paper at ICLR 2026
Table 19: A list symbols used the paper. The first part focuses on RL-specific symbols. The second part details Wasserstein-speciffe notation. The third part summarizes ICVF-specifi symbols (Sec. 3.2).
Name
Meaning
Note
State space State
SES
Action space
Action
a A
Time step
E 1,,1,2....
Discount factor
Y E (0,
Reward function
r(s, a) for single sate-action pair P(s'Is, a) A(S)
P
Transition
E
Expert dataset
state-only expert demonstrations
Random dataset
state-action trajectories of very low quality
Learner policy
The policy we aim to optimize
d
State occupancy of
d) (1 7) E
Pr(s: = =), where 8: is the i-th state
in trajectory
d.
State-pair occupancy of
&i (s, 8¹) (1 - 7) Σ 7 7'Pr(s; = 8,8:41 s'), where s; is the i-th state in trajectory
d
State pair occupancy of expert policy
The expert policy here empirically induced from
Underlying metric for Wasserstein distance Dual function Discriminator
Dual function Rubinstein dual of I-Wasserstein dis- tance: also a discriminator from adversarial perspective and a reward model from IRL perspective In our case, E.ss ll(s,s') (s)).. Eves ll(s, di (s)
Wasserstein matching variable
Wi
1- Wasserstein distance
Outcome state
Latent intention
Value function
takes 8, 84, z as input in ICVF; only takes normal target value function in the critic objective of RL [[condition] 1 if the condition is true, and = otherwise embedding function we use f: ((s) Rd
Vinger
Target value
indicator function
State representation (embedding)
T
Counterfactual intention
T(=) E 1dd ((+.) ) E Rd a E (0.5, 1]
Outcome representation
ICVF constant
Sigmoid function
31


Reviews:
[{"review_id": "lncMiAxNXK", "reviewer": "Reviewer_XUbE", "summary": "The paper proposes a state-only imitation learning approach that combines a Wasserstein Adversarial Imitation Learning objective with a learned latent representation (ICVF). The motivation is that Wasserstein distances in the original state space can behave poorly for high-dimensional or redundant dynamics, and that learning an approximately linear latent embedding can yield a geometry where Euclidean distances better reflect transition similarity. The authors provide empirical results on MuJoCo control tasks, demonstrating that training in such a latent space improves sample efficiency and stability in low-demonstration regimes.", "questions": "To make the contribution clearer and more convincing, I would recommend the following revisions before acceptance:\n\n- Clarify contribution framing. Recast the paper as an investigation of Wasserstein AIL in learned latent metricsrather than as a new algorithm.\n- Systematic baseline parity. Add a study where the same ICVF embedding is used across strong baselines (IQ-Learn, GAIfO, OPOLO) under searched hyperparameters and matched general setups. This would clarify whether the Wasserstein component itself is crucial. Describe how hyperparameters for baselines and variants were chosen, especially when embedding features are shared.\n- Why was the humanoid environment left out? Why is there not at least one more complex or high-dimensional task (e.g., Humanoid or a visual D4RL environment). A test for scalability and the generality of the embedding is missing. Is the claim that the method can learn from on expert trajectory legitimate when only the “simple” mujoco tasks where tested?\n- Highlight experimental design. Why are the key ablations (embedding vs. no-embedding, Wasserstein vs. f-divergence, baselines with latent embeddings) not all in the main paper? Why are there no clear performance summaries and standard deviations over multiple seeds for the ablations?\n- Theory scope and assumptions. Discuss more explicitly when the linear-in-embedding assumption of Theorem 3.1 holds or fails. Does it fail in a truly stochastic environment?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "-The idea of combining a latent, dynamics-aware embedding with a Wasserstein AIL objective is intuitive and well-motivated.\n-The results show consistent gains over baselines in low-data settings, and the ablations suggest that learning a latent embedding indeed facilitates imitation learning.\n-The paper provides a wide range of ablations (different embeddings, action noise, dynamics mismatch), which adds credibility to the empirical findings.", "weaknesses": "-Overstated novelty. Methodologically, the approach is largely a combination of known components (Wasserstein AIL and ICVF-based embeddings) rather than a fundamentally new algorithm. The paper currently presents it as a new framework rather than as a targeted study of how these parts interact.\n- Fairness of comparisons. While the paper shows LWAIL with and without the latent, it does not systematically test whether other IL algorithms (e.g., IQ-Learn, OPOLO, GAIfO) would also benefit from the same latent embedding. While they have a short comparison in the appendix details about the comparison are missing (e.g. how where the hyperparameters selected). This makes it unclear whether the improvement comes from the Wasserstein loss or simply from using ICVF features with the other improvements like the adapted reward.\n- Experimental focus and clarity. Important comparisons are buried in the appendix, and it is difficult to attribute which component (embedding, loss type, or regularization) drives the gains. The main tables mostly cover simple MuJoCo tasks, with only a limited exploration of harder or higher-dimensional environments.\n- Limited evaluation breadth. The experiments remain confined to low-dimensional, low-contact locomotion tasks (Hopper, Walker2D, HalfCheetah, Ant). The approach’s claimed robustness would be more convincing with higher-dimensional or vision-based control tasks. For example why was the humanoid environment left out?\n- Presentation imbalance. Some theoretical and background sections could be shortened to leave space for clearer experimental structure and discussion of assumptions. Key ablations and tuning details should be surfaced from the appendix.", "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": "", "extra_scores": {}}, {"review_id": "aqtOwQV6VN", "reviewer": "Reviewer_EztW", "summary": "This paper proposes a novel Latent Wasserstein Adversarial Imitation Learning (LWAIL) algorithm for state-only imitation. The authors identify that in Wasserstein IL, the gradient penalty commonly used to stabilize adversarial optimization enforces a 1-Lipschitz constraint, which implicitly restricts the underlying distance metric to be Euclidean. To overcome this limitation, LWAIL introduces a dynamics-aware latent space learned via the Intention-Conditioned Value Function, which captures reachability between states and provides a more meaningful ground metric for Wasserstein distance.", "questions": "1. It is unclear what the expert performance level is in Figure 4. Could you provide a reference line or explicit numbers for comparison?\n\n2. In practice, many AIfO algorithms minimize a divergence over state transitions (as in Eq. (5)) rather than over individual states. Wouldn’t this already mitigate the issue highlighted in Fig. 1(a), since the optimization objective ensures that the agent’s transitions (s,s’) are aligned with the expert’s, regardless of the absolute distance between individual state embeddings? How does the proposed embedding provide an additional benefit beyond what is already achieved by transition-level matching?", "limitations": "", "rating": 2, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The paper identifies and addresses the weakness of using Euclidean metrics in Wasserstein IL by introducing a dynamics-aware latent space. Overall, the presentation is clear, well-structured, and easy to follow.", "weaknesses": "1. **Limited novelty**: The contribution is primarily a combination of existing ideas (Wasserstein IL, adversarial learning, and latent representations), with the main advance being the integration of ICVF embeddings.\n\n2. **Narrow experimental scope**: The experiments focus mainly on locomotion tasks. Given the illustrative example in Fig. 1, navigation-style tasks (such as those in Fig. 3) would have been more suitable to highlight the strengths of the proposed approach.\n\n3. **Questionable necessity of state embeddings**: A central motivation of the paper is that Euclidean distances between states are not meaningful. However, in many AIL approaches the divergence is minimized over state transitions rather than individual states. Under the standard expert optimality assumption, these transition distributions already capture the structure of optimal paths, potentially mitigating the issue highlighted in Fig. 1(a). This raises the question of whether a dedicated dynamics-aware embedding is strictly necessary or whether the problem is indirectly addressed by existing formulations.\n\nMinor issues:\n1. DACfO in line 339 outperforms LWAIL.\n2. The Related Work section is missing relevant references, such as: \nGiammarino V, Queeney J, Paschalidis I. Adversarial Imitation Learning from Visual Observations using Latent Information. Transactions on Machine Learning Research.", "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": "", "extra_scores": {}}, {"review_id": "WoFBsLXeEa", "reviewer": "Reviewer_mxsi", "summary": "Previous Wasserstein imitation learning methods couldn’t capture the environment’s dynamics accurately. The paper proposes a novel adversarial imitation learning framework, Latent Wasserstein Adversarial Imitation Learning (LWAIL), that integrates the Intention Conditioned Value Function (ICVF) to capture these dynamics to resolve the issue. By using a two-stage process, first pretraining the ICVF with low-quality data and, secondly, running a standard Wasserstein AIL framework in this new latent space, the paper shows expert-level results in multiple locomotion tasks with a single state-only expert trajectory.", "questions": "- The paper convincingly argues that an ICVF-learned latent space is beneficial for tasks like `umaze`, where the Euclidean distance is misleading due to the environment's dynamics (e.g., a wall). However, it is less clear how this specific benefit translates to the MuJoCo locomotion tasks. What is the intuition for why a dynamics-aware latent space is superior to the standard state space for tasks like Hopper or Walker2D?\n- To better validate the core contribution, have you considered testing LWAIL and key baselines on more complex environments that explicitly share the properties of Figure 1(a), such as the `ant-maze` or other complex navigation tasks from the D4RL benchmark?\n- The performance plot in Figure 4 is very dense, and many of the colors are similar, making it difficult to distinguish the learning curve for LWAIL from the baselines. Could you please consider using a more distinct color (e.g., a thick, bright red) or a different line style for your method in the final version to improve legibility?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is clearly written, well-structured, and comprehensive, making the authors' contributions and methodology easy to follow.\n- The paper's primary contribution, LWAIL, is novel in its integration of an ICVF-learned latent space into a Wasserstein AIL framework to create a dynamics-aware distance metric. This is significant as it demonstrates expert-level performance on several locomotion tasks using only a single state-only expert trajectory, a challenging and practical problem setting.\n- The empirical investigation is of high quality. The authors performed a thorough comparison against numerous baselines and provided an extensive set of ablation studies that validate the key components of their approach, such as the choice of embedding and robustness to noise.", "weaknesses": "- The empirical evaluation, while thorough on the chosen tasks, is limited to locomotion (MuJoCo) and a simple maze environment. These environments are state-based and have relatively stable dynamics. The applicability of LWAIL to more complex, high-dimensional problems (e.g., vision-based tasks) or environments with more stochastic dynamics is not explored. \n- The paper does not include an ablation study on the sensitivity of key hyperparameters. Adversarial Imitation Learning frameworks are often sensitive to settings like the discriminator update frequency (which the authors note must be balanced ) or the gradient penalty coefficient. An analysis of how performance changes with these hyperparameters would strengthen the paper's claims of robustness.\n- The paper's core motivation, powerfully illustrated in Figure 1(a), is that LWAIL's dynamics-aware metric can solve problems where simple Euclidean distance is misleading (e.g., a state near a goal is \"worse\" than a farther state with a valid path). However, the chosen experiments, while standard, do not fully stress-test this specific claim. The `umaze` environment demonstrates this principle, but is very simple. The MuJoCo locomotion tasks are complex in terms of control, but they do not necessarily exhibit the challenging topological structure implied by the motivating example. The paper would be much stronger if LWAIL and key baselines were evaluated on environments specifically designed to have this deceptive property, such as more complex mazes with traps, one-way paths, or regions that require non-Euclidean navigation to succeed.", "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": "", "extra_scores": {}}, {"review_id": "B2vWcoKNxx", "reviewer": "Reviewer_Cuqz", "summary": "The paper introduces Latent Wasserstein Adversarial Imitation Learning (LWAIL), a imitation learning framework that learns from state-only expert demonstrations without requiring expert actions or large datasets. LWAIL leverages the Wasserstein distance computed in a dynamics-aware latent embedding, which is obtained through a pre-training stage using an Intention Conditioned Value Function (ICVF) trained on a small set of random state-only data. Experiments on multiple MuJoCo tasks show that LWAIL outperforms prior imitation learning methods.", "questions": "1. **Justification of Dynamics-aware Embeddings**: My primary concern is the lack of justification for the effectiveness of the proposed dynamics-aware embeddings. Could the authors provide stronger theoretical insights or additional empirical evidence, such as a more detailed analysis of Figure 2 or an ablation study using original observations as inputs to WAIL, to demonstrate that the dynamics-aware embedding is indeed more effective than using raw observations?\n\n2. **Comparison with Latent World Models**: How does the proposed dynamics-aware latent space improve upon the latent representations learned in existing imitation learning methods that employ latent world models?\n\n3. **Sensitivity to Downstream RL Algorithms**: The results show significant performance degradation when changing the downstream RL algorithm. Could the authors provide more justification for this sensitivity? Are there potential robustness issues with the learned embeddings that may contribute to this behavior?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The paper is well-written and easy to follow.  \n- The experimental results show strong performance compared to both classic imitation learning and learning-from-observation (LfO) methods.  \n- The ablation studies and experimental details seem to be thorough and sufficient.", "weaknesses": "- **Novelty**: The paper primarily combines existing methods, as both WAIL and ICVF are established approaches. This limits the overall technical contribution and originality of the work.\n\n- **Lack of Justifications**: The paper does not sufficiently justify why using the ICVF dynamics-aware embedding space should improve performance. From my understanding, standard adversarial imitation learning (AIL) methods based on occupancy matching are already dynamics-aware—the reward or Q-function trained via the adversarial objective implicitly captures the dynamics, as noted in the original IQ-Learn paper [1]. It remains unclear why introducing an additional dynamics-aware latent embedding would further enhance performance. Moreover, while the t-SNE plot in Figure 2 shows some structure in the learned embedding, its connection to the actual system dynamics is not well established, making it difficult to support the claim that the embedding provides a more dynamics-aware metric.\n\n- **Lack of Comparison with Model-based Approaches**: Several existing model-based imitation learning methods also learn latent embeddings through a latent world model, such as [2, 3, 4]. The paper does not provide sufficient discussion or comparison regarding how the proposed dynamics-aware latent space differs from or improves upon these approaches.\n\n- **Potential Misalignment Issue**: The reward function is trained by discriminating between state-action transitions from expert and random policies. During the RL training phase, however, the policy gradually deviates from the random policy, leading to a potential distribution mismatch. The paper appears to overlook this issue and could benefit from an off-policy correction mechanism, such as importance sampling.\n\n- **Choice of Downstream RL Method**: The ablation study in Table 10 shows that changing the RL algorithm significantly degrades performance. For PPO, the poor results might be due to the aforementioned policy misalignment issue, as PPO is an on-policy method. However, even with off-policy algorithms like DDPG, performance remains unsatisfactory, raising concerns about the robustness and generalizability of the proposed embedding approach.\n\n\n[1] Garg, D., Chakraborty, S., Cundy, C., Song, J., & Ermon, S. (2021). Iq-learn: Inverse soft-q learning for imitation. Advances in Neural Information Processing Systems, 34, 4028-4039.\n\n[2] Yin, Z. H., Ye, W., Chen, Q., & Gao, Y. (2022). Planning for sample efficient imitation learning. Advances in Neural Information Processing Systems, 35, 2577-2589.\n\n[3] Li, S., Huang, Z., & Su, H. (2025). Reward-free World Models for Online Imitation Learning. In Forty-second International Conference on Machine Learning.\n\n[4] Kolev, V., Rafailov, R., Hatch, K., Wu, J., & Finn, C. (2024, June). Efficient imitation learning with conservative world models. In 6th Annual Learning for Dynamics & Control Conference (pp. 1777-1790). PMLR.", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)