Optimal Transport for Offline Imitation Learning
Authors: Yicheng Luo
zhengyao jiang
Samuel Cohen
Edward Grefenstette
Marc Peter Deisenroth
Venue: ICLR 2023 notable top 25% (2023)
Keywords: offline reinforcement learning
optimal transport
imitation learning

Abstract: With the advent of large datasets, offline reinforcement learning is a promising framework for learning good decision-making policies without the need to interact with the real environment.
However, offline RL requires the dataset to be reward-annotated, which presents practical challenges when reward engineering is difficult or when obtaining reward annotations is labor-intensive.
In this paper, we introduce Optimal Transport Relabeling (OTR), an imitation learning algorithm that can automatically relabel offline data of mixed and unknown quality with rewards from a few good demonstrations. OTR's key idea is to use optimal transport to compute an optimal alignment between an unlabeled trajectory in the dataset and an expert demonstration to obtain a similarity measure that can be interpreted as a reward, which can then be used by an offline RL algorithm to learn the policy. OTR is easy to implement and computationally efficient. On D4RL benchmarks, we demonstrate that OTR with a single demonstration can consistently match the performance of offline RL with ground-truth rewards.


Published as a conference paper at ICLR 2023
OPTIMAL TRANSPORT FOR OFFLINE IMITATION LEARNING
Yicheng Luo University College London
Zhengyao Jiang University College London
Samuel Cohen
University College London
Edward Grefenstette University College London Cohere
Mare Peter Deisenroth University College London
ABSTRACT
With the advent of large datasets, offline reinforcement learning (RL.) is promis- ing framework for learning good decision-making policies without the need to interact with the real environment. However, offline RL requires the dataset be reward-ane reward-annotated, which presents practical challenges when reward engineering is difficult or when obtaining reward annotations is labor-intensive. In this paper, we introduce Optimal Transport Reward labeling (OTR), an algorithm that assigns rewards to offline trajectories, with a few high-quality demonstrations. OTR's key idea is to use optimal transport to compute an optimal alignment between an un- labeled trajectory in the dataset and an expert demonstration to obtain a similarity measure that can be interpreted RL algorithm to learn the policy. OTR is easy to implement and computationally efficient. On D4RL benchmarks, show that OTR with single demonstration can consistently match the performance of offline RL with ground-truth rewards"]
reward, which then be used by offline
INTRODUCTION
Offline Reinforcement Learning (RL) has made significant progress recently, enabling learning policies from logged experience without any interaction with the environment. Offline RL is relevant when online data collection can be expensive or slow, e.g., robotics, financial trading and autonomous driving. A key feature of offline RL is that it can learn an improved policy that goes beyond the behavior policy that generated the data. However, offline RL requires the existence of a reward function for labeling the logged experience, making direct applications offline RL methods impractical for applications where rewards Even if it is possible to label the trajectories with human preferences, such a procedure to generate reward signals can be expensive. Therefore, enabling offline RL to leverage unlabeled data is an open question of significant practical value.
hard to specify with hand-crafted rules.
Besides labeling every single trajectory, an alternative way to inform the agent about human prefer- ence is to provide expert demonstrations. For many applications, providing expert demonstrations is more natural for practitioners compared to specifying reward function. In robotics, providing ex- pert demonstrations is fairly common, and in the absence of natural reward functions, 'learning from demonstration' has been used for decades to find good policies for robotic systems; see, e.g., (Atke- son & Schaal, 1997: Abbeel & Ng. 20044, [Calinon et al., 2007, |Englert et al.] 2013). One such framework for learning policies from demonstrations is imitation learning (IL). IL aims at learning policies that imitate the behavior of expert demonstrations without explicit reward function.
There are two popular approaches to IL: Behavior Cloning (BC) (Pomerleau) 1988) and Inverse Reinforcement Learning (IRL) (Ng & Russell 2000). BC aims to recover the demonstrator's be- havior directly by setting up an offline supervised learning problem. If demonstrations are of high quality and actions of the demonstrations
recorded, BC can work very well as demonstrated by Pomerleau (1988), but generalization to new situations typically does not work well. IRL learns an intermediate reward function that aims to capture the demonstrator's intent. Current algorithms
Code available at att.:ggg//gghhbb..cc////aaaagggg
ma1 mat___ansspprptreeeeeeaa

---

