AI Продвинутые RL алгоритмы: NPG, TRPO, PPO

[AI]

Продолжение постов про RL:

1) Intro Reinforcement Learning
2) Reinforcement Learning: Model-free & Deep RL
3) Reinforcement Learning: Policy gradient methods

Если вы заметите ошибки в формулах или опечатки — буду очень благодарна, если укажете на них 🙈.
Если вдруг формулы покажутся сложными, то отмечу, что главная ценность этой статьи – в объяснении того, откуда эти формулы берутся и что они означают. Однако если какие-то моменты останутся непонятными, то я рада доработать текст, чтобы сделать его более доступным.

Мой тг канал: not magic neural networks

Полезные ссылки:

1) Practical_RL/week09_policy_II
2) Policy gradient method
3) A Natural Policy Gradient
4) Метод сопряженных градиентов
5) Trust Region Policy Optimization (TRPO)
6) Proximal Policy Optimization (PPO)
7) DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models

1. Прежде всего вспомним Policy Gradient методы

1.1. Policy Gradient Theorem


Пусть мы имеем траектории

. Мы хотим максимизировать математическое ожидание дисконтированной суммы наград (objective function) по политикам

:

Траектории

приходят из

– распределение над траекториями

при условии политики

:

Заметим, что

можно записать как математическое ожидание по состояниям

из

-функции (позже данная запись пригодится):

Мы научились максимизировать функционал

с помощью градиентного подъема. Градиент

находим с помощью Монте-Карло оценки с применением logderivative trick

:

Распишем Монте-Карло оценку:

Перемножим скобки

и

между собой (занесем под одну скобку):

Разделим награду на «прошлое» и «будущее» относительно

:

Награда

к моменту времени

уже произошла и не зависит от того, какое

семплируется. Поэтому

– константа, при фиксированном

.

Рассмотрим часть формулы [1.3]:

Распишем математическое ожидание по определению:

Применим "обратный" logderivative trick

По определению политики

Таким образом:

Возвращаясь к формуле

, обнулим слагаемое с "прошлой" наградой:

Вынесем

из суммы

Применим закон полного математического ожидания

:

где

– условное математическое ожидание возврата при заданных текущих состоянии

и действии

.

Заметим, что это определение

-функции:

. Тогда:

Сформулировали Policy Gradient Theorem:

На практике оказалось, что данная формула имеет большую дисперсию. Для того чтобы снизить дисперсию решили вычитать некий

из

не зависящий от действий

.

1.2. Baseline и Advantage


Можно доказать, что если брать в качестве

использовать

-функцию, то дисперсия Монте-Карло оценки будет падать.

Таким образом:

где

– обучаем, а

будем получать из

по формуле

.

называется advantage функция. Она показывает насколько действие

лучше или хуже, чем среднее действие в этом состоянии при политике

. Если

, то действие

лучше среднего выбираемого политикой, иначе - хуже.

Однако, на практике оказывается удобнее эквивалентная

формула:

Математическое ожидание берётся по данным, сгенерированным самой

.

– discounted state visitation distribution политики. Это распределение состояний, которое показывает как часто политика

посещает состояние

, с учётом дисконтирования по времени.

• 
  
  
  — вероятность быть в состоянии
  
  на шаге
  
  .
  
  
• 
  
  
  – нормировка, чтобы это было распределение
  
  
  

Описанные выше алгоритмы REINFORCE и A2C (Advantage Actor–Critic) являются on-policy, что делает их неэффективными с точки зрения использования данных (градиент оптимизируемой политики

корректно оценивается только на данных, сгенерированных этой же политикой).

При этом, сам градиентный спуск не такой дорогой как сбор данных (траекторий). Естественным образом возникает желание оптимизировать полезность политики

, используя данные, собранные старой политикой

. В этом заключается основная цель дальнейших методов (TRPO, PPO и др.)

2. Как оптимизировать новую политики по данным из старой политики


Для этого в формуле

надо придумать как брать математическое ожидание и advantage по старой политике

.

2.1. Заменяем advantage


Рассмотрим новый функционал:

Теперь мы хотим чтобы новая политика была как можно лучше старой.

И прежде чем расписывать новый функционал, проделаем "подготовительную" работу:
1) Представим

в виде телескопические суммы.
2) Распишем разность

.
3) Распишем новый функционал

.

2.1.1. Телескопическая сумма


Телескопическая сумма — это сумма вида:

где каждое слагаемое ряда представляется в виде разности и поэтому частичная сумма ряда упрощается.

Название "телескопическая сумма" дано по аналогии с трубой телескопа, который может уменьшить свою длину, сложившись несколько раз.

Мы будем пользоваться телескопической суммой для любой последовательности

, такой что

:

По аналогии с формулой выше: для произвольной траектории

