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）。这种学习到的“进度信号”有多种用途，其中两个有前景的应用是：(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）进行训练，该处理器：
- 使用 CLIP (ViT-B/32) 对 Encodes 图像和任务文本进行编码，生成
video_features和text_features - 填充/截断机器人状态至
state_features（最多max_state_dim） - 使用 stage+tau 编码
y = k + τ将 构建目标 作为sparse_targets（以及在dense_only/dual中的dense_targets） - 掩码回退帧使用按样本划分的
lengths张量（回退是一种训练时增强）
至少，每个训练样本需要：
task（字符串）：任务描述- 数据集中的
reward_model.image_key图像和reward_model.state_key状态
标注模式
您可以从 3 种标注模式 中进行选择，这些模式决定了进度标签的计算方式：
| 模式 | 需要注释 | 注意力头 | 用例 |
|---|---|---|---|
single_stage | 无 | 仅稀疏 | 简单任务、快速实验，无需 VLM |
dense_only | 稠密（VLM） | 双模式（稀疏自动生成） | 在不定义高级阶段的情况下进行详细的子任务跟踪 |
dual | 稀疏 + 稠密（VLM） | Dual | 包含两种粒度的完整 SARM 论文设置 |
模式详情
single_stage
dense_only
dual
无需标注。整个 episode 被视为一个名为 "task" 的单一阶段，进度在 episode 持续时间内从 0 线性变化到 1。
- 稀疏头：1 个阶段（“任务”），线性进度
- Dense head：未使用
- 最适用于：简单任务、快速实验，或无法使用 VLM 进行标注的情况
设置您的环境
- 按照我们的安装指南安装 LeRobot。
- 运行以下命令安装 SARM 依赖项：
pip install -e ".[sarm]"工作流：
1. Train SARM → 2. Visualize predictions → 3. (Optional) Train policy with RA-BC仅来自 VLM 的稠密（细粒度）标注。稀疏头自动使用覆盖整个 episode 的单个 "task" 阶段，而稠密头则学习详细的子任务进展。
- 稀疏头：1 个阶段（“任务”），线性进度（自动生成）
- 密集头：来自 VLM 注释的多个细粒度阶段
- 最适用于：当你需要详细的子任务跟踪，但不需要定义高层级阶段时
工作流：
1. Annotate (dense) → 2. Verify → 3. Train SARM → 4. Visualize → 5. (Optional) Train policy with RA-BC来自 VLM 的稀疏和稠密标注。 如 SARM 论文中所述，采用完整的双头模式，同时包含高层级（稀疏）和细粒度（稠密）的阶段预测。
- 稀疏头：来自 VLM 标注的高层阶段
- 密集头：来自 VLM 标注的细粒度阶段
- 最适用于：需要同时利用两种粒度来完成的复杂多阶段任务
工作流：
1. Annotate (sparse+dense) → 2. Verify → 3. Train SARM → 4. Visualize → 5. (Optional) Train policy with RA-BC步骤 1：子任务标注
single_stage
dense_only
dual
无需标注！ 完全跳过此步骤。模型将使用剧集的任务描述并自动计算线性进度。
使用 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- 密集时间比例- 每个 episode 的列位于
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 个 episode。使用
--num-visualizations 0可跳过此步骤。
步骤 2：验证注释
single_stage
dense_only
dual
无需验证！ 跳过此步骤。
使用 --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
single_stage
dense_only
dual
使用无标注进行训练 - 使用从 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 文件的情况下可视化模型预测结果（以及，如果可用的话，基于标注生成的目标值）。
single_stage
dense_only
dual
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 训练策略
奖励对齐行为克隆（RA-BC）使用训练好的 SARM 模型，根据预测的进度提升来对训练样本进行加权。这需要两个步骤：
- 使用训练好的 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 | Episode 内的本地帧索引 |
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 表示 kappa 值过低 |
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 恤折叠任务（约 90 秒的回合，30fps）调优的。对于您的数据集，请检查记录的 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}
}