Обучение с подкреплением внутри моделей мира

Источник: AMD•

Обучение с подкреплением внутри моделей мира

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

Аннотация

Когда модель мира (WM) служит средой для обучения с подкреплением, политика движется в сторону того, что модель мира считает успехом, поэтому любая систематическая склонность модели мира воспринимать неудачи как успехи усиливается. В задаче RoboTwin adjust_bottle, где в качестве среды используется видеомодель с условными действиями BWM, а в качестве политики — π0.5, данное исследование задается вопросом, при каких условиях прирост от ОП внутри модели мира является надежным. Базовая модель BWM воспринимает 82,5% неудач симулятора как успехи — разрыв, который пиксельные метрики едва фиксируют. Дообучение модели мира на собственных траекториях неудач политики снижает количество ложных успехов до 26%; контрольная выборка, состоящая только из успехов, вообще не дает улучшений. Внутри дообученной модели мира рецепт PPO, устраняющий три дефекта (бутстрэппинг при усечении, место подключения критика и скорость обучения ценности), значительно превосходит метод SFT на всех пяти сидах обучения при оценке обратно в симуляторе: средний прирост составил +9,7 процентных пункта, а стандартное отклонение — 4,0. У исходного рецепта стандартное отклонение составляло 14,3, а худший сид достигал −19,0. Повторный запуск в замкнутом контуре показывает, что деградировавший сид при использовании исходного рецепта эксплуатировал ложные успехи модели мира, в то время как исправленный рецепт остается на уровне несоответствия SFT. Абляция показывает, что оставшаяся ошибка кроется в неудачах, исход которых зависит от деталей действия из общего исходного состояния: модель мира считает 29–40% таких случаев успешными даже после обучения на них, добавление данных, собранных on-policy, это не исправляет, и именно из таких состояний поступают градиенты обучения с подкреплением.

Введение

В модели мира, используемой в качестве среды ОП, политика выдает фрагмент действий, модель мира рендерит соответствующее видео, модель вознаграждения считывает видео и оценивает успех, после чего политика обновляется на основе этого вознаграждения. Модель мира, использующаяся в качестве оценщика политики, может допускать ошибки до тех пор, пока они примерно одинаково затрагивают все варианты политики, поскольку ранжирование сохраняется. Среда обучения с подкреплением не может себе этого позволить: политика активно ищет действия, приносящие вознаграждение. Если модель мира систематически ошибается в определенном классе действий, например, считает захват успешным всякий раз, когда схватка закрывается, политика будет все чаще совершать действия этого класса. Следовательно, к модели мира в качестве среды ОП предъявляются более строгие требования, чем в качестве оценщика: исход каждого запуска должен определяться правильно.

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

Данное исследование дает ответы на оба вопроса в рамках одной задачи. Основные результаты:

  • Измерение точности исходов. Одна и та же последовательность действий управляет как симулятором, так и моделью мира, а одна замороженная модель вознаграждения оценивает оба набора кадров, выдавая показатель частоты ложных успехов модели мира. Затем повторный запуск в замкнутом контуре позволяет оценить, использует ли обновленная политика ошибки модели мира в своих интересах.
  • Дообучение на данных о неудачах. Дообучение модели мира на собственных траекториях неудач политики снижает частоту ложных успехов с 82,5% до 26%, причем этот эффект достигается именно за счет данных о неудачах.
  • Воспроизводимый PPO внутри модели мира. После устранения трех дефектов рецепта все пять сидов обучения значительно превосходят метод SFT, а уровень несоответствия находится на уровне SFT.
  • Локализация оставшейся ошибки. Абляция показывает, что оценки исходов модели мира при разветвлении действий из общего начального состояния остаются ненадежными и не реагируют на добавление данных on-policy. Эту проблему необходимо решить до того, как ОП внутри модели мира сможет масштабироваться дальше.

2 Связанные работы

Обучение VLA внутри моделей мира. В проектах WoVR, WMPO и World-Gymnast выполняется пост-обучение VLA с помощью ОП внутри видеомоделей с условными действиями, и сообщается об однократном приросте производительности от +5 до +29 процентных пунктов. Настоящая установка наиболее близка к WoVR, а многокадровое окно обучения следует ее рецепту. Разница заключается в валидации: каждый рецепт здесь обучается с пятью сидами, а повторный запуск в замкнутом контуре проверяет, не использует ли политика ошибки модели мира.

