기술 요약: 몬테카를로 트리 탐색을 위한 인과적 객체 중심 모델 (Causal Object-Centric Models for Planning with Monte Carlo Tree Search)

문제 정의

모델 기반 강화 학습(MBRL)은 환경 역학을 학습하여 생성된 상상된 경험을 통해 실제 세계와의 상호작용 필요성을 줄임으로써 상당한 성공을 거두었습니다. 그러나 고차원적이고 비정상성(non-stationary)이 있는 다수의 상호작용하는 객체들로 구성된 환경에서 정확한 월드 모델을 학습하는 것은 여전히 어려운 과제입니다.

기존의 시각적 RL 접근 방식은 일반적으로 입력 이미지를 하나의 통합된 표현으로 생성하는 합성곱 신경망(CNN) 인코더에 의존합니다. 이러한 단일 구조의 표현은 의사결정에 중요한 객체 수준의 구조와 상호작용을 포착하는 데 실패하는 경우가 많습니다. 복잡한 장면에서 동적인 배경, 무관한 엔티티, 또는 작지만 작업에 중요한 객체의 존재는 에이전트의 성능을 크게 저하시킬 수 있습니다. 인간은 세상을 효율적인 추론을 가능하게 하는 개별적인 엔티티들의 집합으로 인식하지만, 표준 RL 에이전트는 명시적인 감독이나 인스턴스 분할 마스크 없이 이러한 엔티티들을 격리하는 데 어려움을 겪습니다.

방법론: COMET

저자들은 슬롯 구조화된 잠재 공간(slot-structured latent space) 내에서 몬테카를로 트리 탐색(MCTS)을 수행하는 모델 기반 강화 학습 알고리즘인 COMET(Causal Object-centric Model for Efficient Tree search)을 소개합니다. COMET은 세 가지 주요 구성 요소를 통합합니다:

1. 동결된 객체 중심 인코더 (Frozen Object-Centric Encoder)

COMET은 시각적 관측치를 일련의 객체 중심 잠재 표현 또는 "슬롯"($\bar{s}_t = \{s^1_t, \dots, s^n_t\}$)으로 매핑하기 위해 동결된 비지도 객체 중심 인코더(예: SLATE, DINOSAUR 또는 Slot Contrast)를 사용합니다.

• 슬롯 초기화: 정적 이미지 모델에서의 순열 불변성(permutation-invariance) 및 확률적 특성을 해결하기 위해, 저자들은 시간 $t+1$의 슬롯을 시간 $t$의 슬롯을 사용하여 초기화합니다. 이는 에피소드 내에서 객체 할당의 시간적 일관성을 보장합니다.
• 동결: 인코더는 무작위 정책을 통해 수집된 데이터로 사전 학습되며, RL 단계 동안 안정적인 객체 수준 특징을 제공하기 위해 동결된 상태로 유지됩니다.

2. 액션-슬롯 융합을 갖춘 트랜스포머 기반 월드 모델 (Transformer-Based World Model with Action-Slot Fusion)

COMET의 핵심은 이러한 슬롯 위에서 작동하는 트랜스포머 기반 월드 모델입니다. 단일한 상태 벡터를 처리하는 표준 MuZero 스타일 방식과 달리, COMET은 객체 슬롯의 시퀀스를 처리합니다.

• 액션-객체 바인딩: 액션을 객체 슬롯과 융합하는 새로운 메커니즘을 도입합니다. 모든 객체 잠재 변수를 하나의 벡터로 압축하여 다음 상태를 예측하는 대신(이는 병목 현상을 유발함), 액션 임베딩을 각 슬롯에 독립적으로 결합합니다. 이러한 슬롯 조건부 액션 임베딩은 공유 MLP 프로젝터를 통과한 후 트랜스포머 백본을 거쳐 다음 상태 슬롯($\hat{s}_{t+1}$)과 보상을 예측합니다.
• 아키텍처: 모델은 nanoGPT 기반의 트랜스포머 백본을 사용하는 UniZero 스타일의 훈련 파이프라인을 따릅니다. 이는 현재의 슬롯 집합과 취해진 액션을 바탕으로 미래의 슬롯과 보상을 예측합니다.

3. 정책 및 가치를 위한 객체-인과적 어텐션 (Object-Causal Attention for Policy and Value)

정책 및 가치 헤드는 **객체-인과적 어텐션(object-causal attention)**을 갖춘 수정된 트랜스포머 아키텍처를 활용합니다.

