SARM：阶段感知奖励建模

SARM（阶段感知奖励建模，Stage-Aware Reward Modeling）是一个面向长时程机器人操作任务、基于视频的奖励建模框架。本指南介绍如何训练 SARM 奖励模型，以及如何选择性地将其与奖励对齐行为克隆（RA-BC）配合使用。

论文：SARM：面向长时程机器人操作的阶段感知奖励建模

为什么要用奖励模型？

标准的行为克隆会将所有 demonstration 帧一视同仁，但现实世界的机器人 dataset 往往良莠不齐，其中包含迟疑、纠错以及质量参差不齐的轨迹。奖励模型通过从 demonstration 中学习一种可泛化的任务进度概念来解决这一问题：给定视频帧和任务描述，它们预测机器人距离完成任务还有多近（0→1）。这种学习到的“进度信号”可以有多种用途，其中两个有前景的应用是：（1）加权模仿学习（RA-BC），即高进度帧在 policy 训练期间获得更高权重；（2）强化学习，即奖励模型为在线或离线的 policy 改进提供稠密奖励。

概览

SARM 具有以下特性：

1. 阶段感知架构：联合预测高层级的任务阶段以及每个阶段内的细粒度进度
2. 子任务标注：使用自然语言的子任务标注来推导一致的进度标签
3. 时间比例：为每个子任务计算 dataset 级别的先验（α̅_k），以在不同长度的 demonstration 之间归一化进度

SARM 在每一帧上针对紧凑的 stage+tau 目标进行训练：

• stage：整数阶段索引 k ∈ {0, ..., K-1}
• τ（tau）：阶段内进度 τ ∈ [0, 1]
• 目标编码：y = k + τ（这是 dataset 处理器生成的）

在 inference 时（以及下游的 RA-BC 中），SARM 使用 dataset 级别的时间比例 α̅_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 的时间比例

这确保了无论 demonstration 长度如何，相同的任务 state 都映射到一致的进度值。

输入与目标（新代码期望的内容）

SARM 通过其处理器（src/lerobot/rewards/sarm/processor_sarm.py）进行训练，该处理器负责：

• 编码：使用 CLIP（ViT-B/32）将图像和任务文本编码为 video_features 和 text_features
• 填充/截断：将机器人 state 填充/截断为 state_features（最多 max_state_dim）
• 构建目标：使用 stage+tau 编码 y = k + τ 构建 sparse_targets（以及 dense_only/dual 模式下的 dense_targets）
• 掩码回退帧：使用每个样本的 lengths 张量（回退是一种训练时的数据增强）

每个训练样本至少需要：

• task（字符串）：任务描述
• 来自 dataset 的 policy.image_key 图像和 policy.state_key state

标注模式

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

模式详情

pip install -e ".[sarm]"

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

第 1 步：子任务标注

标注参数

注意：标注完成后，默认会自动可视化 5 个 episode。使用 --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 步：可视化预测

使用 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

可视化会显示：

• 进度图：预测的进度（当可用且设置 --stride 1 时，还包括由标注推导出的“GT”）
• 阶段概率：预测阶段概率的堆叠面积图
• 示例帧：episode 中带有进度/阶段标签的关键帧

可视化参数

第 5 步（可选）：使用 RA-BC 训练 policy

奖励对齐行为克隆（RA-BC）使用训练好的 SARM 模型，根据预测的进度提升程度对训练样本进行加权。这需要两个步骤：

1. 预先计算进度值：使用训练好的 SARM 模型为所有帧计算进度值
2. 训练 policy：使用预先计算的值以 RA-BC 加权方式训练 policy

RA-BC 的工作原理

对于每个训练样本，RA-BC 计算进度增量：

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

其中 φ 是 SARM 的进度预测，Δ 是 policy 的 chunk_size。进度为正（即好的 demonstration）的样本获得更高权重，而进度为负或零的样本则被降低权重。

权重计算遵循论文中的公式 8-9：

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

第 5a 步：计算 SARM 进度值

首先，在 dataset 的所有帧上运行 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

该脚本会：

• 处理所有帧并计算进度值
• 将进度值保存到磁盘上 dataset 旁边的 parquet 文件中（默认位置为 <dataset_root>/sarm_progress.parquet）
• 生成前 N 个 episode 的可视化（默认：5）

参数：

输出格式（sarm_progress.parquet）：

第 5b 步：使用 RA-BC 训练 policy

获得进度文件后，即可使用 RA-BC 加权训练 policy。如果未明确指定，系统会从 dataset 路径自动检测进度文件（sarm_progress.parquet）。目前 RA-BC 支持 PI0、PI0.5 和 SmolVLA：

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 文件中加载预先计算的进度值
• 使用 policy 的 chunk_size 计算进度增量（Δ）
• 根据进度提升计算样本权重
• 在训练期间应用加权损失

RA-BC 参数：

调整 RA-BC 的 Kappa

kappa 参数是决定哪些样本获得完整权重（w=1）的阈值。了解如何调整它对于 RA-BC 的有效工作至关重要。

加权是如何工作的：

诊断 kappa 问题：

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

如果 sample_weight_mean_weight ≈ 1.0： 你的 kappa 过低。大多数样本的 delta > kappa，完全绕过了软加权。RA-BC 将等同于普通的 BC。

根据你的数据设置 kappa：

默认的 kappa=0.01 是针对论文中 T 恤折叠任务（30fps 下约 90 秒的 episode）调优的。对于你的 dataset，请检查记录的 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.07

RA-BC 可能无法提供帮助的情况：

如果你的 dataset 已经是高质量的（所有 demonstration 的进度一致），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}
}