SARM: Modelado de recompensas consciente de etapas

SARM (Stage-Aware Reward Modeling) es un marco de modelado de recompensas basado en vídeo para tareas de manipulación robótica de largo horizonte. Esta guía explica cómo entrenar modelos de recompensa SARM y, opcionalmente, cómo usarlos con Reward-Aligned Behavior Cloning (RA-BC).

Artículo: SARM: Stage-Aware Reward Modeling for Long Horizon Robot Manipulation

¿Por qué modelos de recompensa?

El clonado de comportamiento estándar trata todos los fotogramas de demostración por igual, pero los conjuntos de datos de robots del mundo real son desordenados. Contienen vacilaciones, correcciones y trayectorias de calidad variable. Los modelos de recompensa resuelven esto aprendiendo una noción generalizable de progreso de la tarea a partir de las demostraciones: dados fotogramas de video y una descripción de la tarea, predicen qué tan cerca está el robot de completar la tarea (0→1). Esta “señal de progreso” aprendida puede utilizarse de varias maneras; dos aplicaciones prometedoras son: (1) aprendizaje por imitación ponderado (RA-BC), donde los fotogramas con alto progreso reciben más peso durante el entrenamiento de la política, y (2) aprendizaje por refuerzo, donde el modelo de recompensa proporciona recompensas densas para la mejora de la política en línea o fuera de línea.

Resumen

SARM tiene las siguientes características:

1. Arquitectura consciente de la etapa: Predice conjuntamente la etapa de alto nivel de la tarea y el progreso detallado dentro de cada etapa
2. Anotaciones de subtareas: Utiliza anotaciones de subtareas en lenguaje natural para derivar etiquetas de progreso consistentes
3. Proporciones temporales: Calcula los a priori a nivel de conjunto de datos (α̅_k) para cada subtarea a fin de normalizar el progreso entre demostraciones de longitud variable

SARM entrena en un objetivo compacto stage+tau para cada fotograma:

• stage: índice de etapa entero k ∈ {0, ..., K-1}
• τ (tau): progreso dentro de la etapa τ ∈ [0, 1]
• codificación de objetivo: y = k + τ (esto es lo que produce el procesador de datos)

En el momento de la inferencia (y en el RA-BC de niveles inferiores), SARM convierte el valor crudo k + τ en un progreso normalizado en [0, 1] utilizando las proporciones temporales a nivel de conjunto de datos α̅_k (almacenadas en meta/temporal_proportions_*.json).

Esto coincide con la ecuación (2) del artículo:

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

Donde:

• τ_t = (t - s_k) / (e_k - s_k) es el tiempo normalizado dentro de la subtarea
• P_{k-1} es el previo acumulado (suma de las proporciones de las subtareas anteriores)
• α̅_k es la proporción temporal para la subtarea k

Esto garantiza que los estados de tarea idénticos se mapeen a valores de progreso consistentes, incluso entre demostraciones de diferentes longitudes.

Entradas y objetivos (lo que el nuevo código espera)

SARM se entrena a través de su procesador (src/lerobot/rewards/sarm/processor_sarm.py), que:

• Codifica imágenes y texto de la tarea con CLIP (ViT-B/32) en video_features y text_features
• Rellena/trunca el estado del robot hasta state_features (hasta max_state_dim)
• Construye los objetivos como sparse_targets (y dense_targets en dense_only/dual) utilizando la codificación stage+tau y = k + τ
• Las máscaras retroceden fotogramas mediante un tensor lengths por muestra (el retroceso es una aumentación en tiempo de entrenamiento)

Como mínimo, cada muestra de entrenamiento necesita:

• task (string): descripción de la tarea
• las imágenes de reward_model.image_key y los estados de reward_model.state_key del conjunto de datos

Modos de anotación

Puede elegir entre 3 modos de anotación que determinan cómo se calculan las etiquetas de progreso:

Detalles del modo

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

Paso 1: Anotación de subtareas

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

Argumentos de anotación

Nota: Una vez completada la anotación, se visualizan automáticamente 5 episodios de forma predeterminada. Utiliza --num-visualizations 0 para omitir este paso.

Paso 2: Verificar anotaciones

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

Esto genera visualizaciones que muestran fotogramas de video con los límites de subtareas superpuestos y una línea de tiempo de subtareas.

Argumentos de visualización

Consejo: Si las anotaciones no son precisas, ajusta las descripciones de tus subtareas para que sean más específicas y vuelve a ejecutarlas.

Paso 3: Entrenar 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

Entrenamiento con múltiples GPU

Agregue accelerate launch --multi_gpu --num_processes=4 para usar varias GPU para el entrenamiento.

Argumentos de entrenamiento

Paso 4: Visualizar predicciones

Use compute_rabc_weights.py con --visualize-only para visualizar las predicciones del modelo (y, si están disponibles, los objetivos derivados de las anotaciones) sin escribir un archivo 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

