기술 요약: DataMIL – 데이터모델을 이용한 로봇 모방 학습용 데이터 선택

1. 문제 정의

로봇 커뮤니티는 최근 거대하고 다양한 데이터셋(예: Open-X Embodiment)을 사용하여 범용 파운데이션 정책(generalist foundation policies)을 훈련하는 방향으로 전환했습니다. 이러한 정책들은 다양한 작업에 대해 강력한 평균 성능을 달ей지만, 개별적이고 특화된 작업에서는 성능이 저하되는 경우가 빈번하며, 이에 따라 새롭게 습득한 작업 특정적 데이터에 대한 미세 조정(fine-tuning)이 필요합니다.

본 연구가 다루는 핵심 과제는 **데이터 선택(data selection)**입니다. 즉, 제한된 작업 특정적 데모와 결합했을 때 가장 좋은 전문화된 정책을 생성할 수 있는 대규모 사전 데이터셋의 하위 집합을 식별하는 것입니다.

• 기존 방법의 한계: 인간의 개념(예: 의미론적 유사성, 시각적 유사성 또는 상태-행동 근접성)에 기반한 나이브한 선택은 종 often 실패합니다. 이러한 휴리스틱은 가장 "유사한" 데이터가 가장 "유용할" 것이라고 가정하지만, 실제 다운스트림 정책 성능에 미치는 영향은 무시합니다.
• 최적화 장벽: 이상적으로는 하위 집합을 재학습 및 평가하여 최적의 데이터셋을 선택해야 합니다. 그러나 로보틱스에서 정책 평가는 실세계 롤아웃(rollout)을 필요로 하며, 이는 대규모 탐색을 수행하기에는 시간 소모적이고 위험하며 계산적으로 불가능합니다.
• 공백: 데이터모델(datamodels)(모델 성능을 예측하는 데이터모델)과 같은 기존의 데이터 귀속 프레임워크는 NLP 및 컴퓨터 비전에서는 성공적이었으나, 비미분적이고 비용이 많이 드는 실세계 정책 평가의 특성 때문에 로보틱스에 직접 적용하기 어려웠습니다.

2. 방법론: DataMIL

저자들은 데이터모델 패러다임을 로보틱스로 확장한 프레임워크인 DataMIL(Datamodels for Imitation Learning)을 소개합니다. DataMIL은 정적인 휴리스틱에 의존하는 대신, 정책 자체를 사용하여 성능을 개선하는 데이터 포인트를 식별함으로써 엔드투엔드 방식으로 데이터 선택을 추론합니다.

2.1 핵심 프레임워크

DataMIL은 정책 학습 알고리즘을 "블랙박스"로 취급하며, 특정 데이터 하위 집합으로 훈련된 정책의 성능을 예측하는 추정기(데이터모델)를 훈련합니다.

