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]
• 目标编码：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）
• 使用 stage+tau 编码 y = k + τ 构建 sparse_targets（以及 dense_only/dual 模式下的 dense_targets）
• 使用逐样本的 lengths 张量对回退帧进行掩码（rewind 是训练时的数据增强）

每个训练样本至少需要：

• task（字符串）：任务描述
• 数据集中的 policy.image_key 图像和 policy.state_key 状态

---

标注模式

你可以从 3 种标注模式中进行选择，它们决定了如何计算进度标签：

Mode	所需标注	Heads	适用场景
single_stage	无	仅 sparse	简单任务、快速实验、不需要 VLM
dense_only	Dense（VLM）	Dual（sparse 自动生成）	需要细粒度子任务跟踪但无需定义高层阶段
dual	Sparse + Dense（VLM）	Dual	完整复现 SARM 论文设置，兼顾两种粒度

模式详情

single_stagedense_onlydual

无需任何标注。 整个 episode 会被视为一个名为 "task" 的单一阶段，进度会在 episode 持续时间内从 0 线性变化到 1。

• Sparse head：1 个阶段（"task"），线性进度
• Dense head：不使用
• 最适合：简单任务、快速实验，或无法使用 VLM 标注时

设置环境

1. 按照我们的安装指南安装 LeRobot。
2. 运行以下命令安装 SARM 依赖：

pip install -e ".[sarm]"

工作流：

1. 训练 SARM → 2. 可视化预测 → 3.（可选）使用 RA-BC 训练策略

仅使用来自 VLM 的 dense（细粒度）标注。 sparse head 会自动使用覆盖整个 episode 的单一 "task" 阶段，而 dense head 则学习详细的子任务进展。

• Sparse head：1 个阶段（"task"），线性进度（自动生成）
• Dense head：来自 VLM 标注的多个细粒度阶段
• 最适合：希望获得细粒度子任务跟踪，但不需要定义高层阶段时

工作流：

1. 标注（dense）→ 2. 验证 → 3. 训练 SARM → 4. 可视化 → 5.（可选）使用 RA-BC 训练策略

同时使用来自 VLM 的 sparse 和 dense 标注。 这是论文中描述的完整双头模式，同时预测高层（sparse）和细粒度（dense）阶段。

• Sparse head：来自 VLM 标注的高层阶段
• Dense head：来自 VLM 标注的细粒度阶段
• 最适合：复杂多阶段任务，且两种粒度都很有价值时

工作流：

1. 标注（sparse+dense）→ 2. 验证 → 3. 训练 SARM → 4. 可视化 → 5.（可选）使用 RA-BC 训练策略

---

步骤 1：子任务标注

single_stagedense_onlydual

无需标注！ 可以完全跳过此步骤。模型会使用 episode 的任务描述并自动计算线性进度。

使用 VLM 仅生成 dense（细粒度）标注。sparse 阶段会自动生成。

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 - 自动生成的 sparse 比例（{"task": 1.0}）
• meta/temporal_proportions_dense.json - dense 时间比例
• episodes/*.parquet 中的逐 episode 列：
• dense_subtask_names, dense_subtask_start_frames, dense_subtask_end_frames
• （以及基于时间的列：dense_subtask_start_times, dense_subtask_end_times）

使用 VLM 同时生成 sparse（高层）和 dense（细粒度）标注。

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 - sparse 时间比例
• meta/temporal_proportions_dense.json - dense 时间比例
• episodes/*.parquet 中的逐 episode 列：
• sparse_subtask_names, sparse_subtask_start_frames, sparse_subtask_end_frames
• dense_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	仅生成 dense 标注（自动创建 sparse "task" 阶段）
--video-key	使用的相机/视频键（例如 observation.images.top）
--num-workers	并行 GPU worker 数量（默认：1）
--episodes	要标注的特定 episode 索引（默认：全部）
--skip-existing	跳过已有标注的 episode
--model	VLM 模型（默认：claude/claude3-VL-30B-A3B-Instruct）
--num-visualizations	标注后可视化的 episode 数量（默认：5，设为 0 可跳过）

注意：标注完成后，默认会自动可视化 5 个 episode。使用 --num-visualizations 0 可以跳过此步骤。

---

步骤 2：验证标注

single_stagedense_onlydual

无需验证！ 跳过此步骤即可。

使用 --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	要可视化的 episode 数量（默认：5）
--visualize-type	要可视化的标注类型：sparse、dense 或 both

提示：如果标注不准确，请将子任务描述写得更具体，然后重新运行。

---

步骤 3：训练 SARM

single_stagedense_onlydual

在无标注情况下训练——使用从 0 到 1 的线性进度：

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

仅使用dense 标注训练（sparse 自动生成）：

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

同时使用sparse 和 dense 标注训练：

lerobot-train \
  --dataset.repo_id=your-username/your-dataset \
  --policy.type=sarm \
  --policy.annotation_mode=dual \
  --policy.image_key=observation.images.base \
  --output_dir=outputs/train/sarm_dual \
  --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。

训练参数

参数	描述	默认值
--policy.annotation_mode	single_stage、dense_only 或 dual	single_stage
--policy.image_key	图像对应的相机键	observation.images.top
--policy.state_key	关节状态对应的键	observation.state
--policy.n_obs_steps	观测历史步数（总观测帧数 = n_obs_steps + 1）	8
--policy.frame_gap	采样观测之间的帧间隔（30fps 下，30 ≈ 1 秒）	30

---

步骤 4：可视化预测

使用带 --visualize-only 的 compute_rabc_weights.py 可以可视化模型预测（如果可用，还会显示由标注推导出的目标值），且不会写入 parquet 文件。

single_stagedense_onlydual

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	要可视化的 episode 数量（默认：5）
--head-mode	使用哪一个 SARM head：sparse、dense 或 both
--stride	每 N 帧计算一次，其余帧插值（默认：1）

---

步骤 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。进度为正的样本（好的演示）会获得更高权重，而进度为负或零的样本则会被降低权重。

权重计算遵循论文中的公式 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 个 episode 生成可视化（默认：5）

参数：

参数	描述	默认值
--reward-model-path	训练好的 SARM 模型路径	（必需）
--head-mode	使用哪一个 SARM head：sparse、dense 或 both	sparse
--device	推理所用设备	cuda
--visualize-only	仅可视化预测（不执行 RA-BC 计算）	false
--num-visualizations	要可视化的 episode 数量（默认：5，设为 0 可跳过）	5

输出格式（sarm_progress.parquet）：

列名	描述
index	数据集中的全局帧索引
episode_index	Episode 编号
frame_index	Episode 内的局部帧索引
progress_sparse	sparse head 的进度值 [0, 1]
progress_dense	dense head 的进度值 [0, 1]（如果已计算）

步骤 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 \
  --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 head 的进度：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 问题：

在训练期间监控这些 WandB 指标：

指标	健康范围	问题信号
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 秒的 episode，30fps）。对于你的数据集，请查看日志中的 sample_weighting/delta_mean 和 sample_weighting/delta_std：

# 如果 delta_mean ≈ 0.03 且 delta_std ≈ 0.02：
# 大多数 delta 落在 [0.01, 0.05] 范围内

# 方案 1：设置 kappa = delta_mean（中等选择性）
--sample_weighting.kappa=0.03

# 方案 2：设置 kappa = delta_mean + delta_std（高选择性）
--sample_weighting.kappa=0.05

# 方案 3：设置 kappa = delta_mean + 2*delta_std（非常严格）
--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. 通过可视化进行验证：训练前务必检查几个 episode
3. 命名保持一致：所有 episode 中都使用相同的子任务名称

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