Полезные ссылки:
Тренировки. Лекция 5: Современные методы обучения с подкреплением. Advantage actor critic, RLHF
Practical RL: Policy gradient methods
Policy Gradient – Федор Ратников

Тест ниже представляет собой агрегацию лекций, на которые даны ссылки выше.

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

Перед прочтением можно ознакомиться с предыдущими статьями Intro Reinforcement Learning и Reinforcement Learning: Model-free & Deep RL.

В первой статье введены основные понятия Reinforcement Learning и описаны - и -функцийи. Напомню, что V-функция – это математическое ожидание награды для текущего состояния, а Q-функция – это математическое ожидание награды для текущего состояния и выбранного действия. Значения V-функции мы можем найти либо решив матричное уравнение (выразив из уравнения Беллмана), либо используя метод простой итерации. Q-функцию можно выразить из V-функции через то же уравнение Беллмана. А оптимальная политика – это действия ведущие в максимальное значение Q-функции.

Вторая статья является продолжением Intro RL.
Во-первых, усложняется среда или просто говорится, что мы не имеем к ней доступа. Тогда, чтобы найти оптимальную политику было предложено учиться на целых траекториях "состояние->действие->...->конец" (метод Монте-Карло) или на кусочках траекторий (Temporal Difference learning). Если траектории минимальны , то такое обучение называется Q-learning (потому что мы предсказываем значение Q-функции).
Во-вторых, было замечено, что вместо табличного хранения Q-функции можно обучать параметрическую функцию $Q_\theta(s, a)$ (линейную регрессию или нейронную сеть) и предложен алгоритм обучения.
Заметим, что оптимальная политика, как и в первой части, это действия ведущие в максимальное значение Q-функции.

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

1. Политика

1.1. Детерминированная и стохастическая политика

Политика — это “способ выбирать действие”.

$\pi(a|s) = \text{вероятность выбрать действие } a \text{ в состоянии } s$

Когда политика выбирает только одно действие (например, $a = \arg \max Q(s,a)$ ) она называется детерминированной.

Стохастическая политика $a \sim \pi_{\theta}(a|s)$ задаёт распределение действий (например, $\varepsilon$ -greedy или softmax).

Можно заметить, что:

Cтохастическая политика обобщает детерменистическую.
Cуществуют задачи, где стохастическая политика оказывается оптимальнее детерменистической (например, в игре "Камень, ножницы, бумага" оптимальной окажется случайная равновероятная политика).
Иногда политика в виде вероятностей над действиями просто может оказаться более удобной.

По этим причинам далее будем использовать именно cтохастическую политику.

1.2. Как выучиваем стохастическую политику

Когда количество действий ограничено, очевидно, что выходной слой нейронной сети должен возвращать вектор логитов размером (число возможных действий), и применять softmax для преобразования в распределение вероятностей: $\pi_{\theta}(a|s) = softmax(Wh_{\theta}(s))$ .

В случае, когда количество действий не ограничено, политика должна выучивать параметры какого-то вероятностного распределения (например, нормального).

Ранее мы выучивали оптимальные действия косвенно через полезности $Q_{\theta}(s,a)$ или $V_{\theta}(s)$ . Такие методы называются Value based. Сейчас мы попытаемся явно выучивать политику, минуя промежуточные решения. Такие методы называют Policy based.

2. Оптимизация политики (или максимизация ожидаемой награды)

Оптимизируя политику, мы пытаемся максимизировать один из следующих функционалов:

$J = \mathbb{E}_{\substack{s \sim p(s) \\ a \sim \pi_{\theta}(a|s)}} R(s,a,s',a',...)$ – математическое ожидание награды.
$J = \mathbb{E}_{\substack{s \sim p(s) \\ a \sim \pi_{\theta}(a|s)}} G(s,a)$ , где $G(s,a) = r + \gamma \cdot G(s', a')$ – математическое ожидание дисконтированной награды.

Для простоты рассуждений, мы (пока что) будем рассматривать первый случай (математическое ожидание награды) и обойдемся одношаговым процессом $s \rightarrow a \rightarrow r$ (когда одно действие нас сразу приведет к результату).

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

