Align Your Intents: Offline Imitation Learning via Optimal Transport

## Abstract

Abstract Offline Reinforcement Learning (RL) addresses the problem of sequential decision-making by learning optimal policy through pre-collected data, without interacting with the environment.
As yet, it has remained somewhat impractical, because one rarely knows the reward explicitly and it is hard to distill it retrospectively.
Here, we show that an imitating agent can still learn the desired behavior merely from observing the expert, despite the absence of explicit rewards or action labels.
In our method, AILOT (Aligned Imitation Learning via Optimal Transport), we involve special representation of states in a form of intents that incorporate pairwise spatial distances within the data.
Given such representations, we define intrinsic reward function via optimal transport distance between the expert’s and the agent’s trajectories.
We report that AILOT outperforms state-of-the art offline imitation learning algorithms on D4RL benchmarks and improves the performance of other offline RL algorithms by dense reward relabelling in the sparse-reward tasks.

## 1 Introduction

Over the past years, offline learning has remained both the most logical and the most ambitious avenue for the development of RL.
On the one hand, there is an ever-growing reservoir of sequential data, such as videos, becoming available for training the decision-making of RL agents.
On the other hand, these immense data remain largely unlabeled and unstructured for gaining any valuable guidance in the form of a tractable objective for learning the underlying policy, stimulating the development of unsupervised and self-supervised methods (Singh et al., ; Li et al., ; Sinha et al., ; Yu et al., ; Eysenbach et al., ; Park et al., ).

Following the success of language-based foundational models, the offline RL community have also begun to leverage the weakly-labeled data.
However, incorporating such offline data into RL frameworks efficiently remains a challenge (Levine et al., )).
Several factors, such as distribution shift, slow convergence when the labels are missing, and the absence of known rewards or task-specific objectives, hinder the development of offline RL (Kumar et al., ; Fujimoto et al., ).

A potential remedy to these issues is found in Imitation Learning (IL), where the explicit reward function is not needed. Instead, an imitating agent is trained to replicate the behavior of the expert. Behavior cloning (BC) (Ross & Bagnell, ) frames the IL problem akin to classical supervised learning, seeking to maximize the likelihood of the actions provided under the learner’s policy. Despite working well in simple environments, BC is prone to accumulating errors in states, coming from different distributions other than that of the expert. Other notable methods include Inverse RL (Torabi et al., ; Zolna et al., ) and action pseudo-labeling (Kumar et al., ), however they are rarely used in practice due to the introduced overhead. Distribution matching is another promising IL paradigm, where the approaches such as DIstribution Correction Estimation ("DICE"), including SMODICE (Ma et al., ), attempt to match the state-occupancy measures between the imitator’s and the expert’s policies.
However, these methods posses several limitations: 1) they require a non-zero overlap between the supports of the agent and the expert and 2) they mostly use KL-divergence (or, generally, f-divergences), which ignores the underlying geometry of the spaces where the distributions are defined.

Another line of development is associated with Computational Optimal Transport, a domain that has garnered significant popularity for addressing diverse machine learning tasks (Peyré et al., ). Its applications span from domain adaptation to generative modeling (Salimans et al., ; Rout et al., ; Korotin et al., ), including weakly-supervised and unsupervised methods (Bespalov et al., ; ).
Optimal transport theory was proposed to alleviate the necessity for manual reward engineering (Luo et al., ) via establishing an optimal coupling between a few high-quality expert demonstrations and the trajectories of the learning agent, which proved to be useful for Imitation Learning tasks (Haldar et al., ; ). However, performing OT matching in the high-dimensional unstructured space is difficult, thus limiting the approach to low-dimensional tasks.

The absence of reward labels is not the sole challenge; obtaining access to the expert actions can also be problematic.
When no rewards or action labels accompany the abundant demonstrations, one potential remedy is to create a simulator (Krylov et al., ; Saboo et al., ) that enables access to both actions and rewards. However, adopting such an approach entails substantial additional work.
Here, we eliminate the requirement for knowing the rewards or the action labels of the expert altogether.
We propose to map the initial state space to the space of intentions, where there is some high-level semantic imaginary goal that the agent needs to achieve to imitate the expert.
Aligning intents of the agent with those of the expert via Optimal Transport represents a new pathway for training efficient offline RL models.

Our contribution are as follows:

- •
A new intrinsic dense reward relabelling algorithm that ‘comprehends’ the demonstrated dynamics in the environment, on top of which any offline RL method can be used.
- •
We outperform the state-of-the-art models in the majority of Offline RL benchmarks. We do so even without knowing the expert’s action labels and the ground truth rewards.
- •
We report extensive comparison with previous works. We show that our approach enables custom imitation even if the agent’s data are a mix of random policies.

## 2 Related Works

In this study, we broaden the application of offline Reinforcement Learning (RL) to datasets that lack both the rewards and the action labels.
Despite all the research effort on learning the intrinsic rewards for RL, most works assume either online RL setting (Brown & Niekum, ; Yu et al., ; Ibarz et al., ) or the unrealistic setup of possessing some annotated prior data. Moreover, the problem of learning and exploration in sparse-reward environments can be deemed as solved for online RL (Li et al., ; Eysenbach et al., ; Lee et al., ); whereas, only a few works exist that consider the offline goal-conditioned setting with reward-free data (Park et al., ; Zheng et al., ; Wang et al., ).
The goal of our work is to extract guidance from the expert by forcing optimal alignment of information-rich representations of imaginary goals between the expert and the agent in a shared distance-aware latent space.

We build upon the approach introduced in OTR (Luo et al., ), extending it by finding a representative distance-preserving isometry to the shared latent space and forcing alignment via optimal transport in this metric-aware space.
In OTR, the authors estimate the optimal coupling between expert and agent trajectories, which enables rewards relabeling for agent’s state transitions by computed Wasserstein distance. However, performing optimal transportation matching in this way is sensitive to underlying cost function (e.g Cosine or Euclidean) and final results can vary drastically. Choosing right distance function requires additional knowledge of the environment and extensive search, limiting application to simple tasks.

An advantage of our method is in its ability to perform distribution alignment in the space, which captures temporally structured dependencies in the provided datasets, which is completely ignored in previous works. For example, a lot of methods focus on KL divergence, which is agnostic to the distance metric Such temporal structure induces spatial closeness between temporally similiar states.Additionally, the method seamlessly integrates with various downstream RL algorithms, providing flexibility in selecting the most suitable training approach.

Another recent method called Calibrated Latent gUidancE (CLUE) (Liu et al., ), introduces a parallel approach for deriving intrinsic rewards. In this study, the authors employ a conditional variational auto-encoder trained on both expert and agent transitions and compute the euclidean distance between the collapsed expert embedding and the agent trajectory.
The expert embedding might not collapse into a single point in a multimodal expert dataset, requiring clustering to handle different skills. Unlike the CLUE method, our approach doesn’t need state-action paired labeled data, eliminating the need for action annotations.

## 3 Preliminaries

### 3.1 Problem Formulation

