One Model for All Tasks: Leveraging Efficient World Models in Multi-Task Planning

Ce papier présente ScaleZero, un modèle unique pour la planification multi-tâches qui surpasse les approches conventionnelles en résolvant les conflits de gradients grâce à une architecture Mixture-of-Experts et en optimisant l'efficacité de l'apprentissage via une stratégie de mise à l'échelle dynamique des paramètres.

L'essentiel

🌍 Le Grand Défi : Un seul cerveau pour tout apprendre ?

Imaginez que vous vouliez créer un robot super-intelligent capable de jouer à tous les jeux vidéo du monde en même temps.

• Certains jeux sont simples : comme faire du ping-pong (réagir vite).
• D'autres sont complexes : comme explorer un monde ouvert avec des énigmes (penser loin).
• D'autres encore sont très différents : comme des jeux de texte où il faut comprendre des histoires.

Le problème, c'est que si vous essayez d'entraîner un seul cerveau (un seul modèle d'IA) sur tout cela, il se retrouve souvent perdu. C'est comme essayer d'apprendre à jouer de la batterie, à faire de la chirurgie et à cuisiner un soufflé en même temps, avec les mêmes mains. Le cerveau se mélange, oublie ce qu'il savait, et finit par ne plus rien faire de bien. C'est ce que les chercheurs appellent le "collapsus de la plasticité" (le cerveau devient rigide et ne peut plus apprendre).

🚀 La Solution : ScaleZero

Les auteurs de ce papier (publié à ICLR 2026) ont créé un nouveau modèle appelé ScaleZero. Ils ont résolu le problème en utilisant deux astuces géniales, que l'on peut comparer à une école très intelligente et une boîte à outils magique.

Astuce 1 : L'École des Experts (MoE - Mixture of Experts)

Au lieu d'avoir un seul professeur qui essaie d'enseigner tout à tout le monde, ScaleZero est une école avec des centaines de spécialistes.

• L'analogie : Imaginez une grande classe. Quand un élève pose une question sur le ping-pong, seul l'expert en ping-pong intervient. Quand on parle de stratégie de jeu, c'est l'expert en stratégie qui prend la parole.
• Le résultat : L'IA ne mélange plus les idées. Elle envoie chaque tâche vers le "cerveau" spécialisé qui sait exactement quoi faire. Cela évite que les tâches simples (comme le ping-pong) étouffent les tâches complexes (comme l'exploration). C'est comme si chaque élève avait son propre tuteur personnel, mais ils partagent tous la même salle de classe.

Astuce 2 : La Boîte à Outils Dynamique (DPS - Dynamic Parameter Scaling)

Même avec des experts, apprendre tout en même temps est épuisant et coûteux en temps. C'est ici qu'intervient la deuxième astuce : l'adaptation dynamique.

• L'analogie : Imaginez que vous apprenez à conduire. Au début, vous avez besoin de beaucoup d'aide (freins, volant, pédale). Une fois que vous savez conduire, vous n'avez plus besoin de quelqu'un qui vous tient la main.
• La méthode DPS : ScaleZero observe en temps réel.
  1. Si un jeu est facile et que l'IA le maîtrise, elle arrête de gaspiller de l'énergie dessus (elle "gèle" les connaissances acquises).
  2. Si un jeu est difficile et que l'IA bloque, elle ajoute instantanément de nouveaux outils (de nouveaux "experts" ou modules) spécifiquement pour ce problème.
• Le résultat : L'IA apprend beaucoup plus vite et avec beaucoup moins d'essais. Au lieu de répéter 100 fois ce qu'elle sait déjà, elle se concentre uniquement sur ce qu'elle ne sait pas encore.

🏆 Les Résultats : Un champion polyvalent

Les chercheurs ont testé ScaleZero sur trois mondes très différents :

1. Les jeux vidéo classiques (Atari) : Des jeux comme Pong ou Seaquest.
2. La robotique (DMC) : Faire marcher un robot ou lancer une balle.
3. Les jeux d'aventure textuels (Jericho) : Résoudre des énigmes en lisant et en écrivant.

Le verdict ?

