SARM: 단계 인식 보상 모델링

SARM(Stage-Aware Reward Modeling)은 장기 로봇 조작 작업을 위한 비디오 기반 보상 모델링 프레임워크입니다. 이 가이드에서는 SARM 보상 모델을 학습하는 방법과 선택적으로 Reward-Aligned Behavior Cloning(RA-BC)과 함께 사용하는 방법을 다룹니다.

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

왜 보상 모델인가?

표준 행동 클로닝은 모든 데모스트레이션 프레임을 동일하게 취급하지만, 실제 로봇 데이터셋은 불완전합니다. 망설임, 수정, 품질이 상이한 궤적이 포함되어 있기 때문입니다. 보상 모델은 데모스트레이션에서 작업 진행도라는 일반화 가능한 개념을 학습하여 이 문제를 해결합니다. 비디오 프레임과 작업 설명을 입력으로 받아 로봇이 작업을 완료하기까지 얼마나 가까운지(0→1)를 예측합니다. 이렇게 학습된 “진행 신호”는 여러 방식으로 활용될 수 있으며, 두 가지 유망한 응용 분야는 다음과 같습니다. (1) 가중치 모방 학습(RA-BC)은 정책 학습 시 높은 진행도를 가진 프레임에 더 큰 가중치를 부여하고, (2) 강화 학습에서는 보상 모델이 온라인 또는 오프라인 정책 개선에 대해 밀도 높은 보상을 제공합니다.

개요

SARM에는 다음과 같은 기능이 있습니다:

1. 단계 인식 아키텍처: 고수준 작업 단계와 각 단계 내의 세밀한 진행도를 함께 예측합니다
2. 하위 작업 주석: 자연어 하위 작업 주석을 사용하여 일관된 진행 라벨을 유도
3. 시간 비율: 가변 길이의 데모스트레이션에 걸쳐 진행도를 정규화하기 위해 각 하위 작업에 대한 데이터셋 수준의 사전 분포(α̅_k)를 계산합니다

SARM은 각 프레임마다 컴팩트한 stage+tau 타겟을 학습합니다:

• stage: 정수 단계 인덱스 k ∈ {0, ..., K-1}
• τ (tau): 단계 내 진행률 τ ∈ [0, 1]
• target encoding: y = k + τ (데이터셋 프로세서가 생성하는 값)

추론 시(및 하위 RA-BC에서) SARM은 데이터셋 수준의 시간 비율 α̅_k(meta/temporal_proportions_*.json에 저장됨)를 사용하여 원시 k + τ 값을 [0, 1] 범위의 정규화된 진행도로 변환합니다.

이것은 논문에서 수식 (2)와 일치합니다:

progress_t = P_{k-1} + α̅_k × τ_t

여기서:

• τ_t = (t - s_k) / (e_k - s_k)는 서브태스크 내 정규화된 시간입니다
• P_{k-1}은 누적 사전 확률(이전 하위 작업 비율의 합)입니다
• α̅_k는 서브태스크 k의 시간 비율입니다

이로 인해 서로 다른 길이의 데모에서도 동일한 작업 상태가 일관된 진행 값으로 매핑됩니다.

입력과 대상(새로운 코드가 기대하는 것)

SARM은 프로세서(src/lerobot/rewards/sarm/processor_sarm.py)를 통해 학습되며, 이 프로세서는 다음과 같은 역할을 수행합니다:

• CLIP(ViT-B/32)로 이미지와 작업 텍스트를 video_features와 text_features로 인코딩합니다
• 패딩/잘라 로봇 상태를 state_features로 (max_state_dim까지)
• Builds targets를 sparse_targets로(그리고 dense_only/dual의 경우 dense_targets로) stage+tau 인코딩 y = k + τ를 사용하여 생성합니다
• 마스크가 프레임 되감기를 샘플별 lengths 텐서를 사용하여 수행합니다(되감기는 학습 시 증강 기법입니다)

최소한, 각 학습 샘플에는 다음이 필요합니다:

• task (문자열): 작업 설명
• 데이터셋의 reward_model.image_key 이미지와 reward_model.state_key 상태

