Мой блог

Листай вниз

Как создать стриминговый пайплайн обучения роботов на NVIDIA Cosmos3-DROID без скачивания гигабайтов

Теги :

Большие массивы данных вроде NVIDIA Cosmos3-DROID весят сотни гигабайт, что сильно затрудняет их локальную загрузку и предварительную подготовку перед обучением моделей. Я покажу, как спроектировать сквозной пайплайн обучения робототехники в режиме реального времени, обращаясь к репозиторию удаленно. В этом материале мы разберем работу со структурой LeRobotDataset v3.0, применение HTTP-доступа к нужным байтам файлов Parquet и потоковое извлечение видеофрагментов без скачивания всего архива целиком.

Для работы нам понадобится настроить окружение Python и установить ключевые библиотеки, включая клиент Hugging Face, PyArrow для работы с таблицами и PyAV для декодирования видеопотоков. Вместо того чтобы сохранять весь тяжеловесный репозиторий на диск, я настрою выборочное чтение метаданных, индексов эпизодов и параметров обучения, что позволит сразу перейти к анализу траекторий.

Введение в потоковую обработку датасетов

Интроспекция репозитория и метаданные

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

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

Байтовый доступ к Parquet-файлам

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

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

Аналитика траекторий и визуализация

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

Дополнительно я анализирую спектры скоростей суставов с помощью быстрого преобразования Фурье, оценивая распределение частот вплоть до частоты Найквиста, и строю гистограмму распределения длительности эпизодов. Это помогает отсечь аномалии и лучше понять динамику демонстраций перед подачей данных в модель.

Потоковое декодирование видео в обход полной загрузки

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

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

Нормализация и построение обучающего набора

Перед передачей данных в модель необходимо привести их к единому масштабу. Я загружаю официальные статистические показатели из файла stats.json, содержащего средние значения и стандартные отклонения для состояний и действий, а в случае их отсутствия вычисляю эти параметры эмпирически на основе загруженной выборки.

Далее я создаю датасет в стиле архитектуры ACT на базе PyTorch, который объединяет историю наблюдений, опциональные изображения с камер и нормализованные будущие действия в виде чанков фиксированной длины. Для ускорения обучения я кэширую синхронизированные визуальные наблюдения для небольшой подвыборки эпизодов.

Архитектура политики и процесс обучения

Для предсказания последовательности будущих действий я собираю составную нейросетевую модель, объединяющую полносвязный энкодер состояний и опциональный сверточный CNN-энкодер для обработки изображений. Головная часть сети выдает батч прогнозов на заданный горизонт планирования вперед.

В качестве оптимизатора я использую AdamW с планировщиком скорости обучения OneCycleLR, а также задействую смешанную точность вычислений и клиппинг градиентов для стабильности. Функция потерь Smooth L1 делает процесс клонирования поведения более устойчивым к шумам в демонстрациях оператора-человека, а в процессе эпох я отслеживаю динамику ошибок на обучающей и валидационной выборках.

Оценка модели и временной ансамбль

На финальном этапе я оцениваю обученную политику в режиме разомкнутого контура (open-loop rollout) и применяю экспоненциально взвешенный временной ансамбль для сглаживания пересекающихся предсказаний действий. Расчет среднеквадратичной ошибки для каждого сустава и коэффициента детерминации $R^2$ в сравнении с базовой линией позволяет наглядно оценить качество работы модели.

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

01.