Title: Robot Policy Learning with Temporal Optimal Transport Reward
Venue: NeurIPS 2024
Accepted: yes
Keywords: Reinforcement Learning

Imitation Learning

Optimal Transport
Abstract: Reward specification is one of the most tricky problems in Reinforcement Learning, which usually requires tedious hand engineering in practice. One promising approach to tackle this challenge is to adopt existing expert video demonstrations for policy learning. Some recent work investigates how to learn robot policies from only a single/few expert video demonstrations. For example, reward labeling via Optimal Transport (OT) has been shown to be an effective strategy to generate a proxy reward by measuring the alignment between the robot trajectory and the expert demonstrations. However, previous work mostly overlooks that the OT reward is invariant to temporal order information, which could bring extra noise to the reward signal. To address this issue, in this paper, we introduce the Temporal Optimal Transport (TemporalOT) reward to incorporate temporal order information for learning a more accurate OT-based proxy reward. Extensive experiments on the Meta-world benchmark tasks validate the efficacy of the proposed method. Our code is available at: https://github.com/fuyw/TemporalOT.

Robot Policy Learning with Temporal Optimal Transport Reward

Abstract

Reward specification is one of the most tricky problems in Reinforcement Learning, which usually requires tedious hand engineering in practice. One promising approach to tackle this challenge is to adopt existing expert video demonstrations for policy learning. Some recent work investigates how to learn robot policies from only a single/few expert video demonstrations. For example, reward labeling via Optimal Transport (OT) has been shown to be an effective strategy to generate a proxy reward by measuring the alignment between the robot trajectory and the expert demonstrations. However, previous work mostly overlooks that the OT reward is invariant to temporal order information, which could bring extra noise to the reward signal. To address this issue, in this paper, we introduce the Temporal Optimal Transport (TemporalOT) reward to incorporate temporal order information for learning a more accurate OT-based proxy reward. Extensive experiments on the Meta-world benchmark tasks validate the efficacy of the proposed method. Our code is available at: https://github.com/fuyw/TemporalOT .

1 Introduction

Reinforcement Learning (RL) [51] has achieved great success across a wide array of applications [40]. However, it typically requires a large number of interactions with the environment [26, 35], which limits its practical application in the robotic control [9, 49]. A large body of work has been developed to address this issue from different aspects [52], i.e. , using curiosity-based intrinsic reward to encourage exploration [2, 47], leveraging better representation pretrained on large scale robotics datasets [33, 38], incorporating external knowledge from the Vision-Language Models (VLMs) [13, 32], and imitating the behaviors from pre-collected expert demonstrations [50, 62].

Reward specification plays a central role in RL [6]. Since the goal of the RL agent is to maximize the expected cumulative rewards, the reward signal directly influences the learned behaviors [8]. One major challenge in applying RL to real-world problems is how to design the reward functions [10]. A well-designed reward function can guide the agent towards desirable behaviors more efficiently, while a poorly designed one could lead to sub-optimal behaviors [25]. However, designing a good reward function is a nontrivial task [11], which requires related expert domain knowledge and (or) time-consuming hand reward engineering [48]. The lack of a good reward function is one of the main bottlenecks for the low sample-efficiency issue in RL [58].

Imitation Learning (IL) has been proven to be an effective technique to learn control policies without the oracle task reward [24]. Given an expert demonstration dataset, IL formulates the policy learning as a supervised learning paradigm [27]. The IL objective aims to learn a policy that mimics the expert behaviors via minimizing a distance measure of the learned policy and an approximated expert policy [23, 28]. Depending on if the RL agent can learn from further online interactions, IL can be crudely classified as offline IL and online IL [29]. In offline IL, the RL agent purely learns from a

fixed dataset of collected expert experiences. In online IL, the RL agent usually learns a proxy reward function to relabel the collected online trajectories with respect to the expert demonstrations [53].

One notable weakness of IL is that it generally requires a diverse and high-quality demonstration dataset to achieve desired performances [5]. Recently, some researchers found that Optimal Transport (OT) [18, 44] based proxy reward function enables us to learn effective robot policies with only a few expert demonstrations [30, 31]. In this paper, we follow this line of research in applying OT-based proxy rewards to online IL without using any task reward information. In particular, we first revisit the efficacy of OT-based proxy reward in RL and then discuss some challenges of the existing methods due to the overlook of temporal order information. To mitigate this issue, we introduced the Temporal O ptimal T ransport (TemporalOT) reward, which incorporates temporal order information to the OT-based proxy reward via using context embeddings and a mask mechanism.

The primary contributions of this work can be summarized as follows:

we pointed out a weakness of existing OT-based proxy reward methods for imitation learning due to the overlook of temporal order information;

we designed a simple yet effective algorithm to incorporate temporal order information into OT-based proxy reward via using context embeddings and a mask mechanism;

experiments show that the proposed method outperforms other SOTA algorithms.

2 Background

2.1 Reinforcement Learning

In this work, we consider the standard Markov Decision Process (MDP) [46] setting M = ( S , A , R, P, ρ 0 , γ ) , where S and A are state and action spaces, R : S × A → R is a reward function, P : S × A → ∆( S ) is the state-transition probability function, ρ 0 : S → R + is the initial state distribution and γ ∈ [0 , 1) is a discount factor. Our goal is to learn a policy π ( a | s ) : S → ∆( A ) that maximizes the expected cumulative discounted rewards E π [ ∑ ∞ t =0 γ t r ( s t , a t )] where s 0 ∼ ρ 0 , s t +1 ∼ P ( ·| s t , a t ) and a t ∼ π ( ·| s t ) . In partially observable MDP (POMDP) [19], we can only receive an observation o i ∈ O , i.e. , image observation, of the current state s i .

To solve this optimization problem, value-based RL methods typically learn a state-action value function Q π ( s, a ) := E π [ ∑ ∞ t =0 γ t r t | s 0 = s, a 0 = a ] , which is defined as the expected return under policy π . For convenience, we adopt the vector notation Q ∈ R S×A , and define the one-step Bellman operator T π : R S×A → R S×A such that T π Q ( s, a ) := r ( s, a ) + γ E s ′ ∼ P,a ′ ∼ π [ Q ( s ′ , a ′ )] . The Q -function Q π is the fixed point of T π such that Q π = T π Q π [51]. Similarly, we define the optimality Bellman operator as follows T Q ( s, a ) := r ( s, a ) + γ E s ′ ∼ P [max a ′ Q ( s ′ , a ′ )] and the optimal Q -value function Q ∗ is the fixed point of T Q ∗ = Q ∗ . In deep RL, we use neural networks Q θ ( s, a ) to approximate the Q -functions by minimizing the empirical Bellman error:

where we sample transitions ( s, a, r, s ′ ) from a replay buffer and Q π ˆ θ ( s, a ) is the target network.

2.2 Inverse Reinforcement Learning

