NoticiasMacroOpen Dreamer: Reactor publica una reproducción de código abierto en JAX/Flax del pipeline del modelo mundial Dreamer 4 con la receta de entrenamiento completa

Open Dreamer: Reactor publica una reproducción de código abierto en JAX/Flax del pipeline del modelo mundial Dreamer 4 con la receta de entrenamiento completa

Autor: MarkTechPost·

Puntos clave

  • Open Dreamer proporciona una implementación de código abierto del pipeline de Dreamer 4 con un tokenizador de video causal, un modelo dinámico latente condicionado por acciones, utilidades de rollout y código de evaluación FVD.
  • El modelo dinámico de Minecraft utiliza 1.600 millones de parámetros distribuidos en 30 capas de bloques causales y está configurado para 200.000 pasos de entrenamiento con el optimizador Muon.
  • Reactor reporta un 57% a 58% de utilización de FLOPs del modelo en GPUs NVIDIA B200, identificando las activaciones como el principal costo de memoria en lugar del estado del modelo.
  • El equipo documentó múltiples medidas de estabilidad, incluyendo el uso de EMA, límites de precisión mixta, cambios de optimizador, ponderación de pérdidas y transporte óptimo por minilote.
  • La publicación no incluye bucles de entrenamiento de clonación de comportamiento ni de aprendizaje por refuerzo, y el proyecto no ha publicado puntuaciones FVD.
Open Dreamer: Reactor publica una reproducción de código abierto en JAX/Flax del pipeline del modelo mundial Dreamer 4 con la receta de entrenamiento completa

Un grupo de investigación que opera bajo el nombre de Reactor ha publicado Open Dreamer, una implementación de código abierto del pipeline del modelo mundial Dreamer 4 desarrollada en JAX y Flax NNX. Los modelos mundiales aprenden a predecir cómo evoluciona un entorno a partir de datos de observación, lo que permite a los agentes planificar, aprender y generar video sin interactuar con un sistema real. La serie Dreamer, desarrollada por Danijar Hafner, ha sido una de las líneas de trabajo más estudiadas en el aprendizaje por refuerzo basado en modelos, y Dreamer 4 amplía el enfoque hacia modelos mundiales de video a gran escala. El proyecto busca reproducir fielmente la metodología de investigación de Dreamer 4 y hacer que el pipeline de entrenamiento completo —incluidas las correcciones de estabilidad y las configuraciones de cómputo— sea accesible al público, un nivel de detalle que las publicaciones de modelos mundiales a gran escala han omitido habitualmente.

Artefactos publicados

Se publicaron dos repositorios de código. El primero, next-state/open-dreamer, contiene el pipeline de entrenamiento completo: un tokenizador de video causal, un modelo dinámico latente condicionado por acciones, utilidades de generación de rollouts y evaluación FVD (Fréchet Video Distance). El segundo, reactor-team/open-dreamer, proporciona un harness mínimo de rollout local capaz de generar fotogramas de video a partir de un archivo MP4 de entrada junto con su archivo de acciones correspondiente.

Un tercer entregable es una demostración basada en navegador que se ejecuta en el runtime de Reactor. Transmite en tiempo real un entorno generado de Minecraft e incluye un conmutador Game ⟷ Dream que transiciona la transmisión de video entre el juego real y la salida del modelo mundial fotograma a fotograma.

El equipo declaró que su objetivo era reproducir deliberadamente la investigación de Dreamer 4, evitando técnicas fuera del artículo original para mantener un espacio de búsqueda reducido. El desarrollo comenzó en CoinRun, un plataformas 2D generado proceduralmente que puede entrenarse en una sola GPU, antes de escalar el pipeline funcional a video de jugabilidad estilo Minecraft/VPT. Minecraft ha servido como dominio de referencia para el aprendizaje de agentes a partir de demostraciones de jugabilidad desde el proyecto Video PreTraining (VPT) de OpenAI, el cual entrenó un modelo de dinámica inversa con metraje de YouTube etiquetado por contratistas para producir datos de jugabilidad etiquetados con acciones a gran escala.

