SARM: Stage-Aware Reward Modeling for Long Horizon Robot Manipulation

"Stage-Aware Reward Modeling for Long Horizon Robot Manipulation"(ICLR 2026) 논문의 상세 기술 요약입니다.

1. 문제 정의

본 논문은 변형 가능한 물체(예: 티셔츠 접기) 를 다루는 장기적, 접촉이 풍부한 로봇 조작의 핵심 과제를 해결합니다. 최근 로봇 행동 모델 (RBM) 은 유망한 성과를 보였으나, 다음과 같은 어려움에 직면해 있습니다:

• 일관성 없는 시연 품질: 대규모 데이터셋에는 비전문가 운영자로부터 수집된 노이즈가 많거나, 비최적, 혹은 실패한 궤적이 종종 포함됩니다.
• 취약한 라벨링: 기존 방법들은 주로 프레임 인덱스나 단순한 시간 기반 진행 라벨에 의존합니다. 지속 시간이 가변적이고 진행이 비선형적인 작업 (예: 천의 종류에 따라 셔츠를 펴는 데 걸리는 시간이 다름) 의 경우, 프레임 기반 라벨링은 심각한 노이즈를 유발하여 "느리지만 올바른" 단계와 "실패한" 단계를 구분하는 것을 불가능하게 만듭니다.
• 일반화: 큐레이션된 데이터로 훈련된 모델은 오류 복구나 변형 가능한 물체의 가림 현상 처리와 같은 분포 외 (OOD) 시나리오로 일반화되지 못합니다.

핵심 문제는 완벽한 수동 제작 보상 함수 없이 정책 학습을 효과적으로 안내하기 위해, 다양하고 가변적인 길이의 시연으로부터 일관되고 세분화된 보상 신호를 도출하는 방법입니다.

2. 방법론

저자들은 보상을 생성하기 위한 SARM(Stage-Aware Reward Modeling) 과 이러한 보상을 활용하여 정책을 훈련하기 위한 RA-BC(Reward-Aligned Behavior Cloning) 라는 두 부분으로 구성된 프레임워크를 제안합니다.

A. SARM: Stage-Aware Reward Modeling

SARM 은 작업 진행을 강건하게 예측하도록 설계된 비디오 기반 보상 모델링 프레임워크입니다.

• 자연어 주석: 프레임 인덱스 대신 자연어 서브태스크 주석 (예: "잡기", "펴기", "접기") 을 사용하여 작업을 의미론적 단계로 분해합니다.
• 듀얼 헤드 아키텍처:
  1. 단계 추정기 (Stage Estimator): 시각적 관찰 (RGB 프레임) 과 관절 상태를 기반으로 현재 고수준 작업 단계 (예: "펴기 단계") 를 예측하는 분류 헤드입니다.
  2. 서브태스크 추정기 (Subtask Estimator): 예측된 단계에 조건부 (conditioned) 로 세분화된 진행도 ($\in [0, 1]$) 를 예측하는 회귀 헤드입니다.
• 라벨링 전략:
  • 저자들은 데이터셋 전체의 각 서브태스크에 대한 평균 시간 지속 기간을 기반으로 데이터셋 수준의 사전 비율 ($\bar{\alpha}_k$) 을 계산합니다.
  • 진행도 라벨은 서브태스크 내에서 선형 보간하여 생성되므로, 해당 단계에 도달하는 데 걸린 프레임 수와 관계없이 동일한 의미론적 상태 (예: 완전히 펴진 셔츠) 는 일관된 진행도 값을 받습니다.
• 입력 처리: 이 모델은 시각적 특징을 위해 고정된 (frozen) CLIP 인코더를 사용하고, 시간적 의존성을 포착하기 위해 Transformer 인코더를 사용합니다. 또한 실패 사례와 OOD 시나리오에 대한 강건성을 높이기 위해 되감기 증강 (rewind augmentation) (이전 타임스탬프의 역순 프레임을 추가) 을 적용합니다.

B. RA-BC: Reward-Aligned Behavior Cloning

RA-BC 는 학습된 SARM 보상 모델을 활용하여 정책 훈련 중 시연 데이터를 필터링하고 재가중치하는 훈련 프레임워크입니다.

• 진행도 델타 계산: 각 훈련 샘플에 대해 SARM 보상 모델을 사용하여 현재 상태와 다음 상태 간의 진행도 개선 ($\Delta \phi$) 을 계산합니다.
• 동적 가중치 부여:
  • 시스템은 진행도 델타의 온라인 실행 통계 (평균 $\mu$ 및 표준 편차 $\sigma$) 를 유지합니다.
  • 양의 진행 (임계값 이상) 을 보이는 샘플은 가중치를 높이고, 음의 진행 (후퇴/실패) 을 보이는 샘플은 가중치를 낮추거나 필터링합니다.
  • "사전 재정의 (prior override)" 메커니즘은 명확히 좋은 샘플에는 가중치 1 을, 명확히 나쁜 샘플에는 0 을 부여하여 결정적인 필터링을 보장합니다.