Inverse Reinforcement Learning (IRL) aims to infer the underlying reward function from expert demonstrations [39], which further facilitates an RL agent to learn the policy. One key assumption of IRL is that the observed behaviors are optimal such that the observed trajectories maximize the cumulative rewards [17]. Due to the ability to avoid the manual reward specification, IRL holds the promise for practical real-world RL applications. Denote M as an MDP and π E as an expert policy, the IRL problem is to find an optimal reward function R ∗ such that:

where Π is the feasible policy set. That is, the expert policy π E will achieve the maximum expected cumulative discounted reward than any other policy.

Demo

面

Cost

Transport

OT

D

Matrix

Plan

Reward

五

b

0

2.3 Optimal Transport

Optimal Transport is an optimization problem which aims to find an optimal mapping that transforms one probability distribution into another with the least cost. OT has a wide application in various domains such as economics [14], physics [15], and machine learning [1]. Consider two probability distributions p ∈ R n , q ∈ R m and a joint distribution µ ( p, q ) on product space X×Y , the Wasserstein distance [54] between p and q is defined as:

where c ( x, y ) is the cost function for moving mass form x to y . In the RL scenario, p and q are usually in the state space S or the observation space O [21]. For example, given an expert trajectory τ E = ( o E 1 , · · · , o E T ) and an agent trajectory τ = ( o 1 , · · · , o T ) where o i is the image observation at step i , the Wasserstein distance between τ E and τ is defined in the following discrete form:

where µ ∈ R T × T is called the transport plan, and we denote the optimal transport plan as µ ∗ .

3 Method

In this section, we first revisit the application of OT-based proxy reward in RL. In particular, we point out the influence of temporal order information, which has been overlooked in most prior work. Next, we introduce the main idea and formulation of the proposed method based on these observations.

3.1 A Recap of OT Reward in RL

3.1.1 OT reward helps to rank states and actions

In RL, we usually adopt the Wasserstein distance to measure the similarity of two trajectories [31], as illustrated in Figure 1. Given an agent trajectory τ = ( o 1 , · · · , o T ) and an expert trajectory

Traj a

0.02

00

Traj b

0

02

-0.04

-0.06

OT Reward

-0.08

μ*(3,)

μ*(3,)

0.10

D

05

a

0.12

r2(02,a2) = rT(0,tE)

r2(02, a2) = r0T(0,tE)

0.14

0

20

40

60

80

100

Step

τ E = ( o E 1 , · · · , o E T ) , we first compute the optimal transport plan µ ∗ in Eqn.(4) using some iterative optimization algorithms, i.e. , Sinkhorn algorithm [4]. Then, the OT-based proxy reward at the i -th step is defined as follows:

where c ( o i , o E j ) is a cost function that measures the similarity between o i and o E j . One popular choice in prior work is the cosine similarity based cost function c ( o i , o E j ) = 1 -⟨ f ( o i ) ,f ( o E j ) ⟩ ∥ f ( o i ) ∥∥ f ( o E j ) ∥ , where f ( o i ) is the latent representation of observation o i extracted by a visual encoder [3].

Similar to the curiosity-based exploration bonus [2], OT reward is used to distinguish the goodness of different states. We consider the toy example in Figure 1, two agents start from the same state with observation o 2 and take different actions a a 2 and a b 2 , respectively. Then the goodness of a a 2 and a b 2 at o 2 is measured by the OT reward computed w.r.t. the observation o a 3 and o b 3 at the next step. As long as the OT reward can rank r OT 3 ( o a 3 , τ E ) and r OT 3 ( o b 3 , τ E ) correctly, then the policy will be able to learn a better action at o 2 . From Figure 2 (right), we can observe that the OT reward for the better trajectory b is generally larger, which validates the previous explanations.

3.1.2 Two Key Observations

Notably, there are two key observations of OT reward that have been less discussed in prior work:

The OT reward is order invariant.

The OT reward at step i is influenced by the later steps so that two transitions with the same state-action pair could have different OT rewards.

Our first observation is that the standard OT-reward is order invariant. As shown in Eqn.(5), the order information is discarded and the frames from the demo trajectory are treated as bag-of-temporallycollapsed frames. In our view, collapsing the temporal axis drops arguably one of the most important characteristic features of temporal order information. More concretely, consider a demo trajectory of τ 1 = ( o 1 , o 1 , o 2 ) , meaning the agent first stays in the first state and then moves to the second state. Our goal is to imitate this behavior. However, if we discard the order information as in Eqn.(5), from the perspective of OT reward, there is no ability to differentiate between τ 1 and some other undesired trajectories, i.e. , τ 2 = ( o 1 , o 2 , o 1 ) which first moves to the second state and then moves back to the first state. Therefore, discarding the temporal order information in reward calculation

makes the reward on top of it under-constrained, thereby increasing the likelihood of convergence toward undesired solutions.

Our second observation is that the OT reward is non-stationary during the training. As the Eqn.(5) shows that OT reward r OT i depends on the optimal transport plan µ ∗ ( i, j ) , which depends on the entire trajectory. Therefore, the OT reward for each state is a function of the entire agent trajectory. As illustrated in Figure 1, even though trajectory τ a and trajectory τ b have the same first two transitions, their OT rewards have different values. This is very different from the standard RL setting, where reward is usually determined with a fixed state-action pair. Such a non-stationary OT reward could have a pitfall that makes the optimization of RL objective in Eqn.(1) to be less stable.

Inspired by these two observations, we found that the current OT-based RL methods usually overlook the temporal order information of the trajectories. In this work, we aim to investigate how to improve the current OT-based RL methods by incorporating temporal order information.

3.2 Temporal Optimal Transport Reward (TemporalOT)

In this subsection, we present the Temporal O ptimal T ransport (TemporalOT) reward. We first explain our motivations for the model design, and then introduce the details of the proposed method.

3.2.1 Motivation for the Model Design

The pipeline of a standard OT-based reward calculation usually consists of two stages:

(Stage-1) first define a transport cost function c ( · , · ) between two states;

(Stage-2) then solve an OT optimization problem in Eqn.(4) to approximate the optimal transport plan µ ∗ and compute the OT reward r OT in Eqn.(5) for each state in a trajectory τ .

After the OT-reward calculation step, the transition will be relabeled with the OT-reward for training an RL agent as in Eqn.(1). Our method aims to improve both stages of OT-reward calculation. Firstly, previous methods usually use a pair-wise cosine similarity based cost function in Stage-1, which sometimes could be inaccurate and noisy. Secondly, previous methods ignore the temporal order information in Stage-2 as discussed in Section 3.1.2. We will introduce two simple solutions to address these two points, respectively.

3.2.2 Context Embedding-based Cost Matrix for Improving Stage-1

To learn a more accurate transport cost function, we introduce a context embedding based cost matrix. Unlike previous methods that use a pair-wise cosine similarity as the transport cost, we adopt a group-wise cosine similarity that we define the transport cost between agent observation o i and expert observation o E j as following:

