Введение

Статья написана под впечатлением и при поддержке "Большой математической мастерской", проходившей летом 2022 г. в Академгородке на базе НГУ.

Теорема о градиенте стратегии является одним из краеугольных камней современного RL. Удивительно, что некоторые авторы книг, посвященных обучению с подкреплением, стараются обойти тему стороной и неохотно делятся информацией откуда и каким образом у стратегии растут градиенты. В этой заметке, мы разберём детали непростой анатомии стратегий во всех подробностях. За дополнительным сведениями об RL рекомендую обращаться к произведению Р. Саттона и А. Барто “Reinforcement Learning. An Introduction“

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

Некоторые определения и задача оптимизации RL

Начнём разговор с описания окружающего ландшафта, для этого сформулируем некоторые важные определения из области обучения с подкреплением.

Траектория $\tau = s_0, a_0, r_0, \ldots$ – последовательность состояний , между которыми перемещается агент, совершая действияи получая вознаграждение . Конечную траекторию будем также называть эпизодом.

При выборе очередного действия, агент использует стратегию $\pi_{\theta}$ .

Стратегия $\pi_{\theta}$ – отображение пространства состояний на пространство действий. Для каждого состояния стратегия возвращает распределение вероятностей выбора того или иного действия $\pi_{\theta} (s_t)$ .

В роли стратегии $\pi_{\theta}$ выступает какая-то хитрая функция с набором параметров $\theta$ , покрутив которые можно обучить её нужному поведению. Для удобства рассуждений мы будем считать, что это нейросеть.

Эффективность стратегии $\pi_{\theta}$ определяется вознаграждением, которое агент получил при движении вдоль траектории $\tau$ , порожденной $\pi_{\theta}$ .

Функция вознаграждения $R_t(\tau)$ представляет собой сумму всех вознаграждений , которые получил агент, начиная с момента времени и до конца эпизода . Коэффициент обесценивания $\gamma$ регулирует влияние отдаленных шагов на текущий момент времени

$R_t(\tau) = \sum_{t' = t}^{T} \gamma^{t' - t} r_t$

Если траекторию рассматривать целиком, то обозначение слегка упрощается (исчезает индекс ):

$R(\tau) := R_0(\tau) = \sum_{t = 0}^{T} \gamma^{t} r_t$

В сложных средах траектории могут сильно различаться даже для одной и той же стратегии $\pi_{\theta}$ , поэтому для оценки поведения агента удобнее усреднять вознаграждение по набору траекторий. Возникает

Целевая функция $J(\pi_{\theta})$ – среднее вознаграждение агента по всем траекториям, порожденным с помощью стратегии $\pi_{\theta}$ .

$J(\pi_{\theta}) = E_{\tau \sim \pi_{\theta}} \left[ R(\tau) \right] = E_{\tau \sim \pi_{\theta}} \left[ \sum_{t = 0}^{T} \gamma^{t} r_t \right]$

Возникает естественное желание, изменяя параметры $\theta$ стратегии $\pi_{\theta}$ , максимизировать целевую функцию, чтобы добиться наибольшего вознаграждения для агента.

$\max_{\theta} J(\pi_{\theta}) = E_{\tau \sim \pi_{\theta}} \left[ R(\tau) \right]$

Т.к. мы договорились считать $\pi_{\theta}$ нейросетью, то набором параметров $\theta$ являются веса сети.

Вычислительные затруднения

В первом приближении карта местности определена и сформулирована понятная цель: увеличить среднее вознаграждение, которое получает агент в процессе взаимодействия со средой, руководствуясь стратегией $\pi_{\theta}$ . Математическим языком эту цель можно выразить в виде задачи оптимизации

$\max_{\theta} J(\pi_{\theta}) = E_{\tau \sim \pi_{\theta}} \left[ R(\tau) \right] \quad \quad (1)$

Как это можно сделать? Первый ответ, который приходит в голову специалистам по ML, воспользоваться градиентом целевой функции $J(\pi_{\theta})$ по набору параметров $\theta$ : $\nabla_\theta J(\pi_{\theta})$ ! Схему градиентного подъема можно записать вот так

$\theta \leftarrow \theta + \alpha \nabla_\theta J(\pi_{\theta})$