A standard Reinforcement Learning problem is defined as a Markov Decision Process (MDP) with tuple $\mathcal{M}=(\mathcal{S},\mathcal{A},p,r,\rho_{0},\gamma)$, where $\mathcal{S}\subset\mathbb{R}^{n}$ is the state space, $\mathcal{A}\subset\mathbb{R}^{m}$ is the action space, $p$ is a function describing transition dynamics in the environment $p$: $\mathcal{S}\times\mathcal{A}\rightarrow\mathcal{P}(\mathcal{S})$, $r:\mathcal{S}\times\mathcal{A}\rightarrow\mathbb{R}$ is a predefined extrinsic reward function, $\rho_{0}$ is an initial state distribution and $\gamma\in(0,1]$ is the discount factor. The objective is to learn policy $\pi_{\theta}(a|s):\mathcal{S}\rightarrow\mathcal{P}(\mathcal{A})$ that maximizes discounted cumulative return $\mathbb{E}[\sum^{\infty}_{t=0}\gamma^{t}r(s_{t},a_{t})]$. In contrast, offline RL assumes access to a pre-collected static dataset of transitions $\mathcal{D}=\{(s_{i},a_{i},s^{\prime}_{i},r_{i})\}^{n}_{i=1}$ while prohibiting additional interaction with the environment. In this study, we assume access to a dataset of transitions without reward labels $\mathcal{D}_{a}=\{(s_{i},a_{i},s^{\prime}_{i})\}^{n}_{i=1}$, collected from $\mathcal{M}$, and a limited number of ground truth demonstrations from the expert policy without reward and action labels $\mathcal{D}_{e}=\{(s_{i},s^{\prime}_{i})\}^{m}_{i=1}$. These demonstrations are collected from the same MDP ($\mathcal{M}$) and we assume that several trajectories from $\mathcal{D}_{e}$ have high cumulative return. The primary objective is to determine a policy that closely emulates the behavior of the expert and thereby maximises the cumulative return.

### 3.2 Reward Relabelling through Optimal Transport

Optimal Transport provides a reasonable and efficient way of comparing two probability measures supported on high-dimensional measure spaces. Given a defined cost function $c(\cdot,\cdot)$ over $\mathcal{S}\times\mathcal{S}$, the Wasserstein distance between the agent trajectory $\tau^{a}=\{s^{a}_{1},s^{a}_{2},\cdots,s^{a}_{T}\}\subset\mathcal{D}_{a}$ and the expert trajectory $\tau^{e}=\{s^{e}_{1},s^{e}_{2},\cdots,s^{e}_{T}\}\subset\mathcal{D}_{e}$ is defined as:

$$ $W(\tau^{a},\tau^{e})=\min_{P\in\mathbb{R}^{T\times T}}\sum^{T}_{i=1}\sum^{T}_{ j=1}c(s^{a}_{i},s^{e}_{j})P_{ij},$ (1) $$

satisfying the marginal constraints, with $p^{a},p^{e}$ being the state occupancy of the agent data and the expert dataset respectively:

$$ $\sum^{T}_{i=1}P_{ij}=p^{e}(s_{j});\quad\sum^{T}_{j=1}P_{ij}=p^{a}(s_{i}).$ (2) $$

Then, a proxy reward function, given the optimal transport plan $P^{*}(\tau_{a},\tau_{e})$ to transform the agent state $s^{a}_{i}$ into the closest state of the expert, is:

$$ $r_{i}=-\sum^{T}_{j=1}c(s^{a}_{i},s^{e}_{j})P^{*}_{ij}(\tau_{a},\tau_{e}).$ (3) $$

In this work, instead of computing the optimal transport between the initial states from the datasets, we propose to extract intentions of the expert in the form of an isometry map to the latent space, which enables a more accurate modelling of the distribution of the similiar states (based on temporal distances between them).

Figure: Figure 1: AILOT: Aligned Imitation Learning via Optimal Transport. Principal diagram for intents alignment of states in two different trajectories (in blue and red). Left: Stage I: projection into intent space (denoted by $\psi$ in text) (first $3$ principal components are shown). Middle: Stage II: computation of intrinsic rewards for offline RL, where $s_{i}^{a}$ and $s_{j}^{e}$ are the expert’s and the agent’s states, $P$ is the optimal coupling matrix; corresponding intrinsic reward $r(s_{i}^{a})$ is a scaling transform of the product $\sum_{j}P_{ij}C_{ij}$ with some cost function $C$ defined on the intents pairs. Right: squared norm vs. steps count in the same trajectory for the states and the intents differences; It demonstrates that distance between intents is proportional to the total path length (steps count) between the states (AntMaze example is shown).
Refer to caption: x1.png

### 3.3 Pretraining of Distance Preserving Representation

Using pre-collected large unlabeled datasets in the form of prior data enables learning diversified behaviors in the form of skills or general representations. Even in the absence of the ground truth actions or rewards, the policies can still acquire valuable insights about the dynamics in the environment.
Several works tried to define valuable objective for learning from the unstructured data, including offline GCRL (Eysenbach et al., ; Park et al., ) and skill discovery (Eysenbach et al., ; Sharma et al., ).
Finding general mapping $\psi:\mathcal{S}\rightarrow\mathbb{R}^{d}$ which is a lossless compression of a state was addressed in several formulations including successor features (SFs) (Dayan, ; Barreto et al., ; Borsa et al., ) and Forward-Backward (FB) representations (Touati & Ollivier, ) with their extensions to the generalized value functions (Hansen et al., ; Touati et al., ; Ghosh et al., ; Bhateja et al., ). In the same vein, we learn successor features but enforce additional constraint on temporal structure.
Finding optimal temporal goal-conditioned value function $V^{*}$ (Wang et al., ) which depicts the minimal number of steps, required to arrive at certain goal $g$ from state $s$ to $s_{+}$ can be viewed as goal conditioned RL maximizing $r(s,g)=-\mathbbm{1}(s\neq g)$. ICVF (Ghosh et al., ) exchanges $g$ with intent $z$, generalizing previous definition of optimal goal-conditioned value function as

$$ $V^{*}(s,s_{+},z)=\mathbb{E}_{s_{t+1}\sim P_{z}(\cdot|s_{t})}\left[\sum_{t\geq 0 }-\mathbbm{1}(s_{t}\neq s_{+})\bigg{|}s_{0}=s\right],$ (4) $$

where intent $z$ completely specifies the state-occupancy dynamics through the transition function $P_{z}$.
In practice, $V$ is parametrized by three neural networks with $\phi,\psi:\mathcal{S}\to\mathbb{R}^{d}$ and $T:\mathbb{R}^{d}\to\mathbb{R}^{d\times d}$ being a transition matrix, incorporating all possible transition dynamics of each $z$:

$$ $V(s,s_{+},z)=\phi(s)^{T}T(z)\psi(s_{+}).$ (5) $$

During training, the following temporal distance loss is minimized:

$$ $\mathcal{L}^{2}_{\tau}(-\mathbbm{1}(s\neq s_{+})+\gamma\bar{V}(s^{\prime},s_{+ },z)-V(s,s_{+},z)),$ (6) $$

with $\mathcal{L}^{2}_{\tau}(x)=|\tau-\mathbbm{1}(A<0)|x^{2}$, $A=r_{z}(s)+\gamma\bar{V}(s^{\prime},s_{+},z)-\bar{V}(s,s_{+},z)$ being an expectile loss (Kostrikov et al., ), $A$ is an advantage of acting according to $z$, and $\bar{V}$ is a ‘delayed’ copy (target network) of $V$.
Such a representation estimates an average path length between the states $s$ and $s_{+}$, conditioning on guidance of all possible choices of $z$. The pretraining procedure enables a downstream extraction of learned representations $\psi,\phi$ useful for the other tasks. In our experiments, we set intent of a state $s$ as $z(s)=\psi(s)$ and $r_{z}(s)=-\mathbbm{1}(s\neq s_{z})$ with a randomly sampled future state $s_{z}$ from the same trajectory.
The $\psi$ mapping itself (without $\phi$ and $T$) serves as a good estimate for the distances between the states of the environment. In the Experiments section below, we will empirically support this, showing that a squared distance $d(s_{t+k},s_{t})=\|\psi(s_{t+k})-\psi(s_{t})\|^{2}$ is roughly a linear function of the steps count between the states $k$, which maps temporally similar states in the initial state space to the spatially closest in the intent space. Thus, applying optimal transport tools in this space makes it robust to noise or inaccuracies in the initial state space of the environment. Extracting the intents from the pretrained general-purpose value function enables capturing actual behavior from the trajectory.

