В мире генеративного искусственного интеллекта и, в частности, в области обучения с подкреплением (RL) и моделирования мира, появление качественных открытых реализаций сложных архитектур — это всегда событие. Недавно группа исследователей из Reactor выпустила Open Dreamer, проект, который детально воспроизводит пайплайн модели мира Dreamer 4. Это не просто «обертка» или демонстрация, а полноценная реализация на фреймворках JAX и Flax NNX, включающая полный рецепт обучения. Для инженеров и исследователей, работающих с физическим ИИ (Physical AI) и симуляциями, это открывает новые горизонты для локального запуска и экспериментов.
Что делает этот релиз особенным? В отличие от многих проектов, где код является лишь тенью оригинальной статьи, здесь предоставлены рабочие репозитории для обучения и инференса, а также браузерная демонстрация. Исследователи пошли по пути строгой репродукции: они намеренно избегали сторонних методов, не описанных в оригинальной статье Dreamer 4, чтобы сохранить чистоту эксперимента. Начав с простой игры CoinRun, они масштабировали пайплайн до сложного геймплея в стиле Minecraft/VPT, что демонстрирует robustность подхода.
В этой статье мы подробно разберем архитектуру Open Dreamer, технические детали обучения, проблемы с памятью и стабильностью, а также практические выводы, которые можно извлечь для своих проектов. Мы рассмотрим, как работает токенизатор, почему команда выбрала конкретные оптимизаторы и как им удалось достичь высокой эффективности использования вычислительных ресурсов на современных GPU.
01Архитектура: один бэкбон, две модели
Сердцем Open Dreamer является единый бэкбон на основе трансформера с блочно-каузальным вниманием. Эта архитектура универсальна и используется как для токенизатора, так и для модели динамики. Ключевая особенность заключается в чередовании двух типов слоев внимания:
- Пространственные слои (Space layers): Они распространяют информацию среди элементов внутри одного кадра. Это позволяет модели понимать структуру сцены, объекты и их взаимное расположение в данный момент времени.
- Каузальные временные слои (Causal time layers): Эти слои передают информацию между кадрами, обеспечивая понимание движения, изменений и причинно-следственных связей во времени.
Такой подход позволяет модели эффективно обрабатывать как пространственные, так и временные зависимости, что критически важно для предсказания будущего состояния среды.
Токенизатор: MAE вместо VAE
Одним из самых интересных архитектурных решений в Open Dreamer является использование трансформерного Masked Autoencoder (MAE) в качестве токенизатора, а не вариационного автоэнкодера (VAE), который часто встречается в подобных задачах. Команда сообщает о сжатии данных примерно в 100 раз. Важное преимущество MAE заключается в том, что ему не требуются потери KL (Kullback-Leibler) или состязательные потери (adversarial loss), которые часто усложняют обучение и могут приводить к нестабильности.
Исследователи утверждают, что маскирование делает латентное пространство более «диффузионным» (diffusible), что упрощает последующее моделирование динамики. Токенизатор генерирует 512 латентных токенов на кадр при ширине бутылочного горлышка 16. Исходные кадры размером 360x640 пикселей дополняются до 368x640, чтобы оба измерения делились на патчи 16x16. Глубина энкодера составляет 12 слоев с d_model=1536, а декодера — 8 слоев с d_model=1024. Вероятность маскировки достигает 0.9, а метрика LPIPS применяется с весом 0.2 на половине временных шагов.
Модель динамики: предсказание следующего состояния
Модель динамики отвечает за предсказание следующего кадра и следующего действия агента. Она обучается с использованием диффузионного принуждения (diffusion forcing), flow matching и shortcut-моделей. В отличие от традиционных подходов, где модуль перехода (transition module) и политика (policy) разделены, здесь процесс развертывания (rollout) встроен в блоки на каждом временном шаге, состоящие из (предыдущего действия, состояния, политики).
Важно отметить критическое ограничение: токены модели мира не могут «читать» токен агента. Информация о задаче и политике может влиять на будущие состояния только через следующее действие. Это создает четкую границу между восприятием мира и принятием решений, что упрощает обучение и делает модель более интерпретируемой.
02Рецепт обучения: детали конфигурации
Команда предоставила конкретные конфигурации для обучения на Minecraft, которые делают рецепт воспроизводимым. Давайте разберем ключевые параметры модели динамики и процесса обучения.
Модель динамики содержит 1.6 миллиарда параметров. Она состоит из 30 слоев блочно-каузального трансформера с d_model=1920, 30 головами внимания и 3 головами KV для группированного внимания (Grouped-Query Attention). Каждое четвертое слое является слоем временного внимания. Каждый временной шаг содержит 32 обучаемых регистрационных токена (register tokens), а packing_factor=2 упаковывает соседние латентные токены токенизатора в каждый пространственный токен динамики. Временное внимание использует скользящее окно размером 192 шага.
Обучение длится 200 000 шагов с использованием оптимизатора Muon, расписания WSD (Weighted Step Decay) и пиковой скорости обучения 3e-4. Shortcut/bootstrap-выборки включаются на шаге 100 000 с долей батча 0.25. Скорость затухания EMA (Exponential Moving Average) составляет 0.999.
Действия VPT (Video PreTraining) парсятся в 27 бинарных каналов действий и 121 категориальный класс мыши, без непрерывных каналов. Это упрощает задачу предсказания действий, превращая ее в классификацию, а не регрессию.
03Вычислительная мощность и проблема памяти
Одной из самых сложных частей при обучении больших моделей является эффективное использование вычислительных ресурсов. Команда Open Dreamer сообщает о использовании FLOPs модели (MFU) на уровне 57–58%, что близко к заявленному бенчмарку в 60% для здорового обучения трансформеров.
Объяснение лежит в области «roofline» (крыши производительности). На GPU B200 точка пересечения между ограничением пропускной способности (bandwidth-bound) и вычислительной мощностью (compute-bound) находится на уровне 292 FLOP/byte. Подача 256 кадров на GPU выводит рабочую нагрузку за эту точку, делая процесс вычислительно ограниченным.
Шардинг (sharding) пошел неожиданным путем. При 1.6 миллиардах параметров состояние модели — параметры, градиенты, состояние оптимизатора и EMA — занимало около 24 ГБ, что легко помещается на B200. Настоящей проблемой были активации. Команда попробовала параллелизм данных, FSDP, тензорный параллелизм и параллелизм последовательностей, но в итоге остановилась на простом параллелизме данных плюс чекпоинтинг активаций (activation checkpointing).
Загрузка данных была решена путем предварительного токенизации всего набора данных в файлы .arrayrecord, а затем использования Grain с буфером предварительной выборки на стороне GPU. Декодирование ffmpeg оказалось недостаточно быстрым, чтобы поддерживать загрузку GPU.
04Стабильность: главный вызов
Исследователи прямо заявляют, что стабильность заняла самую большую долю их времени. Их ключевое наблюдение: большинство проблем со стабильностью возникают несмотря на снижение функции потерь. MSE (Mean Squared Error) улучшается плавно, но качество генерации ухудшается. Это классическая проблема, когда модель учится предсказывать среднее значение, а не реалистичные кадры.
Документировано шесть основных исправлений:
- Замена оптимизатора: Muon заменил LaProp, который вызывал случайные и все более частые всплески, особенно в двух запусках по 400 часов на B200 каждый.
- EMA веса: Веса EMA считаются обязательными для инференса с диффузией. Без них качество генерации резко падает.
- Смешанная точность: Она чувствительна к границам. Параметры остаются float32, BF16 покрывает большинство matmul-активаций и входов внимания, а float32 сохраняется для нормализации и выходной головки flow динамики.
- Взвешивание потерь: Используется x-prediction с потерей в v-пространстве, что сводится к члену взвешивания, аналогичному Dreamer 4, но с квадратом в знаменателе. Это дало небольшое, но заметное улучшение.
- Оптимальный транспорт: Минибатчный барицентрический оптимальный транспорт между шумом и латентными последовательностями сделал развертывание генерации более стабильным.
- μ-параметризация: Была протестирована, но признана ненужной, частично потому, что Muon удерживает гиперпараметры более стабильными при изменении размера модели.
Еще один результат из фазы CoinRun: сканирование iso-FLOPs показало, что оптимальное масштабирование вычислений составляет примерно 0.56 и 0.44. Это указывает на то, что существует оптимальный баланс между размером модели и количеством вычислений, который не всегда соответствует простому увеличению параметров.
05Что отсутствует в коробке
Важно понимать, что Open Dreamer — это не полный агент. Репозиторий не включает цикл обучения по поведению (behaviour-cloning) или RL. Полный цикл агента Dreamer 4 BC/RL указан как открытый элемент дорожной карты. Работа с политикой CoinRun, описанная в посте, не использовалась для Minecraft и не была выпущена.
Также в посте не публикуются оценки FVD (Fréchet Video Distance), хотя скрипт scripts/eval_fvd.py поставляется сHarness на основе I3D, настроенным на 6 кадров контекста и горизонт 240 кадров. Это позволяет пользователям самостоятельно оценивать качество генерации видео.
06Что это значит на практике
Для разработчиков и исследователей в России и СНГ Open Dreamer представляет значительный интерес по нескольким причинам. Во-первых, использование JAX и Flax позволяет эффективно использовать современные GPU, включая NVIDIA B200, которые становятся все более доступными через облачные провайдеры. Во-вторых, открытый код и полный рецепт обучения позволяют воспроизвести результаты без необходимости доступа к закрытым ресурсам оригинальной команды.
Практическое применение включает:
- Локальное обучение: Благодаря оптимизациям памяти и загрузки данных, запуск на мощных потребительских GPU (например, RTX 4090/5090) становится возможным, хотя и требует тщательной настройки.
- Исследования в области Physical AI: Модель мира может использоваться для симуляции физических процессов, что критически важно для робототехники и автономных систем.
- Генерация видео: Архитектура токенизатора и динамики может быть адаптирована для задач генерации видео с высоким уровнем детализации.
Команда Reactor проделала огромную работу, документировав не только успешные шаги, но и проблемы со стабильностью. Это экономит время других исследователей, позволяя им избежать типичных ошибок. Если вы работаете с моделями мира или генеративным ИИ, Open Dreamer — это обязательный к изучению проект.
Вы можете ознакомиться с блог-постом и демонстрацией, репозиторием обучения, репозиторием инференса и Reactor в X. Весь кредит за это исследование принадлежит исследователям этого проекта.
Источник: MarkTechPost ↗