SARM: Stage-Aware Reward Modeling

← 教程资料库 · 策略与模型

本节介绍 LeRobot 的通用功能。请按目标机器人、数据集和运行环境配置示例；用于 AlohaMini 2 / 2 Pro 前，需要确认硬件与接口兼容性。

---

SARM (Stage-Aware Reward Modeling) is a video-based reward modeling framework for long-horizon robot manipulation tasks. This guide covers how to train SARM reward models and optionally use them with Reward-Aligned Behavior Cloning (RA-BC).

Paper: SARM: Stage-Aware Reward Modeling for Long Horizon Robot Manipulation

Why Reward Models?

Standard behavior cloning treats all demonstration frames equally, but real-world robot datasets are messy. They contain hesitations, corrections, and variable-quality trajectories. Reward models solve this by learning a generalizable notion of task progress from demonstrations: given video frames and a task description, they predict how close the robot is to completing the task (0→1). This learned "progress signal" can be used in multiple ways, two promising applications are: (1) weighted imitation learning (RA-BC), where high-progress frames receive more weight during policy training, and (2) reinforcement learning, where the reward model provides dense rewards for online or offline policy improvement.

Overview

SARM has following features:

1. Stage-aware architecture: Jointly predicts the high-level task stage and fine-grained progress within each stage

2. Subtask annotations: Uses natural language subtask annotations to derive consistent progress labels

3. Temporal proportions: Computes dataset-level priors (α̅_k) for each subtask to normalize progress across variable-length demonstrations

SARM trains on a compact stage+tau target for each frame:

• stage: integer stage index k ∈ {0, ..., K-1}

• τ (tau): within-stage progress τ ∈ [0, 1]

• target encoding: y = k + τ (this is what the dataset processor produces)

At inference time (and in downstream RA-BC), SARM converts the raw k + τ value into a normalized progress in [0, 1] using dataset-level temporal proportions α̅_k (stored in meta/temporal_proportions_*.json).

This matches Formula (2) from the paper:

progress_t = P_{k-1} + α̅_k × τ_t

Where:

• τ_t = (t - s_k) / (e_k - s_k) is within-subtask normalized time

• P_{k-1} is cumulative prior (sum of previous subtask proportions)

• α̅_k is the temporal proportion for subtask k

This ensures identical task states map to consistent progress values, even across demonstrations of different lengths.

Inputs and Targets (What the new code expects)

SARM is trained through its processor (src/lerobot/rewards/sarm/processor_sarm.py), which:

• Encodes images and task text with CLIP (ViT-B/32) into video_features and text_features

• Pads/truncates robot state into state_features (up to max_state_dim)

• Builds targets as sparse_targets (and dense_targets in dense_only/dual) using the stage+tau encoding y = k + τ

• Masks rewind frames using a per-sample lengths tensor (rewind is a training-time augmentation)

At minimum, each training sample needs:

• task (string): task description

• policy.image_key images and policy.state_key states from the dataset

---

Annotation Modes

You can choose from 3 annotation modes that determine how progress labels are computed:

Mode	Annotations Required	Heads	Use Case
single_stage	None	Sparse only	Simple tasks, quick experiments, no VLM needed
dense_only	Dense (VLM)	Dual (sparse auto-generated)	Detailed subtask tracking without defining high-level stages
dual	Sparse + Dense (VLM)	Dual	Full SARM paper setup with both granularities

Mode Details

single_stage

No annotations required. The entire episode is treated as a single stage called "task", and progress is linear from 0 to 1 over the episode duration.

• Sparse head: 1 stage ("task"), linear progress

• Dense head: Not used

• Best for: Simple tasks, quick experiments, or when VLM annotation is not available

Set Up Your Environment

1. Install LeRobot by following our Installation Guide.

2. Install SARM dependencies by running:

pip install -e ".[sarm]"

Workflow:

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

dense_only

Only dense (fine-grained) annotations from a VLM. The sparse head automatically uses a single "task" stage covering the full episode, while the dense head learns detailed subtask progression.

• Sparse head: 1 stage ("task"), linear progress (auto-generated)

• Dense head: Multiple fine-grained stages from VLM annotations

• Best for: When you want detailed subtask tracking but don't need to define high-level stages

Workflow:

1. Annotate (dense) → 2. Verify → 3. Train SARM → 4. Visualize → 5. (Optional) Train policy with RA-BC

dual

Both sparse and dense annotations from VLM. Full dual-head mode as described in the SARM paper, with both high-level (sparse) and fine-grained (dense) stage predictions.

• Sparse head: High-level stages from VLM annotations

• Dense head: Fine-grained stages from VLM annotations

