🤖 machine learning

JEDI: Joint Embedding Diffusion World Model for Online Model-Based Reinforcement Learning

JEDI introduit le premier modèle de monde latent de diffusion en bout en bout en ligne qui unifie l'apprentissage de représentations prédictives de type JEPA avec des objectifs de diffusion pour atteindre une efficacité supérieure et des performances compétitives en apprentissage par renforcement basé sur le modèle par rapport aux approches basées sur les pixels et aux approches latentes entraînées séparément.

Auteurs originaux : Jing Yu Lim, Rushi Shah, Zarif Ikram, Samson Yu, Haozhe Ma, Tze-Yun Leong, Dianbo Liu

Publié 2026-05-14

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Jing Yu Lim, Rushi Shah, Zarif Ikram, Samson Yu, Haozhe Ma, Tze-Yun Leong, Dianbo Liu

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous essayez d'enseigner à un robot comment jouer à des jeux vidéo. Pour ce faire efficacement, le robot a besoin d'un « modèle du monde » — une simulation mentale du fonctionnement du jeu afin qu'il puisse s'entraîner dans sa tête avant de jouer réellement.

Pendant longtemps, la meilleure façon de construire cette simulation mentale consistait à demander au robot de tenter de reconstruire l'écran du jeu pixel par pixel, comme un peintre essayant de copier une photo à l'identique. C'est précis, mais très lent et cela nécessite une quantité massive de mémoire informatique (comme essayer de transporter une bibliothèque dans son sac à dos).

Récemment, une nouvelle technique appelée Diffusion est devenue populaire. Imaginez la diffusion comme un sculpteur qui commence avec un bloc d'argile rempli de bruit et de statique, et qui élimine lentement ce bruit pour révéler une statue claire. C'est excellent pour comprendre des scènes complexes, mais le faire pixel par pixel reste trop lourd pour qu'un robot puisse l'exécuter en temps réel.

Une autre approche a tenté de compresser le jeu dans un « résumé » minuscule et abstrait (un espace latent) au préalable, mais ces résumés étaient souvent entraînés séparément et n'apprenaient pas suffisamment bien les règles du jeu par eux-mêmes.

Voici JEDI (Joint Embedding Diffusion).

Les auteurs de cet article ont créé une nouvelle méthode qui combine le meilleur des deux mondes. Voici comment JEDI fonctionne, en utilisant des analogies simples :

1. La « Photo Mentale » contre la « Peinture »

L'Ancienne Méthode (Diffusion de Pixels) : Imaginez le robot essayant de mémoriser chaque grain de sable d'une plage pour comprendre l'océan. C'est précis, mais cela prend une éternité et épuise toute sa puissance de calcul.
La Méthode de JEDI : Au lieu de mémoriser le sable, JEDI apprend au robot à prendre une photo mentale de l'essence de l'océan (les vagues, la couleur, le mouvement). Il ne se soucie pas des grains de sable individuels. Il compresse l'écran du jeu en un « résumé » minuscule et efficace qui ne capture que ce qui compte pour gagner.

2. L'Entraînement par le « Jeu de Devinettes »

Comment le robot apprend-il à prendre ces photos ?

L'Ancienne Méthode : Le robot était souvent entraîné en se voyant montrer une image et en lui demandant de la repeindre exactement. S'il manquait un détail, il était pénalisé. C'est comme si un élève était noté sur son écriture plutôt que sur sa compréhension de l'histoire.
La Méthode de JEDI : JEDI utilise un jeu de devinettes prédictif.
1. Le robot voit l'état actuel du jeu.
2. On lui demande de deviner à quoi ressemblera la prochaine « photo mentale ».
3. Pour rendre cela plus difficile (et plus intelligent), l'ordinateur prend la prochaine photo réelle, y ajoute du « bruit statique » (comme un flou), et demande au robot de le dénouer pour retrouver la clarté.
4. Crucialement, le robot apprend cela pendant qu'il joue. Il n'a pas besoin d'un enseignant séparé pour lui montrer comment peindre ; il apprend les règles du jeu en essayant de prédire l'avenir.

3. Pourquoi c'est une Grande Nouvelle (Les Résultats)

L'article affirme que JEDI est une mise à niveau majeure pour trois raisons principales :

C'est plus léger : Parce que JEDI fonctionne avec de minuscules « photos mentales » plutôt qu'avec des images vidéo complètes, il utilise 43 % de mémoire informatique en moins. C'est comme passer du transport d'un lourd sac à dos rempli de livres au transport d'un seul carnet intelligent.
C'est plus rapide : Le robot peut penser (échantillonner) 3 fois plus vite et s'entraîner 2,5 fois plus vite que les meilleures méthodes précédentes basées sur les pixels.
Il joue différemment : C'est la partie la plus surprenante. L'article a découvert que JEDI ne joue pas seulement plus vite ; il joue différemment.
- Sur des jeux qui étaient déjà faciles pour d'autres robots, JEDI était bon mais pas toujours le meilleur absolu.
- Cependant, sur les jeux les plus difficiles et les plus chaotiques (comme les jeux de tir avec de nombreuses cibles en mouvement), JEDI a brillé. Il semblait mieux gérer le chaos parce que sa « photo mentale » se concentrait sur la stratégie plutôt que de se laisser distraire par le bruit visuel.

