Grounding Bodily Awareness in Visual Representations for Efficient Policy Learning
This work addresses the problem of efficient policy learning for robotic manipulation by providing agent-centric visual representations, though it is incremental as it builds on existing contrastive learning and Vision Transformer methods.
The paper tackles the challenge of learning effective visual representations for robotic manipulation by introducing ICon, a contrastive learning method that separates agent-specific and environment-specific tokens in Vision Transformers, resulting in improved policy performance and transfer across different robots.
Learning effective visual representations for robotic manipulation remains a fundamental challenge due to the complex body dynamics involved in action execution. In this paper, we study how visual representations that carry body-relevant cues can enable efficient policy learning for downstream robotic manipulation tasks. We present $\textbf{I}$nter-token $\textbf{Con}$trast ($\textbf{ICon}$), a contrastive learning method applied to the token-level representations of Vision Transformers (ViTs). ICon enforces a separation in the feature space between agent-specific and environment-specific tokens, resulting in agent-centric visual representations that embed body-specific inductive biases. This framework can be seamlessly integrated into end-to-end policy learning by incorporating the contrastive loss as an auxiliary objective. Our experiments show that ICon not only improves policy performance across various manipulation tasks but also facilitates policy transfer across different robots. The project website: https://github.com/HenryWJL/icon