Attentive Feature Aggregation or: How Policies Learn to Stop Worrying about Robustness and Attend to Task-Relevant Visual Cues

TL;DR

This paper introduces Attentive Feature Aggregation (AFA), a lightweight trainable pooling method that enables visuomotor policies to focus on task-relevant visual cues, significantly improving robustness against visual perturbations without extra data or fine-tuning.

Contribution

The paper proposes AFA, a novel feature pooling mechanism that enhances robustness of visuomotor policies by attending to relevant cues and ignoring distractors, without requiring dataset augmentation.

Findings

AFA outperforms standard pooling in perturbed scenes

Policies with AFA maintain robustness without dataset augmentation

AFA is effective in both simulation and real-world environments

Abstract

The adoption of pre-trained visual representations (PVRs), leveraging features from large-scale vision models, has become a popular paradigm for training visuomotor policies. However, these powerful representations can encode a broad range of task-irrelevant scene information, making the resulting trained policies vulnerable to out-of-domain visual changes and distractors. In this work we address visuomotor policy feature pooling as a solution to the observed lack of robustness in perturbed scenes. We achieve this via Attentive Feature Aggregation (AFA), a lightweight, trainable pooling mechanism that learns to naturally attend to task-relevant visual cues, ignoring even semantically rich scene distractors. Through extensive experiments in both simulation and the real world, we demonstrate that policies trained with AFA significantly outperform standard pooling approaches in the presence of visual perturbations, without requiring expensive dataset augmentation or fine-tuning of the PVR. Our findings show that ignoring extraneous visual information is a crucial step towards deploying robust and generalisable visuomotor policies. Project Page: tsagkas.github.io/afa