Dear Reviewer dSMV,

Thanks for the constructive comments and suggestions. We also appreciate your acknowledgment of the principledness of our proposed metric, as well as the quality of the presentation.

We also thank you for many insightful questions that promote discussions to further reveal the scientific merits of the proposed metric. Please find our answers and clarifications below in detail.

We hope the answers can help finalize your assessment and the rating of our paper. Please also let us know if you have any further questions that we need to provide additional clarifications.

---

W1: The key states for some tasks are hard to describe using sentences, for example, in the task “hold the tube and align it with the table”, it is hard to describe the key state for “align” if not know the dimension, size, or relative orientation/positions for the tube and table. The proposed method would be more practical if it could identify key states without the state-by-state instruction generated by LLMs.

A: Thanks for the insightful comment. We agree that it would be nice if the agent could automatically discover key states for policy learning. And this is actually the direction we are heading to. In the sense that, manually defined key states (e.g., the one obtained by predefined rules in CoTPC) are too limited in terms of diversity.

For example, some key states are semantically meaningful but hard to be described by a set of predefined rules. We alleviate this constraint by employing LLMs to propose the key states (instead of constructing more rules). We admit that LLMs have their own limitations, which is out of the scope of our current paper. We focus on the grounding in this work but will study the limitation of LLMs in proposing key states in future research.

However, we would also like to point out that there is a benefit of leveraging LLMs for proposing the states, as now the meaning of the key states is aligned with human understanding, which can help promote transparency of the grounded key states and the explainability of the trained policies.

---

W2: Most tasks in the experiment do not really need to understand key states. For example, tasks such as pick and place cubes can be learned with pure behavior cloning (it is a bit surprising the performance reported in this paper was not that good), but it is harder to see the effect of key states grounding with short-horizon pick-and-place experiments.

A: Thanks for the question. Even though Pick & Place Cube seems to be a short-horizon task, our experiments still show that key states help policy learning.

The comparison (on Pick & Place Cube) between the baselines not using key states and CoTPC already demonstrates the effectiveness of the key states. Also, the comparison between CoTPC and our results shows that there are even more key states that are useful even within such a short-horizon P&P Cube task.

For example, in CoPTC, only one key state is defined, which is “grasp.” However, our method employs an LLM to propose many more (semantically meaningful) key states and ground them with the proposed information-theoretic criterion. For example, the key states proposed for Pick & Place Cube in our method are:

{

• KS 1: Robotic arm positioned above the cube.
• KS 2: Robotic arm's gripper aligned with the cube.
• KS 3: Robotic arm has securely grasped the cube.
• KS 4: Robotic arm lifted, holding the cube.
• KS 5: Robotic arm positioned above the goal position.

}

The flexibility of employing LLMs to propose more semantically meaningful key states enables the better performance of our trained policies. It is also evidence that short-horizon tasks can also benefit from fine-grained (but semantically meaningful) key states.

Additionally, we also conducted experiments on the more complex long-horizon task, Franka Kitchen, to verify the effectiveness of key state grounding. The basic neural network architecture is Decision Transformer (DT). The results are shown below. We can find that in this task, DT with key states can achieve superior performance over the one without key states. In the unseen environment, the relative improvement was as high as 73.7%.

---

Dear Reviewer dSMV,

As time passes by, could you take a look at our rebuttal and let us know whether it resolves your concerns? We are also here to answer any further questions you may have.

Best,

The Authors

Dear Reviewer 4ns7,

Thanks for the constructive feedback and suggestions. We appreciate that you found our proposed metric novel and interesting.

We understand that more explanations and comparisons would help your re-evaluation of our paper. Especially, on the intuition of the proposed metric, as well as the comparison with more baselines. Please find our answers and clarifications below in detail.

We hope the answers can help finalize your assessment and the rating of our paper. Please also let us know if you have any further questions that we need to provide additional clarifications.

---

W1: The writing of the paper is not very clear. For example, the majority of the paper describes how we can use MIB to segment trajectories, but there is no description of how this MIB-derived segmentation is actually used in training the policy network. I believe this part should be at least briefly described in the paper. (I found it in the supplementary material, but I would like to mention that reviewers are not subject to reading the full supplementary material).

