Chain-of-Thought Predictive Control

Zhiwei Jia, Vineet Thumuluri, Fangchen Liu, Linghao Chen, Zhiao Huang, Hao Su

ICML 2024

[arxiv] [slides] [poster] [code] [video (old)]

CoTPC leverages suboptimal demonstrations to learn policies for tasks requiring highly-precise and dynamic controls with (geometry-level) generalizability and smooth trajectories.

Key Ideas

Part-I: decompose a task into a chain of subskill segments (CoT) from the demos

Part-II: learn to model actions jointly with dynamic predictions of subskills (replanning) supervised by the demos

Details of Part-I: Observation Space-agnostic Unsupervised Discovery of CoT

CoTPC performs grouping-based segmentation on the action sequences from the demos by grouping temporarily close and functionally similar actions (measured in Cosine dist.) into the same subskills. This approach is effective across different action spaces while robust against the observation space setups (camera angles, RGB vs. point clouds, etc.).

Details of Part-II: Joint Prediction of Actions & CoTs with Masked Transformer

CoTPC jointly predicts actions and CoTs where low-level action predictions rely on extracted trajectory features predictive of the future subskills (hence the name CoTPC). CoTPC is the first work that successfully adopts prompt tokens for low-level controls. We carefully studied its Transformer architecture and masking schemes.

Main Results

CoTPC outperforms SoTA methods for state-based (top) and visual-based (bottom, w\ point cloud observations) policies on challenging low-level object manipulation tasks, while being much faster and (potentially) more data-efficient than diffusion-based methods.

CoTPC also transfers to real robots well, with visual results of Stack Cube and Peg Insertion (both simplified) on the right.

The associated technical paper of this project is accepted to ICML 2024 with the following BibTeX

@InProceedings{jia2023chain,

title={Chain-of-Thought Predictive Control},

author={Jia, Zhiwei and Thumuluri, Vineet and Liu, Fangchen and Chen, Linghao and Huang, Zhiao and Su, Hao},

booktitle={Proceedings of the 41st International Conference on Machine Learning},

pages={21768--21790},

year={2024},

}

Page updated

Google Sites

Report abuse