Предложен SQAM: дообучение flow-политик через RL без векторно-якобиевых произведений
Flow-политики хорошо передают богатые и разнообразные распределения действий, и всё больше интереса вызывает их дообучение с помощью off-policy RL (обучения с подкреплением на данных, собранных другой политикой), чтобы превзойти исходные демонстрации. Сложность в том, что дообучать flow-политику по выученной функции ценности непросто: действие порождается за много шагов потока (flow steps).
Существующий подход, adjoint matching (сопоставление сопряжённых), обновляет саму flow-модель, передавая информацию о ценности от финального действия назад к каждому шагу потока. Но для этого на каждом шаге нужно вычислять векторно-якобиево произведение через политику, а стоимость этого растёт с числом шагов потока и размером политики.
Авторы заметили, что усреднённый по батчу якобиан скорости предобученных flow-политик концентрируется на диагонали. Опираясь на это наблюдение, они вывели замкнутый скалярный adjoint: он масштабирует градиент ценности в финальном действии на время потока, и векторно-якобиевы произведения на каждом шаге становятся не нужны. Кроме того, они обнаружили, что при скалярном adjoint особенно важно контролировать значение критика в действиях, порождённых самой политикой. Из этих двух находок сложился метод Q-learning with Scalar Adjoint Matching (SQAM): скалярный adjoint плюс штраф по значению в таких действиях.
Результаты: выигрыш SQAM сосредоточен на четырёх самых сложных доменах OGBench, где доля успехов превышает показатель сильнейшего базового метода в каждом домене на 18-35 процентных пунктов. Чтобы проверить, переносится ли метод на крупные предобученные политики, авторы также дообучили vision-language-action политику на реальном двуруком роботе. SQAM улучшил результат по сравнению с обучением с учителем (supervised fine-tuning) во всех трёх задачах.
Ключевые факты
- SQAM (Q-learning with Scalar Adjoint Matching), метод off-policy дообучения flow-политик: скалярный adjoint плюс штраф по значению критика в действиях самой политики.
- Идея упрощения: усреднённый по батчу якобиан скорости предобученных flow-политик концентрируется на диагонали, поэтому adjoint выводится в замкнутом виде и масштабирует градиент ценности в финальном действии на время потока.
- Это устраняет векторно-якобиевы произведения на каждом шаге, стоимость которых в обычном adjoint matching растёт с числом шагов потока и размером политики.
- На четырёх самых сложных доменах OGBench доля успехов выше сильнейшего базового метода в каждом домене на 18, 35 процентных пунктов.
- На реальном двуруком роботе vision-language-action политика, дообученная SQAM, лучше варианта с обучением с учителем во всех трёх задачах.
Почему это важно
Дообучать flow-политики по выученной функции ценности непросто: действие строится за много шагов. Adjoint matching решает задачу принципиально, но его стоимость растёт с числом шагов потока и размером политики. SQAM убирает векторно-якобиевы произведения на каждом шаге за счёт замкнутого скалярного adjoint и, по заявлению авторов, даёт заметный прирост на самых сложных доменах OGBench.
Кому это важно
Исследователям обучения с подкреплением и робототехники, которые дообучают flow-политики или vision-language-action политики по данным вне политики, а также тем, кого ограничивает вычислительная стоимость adjoint matching.
Как это применить
В доступном тексте (аннотации) описан только сам метод: скалярный adjoint, масштабирующий градиент ценности в финальном действии на время потока, и штраф по значению критика в действиях политики. Вид штрафа и его коэффициент не приведены, поэтому для воспроизведения понадобится полный текст статьи.
Можно ли доверять
Это препринт на странице Hugging Face papers, доступна лишь аннотация. Все утверждения, слова самих авторов. Цифра «18-35 процентных пунктов», отрыв от сильнейшего базового метода в каждом из четырёх сложных доменов; абсолютные показатели не приведены. Для реального робота сказано лишь, что результат лучше обучения с учителем во всех трёх задачах, без величины улучшения.
Риски и подводные камни
Прирост сосредоточен на четырёх самых сложных доменах OGBench, а не на всех. Проверка на реальном двуруком роботе ограничена тремя задачами, размер улучшения не указан. Метод опирается на наблюдение о диагональной концентрации якобиана, сделанное для предобученных flow-политик. Для скалярного adjoint, по словам авторов, особенно важен контроль значения критика в действиях политики, отсюда штраф в составе метода.