ActualitésMacroOpen Dreamer : Reactor publie une reproduction open source en JAX/Flax du pipeline de modèle du monde Dreamer 4 avec la recette complète d’entraînement

Open Dreamer : Reactor publie une reproduction open source en JAX/Flax du pipeline de modèle du monde Dreamer 4 avec la recette complète d’entraînement

Auteur: MarkTechPost·

Points clés

  • Open Dreamer fournit une implémentation open source du pipeline Dreamer 4 avec un tokenizer vidéo causal, un modèle dynamique latent conditionné par les actions, des utilitaires de rollout et du code d’évaluation FVD.
  • Le modèle dynamique Minecraft utilise 1,6 milliard de paramètres sur 30 couches block-causal et est configuré pour 200 000 étapes d’entraînement avec l’optimiseur Muon.
  • Reactor rapporte une utilisation des FLOPs du modèle de 57 % à 58 % sur des GPU NVIDIA B200, les activations étant identifiées comme le principal coût mémoire plutôt que l’état du modèle.
  • L’équipe a documenté plusieurs mesures de stabilité, notamment l’utilisation d’EMA, des frontières de précision mixte, des changements d’optimiseur, la pondération de la perte et le transport optimal par minibatch.
  • La publication n’inclut pas les boucles d’entraînement de clonage comportemental ou d’apprentissage par renforcement, et le projet n’a pas publié de scores FVD.
Open Dreamer : Reactor publie une reproduction open source en JAX/Flax du pipeline de modèle du monde Dreamer 4 avec la recette complète d’entraînement

Un groupe de recherche opérant sous le nom Reactor a publié Open Dreamer, une implémentation open source du pipeline de modèle du monde Dreamer 4 construite avec JAX et Flax NNX. Les modèles du monde apprennent à prédire l’évolution d’un environnement à partir de données d’observation, ce qui permet aux agents de planifier, d’apprendre et de générer de la vidéo sans interagir avec un système réel. La série Dreamer, développée par Danijar Hafner, figure parmi les axes de recherche les plus étudiés en apprentissage par renforcement fondé sur des modèles, Dreamer 4 étendant l’approche vers des modèles du monde vidéo à grande échelle. Le projet vise à reproduire fidèlement la méthodologie de recherche de Dreamer 4 tout en rendant publiquement accessible l’ensemble du pipeline d’entraînement — y compris les correctifs de stabilité et les configurations de calcul — un niveau de détail que les publications de modèles du monde à grande échelle ont généralement omis.

Artefacts publiés

Deux dépôts de code ont été publiés. Le premier, next-state/open-dreamer, contient le pipeline complet d’entraînement : un tokenizer vidéo causal, un modèle dynamique latent conditionné par les actions, des utilitaires de génération de rollouts et une évaluation FVD (Fréchet Video Distance). Le second, reactor-team/open-dreamer, fournit un harnais minimal de rollout local capable de générer des images vidéo à partir d’un fichier MP4 d’entrée associé à un fichier d’actions correspondant.

Un troisième livrable est une démonstration basée sur navigateur exécutée sur le runtime Reactor. Elle diffuse en temps réel un environnement Minecraft généré et inclut un commutateur Game ⟷ Dream qui fait passer le flux vidéo, image par image, entre le jeu réel et la sortie du modèle du monde.

L’équipe a indiqué que son objectif était de reproduire délibérément la recherche Dreamer 4, en évitant les techniques extérieures à l’article original afin de maintenir un espace de recherche restreint. Le développement a commencé sur CoinRun, un jeu de plateforme 2D généré procéduralement pouvant être entraîné sur un seul GPU, avant que le pipeline fonctionnel ne soit étendu à des vidéos de gameplay de type Minecraft/VPT. Minecraft sert de domaine de référence pour l’apprentissage d’agents à partir de démonstrations de gameplay depuis le projet Video PreTraining (VPT) d’OpenAI, qui a entraîné un modèle de dynamique inverse sur des vidéos YouTube annotées par des prestataires afin de produire des données de gameplay à grande échelle étiquetées par actions.

