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

NVIDIA показала постобучение Cosmos 3 Edge для роботов на Jetson Thor

Что умеет Cosmos 3 Edge

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

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

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

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

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

Запуск политики на роботе

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

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

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

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

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

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

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

Проверка в RoboLab

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

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

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

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

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

Синтетические данные для роботов

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

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

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