SegDAC: Improving Visual Reinforcement Learning by Extracting Dynamic Objectc-Centric Representations from Pretrained Vision Models

TL;DR

SegDAC addresses the challenge of learning robust visual policies from high-dimensional images by grounding dynamic object-centric representations in text-guided segmentation derived from pretrained detectors. It combines a ground segmentation module (YOLO-World + EfficientViT-SAM) with a transformer-based actor-critic that consumes a variable-length sequence of segment embeddings, enabling online RL without human labels or data augmentation. On ManiSkill3, SegDAC achieves superior visual generalization under strong perturbations—up to about 2x IQM gains in the hardest settings—while maintaining competitive sample efficiency. This approach advances practical visual RL by leveraging open-world perception to adaptively focus computation on semantically relevant objects.

Abstract

Visual reinforcement learning (RL) is challenging due to the need to extract useful representations from high-dimensional inputs while learning effective control from sparse and noisy rewards. Although large perception models exist, integrating them effectively into RL for visual generalization and improved sample efficiency remains difficult. We propose SegDAC, a Segmentation-Driven Actor-Critic method. SegDAC uses Segment Anything (SAM) for object-centric decomposition and YOLO-World to ground the image segmentation process via text inputs. It includes a novel transformer-based architecture that supports a dynamic number of segments at each time step and effectively learns which segments to focus on using online RL, without using human labels. By evaluating SegDAC over a challenging visual generalization benchmark using Maniskill3, which covers diverse manipulation tasks under strong visual perturbations, we demonstrate that SegDAC achieves significantly better visual generalization, doubling prior performance on the hardest setting and matching or surpassing prior methods in sample efficiency across all evaluated tasks.

Figures (38)

• Figure 1: SegDAC overview and components.
• Figure 2: Our grounded segmentation module leverages the strengths of both YOLO-World and EfficientViT-SAM models to efficiently detect segments from the image.
• Figure 3: Examples of grounded segments for the PushCube task. Left: image with YOLO-World bounding boxes. Right: segments extracted by EfficientViT-SAM using boxes.
• Figure 4: Segment embedding extraction: For each SAM-predicted mask, we select patch embeddings with sufficient overlap ($\geq 4$ pixels), then apply global average pooling to obtain one embedding per segment.
• Figure 5: SegDAC actor critic networks architectures.