Coupled Local and Global World Models for Efficient First Order RL

Résumé Technique : Couplage de Modèles de Monde Locaux et Globaux pour un Apprentissage par Renforcement du Premier Ordre Efficace

Énoncé du Problème

L'apprentissage par renforcement (RL) a connu un succès significatif dans les tâches de locomotion via des simulateurs GPU parallélisés, mais il se heurte à des obstacles substantiels en matière de contrôle perceptuel et de manipulation. Deux problèmes principaux entravent la progression :

1. Inefficacité d'échantillonnage : L'apprentissage en ligne pratique directement à partir d'interactions avec de vrais robots est souvent irréalisable en raison du coût élevé de la collecte de données.
2. Fidélité de la simulation : Simuler des tâches de manipulation est intrinsèquement plus complexe que la locomotion. L'état de l'environnement change dynamiquement avec les actions de l'agent (ex: objets déformables, dynamique granulaire), introduisant une variabilité infinie que les simulateurs physiques actuels peinent à modéliser avec une fidélité suffisante pour le transfert sim-to-real.

Bien que l'apprentissage d'agents à l'intérieur de modèles de monde différentiables via des méthodes de Gradient du Premier Ordre (FoG) offre une voie pour résoudre l'inefficacité d'échantillonnage, les approches existantes font face à un goulot d'étranglement de scalabilité. Les modèles de monde à grande échelle (ex: modèles de diffusion) requis pour une dynamique du monde réel de haute fidélité sont coûteux à différencier, rendant la rétropropagation standard par backpropagation intraitable pour l'entraînement des politiques FoG. Inversement, les modèles légers adaptés à la différenciation manquent souvent d'expressivité pour capturer les dynamiques complexes du monde réel.

Méthodologie

Le papier propose un nouveau cadre qui découple la simulation directe de la différenciation indirecte pour permettre un apprentissage par renforcement basé sur des modèles de premier ordre (FoG-MBRL) efficace utilisant des modèles de monde à grande échelle. L'architecture centrale se compose de deux composants couplés :

1. Modèle de Monde Global (Passage Avant) :

   • Un modèle à grande échelle, entraîné hors ligne (soit un modèle basé sur la diffusion comme DIAMOND, soit un modèle de diffusion latente de type transformer comme DreamerV4), agit comme un simulateur « global ».
   • Ce modèle opère dans l'espace des pixels (ou un espace latent de haute dimension) pour générer des trajectoires directes ($s^{\rightarrow}_{t+1}$) de haute fidélité qui capturent précisément les dynamiques visuelles et de contact complexes.
   • Il est entraîné sur des données de jeu non structurées et à haute entropie collectées sur de vrais robots, éliminant ainsi le besoin de simulateurs physiques conçus à la main.

2. Modèle de Monde Local (Passage Arrière) :

   • Un modèle léger de type RSSM (Recurrent State Space Model) agit comme un substitut « local » dans un espace latent de faible dimension.
   • Ce modèle approxime la dynamique locale et fournit les gradients nécessaires pour l'optimisation de la politique.
   • Crucialement, les gradients sont calculés en évaluant les Jacobiens du modèle local aux points d'état générés par le modèle global ($s^{\rightarrow}_{t+1}$), plutôt qu'aux prédictions du modèle local lui-même. Ce découplage permet au système de tirer parti de la précision du modèle global pour la génération de trajectoires tout en utilisant le modèle local pour une estimation de gradient calculable.

3. Protocole d'Entraînement :

   • Collecte de Données : Des données de jeu non structurées sont collectées sur de vrais robots sans segmentation stricte des tâches.
   • Modélisation de la Récompense : Selon la tâche, les signaux de récompense sont dérivés de fonctions d'énergie contrastives (entraînées sur des données de jeu/démonstration) ou de têtes de classification légères identifiant les comportements pertinents de la tâche.
   • Optimisation : La politique est optimisée entièrement à l'intérieur du modèle de monde en utilisant l'algorithme de l'optimisation de politique basée sur des modèles déconnectés (DMO - Decoupled forward-backward Model-based policy Optimization). Le modèle local est continuellement affiné pendant l'entraînement RL pour maintenir sa précision par rapport à la distribution d'actions de la politique actuelle.

