Paper: InfoGround: Ground Manipulation Concepts with Maximal Information Boost (ICLR-2024)

Reviewer comment (Theory):
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.

Author response:
**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.

Reviewer comment (Experiments):
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.

Author response:
**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%. 
| Franka Kitchen (Success rate) | Seen environment | Unseen environment |
|-------------------------------|------------------|--------------------|
| Without key states            | 78.4             | 13.7               |
| With key states               | 81.8             | 23.8               |

Reviewer comment (Theory):
If the key states have the maximum mutual information boost, why there is a need to set “n” (the margin) to ensure any two key states are not too close to each other? It is possible that the two key states are quite close, for example, for the task “inserting the key to the keyhole, aligning the key with the keyhole, and moving the key into the keyhole are both key states and can be quite close to each other temporally.

Author response:
**A**: Thank you for your insightful observation. In fact, in our experiments, we typically set "n" to 1 or 2, which are relatively small values. The primary significance of introducing "n" is to ensure the chronological correctness of key state grounding, rather than to see it as a relatively large margin to prevent key states from being too close to each other. 

Our experiments have shown that setting "n" to 1 or 2 has minimal variation in the obtained key states and that using these values can stabilize the training process, resulting in faster convergence. In our experience, setting n to 1 or 2 speeds up convergence by about 30%.

Reviewer comment (Writing):
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. This part should be at least briefly described in the paper. Although this information is in the supplementary material, reviewers are not required to read the full supplementary material.

Author response:
**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.

Reviewer comment (Writing):
While the introduced metric is novel and interesting, its intuition is not fully clear. In particular, is it simply the case that the speed of the robot arm will be very low when it is close to a "key state"? This uncertainty is essentially defined as s_t | s_{t-dt}.

Author response:
**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.

Reviewer comment (Theory):
The proposed method seems to have the limitation that each key state can only be visited once in a demonstration trajectory (due to the Chronological assumption). Therefore, it cannot model task descriptions that require loops (e.g., scoop all sugar, possibly multiple rounds, to transfer them from one bowl into another bowl).

Author response:
**A**: Thanks for the insightful question. The proposed grounding method aims at connecting LLM-generated key state descriptions with the physical states, while the understanding of the (repetitive) structure of the task is mainly dependent on the capacity of the LLM itself, which is currently not our focus. 

Actually, our InfoGround method doesn't confine a key state to a single occurrence within a demonstration trajectory. Provided a key state recurs in the trajectory, our method is able to identify it as a key state each time, by leveraging the maximal information boost detailed in Section 2.3. 

However, the proposed task is an interesting scenario to examine the limitation of using LLMs to propose key states useful for achieving a physical manipulation task. 

To proceed with the discussion, we have prompted GPT4 with the following:

Prompt: I am a robotic arm, please describe the step-by-step instructions for scooping all sugar.

GPT:
1. Position: Move the robotic arm above the container holding the sugar. 
2. Extend: Lower the arm until the scoop or shovel attachment is inside the sugar. 
3. Angle: Tilt the scoop to a 45-degree angle to ensure maximum sugar collection. 
4. Scoop: Close the scoop or shovel attachment to collect sugar, making sure it is full. 
5. Lift: Raise the robotic arm, ensuring the scoop remains level to prevent sugar spillage. 
6. Transport: Move the arm to the desired location for dispensing the sugar. 
7. Release: Open the scoop or shovel attachment, allowing the sugar to pour out. 
8. Repeat: Continue steps 1-7 until all sugar has been scooped and transferred.

From the output, we can still see that LLM has a good understanding of the repetitive structure of the task (step 8), and can propose meaningful key state descriptions. This observation should be evidence that the proposed pipeline can still work for the proposed task of scooping all sugar. However, since there are no existing demonstrations for training the grounding, we would like to check this aspect in the future after collecting more data regarding tasks with repetitive patterns.

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

Author response:
**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.

Reviewer comment (Experiments):
It is necessary to manually create various paraphrases that are semantically consistent depending on the tasks.

Author response:
**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%.

Reviewer comment (Experiments):
The paper does not address how well the proposed approach can generalize to unseen tasks.

Author response:
**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.

Reviewer comment (Experiments):
Do you need to train a separate model for each task? Given the simulated environments are relatively simple, how can you justify that the model is not overfitting the environment?

Author response:
**A**: Thanks for the question. It actually reveals a nice point of our grounding framework. It is true that we train a separate grounding network (the KSL-Net) for each task to get the correct groundings, and then use the groundings to train a policy network to perform the task.

With the discussion on the generalization of the policy networks and the KSL-Net in W3, we would like to point out that, instead of generalization, underfitting is actually a problem if we do not use a separate network for different tasks to derive the accurate grounding. 

First, our experiments show that the global minima of the loss function for a task gives us the best grounding result for that task. Thus it is critical to ensure that the grounding loss for each task is fully optimized. However, unifying the optimization of multiple tasks by training a single network can suffer from under-optimization, i.e., the groundings after training for a certain task may not be the global minima, compared to the one obtained by training a separate KSL-Net using the corresponding loss of this task. In a nutshell, we would like to emphasize that overfitting a separate KSL-Net to the grounding objective of a task can ensure global optimality of the grounding and is critical for the quality of the obtained grounding.

Reviewer comment (Experiments):
The environments presented in the paper are clear scenarios, involving only a limited number of objects relevant to specific tasks. I'm curious about how the approach would handle complex scenes with multiple objects and some of them are unrelated to the task, as this might introduce ambiguity.