Смещение в сторону оптимизма и данные о неудачах. В MiraBench задокументирована тенденция видеомоделей мира принимать неудачи за успехи; ни одна из 12 оцениваемых моделей не решает эту проблему. В BWM, FACT, 1X и параллельном исследовании CureWM в обучение моделей мира добавляются данные о неудачах или контрфактические данные. Данные о неудачах в данной работе представляют собой собственные естественные неудачи политики, а влияние исправления прослеживается вплоть до последующего прироста ОП и показателей несоответствия. Контрфактические данные с одинаковым началом из CureWM нацелены именно на ту ошибку разветвления действий, которая остается в разделе 5.5.

PPO для VLA на основе сопоставления потоков (flow-matching). В πRL запускается онлайн-ОП на π0 / π0.5 и рекомендуется подключать критик к выходу VLM. Данное исследование следует этому совету. Бутстрэппинг усечения на пределе эпизода исправляется отдельно в среде модели мира RLinf; рецепт приведен в разделе 3.3.

3 Метод

3.1 Установка

Компонент

Выбор

Задача

RoboTwin 2.0 adjust_bottle, aloha-agilex с двумя манипуляторами, одна головная камера

Политика

Контрольная точка π0.5 SFT из WorldArena 2.0 Track 2

BWM step-12000, видеомодель с условными действиями на базе Wan2.2-TI2V-5B

Награда

Классификатор кадров, обученный на метках успеха симулятора, а затем замороженный; он оценивает кадры 51–200, находит пик и выдает единицу на шаге 200, если пик превышает пороговое значение

PPO в RLinf, 96 эпизодов на шаг, 10 обновлений политики

Оценка

Обратно в симуляторе RoboTwin, на 144 начальных состояниях, зафиксированных до обучения, три сида оценки на состояние, сопоставленных попарно с SFT для каждого состояния, с 90%-ными интервалами на основе бутстрэппинга для каждого состояния

Повторная оценка одной политики с другим случайным потоком оценки изменяет показатель ее успеха на 1,8 процентных пункта (стандартное отклонение), что значительно ниже различий между сидами, указанных ниже. Метод SFT показывает 0,766 и 0,773 на двух использованных машинах; каждое сравнение выполняется с SFT на той же машине.

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

3.2 Дообучение на данных о неудачах

Траектории обучения BWM для RoboTwin получены с помощью скриптового планировщика, и сохранялись только те сиды, в которых планировщик преуспевал. Модель никогда не видела захваты в стиле π0.5, которые смыкаются в воздухе в нескольких сантиметрах от цели, поэтому она все равно засчитывает подъем. Исправление затрагивает только данные, а не архитектуру или функцию потерь:

  • Данные: 677 запусков в RoboTwin от базовой модели π0.5 и нескольких контрольных точек симуляторного PPO, 165 из них — неудачи. Неудачи происходят естественным образом, без скриптовых возмущений. Метки успехов используются только для коэффициентов выборки, разделения и оценки, но никогда в функции потерь.
  • Окно: обучение и инференс используют общее окно из 61 кадра: 1 начальный кадр, 8 кадров истории и 52 будущих кадра, что согласуется с 50-шаговыми фрагментами действий π0.5.
  • Потери: чистый первый кадр, слегка зашумленные кадры истории, потери сопоставления потоков (flow-matching) только для будущих кадров. Публичный каркас обучения BWM поддерживает только однокадровое кондиционирование; данный шаг обобщает его на многокадровое отслеживание истории в соответствии с WoVR.
  • Обучение: начать с шага step-12000, смешать успехи и неудачи, 750 шагов.

Исправление затрагивает WM, а не вознаграждение. Когда WM допускает ошибку, на кадрах действительно видно успешный захват, поэтому любое вознаграждение, считывающее только кадры, засчитает успех, а при реальном развертывании WM у нее нет состояния симулятора, к которому можно было бы вернуться. Максимум по кадрам для замороженного классификатора, временной гейт, требующий k последовательных положительных кадров, и признаки, построенные только на основе геометрии действий и начальной позы, — все они не справляются с разделением успеха и неудачи на кадрах WM.

3.3 Рецепт PPO внутри WM

