GRASP: новый градиентный планировщик для долгосрочного планирования с мировыми моделями

Исследователи из BAIR (Berkeley AI) представили GRASP (Gradient RelAxed Stochastic Planner) — новый градиентный планировщик для мировых моделей, который решает проблемы долгосрочного планирования. Мировые модели (world models) — это обученные динамические модели, предсказывающие последовательности наблюдений. Планирование с их помощью становится сложным для длинных горизонтов из-за глубоких вычислительных графов, приводящих к взрыву/затуханию градиентов, и нежадного ландшафта с локальными минимумами. Стандартный метод развёртки (rollout) через модель приводит к плохой обусловленности. Авторы предлагают подъём (lifting) — переход к мягкому ограничению динамики, где оптимизация ведётся по состояниям и действиям, что позволяет распараллеливать вычисления и избегать произведения якобианов. Однако для глубоких нейросетевых мировых моделей прямое дифференцирование по состояниям проблематично из-за чувствительности к состязательным возмущениям (adversarial robustness): модель обучается на многообразии низкой размерности и резко меняется в ортогональных направлениях. GRASP решает эту проблему, используя только градиенты по действиям $D_a F_\theta$, которые надёжны, так как пространство действий обычно низкоразмерно и хорошо обучено. Таким образом, GRASP строит планировщик на основе коллокации, зависящий только от якобианов действий, что делает долгосрочное планирование робастным.