Open Dreamer : une reproduction JAX/Flax du modèle Dreamer 4, avec la recette d'entraînement complète publiée
Un petit groupe de chercheurs en IA, l'équipe Reactor, a publié Open Dreamer, une implémentation ouverte du pipeline de modèle du monde Dreamer 4, codée en JAX et Flax NNX. Deux dépôts ont été mis en ligne: next-state/open-dreamer contient le pipeline d'entraînement complet, avec un tokenizer vidéo causal, un modèle de dynamique latente conditionné par l'action, la génération de rollouts et le calcul du score FVD, tandis que reactor-team/open-dreamer propose un harnais de rollout local minimal qui génère des images à partir d'un fichier MP4 et d'un fichier d'actions associé. Un troisième élément, une démo dans le navigateur hébergée sur l'infrastructure Reactor, diffuse en temps réel un monde Minecraft généré et propose un bouton pour basculer entre le jeu réel et le rêve, image par image. L'équipe a délibérément limité ses méthodes à celles décrites dans l'article de recherche original sur Dreamer 4, en commençant par CoinRun, un jeu de plateforme 2D procédural entraînable sur un seul GPU, avant d'étendre le pipeline à des séquences de jeu Minecraft au format VPT.
Résumé et traduction réalisés par Le Fil IA à partir de MarkTechPost. Lire l'article original →
L'architecture repose sur un même squelette transformeur causal par blocs, utilisé à la fois pour le tokenizer et le modèle de dynamique, alternant des couches spatiales qui propagent l'information au sein d'une image et des couches temporelles causales qui relient les images entre elles. Le tokenizer, un auto-encodeur masqué plutôt qu'un VAE classique, atteint une compression d'environ 100 fois sans perte de qualité liée aux pertes KL ou adversariales, le masquage rendant selon l'équipe l'espace latent plus facile à diffuser. Le modèle de dynamique prédit l'image suivante via diffusion forcing, flow matching et modèles shortcut, tout en anticipant la prochaine action; les jetons du modèle du monde n'ont toutefois pas accès au jeton de l'agent, si bien que la politique n'influence le futur qu'à travers l'action suivante. Concrètement, le modèle de dynamique Minecraft compte 1,6 milliard de paramètres répartis sur 30 couches, avec 30 têtes d'attention et une attention par requêtes groupées à 3 têtes clé-valeur, entraîné pendant 200 000 étapes avec l'optimiseur Muon et un taux d'apprentissage de pointe de 3e-4.
Cette publication détaillée offre à la communauté open source une reproduction fidèle et documentée d'une architecture jusque-là propriétaire, avec des choix d'ingénierie précis sur le parallélisme et l'utilisation du matériel. L'équipe rapporte une utilisation des FLOPs du modèle de 57 à 58 %, proche du seuil de 60 % jugé sain pour l'entraînement de transformeurs, et explique avoir dû composer avec un mur mémoire sur les activations plutôt que sur l'état du modèle, qui tient en environ 24 Gio sur un GPU B200. Ce travail s'inscrit dans une course plus large à la reproduction ouverte des modèles du monde utilisés pour la génération d'environnements interactifs et l'entraînement d'agents.
Pas d'impact direct sur la France/UE
Le vrai signal ici, c'est que Dreamer 4 sort du placard propriétaire avec la recette d'entraînement complète, pas juste des poids. Ça compte, parce que jusqu'ici les modèles du monde restaient des boîtes noires qu'on citait sans pouvoir vérifier. 57-58% d'utilisation FLOPs sur B200, c'est du sérieux, pas du proof-of-concept bricolé. Reste à voir si ça tient quand quelqu'un d'autre que l'équipe Reactor essaie de le reproduire sur son propre matériel.