• Best for: Complex multi-stage tasks where both granularities are useful

Workflow:

1. Annotate (sparse+dense) → 2. Verify → 3. Train SARM → 4. Visualize → 5. (Optional) Train policy with RA-BC

---

Step 1: Subtask Annotation

single_stage

No annotation required! Skip this step entirely. The model will use the episode's task description and compute linear progress automatically.

dense_only

Generate dense (fine-grained) annotations only using a VLM. The sparse stage will be auto-generated.

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

What gets saved:

• meta/temporal_proportions_sparse.json - Auto-generated sparse proportions ({"task": 1.0})

• meta/temporal_proportions_dense.json - Dense temporal proportions

• Per-episode columns in episodes/*.parquet:

  • dense_subtask_names, dense_subtask_start_frames, dense_subtask_end_frames

  • (also time-based columns: dense_subtask_start_times, dense_subtask_end_times)

dual

Generate both sparse (high-level) and dense (fine-grained) annotations using a 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

What gets saved:

• meta/temporal_proportions_sparse.json - Sparse temporal proportions

• meta/temporal_proportions_dense.json - Dense temporal proportions

• Per-episode columns in episodes/*.parquet:

  • sparse_subtask_names, sparse_subtask_start_frames, sparse_subtask_end_frames

  • dense_subtask_names, dense_subtask_start_frames, dense_subtask_end_frames

  • (also time-based columns: *_subtask_start_times, *_subtask_end_times)

Annotation Arguments

Argument	Description
--repo-id	HuggingFace dataset repository ID
--sparse-subtasks	Comma-separated list of high-level subtask names
--dense-subtasks	Comma-separated list of fine-grained subtask names
--dense-only	Generate only dense annotations (auto-creates sparse "task" stage)
--video-key	Camera/video key to use (e.g., observation.images.top)
--num-workers	Number of parallel GPU workers (default: 1)
--episodes	Specific episode indices to annotate (default: all)
--skip-existing	Skip episodes that already have annotations
--model	VLM model (default: Qwen/Qwen3-VL-30B-A3B-Instruct)
--num-visualizations	Number of episodes to visualize after annotation (default: 5, set to 0 to skip)

Note: After annotation completes, 5 episodes are automatically visualized by default. Use --num-visualizations 0 to skip this step.

---

Step 2: Verify Annotations

single_stage

No verification needed! Skip this step.

dense_only

Visualize annotations using the --visualize-only flag:

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

dual

Visualize annotations using the --visualize-only flag:

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

This generates visualizations showing video frames with subtask boundaries overlaid and timeline of subtasks.

Visualization Arguments

Argument	Description
--visualize-only	Only visualize existing annotations (no generation)
--num-visualizations	Number of episodes to visualize (default: 5)
--visualize-type	Type of annotations to visualize: sparse, dense, or both

Tip: If annotations are inaccurate, adjust your subtask descriptions to be more specific and re-run.

---

Step 3: Train SARM

single_stage

Train with no annotations - uses linear progress from 0 to 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_only

Train with dense annotations only (sparse auto-generated):

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

dual

Train with both sparse and dense annotations:

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

Multi-GPU Training

Add accelerate launch --multi_gpu --num_processes=4 to use multiple GPUs for training.

Training Arguments

Argument	Description	Default
--policy.annotation_mode	single_stage, dense_only, or dual	single_stage
--policy.image_key	Camera key for images	observation.images.top
--policy.state_key	Key for joint states	observation.state
--policy.n_obs_steps	Observation history steps (total obs frames = n_obs_steps + 1)	8
--policy.frame_gap	Gap (in frames) between sampled observations (at 30 fps: 30 ≈ 1s)	30

---

Step 4: Visualize Predictions

Use compute_rabc_weights.py with --visualize-only to visualize model predictions (and, if available, annotation-derived targets) without writing a parquet file.

single_stage

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

dense_only

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

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 both \
  --output-dir ./sarm_viz

The visualization shows:

• Progress plot: Predicted progress (and optional annotation-derived “GT” when available and --stride 1)

• Stage probabilities: Stacked area plot of predicted stage probabilities

• Sample frames: Key frames from the episode with progress/stage labels

Visualization Arguments

Argument	Description
--visualize-only	Only visualize predictions (no RABC computation)
--num-visualizations	Number of episodes to visualize (default: 5)
--head-mode	SARM head to use: sparse, dense, or both
--stride	Compute every N frames, interpolate the rest (default: 1)

---

Step 5 (Optional): Train Policy with RA-BC

Reward-Aligned Behavior Cloning (RA-BC) uses the trained SARM model to weight training samples based on predicted progress improvement. This requires two steps:

1. Precompute progress values for all frames using the trained SARM model

2. Train policy with RA-BC weighting using the precomputed values

How RA-BC Works

For each training sample, RA-BC computes the progress delta:

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

Where φ is the SARM progress prediction and Δ is the policy's chunk_size. Samples with positive progress (good demonstrations) get higher weights, while samples with negative or zero progress get down-weighted.

The weighting follows Equations 8-9 from the paper:

• Soft weight: w̃_i = clip((r_i − (μ − 2σ)) / (4σ + ε), 0, 1)

• Final weight: w_i = 𝟙{r_i > κ} + 𝟙{0 ≤ r_i ≤ κ} × w̃_i

Step 5a: Compute SARM Progress Values

First, run the SARM model on all frames in your dataset to compute progress values:

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

This script:

• Processes all frames and computes progress values

• Saves progress values to a parquet file next to the dataset on disk (defaults to <dataset_root>/sarm_progress.parquet)

• Generates visualizations of the first N episodes (default: 5)

Arguments:

Argument	Description	Default
--reward-model-path	Path to trained SARM model	(required)
--head-mode	SARM head to use: sparse, dense, or both	sparse
--device	Device for inference	cuda
--visualize-only	Only visualize predictions (no RA-BC computation)	false
--num-visualizations	Number of episodes to visualize (default: 5, set to 0 to skip)	5

Output format (sarm_progress.parquet):

Column	Description
index	Global frame index in dataset
episode_index	Episode number
frame_index	Local frame index within episode
progress_sparse	Sparse head progress value [0, 1]
progress_dense	Dense head progress value [0, 1] (if computed)

Step 5b: Train Policy with RA-BC

Once you have the progress file, train your policy with RA-BC weighting. The progress file is auto-detected from the dataset path (sarm_progress.parquet) if not explicitly provided. Currently PI0, PI0.5 and SmolVLA are supported with 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

The training script automatically:

• Loads the precomputed progress values from the parquet file

• Uses the policy's chunk_size to compute progress deltas (Δ)

• Computes sample weights based on progress improvement

• Applies weighted loss during training

RA-BC Arguments:

Argument	Description	Default
--sample_weighting.type	Weighting strategy type (rabc or uniform)	rabc
--sample_weighting.progress_path	Path to progress parquet file	sarm_progress.parquet
--sample_weighting.head_mode	Which SARM head's progress to use: sparse or dense	sparse
--sample_weighting.kappa	Threshold κ for high-quality samples	0.01
--sample_weighting.epsilon	Small constant for numerical stability	1e-6

Tuning RA-BC Kappa

The kappa parameter is the threshold that determines which samples get full weight (w=1). Understanding how to tune it is critical for RA-BC to work effectively.

How the weighting works:

Condition	Weight
delta > kappa	1.0 (hard threshold)
0 ≤ delta ≤ kappa	Soft weight from Eq. 8
delta < 0	0.0 (negative progress)

Diagnosing kappa issues:

Monitor these WandB metrics during training:

Metric	Healthy Range	Problem Indicator
sample_weight_mean_weight	0.3 - 0.8	≈ 1.0 means kappa too low
sample_weighting/delta_mean	> 0	Should be positive
sample_weighting/delta_std	> 0	Variance in data quality

If sample_weight_mean_weight ≈ 1.0: Your kappa is too low. Most samples have delta > kappa and bypass the soft-weighting entirely. RA-BC becomes equivalent to vanilla BC.

Setting kappa based on your data:

The default kappa=0.01 was tuned for the paper's T-shirt folding task (~90s episodes at 30fps). For your dataset, check the logged sample_weighting/delta_mean and 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

When RA-BC may not help:

If your dataset is already high quality (consistent progress across all demonstrations), RA-BC won't provide much benefit since there's nothing to filter.

Multi-GPU Training with RA-BC

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

---

Tips & Best Practices

Choosing a Mode

• Start with single_stage for quick experiments - no annotation overhead

• Use dense_only when you want detailed progress tracking but tasks don't have clear high-level stages

• Use dual for complex tasks where both coarse and fine-grained progress is meaningful

Annotation Quality

1. Be specific with subtask names: Instead of "fold", use "grab near side and fold toward center"

2. Verify with visualization: Always check a few episodes before training

3. Consistent naming: Use the same subtask names across all episodes

RA-BC

1. Train SARM first: RA-BC quality depends entirely on SARM quality

2. Monitor sample_weight_mean_weight: If it's ≈ 1.0, increase kappa (see Tuning RA-BC Kappa)

---

Citation

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