и произвольной политики

справедливо равенство:

при условии, что

.

2.1.2. Распишем разность V-функций

по определению:

Занесем

под математическое ожидание:

Запишем

в виде телескопической суммы по формуле

:

Вынесем

за скобки

В силу равенства

, внесём матожидание по будущим состояниям

, условно на

.

Заметим, что

:

по определению advantage функции:

Итоговая формула:

2.1.3. Вернемся к разности функционалов J(new) - J(old)


Вспомним формулу для

:

Распишем новый функционал:

Подставим формулу

, получим:

По аналогии с

перепишем эквивалентную формулу

:

2.2. Заменяем математическое ожидание


В получившийся формуле

остается понять как заменить математические ожидания по новой политике

на математические ожидания по старой политике

.

2.2.1. Importance Sampling Trick / Трюк важностного семплирования


Пусть у нас есть математическое ожидание функции

по распределению

:

Теперь мы хотим выразить

через другое распределение

. Чтобы это сделать, просто умножаем и делим под интегралом на

.

С помощью данного трюка можно заменить математическое ожидание по действиям в

:

2.2.2. Суррогатная функция


Для

Importance Sampling сделать не получится, поскольку мы не знаем

.

Просто заменим

на

и назовем получившийся функционал суррогатной функцией

.

Однако, максимизируя данный функционал, мы будем увеличивать вероятность действий, которые максимизируют значение advantage-функции

. Решение будет сходится к

, где

будет у действия приводящего advantage к максимуму, а для остальных действий

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

Если пытаться повышать вероятности маловероятных действий, отношение

. Градиенты будут очень большими.

Аналогично, если пытаться уменьшать вероятность популярных действий, то

. Градиенты будут очень маленькими

2.3. Теоретическая граница ошибки для нового функционала (theoretic error bound)


В статье Trust Region Policy Optimization (2015) важным результатом является следующая теорема (Theorem 1):

Модуль разницы между желаемым оптимизируемым функционалом

и суррогатной функцией

ограничен максимумом по состояниям

-дивергенции между старой и новой политикой.

Другими словами, если политики

и

близки, то разница

не велика.

Введем вариационную низкую оценку

(variational lower bound):

Будем улучшать политику в два шага:

1. 
   Minorization step: фиксируем текущую политику и строим нижнюю оценку
   
   , совпадающую с истинной целью
   
   в текущей точке.
   
   
2. 
   Maximization step: максимизируем нижнюю вариационную оценку
   
   в надежде, что максимизируется желаемый функционал
   
   
   

будет оценивать усредненным

между двумя политиками по состояниям.

А константу

объявим гиперпараметром.

Таким образом, будем оптимизировать

с ограничением на разницу между старой и новой политиками.

3. Проблема оптимизации


На недолгое время абстрагируемся, оставим все результаты выше и рассмотрим задачу градиентного спуска как задачу условной оптимизации.

Пусть

– оптимизируемая функция. Будем оптимизировать

в ее локальной окрестности. Другими словами, рассмотрим задачу обновления параметров

, начиная с точки

. Мы хотим сделать шаг

, который минимизирует линейную аппроксимацию функции

в окрестности

, но при этом величина шага не должна быть слишком большой.

Приблизим

линейной аппроксимацией с помощью разложения в ряд Тейлора:

Поскольку

не зависит от параметров

, то константу можно опустить.
Таким образом, будем оптимизировать

.
Обозначим расстояние между старыми и новыми параметрами:

. Тогда, мы хотим найти

такое что

с квадратичным ограничением

.
Аналитическое решение методом множителей Лагранжа:

Оптимальная дистанция

пропорциональна антиградиенту оптимизируемой функции

.

Немного усложним задачу. Будем минимизировать ту же функцию

, но ограничим расстояние между старыми и новыми параметрами по неевклидовой метрике

, где

– положительно определённая матрица.

Аналитическое решение для такого ограничения:

Таким образом,

пропорционально

3.1. Natural Policy Gradient, NPG (2001)


Статья: A Natural Policy Gradient

Вернемся к задаче обучения с подкреплением.

Пусть имеется политика с параметрами

. Хочется найти новые параметры

, которые находились бы в окрестности параметров

.

Квадратичное расстояние между весами

, где

будет задавать окружность в пространстве весов. Однако, хотелось бы измерять расстояние между политиками

и

какими-то более умными способами.

Будем использовать расхождение Кульбака-Лейблера (

-дивергенцию) между распределениями над действиями в качестве метрики между политиками

и

.

Обозначим

как функцию параметров

.
Разложим

около точки

в ряд Тейлора до второго порядка производной:

1. 
   
   
   – из определения
   
   -дивергенции.
   
   
2. 
   
   
   поскольку
   
   -дивергенция при совпадении распределений достигает своего минимума.
   
   
