Что такое Open Dreamer и зачем это нужно
Исследователи из группы Reactor опубликовали Open Dreamer, точную репродукцию пайплайна модели мира Dreamer 4, написанную на JAX и Flax NNX. В отличие от многих open-source проектов, здесь предоставлен не только инференс, но и полный рецепт обучения (training recipe), что позволяет исследователям воспроизвести результаты и провести собственные эксперименты.
Проект состоит из двух основных репозиториев: next-state/open-dreamer для обучения (включая causal video tokenizer и модель динамики) и reactor-team/open-dreamer для локального роулта. Также доступна браузерная демо-версия, где можно переключаться между реальным геймплеем Minecraft и генерацией мира моделью в реальном времени.
Архитектура: один бэкбон, две задачи
Ключевая особенность архитектуры — использование единого трансформерного бэкбона для токенизатора и модели динамики. Бэкбон чередует два типа внимания:
- Space layers: распространяют информацию между элементами одного кадра.
- Causal time layers: передают информацию между кадрами во времени.
Токенизатор реализован как трансформерный Masked Autoencoder (MAE), а не VAE, что обеспечивает сжатие примерно в 100 раз без необходимости в KL- или adversarial-потерях. Модель динамики предсказывает следующее состояние и действие, используя диффузионное принуждение (diffusion forcing) и flow matching.
Технические спецификации и гиперпараметры
Модель динамики имеет 1.6 миллиарда параметров. Обучение проводилось на 200 000 шагов с использованием оптимизатора Muon. Ниже приведена сводка конфигураций:
| Компонент | Параметры / Конфигурация |
|---|---|
| Модель динамики | 1.6B параметров, 30 слоев, d_model=1920, 30 attention heads, GQA (3 KV heads) |
| Токенизатор | 512 латентных токенов на кадр, bottleneck width=16, MAE masking до 0.9 |
| Обучение | 200k шагов, оптимизатор Muon, LR peak=3e-4, EMA decay=0.999 |
| Входные данные | Кадры 360x640 (пэд до 368x640), 27 бинарных + 121 категориальный канал действий |
| Hardware | NVIDIA B200, 256 кадров на GPU, ~24 GiB состояния модели |
Инженерные вызовы: стабильность важнее пропускной способности
Исследователи отмечают, что главным препятствием стала не вычислительная мощность, а стабильность обучения. Была достигнута MFU (Model FLOPs Utilization) на уровне 57–58% на B200, что близко к теоретическому пределу в 60% для трансформеров.
Ключевые инженерные решения:
- Оптимизатор: Muon заменил LaProp, который вызывал случайные всплески потерь.
- Точность: Параметры хранятся в float32, BF16 используется для matmul и attention, float32 — для нормализации и выхода динамики.
- Параллелизм: Отказ от FSDP и tensor parallelism в пользу простого data parallelism + activation checkpointing из-за высокой стоимости активаций.
- Данные: Предварительное токенизирование датасета в .arrayrecord файлы с GPU-side prefetch через Grain, так как ffmpeg не успевал за GPU.
Что осталось за бортом
В релиз не вошли циклы обучения по поведению (behavior-cloning) и RL-агент. Также не опубликованы финальные метрики FVD (Fréchet Video Distance), хотя скрипт для их оценки (scripts/eval_fvd.py) доступен. Команда планирует выпустить полный RL-пайплайн в будущем, а текущая работа с CoinRun не использовалась для Minecraft и не опубликована.
Источник: MarkTechPost ↗