KOI: Accelerating Online Imitation Learning via Hybrid Key-state Guidance

Authors: jingxian  lu
wenke  xia
dong  wang
zhigang  wang
bin  zhao
di  hu
xuelong  li
Venue: robot-learning_org_CoRL_2024_Conference
Published: 2024-06-06 20:36:51

Abstract: Online Imitation Learning methods struggle with the gap between extensive online exploration space and limited expert trajectories, which hinder efficient exploration due to inaccurate task-aware reward estimation.
    Inspired by the findings from cognitive neuroscience that task decomposition could facilitate cognitive processing for efficient learning, we hypothesize that an agent could estimate precise task-aware imitation rewards for efficient online exploration by decomposing the target task into the objectives of "what to do" and the mechanisms of "how to do".
    In this work, we introduce the hybrid Key-state guided Online Imitation (KOI) learning approach, which leverages the integration of semantic and motion key states as guidance for task-aware reward estimation.
    Initially, we utilize the visual-language models to segment the expert trajectory into semantic key states, indicating the objectives of "what to do". 
    Within the intervals between semantic key states, optical flow is employed to capture motion key states to understand the process of "how to do".
    By integrating a thorough grasp of both semantic and motion key states, we refine the trajectory-matching reward computation, encouraging task-aware exploration for efficient online imitation learning.
    Our experiment results prove that our method is more sample efficient than previous state-of-the-art approaches in the Meta-World and LIBERO environments. We also conduct real-world robotic manipulation experiments to validate the efficacy of our method, demonstrating the practical applicability of our KOI method.

Main paper: 
 introduction: 
 Imitation Learning has achieved significant success in various robotic manipulation tasks through offline and online modes [1, 2, 3, 4] . The offline approach [5, 6, 7, 8, 9] seeks to learn policy by mimicking state-action pairs from expert demonstrations through supervised learning. Due to its independence from interaction with environments, offline imitation learning necessitates extensive robotic data, which requires expensive costs associated with real-world data acquisition [10, 11, 12] . In contrast, online imitation learning [13, 14] , driven by a reward function informed by expert demonstrations, is able to formulate policies via exploration. Although it has the potential to refine policies with minimal expert input, the considerable gap between substantial online exploration space and limited expert trajectories challenges the estimation of exploration reward, which hinders the efficiency of online imitation learning. To estimate fine-grained reward, Optimal Transport (OT) [15, 16, 17, 18] is proposed to facilitate trajectory-matching reward by measuring the distance between exploration trajectories and expert demonstrations. However, their insufficient grasp of the semantic comprehension of tasks makes them easily converge on sub-optimal solutions [15] .Pivotal research in cognitive neuroscience [19, 20] suggests that humans could enhance cognitive processing and learning efficiency via task decomposition. Inspired by these findings, we hypothesize that agents could optimize online imitation learning efficiency by task decomposition. Specifically, by gaining a holistic grasp of both the objectives ("what to do") and the mechanisms ("how to do") of the task, an agent could be encouraged to conduct task-aware exploration.To this end, we propose the hybrid Key-state guided Online Imitation (KOI) learning method, which effectively extracts the semantic and motion key states from expert trajectory to refine the reward estimation. We initially utilize the rich world knowledge of visual-language models to extract semantic key states from expert trajectory, clarifying the objectives of "what to do". Within intervals between semantic key states, optical flow is employed to identify essential motion key states to comprehend the dynamic transition to the subsequent semantic key state, indicating "how to do" the target task. By integrating both types of key states, we adjust the importance weight of expert trajectory states in OT-based reward estimation to empower efficient online imitation learning.To validate the efficiency of our KOI method, we conducted comprehensive experiments across 6 Meta-World manipulation tasks and 3 long-horizon LIBERO tasks. The analysis of experiments highlights the strength of our semantic and motion key state guidance, which enhances the sample efficiency of KOI compared to previous methods. In the end, real-world experiments are conducted to prove that our method could also generalize to real complex scenarios. 
related work: 
 Online Imitation Learning. Online Imitation Learning seeks to develop policies through online exploration, guided by a reward function informed by expert demonstrations [21, 22, 23, 24] . Recognizing the difficulty of obtaining state information in real-world settings, some researchers [25, 26] explored the use of visual and language cues to inform task-aware reward estimation. However, these works mainly focused on task semantic knowledge, overlooking detailed trajectory information. In this study, we adopt the trajectory-matching reward estimation method, integrating task semantic comprehension with fine-grained trajectory motion for efficient online exploration.Optimal Transport for Imitation Learning. Optimal Transport (OT) is a computational framework that addresses the problem of finding the most efficient way to transform one distribution into another, minimizing a predefined cost for this transformation [27, 28, 29] . Prior works borrow various distance metrics to measure the distance between the online exploration trajectories and expert demonstration for fine-grained imitation reward estimation [17, 18, 30, 15] . However, they regrettably ignore the semantic information of manipulation tasks, leading to convergence on suboptimal solutions [15] . In this work, we propose to combine the fine-grained trajectory-matching reward with semantic and motion key states, facilitating the efficiency of online imitation learning. 