• ScaleZero, avec un seul modèle, joue aussi bien (voire mieux) que des experts qui ne connaissent qu'un seul jeu chacun.
• Grâce à son système dynamique (DPS), elle a besoin de 28,5 % d'essais en moins pour atteindre le même niveau de performance. C'est comme si elle apprenait à conduire en 70 % du temps habituel !

💡 En résumé

Ce papier nous dit qu'il est possible de créer une IA généraliste (un "couteau suisse" de l'intelligence) qui ne se perd pas quand on lui donne trop de tâches.

• Le secret ? Ne pas tout faire avec le même cerveau, mais utiliser une équipe d'experts qui travaillent chacun sur leur spécialité.
• L'efficacité ? Ne pas gaspiller d'énergie sur ce qui est déjà appris, mais ajuster la taille de l'équipe en fonction de la difficulté du moment.

C'est un pas de géant vers des agents intelligents capables de s'adapter à n'importe quel environnement, du jeu vidéo à la robotique, sans jamais avoir besoin d'être réentraînés de zéro.

Résumé technique

1. Problématique : L'effondrement de la plasticité dans les modèles du monde unifiés

L'article aborde le défi de l'apprentissage par renforcement multi-tâches (MTRL) hétérogène, où un agent unique doit maîtriser des tâches aux espaces d'observation, d'action et de complexité fondamentalement différents (ex. : jeux Atari, contrôle continu, aventures textuelles).

Les modèles du monde unifiés existants (comme UniZero), bien que performants sur des tâches individuelles, échouent souvent lorsqu'ils sont entraînés sur un large éventail de tâches simultanément. Les auteurs identifient deux obstacles critiques :

• Conflits de gradients et interférence représentationnelle : Dans une architecture monolithique, les gradients des tâches simples (convergeant rapidement) dominent ceux des tâches complexes, supprimant les signaux d'apprentissage nécessaires à ces dernières.
• Effondrement de la plasticité (Plasticity Collapse) : Ce déséquilibre conduit à une perte progressive de la capacité du réseau à s'adapter aux nouvelles données. Les auteurs observent empiriquement une augmentation drastique du ratio de neurones dormants et une inflation incontrôlée de la norme des états latents, signalant que le modèle a perdu sa capacité à apprendre de nouvelles dynamiques, entraînant un effondrement catastrophique des performances sur les tâches difficiles.
• Allocation statique des ressources : Les architectures conventionnelles appliquent les mêmes ressources d'apprentissage à toutes les tâches, indépendamment de leur difficulté ou de leur état de convergence, ce qui est inefficace.

2. Méthodologie : ScaleZero et la Mise à l'Échelle Dynamique des Paramètres (DPS)

Pour surmonter ces limitations, les auteurs proposent une approche en deux volets : une solution architecturale interne et une stratégie d'allocation de ressources externe.

A. Architecture : ScaleZero

ScaleZero est un modèle du monde unifié conçu pour atténuer les conflits de gradients et maintenir la plasticité. Il résulte d'une exploration systématique de l'espace de conception d'UniZero sur cinq axes (conditionnement des tâches, encodeur, normalisation, architecture du backbone, optimisation). Les choix clés sont :

• Backbone à Mélange d'Experts (MoE) : C'est la modification la plus impactante. Au lieu d'un réseau de neurones dense, ScaleZero utilise des couches MoE (Mixture-of-Experts) avec un routage sparse (Top-1). Cela permet de créer des chemins de calcul conditionnels où différentes tâches sont traitées par des sous-réseaux spécialisés ("experts"), réduisant ainsi l'interférence et les conflits de gradients.
• Encodeur ViT (Vision Transformer) : Utilisation d'un encodeur ViT entraîné de zéro pour une extraction de caractéristiques visuelle puissante et évolutive.
• Normalisation par Couche (LayerNorm) : Remplacement de la normalisation SimNorm par une LayerNorm standard pour mieux équilibrer stabilité et expressivité représentationnelle.

B. Stratégie d'Entraînement : Dynamic Parameter Scaling (DPS)

