Policy Contrastive Decoding for Robotic Foundation Models

1. Énoncé du problème

Les modèles de base robotiques (politiques robotiques généralistes) ont montré un grand potentiel pour la manipulation flexible et polyvalente. Cependant, les auteurs identifient un défaut critique : ces modèles apprennent souvent des corrélations fallacieuses à partir des trajectoires de pré-entraînement. Au lieu de se concentrer sur les caractéristiques des objets pertinentes pour la tâche, les politiques s'appuient sur des indices de fond non pertinents (par exemple, l'éclairage, les textures ou des objets de fond spécifiques) pour prédire les actions.

• Conséquence : Lorsque l'environnement de déploiement change (par exemple, modification du fond, ajustements de l'éclairage ou repositionnement des objets), les performances de la politique se dégradent considérablement car les « raccourcis » appris ne sont plus valables.
• Limitations actuelles : Les solutions existantes nécessitent souvent un réentraînement ou un ajustement fin, ce qui est coûteux en calcul et exige l'accès aux poids du modèle. Il manque des méthodes sans réentraînement, plug-and-play pour corriger ces échecs au moment de l'inférence.

2. Méthodologie : Policy Contrastive Decoding (PCD)

Les auteurs proposent le Policy Contrastive Decoding (PCD), un plugin d'inférence sans réentraînement conçu pour rediriger l'attention d'une politique robotique des caractéristiques fallacieuses vers des indices visuels pertinents pour les objets.

Mécanisme central

Le PCD fonctionne en contrastant les distributions de probabilité d'action générées à partir de deux entrées visuelles différentes :

1. Observation originale ($o_i$) : L'entrée visuelle standard.
2. Observation masquée par l'objet ($\hat{o}_i$) : La même entrée visuelle où les objets cibles pertinents pour la tâche sont masqués (supprimés/peints).

La méthode calcule une nouvelle distribution de probabilité d'action contrastive ($\pi^*$) en utilisant la formule suivante :
$$\pi^*_\theta(a_i | \ell, o_i) = \frac{1}{C} \cdot \pi_\theta(a_i | \ell, o_i) \cdot \left( \frac{\pi_\theta(a_i | \ell, o_i)}{\pi_\theta(a_i | \ell, \hat{o}_i)} \right)^\alpha$$

• Logique : Si la politique s'appuie sur des caractéristiques fallacieuses (le fond), la distribution de probabilité sera similaire pour $o_i$ et $\hat{o}_i$, ce qui donnera un ratio proche de 1 (aucun changement). Si la politique s'appuie sur l'objet, la distribution pour $\hat{o}_i$ différera considérablement, amplifiant la probabilité de l'action correcte dans la sortie finale.
• Hyperparamètre ($\alpha$) : Contrôle l'intensité du contraste. $\alpha=0$ récupère la politique de base.

Composants techniques clés

Pour rendre le PCD applicable à diverses politiques robotiques, les auteurs introduisent deux stratégies spécifiques :

1. Tracking-to-Mask (Track2Mask) :

   • Objectif : Générer automatiquement des observations masquées par l'objet pour chaque image d'une trajectoire sans intervention manuelle.
   • Processus :
     • Annotation : Les objets cibles dans l'image initiale sont identifiés à l'aide de prompts humains (Point/Boîte) ou de détecteurs open-vocabulary hors de la boîte (par exemple, Grounding DINO).
     • Suivi : Le SAM2 (Segment Anything Model 2) suit et segmente ces objets à travers les images suivantes.
     • Inpainting : Les objets segmentés sont peints (supprimés) pour créer l'observation masquée $\hat{o}_i$.

