NotizieMacroOpen Dreamer: Reactor rilascia una riproduzione open-source in JAX/Flax della pipeline di world model Dreamer 4 con ricetta di training completa

Open Dreamer: Reactor rilascia una riproduzione open-source in JAX/Flax della pipeline di world model Dreamer 4 con ricetta di training completa

Autore: MarkTechPost·

Punti chiave

  • Open Dreamer offre un’implementazione open-source della pipeline Dreamer 4 con un tokenizer video causale, un modello dinamico latente condizionato dalle azioni, utility di rollout e codice di valutazione FVD.
  • Il modello dinamico per Minecraft usa 1,6 miliardi di parametri su 30 layer block-causali ed è configurato per 200.000 step di training con l’ottimizzatore Muon.
  • Reactor riporta un utilizzo dei FLOP del modello tra il 57% e il 58% su GPU NVIDIA B200, con le activation identificate come principale costo di memoria anziché lo stato del modello.
  • Il team ha documentato diverse misure di stabilità, tra cui l’uso di EMA, confini di precisione mista, cambi di ottimizzatore, pesatura della loss e trasporto ottimale su minibatch.
  • Il rilascio non include cicli di training per behavior cloning o reinforcement learning, e il progetto non ha pubblicato punteggi FVD.
Open Dreamer: Reactor rilascia una riproduzione open-source in JAX/Flax della pipeline di world model Dreamer 4 con ricetta di training completa

Un gruppo di ricerca che opera con il nome Reactor ha rilasciato Open Dreamer, un’implementazione open-source della pipeline di world model Dreamer 4 realizzata in JAX e Flax NNX. I world model imparano a prevedere come evolve un ambiente a partire dai dati di osservazione, consentendo agli agenti di pianificare, apprendere e generare video senza interagire con un sistema reale. La serie Dreamer, sviluppata da Danijar Hafner, è stata una delle linee di ricerca più studiate nel reinforcement learning basato su modelli, con Dreamer 4 che estende l’approccio verso world model video su larga scala. Il progetto punta a riprodurre fedelmente la metodologia di ricerca di Dreamer 4 rendendo al tempo stesso pubblica l’intera pipeline di training — incluse le correzioni di stabilità e le configurazioni di calcolo — con un livello di dettaglio che le release di world model su larga scala hanno generalmente omesso.

Artefatti rilasciati

Sono stati pubblicati due repository di codice. Il primo, next-state/open-dreamer, contiene la pipeline di training completa: un tokenizer video causale, un modello dinamico latente condizionato dalle azioni, utility per la generazione di rollout e scoring FVD (Fréchet Video Distance). Il secondo, reactor-team/open-dreamer, fornisce un harness locale minimale per rollout in grado di generare fotogrammi video da un file MP4 di input abbinato a un file di azioni corrispondente.

Un terzo deliverable è una demo basata su browser eseguita sul runtime Reactor. Trasmette in tempo reale un ambiente Minecraft generato e include un toggle Game ⟷ Dream che trasferisce il flusso video tra il gioco reale e l’output del world model fotogramma per fotogramma.

Il team ha dichiarato che il suo obiettivo era riprodurre deliberatamente la ricerca Dreamer 4, evitando tecniche esterne al paper originale per mantenere ristretto lo spazio di ricerca. Lo sviluppo è iniziato su CoinRun, un platform 2D generato proceduralmente che può essere addestrato su una singola GPU, prima di scalare la pipeline funzionante a video di gameplay in stile Minecraft/VPT. Minecraft è stato un dominio di benchmark per l’apprendimento di agenti da dimostrazioni di gameplay sin dal progetto Video PreTraining (VPT) di OpenAI, che ha addestrato un modello di dinamica inversa su filmati YouTube etichettati da contractor per produrre dati di gameplay etichettati con azioni su larga scala.

Ulteriori dettagli sono disponibili nel blog del progetto, nel paper su arXiv e nell’annuncio di Reactor su X.

Architettura: un backbone, due modelli

Sia il tokenizer sia il modello dinamico condividono lo stesso backbone transformer block-causale, che alterna due tipi di attention. I layer spaziali propagano informazioni tra gli elementi all’interno di un singolo frame, mentre i layer temporali causali propagano informazioni tra frame.

Il tokenizer è progettato come Masked Autoencoder (MAE) basato su transformer, invece che come un tradizionale variational autoencoder (VAE). Il team riporta una compressione di circa 100× e osserva che questo design elimina la necessità di divergenza KL o loss avversarie. Secondo il team, l’approccio basato su masking rende lo spazio latente più adatto alla generazione basata su diffusione.

Il modello dinamico esegue la previsione del frame successivo utilizzando diffusion forcing, flow matching e shortcut models. Prevede anche l’azione successiva. Invece di alternare tra un modulo di transizione separato e un modulo di policy, il rollout è integrato in blocchi per timestep strutturati come (azione precedente, stato, policy). L’attention spaziale opera all’interno di ciascun blocco, mentre l’attention temporale causale collega i blocchi nel tempo.

Un vincolo progettuale critico: i token del world model non possono leggere il token dell’agente. Di conseguenza, le informazioni su task e policy possono influenzare gli stati futuri solo tramite l’azione successiva prevista.

Configurazione di training

I file di configurazione Minecraft inclusi specificano in dettaglio la ricetta di training.