Скорректированный рецепт меняет три вещи по сравнению с PPO по умолчанию в RLinf. Сам рецепт вместе с бэкэндом world-model BWM был представлен по адресу RLinf #1714:

  • Сообщать о завершении по лимиту эпизода. Эпизоды WM всегда завершаются на шаге 200, а вознаграждение представляет собой оценку 0/1 для всего эпизода, однако среда сообщает о шаге 200 как осечке (truncation). Стандартный подход добавляет $\gamma V(s_T)$ к последнему вознаграждению в предположении, что эпизод продолжится; при $\gamma = 1$ целевое значение критика становится «вознаграждением плюс собственная оценка», и отдачи дрейфуют вместе со значением. Исправление заставляет среду world-model сообщать о завершении на max_episode_steps. Эта опция отключена по умолчанию.
  • Подключите критик к выходу VLM и выполните его прогрев (warmup). Следуя совету $\pi$RL для $\pi 0.5$, критик считывает выход VLM, а не эксперта по действиям, входы которого смешаны с зашумленными действиями. Первые 4 шага обучения обучают только критик, а следующие 10 обновляются совместно, поэтому количество обновлений политики совпадает с исходным рецептом.
  • Уменьшите скорость обучения ценности в 10 раз. При исходной скорости примерно половина обновлений критика превышает порог обрезки (clip) ценности.

3.4 Несоответствие при замкнутом контуре (closed-loop) воспроизведения

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

4 Основные результаты

4.1 Тонкая настройка на данных неудач исправляет грубые ложные успехи

При запуске базовой BWM в разомкнутом контуре с действиями неудачных попыток в симуляторе 3 из 4 неудачных захватов отображаются как захваченные и установленные вертикально:

На 85 эпизодах (57 неудач, 28 успехов) одни и те же команды поступают как в симулятор, так и в WM, а один и тот же замороженный механизм вознаграждения считывает оба набора кадров:

Источник кадров

Неудача, считающаяся успехом

Полнота (recall) успехов

Кадры симулятора

0.07

1.00

Базовый BWM, шаг 12000

0.825

1.00

Тонкая настройка на данных неудач, 750 шагов

0.263

0.964

Вознаграждение точно работает для обеих сторон на кадрах симулятора; ошибка заключается в кадрах WM. Пиксельные метрики не показывают этого разрыва: PSNR базовой BWM на неудачных попытках всего примерно на 1 дБ ниже ее PSNR на успешных, в то время как 82,5% неудач преподносятся как успехи. После тонкой настройки число ложных успехов падает до 26,3% при полноте успехов 0,96.

4.2 При использовании скорректированного рецепта PPO внутри WM превосходит SFT на всех сидах обучения

Внутри дообученной WM исходный и скорректированный рецепты обучаются с пятью сидами каждый при прочих равных условиях. Критерий был зафиксирован до запусков: все пять приростов по сравнению с SFT больше 0, каждый 90-процентный нижний предел больше 0, а стандартное отклонение между запусками меньше 7 процентных пунктов.

Сид

Среднее

Исходный рецепт (п.п.)

+17.4

+4.6

+13.0

+10.4

−19.0

+5.3

14.3

Скорректированный рецепт (п.п.)

+11.6

+11.1

+7.4

+14.4

+3.9

+9.7

4.0

Несоответствие для скорректированного рецепта

0.18

0.16

0.11

0.20

0.25

Все пять запусков со скорректированным рецептом проходят критерий с односторонним 95-процентным нижним пределом +5,8; даже худший, сид 5, дает прирост +3,9, а его интервал находится выше 0. Средний интервал исходного рецепта составляет [−8,4, +18,9], который пересекает ноль. Несоответствие для скорректированного рецепта колеблется от 0,11 до 0,25, что соответствует уровню SFT (0,17). Частота успехов WM во время обучения меняется в зависимости от прироста в симуляторе: у сида 3 самая низкая частота во время обучения (0,48–0,66) и наименьший прирост, в то время как у сида 4 самая высокая частота за последние два шага (0,89, 0,90) и наибольший прирост. Пять сидов заняли около 43 часов на 8 графических процессорах AMD Instinct™ MI300X.

5 Абляции и анализ

5.1 Данные о неудачах выполняют работу, а не шаги обучения

При тех же 500 шагах вариант управления только по успехам имеет 86% ложных успехов, что хуже базового BWM, в то время как смешанный рецепт имеет 33%. При кластеризации по исходному состоянию 11 из 14 состояний неудач ухудшаются при обучении только на успехах и ни одно не улучшается. Оба рецепта улучшают PSNR по сравнению с базовым BWM примерно на 2 дБ, поэтому пиксельные метрики также не могут их различить.

У шагов обучения есть предел. На 1000 шагов частота ложных успехов на валидации падает до 0,286, но полнота успехов падает до 0,73, поскольку модель начинает воспринимать успехи как неудачи. Ложные успехи и полноту необходимо оценивать в совокупности; рассмотрение только первого приводит к тому, что пессимистичная в целом WM принимается за отремонтированную.

