SegDAC: Visual Generalization in Reinforcement Learning via Dynamic Object Tokens

이 논문은 고정된 크기의 슬롯이나 보조 손실 없이 텍스트 기반 분할을 통해 동적 객체 토큰을 학습하고, 이를 변형 가능한 길이의 트랜스포머 기반 액터-크리틱에 적용하여 시각적 변화에 대한 강화학습의 일반화 성능을 획기적으로 개선한 'SegDAC'을 제안합니다.

핵심

1. 문제: "로봇이 눈이 나빠진 것 같아!"

기존의 로봇 AI 는 카메라로 찍은 화소 (픽셀) 단위의 이미지를 그대로 보고 학습합니다. 마치 우리가 사진을 볼 때, 사물의 모양보다는 배경의 색이나 빛의 반사 같은 '세부적인 픽셀'에 너무 집중하는 것과 비슷합니다.

• 상황: 로봇이 빨간 공을 밀어내는 일을 배웠다고 칩시다.
• 문제: 시험 시간 (실제 작동) 에 배경 벽의 색이 바뀌거나, 공의 재질이 반짝이는 플라스틱에서 매트한 고무로 바뀌면, 로봇은 **"이건 내가 배운 공이 아니야!"**라고 착각해서 아예 움직이지 못하거나 엉뚱한 행동을 합니다.
• 기존 해결책: 다양한 빛과 색을 인위적으로 만들어 로봇에게 보여주고 훈련시키는 '데이터 증강 (Data Augmentation)'이라는 방법을 썼는데, 이는 마치 로봇에게 "모든 경우의 수를 다 겪어봐"라고 시키는 것과 같아 비효율적이고, 여전히 완벽하지 않았습니다.

2. 해결책: SegDAC - "사물 중심의 눈"

이 논문은 로봇에게 픽셀을 보는 눈이 아니라, 사물 (Object) 을 인식하는 눈을 심어주었습니다. 이를 SegDAC이라고 부릅니다.

🍎 비유: "레시피 vs 재료 목록"

• 기존 방법 (픽셀 기반): 요리사가 "이 그릇 안의 모든 픽셀이 붉은색이야"라고 외우며 요리합니다. 그릇 색이 조금만 변해도 당황합니다.
• SegDAC (사물 중심): 요리사가 **"사과 1 개, 숟가락 1 개, 그릇 1 개"**라는 **재료 목록 (토큰)**만 보고 요리합니다. 사과가 빨간색이든 초록색이든, 그릇이 크든 작든 상관없이 "아, 사과가 있구나"라고 바로 파악합니다.

3. SegDAC 의 핵심 작동 원리 (3 단계)

이 기술은 세 가지 마법 같은 단계를 거칩니다.

① 텍스트로 사물을 찾아라 (Text-Grounded Segmentation)

로봇은 먼저 "로봇", "상자", "배경" 같은 간단한 **단어 (텍스트)**를 입력받습니다. AI 는 이 단어들을 보고 화면에서 해당 사물을 찾아냅니다.

• 비유: 요리사가 "사과 찾아줘!"라고 말하면, AI 가 화면에서 사과 모양을 딱 찾아내서 "여기 있어요!"라고 표시해 주는 것입니다.

② 동적인 '토큰'으로 변환 (Dynamic Object Tokens)

찾아낸 사물들을 **작은 정보 덩어리 (토큰)**로 만듭니다. 여기서 중요한 점은 개수가 변한다는 것입니다.

• 상황: 로봇이 사과를 집어 들면 사과가 사라지고, 새로운 사물이 들어오면 개수가 늘어납니다.
• 기존 기술의 한계: 대부분의 AI 는 "항상 5 개의 사물만 처리해"라고 고정되어 있어, 개수가 변하면 당황합니다.
• SegDAC 의 장점: "오늘은 사과 1 개, 내일은 사과 3 개, 모레는 사과 5 개"처럼 개수가 매일 변해도 자연스럽게 처리할 수 있습니다. 마치 유동적인 팀원 수에 맞춰 업무를 분배하는 유연한 팀장처럼요.

③ 공간 감각을 잊지 마세요 (Segment Positional Encoding)

사물만 인식하면 "어디에 있는 건지"를 모를 수 있습니다. 그래서 SegDAC 는 각 사물의 위치 정보를 토큰에 추가합니다.

• 비유: "사과가 있어"라고만 하면 어디에 있는지 모릅니다. 하지만 **"왼쪽 구석에 사과가 있어"**라고 하면 바로 찾을 수 있죠. SegDAC 는 이 위치 정보를 토큰에 박아두어, 로봇이 사물의 위치를 정확히 파악하게 합니다.

4. 왜 이것이 혁신적인가? (결과)

이 논문의 실험 결과는 매우 놀라웠습니다.