method: 
 During online imitation learning, a significant challenge lies in the gap between extensive online exploration space and limited expert trajectories, resulting in inefficient exploration due to imprecise imitation reward estimations. In this work, to precisely estimate task-aware reward, we extract the hybrid key states from semantic and motion aspects by the Semantic Decomposition Module (SDM) and Motion Capture Module (MCM), indicating "what to do" and "how to do" for the target task. By integrating both types of key states, we adjust the importance weights of expert trajectory within the OT-based reward estimation method, enhancing efficiency of online imitation learning.The goal of online imitation learning is to train a policy π e through exploration with the reward r e informed by the expert demonstrations T d . The i-th state from exploration trajectory T e is denoted as s e i , and the j-th state from expert demonstration T d is denoted as s d j . The similarity between them We first extract semantic key states with the Semantic Decomposition Module. Within intervals between semantic key states, the Motion Capture Module is proposed to identify motion key states. Further, we adjust the importance weight in OT-based reward estimation with these hybrid key states, to enable task-aware exploration for efficient online imitation learning. The color intensity of OT matrix represents the value of estimated reward.is measured by function C, which is defined as cosine similarity in this paper. Accordingly, the cost of transporting from source s e i to destination s d j is calculated as follows:EQUATIONGiven the cost above, methods based on Optimal Transport (OT) [15, 16] aim to find the optimal transport matrix η * to estimate reward r e of each state. For example, the reward of exploration state at timestep i is estimated using the subsequent formula:EQUATIONwhere η ∈ N N e ×N d and each element η ij represents the amount of data transported from i to j. N e and N d are lengths of exploration trajectory T e and expert demonstration T d , respectively. The optimal transport matrix η * is obtained by measuring the distance between importance distribution µ of exploration trajectory T e and importance distribution ν of expert demonstration T d , with Sinkhorn algorithm [27] employing an entropy regularization H to enhance computational efficiency:EQUATIONHowever, previous OT-based methods assume the distribution of expert demonstration as a uniform distribution for each state, i.e., ν = 1 N d 1. While this assumption facilitates fine-grained rewards, it overlooks the importance among states. Such oversight hampers exploration efficiency for lack of task-aware state comprehension, as revealed in the qualitative analysis of Section 4.4.To integrate trajectory-matching reward with task-aware information, we extract the semantic key states, highlighting the objectives of "what to do", and identify the motion key states, revealing the mechanisms of "how to do'. With the guidance of both types of key states, we adjust the importance distribution of expert trajectory ν for accurate reward estimation as described in Section 3.4, which encourages the agent to conduct task-aware exploration to facilitate online imitation learning.Foundation Models have demonstrated their rich world knowledge and strong generalization capabilities in various visual-language tasks [31, 32] . In this work, we propose utilizing GPT-4v to design the Semantic Decomposition Module (SDM) for task-specific subgoal extraction.Concretely, we initially instruct GPT-4 to decompose the task description L t into K subgoals, denoted as L s = {l s 1 , l s 2 , .., l s K }. Subsequently, GPT-4v is employed to locate corresponding semantic key states of the expert demonstration. To manage the number of input images, we sample states every T step from expert observations, forming the query sets O. To ensure the temporal consistency of located key states, both sampled expert observation O and subgoal language descriptions L s are fed into GPT-4v to get the semantic key state index set I s :EQUATIONConsequently, SDM produces a temporally coherent sequence of indices I s that precisely match the semantic key states to subgoal descriptions, clearly defining "what to do" at each stage of the task. Due to the length constraints of the paper, please refer to Appendix A for detailed implementation.While the SDM effectively recognizes semantic key states, they primarily outline the task objectives, without delving into manipulation details required for transition to these target states. To enhance the grasp of mechanisms of "how to do" the task, we propose the Motion Capture Module (MCM) to identify the motion key states within the interval defined by two semantic key states.Specifically, we measure the intensity of motion between two consecutive states by the absolute value of optical flow of their corresponding observations, which is captured by the Farneback [33] optical flow estimation algorithm F. The p-th motion key state represents the critical transition from semantic key state s I s p-1 to s I s p , whose index I m p ∈ (I s p-1 , I s p ) is identified by following formula:EQUATIONBy identifying the most intense optical flow, we select the related states as the motion key states, with index set I m , which indicate "how to do" the target task.Based on the extracted hybrid semantic key states and motion key states, we develop the Importance Weight Adaption method to guide the agent in online exploration. Previous studies [15] typically assumed that the importance weights assigned to each state in expert trajectory ν in Equation 3 are equal, namely, ν = 1 N d 1. While it aids trajectory-matching, the estimated reward under this assumption is disarmed of task awareness. In this module, we enhance task-aware reward by emphasizing the guidance of semantic and motion key states. Thus, we design the importance weight of expert demonstration, which was decomposed into K subgoals, as a Gaussian Mixture Distribution D:EQUATIONwhere we adjust the semantic variance σ s to be smaller than the motion variance σ m , guiding the agent to pay more attention to semantic information. The weights of these two types of key states A s i and A m i are introduced to ensure the dominance of task completion status when reward estimation. The choice of hyperparameters is detailed in Appendix B.2 Additionally, we replace the uniform distribution assumption of expert demonstrations ν, with the normalized distribution D, enhancing task-aware exploration during online imitation learning.To fully utilize the expert demonstrations, we employ an offline-to-online methodology, initially training a behavior cloning policy π b using the state-action pair (s d , a d ) in expert trajectory T d :EQUATIONInitialized our exploration policy π e as the behavior cloning policy π b , the agent could explore the environments with task-aware knowledge and measure the similarity between exploration trajectory and expert demonstrations. Further, we utilize the n-step Deep Deterministic Policy Gradient (DDPG) [34] with our estimated exploration reward r e to update our policy π e , and Q-function Q θ :EQUATION)In this step, we employ the adaptive function λ following the Regularized Optimal Transport (ROT) [15] to effectively balance the weight between BC regularization and online exploration:EQUATIONHowever, the limited expert demonstrations commonly fail to offer task-aware reward r e for extensive online exploration space. The KOI method we proposed release this issue via the guidance of hybrid semantic key states and motion key states, accelerating the online imitation learning. 