3. 
   
   
   – пренебрежимо мало.
   

Остается только член второго порядка:

Сформулируем метод Natural Policy Gradient.

Пока что, пользуемся функционалом

.
Максимизируем

с ограничениями

, где

– матрица вторых производных

-дивергенции.

В таком случае, оптимальный вектор изменения

будет пропорционален

.
Обновление политики происходит по формуле:

Однако, обращение матрицы

очень дорогая операция, потому данный метод будет работать только на игрушечных средах. Хотелось бы этот метод как-то улучшить.

Заметим, что можно не искать отдельно матрицу

, а искать сразу значение выражения

.

Тогда вместо обращения матрицы можно решить линейную систему:

Поскольку матрица

положительно определённая, то мы можем воспользоваться эффективным методом решения линейных систем — методом сопряжённых градиентов (conjugate gradients).

В результате получаем вектор изменения весов:

Параметр

подбирают линейным поиском шагая от

по направлению

:

Остается обновить политику:

3.2. Trust Region Policy Optimization, TRPO (2015)


Статья: Trust Region Policy Optimization

Вернемся к ранее выведенному функционалу:

Сформулируем задачу оптимизации с ограничениями:

Метод Trust Region Policy Optimization:

1. 
   Инициализируем политику
   
   
   
   1. 
      Loop:
      
      1. 
         
         
         
         
         
      2. 
         Генерируем траектории, используя текущую политику
         
         .
         
         
      3. 
         Считаем
         
         
         
         
      4. 
         Цикл по эпохам:
         
         1. 
            Цикл по батчам:
            
            1. 
               Считаем градиент суррогатной функции:
               
               
               
               
            2. 
               Аналитически считаем матрицу вторых производных KL-дивергенции:
               
               
               
               
            3. 
               Методом сопряженных градиентов ищем направление обновления весов, решая систему
               
               методом сопряжённых градиентов:
               
               
               
               
            4. 
               Линейным поиском подбираем параметр
               
               , в направлении
               
               с ограничением:
               
               и одновременной проверкой значения:
               
               
               
               
            5. 
               Обновляем параметры политики
               
               
               

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

В современных условиях классический TRPO практически не применяется — он слишком требователен к памяти и плохо масштабируется на большие модели. На смену ему пришла его усовершенствованная модификация — PPO (Proximal Policy Optimization).

3.3. Proximal Policy Optimization (PPO), 2017


Proximal Policy Optimization Algorithms

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

Добавим регуляризацию прям в максимизируемый функционал:

Однако, данный функционал не гарантирует сходимости алгоритма.

не достаточно штрафует за разницу между политиками и выражение

часто уходит либо в беcконечность, либо в ноль и обучение разваливается.

Ограничиваем изменение политики, "подрезая" (clipping) слишком большие значения importance ratio

.

Если

выходит за границы

, то считаем его равным

.Таким образом, новая политика не сможет выйти за регион

.

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

В статье PPO были представлены следующие графики. По вертикальной оси находятся значения функции ошибки, а по горизонтальной значение функции

.
В красной точке

совпадает с

, соответсвенно

.

Допустим

, то тогда и

должно увеличиваться, поскольку мы хотим увеличить вероятность этого действия. Однако, когда

доходит до числа

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

, то

должно уменьшаться. Но, при достижении границы

выгода с точки зрения функции потерь теряется.

На графике ниже представлена линейная интерполяция изменения весов. Слева расположена

, справа последняя версия

.

• 
  Синяя функция показывает изменение
  
  -дивергенции между старой и новой политикой.
  
  
• 
  Желтая функция – это обычная суррогатная функция без клиппинга. Однако она является лишь приближением целевого функционала
  
  и не гарантирует его улучшения.
  
  
• 
  Зеленая функция – это клипнутый importance ratio (
  
  ). На графике видно, что при опредленном количестве обновлений, она достигает своего плато, что логично из формулы клиппинга.
  
  
• 
  Красная функция – это минимум из клипнутой функции и обычной суррогатная функции. Видно, что ДО того как функция достигнет значения
  
  , она просто улучшает политику. По достижении
  
  хорошие политики больше не улучшаются и плохие начинают портить красную функцию. Значит отходить дальше от
  
  не имеет смысла.
  

Сформулируем метод Proximal Policy Optimization:

1. 
   Инициализируем параметры политики
   
   
   
   
2. 
   Loop:
   
   1. 
      Сохраняем старую политику
      
      
      
      
   2. 
      Генерируем траектории c
      
      
      
      
   3. 
      Считаем advantage estimate
      
      
      
      
   4. 
      Цикл по эпохам:
      
      1. 
         Цикл по батчам:
         
         1. 
            Считаем importance ratio:
            
            
            
            
         2. 
            Считаем clipped importance ratio:
            
            
            
            
         3. 
            Считаем Сlipped objective:
            
            
            
            
         4. 
            Обновляем параметры
            
            
            