5.2 Ухудшенный сид в исходном рецепте использует ложные успехи WM

Значение −19,0 для сида 5 с исходным рецептом на рисунке 2 имеет две возможные причины: нестабильное обновление PPO по правильному сигналу или переход политики в области, где WM формирует неверные результаты. Воспроизведение в замкнутом контуре позволяет разделить эти факторы:

Из 28 эпизодов, в которых сид 5 с исходным рецептом преуспевает внутри WM, 12 завершаются неудачей при воспроизведении в симуляторе (0,43). У SFT этот показатель составляет 5/30 (0,17), избыточное несоответствие равно +0,26 с 90-процентным интервалом [+0,08, +0,43]. Лучший сид, сид 1, имеет показатель 0,18, ничем не отличающийся от SFT. Замороженное вознаграждение считывает каждое из таких воспроизведений «успех в WM, неудача в симуляторе» как неудачу, поэтому ошибка заключается в том, что генерирует WM. Ни один из пяти запусков со скорректированным рецептом не демонстрирует такого несоответствия (рисунок 2, справа); их обновления не были направлены на ошибки WM.

Вариант базовой групповой модели без критика достигает +11,1 и +4,2 на двух сидах, что не более стабильно, чем исходный рецепт; дисперсия между сидами не исчезает просто за счет обхода критика.

5.3 Влияние компонентов рецепта на критик

В рамках исходного рецепта объясненная дисперсия критика никогда не превышает 0,13 и в некоторых запусках остается отрицательной. Если разделить дисперсию успехов от повторной оценки симулятора на межсостояния (between-state) и внутрисостояния (within-state), доля межсостояний составляет около 0,6–0,8. Это предел для идеального критика, который видит только начальное состояние, поэтому у критика достаточно сигнала для обучения. При добавлении исправлений по одному и чтении медианы объясненной дисперсии за 24 шага обучения:

Конфигурация

Медиана объясненной дисперсии

Наблюдение

+ критик на VLM, 4-шаговый прогрев

−0.138

Около половины обновлений превышают порог обрезки ценности

+ скорость обучения ценности ÷ 10

−0.038

Обрезка находится под контролем, но отдачи преследуют ценность: вознаграждения 0/1 дают отдачи до 2,25

+ завершение по лимиту эпизода

−0.016

Дрейф исчезает; критик изучает только среднюю отдачу

При исправленном рецепте критик по-прежнему определяет лишь среднее значение, а его объясненная дисперсия остается близкой к 0 во всех пяти финальных запусках. Линейный зонд на входах критика восстанавливает успешность с R² 0,283, поэтому информация присутствует, но онлайн-обучение ее не извлекает. Таким образом, стабильность исправленного рецепта проистекает не из более точного критика; вклад каждого из трех исправлений в прирост симулятора не исследовался методом абляции по отдельности.

5.4 Ошибки интерфейсного контракта сами по себе способны вызвать регресс RL внутри WM

При наличии трех ошибок интерфейса (переставлены блоки действий двух рук; WM выдает эхо предыдущей команды вместо начального состояния при сбросе; эпизод завершается при первом положительном вознаграждении) PPO внутри базового BWM снижает успешность в симуляторе с 93/144 до 64/144, в то время как PPO, запущенный напрямую в симуляторе, меняет показатель с 93/144 до 114/144. После исправления интерфейса три одиночных запуска внутри базового BWM показывают значения −4,9, +5,6 и +16,0 процентных пункта по сравнению с SFT, без систематического регресса. Отрицательный результат RL внутри WM нельзя относить на счет WM до тех пор, пока не исключены ошибки интерфейса.

5.5 Больший объем on-policy данных не устраняет развилки действий

В разделе 5.2 указывается прямое решение: добавить траектории обновленных политик с метками симулятора в дообучение WM. Пять чекпоинтов исходного рецепта запускаются по 4 раза на 48 начальных состояниях обучения, что дает 960 траекторий (167 сбоев). Они объединяются с 677 траекториями из раздела 3.2, и WM дообучается с шага 12000 в течение 750 шагов по тому же рецепту, причем данные являются единственной переменной. Стоп-фильтр был заморожен заранее: полнота (recall) успешности валидации ниже 0,80 считается сбоем.

Полнота успешности нового WM падает с 0,933 до 0,767, не проходя стоп-фильтр, а ложная успешность возрастает с 0,418 до 0,509. Шаги обучения здесь ни при чем: дискриминация (полнота минус ложная успешность) достигает пика на 500 шагах, а затем снижается. Все начальные состояния новых данных уже присутствовали в исходном обучающем наборе, 85% повторных запусков одной и той же политики на одном и том же состоянии дают одинаковый результат, а 45% сбоев приходятся на деградировавший сид 5.