experiments setting: 
 To comprehensively evaluate our key-state guided method, we initially conducted experiments across 6 tasks from the Meta-World suite [35] to validate the efficacy of our method. Further, we assess KOI's performance in more complex manipulation tasks, we conduct 3 tasks from the LIBERO suite [36] -notable for its complex scenarios, large action spaces, and long task sequences.In offline imitation phase, we collected 1 demonstration and 50 demonstrations per task in the Meta-World suite and LIBERO suite to train BC policy, respectively. For online imitation phase, we loaded pretrained BC policy and initialized critic network. Details of settings are attached at Appendix B.To evaluate our methods, we conducted comparison experiments with various methods, all methods are initialized with pretrained behavior cloning policy π b to ensure a fair comparison.:• BC represents the pretrained behavior cloning policy π b .• UVD [37] utilize pretrained visual encoder R3M [38] for task decomposition in image-goal reinforcement learning.• RoboCLIP [25] leverages the pretrained video encoder S3D [39] to estimate the task-aware reward function instead of the reward of each state.• ROT [15] employs Optimal Transport to estimate the distance between exploration trajectory and expert demonstrations, which focus on fine-grained trajectory matching.As the results presented in Figure 2 shown, the UVD method, which emphasizes semantic subgoals rather than the mechanisms of "how to do", often fails to accomplish the target task. Similarly, RoboCLIP leverages a video encoder to capture trajectory motion, but the absence of per-state reward hinders its ability to guide online imitation. Although ROT employs fine-grained reward and shows commendable performance, it sadly neglects task-aware reward, thereby impeding efficient exploration. This issue is highlighted in the qualitative analysis detailed in Section 4.4. In contrast, our KOI method achieves more efficient online imitation learning in simple Meta-World tasks such as "drawer open". For the challenging "bin picking" task, characterized by multiple manipulation subgoals, only the fine-grained reward estimation methods (our method and ROT) have achieved success. Furthermore, our method outperforms ROT by leveraging a thorough comprehension of both semantic and motion key states, providing a distinct advantage.The tasks in the LIBERO suite present increased difficulties, due to longer task sequences and larger action spaces. The results demonstrate that KOI not only yields more efficient learning but also maintains greater stability, benefiting from the comprehensive guidance from the hybrid key states. We further discuss the trend of variance over time. As shown in Figure 3 , the variance of KOI decreases continuously during exploration while the performance increases, which indicates that our method can enhance model performance across a diverse range of environments, thereby demonstrating its robustness and stability.Our method proposes the Semantic Decomposition Module (SDM) and Motion Capture Module (MCM) to provide task-aware knowledge for trajectory-matching reward estimation. To evaluate the effectiveness of each module, we conducted ablation experiment with the following three settings:• w/o. Both: This setting selects the state every 10 timesteps as the key state and utilizes them to adjust the importance weight of expert demonstrations.• w/o. MCM: This setting only uses states selected by visual-language models to adjust the importance weight of expert demonstrations.• w/o. SDM: This setting removes the semantic subgoals inferred by visual-language models and selects motion key states within every 10-state intervals.As shown in Figure 4 , the absence of either semantic or motion key states detrimentally affects the efficiency of online exploration, particularly in the complex "bin picking" manipulation task. No- Besides, the performance of the "w/o. MCM" setting is comparable to the "w/o. SDM" setting, demonstrating the equal importance of both types of key states. Additionally, the shaded region for the "w/o. MCM" setting is larger than the other methods, indicating that motion key states contribute to stabilizing online exploration by providing detailed guidance on "how to do" tasks. We visualize the semantic and motion key states selected by our method. As shown in Figure 5 (a), our semantic key states effectively extract subgoals of the task. Our motion key states, capture the most intensive optical flow between semantic subgoals, allowing agents to grasp motion dynamics.Furthermore, we present three representative cases in the "bin picking" task. As shown in Figure 5(b) , when the agent fails to complete any subgoals, our method estimates a lower reward to discourage unsuccessful exploration. In cases where the agent completes only the first subgoal of Figure 5 (c), the estimated reward is high during the initial stage but lower at the final stage. In contrast, ROT incorrectly estimates a high reward for the final stage. The last case, shown in Figure 5(d) , demonstrates that our method provides a high reward estimation when all subgoals of this task are completed, thereby facilitating effective online exploration.As shown in Table 1 , We conduct real-world experiments using an XARM robot equipped with a Robotiq gripper. To comprehensively evaluate the performance, the entire imitation learning is divided into three stages: offline, online, and test. For each task, we collect 10 human expert demonstrations to train the Behavior Cloning (BC) policy during the offline phase and use only 1 expert demonstration to estimate reward during the online phase. All the evaluations are conducted at varying initial object positions. Implement details can be found in Appendix C.We conduct comparative experiments between our method and previous trajectory-matching method [15] , reporting the success rates of each phase in Table 1 . As the results show, the ROT method commonly struggle to provide effective reward estimation for unseen initial positions. However, our method utilizes task-aware information to improve the reward estimation, thereby enhancing the agent's exploration efficiency and performance. The differences are most significant in the task "put the cube on the scale", which requires the highest precision in manipulation, making it difficult to complete the task by simply matching trajectories with expert demonstrations. Table 1 : The demonstrations and comparative results of our 3 real-world robotic manipulation tasks. The initial state of each task is represented in a lighter shade, while the completed state is depicted in a darker shade. The results indicate that our method achieves more efficient exploration and better overall performance. 
