CLASS: Contrastive Learning via Action Sequence Supervision for Robot Manipulation

Sung-Wook Lee, Xuhui Kang, Brandon Y. Yang, Yen-Ling Kuo

Keywords: Supervised Contrastive Learning, Imitation Learning, Robot Manipulation, Action Chunking

TL;DR: CLASS improves Behavior Cloning by using DTW-based soft contrastive learning to learn robust action representations, boosting generalization across visual shifts in both simulated and real robotic tasks.

Abstract: Recent advances in Behavior Cloning (BC) have led to strong performance in robotic manipulation, driven by expressive models, sequence modeling of actions, and large-scale demonstration data. However, BC faces significant challenges when applied to heterogeneous datasets, such as visual shift with different camera poses or object appearances, where performance degrades despite the benefits of learning at scale. This stems from BC's tendency to overfit individual demonstrations rather than capture shared structure, limiting generalization. To address this, we introduce Contrastive Learning via Action Sequence Supervision (CLASS), a method for learning behavioral representations from demonstrations using supervised contrastive learning. CLASS leverages weak supervision from similar action sequences identified via Dynamic Time Warping (DTW) and optimizes a soft InfoNCE loss with similarity-weighted positive pairs. We evaluate CLASS on 5 simulation benchmarks and 3 real-world tasks to achieve competitive results using retrieval-based control with representations only. Most notably, for downstream policy learning under significant visual shifts, CLASS achieves an average success rate of 70% with Diffusion Policy, while all other baseline methods fail to perform competitively.

Submission Number: 803