Paper: Chain-of-Thought Predictive Control (ICLR-2024)

Reviewer comment (Novelty):
I think CoT is not very novel in this context and I don't think the proposed approach can be regarded as CoT. The proposed approach refers to a demonstration split methodology (or a subgoal discovery mechanism in hierarchical RL). Since no language is involved, I don't think this method can be related to CoT.

Author response:
**Since no language is involved, I don't think this method can be related to CoT.**

We name our work with chain-of-though similar to [7], which suggests that CoT generally refers to the key intermediate steps in sequential decision-making. Nevertheless, we are willing to change the name to others (e.g., subskill-augmented or subskill-predictive control) if all of the reviewers think it is necessary (as the reviewer iFzy also raised this point).

Reviewer comment (Novelty):
This paper is not positioned well in the context of hierarchical RL + demonstration, so the novelty is not well-stated. The authors should mention, discuss, or compare the following works [1-5]. Disclaim: I am not an author of any of these works. The idea of learning subgoals by similarity or diversity is not novel.

Author response:
**Not positioned well in the context of hierarchical RL + demonstration regarding subgoal decomposition. See [1-5]**

While the high-level idea of utilizing the hierarchical principles for policy learning is shared with the approaches in the hierarchical RL + demonstration literature (e.g., [1-4]), all of these work relies on online interaction with the environments either by using online RL to learn to utilize the learned (sub)skills in a transfer learning setup (e.g, [1]) or by using online exploration to acquire the (sub)skills in the first place (e.g., [2]). On the contrary, ours is a completely offline imitation learning method. 