conclusion and limitation: 
 In this work, we propose the hybrid Key-state guided Online Imitation (KOI) learning method, which extracts the semantic and motion key states for trajectory-matching reward estimation for online imitation learning. By decomposing the target task into the objectives of "what to do" and the mechanisms of "how to do", we refine the trajectory-matching reward estimation to encourage taskaware exploration for efficient online imitation learning. The results from simulation environments and real-world scenarios prove the efficiency of our method.Limitations. The online imitation learning in this study initializes an critic, as shown in our results, which negatively impacts early exploration. In future work, offline reinforcement learning could be employed to pretrain both the actor and critic, potentially enhancing initial online learning.During imitation learning phase, each task was trained using a batch size of 256 over 50k epochs. For the online imitation learning phase, we loaded the pretrained BC policy and initialized the critic network. The model then interacted with the environment over 500k timesteps in the Meta-World suite and 200k timesteps in the LIBERO suite, allowing us to observe the efficiency of different online imitation learning methods. For a fair comparison to ROT [15] , we train the policy using a stack of 3 consecutive RGB frames in Meta-World suite, and each action in the environment is repeated 2 times. 
 
 Appendix: Image size 84 × 84Action shape 4Frame stack 3Action repeat 2Task finish reward 0 Demonstration(s) for BC 1 Demonstration(s) for online finetuning 1Image size 128 × 128Frame stack 1Action repeat 1Task finish reward 5Demonstration(s) for BC 50 Demonstration(s) for online finetuning 5Table 2 : Details of environment settings.The complete list of hyperparameters is provided in Table 3 . As shown, all the methods differ only in reward estimation, i.e., components like encoder and RL backbone are the same. When modeling the importance weight of expert demonstration, we apply distinct weights and standard deviations for semantic and motion key states, respectively. The weight assigned to the ultimate goal of the task is set higher than that of the subgoals, thereby enhancing the agent's awareness of task completion.In addition to the results provided in Section 4.2, Figure 7 shows the performance of KOI on all tasks -6 from the Meta-World suite [35] and 3 from the LIBERO suite [36] . On all tasks, KOI has shown its sample efficiency compared to other reward estimation methods [15, 37, 25] . However, the "door unlock" task, as shown in Figure 12 , requires the robot arm to rotate the handle to unlock the door. However, during online exploration, the agent could complete this task by using its body instead of its gripper to rotate the handle, due to the imprecise physical simulation. The tricky policy can complete the task without imitating the expert's trajectory, resulting in the degradation of performance during online finetuning with any reward estimation method as compared to the pretrained policy.We provide the ablation study of the Semantic Decomposition Module, which select the semantic key states with different methods:• Human Annotation: We manually annotate the key states in the expert demonstration.• UVD R3M: Utilizing the pretrained R3M visual model as a visual encoder for task decomposition, we select semantic key states to adjust the importance weights in the expert demonstration.• UVD trained: We utilize the visual encoder from the Behavior Cloning model, which is trained on the expert demonstrations, to select semantic key states to adjust the importance weights of the expert demonstration.The results in Figure 8 demonstrate that our model, designed based on GPT-4v, can fully leverage the comprehension capabilities of large multimodal models, thereby achieving better decomposition effects than pre-trained models and trained sequential models. Additionally, our model achieves results comparable to manually selected key frames, while reducing labor costs and enabling easier scalability in future processes.Ours UVD_R3M We provide the ablation study of the Motion Capture Module (MCM), which selects the motion key states with different methods:• w/o. MCM: This setting only uses states selected by visual-language models to adjust the importance weight of expert demonstrations.• uniform Motion Capture: This setting selects the motion key states every 5 steps within the interval defined by two semantic key states.As shown in Figure 9 , the uniform motion selection method would choose noisy keyframes, affecting the reward estimation. However, our motion capture module can extract useful information to facilitate online exploration.uniform motion Ours w/o. motion We also provide the ablation study of the Both Module:• w/o. Both: This setting selects the state every 10 timesteps as the key state and utilizes them to adjust the importance weight of expert demonstrations.• w/o. MCM: This setting only uses states selected by visual-language models to adjust the importance weight of expert demonstrations.• w/o. SDM: This setting removes the semantic subgoals inferred by visual-language models and selects motion key states within every 10-state intervals.As shown in Figure 10 , the absence of either semantic or motion key states detrimentally affects the efficiency of online exploration, particularly in the complex "bin picking" manipulation task.Besides, the performance of the "w/o. MCM" setting is comparable to the "w/o. SDM" setting, demonstrating the equal importance of both types of key states.Additionally, the shaded region for the "w/o. MCM" setting is larger than the other methods, indicating that motion key states contribute to stabilizing online exploration by providing detailed guidance on "how to do" tasks.Notably, both the Semantic Decomposition Module and Motion Capture Module we proposed surpass "w/o. Both" setting in most tasks, demonstrating the effectiveness of each module.In this work, we conduct 3 real-world robotic manipulation tasks, as shown in Figure 1 . We conduct our real-world experiments using an XARM robot equipped with a Robotiq gripper. To facilitate efficient policy learning, we limit the robot's action space to 4 dimensions (x, y, z, gripper) and utilize both image and proprioceptive inputs for the policy. To ensure the safety of the exploration, we restrict the robot arm's gripper to operate within a certain safety range.For each task, we collect 10 human expert demonstrations to train the Behavior Cloning (BC) policy during the offline phase and use a single expert demonstration to estimate online imitation reward during the online phase. To evaluate the pretrained BC policy, we conduct 10 experiments at varying initial object positions to calculate the average success rate. We then conduct 20 times of online explorations, recording their average success rates. Finally, we load the updated policy and conduct 10 offline tests to assess its performance.Offline: For image data, we resize them to 120 × 160 and extract 128-dimensional representations through a shallow CNN network. For proprioceptive data, we collect the 6D pose information of the robot arm's end-effector, the 7-axis degrees of freedom joint information, and the gripper's angle. We concatenate the proprioceptive and then map them to the 128-dimensional representation through a fully connected layer. We train the BC policy π bc with 10 expert demonstrations over 200 epochs and evaluate the performance at 10 various initial object positions.Online: We load the pretrained policy π bc and interact with environments within 20 times during online exploration. To ensure the safety of the robot, we restrict the operation range of the gripper to prevent it from colliding with the table . As mentioned in the limitation, the mismatch between pretrained BC policy and initial critic would negatively impact early exploration learning. Consequently, to ensure both safety and efficiency during real-world exploration, we substitute the adaptive function λ in Equation 8 with a constant value of 0.9 and provide a sparse reward to indicate whether the task has been completed. Through this modification, we enable the robot to explore the environment stably and develop the exploration policy π e .To ensure the fairness of the experiments, we also set the adaptive weights of ROT in the comparative experiments as the constant value of 0.9.Test: We load the updated exploration policy π e to evaluate the performance at 10 various initial object positions.  