Contributions Clés

• Cadre FoG-MBRL Découplé : Les auteurs introduisent une méthode qui sépare la simulation directe de haute fidélité (via de grands modèles de diffusion/flux) de la différenciation indirecte efficace (via des RSSM légers). Cela résout le compromis entre la complexité de la modélisation et le coût du calcul des gradients.
• Apprentissage Sans Simulateur à Partir de Données Réelles : L'approche entraîne des agents entièrement à partir de données de robots réels, contournant le besoin de simulateurs physiques et traitant directement le fossé sim-to-real.
• Déploiement Zero-Shot : Le cadre démontre la capacité d'entraîner des politiques dans le modèle de monde et de les déployer sur de vrais robots sans ajustement supplémentaire ou guidage par imitation de comportement (behavior cloning).
• Robustesse aux Dynamiques Complexes : La méthode gère avec succès des tâches impliquant une observabilité partielle, des dynamiques de contact complexes et des degrés de liberté élevés (ex: manipulation humanoïde).

Résultats Expérimentaux

Les auteurs ont évalué leur approche sur trois tâches robotiques réelles de complexité croissante :

1. Push-T : Une tâche de manipulation sur table consistant à pousser un objet en forme de T.
2. Ego-Centric Push Cube : Un robot quadrupède poussant un cube vers un objectif, impliquant des dynamiques de marche et des occlusions partielles.
3. Humanoid Grasp and Lift : Un humanoïde à haut degré de liberté saisissant et soulevant une boîte avec une main dextre.

Principaux Résultats :

• Efficacité d'Échantillonnage et Temporelle : La méthode DMO proposée a nettement surpassé PPO (Proximal Policy Optimization) en termes d'efficacité d'échantillonnage et de temps de calcul (wall-clock time) sur toutes les tâches. PPO a eu du mal à converger ou a nécessité un nombre excessif d'échantillons, tombant souvent dans des solutions locales sous-optimales.
• Nécessité des Modèles Globaux : Les études d'ablation (« No Diffusion »), remplaçant le modèle de diffusion global par un RSSM standard pour les passages directs, ont entraîné un échec total de la tâche (taux de succès de 0 %), confirmant que les modèles directs de haute fidélité sont essentiels pour la manipulation complexe.
• Supériorité sur le Behavior Cloning : Les politiques apprises ont surpassé les bases de comparaison de Behavior Cloning (Action Chunking Transformer), démontrant la capacité à découvrir des stratégies à long terme et à gérer des régions de l'espace d'état mal couvertes par les démonstrations.
• Succès Zero-Shot : Les politiques entraînées entièrement dans le modèle de monde ont atteint des taux de succès élevés (ex: 9/10 sur Push-T, 8/10 sur Humanoid Grasp) lors de leur déploiement sur de vrais robots sans entraînement réel supplémentaire.

Signification et Revendications

Le papier affirme que l'apprentissage d'agents entièrement à l'intérieur de modèles de monde basés sur les données est une voie réalisable et prometteuse pour résoudre des tâches difficiles à modéliser dans l'espace image. En découplant les exigences de précision directe et de tractabilité indirecte, les auteurs démontrent que les modèles de monde de grande échelle et coûteux en calcul peuvent être efficacement utilisés pour l'optimisation de politiques efficaces.

Ce travail se positionne comme une solution à la dépendance envers les simulateurs physiques faits main, offrant une méthode capable de passer à l'échelle pour des environnements complexes et non structurés grâce à l'acquisition de données réelles. Les auteurs notent que, bien que leur implémentation actuelle nécessite l'entraînement de modèles globaux à partir de zéro et repose sur des récompenses étiquetées manuellement pour certaines tâches, le cadre constitue un proxy sûr pour l'entraînement de robots réels, permettant une optimisation de politique scalable sans risque de dommages matériels. Les travaux futurs visent à exploiter les modèles vidéo pré-entraînés à l'échelle d'Internet pour réduire davantage les besoins en données et automatiser l'étiquetage des récompenses.