$J = \underset{\substack{s \sim p(s) \ a \sim \pi_{\theta}(s|a)}}{\mathbb{E}} R(s,a,s',a',...) = \int_s p(s) \int_a \pi_{\theta}(a|s) R(s,a) \mathrm{d}a \mathrm{d}s$

где – вероятность посещений состояния (может зависеть от политики), $\pi_{\theta}(a|s)$ – вероятность выбрать действие в состоянии (политика), – награда одношаговой траектории.

Мы хотим оптимизировать $\pi_{\theta}(a|s)$ чтобы максимизировать . Оптимизировать, конечно, хотелось бы градиентным спуском и для этого надо понять как продифференцировать . Ниже рассмотрим разные способы.

Способ 1: метод Монте-Карло.

Можно попробовать оценить случайную величину методом Монте-Карло (проиграть сессий и усреднить награду) и продифферинцировать Монте-Карло по $\theta$ :

$J \approx \frac{1}{N} \sum_{i=0}^{N} \sum_{s,a \in z_i} R(s,a)\underset{\substack{\text{1-step} \\ \text{process}}}{=}\frac{1}{N} \sum_{i=0}^{N} R(s,a)$

Однако, заметим, что при замене интеграла на выборочное среднее пропадают параметры политики $\theta$ . А значит, невозможно посчитать $\nabla_{\theta} J$ .

Способ 2: метод конечных разностей.

Можно попытаться использовать метод конечных разностей: аппроксимировать градиент путем небольшого возмущения каждого параметра и оценки изменения функции:

$\nabla J = \frac{J_{\theta + \varepsilon} - J_{\theta}}{\varepsilon}$

Этот подход может работать на достаточно гладких функциях, однако для негладких (например, не удовлетворяющих условию Липшица) или зашумленных функций оценка становится крайне ненадежной. Так же, его вычислительная сложность растет линейно с размерностью пространства параметров, что делает его непрактичным для современных моделей с миллионами весов.

Способ 3: стохастическая оптимизация (метод кросс-энтропии)

В первой статье Intro RL был рассмотрен простой эволюционный способ прямой оптимизации политики — метод кросс энтропии (CEM). Его алгоритм выглядит так:

Инициализируем веса нейронной сети (политики) случайным образом $\theta_0 \leftarrow \text{random}$ .
Основной цикл обучения:
1. Генерируем траекторий (сессий), играя эпизоды с текущей политикой.
2. Выбираем лучших сессий с наибольшей наградой $\text{Elite} = [(s_0, a_0), (s_1, a_1), \ldots, (s_M, a_M)]$
3. Обновляем параметры политики, чтобы она повторяла действия, которые встречались в лучших эпизодах: $\theta_{i+1} = \theta_{i} + \alpha \nabla \sum_i \log \pi_{\theta_i}(a_i|s_i) \cdot [s_i, a_i \in \text{Elite}]$

В CEM элита выбирается по эпизодам.

Существует также TD-версия Cross-Entropy Method (CEM), в которой элита выбирается не по целым эпизодам, а по состояниям. В этом подходе “элитной” считается не вся траектория, а лишь конкретная пара (s, a), если она приводит к высокому будущему возврату G(s, a). Такой способ отбора оказывается особенно полезным в ситуациях, когда эпизод в целом оказался неудачным, но отдельные решения внутри него были оптимальными. В отличие от классического CEM, где плохая траектория отбрасывается целиком, TD-CEM сохраняет ценные фрагменты поведения. Благодаря этому обучение становится быстрее, стабильнее и менее зависимым от общей удачи/неудачи эпизода, поскольку полезная информация не теряется.

Однако у стохастической оптимизации есть серьезные недостатки:

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

Способ 4: дифференцирование интеграла

Выпишем еще раз оптимизируемый функционал:

$J = \int_s p(s) \int_a \pi_{\theta}(a|s) R(s,a) \mathrm{d}a \mathrm{d}s \rightarrow \max$

Хочется аналитически выписать его градиент, чтобы затем заменить интегралы на выборочные средние (т.е. приблизить методом Монте-Карло):