Hay detalles adicionales disponibles en el blog del proyecto, el artículo de arXiv y el anuncio de Reactor en X.

Arquitectura: una columna vertebral, dos modelos

Tanto el tokenizador como el modelo dinámico comparten la misma columna vertebral transformadora de bloques causales, que alterna entre dos tipos de atención. Las capas espaciales propagan información entre los elementos dentro de un mismo fotograma, mientras que las capas temporales causales propagan información entre fotogramas.

El tokenizador está diseñado como un Autoencoder Enmascarado (MAE) basado en transformador, en lugar de un autoencoder variacional (VAE) tradicional. El equipo reporta una compresión de aproximadamente 100× y señala que este diseño elimina la necesidad de divergencia KL o pérdidas adversariales. Argumentan que el enfoque de enmascaramiento hace que el espacio latente sea más adecuado para la generación basada en difusión.

El modelo dinámico realiza predicción del siguiente fotograma mediante difusión forzada (diffusion forcing), flow matching y modelos de atajo (shortcut models). También predice la siguiente acción. En lugar de alternar entre un módulo de transición separado y un módulo de política, el rollout se integra en bloques por paso de tiempo estructurados como (acción anterior, estado, política). La atención espacial opera dentro de cada bloque, y la atención temporal causal conecta los bloques a lo largo del tiempo.

Una restricción de diseño crítica: los tokens del modelo mundial no pueden leer el token del agente. En consecuencia, la información de tarea y política solo puede influir en los estados futuros a través de la siguiente acción predicha.

Configuración de entrenamiento

Los archivos de configuración de Minecraft publicados especifican la receta de entrenamiento en detalle.

El modelo dinámico comprende 1.600 millones de parámetros distribuidos en 30 capas de bloques causales, con d_model de 1920, 30 cabezas de atención y 3 cabezas KV mediante atención de consulta agrupada (grouped-query attention). Cada cuarta capa funciona como capa de atención temporal. Cada paso de tiempo incluye 32 tokens de registro aprendidos, y un factor de empaquetado (packing factor) de 2 consolida los latentes vecinos del tokenizador en cada token espacial del modelo dinámico. La atención temporal opera sobre una ventana deslizante de 192 pasos.

El entrenamiento abarca 200.000 pasos con el optimizador Muon, un programa WSD (warmup-stable-decay) y una tasa de aprendizaje máxima de 3e-4. El muestreo de atajo y bootstrap se activa en el paso 100.000 con una fracción de lote de 0,25. La descomposición de media móvil exponencial (EMA) se establece en 0,999.

La configuración del tokenizador produce 512 tokens latentes por fotograma con un ancho de cuello de botella de 16. Los fotogramas originales de 360×640 se rellenan a 368×640 para que ambas dimensiones espaciales se dividan uniformemente en parches de 16×16. La profundidad del codificador es 12 con d_model 1536; la profundidad del decodificador es 8 con d_model 1024. La probabilidad de enmascaramiento MAE alcanza un máximo de 0,9, y la pérdida LPIPS se aplica con un peso de 0,2 en la mitad de los pasos de tiempo.

Las acciones VPT se analizan en 27 canales de acción binarios y 121 clases categóricas de mouse, sin canales continuos.

Rendimiento de cómputo y estrategia de memoria

El equipo reporta un 57–58% de utilización de FLOPs del modelo (MFU), en comparación con un punto de referencia comúnmente citado del 60% para un entrenamiento saludable de transformadores. Su análisis utiliza un argumento de roofline: en una NVIDIA B200, el cruce entre una operación limitada por ancho de banda y una limitada por cómputo se sitúa en 292 FLOP/byte. Procesar 256 fotogramas por GPU empuja la carga de trabajo más allá de ese punto de inflexión hacia el régimen limitado por cómputo.