A: We appreciate your careful review and constructive feedback. We placed the policy network training details in the supplementary material due to page restrictions in the main manuscript. We thought the experimental setup could serve a similar purpose, but apparently, it causes a gap in the information flow. In the revised version, we have added a new subsection 2.5 to elaborate on the training of the policy networks.

---

W2: While I found the introduced metric novel and interesting, I did not fully get its intuition. In particular, is it just the case that the speed of the robot arm will be very low when it's close to a "key state?" Because the uncertainty is essentially defined as s_t | s_{t-dt}

A: We appreciate your observations and would like to clarify that a key state does not necessarily correlate with a very low speed of the robot arm; it's more about the reduction of uncertainty in the latent state space.

An intuitive example is outlined in Appendix A.2. To clarify and avoid any confusion regarding speed, consider another example: We are in the source room aiming to get to the target room, which is connected solely by an open door. Regardless of our path within the source room, accessing the target room necessitates passing through this door. This action of passing through the door symbolizes a key state, which is independent of the speed we traverse it at—we could do so either quickly or slowly. In the semantic space, however, the uncertainty s_t​∣s_{t−dt}​ is significantly reduced.

Moreover, we would like to point out that since mutual information is invariant under scaling, the proposed metric is not affected by the speed, e.g., multiplying the state with a scalar can change the speed but not the mutual information. Thus, the proposed information-theoretic criterion is more than just determining key state by speed.

---

Dear Reviewer 4ns7,

As time passes by, could you take a look at our rebuttal and let us know whether it resolves your concerns? We are also here to answer any further questions you may have.

Best,

The Authors

Dear Reviewer KD76,

Thank you for the constructive comments and suggestions. We also appreciate your acknowledgment of the soundness, presentation, and effectiveness of our work.

We understand that clarifications on the usage of LLMs in our pipeline, the role of paraphrasing with LLMs, and discussion on generalization are needed for your re-evaluation of the proposed key state grounding method. We also provide an experiment on more complex scenes to demonstrate the effectiveness of our algorithm when facing distractions or ambiguities. Please find details below.

We hope the answers can help finalize your assessment and the rating of our paper. Please also let us know if you have any further questions that we need to provide additional clarifications.

---

W1: The LLM lacks visual information or constrains about the environment, which restricts its ability to make geometry-based decisions like collision avoidance.

A: We agree that the LLM lacks visual information to make geometry-based decisions like collision avoidance. However, this is not a critical issue for our method, since we are not relying on LLM to make low-level control decisions.

Our method only leverages an LLM to provide potential key states involved in achieving a task, which does not need accurate geometric information. Our method then focuses on grounding these key states in the demonstrations, which then provides the grounded physical states for better policy learning and generalization. The experiments have shown that by grounding the LLM-generated key states, the trained policies can surpass human-provided key states, evidencing the effectiveness of the proposed grounding method.

Again, thanks for the comment. In the future, we would like to see how the key states provided by an enhanced LLM (e.g., GPT-4V) can help with the learned policy.

---

W2: It is necessary to manually create various paraphrases that are semantically consistent depending on the tasks.

A: It appears there may be a slight misunderstanding. In our experiments, various paraphrases are generated automatically by LLM, eliminating the need for manual modifications. The first stage of the process, instruction generation for getting key state descriptions, relies on minimal task priors such as specified task description. The second stage, paraphrase generation, is solely dependent on the instructions generated in the first stage and does not require any task priors.

The necessity of paraphrases lies in the fact that pretrained visual-language encoders may be subject to domain gaps and insufficient granularity in their training data and thus introduce noise in the compatibility function (between the image and the key state description). Therefore, the primary purpose of paraphrasing the instruction is to create semantically equivalent versions of the key state description and use them to form an averaged compatibility function, so that the noise in the encoders can be canceled out.

This operation is observed to stabilize the training (optimization) process. In our experience, the incorporation of paraphrasing has been observed to speed up the training convergence by 20%.

---

W3: The paper does not address how well the proposed approach can generalize to unseen tasks.