La Conclusion

L'article soutient que vous n'avez pas besoin d'être un peintre parfait (reconstruisant chaque pixel) pour être un grand stratège. En apprenant au robot à prédire l'« essence » du futur en utilisant un jeu de suppression de bruit (diffusion), vous obtenez un système qui est plus rapide, moins cher à exécuter et étonnamment meilleur pour gérer des situations difficiles et chaotiques.

Les auteurs appellent cela JEDI, et ils affirment que c'est la première fois que cette méthode spécifique de « devinettes prédictives » est combinée avec succès à la « suppression de bruit » (diffusion) pour enseigner à un robot à jouer à des jeux en ligne, sans avoir besoin d'enseignants pré-entraînés.

Résumé Technique : JEDI : Modèle de Monde par Diffusion d'Embedding Joint pour l'Apprentissage par Renforcement Basé sur un Modèle en Ligne

Énoncé du Problème

L'apprentissage par renforcement basé sur un modèle en ligne (MBRL) repose sur des modèles de monde appris pour prédire les résultats futurs et améliorer l'efficacité des échantillons. Bien que les modèles de diffusion aient émergé comme des architectures puissantes pour la modélisation du monde dans des domaines visuellement complexes, les approches actuelles font face à une tension significative :

Diffusion Pixelique : Des modèles comme DIAMOND opèrent directement dans l'espace des pixels. Bien qu'efficaces, ils sont coûteux en termes de ressources informatiques (mémoire VRAM), de temps d'échantillonnage et de durée d'entraînement.
Diffusion Latente : Le déplacement de la diffusion vers l'espace latent améliore l'efficacité mais repose souvent sur des latents entraînés séparément (par exemple, via des objectifs de reconstruction) plutôt que sur des objectifs de modèle de monde de bout en bout. Ces approches, telles que Horizon Imagination (HI), n'ont pas encore égalé les performances des modèles de diffusion basés sur les pixels et manquent de l'apprentissage de représentation prédictive de bout en bout qui a propulsé les progrès modernes du MBRL.

La question centrale abordée est de savoir si un modèle de monde par diffusion peut apprendre sa représentation latente de bout en bout à partir de l'objectif de débruitage en utilisant un cadre d'Architecture Prédictive d'Embedding Joint (JEPA), combinant ainsi l'efficacité de la diffusion latente avec les performances de l'apprentissage prédictif de bout en bout.

Méthodologie : JEDI

Les auteurs introduisent JEDI (Joint Embedding Diffusion), le premier modèle de monde latent par diffusion en ligne de bout en bout construit sur un cadre JEPA.

Architecture Principale

JEDI apprend simultanément un encodeur d'espace latent ( $E_\phi$ ) et un modèle de dynamique. Le système se compose de :

Encodeur : Compresse les observations $x_t$ en latents propres $z^0_t$ .
Dynamique de Diffusion Latente : Prédit l'état latent suivant par débruitage conditionnel. Au lieu d'une prédiction en une seule étape, elle utilise une chaîne de diffusion inverse conditionnée par le latent courant $z^0_t$ et les actions passées pour prédire le latent futur propre $z^0_{t+1}$ .
Têtes de Récompense et de Terminaison : Prédisent les récompenses et les signaux de terminaison directement à partir de l'état latent.

Mécanisme d'Entraînement

Contrairement aux modèles de diffusion latente traditionnels qui s'appuient sur des pertes de reconstruction (par exemple, les VAE) pour entraîner l'encodeur, JEDI entraîne l'encodeur de bout en bout en utilisant la perte de débruitage par diffusion.

Objectif d'Embedding Joint : L'encodeur est mis à jour via les gradients de l'objectif de débruitage conditionnel. Le futur latent cible est traité comme une « étiquette » dérivée de l'environnement, et le modèle apprend à débruitier une version bruitée de ce futur latent pour revenir à l'état propre, conditionné par l'état courant.
Stop-Gradient : Pour éviter l'effondrement de la représentation et stabiliser l'entraînement, un stop-gradient est appliqué au futur latent cible lors des mises à jour de l'encodeur, suivant les principes JEPA et TDMPC2.
Techniques de Stabilisation : Le modèle emploie un clampage tanh différentiable sur les activations latentes et un mécanisme de commutation aléatoire entre la sortie du débruiteur et la sortie de l'encodeur pour la conditionnement de l'étape suivante, afin d'équilibrer la cohérence et la supervision.

Motivation Théorique

L'article fournit une justification théorique reliant JEDI au principe du Goulot d'Étranglement Informationnel.

