04. Перенос SmolVLA на ROCm и взвешивание выборки

Эта задача посвящена SmolVLA. По сравнению с ACT SmolVLA сильнее опирается на языковые условия и визуальные признаки, поэтому она особенно подходит для проверки на задаче с красной и синей кружкой: понимает ли модель инструкцию, а не просто запоминает распределение данных.

Практический Notebook к этой задаче: 04_smolvla_weighted_sampling.ipynb.

Фиксированная оценка красной и синей кружки

Не выбирайте задачи случайно. Используйте один и тот же набор seed и для каждого принудительно задайте обе инструкции:

Place the red mug on the plate.
Place the blue mug on the plate.

Затем сравните:

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

Рисунок 1: сравнение physical_success у SmolVLA при фиксированных инструкциях для красной и синей кружки. Важно не добиться максимальной успешности для одного цвета, а проверить, не ухудшает ли рост синей кружки результат для красной.

Почему нельзя просто копировать данные

Самый очевидный способ — продублировать эпизоды с синей кружкой. Но копирование меняет статистику набора данных и распределение эпизодов: модель может переобучиться на синюю кружку и ухудшить работу с красной.

Более мягкий вариант — взвешенная выборка по кадрам или эпизодам. Например, WeightedRandomSampler чаще выбирает кадры с синей кружкой, но не меняет исходные файлы parquet.

Сравните два метода:

Что на самом деле изменено в этом раунде SmolVLA

Главная проблема SmolVLA была не в том, что она «совсем не умеет захватывать кружку», а в явном смещении к красной кружке. При фиксированных seed и принудительных инструкциях для красной и синей кружки baseline даёт 8/10 для красной и 0/10 для синей. Это значит, что модель умеет захватывать, но языковые условия и распределение цветов не выучены сбалансированно.

Сначала мы проверили самый очевидный путь — копирование эпизодов с синей кружкой в пропорции 1.5x, 2x и 3x. Синяя кружка действительно улучшилась, но красная заметно просела:

Поэтому основной линией выбран WeightedRandomSampler. Он не копирует эпизоды parquet, не меняет исходный набор данных LeRobot и не изменяет статистику нормализации, а лишь повышает вероятность выборки кадров с синей кружкой во время обучения. Используются две переменные окружения:

export LEROBOT_TASK_WEIGHT_SUBSTRING=blue
export LEROBOT_TASK_WEIGHT=2.0

Преимущество такого подхода в простом откате: набор данных demo_data_language остаётся прежним, и достаточно убрать веса sampler, чтобы вернуться к обычной выборке. Кроме того, статистика набора не переписывается, как при копировании эпизодов.

Лучшим checkpoint в этом раунде оказался не финальный шаг запуска, а промежуточный 000500:

ckpt/smolvla_weighted_blue2_from5000_1000_20260629_143123/checkpoints/000500/pretrained_model

Loss у 001000 продолжал снижаться, но баланс между красной и синей кружкой был хуже, чем у 000500. Поэтому в теме 000500 назначен защищённым checkpoint. Любой новый вариант SmolVLA должен превзойти его на том же наборе фиксированных seed и при том же определении physical_success, иначе замены не происходит.

Итоговые результаты приводятся на двух уровнях:

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

Запись обучения на ROCm

При обучении SmolVLA на ROCm фиксируйте:

• начальный checkpoint;
• корневой каталог данных;
• число шагов обучения;
• размер батча;
• подстроку задачи для веса;
• вес задачи;
• загрузку GPU;
• диапазон температур;
• использование видеопамяти;
• шаги, на которых сохранялись checkpoint.

Не ограничивайтесь итоговым checkpoint. Промежуточный checkpoint может оказаться лучше: например, на 500 шаге баланс может быть лучше, чем на 1000.

Видео и повторный запуск неудачных seed

Если некоторые seed падают в пакетной оценке, запустите их ещё три раза, чтобы понять, стабильна ли неудача:

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

Пример результатов воспроизведения

В примере учебника лучше всего показал себя checkpoint SmolVLA weighted500, восстановленный в сквозном Notebook. На небольшой выборке он даёт красную 8/10 и синюю 10/10. В расширенной защищённой оценке на 30 seed результат: красная 27/30, синяя 30/30, всего 57/60. Это показывает, что основная линия SmolVLA на устройстве ROCm воспроизвела достаточно стабильное захватывание кружки по фиксированной инструкции.

Здесь этот checkpoint назван «защищённым базовым»: все новые эксперименты должны сравниваться с ним. Новый checkpoint заменит его только при превышении результата на том же наборе seed и при том же определении physical_success. Так можно не принять за улучшение модель, у которой просто ниже loss на одном запуске.

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

Рисунок 3: успешная последовательность для синей кружки после взвешенной выборки. Её удобно использовать в учебнике, чтобы объяснить, почему оценку нужно разделять по цвету кружки в инструкции.

Checkpoint

После выполнения задачи сохраните:

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