Decision: The presented approach to improve online imitation learning is very interesting and generally well presented. However, some design choices could be better motivated and discussed in more detail. The reviewers highlighted several issues that should be addressed to further improve the paper. For instance, additional clarification regarding the large variance in KOI’s performance would improve the discussion of the evaluations. Furthermore, the real-world experiments are not described very well and raise some concerns regarding the claims made. The paper also lacks comparison to other hierarchical or subgoal-based methods.

## Post Rebuttal
The rebuttal addressed all major concerns and further improved the already strong work. The reviewers unanimously agree that this work should be presented at CoRL 2024.

Reviews: review: title: title: Review of submission 336
paper_summary: paper_summary: This paper introduces a hybrid Key-state guided Online Imitation (KOI) learning approach to improve the efficiency of online imitation learning. The authors claim that this addresses the challenge of limited expert trajectories in extensive exploration spaces. The KOI method decomposes tasks into "what to do" and "how to do" aspects, utilizing visual-language models to identify semantic key states for the former and optical flow analysis to determine motion key states for the latter. This dual approach enhances task-aware exploration by refining trajectory-matching reward computations. Experiments in Meta-World and LIBERO environments show that KOI significantly boosts sample efficiency compared to previous methods, and its effectiveness is further validated in real-world robotic manipulation tasks.
main_review: main_review: # Clarity