• 극한의 환경에서도 승리: 배경이 완전히 바뀌거나, 사물의 색과 질감이 엉망이 되어도 (가장 어려운 난이도), SegDAC 는 기존 기술들보다 88% 더 높은 성능을 보였습니다.
• 학습 속도도 빠름: 보통 시각적 변화에 강한 AI 는 학습이 느린데, SegDAC 는 가장 빠른 AI 와도 동일한 학습 속도를 보여주었습니다.
• 불필요한 노력 제거: 복잡한 데이터 증강이나 추가적인 학습 과정 없이, 오직 사물 중심의 인식만으로 이 모든 것을 달성했습니다.

5. 한 줄 요약

"SegDAC 는 로봇에게 픽셀의 색이나 빛에 휘둘리지 않고, '사물이 무엇이고 어디에 있는지'라는 본질적인 정보만 보고 유연하게 행동하는 능력을 심어주었습니다. 마치 요리사가 재료의 이름과 위치만 기억하고, 그릇 색이 바뀌어도 요리를 완벽하게 해내는 것과 같습니다."

이 기술은 앞으로 실제 세상 (실제 조명, 다양한 배경, 예측 불가능한 상황) 에서 로봇이 더 똑똑하고 안정적으로 일할 수 있는 기반이 될 것입니다.

기술 요약

1. 문제 정의 (Problem)

시각 강화학습 (Visual RL) 은 픽셀 관측을 기반으로 정책을 학습하지만, 테스트 시 시각적 조건 (배경, 조명, 물체 색상 등) 이 변경되면 성능이 급격히 저하되는 시각적 일반화 (Visual Generalization) 문제가 있습니다.

• 기존 접근법의 한계:
  • 픽셀 기반: 데이터 증강 (Data Augmentation) 은 유용하지만, 작업 관련 정보와 무관한 정보가 섞인 원시 픽셀에서 작동하여 분포 이동 (Distribution Shift) 에 취약합니다.
  • 객체 중심 (Object-Centric) 표현: Slot Attention 등의 방법은 고정된 수의 슬롯 (Slot) 을 사용하거나, 이미지 재구성 (Reconstruction) 목적 함수, 보조 손실 (Auxiliary Loss) 을 필요로 합니다. 이는 학습을 불안정하게 하거나 작업 관련성이 낮은 시각적 충실도에 치중하게 만듭니다.
  • 분할 기반 방법: 기존 분할 기반 RL 은 정답 마스크 (Ground-truth masks) 가 필요하거나, 재구성된 픽셀 이미지를 입력으로 사용하여 진정한 객체 수준의 추론을 하지 못합니다.

핵심 질문: 재구성, 보조 손실, 데이터 증강, 정답 마스크 없이, **가변적인 길이의 객체 토큰 (Variable-length Object Tokens)**으로부터 직접 RL 정책을 학습하여 강력한 일반화 성능을 낼 수 있는가?

2. 제안 방법: SegDAC (Methodology)

저자들은 **SegDAC (Segmentation-Driven Actor-Critic)**을 제안합니다. 이는 고정된 시각 모델과 가변적인 토큰 처리를 결합한 엔드 - 투 - 엔드 파이프라인입니다.

A. 동적 객체 토큰 생성 (Dynamic Object Token Construction)

1. 텍스트 기반 분할 (Text-Grounded Segmentation):
   • YOLO-World(개체 감지) 와 EfficientViT-SAM(분할) 을 사용하여 프레임을 분할합니다.
   • "robot", "cube", "background"와 같은 간단한 텍스트 개념 (Text Tags) 을 입력받아 개체 마스크를 생성합니다.
   • 장점: 정답 마스크가 필요 없으며, 텍스트 프롬프트가 분할의 정확도를 높입니다.
2. 경량 마스크 정제:
   • 반복적인 정제 대신 **형태학적 연산 (Morphological Opening & Closing)**을 사용하여 노이즈를 제거하고 경계를 채웁니다. 이는 온라인 RL 의 실시간성 (Low Latency) 을 보장합니다.
3. 맥락적 토큰 추출:
   • 분할된 마스크 영역에 해당하는 고정된 비전 인코더 (ViT) 의 패치 임베딩 (Patch Embeddings) 을 선택합니다.
   • 해당 패치들에 대해 **전역 평균 풀링 (Global Average Pooling)**을 수행하여 각 객체마다 하나의 임베딩 (Token) 을 생성합니다.
   • 핵심: 객체 토큰은 전역 컨텍스트를 이미 포함하고 있는 ViT 의 특징을 재사용하므로, 추가적인 어텐션 계산 없이도 장면 전체의 맥락을 유지합니다.

B. Transformer 기반 Actor-Critic 아키텍처

