NVIDIA Cosmos 3 Edge: руководство по дообучению модели

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

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

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

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

В руководстве показано, как получить пост-обученную политику для манипуляций на базе Cosmos 3 Edge, запустить её на Jetson Thor и проверить в симуляции с замкнутым контуром.

Все шаги можно воспроизвести из открытого репозитория cosmos-framework. Готовая контрольная точка опубликована на Hugging Face.

Что дают модели мира

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

Cosmos умеет генерировать действия на основе понимания и прогнозирования физических процессов. Эти знания становятся отправной точкой для обучения роботизированной политики: ей не приходится заново осваивать все физические зависимости только по примерам конкретной задачи.

При запуске таких моделей на физическом роботе возникают два практических ограничения. Пост-обучение Cosmos 3 Edge устраняет их: итоговая политика помещается в память Jetson Thor, а инференс выполняется непосредственно на роботе, без передачи вычислений на GPU в дата-центре.

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

Следующий фрагмент готов до завершения текущего, поэтому манипулятор движется непрерывно, а GPU в дата-центре не участвует. Политика поддерживает потоковую работу на устройстве: после каждого цикла инференса она создаёт новый фрагмент действий и заново планирует движение. При этом перепланирование происходит после каждого цикла инференса, а не после каждого наблюдения.

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

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

Политика из руководства обучается на наборе данных 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 узла по 4× GB200, 60 тыс. итераций и занимает около 68 часов — примерно 17,4 тыс. часов работы GB200. Вычислительные ресурсы нужно планировать с учётом этих требований.

Верхнеуровнево пост-обучение разделено на четыре шага. Исходный загрузчик регистрирует сценарий для Cosmos 3 Nano. Чтобы обучать Cosmos 3 Edge, перед запуском нужно внести три изменения.

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

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

Запуск политики

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

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

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

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

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

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

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

Проверка в RoboLab

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

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

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

В замкнутой оценке RoboLab пост-обученная политика Edge достигла 22,9% успеха на всём наборе задач. При этом она требует меньших вычислительных затрат на инференс, чем более крупные варианты Cosmos3: результат Nano составляет 36,8%. Такой обмен характеристиками заложен в концепцию Edge — автономная работа в реальном времени полностью на роботе при сопоставимой успешности.

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

Другие применения

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

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

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

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