Об этом проекте

LeWorldModel (LeWM) — это официальный репозиторий кода для статьи «LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels» авторов Lucas Maes, Quentin Le Lidec, Damien Scieur, Yann LeCun и Randall Balestriero. Он реализует архитектуру Joint-Embedding Predictive Architecture (JEPA), которая, согласно аннотации, стабильно обучается сквозным образом на сырых пикселях, используя только два члена функции потерь: потерю предсказания следующего эмбеддинга и регуляризатор, обеспечивающий гауссово распределение латентных эмбеддингов. Авторы утверждают, что это сокращает количество настраиваемых гиперпараметров потерь с шести до одного по сравнению с единственной существующей сквозной альтернативой, и что модель имеет примерно 15 миллионов параметров и обучается на одной GPU за несколько часов. Репозиторий основан на двух внешних проектах: stable-worldmodel для управления окружением, планирования и оценки, и stable-pretraining для обучения. Заявленная цель — свести репозиторий к его основной части, архитектуре модели и цели обучения, реализованным в jepa.py как модуль PyTorch. Установка использует uv с Python 3.10 и устанавливает stable-worldmodel с дополнительными пакетами train и env. Наборы данных распространяются в формате HDF5 через Hugging Face и распаковываются с помощью tar --zstd; извлеченные файлы .h5 помещаются в $STABLEWM_HOME (по умолчанию ~/.stable-wm/), который можно переопределить переменной окружения. Имена наборов данных указываются без расширения .h5 в конфигурационных файлах Hydra в config/train/. Обучение запускается командой python train.py data=pusht после установки сущности и проекта WandB в config/train/lewm.yaml. Чекпоинты сохраняются в $STABLEWM_HOME по завершении. Конфигурации оценки находятся в config/eval/, и поле policy должно указывать на путь к чекпоинту относительно $STABLEWM_HOME без суффикса _object.ckpt. Предобученные чекпоинты LeWM зеркалируются на Hugging Face Hub для окружений pusht, cube, tworooms и reacher, вместе с наборами данных в той же коллекции. Более широкий набор базовых чекпоинтов (PLDM, LeJEPA, IVL, IQL, GCBC, DINO-WM, DINO-WM-noprop) доступен через Google Drive, с таблицей, указывающей, какие методы покрывают какие окружения. Чекпоинты представлены в двух формах: _object.ckpt — сериализованный объект Python, используемый eval.py и API stable_worldmodel, и _weight.ckpt — словарь состояний только с весами. Загрузка выполняется через swm.policy.AutoCostModel с именем запуска относительно $STABLEWM_HOME; возвращаемый модуль находится в режиме оценки и предоставляет .state_dict(). Для зеркал Hugging Face в README приведен скрипт преобразования, который восстанавливает модель JEPA из config.json и weights.pt и сохраняет объектный чекпоинт. README также содержит ссылки на статью, сайт проекта и коллекцию Hugging Face, а также приглашает к обсуждению и сотрудничеству по электронной почте. Сообщается, что LeWM планирует до 48 раз быстрее, чем мировые модели на основе фундаментальных моделей, оставаясь конкурентоспособным в разнообразных задачах управления 2D и 3D, что его латентное пространство кодирует физическую структуру, как показано зондированием физических величин, и что оценка неожиданности обнаруживает физически неправдоподобные события. Это утверждения из аннотации авторов, не проверенные независимо.