• 목표: $D' \subset D$인 하위 집합 $D'$를 찾아 $M(A(D'))$인 타겟 메트릭을 최대화합니다. 여기서 $A$는 학습 알고리즘이고 $M$은 성능 메트릭입니다.
• 데이터모델 공식화: 정책을 직접 훈련하는 대신, DataMIL은 데이터의 존재 여부에 기반하여 성능을 예측하는 함수 $\hat{f}(D') \approx M(A(D'))$를 학습합니다. 선형 데이터모델의 경우, 이는 각 데이터 포인트 $z_i$에 스칼라 점수 $\tau(z_i)$를 할당하는 것으로 축소되며, 예측된 성능은 선택된 하위 집합의 점수 합계가 됩니다.

2.2 로보틱스를 위한 주요 적응

데이터모델을 로보틱스에서 실행 가능하게 만들기 위해, 저자들은 세 가지 중요한 수정을 도입했습니다.

1. 대리 손실 함수 (롤아웃 회피):
   실세계 롤아웃은 미분 불가능하고 비용이 많이 듭니다. DataMIL은 실제 성공률 메트릭을 홀드아웃된 타겟 데모 세트에 대한 대리 메트릭(proxy metric) $\tilde{M}$으로 대체합니다. 이는 다음과 같이 정의됩니다:
   $$\tilde{M}(\pi, D_{target}) = \frac{1}{|D_{target}|} \sum_{(s,a) \in D_{target}} -L_{BC}(\pi(s), a)$$
   이 메트릭은 완전히 미분 가능하며 환경과의 상호작용을 필요로 하지 않으므로, 경사 기반 추정법을 사용할 수 있게 합니다.

2. 데이터모델 추정기:
   논문은 영향력 점수 $\tau(z_i)$를 추정하기 위해 두 가지 방법을 사용합니다.

• 회귀 추정기(Regression Estimator): 무작위 데이터 하위 집합을 샘플링하고, 정책을 훈련하며, 대리 메트릭을 평가한 후, 데이터 존재 여부에 따라 성능을 예측하도록 선형 모델을 적합시킵니다. 이는 정확하지만 계산량이 많습니다.
• 메타그레이디언트 기반 추정기(Metagradient-Based Estimator): 영향력 함수(influence functions)와 메타그레이디언트를 사용하여 타겟 메트릭의 데이터 가중치에 대한 미분값을 계산합니다. 이 방법은 대규모 모델(예: Octo)에 더 효율적이며, 모든 하위 집합에 대해 정책을 재훈련할 필요가 없습니다.

3. 클러스터링 및 분포 변화 완화:

• 클러스터링: 영향력 추정의 노이즈를 줄이기 위해, 개별 상태-행동 쌍을 시간적 클러스터(하위 궤적 또는 전체 궤적)로 그룹화합니다. 최적의 클러스터 입도(granularity)는 데이터셋 크기에 따라 달라지며, OXE와 같은 대규모 데이터셋은 궤적 수준의 집계가 유리합니다.
• 분포 변화(Distribution Shift): 사전 도메인과 타겟 도메인 간의 차이로 인한 잘못된 추정을 방지하기 위해, 데이터모델 추정 단계에서 타겟 작업 데이터의 작은 비율을 포함시켜 학습 과정을 타겟 도메인에 정렬합니다.

2.3 훈련 파이프라인

영향력 점수가 추정되면 다음 단계를 따릅니다:

1. 상위 $x\%$의 사전 데이터(가장 높은 양의 영향력)를 선택하여 $D_{sel}$을 구성합니다.
2. 최종 정책은 타겟 데이터 $D_{target}$과 선택된 데이터 $D_{sel}$에 대해 **공동 훈련(co-training)**을 통해 훈련됩니다. 각 단계에서 $D_{target}$에서 확률 $\alpha$로, $D_{sel}$에서 확률 $1-\alpha$로 배치를 샘플링합니다.

3. 주요 기여

• 로보틱스로의 데이터모델 확장: DataMIL은 미분 가능한 대리 손실을 통해 비용이 많이 드는 실세계 롤아웃 문제를 극복하고, 로봇 모방 학습에 데이터모델을 성공적으로 적용한 첫 번째 프레임워크입니다.
• 엔드투엔드 성능 인식: 시각적, 동작적, 또는 상태 유사성과 같은 휴리스틱 방식과 달리, DataMIL은 성능에 미치는 영향을 직접 최적화합니다. 이를 통해 시각적으로는 유사해 보일지라도 성능을 개선하는 데이터를 선택하고 성능을 저하시키는 데이터를 버릴 수 있습니다.
• 확장 가능한 추정기: 메타그레이디언트 기반 추정기의 도입은 회귀 기반 접근 방식이 계산적으로 불가능한 대규모의 복잡한 정책(예: Octo)과 거대한 데이터셋(예: OXE)으로 프레임워크를 확장할 수 있게 합니다.
• 교차 구현 및 다중 작업 선택: 사전 데이터셋에 존재하지 않는 작업 및 로봇 구현체에 대해 유용한 데이터를 선택하는 능력과, 여러 다운스트림 목적을 동시에 지원하는 데이터셋을 큐레이션하는 능력을 입증했습니다.

4. 실험 결과

저자들은 시뮬레이션 및 실세계 환경에서 60개 이상의 작업을 통해 DataMIL을 검증했습니다.

• MetaWorld (50개 작업): DataMIL은 최신 베이스라인(Behavior Retrieval, Flow Retrieval, STRAP) 대비 10%의 성능 향상을 달성했습니다. 자율 탐사 데이터로부터 노이즈가 심하고 최적화되지 않은 데모를 성공적으로 필터링했으며, 반면 유사성 기반 베이스라인들은 무관하거나 저품질의 데이터를 검색하는 경우가 많았습니다.
• LIBERO (10개 작업): Octo 트랜스포머 정책을 사용하여, DataMIL은 다양한 롱 호라이즌(long-horizon) 작업에서 베이스라인을 일관되게 능가했습니다. 시각적 유사성 베이스라인이 일부 특정 작업에서는 잘 작동하기도 했으나, DataMIL이 가장 견고한 평균 성능을 제공했습니다.
• Open-X Embodiment (OXE) (실세계):
  • Franka-Ball & Franka-Pouch: DataMIL은 무작위 선택 및 휴리스틱 베이스라인에 비해 성공률을 크게 높이는 데이터를 선택했습니다.
  • Tiago-Sink: Tiago 데이터가 없는 데이터셋에서 Tiago 로봇을 위한 데이터를 선택하는 도전적인 교차 구현 설정에서, DataMIL은 작업의 "본질"(테이블탑 조작)을 공유하는 다른 로봇(예: RT-1, BC-Z)의 관련 데모를 성공적으로 식별한 반면, 베이스라인들은 실패했습니다.
  • Droid-Multitask: DataMIL은 세 가지 서로 다른 작업을 동시에 개선하는 단일 데이터셋을 큐레이션하여 베이스라인보다 높은 평균 성공률을 보였습니다.

정성적 발견:

• DataMIL은 하나의 데이터 소스에 과도하게 의존하기보다(하나의 데이터셋에서 80% 이상을 가져오는 베이스라인과 대조적) 다양한 데이터셋의 혼합을 선택합니다.
• 시각적으로 동일해 보일 때라도(예: 동일한 상태이지만 다른 행동 분포를 가짐) 도움이 되는 데이터와 해로운 데이터를 효과적으로 구분해 냅니다. 이는 유사성 기반 검색에서 흔히 발생하는 실패 모드입니다.

5. 의의 및 주장

본 논문은 엔드투엔드 방식의 성능 인식 데이터 선택이 로보틱스에서 거대 사전 데이터셋의 잠재력을 끌어내는 데 필수적이라고 주장합니다. 저자들은 거대 파운데이션 모델이 강력한 기초를 제공하지만, 그 전문화는 미세 조정 데이터의 품질에 크게 의존한다고 강조합니다.

• 확장성에 대한 신중한 주장: 저자들은 데이터모델을 추정하는 것이 여전히 계산적으로 비용이 많이 든다는 점(전체 데이터에 대한 정책 훈련 비용의 몇 배)을 인정하지만, 메타그레이디언트 접근 방식이 이를 완화한다고 설명합니다. 또한 하이퍼파라미터 튜닝(예: 클러스터 크기, 선택 비율)에 현재 강력한 이론적 직관이 부족하다는 점을 언급합니다.
• 범위: 본 연구는 주로 단일 작업 전문화에 초점을 맞추고 있으며, Droid-Multitask 설정은 다중 작업 큐레이션을 향한 예비 단계 역할을 합니다.
• 결론: DataMIL은 인간이 정의한 휴리스틱을 넘어 데이터 중심의 모델 인식 선택으로 나아가는 것이, 전통적인 검색 방법이 무너지는 복잡하고 이질적인 실세계 시나리오에서 성공률을 일관되게 높이는 데 효과적임을 보여줍니다.