Grounding Bodily Awareness in Visual Representations for Efficient Policy Learning

TL;DR

This paper introduces ICon, a contrastive learning method for Vision Transformers that enhances robotic manipulation by embedding body-specific cues into visual representations, leading to improved policy learning and transfer.

Contribution

We propose ICon, a novel contrastive learning approach that separates agent-specific and environment-specific features in ViTs, improving manipulation policy efficiency and transferability.

Findings

Enhanced manipulation policy performance across tasks

Facilitated transfer of policies between robots

Improved agent-centric visual representations

Abstract

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://inter-token-contrast.github.io/icon/