SARM : Modélisation de la récompense tenant compte des étapes

SARM (Stage-Aware Reward Modeling) est un cadre de modélisation de récompense basé sur la vidéo pour les tâches de manipulation robotique à long horizon. Ce guide explique comment entraîner des modèles de récompense SARM et, le cas échéant, les utiliser avec le clonage de comportement aligné sur la récompense (RA-BC).

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

Pourquoi les modèles de récompense ?

Le clonage de comportement standard traite tous les cadres de démonstration de manière égale, mais les jeux de données robotiques réels sont hétérogènes. Ils contiennent des hésitations, des corrections et des trajectoires de qualité variable. Les modèles de récompense résolvent ce problème en apprenant une notion généralisable de progression de la tâche à partir des démonstrations : étant donné des cadres vidéo et une description de la tâche, ils prédisent à quel point le robot est proche de l'accomplissement de la tâche (0→1). Ce « signal de progression » appris peut être utilisé de plusieurs manières, deux applications prometteuses étant : (1) l'apprentissage par imitation pondérée (RA-BC), où les cadres à forte progression reçoivent un poids plus élevé lors de l'entraînement de la politique, et (2) l'apprentissage par renforcement, où le modèle de récompense fournit des récompenses denses pour l'amélioration en ligne ou hors ligne de la politique.

Aperçu

SARM dispose des fonctionnalités suivantes :

1. Architecture consciente des étapes : Prédit conjointement l'étape de haut niveau de la tâche et la progression fine au sein de chaque étape
2. Annotations de sous-tâches : Utilise des annotations de sous-tâches en langage naturel pour dériver des étiquettes de progression cohérentes
3. Proportions temporelles : Calcule les a priori au niveau du jeu de données (α̅_k) pour chaque sous-tâche afin de normaliser la progression à travers des démonstrations de longueurs variables

SARM s’entraîne sur une cible compacte stage+tau pour chaque image :