Author response:
**A**: Thanks for the suggestion. We have conducted experiments on the Franka Kitchen environment, which is a complex robotic task/environment, where a scene involves many other objects that may introduce ambiguity. With the same decision transformer architecture, our method (with key states) has improved from the baseline(without key states) for both seen and unseen environments. Especially in the unseen environment, the relative boost is 73.7%. The results are listed below.
| Franka Kitchen (Success rate) | Seen environment | Unseen environment |
|-------------------------------|------------------|--------------------|
| Without key states            | 78.4             | 13.7               |
| With key states               | 81.8             | 23.8               |

The results demonstrate the importance of leveraging key states in improving the training efficiency and the generalization capability of the policy. Moreover, it shows that our grounding network can generalize to more complex scenarios.

Reviewer comment (Experiments):
Do you still need to do additional planning in your setup to complete the task with key frames?

Author response:
**A**: To ensure a fair comparison, no additional planning is employed. The framework is end-to-end, with the policy network concurrently undertaking action prediction and dynamic key state estimation during the training phase. This streamlined approach facilitates easy comprehension and seamless integration with other imitation learning pipelines.

Reviewer comment (Experiments):
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. Key state grounding is useful for long-horizon tasks. I would recommend environments/tasks like [franka Kitchen](https://robotics.farama.org/envs/franka_kitchen/franka_kitchen/) used by BeT and by real-world tasks in [MimicPlay](https://www.notion.so/CoRL2023-reviews-92631182f55f47979d5f15dbdf7c0f7c?pvs=21). I will change my view if I see more complex robotic applications.

Author response:
**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.
| Algorithm         | P&P Cube | Stack Cube | Turn Faucet | Peg Insertion | Mean |
|-------------------|----------|------------|-------------|---------------|------|
| BC-RNN [6]        | 62.7     | 7.8        | 43.2        | 3.2           | 29.2 |
| Decision Transformer | 65.4  | 13.0       | 39.4        | 5.6           | 30.9 |
| CoTPC             | 75.2     | 58.8       | 56.4        | 52.8          | 60.8 |
| Ours              | 89.4     | 83.2       | 67.6        | 75.4          | 78.9 |


**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.

| Franka Kitchen (Success rate) | Seen environment | Unseen environment |
|-------------------------------|------------------|--------------------|
| Without key states            | 78.4             | 13.7               |
| With key states               | 81.8             | 23.8               |


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.

Reviewer comment (Novelty):
If I understand correctly, the LLM in your setup requires task-specific prompts, as described in Appendix A.9. If the prompt already contains the information needed for the LLM to respond, I don’t see why I should use an LLM. I could directly provide the information to the robot.

Author response:
**W2: If I understand correctly, the LLM in your setup requires task-specific prompts. As described in Appendix A.9. If the prompt already contains the information needed for the LLM to respond, I don't see why I should use an LLM. I could directly provide the information to the robot.**

**A**: We would like to clarify that the LLM in our setup does not require task-specific prompts. In other words, throughout the whole experiment, a single in-context learning example is provided for all tasks to bias the LLM toward describing the key state using more physically relevant language. So we do not need to change the example for a different task.

In Appendix A.9, the prompt consists of a task description (e.g., what this task is, specified with a single sentence), while the other parts are fixed and provide basic contextual information to the robot. These fixed parts do not require any additional modification for different robot tasks; the only part that needs modification is the task description.

Specifying the task description is crucial because, within the same environment, the robot's actions vary depending on different objectives. Recent research [2,3] also validates this point. However, we could improve this part by asking a more powerful foundation model to automatically come up with the tasks to perform in a scene, e.g., GPT-4V, which introduces a visual perception module and could utilize its visual perception capability to formulate an achievable task. This is out of the scope of our current work as we mainly focus on grounding key states related to a task but not proposing tasks in the scene. However, we would like to explore a fully automatic pipeline for proposing tasks by large models in the future.

Reviewer comment (Writing):
Is the Exemplar part of the prompt? For example, on page 19: {Exemplar: For a task that involves picking up a cube and moving it to a goal position…..}

Author response:
**Q1: Is the Exemplar part of the prompt? For example, on page 19: {Exemplar: For a task that involves picking up a cube and moving it to a goal position…..}**

**A**: Yes, the (in-context learning) exemplar is part of the prompt. In-context exemplar is beneficial, providing effective context to bias more accurate results from Large Language Models (LLMs) to enhance instruction performance. Recent research has also indicated the necessity of in-context examples for LLMs [3,4,5]. Meanwhile, it's important to note that the (single) exemplar part remains fixed across different task settings, eliminating the need for manual modifications.

Reviewer comment (Experiments):
Doesn’t the prompt already contain all the information needed for the task? Why do I need to use an LLM? What's the key difference between the exemplar and LLM response? How many attempts did you make to ask ChatGPT to get the answer you wanted?

Author response:
**Q2: Doesn’t the prompt already contain all the information needed for the task? Why do I need to use an LLM? What's the key difference between the exemplar and LLM response? How many attempts did you make to ask ChatGPT to get the answer you wanted?**

**A**: It seems there may be a slight misunderstanding. To clarify, the prompt does not contain all the information needed for the task. It uses a single in-context example to enhance the instruction quality, but this single example is fixed across all the tasks we have trained on, thus it does not give all the information needed. The only task-specific thing in the prompt is the task description (i.e., tell the robot what to do), which we must have since the LLM needs to know the task to propose relevant key state descriptions. Moreover, LLM acts in the role of proposing relevant key states for our pipeline to ground.

The key difference between the exemplar and the LLM response is that the exemplar serves merely as a fixed template for the output across different tasks, whereas LLM responses can vary depending on the task at hand in proposing task-relevant key states. Our experiments showed that it **only requires one attempt** for GPT to generate the answer we wanted.