The paper is generally well-written and easy to follow. The authors clearly articulate the motivation behind their work and provide a comprehensive overview of the problem they are addressing. The explanation of the KOI method, including the use of visual-language models and optical flow for extracting semantic and motion key states, is particularly well-structured. The figures and diagrams included in the paper are helpful in illustrating the proposed methodology and results. Overall, the clarity of the paper makes it accessible to a broad audience, including those who may not be experts in the field. Minor correction: In line 101 expert demonstration should be T^d and not T^e.

# Strengths

The paper presents an interesting and promising approach to improving online imitation learning by integrating visual-language models to extract semantically relevant information and combining this with optical flow to understand the necessary actions for task completion. This hybrid approach, which decomposes tasks into "what to do" and "how to do" components, is an interesting idea that leverages recent advances in multi-modal language models. The method shows promise in enhancing task-aware exploration and has potential applications in various domains requiring efficient learning from limited expert demonstrations. 

# Weaknesses

The primary weakness of the paper lies in the evaluations provided by the authors. In Figure 2, there appears to be a significant amount of variance in the results. In several cases, the lower bound for KOI is lower than the average performance of baseline methods, particularly ROT and RoboCLIP. This raises questions about the consistency and reliability of the KOI method's performance. The authors should address the reasoning behind this high variability. Is it due to having only three trials? If so, I recommend conducting more trials to obtain more robust and reliable bounds for their experimental results. Additionally, the lack of a comparative study with the real-world robot arm is disappointing. The authors claim their method works in real-world settings but provide only a figure with the experimental setup and no detailed results or comparisons. This omission is particularly concerning given the variability observed in the simulation results. I recommend the authors perform a full comparative study using the same baselines with the real robot and include these findings in the paper to strengthen their claims.
questions: questions: My primary issue is with the evaluations and I would recommend the authors perform more trials in sim to obtain more reasonable variance bounds and also perform a comparative study with the real robot to strengthen their claims.
limitations: limitations: 3
review_summary: review_summary: I believe the authors' idea is interesting. However, the experiments have to be more extensive to provide clear results.
title: title: Paper applying abstraction to imitation learning
paper_summary: paper_summary: This paper presents a novel approach called hybrid Key-state guided Online Imitation (KOI) learning to improve online imitation learning. The method addresses the challenge of efficiently exploring large state spaces with limited expert demonstrations by decomposing tasks into "what to do" (objectives) and "how to do" (mechanisms).
main_review: main_review: ## Summary

This paper introduces the hybrid Key-state guided Online Imitation (KOI) learning approach, which aims to improve online imitation learning by addressing the gap between extensive online exploration space and limited expert trajectories. The authors propose decomposing tasks into "what to do" (objectives) and "how to do" (mechanisms) to estimate more precise task-aware imitation rewards for efficient online exploration.
The KOI method consists of three main components:

* Semantic Decomposition Module (SDM): Uses GPT-4v to decompose task descriptions into subgoals and identify corresponding key states.
* Motion Capture Module (MCM): Employs optical flow to capture motion key states between semantic key states.
* Importance Weight Adaptation: Adjusts the importance weights of expert demonstration states using a Gaussian Mixture Distribution based on the extracted key states.

The authors evaluate their method on 6 Meta-world tasks and 3 LIBERO tasks, comparing it to baseline methods such as Behavior Cloning (BC), ROT, RoboCLIP, and UVD. They also conduct ablation studies and real-world experiments using an Xarm robot.

## Strengths