• 목적: 표준 행동 복제 (BC) 손실은 가중 손실로 대체되어, 정책이 다양하고 노이즈가 많은 데이터셋의 고품질 진행 구간만을 모방하도록 효과적으로 가르칩니다.

3. 주요 기여

1. SARM 프레임워크: 자연어 주석을 활용하여 작업 단계와 세분화된 진행도를 동시에 예측하는 새로운 단계 인식형 비디오 기반 보상 모델입니다. 이는 프레임 인덱스 라벨링의 취약성을 극복하고 OOD 시나리오에 잘 일반화됩니다.
2. RA-BC 프레임워크: 학습된 보상 모델을 사용하여 고품질 시연을 자동으로 식별하고 재가중치하는 정책 훈련 방법입니다. 수동 필터링 없이 노이즈가 많고 다양한 데이터셋으로부터 효과적으로 학습할 수 있게 합니다.
3. 변형 가능한 물체에 대한 경험적 검증: 매우 까다로운 티셔츠 접기 작업에 성공적으로 적용되어, 표준 행동 복제 및 기타 보상 모델링 베이스라인 대비 상당한 개선을 입증했습니다.

4. 실험 결과

저자들은 티셔츠 접기(장기적, 변형 가능) 와 접시 내리기(다단계, 강체) 라는 두 가지 작업에서 접근법을 평가했습니다.

• 보상 모델 성능:

  • SARM 은 인간 시연과 실제 로봇 롤아웃에서 진행도를 예측하는 데 있어 베이스라인 (LIV, VLC, GVL, REDS, ReWiND) 보다 현저히 우수한 성능을 보였습니다.
  • 티셔츠 접기 작업에서 SARM 은 검증 데이터에서 0.008 MSE를 달성했으며 (다음으로 좋은 베이스라인은 0.019), 실제 로봇 롤아웃에서는 0.94 분류 점수를 기록했습니다 (ReWiND 는 0.78).
  • SARM 은 조명 변화와 다른 모델들이 노이즈가 많거나 낙관적인 신호를 생성하는 실패 모드에서도 강건했습니다.

• 정책 성능 (티셔츠 접기):

  • 베이스라인: 전체 다양 데이터셋 ($D_{all}$) 으로 훈련된 일반 BC 는 "Hard" 작업 (구겨진 상태에서 접기) 에서 **0%**의 성공률을 보였습니다.
  • 단순 필터링: 지속 시간으로 필터링 ($D_{2min}$) 한 경우 성능이 개선되었으나 Hard 작업에서는 여전히 실패했습니다 (0% 성공).
  • SARM 을 사용한 RA-BC:
    • Medium 작업 (펴진 상태에서 접힌 상태): 83% 성공률.
    • Hard 작업 (구겨진 상태에서 접힌 상태): 67% 성공률.
  • 비교: SARM 을 사용한 RA-BC 는 ReWiND 보상 모델을 사용한 RA-BC 보다 현저히 우수했습니다 (Hard 작업에서 67% 대 25%). 이는 보상 모델의 품질이 RA-BC 의 병목 현상임을 입증합니다.

• 절대 실험 (Ablation Studies):

  • 단계 인식: 단계 추정기를 제거하면 성능이 저하되어 계층적 모델링의 필요성이 확인되었습니다.
  • 되감기 증강: 실제 세계의 실패 사례로 일반화하는 데 필수적이며, 이를 사용하지 않으면 모델이 지나치게 낙관적이 되었습니다.
  • 모델 규모: 60M 파라미터 모델이 용량과 과적합 사이의 최적 균형을 제공하는 것으로 확인되었습니다.

5. 의의

본 논문은 장기적 로봇 조작을 위한 확장 가능하고 주석 효율적인 솔루션을 제공합니다.

• 데이터 효율성: 보상 신호가 "느린 진행"과 "실패"를 정확하게 구분할 수 있다면, 노이즈가 많고 다양한 데이터셋으로도 고품질 정책 학습이 가능함을 입증했습니다.
• 확장성: 밀집 보상 함수보다 획득이 쉬운 자연어 주석을 활용하고 보상 도출을 자동화함으로써, 변형 가능한 물체 조작과 같은 복잡한 작업으로 확장 가능합니다.
• 실용적 영향: 구겨진 상태에서 티셔츠를 접는 67% 의 성공률은 접촉이 풍부하고 변형 가능한 상호작용의 복잡성으로 인해 기존 모델이 일반적으로 실패하는 가정용 로봇 분야에서 중요한 도약을 의미합니다.

이 연구는 보상 모델링이 모방 학습의 핵심 병목 현상임을 확립하며, 단계 인식형 의미론적 진행 추정이 견고한 장기 로봇 행동을 unlocking 하는 열쇠임을 보여줍니다.