where k c is the parameter of the context length and f ( · ) is a fixed visual encoder. The goal of the context cost matrix ˆ C is to facilitate expert progress estimation by taking nearby information into consideration. For example, we use k c = 3 in Figure 3 and the transport cost between o 1 and o E 2 is ˆ c ( o 1 , o E 2 ) = 1 -[cos( f ( o 1 ) , f ( o E 2 )) + cos( f ( o 2 ) , f ( o E 3 )) + cos( f ( o 3 ) , f ( o E 4 ))] / 3 .

3.2.3 Temporal-masked Optimal Transport Objective for Improving Stage-2

The prevalent OT reward ignores the temporal order information and takes the information of every step in the trajectory into consideration, as shown in the Eqn.(5). As pointed out by some previous work, the OT reward is not always correct where noisy OT rewards could distract the agent from learning some key early behaviors [30]. To mitigate this issue, we introduce a concise solution by adding a temporal mask to the cost matrix.

For an agent trajectory τ = ( o 1 , · · · , o T ) and an expert trajectory τ E = ( o E 1 , · · · , o E T ) , we denote the context cost matrix as ˆ C ∈ R T × T and the transport plan as µ ∈ R T × T . The row sum and column

ContextEmbedding

ContextCostMatrix

TemporalMatrix

ContextCostMatrix

TransportPlan

OTReward

C

M

C

OT

(00,01,02)

(,,)

(01,02,03)

(,，)

(02,03,04)

(0,0，0)

