NachrichtenMakroOpen Dreamer: Reactor veröffentlicht quelloffene JAX/Flax-Reproduktion der Dreamer-4-World-Model-Pipeline mit vollständigem Trainingsrezept

Open Dreamer: Reactor veröffentlicht quelloffene JAX/Flax-Reproduktion der Dreamer-4-World-Model-Pipeline mit vollständigem Trainingsrezept

Autor: MarkTechPost·

Wichtige Erkenntnisse

  • Open Dreamer stellt eine quelloffene Implementierung der Dreamer-4-Pipeline mit kausalem Video-Tokenizer, aktionskonditioniertem latentem Dynamikmodell, Rollout-Werkzeugen und FVD-Evaluierungscode bereit.
  • Das Minecraft-Dynamikmodell nutzt 1,6 Milliarden Parameter über 30 blockkausale Schichten hinweg und ist für 200.000 Trainingsschritte mit dem Muon-Optimizer konfiguriert.
  • Reactor berichtet von 57% bis 58% Model FLOPs Utilization auf NVIDIA B200 GPUs, wobei Aktivierungen und nicht der Modellzustand als Hauptspeicherkosten identifiziert wurden.
  • Das Team dokumentierte mehrere Stabilitätsmaßnahmen, darunter EMA-Nutzung, Grenzen gemischter Präzision, Optimizer-Wechsel, Loss-Gewichtung und Minibatch Optimal Transport.
  • Die Veröffentlichung enthält keine Trainingsschleifen für Behaviour Cloning oder Reinforcement Learning, und das Projekt hat keine FVD-Scores veröffentlicht.
Open Dreamer: Reactor veröffentlicht quelloffene JAX/Flax-Reproduktion der Dreamer-4-World-Model-Pipeline mit vollständigem Trainingsrezept

Eine Forschungsgruppe unter dem Namen Reactor hat Open Dreamer veröffentlicht, eine quelloffene Implementierung der Dreamer-4-World-Model-Pipeline, die in JAX und Flax NNX entwickelt wurde. World Models lernen, aus Beobachtungsdaten vorherzusagen, wie sich eine Umgebung entwickelt, sodass Agenten planen, lernen und Videos erzeugen können, ohne mit einem realen System zu interagieren. Die von Danijar Hafner entwickelte Dreamer-Reihe gehört zu den meistuntersuchten Arbeiten im modellbasierten Reinforcement Learning; Dreamer 4 erweitert den Ansatz in Richtung großskaliger Video-World-Models. Das Projekt soll die Forschungsmethodik von Dreamer 4 originalgetreu reproduzieren und zugleich die vollständige Trainingspipeline — einschließlich Stabilitätskorrekturen und Compute-Konfigurationen — öffentlich zugänglich machen, ein Detailgrad, der bei Veröffentlichungen großskaliger World Models bislang typischerweise fehlte.

Veröffentlichte Artefakte

Es wurden zwei Code-Repositories veröffentlicht. Das erste, next-state/open-dreamer, enthält die vollständige Trainingspipeline: einen kausalen Video-Tokenizer, ein aktionskonditioniertes latentes Dynamikmodell, Werkzeuge zur Rollout-Generierung und FVD-Scoring (Fréchet Video Distance). Das zweite, reactor-team/open-dreamer, stellt ein minimales lokales Rollout-Harness bereit, das Videoframes aus einer Eingabe-MP4-Datei und einer zugehörigen Aktionsdatei erzeugen kann.

Ein drittes Ergebnis ist eine browserbasierte Demo, die auf der Reactor-Runtime läuft. Sie streamt eine generierte Minecraft-Umgebung in Echtzeit und enthält einen Game-⟷-Dream-Schalter, der den Videostream bildweise zwischen dem tatsächlichen Spiel und der Ausgabe des World Models überführt.

Das Team erklärte, sein Ziel sei die gezielte Reproduktion der Dreamer 4-Forschung gewesen, wobei Techniken außerhalb des ursprünglichen Papers vermieden wurden, um den Suchraum eng zu halten. Die Entwicklung begann mit CoinRun, einem prozedural generierten 2D-Plattformer, der auf einer einzelnen GPU trainiert werden kann, bevor die funktionierende Pipeline auf Gameplay-Videos im Stil von Minecraft/VPT skaliert wurde. Minecraft dient seit OpenAIs Video-PreTraining-(VPT)-Projekt als Benchmark-Domäne für das Lernen von Agenten aus Gameplay-Demonstrationen; dort wurde ein inverses Dynamikmodell auf von Auftragnehmern gelabeltem YouTube-Material trainiert, um großskalige, aktionsgelabelte Gameplay-Daten zu erzeugen.

