Open Dreamer: Dreamer 4 в открытом коде на JAX — полный рецепт обучения

Исследовательская группа Reactor опубликовала Open Dreamer — открытую реализацию пайплайна Dreamer 4 на JAX и Flax NNX. В двух репозиториях выложены полный код обучения (каузальный видео-токенизатор, латентная динамическая модель с учётом действий, генерация роллаутов и FVD-оценка) и минимальный локальный инструмент для генерации кадров из MP4 с файлом действий. Также доступно браузерное демо, которое стримит сгенерированный мир Minecraft в реальном времени с переключателем Game ↔ Dream.

Архитектура построена на едином block-causal transformer backbone, который чередует два типа внимания: пространственные слои обрабатывают элементы внутри одного кадра, а каузальные временные — между кадрами. Токенизатор — это transformer-based Masked Autoencoder (не VAE), обеспечивающий примерно 100-кратное сжатие без KL- или adversarial-лосса. Маскирование, по утверждению авторов, делает латентное пространство более пригодным для диффузии.

Динамическая модель предсказывает следующий кадр и обучается с помощью diffusion forcing, flow matching и shortcut models. Также она предсказывает следующее действие. Роллаут встроен в per-timestep блоки (предыдущее действие, состояние, политика) — пространственное внимание внутри блока, каузальное временное между блоками. Критически важно: токены world model не могут читать токен агента, поэтому информация о задаче и политике влияет на будущие состояния только через следующее действие.

Команда намеренно воспроизводила именно Dreamer 4, не используя методы за пределами этой статьи, чтобы сохранить чистоту эксперимента. Сначала пайплайн обкатали на CoinRun (2D платформер, обучаемый на одной GPU), затем масштабировали на Minecraft/VPT-видео.