Résumé Technique : Apprentissage de Représentations Structurées avec l'Embedding Linéaire Local et la Fusion de Caractéristiques Adaptative

Énoncé du Problème

L'apprentissage par renforcement (RL) a connu un grand succès dans des domaines tels que la robotique et les jeux, pourtant un défi persistant demeure : le développement de représentations d'états qui maximisent simultanément la récompense et capturent fidèlement la dynamique intrinsèque de l'environnement. Les méthodes traditionnelles optimisent souvent un objectif unique qui confond les signaux de récompense avec les transitions d'état, menant à des représentations qui peuvent négliger la structure géométrique sous-jacente régissant les changements d'état. Bien que l'apprentissage de représentations non supervisé offre une alternative prometteuse, l'application de techniques comme l'Embedding Linéaire Local (LLE) directement au RL en ligne est difficile en raison des distributions d'états non stationnaires. De plus, les méthodes de perte auxiliaire existantes reposent souvent sur l'apprentissage contrastif ou l'augmentation de données conçus pour des entrées visuelles de haute dimension, ce qui peut ne pas exploiter la lissité locale inhérente aux espaces d'états robotiques basés sur la physique.

Méthodologie

Les auteurs proposent un cadre de double représentation fondé sur des principes neuroscientifiques, plus précisément la séparation des flux de traitement spécifiques à la dynamique et de ceux spécifiques à la récompense, tels qu'observés dans les systèmes biologiques. L'architecture se compose de trois composants principaux :

1. Représentation Spécifique à la Dynamique (Embedding Linéaire Local)

Au lieu d'apprendre un modèle de transition global, le cadre extrait des caractéristiques spécifiques à la dynamique ($z^{LLE}$) en utilisant un encodeur paramétré entraîné pour préserver les structures de voisinage locales.

• Mécanisme : Pour un état $x_t$ donné, l'algorithme identifie $K$ voisins temporels les plus proches au sein d'une fenêtre de trajectoire. Il calcule les poids de reconstruction $w_{ij}$ pour minimiser l'erreur de reconstruction de $x_i$ à partir de ses voisins.
• Embedding : Un embedding de faible dimension $z^{LLE}_i$ est appris de telle sorte qu'il préserve ces mêmes poids de reconstruction, imposant une linéarité locale.
• Analogie Neuroscientifique : Cette branche reflète les cortex moteurs et associatifs, qui encodent des trajectoires lisses et de faible dimension de l'activité neuronale.
• Entraînement : L'encodeur est mis à jour périodiquement à l'aide d'un minibatch d'états provenant du tampon de relecture (replay buffer), en minimisant la perte d'embedding $L_z$ et un régularisateur de reconstruction $L_{rec}$ pour éviter l'effondrement de la représentation.

2. Représentation Spécifique à la Récompense

Parallèlement à la branche LLE, une branche spécifique à la récompense apprend des caractéristiques ($z^{Rew}$) en utilisant un réseau de RL profond standard (par exemple, Soft Actor-Critic). Cette branche est optimisée uniquement pour maximiser la récompense cumulée, reflétant les circuits sensibles à la récompense comme les voies dopaminergiques.

3. Fusion de Caractéristiques Adaptative via l'Auto-Attention

Pour intégrer ces représentations complémentaires, le cadre utilise un mécanisme d'auto-attention.

• Processus : Les caractéristiques spécifiques à la dynamique ($z^{LLE}$) et celles spécifiques à la récompense ($z^{Rew}$) sont concaténées et passées à travers une couche d'auto-attention.
• Fonction : Le mécanisme d'attention fusionne de manière adaptative les caractéristiques selon l'état, pondérant dynamiquement l'importance de la structure géométrique locale par rapport aux signaux de récompense en fonction du contexte actuel.
• Analogie Neuroscientifique : Cela agit comme un mécanisme de « portage cortical » (cortical gating), similaire à la façon dont le cortex frontal amplifie sélectivement les informations pertinentes et intègre plusieurs flux lors de la prise de décision.

Procédure d'Entraînement

La boucle d'entraînement entrelace les mises à jour standard de SAC avec des mises à jour périodiques de LLE :

1. Collecte de Données : Les transitions sont stockées dans un tampon de relecture.
2. Mise à jour SAC : La politique et le critique sont mis à jour à l'aide des pertes SAC standard ; l'encodeur de dynamique n'est pas mis à jour par les gradients de récompense lors de cette étape.
3. Mise à jour LLE : Périodiquement, l'encodeur de dynamique est mis à jour en utilisant les objectifs LLE (reconstruction de voisinage et régularisation de reconstruction) sur un minibatch d'états. Ce découplage empêche la tâche auxiliaire d'interférer avec la dynamique d'apprentissage de la politique primaire.

