Chain-of-Thought CheckList

• The paper introduces a hierarchical imitation learning framework that decomposes complex manipulation tasks into key subgoal states using sub-optimal demonstrations.
• It employs an unsupervised subskill segmentation method and a custom Transformer architecture with learnable prompt tokens and hybrid masking to jointly predict actions and key states.
• Empirical results on tasks like peg insertion and cube stacking demonstrate superior task success rates and robust generalization compared to state-of-the-art baselines.

Chain-of-Thought Predictive Control (CoTPC) is a hierarchical imitation learning approach that leverages sub-optimal demonstrations to learn generalizable low-level control policies for complex manipulation tasks. The method introduces a chain-of-thought (CoT) abstraction by decomposing demonstration trajectories into sequences of key states that mark subgoal completions. Through an observation space-agnostic, unsupervised subskill extraction and a custom Transformer-based architecture, CoTPC generates policies that are robust to the noise and suboptimality inherent in real-world demonstration data, consistently surpassing strong baselines in both task success rates and generalization across challenging manipulation environments.

1. Hierarchical Imitation Learning via Key State Abstraction

CoTPC is grounded in the idea that many challenging robotic and manipulation tasks are naturally decomposable into a sequence of lower-variance subskills or subgoals—so-called “key states”—which persist across diverse (including sub-optimal) demonstrations. Rather than modeling the entire task as a flat, monolithic control process, CoTPC extracts these states at segment boundaries where subgoals are completed (e.g., grasping, aligning, inserting in a peg insertion task).

By leveraging the inherent multi-stage structure, CoTPC exploits even noisy and non-Markovian trajectories: subgoal boundary transitions are statistically more stable, admitting an effective hierarchical representation. Sub-optimality, typically problematic for standard behavioral cloning due to compounding errors, is mitigated by focusing policy learning at these critical waypoints.

2. Observation Space-Agnostic Subskill Decomposition

The CoTPC pipeline begins with a robust, observation-invariant approach to unsupervised subskill segmentation:

• Temporally adjacent and functionally similar actions are clustered into trajectory segments, with segment boundaries corresponding to completion of subskills.
• Extraction mechanisms are intentionally simple yet flexible: in simulation, information such as contact forces, object states, or even vision-language based criteria (for zero-shot transfer) serve to detect subgoal completions.
• This process does not require hand-crafted features or modality-specific engineering, rendering CoTPC broadly applicable to various input representations.
• Across demonstrations, key state sequences are identified as “chains of planning steps,” forming the CoT that captures shared task decomposition patterns.

By focusing on the temporal structure of state transitions, CoTPC sidesteps the need for domain-dependent feature engineering and enables broad generalization capabilities.

3. Chain-of-Thought as Structured Multi-Step Planning

Within the CoTPC framework, the chain-of-thought (CoT) is conceptualized as the sequence of extracted key states $[s_0^{cot}, s_1^{cot}, \ldots, s_{K-1}^{cot}]$ for each trajectory, where each $s_k^{cot}$ represents the state at the $k$-th subgoal boundary. This ordered list reflects the “thought process” of the demonstrator.

Unlike classical hierarchical models that might only condition on the next subgoal, CoTPC predicts the full sequence of key state milestones for the entire planning horizon. Joint prediction of low-level actions and the overarching CoT enables closed-loop correction and dynamic policy adjustment at test time, improving the generalization and robustness of the learned controller.

4. Transformer Architecture with Learnable Prompt Tokens and Hybrid Masking

The CoTPC policy architecture employs a specialized Transformer network designed for the joint prediction of (a) immediate actions, and (b) the full chain-of-thought key states, using the following design:

• Learnable Prompt Tokens: $K$ trainable “prompt” vectors are prepended to the input sequence, each corresponding to a subgoal boundary. The number of prompts $K$ matches the number of key states considered for the task.
• Hybrid Masking: Standard state–action tokens use a causal mask (autoregressive sequence), while prompt tokens use an all-to-all mask, enabling them to attend to the full context, thus capturing global plan information.
• Action and Key State Decoders:
  • The action decoder $g_a$ maps the representation of the last token at the output of the final layer to the predicted action.
  • The key state decoder $g_{cot}$ projects the representations of prompt tokens (from an earlier layer) to the predicted key states.
