Open Dreamer: Pesquisadores Lançam Reprodução Aberta do Dreamer 4 em JAX/Flax com Pipeline de Treinamento Completo

O que é o Open Dreamer

Um pequeno grupo de pesquisadores do coletivo Reactor lançou o Open Dreamer, uma implementação de código aberto do pipeline de modelos de mundo Dreamer 4, escrita inteiramente em JAX e Flax NNX. O projeto é uma reprodução fiel do artigo original do Dreamer 4, evitando deliberadamente qualquer método não descrito no paper para manter o escopo controlado.

O que foi publicado

Dois repositórios foram disponibilizados. O primeiro, next-state/open-dreamer, contém o pipeline de treinamento completo: tokenizador causal de vídeo, modelo de dinâmica latente condicionado por ações, geração de rollouts e pontuação FVD. O segundo, reactor-team/open-dreamer, oferece um harness mínimo de inferência local que gera quadros a partir de um arquivo MP4 e ações correspondentes.

Um terceiro artefato é a demo no navegador hospedada no runtime Reactor. Ela transmite um mundo de Minecraft gerado em tempo real e oferece um toggle Game ↔ Dream que entrega o fluxo do jogo real para o modelo de mundo, quadro a quadro.

A equipe começou com o CoinRun, um jogo de plataforma 2D gerado proceduralmente que pode ser treinado em uma única GPU, e depois escalou o pipeline funcional para vídeos de gameplay no estilo Minecraft/VPT.

Arquitetura: um backbone, dois modelos

Tanto o tokenizador quanto o modelo de dinâmica usam o mesmo backbone transformer block-causal. Esse backbone alterna dois tipos de atenção: camadas espaciais propagam informações entre elementos de um único quadro, e camadas temporais causais propagam informações entre quadros.

O tokenizador é um Masked Autoencoder (MAE) baseado em transformer, não um VAE. A equipe relata compressão de aproximadamente 100× e destaca que o design não precisa de perda adversarial ou KL. O mascaramento, segundo os pesquisadores, torna o espaço latente mais “difusível”.

O modelo de dinâmica realiza predição do próximo quadro e é treinado com diffusion forcing, flow matching e shortcut models. Ele também prevê a próxima ação. Tokens do modelo de mundo não podem ler o token do agente, de modo que informações de tarefa e política só influenciam estados futuros através da próxima ação.

A receita de treinamento

O modelo de dinâmica tem 1,6 bilhão de parâmetros: 30 camadas block-causais, d_model 1920, 30 cabeças de atenção e 3 cabeças KV para grouped-query attention. Cada quarta camada é de atenção temporal. O treinamento roda por 200.000 passos com o otimizador Muon, um schedule WSD e learning rate de pico de 3e-4.

O tokenizador emite 512 tokens latentes por quadro com largura de bottleneck de 16. Quadros de 360×640 são preenchidos para 368×640 para divisão em patches 16×16. A profundidade do encoder é 12 com d_model 1536; o decoder tem profundidade 8 com d_model 1024.

O muro da memória e a engenharia de GPU

A equipe relata 57-58% de utilização de FLOPs do modelo em GPUs B200, contra um benchmark de 60% para treinamento saudável de transformers. Com 256 quadros por GPU, a carga de trabalho ultrapassa o ponto de crossover entre operações limitadas por largura de banda e por computação.

O estado do modelo (parâmetros, gradientes, estado do otimizador e EMA) ocupa aproximadamente 24 GiB, cabendo em uma B200. O verdadeiro custo foram as ativações. A equipe testou data parallelism, FSDP, tensor parallelism e sequence parallelism, optando por data parallelism simples com checkpointing de ativação.

Estabilidade: o verdadeiro desafio

Os pesquisadores são explícitos: estabilidade consumiu a maior parte do tempo, e a maioria dos problemas ocorre apesar de a perda (MSE) estar caindo — o loss melhora suavemente enquanto a qualidade da geração se degrada. Seis correções são documentadas:

• Muon substituiu LaProp, que apresentava spikes aleatórios cada vez mais frequentes.
• Pesos EMA são tratados como obrigatórios para inferência por difusão.
• Precisão mista com parâmetros em float32, BF16 para a maioria das ativações de matmul, e float32 para normalização e cabeça de saída do fluxo de dinâmica.
• Ponderação de perda usando x-prediction com v-space loss, que reduz a um termo similar ao do Dreamer 4 com denominador ao quadrado.
• Transporte ótimo baricêntrico em minibatch entre ruído e sequências latentes tornou a geração de rollouts mais estável.
• μ-parametrização foi testada e considerada desnecessária, já que Muon mantém hiperparâmetros estáveis entre tamanhos de modelo.

O que ainda não está incluído

O repositório não inclui o loop de treinamento de behavior cloning ou RL — um loop completo de agente BC/RL do Dreamer 4 está listado como item do roadmap. O trabalho com política no CoinRun descrito no post não foi usado para Minecraft e não foi publicado. Scores FVD também não foram divulgados, embora o script de avaliação eval_fvd.py esteja incluído.

Por que isso importa

O Open Dreamer é a primeira reprodução aberta completa do Dreamer 4 com código de treinamento funcional, documentação detalhada de engenharia e uma demo interativa no navegador. Para times que trabalham com world models, agentes em ambientes 3D ou pesquisa em reinforcement learning, o projeto oferece um ponto de partida sólido — com lições reais de estabilidade e engenharia de GPU que vão muito além do paper original.