→Avg(

(03,04,05)

(05,，)

(04,05,05)

(,,)

(05,05,05)

(,,)

sum of µ equal to the constraint s = [ 1 T , · · · , 1 T ] ∈ R T . We proposed to introduce a temporal mask M ∈ R T × T to the transport plan, where M ( i, j ) ∈ [0 , 1] . We can express the masked optimal transport objective in the following vector form [16]:

where ⟨· , ·⟩ F is the Frobenius norm and we add an entropy regularizer H ( · ) of the masked transport plan M ⊙ µ . We can solve Eqn.(7) by the Lagrangian:

where α and β are two Lagrangian multipliers. By using different temporal mask, we can control what kind of temporal order information we use in the OT reward. For example, M = 1 degrades to the original OT reward without temporal order information, and a lower triangle matrix corresponds to the causal mask in the Transformer decoder [55], which indicates that we only concern the past steps observations. In our method, we use a variant of the diagonal matrix:

where k m is a window size parameter that controls the scope we use in the masked OT rewards. A smaller mask window size k m refers to a closer match w.r.t. to the expert demonstration. We select a diagonal-like matrix because we follow previous learning from demonstration literature to assume that the agent has a similar movement speed as the expert [30]. Under this assumption, we adopt the distance between the time step indexes to represent temporal affinity information. Figure 3 illustrates the main ideas of the proposed TemporalOT method.

4 Experiments

In this section, we aim to answer the following questions: (1) How does the proposed TemporalOT method perform compared with other baselines? (2) Are the proposed context-embedding based cost function and temporal mask useful? (3) How do the key parameters influence the performances? (4) Is TemporalOT effective with both state-based and pixel-based observations?

4.1 Experimental Setup

We implement TemporalOT-RL in PyTorch [42] based on the official ADS implementation 1 . We use a pretrained ResNet50 [20] network as the fixed visual encoder to extract the image embedding for each pixel observation. Unlike the original ADS experiments which use a fixed goal in each task, we adopt a more challenging setting where the goal position changes for each episode. Moreover, we only provide two expert video demonstrations to the RL agent. For the experiment results, we evaluate the RL agent for 100 trajectories every 20000 steps. We report the mean and standard deviation of the evaluation success rate across 5 random seeds. We define one trajectory to be successful if the RL agent solves the task at the last step. More detailed information is available in the Appendix B.

Door-open-v2

Hand-insert-v2

Stick-push-v2

100

125

no-mask

80

no-context

40

100

60

TemporalOT

30

75

R

S

40

20

50

Suc

S

只

20

10

25

S

0

0

-20

0

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

T

0.4

0.6

0.8

1.0

Step (1e6)

Step (1e6)

Step (1e6)

4.2 Baselines

We compare the following baseline methods. (1) TaskReward: training a backbone RL agent from DrQ-v2 [56] with the oracle task reward r task i = δ success. The reward is 1 when the task is solved and otherwise the reward is 0. (2) BC: a naive behavior cloning agent which has the access of the expert action. (3) GAIfO: another IL baseline which learns a discriminator to provide proxy reward [53]. (4) OT: we use an online version of OTR [31] agent where we first rollout the RL agent to collect online trajectories and then relabel the reward with OTR for RL training. (5) ADS: a variant of OT baseline which adaptively adjusts the discount factor w.r.t. a progress tracker [30].

4.3 Results on the Meta-world Benchmark Tasks

We first validate the effectiveness of TemporalOT on nine Meta-world [57] tasks. Experiment results are shown in Table 1. We can observe that TemporalOT generally outperforms the other baselines without using the task rewards. Moreover, the TaskReward baseline only shows good performance on the Door-lock and Window-open tasks. The main reason is that the RL agent fails to collect the first successful trajectory. For example, the goal of the Basketball task as shown in Figure 1 is to pick up the basketball and move to a target position above the rim. With the oracle sparse task reward, the RL agent only receives a nonzero reward until it first successfully solves the task. Under such circumstances, it is particularly challenging to collect the first successful trajectory with only zero task rewards and random action explorations. On the other hand, we can observe that the two IL baselines, BC and GAIfO [53], perform much worse than other OT-reward based baselines. This is because we only provide two expert video demonstrations with a few hundred samples, where the IL-based methods suffer from an over-fitting issue. We further compare two OT agent baselines, where OT0.99 uses γ = 0 . 99 , and OT0.9 uses γ = 0 . 9 . We have a similar conclusion as in ADS that using a smaller discount factor is helpful to learn early behaviors in some tasks that strongly rely on the progress dependency, i.e. , Basketball . TemporalOT outperforms the recent SOTA baseline ADS in 8 out of 9 tasks, which proves the effectiveness of the proposed method.

4.4 Ablation Studies on Different Model Components

We further conduct ablation studies to validate the effectiveness of the proposed context cost matrix and temporal mask in TemporalOT. In Figure 4, no-mask refers to a variant of TemporalOT without mask and no-context refers to a variant of TemporalOT without the context cost matrix. We can

Door-open-v2

Door-open-v2

Push-v2

100

100

80

kc=2

80

Km=5

Km=10

25

Ne=2

N=5

60

kc=6

60

20

R

40

kc=8

40

km=40

15

S

20

20

10

Suc

S

0

0

5

-20

-20

0

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

Step (1e6)

Step (1e6)

Step (1e6)

Basketball-v3

Door-lock-v2

Door-open-v2

100

40

(%)

80

80

30

Rate

60

60

20

Success

40

40

20

10

20

ADS

0

TemporalOT

0

0

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

Step (1e6)

Step (1e6)

Step (1e6)

observe that removing any of the two components will lead to a degraded performance. Moreover, the more important component varies depending on the task. For example, the temporal mask is more important in the Door-open task, and the context cost matrix is more important in the Stick-push task.

4.5 Ablation Studies on Different Key Parameters

We then validate the efficacy of different key parameters, i.e. , the context length k c for the context embedding, window size k m for the temporal mask, and the demonstration number N E . From Figure 5, we can observe that a medium number of k c and k m performs the best and a larger N E improves the performances. A large context length k c does not perform well because it will distract the OT reward from the current step and introduce extra reward noise. A smaller mask window size k m makes the learning more difficult because it only receives information from nearby observations, and a larger k c will gradually degrade to the naive OT reward. Further, having more expert demonstrations is helpful in mitigating the potential over-fitting issue and improving the final performance.

4.6 Results with Pixel-based Observations

We also evaluate the proposed method with pixel-based observations, where we follow the same DrQv2 model setting as the ADS baseline. Figure 6 shows the results of the comparison of TemporalOT with ADS. We can observe similar conclusions as in Table 1 that our proposed TemporalOT method also outperforms the ADS baseline with pixel-based inputs, where TemporalOT usually converges faster than ADS and (or) achieves a higher final success rate. Moreover, we can observe that sometimes the pixel-based agent learns faster than its dense state-based counterpart, which indicates that the agent can extract more effective representations from the pixel inputs.

4.7 Visualization for Bad Cases

In this subsection, we visualize some bad cases of OT-based RL agents to provide readers more insights about when OT-based RL agents are less useful. Figure 7 plots a typical bad case for the OT/ADS/TemporalOT agents in the Hand-insert task. The top row is the expert trajectory, and the second row is the agent trajectory. The goal of the Hand-insert task is to move the brown block to a target position in the hole. We can observe that the RL agent mainly focuses on imitating the arm behaviors which ends in the target position, but it ignores the brown block. The main reason for the this bad case is that the color of brown box is very close to the table background which sometimes

make it difficult for the pretrained visual encoder to capture the subtle information. More bad case analyses are available in the Appendix A.2.

5 Related Work

Learning with a Few Demonstrations. There is a large body of work on leveraging demonstrations for policy learning, ranging from the basic behavior cloning [12, 45] to demonstration-aided RL [43]. There is also work on leveraging demonstration data for offline pre-training [60, 36, 61], to either warm-start the policy [36] or help with exploration [61, 37, 22]. However, the amount of demonstrations required for a high-quality pre-training is typically large. In this work, we focus on the setting where only a small number of demonstrations are provided [5], thus greatly relieving the burden of generating demonstrations. Optimal Transport based imitation is a recently emerged approach in this direction, which will be reviewed in the subsequent section.

Optimal Transport-based Reward for Imitation and RL. Optimal Transport (OT) has been shown to be effective for imitation learning [18, 31]. Optimal Transport Reward Labeling (OTR) [31] uses Sinkhorn distance [4] to compute a similarity metric for a trajectory w.r.t. an expert demonstration and uses this metric as rewards for offline RL datasets without rewards [31]. Automatic Discount Scheduling (ADS) [30] uses a similar OT-based approach for reward calculation. The core idea of ADS is to incorporate a scheduling of the discount factor for online RL to mitigate the potentially distracting OT reward from temporally distant states. Our work aligns with previous work in this category, and addresses some common issues that are shared by previous methods.

6 Limitations

Since our work is closely related to IL, our method shares some common limitations of IL. For example, the success of our method heavily depends on high-quality expert video demonstrations. If we are facing a new task without any available expert demonstrations, our method will be less useful. Moreover, if the given demonstrations are sub-optimal or biased, the learned policies will inherit these flaws as well. Moreover, the performance of the proposed method relies on the quality of the pretrained visual encoder. If the pretrained visual encoder fails to capture some key information in the pixel observation, then our method will also fail to take such key details into consideration. Another limitation of our work is that the computation cost of the proposed method is related to the number of the given expert video demonstrations. A larger number of expert demonstrations will increase the computation cost when we compute the optimal transport plan.

7 Conclusion

This paper studies the problem of learning effective robot policies with expert video demonstrations. We focus on a challenging setting where there are only two demonstrations available and the environment does not provide any task reward. Following the line of research of OT-based proxy reward, we first discuss some challenges of the existing methods due to the overlook of temporal information. Further, we introduced a new method named TemporalOT, which incorporates temporal information to existing baseline by using a context-embedding based cost matrix and a mask mechanism. Experiments on nine Meta-world benchmark tasks showcase the effectiveness of the proposed method. One interesting future direction is to extend the current method to a camera-view invariant agent, where we can learn policies w.r.t. expert video demonstrations from different camera views.

References

Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International Conference on Machine Learning , 2017.

Yuri Burda, Harrison Edwards, Amos Storkey, and Oleg Klimov. Exploration by random network distillation. In International Conference on Learning Representations , 2019.

Samuel Cohen, Brandon Amos, Marc Peter Deisenroth, Mikael Henaff, Eugene Vinitsky, and Denis Yarats. Imitation learning from pixel observations for continuous control. 2021.

Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in Neural Information Processing Systems , 2013.

Christopher R. Dance, Julien Perez, and Théo Cachet. Demonstration-conditioned reinforcement learning for few-shot imitation. In International Conference on Machine Learning , 2021.

Peter Dayan and Bernard W Balleine. Reward, motivation, and reinforcement learning. Neuron , 36(2):285-298, 2002.

Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In IEEE Conference on Computer Vision and Pattern Recognition , 2009.

Rati Devidze, Goran Radanovic, Parameswaran Kamalaruban, and Adish Singla. Explicable reward design for reinforcement learning agents. In Advances in Neural Information Processing Systems , 2021.

Pierluca D'Oro, Max Schwarzer, Evgenii Nikishin, Pierre-Luc Bacon, Marc G Bellemare, and Aaron Courville. Sample-efficient reinforcement learning by breaking the replay ratio barrier. In International Conference on Learning Representations , 2023.

Gabriel Dulac-Arnold, Nir Levine, Daniel J Mankowitz, Jerry Li, Cosmin Paduraru, Sven Gowal, and Todd Hester. Challenges of real-world reinforcement learning: definitions, benchmarks and analysis. Machine Learning , 110(9):2419-2468, 2021.

Jonas Eschmann. Reward function design in reinforcement learning. Reinforcement learning algorithms: Analysis and Applications , pages 25-33, 2021.

Pete Florence, Corey Lynch, Andy Zeng, Oscar A Ramirez, Ayzaan Wahid, Laura Downs, Adrian Wong, Johnny Lee, Igor Mordatch, and Jonathan Tompson. Implicit behavioral cloning. In Annual Conference on Robot Learning , 2022.

Yuwei Fu, Haichao Zhang, Di Wu, Wei Xu, and Benoit Boulet. FuRL: Visual-language models as fuzzy rewards for reinforcement learning. In International Conference on Machine Learning , 2024.

Alfred Galichon. Optimal transport methods in economics . Princeton University Press, 2018.

Wilfrid Gangbo, Wuchen Li, Stanley Osher, and Michael Puthawala. Unnormalized optimal transport. Journal of Computational Physics , 399:108940, 2019.

Xiang Gu, Yucheng Yang, Wei Zeng, Jian Sun, and Zongben Xu. Keypoint-guided optimal transport with applications in heterogeneous domain adaptation. In Advances in Neural Information Processing Systems , 2022.

Dylan Hadfield-Menell, Stuart J Russell, Pieter Abbeel, and Anca Dragan. Cooperative inverse reinforcement learning. In Advances in Neural Information Processing Systems , 2016.

Siddhant Haldar, Vaibhav Mathur, Denis Yarats, and Lerrel Pinto. Watch and match: Supercharging imitation with regularized optimal transport. In Annual Conference on Robot Learning , 2022.

Matthew Hausknecht and Peter Stone. Deep recurrent q-learning for partially observable MDPs. In 2015 AAAI fall symposium series , 2015.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In IEEE Conference on Computer Vision and Pattern Recognition , 2016.

Shuncheng He, Yuhang Jiang, Hongchang Zhang, Jianzhun Shao, and Xiangyang Ji. Wasserstein unsupervised reinforcement learning. In AAAI Conference on Artificial Intelligence , 2022.

Todd Hester, Matej Vecerík, Olivier Pietquin, Marc Lanctot, Tom Schaul, Bilal Piot, Andrew Sendonaris, Gabriel Dulac-Arnold, Ian Osband, John P. Agapiou, Joel Z. Leibo, and Audrunas Gruslys. Deep q-learning from demonstrations. CoRR , abs/1704.03732, 2017.

Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. In Advances in neural information processing systems , 2016.

Ahmed Hussein, Mohamed Medhat Gaber, Eyad Elyan, and Chrisina Jayne. Imitation learning: A survey of learning methods. ACM Computing Surveys (CSUR) , 50(2):1-35, 2017.

Rodrigo Toro Icarte, Toryn Q Klassen, Richard Valenzano, and Sheila A McIlraith. Reward machines: Exploiting reward function structure in reinforcement learning. Journal of Artificial Intelligence Research , 73:173-208, 2022.

Steven Kapturowski, Víctor Campos, Ray Jiang, Nemanja Rakicevic, Hado van Hasselt, Charles Blundell, and Adria Puigdomenech Badia. Human-level Atari 200x faster. In International Conference on Learning Representations , 2023.

Beomjoon Kim, Amir-massoud Farahmand, Joelle Pineau, and Doina Precup. Learning from limited demonstrations. In Advances in Neural Information Processing Systems , 2013.

Anqi Li, Byron Boots, and Ching-An Cheng. Mahalo: Unifying offline reinforcement learning and imitation learning from observations. In International Conference on Machine Learning , 2023.

Jinxin Liu, Li He, Yachen Kang, Zifeng Zhuang, Donglin Wang, and Huazhe Xu. Ceil: Generalized contextual imitation learning. In Advances in Neural Information Processing Systems , 2024.

Yuyang Liu, Weijun Dong, Yingdong Hu, Chuan Wen, Zhao-Heng Yin, Chongjie Zhang, and Yang Gao. Imitation learning from observation with automatic discount scheduling. In International Conference on Learning Representations , 2024.

Yicheng Luo, Zhengyao Jiang, Samuel Cohen, Edward Grefenstette, and Marc Peter Deisenroth. Optimal transport for offline imitation learning. arXiv preprint arXiv:2303.13971 , 2023.

Yecheng Jason Ma, William Liang, Vaidehi Som, Vikash Kumar, Amy Zhang, Osbert Bastani, and Dinesh Jayaraman. LIV: Language-image representations and rewards for robotic control. In International Conference on Machine Learning , 2023.

Yecheng Jason Ma, Shagun Sodhani, Dinesh Jayaraman, Osbert Bastani, Vikash Kumar, and Amy Zhang. VIP: Towards universal visual reward and representation via value-implicit pre-training. In International Conference on Learning Representations , 2023.

Sébastien Marcel and Yann Rodriguez. Torchvision the machine-vision package of torch. In ACM international conference on Multimedia , 2010.

Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample-efficient world models. In International Conference on Learning Representations , 2023.

Ashvin Nair, Murtaza Dalal, Abhishek Gupta, and Sergey Levine. Awac: Accelerating online reinforcement learning with offline datasets. CoRR , abs/2006.09359, 2020.

Ashvin Nair, Bob McGrew, Marcin Andrychowicz, Wojciech Zaremba, and Pieter Abbeel. Overcoming exploration in reinforcement learning with demonstrations. CoRR , abs/1709.10089, 2017.

Suraj Nair, Aravind Rajeswaran, Vikash Kumar, Chelsea Finn, and Abhinav Gupta. R3M: A universal visual representation for robot manipulation. In Conference on Robot Learning , 2023.

Andrew Y Ng, Stuart Russell, et al. Algorithms for inverse reinforcement learning. In International Conference on Machine Learning , 2000.

Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. In Advances in Neural Information Processing Systems , 2022.

Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.

Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. PyTorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems , 2019.

Karl Pertsch, Youngwoon Lee, Yue Wu, and Joseph J Lim. Demonstration-guided reinforcement learning with learned skills. In Annual Conference on Robot Learning , 2021.

Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport: With applications to data science. Foundations and Trends® in Machine Learning , 11(5-6):355-607, 2019.

Dean A. Pomerleau. ALVINN: An autonomous land vehicle in a neural network. In Advances in Neural Information Processing Systems , 1988.

Martin L Puterman. Markov decision processes: discrete stochastic dynamic programming . John Wiley & Sons, 2014.

Younggyo Seo, Lili Chen, Jinwoo Shin, Honglak Lee, Pieter Abbeel, and Kimin Lee. State entropy maximization with random encoders for efficient exploration. In International Conference on Machine Learning , 2021.

Avi Singh, Larry Yang, Kristian Hartikainen, Chelsea Finn, and Sergey Levine. End-to-end robotic reinforcement learning without reward engineering. arXiv preprint arXiv:1904.07854 , 2019.

Laura Smith, Ilya Kostrikov, and Sergey Levine. A walk in the park: Learning to walk in 20 minutes with model-free reinforcement learning. arXiv preprint arXiv:2208.07860 , 2022.

Sumedh Sontakke, Jesse Zhang, Séb Arnold, Karl Pertsch, Erdem Bıyık, Dorsa Sadigh, Chelsea Finn, and Laurent Itti. Roboclip: One demonstration is enough to learn robot policies. In Advances in Neural Information Processing Systems , 2024.

Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction . MIT press, 2018.

Daniil Tiapkin, Denis Belomestny, Daniele Calandriello, Eric Moulines, Remi Munos, Alexey Naumov, Pierre Perrault, Yunhao Tang, Michal Valko, and Pierre Menard. Fast rates for maximum entropy exploration. arXiv preprint arXiv:2303.08059 , 2023.

Faraz Torabi, Garrett Warnell, and Peter Stone. Generative adversarial imitation from observation. arXiv preprint arXiv:1807.06158 , 2018.

Luis Caicedo Torres, Luiz Manella Pereira, and M Hadi Amini. A survey on optimal transport for machine learning: Theory and applications. arXiv preprint arXiv:2106.01963 , 2021.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems , 2017.

Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control: Improved data-augmented reinforcement learning. arXiv preprint arXiv:2107.09645 , 2021.

Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on robot learning , 2020.

Yang Yu. Towards sample efficient reinforcement learning. In International Joint Conference on Artificial Intelligence , 2018.

Zishun Yu and Xinhua Zhang. Actor-critic alignment for offline-to-online reinforcement learning. In International Conference on Machine Learning , 2023.

Haichao Zhang, Wei Xu, and Haonan Yu. Policy expansion for bridging offline-to-online reinforcement learning. In International Conference on Learning Representations , 2023.

Bohan Zhou, Ke Li, Jiechuan Jiang, and Zongqing Lu. Learning from visual observation via offline pretrained state-to-go transformer. In Advances in Neural Information Processing Systems , 2023.

Zhuangdi Zhu, Kaixiang Lin, Bo Dai, and Jiayu Zhou. Self-adaptive imitation learning: Learning tasks with delayed rewards from sub-optimal demonstrations. In AAAI Conference on Artificial Intelligence , 2022.

Appendix

A Additional Experiment Results

A.1 Evaluation Curves

Figure 8 shows the evaluation curves corresponding to the results in Table 1. We can observe that TemporalOT generally outperforms the other baselines which do no use the task reward.

TaskReward

GAIfO

OT0.99

OT0.9

ADS

—TemporalOT

Basketball

Button-press

Door-lock

100

100

100

Success Rate (%)

0 00 4

08

80

60

60

40

40

20-

20

20

0

20

20

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

Door-open

Hand-insert

Lever-pull

100

50

60

Success Rate (%)

00 00 00

40

50

40

30

30

20

20

W

10

10

-20

0

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

Push

Stick-push

Window-open

35

100

100

Success Rate (%)

80

80

60

60

40

40

20

2

20

0

0

20

20

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

Step (1e6)

Step (1e6)

Step (1e6)

A.2 Bad Case Analysis

In this subsection, we analyze more bad cases in different tasks. Similar to Figure 7, we visualize the expert trajectory (top) and an agent trajectory (bottom) of some typical bad cases in different tasks. We can observe that the robot arm generally displays a similar behavior as demonstrated in the expert trajectory. The agent did not solve these tasks because it - (A) failed to grab the handle in the Door-open task; (B) failed to touch the knob in the Door-unlock task; (C) failed to pick the red block in the Push task; (D) failed to pick up the blue stick in the Stick-push task.

A.3 Pretraining with Expert Data

In this subsection, we validate the effectiveness of using imitation learning, i.e. , behavior cloning (BC), to first initialize the robot policy and then fine-tune it with TemporalOT. In this experiment, we use action-inclusive expert demonstrations to pretrain the robot policy with behavior cloning loss. Table 2 shows the results on the Door-open task. The BC baseline is a pure offline method where the parameters are fixed after pretraining. TemporalOT-P is the variant that fine-tunes a pretrained BC policy using TemporalOT. We can observe that incorporating pretraining helps to improve the sample efficiency. Moreover, we can notice a small success rate drop (which recovers later) at the initial phase when we transit from offline to online training from step 0 to step 4e4. This is an initial-dipping

(A)

(B)

(C)

(D)

phenomenon as also being reported in previous offline RL literature and can be improved by designing more specific offline-to-online RL methods [59, 60].

A.4 Ablation Studies for Expert Demonstration with Different Speed

Our work inherits an implicit assumption from learning from demonstration literature that the agent has a similar movement speed as the expert agent. Under this assumption, we use the distance between the time step indexes to represent temporal affinity information in the temporal mask. If the discrepancy between expert-agent movement speed is large, then the assumption will be broken. To validate how does such discrepancy influence the model performance, we run experiments using expert demonstrations with different speed. We generate an N times faster expert demonstration by sampling the original expert trajectory every N steps. Table 3 shows the results of using expert demo with different speed. We can observe that as the discrepancy between the expert-agent movement speed increases, the final success rate decreases. Using a double speed demonstration in the Basketball task and Door-open task achieves a slightly worse performance than the original demonstration. However, using a three times faster or four times faster expert demonstration performs quite badly. This is because a faster expert trajectory corresponds to a shorter sequence of demonstration frames, where it is more likely that the important information is dropped.

A.5 Ablation Studies for Different Visual Encoders

In this subsection, we compare the effectiveness of using different visual encoders in TemporalOT. In particular, we compare three different ResNet variants (ResNet18, ResNet50, ResNet152) using the checkpoints from torchvision [34]. From Figure 10 (left), we can observe that ResNet50 and ResNet152 encoders show similar final performances. Here, ResNet18 underperforms the other two encoders because it is quite weak, such that sometimes it fails to capture the key information of the image. The results in the Figure 10 (left) indicate that a reasonably good visual encoder is usually enough to extract effective visual embeddings for computing OT rewards in RL.

A.6 Ablation Studies for Different Mask Designs

In this subsection, we try a variant of TemporalOT that uses a dynamic mask in computing the OT rewards. Unlike Eqn.(9) where the masked position M ( i, j ) always follows a fixed rule that j ∈ [ i -k m , i + k m ] , we introduce the following dynamic counterpart,

where the mask window center c = arg min j ˆ C ( i, j ) and j ∈ [0 , i ] . Eqn.(10) means we select an index j in the expert trajectory that has the lowest transport cost w.r.t. the current observation o i . We further add a constraint j ∈ [0 , i ] to avoid looking into distant future steps as pointed in Section 3.2.3. The experiment results in Figure 10 (right) show that the learning-based temporal mask in Eqn.(10) slightly outperforms our previous rule-based temporal mask in Eqn.(9). Since the main focus of this work is to use a simple and easy to understand design to illustrate the benefits of incorporating temporal information into OT rewards in RL, we leave the investigation of more sophisticated dynamic temporal masks to the future work.

Door-open-v2

Door-open-v2

100

ResNet18

100

diagonal

ResNet50

dynamic

ResNet152

80

80

Success Rate (%)

Success Rate (%)

60

60

40

40

20

20

0

0

-20

-20

0.0

0.2

0.4

Step (1e6)

0.6

8'0

1.0

0.0

0.2

0.4

0.6

0.8

1.0

Step (1e6)

B Experimental Setup

B.1 Evaluation Environments

In the experiments, we focus on the Meta-world benchmark tasks [57]. Figure 11 shows the nine selected environments that we use in the experiments. In each environment, the robot arm aims to solve a specific task. We use the same task length parameter from the ADS paper [30]. The task lengths for the Basketball task and Lever-pull task are 175 steps and the task lengths for the other tasks are 125 steps. The goals for the nine selected tasks are as following:

Basketball task: the arm aims to grasp the orange basketball and move to a target position above the rim.

Button-press task: the arm aims to push down the red button.

Door-lock task: the arm aims to rotate the knob to a target angle.

Door-open task: the arm aims to open the door to a target position.

Hand-insert task: the arm aims to move a brown block to a target position in the hole.

Lever-pull task: the arm aims to move the lever to a target height.

Push task: the arm aims to move the red cylinder to a target position on the table.

Stick push task: the arm aims to grab the blue stick and push the bottle to a target position.

Window open task: the arm aims to open the window to a target position.

B.2 Implementation Details

Our code is based on the official implementation of ADS in PyTorch [41]. We also use the Metaworld environment [57] provided in the ADS codebase. For the fixed pretrained visual encoder, we use the Resnet50 [20] network trained on the ImageNet dataset [7]. In particular, we adopt the official Resnet50 checkpoint provided by Torchvision [34]. For the other main softwares, we use the following package versions:

Python 3.9.19

numpy 1.26.4

torch 2.2.2

torchvision 0.17.2

pot 0.9.3

dm-control 1.0.17

dm_env 1.6

mujoco-py 2.1.2.14

cython 3.0.0a10

gym 0.22.0

The pseudo-code of TemporalOT is summarized in Algorithm 1.

Basketball

Button-press

Door-lock

Door-open

Hand-insert

Lever-pull

Push

Stick-push

Window-open

Algorithm 1

Temporal Optimal Transport Reward (TemporalOT) Input: a fixed pretrained visual encoder f , N E expert video demonstrations, temporal mask window size k m , context embedding length k c , trajectory length T , total trajectory number N , experience replay buffer D . Output: trained RL agent π ( a | s ) . for i = 1 to N do Unroll policy π ( a | s ) to collect a trajectory τ = ( o 1 , · · · , o T ) . Compute visual embeddings using of the trajectory ( f ( o 1 ) , · · · , f ( o T )) . for j = 1 to N E do Compute the context cost matrix ˆ C as in Eqn.(6) for the j -th expert demonstration. Compute the masked OT reward R j τ = ( r j 0 , · · · , r j T ) with Eqn.(9). end for Select the expert demonstration with the largest trajectory OT reward sum as the final OT reward r OT τ = ( r OT 0 , · · · , r OT T ) . Save the labeled transition ( o i , a i , r OT i , o i +1 ) to the replay buffer. Update the RL agent with sampled transitions as in Eqn.(1). end for

B.3 Parameter Settings

In the experiments, we mainly follow the parameter settings as in the ADS baseline. Some of the key parameters are summarized in the Table 4.

B.4 Computation Resources

We run our experiments on a workstation with an NVIDIA GeForce RTX 3090 GPU and a 12th Gen Intel(R) Core(TM) i9-12900KF CPU. The average wall-clock running time for TemporalOT on the state-based experiment and pixel-based experiment are 1 hour and 3 hours, respectively.

C Solution of the Masked OT Objective

In this section, we provide the solution of the masked OT objective in the Eqn.(7).

To compute the optimal transport plan µ ∗ , we first write its Lagrangian as follows:

We set the partial derivative of the Lagrangian to zero:

If M i,j = 1 , then we have:

We can express it in the matrix form:

where u = e α/ϵ , K = M ⊙ e -C/ϵ , and v = e β/ϵ . The constraints can be expressed as:

We can later use the Sinkhorn algorithm [4] to iteratively compute u and v :

D Broader Societal Impacts

In this work, we investigate the application of OT-based proxy reward in learning effective robot policies with a few expert video demonstrations. Since there are usually no available task rewards in real-world scenarios, our method holds promise for the advancement of real-world RL applications. On the other hand, the proposed method aims to learn a policy that behaves similarly to the given demonstrations. Once the given expert demonstration contains some dangerous behaviors, our method is also likely to learn the dangerous behaviors, which will lead to some negative societal impacts. To mitigate such potential negative societal impacts, we can introduce an extra safety reward to prevent the agent from learning dangerous behaviors.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: We summarized the contributions and scope in the abstract and introduction, and the claims match the experimental results.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We discussed about the limitations of this work in Section 6, where we discussed about the assumptions, factors that influence the performance of the approach, and the computation cost.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We provided some theoretical result for Eqn.(7) in the Appendix C.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We provide details of the experiments in the Appendix B.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We will release the code and data.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We provided details of the experiments in the Appendix B.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We reported the mean and standard deviation of the results.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We describe the computational cost of the experiments in Appendix B.4.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We followed the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: We discussed about both potential positive societal impacts and negative societal impacts of this work in the Appendix D.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: Our paper poses no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: we have cited the original paper and our code follows the MIT License.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: We will release the code with documentation.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA] .

