Robot Policy Learning with Temporal Optimal Transport Reward
Authors: Yuwei Fu
Haichao Zhang
Di Wu
Wei Xu
Benoit Boulet
Venue: NeurIPS 2024 poster (2024)
Keywords: Reinforcement Learning
Imitation Learning
Optimal Transport
Primary area: reinforcement_learning
TL;DR: Temporal optimal transport reward for policy learning

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
Yuwei Fu'
Haichao Zhang²
Di Wu
Wei Xu²
Benoit Boulet
McGill University
Horizon Robotics
guwei.uummate megilll...mi...c
haichao haichao..hannghorrreee wwwgggorrzzo...ccccc..
Abstract
Reward specification is one of the most tricky problems in Reinforcement Learn- ing, 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 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 temporal order information, which could bring extra noise the reward signal. To address this issue, in this paper, we introduce the Temporal Optimal Transport (TemporalOT) reward 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: htpps:///itthll http/fuuy//eeppraa107
Introduction
Reinforcement Learning (RL) [11] has achieved great success across a wide array of applications 440]. However, it typically requires 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 [22 iee, using curiosity-based intrinsic reward to
encourage exploration 22] 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)
Reward specification plays central role in RL [6]. Since the goal of the RL agent is maximize the
expected cumulative rewards, the reward signal directly influences the learned behaviors (8]. One
major challenge in applying RL to real-worke problems is how to design the reward functions [10]. well-designed reward function can guide the agent towards desirable behaviors more efficiently, while poorly designed one could lead sub-optimal behaviors 225 However, designing 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 good reward function is one of the main
bottlenecks for the low sample-efficieney issue in RL [58).
Imitation Learning (IL) has been proven be an effective technique to learn control policies without the oracle task reward (24). Given an expert demonstration dataset, IL formulates the policy learning supervised learning paradigm
(27). The IL objective aims to learn 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, II. can be crudely classified as offline II. and online II. 29 In offline II, the RI. agent purely learns from a
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

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 II is that it generally requires a diverse and high-quality demonstration dataset to achieve desired performances [3]. Recently, some researchers found that Optimal Transport
(OT) 18, 44] based proxy reward function enables us to learn effective robot policies with only few expert demonstrations [30] 311. 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 Optimal Transport (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 be summarized 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 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, P. 7), where and A are state and action spaces, R
A R is reward
function, S x A + A(S) is the state-transition probability function, Po S + R. is the initial state distribution and y E [0, 1) is a discount factor. Our goal to learn policy (als)) : - A(A) that maximizes the expected cumulative discounted rewards E,, \ t=0 711(8t, ax)) where 80 po. 8141 ~ P(.|st. 0z) and a, ~ T(-(s)). In partially observable MDP (POMDP) 19] we can only receive an observation o O, image observation, of the current state s,
00
To solve this optimization problem, value-based RL methods typically learn state-action value function Q" ("((sa)= == E. oo 1 7're|so = 8, = a). which is defined as the expected return under policy æ. For convenience, we adopt the vector notation Q E S and define the one-step Bellman operator T SRS + R.SA such that T=Q(s, a) r(s, a) + E..'' Q(s', a') a²)]. The Q-function Qa is the fixed point of T such that Q" optimality Bellman operator as follows TQ(s, a) == r(s, a) + ? (e.rP((max Q(s', a')] and the optimal Q-value function Q* is the fixed point of TQ" = =Q" In deep RL, we use neural networks (s,s, a) to approximate the Q-functions by minimizing the empirical Bellman error:
TTQ" [51] Similarly, we define the
((sa..))))) ((r y max 07((' a') 17 a))²
(I)
where we sample transitions (s,a, r,),),))) s') from replay buffer and 07 (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 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:
E
(',
2
Σ
7 R (st, ar)
VT II,
(2)
where II is the feasible policy set. That is, the expert policy n E will achieve the maximum expected cumulative discounted reward than any other policy.

---

Demo of
Cost Matrix
Transport
Plan
Reward
of
G
Kt
An illustration of the pipeline of applying OT-based reward in RL. In this toy ex- ample. we rollout two agent for five steps of transitions. Both agents start from the initial state and take same actions an and a1 at the first two states. Then the two agents take different ac- tions aß and a generate different trajectories Ta (op. ag: 01, a1 02. az oß. ag. of, aq of) and To (Oo, do, 01, 01, 02, o², o3, ag, of a4, of) The OT rewards for (oo, ag) and (01, a1) 7a and are different even though the state-action pairs are exactly the same.
Figure I:
2.3 Optimal Transport
Optimal Transport is optimization problem which aims find optimal mapping that transforms one probability distribution into another with the least cost. OT has a wide application various domains such economics [14], physics [15], and machine learning [1. Consider two probability distributions E R" IR and joint distribution p(p, q) on product space X' x V/ the Wasserstein distance [54] between p and g is defined
W(p. q) inf inf J c(x, y)du.
_xxy
where c(x, y) is the cost function for moving mass form to y. In the RL scenario, p and usually the state space S or observation space O (2T]. For example, given an expert trajectory T (of .... of) and agent trajectory == (01, ... ,or) where o, is the image observation at step i, the Wasserstein distance between E and r is defined in the following discrete form:
T
T
W(T,
min PERTXT
ΣΣ
c(or, of )n(,s))
i=1
j=1
(4)
((i,)) Σ p(i.)) T
s.t.
where j E RT'T is called the transport plan, and we denote the optimal transport plan as n².
Method
In this section, first revisit the application of OT-based proxy reward in RL. In particular, point out the influence of temporal order information, which has been overlooked in most prior work. Next, introduce the main idea and formulation of the proposed method based these observations.
3.1 A Recap of OT Reward in RL
3.1.1
OT reward helps rank states and actions
In RI., we usually adopt the Wasserstein distance to measure the similarity of two trajectories (i) illustrated in Figure Given an agent trajectory (01,--. ,oT) and an expert trajectory

---

o5
of
Taj
Taj
-000
-0.06
OT Reward
P
-0.10
1/0% = (og.
220 2 = (d.
Step
Figure 2: Why OT reward could be useful? When the OT reward is generally correct, it helps to rank the goodness different states and induce the policy to take better actions. (left) In the toy example, two agents takes different action aB and ag O2 and thereafter. The goodness of ag and al is measured by the OT reward computed w.r.t. to the observation of the next state oğ and og. (right) A comparison of the true OT reward curves for trajectory T" and 7", where og/or/oz/o//o// correspond to observations at the 0/20/40/60/80/ 100tth step. We can observe that the OT reward for trajectory b is generally larger, which shows that the OT reward is generally correct.
-F= (of..... of), we first compute the optimal transport plan μ in Eqn. using some iterative optimization algorithms, i.e., Sinkhorn algorithm step is defined as follows:
Then, the OT-based proxy reward at the i-th
OTT - of
(5)
j=1
where c(a,, of!) is a cost function that measures the similarity between o and of. One popular choice (f(o)))((a))
prior work is the cosine similarity based cost function c(o,, off) the latent representation of observation o, extracted by visual encoder (3.
where fo;) f fll
Similar to the curiosity- exploration bonus 2 OT reward is used to distinguish the goodness of different states. We consider the toy example in Figure T two agents start from the same state with observation 02 and take different actions aß and a², respectively. Then the goodness aj and al at 02 measured by the OT reward computed w.r.t. the observation og and of the next step. As long as the OT reward can rank O''' (og, ( -F) and 1'OT' r
(o =5) correctly, then the policy will be able to
learn better action at 02. From Figure 2 (right), we can observe that the OT reward for the better trajectory is generally larger, which validates the previous explanations.
3.1.2 Two Key Observations
Notably, there are two key observations OT reward that have been less discussed prior work:
1. The OT reward is order invariant.
2. The OT reward at step 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. the order information is discarded and the frames from the demo trajectory are treated as teepoollyyy collapsed frames. In our view, collapsing the temporal axis drops arguably one of the most important characteristic features of temporal order information. More concretely, consider demo trajectory of (01, 01, 02), 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 Eqn. ( from the perspective of OT reward, there is no ability to differentiate between and some other undesired trajectories, iee, 72 = (01, 02, 01) 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. ( shows that OT reward r depends the optimal transport plan j (i, j), which depends on the entire trajectory. Therefore, the OT reward for each state is function of the entire agent trajectory. As illustrated in Figure | even though trajectory r" and trajectory 7 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 -on-stationary OT reward could have a pitfall that makes the optimization of RL objective in Eqn. ① to be less stable.
Inspired by these two observations, found that the current OT-based RL methods usually overlook the temporal order information of the trajectories. 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 Optimal Transport (TemporaIOT) 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 standard OT-based reward calculation usually consists of two stages:
(Stage-1) 1)
first define 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 OT in Eqn. for each state in a trajectory 7.
After the OT-reward calculation step, the transition will be relabeled with the OT-reward for training an RL agent as in (I.. 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 address these two points, respectively.
We will introduce two simple solutions to
3.2.2 Context Embedding- 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 the transport cost, adopt group-wise cosine similarity that we define the transport cost between agent observation o, and expert observation of as following:
(((o.))) )
a(oi, of
Σ lic h=0
1
(6)
where k is the parameter of the context length and f(.) is fixed visual encoder. The goal of the context cost matrix Ô is to facilitate expert progress estimation by taking nearby information into consideration. For example,
ke in Figure 3 and the transport cost between 0₁ and of is a(o,, o = [c (( + cos([(o2)) [(2) )) + cos(s(os), ()]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 (3n. 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 [00. To mitigate this issue, we introduce concise solution by adding a temporal mask to the cost matrix.
For an agent trajectory 7 (o),-.r, oT) and an expert trajectory , = -off of) we denote the context cost matrix Ó e TT and the transport plan μ RXT The row sum and column
5

---

Context Embedding
Context Cost Matrix
Temporal Matrix
Context Cost Matrrx
Transport
OT Reward por
C
M
é
(ex, [s)
< 4f 42)
(0
(92.9)99)
48
Avg
(u
(n 48 .00
Figure 3: An illustration of the proposed TemporalOT method. (left) Instead of using pair- pair-wise cosine similarity as the transport cost, we use group-wise cosine similarity to learn a more accurate cost matrix. (right) We use temporal mask to enforce the OT reward to focus narrow scope to avoid potential distractions from observations outside the mask window.
sum of μ equal to the constraint s = I ... 71 E R We proposed to introduce a temporal mask M E T to the transport plan, where M(i.) j) E [0, 1]. We can express the masked optimal transport objective the following vector form [16]
1 arg min (M & 14 FF <H(M @ (1), s.t. p1 1=1=s, =s
where ( is the Frobenius norm and we add an entropy regularizer H-(-) of the masked transport plan M 3 J.. We can solve Eqn. 7 by the Lagrangian:
L(, (M @ n. Ör + ((M @ p. log @ p))r 12 ④ p)1)
(8)
(o, (M @ (M)))1111 111-111- s)F (M, () s)F,
where a and 3 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 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 the diagonal matrix:
M(i (0. 0, otherwise,
if E [i km. km],
(9)
where kim is window size parameter that controls the scope we in the masked OT rewards. A smaller mask window size km 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 assume that the agent has similar movement speed as the expert [30]. Under this assumption, we adopt the distance between the time step indexes represent temporal affinity information. Figure illluttrates 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-RU in PyTorch 442 based on the official ADS implementation We use 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 fixed goal in each task, we adopt more challenging setting where the goal position changes for each episode. Moreover, we only provide two expert video demonstrations to the RI. 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

---

Table 1: Experiment results of success rate on the Meta-world benchmark.
Environment
TaskRRearrd
BC
GAIfO
OTT0999
OT0.9
ADX
TemporalOT
Basketball
0.0 (0.0)
0.2 (0.4)
(..0)
0.0 (0.00
76.6 (27.4)
42.2 (44.5)
94.4 (4.7)
Button-press
14.0 (18.5)
1.7 (2.4)
(1.1)
88.8 (2.5)
(3.3)
89.0 (3.8)
92.4 (3.6)
Dour-lock
86.2 (12.4)
(7.0)
(12.))
(5.5)
(200)
(2.7)
33.4 (2.8)
Door-open
0.0 (0.0)
10.7 (10.3)
21.71.7)
46.2
30.25)4.5)
(42.7)
78.4 (12.4)
Hand-insert
0.8 (1.6)
(2.1)
(4.4)
(9.7)
11.2
35.0
36.8 (6.6)
Lever-pull
0.0 (0.0)
(1.6).))
(1.9)
15.4 (15.5)
312.8)
212.
53.6 (7.7)
Push
1.0 (0.7)
(0.8)
(0.0)
14.2 (7.5)
(2.6)
(5.6)
8.4 (1.7)
Stick-push
0.0 (0.0)
(0.0)
18.8 (22.9)
0.0 (0.0)
48.8 (41.5)
20.0 (40.0)
97.6 (2.6)
Window-open
85.6 (12.2)
(2.))
(4.7)
54.0 (28.0)
22.4 (22.9)
43.6 (20.5)
55.2 (2.3)
Average
20.8
2.5
5.2
27.8
35,5
35.9
61.1
Door-open-v2
Hand-insert-v2
Stick-push-v2
100
Success Rate (9)
80
no-context
TemporalOT
75
40
20
50
20
10
25
Step (1e6)
Step (1e6)
Step (1e6)
Figure 4: Ablation for model components. Both proposed components are useful.
4.2 Baselines
We compare the following baseline methods. (1) ReskR: training backbone RL agent from DrQ-v2 [56] with the oracle task reward r 6 success. The reward is when the task is solved and otherwise the reward is O. (2) BC: naive behavior cloning agent which has the access of the expert action. (3) GAIfO: another IL baseline which learns discriminator to provide proxy reward [53). (4) OT: we use an online version of OTR [T agent where we first rollout the RL agent to collect online trajectories and then relabel the reward with OTR for RL training. (5) ADS: variant of OT baseline which adaptively adjusts the discount factor w.r.t. 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 shown in Table T. 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 that the RL agent fails to collect the first successful trajectory. For example, the goal of the Basketball task as shown in Figure i to pick up the basketball and move 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, is particularly challenging collect the first successful trajectory with only zero task rewards and random action explorations. On the other hand, we 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 OTO.99 uses 0.99, and OTO.9 uses 7 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 out 9 tasks, which proves the effectiveness 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 variant of TemporalOT without the context cost matrix. We can
7

---

Door-open-v2
Door-open-v2
Push-v2
Success Rate (%)
80
k--2
90
11
kn-5
Nr=2
N
60
km
60
K
ka-40
Nç=20
15
20
20
10
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)
Figure 5: Influences of key parameters. A medium number of context length kg or mask length km performs the best. The agent performs better with more expert demonstrations.
Basketball-v3
Door-lock-v2
Door-open-v2
100
Success Rate (%)
80
80
30
20
40
10-
20
TemporalO1
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)
Figure 6: Results with pixel-based inputs. TemporalOT is also effective with pixel-based inputs.
observe that removing any of the two components will lead to 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 ke for the context embedding, window size km for the temporal mask, and the demonstration number Ng. From Figure we can observe that medium number of ke and km performs the best and larger Ng improves the performances. A large context length ke does not perform well because will distract the OT reward from the current step and introduce extra reward noise. A smaller mask window size k.. makes the learning more difficult because it only receives information from nearby observations, and larger ke will gradually degrade the naive OT reward. Further, having expert demonstrations helpful 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 DrQ- v2 model setting as the ADS baseline. Figure T shows the results of the comparison of TemporalOT with ADS. We can observe similar conclusions as in T 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- bT-ba RL agents are less useful. Figure plotss a typical bad case for the OT/ADS/TemporalOY 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 target position in the hole. We can observe that the RL agent mainly focuses on imitating the arm behaviors which ends 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 the table background which sometimes

---

Figure 7:
Bad case analysis. Compared with the expert trajectory (top), the agent focused on imitating the arm behavior (bottom) and missed the details, i.e., grasping the block.
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 large body work on leveraging demonstrations for policy learning, ranging from the basic behavior cloning [12 45 to remoontration-aidee RL [43). There is also work on leveraging demonstration data for offline pre-training [60] 36 [1], to either warm-start the policy [36] or help with exploration [6] 37 22]. However, the amount of demonstra- tions required for a high-quality pre-training is typically large. In this work, we focus on the setting where only small number of demonstrations generating demonstrations. Optimal Transport based imitation is a recently emerged approach in this direction, which will be reviewed in the subsequent section.
provided [5), thus greatly relieving the burden of
Optimal Transport-basee Reward for Imitation and RL. Optimal Transport (OT) has been shown to be effective for imitation learning [18 31.. Optimal Transport Reward Labeling (OTR) [1 uses Sinkhorn distance
[ to compute a similarity metric for a trajectory w.r.t an expert demonstration and uses this metric rewards for offline RL datasets without rewards (31]. Automatic Discount Scheduling (ADS) 33 uses a similar OT-based approach for reward calculation. The core idea of ADS is to incorporate scheduling of the discount factor for online RI 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 II., our method shares some common limitations of IL. For example, the success of our method heavily depends 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 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. 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 our work is that the computation cost of the proposed method related 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.
Conclusion
This paper studies the problem of learning effective robot policies with expert video demonstrations. We focus on challenging setting where there are only two demonstrations available and the environ- ment does not provide any task reward. Following the line of research of OTT-base proxy reward, we first discuss some challenges the existing methods due 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. Experi- ments 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 can learn policies w.r.t. expert video demonstrations from different camera views.

---

References
[1] Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International Conference Machine Learning, 2017.
[2] Yuri Burda, Harrison Edwards, Amos Storkey, and Oleg Klimov. Exploration by random network distillation. In International Conference on Learning Representations, 2019.
[3] Samuel Cohen, Brandon Amos, Mare Peter Deisenroth, Mikael Henaff, Eugene Vinitsky, and
Denis Yarats. Imitation learning from pixel observations for continuous control. 2021.
[4] Marco Cuturi. Sinkhorn distances: Lightspeed computation optimal transport. In Advances in Neural Information Processing Systems, 2013.
[5] Christopher R. Dance, Julien Perez, and Théo Cachet. cetponstrationcconiiiiii reinforcement learning for few-shot imitation. In International Conference on Machine Learning. 2021.
[6] Peter Dayan and Bernard W Balleine. Reward, motivation, and reinforcement learning. Neuron, 36(2):285-22 2988 2002.
[7] 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.
[8] Rati Devidze, Goran Radanovic, Parameswaran Kamalaruban, and Adish Singla. Explicable reward design for reinforcement learning agents. In Advances in Neural Information Processing Systems, 2021.
[9] 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.
[10] Gabriel Dulac-Amold, 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, 11099)2411-2468, 2021.
[11] Jonas Eschmann. Reward function design in reinforcement learning. Reinforcement learning algorithmss: Analysis and Applications, pages 25-33, 2021.
[12] 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.
[13] 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.
[14] Alfred Galichon. Optimal transport methods in economics. Princeton University Press, 2018.
[15] Wilfrid Gangbo, Wuchen Li, Stanley Osher, and Michael Puthawala. Unnormalized optimal transport. Journal of Computational Physics, 399:108940, 2019.
[16] Xiang Gu, Yucheng Yang, Wei Zeng, Jian Sun, and Zongben Xu. Keypoint-guided optimal trans- port with applications in heterogeneous domain adaptation. In Advances in Neural Information Processing Systems, 2022.
[17] Dylan Hadfield-Menell, Stuart Russell, Pieter Abbeel, and Anca Dragan. Cooperative inverse reinforcement learning. In Advances in Neural Information Processing Systems, 2016.
[18] Siddhant Haldar, Vaibhav Mathur, Denis Yarats, and Lerrel Pinto. Watch and match: Super- charging imitation with regularized optimal transport. In Anmual Conference on Robot Learning. 2022.
[19] Matthew Hausknecht and Peter Stone. Deep recurrent q-learning for partially observable MDPs. In 2015 AAAI fall symposium series, 2015.
10

---

[20] 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.
[21] Shuncheng He, Yuhang Jiang. Hongchang Zhang, Jianzhun Shao, and Xiangyang Ji. Wasserstein unsupervised reinforcement learning. In AAAI Conference on Artificial Intelligence, 2022.
[22] Todd Hester, Matej Vecerik, Olivier Pictquin, Marc Lanctot, Tom Schaul, Bilal Piot, Andrew Sendonaris, Gabriel Dulac-Arolld lan Osband, John Agapiou, Joel Z. Leibo, and Audrunas Gruslys. Deep q q-learning from demonstrations. CoRR, abs/1704.03732, 2017.
[23] Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. In Advances in neural information processing systems, 2016.
[24] Ahmed Hussein, Mohamed Medhat Gaber, Eyad Elyan, and Chrisina Jayne. Imitation learning:
A survey of learning methods. ACM Computing Surveys (CSUR), 10(2):1-35, 2017.
[25] Rodrigo Toro Icarte, Toryn Q Klassen, Richard Valenzano, and Sheila A Mcllraith. Reward machines: Exploiting reward function structure in reinforcement learning. Journal of Artificial Intelligence Research, 73:173 1782208, 2022.
[26] 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.
[27] Beomjoon Kim, Amir-massout Farahmand, Joelle Pineau, and Doina Precup. Learning from limited demonstrations. In Advances in Neural Information Processing Systems, 2013.
[28] 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.
[29] 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.
[30] 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.
[31] Yicheng Luo, Zhengyao Jiang, Samuel Cohen, Edward Grefenstette, and Mare Peter Deisenroth. Optimal transport for offline imitation learning. arXiv preprint arXiv:2303. /3971, 2023.
[32] 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.
[33] 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.
[34] Sébastien Marcel and Yann Rodriguez. Torchvision the machine-vision package of torch. In ACM international conference on Multimedia, 2010.
[35] Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample-efficient world models. In International Conference on Learning Representations, 2023.
[36] Ashvin Nair, Murtaza Dalal, Abhishek Gupta, and Sergey Levine. Awac: Accelerating online reinforcement learning with offline datasets. CoRR, 0bs/22006.09399) 2020.
[37] Ashvin Nair, Bob McGrew, Marcin Andrychowiez, Wojeiech Zaremba, and Pieter Abbeel. Overcoming exploration in reinforcement learning with demonstrations. CoRR, abs/1770.10008) 2017.
[38] Suraj Nair, Aravind Rajeswaran, Vikash Kumar, Chelsea Finn, and Abhinav Gupta. R3M: universal visual representation for robot manipulation. In Conference on Robot Learning, 2023.
11

---

[39] Andrew Y Ng. Stuart Russell, et al. Algorithms for inverse reinforcement learning. In Interna- tional Conference on Machine Learning, 2000.
[40] 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.
[41] Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary De Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.
[42] 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, nigh-performance deep learning library. In Advances in Neural Information Processing Systems, 2019.
[43] Karl Pertsch, Youngwoon Lee, Yue Wu, and Joseph Lim. Demonstration--uuin reinforcement learning with learned skills. In Annual Conference on Robot Learning, 2021.
[44] Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport: With applications to data science. Foundations and Trendy@ in Machine Learning, 11(5-6):35-6677 2019.
[45] Dean Pomerleau. ALVINN: An autonomous land vehicle in a neural network. In Advances in Neural Information Processing Systems, 1988.
[46] Martin L Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 2014.
[47] Younggyo Seo, Lili Chen, Jinwoo Shin, Honglak Lee, Pieter Abbeel, and Kimin Lee. State en- tropy maximization with random encoders for efficient exploration. In International Conference on Machine Learning, 2021.
[48] Avi Singh, Larry Yang, Kristian Hartikainen, Chelsea Finn, and Sergey Levine. End-to-end robotic reinforcement learning without reward engineering, arXiv preprint 1rXiv:1904.07854 2019.
[49] 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.
[50] Sumedh Sontakke, Jesse Zhang, Séb Arnold, Karl Pertsch, Erdem BIyIk, Dorsa Sadigh, Chelsea Finn, and Laurent Itti. Roboclip: One demonstration is enough to learn robot policies. In Advances in Neural Information Processing Systems, 2024.
[51] Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
[52] 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.088. 08059, 2023.
[53] Faraz Torabi, Garrett Warnell, and Peter Stone. Generative adversarial imitation from observa- tion. arXiv preprint arXiv: 18Xi::1807.06588 2018.
[54] 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.01996 2021.
[55] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Informa- tion Processing Systems, 2017.
[56] Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control: Improved data-augmented reinforcement learning. arXiv preprint arXiv:2 107.000000000 2021.
12

---

[57] 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 robol learning, 2020.
[58] Yang Yu. Towards sample efficient reinforcement learning. In International Joint Conference on Artificial Intelligence, 2018.
[59] Zishun Yu and Xinhua Zhang. Actor-critic alignment for offline-to-online reinforcement learning. In International Conference Machine Learning, 2023.
[60] Haichao Zhang, Wei Xu, and Haonan Yu. Policy expansion for bridging oflline-to-online reinforcement learning. In International Conference on Learning Representations, 2023.
[61] Bohan Zhou, Ke Li, Jiechuan Jiang, and Zongqing Lu. Learning from visual observation via offline pretrained state-to-go transformer. In Advances Neural Information Processing Systems, 2023.
[62] 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.
13

---

Appendix
A
Additional Experiment Results
A.1 Evaluation Curves
Figure shows the evaluation curves corresponding to the results in Table I We can observe that Tempora OT generally outperforms the other baselines which do no the task reward.
TaskReward
GAITO
OT0.99
OTO.9
ADS
TemporalOT
Basketball
Button-press
Door-lock
100
100
100
Success Rate (%)
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.0
Door-open
Hand-insert
Lever-puil
100
50-
Success Rate (%)
30-
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
Push
Stick-push
Window-open
35-
100
Success Rate (%)
30-
60-
25
60
20-
15-
40
10-
-20
0.0
0.4
0.65
0.8
1.0
0.0
0.2
0.6
1.0
0.0
0.4
0.6
0.8
Step (1e6)
Step (1e6)
Step
(1e6)
Figure 8: Evaluation curves the Meta-world benchmark tasks.
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 observe that the robot expert trajectory. The agent did not solve these tasks because it (A) failed to grab the handle the Door-open task; (B) failed to touch the knob in the Door-unlock task; (C) failed pick the red block in the Push task; (D) failed to pick up the blue stick in the Stick- trick-push task.
generally displays similar behavior demonstrated in the
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 with TemporalOT. In this experiment, we use action-inclusio expert demonstrations Table (eeossswww the results the Door-open task. The BC baseline is pure offline method where the parameters are fixed after pretraining. TemporalOT-PP is the variant that fine-tunes pretrained BC policy using TemporalOT. We can observe that incorporating pretraining helps improve the sample efficiency. Moreover, we can notice small success rate drop (which recovers later) at the initial phase when transit from offline to online training from step to step 4c4. This is an initial- -nitial-dipping
pretrain the robot policy with behavior cloning loss.
14

---

(A)
(B)
(C)
(D)
Figure 9: Visualization of more bad cases in the (A) Door-open task, (B) Door-unlock task, (C) Push task and (D) Stick-push task.
phenomenon as also being reported previous offline RI. literature and can improved by designing more specific offline-to-omi -oniinee RL methods [59 60).
15

---

Table 2: Pretraining with action-inclusive demonstrations is helpful.
0
2e4
4e4
6e4
8e4
le5
5e5
le6
BC
10.8
TemporalOT
0
2.0
0
5.0
16.6
57.8
78.4
TemporalOT-P
10.8
6.8
25.0
42.8
48.6
55.4
70.8
82.0
Table 3: Experiment results of success rate for demonstrations with different speed.
1x
2x
3x
4x
Basketball
94,4 (4.7)
91.8 (8.9)
77.2 (8.4)
43.8 (29.5)
Button-press
89.0 (3.8)
(3.0)
65.6 (5.4)
57.8 (7.4)
Door-open
(12.4)
76.2 (10.4)
49.4 (20.3)
20.4 (9.2)
A.4 Ablation Studies for Expert Demonstration with Different Speed
Our work inherits implicit assumption from learning from demonstration literature that the agent has 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 times faster expert demonstration by sampling the original expert trajectory every N steps. Table 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 double speed demonstration in the Basketball task and Door-open task achieves a slightly worse performance than the original demonstration. However, using three times faster four times faster expert demonstration performs quite badly. This is because faster expert trajectory corresponds to a shorter sequence of demonstration frames, where it 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 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 55 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 I0 (left) indicate that reasonably good visual encoder is usually enough to extract effective visual embeddings for computing OT rewards RL.
A.6 Ablation Studies for Different Mask Designs
In this subsection, we try variant of TemporalOT that uses dynamic mask in computing the OT rewards. Unlike Eqn. (9) where the masked position (i, always follows a fixed rule that j E ji km, i + km], we introduce the following dynamic counterpart,
[1, M(I.) - 0, otherwise,
j E [e km. c km]]
(10)
where the mask window center c arg min, C(i,j) and j [0, Eqn.q10 means we select an index j in the expert trajectory that has the lowest transport cost w.r.t. the current observation og- We further add constraint j E [0, i] to avoid looking into distant future steps as pointed in Section 3.2.3 The experiment results in Figure IO (right) show that the baaning-based temporal mask in Eqn. (10) slightly outperforms our previous rule-based temporal mask in Eqn.O). Since the main focus of this work is to use simple and easy to understand design to illustrate the benefits of incorporating temporal information into OT rewards in RI., we leave the investigation more sophisticated dynamic temporal masks to the future work.
16

---

Doar-open-22
Decr-opetwe
Nasherd
Gynare
Success Sate thi
Step (led)
(3e5)
Figure 10: More ablation studies. (left) A comparison of different visual encoders. (right) Using dynamic temporal mask slightly improves the performance.
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 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
following:
Basketball task: the arm aims grasp the orange basketball and move target position above the rim.
Button-press task: the arm aims to push down the red button.
Door-lock task: the aims rotate the knob target angle.
Door-open task: the arm aims to open the door to a target position.
Hand-insert task: the arm aims to move brown block to a target position in the hole.
Lever-pull task: the aims to move the lever to 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 open the window target position.
B.2
Implementation Details
We also use the Meta- world environment [57] provided in the ADS codebase. For the fixed pretrained visual encoder,
Our code based on the official implementation of ADS in PyTorch 41
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:
1. Python 3.9.19
2. numpy 1.26.4
3. torch 2.2.2
4. torchvision 0.17.2
5. pot 0.9.3
6. dm-control 1.0.17
7. dm_env 1.6
8. mujoco-py 2.1.2.14
9. cython 3.0.0a10
10. gym 0.22.0
The pseudo-code of TemporalOT is summarized in Algorithm I
17

---

Bosketball
Button-press
Door-lock
Deor-open
Hand-insert
Lever-pull
Push
Stick-push
Window-open
Figure 11: Visualization of the nine evaluation tasks in the experiments.
Algorithm Temporal Optimal Transport Reward (TemporalOT)
Input: a fixed pretrained visual encoder f. Ngg expert video demonstrations, temporal mask window size km context embedding length lics trajectory length T, total trajectory number N, experience replay buffer DD
Output: trained RL. agent (aass).
for i to N do
Unroll policy (Taas) to collect trajectory t = .T).
Compute visual embeddings using of the trajectory (f(oi). ((OT)).
for j 1 to Ng
Compute the context cost matrix Ô as in Equ. 6 for the j-th expert demonstration.
Compute the masked OT reward R (ro. -. mry with Eqn.
end for
Select the expert demonstration with the largest trajectory OT reward sum as the final OT reward ,0T (r ..... -,7")
Save the labeled transition (o,, a,, ,0T 0.41) to the replay buffer.
Update the RL agent with sampled transitions as Eqn.
end for
B.3 Parameter Settings
In the experiments, we mainly follow the parameter settings in the ADS baseline. Some of the key parameters summarized in the Table
18

---

Table 4: Summarization of hyper-paramelers.
Parameter
Value
Total environment step
le6
Adam learning rate
le-4
Batch size
Target network
0.005
Discount factor y
0.9
Expert demo number Ng
2
Context length &cc
Mask window size kim
DrQ buffer size
1.5c5
DrQ action repeat
DrQ frame stack
DrQ image size
(84, 84, 3)
DrQ embedding dimension
DrQ CNN features
(32, 32, 32)
DrQ CNN kernels
(3, 3, 3, 3)
DrQ CNN strides
(2, I,
DrQ CNN padding
VALID (1024, 1024, 1024)
DrQ actor network
DrQ critic network
(1024, 1024, 1024)
B.4
Computation Resources
We run experiments workstation with NVIDIA GeForce RTX 3090 GPU and a 12th Gen Intel(R) Core(TM) 19-12900KF CPU. The average wall-clock running time for TemporalOT on the state-based experiment and pixel-based experiment are 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
arg (i @ μr C)F eH(M @ n)), s.t. p1 11 s,
To compute the optimal transport plan μu, we first write its Lagrangian as follows:
Lp, (M F) + ((M @ p. log(M @ p)) 1 (M @ p)1) (a, (M @ 111 - s)F - (8, (M o 11) T 1)
We set the partial derivative of the Lagrangian to zero:
aL
Mis + eMi,, log(Mr, Mi Mr,883 = 0.
If Maj 1, then have:
log(uis)
Oi +
3
Cij
can express it in the matrix form:
M O μt = diag( u)) Kdiag(v)) (v),
where / K M -0-e-- C/e and eB The constraints can be expressed as:
diag(u) ()) K diag(v) = (diag(u) (1) (diia(v))) s.
We can later use the Sinkhorn algorithm
to iteratively compute and v:
(1+))
K kv
(++1)
K u (+1)
19

---

D Broader Societal Impacts
In this work, we investigate the application of OT-based proxy reward in learning effective robot policies with 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 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.
20

---

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 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 is clear that these goals attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: We discussed about the limitations of this work in Section 6 where 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 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 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 the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or 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 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 used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant 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 complete (and correct) proof?
21

---

Answer: [Yes]
Justification:
We provided some theoretical result for Eqn. 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 cross- referenced.
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 provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix 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 ex- perimental results of the paper to the extent that 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 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 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 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 large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily 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 new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with open-source dataset instructions for how construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome 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., 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
22

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
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 NeurlPS code and data submission guidelines (https://nips... cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "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
/ "cippscc/puullc//uuu /ttee//ggudddeeeedSuuuummbibbi ssionPolicy) 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. only subset of experiments are reproducible, they should state which ones omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, 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 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, confi- dence 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.)
23

---

The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
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 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 range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
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, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't 't make into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics
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.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration 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 no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact 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 particular applications, let alone deployments. However, if there is 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 being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there 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 monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our paper poses no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have 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 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 best faith effort.
12. Lieenses for existing assets
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?
Question:
Answer: [Yes]
Justification: we have cited the original paper and our code follows the MIT License.
Guidelines:
existing assets. The authors should cite the original paper that produced the code package or dataset.
The answer NA means that the paper does
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.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, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of 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.
25

---

If this information is not available online, the authors are encouraged to reach out to the asset's '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/moodl 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 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 crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, 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 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 NeurlPS 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.
26


Reviews:
[{"review_id": "8eir6v11iX", "reviewer": "Reviewer_hYj8", "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.", "questions": "1. Could the author analyze what potential property of the push task makes it different that ADS outperforms TemporalOT?\n2. 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.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The motivation of incorporating temporal information into OT-reward is clear.\n2. The paper is well-written and easy to follow.\n3. The experiments are thorough and convincing\n4. The ablation study of no-mask, no-context, discount factor, number of demonstrations is very sufficient.", "weaknesses": "See Questions", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for the positive review and for the efforts in reviewing our work!\n\nYour comments are valuable in helping us to further improve the quality of the paper.\n\nResponses to the questions are below:\n\n| **Q1: Could the author analyze what potential property of the push task makes it different that ADS outperforms TemporalOT?**\n\nThanks for raising this point. The goal of the push task is to move the red cylinder to a target position (green dot) on the surface of the table, as illustrated in Figure10 (C) in the appendix. Usually, we can solve a task by focusing on the robot arm trajectory in the demo, i.e., open a door. However, we need to pay special attention to the red cylinder in the *push task* and move it to the exact target position.\n\nADS performs better in this task because:\n- Firstly, the red cylinder is quite small compared to the red robot arm. Therefore, the cosine similarity based OT reward usually contains more information about the robot arm than the red cylinder. The agent is likely to focus on imitating the robot arm movement and ignores the small red cylinder.\n- Secondly, ADS uses a larger 𝛾 than TemporalOT due to its adaptive discount scheduler. From the comparison of OT0.99 and OT0.9 in Table 1 of the main paper, we can observe that using a larger discount factor indeed helps to improve performance.\n\nThe benefit of using the adaptive discount scheduler in ADS is significant for the *push task* because the larger 𝛾 encourages the RL agent to look into more future steps via TD-learning and assign higher weight for future rewards. This is important because the OT reward is generally more noisy in the *push task* than the other tasks due to the small size of the red cylinder. For example, the OT reward $r_t$ at step $t$ is sometimes not that accurate but the OT reward $r_{t+3}$ at step $t+3$ is more accurate and contains more useful information about the red cylinder. Under a such circumstance, a larger 𝛾 will let the agent pay more attention to $r_{t+3}$ and learn a more accurate Q-value function and policy accordingly.\n\nIn the experiment, we always use 𝛾=0.9 to match the initial 𝛾 value in ADS and the second best baseline OT0.9. We added experiments to run TemporalOT with a larger discount factor 𝛾=0.99. From PDF-Fig 1 in the uploaded PDF file, we can observe that TemporalOT with a larger discount factor achieves better performance in the push task.\n\n| **Q2: Does the author intend to incorporate visual encoder training instead of fixed pre-trained visual encoder?**\n\nThanks for the great question. The reason for using a fixed pre-trained visual encoder in our work is mainly following standard settings in literature [e.g. ADS paper], which removes the factor of visual encoder learning and helps us to focus on reward structure design.\n\nAlthough it is not the main focus of this work, we can definitely incorporate visual encoder training instead of using a fixed pre-trained visual encoder. For example, we have the following different strategies to train the encoder:\n\n- In the pixel-based setting, we unfreeze the encoder in the critic and use the Q-learning loss to optimize the encoder.\n- We can also adopt some existing representation learning methods for RL to train the encoder, i.e., ATC [Stooke et al., 2021].\n\nIntuitively, this has the potential to further boost overall performance by reducing the potential domain shift issue, better capturing task-relevant signals, etc. Since the main focus of our current work is to demonstrate the effectiveness of adding temporal information to the OT rewards in RL, we follow standard setting in literature in order to make the comparison clear and consistent and leave the incorporation of visual encoder training as an interesting future work.\n\n| **Q3: Rely on demonstration quality and visual encoder effectiveness.**\n\nThanks for the comment.\n\n**Demonstration quality.** We agree with you that this is a common limitation of imitation learning (IL). Since our method is closely related to IL, TemporalOT also faces this common issue. On the other hand, unlike many other IL methods that require a large diverse dataset of high-quality demonstrations, TemporalOT is designed to work with only a few expert demonstrations. For example, we only provide TemporalOT with two expert trajectories in the experiments. Such a small number of expert demonstrations makes our method more practical than the other IL methods, which require more expert demonstrations.\n\n**Visual encoder.** We leverage the visual encoder to extract embedding for each image and compute the cosine similarity based transport cost. Since this is also a common component shared by baseline methods such as OT and ADS,  any reasonably effective visual encoder can be used to make the comparison fair. In this work, we follow the experiment setup in ADS and use a pre-trained ResNet50 to extract the visual embedding.\n\nOther visual encoders could also be used. Following your suggestion, we also run some experiments quickly to compare different visual encoders in PDF-Fig 2 in the uploaded PDF file. In PDF-Fig 2, we compare three different ResNet variants (ResNet18, ResNet50, ResNet152) using the checkpoints from torchvision. 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. PDF-Fig 2 indicates that a reasonably good visual encoder is usually enough to extract effective visual embeddings for computing OT rewards in RL.\n\nGiven the rapid development of computer vision models, we can easily build our own model with some strong open-sourced vision encoder checkpoints. Moreover, we can incorporate visual encoder training, as discussed in Q2, to further improve performance.\n\n**References**\n\n[Stooke et al., 2021] Decoupling Representation Learning from Reinforcement Learning, ICML 2021", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "XxF833kkou", "reviewer": "Reviewer_FnLe", "summary": "This 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.", "questions": "Please refer to the weakness section.", "limitations": "Yes.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- 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.", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for your positive review and insightful comments!\n\nYour valuable comments helped us to further improve the quality of the paper.\n\nResponses to the questions are below:\n\n| **Q: The proposed method to restrict the temporal information works by masking the transport plan to only include temporarily nearby states. This \"temporarily 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\nThanks for the great question.\n\nOur work indeed inherits an implicit assumption from *learning from demonstration* literature:\n\n- The agent has a similar movement speed as the expert. \n\nUnder this assumption, using the distance between the time step indexes is a simple and effective strategy to represent temporal affinity information.  We will make this clear in our revised paper.\n\n\nAs to your question on wondering how general this is true, we do admit that this assumption can be broken if the discrepancy between expert-agent movement speed / trajectory alignment is large enough.\n\nHere we share some quick empirical results that we were able to finish within the limited time on investigating the performance v.s. discrepancy (more results with the complete transition phase, i.e. the stage that it starts to fail due to increased discrepancy, will be added to the appendix of revised paper).\n\nWe have done experiments to use double speed expert demonstrations, where we get the demos by sampling the original expert trajectory every 2 steps. PDF-Fig 3 in the uploaded PDF file shows the experiment results. Though the expert trajectory length is only half of the agent trajectory length, the proposed TemporalOT method is still effective under a such setting. This experiment provides some positive empirical evidence that there are some tolerance in terms of the deviation from the assumption (discrepancy). \n\nIn addition to these empirical investigations, as a future extension, we may need to further generalize the definition of temporal affinity accordingly, potentially at a more abstract level. For example, we can first learn to infer the intention/latent goal of the agent and then compute the OT reward w.r.t. the intentions [Ghosh et al., 2023]. This is a very interesting direction that is worth exploring in future work.\n\n\n**References**\n\n[Ghosh et al., 2023] Reinforcement learning from passive data via latent intentions, ICML 2023", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DqpXq0rdQN", "reviewer": "Reviewer_M6BJ", "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.", "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.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "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? \n2. 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]).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the detailed, constructive review and insightful comments.\n\nYour valuable comments helped us to further improve the quality of the paper.\n\nResponses to the questions are below:\n\n> **Q1: Try other design choices for the temporal mask or try to make it a part of the learning**\n\nThanks for the great question.\n\nWe used the current design in the paper is because:\n\n- We want to use a simple and easy to understand design to illustrate the benefits of incorporating temporal information into OT rewards; \n- The current design is interpretable, easy to implement and computationally cheap.\n\nTo mitigate the issue that the agent could get stuck at early-stage behaviors, here, we introduce a variant of TemporalOT that uses a dynamic mask for each trajectory. We begin with some notations:\n- The agent trajectory is $\\tau = (o_1, \\cdots, o_T)$ and the expert trajectory is $\\tau^E = (o^E_1, \\cdots, o^E_T)$.\n- We compute a cosine similarity based transport cost matrix $\\hat C\\in\\mathbb R^{T \\times T}$ as Eq 9.\n- For each step $o_i$ in the agent trajectory, we compute a mask window $M_i\\in\\mathbb R^{1 \\times T}$ where $M_i(j)$=0 or 1 for $1\\le j \\le L$.\n\nIn the current paper, we use a variant of the diagonal matrix with a mask window size $k_m$:\n$$\nM_i(j) = \\begin{cases}1, &\\text{if } j \\in [i-k_m, i+k_m] \\\\\\\\0, &\\mathrm{otherwise} \\end{cases}\n$$\n\nIn the added experiment, we learn a dynamic mask window for each step as follows:\n$$\nM_i(j) = \\begin{cases}1, &\\text{if } j \\in [c-k_m, c+k_m] \\\\\\\\0, &\\mathrm{otherwise} \\end{cases}\n$$\nHere, we select the window center index as $c = \\arg\\min_j \\hat C(i, j), j \\in [0, i]$. This means we select an index $j$ in the expert trajectory that has the least transport cost w.r.t $o_i$. We further add a constraint $j \\in [0, i]$ to avoid looking into distant future steps as pointed out in Section 3.1.2.\n\nThe experiment results are shown in PDF-Fig 4 of the attached PDF. We can observe that this learning-based temporal mask slightly outperforms our previous rule-based temporal mask. \n\nWe have been focusing on using a simple and effective approach to illustrate the main point of the paper, and therefore did not try to learn the temporal masks. We do agree with your intuition that learning temporal masks is a very promising direction to pursue. \n\n> **Q2: As a robot learning paper, it would be better to apply the algorithm to real robots similar to previous OT-based methods**\n\nThanks for the suggestion. \n\nWe apologize that simulated environments are mainly used in experiments.\nThe reason is that we currently have no real robots at hand. \nAt the same time, in our humble opinion, benchmarking different methods under the same setting in simulation also provide valuable empirical verifications on the effectiveness of different methods. \n\nWe fully agree with you that it is important to validate the effectiveness of the proposed method in real-world scenarios. \nWe are very willing to deploy our algorithm to real robots once physical robots are available in the future. \n\n> **Q3: High variance of ADS in some environments.**\n\nThanks for the great question.\n\nActually, the ADS baseline also shows high variance in its original paper [ADS]. The main reasons are:\n\n- ADS uses the naive OT reward definition (Eq 5) that depends on the whole trajectory. Such a global scope sometimes introduces noisy information from far away steps. Therefore, we can observe that both the ADS and OT show high variances in Figure 4, 5, 10, 11, 12 from the ADS original paper.\n\n- Moreover, the discount factor changes during the training in ADS. This makes the TD target less stable. Sometimes, such a changing discount factor will lead to a change in the Q-value function and the learned policy. In some tasks, the less stable ADS policy will also lead to high variance.\n\n\n> **Q4: 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.**\n\nThanks for the question.\n\nUsing imitation learning to initialize the robot policy and then fine-tuning it with TemporalOT will very likely to help converge faster. \n\nOne thing to note is that the demos used in our work and previous literature [e.g. ADS] are assumed to be **action-free videos**. Therefore, many of the standard action-based imitation learning approaches are not applicable, preventing a simple imitation-based policy pre-training.\n\nTo do policy pre-training, state-based imitation learning approach should be used. However, the proposed approach is actually a type of state-based policy learning approach already. From this sense, although we can further improve sample efficiency with pre-training, it can be regarded as a factor that can be used by all baselines. Therefore, for fair comparison in this work, we place all the baseline methods on the same footing without policy-pre-training. \n\nOn the other hand, to verify the point raised in your comment and go beyond the settings used in this work, we did some quick experiments with **action-inclusive demonstrations**. In PDF-Fig 5 of the attached PDF file, the BC Pretrain baseline refers to pre-training the policy with imitation learning. In addition, the BC Loss baseline refers to using both imitation learning loss and RL-loss for policy training throughout. \nWe can observe that both BC Pretrain and BC Loss show performance gain in terms of sample efficiency, which confirms your intuition.\n\n**References**\n\n[ADS] Liu et al., Imitation Learning from Observation with Automatic Discount Scheduling, ICLR 2024", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sNCRGL1BDV", "reviewer": "Reviewer_GZSS", "summary": "This 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", "questions": "### 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", "limitations": "- 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", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- the paper is well written and organized, with rich benchmarking experiments and comparisons\n- strong experimental results in comparison to baselines", "weaknesses": "- 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the detailed and constructive comments.\n\n## Questions\n\n> **Q1: Presentation of the central claim**\n\nDue to the importance of this point and limited space, we have clarified the motivation/central claim in the common reply (top) on Clarifications on motivation.\nIn short, the permutation study provides empirical evidence that the conventional OT-reward (Eq 5) is order-invariant (1st obs),  since the frame order in the expert trajectory is fully discarded. \n\n> **Q2: What exactly is meant by overlooking temporal information? How does 1st obs relate to it**\n\nWe meant overlooking the temporal *order* information (apologize and will revise). Since Eq 5 is order-invariant, methods using it will overlook order info. We further detailed on this in C1. \n\n> **Q3: Doesn't the masking scheme weaken ... while context embedding enlarges....**\n\nBoth scheme are for incorporating the temporal *order* information (C2).\nThis seems to be caused by the term \"temporal information\". We replaced it with \"temporal order information\" in the revised paper. \n\n> **Q4: Introducing context embeddings seems to contradict the second observation**\n\nThe context length $k_c$ (i.e., 2-3) is much smaller than the trajectory length (i.e.,100-200). In observation 2, we discuss the distraction from distant states (i.e.,100-200 steps). The context embedding with a small context length $k_c$ will have a small influence on the distraction issue as discussed in observation 2.\n\n> **Q5: Is the context cost matrix in Eq9 used in the optimization of transport plan 𝜇 in Eq4, or just for reward inference in Eq5?**\n\nThe context cost matrix in Eq9 is used in *both* transport plan optimization and reward inference.\n\n> **Q6: In Figure 6, the average of TemporalOT appears less than 7.5 but is reported to be 8.4 in Table 1**\n\nThis is because we plot a smoothed evaluation curve. Table 1 reports the result at the last step: (7+7+11+10+7)/5 = 8.4. Figure 6 plots a smoothed curve: (5+4.33+8+9.33+7)/5 = 6.73. Logs are available in PDF-Table 1 in the uploaded PDF.\n\n> **Q7:  $o_2$ in Figure 2** \n\n$o_2$ is the 40th step. \n\n> **Q8: Misspecified value for TaskReward**\n\nSorry that we made a mistake when plotting Fig 9. We extracted the same number of data samples for TaskReward from the evaluation log as other baselines, and thought it corresponds to1e6 steps. But actually it is 5e5 steps since it uses doubled evaluation frequency. PDF-Fig 7 is the corrected figure matching Table 1\n\n> **Q9: Effect of 𝜖**\n\nWe follow the parameter setting from ADS where 𝜖=0.01. PDF-Fig 6 from the PDF file shows the ablation for different 𝜖. The entropy coefficient 𝜖 controls the sparseness of the transport plan 𝜇. A large value will generate a 𝜇 closer to uniform distribution.\n\n\n> **Q10: Discrepancy: Table 1 shows the opposite comparison between ASD and vanilla OT for basketball and lever-pull**\n\nSorry about the confusion. OT in the ADS paper corresponds to OT0.99 in Table 1 of our paper. By comparing OT0.99 with ADS, we can see that the results are *not opposite*, with ADS outperforming OT0.99 on basketball and lever-pull tasks. The overall performance drop is indeed due to the more challenging experimental setup as you mentioned.\n\n\n> **Q11: Figure 1, five steps ... what does $a^a_2$ differ from $a^b_2$**\n\nSorry that *steps* is a confusing term and we changed it to *steps of transitions*.\n$r_b$ is larger than $r_a$ since  $a^b_2$ leads to a success grasp ($a^a_2$ fails resp.) and matches demo better. \n\nThanks for all the good suggestions and will include the improved figure in revision. \n\n\n> **Q13: Different pretrained visual encoders**\n\nThanks for the suggestion and we will add more discussions in revision.\n\nIn this work, we follow the standard setting [ADS] and use a pre-trained ResNet50 to extract the visual embedding.\nOther visual encoders could also be used. We also run some experiments quickly to compare different visual encoders in PDF-Fig 2 in the uploaded PDF file. In PDF-Fig 2, we compare three different ResNet variants (ResNet18, ResNet50, ResNet152). PDF-Fig 2 indicates that a reasonably good visual encoder is usually enough to extract effective visual embeddings for computing OT rewards in RL. \n\n\n## Weaknesses\n\n> **W1: Novelty**\n\nTo the best of our knowledge, we are the first to discuss how to leverage temporal order information in computing OT rewards for training an RL agent. As an initial attempt on this topic, we aim to address this problem from a minimalist perspective. Though being simple, the method is effective and could be a good starting point to inspire others to reproduce and build upon.\n\n> **W2: Bad performance in door-lock and window-open**\n\nBecause these two tasks have strict success conditions, i.e., the door-lock task is defined as solved only when the knob rotates to 60°. IL proxy reward is not accurate enough for such details, i.e., 55° w.r.t. 60°.\n\n> **W3: Important components vary across tasks**\n\nWhen long distance distraction is the main bottleneck, temporal mask is the key. When inaccurate transport cost is the main bottleneck, context embedding is the key.\n\n> **W4: How to choose of $k_c$ and $k_m$**\n\nA small value of $k_c$=3,4 is effective enough for a context information-augmented OT cost, while a larger value will oversmooth the cost information. $k_m$ relates to the length of the expert trajectory L. Setting $k_m$ to around 0.1L is a good rule of thumb.\n\n> **W5: Example of computing group-wise cosine similarity**\n\nIn Figure 4, we use $k_c=3$. The transport cost between $o_1$ and $o^E_2$ is: $\\hat c(o_1, o^E_2) = 1 - \\left[\\mathrm{cos}(f(o_1),f(o^E_2))+\\mathrm{cos}(f(o_2),f(o^E_3))+\\mathrm{cos}(f(o_3),f(o^E_4))\\right]/3$.\n\n\n## Other\n> **Caption of Tab 1; details about sec 4.6 in appendix; comments on sec 4.7; description of meta-world; more explanation on the 4th baseline ...**\n\nThanks for all the other comments (cannot list all due to space; all received, thanks!) and we have revised paper accordingly.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
# To all the reviewers: Thanks for the reviews and summary of key paper changes

We thank all the reviewers (**R1**-hYj8, **R2**-FnLe, **R3**-M6BJ and **R4**-GZSS) for their time and efforts in reviewing the paper. The reviewers appreciated that:
- The motivation is clear and the problem is important (**R1**-hYj8, **R2**-FnLe, **R3**-M6BJ, **R4**-GZSS).
- The paper is well-written and easy to follow (**R1**-hYj8, **R2**-FnLe, **R4**-GZSS).
- The proposed method is simple and effective (**R1**-hYj8, **R2**-FnLe, **R3**-M6BJ).
- The experiments are thorough (**R1**-hYj8, **R2**-FnLe, **R4**-GZSS).

We have further improved the paper according to the valuable suggestions from all reviewers. Below is a summary of the key changes made to the current paper. We:
- added experiments for $\gamma$=0.99 on the push task (**R1**-hYj8) (c.f. PDF-Fig 1)
- added experiments to compare different visual encoders (**R1**-hYj8, **R4**-GZSS) (c.f. PDF-Fig 2)
- added experiments to use double speed expert demonstrations (**R2**-FnLe) (c.f. PDF-Fig 3)
- added experiments to use different masks (**R3**-M6BJ) (c.f. PDF-Fig 4)
- added experiments to incorporate BC pre-trained policy (**R3**-M6BJ) (c.f. PDF-Fig 5)
- added ablations for $\epsilon$ parameter (**R4**-GZSS) (c.f. PDF-Fig 6)
- re-generated Figure 9 by fixing a plotting issue that caused a mis-match with the results in Table 1 (**R4**-GZSS) (c.f. PDF-Fig 7)
- added clarifications and discussions on the assumptions of TemporalOT (**R3**-M6BJ, **R4**-GZSS).
- added more detailed explanations for the experiment results (**R4**-GZSS).

We greatly appreciate the efforts the reviewers spent on helping to improve our work and we are deeply grateful for all the valuable suggestions. We have carefully addressed each comment and improved our work accordingly. We hope they are helpful to address the concerns the reviewers had previously. We would love to hear feedback from all the reviewers.

# Clarifications on motivation
The motivation of this work has been well-perceived by **R1**-hYj8, **R2**-FnLe and **R3**-M6BJ (e.g. quoted from their reviews respectively: **R1**-hYj8: "the motivation of incorporating temporal information into OT-reward is clear"; **R2**-FnLe: "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)";  **R3**-M6BJ: "the paper is clearly motivated. It points out a major weakness of OT-based reward calculation").

Here we try to make some further clarifications to help resolve any potential confusions **R4**-GZSS seems to have (we apologize to **R4**-GZSS for any potential confusions we caused and we will resolve them below and also in the revised paper). We decide to re-clarify the motivation here since it is one of the most crucial prerequisite to better understand the method and resolve many related questions and confusions. Another reason to list it here is that it could be useful to readers from a broader audience later after the reviews and rebuttals are made public.

> **C1: Is being order-invariant within a trajectory problematic? Why order information matters in OT-based RL?**

Temporal order information is very important and being oder-invariant within a trajectory is problematic.

More specifically, in the standard OT-reward (Eq. 5), the order information is discarded and the frames from the demo sequence are treated as *bag-of-temporally-collapsed* frames. In our view, collapsing the temporal axis drops arguably one of the most important characteristic features of temporal information. 

More concretely, consider a demo trajectory of 𝜏1 = (S1, S1, S2), meaning first to stay at S1 and then move to S2, and we want to imitate this behavior. However, if order information is discarded as in standard OT-reward (Eq. 5), from the perspective of reward, there is no ability to differentiate between 𝜏1 and some other (undesired) trajectories, e.g. 𝜏2 = (S1, S2, S1) (first move to S2 and move back to S1).
Therefore, it is possible for the policy to converge to a solution that generates trajectories like 𝜏2, which is undesirable. 

In summary, discarding the temporal ordering information in reward calculation makes the reward on top of it under-constrained, which in consequence increases the possibility of converging to undesired solutions. This motivates us to investigate how to address this issue by incorporating temporal ordering information.

> **C2: Motivation of the proposed method. Which proposed component (masking scheme or context embeddings) addresses the lacking of ordering issue in the standard OT-reward?**

Both masking scheme and context embeddings contribute to the final overall performance and are incorporated into different stages in the algorithmic pipeline.

As a brief recap, the pipeline of standard OT-based RL is as follows:

(Stage 1) Transport cost matrix => (Stage 2) OT reward => (Stage 3) RL training

**Temporal mask.** We leverage the mask mechanism to incorporate temporal affinity and therefore the order information. This is the component that is designed to address the lacking of ordering issue in the standard OT-reward explicitly (Stage 2).

**Context embedding.**  This is a component for improving the quality of the cost matrix (Stage 1) by incorporating nearby context information. This component contribute to the inclusion of temporal ordering implicitly (and at a coarser temporal scale), by the incorporating of a temporal window as context for cost calculation (e.g. frames outside of the temporal window will not be considered in cost calculation).

We realized that the roles of these two components might not be well clarified in the original submission. We apologize to **R4**-GZSS for any potential confusions we caused and will resolve them in the revised paper.

Meta-review: This paper proposes a method for inverse RL based on optimal transport, building on top of prior work with a new masking mechanism. Reviewers found the paper well written and well motivated, with thorough experiments. While the differences from prior work are small, the paper presents strong empirical results on metaworld. The authors and reviewers had a thorough discussion, which cleared up many questions and concerns about the paper. 

Given that all reviewers scores indicate that the paper should be accepted, I’m recommending that the paper be accepted. I encourage the authors to incorporate the rebuttal experiments and discussion into the final version of the paper.

Minor comments: In Table 1, please bold all entries that are within the standard error of the best method (e.g., some of the numbers in the ADS column should be bolded).
Decision: Accept (poster)