SARM：階段感知獎勵建模

SARM (Stage-Aware Reward Modeling) 是一種基於影像的獎勵建模框架，專為長時程機器人操作任務所設計。本指南涵蓋如何訓練 SARM 獎勵模型，以及如何選擇性地將其與獎勵對齊行為複製 (Reward-Aligned Behavior Cloning, RA-BC) 結合使用。

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

為什麼需要獎勵模型？

標準的行為複製 (Behavior Cloning) 對所有示範影格一視同仁，但現實世界中的機器人資料集往往較混亂，包含猶豫、修正和品質不一的軌跡。獎勵模型透過從示範中學習可推廣的任務進度 (task progress) 概念來解決此問題：給定影片影格與任務描述，它們能預測機器人完成任務的接近程度 (0→1)。這種學習到的「進度訊號」可以有多種用途，兩個有前景的應用為：(1) 加權模仿學習 (RA-BC)，在策略訓練期間為進度高的影格賦予更多權重；以及 (2) 強化學習，獎勵模型可為線上或離線策略改進提供密集的獎勵。

概覽

SARM 具備以下特性

1. 階段感知架構：同時預測高階任務階段與各階段內的細粒度進度
2. 子任務註解：利用自然語言子任務註解來導出一致的進度標籤
3. 時間比例：計算每個子任務的資料集層級先驗 (α̅_k)，以在長度不一的示範中正規化進度

SARM 針對每個影格訓練一個簡潔的 階段+tau (stage+tau) 目標

• 階段 (stage)：整數階段索引 k ∈ {0, ..., K-1}
• τ (tau)：階段內進度 τ ∈ [0, 1]
• 目標編碼：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/policies/sarm/processor_sarm.py) 進行訓練，該處理器會：

• 編碼：使用 CLIP (ViT-B/32) 將影像與任務文字編碼為 video_features 與 text_features
• 填充/截斷：將機器人狀態填充或截斷為 state_features（上限為 max_state_dim）
• 建立目標：使用階段+tau 編碼 y = k + τ 作為 sparse_targets（在 dense_only/dual 模式下也包含 dense_targets）
• 遮蔽回溯影格：使用每個樣本的 lengths 張量來遮蔽回溯影格（回溯是一種訓練時的增強方式）

每個訓練樣本至少需要：

• task (字串)：任務描述
• 來自資料集的 policy.image_key 影像與 policy.state_key 狀態

註解模式

您可以選擇 3 種註解模式，這決定了進度標籤的計算方式：

模式詳細資訊

pip install -e ".[sarm]"

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

第 1 步：子任務註解

註解參數

注意：註解完成後，預設會自動視覺化 5 個片段。使用 --num-visualizations 0 可跳過此步驟。

第 2 步：驗證註解

這會產生視覺化內容，顯示疊加了子任務邊界及子任務時間軸的影片影格。

視覺化參數

提示：若註解不準確，請將子任務描述調整得更具體後重新執行。

第 3 步：訓練 SARM

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

多 GPU 訓練

加入 accelerate launch --multi_gpu --num_processes=4 以使用多個 GPU 進行訓練。

訓練參數

第 4 步：視覺化預測結果

使用帶有 --visualize-only 的 compute_rabc_weights.py 來視覺化模型預測（若可用，包含註解產生的目標），而無需寫入 parquet 檔案。

python src/lerobot/policies/sarm/compute_rabc_weights.py \
  --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

視覺化內容顯示：

• 進度圖：預測的進度（以及若可用且 --stride 1 時，額外顯示的註解導出“GT”進度）
• 階段機率：預測階段機率的堆疊面積圖
• 樣本影格：具有進度/階段標籤的片段關鍵影格

視覺化參數

第 5 步（選擇性）：使用 RA-BC 訓練策略

獎勵對齊行為複製 (RA-BC) 使用已訓練的 SARM 模型，根據預測的進度改進情況來對訓練樣本進行加權。這需要兩個步驟：

1. 預先計算進度值：使用已訓練的 SARM 模型計算所有影格的進度值
2. 訓練策略：使用預先計算的數值透過 RA-BC 加權進行訓練

RA-BC 的運作原理

對於每個訓練樣本，RA-BC 會計算進度增量 (progress delta)：

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

其中 φ 是 SARM 的進度預測，Δ 是策略的 chunk_size。進度為正的樣本（良好的示範）會獲得較高的權重，而進度為負或零的樣本則會被降低權重。

加權方式遵循論文中的方程式 8-9

• 軟權重：w̃_i = clip((r_i − (μ − 2σ)) / (4σ + ε), 0, 1)
• 最終權重：w_i = 𝟙{r_i > κ} + 𝟙{0 ≤ r_i ≤ κ} × w̃_i

第 5a 步：計算 SARM 進度值

首先，在您的資料集所有影格上執行 SARM 模型以計算進度值。

python src/lerobot/policies/sarm/compute_rabc_weights.py \
  --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）

引數 (Arguments)

輸出格式 (sarm_progress.parquet)

第 5b 步：使用 RA-BC 訓練策略

取得進度檔案後，即可使用 RA-BC 加權來訓練您的策略。訓練程式會從資料集路徑自動偵測該進度檔案 (sarm_progress.parquet)。目前 PI0、PI0.5 與 SmolVLA 支援 RA-BC。

lerobot-train \
  --dataset.repo_id=your-username/your-dataset \
  --policy.type=pi0 \
  --use_rabc=true \
  --rabc_head_mode=sparse \
  --rabc_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 的有效運作至關重要。

加權運作方式：

診斷 kappa 問題

訓練期間請監控以下 WandB 指標：

若 rabc_mean_weight ≈ 1.0：您的 kappa 太低。大多數樣本具有 delta > kappa，完全繞過了軟加權。RA-BC 會變得與傳統 BC 無異。

根據您的資料設定 kappa

預設的 kappa=0.01 是針對論文中摺疊 T 恤任務（30fps 下約 90 秒的片段）進行調整的。對於您的資料集，請檢查記錄的 rabc_delta_mean 與 rabc_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)
--rabc_kappa=0.03

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

# Option 3: Set kappa = delta_mean + 2*delta_std (very selective)
--rabc_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 \
  --use_rabc=true \
  --rabc_kappa=0.01 \
  --output_dir=outputs/train/policy_rabc \
  --batch_size=32 \
  --steps=40000

技巧與最佳實踐

選擇模式

• 從 single_stage 開始進行快速實驗 - 無需註解開銷
• 當您需要詳細進度追蹤，但任務沒有明確的高階階段時，請使用 dense_only
• 對於複雜任務，若粗略與細粒度進度皆有意義，請使用 dual

註解品質

1. 子任務名稱要具體：不要只用“抓取”，改用“抓取近側並向中心摺疊”
2. 透過視覺化驗證：訓練前務必檢查幾個片段
3. 命名一致性：在所有片段中使用相同的子任務名稱

RA-BC

1. 先訓練 SARM：RA-BC 的品質完全取決於 SARM 的品質
2. 監控 rabc_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}
}