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

1. Énoncé du problème

Le papier aborde le défi critique de la manipulation robotique à long horizon et riche en contacts, impliquant spécifiquement des objets déformables (par exemple, le pliage d'un t-shirt). Bien que les modèles récents de comportement robotique (RBMs) aient montré des promesses, ils peinent face à :

• Une qualité de démonstration incohérente : Les grands ensembles de données contiennent souvent des trajectoires bruitées, sous-optimales ou échouées provenant d'opérateurs non experts.
• Un étiquetage fragile : Les méthodes existantes reposent souvent sur des indices de trame ou des étiquettes de progression basées sur le temps. Dans des tâches à durées variables et à progression non linéaire (par exemple, aplatir un t-shirt prend des quantités de temps différentes selon le tissu), l'étiquetage basé sur les trames introduit un bruit sévère, rendant impossible la distinction entre une étape "lente mais correcte" et une étape "échouée".
• La généralisation : Les modèles entraînés sur des données curatées échouent à se généraliser à des scénarios hors distribution (OOD), tels que la récupération après erreur ou la gestion des occlusions dans les objets déformables.

Le problème central réside dans la manière de dériver des signaux de récompense cohérents et fins à partir de démonstrations diverses et de longueurs variables pour guider efficacement l'apprentissage de la politique, sans nécessiter de fonctions de récompense parfaites et conçues à la main.

2. Méthodologie

Les auteurs proposent un cadre en deux parties : SARM (Stage-Aware Reward Modeling) pour générer des récompenses, et RA-BC (Reward-Aligned Behavior Cloning) pour utiliser ces récompenses afin d'entraîner des politiques.

A. SARM : Modélisation de la récompense consciente des étapes

SARM est un cadre de modélisation de récompense basé sur la vidéo conçu pour prédire la progression de la tâche de manière robuste.

• Annotations en langage naturel : Au lieu d'indices de trames, la méthode utilise des annotations de sous-tâches en langage naturel (par exemple, "saisir", "aplatir", "plier") pour décomposer les tâches en étapes sémantiques.
• Architecture à double tête :
  1. Estimateur d'étape : Une tête de classification qui prédit l'étape de tâche de haut niveau actuelle (par exemple, "Phase d'aplanissement") basée sur les observations visuelles (trames RGB) et les états articulaires.
  2. Estimateur de sous-tâche : Une tête de régression qui prédit la progression fine ($\in [0, 1]$) conditionnée par l'étape prédite.
• Stratégie d'étiquetage :
  • Les auteurs calculent une proportion a priori au niveau de l'ensemble de données ($\bar{\alpha}_k$) pour chaque sous-tâche, basée sur la durée temporelle moyenne de cette sous-tâche à travers l'ensemble de données.
  • Les étiquettes de progression sont générées par interpolation linéaire au sein des sous-tâches, garantissant que des états sémantiques identiques (par exemple, un t-shirt entièrement aplati) reçoivent des valeurs de progression cohérentes, indépendamment du nombre de trames nécessaires pour y parvenir.
• Traitement des entrées : Le modèle utilise un encodeur CLIP figé pour les caractéristiques visuelles et un encodeur Transformer pour capturer les dépendances temporelles. Il emploie l'augmentation par rebobinage (ajout de trames inversées provenant d'horodatages antérieurs) pour améliorer la robustesse face aux cas d'échec et aux scénarios OOD.

B. RA-BC : Clonage de comportement aligné sur la récompense

RA-BC est un cadre d'entraînement qui exploite le modèle de récompense SARM appris pour filtrer et repondérer les données de démonstration lors de l'entraînement de la politique.

