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}
• τ (тау): прогресс внутри этапа τ ∈ [0, 1]
• целевое кодирование: y = k + τ (это то, что генерирует процессор набора данных)

На этапе вывода (и в downstream RA-BC) SARM преобразует исходное значение k + τ в нормализованный прогресс в диапазоне [0, 1] с использованием временных пропорций на уровне набора данных α̅_k (сохранённых в meta/temporal_proportions_*.json).

Это соответствует формуле (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)
• Строит целевые значения в виде sparse_targets (а также dense_targets в режимах dense_only/dual) с использованием кодирования stage+tau y = k + τ
• Маски перемотки кадров с использованием тензора lengths для каждого примера (перемотка — это аугментация на этапе обучения)

Минимально каждый образец для обучения должен содержать:

• task (строка): описание задачи
• изображения reward_model.image_key и состояния reward_model.state_key из набора данных

Режимы аннотирования

Вы можете выбрать один из 3 режимов аннотирования, которые определяют, как вычисляются метки прогресса:

Детали режима

pip install -e ".[sarm]"

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

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

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

Шаг 1: Аннотация подзадач

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

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

Аргументы аннотации

Примечание: после завершения аннотирования по умолчанию автоматически визуализируются 5 эпизодов. Используйте --num-visualizations 0, чтобы пропустить этот шаг.

Шаг 2: Проверка аннотаций

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

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

Это генерирует визуализации, показывающие кадры видео с наложенными границами подзадач и временную шкалу подзадач.

Аргументы визуализации

Совет: если аннотации неточны, уточните описания подзадач и запустите процесс повторно.

Шаг 3: Обучение SARM

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 для обучения.

Аргументы обучения

Шаг 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

Визуализация показывает:

• Диаграмма прогресса: предсказанный прогресс (и необязательное «GT», полученное из аннотаций, если доступно и --stride 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

Шаг 5а: Вычисление значений прогресса 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)

Аргументы:

Формат вывода (sarm_progress.parquet):

Шаг 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:

Настройка RA-BC Kappa

Параметр kappa — это порог, определяющий, какие примеры получают полный вес (w=1). Понимание того, как его настраивать, критически важно для эффективной работы RA-BC.

Как работает взвешивание:

Диагностика проблем с kappa:

Отслеживайте эти метрики Trackio во время обучения:

Если sample_weight_mean_weight ≈ 1.0: ваш коэффициент каппа слишком низкий. У большинства выборок delta > kappa, и они полностью обходят мягкое взвешивание. RA-BC становится эквивалентным обычному BC.

Настройка kappa на основе ваших данных:

Значение по умолчанию kappa=0.01 было подобрано для задачи складывания футболок из статьи (эпизоды ~90 с при 30 fps). Для вашего набора данных проверьте значения 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 может не помочь:

Если ваш датасет уже высокого качества (последовательный прогресс во всех демонстрациях), RA-BC не принесёт значительной пользы, так как фильтровать нечего.

Обучение на нескольких GPU с 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

Советы и лучшие практики

Выбор режима

• Начните с single_stage для быстрых экспериментов — без затрат на разметку
• Используйте dense_only, если вам нужен детальный контроль прогресса, но у задач нет чётко выраженных высокоуровневых этапов
• Используйте dual для сложных задач, где осмысленно отслеживание прогресса как на грубом, так и на тонком уровне

Качество аннотаций

1. Указывайте конкретные названия подзадач: вместо «fold» используйте «grab near side and fold toward center»
2. Проверка с помощью визуализации: всегда проверяйте несколько эпизодов перед обучением
3. Единое именование: Используйте одинаковые имена подзадач во всех эпизодах

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