Choose What to Observe: Task-Aware Semantic-Geometric Representations for Visuomotor Policy

TL;DR

This paper introduces a task-aware visual observation method that canonicalizes input images into semantic-geometric representations, significantly enhancing the robustness of visuomotor policies against appearance changes without retraining.

Contribution

It proposes a novel observation interface combining semantic segmentation and depth information to improve policy robustness to visual domain shifts.

Findings

Improved robustness to background and object recoloring shifts.

Maintains in-distribution performance across various benchmarks.

Effective with different policy architectures.

Abstract

Visuomotor policies learned from demonstrations often overfit to nuisance visual factors in raw RGB observations, resulting in brittle behavior under appearance shifts such as background changes and object recoloring. We propose a task-aware observation interface that canonicalizes visual input into a shared representation, improving robustness to out-of-distribution (OOD) appearance changes without modifying or fine-tuning the policy. Given an RGB image and an open-vocabulary specification of task-relevant entities, we use SAM3 to segment the target object and robot/gripper. We construct an L0 observation by repainting segmented entities with predefined semantic colors on a constant background. For tasks requiring stronger geometric cues, we further inject monocular depth from Depth Anything 3 into the segmented regions via depth-guided overwrite, yielding a unified semantic--geometric observation (L1) that remains a standard 3-channel, image-like input. We evaluate on RoboMimic (Lift), ManiSkill YCB grasping under clutter, four RLBench tasks under controlled appearance shifts, and two real-world Franka tasks (ReachX and CloseCabinet). Across benchmarks and policy backbones (Flow Matching Policy and SmolVLA), our interface preserves in-distribution performance while substantially improving robustness under OOD visual shifts.