주석 모드

진행 라벨이 계산되는 방식을 결정하는 3가지 어노테이션 모드 중 하나를 선택할 수 있습니다:

모드 세부 정보

pip install -e ".[sarm]"

1. Train SARM → 2. Visualize predictions → 3. (Optional) Train policy with RA-BC

1. Annotate (dense) → 2. Verify → 3. Train SARM → 4. Visualize → 5. (Optional) Train policy with RA-BC

1. Annotate (sparse+dense) → 2. Verify → 3. Train SARM → 4. Visualize → 5. (Optional) Train policy with RA-BC

1단계: 하위 작업 주석

python src/lerobot/data_processing/sarm_annotations/subtask_annotation.py \
  --repo-id your-username/your-dataset \
  --dense-only \
  --dense-subtasks "Bring robot arms up from starting position,Grab near side and do 1st fold,Grab side and do 2nd fold,Grab side and do 3rd fold to finish folding" \
  --video-key observation.images.base \
  --num-workers 4 \
  --push-to-hub

python src/lerobot/data_processing/sarm_annotations/subtask_annotation.py \
  --repo-id your-username/your-dataset \
  --sparse-subtasks "Bring arms up from starting position,Fold the towel (3 folds in total)" \
  --dense-subtasks "Bring robot arms up from starting position,Grab near side and do 1st fold,Grab side and do 2nd fold,Grab side and do 3rd fold to finish folding" \
  --video-key observation.images.base \
  --num-workers 4 \
  --push-to-hub

Annotation Arguments

참고: 주석이 완료되면 기본적으로 5개의 에피소드가 자동으로 시각화됩니다. 이 단계를 건너뛰려면 --num-visualizations 0을 사용하세요.

2단계: 주석 확인

python src/lerobot/data_processing/sarm_annotations/subtask_annotation.py \
  --repo-id your-username/your-dataset \
  --visualize-only \
  --visualize-type dense \
  --num-visualizations 5 \
  --video-key observation.images.base \
  --output-dir ./subtask_viz

python src/lerobot/data_processing/sarm_annotations/subtask_annotation.py \
  --repo-id your-username/your-dataset \
  --visualize-only \
  --visualize-type both \
  --num-visualizations 5 \
  --video-key observation.images.base \
  --output-dir ./subtask_viz

…

시각화 인자

팁: 주석이 부정확한 경우, 하위 작업 설명을 더 구체적으로 조정하고 다시 실행하세요.

3단계: SARM 학습

lerobot-train \
  --dataset.repo_id=your-username/your-dataset \
  --reward_model.type=sarm \
  --reward_model.annotation_mode=single_stage \
  --reward_model.image_key=observation.images.base \
  --output_dir=outputs/train/sarm_single \
  --batch_size=32 \
  --steps=5000 \
  --trackio.enable=true \
  --trackio.project=sarm \
  --reward_model.repo_id=your-username/your-model-name

lerobot-train \
  --dataset.repo_id=your-username/your-dataset \
  --reward_model.type=sarm \
  --reward_model.annotation_mode=dense_only \
  --reward_model.image_key=observation.images.base \
  --output_dir=outputs/train/sarm_dense \
  --batch_size=32 \
  --steps=5000 \
  --trackio.enable=true \
  --trackio.project=sarm \
  --reward_model.repo_id=your-username/your-model-name

lerobot-train \
  --dataset.repo_id=your-username/your-dataset \
  --reward_model.type=sarm \
  --reward_model.annotation_mode=dual \
  --reward_model.image_key=observation.images.base \
  --output_dir=outputs/train/sarm_dual \
  --batch_size=32 \
  --steps=5000 \
  --trackio.enable=true \
  --trackio.project=sarm \
  --reward_model.repo_id=your-username/your-model-name

멀티-GPU 학습

여러 GPU를 사용하여 학습하려면 accelerate launch --multi_gpu --num_processes=4를 추가하세요.

학습 인자

4단계: 예측 시각화

compute_rabc_weights.py에 --visualize-only를 사용하여 parquet 파일을 작성하지 않고 모델 예측(및 사용 가능한 경우, 주석에서 파생된 대상)을 시각화합니다.

