ARM-FM: Automated Reward Machines via Foundation Models for Compositional Reinforcement Learning

Ce papier présente ARM-FM, un cadre qui utilise des modèles de fondation pour générer automatiquement des machines de récompense à partir de spécifications en langage naturel, permettant ainsi une conception de récompenses compositionnelle et une généralisation zéro-shot en apprentissage par renforcement.

L'essentiel

Imagine que vous essayez d'enseigner à un robot comment cuisiner un gâteau complexe, mais au lieu de lui donner une recette étape par étape, vous lui dites simplement : « Fais un gâteau ». Si le robot ne reçoit de récompense (un bon point) qu'une fois le gâteau fini, il va probablement passer des heures à mélanger de la farine, à casser des œufs au sol, ou à manger la pâte, sans jamais comprendre qu'il doit d'abord préchauffer le four. C'est le problème principal de l'apprentissage par renforcement (RL) : comment guider l'agent sans le noyer de détails ?

Voici l'explication de la méthode ARM-FM décrite dans ce papier, imagée pour tout le monde.

1. Le Problème : Le Robot Perdu dans le Brouillard

Dans le monde de l'intelligence artificielle, les robots apprennent souvent par essais et erreurs. Mais si la tâche est longue (comme aller chercher un diamant dans un jeu vidéo type Minecraft), les récompenses sont trop rares. Le robot se perd, comme un touriste sans carte dans une ville immense. Il ne sait pas s'il fait du bon travail ou s'il tourne en rond.

2. La Solution : Le « GPS » Généré par une Machine à Rêver

Les auteurs proposent ARM-FM. C'est un système qui utilise un Modèle de Fondation (une IA très intelligente, comme un grand cerveau qui a lu tout Internet) pour créer un GPS pour le robot.

Ce GPS s'appelle une Machine à Récompense (Reward Machine).

• L'analogie : Imaginez que vous donnez à un enfant un puzzle géant. Au lieu de lui dire « Assemble le puzzle », vous lui donnez un tableau avec des étapes claires : « 1. Trouve les pièces bleues. 2. Assemble le ciel. 3. Trouve les pièces vertes. 4. Assemble l'herbe. »
• Le rôle de l'IA (Modèle de Fondation) : Dans le passé, il fallait qu'un humain expert dessine ce tableau à la main, ce qui prenait du temps et était difficile. Ici, vous demandez simplement à l'IA : « Voici le but : aller chercher un diamant. Crée-moi le plan étape par étape. » L'IA génère automatiquement ce plan (la Machine à Récompense).

3. Le Secret : La « Langue des Robots »

C'est ici que la magie opère. La méthode ARM-FM ne se contente pas de créer un plan sec. Elle donne à chaque étape du plan une étiquette en langage naturel (par exemple : « Prends la clé bleue ») et la transforme en une empreinte digitale mathématique (un vecteur d'incrustation).

• L'analogie de la bibliothèque : Imaginez que le robot a une bibliothèque de compétences. Si l'étape actuelle est « Prends la clé bleue », le robot regarde dans sa bibliothèque. Il voit que « Prends la clé rouge » est très proche dans l'esprit (c'est la même action, juste une couleur différente).
• Le résultat : Grâce à cette connexion, si le robot a appris à prendre une clé rouge dans un jeu, il sait instantanément comment prendre une clé bleue dans un nouveau jeu, sans avoir besoin de réapprendre de zéro. C'est ce qu'on appelle la généralisation zéro-shot (réussir une nouvelle tâche sans entraînement supplémentaire).

4. Comment ça marche en pratique ?

Le processus ressemble à une boucle de perfectionnement :

