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

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

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

Что такое Cosmos 3 Edge

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

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

Размер Cosmos 3 Edge позволяет запускать её на устройстве NVIDIA Jetson Thor. В руководстве NVIDIA показано, как получить дообученную политику для манипуляций, запустить её на 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 000 успешных траекторий телеоперации — около 350 часов данных, собранных для 86 задач и 564 сцен. В экспериментах использовались манипулятор Franka Panda и захват Robotiq.

Набор данных подготовлен в формате LeRobotDataset v3.0 с разрешением 640 × 360.

Подготовка данных начинается с преобразования в формат LeRobot Dataset v3. В нём для каждого кадра хранятся видео с камер, состояния суставов, состояние захвата, действия и инструкция к задаче.

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

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

Дообучение модели

Речь идёт о дообучении фундаментальной модели, а не о файн-тюнинге на одной GPU. Проверенный запуск использует 64 узла с четырьмя GB200 на каждом, 60 000 итераций и занимает около 68 часов — примерно 17 400 часов работы 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` являются настоящими входами модели, а не формальными полями. Если манипулятор не подключён, в качестве заполнителя можно честно передать нули: такой вызов лишь проверяет, что сервер возвращает корректный фрагмент действий.

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

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

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

Проверка в RoboLab

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

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

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

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

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

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

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

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

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

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