Но как применить её на практике? Посмотрим внимательно на уравнение нашей задачи (1) оптимизации. Ключевым элементом выражения является функция вознаграждения $R(\tau)$ . Что же нам известно о ней? Как она зависит от $\theta$ ?

К сожалению, мы не знаем об $R(\tau)$ почти ничего, она выступает в роли загадочного черного ящика. Выражение $R(\tau) = \sum \gamma^{t} r_t$ не получится продифференцировать по $\theta$ . Но связь между $\theta$ и $R(\tau)$ можно почувствовать, семплируя траектории $\tau \sim \pi_{\theta}$ и вычисляя вдоль них вознаграждение $R(\tau)$ . То есть, нужная информация хранится в распределении $\tau \sim \pi_{\theta}$ и нужно найти способ до неё добраться.

Теорема о градиенте

Такой способ существует и называется теоремой о градиенте стратегии. Для удобства восприятия, мы сперва разберем её в общих обозначениях, а затем подставим интересующие выражения и выведем основной результат.

Некоторые полезные математические преобразования

Пусть определены следующие объекты: переменная , функция , условное (оно же параметризованное) распределение вероятностей $p(x | \theta)$ и математическое ожидание $E_{x \sim p(x|\theta)} \left[ f(x)\right]$ (то есть интеграл). Будем искать градиент математического ожидания по параметру $\theta$ .

Сперва выпишем определение мат. ожидания, а затем внесём оператор дифференцирования $\nabla_\theta$ под знак интеграла

$\nabla_\theta E_{x \sim p(x|\theta)} \left[ f(x)\right] = \nabla_\theta \int f(x) p(x | \theta) dx = \int \nabla_\theta \left( f(x) p(x | \theta) \right) dx =$

не зависит от $\theta$ , поэтому выносим её из-под оператора и домножаем всё подынтегральное выражение на единицу в форме $p(x | \theta) / p(x | \theta)$

$= \int f(x) \nabla_\theta p(x | \theta) dx = \int f(x) p(x | \theta) \dfrac{\nabla_\theta p(x | \theta)}{p(x | \theta)} dx =$

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

$= \int f(x) p(x | \theta) \nabla_\theta \log p(x | \theta) dx = E_{x \sim p(x|\theta)} \left[ f(x) \nabla_\theta \log p(x | \theta) \right]$

За счет смены порядка операторов дифференцирования и интегрирования, а также пары математических трюков удалось получить тождество, в котором градиент находится под интегралом и действует только на распределение $p(x | \theta)$ , явным образом зависящее от параметра $\theta$ .

$\nabla_\theta E_{x \sim p(x|\theta)} \left[ f(x)\right] = E_{x \sim p(x|\theta)} \left[ f(x) \nabla_\theta \log p(x | \theta) \right] \quad \quad \quad (2)$

На функцию при этом накладываются минимальные требования: мы хотим, чтобы она была интегрируемой. Тогда интеграл для мат. ожидания можно оценивать численно с помощью выборок $x \sim p(x | \theta)$ .

Градиент стратегии для RL

Пора вернуться обратно к градиенту целевой функции $\nabla_\theta J(\pi_{\theta})$ . Подставим в получившееся равенство (2) выражения для стратегии $\pi_{\theta}$ и переместимся в пространство траекторий. Теперь в роли выступает траектория $\tau$ , – это функция вознаграждения для траектории $R(\tau)$ , а таинственное распределение $p(x | \theta)$ превращается в $p(\tau | \theta)$ .

$\nabla_\theta J(\pi_{\theta}) = E_{\tau \sim \pi_{\theta}} \left[ R(\tau) \nabla_\theta \log p(\tau | \theta) \right]$

В целом выражение выглядит хорошо, но появился множитель $p(\tau | \theta)$ , который выражает вероятность возникновения траектории $\tau$ при условии, что стратегия задана набором параметров $\theta$ . Как его вычислить?

Мы знаем, что для каждого шага вероятность действия определяется стратегией и равна $\pi_{\theta} (a_t | s_t)$ , а вероятность перехода между состояниями и $s_{t+1}$ определяется как $p(s_{t+1} | s_t, a_t)$ . Тогда вероятность осуществления всей траектории $\tau$ равна произведению вероятностей этих переходов.

$p(\tau | \theta) = \prod_{t \ge 0} p(s_{t+1} | s_t, a_t) \pi_\theta(a_t | s_t) \quad \quad \quad (3)$