• Calcul de la différence de progression : Pour chaque échantillon d'entraînement, le modèle calcule l'amélioration de la progression ($\Delta \phi$) entre l'état actuel et l'état suivant en utilisant le modèle de récompense SARM.
• Pondération dynamique :
  • Le système maintient des statistiques de fonctionnement en ligne (moyenne $\mu$ et écart-type $\sigma$) des différences de progression.
  • Les échantillons avec une progression positive (au-dessus d'un seuil) sont surpondérés, tandis que ceux avec une progression négative (régressions/échecs) sont sous-pondérés ou filtrés.
  • Un mécanisme de "remplacement a priori" attribue un poids de 1 aux échantillons clairement bons et 0 aux échantillons clairement mauvais pour assurer un filtrage décisif.
• Objectif : La perte standard de clonage de comportement (BC) est remplacée par une perte pondérée, enseignant efficacement à la politique d'imiter uniquement les segments de haute qualité et progressifs de l'ensemble de données diversifié.

3. Contributions clés

1. Cadre SARM : Un modèle de récompense vidéo nouveau, conscient des étapes, qui prédit conjointement les étapes de tâche et la progression fine en utilisant des annotations en langage naturel. Il surmonte la fragilité de l'étiquetage par indice de trame et se généralise bien aux scénarios OOD.
2. Cadre RA-BC : Une méthode d'entraînement de politique qui utilise le modèle de récompense appris pour identifier et repondérer automatiquement les démonstrations de haute qualité, permettant un apprentissage efficace à partir d'ensembles de données bruités et diversifiés sans filtrage manuel.
3. Validation empirique sur des objets déformables : Application réussie à la tâche hautement difficile du pliage de t-shirt, démontrant des améliorations significatives par rapport au clonage de comportement standard et à d'autres bases de modélisation de récompense.

4. Résultats expérimentaux

Les auteurs ont évalué leur approche sur deux tâches : le pliage de t-shirt (long horizon, déformable) et le déchargement de vaisselle (multi-étapes, rigide).

• Performance du modèle de récompense :

  • SARM a nettement surpassé les bases de comparaison (LIV, VLC, GVL, REDS, ReWiND) dans la prédiction de la progression sur les démonstrations humaines et les déployements réels de robots.
  • Sur le pliage de t-shirt, SARM a atteint un MSE de 0,008 sur les données de validation (contre 0,019 pour la meilleure base de comparaison suivante) et un score de classification de 0,94 sur les déployements réels de robots (contre 0,78 pour ReWiND).
  • SARM était robuste aux variations d'éclairage et aux modes d'échec où d'autres modèles produisaient des signaux bruités ou optimistes.

• Performance de la politique (Pliage de t-shirt) :

  • Bases de comparaison : Le BC vanilla entraîné sur l'ensemble complet de données diversifiées ($D_{all}$) a atteint 0% de réussite sur la tâche "Difficile" (pliage à partir d'un état froissé).
  • Filtrage naïf : Le filtrage par durée ($D_{2min}$) a amélioré les performances mais a toujours échoué sur la tâche Difficile (0% de réussite).
  • RA-BC avec SARM :
    • Tâche Moyenne (Plat à Plié) : 83% de taux de réussite.
    • Tâche Difficile (Froissé à Plié) : 67% de taux de réussite.
  • Comparaison : RA-BC avec SARM a nettement surpassé RA-BC avec le modèle de récompense ReWiND (67% contre 25% sur la tâche Difficile), prouvant que la qualité du modèle de récompense est le goulot d'étranglement pour RA-BC.

• Études d'ablation :

  • Conscience des étapes : La suppression de l'estimateur d'étape a dégradé les performances, confirmant la nécessité d'une modélisation hiérarchique.
  • Augmentation par rebobinage : Essentielle pour la généralisation aux cas d'échec réels ; sans elle, le modèle devenait excessivement optimiste.
  • Échelle du modèle : Un modèle de 60 millions de paramètres a été identifié comme le compromis optimal entre capacité et surapprentissage.

5. Signification

Ce papier fournit une solution évolutive et efficace en termes d'annotation pour la manipulation robotique à long horizon.

• Efficacité des données : Il démontre qu'un apprentissage de politique de haute qualité est possible même avec des ensembles de données bruités et diversifiés si le signal de récompense peut distinguer avec précision entre "progression lente" et "échec".
• Évolutivité : En exploitant des annotations en langage naturel (plus faciles à obtenir que des fonctions de récompense denses) et en automatisant la dérivation de la récompense, la méthode s'étend à des tâches complexes comme la manipulation d'objets déformables.
• Impact pratique : Les résultats sur le pliage de t-shirt (67% de réussite à partir d'un état froissé) représentent un bond en avant significatif pour la robotique domestique, un domaine où les modèles actuels échouent généralement en raison de la complexité des interactions riches en contacts et déformables.

Ce travail établit que la modélisation de la récompense est un goulot d'étranglement critique dans l'apprentissage par imitation et que l'estimation de la progression sémantique consciente des étapes est la clé pour débloquer des comportements robotiques à long horizon robustes.