Дообучение NVIDIA Cosmos 3 Edge для управления роботом на устройстве

NVIDIA показала управление роботом на Cosmos 3 Edge без облака

NVIDIA выпустила руководство по дообучению Cosmos 3 Edge для управления роботом прямо на бортовом компьютере. Модель с 4 млрд параметров, включая модуль рассуждения на базе NVIDIA Nemotron с 2 млрд параметров, запускается на NVIDIA Jetson Thor и формирует действия без обращения к серверной GPU. В замкнутом моделировании RoboLab дообученная политика достигла точности 22,9%. Все шаги можно воспроизвести в открытом репозитории cosmos-framework, а готовая контрольная точка опубликована на Hugging Face.

Что такое Cosmos 3 Edge

Роботам нужны политики управления, которые подстраиваются под сенсоры, окружение и задачу и при этом работают на бортовом вычислительном оборудовании. Модели мира дают основу для обучения взаимодействию с физическими объектами, но их размер часто мешает запуску на самом роботе. Cosmos 3 Edge рассчитана на такой сценарий.

Cosmos 3 Edge — универсальная модель на 4 млрд параметров из семейства Cosmos 3. В неё входит модуль рассуждения на базе NVIDIA Nemotron с 2 млрд параметров. Модель обучали на тех же данных о физическом мире, что и NVIDIA Cosmos 3 Nano и NVIDIA Cosmos 3 Super, поэтому она изначально учитывает, как объекты двигаются и взаимодействуют. При этом Cosmos 3 Edge достаточно мала для запуска на NVIDIA Jetson Thor.

Руководство приводит полный путь до дообученной политики манипуляции, которая работает на Jetson Thor и проверяется в моделировании с замкнутым циклом. Каждый шаг воспроизводится из открытого репозитория cosmos-framework, а готовая контрольная точка доступна на Hugging Face.

Почему помогают модели мира

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

Cosmos умеет генерировать действия, опираясь на понимание физики и прогнозирование. Благодаря этим знаниям политика управления роботом начинает обучение с уже сформированным представлением о физических связях, вместо того чтобы осваивать их только по демонстрациям конкретной задачи.

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

Дообучение Cosmos 3 Edge решает обе проблемы. Политика помещается в память Jetson Thor, поэтому инференс выполняется непосредственно на роботе и не передаётся на GPU в центре обработки данных.

На NVIDIA Jetson AGX Thor T5000 политика DROID работает в реальном времени. При разрешении 640×540 и частоте 15 Гц она формирует каждый фрагмент действий примерно за 1,53 секунды, а один такой фрагмент описывает около 2,13 секунды движения робота. Следующий фрагмент готов до завершения текущего, поэтому манипулятор движется непрерывно.

| Параметр | Значение |

|---|---:|

| Время формирования фрагмента действий | около 1,53 с |

| Длительность движения во фрагменте | около 2,13 с |

| Разрешение и частота | 640×540, 15 Гц |

Политика поддерживает непрерывную потоковую работу на устройстве: после каждого цикла инференса она строит новый план. При этом перепланирование происходит после каждого цикла вывода, а не после каждого наблюдения.

В задачах RoboLab с замкнутым циклом дообученная политика достигла 22,9% успешных выполнений. Это показывает, что базовая модель мира на 4 млрд параметров может работать основой политики управления в реальном времени непосредственно на роботе. Модель помещается в Jetson AGX Thor и полностью выполняется на нём.

Источник: developer.nvidia.com

Данные для обучения

Политика из руководства обучается на наборе данных nvidia/Cosmos3-DROID. В нём 76 тысяч успешных траекторий, записанных оператором вручную: примерно 350 часов данных для 86 задач и 564 сцен. Съёмка проводилась с манипулятором Franka Panda и захватом Robotiq.

Набор упакован в формате LeRobotDataset v3.0 с разрешением 640 × 360. Его подготовка проходит в три этапа. Данные нужно привести к формату LeRobot Dataset v3: покадровое видео с камер, состояния суставов, состояние захвата, действия и инструкция для задачи.