python -m lerobot.rewards.sarm.compute_rabc_weights \
  --dataset-repo-id your-username/your-dataset \
  --reward-model-path your-username/sarm-model \
  --visualize-only \
  --num-visualizations 5 \
  --head-mode sparse \
  --output-dir ./sarm_viz

python -m lerobot.rewards.sarm.compute_rabc_weights \
  --dataset-repo-id your-username/your-dataset \
  --reward-model-path your-username/sarm-model \
  --visualize-only \
  --num-visualizations 5 \
  --head-mode dense \
  --output-dir ./sarm_viz

python -m lerobot.rewards.sarm.compute_rabc_weights \
  --dataset-repo-id your-username/your-dataset \
  --reward-model-path your-username/sarm-model \
  --visualize-only \
  --num-visualizations 5 \
  --head-mode both \
  --output-dir ./sarm_viz

시각화는 다음을 보여줍니다:

• 진행도 플롯: 예측된 진행도(사용 가능한 경우 및 --stride 1일 때 선택적으로 어노테이션에서 파생된 “GT”)
• 단계별 확률: 예측된 단계별 확률의 누적 영역 플롯
• 샘플 프레임: 에피소드의 주요 프레임으로 진행/단계 라벨이 포함됩니다

시각화 인자

5단계 (선택 사항): RA-BC로 정책 학습

Reward-Aligned Behavior Cloning (RA-BC)는 학습된 SARM 모델을 사용하여 예측된 진행도 개선에 따라 학습 샘플의 가중치를 설정합니다. 이를 위해 두 단계가 필요합니다:

1. 학습된 SARM 모델을 사용하여 모든 프레임에 대해 진행 값을 사전 계산합니다
2. 정책 학습을 사전 계산된 값을 사용하여 RA-BC 가중치와 함께 수행합니다

RA-BC 작동 방식

각 학습 샘플에 대해 RA-BC는 진행 델타를 계산합니다:

r_i = φ(o_{t+Δ}) - φ(o_t)

φ는 SARM 진행도 예측값이며, Δ는 정책의 chunk_size입니다. 진행도가 양수인 샘플(좋은 데모)은 더 높은 가중치를, 진행도가 음수 또는 0인 샘플은 가중치가 낮아집니다.

가중치는 논문에서 식 8-9를 따릅니다:

• 소프트 가중치: w̃_i = clip((r_i − (μ − 2σ)) / (4σ + ε), 0, 1)
• 최종 가중치: w_i = 𝟙{r_i > κ} + 𝟙{0 ≤ r_i ≤ κ} × w̃_i

5a단계: SARM 진행 값 계산

먼저, 데이터셋의 모든 프레임에서 SARM 모델을 실행하여 진행 값(progress values)을 계산합니다:

python -m lerobot.rewards.sarm.compute_rabc_weights \
  --dataset-repo-id your-username/your-dataset \
  --reward-model-path your-username/sarm-model \
  --head-mode sparse \
  --num-visualizations 5 \
  --push-to-hub

이 스크립트는:

• 모든 프레임을 처리하고 진행 값을 계산합니다
• 디스크의 데이터셋 옆에 parquet 파일로 진행 값을 저장합니다(기본값은 <dataset_root>/sarm_progress.parquet)
• 첫 N개의 에피소드(기본값: 5)의 시각화를 생성합니다

인수:

출력 형식 (sarm_progress.parquet):

5b단계: RA-BC로 정책 학습

진행 파일이 준비되면 RA-BC 가중치를 사용하여 정책을 학습합니다. 명시적으로 제공되지 않으면 데이터셋 경로(sarm_progress.parquet)에서 진행 파일이 자동으로 감지됩니다. 현재 RA-BC는 PI0, PI0.5, SmolVLA, FastWAM을 지원합니다:

lerobot-train \
  --dataset.repo_id=your-username/your-dataset \
  --policy.type=pi0 \
  --sample_weighting.type=rabc \
  --sample_weighting.head_mode=sparse \
  --sample_weighting.kappa=0.01 \
  --output_dir=outputs/train/policy_rabc \
  --batch_size=32 \
  --steps=40000

