Performance Asymmetry in Model-Based Reinforcement Learning

1. Problématique : L'Asymétrie de Performance

Les auteurs identifient un problème fondamental dans les évaluations récentes de l'Apprentissage par Renforcement Basé sur un Modèle (MBRL) sur le benchmark Atari100k. Bien que les agents MBRL aient atteint un niveau de performance moyen "surhumain" (Human Normalized Score ou HNS > 1.0), cette moyenne masque une asymétrie de performance critique :

• Tâches "Agent-Optimales" (Agent-Optimal) : Les agents surperforment massivement les humains (parfois de plus d'un ordre de grandeur). Ces tâches sont souvent caractérisées par un espace d'actions réduit et une complexité visuelle élevée mais une dynamique de jeu plus simple.
• Tâches "Humain-Optimales" (Human-Optimal) : Les agents sous-performent drastiquement les humains. Ces tâches sont généralement plus complexes, avec un espace d'actions plus grand et une forte présence de tâches de type "Shooter" (nécessitant des actions de tir précises et temporelles).

Le problème des métriques actuelles : Les métriques d'agrégation conventionnelles (comme la moyenne arithmétique) sont insensibles aux valeurs extrêmes. Une performance exceptionnelle sur quelques tâches faciles (faible score humain) peut compenser mathématiquement un échec total sur la majorité des tâches difficiles (fort score humain), rendant l'asymétrie invisible. De plus, les agents basés sur les pixels (pixel-based) récents, bien que SOTA (State-of-the-Art) globalement, montrent la pire asymétrie, échouant catastrophiquement sur les tâches complexes.

2. Méthodologie et Analyse

A. Redéfinition de l'évaluation

Pour révéler ce problème, les auteurs :

1. Partitionnent Atari100k en deux ensembles équilibrés de 13 tâches chacun : Agent-Optimal et Human-Optimal, basés sur le HNS moyen de plusieurs agents de référence.
2. Proposent une nouvelle métrique : Sym-HNS.
   • Au lieu de la moyenne arithmétique, ils utilisent la moyenne harmonique entre le HNS moyen des tâches Agent-Optimales et celui des tâches Humain-Optimales.
   • Formule : $Sym-HNS = \frac{2 \cdot \mu(AO) \cdot \mu(HO)}{\mu(AO) + \mu(HO)}$
   • Cette métrique pénalise les déséquilibres extrêmes et force les modèles à performer de manière équilibrée sur les deux sous-ensembles.

B. Analyse des causes de l'échec

L'analyse révèle deux facteurs clés expliquant la sous-performance des agents basés sur les pixels (comme DIAMOND) sur les tâches Humain-Optimales :

1. La Malédiction de la Dimensionnalité : Les tâches Humain-Optimales ont un espace d'actions plus grand et plus de tâches de type "Shooter". Les réseaux de RL entrant directement par des pixels ($64 \times 64 \times 3$) peinent à apprendre des politiques efficaces dans ces espaces de recherche exponentiels.
2. Complexité des tâches Shooter : Les tâches impliquant des actions de tir (FIRE) ont une variance de retour très élevée (le timing est crucial). Les agents basés sur les pixels ont du mal à attribuer les probabilités correctes à ces actions, contrairement aux modèles latents qui compressent l'information.

C. Proposition : Le modèle JEDI

Pour résoudre ce problème, les auteurs proposent JEDI (Joint Embedding DIffusion), un nouveau modèle du monde (World Model) :

• Architecture Latente End-to-End : Contrairement aux modèles précédents qui utilisent des décodeurs de reconstruction, JEDI apprend un encodeur latent directement pour la tâche de RL.
• Dynamique de Diffusion Latente : Il intègre un modèle de diffusion (inspiré par EDM) opéré dans l'espace latent comprimé, permettant de modéliser efficacement les distributions stochastiques complexes.
• Entraînement JEPA (Joint Embedding Predictive Architecture) :
  • L'encodeur est entraîné via la fonction de perte de débruitage du modèle de diffusion latent.
  • Innovation clé : C'est la première fois qu'une perte de diffusion "hors de la boîte" (off-the-shelf) est utilisée directement avec un cadre JEPA, sans perte de reconstruction, sans extraction d'états cachés et sans distillation.
  • Des techniques de stabilisation (clamp différentiable, commutation aléatoire des entrées) sont utilisées pour éviter l'effondrement de la représentation.

3. Résultats Clés

Les expériences sur Atari100k (5 graines) démontrent que JEDI inverse la tendance de l'asymétrie de performance :

• Performance Équilibrée (Sym-HNS) : JEDI atteint le niveau SOTA sur la métrique Sym-HNS, surpassant tous les agents précédents (y compris DIAMOND).
• Tâches Humain-Optimales : JEDI est le premier à atteindre une performance surhumaine sur la tâche BankHeist et obtient les meilleurs scores agrégés sur l'ensemble des tâches Humain-Optimales.
• Tâches Visuelles (Agent-Optimales) : JEDI maintient une performance compétitive sur les tâches visuellement complexes (comme Breakout), prouvant que la diffusion latente ne sacrifie pas la fidélité visuelle nécessaire.
• Efficacité Computationnelle : Grâce à la compression latente (12x), JEDI est :
  • 2x plus rapide à l'entraînement.
  • 3x plus rapide à l'inférence.
  • Utilise 43% de mémoire GPU en moins que DIAMOND.
• Robustesse Stochastique : JEDI surpasse significativement les agents latents non-diffusion dans des environnements avec des sauts de trames aléatoires (stochastic frameskip).

4. Contributions Principales

1. Révélation de l'Asymétrie : Identification et quantification du problème de sur-ajustement (overfitting) des agents MBRL sur des sous-ensembles de tâches spécifiques, masqué par les métriques moyennes actuelles.
2. Nouvelle Métrique et Partition : Introduction de la partition Agent/Human-Optimal et de la métrique Sym-HNS pour une évaluation plus juste et équilibrée.
3. Analyse Causale : Démonstration que la malédiction de la dimensionnalité et la complexité des tâches "Shooter" sont les causes principales de l'échec des agents basés sur les pixels sur les tâches difficiles.
4. Modèle JEDI : Proposition du premier modèle du monde latent end-to-end basé sur la diffusion, entraîné via un cadre JEPA sans reconstruction, validé empiriquement comme étant à la fois plus performant, plus équilibré et plus efficace.

5. Signification et Impact

Ce travail remet en question la notion de "progrès" dans le MBRL basée uniquement sur les scores moyens. Il démontre que les modèles les plus récents peuvent être en réalité en régression sur des tâches complexes et critiques.

L'introduction de JEDI prouve qu'il est possible de combiner la puissance générative des modèles de diffusion avec l'efficacité des représentations latentes pour surmonter les limitations des approches purement pixelles. Cela ouvre la voie à des agents plus robustes, capables de généraliser sur des tâches variées et complexes, tout en réduisant les coûts computationnels. La méthodologie d'évaluation proposée (Sym-HNS) devrait devenir un standard pour éviter les biais d'évaluation futurs dans le domaine du RL.