A: In terms of the generalization of the policy on unseen environments, Table 2 in our paper presents the required evaluation. To elaborate, Table 1 represents the performance of our proposed method on seen task configurations, while Table 2 demonstrates its performance on unseen task configurations, showcasing our strong generalization capabilities of the trained policies under key state guidance.

In terms of the generalization of the grounding network, we would like to point out that it is generalizable by nature. The full training procedure of the grounding network requires no human annotation of the ground-truth groundings. In other words, the training objective for the grounding network is unsupervised, which contains only the compatibility function, and the mutual information boost (MIB) for enhancing the physical meaningfulness of the grounding. Thus the pipeline can be trained in any situation without concerning the test-time generalization.

---

Dear Reviewer KD76,

As time passes by, could you take a look at our rebuttal and let us know whether it resolves your concerns? We are also here to answer any further questions you may have.

Best,

The Authors

Dear Reviewer S4jr,

Thanks for your questions and constructive suggestions. We also appreciate your acknowledgment of the novelty and effectiveness of the proposed key state grounding algorithm.

We understand that extra baseline comparisons are needed for your re-evaluation of our paper, and some clarifications are needed on the usage of the LLM to alleviate the misunderstanding. Accordingly, we have provided detailed answers to the questions and comments below.

We hope the answers can help finalize your assessment and the rating of our paper. Please also let us know if you have any further questions that we need to provide additional clarifications.

---

W1.1: The robotic environment/task as described by the authors does not necessarily need key-state grounding. I believe 100 clean expert demonstrations with end-to-end BC-RNN could achieve similar performances.

A: Thanks for the comment. We have performed an additional experiment to provide more information on this point. Specifically, we have experimented with a BC-RNN structure [6] using 500 demonstrations per task for 4 tasks from our environment (listed below).

Both BC-RNN and Decision Transformer (DT) do not employ key states during the training of the policy, while CoTPC and Ours utilize key states to improve the policy training efficiency and generalization. We can observe that the BC-RNN structure performs worse than the DT structure when trained using the same 500 demonstrations. Moreover, by employing key states during policy training, the latter two (CoTPC [1] and Ours) outperform BC-RNN and DT by a big margin. Together these comparisons demonstrate that BC-RNN without key state prediction encounters difficulty in efficiently learning the tasks used in our environment. Also, the comparison between CoPTC and Ours demonstrates that the way of using (grounding) key states proposed in our paper is more effective.

Particularly, we would like to clarify that Peg insertion is a complex, multi-stage, long-horizon task, and we have observed that other methods do not achieve high success rates on such long-horizon tasks. Moreover, our method shows about a 30% improvement over the SOTA baseline (CoTPC), both in seen and unseen environments.

W1.2: Key state grounding is useful for long-horizon tasks. I would recommend environments/tasks like Franka Kitchen used by BeT and by real-world tasks in MimicPlay. I will change my view if I see more complex robotic applications.

A: Thanks for your suggestion. We have further experimented with the mentioned environments/tasks to demonstrate that key states aid in policy learning.

For the Franka Kitchen task, where the success rate is measured in completing a specified sequence of multiple subtasks, the Decision Transformers (DT) with and without key states are compared.

We can observe that the DT with key states demonstrates competitive performance in seen environments, showing better learning efficiency of the policy when supplied with the key states. Moreover, we can observe the DT trained with key states performs much better than the DT w/o key states in unseen environments, achieving a relative 73.7% improvement. This indicates that key state guidance has a substantial positive impact on the generalization of the learned policy for complex long-horizon tasks. The results are listed below.

We have also examined the MimicPlay dataset and attempted to perform similar experiments. However, we found that MimicPlay focuses on transferring human policy to robot policy with a large number of human demonstrations, which is currently out of our scope. But we would like to thank the reviewer for pointing this out, which can actually be an interesting future research direction for leveraging key state grounding for policy transfer.

---

Dear Reviewer S4jr,

As time passes by, could you take a look at our rebuttal and let us know whether it resolves your concerns? We are also here to answer any further questions you may have.

Best,

The Authors