Contributions Clés

1. Cadre de Double Représentation : Une nouvelle architecture de RL qui explicite la séparation des caractéristiques spécifiques à la dynamique (via LLE) de celles spécifiques à la récompense, plutôt que de les confondre dans un seul réseau monolithique.
2. Conception Motivée par les Neurosciences : La séparation des flux et l'utilisation du portage par attention sont explicitement alignées sur les principes biologiques de l'organisation des variétés neuronales et du portage cortical, fournissant un biais inductif biologiquement plausible.
3. Mécanisme de Fusion Adaptative : L'utilisation de l'auto-attention pour fusionner dynamiquement les représentations permet à l'agent de basculer entre la structure géométrique et les indices de récompense selon l'état spécifique, offrant à la fois des bénéfices de performance et une interprétabilité via les cartes d'attention.
4. Efficacité dans les Domaines Structurés : La méthode tire parti de la linéarité locale inhérente aux environnements robotiques basés sur la physique, évitant ainsi le besoin d'une augmentation de données étendue ou d'un échantillonnage contrastif requis par d'autres méthodes d'apprentissage de représentation.

Résultats Expérimentaux

Le cadre a été évalué sur Robosuite (tâches de manipulation robotique modulaire comme Lift, NutAssembly, PickPlace, Stack) et Dexterous Gym (tâches de manipulation de la main Shadow).

• Performance : La méthode proposée SAC-LLE a systématiquement surpassé le SAC vanilla et plusieurs bases de perte auxiliaire (incluant SAC-Recon, SAC-Next, SAC-SPR et SAC-DBC).
• Études d'Ablation :
  • Mécanisme de Fusion : Remplacer l'auto-attention par une concaténation statique ou un simple portage (MLP/Scalaire) a entraîné une performance moindre, confirmant que la fusion adaptative conditionnée par l'état est critique.
  • Régularisation : Supprimer le régularisateur de reconstruction a dégradé la performance, soulignant son rôle dans la prévention de l'effondrement de la représentation.
  • Apprentissage Conjoint : Une variante apprenant conjointement les caractéristiques LLE et de récompense sans fusion explicite (SAC-LLE Joint) a mal performé, suggérant que le découplage des objectifs d'apprentissage est essentiel.
• Interprétabilité : Les visualisations de cartes d'attention ont révélé des modèles comportementaux distincts :
  • Dans les tâches simples (ex: Lift), l'agent s'appuie initialement fortement sur les caractéristiques de récompense et intègre progressivement les caractéristiques LLE au fil de l'épisode.
  • Dans les tâches complexes à contacts élevés (ex: Dexterous Gym), l'agent s'appuie sur les caractéristiques LLE de manière constante tout au long de l'épisode, indiquant leur nécessité pour gérer les dynamiques de contact riches.
• Préservation de la Structure : L'analyse quantitative de l'erreur de reconstruction locale a montré que SAC-LLE maintenait des taux d'erreur nettement inférieurs par rapport au SAC standard, fournissant une preuve directe que la méthode préserve avec succès la structure de voisinage locale dans l'espace d'embedding.

Signification et Revendications

L'article affirme qu'un modèle explicite des structures d'état locales et une fusion adaptative avec l'information de récompense conduisent à une amélioration de l'efficacité de l'apprentissage et de la performance globale dans les domaines robotiques simulés. En ancrant l'architecture dans les principes neuroscientifiques — spécifiquement la séparation du traitement de la dynamique et de la récompense et l'utilisation de mécanismes de portage — les auteurs démontrent que les agents de RL peuvent former des modèles internes plus robustes et adaptables.

Les auteurs soulignent que leur approche est particulièrement efficace dans les environnements où les lois physiques imposent des transitions d'état prévisibles et localement linéaires. Ils reconnaissent des limites, notant que l'hypothèse de linéarité locale peut ne pas tenir dans des environnements chaotiques ou basés sur des pixels bruts, et suggèrent que l'extension du cadre à des dynamiques visuelles non linéaires ou de haute dimension est une direction pour les recherches futures. Ce travail se positionne comme une étape vers le rapprochement entre l'intelligence artificielle et biologique en exploitant les biais structurels observés dans les systèmes neuronaux.