[논문 리뷰] SARM: Stage-Aware Reward Modeling for long horizon Robot Manipulation

두부김치·2026년 6월 29일

논문 리뷰

목록 보기

18/26

Title : SARM: Stage-Aware Reward Modeling for long horizon Robot Manipulation(ICLR 2026)
논문 링크 : https://arxiv.org/pdf/2509.25358
blog: https://qianzhong-chen.github.io/sarm.github.io/

1. Introduction

연구 배경
- 대규모 데이터의 중요성: 최근 로봇 조작(Manipulation)분야는 인터넷 규모의 데이터와 인간의 시연(Demonstration)데이터를 대량으로 활용하여 복잡하고 범용적인 작업을 수행하는 방향으로 발전하고 있음.
- Long-horizon task의 도전 과제: 옷접기와 같은 Deformable object를 다루거나, 접촉이 많은(Contact-rich) task와 Long-horizon task는 여전히 해결하기 매우 어려움.
기존 방법론들의 한계점
- 데이터 수집의 노이즈: 대규모 데이터셋(예: OXE)을 구축하는 과정에서는 필연적으로 실패한 시도, 최적화되지 않은 움직임(Suboptimal trajectories), 무의미하게 멈춰 있는 구간 등 상당한 양의 'Noise label'이 포함됨.
- 시간/프레임 기반 label의 취약성: 기존의 비디오 기반 Reward 모델들은 영상의 시작 프레임을 0, 끝 프레임을 1로 두고 Frame Index에 따라 선형적으로 Progress를 labeling함.
  - 하지만 옷을 접느 작업처러머 매 시도마다 걸리는 시간이 다르고 Time warping이 심한 작업에서는, 단순히 영상의 중간 지점이라고 해서 작업이 정확히 절반 진행되었다고 볼 수 없음.
  - 이로 인해 보상 모델이 잘못된 신호를 학습하게 되고, Downstream Policy학습을 방해함.
Proposd Method
- 자연어 서브태스크 주석 활용: 저자들은 고정된 시간 프레임 대신, 영상에 포함된 고수준의 language Subtask Annotation을 활용하여 정교하고 일관된 진행도 label을 자동 생성하는 방식을 제안함.
  - SARM(Stage-Aware Reward Modeling) 프레임워크:
    - 듀얼 구조: 작업의 대단위 단계를 분류하는 '스테이지 모델(Stage Model)'과, 해당 스테이지 내부에서의 미시적인 진척도를 정밀하게 예측하는 '서브태스크 모델(Subtask Model)'이 결합된 형태임.
    - 이를 통해 Time warping이나 시연의 가변성에 영향을 받지 않는 강인하고 안정적인 보상 신호(Reward Signal)를 제공함.
  - RA-BC(Reward-Aligned Behavior Cloning):
    - SARM이 예측한 보상 신호를 기반으로, 학습 데이터 중 영양가가 높고 진행도가 확실한 구간에는 높은 가중치를 부여하고, 실패하거나 정체된 구간은 필터링(가중치 0)하는 데이터 정제 기법을 도입함.
Contribution
- SARM 개발: 시연 데이터의 다양성과 가변성에 기인하며,
- RA-BC 알고리즘 제안: 학습 데이터 내의 보상 정렬(Reward Alignment)을 통해 소프트 가중치 조절 및 필터링을 수행하여 정책 학습 효율을 극대화함.
- 실제 로봇 환경 검증: Bimanual Robot 인프라를 활용하여 난이도가 높은 T-shirt 접기(Folding) 작업 등에서 기존 SOTA 베이스라인 모델들을 압도하는 성과를 증명함.

2. Method

SARM 프레이워크는 크게 두 가지 단계로 구성됨. 첫째는 자연어 Subtask 주석을 활용한 스테이지 기반 모상 모델(SARM)이며, 둘째는 이를 활용해 정책(Policy)의 데이터 품질을 정제하는 보상 정렬 모방 학습(RA-BC)임.

2.1 Problem Formulation & Data Processing

기존의 프레임 인덱스 기반 방식과 달리, SARM은 비디오에 제공된 고수준의 언어 서브태스크 주석을 정밀한 연속적 진행도(Progress) 레이블로 변환함
- Subtask Prior 비율 산정( $\overline{\alpha}_k$ ):
  - 전체 데이터셋 $M$ 개의 trajectory중에서 각 subtask $K(K \in \{1,... ,K\})$ 가 차지하는 평균 시간적 비율을 계산
  - $\overline{\alpha}_k = \frac{1}{M}\sum_{i=1}^M\frac{T_{i,k}}{T_i}$