Загадка множителя $p(\tau | \theta)$ разгадана, но попробуем ещё немного улучшить уравнение градиента, прологарифмировав выражение (3) и вычислив градиент по $\theta$ . Сперва подставляем (3) и превращаем произведение в сумму, пользуясь свойствами логарифма

$\log p(\tau | \theta) = \log \prod_{t \ge 0} p(s_{t+1} | s_t, a_t) \pi_\theta(a_t | s_t) = \sum_{t \ge 0} \left( \log p(s_{t+1} | s_t, a_t) + \log \pi_\theta(a_t | s_t) \right)$

затем избавляемся от слагаемых, не зависящих от $\theta$ , т.к. при дифференцировании они дадут ноль

$\nabla_\theta \log p(\tau | \theta) = \nabla_\theta \sum_{t \ge 0} \left( \log p(s_{t+1} | s_t, a_t) + \log \pi_\theta(a_t | s_t) \right) = \nabla_\theta \sum_{t \ge 0} \log \pi_\theta (a_t | s_t)$ $\nabla_\theta \log p(\tau | \theta) = \nabla_\theta \sum_{t \ge 0} \log \pi_\theta (a_t | s_t) \quad \quad \quad (4)$

Можно сказать, что в равенстве (4) удалось избавиться от вероятностей, связанных с переходами между состояниями, которые не зависят от стратегии, т.е. агент не может на них повлиять. Собрав все множители под знаком суммы, мы получаем формулировку теоремы о градиенте стратегии

$\nabla_\theta J(\pi_{\theta}) = E_{\tau \sim \pi_{\theta}} \left[ \sum_{t \ge 0} R(\tau) \nabla_\theta \log \pi_\theta (a_t | s_t) \right] \quad \quad \quad (5)$

Тонкая настройка вознаграждений

На практике имеет смысл выполнить ещё одно преобразование, хотя оно уже и не будет тождественным. В выражении (5) каждое слагаемое содержит коэффициент $R(\tau)$ , зависящий от полной траектории. Т.е. информация о порядке действий не используется. Это можно исправить, если считать вознаграждения по отдельности для каждого временного шага . Для этого мы заменим $R(\tau)$ на $R_t(\tau)$ .

$\nabla_\theta J(\pi_{\theta}) = E_{\tau \sim \pi_{\theta}} \left[ \sum_{t = 0}^{T} R_t(\tau) \nabla_\theta \log \pi_\theta (a_t | s_t) \right]$

Получившиеся выражения (4) и (5) в общем случае не получится аналитически проинтегрировать, но на самом деле этого и не требуется. Основная ценность уравнений градиента заключается в том, что с их помощью можно определять направление, в котором следует изменить стратегию $\pi_\theta$ , чтобы увеличить совокупное вознаграждение. Нужно просто собирать информацию о пройденных траекториях и пересчитывать значения суммы для всех шагов.

Последнее замечание касается множителя $\nabla_\theta \log \pi_\theta (a_t | s_t)$ . Т.к. $\pi_\theta$ является по сути нейросетью, возвращающей вероятностные распределения для действий на каждом шаге, значит с вычислением градиента способен справиться любой DL-фреймворк.

Пример алгоритма

В настоящее время немало продвинутых RL-алгоритмов в том или ином виде опираются на градиент стратегии. Для примера взглянем на простейший алгоритм REINFORCE, использующий для обучения только актуальный опыт, то есть текущую траекторию.

def reinforce(env, pi, n_episode, gamma=1.0):
   """
   Алгоритм REINFORCE
   @param env: имя среды Gym
   @param pi: сеть, аппроксимирующая стратегию
   @param n_episode: количество эпизодов
   @param gamma: коэффициент обесценивания
   """
   for episode in range(n_episode):
     log_probs = []
     rewards = []
     state = env.reset()
     while True:
       action, log_prob = pi.get_action(state)
       next_state, reward, is_done, _ = env.step(action)
       log_probs.append(log_prob)
       rewards.append(reward)
       if is_done:
          returns = []
          Gt, k = 0, 0
          for reward in rewards[::-1]:
             Gt += gamma ** k * reward
             k += 1
             returns.append(Gt)
       returns = torch.tensor(returns)			
       pi.update(returns, log_probs)

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

Захватывающая история о градиенте стратегии в RL