• 인과성 점수: 각 잠재 슬롯 $h^k_t$에 대해, 네트워크는 해당 객체가 현재 결정에 관련될 확률을 나타내는 인과성 점수 $\alpha^k_t \in [0, 1]$를 추정합니다.
• 어텐션 변조: 이 점수들은 스케일드 닷 프로덕트 어텐션(scaled dot-product attention)을 변조하는 데 사용됩니다. 학습된 가중치 행렬 $W_t$는 정의된 인과 그래프 구조를 토큰 수준으로 끌어올려, 어텐션 메커니즘이 무관한 객체(비-인과적 객체)를 억제하고 작업 관련 엔티티(인과적 객체)에 집중할 수 있도록 합니다. 이를 통해 에이전트는 현재 환경과 상호작용 중인 특정 객체 부분집합에 계획을 집중할 수 있습니다.

주요 기여

1. COMET 알고리즘: 동결된 객체 수준 표현과 구조화된 잠재 공간에서의 계획을 위한 트랜스포르머 기반 월드 모델을 결합한 MCTS 기반 객체 중심 MBRL을 제안합니다.
2. 액션-슬롯 융합 메커니즘: 통합된 트랜스포머 백본 내에서 액션을 객체 중심 슬롯에 바인딩하는 새로운 방법입니다. 이를 통해 객체 정보를 단일 벡터로 붕괴시키지 않고 객체 중심의 월드 모델링 및 정책/가치 예측을 수행할 수 있습니다.
3. 객체-인과적 어텐션: 토큰 간의 상호작용을 변조하기 위해 학습된 인과성 점수를 통합하여, 정책이 의사결정 중에 작업 관련 엔티티에 동적으로 집중할 수 있게 합니다.
4. 경험적 평가: Object-Centric Visual RL 벤치마크, ManiSkill, Robosuite, VizDoom을 포함한 8가지의 시각적 및 동적 다양성을 가진 작업에 대한 포괄적인 평가를 수행했습니다.

실험 결과

COMET은 다음과 같은 베이스라인들과 비교 평가되었습니다:

• 객체 중심 베이스라인: OCRL (객체 인코더를 사용하는 모델 프리 PPO), OCDreamer (객체 중심 RSSM), SOLD (비지도 슬롯 기반 역학).
• 단일 구조 베이스라인: UniZero (표준 MuZero 스타일 잠재 계획).

성능 발견 사항:

• 샘플 효율성: 다양한 작업 세트 전반에서, COMET은 강력한 단일 구조 MCTS 기반 방법 및 다른 객체 중심 RL 베이스라인과 비교하여 훈련 초기 단계에서 더 높은 **평균 정규화 점수(Mean Normalized Score)**를 달s합니다.
• 작업 특이성:
  • 시각적으로 단순한 작업(예: Object Goal, Object Comparison)에서, COMET은 고품질의 객체 표현과 인과적 어텐션을 활용하여 관련 객체를 정확하게 식별함으로써 빠른 수렴을 이끌어냅/냅니다.
  • 연속 제어 작업(예: Block Lifting, Cube Pushing)에서, COMET은 일관된 성능을 보여주지만, 모든 특정 지표에서 보편적으로 모든 베이스라인을 능가하지는 않습니다 (예: Block Lifting에서 경쟁 모델을 능가하지 못함).
• 일반화: 이 방법은 단색 배경 위의 2D 도형부터 복잡한 3D 로봇 조작 및 1인칭 슈팅 게임 시나리오에 이르기까지 다양한 환경의 시각적 복잡성에 걸쳐 강건함을 보여줍니다.

의의 및 주장

본 논문은 의사결정을 위한 개별 객체의 중요성을 명시적으로 모델링하는 것이 객체 지향형 환경에서 정책 학습을 개선한다고 주장합니다. COMET은 MuZero 스타일의 잠재 계획에 명시적인 객체 수준 유도 편향(inductive bias)을 도입함으로써, 장면을 하나의 통합된 덩어리로 취급하는 대신 개별 엔티티와 그 상호작용을 추론할 수 있게 합니다.

저자들은 COMET이 복잡한 다중 객체 환경에서 관측치를 표현하는 문제를 효과적으로 해결한다고 주장합니다. 객체 표현을 분리하고 인과적 어텐션을 사용하여 관련성을 필터링함으로써, 이 알고리즘은 더 높은 샘플 효율성과 더 집중된 의사결정을 달성합니다. 이 연구는 객체 중심 유도 편향이 인스턴스 분할을 사용할 수 없거나 비지도 표현 학습이 필요한 경우를 포함하여, 시각적 RL에서의 모델 기반 계획을 개선하기 위한 실행 가능하고 효과적인 경로임을 시사합니다.