4. Generalized Advantage Estimation (GAE)


Можно заметить, что в методах TRPO и PPO был упущен момент c подсчетом advantage estimate

.

На практике функции

и

неизвестны и аппроксимируются по данным. Обычно обучается только value-функция, а advantage оценивается через возвраты.

4.1. Обучение Value-функции


Функцию ценности обычно аппроксимируют нейронной сетью

с функцией потерь

где

– оценка истинного возврата.

Monte-Carlo возврат

несмещённый, но имеет огромную дисперсию.
Поэтому на практике используют усечённую

-шаговую оценку:

Из определения

. Следовательно, оценка возврата может быть записана как

Тогда:

Обновление параметров градиентным спуском:

4.2. n-шаговый Advantage


В зависимости от метода оценки

-функции, существуют различные оценки для Advantage:

Монте-Карло оценка будет иметь высокий разброс (и низкое смещение).

Можно

-функцию оценить одношаговой оценкой из знаний

:

Такая оценка будет иметь высокое смещение (и низкий разброс).

Возникает задача умного подбора шага

при оценивании функции

:

Переобозначим функционал (актора):

и функционал критика:

4.3. Generalized Advantage Estimation (GAE)


Вместо того чтобы подбирать количество шагов

, можно использовать все

-шаговые оценки Advantage с затухающими весами:

Estimation​
Weight​

Однако, пересчитывать все

слишком долго и не оптимально.

Выпишем двушаговый

через

.

1. 
   Из определения
   
   :
   

1. 
   Выпишем одношаговый
   
   , но для следующего состояния
   
   и следующего действия
   
   :
   

1. 
   Умножим на
   
   :
   

1. 
   Добавим и вычтем
   
   из
   
   :
   

Если обобщить

на

шагов, то получится общая формула

-шаговой ошибки через одношаговые:

Подставив получившуюся формулу

в

можно доказать, что:

Однако

считаем не до конца траектории, а только по сделанным шагам, потому ограничим

до

:

На практике удобнее использовать следующую эквивалентную формулу к

:

А функцию ценности обновляем:

Основная проблема классического

-шагового возврата заключается в необходимости заранее выбрать конкретный горизонт

. Выбрав маленькое

, оценка возврата будет иметь высокий

. Выбрав большое

, оценка возврата будет с огромной дисперсией. На практике это создаёт трудности: оптимальный горизонт

различается для разных состояний, длины эпизодов непостоянны, а обучение становится нестабильным.

Вместо выбора одного

,

строит смесь всех

-шаговых возвратов. Таким образором, выбирается не конкретный горизонт, а скорость затухания памяти. При

алгоритм имеет короткую память и оценка близка к

(низкая дисперсия, но высокое смещение). При

– долгая память, оценка стремится к полному Монте-Карло возврату (низкое смещение, но высокая дисперсия). Эмпирически для многих задач оптимальный баланс достигается в диапазоне

.

4.4. Вернемся к PPO


Напишем полностью алгоритм Proximal Policy Optimization:

1. 
   Инициализируем параметры политики
   
   и параметры критика
   
   .
   
   
2. 
   Loop:
   
   1. 
      Сохраняем политику
      
      
      
      
   2. 
      Генерируем траектории (rollouts) из политики
      
      , сохраняя логиты
      
      для подсчета importance ratio.
      
      
   3. 
      Сохраняем выходы критика
      
      
      
      
   4. 
      Считаем одношаговые оценки для всех состояний:
      if not
      
      :
      
      
      
      else:
      
      
      
      
      
   5. 
      Считаем рекурсивно по формуле
      
      GAE оценки:
      
      
      
      for t in (N-2, 0, -1):
      if not
      
      :
      
      
      
      else:
      
      
      
      
      
   6. 
      Считаем целевые значения критика:
      
      
      
      
   7. 
      for epoch in n_epochs:
      
      1. 
         for batch in batches:
         
         1. 
            Считаем importance ratio:
            
            
            
            
         2. 
            Считаем clipping importance ratio
            
            
            
            
         3. 
            Считаем clipped objective:
            
            
            
            
         4. 
            Обновляем параметры актора:
            
            
            
            
         5. 
            Обновляем лосс критика:
            
            
            
            
         6. 
            И обновляем параметры критика:
            
            
            

Реализацию PPO можно посмотреть здесь: Practical_RL/week09_policy_II/ppo.ipynb at main · anyakangur/Practical_RL https://github.com/anyakangur/Practical_RL/blob/main/week09_policy_II/ppo.ipynb

На GRPO сил не хватило, но будет в следующий раз!