VLA-JEPA

Это порт VLA-JEPA для LeRobot — модели Vision-Language-Action, которая объединяет языковую основу Qwen3-VL, самообучаемую модель мира для видео (V-JEPA2) и головку действий flow-matching DiT.

Обзор архитектуры

VLA-JEPA состоит из трёх основных компонентов:

Поток данных

Обучение:

1. Видеоклип из num_video_frames кадров кодируется V-JEPA2 в токены патчей для каждого кадра.
2. Базовая модель Qwen3-VL обрабатывает изображения с нескольких ракурсов + инструкцию к задаче и формирует последовательность контекстных токенов, включающую специальные токены действий (для настройки мировой модели) и токены воплощения.
3. Голова действий получает эти токены контекста в качестве ключей и значений для кросс-внимания и предсказывает очищенный от шума блок действий с помощью flow matching.
4. Предиктор мировой модели использует токены действий, извлечённые из Qwen, для предсказания будущих эмбеддингов кадров V-JEPA2; к потере действий добавляется регрессионная потеря для этих предсказаний.

Вывод: используются только Qwen и головка действий. Модель мира на этапе вывода не требуется.

Детали action head

Доступные пресеты через action_model_type:

Пресет задаёт только геометрию внимания, и каждую запись можно переопределить с помощью action_num_heads / action_attention_head_dim. Из неё следуют две ширины:

• внутренняя ширина DiT равна heads x head_dim (768 для DiT-B), она вычисляется, а не настраивается;
• Ширина выхода DiT, а также ширина MLP-блоков декодера действий и энкодера состояния, задается параметром action_hidden_size (по умолчанию 1024).

Так DiT-B использует трансформер шириной 768, который проецируется в 1024. Эти два значения независимы.

Детали модели мира

Видео-предиктор представляет собой трансформер в стиле ViT (ActionConditionedVideoPredictor), который принимает:

• Токены кадров: эмбеддинги патчей V-JEPA2, спроецированные в predictor_embed_dim
• Токены действий: эмбеддинги токенов действий Qwen, спроецированные в predictor_embed_dim

Он использует блочно-каузальное внимание, благодаря которому каждый временной шаг может обращать внимание на все предыдущие шаги. Входное значение embed_dim предсказателя равно num_views × video_encoder_hidden_size (например, 2 вида × 1024 = 2048 для предобученных контрольных точек).

Предобученные контрольные точки

Внутри организации LeRobot здесь доступны три контрольные точки: lerobot/VLA-JEPA, сконвертированные из ginwind/VLA-JEPA:

Все контрольные точки используют Qwen/Qwen3-VL-2B-Instruct в качестве языковой основы.

Конфигурация

Ключевые параметры в VLAJEPAConfig:

pre_snap_gripper_action и binarize_gripper_action — это порт цикла оценки LIBERO из starVLA, и они корректны только для конвенции действий LIBERO. pre_snap записывает {0, 1} в пространство normalized, денормализатор отображает эти значения на середину и максимум диапазона, а затем binarize сравнивает это physical значение с gripper_threshold (0.5). Если захват измеряется в градусах, мм или в диапазоне [0, 100], оба значения окажутся выше порога, и управляющий сигнал захвата станет константой. По этой причине по умолчанию они установлены в False; включайте их только для конфигураций в стиле LIBERO, и при этом задавайте gripper_threshold в собственных единицах захвата. Фабрика процессоров выводит предупреждение, если статистика датасета показывает, что диапазон не может работать.

Обучение

Количество шагов обучения может варьироваться в зависимости от размера набора данных и вычислительного бюджета. В исходной статье модель предварительно обучалась на 50k шагов на ssv2 + droid совместно, затем на дополнительные 30k шагов на LIBERO, однако при дообучении из предоставленных предварительно обученных контрольных точек хорошее качество может быть достигнуто за меньшее количество шагов.

Полная обучение с нуля

lerobot-train \
  --policy.type=vla_jepa \
  --policy.repo_id=your_org/your_repo \
  --dataset.repo_id=your_org/your_dataset

Доработка из предварительно обученной контрольной точки

lerobot-train \
  --policy.path=lerobot/VLA-JEPA-Pretrain \
  --policy.repo_id=your_org/your_repo \
  --dataset.repo_id=your_org/your_dataset

Если вы хотите заморозить базовую модель Qwen и обучать только головку действий, установите policy.freeze_qwen=True:

lerobot-train \
  --policy.path=lerobot/VLA-JEPA-Pretrain \
  --policy.repo_id=your_org/your_repo \
  --policy.freeze_qwen=true \
  --dataset.repo_id=your_org/your_dataset

Доподнастройка на другом воплощении

