π0: A Vision-Language-Action Flow Model for General Robot Control

Сегодня разберём статью о важной для self-driving теме из смежной области — как делать large-scale pre-training, если обучающих данных исчезающе мало.

У ребят, которые создают автономных роботов, обучающих датасетов ещё меньше, чем у нас. Авторы из команды Physical Intelligence предлагают решение: прикручивают flow-matching-декодер для экшнов к VLM, натренированной на internet-scale-датасете.

Архитектурно эта система устроена следующим образом:

1. Изображения с камер и текст задания попадают в префикс, проходят через VLM (PaliGemma 3B на базе Gemma 2) — получается KV-кеш.

2. В получившуюся последовательность после префикса добавляют эмбеддинг текущего состояния — вектора Q из 18 float-значений.

3. Траекторию на H=50 шагов, предсказывают с помощью flow-matching-декодера на 300M параметров, используя KV-кеш из 1 и 2.

Если у робота меньше трёх камер, вместо недостающих изображений указывают нули. Аналогично, если размерность стейта меньше 18, дополняют вектор нулями до 18 значений.

При этом у аттеншна блочно-каузальная маска:

🔴 префикс учитывает только собственные данные,
🔴 вектор состояния Q ориентируется и на себя, и на префикс,
🔴 траектория опирается и на собственные значения, и на Q, и на префикс,
🔴 внутри каждого блока маска полная: все смотрят на всех.

Так как префикс лежит в KV-кеше, важно, чтобы у декодера траектории были те же количество слоёв и размерность слоя. Но можно уменьшить размерность в MLP: именно так удаётся получить 300M параметров в декодере траектории при 3B в VLM.

Сначала модель претрейнят на составном датасете:

🔴 9,1% разнообразных опенсорс-данных — Oxe Magic Soup (2 Гц + 10 Гц),
🔴 остальные 90% (903M таймстепов на 68 заданий для single-/dual-arm роботов) — π0-dataset,
🔴 чтобы избежать дисбаланса, каждую task-robot-пару взвешивают с весом n⁰˒⁴³.

Потом обучают все части модели: image-backbone, VLM, Action-Expert. А во время пост-трейнинга каждый Action-Expert файнтюнится на своём task-specific-датасете.

В инференсе 10 шагов интегрирования (денойзинга). Он работает 73 мс на RTX 4090 (86 мс с задержкой сети). Для 20-герцовых роботов инференс делают каждые 0,8 секунды (после 16 шагов), для 50 Гц роботов — каждые 0,5 секунды (после 25 шагов). Для траекторий пробовали добавлять temporal ensembling, но результаты получились хуже чем open-loop.

Авторы утверждают, что на момент публикации были в топе всех популярных бенчмарков.

Разбор подготовил ❣️ Сергей Репьевский
404 driver not found