Justification: This work does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA] .

Justification: This work does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper extends the vanilla optimal transport-based proxy reward method in imitation learning by 1) masking out distant steps within a trajectory in the optimization of the transport plan and 2) considering neighbouring steps in the reward estimate\n\nSTRENGTHS:\n- the paper is well written and organized, with rich benchmarking experiments and comparisons\n- strong experimental results in comparison to baselines\n\nWEAKNESSES:\n- the novelty is incremental, focusing on two additions to vanilla OT\n- the presentation of the central claim seems to contradict the rationale behind the two additions to OT (see the question below).\n- some result analyses contain plain observations, lacking in-depth insights\n\t- missing discussion on why learning from demonstration performs worse in door-lock and window-open tasks (sec. 4.3), important for understanding the efficacy scope of the proposed extension; in contrast, RL with simply binary reward performs well in the twos\n\t- no insights on why important components vary across tasks (sec. 4.4), beneficial for understanding the sensitivity of the proposed extension\n   - ablation of $k_c$ and $k_m$ in sec. 4.4 is good, but a high-level guide for choosing them under different task conditions would benefit others following this work and be important for understanding the generalisation capability of the two additions\n   - as to sec. 4.6, the description of the pixel-based setting is a bit over-simplified, the author can instead put detailed explanations in the appendix\t\n   - sec. 4.7 feels like a case-by-case study; it does not relate to the method's efficacy and the paper's central claim\n\t\n- minor comment: \n     - a brief description of the nine meta-world tasks (and task length) would provide an overview of task complexity\n     - a detailed example of computing group-wise cosine similarity is missing. Including this would help clarify the value of introducing $k_c$ in the cost estimate, especially concerning the example on the right side of Figure 4\n\nQUESTIONS:\n### Presentation of the central claim\n- in section 3.1.2, how does the permutation study with ADS connect to the first observation; is being oder-invariant within a trajectory problematic? Which proposed component (masking scheme or context embeddings) addresses this? Further, how does the 1st observation relate to the claim that existing methods overlook temporal information (e.g., lines 44-45)?\n- what exactly is meant by *overlooking temporal information*? eq. 5 already counts all future steps in calculating the original OT proxy reward for current state. Why incorporate temporal information?\n- the paper claims to incorporate this information (lines 144-145), but doesn't the masking scheme *weaken* the influence of future/distant steps (lines 149-152), while context embedding *enlarges* it (lines 168-169)? This divergence causes confusion about the central claim of mitigating the issue of overlooking temporal information (e.g., lines 8-12, 45-47)\n\t- Introducing context embeddings (section 3.2.2) seems to contradict the second observation in section 3.1.2 (lines 116-117) as considering later steps could add noise to the reward estimate of the current state-action pair.\n\n### Methodology\n- is the context cost matrix in eq. 9 used in the optimisation of transport plan $\\mu$ in eq 4, or just for reward inference in eq. 5?\n\n### Experimental results\n- in Figure 6, the average of TemporalOT appears less than 7.5 but is reported to be 8.4 in Table 1.\n- in Figure 2, which step does $o_2$ correspond to? better to add x-y labels and mark $o_2$, $o_3$ to the reward curve on the right\n- misspecified value for TaskReward in Door-lock env, according to Figure 9, the mean value is around 60. Similarly, the Window-open env mean value should be around 30\n- which $\\epsilon$ values were used in the experiments (or are they determined after optimizing equation 7), and how do they affect the task performance?\n- the 4th baseline OT needs more explanation, though I can roughly guess how online OT differ from offline OT\n\n### Discrepancy in ASD vs. OT performance\n\n- unsure about why ASD performs worse than vanilla OT in some tasks as shown in Table 1? I understand this paper adopts a more challenging experimental setup, with varying goals per episode and only two demonstration videos provided. However, in the benchmarking results of the ASD paper, ASD achieves over 80% success in the basketball task within 1M steps, significantly higher than OT's performance (around 10%), as shown in Figure 11 of Appendix C.2. This trend is also observed in the lever-pull task. Yet, Table 1 in this paper shows the opposite comparison between ASD and vanilla OT for these tasks. Could you explain this discrepancy?\n\n### Minor points\n\n- Figure 1, five steps should be -> six steps? the scheme is not indicative of the magnitude of reward $r_a$ or $r_b$; what does $a_2^a$ differ from $a_2^b$?\n- It would be better to indicate in the caption of Table 1 that the values represent success rates, even though this is mentioned in section 4.1, as tables should be self-contained\n\nLIMITATIONS:\n- authors acknowledge the limitations of the work, and their comments are fair\n- it's unclear how different pretrained visual encoders will influence OT's efficacy under the same conditions; a discussion on this topic would be helpful\n- nonetheless, the authors discuss the broader impact of their work, including immediate societal implications, and suggest a straightforward solution, which is appreciated", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-04"}

