LingBot-VA

LingBot-VA — это авто регрессионная политика-модель мира для видео и действий, построенная на стеке видео-диффузии Wan2.2. В одной авторегрессионной последовательности она чередует предсказание будущих видео-латентов и действий робота («VA» = Video-Action). Интеграция с LeRobot подключает LingBot-VA к стандартным интерфейсам обучения, оценки и процессоров.

Обзор модели

LingBot-VA представляет собой двухпоточную «смесь трансформеров»: видеопоток/поток латентов (patch_embedding_mlp → blocks → proj_out) и поток действий (action_embedder → blocks → action_proj_out) используют одни и те же 30 блоков трансформера и текстовое условие.

На этапе инференса политика выполняет авторегрессионный цикл для каждого чанка: она денойзирует поток латентов видео (CFG, ~20 шагов) и поток действий (~50 шагов) с помощью двух независимых планировщиков flow-matching, поддерживая KV-кэш между чанками. Реально наблюдаемые ключевые кадры подаются обратно в KV-кэш по мере выполнения чанка (замкнутый цикл моделирования мира).

Что охватывает интеграция LeRobot

• Стандартная конфигурация policy.type=lingbot_va через LeRobot.
• Готовые к использованию контрольные точки в формате LeRobot на Hub (конвертированные из опубликованных исходных).
• Авторегрессионный инференс с двумя потоками за стандартным интерфейсом select_action (оценка в одном окружении, --eval.batch_size=1).
• Опциональное сохранение предсказанных (воображаемых) видео политики во время оценки / обучения.
• Оценка с помощью lerobot-eval на LIBERO и RoboTwin.
• Обучение / дообучение с помощью функции потерь dual-stream flow-matching (policy.forward), см. ниже.

Установка

1. Установите LeRobot, следуя Руководству по установке.
2. Установите дополнительный модуль LingBot-VA:

pip install -e ".[lingbot_va]"

Чекпоинты

Опубликованные контрольные точки из upstream-репозитория были преобразованы в формат LeRobot и загружены на Hub:

