НовостиМакроOpen Dreamer: Reactor выпустила open-source-воспроизведение конвейера мировой модели Dreamer 4 на JAX/Flax с полным рецептом обучения

Open Dreamer: Reactor выпустила open-source-воспроизведение конвейера мировой модели Dreamer 4 на JAX/Flax с полным рецептом обучения

Автор: MarkTechPost·

Ключевые выводы

  • Open Dreamer предоставляет open-source-реализацию конвейера Dreamer 4 с causal video tokenizer, латентной динамической моделью, обусловленной действиями, rollout-утилитами и кодом оценки FVD.
  • Динамическая модель Minecraft использует 1.6 млрд параметров в 30 block-causal layers и настроена на 200,000 шагов обучения с оптимизатором Muon.
  • Reactor сообщает о 57%–58% model FLOPs utilization на GPU NVIDIA B200, при этом activations названы основной стоимостью памяти, а не состояние модели.
  • Команда задокументировала несколько мер стабильности, включая использование EMA, границы mixed precision, смену оптимизатора, взвешивание loss и minibatch optimal transport.
  • Релиз не включает training loops для behavior cloning или reinforcement learning, а проект не опубликовал FVD scores.
Open Dreamer: Reactor выпустила open-source-воспроизведение конвейера мировой модели Dreamer 4 на JAX/Flax с полным рецептом обучения

Исследовательская группа, работающая под названием Reactor, выпустила Open Dreamer — open-source-реализацию конвейера мировой модели Dreamer 4, построенную на JAX и Flax NNX. Мировые модели учатся предсказывать, как среда развивается на основе данных наблюдений, позволяя агентам планировать, обучаться и генерировать видео без взаимодействия с реальной системой. Серия Dreamer, разработанная Danijar Hafner, стала одним из наиболее изучаемых направлений в reinforcement learning на основе моделей, а Dreamer 4 расширяет этот подход в сторону крупномасштабных видео-мировых моделей. Проект нацелен на точное воспроизведение исследовательской методологии Dreamer 4 и одновременно делает полный конвейер обучения — включая исправления стабильности и вычислительные конфигурации — публично доступным, что обычно отсутствовало в релизах крупномасштабных мировых моделей.

Выпущенные артефакты

Были опубликованы два репозитория с кодом. Первый, next-state/open-dreamer, содержит полный конвейер обучения: causal video tokenizer, латентную динамическую модель, обусловленную действиями, утилиты генерации rollout и оценку FVD (Fréchet Video Distance). Второй, reactor-team/open-dreamer, предоставляет минимальный локальный rollout harness, способный генерировать видеокадры из входного MP4-файла в паре с соответствующим файлом действий.

Третьим результатом стала браузерная демоверсия, работающая на рантайме Reactor. Она транслирует сгенерированную среду Minecraft в реальном времени и включает переключатель Game ⟷ Dream, который покадрово переводит видеопоток между реальной игрой и выводом мировой модели.

Команда заявила, что ее цель состояла в намеренном воспроизведении исследования Dreamer 4 без использования техник за пределами оригинальной статьи, чтобы сохранить узкое пространство поиска. Разработка началась с CoinRun, процедурно генерируемого 2D-платформера, который можно обучать на одном GPU, после чего рабочий конвейер был масштабирован до видео игрового процесса в стиле Minecraft/VPT. Minecraft служит эталонной областью для обучения агентов на демонстрациях игрового процесса со времени проекта OpenAI Video PreTraining (VPT), в котором inverse dynamics model обучалась на размеченных подрядчиками YouTube-роликах для получения крупномасштабных данных игрового процесса с метками действий.

Дополнительные сведения доступны в блоге проекта, статье на arXiv и анонсе Reactor в X.

Архитектура: один backbone, две модели

И tokenizer, и динамическая модель используют один и тот же block-causal transformer backbone, который чередует два типа внимания. Space layers распространяют информацию между элементами внутри одного кадра, тогда как causal time layers передают информацию между кадрами.

Tokenizer спроектирован как Masked Autoencoder (MAE) на базе transformer, а не как традиционный variational autoencoder (VAE). Команда сообщает примерно о 100× сжатии и отмечает, что такая конструкция устраняет необходимость в KL divergence или adversarial losses. По их утверждению, подход с masking делает латентное пространство более пригодным для генерации на основе diffusion.

Динамическая модель выполняет предсказание следующего кадра с использованием diffusion forcing, flow matching и shortcut models. Она также предсказывает следующее действие. Вместо чередования отдельного transition module и policy module rollout интегрирован в блоки на каждом временном шаге со структурой (previous action, state, policy). Spatial attention работает внутри каждого блока, а causal temporal attention связывает блоки во времени.

Ключевое проектное ограничение: токены мировой модели не могут читать agent token. Следовательно, информация о задаче и политике может влиять на будущие состояния только через предсказанное следующее действие.

Конфигурация обучения

Поставляемые конфигурационные файлы Minecraft подробно задают рецепт обучения.

