SegDAC: Visual Generalization in Reinforcement Learning via Dynamic Object Tokens

SegDAC est une méthode d'apprentissage par renforcement qui améliore la généralisation visuelle en utilisant des représentations centrées sur des objets dynamiques et textuellement ancrés, permettant aux politiques d'agir directement sur des tokens d'objets sans nécessiter de reconstruction d'image ni de pertes auxiliaires, tout en surpassant significativement les méthodes existantes sur des tâches de manipulation sous diverses perturbations visuelles.

L'essentiel

🎨 Le Problème : Le Robot qui a peur du changement

Imaginez que vous apprenez à un robot à ranger une chambre en lui montrant des milliers de photos. Il devient très fort pour ranger cette chambre précise, avec ce tapis rouge, cette lumière du matin et ces murs blancs.

Mais dès que vous changez la couleur du tapis en bleu, ou que vous allumez une lampe différente, le robot panique. Il ne sait plus quoi faire. Pourquoi ? Parce qu'il a appris à "lire" les pixels de l'image (les petits points de couleur) comme un tableau de bord, au lieu de comprendre les objets qui sont dans la pièce. Pour lui, un cube rouge sur une table bleue est une chose, et un cube rouge sur une table verte est une chose totalement différente, même si la tâche est la même.

C'est ce qu'on appelle un manque de généralisation visuelle.

🧩 La Solution : SegDAC (Le Robot qui voit par "objets")

Les chercheurs ont créé une nouvelle méthode appelée SegDAC. Au lieu de regarder l'image comme une immense mosaïque de pixels, SegDAC apprend à voir le monde comme un ensemble d'objets distincts, un peu comme un enfant qui joue avec des blocs de construction.

Voici comment cela fonctionne, étape par étape, avec des analogies simples :

1. Le Détective Textuel (La Segmentation)

Imaginez que vous donnez au robot une liste de mots-clés, comme une liste de courses : "Robot, Cube, Table, Fond".
Au lieu de chercher à deviner ce qu'il voit, le robot utilise cette liste pour demander à un expert (un modèle d'intelligence artificielle pré-entraîné) : "Montre-moi où est le cube, et où est la table."
Le robot découpe alors l'image en plusieurs morceaux (des masques), un pour chaque objet. Il ne regarde plus l'image entière, mais il a une liste d'objets clairs.

2. Les Cartes de Visite (Les "Tokens")

Pour chaque objet détecté (le cube, le robot, la table), le système crée une petite "carte de visite" numérique.

• L'astuce géniale : Le nombre de cartes n'est pas fixe. Si le robot voit 3 objets, il a 3 cartes. S'il voit 5 objets (parce qu'un nouveau jouet est apparu), il a 5 cartes.
• L'analogie : Imaginez un chef d'orchestre. Si un musicien arrive en retard, le chef ne s'arrête pas. Il s'adapte instantanément au nombre de musiciens présents. SegDAC fait pareil : il s'adapte au nombre d'objets dans la scène, sans être bloqué par un nombre fixe prédéfini.

3. Le Chef d'Orchestre (Le Transformer)