Правая панель Рисунка 5 показывает более глубокую причину. В добавленных данных исходы из общего начального состояния в основном зависят от действий: только 32% начальных состояний имеют единственный исход по сравнению с 71% в исходном обучающем наборе. На этих сбоях, даже когда сами траектории присутствуют в обучающем наборе, исходный WM все равно относит 29% к успешным, а новый WM — 40% (по сравнению с 12% для собственных кадров симулятора); на неизученных начальных состояниях новый WM показывает 56%. Данные о сбоях исправляют грубые ошибки, определяемые начальным состоянием; сбои, определяемые деталями действий, не уменьшаются при увеличении объема данных или шагов. Невозможно текущими данными разделить, не научился ли WM распознавать развилку действий или генерация с разомкнутым контуром отклоняется все дальше и дальше.

6 Обсуждение и ограничения

Результаты разделяют надежность RL внутри WM на два уровня. Первый — это рецепт обучения: на одной и той же WM одни лишь дефекты рецепта переводят межсидовое стандартное отклонение с 4,0 до 14,3 и позволяют отдельным сидам эксплуатировать ошибки WM. Одиночный прирост мало говорит о том, пригоден ли WM к использованию; для валидации требуется как минимум несколько сидов обучения и проверка рассогласования при воспроизведении в замкнутом контуре. Второй уровень — это точность исходов WM: данные о сбоях устраняют грубую ложную успешность, но развилки действий из общего начального состояния по-прежнему определяются неверно в 29–40% случаев.

Второй уровень приобретает большую значимость по мере масштабирования обучения. Градиенты RL поступают из начальных состояний со смешанными исходами: когда все запуски из состояния успешны или все неуспешны, групповое преимущество примерно равно нулю. Ложные успехи WM концентрируются ровно в таких состояниях. В рамках использованного здесь бюджета в 10 обновлений политики исправленный рецепт не утягивается ими вниз; более длительное обучение и более крупные групповые базисы еще сильнее сосредоточат градиенты на этих состояниях, и точность на развилках действий, вероятно, установит верхний предел прироста RL внутри WM. Можно напрямую протестировать два направления: принимать WM, сравнивая вердикты WM и симулятора на траекториях обновленной политики из начальных состояний со смешанными исходами, и обучать WM на парных контрфактуалах из одного старта (номинальное действие и возмущенное действие, приводящее к сбою), что и является направлением CureWM. Большее количество запусков на шаг снижает дисперсию градиента, не меняя смещение WM; исходя из межсидовой дисперсии исправленного рецепта, ожидаемый прирост составляет около 0–3 процентных пунктов, что ниже примерно 4 пунктов, которые могут обнаружить пять сидов, поэтому данный вариант не запускался.

Ограничения:

  • Одна задача, одна политика, один WM и всего 10 обновлений политики за запуск; неизвестно, сохраняются ли выводы для других задач или более длительного обучения.
  • Три изменения в исправленном рецепте не исследовались методом индивидуальной абляции относительно прироста симулятора.
  • Новый WM в разделе 5.5 обучался единожды с единственным сидом генерации; воспроизведение в замкнутом контуре охватывает лишь 24 начальных состояния на чекпоинт.
  • Для нескольких начальных состояний исход зависит от состояний, невидимых во входах WM (субпиксельные возмущения меняют его), поэтому сама наземная истина (ground truth) симулятора там нестабильна.

Данные, веса и код

  • 677 запусков π0.5 из раздела 3.2 (165 сбоев, с метками успешности симулятора): lidavidsh/robotwin-adjust-bottle-pi05-rollouts
  • Веса BWM после 750 шагов дообучения на данных сбоев (среда RL из раздела 4.2): lidavidsh/bwm-robotwin-adjust-bottle-failure-ft
  • Исправленный рецепт из раздела 3.3 и бэкенд мир-модели BWM: , пока не объединены.

Использование ИИ

Выполнение экспериментов, анализ логов и первый вариант этого поста были выполнены с помощью ассистентов кодирования на базе ИИ (Claude в Cursor, с Codex для проектирования и обзора результатов); экспериментальные компромиссы, выводы и правки перед публикацией принадлежат автору. Рисунки построены по скриптам из результатов экспериментов или захвачены из кадров экспериментов; ни один из них не сгенерирован ИИ. Эксперименты проводились на AMD MI300X.

Ссылки

О чём эта статья