SARM: Pemodelan Hadiah yang Mengetahui Tahap

SARM (Stage-Aware Reward Modeling) adalah kerangka pemodelan reward berbasis video untuk tugas manipulasi robot jangka panjang. Panduan ini membahas cara melatih model reward SARM dan, secara opsional, menggunakannya dengan Reward-Aligned Behavior Cloning (RA-BC).

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

Mengapa Model Reward?

Behavior cloning standar memperlakukan semua frame demonstrasi secara setara, tetapi dataset robot di dunia nyata bersifat berantakan. Dataset tersebut mengandung keraguan, koreksi, dan trajektori dengan kualitas yang bervariasi. Model reward menyelesaikan masalah ini dengan mempelajari konsep progres tugas yang dapat digeneralisasi dari demonstrasi: diberikan frame video dan deskripsi tugas, model memprediksi seberapa dekat robot dalam menyelesaikan tugas (0→1). Sinyal “progres” yang dipelajari ini dapat digunakan dengan berbagai cara, dua aplikasi yang menjanjikan adalah: (1) weighted imitation learning (RA-BC), di mana frame dengan progres tinggi menerima bobot lebih besar selama pelatihan kebijakan, dan (2) reinforcement learning, di mana model reward menyediakan reward padat untuk perbaikan kebijakan secara online atau offline.

Ikhtisar

SARM memiliki fitur-fitur berikut:

1. Arsitektur yang sadar tahap: Memprediksi secara bersama-sama tahap tugas tingkat tinggi dan kemajuan terperinci di dalam setiap tahap
2. Anotasi subtask: Menggunakan anotasi subtask dalam bahasa alami untuk menurunkan label progres yang konsisten
3. Proporsi temporal: Menghitung prior pada tingkat dataset (α̅_k) untuk setiap subtask guna menormalisasi kemajuan di antara demonstrasi dengan panjang yang bervariasi

SARM melatih pada target stage+tau yang ringkas untuk setiap frame:

• stage: indeks stage integer k ∈ {0, ..., K-1}
• τ (tau): kemajuan dalam tahap τ ∈ [0, 1]
• target encoding: y = k + τ (ini adalah apa yang dihasilkan oleh pemroses dataset)

Pada saat inferensi (dan dalam RA-BC turunan), SARM mengonversi nilai mentah k + τ menjadi progres ternormalisasi dalam [0, 1] menggunakan proporsi temporal tingkat dataset α̅_k (disimpan di meta/temporal_proportions_*.json).

Ini sesuai dengan Rumus (2) dari makalah:

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

Di mana:

• τ_t = (t - s_k) / (e_k - s_k) adalah waktu ternormalisasi dalam subtask
• P_{k-1} adalah kumulatif sebelumnya (jumlah proporsi subtask sebelumnya)
• α̅_k adalah proporsi temporal untuk subtask k

Hal ini memastikan bahwa status tugas yang identik dipetakan ke nilai progres yang konsisten, bahkan di antara demonstrasi dengan panjang yang berbeda.

Input dan Target (Apa yang diharapkan kode baru)

SARM dilatih melalui prosesor-nya (src/lerobot/rewards/sarm/processor_sarm.py), yang:

• Mengodekan gambar dan teks tugas dengan CLIP (ViT-B/32) menjadi video_features dan text_features
• Menambahkan padding/memotong status robot menjadi state_features (hingga max_state_dim)
• Membangun target sebagai sparse_targets (dan dense_targets dalam dense_only/dual) menggunakan enkoding stage+tau y = k + τ
• Masks memutar ulang frame menggunakan tensor lengths per sampel (rewind adalah augmentasi saat pelatihan)

Setidaknya, setiap sampel pelatihan membutuhkan:

• task (string): deskripsi tugas
• gambar reward_model.image_key dan state reward_model.state_key dari dataset

Mode Anotasi

Anda dapat memilih dari 3 mode anotasi yang menentukan bagaimana label progres dihitung:

Mode Detail

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

Langkah 1: Anotasi Subtugas

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

Argumen Anotasi

Catatan: Setelah anotasi selesai, 5 episode secara otomatis divisualisasikan secara default. Gunakan --num-visualizations 0 untuk melewati langkah ini.

Langkah 2: Verifikasi Anotasi

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

Ini menghasilkan visualisasi yang menampilkan bingkai video dengan batas subtask yang ditumpangkan serta linimasa subtask.

Argumen Visualisasi

Tips: Jika anotasi tidak akurat, sesuaikan deskripsi subtask Anda agar lebih spesifik dan jalankan ulang.

Langkah 3: Latih 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

Pelatihan Multi-GPU

Tambahkan accelerate launch --multi_gpu --num_processes=4 untuk menggunakan beberapa GPU dalam pelatihan.

Argumen Pelatihan

Langkah 4: Visualisasi Prediksi

Gunakan compute_rabc_weights.py dengan --visualize-only untuk memvisualisasikan prediksi model (dan, jika tersedia, target yang diturunkan dari anotasi) tanpa menulis file 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

Visualisasi menunjukkan:

