Open Dreamer: Implementación de Código Abierto del Pipeline Dreamer 4 en JAX/Flax con Receta de Entrenamiento Completa
Un equipo de investigadores de IA lanzó Open Dreamer, una implementación de código abierto del pipeline del modelo de mundo Dreamer 4 construido con JAX y Flax NNX. El proyecto incluye un pipeline de entrenamiento, un arnés de inferencia local y una demostración en navegador que genera un mundo de Minecraft en tiempo real. El modelo utiliza un backbone de transformador con causalidad de bloques, alcanza un 57-58% de utilización de FLOPs del modelo, y el equipo enfatiza que la estabilidad fue el principal desafío, no el rendimiento.
Los investigadores Francesco Sacco, Diego Martí y Edward Hu, patrocinados por Reactor, lanzaron Open Dreamer, una implementación abierta de Dreamer 4 en JAX/Flax NNX. Se publicaron dos repositorios: uno para entrenamiento (tokenizador, modelo de dinámica, rollout, puntuación FVD) y otro para inferencia local. Una demo en el navegador transmite un mundo de Minecraft generado y permite alternar entre el juego real y el modelo del mundo. La arquitectura utiliza un transformer de causalidad de bloques como columna vertebral tanto para el tokenizador como para el modelo de dinámica. El tokenizador es un autoencoder enmascarado basado en transformer que logra una compresión de aproximadamente 100 veces sin pérdida KL ni adversaria. El modelo de dinámica (1.6 mil millones de parámetros) realiza predicción del siguiente fotograma utilizando forzado de difusión, flujo matching y modelos de atajo, y también predice la siguiente acción. El entrenamiento se ejecuta durante 200,000 pasos con el optimizador Muon, programación WSD, tasa de aprendizaje máxima de 3e-4 y decaimiento de EMA de 0.999. El equipo logró una utilización de FLOPs del modelo del 57-58% en GPUs B200, con el estado del modelo ocupando aproximadamente 24 GiB. Las activaciones fueron el principal costo de memoria, resuelto mediante paralelismo de datos y checkpointing de activaciones. Los problemas de estabilidad fueron el mayor desafío; seis soluciones incluyen el optimizador Muon, pesos EMA obligatorios, precisión mixta cuidadosa, ponderación de pérdida de espacio v, transporte óptimo barycéntrico con minilotes y restar énfasis a la parametrización μ. El repositorio no incluye el bucle de entrenamiento de clonación de comportamiento o refuerzo, y las puntuaciones FVD no se publican.
Fuente: MarkTechPost —
original