JEPA comme Goulot d'Étranglement : L'objectif JEPA standard est montré pour induire un goulot d'étranglement informationnel variationnel, récompensant les représentations qui conservent les informations utiles pour prédire l'avenir tout en compressant les détails non pertinents.
Décomposition Diffusion-JEPA : Les auteurs démontrent que la diffusion conditionnelle latente admet une décomposition similaire de goulot d'étranglement informationnel variationnel. La perte de débruitage conditionnel agit comme le terme prédictif au sein de cette structure de goulot. Cela suggère que le débruitage par diffusion fournit un signal d'apprentissage prédictif de bout en bout, principiel, pour l'encodeur latent, plutôt que de servir simplement de substitut génératif.

Contributions Clés

Premier Modèle de Monde Latent par Diffusion de Bout en Bout : JEDI est le premier modèle MBRL en ligne à entraîner son encodeur latent de bout en bout en utilisant un objectif de débruitage par diffusion dans un cadre JEPA, éliminant le besoin d'un préentraînement séparé ou de pertes de reconstruction.
Lien Théorique : L'article établit un lien mathématique entre les objectifs de diffusion conditionnelle et le goulot d'étranglement informationnel variationnel, expliquant pourquoi le débruitage peut efficacement soutenir l'apprentissage de représentations prédictives.
Gains d'Efficacité : En passant de l'espace des pixels à un espace latent appris compact ( $16 \times 8 \times 8$ ) et en entraînant de bout en bout, JEDI réduit considérablement la surcharge computationnelle par rapport à la diffusion basée sur les pixels.

Résultats Expérimentaux

JEDI a été évalué sur le benchmark Atari100k (26 tâches, 5 graines) et l'environnement Craftium (prise de décision incarnée en 3D).

Performance

Atari100k : JEDI atteint des performances compétitives, égalant ou dépassant de solides références (DIAMOND, DreamerV3, STORM) sur les métriques de Moyenne Interquartile (IQM) et d'Écart d'Optimalité. Il obtient des résultats de l'état de l'art (SOTA) dans 7 jeux et une performance de deuxième place en nombre de tâches surhumaines.
Craftium : JEDI surpasse Horizon Imagination (HI) sur les quatre tâches testées, démontrant que des latents prédictifs de bout en bout peuvent soutenir de fortes performances dans des environnements 3D sans entraînement latent basé sur la reconstruction.
Stochasticité : Dans les expériences avec saut de trames stochastique, JEDI surpasse systématiquement DreamerV3, suggérant que les modèles de diffusion peuvent être plus robustes face à une forte incertitude aléatoire.

Efficacité

Par rapport à la référence basée sur les pixels DIAMOND, JEDI démontre des améliorations substantielles en efficacité :

Utilisation VRAM : Utilise 43 % de moins de VRAM.
Vitesse d'Échantillonnage : Réalise un échantillonnage de modèle de monde >3× plus rapide.
Vitesse d'Entraînement : S'entraîne >2,5× plus vite (38 heures contre 98 heures sur A100 pour Atari100k).
Taille du Modèle : Maintient un nombre de paramètres similaire (13,5 M) à DIAMOND mais surpasse significativement HI (97 M de paramètres) en efficacité tout en conservant des performances compétitives.

Décalage du Profil de Performance

Une découverte distincte est que JEDI présente un profil de performance au niveau des tâches différent de celui de DIAMOND. Alors que DIAMOND excelle dans les régimes à fort HNS (Human Normalized Score), JEDI montre des améliorations marquées dans les régimes à faible HNS (jeux difficiles). Cela est capturé par la métrique Har-HNS, où JEDI surpasse DIAMOND, indiquant qu'il échoue moins sévèrement sur les tâches les plus difficiles. Les auteurs attribuent cela à l'interface latente de bout en bout réduisant la dimensionnalité de l'entrée pour l'acteur-critique, permettant une meilleure optimisation dans des jeux complexes à grand espace d'actions.

Signification et Revendications

L'article revendique que JEDI démontre que la diffusion latente prédictive est une alternative viable et efficace à la fois à la diffusion dans l'espace des pixels et aux modèles latents entraînés séparément pour le MBRL en ligne.

Au-delà du Calcul : Les résultats suggèrent que les latents prédictifs de bout en bout modifient le comportement du modèle de monde lui-même, et pas seulement le coût computationnel. Le décalage des profils de performance indique que la représentation apprise modifie fondamentalement la façon dont l'agent interagit avec des environnements complexes.
Approche Principielle : En ancrant la méthode dans la théorie du goulot d'étranglement informationnel, l'article positionne la diffusion latente prédictive comme une direction principielle pour la mise à l'échelle du MBRL en ligne.
Limites : Les auteurs notent modestement que les preuves sont actuellement limitées à Atari100k et Craftium, et que la méthode n'a pas encore été testée sur des domaines comme DMControl ou Procgen. Ils reconnaissent également que les performances dépendent de choix de conception spécifiques (par exemple, le clampage, la commutation aléatoire) et que la dérivation théorique motive les propriétés de représentation plutôt que de garantir l'optimisation de l'objectif variationnel complet.