Purely offline subskill discovery that can be used for downstream imitation learning is very challenging, and recent work mostly relies on RL to provide extra supervision [6]. (Also see the author's response to reviewer HDpH). Work that is similar to ours in this setup includes CompILE [5] and PC [7]. CompILE adopts a VAE to perform soft-segmentation of the demo trajectories. However, its evaluation was carried out only on a limited set of relatively simple control tasks. We try our best to find [8] as a relevant follow-up work, yet its evaluation only involves grid worlds. On the other hand, [7] performs hierarchical policy learning with extra procedure-level supervision that limits its applicability and is hard to be directly compared with (also see a discussion in our Appendix). Notice that due to the relatively small previous work in this direction (subskill discovery + hierarchical offline policy learning), both [5] and [7] only use non-hierarchical imitation learning methods as baselines. Similarly in our comparison, we include several non-hierarchical strong baselines.

We will include a discussion regarding this in the appendix.

Reviewer comment (Novelty):
Using transformers in policy learning is not novel as well, and I don't think the major goal of this paper is related to architecture design. If the goal is to claim the novelty of CoT control, the authors should try other architectures as well. Since nowadays transformers are very common architectures, I don't think this can be claimed as a major novelty.

Author response:
**Since nowadays transformers are very common architectures, I don't think the architecture design can be claimed as a major novelty.**

Our contributions are two-fold: a novel formulation of observation-agnostic subskill from offline demonstrations based on discovery based on breakpoint detection and a customized Transformer architecture for the subskill-augmented hierarchical imitation learning method.

Our specific architecture design makes our framework suitable for leveraging the hierarchical information discovered in an offline manner. Granted, Transformers are common standard architectures nowadays. We do not aim for a major overhaul of Transformers (neither does existing work such as BeT). Nevertheless, our design to enhance Transformers for hierarchical imitation learning is shown to be effective and important with several ablations studies. Specifically, we add a diagram to the Appendix to better illustrate the differences in the network design in our third ablation study. We add an ablation study in the author's response to reviewer 2mSo in this direction.

Reviewer comment (Experiments):
It's unclear why the authors use a single model to learn CoT and learn the action. The authors should try to compare to architectures like [6].

Author response:
**Why do the authors use a single model to learn CoT and learn the action? Should try to compare to architectures like [9].**

Here we add another ablation study to verify our architecture design. We report unseen SR for Push Chair and 0-shot SR for Peg Insertion. In this study, we show that using two separate models for predicting the CoT and the actions respectively is outperformed by the joint modeling approach. Note that the reviewer mentioned [9], yet it is not directly comparable since the setup of [9] takes target images as inputs rather than trying to predict them as outputs. Intuitively, the auxiliary loss we propose together with the joint learning approach improves sequential feature learning by encouraging the model to extract information predictive of not only the actions but also the hierarchical structure of the task.

|    | CoTPC (sep) | CoTPC |
| ------- | ------- | ------- |
| Peg Insertion | 49.0 | 59.3 |
| Push Chair | 35.7 | 41.0 |

Reviewer comment (Presentation):
The function of the model is not discussed and section 4.2 looks confusing as a result. I think the function and signature of the model should be discussed prior to the model details.

Author response:
**Need better writing for section 4.2 to introduce the details of the model**

We updated Section 4.2 (mostly 4.2.2) to make it clearer (also see the author's response to the reviewer 2mSo).

Reviewer comment (Reproducibility):
The details of the CoT algorithm are unclear. Particularly, these two sentences look very confusing to me.

"therefore, propose to group contiguous actions into segments, using a similarity-based heuristic to find these subskills." How does the group algorithm work? How to discretize continuous actions?

"We then utilize the Pruned Exact Linear Time (PELT) method [38] with cosine similarity as the cost metric to generate the changepoints in a per-trajectory manner." How does this work? I'm not familiar with the PELT method and this should be discussed in detail.

Author response:
**The paper utilizes PELT for breakpoint detection. Might need to explain it in more detail as background knowledge. Should include an introduction to BeT as well.**

We added a brief introduction paragraph to BeT to the appendix. We will add another short intro to the main paper for PELT. In short, PELT is an optimized version of a more classic breakpoint detection method that is based on the minimization of certain costs. A time series can be partitioned into consecutive groups where each group is associated with a cost term similar to the within-cluster distance in k-means. The total cost is the sum of all individual cost terms plus some penalty to control the number of groups. The classic breakpoint detection method performs dynamic programming to find the exact optimal partitioning that minimizes the total cost. PELT is an optimized version of the method.

Reviewer comment (Writing):
The authors assume the readers know the decision transformer and detection transformer in advance, which is not very good.

Author response:
**The paper utilizes PELT for breakpoint detection. Might need to explain it in more detail as background knowledge. Should include an introduction to BeT as well.**

We added a brief introduction paragraph to BeT to the appendix. We will add another short intro to the main paper for PELT. In short, PELT is an optimized version of a more classic breakpoint detection method that is based on the minimization of certain costs. A time series can be partitioned into consecutive groups where each group is associated with a cost term similar to the within-cluster distance in k-means. The total cost is the sum of all individual cost terms plus some penalty to control the number of groups. The classic breakpoint detection method performs dynamic programming to find the exact optimal partitioning that minimizes the total cost. PELT is an optimized version of the method.

Reviewer comment (Theory):
My main concern with the method is their motivation for how they perform sub-task decomposition. Cosine similarity metric seems heuristical, not well motivated, and anecdotal.

Author response:
Our subskill discovery approach by decomposing the action sequences is motivated by the intuition that functionally similar actions that are temporally close should be grouped into the same subskill. The choice of the metric to measure such functional similarity is an open problem. We tried cosine similarity, L2 distance, and a more complicated Hausdorff distance. Among these, we find that the cosine similarity exhibits both simplicity and generalizability. In our experiments, we find it works across different action spaces (e.g., delta position in Moving Maze, delta joint pose in Peg Insertion, and delta joint velocity in Push Chair) with different action statistics based on how they are generated (sampling-based methods such as RRT, heuristics, etc.) and scale well to higher action dimensions (e.g., Moving Maze uses 2-d actions, Push Chair uses 19-d actions with a dual-arm mobile robot). In practice, action spaces of much higher dimensions are relatively uncommon in robotic tasks.

Reviewer comment (Experiments):
Whilst the results are promising, it is unclear whether more principled decompositions would lead to better results: e.g. obtained via bottleneck options [1] or gaussian processes [2]. The paper would benefit from a greater discussion/comparison on this front.

Author response:
Skill or sub-skill discovery purely from offline demonstration sets is very challenging since there is barely any useful supervision. In the related literature, the option-based approach is a popular strategy. For instance, MO2 [1], an offline option learning framework using the bottleneck state principle, is shown to work well on continuous control problems for learning the options. However, this line of work has the shortcoming of relying on good state space representation. It is unclear if MO2 can work well for high-dimensional visual observations (e.g., the Pour task only supports visual observation due to soft-body manipulation). Our action-based approach is observation-agnostic and thus avoids this issue. Moreover, the option approach requires learning good initial and termination conditions for the options, which is a hard problem itself [2]. Usually, it requires further online learning for skill chaining to compensate for the suboptimal options (MO2 still requires online learning after the options learned from offline datasets are fixed).

The alternative, as we propose, is to utilize methods from the breakpoint detection community for time series modeling to perform action segmentation. Breakpoint detection methods can roughly be divided into predictive model-based approaches (e.g., [3]) and optimization-based ones (e.g., [4]). The former requires an underlying predictive model (UPM) where [3] chooses Gaussian Processes. However, it is unclear if it can model high-dimensional and complex control signals as the action sequences can be hard to model in the first place ([3] suggests it might require a large set of sensitive hyper-parameters). The latter does not model the time series directly but finds breakpoints by optimizing a cost function instead. In this case, a good choice of a metric can be critical. 

We will add a discussion to the appendix.

Reviewer comment (Experiments):
It is unclear to me whether their decomposition approach would favor different tasks, with distinct action-space statistics.

Author response:
Our subskill discovery approach by decomposing the action sequences is motivated by the intuition that functionally similar actions that are temporally close should be grouped into the same subskill. The choice of the metric to measure such functional similarity is an open problem. We tried cosine similarity, L2 distance, and a more complicated Hausdorff distance. Among these, we find that the cosine similarity exhibits both simplicity and generalizability. In our experiments, we find it works across different action spaces (e.g., delta position in Moving Maze, delta joint pose in Peg Insertion, and delta joint velocity in Push Chair) with different action statistics based on how they are generated (sampling-based methods such as RRT, heuristics, etc.) and scale well to higher action dimensions (e.g., Moving Maze uses 2-d actions, Push Chair uses 19-d actions with a dual-arm mobile robot). In practice, action spaces of much higher dimensions are relatively uncommon in robotic tasks.

Reviewer comment (Evaluation):
Results are lacking confidence intervals (how many runs/model seeds).

Author response:
For all results on ManiSkill2, we report the best result among the three training runs in the original paper. We report the additional mean and std below. The format is best (mean ± std).

|    | Stack Cube (unseen) | Peg Insertion (0-shot) | Turn Faucet (0-shot) | Push Chair (0-shot) |
| ------- | ------- | ------- | ------- | ------- |
| BeT | 73.0 (68.7±4.0)  | 42.5 (39.5±2.7) | 32.5 (31.7±0.8) | 33.4 (32.3±1.0) |
| CoTPC | 86.0 (84.0±2.6) | 59.3 (51.5±7.4) | 39.3 (35.4±3.6) | 41.0 (36.4±4.0) |

Reviewer comment (Experiments):
How well does this approach to sub-task decomposition scale to larger action-spaces?

Author response:
Our subskill discovery approach by decomposing the action sequences is motivated by the intuition that functionally similar actions that are temporally close should be grouped into the same subskill. The choice of the metric to measure such functional similarity is an open problem. We tried cosine similarity, L2 distance, and a more complicated Hausdorff distance. Among these, we find that the cosine similarity exhibits both simplicity and generalizability. In our experiments, we find it works across different action spaces (e.g., delta position in Moving Maze, delta joint pose in Peg Insertion, and delta joint velocity in Push Chair) with different action statistics based on how they are generated (sampling-based methods such as RRT, heuristics, etc.) and scale well to higher action dimensions (e.g., Moving Maze uses 2-d actions, Push Chair uses 19-d actions with a dual-arm mobile robot). In practice, action spaces of much higher dimensions are relatively uncommon in robotic tasks.

Reviewer comment (Experiments):
How sensitive is their approach to the beta parameter that controls the number of detected changepoints?

Author response:
We perform experiments on the Peg Insertion task (state observation policies) and find that shorter CoT (approximately -1 length on average) leads to slightly decreased performance (59.3 vs. 54.0 in terms of 0-shot SR); we find longer CoT (approximately +1 length on average) have very similar performance.

Reviewer comment (Presentation):
Fig 2 - Can the authors comment on what the action groupings correspond to intuitively for these examples?

Author response:
Figure 6 in the appendix (updated pdf version) illustrates semantically what each subskill corresponds to.

Reviewer comment (Writing):
I am still not fully convinced by using the term 'Chain-of-Thought'...

Author response:
Thanks for the valuable review (again)!

Regarding the term CoT, we are willing to change the name to others (e.g., subskill-augmented or subskill-predictive control) if all of the reviewers think it is necessary (as the reviewer QpAT raised this point as well).

Reviewer comment (Writing):
The writing of section 4 needs improvements. Specifically, Section 4.2.2 is quite hard to parse and easy to get lost in. The authors are encouraged to re-write this section and redo Figure 1 so that things are clearer.

Author response:
**Section 4.2.2 writing needs improvement**

1. We’ve updated Section 4.2.2 to make it clearer. 
2. We reword all “predictors” to “decoders” to eliminate ambiguities.
3. We describe the inputs of the CoT decoders in more detail.
4. Regarding why not use action tokens for predictions: The action tokens are not directly used as inputs to all the decoders since during inference there is no input for the action token corresponding to the current timestep (this will be a_t, exactly what we aim to predict). The action tokens in the past are used in the attention layers, though. The alternative is to add an extra prompt token for the current actions. We stick to the standard approach of using the state tokens for inputs to the decoders as in DT/BeT.
5. The content of the CoT tokens is the learnable prompt embeddings. Each CoT token has its embeddings (not shared since we do not use position embedding for the CoT tokens). They are fixed during inference.
6. The T CoT tokens correspond to the T timesteps in the context length of the Transformer. Each token is learned to predict the same CoT content (next and last subgoal). There is no alignment issue between subgoals and CoT tokens in this setup. Because we couple action prediction and CoT prediction, we need T branches for T action offset predictions in the first place. We find the alternative producing less generalizable policies.

Reviewer comment (Writing):
There is confusion about what the inputs of \( g_\text{CoT} \) are.

Author response:
**Section 4.2.2 writing needs improvement**

1. We’ve updated Section 4.2.2 to make it clearer. 
2. We reword all “predictors” to “decoders” to eliminate ambiguities.
3. We describe the inputs of the CoT decoders in more detail.
4. Regarding why not use action tokens for predictions: The action tokens are not directly used as inputs to all the decoders since during inference there is no input for the action token corresponding to the current timestep (this will be a_t, exactly what we aim to predict). The action tokens in the past are used in the attention layers, though. The alternative is to add an extra prompt token for the current actions. We stick to the standard approach of using the state tokens for inputs to the decoders as in DT/BeT.
5. The content of the CoT tokens is the learnable prompt embeddings. Each CoT token has its embeddings (not shared since we do not use position embedding for the CoT tokens). They are fixed during inference.
6. The T CoT tokens correspond to the T timesteps in the context length of the Transformer. Each token is learned to predict the same CoT content (next and last subgoal). There is no alignment issue between subgoals and CoT tokens in this setup. Because we couple action prediction and CoT prediction, we need T branches for T action offset predictions in the first place. We find the alternative producing less generalizable policies.

Reviewer comment (Writing):
The role and definition of the CoT predictor are unclear; it appears only once in the last paragraph of section 4.2.2.

Author response:
**Section 4.2.2 writing needs improvement**

1. We’ve updated Section 4.2.2 to make it clearer. 
2. We reword all “predictors” to “decoders” to eliminate ambiguities.
3. We describe the inputs of the CoT decoders in more detail.
4. Regarding why not use action tokens for predictions: The action tokens are not directly used as inputs to all the decoders since during inference there is no input for the action token corresponding to the current timestep (this will be a_t, exactly what we aim to predict). The action tokens in the past are used in the attention layers, though. The alternative is to add an extra prompt token for the current actions. We stick to the standard approach of using the state tokens for inputs to the decoders as in DT/BeT.
5. The content of the CoT tokens is the learnable prompt embeddings. Each CoT token has its embeddings (not shared since we do not use position embedding for the CoT tokens). They are fixed during inference.
6. The T CoT tokens correspond to the T timesteps in the context length of the Transformer. Each token is learned to predict the same CoT content (next and last subgoal). There is no alignment issue between subgoals and CoT tokens in this setup. Because we couple action prediction and CoT prediction, we need T branches for T action offset predictions in the first place. We find the alternative producing less generalizable policies.

Reviewer comment (Theory):
It is unclear why actions are not used as inputs to \( g^x(.) \) functions.

Author response:
**Section 4.2.2 writing needs improvement**

1. We’ve updated Section 4.2.2 to make it clearer. 
2. We reword all “predictors” to “decoders” to eliminate ambiguities.
3. We describe the inputs of the CoT decoders in more detail.
4. Regarding why not use action tokens for predictions: The action tokens are not directly used as inputs to all the decoders since during inference there is no input for the action token corresponding to the current timestep (this will be a_t, exactly what we aim to predict). The action tokens in the past are used in the attention layers, though. The alternative is to add an extra prompt token for the current actions. We stick to the standard approach of using the state tokens for inputs to the decoders as in DT/BeT.
5. The content of the CoT tokens is the learnable prompt embeddings. Each CoT token has its embeddings (not shared since we do not use position embedding for the CoT tokens). They are fixed during inference.
6. The T CoT tokens correspond to the T timesteps in the context length of the Transformer. Each token is learned to predict the same CoT content (next and last subgoal). There is no alignment issue between subgoals and CoT tokens in this setup. Because we couple action prediction and CoT prediction, we need T branches for T action offset predictions in the first place. We find the alternative producing less generalizable policies.

Reviewer comment (Writing):
The exact contents of \( \{ \mathbf{S}^\text{CoT}_{...}\} \) are unclear. Do they change with time? What is the (...) subscript?

Author response:
**Section 4.2.2 writing needs improvement**

1. We’ve updated Section 4.2.2 to make it clearer. 
2. We reword all “predictors” to “decoders” to eliminate ambiguities.
3. We describe the inputs of the CoT decoders in more detail.
4. Regarding why not use action tokens for predictions: The action tokens are not directly used as inputs to all the decoders since during inference there is no input for the action token corresponding to the current timestep (this will be a_t, exactly what we aim to predict). The action tokens in the past are used in the attention layers, though. The alternative is to add an extra prompt token for the current actions. We stick to the standard approach of using the state tokens for inputs to the decoders as in DT/BeT.
5. The content of the CoT tokens is the learnable prompt embeddings. Each CoT token has its embeddings (not shared since we do not use position embedding for the CoT tokens). They are fixed during inference.
6. The T CoT tokens correspond to the T timesteps in the context length of the Transformer. Each token is learned to predict the same CoT content (next and last subgoal). There is no alignment issue between subgoals and CoT tokens in this setup. Because we couple action prediction and CoT prediction, we need T branches for T action offset predictions in the first place. We find the alternative producing less generalizable policies.

Reviewer comment (Experiments):
The ablation study has limited value. It should aim to provide more intuition on what exactly is learned by CoT embeddings. This could possibly be shown in the maze environment to illustrate subgoal discovery clearly. Figure 5 in the appendix is a step in the right direction, but the maze environment would be easier and more informative.

Author response:
In the Maze task, since the state observation is essentially a vector of the position of the agent as well as the moving blocks, we find that the CoT tokens and the CoT decoder learn to predict the positions of the moving blocks relatively precisely from which the agent will reach them. Since it is a relatively simple task with very clear subgoals, the illustration provides less information than that of the ManiSkill2 tasks. We will add it to the Appendix if required.

Reviewer comment (Theory):
The authors should express the limitations of the approach. For example, is it always possible to assume that actions similar or close temporally belong to the same subskill?

Author response:
**Limitations of the approach**

We will add a paragraph to discuss the limitation of action-based subskill discovery. Specifically, the assumption that similar actions that are also temporarily close should be grouped into the same subskill can be violated with complex manipulations with many micro-adjustments (e.g., standing and balancing on a rope). In this case, extra force feedback might be required to distinguish subskills that are only subtly different.

Reviewer comment (Presentation):
BeT should perhaps have been explained in section 3 since the method seems to be heavily based on it. Including it there could ease exposition, or alternatively, it could be placed in the appendix similar to the Rt-1 paragraph.

Author response:
**BeT should have been explained more**

We’ve added a brief intro to the BeT architecture in the appendix.

Reviewer comment (Presentation):
In the last paragraph of section 5.4 ("Shared tokens for CoT and action predictions"), the different variants should be shown as a three-part diagram to help readers understand the differences in inputs. Currently, the differences are not entirely clear.

Author response:
**The third ablation study needs a diagram to illustrate the different variants of the architecture**

We’ve added a diagram in the appendix to illustrate the differences.

Reviewer comment (Experiments):
It would be interesting to see the effect of setting the component of the auxiliary loss term to 0 while keeping learnable prompt tokens. This seems equivalent to BC but with added learnable tokens, maintaining the same capacity as the CoTPC architecture.

Author response:
**Curious to see what the effect of setting the component of the auxiliary loss term to 0, but keeping learnable prompt tokens (maintain the same capacity as the CoTPC architecture)**

We perform an ablation study on Peg Insertion and Push Chair with state observation (similar to our other ablation studies) where we simply add additional prompt tokens to the BeT baseline (no auxiliary loss and thus no CoT decoder). This variant (named BeT+prompt) has the same capacity as CoTPC yet performs very similarly to the BeT baseline (shown below). We report unseen SR for Push Chair and 0-shot SR for Peg Insertion (as in our other ablation studies).

|    | BeT | BeT+prompt | CoTPC |
| -------- | ------- | ------- | ------- |
| Peg Insertion | 42.5 | 43.0 | 59.3 |
| Push Chair | 33.4 | 32.7 | 41.0 |