• stage : index de l'étape entière k ∈ {0, ..., K-1}
• τ (tau) : progression au sein de l'étape τ ∈ [0, 1]
• encodage de la cible : y = k + τ (c'est ce que produit le processeur de jeu de données)

Au moment de l'inférence (et dans les RA-BC en aval), SARM convertit la valeur brute k + τ en progression normalisée dans [0, 1] à l'aide des proportions temporelles au niveau du jeu de données α̅_k (stockées dans meta/temporal_proportions_*.json).

Cela correspond à la Formule (2) de l'article :

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

Où :

• τ_t = (t - s_k) / (e_k - s_k) est le temps normalisé au sein de la sous-tâche
• P_{k-1} est la probabilité a priori cumulée (somme des proportions des sous-tâches précédentes)
• α̅_k est la proportion temporelle pour la sous-tâche k

Cela garantit que des états de tâche identiques correspondent à des valeurs de progression cohérentes, même entre des démonstrations de longueurs différentes.

Entrées et cibles (ce que le nouveau code attend)

SARM est entraîné par son processeur (src/lerobot/rewards/sarm/processor_sarm.py), qui :

• Encode les images et le texte de la tâche avec CLIP (ViT-B/32) en video_features et text_features
• Remplit/tronque l'état du robot en state_features (jusqu'à max_state_dim)
• Construit les cibles sous forme de sparse_targets (et dense_targets en dense_only/dual) en utilisant l'encodage étape+tau y = k + τ
• Les masques rétrogradent les images à l'aide d'un tenseur lengths par échantillon (le rétrogradage est une augmentation appliquée lors de l'entraînement)

Au minimum, chaque exemple d'entraînement nécessite :

• task (string) : description de la tâche
• les images reward_model.image_key et les états reward_model.state_key du jeu de données

Modes d'annotation

Vous pouvez choisir parmi 3 modes d’annotation qui déterminent la manière dont les étiquettes de progression sont calculées :

Détails du mode

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

Étape 1 : Annotation des sous-tâches

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

Arguments d'annotation

Note : une fois l'annotation terminée, 5 épisodes sont visualisés automatiquement par défaut. Utilisez --num-visualizations 0 pour passer cette étape.

Étape 2 : Vérifier les annotations

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

Cela génère des visualisations montrant les images de la vidéo avec les limites des sous-tâches superposées et une chronologie des sous-tâches.

Arguments de visualisation

Astuce : Si les annotations sont imprécises, ajustez vos descriptions de sous-tâches pour qu'elles soient plus spécifiques et relancez l'exécution.

Étape 3 : Entraîner 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

Entraînement multi-GPU

Ajoutez accelerate launch --multi_gpu --num_processes=4 pour utiliser plusieurs GPU pour l'entraînement.

Arguments d'entraînement

Étape 4 : Visualiser les prédictions

Utilisez compute_rabc_weights.py avec --visualize-only pour visualiser les prédictions du modèle (et, si disponibles, les cibles dérivées des annotations) sans écrire de fichier 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 visualisation montre :

• Graph de progression : progression prédite (et « GT » dérivé des annotations facultatives, lorsqu’il est disponible et --stride 1)
• Probabilités par étape : Graphique en aires empilées des probabilités d'étape prédites
• Images d'exemple : Images clés de l'épisode avec des étiquettes de progression/étape

Arguments de visualisation

Étape 5 (Facultatif) : Entraîner la politique avec RA-BC

Le Behavior Cloning aligné sur la récompense (RA-BC) utilise le modèle SARM entraîné pour pondérer les échantillons d'entraînement en fonction de l'amélioration du progrès prédite. Cela nécessite deux étapes :

1. Précalculer les valeurs de progression pour toutes les images à l'aide du modèle SARM entraîné
2. Entraîner la politique avec un pondération RA-BC en utilisant les valeurs précalculées

Fonctionnement de RA-BC

Pour chaque exemple d'entraînement, RA-BC calcule la variation de progression :

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

Où φ est la prédiction de progression SARM et Δ est le chunk_size de la politique. Les échantillons avec une progression positive (bonnes démonstrations) reçoivent des poids plus élevés, tandis que les échantillons avec une progression négative ou nulle voient leur poids diminué.

Le pondération suit les équations 8-9 de l'article :

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

Étape 5a : Calculer les valeurs de progression SARM

Tout d'abord, exécutez le modèle SARM sur toutes les images de votre jeu de données pour calculer les valeurs de progression :

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

Ce script :

• Traite toutes les images et calcule les valeurs de progression
• Enregistre les valeurs de progression dans un fichier parquet à côté du jeu de données sur le disque (par défaut <dataset_root>/sarm_progress.parquet)
• Génère des visualisations des N premiers épisodes (par défaut : 5)

Arguments :

Format de sortie (sarm_progress.parquet) :

Étape 5b : Entraîner la politique avec RA-BC

Une fois que vous disposez du fichier de progression, entraînez votre politique avec le pondération RA-BC. Le fichier de progression est détecté automatiquement à partir du chemin du jeu de données (sarm_progress.parquet) s'il n'est pas fourni explicitement. Actuellement, PI0, PI0.5, SmolVLA et FastWAM sont pris en charge avec 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

Le script d'entraînement effectue automatiquement les opérations suivantes :

• Charge les valeurs de progression précalculées depuis le fichier parquet
• Utilise le chunk_size de la politique pour calculer les deltas de progression (Δ)
• Calcule les poids d'échantillonnage en fonction de l'amélioration du progrès
• Applique une perte pondérée pendant l'entraînement

Arguments RA-BC :

Réglage du Kappa RA-BC

Le paramètre kappa est le seuil qui détermine quels échantillons reçoivent un poids complet (w=1). Comprendre comment le régler est essentiel pour que RA-BC fonctionne efficacement.

Fonctionnement du pondération :

Diagnostic des problèmes de kappa :

Surveillez ces métriques Trackio pendant l’entraînement :

Si sample_weight_mean_weight ≈ 1.0 : votre kappa est trop faible. La plupart des échantillons ont delta > kappa et contournent entièrement le pondération douce. RA-BC devient équivalent à un BC classique.

Définir kappa en fonction de vos données :

La valeur par défaut kappa=0.01 a été ajustée pour la tâche de pliage de T-shirt de l’article (épisodes d’environ 90 s à 30 fps). Pour votre jeu de données, vérifiez les valeurs journalisées de sample_weighting/delta_mean et 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

Cas où RA-BC peut ne pas être utile :

Si votre jeu de données est déjà de haute qualité (progression cohérente à travers toutes les démonstrations), RA-BC n’apportera pas grand-chose, car il n’y a rien à filtrer.

Entraînement multi-GPU avec 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

Astuces et bonnes pratiques

Choisir un mode

• Commencez par single_stage pour des expériences rapides - sans surcharge d'annotation
• Utilisez dense_only lorsque vous souhaitez un suivi détaillé des progrès mais que les tâches ne comportent pas de stades de haut niveau clairement définis
• Utilisez dual pour les tâches complexes où le suivi de la progression à la fois grossière et fine est pertinent

Qualité de l'annotation

1. Soyez précis dans les noms de sous-tâches : au lieu de « plier », utilisez « saisir le côté proche et plier vers le centre »
2. Vérifier par visualisation : vérifiez toujours quelques épisodes avant l'entraînement
3. Nommage cohérent : Utilisez les mêmes noms de sous-tâches dans tous les épisodes

RA-BC

1. Entraînez d'abord SARM : la qualité de RA-BC dépend entièrement de la qualité de SARM
2. Surveiller sample_weight_mean_weight : s’il est ≈ 1,0, augmenter kappa (voir Réglage du kappa RA-BC)

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