GRASP: градиентное планирование для мировых моделей
Современная мировая модель может предсказать, как будет выглядеть сцена через десятки шагов вперёд. Но попросите её спланировать, как туда добраться, и всё разваливается: градиенты взрываются, оптимизация застревает в локальных минимумах, а сама модель оказывается уязвимой к чему-то очень похожему на adversarial-атаки. Исследователи из BAIR (Berkeley AI Research) совместно с Яном Лекуном предложили GRASP, планировщик, который делает градиентное планирование с мировыми моделями практичным на длинных горизонтах. Разбираем, как это работает.
Что такое мировая модель
Мировая модель (world model) это обученная модель, которая по текущему состоянию среды и последовательности будущих действий предсказывает, что произойдёт дальше. Формально она задаёт распределение P(s_{t+1} | s_{t-h:t}, a_t), аппроксимирующее истинную динамику среды. В детерминированном случае это просто отображение s_{t+1} = F(s_t, a_t).
Состояние здесь обычно не пиксели, а выученное латентное представление, компактное и дифференцируемое. Ключевое свойство: мировая модель даёт вам дифференцируемый симулятор. Раз всё дифференцируемо, возникает соблазн: почему бы не найти оптимальную последовательность действий обычным градиентным спуском, минимизируя расстояние от финального состояния до цели? Именно так устроен классический подход, и именно он ломается на длинных горизонтах.
Почему планирование на длинных горизонтах не работает
Авторы выделяют три независимые причины, и каждая сама по себе способна убить планировщик.
Первая причина знакома всем, кто сталкивался с backpropagation through time. Когда вы дифференцируете через модель, применённую к самой себе T раз, производная по ранним действиям превращается в произведение T якобианов. Обусловленность этого произведения растёт экспоненциально с горизонтом: градиенты либо взрываются, либо затухают до нуля. На горизонте в 5 шагов это терпимо, на горизонте в 60 шагов сигнал до первых действий просто не доходит.
Вторая причина связана со структурой самой задачи. На коротких горизонтах жадное решение, двигаться к цели напрямик на каждом шаге, почти оптимально. На длинных горизонтах оптимальная траектория принципиально не жадная: обойти стену, отойти назад, чтобы потом толкнуть объект, сменить позицию перед рывком. Расстояние до цели вдоль оптимального пути немонотонно, и ландшафт потерь покрывается локальными минимумами. Дополнительно пространство оптимизации растёт линейно с горизонтом: T умножить на размерность действий. Локальных ловушек становится всё больше.
Третья причина специфична именно для глубоких мировых моделей, и она самая коварная.
Проблема, о которой мало говорят: adversarial-хрупкость градиентов по состоянию
Вспомните классические adversarial-примеры в компьютерном зрении. Работы Szegedy et al. (2014) и Goodfellow et al. (2015) показали: достаточно сдвинуть изображение на малую величину вдоль градиента логита, и классификатор уверенно выдаст неверный класс. Позднее Stutz et al. (2019) дали этому геометрическое объяснение: обучение контролирует поведение модели вдоль касательных направлений к многообразию данных, но никак не регуляризует её в ортогональных направлениях. Модель остаётся гладкой вдоль данных и патологически острой поперёк них. Более того, известен и трейд-офф между точностью и робастностью (Tsipras et al., 2019): чтобы точнее фитить сложные функции, модели выгодно быть острее в этих направлениях. Так что эту проблему нельзя просто «выучить наружу».
Как это бьёт по планированию? Мировая модель обучается на траекториях вида (s_1, a_1, s_2, ...), поэтому многообразие состояний, которые она реально видела, имеет размерность, ограниченную сверху размерностью пространства действий. Оно несравнимо уже, чем всё пространство состояний. Следствие: найти adversarial-сдвиг состояния, который обманет модель, тривиально просто.
Когда вы оптимизируете состояния напрямую через F, градиент по входу-состоянию ведёт вас именно в эти хрупкие ортогональные направления. Оптимизация становится «липкой»: состояние сдвигается на микрон в сторону, модель уверена, что локальная цель достигнута, хотя физически ничего не произошло. Планировщик взламывает собственную мировую модель вместо того, чтобы планировать в реальности.
Важная деталь: это бьёт не только по lifted-подходам. Даже при последовательной оптимизации через полный rollout произведение якобианов легко выносит промежуточные состояния за пределы многообразия данных, и дальше включаются те же хрупкие направления.
Коллокация: классический приём против экспоненциального взрыва
Первый шаг к решению пришёл из классической литературы по планированию и робототехнике. Вместо того чтобы жёстко требовать s_{t+1} = F(s_t, a_t) и прокатывать всю траекторию через модель, ограничение динамики смягчают до штрафа: минимизируем сумму квадратов ||F(s_t, a_t) - s_{t+1}|| по всем t, оптимизируя и действия, и состояния одновременно. Это называется коллокацией (collocation) или lifting.
У такой формулировки те же глобальные минимумы, что и у исходной задачи (обе равны нулю ровно на физически допустимых траекториях), но ландшафт оптимизации совершенно другой, и выгоды две. Во-первых, каждое вычисление F зависит только от локальных переменных, поэтому все T слагаемых считаются параллельно по времени. На длинных горизонтах это огромное ускорение. Во-вторых, произведение якобианов превращается в сумму: градиент по a_0 больше не проходит через глубокую T-шаговую композицию. Проблема взрывающихся градиентов исчезает.
Бонус: раз состояния оптимизируются напрямую, траектория может временно проходить через «нефизичные» области, исследуя промежуточные точки, которые rollout-планировщик никогда бы не посетил.
Но бесплатных обедов не бывает. Коллокация усиливает именно ту проблему, о которой шла речь выше: теперь состояния шевелятся свободно, и градиенты по входу-состоянию модели ведут прямиком в adversarial-области. Нужен следующий слой решения.
GRASP: как это чинится
Ключевое наблюдение авторов красиво своей асимметрией. Градиент по состоянию D_s F ненадёжен и adversarial-хрупок. А вот пространство действий обычно низкоразмерное и покрыто обучающими данными плотно: модель видела каждое направление действий. Значит, градиент по действию D_a F ведёт себя прилично, и через него оптимизировать можно.
Отсюда архитектура GRASP (Gradient Relaxed Stochastic Planner): lifted-планировщик первого порядка, который зависит только от якобианов по действиям и никогда не доверяет якобианам по состояниям. Три ингредиента.
Первый ингредиент: исследование через зашумление состояний. Планирование остаётся невыпуклым, поэтому в обновления виртуальных состояний добавляют гауссов шум: s_t сдвигается на градиентный шаг плюс случайная компонента с амплитудой sigma. Действия при этом обновляются обычным детерминированным спуском. Шум в состояниях позволяет «перепрыгивать» между бассейнами притяжения в lifted-пространстве, а чистые действия удерживают точность сходимости. Авторы специально отмечают: зашумлять именно состояния, а не действия, оказалось лучшим балансом между исследованием и нахождением острых минимумов.
Второй ингредиент: хирургия градиентов. Хрупкий путь, это градиент, текущий во вход-состояние модели. Простейшее решение, stop-gradient на состояние при вычислении F, само по себе не работает: сигнал цели перестаёт доходить до ранних состояний, и оптимизация вырождается в тривиальный минимум, где всё стоит на месте, а последнее действие пытается добраться до цели одним прыжком. Лечится добавлением плотного целевого члена: сумма ||F(s_t, a_t) - g||^2 по всем t. В обычной постановке такой член смещал бы решение к жадному движению напролом, но здесь он уравновешивается stop-gradient членом динамики, который тянет траекторию к физической допустимости. Итоговый лосс, сумма этих двух компонент с коэффициентом gamma, не содержит зависимости от градиентов по состоянию вообще.
Третий ингредиент: периодическая синхронизация. Lifted-цель это всё же аппроксимация исходной последовательной задачи. Поэтому каждые K_sync итераций GRASP делает короткую фазу уточнения: прокатывает текущие действия от s_0 через модель и делает несколько маленьких шагов по исходному rollout-лоссу ||s_T(a) - g||^2. Это возвращает траекторию к реальным состояниям и компенсирует накопленную неточность аппроксимации. При желании этот шаг можно заменить любым последовательным планировщиком, например CEM: суть в том, чтобы получить немного пользы от полной синхронизации пути, не теряя преимуществ lifted-оптимизации.
Результаты на Push-T: чем длиннее горизонт, тем сильнее отрыв
Проверяли на задаче Push-T, классическом бенчмарке, где нужно толкать T-образный блок в целевую позицию. Сравнение с CEM (cross-entropy method), обычным градиентным спуском (GD) и lifted-планировщиком LatCo. Формат чисел: процент успеха и медианное время до успеха.
На горизонте H=40 GRASP показывает 59.0% успеха за 8.5 секунды против 61.4% у CEM за 35.3 секунды и 51.0% у GD за 18 секунд. LatCo заметно хуже: 15.0% за почти 600 секунд. То есть на коротком горизонте GRASP сопоставим по качеству с лидером, но вчетверо быстрее.
Дальше начинается главное. На H=50 GRASP выходит в лидеры по качеству: 43.4% против 30.2% у CEM, и при этом 15.2 секунды против 96.2. На H=60 отрыв растёт: 26.2% у GRASP против 16.4% у GD и 7.2% у CEM. На H=70 у GRASP ещё 16.0% успеха, у CEM 7.8%, у LatCo уже ноль. На H=80 GRASP держит 10.4% и 58.9 секунд, тогда как GD даёт 6.4%, CEM 2.8%, а LatCo полностью сдаётся.
Картина однозначная: с ростом горизонта все методы деградируют, но GRASP деградирует медленнее всех и остаётся самым быстрым. Отдельная тонкость, которую честно отмечают авторы: медианное время успеха смещено вверх у методов с высоким процентом успеха (они «успевают» и на более тяжёлых эпизодах). GRASP умудряется быть быстрее конкурентов даже с этой поправкой.
Что дальше
Авторы называют GRASP первой итерацией нового класса планировщиков. Среди естественных продолжений: перенос на диффузионные мировые модели, где глубокие латентные шаги можно трактовать как сглаженные версии самой модели, более изощрённые оптимизаторы и стратегии зашумления, интеграция в замкнутый контур управления или в обучение RL-политик для адаптивного длинно-горизонтного планирования.
Финальное наблюдение из статьи звучит оптимистично: область планирования и управления невероятно зрелая, десятилетия накопленной теории, а вот применение этой теории к современным крупным мировым моделям почти не исследовано. Когда правильные идеи найдутся, планировщики по мировым моделям станут таким же стандартным инструментом, каким сегодня является RL. Судя по тому, как GRASP переиспользует коллокацию из робототехники семидесятых для моделей 2026 года, этот процесс уже идёт.
Часто задаваемые вопросы
Чем GRASP отличается от CEM?
CEM это сэмплирующий метод без градиентов: он генерирует популяцию последовательностей действий и отбирает лучшие. GRASP использует градиенты мировой модели по действиям, поэтому находит решения заметно быстрее: на H=50 он в шесть раз быстрее CEM при более высоком проценте успеха. CEM при этом можно встроить внутрь GRASP как фазу синхронизации.
Почему нельзя просто сделать мировую модель робастной к adversarial-сдвигам?
Известен фундаментальный трейд-офф между точностью и робастностью: модель, сглаженная в ортогональных направлениях, хуже фитит сложные функции. Пока глубокое обучение не сменит парадигму, хрупкость градиентов по состоянию останется. GRASP обходит проблему, а не решает её: он вообще не использует эти градиенты.
Где применять планирование с мировыми моделями?
Везде, где нужно многошаговое управление без проб и ошибок в реальном мире: манипуляция объектами, навигация роботов, автономное вождение. Длинные горизонты критичны именно там, где жадные стратегии не работают, а обучение RL-политики на реальном железе дорого или опасно.
Итог
Планирование с мировыми моделями упиралось в три стены: экспоненциально плохую обусловленность градиентов на длинных rollout, не жадный ландшафт с локальными минимумами и adversarial-хрупкость градиентов по состоянию. GRASP отвечает на все три: коллокация снимает первую, зашумление состояний вторую, stop-gradient с плотным целевым членом третью. Результат: единственный метод, который на горизонте 80 шагов ещё решает задачу, и делает это быстрее всех конкурентов. Если вы строите агентов поверх мировых моделей, посмотрите на статью Psenka et al. (arXiv 2602.00475): вероятно, ваш планировщик можно ускорить в разы, просто перестав доверять не тем градиентам.