В файл model.safetensors LeRobot сохраняется только обучаемый трансформер объёмом ~5B. Замороженные компоненты VAE + UMT5 + токенизатор (~20 ГБ) загружаются из config.wan_pretrained_path во время загрузки (по умолчанию используется исходный репозиторий robbyant/*). Текстовый энкодер UMT5-XXL по умолчанию работает на CPU (config.text_encoder_device), благодаря чему трансформер 5B и VAE помещаются на один GPU объёмом 24–32 ГБ.

Оценка (LIBERO)

lerobot-eval \
    --policy.path=lerobot/lingbot_va_libero_long \
    --policy.device=cuda \
    --env.type=libero --env.task=libero_10 \
    --env.observation_height=128 --env.observation_width=128 \
    --eval.n_episodes=50 --eval.batch_size=1 \
    --output_dir=outputs/eval/lingbot_va_libero

Стриминговый вывод LingBot-VA (кэш KV + обратная связь по наблюдаемым ключевым кадрам) реализован для оценки в одной среде; используйте --eval.batch_size=1.

Оценка (RoboTwin)

Для RoboTwin 2.0 необходим стек симуляторов SAPIEN + CuRobo. Можно использовать Docker-образ бенчмарка (docker/Dockerfile.benchmark.robotwin, для которого также требуется warp-lang==1.3.1 и CuRobo, собранный с учётом вычислительной возможности GPU, указанной в TORCH_CUDA_ARCH_LIST). RoboTwin использует управление по позе конечного эффектора, поэтому запускайте с флагом --env.action_mode=ee: политика предсказывает дельты xyz+quaternion+gripper для каждой руки (латентная раскладка robotwin_tshape), которые накладываются на начальную позу eef эпизода и выполняются через IK в CuRobo.

lerobot-eval \
    --policy.path=lerobot/lingbot_va_robotwin \
    --policy.device=cuda \
    --env.type=robotwin --env.task=beat_block_hammer --env.action_mode=ee \
    --eval.n_episodes=10 --eval.batch_size=1 \
    --output_dir=outputs/eval/lingbot_va_robotwin

Сохранение предсказанных (воображаемых) видео

Установите --policy.save_predicted_video=true, чтобы дополнительно декодировать предсказанные латенты видео с помощью VAE и записать pred_episode_*.mp4 рядом с видео eval_episode_*.mp4, отрендеренными средой. Этот же флаг работает для периодической оценки во время lerobot-train.

Наблюдение за воображаемым видео в реальном времени

Установите --policy.return_predicted_video=true при запуске lerobot-rollout --display_data=true, чтобы декодировать предсказанное видео для каждого чанка и отображать воображаемые кадры камеры, к которым относится каждое выполняемое действие (prediction.images.<camera>, рядом с каждой реальной камерой). Это требует одного декодирования VAE на чанк.

Обучение / дообучение

LingBotVAPolicy.forward(batch) реализует двухпоточную функцию потерь flow-matching (latent_loss + action_loss, с весами по временным шагам и маской действий) из статьи: метод кодирует видеоклипы с камеры в латентные представления с помощью VAE, кодирует задачу с помощью UMT5, добавляет шум в оба потока, выполняет блок-каузальный проход обучения трансформера и возвращает (loss, metrics). Пресет оптимизатора — AdamW с расписанием линейного прогрева, за которым следует постоянное значение (в соответствии с исходной реализацией).

Требования:

• Блоково-каузальные маски используют flex-attention из PyTorch, поэтому для обучения создавайте политику с флагом --policy.attn_mode=flex (по умолчанию используется SDPA torch, который предназначен только для инференса).
• Полная модель DiT на 5B параметров не помещается в одну GPU объёмом 24–32 ГБ при использовании AdamW; для дообучения используйте LoRA (--peft.method_type=LORA) и/или выгрузку оптимизатора. get_optim_params возвращает только обучаемые (например, адаптерные) параметры; VAE и текстовый энкодер UMT5 остаются замороженными. Установите дополнительный пакет lerobot[peft], чтобы включить поддержку PEFT (см. руководство по обучению с PEFT).

lerobot-train \
  --policy.path=lerobot/lingbot_va_libero_long --policy.attn_mode=flex \
  --peft.method_type=LORA --peft.r=32 --peft.lora_alpha=32 \
  --peft.target_modules='transformer\.blocks\.\d+\.attn[12]\.(to_q|to_v)' \
  --dataset.repo_id=<your LeRobot-format dataset> \
  --batch_size=1 --steps=... --output_dir=outputs/train/lingbot_va

В отличие от SmolVLA / π₀, LingBot-VA не содержит встроенных целевых модулей LoRA по умолчанию, поэтому параметр --peft.target_modules необходимо передавать явно. Обучаемой является только self.transformer (двухпоточный трансформер Wan); в примере выше адаптируются проекции query/value как в блоках self-attention (attn1), так и в блоках cross-attention (attn2), что является стандартным набором целевых модулей для LoRA. При необходимости более ёмкого адаптера расширьте этот набор (например, добавьте to_k/to_out или слои ffn). Передача параметра --peft.method_type подразумевает использование PEFT, поэтому указывать --policy.use_peft=true не требуется.

Набор данных должен предоставлять клипы с камер (временное окно для каждой камеры, закодированное VAE в frame_chunk_size латентных кадров) и frame_chunk_size * action_per_frame шагов действий на элемент.

Формат данных (каналы действий & порядок камер)

LingBot-VA — это политика для позы концевой эффектора (декартова): она предсказывает позы EEF и положение захвата, а не позиции суставов. Действия размещаются в фиксированной многоэпизодной 30-мерной структуре; отобразите размеры действий вашего робота в эти каналы, а остальные заполните 0 (used_action_channel_ids выбирает каналы, которые действительно используются в данном контрольной точке):

• LIBERO использует каналы 0–6: дельта 6-DoF EEF (xyz + вращение) + захват (одна рука).
• RoboTwin использует каналы [0–6, 28, 7–13, 29]: левый EEF (xyz + кватернион) + левый захват + правый EEF + правый захват (16 измерений). Окружение преобразует эти позы в траектории суставов с помощью CuRobo IK — суставы никогда не предсказываются.

Датасеты в пространстве суставов (или с другой конвенцией EEF) необходимо переместить в эту схему перед дообучением этих контрольных точек.

Порядок камер фиксирован и чувствителен к последовательности: латенты для каждой камеры конкатенируются в пространственном направлении в порядке obs_cam_keys, поэтому физическое соответствие «камера → слот» должно совпадать с тем, что использовалось при обучении:

Первая камера — это внешний вид/вид спереди, остальные — вид с запястья.

Гиперпараметры инференса (LIBERO)

Это значения по умолчанию для LingBotVAConfig; любое из них можно переопределить с помощью --policy.<name>=....

Примечания

• Бэкенд внимания: для инференса используется бэкенд torch SDPA (всегда доступен). Бэкенды flashattn и flex являются опциональными; flex требуется только для обучения.
• Размер модели: DiT содержит ~5 млрд параметров, а замороженные VAE и UMT5 добавляют ~20 ГБ; для инференса требуется примерно 18–24 ГБ видеопамяти.

Лицензия

LingBot-VA распространяется по лицензии Apache-2.0. См. исходный репозиторий.