2. Modélisation probabiliste basée sur KDE (KDE-PM) :

   • Défi : Les politiques autorégressives (par exemple, OpenVLA) sortent directement les probabilités de tokens, mais les politiques basées sur la diffusion (par exemple, Octo, $\pi_0$) génèrent des actions via un processus de débruitage et ne sortent pas nativement une distribution de probabilité $\pi_\theta(a|o)$.
   • Solution : Les auteurs utilisent l'Estimation de Densité par Noyaux (KDE) pour approximer la distribution de probabilité d'action. Ils échantillonnent $N$ actions candidates à partir de la politique de diffusion et ajustent une densité de noyau gaussien à ces échantillons pour estimer $\pi_\theta(a|o)$ et $\pi_\theta(a|\hat{o})$, permettant ainsi à la formule de décodage contrastif de fonctionner avec les modèles de diffusion.

3. Contributions clés

1. Première solution sans réentraînement : Le PCD est le premier travail à proposer une approche plug-and-play, sans réentraînement, pour atténuer les corrélations fallacieuses dans les politiques robotiques lors de l'inférence.
2. Indépendant du modèle : Il fonctionne comme un plugin boîte noire compatible avec les politiques autorégressives (OpenVLA) et basées sur la diffusion (Octo, $\pi_0$) sans nécessiter l'accès aux poids pré-entraînés ni d'ajustement fin.
3. Robustesse au décalage de distribution : En contrastant explicitement les entrées pertinentes pour l'objet par rapport aux entrées non pertinentes, le PCD force la politique à se concentrer sur des caractéristiques causales, améliorant considérablement la généralisation à des scénarios non vus.

4. Résultats expérimentaux

Les auteurs ont évalué le PCD sur 15 tâches diverses dans des environnements de simulation (SIMPLER) et réels en utilisant trois politiques de pointe : OpenVLA, Octo et $\pi_0$.

• Performances en simulation :
  • OpenVLA : Amélioration du taux de réussite de 50,6 % en moyenne.
  • Octo : Amélioration du taux de réussite de 29,7 % en moyenne.
  • $\pi_0$ : Amélioration du taux de réussite de 8,9 % en moyenne (malgré le fait que $\pi_0$ soit déjà une base solide).
• Performances dans le monde réel :
  • Sur des tâches de manipulation dans le monde réel utilisant la politique $\pi_0$, le PCD a atteint une amélioration massive de 108 % du taux de réussite par rapport à la base.
  • La méthode s'est révélée robuste face à divers distracteurs (changements d'éclairage, encombrement du fond, déplacements de position des objets).
• Études d'ablation :
  • La méthode est robuste face à différentes stratégies de détection d'objets (prompts humains vs Grounding DINO vs YOLO World).
  • Elle est robuste face à différents algorithmes d'inpainting (LaMa a donné les meilleurs résultats).
  • L'hyperparamètre $\alpha$ montre un impact positif clair lorsqu'il est $>0$.

5. Importance et limites

• Importance : Le PCD aborde une faiblesse fondamentale des modèles de base robotiques actuels (dépendance aux corrélations fallacieuses) sans le coût lourd du réentraînement. Il démontre que de simples interventions au moment de l'inférence peuvent améliorer considérablement la généralisation et la sécurité dans le déploiement réel.
• Limites :
  • Surcharge computationnelle : Le PCD nécessite un deuxième passage avant (pour l'image masquée) ainsi que le suivi et l'inpainting des objets, augmentant la latence d'inférence d'environ 24 à 100 % (selon la politique et le parallélisme). Cependant, les auteurs soutiennent que cela est acceptable étant donné les gains massifs en taux de réussite.
  • Portée : Se concentre actuellement sur la correction au moment de l'inférence ; les travaux futurs visent à traiter les corrélations fallacieuses pendant la phase d'entraînement.

Conclusion : L'article présente un cadre hautement efficace et polyvalent (PCD) qui améliore considérablement la fiabilité et la généralisation des modèles de base robotiques en les forçant à prêter attention aux caractéristiques visuelles correctes, comblant ainsi le fossé entre les performances en simulation et la robustesse dans le monde réel.