Des détails supplémentaires sont disponibles sur le blog du projet, l’article arXiv et l’annonce de Reactor sur X.

Architecture : un backbone, deux modèles

Le tokenizer et le modèle dynamique partagent le même backbone Transformer block-causal, qui alterne entre deux types d’attention. Les couches spatiales propagent l’information entre les éléments d’une même image, tandis que les couches temporelles causales propagent l’information entre les images.

Le tokenizer est conçu comme un Masked Autoencoder (MAE) basé sur Transformer plutôt que comme un autoencodeur variationnel (VAE) traditionnel. L’équipe rapporte une compression d’environ 100× et note que cette conception supprime le besoin de divergence KL ou de pertes adversariales. Elle soutient que l’approche par masquage rend l’espace latent plus adapté à une génération fondée sur la diffusion.

Le modèle dynamique effectue une prédiction de l’image suivante à l’aide de diffusion forcing, de flow matching et de shortcut models. Il prédit également l’action suivante. Au lieu d’alterner entre un module de transition séparé et un module de politique, le rollout est intégré dans des blocs par pas de temps structurés sous la forme (action précédente, état, politique). L’attention spatiale opère à l’intérieur de chaque bloc, et l’attention temporelle causale relie les blocs au fil du temps.

Une contrainte de conception critique : les tokens du modèle du monde ne peuvent pas lire le token de l’agent. Par conséquent, les informations de tâche et de politique ne peuvent influencer les états futurs qu’au moyen de l’action suivante prédite.

Configuration d’entraînement

Les fichiers de configuration Minecraft publiés précisent la recette d’entraînement en détail.

Le modèle dynamique comprend 1,6 milliard de paramètres répartis sur 30 couches block-causal, avec un d_model de 1920, 30 têtes d’attention et 3 têtes KV utilisant une attention grouped-query. Chaque quatrième couche fonctionne comme une couche d’attention temporelle. Chaque pas de temps contient 32 tokens de registre appris, et un facteur de packing de 2 consolide les latents voisins du tokenizer dans chaque token spatial dynamique. L’attention temporelle opère sur une fenêtre glissante de 192 pas.

L’entraînement s’étend sur 200 000 étapes avec l’optimiseur Muon, un calendrier WSD (warmup-stable-decay) et un taux d’apprentissage maximal de 3e-4. L’échantillonnage shortcut et bootstrap s’active à l’étape 100 000 avec une fraction de batch de 0,25. La décroissance de la moyenne mobile exponentielle (EMA) est fixée à 0,999.

La configuration du tokenizer produit 512 tokens latents par image avec une largeur de goulot d’étranglement de 16. Les images brutes 360×640 sont remplies jusqu’à 368×640 afin que les deux dimensions spatiales se divisent uniformément en patchs de 16×16. La profondeur de l’encodeur est de 12 avec d_model 1536 ; la profondeur du décodeur est de 8 avec d_model 1024. La probabilité de masquage MAE atteint au maximum 0,9, et la perte LPIPS est appliquée avec un poids de 0,2 sur la moitié des pas de temps.

Les actions VPT sont analysées en 27 canaux d’action binaires et 121 classes catégorielles de souris, sans canaux continus.

Performance de calcul et stratégie mémoire

L’équipe rapporte une utilisation des FLOPs du modèle (MFU) de 57–58 %, contre un repère couramment cité de 60 % pour un entraînement Transformer sain. Son analyse repose sur un argument de type roofline : sur un NVIDIA B200, le point de bascule entre une opération limitée par la bande passante et une opération limitée par le calcul se situe à 292 FLOP/byte. Le traitement de 256 images par GPU pousse la charge de travail au-delà de ce point de crête, dans le régime limité par le calcul.