Las decisiones de fragmentación (sharding) desafiaron las expectativas iniciales. Con 1.600 millones de parámetros, el estado completo del modelo —parámetros, gradientes, estado del optimizador y EMA— ocupó aproximadamente 24 GiB, lo cual cabe en una sola B200. Las activaciones, y no el estado del modelo, resultaron ser el principal costo de memoria. El equipo experimentó con paralelismo de datos, paralelismo de datos completamente fragmentado (FSDP), paralelismo de tensores y paralelismo de secuencias antes de optar por un paralelismo de datos simple combinado con checkpointing de activaciones.

Para la carga de datos, el equipo pre-tokenizó todo el conjunto de datos en archivos .arrayrecord y utilizó Grain con un búfer de prefetch del lado de la GPU. La decodificación estándar basada en ffmpeg no fue suficiente para mantener las GPUs saturadas.

Ingeniería de estabilidad

El equipo de investigación declara explícitamente que los problemas de estabilidad consumieron la mayor parte de su tiempo de desarrollo. Su observación central: la mayoría de los problemas de estabilidad surgen incluso cuando la pérdida sigue disminuyendo. El error cuadrático medio (MSE) mejora de manera fluida mientras que la calidad de generación se degrada simultáneamente, un fenómeno que hace que la monitorización convencional basada en pérdidas no sea confiable para modelos mundiales basados en difusión.

Se documentan seis correcciones específicas:

  1. Cambio de optimizador: Muon reemplazó a LaProp, que presentaba picos aleatorios y cada vez más frecuentes en dos ejecuciones de entrenamiento separadas de aproximadamente 400 horas B200 cada una.

  2. EMA como obligatorio: Los pesos EMA se tratan como esenciales para la inferencia de difusión, no como opcionales.

  3. Límites de precisión mixta: Los parámetros se mantienen en float32, BF16 cubre la mayoría de las activaciones de matmul y las entradas de atención, y float32 se conserva para las capas de normalización y la cabeza de salida de flujo dinámico.

  4. Ponderación de pérdidas: El equipo utiliza x-prediction con una pérdida en espacio v, que se reduce a un término de ponderación similar a la formulación de Dreamer 4 pero con un denominador al cuadrado. Reportan una mejora pequeña pero perceptible.

  5. Transporte óptimo: El transporte óptimo baricéntrico por minilote aplicado entre secuencias de ruido y latentes mejoró la estabilidad de la generación de rollouts.

  6. μ-parametrización: Se probó y se consideró innecesaria, en parte porque Muon mantiene la estabilidad de los hiperparámetros de manera más efectiva en diferentes tamaños de modelo.

Un hallazgo adicional de la fase de desarrollo en CoinRun: un barrido iso-FLOPs estimó el escalado óptimo en cómputo en aproximadamente N ∝ C^0.56 y D ∝ C^0.44.

Lo que no se incluye

El repositorio no contiene el bucle de entrenamiento de clonación de comportamiento (BC) ni de aprendizaje por refuerzo (RL). Un bucle completo de agente BC/RL de Dreamer 4 aparece como un elemento pendiente en la hoja de ruta. El trabajo de política de CoinRun descrito en la documentación del proyecto no se utilizó para la implementación de Minecraft y no fue publicado.

El proyecto no publica puntuaciones FVD, aunque el repositorio incluye scripts/eval_fvd.py, un harness de evaluación basado en I3D configurado para 4 fotogramas de contexto y un horizonte de 240 fotogramas.

Resumen de cifras clave de ingeniería

  • Modelo dinámico: 1.600 millones de parámetros, 30 capas, d_model 1920, entrenado durante 200.000 pasos con Muon
  • Eficiencia de hardware: 57–58% MFU en GPUs NVIDIA B200, 256 fotogramas por GPU, aproximadamente 24 GiB de estado del modelo
  • Desafío principal: estabilidad, no rendimiento — las curvas de pérdida ocultaron la mayoría de las regresiones en calidad de generación