학습 스크립트는 자동으로 다음을 수행합니다:

• parquet 파일에서 사전 계산된 진행 값을 로드합니다
• 정책의 chunk_size를 사용하여 진행 증분(Δ)을 계산합니다
• 진행도 개선에 기반하여 샘플 가중치를 계산합니다
• 학습 중 가중 손실 적용

RA-BC 인자:

RA-BC Kappa 조정

kappa 매개변수는 어떤 샘플이 전체 가중치(w=1)를 부여받는지 결정하는 임계값입니다. 이를 어떻게 조정하는지 이해하는 것은 RA-BC가 효과적으로 작동하는 데 매우 중요합니다.

가중치가 작동하는 방식:

카파 문제 진단:

학습 중에 다음 Trackio 메트릭을 모니터링하세요:

sample_weight_mean_weight ≈ 1.0인 경우: kappa 값이 너무 낮습니다. 대부분의 샘플에서 delta > kappa가 되어 소프트 가중치가 완전히 무시됩니다. RA-BC는 일반 BC와 동일해집니다.

데이터에 기반하여 kappa 설정:

기본값 kappa=0.01은 해당 논문의 T-셔츠 접기 작업(30fps에서 약 90초 에피소드)에 맞게 조정되었습니다. 사용자의 데이터셋에 대해서는 기록된 sample_weighting/delta_mean 및 sample_weighting/delta_std 값을 확인하세요:

# If delta_mean ≈ 0.03 and delta_std ≈ 0.02:
# Most deltas fall in range [0.01, 0.05]

# Option 1: Set kappa = delta_mean (medium selectivity)
--sample_weighting.kappa=0.03

# Option 2: Set kappa = delta_mean + delta_std (high selectivity)
--sample_weighting.kappa=0.05

# Option 3: Set kappa = delta_mean + 2*delta_std (very selective)
--sample_weighting.kappa=0.07

RA-BC가 도움이 되지 않을 수 있는 경우:

데이터셋이 이미 고품질(모든 데모에서 일관된 진행)이라면, 필터링할 대상이 없으므로 RA-BC는 큰 이점을 제공하지 않습니다.

RA-BC를 사용한 멀티 GPU 학습

accelerate launch \
  --multi_gpu \
  --num_processes=4 \
  src/lerobot/scripts/lerobot_train.py \
  --dataset.repo_id=your-username/your-dataset \
  --policy.type=pi0 \
  --sample_weighting.type=rabc \
  --sample_weighting.kappa=0.01 \
  --output_dir=outputs/train/policy_rabc \
  --batch_size=32 \
  --steps=40000

팁 & 모범 사례

모드 선택

• 빠른 실험을 위해 single_stage로 시작하세요 - 어노테이션 오버헤드가 없습니다
• 상세한 진행 상황 추적이 필요하지만 작업에 명확한 상위 단계가 없는 경우 dense_only를 사용하세요
• 와 의 세밀한 진행 상황 모두 의미 있는 복잡한 작업에는 dual을 사용하세요

주석 품질

1. 하위 작업 이름을 구체적으로 지정: “fold” 대신 “grab near side and fold toward center”를 사용
2. 시각화로 검증: 학습 전에 몇 에피소드를 항상 확인하세요
3. 일관된 이름: 모든 에피소드에서 동일한 하위 작업 이름을 사용합니다

RA-BC

1. 먼저 SARM을 학습합니다: RA-BC의 품질은 SARM의 품질에 전적으로 의존합니다
2. sample_weight_mean_weight 모니터링: 값이 약 1.0이라면 kappa를 증가시키세요(RA-BC Kappa 튜닝 참고)

인용

@article{chen2025sarm,
  title={SARM: Stage-Aware Reward Modeling for Long Horizon Robot Manipulation},
  author={Chen, Qianzhong and Yu, Justin and Schwager, Mac and Abbeel, Pieter and Shentu, Yide and Wu, Philipp},
  journal={arXiv preprint arXiv:2509.25358},
  year={2025}
}