Les décisions de sharding ont contredit les attentes initiales. Avec 1,6 milliard de paramètres, l’état complet du modèle — paramètres, gradients, état de l’optimiseur et EMA — occupait environ 24 GiB, ce qui tient dans un seul B200. Les activations, plutôt que l’état du modèle, se sont révélées être le principal coût mémoire. L’équipe a expérimenté le parallélisme de données, le fully sharded data parallelism (FSDP), le parallélisme tensoriel et le parallélisme de séquence avant de retenir un parallélisme de données simple combiné au checkpointing des activations.

Pour le chargement des données, l’équipe a pré-tokenisé l’ensemble du jeu de données dans des fichiers .arrayrecord et utilisé Grain avec un tampon de préchargement côté GPU. Le décodage standard basé sur ffmpeg était insuffisant pour maintenir les GPU saturés.

Ingénierie de la stabilité

L’équipe de recherche indique explicitement que les problèmes de stabilité ont absorbé la plus grande part du temps de développement. Son observation centrale : la plupart des problèmes de stabilité surviennent alors même que la perte continue de diminuer. L’erreur quadratique moyenne (MSE) s’améliore régulièrement tandis que la qualité de génération se dégrade simultanément — un phénomène qui rend la surveillance conventionnelle fondée sur la perte peu fiable pour les modèles du monde basés sur la diffusion.

Six correctifs spécifiques sont documentés :

  1. Changement d’optimiseur : Muon a remplacé LaProp, qui présentait des pics aléatoires et de plus en plus fréquents sur deux exécutions d’entraînement distinctes d’environ 400 heures B200 chacune.

  2. EMA obligatoire : les poids EMA sont considérés comme essentiels pour l’inférence par diffusion, et non comme optionnels.

  3. Frontières de précision mixte : les paramètres restent en float32, BF16 couvre la plupart des activations de matmul et des entrées d’attention, et float32 est conservé pour les couches de normalisation et la tête de sortie du flux dynamique.

  4. Pondération de la perte : l’équipe utilise la x-prediction avec une perte en v-space, ce qui se réduit à un terme de pondération similaire à la formulation de Dreamer 4, mais avec un dénominateur au carré. Elle rapporte une amélioration faible mais perceptible.

  5. Transport optimal : le transport optimal barycentrique par minibatch appliqué entre le bruit et les séquences latentes a amélioré la stabilité de la génération de rollouts.

  6. μ-paramétrisation : testée et jugée inutile, en partie parce que Muon maintient plus efficacement la stabilité des hyperparamètres entre différentes tailles de modèles.

Un résultat supplémentaire issu de la phase de développement sur CoinRun : un balayage iso-FLOPs a estimé l’échelle optimale en calcul à environ N ∝ C^0.56 et D ∝ C^0.44.

Ce qui n’est pas inclus

Le dépôt ne contient pas la boucle d’entraînement de clonage comportemental (BC) ni celle d’apprentissage par renforcement (RL). Une boucle complète d’agent Dreamer 4 BC/RL est listée comme élément ouvert de la feuille de route. Le travail de politique sur CoinRun décrit dans la documentation du projet n’a pas été utilisé pour l’implémentation Minecraft et n’a pas été publié.

Le projet ne publie pas de scores FVD, bien que le dépôt fournisse scripts/eval_fvd.py, un harnais d’évaluation basé sur I3D configuré pour 4 images de contexte et un horizon de 240 images.

Synthèse des principaux chiffres d’ingénierie

  • Modèle dynamique : 1,6B de paramètres, 30 couches, d_model 1920, entraîné pendant 200 000 étapes avec Muon
  • Efficacité matérielle : 57–58 % de MFU sur des GPU NVIDIA B200, 256 images par GPU, environ 24 GiB d’état du modèle
  • Principal défi : la stabilité, et non le débit — les courbes de perte masquaient la plupart des régressions de qualité de génération