We also provide a theoretical confirmation of the convergence of extracted intents to the temporal distance function between states (ref. A.1), showing
explicitly that if the intentions of states $s$ and $s_{+}$ converges to each other, then the function $V(s,s_{+})$ approaches zero in the limit.

## 4 Method

Opposite to Luo et al. (Luo et al., ), we change perspective of aligning occupancy measures from the state space into the alignment of intents in the metric-aware latent space by performing isometric transformation by terms of temporally preserving function (Section 3.3). First, we briefly outline how the optimal transportation is computed in the state space.

For two trajectories $\{s^{a}_{i}\}_{i=1}^{T_{a}}$ and $\{s^{e}_{j}\}_{j=1}^{T_{e}}$, the optimal transition matrix is obtained by solving the entropy-regularized (Sinkhorn) (Cuturi, ) OT problem:

$$ $P^{*}=\underset{P\in\Pi[T_{a},T_{e}]}{\arg\min}\left\{\sum_{ij}P_{ij}C_{ij}+ \varepsilon\sum_{ij}P_{ij}\log P_{ij}\right\},$ (7) $$

where $C_{ij}=c(s^{a}_{i},s^{a}_{\min(i+k,T_{a})},s^{e}_{j},s^{e}_{\min(j+k,T_{e})})$ is some cost function with the parameter $k>0$ and $P$ has the following marginal distributions $\forall i,j$:
$\sum_{j=1}^{T_{e}}P_{ij}=1/T_{a},\quad\sum_{i=1}^{T_{a}}P_{ij}=1/T_{e}$, which we take to be uniform across both trajectories.
By finding the optimal matrix $P^{*}$, we can determine the reward for each state of the trajectory $\{s^{a}_{i}\}_{i=1}^{T_{a}}$ as

$$ $r(s_{i}^{a})=\alpha\exp\left(-\tau T_{a}\sum_{j=j_{1}}^{T_{e}}P^{*}_{ij}C_{ij} \right),$ (8) $$

where

$$ $j_{1}=\arg\min_{j}C_{1j},$ (9) $$

and the exponent function with a scalar hyperparameters $\alpha$ and $\tau$ serves as an additional scaling factor to diminish the impact of the states with a large total cost.
The negative sign ensures that, in the process of maximizing the sum of rewards, we are effectively minimizing the optimal transport (OT) distance.

Figure: Algorithm 1 AILOT Training

We proceed to estimate the transition matrix $P^{*}$ between each trajectory of the agent and the expert’s demonstrations (which may consist of one or more trajectories) as described in Section 3.2.
For enhanced alignment, we selectively consider the tail of the expert’s trajectory $j_{1},\ldots,T_{e}$, focusing on the nearest first states to the agent’s starting position according to the cost matrix $C$.
The maximum reward is selected across different trajectories when the expert provides multiple trajectories.
The aggregated rewards for each trajectory of the agent are then incorporated into the offline dataset for subsequent RL training, with the policy trained by the IQL offline RL algorithm (Kostrikov et al., ). We use ICVF representations only in rewards definition and employ IQL without any modifications. A comprehensive recipe for the rewards computation is shown in Algorithm 1.

Unlike the OTR method (Luo et al., ) which computes distances directly between the states, our approach measures differences between intentions $\psi(s)$.
Specifically, we propose to incorporate the following cost matrix:

$$ $C_{ij}=\|\psi(s^{a}_{i})-\psi(s^{e}_{j})\|^{2}+\|\psi(s^{a}_{\min(i+k,T_{a})}) -\psi(s^{e}_{\min(j+k,T_{e})})\|^{2}.$ (10) $$

The second term in the cost function is necessary for an ordered comparison of the trajectories.
In imitation learning, we always compare distributions of pairs
of the expert and the agent states. This is so because we want to be not only in the same state, but we also want to act in a manner similar to the expert.
Hence, during training, we enforce the distribution of the agent’s intent pairs to converge to the empirical measure of the intent pairs of the expert. The motivation behind the extraction of the behaviors comes from the observation that the time-step differences in the state space provide erroneous values, not capturing temporal dependencies between the dataset states, which we show in Figure 3, and thus, the geometry of the environment is not represented properly, making it hard for the optimal transport to scale in high-dimensional problems.

## 5 Experiments

In this Section, we demonstrate the performance of AILOT on several benchmark tasks, report an ablation study on varying the number of provided expert demonstrations, and provide the implementation details.
First, we empirically show the ability of proposed method to efficiently utilize expert demonstrations in sparse-reward tasks (such as Antmaze and Adroit environments) and improve learning ability of Offline RL algorithms by providing geometrically aware dense intrinsic reward signal to agent’s transitions. Second, we evaluate AILOT in the offline Imitation Learning (IL) setting, outperforming the state-of-the art offline IL algorithms on MuJoco locomotion tasks. Also, we include additional experiments with custom behaviors, e.g., the expert hopper performing a backflip, when the agent dataset consists of completely random behaviors.

### 5.1 Implementation Details

Dense reward relabelling by AILOT is completely decoupled from the offline policy training. In all our experiments, we endow AILOT with Implicit Q-Learning (IQL) (Kostrikov et al., ), which is a simple and a robust offline RL algorithm. In our additional experiments, we also test AILOT+IQL against Diffusion-QL (Wang et al., ), which is a recent state-of-the-art approach for offline RL.