1. Le Chef (L'IA Génératrice) : Vous lui donnez une description en français (« Va chercher le diamant »). Il dessine le plan (la Machine à Récompense) et écrit le code pour vérifier si l'étape est réussie.
2. Le Critique (L'IA Vérificatrice) : Un autre cerveau IA relit le plan pour voir s'il y a des erreurs (ex: « Attends, si le robot perd la clé, il doit revenir en arrière ! »).
3. L'Apprentissage : Le robot joue le jeu. À chaque fois qu'il fait une bonne action (comme trouver la clé), la Machine à Récompense lui donne un petit bonbon (récompense dense) au lieu d'attendre la fin du jeu.
4. Le Guide : Pendant qu'il joue, le robot regarde l'étiquette de l'étape actuelle (« Je dois prendre la clé ») et utilise sa mémoire des actions similaires pour savoir quoi faire.

5. Les Résultats : Des Super-Héros du Jeu Vidéo

Les auteurs ont testé leur méthode sur des environnements très difficiles :

• Des labyrinthes (MiniGrid) : Là où les autres robots se perdaient, le leur trouvait la clé, ouvrait la porte et gagnait.
• Un monde 3D type Minecraft (Craftium) : Le robot a appris à couper du bois, miner de la pierre, forger du fer et enfin trouver le diamant, tout seul, juste en suivant le plan généré par l'IA.
• Des robots réels (Meta-World) : Même avec des bras robotiques complexes, la méthode a permis de résoudre des tâches de manipulation qui étaient jusque-là impossibles.

En Résumé

ARM-FM, c'est comme donner à un robot un tuteur personnel généré par une IA.
Au lieu de le laisser errer dans le noir en attendant une récompense lointaine, l'IA lui donne :

1. Un plan d'action clair (les étapes du jeu).
2. Des petites félicitations à chaque étape réussie.
3. Une mémoire sémantique qui lui permet de transférer ce qu'il a appris sur une tâche à une autre (si je sais ouvrir une porte rouge, je sais ouvrir une porte bleue).

C'est une avancée majeure car cela permet de transformer des objectifs humains complexes (« Construis une maison ») en instructions précises et apprenables pour des machines, sans avoir besoin d'un ingénieur humain pour coder chaque détail.

Résumé technique

1. Problématique

L'apprentissage par renforcement (RL) souffre d'une sensibilité critique à la définition des fonctions de récompense. La conception manuelle de récompenses efficaces pour des tâches complexes est un défi majeur :

• Récompenses éparses : Elles fournissent un signal d'apprentissage insuffisant, empêchant l'agent de progresser (ex: atteindre un but lointain sans feedback intermédiaire).
• Récompenses denses mal conçues : Même conçues à la main, elles sont sujettes à des "failles" (reward hacking) où l'agent exploite la spécification sans atteindre l'objectif réel.
• Gap sémantique : Bien que les Modèles Fondamentaux (FM) excellent à décomposer des tâches à partir de descriptions en langage naturel, il existe un fossé entre cette compréhension abstraite et la génération de signaux de récompense structurés et exécutables nécessaires au RL.

2. Méthodologie : ARM-FM

Le cadre proposé, ARM-FM, comble ce fossé en automatisant la conception de Machines de Récompense (Reward Machines - RMs) via des Modèles Fondamentaux.

A. Machines de Récompense Alignées sur le Langage (LARMs)

Une RM est un automate fini qui décompose une tâche complexe en sous-objectifs séquentiels. ARM-FM introduit les LARMs, qui enrichissent les RMs classiques avec :

1. Instructions en langage naturel : Chaque état de l'automate est associé à une description textuelle du sous-objectif.
2. Embeddings sémantiques : Une fonction $\phi(\cdot)$ transforme ces descriptions textuelles en vecteurs d'embedding ($z_u$). Cela permet de créer un espace de compétences partagé où des sous-tâches sémantiquement similaires (ex: "prendre une clé bleue" vs "prendre une clé rouge") sont proches dans l'espace latent.

B. Processus de Génération Automatisée

La construction des LARMs repose sur une boucle d'auto-amélioration (self-improvement) impliquant deux FM :

• Générateur : Produit la spécification de la RM (structure de l'automate, fonctions de labellisation en Python, instructions).
• Critique : Évalue la correction, la compacité et la logique de la RM générée.
• Boucle itérative : Le critique fournit un feedback pour affiner la RM sur plusieurs tours. Une vérification humaine peut intervenir optionnellement pour corriger des cas limites, mais le système fonctionne souvent sans intervention humaine.

C. Entraînement du RL

L'agent RL est entraîné sur un MDP augmenté ($S \times U$) où l'état inclut l'état de la RM.

• Conditionnement par le langage : La politique de l'agent $\pi(s_t, z_{u_t})$ est conditionnée par l'embedding de l'état courant de la RM. Cela permet à l'agent de comprendre son sous-objectif actif et de transférer des connaissances entre tâches similaires.
• Signal de récompense : La récompense totale est la somme de la récompense de l'environnement (souvent éparse) et de la récompense dense fournie par la RM ($R_{total} = R_{env} + R_{RM}$).

3. Contributions Clés

1. Cadre de génération automatique : Développement d'une méthode pour générer des spécifications de tâches complètes (automate, code de labellisation, instructions) directement à partir de descriptions en langage naturel, sans démonstrations expertes.
2. Espace de compétences sémantique : Introduction d'un mécanisme permettant le partage de connaissances et le transfert de politiques entre tâches via l'alignement des états de la RM avec des embeddings de langage.
3. Validation empirique étendue : Démonstration de l'efficacité sur une suite diversifiée d'environnements, prouvant la capacité du système à gérer des tâches à horizon long et des récompenses éparses là où les méthodes classiques échouent.

4. Résultats Expérimentaux

Les expériences ont été menées sur quatre environnements principaux :

• MiniGrid & BabyAI (Tâches à récompenses éparses) :
  • Sur des tâches complexes comme UnlockToUnlock et KeyCorridor, les agents baselines (DQN standard, DQN+ICM, ReAct) échouent totalement.
  • L'agent ARM-FM résout systématiquement ces tâches avec une efficacité d'échantillonnage bien supérieure, grâce à la densification du signal de récompense.
• Craftium (Environnement 3D procédural type Minecraft) :
  • Dans un monde 3D complexe où l'agent doit miner un diamant (nécessitant d'abord du bois, de la pierre, du fer), un agent PPO standard ne progresse pas.
  • L'agent guidé par une LARM générée automatiquement réussit à accomplir la séquence complète de sous-tâches, démontrant la scalabilité du cadre vers des environnements visuellement complexes.
• Meta-World (Manipulation robotique en contrôle continu) :
  • Sur des tâches de manipulation (ex: Shelf-Place, Stick-Push), l'approche surpasse les agents utilisant uniquement des récompenses éparses, évitant le besoin de conception manuelle complexe de signaux de récompense denses.
• Généralisation Zero-Shot (XLand-MiniGrid) :
  • Ablation : Les résultats montrent que les récompenses structurées ET les embeddings d'état sont tous deux nécessaires pour un apprentissage multi-tâches robuste.
  • Généralisation : Un agent entraîné sur un ensemble de tâches (A, B) peut résoudre une nouvelle tâche composite (C) sans entraînement supplémentaire, à condition que les sous-objectifs de C soient sémantiquement familiers (leurs embeddings sont proches de ceux vus durant l'entraînement).

5. Signification et Impact

ARM-FM représente une avancée significative en connectant le raisonnement sémantique des Modèles Fondamentaux à la rigueur formelle des Machines de Récompense.

• Démocratisation du RL : Il permet de spécifier des objectifs complexes en langage naturel, rendant le RL accessible sans expertise en ingénierie de récompenses.
• Interprétabilité et Modularité : Contrairement aux modèles de récompense "boîte noire" (comme les modèles de récompense appris par RLHF), les LARMs offrent une structure explicite, vérifiable et modifiable par l'humain.
• Apprentissage Compositionnel : En ancrant les politiques dans un espace sémantique, le cadre permet un véritable transfert de compétences et une généralisation zero-shot, ouvrant la voie à des agents capables de comprendre et d'exécuter des plans humains complexes dans des environnements dynamiques.

En résumé, ARM-FM transforme les machines de récompense d'outils statiques en bibliothèques dynamiques de compétences réutilisables, pilotées par l'intelligence artificielle générative.