Соревнование 1X World Model

Прогресс в генерации видео может скоро сделать возможной оценку политик роботов в полностью обученных моделях мира. Сквозной обучаемый симулятор, который моделирует миллионы сред для роботов, заметно ускорит прогресс в универсальной робототехнике и даст полезный сигнал для масштабирования данных и вычислений.

Чтобы ускорить прогресс в симуляторах для обучения роботов, мы объявляем соревнование 1X World Model: нужно предсказывать будущие наблюдения от первого лица для EVE Android. Мы предоставляем более 100 часов токенов изображений векторного квантования и исходные данные действий, собранные при управлении EVE, базовую модель мира (в стиле GENIE) и покадровый автоэнкодер MAGVIT2, который сжимает изображения в токены 16×16 и восстанавливает из них изображения.

Мы надеемся, что этот набор данных поможет исследователям робототехники, которые хотят работать в средах, рассчитанных на людей. Достаточно сильная модель мира позволит любому получить доступ к «нейросимулированной EVE». Конечная цель — задача оценки; промежуточные цели мы награждаем денежными призами, например за хорошее согласование с данными (соревнование по сжатию) и за генерацию правдоподобных видео (соревнование по сэмплированию).

Набор данных на Hugging Face

Присоединиться к Discord

Блог первого этапа, блог второго этапа

Следите за обновлениями второго этапа соревнования 1X World Model!

В исходном документе здесь была сетка анимаций с примерами генерации; в этой версии она опущена.

Задачи

Каждый пример — последовательность из 16 кадров от первого лица с робота, частота 2 Гц (всего 8 секунд). Нужно по предыдущим изображениям предсказать следующее изображение.

Соревнование по сжатию (призовой фонд $10 k) * * : предсказатьдискретноераспределениетокеновследующегоизображения . - Критерий : нанашемзакрытомтестовомнаборесначаладобиться * * потерьвременного t e a c h er f orc in g ниже 8.0 * * . - * * Соревнованиепосэмплированию (призовойфонд$ 10k): методы предсказания будущего не обязаны ограничиваться предсказанием следующего логита. Например, можно генерировать будущие изображения с помощью GAN, диффузионных моделей или MaskGIT. Критерии опубликуем позже.
Соревнование по оценке (скоро): дан набор из N политик $π_{1}, π_{2}, \dots π_{N}$ , где каждая политика $π_{i} (a_{t} ∣ z_{t})$ предсказывает токены действий по токенам изображений. Можно ли оценить все политики в «модели мира» $p (z_{t + 1} ∣ z_{t}, a_{t})$ и сказать, какие политики окажутся лучшими?

Эти соревнования во многом вдохновлены соревнованием по сжатию commavq. Подробности см. в разделе «Дополнительные детали соревнования».

Начало работы

Требуется Python 3.10 или новее. Код проверен на Python 3.10.12.

# Установить зависимости и загрузить данные
./build.sh

# Активировать окружение Python
source venv/bin/activate

conda create -n your_environment_name python=3.10.12
conda activate your_environment_name

# Установка зависимостей
conda install -c conda-forge accelerate=0.30.1 torchvision=0.18.0 lpips=0.1.4 matplotlib tqdm wandb xformers=0.0.26.post1 wheel packaging ninja einops

# Для transformers и torch (проверьте совместимость с вашей версией CUDA)
pip install transformers==4.41.0 lightning>2.3.1 git+https://github.com/janEbert/mup.git@fsdp-fix

pip install triton
FLASH_ATTENTION_SKIP_CUDA_BUILD=TRUE python -m pip install flash-attn==2.5.8 --no-build-isolation
huggingface-cli download 1x-technologies/worldmodel --repo-type dataset --local-dir data

Нужно выполнить:

export XFORMERS_FORCE_DISABLE_TRITON=1

GENIE

В этом репозитории реализованы пространственно-временной трансформер и сэмплер MaskGIT, описанные в статье Genie: Generative Interactive Environments. Обратите внимание: эта реализация обучается только на видеопоследовательностях, без действий (хотя добавить действия несложно, добавив эмбеддинги).

# Обучение модели GENIE
python train.py --genie_config genie/configs/magvit_n32_h8_d256.json --output_dir data/genie_model --max_eval_steps 10

# Генерация кадров из обученной модели
python genie/generate.py --checkpoint_dir data/genie_model/final_checkpt

# Визуализация сгенерированных кадров
python visualize.py --token_dir data/genie_generated

# Оценка обученной модели
python genie/evaluate.py --checkpoint_dir data/genie_model/final_checkpt

Базовая линия 1X GENIE

Мы предоставляем две предобученные модели GENIE; ссылки есть в таблице лидеров.

# Генерация и визуализация
output_dir='data/genie_baseline_generated'
for i in {0..240..10}; do
    python genie/generate.py --checkpoint_dir 1x-technologies/GENIE_138M \
        --output_dir $output_dir --example_ind $i --maskgit_steps 2 --temperature 0
    python visualize.py --token_dir $output_dir
    mv $output_dir/generated_offset0.gif $output_dir/example_$i.gif
    mv $output_dir/generated_comic_offset0.png $output_dir/example_$i.png
done

# Оценка
python genie/evaluate.py --checkpoint_dir 1x-technologies/GENIE_138M --maskgit_steps 2

Описание данных

См. карточку набора данных на Hugging Face.

Обучающий набор хранится в каталоге data/train_v1.1.

Участие в соревновании

Сначала прочитайте дополнительные детали соревнования, чтобы понять правила.