Published as a conference paper at ICLR 2023
Expert Demonstration
Offline Dataset
$
sy
OTR
$1
22
$2
42
Offline RL
F2
FT
Figure 1: Illustration of Optimal Transport Reward Labeling (OTR). Given expert demonstrations (left) and an offline dataset without reward labels (center), OTR adds reward labels to the offline dataset by means of optimal transport (orange, center). The labeled dataset can then be used by an offline RI. algorithm (right) to learn policies.
have demonstrated very strong empirical performance, requiring only few expert demonstrations to obtain good performance (Ho & Ermon 2016;; Dadashi et al.] (2022) While these IRL methods do not require many demonstrations, they typically focus on the online RL setting and require a large number of environment interactions to learn imitating policy, i.e., these methods are suitable for offline learning.
In this paper, we introduce Optimal Transport Reward labeling (OTR), an algorithm that uses op- timal transport theory to automatically assign reward labels to unlabeled trajectories in an offline dataset, given one or more expert demonstrations. This reward-annotated dataset can then be used by offline RL algorithms to find good policies that imitate demonstrated behavior. Specifically, OTR uses optimal transport to find optimal alignments between unlabeled trajectories in the dataset and expert demonstrations. The similarity measure between state in unlabeled trajectory and that of an expert trajectory is then treated as a reward label. These rewards can be used by any offline RL algo- rithm for learning policies from small number of expert demonstrations and a large offline dataset. Figure fillustrates how OTR uses expert demonstrations to add reward labels to an offline dataset, which then be used by an offline RL. algorithm to find good policy that imitates demonstrated behavior. Empirical evaluations on the D4RL (Fu et al.) (021) datasets demonstrate that OTR recov- ers the performance of offline RL methods with ground truth rewards with only a single demonstra- tion. Compared to previous reward learning and imitation learning approaches, our approach also achieves consistently better performance across a wide range of offline datasets and tasks.
OFFLINE REINFORCEMENT LEARNING AND IMITATION LEARNING
Offline Reinforcement Learning In offline/batch reinforcement learning, we are interested in learning policies directly from fixed offline datasets (Lange et al. 2012 Levine al. 2020). i.e., the
agent is not permitted any additional interaction with the environment. Offline RL research typically assumes access to an offline dataset of observed transitions D {s),, at, rf, 1) This setting is particularly attractive for applications where there is previous logged experience available but online data collection is expensive (e.g., robotics, healthcare). Recently, the field of offline RL has made significant progress, and many offline RL algorithms have been proposed to learn improved policies from diverse and sub-optimal offline demonstrations (Levine al., 2020 Fujimoto Gu
2021 Kumar et al.) 2020: Kostrikov et al. 2022c Wang et al. 2020).
Offline RL research typically assumes that the offline dataset is anwart-annotated. That is, each transition (si, at, rf. s(+ in the dataset is labeled with reward rř. One common approach to offline RL is to leverage the reward signals in the offline dataset and learn a policy with an actor-critic algorithm purely from offline transitions (Levine et al., 2020). However, in practice having per- transition reward annotations for the offline dataset may be difficult due to the challenges in design- ing good reward function. Zolna et al. (2020) propose ORIL which learns reward function based on positive- vositive-unlabeled unlabeled (PU) learning (Elkan & Noto (008) that can be used to add reward labels to offline datasets, allowing for unlabeled datasets be used by offline RL. algorithms.
2

---

Published as a conference paper at ICLR 2023
Imitation Learning Imitation Learning (IL) aims at learning policies that mimic the behavior of an expert demonstrator. Unlike the typical RL setting, a reward function is not required. Instead, IL methods assume access an expert that can provide demonstrations of desired behavior.
There are two commonly used approaches to imitation learning. Behavior Cloning (BC) (Pomerleau, 1988) learns a policy by learning mapping from states to actions in the trajectory demonstrated by an expert. By formulating policy learning as a supervised learning problem, BC can suffer from classical problems of supervised regression: generalizing the policy unseen states may not be successful due to overfitting; multiple expert trajectories that follow different paths (bifurcation) will cause challenges. Inverse Reinforcement Learning (IRL) considers first learning reward functionn based on expert demonstrations. The learned reward function can then be used to train policy using an RI. algorithm. Previous work shows that IRL.based methods can learn good policies with a small number of expert demonstrations. However, popular recent IRL methods, such as GAIL (Ho & Ermon 2016). require a potentially large number of online samples during training. resulting poor sample efficiency. Moreover, algorithms, such as GAIL, follow a training paradigm that is similar to Generative Adversarial Networks (GANs) [Goodfellow et al. 2014), by formulating the problem of IRL minimax optimization problem to learn a discriminator that implicitly minimizes an f- divergence. It was found, however, that Adversarial imitation learning (AIL) methods such as GAIL can be very difficult to optimize, requiring careful tuning of hyperparameters (Orsini et al.] (021). Another popular approach is to perform distribution matching between the policy and the expert demonstrations (Englert al. 2013 Kostrikov et al. 20226).
More recently, IRL methods based on Optimal Transport (OT) have demonstrated success as an alternative method for IRL compared AIL approaches. Unlike AIL approaches, OT methods minimize the Wasserstein distance between the expert's and the agent's state-action distributions. Building on this formulation, Xiao et al. (2019) propose to minimize the Wasserstein distance via its dual formulation, which may lead to potential optimization issues. More recently. Dadashi et al.) (022)) introduce PWIL, which instead minimizes the Wasserstein distance via its primal formula- tion, avoiding the potential optimization issues in the dual formulation. Along this line of work, [o- hen et al./ (202I) propose a series of improvements to the primal Wasserstein formulation used in PWIL and demonstrate strong empirical results in terms of both sample efficiency and asymptotic performance. Still, these approaches require a large number of online samples to learn good poli- cies. While progress has been made to improve the sample efficiency of these approaches (Kostrikov) jet al. (022a)), imitation learning without any online interaction remains an active research area. In this paper, we provide constructive algorithm address the issue of offiine imitation learning by using optimal transport to annotate offline datasets with suitable rewards.
3 OFFLINE IMITATION LEARNING WITH OPTIMAL TRANSPORT
We consider learning in an episodic, finite-horizon Markov Decision Process (MDP) (S, A. p, r. Po. I') where S is the state space, A is the action space, p is the transition function, is the reward function, 7 is the discount factor, Po is the initial state distribution and T is the episode horizon. A policy is function from state distribution over actions. The goal of RL is to find policies that maximize episodic return. Running a policy π in the MDP generates a state-action episode/trajectory (s1,a), a,, 82, .2,
We consider the problem of imitation learning purely from offline datasets. Unlike in the stan- dard RL setting, no explicit reward function available. Let 7 (s1, 82, d2y....88p 8T) de- note an episode of interaction with the MDP using policy that selects actions a, at time steps t L....,T 1. Instead of reward function, we have access a dataset of expert demonstra-
tions D= = (Te (m) 1 Jn=1 generated by an expert policy म. and a large dataset of unlabeled trajectories D (T(m)) - generated by an arbitrary behavior policy т3- We are interested in learning an offline policy combining information from the expert demonstrations and unlabeled experience, without any interaction with the environment. We will address this problem by using optimal trans- port, which will provide a way to efficiently annotate large offline RL datasets with rewards.
3

---

Published as a conference paper at ICLR 2023
Px P.
C
For
C
Figure 2: Illustration of the computations performed by OTR. In this example, consider an MDP with a vvo-dimensional state-space ([S) = 2). We have two empirical state distributions from an expert Pe with samples { 2 (o) and policy Pa with samples {s7 3 (+) as denoted by points in the leftmost figure. OTR assigns rewards Tor (blue) to each sample in the policy's empirical state dis- tribution as follows: (i) Compute the pairwise cost matrix C' (purple) between expert trajectories and trajectories generated by behavior policy: (ii) Solve for the optimal coupling matrix μ (green) be- tween P.7 and Pe: (iii) Compute the reward for of Toi (s]) = -G7 14 Consider for example state of px: the row C, in the cost matrix corresponds to the costs between Sf and {sf: J iy repre-
sents the optimal coupling between S. and the expert samples. The optimal coupling moves most of the probability mass sj and small fraction of the mass si (green lines in the leftmost figure).
3.1 REWARD LABELING VIA WASSERSTEIN DISTANCE
Optimal Transport (OT) (Cuturi 2013 Peyré & Cuturi 2020 is principled approach for comparing probability measures. The (squared) Wasserstein distance between two discrete measures μıx + + 6n and By = + Ei= 5y. is
T Σ>
min M =====
c(x,, ye Pe,e
where M Epr E RTXT", μ1l 11, A. 1) is the set of coupling matrices, is cost function, and 8g refers to the Dirac measure for x. The optimal coupling μr" provides an alignment between the samples in Pzz and My. Unlike other divergence measures (e.g.. KL-divergence), the Wasserstein distance is metric and incorporates the geometry of the space.
Let p. + Et= deg and P. + ΣT= 6,77 denote the empirical state distribution of an expert policy ग. and behavior policy respectively. Then the (squared) Wasserstein distance
(p..
min ΣΣ μEM
11=11
can be used to measure the distance between expert policy and behavior policy. Let μ² denote the optimal coupling for the optimization problem above, then eq. (2) provides reward signal
7"
Fa (s7)
Σ cl87
(3)
P=1
which can be used for learning policy in an imitation learning setting.
3.2
IMITATION LEARNING USING REWARD LABELS FROM OPTIMAL TRANSPORT
We leverage the reward function from eq. 3 to annotate the unlabeled dataset with reward sig- nals. Computing the optimal alignment between the expert demonstration with trajectories in the unlabeled dataset allows us to assign a reward for each step in the unlabeled trajectory. Figure 2 illustrates the computation performed by OTR annotate an unlabeled dataset with rewards using demonstrations from expert.
The pseudo-code for our approach given algorithm I OTR takes the unlabeled dataset D and expert demonstration De as input. For each unlabeled trajectory ((m) Du, OTR solves the optimal
4

---

Published as a conference paper at ICLR 2023
Algorithm 1: Pseudo-code for Optimal Transport Reward labeling (OTR)
Input: unlabeled dataset D expert dataset D
Output: labeled dataset Dilael Dlabel
foreach 7(m) in D do C(m), 1 [m(m) t SolveOI(De, r(m))
// Label each episode in the unlabeled dataset
// Compute the optimal alignment with eq. ?
for /. to T do
(sr (m)) t - Ev= O Hipe
a(m)
5
Compute the per-step rewards with eq.
end
Dllbel
Dlabel U (m) .01
(m) oT
(m)
// Append labeled episode
end
return Dinhel
transport problem for each, obtaining the cost matrix ((m) and optimal aligmment f f((m)
(line
OTR then computes the per-step reward label following eq. (3) (line
5) The rewand-annotated
trajectories are then combined, forming a reward-labeled dataset Dlabel
Solving eq. (2) requires solving the OT problem to obtain the optimal coupling matrix /²" This amounts to solving linear program (LP), which may be prohibitively expensive with standard LP solver. In practice, we solve the entropy-regularizee OT problem with Sinkhorn's algorithm (Cuturi) 2013). We leverage the Sinkhorn solver in OTT-JAX (Cuturi et al. (022) for this computation. Once OTR has annotated the unlabeled offline dataset with intrinsic rewards, we can use an offline RL algorithm for learning policy. Since we are working in the pure offline setting, it is important use an offline RI. algorithm that can minimize the distribution shifts typically encountered in the offline setting.
Unlike prior works that compute rewards using online samples (Dadashi et al., 2022 Cohen et al. 2021) ), we compute the rewards entirely offline, prior running offline RL training, avoiding the need to modify any part of the downstream offline RL pipeline. Therefore, our approach can be combined with any offline RL algorithms, providing dense reward annotations that required by the downstream algorithms. Figure I illlusratees the entire pipeline of using OTR for relabeling and running offline RI. algorithm using the reward annotated datasets.
Compared to previous work that aims at solving offline imitation learning with a single algorithm, OTR focuses generating high-quality reward annotations for downstream offline RL algorithms. In addition, our approach enjoys several advantages:
Our approach does not require training separate reward models or discriminators, which may incur higher runtime overhead. By not having train separate parametric model, avoid byper-parameter tuning on the discriminator network architectures.
Unlike other approaches, such GAIL DemoDICE, our approach does not require solving a minimax optimization problem, which can suffer from training instability (Orsini et al. 2021)
Our approach is agnostic to the offline RL methods for learning the policy since OTR computes reward signals independently of the offline RL algorithm.
4 EXPERIMENTS
In this section, we evaluate OTR on D4RL. Locomotion, Antmaze, and Adroit benchmark tasks (see also Figure 3 with the goal of answering the following research questions:
1. Can OTR recover the performance of offline RL algorithms that has access to well-engineered reward function (i.e., ground-truth rewards provided by the environment)?
2. Can OTR handle unlabeled datasets with behaviors of unknown and mixed quality?
3. How does OTR perform with varying number of expert demonstrations?
5

---

Published as a conference paper at ICLR 2023
Figure Benchmark tasks: D4RL Locomotion, Antmaze, and Adroit.
4. How does OTR compare with previous work on offline imitation learning in terms of perfor- mance and runtime complexity?
We demonstrate that OTR can be effectively combined with an offline RL algorithm to learn policies from a large dataset of unlabeled episodes and small number high-quality demonstrations. Since OTR is only method for reward learning, it can be combined with any offline RL algorithm that requires peward-annotate data for offline learning. In this paper, we combine OTR with the Implicit Q-Learning (IQL.) algorithm (Kostrikov et al. 2022c).
4.1 SETUP
We evaluate the performance OTR+IQL on the D4RL locomotion benchmark (Fu al. 2021]. We start by evaluating OTR on three environments (Hal ECheetah-2,, Walker-v2. Hopper-v2) from the OpenAl Gym MuJoCo locomotion tasks. For each environment, we use the medium-v2. medi wwwuumppllle rem-eepay--2 and expert demonstrations and the unlabeled dataset. For the expert demonstrations, we choose the best episodes from the D4RL dataset based on the episodic return P discard the original reward information the dataset. We then run OTR to label the dataset based on the optimal coupling between the unlabeled episodes and the chosen expert demonstrations. Afterward, we proceed with running the offline RL algorithm.
datasets to construct the
To obtain the unlabeled dataset,
For OTR, follow the recommendation by Cohen al. (2021) and the cosine distance the cost function. When there are more than one episode of expert demonstrations, we compute the optimal transport with respect to each episode independently and the rewards from the expert trajectory that gives the best episodic return. Similar to (Dadashi et al.) 2022 Cohen et al.) 2021). we squash the rewards computed by 5ith exponential function s s advantage of ensuring that the rewards consumed by the offline RL algorithm have an appropriate range since many offline RL algorithms can be sensitive to the scale of the reward values. We refer to appendix A.1 for additional experimental details and hyperparameters.
exp(Brr This has the
Implementation We implement OTR in JAX (Bradbury et al. 2018). For computing the optimal coupling. OTT-JAX (Cuturi et al. 2022). library for optimal transport that includes scalable and efficient implementation of the Sinkhorn algorithm that can leverage accelerators, such GPU TPU, for speeding up computations. Our IQL implementation is adapted from (Kostrikov et al. 2022c and we set all hyperparameters to the ones recommended in the original paper. All of our algorithms and baselines are implemented in Acme (Hoffman al. 2022))
II practice, the expert demonstrations can he provided separately; we only select expert demonstration in this way for case of evaluation.
https: /i thub. com/ kttm///kottrkk/////ciiiiiii

---

Published as a conference paper at ICLR 2023
Our implementation is computationally efficient requiring only about minute label dataset with 1 million transitions (or 1000 episodes of length 1000 For larger-scale problems, OTR can be scaled up further by processing the episodes in the dataset in parallel. Our implementation of OTR and re-implementation of baselines are computationally efficient. Even so, the training time for IQL. is about 20 minutes, so that OTR adds relatively small amount of overhead for reward annotation to an existing offline RL algorithm.
Baselines
We compare OTR+IQL with the following baselines:
IQL (oracle): this the original Implicit Q-learning (Kostrikov et al. 2022c algorithm using the ground-trute rewards provided by the D4RL datasets.
DemoDICE: an offline imitation learning algorithm proposed by Kim moDICE was found to perform better than previous imitation learning algorithms (e.g.. Val- ueDICE (Kostrikov
2022). De-
al., 2022bb) that can also (in principle) work in the pure offline setting- We ran the original implementation under the same experimental setting as used for the other algorithms in the paper.
ORIL: reward function learning algorithm from Zolna et al., 2020). ORIL learns reward function by contrasting the expert demonstrations with the unlabeled episodes. The learned reward function is then used for labeling the unlabeled dataset. We implement ORIL in com- parable setting to the other baselines.
we keep the rewards from the expert demonstrations and relabel all rewards from the unla- beled datasets with the minimum rewards from the environment. This was found to perform well in Yu et al. 2022). This also resembles online imitation learning methods such as SQIL (Reddy ct al. 2022 or offline RL algorithms, such as COG (Singh et al., (022)).
UDS:
Since UDS and ORII. are also agnostic about the underlying offline RL algorithm used, we combined these algorithms with IQL so that we can focus on comparing the performance difference due to different approaches used in generating reward labels.
For all algorithms, we repeat experiments with 10 random seeds and report the mean and standard deviation of the normalized performance of the last 10 episodes of evaluation. We compare all algorithms by using either K = 1 or K = 10 expert demonstrations, Obtaining the results on the locomotion datasets took approximately 500 GPU hours.
4.2 RESULTS
MuJoCo Locomotion Table DaleeJcomme compares the performance between OTR+IQL with the other base- lines. Overall, OTR+IQL performs best compared with the other baselines terms aggregate score over all of the datasets we used, recovering the performance of IQL with ground-truth rewards provided by the dataset. While we found that other baselines can perform well on some datasets, the performance is not consistent across the entire dataset and can deteriorate significantly some datasets. In contrast, OTR+IQL is the only method that consistently performs well for all datasets of different compositions.
"An efficient implementation of OTR is facilitated by JAX, which includes useful functionality that allows to casily paralletize computations. Concretely, JIT-compile the computation of rewards for episode and further leverage the vmap function compute the optimal coupling between an unlabeled episode with all the expert episodes in parallel. Efficiently parallelizing the computation of optimal coupling requires that all the episades share the same length. This is necessary both for parallelizing the computation across multiple expert demonstrations as well as for avoiding recompilation by XLA due to changes in the shape of the input arrays. To achieve high throughput for datasets with varying episodic length, maximum episode length allowed by the environment (which 1000 for the OpenAl Gym environments) but set the weights of the observations to zero. Padding the episodes this way does not change the solution to the optimal coupling problem. Note that padding means that a IM transition dataset may create more than 1000 episodes of experience, in this case, the runtime for our OTR implementation may be higher effectively due to having to process a larger number of padded episodes.
pad all observations to the
RRnntime measured on hal fcheet -medi um-v2 with an NVIDIA GeForce RTX 3080 GPU.
htttps:/ /i thub. com www.:/cww.oo..eeeeeeeeeeeee.ee haweoonnnyyyooo////ee tat on-dice
7

---

Published as a conference paper at ICLR 2023
IQL (oracle)
DemoDICE
IQL+ORIL
IQL+UDS
OTR+IQL (ours)
halfchoctal-modiun-222
42.6
42.5
47.4 0.2
42.5
49.0 0.2
42.4
43.3 0.2
hopper meeppur-edium-22
52.9
56.9
66.2 a
5..1
4,00
2..5
78.7 2 5.5
walker2d medium v2
75.3
75.0
78.3
73
61.9
68.9
79.4 1.4
halfchhetall- medium replay-v2
44.2
hopper mediumm replay +22
18.1
75.9
94.7 8.6
39.0 15.4
82.4
49.3
84.8 2.6
walkkrrdd medium weediii---pply-yee replay-v2
52.2
salfehoctah-meding ceftim-med--ppet-eeee
55.2
92.9
86.7 16.7
85.
87,5
63.0
89.6 3.0
hopper- meddim--eepeesttttt
52.5
110.9
91.3 14.3
92.3 14.2
29.7 22.2
53.9 2.5
93.2 20.6
walker2d medium expert
107.5
109.0
109.6 1.0
106.9
110.6
107.5
109.3 0.8
locomotion-v2-40
666.2
692.4
494.9
685.5
runtime
10m
10m
20m
100m
30mh
20mn
22m
The nuntime is measured with the original PyTorch implementation.
Table I: D4RL performance comparison between IQL with ground-truth rewards and OTR+IQL with a single expert demonstration (K 1). We report mean 4 standard deviation per task and ag- gregate performance and highlight near-optimal performance in bold and extreme negative outliers in red. OTR+IQL is the only algorithm that performs consistently well across all domains.
Runtime Despite applying optimal transport, we found that with GPU-aceelerated Sinkhorn solver (Cuturi et al. 2022), combined with our efficient implementation in JAX, OTR achieves faster runtime compared to algorithms that learn additional neural networks as discriminators (DemoDICE (Kim et al.) 2022)) or reward models (ORIL. Zolna et al.) 2020)). For methods that learn a neural network reward functiinn an overhead of least 10 minutes is incurred, whereas OTR only incurs approximately 2 minutes of overheads when compared with the same amount of expert demonstrations.
Effect of the number of demonstrations We investi- gate if the performance of the baselines can be improved by increasing the number of expert demonstrations used. Table 2 compares the aggregate performance on the lo- comotion datasets between OTR and the baselines when we increase the number of demonstrations from K = 10. DemoDICE's performance improves little with the additional amount of expert demonstrates. While ORIL and UDS enjoy a relatively larger improvement, they are still unable match the performance of IQL (or- acle) or OTR in terms of aggregate performance despite using the same offline RL backbone. OTR's performance close to IQL (oracle) even when = 1 and matches the performance of IQL (oracle) with K 10.
locomotion-v2-total
10
DemoDICE
IQI+
588.5
618.3
IQL+UDS OTR-IQL
494.9
575.8
694.3
IQL (oracle)
692.4
Table 2: Aggregate performances of dif- ferent reward labeling algorithms with different numbers of expert demonstra- tions. OTR is the only algorithm that leads to an offline RL performance close to using ground-trui -truth rewards.
Antmaze & Adroit We additionally evaluate OTR+IQL on the antmaze-vO and adroit-ia --vo datasets. Table 3 shows that OTR+IQL again recovers the performance of IQL with ground-truth rewards. This suggests that OTR+IQL can learn from datasets with diverse behavior and human demonstrations even without ground-truth reward annotation; additional results appendix A.2
Qualitative comparison of the reward predictions Figure 4a provides qualitative comparison of the reward predicted by OTR, ORIL, and UDS. UDS annotates all transitions the unlabeled dataset with the minimum reward from the environment. Thus, the episodes with non-zero rewards are expert demonstrations. This means that UDS is unable distinguish between episodes in the unlabeled datasets. Compared to reward learning algorithms, such as ORIL, OTR's 's reward predic- tion more strongly correlates with the ground-truth rewards from the environment, which a good precondition for successful policy learning by downstream offline RL algorithm. We also evaluate OTR's s reward prediction on more diverse datasets, such as those in antmaze. Figure 4b shows the expert demonstrations received the best OTR reward labels in the unlabeled dataset (right). OTR correctly assigns more rewards trajectories that are closer to the expert demonstrations.
used in antmaze med um-play-vo (left) and the trajectories that
8

---

Published as a conference paper at ICLR 2023
(oracle)
(oracle)
OIR+IQL
antmaze-lage-daae large play-vD
47.5 9.5
45.5
doce-cloned-v0
1.60
0.01 0.01
39.6 5.8
6,9
dove human-vu
4.30
5.92
antmaee medium diverse
70.0 10.9
70.4 4.8
hummer cloned vO
2.10
0.88 = 0.30
(elmaaze-edium-ppee play----ppp animaze umaze diverse-vt vt)
hummer- human--O
1.40
62.2 13.8
68.9 13.6
pen cloned y pen-human-uu
37.30
46.87 +: 20.85
aummze-uuaze-v0
relocate-clond--10
-0.20
-0.24 0.03
animaze v0 total
378.0
384.0
relocat-huuann-v0
0.10
0.11 2 0.10
adroit-vO-total
118.1
122.16
Table 3: Performance of OTR+IQL antmaze and adroit with single expert demonstration. Similar to the results on locomotion, OTR+IQL recovers the performance offline RL with ground-truth rewards. We report mean + standard deviation per task and aggregate performance.
OTR
ORIL
UDS
Labeled Return
5000
1000
3000
2500
500
1500
10002000
3000
1000
20003000
30000003000
Environment Return
(a) Ground-truth return vs labeled return.
(h) Top trajectories selected by OTR.
Figure 4: Visualization of rewards predicted by OTR and baselines.
Qualitative differences be-
tween the rewards predicted by OTR, ORIL and UDS on hoppe r r- um-v2. (b) Visualization of top trajectories selected OTR antmaze-medl um-plmm--oo p-pppla v-v Left: Expert demonstra-
tions. Right: ranking of trajectories according to rewards per step computed by OTR. Trajectories with lighter colors have higher rewards per step.
Summary Our empirical evaluation demonstrates that OTR+IQL can recover the performance of offline RL with ground-truth rewards given only single episode of expert demonstration. Compared to previous work, it achieves better performance even with larger number of demonstrations, Our qualitative comparison shows that OTR assigns better reward estimates that correlate strongly with a hand-crafted reward function.
DISCUSSION
OTR can be effective in augmenting unlabeled datasets with rewards for by downstream offline RI. algorittms (see section 4). Compared to prior imitation learning and reward learning approaches, OTR enjoys better and more robust performance paper. Furthermore, our approach is easy to implement. Most of the complexity of our approach arises from solving the optimal transport problem. However, there libraries for solving OT efficiently in different frameworks or programming languages
the benchmarks
have considered in this
Unlike prior works, such as ValueDICE DemoDICE, we split offline imitation learning into the two distinct phases of reward modeling via OTR and subsequent offline RL. This improves mod- ularity, allows for improvements in the two phases to be made independently. and adds flexibility. However, this modularity means that a practitioner needs to make a separate choice for the down- stream offline RL algorithm.
We believe that OTR is most useful in situations where providing ground-truth rewards is difficult but providing good demonstrations and unlabeled data is feasible. However, will not be feasible in cases where collecting expert demonstrations is difficult.
For example, see Python Optimal Transport (POT) Fllaaay et al. 2021 which supports PyTorch, JAX, TensorFlow.

---

Published as a conference paper at ICLR 2023
The Wasserstein distance formulation used can be extended to perform cross-domain imitation learning by using the Gromov- Wasserstein distance to align expert demonstrations and offline tra- jectories from different spaces (Fickinger et al. 2022).
CONCLUSION
We introduced Optimal Transport Reward labeling (OTR), a method for adding reward labels to offline dataset, given more expert demonstrations. OTR computes Wasserstein distances between expert demonstrations and trajectories in a dataset without reward labels, which then are turned into reward signal. The reward-annotated offline dataset can then be used by an(y) of- fline RL algorithm to determine good policies. OTR adds minimal overhead to existing offline RL algorithms while providing the same performance as learning with pre-specified reward function.
ACKNOWLEDGMENTS
Yicheng Luo is supported by the UCL Overseas Research Scholarship (UCLORS). Samuel Cohen is supported by the UKRI AI Centre for Doctoral Training in Foundational Artificial Intelligence (EP/8021566/1).
REFERENCES
Pieter Abbeel and Andrew Y. Ng. Apprenticeship Learning via Inverse Reinforcement Learning. In International Conference on Machine Learning, 2004.
Christopher G. Atkeson and Stefan Schaal. Robot Learning From Demonstration, In International
Conference Machine Learning, 1997.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- and Qiao Zhang. JAX: Composable Transformations of Python+NumPy Programs, 2018. URL http:
7/github. com/ google/jax
Sylvain Calinon, Florent Guenter, and Aude Billard. On Learning, Representing, and Generalizing
Task a Humanoid Robot. IEEE Transactions on Systems, Man, and Cybernetics, Part (Cybernetics), 2007.
Samuel Cohen, Brandon Amos, Mare Peter Deisenroth, Mikael Henaff, Eugene Vinitsky, and Denis Yarats. Imitation Learning from Pixel Observations for Continuous Control. 2021. URL hthhsts::
7/ //wpeneeview.nee Ltte://oooummmddd-Lbb........
Marco Cuturi. Sinkhorn Distances: Lightspeed Computation of Optimal Transport. Advances in Neural Information Processing Systems, 2013.
Marco Cuturi, Laetitia Meng- Pang-Papaxsnttos, Yingtao Tian, Charlotte Bunne, Geoff Davis, and Olivier Teboul. Optimal Transport Tools (OTT); A JAX Toolbox for all things Wasserstein. arXiv:2201. 12324, 2022.
Robert Dadashi, Leonard Hussenot, Matthieu Geist, and Olivier Pietquin. Primal Wasserstein Imi- tation Learning. In International Conference Learning Representutions, 2022.
Charles Elkan and Keith Noto. Learning Classifiers from Only Positive and Unlabeled Data. In International Conference Knowledge Discovery and Data Mining, 2008
Peter Englert, Alexandros Paraschos, Jan Peters, and Mare Peter Deisenroth. Model-based Imitation Learning by Probabilistic Trajectory Matching. IEEE International Conference on Robotics and Automation, 2013.
Arnaud Fickinger, Samuel Cohen, Stuart Russell, and Brandon Amos. Mooss-Domain Imitation Learning via Optimal Transport. In International Conference Learning Representations, 2022.
10

---

Published as a conference paper at ICLR 2023
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanis- las Chambon, Lactitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalic T. H. Gayraud, Hicham Janati, Alain Rakotomamoniy. levgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. POT: Python Optimal Transport. Journal of Machine Learning Research, 2021.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4RL: Datasets for
Deep Data-Driven Reinforcement Learning. arXiv:2004.07219, 2021.
Scott Fujimoto and Shixiang Gu. A Minimalist Approach to Offline Reinforcement Learning. In Advances in Neural Information Processing Systems, 2021.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative Adversarial Nets. In Advances in Neural Infor- mation Processing Systems, 2014.
Jonathan Ho and Stefano Ermon. Generative Adversarial Imitation Learning. In Advances Neural Information Processing Systems, 2016.
Matthew W. Hoffman, Bobak Shahriari, John Aslanides, Gabriel Barth-Maron, Nikola Momehev, Danila Sinopalnikov, Piotr Stańczyk, Sabela Ramos, Anton Raichuk, Damien Vincent, Léonard Hussenot, Robert Dadashi, Gabriel Dulac-Amold, Manu Orsini, Alexis Jacq, Johan Ferret, Nino Vieillard, Seyed Kamyar Seyed Ghasemipour, Sertan Girgin, Olivier Pietquin, Feryal Behba- hani, Tamara Norman, Abbas Abdolmalcki, Albin Cassirer, Fan Yang, Kate Baumli. Sarah Hen- derson, Abe Friesen, Ruba Haroun, Alex Novikov, Sergio Gómez Colmenarejo, Serkan Cabi, Caglar Gulcehre, Tom Le Paine, Srivatsan Srinivasan, Andrew Cowie, Ziyu Wang, Bilal Piot, and Nando de Freitas. Acme: A Research Framework for Distributed Reinforcement Learning. arXiv:2006.00979, 2022.
Geon-Hyeong Kim, Seokin Seo, Jongmin Lee, Wonseok Jeon, HyeongJoo Hwang, Hongseok Yang, and Kee-Eung Kim. DemoDICE: Offline Imitation Learning with Supplementary Imperfect Demonstrations. In International Conference on Learning Representations, 2022.
Ilya Kostrikov, Kumar Krishna Agrawal, Debidatta Dwibedi, Sergey Levine, and Jonathan Tomp- son. Discriminato--Actrr Cacioior---oo---iittiee Addressing Sample Inefficiency and Reward Bias in Adversarial Imitation Learning. In International Conference Learning Representations, 2022a.
Ilya Kostrikov, Ofir Nachum, and Jonathan Tompson. Imitation Learning via Off-Policy Distribution Matching. In International Conference on Learning Representations, 2022b.
Nya Kostrikov, Ashvin Nair, and Sergey Levine. Offline Reinforcement Learning with Implicit Q-Learning In International Conference on Learning Representations, 2022c.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative Q-Learning for Of-
fline Reinforcement Learning. In Advances Neural Information Processing Systems, 2020.
Sascha Lange, Thomas Gabel, and Martin Riedmiller. Batch Reinforcement Learning. In Reinforce- ment Learning. Springer Berlin Heidelberg, 2012.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline Reinforcement Learning: Tutorial, Review, and Perspectives Open Problems. arXiv:2005.01643, 2020.
Andrew Y. Ng and Stuart Russell. Algorithms for Inverse Reinforcement Learning. In Interna- tional Conference on Machine Learning, 2000.
Manu Orsini, Anton Raichuk, Leonard Hussenot, Damien Vincent, Robert Dadashi, Sertan Girgin, Matthieu Geist, Olivier Bachem, Olivier Pietquin, and Marcin Andrychowicz. What Matters for Adversarial Imitation Learning? In Advances in Neural Information Processing Systems, 2021.
Gabriel Peyré and Marco Cuturi. Computational Optimal Transport. arXiv:1803. 103.00803.00567, 2020.
Dean A. Pomerleau. ALVINN: An Autonomous Land Vehicle in a Neural Network. In Advances in Neural Information Processing Systems, 1988.
11

---

Published as a conference paper at ICLR 2023
Siddharth Reddy, Anca D. Dragan, and Sergey Levine. SQIL: Imitation Learning via Reinforcement Learning with Sparse Rewards. International Conference on Learning Representations. 2022.
Avi Singh, Albert Yu, Jonathan Yang, Jesse Zhang, Aviral Kumar, and Sergey Levine. Chaining Be- haviors from Data with Model-Free Reinforcement Learning. In Conference on Robot Learning, 2021.
Ziyu Wang, Alexander Novikov, Konrad Zolna, Josh S Merel, Jost Tobias Springenberg, Scott E Reed, Bobak Shahriari, Noah Siegel, Caglar Gulcehre, Nicolas Heess, and Nando de Freitas. Critic Regularized Regression. In Advances in Neural Information Processing Systems, 2020.
Huang Xiao, Michael Herman, Joerg Wagner, Sebastian Ziesche, Jalal Etesami, and Thai Hong Linh. Wasserstein Adversarial Imitation Learning. arXiv:1906.08 13, 2019.
Tianhe Yu, Aviral Kumar, Yevgen Chebotar, Karol Hausman, Chelsea Finn, and Sergey Levine. How to Leverage Unlabeled Data in Offline Reinforcement Learning. In International Conference on Machine Learning, 2022.
Konrad Zolna, Alexander Novikov, Ksenia Konyushkova, Caglar Gulcehre, Ziyu Wang, Yusuf Aytar, Misha Denil, Nando de Freitas, and Scott Reed. Offline Learning from Demonstrations and Unlabeled Experience, LtXiv:2011.13885 2020.
12

---

Published as a conference paper at ICLR 2023
A APPENDIX
A.1 HYPERPARAMETERS
Table 4 lists the hyperparameters used by OTR and IQL on the locomotion datasets. For Antmaze and Adroit, unless otherwise specified by table or table the hyperparameters follows from those used in the locomotion datasets.
The IQL hyperparameters are kept the same those used in (Kostrikov al. 2022c). Note that IQL rescales the rewards in the dataset so that the same set of hyperparameterss can be used for datasets of different qualities. Since OTR computes rewards offline, we also apply reward scaling as in IQL. For the locomotion datasets, the rewards are rescaled by max min_retum.. while for antmaze we subtract 2 to the rewards computed by OTR. The reward processing in antmaze is different from the one used by the original IQL paper (which subtracts 1) since the rewards computed by OTR have different range.
1000
The squashing function used by OTR based on the one used (Dadashi et al. 2022). The antmaze squashing differs slightly from the one used in locomotion and adroit due use of an carlier configuration. In practice, this should have minimal effect on the performance.
Hyperparameter
Value
Discount
0.99
Hidden layers
(256,
Network Architectures
Dropout
none
Network initialization
orthogonal
Optimizer
Adam
IQL
Policy learning rate
3e
cosine decay to
Critic learning rate
3e
Value learning rate
3e-1
Target network update rate
5
Temperature
3.0
Expectile
0.7
Episode length T Cost function Squashing function
1000
OTR
cosine
s(r) = 5.0 exp(5.0 T r//lA)
Table 4: OTR hyperparameters for D4RL Locomotion.
Hyperparameter
Value
Temperature
10.0
IQL
Expectile
0.9
OTR
Squashing function s(r) = 5.0 ex00eexp[ (Spp(T-)) r)
Table 5: OTR hyperparameters for D4RL Antmaze.
Hyperparameter
Value
Network Architectures
Dropout
0.1
Temperature
0.5
IQL
Expectile
0.7
Table 6: OTR hyperparameters for D4RL Adroit.
13

---

Published as a conference paper at ICLR 2023
A.2 ADDITIONAL EXPERIMENTAL RESULTS ON ADROIT AND ANTMAZE
We evaluate OTR on additional datasets from the antmaze and adroit domains with varying number of expert demonstrations. The results are presented in table aable table OTR consistently recovers the performance of IQL with ground-truth rewards on these datasets, largely independent of the number of expert demonstrations provided.
K OTR+IQL
10
Dataset
IQL (oracle)
OTR+IQL
door-cloned-v0)
1.60
0.01+00001
0.01+00.01
door-human-v0
4.30
5,92+2272
4.1542.08
hammer-cloned-v0
2.10
0.88
1.31
hammer-human-v0
1.40
1.79:11.43
1.36:10.22
pen-cloned-v0
37.30
46.87+220.85
42.68++24.98
pen-human-vo
71.50
66.82+21.18
69.41+21550
relocate-cloned-vd
-0.20
-0.2440..33
-0.24+00.3
relocate-human-v0
0.10
0.11+0.11
0.10+00.7
adroit-vO-total
118.1
122.16
118.78
Table 7: OTR+IQL Results on Adroit. The standard deviations for IQL (oracle) are not available
from Kostrikov al. 2022c).
Dataset
IQL (oracle)
OTR+IQL(K = 1)
OTR+IQL (K = 10)
mutmaze-lrre--divese--y
47.5+
4..5+6.2
50.7+669
antmaze-large-play
39.6+
45.3++6.9
51.2+
untmaze-mediium-divee-----
70.0=10.9
70.4
70.5
autmaze-meddimmpllay-oo
71.2+7.3
70.5
72.716.2
untiaaze-uaaze-divrre---0
63.2
18.9
64.4
antmaze-umaze-vO
87.5+2.6
83.4+3.3
88.743.5
antmaze-vO-total
378.0
384.0
398.2
Table 8: OTR+IQL Results on Antmaze.
14

---

Published as a conference paper at ICLR 2023
A.3 COMBINING OTR WITH DIFFERENT OFFLINE RL ALGORITTMS
In the main experiments, we evaluated OTR by pairing it with the IQL algorithm. In this sec- tion, investigate if OTR can recover the performance of different offline RL algorithm (TD3- BC) (Fujimoto & Gu 202T) using ground-truth rewards. We observe that (i) the performance from OTR+TD3-BC mostly matches those using the ground-truth rewards; (ii) the performance fairly robust to the choice of the number of expert trajectories (K = 1 and K = 10 many expert demon- strations provide comparable performance). However, There are more variances on some datasets Nevertheless, the differences are smaller com- pared to the baselines and OTR+TD3-BC still performs better than the baselines presented in sec- tion 4 in terms of the aggregate performance.
(e.g., hal fcheet -tt-eteeeet-e---eeee-e---------e um-expert--2).
Dataset
TD3-BC (oracle)
OTR+TD3-BC (K=1)
OTR+TD3-BC (K=10)
malfcheetah-medium-eee expert-v
93.5+2200
74.8=20.1
71.6=23.1
nalfcheetab-mecium-rrppoe
44.4
39.4+1.3
38.9
halfcheetab-mmedium-
48.0+10.7
42.6+11.0
42.7+1.1
hopper-mediium-experryy
100.2++20.0
103.2+13.9
98.9+19.7
hopper-meddium-rplla--.
64.8+25.5
74.9-28.8
80.2+23.1
hopper--medium-v2
60.7+12.5
66.4=10.3
69.8=13.9
walleer2d-medium-exxprrr
109.5+0.55
109.0=0.66
108.8=0.8
wulkerzddmmedium-repleee
87.4+8.4
69.7=16.4
27.4
mallee22d-mediumv22
83.745.3
76.9+5.4
78.0+2.6
locomotion-v2-total
692.3
656.9
656.4
Table 9: OTR+TD3-BC Results on MuJoCo.
15

---

Published as a conference paper at ICLR 2023
A.4 IMPORTANCE OF USING THE OPTIMAL TRANSPORT PLAN
In the main experiments, we compute the rewards based on the optimal coupling computed by the Sinkhorn solver. The optimal transport plan is sparse and transports most of the probability masses to only a few expert samples. In this section, we investigate what happens if we use a suboptimal transport plan where each sample from the policy's trajectory transported equally each sample in the expert's 's trajectory. In this case, the reward function essentially boils down to computing the average costs with respect to all of the states the expert's trajectory.
Dataset
OTR+IQL
OTR+IQL (Uniform Plan)
malichectah-medium--y
43.3::022
43.5#:..3
hopper-medium--v2
78.7+55.5
80.5+2..3
walker2dd-medium--22
1..4
77.6
halfchhetaa--mcdiuum-rrpppaaa
41.3=0.66
41.6=10.8
hopper-meciuum-rellayy
84.8+2.6
69.8+10.1
walleer22-medduun--eplaaay
66.0+6.7
62.2+14.4
haffchhetal-mmddiu--eeeerrrv
89.6+3.0
90.6+2299
hopper-mcdiuun-exper--22
93.2420..6
89.2
walker2d-mmcdium-exprrre
109.3=0.8
106.015.9
Table 10: OTR with Uniform Transport Plan
Table 10 compares the performance of OTR+IQL using the optimal transport plan and uniform trans- port plan. We find that for many datasets, using the suboptimal uniform transport plan is sufficient for reaching good performance. This indicates that using a reward function based on the similarity of states from the policy and the expert can be a simple and effective method for reward labeling. However, note that the uniform transport plan can still underperform compared to using the opti- mal transport plan (e.g., hopper peeemmeiie medium-- replay-v2). This shows that the optimal transport formulation enables better and more consistent performance.
A.5 COMPARISON TO PWIL
In this section, we investigate if the online imitation learning algorithm PWIL (Dadashi al.) 2022) can be used in the offline setting with change from using an online RL algorithm to an offline RT. algorithm. We ran PWIL with IQI. similar to what we did for OTR the main paper. We use the PWIL implementation from Acme (Hoffman et al. 2022
Note that although OTR is similar to PWIL in using the Wasserstein distance to construct RL reward signals, OTR differs from PWIL in the choices of OT solver, the cost function well as the approach used for aggregating results multiple expert demonstrations. Also note that for experiments paper consider learning only from expert state instead of state-ache action pairs. This is both more general and challenging setting. It was found in (Dadashi et al. 2022) that PWIL sometimes perform badly without expert actions. We OTR and PWIL. using only expert observations (denoted as OTR-state and PWIL-state) and OTR and PWIL using state-action pairs (denoted as OTR-action and PWIL-action). The results are illustrated in table We found that we are unable to get good results when running PWIL using only expert state sequences. This is possibly due to difference choices of OT solver and cost functions. PWIL can perform well when combined with IQL to learn in the offline setting although sometimes performance is significantly worse compared to IQL oracle or OTR (e.g., www.hhopppeseeee wwwperr-meiin--expppeeee un-exper--xx)
In addition, Dadashi et al. (2022) found that PWIL's performance may deteriorate when learning from demonstrations consisting of only expert observations (i.e., no actions are present in the expert demonstrations).
16

---

Published as a conference paper at ICLR 2023
method
OTR-state
stae-state
OTR-action
PWIL-actiin
halfcheetah-meddiun-z
43.1+0.3
1.6+1.2
43.4++0.3
47.5++022
hopper-medium-v2
80.0+:5.2
2.1+1.3
75.4+4.6
70.4+44.2
walker2d-medium--2
79.2
0.9
79.7#1.2
81.9
halichetal-mmddii--rpla---
41.6:10.3
-2.3
41.9410.3
44.6111.1
84.4+1.8
1.4+1.2
85.3+1.1
89.7#4.9
walker2d-medium-rrpph
71,8+3.8
-0.1
69.1#4.6
72.2+10.6
halicheccah-nedddu-eexyrtzrz2Z
87.9+13.4
-0.3+1.5
88.3+5.1
88.6+4.33
hopper-meddumm-exppet-2
96.6++21.5
1.5+0.6
86.6+22.99
32.9+25.0
walker22-medum-eexpertt2
109.6+00.5
1.0+199
109.2+0..5
110.2+0..
Table 11: Comparison between OTR and PWIL with IQL as offline RL backbone.
A.6
PYPER-PARAMETER SENSITIVITY
K
10
method
OTR (a 3 = 5)
OTR (a $ 1)
hallchecah--mediu---epprrrz-
87.9+3.4
86.9+4.0
hallcheccah-meddii---eey---
41.6
1..
halfcheetall-mediuu--22
43.1
42.7
hopper- seetee--mdiiuxppeeert
96.6#21.5
82.6+99.
hopper-medium medium-replay--v
84.4
71.2=15.2
hopper- mepper--edium-22
80.0
75.7
medium-experr--v
109. 0.5
106.3+8.2
walker2d-med mediam---miiuii-m-------
71,8
63.2+t5.7
walker2d-medium--2
79.2
77.4
Table 12: Effect of and / in the squashing function.
For the main results, the hyper-parameters for the squashing function (a and (3) was chosen to be consistent with those used in (Dadashi et al. 2022). In this section compare the differences in the choices these hyper- ppprm-paraetter by running OTR with a = = 1. This reduces to simply applying exponential transformation the optimal transport costs. The results illustrated in [2 We find that OTR still performs well, demonstrating that is not sensitive to the choices of these hyper-parameters.
17


Reviews:
[{"review_id": "yR0MBDgWvw", "reviewer": "Reviewer_eZXx", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "Authors propose an offline imitation learning algorithm that annotates rewards by using Optimal Transport (w.r.t. expert trajectories) solver and uses those reward with the offline RL to mimic expert behavior without environment interactions. The algorithm is evaluated on D4RL benchmarks (which is a widely adopted benchmark for offline RL and IL algorithms) and is shown to outperform existing offline IL algorithms. ", "summary_of_the_review": "Although the algorithm empirically outperforms baselines, the idea doesn't seem novel, and the contribution seems incremental. ", "strength_and_weaknesses": "### Strength\n\n- The proposed algorithm outperforms baselines. \n\n### Weaknesses\n\n- The contribution seems incremental. Using OT to acquire IL rewards was proposed already by PWIL (by Dadashi et al) although the idea was for online imitation learning. The submission extends PWIL's idea to offline IL by using offline RL instead. \n- The complexity of using OT is not well-described. When the number of expert demonstrations is large, the reward estimation may not be scalable due to solver's complexity. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "ii_TSs_EfH", "reviewer": "Reviewer_heNC", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper proposes a new algorithm for reward-free offline reinforcement learning when a few expert demonstrations are available. The core idea is to assign rewards to the unlabelled trajectories based on their distance to the expert demonstrations such that the closer the trajectories they are to the expert demonstrations, the higher rewards they get assigned. The proposed approach (OTR) uses an optimal transport formulation where each trajectory is treated as a discrete distribution with uniform density over the states in the trajectory and the distance between two trajectories is defined as the Wasserstein distance between these two discrete distributions. The paper demonstrates that OTR can be combined with existing offline RL methods with little implementation and computation overhead to tackle reward-free offline RL with demonstrations and often matches the performance of offline RL algorithms when rewards are given on a range of MuJoCo locomotion tasks and Adroit manipulation tasks.", "summary_of_the_review": "The paper is generally well-written with strong empirical results. Although I have some concerns over the reproducibility of the work and empirical comparisons with respect to the baselines (see above), they do not occur to me to be major issues as long as they are addressed appropriately. Overall, I would recommend acceptance of the paper (borderline).", "strength_and_weaknesses": "*Strength*\n- The proposed approach is conceptually simple and elegant with strong empirical results on a range of tasks.\n\n*Weaknesses*\n- OTR introduces additional hyper-parameters such as $\\alpha, \\beta$ in the squashing function (for reward labeling). The paper currently has no sensitivity analysis on how these parameters could affect the performance.\n- There is uncertainty over how the cost between states is defined in the optimal transport formulation (see the section below), which could impact the reproducibility of the paper. \n- I found the claim that the proposed approach (OTR) can consistently match the performance of offline RL with ground-truth rewards is not well-supported by the current set of experiments (see the section below). ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "O8g6Gcgquv", "reviewer": "Reviewer_C2wE", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "In this paper, the authors focused on the offline reinforcement learning setting where a small set of expert demonstrations and a larger dataset without reward are available. \nThe proposed method labels the reward for the larger dataset with optimal transport theories. Specifically, the authors considered the distribution over the state space and use the optimal coupling to decide the closeness of states. \nThe unlabeled trajectory will be assigned a higher reward if its state distribution is closer to the state distribution from experts' data. The authors demonstrated the effectiveness of this method in several environments\n", "summary_of_the_review": "The paper is overall well-written. However, I would suggest the authors elaborate more on the benefits of labeling rewards with optimal transport coupling. ", "strength_and_weaknesses": "Pros:\n- Overall, the paper is well-written and quite easy to follow.\n- The idea of combining optimal transport with reinforcement learning is pretty novel. Given that there are only limited previous studies.\n- The experimental results are encouraging. \n\nCons:\n\n- The novelty is limited, considering some previous studies. I have included more detailed questions in the next section.\n\n- In Figure 2, the cost matrix and the optimal coupling do not seem to correspond. In addition, the authors may want to annotate $s^e_3$ and $s^e_4 $in the leftmost figure.\n\n- The authors used cosine similarity as the ground metric. Is there a principle to choose the ground cost since this is very important for the OT objective? The authors may want to refer to bisimulation metrics.\n\n- One of the pitfalls of offline RL approach is the distribution shift problem. However, the authors did not discuss how the proposed method can handle this issue. When the expert trajectories state distribution is far away from the unlabeled ones, I am wondering whether the proposed method will still annotate the closest states by high rewards. \n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "8iY1LU6eTsc", "reviewer": "Reviewer_bBjs", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper presents a method for performing imitation learning by first computing rewards using a method called Optimal Transport Reward labeling (OTR) that computes the Wasserstein distance between an expert demonstration and an unlabeled one. The rewards are then used in combination with existing offline RL algorithms, for example, implicit Q-learning (IQL). The process is made computationally efficient through the use of GPU-accelerated optimal transport solvers. The method is tested on locomotion tasks from the D4RL benchmark, and compared to DemoDICE, and ORIL and UDS for computing rewards. Through additional experiments, the authors show that OTR is able to effectively use up to 10 demonstrations to boost performance.", "summary_of_the_review": "While there are related works with a similar reward relabeling idea such as [1], this paper tackles a challenging and (as far as I am aware) novel situation of learning from a very small number of demonstrations. It demonstrates promising empirical results of being able to match the performance of offline RL algorithms when given ground truth reward functions. The paper is well written and clearly presented. Therefore I vote weak accept.\n\n[1]: Imitation Learning from Pixel Observations for Continuous Control (Cohen et. al)", "strength_and_weaknesses": "Strengths:\n- The method is conceptually simple, and the quantitative results are encouraging – OTR+IQL is able to recover the performance of IQL using oracle rewards for many of the D4RL tasks.\n- It’s quite surprising and impressive that a single good expert demonstration can be used for the expert trajectory in this setup.\n- The empirical evaluation compares to sensible baselines on the D4RL benchmark.\n- The paper is particularly clearly written and easy to understand! \n\nWeaknesses:\n- The introduction states that for behavior cloning, “generalization to new situations typically does not work well”, but I’m not sure I found the authors’ intuition in the paper for which aspect of OTR will improve generalization to new situations? It seems that labeling rewards based on a single demonstration will always make it more challenging to develop robust policies, as a single demonstration cannot provide signal about how to detect or recover from bad states. I’m curious if the authors could elaborate or if this would continue to work well in more challenging environments.\n- I think it would be helpful to put the performance of the method in context by comparing the performance to behavior cloning on the top K% of dataset trajectories (like 1% or 10% that previous works like IQL report).\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-25%