Il modello dinamico comprende 1,6 miliardi di parametri su 30 layer block-causali, con d_model pari a 1920, 30 attention head e 3 KV head che usano grouped-query attention. Ogni quarto layer funziona come layer di time-attention. Ogni timestep contiene 32 register token appresi e un packing factor di 2 consolida i latenti vicini del tokenizer in ciascun token spaziale della dinamica. La time attention opera su una finestra scorrevole di 192 step.

Il training copre 200.000 step usando l’ottimizzatore Muon con una schedule WSD (warmup-stable-decay) e un learning rate di picco di 3e-4. Il campionamento shortcut e bootstrap si attiva allo step 100.000 con una frazione batch di 0,25. Il decadimento della media mobile esponenziale (EMA) è impostato a 0,999.

La configurazione del tokenizer produce 512 token latenti per frame con una larghezza di bottleneck pari a 16. I frame grezzi 360×640 vengono sottoposti a padding fino a 368×640 affinché entrambe le dimensioni spaziali siano divisibili in patch 16×16. La profondità dell’encoder è 12 con d_model 1536; la profondità del decoder è 8 con d_model 1024. La probabilità di masking MAE raggiunge al massimo 0,9 e la loss LPIPS viene applicata con peso 0,2 su metà dei timestep.

Le azioni VPT sono analizzate in 27 canali di azione binari e 121 classi mouse categoriche, senza canali continui.

Prestazioni di calcolo e strategia di memoria

Il team riporta un utilizzo dei FLOP del modello (MFU) del 57–58%, rispetto a un benchmark comunemente citato del 60% per un training transformer sano. La loro analisi utilizza un argomento roofline: su una NVIDIA B200, il punto di crossover tra operazioni limitate dalla bandwidth e operazioni limitate dal compute si verifica a 292 FLOP/byte. L’elaborazione di 256 frame per GPU spinge il workload oltre quel punto di cresta, nel regime compute-bound.

Le decisioni di sharding hanno smentito le aspettative iniziali. Con 1,6 miliardi di parametri, l’intero stato del modello — parametri, gradienti, stato dell’ottimizzatore ed EMA — occupava circa 24 GiB, rientrando in una singola B200. Le activation, più che lo stato del modello, si sono rivelate il principale costo di memoria. Il team ha sperimentato data parallelism, fully sharded data parallelism (FSDP), tensor parallelism e sequence parallelism prima di optare per un semplice data parallelism combinato con activation checkpointing.

Per il dataloading, il team ha pre-tokenizzato l’intero dataset in file .arrayrecord e ha usato Grain con un buffer di prefetch lato GPU. La decodifica standard basata su ffmpeg non era sufficiente a mantenere le GPU sature.

Ingegneria della stabilità

Il team di ricerca afferma esplicitamente che i problemi di stabilità hanno assorbito la quota maggiore del tempo di sviluppo. La loro osservazione centrale: la maggior parte dei problemi di stabilità emerge anche mentre la loss continua a diminuire. Il mean squared error (MSE) migliora in modo regolare mentre la qualità di generazione degrada simultaneamente — un fenomeno che rende inaffidabile il monitoraggio convenzionale basato sulla loss per i world model basati su diffusione.

Sono documentate sei correzioni specifiche:

  1. Cambio di ottimizzatore: Muon ha sostituito LaProp, che mostrava spike casuali e sempre più frequenti in due run di training separate da circa 400 ore B200 ciascuna.

  2. EMA come requisito obbligatorio: i pesi EMA sono trattati come essenziali per l’inferenza con diffusione, non opzionali.

  3. Confini della precisione mista: i parametri restano in float32, BF16 copre la maggior parte delle activation matmul e degli input di attention, e float32 viene mantenuto per i layer di normalizzazione e la testa di output del dynamics flow.

  4. Pesatura della loss: il team usa x-prediction con una loss in v-space, che si riduce a un termine di pesatura simile alla formulazione di Dreamer 4 ma con un denominatore al quadrato. Riporta un miglioramento piccolo ma percepibile.

  5. Trasporto ottimale: il trasporto ottimale baricentrico su minibatch applicato tra rumore e sequenze latenti ha migliorato la stabilità della generazione di rollout.

  6. μ-parametrization: testata e ritenuta non necessaria, in parte perché Muon mantiene la stabilità degli iperparametri in modo più efficace tra dimensioni di modello diverse.

Un ulteriore risultato dalla fase di sviluppo su CoinRun: uno sweep iso-FLOPs ha stimato lo scaling compute-optimal a circa N ∝ C^0.56 e D ∝ C^0.44.

Cosa non è incluso

Il repository non contiene il ciclo di training per behaviour cloning (BC) o reinforcement learning (RL). Un ciclo agente completo Dreamer 4 BC/RL è indicato come elemento aperto della roadmap. Il lavoro sulla policy CoinRun descritto nella documentazione del progetto non è stato usato per l’implementazione Minecraft e non è stato rilasciato.

Il progetto non pubblica punteggi FVD, anche se il repository include scripts/eval_fvd.py, un harness di valutazione basato su I3D configurato per 4 frame di contesto e un orizzonte di 240 frame.

Sintesi delle principali metriche ingegneristiche

  • Modello dinamico: 1,6B parametri, 30 layer, d_model 1920, addestrato per 200.000 step con Muon
  • Efficienza hardware: 57–58% MFU su GPU NVIDIA B200, 256 frame per GPU, circa 24 GiB di stato del modello
  • Sfida principale: stabilità, non throughput — le curve di loss nascondevano la maggior parte delle regressioni nella qualità di generazione