Динамическая модель включает 1.6 млрд параметров в 30 block-causal layers, с d_model 1920, 30 attention heads и 3 KV heads с использованием grouped-query attention. Каждый четвертый слой работает как time-attention layer. Каждый временной шаг содержит 32 обучаемых register tokens, а packing factor 2 объединяет соседние латенты tokenizer в каждый dynamics spatial token. Time attention работает по скользящему окну длиной 192 шага.

Обучение длится 200,000 шагов с оптимизатором Muon, расписанием WSD (warmup-stable-decay) и пиковым learning rate 3e-4. Shortcut- и bootstrap-сэмплирование активируется на шаге 100,000 с долей batch 0.25. Exponential moving average (EMA) decay установлен на 0.999.

Конфигурация tokenizer выдает 512 латентных токенов на кадр при ширине bottleneck 16. Исходные кадры 360×640 дополняются до 368×640, чтобы оба пространственных измерения ровно делились на патчи 16×16. Глубина encoder составляет 12 при d_model 1536; глубина decoder — 8 при d_model 1024. Вероятность MAE masking достигает максимума 0.9, а LPIPS loss применяется с весом 0.2 на половине временных шагов.

Действия VPT разбираются на 27 бинарных каналов действий и 121 категориальный класс мыши, без непрерывных каналов.

Вычислительная производительность и стратегия памяти

Команда сообщает о 57–58% model FLOPs utilization (MFU) по сравнению с часто упоминаемым ориентиром 60% для устойчивого обучения transformer. Их анализ использует roofline-аргумент: на NVIDIA B200 переход между режимами, ограниченными пропускной способностью и вычислениями, происходит на уровне 292 FLOP/byte. Обработка 256 кадров на GPU переводит нагрузку за эту пороговую точку в вычислительно ограниченный режим.

Решения по sharding оказались не такими, как ожидалось изначально. При 1.6 млрд параметров полное состояние модели — параметры, градиенты, состояние оптимизатора и EMA — занимало примерно 24 GiB, что помещается в один B200. Основной стоимостью памяти оказались activations, а не состояние модели. Команда экспериментировала с data parallelism, fully sharded data parallelism (FSDP), tensor parallelism и sequence parallelism, после чего остановилась на обычном data parallelism в сочетании с activation checkpointing.

Для dataloading команда предварительно токенизировала весь набор данных в файлы .arrayrecord и использовала Grain с prefetch buffer на стороне GPU. Стандартного декодирования на базе ffmpeg оказалось недостаточно, чтобы полностью загружать GPU.

Инженерия стабильности

Исследовательская команда прямо указывает, что проблемы стабильности заняли крупнейшую часть времени разработки. Их центральное наблюдение: большинство проблем стабильности возникает даже тогда, когда loss продолжает снижаться. Mean squared error (MSE) плавно улучшается, в то время как качество генерации одновременно ухудшается — явление, которое делает традиционный мониторинг на основе loss ненадежным для мировых моделей на базе diffusion.

Документированы шесть конкретных исправлений:

  1. Смена оптимизатора: Muon заменил LaProp, у которого наблюдались случайные и все более частые всплески в двух отдельных обучающих запусках примерно по 400 B200-часов каждый.

  2. EMA как обязательный компонент: веса EMA рассматриваются как необходимые для diffusion inference, а не как опциональные.

  3. Границы mixed precision: параметры остаются в float32, BF16 покрывает большинство matmul activations и attention inputs, а float32 сохраняется для слоев normalization и выходной головы dynamics flow.

  4. Взвешивание loss: команда использует x-prediction с v-space loss, что сводится к весовому члену, похожему на формулировку Dreamer 4, но с квадратом в знаменателе. Они сообщают о небольшом, но заметном улучшении.

  5. Optimal transport: minibatch barycentric optimal transport, примененный между шумом и латентными последовательностями, повысил стабильность генерации rollout.

  6. μ-parametrization: была протестирована и признана ненужной, отчасти потому что Muon эффективнее сохраняет стабильность гиперпараметров при изменении размеров моделей.

Еще один результат этапа разработки на CoinRun: iso-FLOPs sweep оценил compute-optimal scaling примерно как N ∝ C^0.56 и D ∝ C^0.44.

Что не включено

Репозиторий не содержит training loop для behaviour-cloning (BC) или reinforcement learning (RL). Полный цикл агента Dreamer 4 BC/RL указан как открытый пункт дорожной карты. Работа над policy для CoinRun, описанная в документации проекта, не использовалась для реализации Minecraft и не была выпущена.

Проект не публикует FVD scores, хотя в репозитории поставляется scripts/eval_fvd.py — evaluation harness на базе I3D, настроенный на 4 context frames и горизонт 240 кадров.

Сводка ключевых инженерных показателей

  • Динамическая модель: 1.6B параметров, 30 слоев, d_model 1920, обучение в течение 200,000 шагов с Muon
  • Эффективность аппаратного обеспечения: 57–58% MFU на GPU NVIDIA B200, 256 кадров на GPU, примерно 24 GiB состояния модели
  • Основная проблема: стабильность, а не пропускная способность — кривые loss скрывали большинство ухудшений качества генерации