1. 가변 길이 처리 (Variable-Length Processing):
   • 각 타임스텝마다 감지된 객체 수 ($N$) 가 달라질 수 있습니다.
   • 패딩 (Padding) 이나 자르기 (Truncation) 를 하지 않고, 모든 토큰을 하나의 긴 시퀀스로 **패킹 (Packing)**하여 배치합니다.
   • 어텐션 마스크를 사용하여 각 타임스텝의 토큰이 서로 간섭하지 않도록 제한합니다. 이는 GPU 메모리 효율성을 극대화합니다.
2. 세그먼트 위치 인코딩 (Segment Positional Encoding):
   • 객체의 공간적 위치 정보를 보존하기 위해, 객체 마스크의 바운딩 박스 좌표를 기반으로 학습 가능한 위치 인코딩을 토큰에 추가합니다.
   • 이는 객체 간의 공간적 관계를 이해하는 데 필수적입니다.
3. 모달리티 임베딩:
   • 객체 토큰, proprioception (자기 감각), 학습된 쿼리 토큰을 구별하기 위해 모달리티 임베딩을 추가합니다.

3. 주요 기여 (Key Contributions)

1. 재구성과 보조 손실 없는 안정적 학습:
   • 고정된 슬롯 수, 이미지 재구성, 보조 손실 없이, 가변 길이의 객체 토큰으로부터 직접 안정적인 모델 프리 (Model-free) 정책을 학습하는 최초의 방법입니다.
2. 프리트레인 비전 모델 활용 및 공간적 정합:
   • 정답 마스크나 파인튜닝 없이 프리트레인된 비전 모델을 사용하여 컨텍스트가 포함된 객체 토큰을 생성하고, 세그먼트 위치 인코딩을 통해 공간적 정보를 보존합니다.
3. 샘플 효율성과 일반화의 동시 달성:
   • 기존 시각 일반화 방법들은 샘플 효율성이 낮거나 데이터 증강에 의존했으나, SegDAC 는 DrQ-v2 와 동급의 샘플 효율성을 유지하면서 뛰어난 일반화 성능을 달성했습니다.
4. 새로운 벤치마크 (ManiSkill3 Visual Generalization Benchmark):
   • 8 가지 조작 작업, 12 가지 시각적 교란 유형, 3 가지 난이도 (Easy, Medium, Hard) 를 포함하는 새로운 벤치마크를 제시했습니다. 특히 Hard 난이도는 시각적/의미적 충돌을 유발하여 기존 방법들의 한계를 명확히 드러냅니다.

4. 실험 결과 (Results)

• 벤치마크 성능: ManiSkill3 의 8 가지 작업에서 12 가지 교란 유형을 평가했습니다.
  • Easy: 기존 방법 대비 15% 향상.
  • Medium: 기존 방법 대비 66% 향상.
  • Hard: 기존 방법 대비 88% 향상 (가장 극적인 개선).
  • 특히 DrQ-v2 는 Hard 설정에서 성능이 90% 이상 하락한 반면, SegDAC 는 높은 성능을 유지했습니다.
• 샘플 효율성: 데이터 증강 없이 DrQ-v2 와 유사한 학습 속도와 효율성을 보였습니다.
• Ablation Study:
  • 위치 인코딩 제거: 학습 불안정 및 성능 저하 발생 (공간 정보의 중요성 입증).
  • 고정 토큰 수 (N=5): 복잡한 작업에서 성능이 급격히 떨어짐 (가변 길이 처리의 필요성 입증).
  • 전역 토큰 (Global Token) 사용: 객체 수준의 구조가 없으면 성능이 크게 떨어짐.
• 실패 분석: 교란이 심할 때 MaDi 와 같은 방법은 비논리적인 행동을 보이지만, SegDAC 는 작업 의도를 유지하며 구조화된 실패 (Graceful Degradation) 를 보입니다.

5. 의의 및 결론 (Significance)

SegDAC 는 시각 RL 의 핵심 난제인 시각적 일반화와 샘플 효율성이라는 상충되는 목표를 동시에 달성하는 새로운 패러다임을 제시합니다.

• 표현의 본질: 픽셀 단위가 아닌 **객체 단위 (Object-centric)**로 추론하는 것이 시각적 변화에 훨씬 강건함을 증명했습니다.
• 실용성: 데이터 증강, 보조 손실, 정답 마스크 없이도 실시간 온라인 RL 이 가능하도록 파이프라인을 최적화하여 실제 로봇 적용 가능성을 높였습니다.
• 미래 방향: 프리트레인된 비전 모델의 힘을 RL 에 효과적으로 통합하는 방법론을 제시함으로써, 향후 더 복잡한 환경에서의 지능형 에이전트 개발에 중요한 토대를 마련했습니다.

이 연구는 **"객체 중심의 동적 토큰 처리"**가 시각 RL 의 일반화 문제를 해결하는 핵심 열쇠임을 입증했습니다.