• Plot progres: Progres yang diprediksi (dan “GT” opsional yang diturunkan dari anotasi jika tersedia dan --stride 1)
• Probabilitas tahap: Plot area bertumpuk dari probabilitas tahap yang diprediksi
• Frame sampel: Frame kunci dari episode dengan label progres/tahap

Argumen Visualisasi

Langkah 5 (Opsional): Latih Policy dengan RA-BC

Reward-Aligned Behavior Cloning (RA-BC) menggunakan model SARM yang telah dilatih untuk memberikan bobot pada sampel pelatihan berdasarkan peningkatan progres yang diprediksi. Hal ini memerlukan dua langkah:

1. Hitung nilai progres untuk semua frame menggunakan model SARM yang telah dilatih
2. Latih kebijakan dengan pembobotan RA-BC menggunakan nilai yang telah dihitung sebelumnya

Cara Kerja RA-BC

Untuk setiap sampel pelatihan, RA-BC menghitung delta progres:

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

Di mana φ adalah prediksi kemajuan SARM dan Δ adalah chunk_size dari kebijakan. Sampel dengan kemajuan positif (demonstrasi yang baik) mendapatkan bobot lebih tinggi, sedangkan sampel dengan kemajuan negatif atau nol mendapatkan bobot yang lebih rendah.

Penimbangan mengikuti Persamaan 8-9 dari makalah:

• Berat lunak: w̃_i = clip((r_i − (μ − 2σ)) / (4σ + ε), 0, 1)
• Bobot akhir: w_i = 𝟙{r_i > κ} + 𝟙{0 ≤ r_i ≤ κ} × w̃_i

Langkah 5a: Hitung Nilai Progres SARM

Pertama, jalankan model SARM pada semua frame dalam dataset Anda untuk menghitung nilai progres:

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

Skrip ini:

• Memproses semua frame dan menghitung nilai progres
• Menyimpan nilai progres ke file parquet di samping dataset di disk (default ke <dataset_root>/sarm_progress.parquet)
• Membuat visualisasi dari N episode pertama (default: 5)

Argumen:

Format output (sarm_progress.parquet):

Langkah 5b: Latih Kebijakan dengan RA-BC

Setelah Anda memiliki file progres, latih kebijakan Anda dengan pembobotan RA-BC. File progres akan terdeteksi secara otomatis dari path dataset (sarm_progress.parquet) jika tidak diberikan secara eksplisit. Saat ini PI0, PI0.5 dan SmolVLA didukung dengan RA-BC:

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

Skrip pelatihan secara otomatis:

• Memuat nilai progres yang telah dihitung sebelumnya dari file parquet
• Menggunakan chunk_size kebijakan untuk menghitung delta kemajuan (Δ)
• Menghitung bobot sampel berdasarkan peningkatan progres
• Menerapkan kerugian berbobot selama pelatihan

Argumen RA-BC:

Menyetel RA-BC Kappa

Parameter kappa adalah ambang batas yang menentukan sampel mana yang mendapatkan bobot penuh (w=1). Memahami cara menyetelnya sangat penting agar RA-BC dapat bekerja secara efektif.

Cara kerja pembobotan:

Mendiagnosis masalah kappa:

Pantau metrik Trackio ini selama pelatihan:

Jika sample_weight_mean_weight ≈ 1.0: kappa Anda terlalu rendah. Sebagian besar sampel memiliki delta > kappa dan melewati soft-weighting sepenuhnya. RA-BC menjadi setara dengan vanilla BC.

Menetapkan kappa berdasarkan data Anda:

Nilai default kappa=0.01 telah disetel untuk tugas melipat kemeja T dalam makalah tersebut (episode ~90 detik pada 30fps). Untuk dataset Anda, periksa nilai sample_weighting/delta_mean dan sample_weighting/delta_std yang tercatat:

# 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

Kapan RA-BC mungkin tidak membantu:

Jika dataset Anda sudah berkualitas tinggi (kemajuan yang konsisten di seluruh demonstrasi), RA-BC tidak akan memberikan banyak manfaat karena tidak ada yang perlu difilter.

Pelatihan Multi-GPU dengan RA-BC

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

Tips & Praktik Terbaik

Memilih Mode

• Mulai dengan single_stage untuk eksperimen cepat - tanpa beban anotasi
• Gunakan dense_only ketika Anda ingin pelacakan progres yang detail tetapi tugas tidak memiliki tahap tingkat tinggi yang jelas
• Gunakan dual untuk tugas kompleks di mana kemajuan kasar dan halus sama-sama bermakna

Kualitas Anotasi

1. Jadilah spesifik dengan nama subtask: Alih-alih “fold”, gunakan “grab near side and fold toward center”
2. Verifikasi dengan visualisasi: Selalu periksa beberapa episode sebelum pelatihan
3. Penamaan yang konsisten: Gunakan nama subtask yang sama di semua episode

RA-BC

1. Latih SARM terlebih dahulu: kualitas RA-BC sepenuhnya bergantung pada kualitas SARM
2. Pantau sample_weight_mean_weight: Jika nilainya ≈ 1.0, tingkatkan kappa (lihat Penyetelan Kappa RA-BC)

Sitasi

@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}
}