Pour optimiser l'allocation des ressources sur l'ensemble du processus d'apprentissage, les auteurs introduisent le DPS, une stratégie en ligne qui couple la capacité du modèle à la progression de l'apprentissage :

• Curriculum de complexité du modèle : Le modèle commence par un "backbone" de base (Stage 0) entraîné sur toutes les tâches.
• Expansion progressive via LoRA : À mesure que les tâches sont résolues ou que de nouvelles difficultés apparaissent, de nouveaux modules LoRA (Low-Rank Adaptation) sont injectés dynamiquement.
• Sélection adaptative des tâches : Seules les tâches non résolues ("unsolved") sont utilisées pour la collecte de données et les mises à jour de gradients. Une fois une tâche résolue (définie par un seuil de performance), elle est retirée de l'ensemble actif pour économiser des ressources.
• Isolation des paramètres : Les paramètres du backbone et des adaptateurs précédents sont gelés lors de l'ajout de nouveaux modules LoRA, empêchant l'oubli catastrophique (catastrophic forgetting) et permettant une plasticité ciblée.

3. Contributions Clés

1. Diagnostic Quantitatif de l'Effondrement de Plasticité : Les auteurs établissent un lien concret entre la dégradation des performances dans les modèles du monde unifiés et des métriques internes spécifiques (ratio de neurones dormants, norme des états latents, rang effectif des représentations).
2. ScaleZero : Une nouvelle architecture unifiée intégrant un backbone MoE sparse, démontrant des performances exceptionnelles et une généralisation robuste sur des benchmarks hétérogènes.
3. Stratégie DPS : Une méthode d'apprentissage adaptatif qui réduit considérablement le coût computationnel et les interactions avec l'environnement tout en maintenant des performances de pointe, en allouant dynamiquement la capacité du modèle aux tâches qui en ont le plus besoin.

4. Résultats Expérimentaux

Les évaluations ont été menées sur 48 tâches distinctes provenant de trois benchmarks majeurs, en utilisant uniquement l'apprentissage par renforcement en ligne (sans données d'experts) :

• Atari (26 jeux) : ScaleZero (MT) atteint un score moyen normalisé par rapport à l'humain supérieur à celui d'une collection de 26 agents spécialisés (UniZero ST), prouvant un transfert de connaissances positif. Il surpasse particulièrement les tâches d'exploration difficiles comme Seaquest.
• DeepMind Control Suite (18 tâches) : ScaleZero obtient un score médian supérieur à celui des agents spécialisés, démontrant une robustesse dans le contrôle continu.
• Jericho (4 jeux textuels) : L'agent atteint des performances comparables aux agents spécialisés et au modèle de base CALM+OC, validant l'efficacité du modèle sur des tâches basées sur le langage.
• Efficacité (DPS) : Sur le benchmark DMC, l'utilisation de la stratégie DPS permet d'atteindre des performances équivalentes aux agents spécialisés en utilisant seulement 71,5 % des interactions avec l'environnement (soit une réduction de 28,5 % des coûts de données et d'entraînement).

5. Signification et Impact

Ce travail représente une avancée significative vers la création d'agents généralistes capables de planifier dans des environnements hétérogènes.

• Théorique : Il démontre que les architectures MoE, souvent associées aux LLM, sont cruciales pour résoudre les problèmes de plasticité et de conflit de gradients dans l'apprentissage par renforcement multi-tâches. La preuve théorique fournie montre que le conflit de gradients dans une couche MoE a une borne supérieure strictement inférieure à celle d'une couche dense.
• Pratique : La combinaison de ScaleZero et de DPS offre une voie viable pour entraîner des agents unifiés de manière efficace, réduisant la barrière de l'entrée computationnelle tout en maintenant une haute performance. Cela ouvre la voie à des systèmes d'IA plus adaptatifs capables de gérer une diversité de tâches sans nécessiter de réentraînement complet ou de modèles séparés pour chaque domaine.

En résumé, l'article propose une solution architecturale (MoE) et algorithmique (DPS) pour transformer les modèles du monde unifiés, autrefois limités par l'interférence des tâches, en des agents robustes et efficaces capables de maîtriser un large éventail de défis de décision séquentielle.