Ces "cartes de visite" (les objets) sont envoyées à un cerveau artificiel très puissant (un Transformer, la même technologie que derrière les chatbots intelligents).
Ce cerveau ne se contente pas de lire les cartes. Il sait aussi où se trouve chaque objet grâce à une étiquette spéciale (l'encodage de position).

• L'analogie : C'est comme si vous aviez une carte au trésor où chaque objet est un point. Le cerveau sait que le "Cube" est à gauche et la "Table" est à droite. Même si le décor change (le tapis devient bleu), la relation entre le Cube et la Table reste la même. Le robot comprend la structure de la scène, pas juste la couleur.

🚀 Pourquoi c'est révolutionnaire ?

Jusqu'à présent, pour rendre les robots robustes, on utilisait deux méthodes qui avaient des défauts :

1. L'augmentation de données : On montrait au robot des milliers de versions de la même image (floue, noire, colorée) pour qu'il s'habitue. C'est lent et ça ne marche pas toujours.
2. Les slots fixes : On forçait le robot à toujours chercher exactement 5 objets, même s'il n'y en avait que 2. C'est rigide et inefficace.

SegDAC change la donne :

• Pas de triche : Il n'a pas besoin de voir des milliers d'images modifiées. Il apprend directement la logique des objets.
• Flexibilité totale : Il gère aussi bien une scène vide qu'une scène encombrée.
• Résultats impressionnants : Sur des tests très difficiles (où l'on change la couleur, la texture, la lumière de manière extrême), SegDAC a surpassé les meilleures méthodes actuelles de 15 % à 88 %. C'est comme si un élève qui apprenait à conduire sur une route sèche pouvait, du jour au lendemain, conduire parfaitement sous la pluie, dans le brouillard et sur une route de gravier, sans jamais avoir pratiqué ces conditions.

🏁 En résumé

SegDAC, c'est comme donner au robot des lunettes qui lui permettent de voir le monde non pas comme une image floue et changeante, mais comme un ensemble d'objets intelligents et connectés.

Au lieu de mémoriser "l'image d'une pièce rouge", il apprend "comment manipuler un cube sur une table". Peu importe que la pièce soit rouge, bleue ou verte, la logique reste la même. C'est une étape majeure pour rendre les robots intelligents capables de fonctionner dans notre monde réel, imprévisible et changeant.

Résumé technique

1. Problématique

Les politiques d'apprentissage par renforcement (RL) visuel, entraînées sur des observations de pixels bruts, souffrent souvent d'un manque de robustesse face aux changements de conditions visuelles lors du déploiement (test). Des variations mineures dans les textures de fond, l'éclairage ou les couleurs des objets peuvent entraîner une chute drastique des performances, même si la structure de la tâche reste inchangée.

Bien que les représentations centrées sur les objets (object-centric) soient une alternative prometteuse pour séparer la structure pertinente de la tâche du bruit visuel, les approches existantes présentent plusieurs limitations :

• Elles utilisent souvent des représentations à taille fixe (slots), ce qui est inadapté aux scènes où le nombre d'objets varie dynamiquement.
• Elles nécessitent des objectifs de reconstruction d'image ou des pertes auxiliaires pour apprendre la décomposition des objets.
• Elles dépendent fréquemment de masques de segmentation ground-truth (annotation manuelle) ou d'augmentations de données massives.

L'objectif est donc de concevoir une méthode capable d'apprendre directement à partir d'entrées de niveau objet, sans ces contraintes, tout en maintenant une efficacité d'échantillonnage (sample efficiency) compétitive.

2. Méthodologie : SegDAC

Les auteurs proposent SegDAC (Segmentation-Driven Actor-Critic), une architecture qui opère sur un ensemble de tokens d'objets de longueur variable. Le pipeline se décompose en trois étapes principales :

A. Segmentation Centrée sur les Objets (Object-Centric Segmentation)

Au lieu d'apprendre la segmentation ou d'utiliser des masques ground-truth, SegDAC utilise un pipeline de segmentation pré-entraîné et figé (frozen) :

1. Détection guidée par le texte : Un détecteur à vocabulaire ouvert (YOLO-World) propose des boîtes englobantes basées sur une liste courte de concepts textuels (ex: "robot", "cube", "fond").
2. Segmentation : Un modèle de segmentation sémantique (EfficientViT-SAM) génère des masques d'instance à l'intérieur de ces boîtes.
3. Raffinement morphologique : Une étape légère d'ouverture et de fermeture morphologique améliore la qualité des masques sans ajouter de latence significative.

• Résultat : Un nombre variable $N$ de segments par image, sans nécessiter d'annotation manuelle ni de reconstruction d'image.

B. Construction de Tokens d'Objets Contextuels

Pour chaque masque détecté, un token d'objet compact est extrait :

• Les patch embeddings d'un encodeur de vision pré-entraîné (frozen) sont filtrés pour ne garder que ceux qui chevauchent le masque de l'objet.
• Un pooling global moyen (global average pooling) est appliqué sur ces patches pour créer un vecteur unique par objet.
• Avantage clé : Ces tokens capturent à la fois la structure locale de l'objet et le contexte global de la scène (grâce à l'attention du transformateur de l'encodeur), sans nécessiter de calculs d'attention supplémentaires entre les objets. Les tokens sont stockés directement dans le buffer de replay, évitant de réexécuter l'encodeur pendant l'entraînement.

C. Actor-Critic basé sur les Transformers

Le cœur de l'apprentissage est un Actor-Critic basé sur un transformateur qui traite la séquence variable de tokens d'objets :

• Encodage de position des segments : Des encodages de position appris, dérivés des coordonnées des boîtes englobantes, sont injectés dans les tokens. Cela permet au modèle de comprendre la localisation spatiale des objets, information souvent perdue dans les features pré-entraînées brutes.
• Traitement de longueur variable : Contrairement aux méthodes qui padding (remplissage) ou tronquent les séquences, SegDAC utilise une stratégie de "packing" (empaquetage). Tous les tokens d'un batch sont concaténés en une seule séquence avec un masque d'attention qui empêche les transitions de différents timesteps de s'interférer. Cela permet de gérer dynamiquement tout nombre d'objets.
• Architecture : Un décodeur de transformateur fusionne les tokens d'objets avec les informations proprioceptives pour prédire les actions (Actor) et les valeurs Q (Critic).
• Entraînement : Uniquement avec la perte SAC (Soft Actor-Critic) standard, sans pertes auxiliaires, reconstruction ou augmentation de données.

3. Contributions Clés

1. Actor-Critic à tokens dynamiques : Une méthode capable d'apprendre des politiques stables à partir d'un ensemble de tokens d'objets dont la taille et le contenu changent à chaque étape, sans reconstruction ni pertes auxiliaires.
2. Tokens contextuels sans fine-tuning : Une méthode pour construire des tokens d'objets à partir de modèles de vision pré-entraînés, utilisant un encodage de position des segments pour préserver l'ancrage spatial, sans masques ground-truth ni ajustement des poids du modèle de vision.
3. Évaluation empirique robuste : Une démonstration que SegDAC améliore la généralisation visuelle par rapport aux méthodes précédentes (jusqu'à +88% sur les scénarios les plus difficiles) tout en égalisant l'efficacité d'échantillonnage de l'état de l'art (DrQ-v2), et ce, sans augmentation de données.
4. Nouveau Benchmark : Introduction d'un benchmark de généralisation visuelle sur 8 tâches de manipulation ManiSkill3, avec 12 types de perturbations visuelles et 3 niveaux de difficulté, organisés selon une taxonomie d'entités de scène.

4. Résultats Expérimentaux

Les expériences ont été menées sur 8 tâches de manipulation (ManiSkill3) avec des perturbations visuelles (caméra, éclairage, couleur, texture) à trois niveaux de difficulté (Facile, Moyen, Difficile).

• Généralisation Visuelle :

  • Sur les paramètres Faciles, SegDAC améliore les méthodes précédentes de 15%.
  • Sur les paramètres Moyens, l'amélioration est de 66%.
  • Sur les paramètres Difficiles (chocs sémantiques et visuels), SegDAC surpasse les méthodes antérieures de 88%.
  • Les méthodes basées sur les pixels (DrQ-v2, SAC-AE) s'effondrent souvent (chute de >90% de performance) sur les scénarios difficiles, tandis que SegDAC maintient une performance robuste.

• Efficacité d'Échantillonnage (Sample Efficiency) :

  • SegDAC atteint une efficacité d'apprentissage comparable à DrQ-v2 (la référence actuelle en efficacité), ce qui est remarquable car les méthodes de généralisation visuelle sacrifient souvent cette efficacité.
  • Contrairement à DrQ-v2 qui dépend fortement de l'augmentation de données, SegDAC atteint ces résultats sans aucune augmentation de données pendant l'entraînement RL.

• Robustesse et Ablations :

  • L'ablation montre que l'encodage de position des segments et le traitement de longueur variable sont tous deux individuellement nécessaires pour des performances élevées.
  • Le modèle est robuste à la variabilité naturelle du nombre de segments détectés (qui peut varier de 5 à 21 objets par épisode), agissant comme une augmentation structurelle implicite.
  • En cas d'échec sous de fortes perturbations, SegDAC dégrade ses performances de manière "gracieuse" (comportement structuré mais manquant de précision) plutôt que de produire des actions erratiques, ce qui est crucial pour le déploiement réel.

5. Signification et Impact

Ce travail démontre que raisonner au niveau des objets plutôt qu'au niveau des pixels offre un biais inductif favorable pour l'apprentissage par renforcement visuel. En exploitant les capacités de segmentation des modèles pré-entraînés modernes (SAM, YOLO-World) sans les réentraîner, SegDAC résout le compromis traditionnel entre généralisation visuelle et efficacité d'échantillonnage.

La méthode propose une voie pratique pour le RL en environnement réel, où les conditions visuelles sont imprévisibles et où l'annotation manuelle (ground-truth) est impossible. En éliminant le besoin d'augmentations de données complexes et de pertes auxiliaires, SegDAC simplifie le pipeline d'entraînement tout en augmentant considérablement la robustesse des agents robotiques face aux changements d'environnement.