$[1] \quad \nabla J = \int_s p(s) \int_a \nabla \pi_{\theta}(a|s) R(s,a) \mathrm{d}a \mathrm{d}s$

Однако, градиент плотности $\nabla \pi_{\theta}(a|s)$ не будет обладать свойствами вероятностного распределения (например, может стать отрицательным или больше единицы). Поэтому его невозможно использовать как вес при оценке математического ожидания сэмплами. А значит, мы опять не может оценить $\nabla J$ .

Для того чтобы это исправить используют логарифмический трюк (logderivative trick):
Известно, что

$\nabla \log \pi(z) = \frac{1}{\pi(z)} \cdot \nabla \pi(z) \quad \text{,} \quad \pi(z) > 0$

Тогда можно выразить градиент политики:

$[2] \quad \pi(z) \cdot \nabla \log \pi(z) = \nabla \pi(z)$

Это тождество просто переписывает градиент плотности в виде плотность × градиент лог-плотности. И вот это уже важно: в левой части стоит сама плотность $\pi$ , которую можно использовать как вес при семплировании.

Применим логарифмический трюк к нашем :

$\nabla J = \int_s p(s) \int_a {\color{red}{\nabla \pi_{\theta}(a|s)}} R(s,a) \mathrm{d}a \mathrm{d}s = \int_s p(s) \int_a {\color{red}{\pi_{\theta}(a|s) \nabla \log \pi_{\theta}(a|s)}} R(s,a) \mathrm{d}a \mathrm{d}s$

Данная формула носит называется Policy Gradient и с ее помощью можно оценить $\nabla J$ методом Монте-Карло (Monte Carlo Policy Gradient):

$\nabla J = \frac{1}{N} \sum^{N}_{i=0} \nabla \log \pi_{\theta}(a|s) \cdot R(s,a)$

Остается обновить параметры политики градиентным подъемом: $\theta_{i+1} \leftarrow \theta_i + \alpha \nabla J$ .

Так мы получили алгоритм REINFORCE для одношаговой игры:

Инициализируем веса нейронной сети случайным образом θ ← random
Основной цикл обучения: 
  Семплируем N сессий
  Вычисляем Монте-Карло оценку градиента ∇J = 1/N Σ ∇log π(a|s) R(s,a)
  Обновляем параметры θ ← θ + alpha ∇J

3. Алгоритм REINFORCE

Policy Gradient Methods for Reinforcement Learning with Function Approximation

В многошаговой постановке агент получает не одну награду, а последовательность наград. Поэтому, вместо мгновенной награды, рассматривают дисконтированную сумму будущих наград: $G_t = r + \gamma \cdot G(s', a')$ .

Тогда оптимизируемый функционал: $J = \underset{\substack{s \sim p(s) \ a \sim \pi_{\theta}(a|s)}}{\mathbb{E}}G(s,a) \rightarrow max_{\pi}$ .

Аналогично, применим к градиенту логарифмический трюк:

$\nabla J = \int_s p(s) \int_a \nabla \pi_{\theta}(a|s) {\color{red} Q(s,a)} \mathrm{d}a \mathrm{d}s \underset{\substack{\text{logderivative} \\ \text{trick}}}{=}$

$\underset{\substack{\text{logderivative} \\ \text{trick}}}{=} \int_s p(s) \int_a \pi_{\theta}(a|s) \nabla \log \pi_{\theta}(a|s) {\color{red} Q(s,a)} \mathrm{d}a \mathrm{d}s$

Монте-Карло оценка градиента:

$\nabla J = \frac{1}{N} \sum_{i=0}^{N} \sum_{a \in z_i} \nabla \log \pi_{\theta}(a|s) \cdot Q(s,a)$

Получили классический алгоритм REINFORCE (Williams, 1992):

Инициализируем веса нейронной сети случайным образом θ ← random
Основной цикл обучения: 
  Семплируем N эпизодов текущей политикой π(a|s), где каждый запуск возвращает траекторию (эпизод): z_i=(s_0,a_0,r_0,s_1,a_1,r_1,…,s_T,a_T,r_T) 
  Вычисляем Монте-Карло оценку градиента целевой функции ∇J = 1/N Σ ∇log π(a|s) R(s,a)
  Обновляем параметры политики (gradient ascent) θ ← θ + alpha ∇J

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

