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.
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.
Dans nos dossiers
Vu une erreur factuelle dans cet article ? Signalez-la. Toutes les corrections valides sont publiées sur /corrections.