Отправьте исходный код, скрипты сборки и немного информации о вашем методе на адрес challenge@1x.tech. Мы оценим отправку на нашем наборе данных и ответим по почте с результатами.

Пришлите, пожалуйста:

выбранное имя пользователя (это может быть настоящее имя или псевдоним; оно будет однозначно связано с вашей перепиской);
исходный код в виде zip-архива;
примерное число FLOP, использованных при обучении модели;
любые внешние данные, которые использовались при обучении;
ваши результаты оценки на предоставленном валидационном наборе (чтобы мы примерно понимали ваши ожидания от модели).

После ручной проверки кода мы запускаем оценку в песочнице Ubuntu 22.04 + CUDA 12.3 такими командами:

./build.sh # Установка нужных зависимостей и весов модели
./evaluate.py --val_data_dir <PATH-TO-HELD-OUT-DATA>  # Запуск вашей модели на отложенном наборе данных

Дополнительные детали соревнования

Мы предоставляем magvit2.ckpt — веса кодировщика/декодера MAGVIT2. Кодировщик позволяет токенизировать внешние данные, чтобы улучшить метрики.
В отличие от LLM, функция потерь здесь нестандартна: словарь токенов изображений поменялся в версии 1.0 (от 8 июля 2024 года). Вместо кросс-энтропии над логитами для 2^18 классов мы считаем кросс-энтропию предсказаний для 2 × 2^9 классов и суммируем её. Причина в том, что большой словарь (2^18) делает тензор логитов формы (B, 2^18, T, 16, 16) очень тяжёлым для памяти. Поэтому соревнование по сжатию рассматривает семейство моделей с факторизованным PMF вида p(x1, x2) = p(x1)p(x2). Для соревнований по сэмплированию и оценке факторизованный PMF является обязательным требованием.
Для соревнования по сжатию мы намеренно оцениваем отложенные данные при фиксированном факторизованном распределении p(x1, x2) = p(x1)p(x2). Хотя модель без факторизации, например p(x1, x2) = f(x1, x2), должна за счёт недиагональных членов ковариации x1 и x2 получать меньшие потери на тестовых данных, мы хотим поощрять решения, которые получают меньшие потери при фиксированной факторизации.
В соревновании по сжатию отправка может использовать только действия, которые произошли до текущего кадра-подсказки. Отправка может авторегрессивно предсказывать последующие действия для улучшения результата, но эти действия не будут поданы вместе с подсказкой.
Простой поиск ближайшего соседа с выбором следующего кадра из обучающего набора дал бы довольно хорошие потери и результаты сэмплирования на валидационном наборе, так как в обучающей выборке есть похожие последовательности. Однако мы явно запрещаем такое решение (закрытый тестовый набор накажет его).
Если правила соревнования нарушаются, мы не сможем выплатить награду лицам из стран, находящихся под санкциями США. Мы оставляем за собой право не выплачивать призы в соревновании.

Детали метрики

Сценарии оценки различаются в зависимости от того, насколько много реального контекста получает модель. В порядке убывания контекста эти сценарии:

Полностью авторегрессивный: модель получает заранее заданное число реальных кадров и авторегрессивно предсказывает все остальные кадры.
Временной teacher forcing: модель получает все реальные кадры до текущего кадра и авторегрессивно предсказывает все токены текущего кадра.
Полный teacher forcing: модель получает все реальные токены до текущего кадра, включая токены текущего кадра. Применим только для каузальных языковых моделей.

Например, рассмотрим предсказание последнего токена видео, который находится в правом нижнем углу 15-го кадра. В каждом сценарии модель получает такой контекст:

Полностью авторегрессивный: первые $t$ токенов 16×16 — настоящие токены первых $t$ кадров, остальные токены генерируются авторегрессивно; $0 < t < 15$ — заранее заданное число кадров-подсказок.
Временной teacher forcing: первые 15 токенов 16×16 — настоящие токены первых 15 кадров, остальные генерируются авторегрессивно.
Полный teacher forcing: все первые (16×16×16 − 1) токенов — настоящие.

Соревнование по сжатию использует сценарий «временной teacher forcing».

Таблица лидеров

Результаты оценки на data/val_v1.1.

Помогите нам улучшить соревнование!

Помимо соревнования по моделям мира, мы хотим сделать соревнование и набор данных полезными для ваших исследовательских вопросов. Нужны данные о взаимодействии с людьми? Больше примеров задач, требующих безопасности, например аккуратной переноски горячей кофейной чашки? Больше тонкой работы с инструментами? Роботы, которые работают вместе с другими роботами? Роботы, которые одеваются перед зеркалом? Представьте 1X как операционную команду, которая собирает высококачественные человекоподобные данные в самых разных сценариях.

Пришлите свои потребности в данных (и почему, по вашему мнению, эти данные важны) на challenge@1x.tech. Мы постараемся включить их в будущие публикации данных. Обсудить потребности в данных можно и в Discord.

Мы также приветствуем спонсоров, которые помогут увеличить призовой фонд.

Цитирование

Если вы используете это программное обеспечение или набор данных в работе, пожалуйста, цитируйте через кнопку «Cite this repository» на GitHub.

Журнал изменений

v1.1 — опубликованы стандарты соревнования по сжатию; из набора данных удалены видео с паузами и разрывами; более высокая обрезка изображений.
v1.0 — более эффективный токенизатор MAGVIT2 с отображением 16×16 (C=2^18) на изображения 256×256; добавлены исходные данные действий.
v0.0.1 — первоначальный выпуск соревнования с токенизатором изображений 20×20 (C=1000) и отображением на изображения 160×160.

Keyboard shortcuts

Воплощённый ИИ и робототехника