Если целевой робот имеет другую размерность действий или состояний, чем предобученная контрольная точка, слои проекции входа/выхода в action head будут иметь несовпадающие формы и не смогут быть загружены напрямую. reinit_modules позволяет указать префиксы ключей, для которых допустимо несовпадение форм — эти слои будут случайным образом переинициализированы, тогда как все остальные веса будут переиспользованы из контрольной точки. Любое несовпадение форм вне указанных префиксов приводит к ошибке.

Слои, зависящие от action_dim и state_dim:

lerobot-train \
  --policy.path=lerobot/VLA-JEPA-Pretrain \
  --policy.repo_id=your_org/your_repo \
  --policy.freeze_qwen=true \
  --policy.reinit_modules='["model.action_model.action_encoder", "model.action_model.action_decoder", "model.action_model.state_encoder"]' \
  --dataset.repo_id=your_org/your_dataset

Если у вашего робота нет проприоцептивного состояния, опустите model.action_model.state_encoder из списка.

Воспроизведение результатов LIBERO

Обучение на LIBERO: запускает обучение с контрольной точки Pretrain, обучает модель в течение 30k шагов на наборе данных LIBERO. В исходной статье упоминается обучение на 8 GPU с размером батча 32, что означает глобальный размер батча 256.

lerobot-train \
  --policy.path=lerobot/VLA-JEPA-Pretrain \
  --policy.repo_id=your_org/your_repo \
  --dataset.repo_id=HuggingFaceVLA/libero \
  --steps=30000

Оценка предобученного контрольного файла LIBERO-10:

lerobot-eval \
  --policy.path=lerobot/VLA-JEPA-LIBERO \
  --env.type=libero \
  --env.task=libero_spatial,libero_object,libero_goal,libero_10 \
  --eval.n_episodes=10 \
  --eval.batch_size=5

Для оценки только подмножества задач:

lerobot-eval \
  --policy.path=lerobot/VLA-JEPA-LIBERO \
  --env.type=libero \
  --env.task=libero_10 \
  --env.task_ids='[0,1,2]' \
  --eval.n_episodes=10 \
  --eval.batch_size=5

Ожидаемые результаты:

Допобучение на наборах данных с разным количеством камер

Предобученный предиктор мировой модели обучался с embed_dim = world_model_num_views × 1024, то есть с двумя камерными видами.

Количество кадров в этом представлении ранее читалось из jepa_tubelet_size, который также задавал размер временного tubelet в кодере JEPA. Теперь для этого используется поле world_model_num_views; если оставить его равным None, значение будет взято из jepa_tubelet_size, чтобы опубликованные контрольные точки продолжали загружаться без изменений.

Поведение по умолчанию — просмотр отступов / обрезки (действия не требуются)

При дообучении из VLA-JEPA-Pretrain модель автоматически настраивает количество представлений, подаваемых в мировую модель, в соответствии с world_model_num_views:

• Наборы данных с одним видом (например, BridgeV2): латентное представление одного вида дублируется для создания входных данных модели мира с двумя видами, что сохраняет самонастраивающийся сигнал JEPA без какого-либо рассогласования весов.
• >Датасеты с более чем 2 видами (например, DROID с 3 видами): все виды передаются в Qwen backbone (для более богатого контекста), но для world model используются только первые world_model_num_views виды (один вид с запястья и один вид от третьего лица, в соответствии с заданным порядком видов).

Вариант 1 — отключить world model

Установите enable_world_model=False, чтобы полностью пропустить потерю JEPA. Загружаются и обучаются только базовая модель Qwen и головная часть для действий. Этого достаточно для хорошей производительности действий.

lerobot-train \
  --policy.path=lerobot/VLA-JEPA-Pretrain \
  --policy.enable_world_model=false \
  --policy.repo_id=your_org/your_repo \
  --dataset.repo_id=your_org/single_camera_dataset

Вариант 2 — Переинициализация проекции входа предсказателя

Если вы хотите изменить world_model_num_views на значение, отличное от 2, загрузите контрольную точку с strict=False и заново инициализируйте model.video_predictor.predictor_embed для нового embed_dim. Все остальные веса блоков предсказателя (внимание, MLP, нормализация, выходная проекция) не зависят от количества камер и могут быть переиспользованы из предобученной контрольной точки.

Цитирование

@misc{sun2026vlajepaenhancingvisionlanguageactionmodel,
  title         = {VLA-JEPA: Enhancing Vision-Language-Action Model with Latent World Model},
  author        = {Jingwen Sun and Wenyao Zhang and Zekun Qi and Shaojie Ren and Zezhi Liu and Hanxin Zhu and Guangzhong Sun and Xin Jin and Zhibo Chen},
  year          = {2026},
  eprint        = {2602.10098},
  archivePrefix = {arXiv},
  primaryClass  = {cs.RO},
  url           = {https://arxiv.org/abs/2602.10098},
}

Лицензия

Веса распространяются на условиях лицензии исходного репозитория ginwind/VLA-JEPA (Apache 2.0 License). Код интеграции LeRobot распространяется по лицензии Apache 2.0 License.