We implement AILOT in JAX (Bradbury et al., ) and use the official IQL implementation(^1^11[https://github.com/ikostrikov/implicit_q_learning](https://github.com/ikostrikov/implicit_q_learning)). Our code is written using Equinox library (Kidger & Garcia, ). To compute optimal transport, we use implementation of Sinkhorn algorithm (Cuturi, ) from OTT-JAX library (Cuturi et al., ). Additional details on chosen hyperparameters and settings are provided in the Appendix A.2.

Figure: (a)
Refer to caption: extracted/5900934/Figures/rl-video-episode-0_004.jpg

### 5.2 Runtime

All of the experiments were conducted on a single RTX 3090. The runtime overhead of the application of optimal transport is not greater than extra $\sim 10$ minutes (compared to the offline RL algorithm training time). Thus, the total runtime of AILOT+IQL equals $\sim$ 25 minutes.
Finding optimal matrix $P^{*}$ and correspondent internal rewards takes several hundred iterations of Sinkhorn solver, which is relatively fast. Moreover, because the optimal transport is computed in a fixed latent space of intents, the overhead stays constant across broad range of tasks regardless of the state dimensionality. For high-dimensional tasks (e.g., learning from pixels), AILOT is more suitable than OTR due to performing OT in the latent space, accelerating the overall relabelling.

### 5.3 Baselines

AILOT is compared against the following state-of-the-art algorithms:
IQL (Kostrikov et al., ),
OTR (Luo et al., ),
CLUE (Liu et al., ),
IQ-Learn (Garg et al., ),
Diffusion-QL (Wang et al., ),
SQIL (Reddy et al., ),
ORIL (Zolna et al., ),
and
SMODICE (Ma et al., ).
The details are given in Appendix B.

**Table 1: Normalized scores (mean $\pm$ standard deviation) of AILOT on MuJoco locomotion tasks, compared to baselines. The upper sub-table includes methods for imitation learning without any kind of optimal transport. The lower sub-table shows results for the conceptually close approaches – oracle IQL, OTR, and CLUE. For these methods the results are given for $K=1$ number of expert trajectories. The highest scores are highlighted in green.**
| Dataset | IQ-Learn | SQIL | ORIL | SMODICE | AILOT |
| --- | --- | --- | --- | --- | --- |
| halfcheetah-medium-v2 | $21.7$ $\pm\>1.5$ | $24.3$ $\pm\>2.7$ | $56.8$ $\pm\>1.2$ | $42.4$ $\pm\>0.6$ | $47.7$ $\pm\>0.2$ |
| halfcheetah-medium-replay-v2 | $6.7$ $\pm\>1.8$ | $43.8$ $\pm\>1.0$ | $46.2$ $\pm\>1.1$ | $38.3$ $\pm\>2.0$ | $42.4$ $\pm\>0.2$ |
| halfcheetah-medium-expert-v2 | $2.0$ $\pm\>0.4$ | $6.7$ $\pm\>1.2$ | $48.7$ $\pm\>2.4$ | $81.0$ $\pm\>2.3$ | $92.4$ $\pm\>1.5$ |
| hopper-medium-v2 | $29.6$ $\pm\>5.2$ | $66.9$ $\pm\>5.1$ | $96.3$ $\pm\>0.9$ | $54.8$ $\pm\>1.2$ | $82.2$ $\pm\>5.6$ |
| hopper-medium-replay-v2 | $23.0$ $\pm\>9.4$ | $98.6$ $\pm\>0.7$ | $56.7$ $\pm\>12.9$ | $30.4$ $\pm\>1.2$ | $98.7$ $\pm\>0.4$ |
| hopper-medium-expert-v2 | $9.1$ $\pm\>2.2$ | $13.6$ $\pm\>9.6$ | $25.1$ $\pm\>12.8$ | $82.4$ $\pm\>7.7$ | $103.4$ $\pm\>5.3$ |
| walker2d-medium-v2 | $5.7$ $\pm\>4.0$ | $51.9$ $\pm\>11.7$ | $20.4$ $\pm\>13.6$ | $67.8$ $\pm\>6.0$ | $78.3$ $\pm\>0.8$ |
| walker2d-medium-replay-v2 | $17.0$ $\pm\>7.6$ | $42.3$ $\pm\>5.8$ | $71.8$ $\pm\>9.6$ | $49.7$ $\pm\>4.6$ | $77.5$ $\pm\>3.1$ |
| walker2d-medium-expert-v2 | $7.7$ $\pm\>2.4$ | $18.8$ $\pm\>13.1$ | $11.6$ $\pm\>14.7$ | $94.8$ $\pm\>11.1$ | $110.2$ $\pm\>1.2$ |
| D4RL Locomotion total | $122.5$ | $366.9$ | $433.6$ | $541.6$ | $732.8$ |

**Table 2: Normalized scores (mean $\pm$ standard deviation) of AILOT on sparse-reward environments. We compare different dense relabelling methods (OTR, CLUE) and show that AILOT outperforms those approaches and accelerates learning of offline IQL method.**
| Dataset | IQL | OTR | CLUE | AILOT |
| --- | --- | --- | --- | --- |
| umaze-v2 | $88.7$ | $81.6$ $\pm\>7.3$ | $92.1$ $\pm\>3.9$ | $93.5$ $\pm\>4.8$ |
| umaze-diverse-v2 | $67.5$ | $70.4$ $\pm\>8.9$ | $68.0$ $\pm\>11.2$ | $63.4$ $\pm\>7.6$ |
| medium-play-v2 | $72.9$ | $73.9$ $\pm\>6.0$ | $75.3$ $\pm\>6.3$ | $71.3$ $\pm\>5.2$ |
| medium-diverse-v2 | $72.1$ | $72.5$ $\pm\>6.9$ | $74.6$ $\pm\>7.5$ | $75.5$ $\pm\>7.4$ |
| large-play-v2 | $43.2$ | $49.7$ $\pm\>6.9$ | $55.8$$\pm\>7.7$ | $57.6$ $\pm\>6.6$ |
| large-diverse-v2 | $46.9$ | $48.1$ $\pm\>7.9$ | $49.9$ $\pm\>6.9$ | $66.6$ $\pm\>3.1$ |
| AntMaze-v2 total | $391.3$ | $396.2$ | $415.7$ | $427.9$ |

**Table 3: Normalized scores (mean $\pm$ standard deviation) of AILOT on Adroit and AntMaze tasks, compared to baselines. OTR, CLUE, and AILOT use IQL as offline RL baseline algorithm with only a single expert episode. The highest scores are highlighted.**
| Dataset | IQL | OTR | CLUE | AILOT |
| --- | --- | --- | --- | --- |
| door-cloned-v0 | $1.6$ | $0.01$ $\pm\>0.01$ | $0.02$ $\pm\>0.01$ | $0.05$ $\pm\>0.02$ |
| door-human-v0 | $4.3$ | $5.9$$\pm\>2.7$ | $7.7$ $\pm\>3.9$ | $7.9$ $\pm\>3.2$ |
| hammer-cloned-v0 | $2.1$ | $0.9$$\pm\>0.3$ | $1.4$ $\pm\>1.0$ | $1.6$ $\pm\>0.1$ |
| hammer-human-v0 | $1.4$ | $1.8$$\pm\>1.4$ | $1.9$ $\pm\>1.2$ | $1.8$ $\pm 1.3$ |
| pen-cloned-v0 | $37.3$ | $46.9$$\pm\>20.9$ | $59.4$ $\pm\>21.1$ | $61.4$ $\pm\>19.5$ |
| pen-human-v0 | $71.5$ | $66.8$$\pm\>21.2$ | $82.9$ $\pm\>20.2$ | $89.4$ $\pm\>0.1$ |
| relocate-cloned | -$0.2$ | -$0.24$$\pm\>0.03$ | -$0.23$ $\pm\>0.02$ | -$0.20$ $\pm\>0.03$ |
| relocate-human | $0.1$ | $0.1$$\pm\>0.1$ | $0.2$ $\pm\>0.3$ | $0.28$ $\pm\>0.1$ |
| Adroit-v0 total | $118.1$ | $122.2$ | $153.3$ | $162.2$ |

### 5.4 Results & Experiments

We present the results for the following offline RL settings: (1) the imitation setting, where the goal is to mimic the behavior as closely as possible given several expert demonstrations in the form of trajectories; (2) the sparsified reward tasks, where the reward of one is given only when the agent reaches the goal state.
We show that the rewards obtained through AILOT relabelling in both settings are descriptive enough to recover the demonstrated policy. We evaluate proposed method on common D4RL (Fu et al., ) benchmark and consider following domains: Gym, Adroid and Antmaze. Additional details on the environments are included in the Appendix F.

#### Offline Imitation Learning.

We compare AILOT performance on IL tasks in the offline setting. D4RL MuJoco locomotion datasets are used as the main source of offline data with the original reward signal and the action labels discarded, and the expert trajectories chosen for each task as the best episodes achieving maximal return. Results for D4RL locomotion tasks are presented in Table 1. AILOT achieves the best performance in $7$ out of $9$ benchmarks, when compared to OTR and CLUE. Note that, unlike the CLUE method, our algorithm does not rely on the action labels. Empirically, we confirm that the optimal geometrically aware map between the expert and the agent trajectories in the intents latent space gives reasonable guidance for the agent to learn properly. Since representation $\psi(s)$ was pretrained by general objective Eq.5 for all possible $z$, we can extract it to obtain semantic behavior of policy, which produced high return expert trajectories.
We chose Euclidean distance as a measure of similarity in the intent latent space as in Eq.10 and perform exponential scaling of rewards according to Eq.8 in order to maintain them in the appropriate range.

Also, we include additional imitation results when a custom demonstrations are provided, as in Figure 2 and Figure 5, where the Hopper expert makes a backward flip and the HalfCheetah stands upwards respectively.
For those tasks, agent’s dataset initially consists of only the random behaviors and agent learns to recover the desired behavior through the AILOT dense relabelling.

#### Sparse-Reward Offline RL Tasks.

Next, we evaluate AILOT on several sparse D4RL benchmarks (AntMaze-v2 and Adroit-v2). To obtain expert trajectories, we consider only those episodes that accomplish the goal task, dismissing all the others.

Table 3 compares the performance of AILOT+IQL to OTR+IQL,  CLUE+IQL, when only a single demonstration trajectory is available, and to the original IQL. We employed IQL, which acts as a baseline, without any modifications. Given that OTR + IQL always performs better than IQL, we decided to compare directly with OTR, which showcases the importance of performing optimal transport alignment under distance preserving mapping.
We observe that AILOT outperforms the current state-of-the art results, with the hardest antmaze-large-diverse task showing the most remarkable margin.
This proves the ability of AILOT to employ the expressive representations from general pretrained value function (5) through the optimal transport for functional learning in the sparse tasks as well.
In door-cloned-v0 and hammer-cloned-v0 tasks IQL performs best with sparse rewards. The inferior results in those two tasks stem from the fact that the cloned version was obtained from collecting the data from trained imitation learning policy on a mix of human and the expert data. Human data contains many ambiguous trajectories, from which it is hard to extract valuable intents (this is the case for both the door and the hammer-cloned tasks).

### 5.5 Ablation Study

#### Varying the number of expert trajectories.

We investigate whether performance of learned behavior tends to improve with increased number of provided expert trajectories. Table 7 in the Appendix shows overall performance for varying number of expert trajectories for $K=1$ to $K=5$ across OTR and AILOT. However, we observe that performance across both algorithms improves slightly. We make comparison with OTR here because it’s the most similar to ours (it also uses OT for reward labelling and IQL for RL problem solution). Still, AILOT achieves better normalized scores than OTR, thus proving that alignment of intents in geometry-aware space improves intrinsic rewards labelling in comparison with similar rewards but with pairwise distances between original states as done in OTR.

Figure: Figure 3: The squared norm of states and intents differences, depending on the total steps count between the states in the same trajectory. The intents differences $||\psi(s_{t+k})-\psi(s_{t})||^{2}$ has a near-linear dependence on the steps count. The squared norm in the state space is not a monotone function, which is less efficient for training an imitating agent, since it completely ignores global geometric dependencies between states in the dataset.
Refer to caption: x2.png

#### Intents distance dependence on the steps count.

By learning temporally preserving value function (i.e., similar states get mapped to temporally closed points in the intents space), the underlying space becomes more structured. Because the optimal transport takes the geometry of space into account, the task of finding the alignment becomes simpler.
In the Method section, we have already mentioned that $\mathbb{E}_{t}\|\psi(s_{t+k})-\psi(s_{t})\|^{2}$ is near-linearly dependent on the steps count between the states ($k$). This important property plays a constructive role in the cost matrix in the OT problem (10), and ultimately, gives a good estimate of the distance between the trajectories. Empirical evidence of the near-linear dependence is presented in Figure 3.
On this figures we show that the pairwise distances between the original states $\mathbb{E}_{t}\|s_{t+k}-s_{t}\|^{2}$, which is what used in previous OT methods, have no such feature and their direct use as a cost function is less preferable, since they discarding global temporal geometry.
Additionally, we compare AILOT to OTR with the same $\alpha$ and $\beta$ hyperparameters in the Appendix, showcasing that performance gains come from performing OT in temporally grounded latent space.

## 6 Discussion

In our work, we introduced AILOT – a new non state-occupancy estimation method for extracting expert’s behavior in terms of its intentions and guiding the learning agent to them through an intrinsic reward. We empirically show that it surpasses the state-of-the art results both in the sparse-reward RL tasks and in the offline imitation learning setting. AILOT can mimic the expert behaviour without knowing its action labels, and without the ground truth rewards. Moreover, we show that the intrinsic rewards, distilled by AILOT, could be used to efficiently boost the performance of other offline RL algorithms, thanks to the proper alignment to the expert intentions via the optimal transport.

#### Limitations

of AILOT include assumption on sufficient access to a large number of unlabeled trajectories of some acceptable quality. In our work, we have focused on the expert behaviours, typically considered in the offline RL publications: popular synthetic environments with some comprehensible expert movements and, consequently, some sufficiently ‘intuitive’ intentions, which can be easily extracted from the provided datasets.

Another limitation follows from the multi-modality of intents, because the expert can have several goals or performs a vague action. The imitation efficiency could, of course, drop, because the expert intents may no longer be transparent to the agent. While such a trait would be on par with the way humans learn a certain skill by observing an adept, weighing the hierarchy of multiple possible intents in AILOT could prove useful to further regularize the learning dynamics in such uncertain scenarios and is an interesting direction for further research.
Other direction of future work is to venture into the cross-domain imitation.
Based on the results observed here, it should be possible to generalize AILOT to handle the transition shift between the expert and the agent in the presence of larger mismatches pertinent to the different domains.

In conclusion, the development of AILOT sets a robust benchmark for future generalizations, enhancing ongoing research in crafting generalist knowledge distillation agents from the offline data.

## 7 Reproducibility Statement

We included source code in the supplementary material. All of the hyperparameters used, along with
additional dataset preprocessing step, are discussed in Appendix A.2. In our experiments, we
used open-source D4RL datasets (Fu et al., ). In order to ensure reproducibility and account for
randomness, we repeated our experiments over 10 random seeds.

## Appendix A Appendix

Figure: Figure 4: Principal diagram of AILOT approach for imitation learning.
Refer to caption: extracted/5900934/Figures/flow_alg_1_offlinerl.png

### A.1 Value function convergence

###### Proposition 1

Consider all possible pairs of current and goal states ($s,s_{+}$) with $z=\psi(s_{+})$. Assume that functions $\phi(s)$, $\psi(s_{+})$, $T(\psi(s_{+}))$ (decomposition of $V(s,s_{+},z)$ from eq. 5) and expression

$$ $\left\|\phi^{T}(s)\frac{\partial T(\psi_{0})}{\partial\psi}\psi(s_{+})\right\|$ (11) $$

are upper-bounded by some constant value. Here $\psi_{0}$ is some intermediate value between $\psi(s)$ and $\psi(s_{+})$. Then from convergence of intent embedding $\psi(s)\to\psi(s_{+})$ (w.r.t Euclidean norm) it follows

$$ $V(s,s_{+},z)\to 0.$ $$

Proof: Define $V(s,s_{+})=V(s,s_{+},\psi(s_{+}))$ and
$\delta=\|\psi(s)-\psi(s_{+})\|.$
The goal is to find an upper bound $V(s,s_{+})$.
Since $V(s,s)=0$, $-V(s,s_{+})$ can be rewritten as

$$ $-V(s,s_{+})=-V(s,s_{+})+V(s,s).$ (12) $$

Take by definition $V(s,s_{+})=\phi^{T}(s)T(\psi(s_{+}))\psi(s_{+})$ and obtain

$$ $\displaystyle-V(s,s_{+})$ $\displaystyle=\phi^{T}(s)T(\psi(s))\psi(s)-\phi^{T}(s)T(\psi(s_{+}))\psi(s_{+})$ (13) $\displaystyle=\phi^{T}(s)T(\psi(s))[\psi(s)-\psi(s_{+})]+\phi^{T}(s)[T(\psi(s) )-T(\psi(s_{+}))]\psi(s_{+})$ (14) $\displaystyle\leq\|\phi^{T}(s)T(\psi(s))\|\delta+\left\|\phi^{T}(s)\frac{ \partial T(\psi_{0})}{\partial\psi}\psi(s_{+})\right\|\delta.$ (15) $$

Since the norms in the expression above are upper-bounded then from $\delta\to 0$ follows $V(s,s_{+})\to 0$.

###### Proposition 1

### A.2 Hyperparameters

We found that $n=200$ Sinkhorn solver iterations are enough to find optimal coupling and set $\epsilon=0.001$ as entropy regularization parameter. All hyperparameters for IQL on D4RL benchmarks are set to those recommended in original paper to ensure reproducibility. We use default latent dimension of $d=256$ for general value function from Eq. 5 embeddings pretraining across all benchmarks and set other hyperparameters as in original paper. In order to maintain rewards in a reasonable range, they are scaled by exponential function as written in Equation 8 with $\alpha=5,\tau=0.5$ for MuJoco and AntMaze tasks and $\alpha=5,\tau=10$ for Adroit. We found that lookahead parameter in Eq. 10 works best for $k=2$. Results are evaluated across 10 random seeds and 10 evaluation episodes for each seed in order to be consistent with previous works.

It should be noted that IQL incorporates its own rewards rescaling function within the dataset. We apply similar technique, using the reward scaling factor of
$\displaystyle\frac{1000}{\mbox{max_return}-\mbox{min_return}}.$

For AntMaze-v2 tasks, we subtract 1 from the rewards outputted by AILOT. All other parameters for IQL for other tasks are kept intact. Full list of crucial parameters are presented in Tables 4 and 5. For AntMaze tasks, the parameters are the same as for Mujoco, except for the expectile in IQL, which we set to $0.9$. Pretraining of general value function is executed for $250$k steps, which we found to be enough to provide reasonable distance estimation, which coincides with original ICVF implementation.

**Table 4: Hyperparameters for MuJoco tasks.**
| Algorithm | Hyperparameter | Value |
| --- | --- | --- |
| IQL | Temperature | 6 |
| Expectile | 0.7 |  |
| Hidden layers | (256, 256) |  |
| Optimizer | Adam |  |
| Critic learning rate | $3e^{-4}$ |  |
| Value learning rate | $3e^{-4}$ |  |
| Policy learning rate | $3e^{-4}$ |  |
| AILOT | $\tau$ | 0.5 |
| $\alpha$ | 5 |  |
| Sinkhorn $\epsilon$ | 0.001 |  |
| k | 2 |  |

**Table 5: Hyperparameters for Adroit tasks. IQL parameters are the same as for the locomotion tasks.**
| Algorithm | Hyperparameter | Value |
| --- | --- | --- |
| IQL | Temperature | 5 |
| Expectile | 0.7 |  |
| AILOT | $\tau$ | 10 |
| $\alpha$ | 5 |  |
| Sinkhorn $\epsilon$ | 0.001 |  |
| k | 2 |  |

## Appendix B Details about the Baseline Models

Performance of AILOT + IQL is compared to the following algorithms:

- •
IQL (Kostrikov et al. ()) is state-of-the-art offline RL algorithm, which avoids querying out of the distribution actions by viewing value function as a random variable, where upper bound of uncertanity is controlled through expectile of distribution. In our experiments, evaluation is made using ground-truth reward from D4RL tasks.
- •
OTR (Luo et al. ()) is a reward function algorithm, where reward signal is based on optimal transport distance between states of expert demonstration and reward unlabeled dataset.
- •
CLUE (Liu et al. ()) learns VAE calibrated latent space of both expert and agent state-action transitions, where intrinsic rewards can be defined as distance between agent and averaged expert transition representations.
- •
IQ-Learn (Garg et al. ()) is an imitation learning algorithm, which implicitly encodes into learned inverse Q-function rewards and policy from expert data.
- •
Diffusion-QL (Wang et al. ()) is an state-of-the-art offline RL algorithm, which models learning policy as conditional diffusion model in order to effectively increase expressiveness and provide more flexible regularization towards behavior policy, which collected dataset.
- •
SQIL (Reddy et al. ()) proposes to learn soft Q-function by setting expert transitions to one and for non-expert transitions to zero.
- •
ORIL (Zolna et al. ()) utilizes discriminator network which distinguishes between optimal and suboptimal data in mixed dataset to provide reward relabelling through learned discriminator.
- •
SMODICE (Ma et al. ()) offline state occupancy matching algorithm, which solves the problem of IL from observations through state divergence minimization by utilizing dual formulation of value function.

## Appendix C Extra ablation study: Varying Number of Expert Trajectories

In this section, we report an ablation study on varying the number of provided expert trajectories, $K$. Table 7 compares AILOT and OTR for $K=1,5$. We observe that the extraction of intent representations is crucial for a good performance. We also observed a negligible performance improvement when the number of provided expert trajectories exceeds certain threshold ($K\geq 10$).

**Table 6: Normalized scores for D4RL locomotion tasks with varying number of expert trajectories. The highest scores are highlighted.**
| Dataset/Method | OTR | AILOT | OTR | AILOT |
| --- | --- | --- | --- | --- |
| # of expert trajectories | $K=1$ | $K=1$ | $K=5$ | $K=5$ |
| halfcheetah-medium-v2 | $43.3$ $\pm\>0.2$ | $47.7$$\pm\>0.2$ | $45.2$$\pm\>0.2$ | $46.6$$\pm\>0.2$ |
| halfcheetah-medium-replay-v2 | $41.3$$\pm\>0.6$ | $42.4$$\pm\>0.2$ | $41.9$$\pm\>0.3$ | $41.2$ $\pm\>0.5$ |
| halfcheetah-medium-expert-v2 | $89.6$ $\pm\>3.0$ | $92.4$$\pm\>1.5$ | $89.9$ $\pm\>1.9$ | $92.4$ $\pm\>1.0$ |
| hopper-medium-v2 | $78.7$ $\pm\>5.5$ | $82.2$$\pm\>5.6$ | $79.5$ $\pm\>5.3$ | $82.5$$\pm\>3.7$ |
| hopper-medium-replay-v2 | $84.8$ $\pm\>2.6$ | $98.7$$\pm\>0.4$ | $85.4$ $\pm\>1.7$ | $97.4$ $\pm\>0.1$ |
| hopper-medium-expert-v2 | $93.2$ $\pm\>20.6$ | $103.4$$\pm\>5.3$ | $90.4$$\pm\>21.5$ | $107.3$ $\pm\>5.6$ |
| walker2d-medium-v2 | $79.4$$\pm\>1.4$ | $78.3$$\pm\>0.8$ | $79.8$ $\pm\>1.4$ | $80.9$ $\pm\>1.4$ |
| walker2d-medium-replay-v2 | $66.0$$\pm\>1.4$ | $77.5$$\pm\>3.1$ | $71.0$ $\pm\>5.0$ | $76.9$ $\pm\>1.6$ |
| walker2d-medium-expert-v2 | $109.3$ $\pm\>0.8$ | $110.2$$\pm\>1.2$ | $109.4$ $\pm\>0.4$ | $110.3$ $\pm\>0.4$ |
| D4RL Locomotion total | $685.6$ | $732.8$ | $690.6$ | $735.5$ |

Also, we performed several experiments with modified cost function in OTR method, i.e., exactly the same as in our Eq. 10 with varying levels of $k$, which corresponds to intents shift in Algorithm 1. We observe that the lookahead parameter
introduces only negligible improvements to OTR (up to $+0.5$), which is an indicator that the temporal grounding of OT is completely lacking in OTR.

**Table 7: Normalized scores for D4RL locomotion tasks with varying number of expert trajectories. The highest scores are highlighted.**
| Dataset/Method | OTR ($k=2$) | AILOT |
| --- | --- | --- |
| halfcheetah-medium-v2 | $43.3$ $\pm\>0.1$ (+0.1) | $47.7$$\pm\>0.2$ |
| halfcheetah-medium-replay-v2 | $41.6$$\pm\>0.8$ (+0.4) | $42.4$$\pm\>0.2$ |
| halfcheetah-medium-expert-v2 | $89.8$ $\pm\>3.0$ (+0.2) | $92.4$$\pm\>1.5$ |
| hopper-medium-v2 | $79.3$ $\pm\>5.5$ (+0.6) | $82.2$$\pm\>5.6$ |
| hopper-medium-replay-v2 | $85.1$ $\pm\>2.6$ (+0.3) | $98.7$$\pm\>0.4$ |
| hopper-medium-expert-v2 | $93.5$ $\pm\>20.6$ (+0.2) | $103.4$$\pm\>5.3$ |

Figure: Figure 5: Top: HalfCheetah expert sample trajectory, performing an upwards standing which the agent should imitate; Bottom: HalfCheetah agent successfully performs imitation of standing upwards from the observations via AILOT.
Refer to caption: extracted/5900934/Figures/skill12_006.jpg

## Appendix D Additional Comparison Experiments

We also compare the performance of AILOT to Diffusion-QL and Behavior Cloning (10%), which are shown in Table 8. We evaluated Diffusion-QL with recommended parameters from the original paper and tested performance on sparse-reward AntMaze tasks.

| Dataset | DiffusionQL | AILOT |
| --- | --- | --- |
| antmaze-large-diverse | $56.6$ $\pm\>7.6$ | $66.6$ $\pm\>3.1$ |
| antmaze-large-play | $46.4$ $\pm\>8.3$ | $57.6$ $\pm\>6.6$ |

**Table 8: Peformance of AILOT in comparison to Diffusion-QL on sparse-reward AntMaze task and BC-10 results.**
| Dataset | BC-10 | AILOT |
| --- | --- | --- |
| halfcheetah-medium-v2 | $42.5$ | $47.7$ $\pm\>0.2$ |
| halfcheetah-medium-replay-v2 | 40.6 | $42.4$ $\pm\>0.2$ |
| halfcheetah-medium-expert-v2 | 92.9 | $92.4$ $\pm\>1.5$ |
| hopper-medium-v2 | 56.9 | $82.2$ $\pm\>5.6$ |
| hopper-medium-replay-v2 | 75.9 | $98.7$ $\pm\>0.4$ |
| hopper-medium-expert-v2 | 110.9 | $103.4$ $\pm\>5.3$ |
| walker2d-medium-v2 | 75.0 | $78.3$ $\pm\>0.8$ |
| walker2d-medium-replay-v2 | 62.5 | $77.5$ $\pm\>3.1$ |
| walker2d-medium-expert-v2 | 109.0 | $110.2$ $\pm\>1.2$ |

## Appendix E Ablating α 𝛼 \alpha italic_α and τ 𝜏 \tau italic_τ

In the current section we provide expriments on how hyperparameters $\alpha$ and $\tau$ are influencing overall performance. Provided results are similiar to those reported in (Luo et al., ).

**Table 9: How the hyperparameters $\alpha$ and $\tau$ influence the performance.**
| Dataset | AILOT($\alpha=5$, $\tau=0.5$) | AILOT($\alpha=1$, $\tau=1$) |
| --- | --- | --- |
| halfcheetah-medium-v2 | 47.7 $\pm\>0.35$ | 45.6 $\pm\>0.3$ |
| halfcheetah-medium-replay-v2 | 42.4 $\pm\>0.8$ | 40.2 $\pm\>0.5$ |
| halfcheetah-medium-expert-v2 | 92.4 $\pm\>1.5$ | 89.4 $\pm\>0.7$ |
| walker2d-medium-v2 | 78.3 $\pm\>0.8$ | 74.3 $\pm\>0.4$ |
| walker2d-medium-replay-v2 | 77.5 $\pm\>3.1$ | 71.7 $\pm\>2.1$ |
| walker2d-medium-expert-v2 | 110.2 $\pm\>1.2$ | 96.7 $\pm\>1.0$ |

## Appendix F Description of Evaluated Environments

Gym-MuJoCo locomotion are the most commonly used tasks for the evaluation of performance of RL algorithms. D4RL provides precollected datasets of different quality, corresponding to training policies evaluated after training for certain number of steps.

#### Gym-Locomotion.

Locomotion tasks consist of four environments: Hopper, Walker2d, HalfCheetah, and Ant. Provided datasets from D4RL come in "-expert", "-replay", and "-random" splits. Where "-expert" split corresponds to high-reward trajectories, whilst the others contain a mixture of roll-outs from partially trained policies or even some random transitions. All of the datasets are obtained with the SAC algorithm.

#### Manipulation

In addition to the locomotion experiments, we test the proposed approach on Adroit environments. Adroid involves controlling a 24-DoF Hand robot tasked with opening a door, hammering a nail, picking up and moving a ball. Two types of offline data from D4RL are considered: human demonstrations (with "-human" split), which were obtained from logged experiences of humans and data obtained from a sub-optimal policy (with "-cloned" split), obtained from the imitation policy.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 5
presentation: 1
questions: 1. Why is there minimal benefit to scaling the number of expert trajectories? How well would this method handle using expert trajectories that take different behaviors to solve the same problem (for example, the Push-T task from Diffusion Policy)

2. Could you explain better what the two different components of the cost function are doing? The text didn't well-motivate why these were chosen in this way.

3.  How sensitive is the method to `k`?

4. Would be nice to expand the discussion about how this method handles sub-optimal / orthogonal data compared to traditional offline algorithms -- Can this method "stitch" trajectories together?
soundness: 3
strengths: The problem setting is topical, and the method is simple and well-motivated -- Figure 1 in particular illustrates clearly the benefit of parameterizing distances in a latent space instead of the raw state space (or equivalent).

The experiments clearly demonstrate performance improvement over prior approaches (both based on optimal transport, or other imitation learning) in the D4RL suite. Admittedly, these tasks are relatively toy and now saturated, but even so, the results seem convincing. 

The related work throughout the paper (in intro, related work, and method section) contextualizes the contributions of this paper well.

The appendix (and experimental section) thoroughly describes the comparisons and the benchmark setting.
summary: This paper considers learning from offline data in settings where reward may be difficult to specify, but one (or multiple) expert trajectories demonstrating the behavior may be found. The general idea is to assign rewards within a trajectory based on an optimal transport distance between trajectories in the offline data, and this optimal trajectory. The primary innovation is to use a dynamical distance (ICVF) to parameterize a more semantically meaningful cost function for the optimal transport problem. The evaluation demonstrates improvement over prior approaches in this problem setting on all state-based D4RL tasks (including locomotion, antmaze, and adroit).
weaknesses: I found the writing in the paper to be difficult to comprehend at many parts, making it difficult to understand the method exactly and what the exact contributions are relative to prior work in this space (e.g. Luo et al).

For instance, the introduction barely touches on the method being proposed, instead discussing in great detail the motivation for IL methods, for optimal transport, etc. This makes it difficult to understand and contextualize the specific contributions of the method being proposed in the paper.

The paper is most closely related to Luo et al, 2023 (OTR), but within the method section, does not distinguish between what ideas come from Luo et al, and which are newly introduced in this paper. For readers who may not be familiar with this prior work, this can lead to misattribution of ideas. It would be useful (whether in the related work, background, or method) to more clearly lay out what is done in Luo et al, and what new ideas are being considered. 

The novelty of the idea (to my understanding) over Luo et al is relatively low -- this, in itself, is not a bad thing. However, given the simplicity of the idea, it would have been nice to see more thorough ablations and analyses to understand how e.g different dynamical distances perform, what types of data this is most helpful with, the importance of both components of the cost function. Another axis that could improve the thoroughness of the paper is to evaluate on more challenging domains beyond where standard cost metrics succeed (for example, in image-based domains). One other possible avenue of improvement here may be to thoroughly investigate what the actual computed rewards look like between this method and prior work. 

As it stands right now, while the method demonstrates mild improvements on D4RL, the paper could be much improved by expanding the analysis on the axes why the learned representation is much more useful, or by testing on a more difficult suite of tasks.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1.	In the experiments, AILOT is applied with Implicit Q-Learning (IQL) because it is a simple and robust offline RL algorithm. Is there any special reason or motivation for using IQL here, and will AILOT also perform well with any other offline RL algorithm?
soundness: 3
strengths: 1.	The proposed AILOT method eliminates the requirements of expert rewards and actions. Instead of performing Optimal Transport matching, AILOT maps the initial state space to the space of intentions and aligns the intents of the agent with those of the expert via Optimal Transport.  This approach involves several steps: 1) training general-purpose value functions from the expert dataset to learn the metric-aware representations; 2) solving the Optimal Transport alignment to obtain the coupling matrix; 3) reward labeling for the expert observations using the coupling matrix; and 4) training RL using the expert dataset with labeled rewards to obtain the final policy. 
2.	The intent differences between the k-step state representations have a linear dependence on the step count. This near-monotone function reflects the global geometric dependencies between states in the expert dataset. This good property is important for defining the cost function of Optimal Transport alignment learning. 
3.	The dense reward from AILOT can also boost the performance of other offline reinforcement learning methods. The performances of offline imitation learning and offline reinforcement learning have been demonstrated in the extensive experiments on D4RL benchmarks.
summary: This paper focuses on practical offline reinforcement learning tasks with only expert observations, avoiding the requirements for expert actions and reward labels. Specifically, this paper proposed AILOT (Aligned Imitation Learning via Optimal Transport), which defines the intrinsic rewards using optimal transport distance between the intention representations of the expert’s and agent’s trajectories. Through dense reward relabeling, AILOT outperforms state-of-the-art offline imitation learning methods and improves other offline reinforcement learning methods on the D4RL benchmarks.
weaknesses: 1.	AILOT is built on top of OTR, following the idea of performing reward relabeling through optimal transport. The most interesting part of AILOT is to perform Optimal Transport alignment in the space of intention instead of the original state space. However, the intention learning method is an existing work called ICVF, which limits the novelty. 
2.	Optimal Transport introduces additional runtime overhead compared to the offline RL algorithms, with the benefits of reward labeling.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: 1.	**Clarification on Dataset Size and Training Configuration:** Could the authors provide specific details on the number of samples and epochs used during training? This information would help clarify the computational efficiency of the method, beyond the hardware and runtime specifics provided.
2.	**Comparison with Modern State-of-the-Art Methods:** Have the authors considered including comparisons with more recent state-of-the-art methods in imitation learning, such as O-DICE or other recent 2024 approaches? This would offer a more comprehensive view of AILOT’s performance relative to current advancements in the field.
3. **Comprehensive Sensitivity Analysis for Cost Function and Hyperparameters:** While the paper includes a limited ablation study with only two configurations (α=5, τ=0.5 and α=1, τ=1), a broader exploration of these hyperparameters would provide a clearer picture of AILOT’s robustness. Could the authors expand the sensitivity analysis with more variations in these parameters or offer additional insights into how these choices affect the model’s performance across different tasks?
soundness: 3
strengths: 1.	**Novel approach within the scope of existing methods:** While optimal transport has been used in similar imitation learning research, AILOT applies it uniquely by focusing on “intent alignment” in a metric-aware latent space, which differentiates it from other OT-based approaches.
2.	**Demonstrated performance improvement:** AILOT outperforms several baseline models, including those using OT, in various benchmark tasks, indicating it successfully optimizes OT alignment in a manner that strengthens offline RL without explicit reward signals.
3.	**Robust integration with other RL algorithms:** The method is designed to enhance the performance of other offline RL algorithms, making it versatile for broader applications.
summary: This paper addresses how to effectively imitate expert behavior in **offline reinforcement learning** with **sparse rewards** and without action labels or ground truth rewards. Previous methods used optimal transport to measure similarity between agent and expert trajectories as a reward signal, but relied on raw state distances. This paper introduces **AILOT**, which uses "intent alignment" and "optimal transport" in an intent space to calculate intrinsic rewards, enabling the agent to learn expert behavior more effectively and improve performance in sparse reward tasks.
weaknesses: 1.	**Overlap with existing research:** The approach shares similarities with prior work that applies optimal transport to offline imitation learning, such as Optimal Transport for Offline Imitation Learning (arXiv:2303.13971) and Combining Expert Demonstrations in Imitation Learning via Optimal Transport (arXiv:2307.10810). These papers also use OT to create reward signals from expert trajectories, raising concerns about the novelty of AILOT’s contribution. Although AILOT introduces intent alignment as a distinct feature, further justification of how this approach advances beyond these prior works would strengthen the contribution.
2.	**Limited comparison with recent state-of-the-art methods:** The paper does not include comparisons with more recent imitation learning algorithms, such as O-DICE (ODICE: Revealing the Mystery of Distribution Correction Estimation via Orthogonal-gradient Update, arXiv:2402.00348), which have demonstrated strong performance in offline imitation learning tasks. Including such comparisons would provide a clearer understanding of AILOT’s relative performance and contributions.
3.	**Dependency on well-defined intents:** AILOT’s performance may be compromised if the expert’s behavior is ambiguous or multi-modal, making alignment challenging. This is especially relevant when handling multi-intent expert demonstrations, an area where existing OT-based methods may also encounter limitations.
4.	**Lack of clarity in training configuration:** While the paper provides an estimated runtime on an NVIDIA RTX 3090 GPU (10-25 minutes), it lacks specific details on training configurations, such as the number of samples or epochs used. Including this information would improve reproducibility and allow readers to better assess the computational efficiency of AILOT.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: - The paper argues that learning expert states alone (without actions nor rewards) is one strong contribution of the paper. Are previous works incapable of doing that (e.g., (Luo et al., 2023))?
soundness: 2
strengths: - Strong empirical performance
summary: This paper introduces Aligned Imitation Learning via Optimal Transport (AILOT), a method for offline reinforcement learning that uses optimal transport to align an agent's behavior with an expert's in a "intent" space. The intent space is learned with some previously suggested method. AILOT outperforms existing methods on benchmark tasks, especially in sparse-reward environments.
weaknesses: - This paper is basically a souped-up version of (Luo et al., 2023), where the optimal transport between state is replaced with the optimal transport between intentions. Since the intention space is also learned with the previously suggested method, the contribution of the paper is  1) idea of using intention instead of state itself, and 2) the design of the cost matrix in Eq. (10). However, the design of the cost matrix Eq. (10) is not well analyzed in the paper, neither theoretically nor empirically. In my opinion, the idea of using intention instead of state is straightforward, and the paper for ICLR should contain more messages than it.