- 내부 정규화 시간 ( $\tau_t$ ) 및 최종 진행도 타겟 ( $y_t$ ) 생성:
  - 특정 서브태스크 $k$ 내의 프레임 $t$ 에 대해, 해당 서브태스크 내부에서의 상대적 위치(시간)를 $\tau_t \in [0, 1]$ 로 정의함.
  - 이를 기반으로 전체 작업 기준의 글로벌 진행도 타겟 $y_t \in [0, 1]$ 를 다음과 같이 선형 보간(Linear Interpolation)하여 할당.
    - $y_t = \sum_{j=1}^{k-1} \overline{\alpha}_j + \tau_t \cdot \overline{\alpha}_k$
  - 시연마다 속도가 다르더라도(Time warping), 고수준 스테이지의 경계면과 내부 진행도가 일관되게 정렬되어 레이블 노이즈가 획기적으로 줄어듬.

2.2 SARM: Stage-Aware Reward Model

SARM은 비디오 프레임을 입력받아 현재 '어떤 stage'에 있고, 그 Stage 내부에서 얼마나 진행되었는지를 Dual Head로 나눠서 동시에 예측함.
Shared Backbone:
- Visual Encoder: 각 비디오 프레임은 사전 학습된 고정된 시각 백본인 Frozen CLIP(ViT-B/32)을 통과하여 임베딩됨.
- Temporal Aggregator: 시간축 정보를 통합하기 위해 Transformer 인코더를 사용함.
- 절대적인 frame index 정보가 유출되어 모델이 shortcut을 배우는 것을 막기 위해, 오직 첫 번째 프레임에만 Positional Embedding을 적용하고 나머지 프레임에는 적용하지 않음.
Stage Classification head:
- 현재 프레임이 K 개의 Subtask 중 어디에 속하는지 분류하는 Discrete Head
- Loss Function: Cross Entropy Loss를 사용
Subtask Progress Head
- 스테이지 분류 헤드의 출력(Stage Classification Head의 출력(예측된 Stage 확률 분포)을 Prior로 입력 받아, 현재 스테이지 내에서의 연속적인 미시적 진행도 $y_t \in [0, 1]$ 를 Regression 예측함.
- MSE Loss function 사용

2.3 Robustness via Rewind Augmentation

단순히 성공한 데이터만 학습하면, 보상 모델이 비디오의 후반부 프레임만 보고 무조건 높은 보상을 주는 OverOptimism에 빠짐. 이를 방지하기 위해 역방향 데이터 증강(Rewind Augmentation)을 도입
- 방법: 정상적인 Forward 비디오 외에, 성공했다가 다시 구겨지거나 뒤로 되돌아가는 Reverse 시퀀스를 의도적으로 생성하여 학습 데이터에 포함함.
- 효과: 로봇이 실패하거나 작업을 망치는 상황(Negative Porgress)이 발생했을 때, 보상 값이 정확하게 감소하도록 만들어 Robusness와 Out-Of-Distribution 일반화 능력을 향상시킴.

2.4 RA-BC: Reward-Aligned Behavior Cloning

학습된 SARM 모델( $\phi$ )을 다운스트림 정책(Policy) 학습에 적용하는 단계, 데이터셋 내의 불완전하고 노이즈 있는 시연 중에서 실질적으로 작업 진척에 기여한 프레임 윈도우에만 가중치를 부여함.
- 보상 델타( $\hat{r}_i$ ) 정의:
  - 특정 학습 윈도우 $i$ 의 시작점( $t$ )과 미래 지점( $t+\tilde{\Delta}$ )사이의 SARM 보상 값의 차이를 측정
  - $\hat{r}_i = \phi(o_i^{t+\tilde{\Delta}}) - \phi(o_i^{t})$
- 실시간 온라인 가중치 산정(Welford's Algorithm):
  - 데이터셋 전체의 보상 델타 분포를 파악하기 위해, 학습 도중 실시간으로 온라인 평균( $\mu$ )과 표쥰편차( $\sigma$ )를 업데이트함.
  - 이를 통해 원시 델타값을 정규화된 소프트 가중치 $\tilde{w}_i$ 로 변환함.
- 임계값 기반 필터링 및 오버라이드(Threshold Override):
  - 확실하게 진전이 있는 구간(정규화된 값이 임계값 $k$ 이상)은 과감하게 가중치 $\tilde{w}_i=1$ 을 부여함.
  - 반대로, 제자리걸음을 하거나 뒤로 역행하는 구간( $\hat{r_t}<0$ )은 실패 혹은 무의미한 움직임으로 간주하여 가중치를 0으로 필터링(제거)함.
- RA-BC Object Function
  - $\mathcal{L}_{\text{RA-BC}} = \sum_{i} w_i \cdot \| \pi(a_i | o_i) - a_i^* \|^2$
  - 이 필터링을 통해 노이즈가 많은 대규모 데이터셋 안에서도 유용한 데이터셋만 골라 학습할 수 있음.

3. Result

해당 연구는 실험에 대해 3가지 질문을 던짐
- Q1: Reward Model로서 SARM의 강인성
- Q2: Policy 학습 시 RA-BC의 효과
- Q3: Reward Model의 품질이 Policy에 미치는 영향

3.1 Q1: Reward Model 평가

SARM의 성능을 검증하기 위해 Long-horizon인 T-shirt folding과 multi-stage dense contact task인 Dish unloading 두 가지 Task에 대해 최신 baseline과 비교를 함.
T-shirt Folding Task 결과
- human demonstration progress estimation(Demo $L$ ):
  - SARM은 단일 단계 MSE Loss에서 0.009를 기록하며 다른 모든 베이스라인을 압도함.
  - 이는 Subtask Prior 레이블링을 통해 프레임 인덱스 방식의 레이블 노이즈를 완벽히 언제했음을 보여줌.
- robot rollout progress estimation(Rollout $\rho$ ):
  - 로봇이 직접 움직임 36개 에피소드(성공 12, 부분 성공 12, 실패 12)에 대한 분류 일치 점수에서 SARM은 0.94를 달성하며 가장 강력한 베이스라인인 ReWiND(0.5)보다 80%이상의 상대적 성능 향상을 기록함.
- OOD(분포 외) 실패 감지 능력(Ablation Study):
  - 역방향 Frame Augmentation을 제거한 모델(w/o R)의 경우, human demonstration은 좋았으나 실제 Robot rollout evaulation은 떨어짐.
  - 로봇이 도중에 실수하거나 뒤로 돌아가는 상황(OOD)에서 모델이 지나치게 낙관적으로 Progress를 판단하기 때문에, 역방향 Augmentation이 필수적임을 증명함.
Dish Unloading Task 결과
- Undeformable한 객체가 아닌 환경에서의 일반화
  - SARM은 접촉 밀집형 Rigid object작업에서도 단일 단계 MSE 0.013, 롤아웃 분류 점수 0.67을 기록하며 베이스라인을 일관되게 능가했음.
  - 8개의 접시를 연속으로 내리는 1.5분 이상의 Long-horzion이상의 sequence에서도 흔들림 없이 안정적인 reward 곡선을 유지함.

3.2 Q2: Policy Learning with RA-BC

GELLO 시스템과 YAM 6-DoF 양팔 로봇을 활용해 수집된 200시간의 T-shirt Folding 대규모 데이터셋( $D_{all}$ )으로 정책을 Fine-tune하여 비교함.
난이도별 Policy 성공률 비교
- Simple Task(상자에서 옷을 꺼내 테이블 중앙에 놓기):
  - 모든 모델(Vanila BC, ReWiND, SARM)이 20K, 40K 스텝에서 12/12(100%) 성공률을 보이며 기본 제어 능력은 모두 확보함을 확인함.
- Medium Task(평평하게 펴진 상태에서 옷 접기):
  - 전체 데이터셋으로 학습한 Vanila BC(BC-All)는 데이터 노이즈로 인해 1/12 수준의 낮은 성공률을 보임.
  - 단순히 시간 단축 기준(2분 이내)으로만 데이터를 필터링한 데이터셋(BC-2min)은 71/12로 개선되었으나 plateau에 갇힘
- Hard task(구겨진 상태에서 펴고 접는 전체 파이프라인 수행)
  - 시간 단축 기반 필터링(BC-2min)은 정밀한 시각 인지 및 양팔 조율 능력을 가종하지 못해 0% 성공률로 완전히 실패
  - RA-BC-SARM은 Hard Task에서도 8/12의 높은 성공률을 기록하며 데이터 정제 및 고품질 프레임 가중치 부여의 강력함을 증명함.

3.3 Effect of Reward Model Quality in RA-BC

잘못된 보상 신호의 위험성:
- 베이스라인인 ReWiND 보상모델을 사용하여 RA-BC 정책을 학습시킨 결과(RA-BC-ReWiND), Medium task 성공률은 50%, Hard Task 성공룰은 25%로 그침.
- 이는 단순히 RA-BC 알고리즘 구조를 사용하는 것 보다, 그 안에 들어가느 Reward Model(SARM)이 얼마나 정확하게 신뢰성 있게 진척도를 짚어내느냐(Quality)ㅏ 전체 시스템 성능을 결정짓는 핵심 기둥임을 보여줌