Для конфигурации, похожей на DROID с манипулятором Franka, основным изменением будет путь к набору данных. Если используется другой тип робота, для него потребуется отдельная конфигурация эксперимента. В ней задаются пространство действий, размерность, схема камер и параметры нормализации.

Cosmos 3 поддерживает несколько конфигураций роботов, включая двухрукого Franka, UR, WidowX 250 и LeRobot SO101. Полный список приведён в карточке модели Cosmos 3 Edge.

Сколько нужно вычислений

Речь идёт о дообучении базовой модели, а не о файн-тюнинге на одной GPU. Проверенный запуск использовал 64 узла с четырьмя GB200 на каждом, 60 тысяч итераций и около 68 часов работы — примерно 17,4 тысячи часов на GB200.

На верхнем уровне дообучение разделено на четыре шага. Исходный запускатель регистрирует рецепт для Cosmos 3 Nano. Перед запуском обучения Cosmos 3 Edge в нём нужно изменить три параметра. Остальные настройки — набор данных, пространство действий, фильтр отбора и расписание обучения — остаются прежними.

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

Как работает политика

Политика запускается через сервер WebSocket, который использует протокол OpenPI. Этот же протокол применяется в экосистеме политик DROID. Клиент отправляет словарь наблюдений, а сервер возвращает фрагмент действий.

Для Cosmos 3 Edge сервер работает непосредственно на Jetson Thor. Веса занимают около 9 ГБ в формате BF16 и помещаются в бортовую память, поэтому сервер политики и клиент управления могут работать на самом роботе без GPU в центре обработки данных.

Поскольку сервер запускается на Thor с адресом `host="localhost"`, запросы не покидают робота. Инференс выполняется в реальном времени на устройстве — именно это обеспечивает Cosmos Edge.

Политика DROID учитывает состояние робота. Позиции суставов `joint_position` и захвата `gripper_position` — настоящие входы модели, а не формальные поля. Если манипулятор не подключён, вместо них используются нулевые значения. Такой вызов нужен только для проверки того, что сервер возвращает корректный фрагмент действий.

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

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

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

Проверка в симуляции

Для проверки политики необязательно использовать физического робота. NVIDIA рекомендует сначала оценивать её в моделировании, чтобы неожиданные действия не повлияли на разработчиков или оборудование.

RoboLab — открытый бенчмарк на базе Isaac Lab-Arena, связанный с рейтингом RoboLab. Его клиент подключается к тому же серверу политики, исполняет каждый фрагмент действий в физическом моделировании и передаёт обратно отрисованные наблюдения. В результате получается замкнутый цикл для 120 задач манипуляции, заданных на естественном языке.

Параметр `--task` можно заменить на любую из 120 задач, чтобы получить более полную картину поведения политики. Каждый запуск создаёт видео из окна моделирования и с камеры робота, а также журнал успешности. Так можно отдельно изучить результат и траекторию, которая к нему привела.

В полном тестировании RoboLab дообученная политика Cosmos 3 Edge достигла 22,9% успешных выполнений по всему набору задач. Для сравнения, Cosmos 3 Nano показала 36,8%. При этом Edge требует меньше вычислений для инференса и рассчитана на работу в реальном времени полностью на роботе.

Для NVIDIA Isaac Sim 5.x требуется драйвер NVIDIA RTX Server версии 580 или новее. Перед сборкой нужно проверить поддерживаемую версию в документации RoboLab и Isaac Sim.

Другие сценарии дообучения

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

Компания Aigen дообучила Cosmos для генерации разнообразных вариантов сельскохозяйственных культур и сорняков. Благодаря этому автономная система прополки показала высокую эффективность, используя для обучения всего 1% данных из реального мира.

Открытые веса Cosmos и программная платформа под лицензией OpenMDW1.1 позволяют дообучать специализированные модели для физических ИИ-систем.