LeRobot documentation
SARM: 段階認識型報酬モデリング
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）。この学習された「進捗シグナル」は、複数の方法で利用でき、有望な応用例として以下の2つが挙げられます。(1) 重み付き模倣学習（RA-BC）では、ポリシーの学習時に高進捗のフレームにより大きな重みが付与されます。(2) 強化学習では、報酬モデルがオンラインまたはオフラインのポリシー改善のために密な報酬を提供します。
概要
SARMには以下の機能があります:
- ステージ認識アーキテクチャ: 高レベルのタスクステージと、各ステージ内の詳細な進捗を同時に予測します
- サブタスクアノテーション: 自然言語のサブタスクアノテーションを使用して、一貫性のある進捗ラベルを導出します
- 時間的比率: 各サブタスクについてデータセットレベルの事前分布（α̅_k）を計算し、長さの異なるデモンストレーション間で進捗を正規化します
SARMは、各フレームに対してコンパクトな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/rewards/sarm/processor_sarm.py)を通じて学習されます。このプロセッサは以下を行います:
- Encodes 画像とタスクのテキストを CLIP (ViT-B/32) でエンコードし、
video_featuresとtext_featuresに変換します - パディング/切り詰めロボット状態を
state_featuresに(max_state_dimまで) - ビルド対象を、ステージ+タウ符号化
y = k + τを用いてsparse_targets（およびdense_only/dualの場合のdense_targets）として構築します - Masks rewind framesは、サンプルごとの
lengthsテンソルを使用してフレームを巻き戻します（rewindは学習時の拡張処理です）
最低限、各トレーニングサンプルには以下が必要です:
task(文字列): タスクの説明- データセットから
reward_model.image_keyの画像とreward_model.state_keyの状態
アノテーションモード
進捗ラベルの計算方法を決定する3つのアノテーションモードから選択できます:
| モード | アノテーションが必要 | ヘッド | ユースケース |
|---|---|---|---|
single_stage | None | スパースのみ | シンプルなタスク、クイックな実験、VLMは不要 |
dense_only | Dense (VLM) | デュアル（スパース自動生成） | 高レベルのステージを定義せずに詳細なサブタスクを追跡 |
dual | Sparse + Dense (VLM) | デュアル | 両方の粒度を含む、SARM論文のフルセットアップ |
モードの詳細
アノテーションは不要です。 エピソード全体が"task"と呼ばれる単一のステージとして扱われ、エピソードの持続時間を通じて進捗が0から1まで線形に進行します。
- Sparse head: 1段階（「task」）、線形進捗
- Dense head: 未使用
- 最適な用途: 単純なタスク、迅速な実験、またはVLMアノテーションが利用できない場合
環境のセットアップ
- LeRobotのインストールは、インストールガイドに従って行ってください。
- …を実行して、SARMの依存関係をインストールします。
pip install -e ".[sarm]"ワークフロー:
1. Train SARM → 2. Visualize predictions → 3. (Optional) Train policy with RA-BCVLMからの高密度（細粒度）アノテーションのみ。スパースヘッドはエピソード全体をカバーする単一の"task"ステージを自動的に使用し、デンスヘッドは詳細なサブタスクの進行を学習します。
- Sparse head: 1段階（“task”）、線形進捗（自動生成）
- Dense head: VLMアノテーションからの複数の細粒度ステージ
- 最適な用途: 詳細なサブタスクの追跡が必要だが、高レベルのステージを定義する必要がない場合
ワークフロー:
1. Annotate (dense) → 2. Verify → 3. Train SARM → 4. Visualize → 5. (Optional) Train policy with RA-BCVLMからのスパースアノテーションとデンスアノテーションの両方。 SARM論文で説明されている完全なデュアルヘッドモードで、高レベル（スパース）と細粒度（デンス）のステージ予測の両方を行います。
- スパースヘッド: VLMアノテーションからの高レベルステージ
- Dense head: VLMアノテーションからの細粒度ステージ
- 最適な用途: 両方の粒度が有用な複雑な多段階タスク
ワークフロー:
1. Annotate (sparse+dense) → 2. Verify → 3. Train SARM → 4. Visualize → 5. (Optional) Train policy with RA-BCステップ1: サブタスクのアノテーション
アノテーションは不要です！ このステップは完全にスキップできます。モデルはエピソードのタスク説明を使用して、線形進捗を自動的に計算します。
VLMを使用して高密度（細粒度）のアノテーションのみを生成します。スパースな段階は自動生成されます。
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保存される内容:
meta/temporal_proportions_sparse.json- 自動生成されたスパースな比率（{"task": 1.0}）meta/temporal_proportions_dense.json- 高密度な時間的比率- エピソードごとのカラムは
episodes/*.parquet:dense_subtask_names,dense_subtask_start_frames,dense_subtask_end_frames- (また、時間ベースの列:
dense_subtask_start_times、dense_subtask_end_times)
VLMを使用して、スパース（高レベル）とデンス（細粒度）の両方のアノテーションを生成します。
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保存される内容:
meta/temporal_proportions_sparse.json- 疎な時間的比率meta/temporal_proportions_dense.json- 高密度な時間的比率- の各エピソードの列は
episodes/*.parquet:sparse_subtask_names,sparse_subtask_start_frames,sparse_subtask_end_framesdense_subtask_names,dense_subtask_start_frames,dense_subtask_end_frames- (また、時間ベースの列:
*_subtask_start_times、*_subtask_end_times)
アノテーション引数
| 引数 | 説明 |
|---|---|
--repo-id | HuggingFace データセットリポジトリID |
--sparse-subtasks | ハイレベルなサブタスク名のコンマ区切りリスト |
--dense-subtasks | 細粒度のサブタスク名のコンマ区切りリスト |
--dense-only | スパースな「task」ステージを自動作成し、デンスなアノテーションのみを生成します |
--video-key | 使用するカメラ/ビデオのキー（例: observation.images.top） |
--num-workers | 並列GPUワーカーの数（デフォルト: 1） |
--episodes | アノテーションする特定のエピソードのインデックス（デフォルト: すべて） |
--skip-existing | アノテーションが既に付与されているエピソードをスキップします |
--model | VLMモデル（デフォルト: Qwen/Qwen3-VL-30B-A3B-Instruct） |
--num-visualizations | アノテーション後に可視化するエピソード数（デフォルト: 5、0に設定するとスキップされます） |
注記: アノテーションが完了すると、デフォルトで5エピソードが自動的に可視化されます。このステップをスキップするには、
--num-visualizations 0を使用してください。
ステップ2: アノテーションの検証
検証は不要です！ この手順をスキップしてください。
--visualize-only フラグを使用してアノテーションを可視化します:
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
--visualize-only フラグを使用してアノテーションを可視化します:
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
これにより、サブタスクの境界線が重ねられた動画フレームと、サブタスクのタイムラインを示す可視化が生成されます。
可視化引数
| 引数 | 説明 |
|---|---|
--visualize-only | 既存のアノテーションのみを可視化します（生成は行いません） |
--num-visualizations | 可視化するエピソード数（デフォルト: 5） |
--visualize-type | 可視化するアノテーションの種類: sparse、dense、またはboth |
ヒント: アノテーションが不正確な場合は、サブタスクの説明をより具体的にして再実行してください。
ステップ3: SARMの学習
アノテーションなしで学習 - 0から1までの線形進捗を使用します:
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トレーニング
accelerate launch --multi_gpu --num_processes=4 を追加して、学習に複数の GPU を使用します。
トレーニング引数
| 引数 | 説明 | デフォルト |
|---|---|---|
--reward_model.annotation_mode | single_stage、dense_only、またはdual | single_stage |
--reward_model.image_key | 画像のカメラキー | observation.images.top |
--reward_model.state_key | 関節状態のキー | observation.state |
--reward_model.n_obs_steps | 観測履歴ステップ（総観測フレーム数 = n_obs_steps + 1） | 8 |
--reward_model.frame_gap | サンプリングされた観測値間のギャップ（フレーム単位）（30 fpsの場合: 30 ≈ 1秒） | 30 |
ステップ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」を含む） - ステージ確率: 予測されたステージ確率の積み上げ面積グラフ
- サンプルフレーム: エピソードのキーフレーム（進捗/ステージラベル付き）
可視化引数
| 引数 | 説明 |
|---|---|
--visualize-only | 予測のみを可視化します（RABCの計算は行いません） |
--num-visualizations | 可視化するエピソード数（デフォルト: 5） |
--head-mode | 使用する SARM ヘッド: sparse、dense、または both |
--stride | Nフレームごとに計算し、残りは補間します（デフォルト: 1） |
ステップ5（オプション）: RA-BCでポリシーを学習
Reward-Aligned Behavior Cloning (RA-BC) は、学習済みの SARM モデルを使用して、予測される進捗の改善に基づいてトレーニングサンプルに重みを付けます。これには2つの手順が必要です:
- 進捗値を事前計算する 学習済みSARMモデルを使用してすべてのフレームに対して
- ポリシーの学習を、事前計算された値を用いたRA-BC重み付けで実行します
RA-BCの仕組み
各トレーニングサンプルについて、RA-BCは進捗の差分を計算します:
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 -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）の可視化を生成します
引数:
| 引数 | 説明 | デフォルト |
|---|---|---|
--reward-model-path | 学習済みSARMモデルへのパス | （必須） |
--head-mode | 使用する SARM ヘッド: sparse、dense、または both | sparse |
--device | 推論用のデバイス | cuda |
--visualize-only | 予測のみを可視化します（RA-BCの計算は行いません） | false |
--num-visualizations | 可視化するエピソード数（デフォルト: 5、0に設定するとスキップされます） | 5 |
出力形式 (sarm_progress.parquet):
| カラム | 説明 |
|---|---|
index | データセット内のグローバルフレームインデックス |
episode_index | エピソード番号 |
frame_index | エピソード内のローカルフレームインデックス |
progress_sparse | スパースヘッドの進捗値 [0, 1] |
progress_dense | Dense head の進捗値 [0, 1]（計算されている場合） |
ステップ 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引数:
| 引数 | 説明 | デフォルト |
|---|---|---|
--sample_weighting.type | 重み付け戦略のタイプ（rabcまたはuniform） | rabc |
--sample_weighting.progress_path | 進捗のparquetファイルへのパス | sarm_progress.parquet |
--sample_weighting.head_mode | 使用するSARMヘッドの進捗: sparseまたはdense | sparse |
--sample_weighting.kappa | 高品質サンプルの閾値κ | 0.01 |
--sample_weighting.epsilon | 数値安定性のための小さな定数 | 1e-6 |
RA-BC Kappaのチューニング
kappa パラメータは、どのサンプルに完全な重み (w=1) が付与されるかを決定する閾値です。これを適切に調整する方法を理解することは、RA-BC が効果的に機能するために不可欠です。
重みの仕組み:
| 条件 | 重み |
|---|---|
delta > kappa | 1.0（ハードしきい値） |
0 ≤ delta ≤ kappa | 式8からのソフトウェイト |
delta < 0 | 0.0（負の進捗） |
kappaの問題の診断:
学習中にこれらのTrackioメトリクスを監視してください:
| メトリック | 健全な範囲 | 問題の指標 |
|---|---|---|
sample_weight_mean_weight | 0.3 - 0.8 | ≈ 1.0 はカッパが低すぎることを意味します |
sample_weighting/delta_mean | > 0 | 正である必要があります |
sample_weighting/delta_std | > 0 | データ品質のばらつき |
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.07RA-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を使用します
アノテーションの品質
- サブタスク名は具体的に: “fold”ではなく、“grab near side and fold toward center”のように使用します
- 可視化で確認: 学習前にいくつかのエピソードを確認してください
- 一貫した命名: すべてのエピソードで同じサブタスク名を使用します
RA-BC
- まずSARMを訓練する: RA-BCの品質は完全にSARMの品質に依存します
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}
}