La visualización muestra:

• Gráfico de progreso: Progreso previsto (y “GT” derivado de anotaciones opcional cuando está disponible y --stride 1)
• Probabilidades de etapa: Gráfico de áreas apiladas de las probabilidades de etapa predichas
• Frames de muestra: Frames clave del episodio con etiquetas de progreso/etapa

Argumentos de visualización

Paso 5 (Opcional): Entrenar la política con RA-BC

El clonado de comportamiento alineado con la recompensa (Reward-Aligned Behavior Cloning, RA-BC) utiliza el modelo SARM entrenado para ponderar las muestras de entrenamiento en función de la mejora del progreso predicho. Esto requiere dos pasos:

1. Precomputar los valores de progreso para todos los fotogramas usando el modelo SARM entrenado
2. Entrenar la política con ponderación RA-BC usando los valores precalculados

Cómo funciona RA-BC

Para cada muestra de entrenamiento, RA-BC calcula el delta de progreso:

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

Donde φ es la predicción de progreso de SARM y Δ es el chunk_size de la política. Las muestras con progreso positivo (buenas demostraciones) reciben pesos más altos, mientras que las muestras con progreso negativo o cero reciben pesos más bajos.

La ponderación sigue las ecuaciones 8-9 del artículo:

• Peso suave: w̃_i = clip((r_i − (μ − 2σ)) / (4σ + ε), 0, 1)
• Peso final: w_i = 𝟙{r_i > κ} + 𝟙{0 ≤ r_i ≤ κ} × w̃_i

Paso 5a: Calcular los valores de progreso de SARM

Primero, ejecute el modelo SARM en todos los fotogramas de su conjunto de datos para calcular los valores de progreso:

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

Este script:

• Procesa todos los fotogramas y calcula los valores de progreso
• Guarda los valores de progreso en un archivo parquet junto al conjunto de datos en el disco (por defecto, <dataset_root>/sarm_progress.parquet)
• Genera visualizaciones de los primeros N episodios (predeterminado: 5)

Argumentos:

Formato de salida (sarm_progress.parquet):

Paso 5b: Entrenar la política con RA-BC

Una vez que tengas el archivo de progreso, entrena tu política con ponderación RA-BC. El archivo de progreso se detecta automáticamente desde la ruta del conjunto de datos (sarm_progress.parquet) si no se proporciona explícitamente. Actualmente, PI0, PI0.5, SmolVLA y FastWAM son compatibles con 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

El script de entrenamiento realiza automáticamente:

• Carga los valores de progreso precalculados desde el archivo parquet
• Utiliza el chunk_size de la política para calcular las diferencias de progreso (Δ)
• Calcula los pesos de las muestras en función de la mejora del progreso
• Aplica una pérdida ponderada durante el entrenamiento

Argumentos de RA-BC:

Ajuste de RA-BC Kappa

El parámetro kappa es el umbral que determina qué muestras reciben peso completo (w=1). Comprender cómo ajustarlo es fundamental para que RA-BC funcione de manera efectiva.

Cómo funciona la ponderación:

Diagnóstico de problemas de kappa:

Monitorea estas métricas de Trackio durante el entrenamiento:

Si sample_weight_mean_weight ≈ 1.0: su kappa es demasiado bajo. La mayoría de las muestras tienen delta > kappa y omiten por completo el ponderado suave. RA-BC se vuelve equivalente a un BC estándar.

Configuración de kappa en función de tus datos:

El valor predeterminado kappa=0.01 se ajustó para la tarea de doblado de camisetas del artículo (episodios de ~90s a 30fps). Para su conjunto de datos, revise los valores registrados de sample_weighting/delta_mean y 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

Cuando RA-BC puede no ayudar:

Si su conjunto de datos ya es de alta calidad (progreso constante en todas las demostraciones), RA-BC no aportará mucha ventaja, ya que no hay nada que filtrar.

Entrenamiento multi-GPU con 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

Consejos y mejores prácticas

Elegir un modo

• Comience con single_stage para experimentos rápidos: sin sobrecarga de anotación
• Usa dense_only cuando quieras un seguimiento detallado del progreso pero las tareas no tengan etapas de alto nivel claras
• Use dual para tareas complejas en las que tanto el progreso a gran escala como el de gran detalle son significativos

Calidad de la anotación

1. Sea específico con los nombres de las subtareas: En lugar de “doblar”, use “agarrar el lado cercano y doblar hacia el centro”
2. Verificar con visualización: Siempre comprueba algunos episodios antes de entrenar
3. Nomenclatura coherente: Utilice los mismos nombres de subtareas en todos los episodios

RA-BC

1. Entrena SARM primero: la calidad de RA-BC depende por completo de la calidad de SARM
2. Monitoree sample_weight_mean_weight: si es ≈ 1.0, aumente kappa (consulte Ajuste de Kappa de RA-BC)

Cita

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