Chain-of-Thought Predictive Control

Author Response to Reviewer QpAT

Thanks for the valuable feedback with so many details! **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). **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. **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. **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 |   [1] Learning Options via Compression [2] Diversity is all you need: Learning skills without a reward function [3] Robot learning from demonstration by constructing skill trees [4] Policyblocks: An algorithm for creating useful macro-actions in reinforcement learning [5] Compile: Compositional imitation learning and execution [6] Hierarchical Imitation Learning with Vector Quantized Models [7] Chain of Thought Imitation with Procedure Cloning [8] Unsupervised Learning of Temporal Abstractions with Slot-based Transformers [9] Target-driven Visual Navigation in Indoor Scenes using Deep Reinforcement Learning

Author Response to Reviewer QpAT (cont.)

**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). **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. **Sensitivity of the length of CoT.** 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. **The variance should be shown for each method, and the learning curves are required.** 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) | Note that the learning curves might not be very meaningful for purely offline imitation learning tasks. Both BeT and DT do not include them for the offline learning setup. **How efficient is the transformer architecture?** We use around 1 million parameters for all of the CoTPC, DT, and BeT models in our experiments. We wonder what additional information is requested here. **If the transformer's weights are tuned, why the CoT can accurately serve as a target signal?** We are not quite sure what is referred to here. The CoT tokens are learned to extract information predictive of the hierarchical structure of the task in the form of subgoals. The Transformer is learned to utilize the extracted information for better action decoding.

Author Response to Reviewer HDpH

Thanks for the valuable review! **Major concern: Cosine similarity metric for sub-task decomposition seems heuristical. How does it scale up to a higher action dimension?** 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. **Discussion about the proposed decomposition approach compared to others in the literature (e.g, [1] and [2])** 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. **Citations should be in ICLR format** We will fix it. **Results are lacking confidence intervals (how many runs/model seeds)** 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) | **How sensitive is their approach to the number of detected changepoints (length of CoT)?** 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. **Fig 2 - Can the author's comment on what the action groupings correspond to intuitively for these examples?** Figure 6 in the appendix (updated pdf version) illustrates semantically what each subskill corresponds to.   [1] Mo2: Model-based offline options [2] Multi-skill Mobile Manipulation for Object Rearrangement [3] Gaussian process change point models [4] Optimal detection of changepoints with a linear computational cost

Author Response to Reviewer iFzy

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

Author Response to Reviewer 2mSo

Thanks for the valuable review! **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. **BeT should have been explained more** We’ve added a brief intro to the BeT architecture in the appendix. **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. **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 | **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.

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