comments: SUMMARY:
This work proposes a reinforcement learning method from a few expert demonstrations based on optimal transport. It incorporates temporal information into the framework of optimal transport reward so that the agent can focus on more relevant information in learning. The work proposes two simple tricks to achieve this, i.e., a masking mechanism and context embedding-based cost matrix. The experiments in simulated robotic control benchmarks show that the proposed temporal OT converges to a higher success rate than imitation learning and inverse reinforcement learning methods that do not consider temporal information.

STRENGTHS:
The paper is clearly motivated. It points out a major weakness of OT-based reward calculation, that is the temporal information not being considered, and proposes a simple solution that applies a temporal mask to the cost matrix.

WEAKNESSES:
1. TemporalOT uses a manually designed diagonal-like matrix as the mask on the transportation plan, which could be ineffective when the learning policy has not been close to the expert or gets stuck at early-stage behaviors. Have the authors tried other design choices for the temporal mask or tried to make it a part of the learning? 
2. As a robot learning paper, it would be better to apply the algorithm to real robots similar to previous OT-based methods (e.g., [18]).

QUESTIONS:
The variance of the baseline ADS is quite large in some environments shown in Fig. 7. Is there any explanation/discussion on it?

LIMITATIONS:
As shown in the learning curves, the proposed algorithm requires ~1e6 timesteps to converge, which may limit its practical use for real robots. I think initializing the policy with imitation learning, then tuning it with TemporalOT might converge faster.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-11

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes several improvements on top of the Optimal Transport reward for inverse reinforcement learning. The key observation is that the traditional OT-based reward does not consider the temporal information, i.e., it is invaraint to the order of the state-action pairs in a trajectory, and that the reward is not defined over state-action pairs but over a whole trajectory, so later transitions in the trajectory would affect the reward of an early state-action pair, which does not follow the standard definiation of reward in RL. This the key technique the paper proposes is to introduce a mask for the transport plan, such that only temporaily neighboring states between expert and the agent will be considered when computing the reward. It futher proposes a technique to futher improve the computation of the cost matrix in the OT reward. Experiments are conducted on 9 metaworld tasks, and the proposed masked OT rewards outperform several previous baselines. Ablation studies are also conducted to help understand some design choices.\n\nSTRENGTHS:\n- This paper studies an important problem in the formulation of the OT-based reward, which is it being non-temporal (invariant to the order of the trajectory). \n- The paper is well-written and easy to follow.\n- The proposed algorithm is simple yet effective, and it achieves better performances than prior methods in the experiments. \n- The experiments are thorough, with results compared on 9 environments, and with ablaiton studies that are very helpful for understanding how each component helped to improve the performance.\n\nWEAKNESSES:\n- The proposed method to restrict the temporal information works by masking the transport plan to only include temporaily nearby states. This \"temporaily nearby\" is measured using the distance between the time step indexes, which assumes that the time interval between each step, or the time scale, is similar between the agent and the expert. I am wondering how general this is true. I encourage the authors to think of cases where this assumption might be broken and add some discussion on how the method would work & can be extended to work in such cases.\n\nQUESTIONS:\nPlease refer to the weakness section.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

comments: SUMMARY:
The authors introduce TemporalOT, a learning-based proxy reward that incorporates temporal information. via using a mask mechamism and context-embeddings based cost matrix. They implemented the TemporalOT-RL based on ADS implementation, conducted thorough experiments on Meta-world benchmark focusing on the a challenge setting where there are only two demonstrations available and outperformed the OT-based reward baseline.

STRENGTHS:
1. The motivation of incorporating temporal information into OT-reward is clear.
2. The paper is well-written and easy to follow.
3. The experiments are thorough and convincing
4. The ablation study of no-mask, no-context, discount factor, number of demonstrations is very sufficient.

WEAKNESSES:
See Questions

QUESTIONS:
1. Could the author analyze what potential property of the push task makes it different that ADS outperforms TemporalOT?
2. Does the author intend to incorporate visual encoder training instead of fixed pre-trained?

LIMITATIONS:
Performance relies heavily on the demonstration quality and visual encoder effectiveness.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-30