Еще несколько особенностей алгоритма REINFORCE:

Это on-policy алгоритм (обучение происходит по текущей политике).
Работает быстрее, чем ванильный -learning.
Градиент политики можно подавать на вход любому современному оптимизатору (SGD, Adam, RMSProp), что обеспечивает устойчивость и скорость сходимости.
Совместим с любыми архитектурами нейросетей (свёрточные сети, рекуррентные сети или трансформеры).
Концептуально похож на обучение с учителем.

Рассмотрим последний (5) пункт подробнее.

4. Обучение с учителем & Policy Gradient

Вспомним, что в задачах обучения с учителем (Supervised Learning), мы пытались максимизировать логарифм функции правдоподобия (log likelihood, llh). То есть найти такое $\theta$ , для которого вероятность выбора действий $a={a_1,…,a_N}$ была бы максимальной (4.1. Вероятностный подход в ML):

$llh = \sum \log \pi_{\theta}(a_{opt}|s, \theta) \rightarrow max_{\theta}$ .

Градиент логарифма правдоподобия: $\nabla llh = \underset{s,a \sim D}{\mathbb{E}}\nabla \log \pi_{\theta}(a_{opt}|s, \theta)$ .

Взглянем еще раз на градиент математического ожидания суммарной дисконтированной награды:

$\nabla J = \underset{\substack{s \sim d(s) \\ a \sim \pi(a|obs(s))}}{\mathbb{E}}\nabla \log \pi(a|s) Q(s,a)$

Не трудно заметиться, что $\nabla llh$ и $\nabla J$ похожи с точностью до .

Можно сказать, что в $\nabla llh$ сидит награда в виде индикаторной функции которая равна единице если объекты принадлежат обучающей выборке, а в противном случае ноль:

$\nabla llh = \underset{s,a \sim D}{\mathbb{E}}\nabla \log \pi_{\theta}(a_{opt}|s, \theta) {\color{red}{Q(s,a)}}$ $Q(s,a) = \mathbb{I}(s, a) = \begin{cases} 1 & \text{if } s, a \in D \\0 & \text{else}\end{cases}$

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

Максимизируя математическое ожидание награды в RL мы можем положить любую интересующую нас функцию награды (метрику или обученную модель rewordа), тем самым придав разную важность примерам.

Таким образом, задачу supervised learning (грубо говоря) можно рассматривать как частный частный случай оптимизации ожидаемой награды, где награда даётся за правильный ответ.

Преимущества supervised learning заключаются в том, что обучение в этой постановке гораздо проще: нет траекторий, нет стохастических политик, и каждая пара «вход–ответ» рассматривается независимо. Благодаря этому обучение быстро сходится, градиенты более стабильны, а оценка ошибки обладает низкой дисперсией. Основные недостатки такого обучения - необходимость в хорошо размеченной выборке и проблема смещения распределений (distribution shift): модель обучается на одном распределении данных, а применять её приходится на другом, что приводит к деградации качества.

В обучении с подкреплением данные разметки не требуются: достаточно иметь состояния , действия и сигнал награды . Агент сам генерирует данные в процессе взаимодействия со средой, поэтому отсутствует distribution shift — распределение данных всегда соответствует текущей политике. К недостаткам относятся проблема холодного старта, когда агент поначалу действует почти случайно и получает слабый обучающий сигнал, а также большой разброс оценок: стохастичность среды и траекторий приводит к высокой дисперсии градиентов и нестабильному обучению.

На практике модель сначала предварительно обучают supervised-методами, чтобы задать хорошую инициализацию, а затем дооптимизируют с помощью RL.

5. Baselines и A2C (Advantage Actor-Critic)