- SInce the usage of intention on OTR is the key contribution of the paper, I would expect the paper to analyze on what aspects do we need to improve from state space OTR methods. However, the paper relies on a single intention learning method and do not discuss on why the used intention learning method improves.

Meta-review: This paper proposes a method for the offline RL setting where rewards are difficult to specify but one (or multiple) expert trajectories demonstrating the behavior may be found. The proposed method computes rewards by compares the optimal transport distance (computed via ICVF) between a new trajectory and the expert trajectory. The proposed method outperforms state-of-the-art offline imitation learning methods and improves other offline reinforcement learning methods on the D4RL benchmarks, especially in sparse-reward environments.

Reviewers appreciated the strong empirical performance, the strong motivation (esp. Fig 1) and the discussion of prior work. They also noted appreciated the thorough experimental details.

One repeated concern was the relationship with Luo et al (OTR), and urged the authors to revise the paper to include a more detailed discussion of the relationship with this prior work to avoid misleading readers into misattributing some of the ideas discussed in this paper. Given the similarity with Luo et al, the reviewers were looking for more thorough experiments (e.g., when and why does using the learned metric help). Reviewers also requested an additional baseline (O-DICE) and more details about the training setup.

Overall, I feel like this has the makings of a strong paper, but the paper needs additional revisions/analysis to address the concerns about similarity with Luo et al.

Decision: Reject