Weitere Details sind im Projektblog, im arXiv-Paper und in Reactors Ankündigung auf X verfügbar.

Architektur: Ein Backbone, zwei Modelle

Sowohl der Tokenizer als auch das Dynamikmodell verwenden denselben blockkausalen Transformer-Backbone, der zwischen zwei Arten von Attention wechselt. Räumliche Schichten propagieren Informationen über Elemente innerhalb eines einzelnen Frames, während kausale Zeitschichten Informationen über Frames hinweg weitergeben.

Der Tokenizer ist als transformerbasierter Masked Autoencoder (MAE) konzipiert und nicht als traditioneller variationaler Autoencoder (VAE). Das Team berichtet von einer Kompression von etwa 100× und weist darauf hin, dass dieses Design weder KL-Divergenz noch adversarielle Verluste erfordert. Nach Einschätzung des Teams macht der Maskierungsansatz den latenten Raum besser für diffusionsbasierte Generierung geeignet.

Das Dynamikmodell führt Next-Frame Prediction mithilfe von Diffusion Forcing, Flow Matching und Shortcut Models aus. Es sagt außerdem die nächste Aktion voraus. Statt zwischen einem separaten Übergangsmodul und einem Policy-Modul zu wechseln, wird der Rollout in Blöcke pro Zeitschritt integriert, die als (vorherige Aktion, Zustand, Policy) strukturiert sind. Spatial Attention arbeitet innerhalb jedes Blocks, während kausale temporale Attention die Blöcke über die Zeit hinweg verbindet.

Eine zentrale Designbeschränkung lautet: World-Model-Tokens können das Agenten-Token nicht lesen. Folglich können Aufgaben- und Policy-Informationen zukünftige Zustände nur über die vorhergesagte nächste Aktion beeinflussen.

Trainingskonfiguration

Die ausgelieferten Minecraft-Konfigurationsdateien spezifizieren das Trainingsrezept im Detail.

Das Dynamikmodell umfasst 1,6 Milliarden Parameter über 30 blockkausale Schichten hinweg, mit d_model 1920, 30 Attention-Heads und 3 KV-Heads unter Verwendung von Grouped-Query Attention. Jede vierte Schicht fungiert als Time-Attention-Schicht. Jeder Zeitschritt enthält 32 gelernte Register-Tokens, und ein Packing-Faktor von 2 fasst benachbarte Tokenizer-Latents in jeweils einem räumlichen Dynamik-Token zusammen. Time Attention arbeitet über ein gleitendes Fenster von 192 Schritten.

Das Training erstreckt sich über 200.000 Schritte mit dem Muon-Optimizer, einem WSD-Zeitplan (warmup-stable-decay) und einer maximalen Lernrate von 3e-4. Shortcut- und Bootstrap-Sampling werden bei Schritt 100.000 mit einem Batch-Anteil von 0,25 aktiviert. Der Decay des exponentiellen gleitenden Durchschnitts (EMA) ist auf 0,999 gesetzt.

Die Tokenizer-Konfiguration erzeugt 512 latente Tokens pro Frame bei einer Bottleneck-Breite von 16. Rohe Frames mit 360×640 werden auf 368×640 aufgefüllt, damit beide räumlichen Dimensionen gleichmäßig in 16×16-Patches aufgeteilt werden können. Die Encoder-Tiefe beträgt 12 bei d_model 1536; die Decoder-Tiefe beträgt 8 bei d_model 1024. Die MAE-Maskierungswahrscheinlichkeit erreicht maximal 0,9, und der LPIPS-Verlust wird auf der Hälfte der Zeitschritte mit einem Gewicht von 0,2 angewendet.

VPT-Aktionen werden in 27 binäre Aktionskanäle und 121 kategoriale Mausklassen geparst, ohne kontinuierliche Kanäle.

Rechenleistung und Speicherstrategie

Das Team berichtet von einer Model FLOPs Utilization (MFU) von 57–58%, verglichen mit einem häufig genannten Richtwert von 60% für gesundes Transformer-Training. Die Analyse basiert auf einem Roofline-Argument: Auf einer NVIDIA B200 liegt der Übergang zwischen bandbreitenlimitierter und rechenlimitierter Ausführung bei 292 FLOP/byte. Die Verarbeitung von 256 Frames pro GPU verschiebt die Arbeitslast über diesen Ridge Point hinaus in den rechenlimitierten Bereich.