Рассмотрим одношаговую игру, в которой агенту доступно всего два действия: выбрать левую или правую кнопку. Каждое действие немедленно приводит к завершению эпизода и выдаёт фиксированную награду: награда за левую кнопку +1 и награда за правую кнопку -1. Таким образом, оптимальная стратегия очевидна: вероятность того что агент выберет левую кнопку будет расти, а вероятность выбора правой кнопки будет падать.

Немного поменяем условия задачи выше: прибавим к награде +100. Теперь награда за левую кнопку +101, а награда за правую кнопку 100. Не смотря на то что условие задачи не поменялось, вероятность выбора и левой и правой кнопки будет возрастать, хоть левой немного сильнее.

И левая и правая кнопка принесут большую награду

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

Можно выучить некоторую и вычесть его из .

$\nabla J = \underset{\substack{s \sim p(s) \\ a \sim \pi(a|s)}}{\mathbb{E}}\nabla \log \pi(a|s) (Q(s,a) - {\color{red}{b(s)}})$

Опустив все формулы, сразу скажу вывод: baseline не меняет математическое ожидание градиента (оно остаётся корректным), но при этом влияет на его дисперсию. Если подобрать удачно, можно значительно уменьшить разброс Монте-Карло оценок и заметно ускорить обучение. При этом, известный практический факт, что не учится. Но можно показать, что если вместо использовать математическое ожидание награды, то дисперсия будет меньше (хоть и не минимальной).

Алгоритмы, которые учат не только вероятность действий $\pi(a|s)$ , но и оценку полезности состояния V(s), называются actor–critic.

В данном месте возникает довольно ироничная ситуация: мы хотели обойтись без обучения v- и q-функций, но для эффективного обучения в итоге приходится снова вводить v(s) - чтобы оценивать преимущество.

В таких моделях оценка градиента политики:

$\nabla J = \frac{1}{N} \sum_{i=0}^{N} \sum_{s,a \in z_i} \nabla \log \pi(a|s) ( {\color{red}{Q(s,a) - V(s)}})$

использует advantage — насколько действие лучше или хуже ожидаемого результата в этом состоянии. Если действие хуже обычного — градиент уменьшает вероятность. Если лучше — увеличивает.

Заметим, что если мы знаем и кортеж можно оценить по формуле:

$Q(s,a) = r + \gamma V(s')$

Обозначим разницу как (от слова advantage):

$A(s,a) = Q(s,a) - V(s) = r + \gamma V(s') - V(s)$

Модель, которая учится делать действия называется Actor.
Ее функция потерь:

$J_{actor} = \int_s p(s) \int_a \pi(s|a) A(s,a) da ds \rightarrow \max$

И градиент:

$\nabla J_{actor} \approx \frac{1}{N} \sum_{i=0}^{N} \sum_{s,a \in z_i} \nabla \log \pi (a|s) A(s,a)$

Модель, которая оценивает качества состояний/действий (критикует действие) называется Critic. Чтобы обучить Critic пользуемся идеей из Q-learning:

$L_{critic} \approx \frac{1}{N} \sum_{i=0}^{N} \sum_{s,a \in z_i} \big( V(s) - \left[ r + \gamma \cdot V(s')\right] \big) ^2 \rightarrow \min$

Немного практических советов:

Важно правильно подобрать коэффициенты перед компонентами loss-функции.
Исследование среды можно делать с помощью регуляризации, а именно максимизации энтропи политики.
Собирайте траектории на параллельных сессиях, в нескольких средах одновременно.
Вместо log(softmax(logits)) используйте log_softmax(logits) для численной стабильности

Резюмируя

Мы хотели избавиться от обучения value-функций, обучая политику напрямую через REINFORCE. Однако REINFORCE страдает от высокой дисперсии градиентов, что замедляет обучение. Для стабилизации пришлось ввести baseline, который по сути является оценкой V(s). В итоге появился Actor-Critic, где Critic учит именно V-функцию для оценки advantage. Таким образом, от value-функций избавиться не удалось — они вернулись как необходимый инструмент для эффективного обучения.

В следующей статье мы продолжим улучшать алгоритмы RL и разберём более продвинутые подходы — Trust Region Policy Optimization (TRPO) и Proximal Policy Optimization (PPO).

Reinforcement Learning: Policy gradient methods