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
Strengths
- the paper is well written and organized, with rich benchmarking experiments and comparisons
- strong experimental results in comparison to baselines
Weaknesses
- the novelty is incremental, focusing on two additions to vanilla OT
- the presentation of the central claim seems to contradict the rationale behind the two additions to OT (see the question below).
- some result analyses contain plain observations, lacking in-depth insights
- 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
- no insights on why important components vary across tasks (sec. 4.4), beneficial for understanding the sensitivity of the proposed extension
- 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
- 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
- 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
- minor comment:
- a brief description of the nine meta-world tasks (and task length) would provide an overview of task complexity
- 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
Questions
### Presentation of the central claim
- 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)?
- 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?
- 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)
- 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.
### Methodology
- 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?
### Experimental results
- in Figure 6, the average of TemporalOT appears less than 7.5 but is reported to be 8.4 in Table 1.
- 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
- 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
- which $\epsilon$ values were used in the experiments (or are they determined after optimizing equation 7), and how do they affect the task performance?
- the 4th baseline OT needs more explanation, though I can roughly guess how online OT differ from offline OT
### Discrepancy in ASD vs. OT performance
- 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?
### Minor points
- 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$?
- 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
- 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
- nonetheless, the authors discuss the broader impact of their work, including immediate societal implications, and suggest a straightforward solution, which is appreciated