Die Sharding-Entscheidungen widersprachen den anfänglichen Erwartungen. Bei 1,6 Milliarden Parametern belegte der vollständige Modellzustand — Parameter, Gradienten, Optimizer-Zustand und EMA — etwa 24 GiB und passte damit auf eine einzelne B200. Die Aktivierungen, nicht der Modellzustand, erwiesen sich als primärer Speicherkostenfaktor. Das Team experimentierte mit Data Parallelism, Fully Sharded Data Parallelism (FSDP), Tensor Parallelism und Sequence Parallelism, bevor es sich für einfaches Data Parallelism in Kombination mit Activation Checkpointing entschied.

Für das Dataloading tokenisierte das Team den gesamten Datensatz vorab in .arrayrecord-Dateien und verwendete Grain mit einem GPU-seitigen Prefetch-Puffer. Standardmäßiges ffmpeg-basiertes Decoding reichte nicht aus, um die GPUs auszulasten.

Stabilitäts-Engineering

Das Forschungsteam erklärt ausdrücklich, dass Stabilitätsprobleme den größten Teil der Entwicklungszeit beanspruchten. Die zentrale Beobachtung: Die meisten Stabilitätsprobleme treten auf, obwohl der Loss weiter sinkt. Der mittlere quadratische Fehler (MSE) verbessert sich gleichmäßig, während sich die Generierungsqualität gleichzeitig verschlechtert — ein Phänomen, das konventionelles, verlustbasiertes Monitoring für diffusionsbasierte World Models unzuverlässig macht.

Sechs konkrete Korrekturen werden dokumentiert:

  1. Optimizer-Wechsel: Muon ersetzte LaProp, das über zwei separate Trainingsläufe von jeweils etwa 400 B200-Stunden zufällige und zunehmend häufige Ausschläge zeigte.

  2. EMA als verpflichtend: EMA-Gewichte werden für die Diffusionsinferenz als essenziell und nicht als optional behandelt.

  3. Grenzen gemischter Präzision: Parameter bleiben in float32, BF16 deckt die meisten Matmul-Aktivierungen und Attention-Eingaben ab, und float32 wird für Normalisierungsschichten sowie den Dynamics-Flow-Ausgabekopf beibehalten.

  4. Loss-Gewichtung: Das Team nutzt x-Prediction mit einem v-space loss, der sich auf einen Gewichtungsterm ähnlich der Dreamer-4-Formulierung reduziert, jedoch mit quadriertem Nenner. Es berichtet von einer kleinen, aber wahrnehmbaren Verbesserung.

  5. Optimal Transport: Minibatch Barycentric Optimal Transport zwischen Rausch- und latenten Sequenzen verbesserte die Stabilität der Rollout-Generierung.

  6. μ-Parametrisierung: Sie wurde getestet und als unnötig bewertet, teilweise weil Muon die Stabilität von Hyperparametern über Modellgrößen hinweg effektiver aufrechterhält.

Ein weiteres Ergebnis aus der CoinRun-Entwicklungsphase: Ein iso-FLOPs-Sweep schätzte die compute-optimale Skalierung auf ungefähr N ∝ C^0.56 und D ∝ C^0.44.

Was nicht enthalten ist

Das Repository enthält keine Trainingsschleife für Behaviour Cloning (BC) oder Reinforcement Learning (RL). Eine vollständige Dreamer-4-BC/RL-Agentenschleife ist als offener Roadmap-Punkt aufgeführt. Die in der Projektdokumentation beschriebene CoinRun-Policy-Arbeit wurde für die Minecraft-Implementierung nicht verwendet und nicht veröffentlicht.

Das Projekt veröffentlicht keine FVD-Scores, obwohl das Repository scripts/eval_fvd.py ausliefert, ein I3D-basiertes Evaluierungs-Harness, das für 4 Kontextframes und einen Horizont von 240 Frames konfiguriert ist.

Zusammenfassung zentraler technischer Kennzahlen

  • Dynamikmodell: 1,6B Parameter, 30 Schichten, d_model 1920, trainiert über 200.000 Schritte mit Muon
  • Hardwareeffizienz: 57–58% MFU auf NVIDIA B200 GPUs, 256 Frames pro GPU, etwa 24 GiB Modellzustand
  • Primäre Herausforderung: Stabilität, nicht Durchsatz — Loss-Kurven verdeckten die meisten Rückgänge der Generierungsqualität