* Novelty: the integration of semantic and motion key states for guiding online imitation learning is an interesting idea, while abstraction is not new, I think this approach is sufficiently novel

* Evaluation: the authors do a good job of evaluating their method across a range of simulated tasks and demonstrating transfer to a real robot system

* Qualitative analysis: I appreciated the author's qualitative discussion of the method and intuition for what is learned

## Weaknesses

* Unclear justification for design choices: Several key design choices are not well-motivated or evaluated. In particular, the way that motion key states are determined was not motivated or evaluated. I also had trouble understanding the motivation behind the method for importance weight adaptation

* Limited real world evaluation: while there is some evaluation on a real robot, it is a bit limited. It's not clear that a reader could fully reproduce the experiment. Similarly, additional tasks on the real robot would help to argue for the practicality of the method.

* Lack of comparison to other hierarchical or subgoal-based methods: for a paper like this, I expect the authors to experiment with and compare to other imitation learning methods that leverage abstraction/hierarchy. 

## Typos and clarity issues

* Line 132: "To extract the important motion key states index set I_m = {I_m_1, .., I_m_j, .., I_m_K}, we compute the optical flow between successive expert observation o_j and o_j+1, selecting the most intense state between two semantic key states, which represents a critical state transition during the manipulation process:" --- The sentence is long and could be split for clarity.

* Line 122: "To manage the number of input images, we select state every T step to form the query expert observation sets O." --- Should be "we select states every T steps to form..."

* Equation 5: The equation needs more explanation and context.

* Line 206: "w/o. SDM: This setting removes the semantic subgoals inferred by visual-language models and selects motion key states within every 10 states interval." Should be "10-state interval" or "10-state intervals"

Update after rebuttal: the authors' response largely addresses my concerns; I have updated my score accordingly.
questions: questions: It seems like the semantic key states are identified entirely from the language model. How would your approach work in a task where there are multiple plausible ways to accomplish it, but the expert demonstrations select one? Would you still learn effectively in this situation?
limitations: limitations: 2
review_summary: review_summary: An interesting idea with room to improve on presentation and evaluation of the method
title: title: Simple but effective improvement over an existing baseline on Online Imitation Learning
paper_summary: paper_summary: The paper proposes a simple modification to an exisiting online Imitation Learning algorithm, by incorporating task decomposition and subsequent estimation of task aware rewards using Optimal Transport. The task decomposition is achieved by leveraging exisitng vision language models and optical flow etimation techniques. The experiments show how the proposed modification with task decomposition as an inductive bias for learning give superior quantitative and qualitative results over valid baselines.
main_review: main_review: These are the strengths and weakness of the paper in my opinion.

Strengths
1. Simple but intuitive idea of utilizing task decomposition of expert trajectories for better reward estimation.
2. Clever use of recent advances in sequential models (vision-language models) to automate this task decomposition task.
3. Convincing performance improvements over valid baselines on Online Imitation Learning.
4. The method's superiority is further emphasized with qualitative and empirical experiments.

Weakness
1. Incremental advancement over an existing algorithm.
2. The algorithm is limited by the generalization capability of vision-language foundation models. There might be many instances beyond the said "naturally" hierarchical manipulation tasks, where the model might fail to decompose the task effectively. 
3. I doubt if the task decomposition achieved by a foundation model on the expert trajectories would outperform a human expert annotating the trajectories with key states. If the expert trajectories can be collected by an expert, it can as well be labelled into subtasks by the same expert or another human.
questions: questions: 1. Is the use of pre-trained vision-language model the only way to achieve this task decomposition. Can a sequential model be explicitly trained to decompose the task in case the pre-trained model fails at this task?

2. What would happen if the tasks are not inherently hierarchical? Let's say the simpler tasks like a cheetah run from dmc, as performed in the original ROT paper?

3. Can't we let a human do the task decomposition instead of the GPT model?
limitations: limitations: 2
review_summary: review_summary: I recommend for a Weak Accept with a lower confidence.
title: title: key insight is to facilitate reward estimation for online imitation learning using task relevant key state information - particularly when there are limited demonstrations available
paper_summary: paper_summary: The paper presents a novel approach called Hybrid Key-state guided Online Imitation (KOI) learning to enhance the efficiency of Online Imitation Learning (OIL). The main idea is to address the challenges of task-aware reward estimation by decomposing tasks into "what to do" (semantic key states) and "how to do" (motion key states). By integrating these key states, the method refines reward computation, encouraging more efficient online exploration. The KOI method leverages visual-language models and optical flow for key state extraction, which is validated through experiments in simulated environments (Meta-World and LIBERO) and real-world robotic manipulation tasks, demonstrating improved sample efficiency and practical applicability.
main_review: main_review: Quality:
The paper is well-structured and presents a clear methodology for improving reward estimation from limited demonstrations for online imitation learning. 