• Formal Architecture:
  • Context window:

  $$\tau_T(t) = \{ s_{t-(T-1)}, a_{t-(T-1)}, \ldots, s_t \}$$ - Multi-head attention with hybrid mask:

  $$h_j(\tau_T(t)) = \text{MHA}_{hmask}[ F_{enc}(\tau_T(t)) ]$$

  recursively for all Transformer layers. - Decoding steps:

  $$\hat{a}_t = g_a(h_J(\tau_T(t))[-1]) \ \hat{s}_{k,t}^{cot} = g_{cot}(h_I(\tau_T(t))[k]), \quad k = 0, \ldots, K-1$$ - Training objective:

  $$\mathcal{L}_{total} = \mathbb{E}_{(s_t, a_t) \in D} [\mathcal{L}_{bc}(\hat{a}_t, a_t)] + \frac{\lambda}{K} \sum_{k=0}^{K-1} \mathbb{E}_\tau \left[ \frac{1}{|\tau|} \sum_t \mathcal{L}_{cot}(\hat{s}_{k, t}^{cot}, s_k^{cot}) \right]$$

This dual-head and prompt-token design enables dynamic trajectory-level reasoning and fine-grained subgoal tracking during both training and deployment.

5. Empirical Performance and Generalization

CoTPC undergoes comprehensive evaluation across challenging low-level robotic manipulation benchmarks: Pick-and-Place Cube, Stack Cube, Turn Faucet, and Peg Insertion Side. Noteworthy findings include:

• Superior Task Success: CoTPC achieves higher success rates than behavioral cloning, Decision Transformer, Behavior Transformer, Decision Diffuser, and a masked prediction baseline using ground-truth key states.
• Enhanced Generalization: The architecture demonstrates robust transfer to unseen task variations (environment seeds), with particularly strong results on tasks requiring fine-grained subgoal sequencing (e.g., distinctly improved rates for subgoal success metrics in peg insertion: grasp, align, insert).
• Closed-Loop Control: Dynamic adjustment based on explicit key state guidance supports more resilient performance in the presence of compounding errors or suboptimal demonstrator behavior.

6. Implementation Considerations and Deployment

When implementing CoTPC in practice, several technical and system-level factors warrant attention:

• Computational Overhead: The Transformer architecture, with hybrid masking and prompt token management, introduces additional computational cost compared to vanilla cloning baselines. However, the policy remains tractable for real-time robotic deployment in simulation.
• Feature Extraction: While the observation space-agnostic segmentation approach simplifies adaptation to new environments, care is needed in designing rule-based or learned detectors for key state identification in real-world (e.g., vision-based) settings.
• Scalability: The reliance on sub-optimal demonstrations, rather than high-quality expert data, allows for scalable data collection; the hierarchical decomposition mechanism gracefully recovers structure from noisy datasets, facilitating domain adaptation.
• Deployment Strategy: CoTPC supports closed-loop prediction, making it suitable for both offline learning and online correction. Prompt tokens can be dynamically updated at test time for interactive or incremental plan adjustment.

7. Summary Table: Core Components of CoTPC

Component	Function	Key Feature
Key State Extraction	Unsupervised decomposition into subgoal boundaries	Observation space-agnostic, simple
Chain-of-Thought (CoT)	Sequence of critical subgoal-completed states	Guides policy at plan-level
Transformer Architecture	Joint prediction of actions and key states (prompts)	Hybrid masking, prompt tokens
Training Objective	Combines BC loss and auxiliary MSE on key state preds	Enables joint, dynamic optimization
Empirical Performance	Surpasses strong baselines on manipulation benchmarks	Robust generalization and control

In summary, Chain-of-Thought Predictive Control demonstrates that extracting and leveraging hierarchical chains of subgoal states from sub-optimal demonstrations—coupled with a prompt-enhanced, multi-headed Transformer architecture—facilitates robust, generalizable policy learning for complex sequential decision-making tasks in robotics and control.