Clarity:
The paper is generally clear and easy to follow. The authors provide a detailed explanation of the problem, the proposed solution, and the results. The use of diagrams and qualitative analyses aids in understanding the methodology and its benefits. However, some sections, such as the mathematical formulations, could benefit from additional explanations to enhance clarity for readers less familiar with the concepts. There are some outputs that are not clearly identified eg what is the representation of I^s and I^m  - please clearly define all the terms, their dimensionality and how they are used to derive the reward. I can see that you utilise these terms to define a normal distribution but and better justification on the choice of sigma here is required. How would this extend to new tasks the choice of these variables - how would one estimate them for a new task? 

Why does UVD fail across all tasks? Why not use it to identify the keypoint states instead if GPT4? - seems like a more relevant use of this method than direct comparison to online imitation learning which it was not designed for.

In figure 3 it would be good to see the graph without both systems as well.

Originality:
The approach of integrating semantic and motion key states for reward estimation in online imitation learning is novel. The use of visual-language models to decompose tasks and optical flow to capture motion dynamics adds a unique dimension to the research. This combination of techniques is innovative and I do like how the authors combined the use of foundation models in a unique manner to improve reward estimation.

Significance:
I am not full convinced on the significance of the results. The major claim is it improves sample efficiency in downstream exploration however based on the key results it seems like ROT and the KOI method are converging at similar times across the majority of the experiments. While KOI does get slightly higher reward I am not convinced why this is a result of the task-dependent key state information for reward estimation - an explanation of why this is should be included. Furthermore, the results are missing a key component/experimental results demonstrating how the number of demonstrations used for reward estimation affects downstream task performance with the proposed semantic key points and without them. Basically, what is 'limited' expert demonstration quantitatively - and at what point is limited no longer relevant that we don't need this semantic decomposition to facilitate reward estimation? 

Strengths:

Novel Approach: The integration of semantic and motion key states for task decomposition is innovative and effective.
Comprehensive Evaluation: The method is validated through extensive experiments in both simulated and real-world environments.
Practical Applicability: Demonstrates success in real-world robotic tasks, indicating the method's robustness and applicability.
Clear Methodology: The step-by-step explanation and visual aids enhance the understanding of the proposed approach.

Weaknesses:

Clarity of formulation: Some sections could benefit from additional explanations and definitions to improve clarity as discussed above.
Results to demonstrate the key claim of limited demonstrations and how many is 'limited' are missing.
How does the performance of the SDM and MCM impact performance - how robust is this approach if these modules decomposed incorrectly? Will it only work on simple tasks where it is straightforward to identify these key states? 
You empirically chose 10 for T. What impact does this have on the system performance?
There seem to be many hyperparameters that are not ablated to understand their impact on the system.

Refer to the questions raised below.
questions: questions: 1. Can you provide more detailed explanations or examples for the mathematical formulations, especially the OT-based reward estimation and the Importance Weight Adaption?
2. The paper mentions a reliance on pretrained behavior cloning policies. How significant is this dependency, and what would be the potential impact on early exploration efficiency without such pretraining?
3. Could you provide more insights or results regarding the choice of variance parameters \( \sigma_1 \) and \( \sigma_2 \) for the Gaussian Mixture Model? How sensitive is the model's performance to these parameters?
4. Can you give more detailed descriptions of the real-world experiments, including the setup, specific tasks, and any encountered challenges?
5. How does the KOI method scale with more complex tasks? Any additional insights or experiments on generalization capabilities would be valuable.
6. How does the semantic decomposition from GPT-4 impact performance. What is the variation in this process? Would something like UVD be more suitable to extend to more complex tasks where the semantic keypoint starts arent as obvious to a vision model? Explain how reliant/robust your method is to the selection of these keystates and comment comment on how general it is to scale to arbitrary tasks.
7. In figure 3, could you include a graph without both the Semantic Decomposition Module (SDM) and the Motion Capture Module (MCM) to better illustrate their individual contributions?
8. Can you clarify the representation, dimensionality, and usage of \( I^s \) and \( I^m \)? How are these terms specifically utilized to derive the reward, and what is the justification for the choice of sigma values?
9. Can you explain why the KOI method shows slightly higher rewards compared to ROT, and how this result is attributed to task-dependent key state information for reward estimation?
10. Could you provide experimental results demonstrating how the number of demonstrations used for reward estimation affects downstream task performance with and without the proposed semantic keypoints? What is considered 'limited' in terms of expert demonstrations?
limitations: limitations: 1
review_summary: review_summary: I believe this is good idea and I do appreciate the unique use of higher level reasoning to faciliate reward estimation. There are some components that are not clear within the paper to fully convince me of the key contributions and impact of this approach on the online imitation learning setting and how this method can scale/generalise to other tasks.