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

Алгоритмы REINFORCE и A2C (Advantage Actor–Critic) относятся к классу on-policy алгоритмов: политика обучается только на данных, собранных текущей версией агента. После каждого обновления старые траектории не могут быть повторно использованы. Естественное желание – обучаться более оптимально, выполняя несколько шагов оптимизации на одном и том же наборе траекторий.

Эта идея была реализована в алгоритме TRPO (Trust Region Policy Optimization). Однако, он требует сложной оптимизации второго порядка и плохо масштабируется. Но его идеи легли в основу алгоритма PPO (Proximal Policy Optimization) – более простого с точки зрения реализации и более эффективного на практике. Благодаря этому PPO стал стандартным инструментом для обучения агентов в задачах управления, игровых средах, робототехнике и в RLHF при обучении больших языковых моделей.

Если вы заметите ошибки в формулах или опечатки — буду очень благодарна, если укажете на них 🙈.

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

В конце-концов будет сформулирован алгоритм PPO в виде псевдокода. Можно попробовать его реализовать в ноутбуке курса Practical RL. Агентом в ноутбуке будет "наполовину гепард", целью которого будет бег вперед (направо) как можно быстрее.

Мой тг канал: 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

Пусть мы имеем траектории $(s_0, a_0, s_1, a_1, \cdots)$ , гдe– состояния,– действия агента.

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

$J(\pi_{\theta}) = \mathbb{E}_{\tau \sim \pi} \sum_{t=0} \gamma^{t} r_t \rightarrow \max \quad [1.1]$

Траектории $\tau$ приходят из $p(\tau|\pi_{\theta})$ – распределение над траекториями $\tau$ при условии политики $\pi_{\theta}$ :

$p(\tau|\pi_{\theta}) = p(s_0) \cdot \pi_{\theta}(a_0, s_0) \cdot p(s_1|s_0, a_0) \cdots = p(s_0) \prod_{t=0}^{\infty} \pi_{\theta}(a_t, s_t) p(s_{t+1}|s_{t}, a_{t})$

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

$J(\pi_{\theta}) = \mathbb{E}_{s_0\sim p(s_0)} V^{\pi_{\theta}}(s_0) \quad [1.2]$

Мы научились максимизировать функционал $J(\pi_{\theta})$ с помощью градиентного подъема. Градиент $\nabla J(\pi_{\theta})$ находим с помощью Монте-Карло оценки с применением logderivative trick $\nabla p(\tau| \pi_{\theta}) = p(\tau|\pi_{\theta}) \cdot \nabla \log p(\tau|\pi_{\theta})$ :

$\nabla J(\pi_{\theta}) = \nabla \int_{\tau} p(\tau|\pi_{\theta}) \left[ \sum_t \gamma^t r_t \right] d\tau =\mathbb{E}_{\tau \sim \pi} \sum_t \gamma^{t} \cdot\nabla \log \pi(a_t|s_t) \cdot Q(s_t,a_t) \quad [1.3]$

Это и есть Policy Gradient Theorem.

Вывод [1.3]

$\nabla J(\pi_{\theta}) = \nabla \int_{\tau} p(\tau|\pi_{\theta}) \left[ \sum_t \gamma^t r_t \right] d\tau =$ $= \int {\color{red}{\nabla p(\tau|\pi_{\theta})}} \left[ \sum_t \gamma^t r_t \right] d\tau \underset{\substack{\text{logderivative} \\ \text{trick}}}{=}$ $\underset{\substack{\text{logderivative} \\ \text{trick}}}{=}\int {\color{red}{p(\tau|\pi_{\theta}) \cdot \nabla \log p(\tau|\pi_{\theta})}} \left[ \sum_t \gamma^t r_t \right] d\tau$

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

$\nabla J(\pi_{\theta}) =\mathbb{E}_{\tau \sim \pi} \nabla \log p(\tau|\pi) \sum_t \gamma^t r_t =$ $= \mathbb{E}_{\tau \sim \pi} \left[ \sum_{t=0} \nabla \log \pi(a_{t}|s_{t}) \right]\left[ \sum_{t=0} \gamma^{t} r_{t} \right]=$

Перемножим скобки $\left[ \sum_{t=0} \nabla \log \pi(a_{t}|s_{t}) \right]$ и $\left[ \sum_{t=0} \gamma^{t} r_{t} \right]$ между собой (занесем под одну скобку):

$= \mathbb{E}_{\tau \sim \pi} \sum_{t=0}^{T} \left[ \nabla \log \pi(a_{t}|s_{t}) \left[ \sum_{t'=0}^{T} \gamma^{t'} r_{t'} \right] \right] =$

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

$= \mathbb{E}_{\tau \sim \pi} \sum_{t} \left[ \nabla \log \pi(a_{t}|s_{t})\left[ {\color{red}{\sum_{t' < t} \gamma^{t'} r_{t'}}} + \sum_{t' \geq t} \gamma^{t'} r_{t'} \right] \right] \quad [1]$

где $\mathbb{E}_{a_t} \nabla \log \pi(a_{t}|s_{t}) \left[ \sum_{t' < t} \gamma^{t'} r_{t'} \right] = 0$

Почему градиент по "прошлым" состояниям зануляется

Награда $R_{t'<t} = {\color{red}{\sum_{t' < t} \gamma^{t'} r_{t'}}}$ к моменту времени уже произошла и не зависит от того, какое семплируется. Поэтому $R_{t'<t}$ – константа, при фиксированном .

Рассмотрим часть формулы c $R_{t'<t}$ :

$\mathbb{E}_{a_t} \nabla \log \pi(a_{t}|s_{t}) \cdot R_{t'<t} = R_{t'<t} \cdot \mathbb{E}_{a_t} \nabla \log \pi(a_{t}|s_{t}) =$

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

$=R_{t'<t} \cdot \sum_{a_t} \pi(a|s_t) \nabla \log \pi(a_{t}|s_{t}) =$

Применим "обратный" logderivative trick: $\nabla \log \pi(a_{t}|s_{t}) = \frac{\nabla \pi(a_{t}|s_{t})}{\pi(a_{t}|s_{t})}$

$= R_{t'<t} \cdot \sum_{a_t} \pi(a|s_t) \frac{\nabla \pi(a_{t}|s_{t})}{\pi(a_{t}|s_{t})} = R_{t'<t} \cdot \sum_{a_t \sim \pi} \nabla \pi(a_{t}|s_{t})$ $= R_{t'<t} \cdot \nabla \sum_{a_t} \pi(a_{t}|s_{t}) =$

По определению политики $\sum_{a_t \sim \pi} \pi(a_{t}|s_{t}) = 1$

$= R_{t'<t} \cdot \nabla \sum_{a_t} \pi(a_{t}|s_{t}) = R_{t'<t} \cdot \nabla 1 = 0$

Таким образом: $\mathbb{E}[\nabla \log \pi(a_t∣s_t) \cdot R_{t'<t}]=0$

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

$[1] = \mathbb{E}_{\tau \sim \pi} \sum_{t} \left[ \nabla \log \pi(a_{t}|s_{t})\left[\sum_{t' \geq t} \gamma^{t'} r_{t'} \right] \right] =$

Вынесем $\gamma^t$ из суммы $\sum_{t' \geq t} \gamma^{t'} r_{t'} = \gamma^t \sum_{t' = 0}^{\infty} \gamma^{t'} r_{t'+t}$

$= \mathbb{E}_{\tau \sim \pi} \sum_{t} \left[ \nabla \log \pi(a_{t}|s_{t}) \cdot \gamma^{t} \cdot \sum_{t'=0}^{\infty} \gamma^{t'} r_{t'+t} \right]$

Согласно равенству $\mathbb{E}[Y] = \mathbb{E}[ \mathbb{E}[Y∣X]]$ :

$= \mathbb{E}_{\tau \sim \pi}\sum_t\gamma^t \nabla \log \pi(a_t|s_t){\color{red}{\mathbb{E}\left[\sum_{t'=0}^{\infty} \gamma^{t'} r_{t'+t}\;\middle|\;s_t, a_t\right]}} =$

где $\mathbb{E} \left[ \sum_{t'=0}^{\infty} \gamma^{t'} r_{t'+t} | s_t, a_t\right]$ – условное математическое ожидание возврата при заданных текущих состоянии и действии .

Заметим, что это определение -функции: $Q(s_t, a_t)= \mathbb{E} \left[ \sum_{t'=0}^{\infty} \gamma^{t'} r_{t'+t} | s_t, a_t \right]$ . Тогда:

$\mathbb{E}_{\tau \sim \pi} \sum_{t} \left[ \gamma^{t} \cdot \nabla \log \pi(a_{t}|s_{t}) \cdot Q(s_{t}, a_{t}) \right]$

Доказали Policy Gradient Theorem:

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

$\nabla J(\pi_{\theta}) =\mathbb{E}_{s, a \sim \pi} \sum_t \gamma^{t} \nabla_{\theta} \log \pi(a|s) \cdot \left( Q (s,a) - {\color{red}{baseline}}\right)$

1.2. Baseline и Advantage

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

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

$\nabla J(\pi_{\theta}) =\mathbb{E}_{s, a \sim \pi} \sum_t \gamma^{t} \nabla_{\theta} \log \pi(a|s) \cdot {\color{red}{ \left( Q (s,a) - V(s)\right)}}$

где – обучаем, а будем получать из по формуле $Q(s_t,a_t) = r_t + \gamma V(s_{t+1})$ .

называется advantage функция. Она показывает насколько действие лучше или хуже, чем среднее действие в этом состоянии при политике $\pi$ . Если , то действие лучше среднего выбираемого политикой, иначе - хуже.

$\nabla J(\pi_{\theta}) =\mathbb{E}_{s, a \sim \pi} \sum_t \gamma^{t} \nabla_{\theta} \log \pi(a|s) \cdot {\color{red}{ A (s,a)}} \quad [1.4]$

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

$\nabla J(\pi_{\theta}) =\frac{1}{1-\gamma} \cdot \underbrace{ \mathbb{E}_{s \sim d_{\pi}(s)} \mathbb{E}_{a \sim \pi(a|s)}}_{\text{data generated by }\pi_\theta}\nabla_{\theta} \log \pi(a|s) \cdot A(s,a) \quad [1.5]$

Математическое ожидание берётся по данным, сгенерированным самой $\pi_\theta$ .
$d_{\pi{_\theta}}(s)$ – discounted state-visitation distribution политики. Это распределение состояний, которое показывает как часто политика $\pi$ посещает состояние , с учётом дисконтирования по времени.

$d_{\pi{_\theta}}(s) = (1−\gamma) \sum_{t=0}^{\infty} \gamma^t P(s_t=s | \pi)$

$P(s_t=s|\pi)$ — вероятность быть в состоянии на шаге .
$(1-\gamma)$ – нормировка, чтобы это было распределение $\sum_s d_\pi(s) = 1$

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

При этом, сам градиентный спуск не такой дорогой как сбор данных (траекторий). Естественным образом возникает желание оптимизировать полезность политики $\pi_{\theta_{new}}$ , используя данные, собранные старой политикой $\pi_{\theta_{old}}$ . В этом заключается основная цель дальнейших методов (TRPO, PPO и др.)

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

Для этого в формуле надо придумать как брать математическое ожидание и advantage по старой политике $\pi_{\theta_{old}}$ .

Дальше будет много формул по их замене. Если хотите их скипуть (или не поняли что произошло ниже), то прочитайте "скрытый текст" и переходите к разделу 3. Проблема оптимизации

Скрытый текст

В тексте ниже, вместо оптимизации $J(\pi_{\theta}) \rightarrow \max$ , будет оптимизироваться $J^{\pi^{new}} - J^{\pi^{old}} \rightarrow \max$ .

В новом функционале удастся взять advantage по старой политике, используя трюк с телескопическими суммами. Удастся заменить математическое ожидание по действиям с помощью Importance Sampling трюка. Но математическое ожидание $d_{\pi{_\theta}}(s)$ по старой политике взять не получится.

Поэтому, вместо оптимизации $J^{\pi^{new}} - J^{\pi^{old}} \rightarrow \max$ будем "не честно" оптимизировать суррогатную функцию $L_{\pi^{old}}(\theta) \rightarrow \max$ .

$L_{\pi^{old}}(\theta) =\frac{1}{1-\gamma} \mathbb{E}_{s \sim d_{\pi^{old}}(s)} \mathbb{E}_{a \sim \pi^{old}(a|s)}\frac{\pi^{new}(a|s)}{\pi^{old}(a|s)}A^{\pi^{old}}(s, a)\rightarrow \max$

При этом, разница функционалов $|J(\pi^{new}) - J(\pi^{old}) - L_{\pi^{old}}(\theta)|$ остается не велика, если разница между политиками $\pi^{new}$ и $\pi^{old}$ ограничена-дивергенцией.

Так возникает задача максимизации $L_{\pi^{old}}(\theta) \rightarrow \max$ с ограничением $KL(\pi(\theta^{old})|| \pi(\theta^{old})) \leq \delta$ .

Как оптимизировать такую задачу, смотрите в проблемах оптимизации.

2.1. Заменяем advantage

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

$J^{\pi^{new}} - J^{\pi^{old}} \rightarrow \max \quad [2.1]$

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

И прежде чем расписывать новый функционал, проделаем "подготовительную" работу:
1) Представим $-V^{\pi}(s_0)$ в виде телескопические суммы.
2) Распишем разность $V^{\pi^{new}}(s) - V^{\pi^{old}}(s)$ .
3) Распишем новый функционал $J^{\pi^{new}} - J^{\pi^{old}}$ .

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

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

$\sum_{k=1}^{n} a_{k} - a_{k+1} = a_1 - a_2 + a_2 - a_3 + a_3 - a_4 + \cdots = a_1 - a_{k+1}$

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

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

Мы будем пользоваться телескопической суммой для любой последовательности , такой что $\lim_{t \rightarrow \infty} a_t = 0$ :

$\sum_{t \geq 0}^{\infty} (a_{t+1} - a_{t}) = - a_0$

По аналогии с формулой выше: для произвольной траектории $\mathcal{T} = s_0, a_0, s_1, a_1, \cdots$ и произвольной политики $\pi$ справедливо равенство:

$\sum_{t \geq 0 } \left[ \gamma^{t+1} V^{\pi}(s_{t+1}) - \gamma^t V^{\pi}(s_t) \right] = -V^{\pi}(s_0) \quad [2.2]$

при условии, что $\lim_{t \rightarrow \infty} \gamma^t V^{\pi}(s_t) = 0$ .

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

$V^{\pi^{new}}(s) - V^{\pi^{old}}(s) =$

$V^{\pi}$ по определению:

$= \mathbb{E}_{\tau \sim \pi^{new} | s_0 = s} \sum_{t \geq 0} \gamma^t r_t - V^{\pi^{old}}(s) =$

Занесем $V^{\pi^{old}}(s)$ под математическое ожидание:

$= \mathbb{E}_{\tau \sim \pi^{new} | s_0 = s} \left[ \sum_{t \geq 0} \gamma^t r_t {\color{red}{ - V^{\pi^{old}}(s_0)}} \right] =$

Запишем $- V^{\pi^{old}}(s_0)$ в виде телескопической суммы по формуле :

$\underset{\text{telescoping sums}}{=}\mathbb{E}_{\tau \sim \pi^{new} | s_0 = s} \left[ \sum_{t \geq 0} \gamma^t r_t + {\color{red}{ \sum_{t \geq 0} \left[ \gamma^{t+1} V^{\pi^{old}}(s_{t+1}) - \gamma^t V^{\pi^{old}}(s_t) \right] }} \right]=$

Вынесем $\gamma^t$ за скобки

$=\mathbb{E}_{\tau \sim \pi^{new} | s_0 = s} \left[ \sum_{t \geq 0} {\color{green}{\gamma^{t}}} r_t + \sum_{t \geq 0} \left[ {\color{green}{\gamma^{t+1}}} V^{\pi^{old}} (s_{t+1}) - {\color{green}{\gamma^{t}}} V^{\pi^{old}}(s_t) \right] \right] =$ $=\mathbb{E}_{\tau \sim \pi^{new} | s_0 = s} \sum_{t \geq 0} {\color{green}{\gamma^t}}\left( r_t + \gamma V^{\pi^{old}} (s_{t+1}) - V^{\pi^{old}}(s_t) \right) =$

В силу равенства $\mathbb{E} f(X)=\mathbb{E}\big[\mathbb{E}[f(X)\mid Y]\big]$ , внес��м матожидание по будущим состояниям $s_{t+1}$ , условно на .

$\color{green}{\mathbb{E}_{\tau \sim \pi^{new}} [V^{\pi^{old}}(s_{t+1})] = \mathbb{E}_{s_t, a_t \sim \pi^{new}} \mathbb{E}_{s_{t+1}} [V^{\pi^{old}}(s_{t+1})]}$ $= \mathbb{E}_{\tau \sim \pi^{new} | s_0 = s} \sum_{t \geq 0} \gamma^t \left( r_t + \gamma {\color{green}{\mathbb{E}_{s_{t+1}}}} V^{\pi^{old}}(s_{t+1}) - V^{\pi^{old}}(s_t) \right) =$

Заметим, что $r_t + \gamma {\mathbb{E}_{s_{t+1}}} V^{\pi^{old}}(s_{t+1}) = Q^{\pi^{old}}(s_t, a_t)$ :

$= \mathbb{E}_{\tau \sim \pi^{new} | s_0 = s} \sum_{t \geq 0} \gamma^t \left( {\color{green}{r_t + \gamma {\mathbb{E}_{s_{t+1}}} V^{\pi^{old}}(s_{t+1})}} - V^{\pi^{old}}(s_t) \right) =$ $=\mathbb{E}_{\mathcal{T} \sim \pi^{new} | s_0 = s} \sum_{t \geq 0} \gamma^t \left( {\color{green}{ Q^{\pi^{old}}(s_t, a_t) }} - V^{\pi^{old}}(s_t) \right) =$

$Q^{\pi^{old}}(s_t, a_t) - V^{\pi^{old}}(s_t) = A^{\pi^{old}}(s_t, a_t)$ по определению advantage функции:

$= \mathbb{E}_{\mathcal{T} \sim \pi^{new} | s_0 = s} \sum_{t \geq 0} \gamma^t \left( {\color{green}{ Q^{\pi^{old}}(s_t, a_t) - V^{\pi^{old}}(s_t) }}\right)=$ $= \mathbb{E}_{\mathcal{T} \sim \pi^{new} | s_0 = s} \sum_{t \geq 0} \gamma^t {\color{green}{ A^{\pi^{old}}(s_t, a_t) }}$

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

$V^{\pi^{new}}(s) - V^{\pi^{old}}(s) = \mathbb{E}_{\mathcal{T} \sim \pi^{new} | s_0 = s} \sum_{t \geq 0} \gamma^t A^{\pi^{old}}(s_t, a_t) \quad [2.3]$

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

Вспомним формулу для $J(\pi_{\theta})$ :

$J(\pi_{\theta}) = \mathbb{E}_{s \sim p(s_0)} \sum_{t} \gamma^{t} r_t = \mathbb{E}_{s_0\sim p(s_0)} V^{\pi_{\theta}}(s_0)$

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

$J(\pi^{new}) - J(\pi^{old}) = \mathbb{E}_{s \sim p(s_0)} \left[ V^{\pi^{new}}(s_0) - V^{\pi^{old}}(s_0) \right]$

Подставим формулу , получим:

$J(\pi^{new}) - J(\pi^{old}) = \mathbb{E}_{\tau \sim \pi^{new} | s_0 = s} \sum_{t} \gamma^t A^{\pi^{old}}(s_t, a_t) \quad [2.4]$

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

$J(\pi^{new}) - J(\pi^{old})=\frac{1}{1 - \gamma}\;\mathbb{E}_{s \sim d_{\pi^{new}}(s)}\mathbb{E}_{a \sim \pi^{new}(a \mid s)}A^{\pi^{old}}(s, a) \quad [2.5]$

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

В получившийся формуле остается понять как заменить математические ожидания по новой политике $\pi^{new}$ на математические ожидания по старой политике $\pi^{old}$ .

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

Пусть у нас есть математическое ожидание функции по распределению :

$\mathbb{E}_{x \sim p(x)} f(x)= \int f(x)p(x)dx \quad [2.6]$

Теперь мы хотим выразить через другое распределение $\mu(x)$ . Чтобы это сделать, просто умножаем и делим под интегралом на $\mu(x)$ .

$\mathbb{E}_{x \sim p(x)} f(x) = \int f(x)p(x)dx = \int f(x) \frac{p(x)}{\color{red}{\mu(x)}} {\color{red}{\mu(x)}} dx = \mathbb{E}_{\color{red}{x \sim \mu(x)}} \frac{p(x)}{{\color{red}{\mu(x)}}}f(x) \quad [2.7]$

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

$\mathbb{E}_{a \sim \pi^{new}(a \mid s)} A^{\pi^{old}}(s, a) =\mathbb{E}_{a \sim \pi^\text{old}(a|s)} \frac{\pi^{new}(a|s)}{\pi^{old} (a|s)} A^{\pi^{old}}(s, a) \quad [2.8]$ $J(\pi^{new}) - J(\pi^{old})=\frac{1}{1 - \gamma} \mathbb{E}_{s \sim d_{\pi^{new}}(s)} \mathbb{E}_{a \sim \pi^\text{old}(a|s)} \frac{\pi^{new}(a|s)}{\pi^{old} (a|s)} A^{\pi^{old}}(s, a)$

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

Для $\mathbb{E}_{s \sim d_{\pi^{new}}(s)}$ Importance Sampling трюк сделать не получится. В таком случае, давайте просто заменим $\mathbb{E}_{s \sim d_{\pi^{new}}(s)}$ на $\mathbb{E}_{s \sim d_{\pi^{old}}(s)}$ и назовем получившийся функционал суррогатной функцией $L_{\pi^\text{old}}(\theta)$ .

$J(\pi^{new}) - J(\pi^{old}) \approx L_{\pi^\text{old}}(\theta)$ $L_{\pi^\text{old}}(\theta) = \frac{1}{1-\gamma} \underbrace{\mathbb{E}_{s \sim d_{\pi^{old}}(s)} \mathbb{E}_{a \sim \pi^{old}(a|s)}}_{\text{data generated by } \pi^{old}} \underbrace{\frac{\pi_\theta(a|s)}{\pi^{old}(a|s)}}_{\text{importance sampling correction}} A^{\pi^{old}}(s, a)$ $L_{\pi^{old}}(\theta) =\frac{1}{1-\gamma} \mathbb{E}_{s \sim d_{\pi^{old}}(s)} \mathbb{E}_{a \sim \pi^{old}(a|s)}\frac{\pi^{new}(a|s)}{\pi^{old}(a|s)}A^{\pi^{old}}(s, a)\rightarrow \max\quad [2.9]$

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

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

$|J(\pi^{new}) - J(\pi^{old}) - L_{\pi^{old}}(\theta) | \leq C \cdot KL^{\max} (\pi^{old}\|\pi^{new})$ $\text{где } \quad KL^{\max}(\pi^{old}\|\pi^{new}) = \max_s KL (\pi^{old} (\cdot|s)\|\pi^{new}(\cdot|s))$ $C = \frac{4 \gamma \max_{s,a} |A^{\pi^{old}}(s,a)|}{(1-\gamma)^2}$

Модуль разницы между желаемым оптимизируемым функцио��алом $J(\pi^{new}) - J(\pi^{old})$ и суррогатной функцией $L_{\pi^{old}}(\theta)$ ограничен максимумом по состояниям -дивергенции между старой и новой политикой.

Другими словами, если политики $\pi^{new}$ и $\pi^{old}$ близки, то разница $|J(\pi^{new}) - J(\pi^{old}) - L_{\pi^{old}}(\theta)|$ не велика.

Введем вариационную низкую оценку $J(\pi^{new}) - J(\pi^{old})$ (variational lower bound):

$J(\pi^{new}) - J(\pi^{old}) \geq L_{\pi^{old}}(\theta) - C \cdot KL^{\max} (\pi^{old}\|\pi^{new})$

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

Minorization step: фиксируем текущую политику и строим нижнюю оценку $L_{\pi^{old}}(\theta) - C \cdot KL^{\max} (\pi^{old}\|\pi^{new})$ , совпадающую с истинной целью $J(\pi^{new}) - J(\pi^{old})$ в текущей точке.
Maximization step: максимизируем нижнюю вариационную оценку $L_{\pi^{old}}(\theta) - C \cdot KL^{\max} (\pi^{old}\|\pi^{new}) \rightarrow \max$ в надежде, что максимизируется желаемый функционал $J(\pi^{new}) - J(\pi^{old}) \rightarrow \max$

$KL^{\max} (\pi^{old}\|\pi^{new})$ будет оценивать усредненным между двумя политиками по состояниям.

$KL^{\max} (\pi^{old}\|\pi^{new}) \approx \mathbb{E}_s KL \left(\pi^{old}(\cdot|s) \; \| \; \pi^{new}(\cdot|s) \right)$

А константу объявим гиперпараметром.

Таким образом, будем оптимизировать $L_{\pi^{old}}(\theta)$ с ограничением на разницу между старой и новой политиками.

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

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

Пусть $F(\theta)$ – оптимизируемая функция. Будем оптимизировать $F(\theta)$ в ее локальной окрестности. Другими словами, рассмотрим задачу обновления параметров $\theta$ , начиная с точки $\theta_0$ . Мы хотим сделать шаг $\theta - \theta_0$ , который минимизирует линейную аппроксимацию функции в окрестности $\theta_0$ , но при этом величина шага не должна быть слишком большой.

Приблизим $F(\theta)$ линейной аппроксимацией с помощью разложения в ряд Тейлора:

$\hat F(\theta) \approx F(\theta_0) + \nabla F(\theta_0)^T (\theta - \theta_0) \rightarrow \min$

Поскольку $F(\theta_0)$ не зависит от параметров $\theta$ , то константу можно опустить.
Таким образом, будем оптимизировать $\nabla F(\theta_0)^T (\theta - \theta_0) \rightarrow \min$ .
Обозначим расстояние между старыми и новыми параметрами: $d = \theta - \theta_0$ . Тогда, мы хотим найти такое что $\nabla F(\theta_0)^T \cdot d \rightarrow min$ с квадратичным ограничением $d^T d < \delta$ .

Аналитическое решение методом множителей Лагранжа: $L(d,\lambda) = \nabla F^T d + \lambda(d^T d − \delta)$
$\nabla F + 2 \lambda d = 0$
Оптимальная дистанция $d_{opt}$ пропорциональна антиградиенту оптимизируемой функции $- \nabla F (\theta)$ .

$d_{opt} \propto - \nabla F (\theta)$

Немного усложним задачу. Будем минимизировать ту же функцию $\nabla F(\theta_0)^T \cdot d \rightarrow min$ , но ограничим расстояние между старыми и новыми параметрами по неевклидовой метрике $d^T K d < \delta$ , где – положительно определённая матрица.

Аналитическое решение для такого ограничения:
$L(d,\lambda) = \nabla F^T d + \lambda(d^T К d − \delta)$
$\nabla F + 2 \lambda К d = 0$
$d_{opt} = − \frac{1}{2\lambda} K^{-1} \nabla F (\theta)$
Таким образом, $d_{opt}$ пропорционально $- K^{-1}\nabla F(\theta)$

$d_{opt} \propto - K^{-1} \nabla F(\theta)$

3.1. Natural Policy Gradient, NPG (2001)

Статья: A Natural Policy Gradient

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

Пусть имеется политика с параметрами $\theta^{old}$ . Хочется найти новые параметры $\theta^{new}$ , которые находились бы в окрестности параметров $\theta^{old}$ .

Квадратичное расстояние между весами $d^T d < \epsilon$ , где $d = \theta^{new} - \theta^{old}$ будет задавать окружность в пространстве весов. Однако, хотелось бы измерять расстояние между политиками $\pi(\theta^{old})$ и $\pi(\theta^{new})$ какими-то более умными способами.

Будем использовать расхождение Кульбака-Лейблера (-дивергенцию) между распределениями над действиями в качестве метрики между политиками $\pi^{old}$ и $\pi^{new}$ .

$\rho = KL( \pi(\theta^{new})|| \pi(\theta^{old})) = \sum \pi(\theta^{old})(a|s) \log \frac{\pi(\theta^{old})(a|s) }{\pi(\theta^{new})(a|s) } \leq \delta$

Обозначим как функцию параметров $f(\theta) =KL(\pi(\theta^{old})|| \pi(\theta^{new}))$ .
Разложим $f(\theta)$ около точки $\theta^{old}$ в ряд Тейлора до второго порядка производной:

$f(\theta) \approx f(\theta^{old}) + \nabla f(\theta^{old})^T (\theta - \theta^{old}) + \frac{1}{2} (\theta - \theta^{old})^T \nabla^2 f(\theta^{old}) (\theta - \theta^{old}) + O(\|\theta - \theta^{old}\|^3)$

$f(\theta^{old}) = KL(\pi(\theta^{old})|| \pi(\theta^{old})) = 0$ – из определения -дивергенции.
$\nabla f(\theta^{old}) = \nabla KL(\pi(\theta^{old})|| \pi(\theta^{old})) = 0$ поскольку -дивергенция при совпадении распределений достигает своего минимума.
$O(\|\theta - \theta^{old}\|^3)$ – пренебрежимо мало.

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

$f(\theta) \approx \frac{1}{2} (\theta - \theta^{old})^T \cdot \nabla^2 f (\theta^{old}) \cdot (\theta - \theta^{old}) = \frac{1}{2} \cdot d^T \cdot \nabla^2 f(\theta^{old}) \cdot d, \quad \text{где} \quad d = \theta - \theta^{old}$ $KL(\pi(\theta^{new})|| \pi(\theta^{old})) \approx \frac{1}{2} d^T \cdot \nabla^2 KL(\pi(\theta^{new})|| \pi(\theta^{old})) \cdot d$

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

Пока что, пользуемся функционалом $J(\pi_{\theta}) = \mathbb{E}_{\tau \sim \pi} \sum_{t=0} \gamma^{t} r_t$ .
Максимизируем $\nabla J^T d \rightarrow \max_{d}$ с ограничениями $d^T K d < \delta$ , где $K = \nabla^2 KL(\pi(\theta^{new})|| \pi(\theta^{old}))$ – матрица вторых производных -дивергенции.

В таком случае, оптимальный вектор изменения $d_{opt}$ будет пропорционален $K^{-1} \nabla J(\theta)$ .
Обновление политики происходит по формуле:

$\theta_{t+1} = \theta_t + \alpha K^{-1} \nabla J(\theta_t)$

Однако, обращение матрицы $K^{-1}$ очень дорогая операция, потому данный метод будет работать только на игрушечных средах. Хотелось бы этот метод как-то улучшить.

Заметим, что можно не искать отдельно матрицу $K^{-1}$ , а искать сразу значение выражения $K^{-1} \nabla J(\theta)$ .

$x = K^{-1} \nabla J(\theta)$

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

$Kx = \nabla J(\theta)$

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

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

$\Delta \theta \approx \alpha \cdot K^{-1} \nabla J(\theta) = \alpha \cdot \text{ConjugateGradient} (Kx = \nabla J(\theta))$

Параметр $\alpha$ подбирают линейным поиском шагая от $\theta^{old}$ по направлению $K^{-1} \nabla J(\theta)$ :

$\frac{1}{N} \sum KL(\pi^{new}(s_i)\| \pi^{old}(s_i)) < \alpha$

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

$\theta_{t+1} = \theta_t + \Delta \theta$

3.2. Trust Region Policy Optimization, TRPO (2015)

Статья: Trust Region Policy Optimization

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

$J(\pi^{new}) - J(\pi^{old}) \approx L(\theta^{new},\theta^{old}) =\mathbb{E}_{s \sim d_{\pi^{old}}} \mathbb{E}_{a \sim \pi^{old}}\frac{\pi^{new}(a|s)}{\pi^{old}(a|s)}A^{\pi^{old}}(s, a)\rightarrow \max$

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

$L(\theta^{new},\theta^{old}) \rightarrow \max$ $KL(\pi(\theta^{new})|| \pi(\theta^{old})) = d^T K d < \delta, \quad \text{где} \;\; K = \nabla^2 KL(\pi(\theta^{new})|| \pi(\theta^{old}))$

Метод Trust Region Policy Optimization:

Инициализируем политику $\pi^{new}$
1. Loop:
  1. $\pi^{old} = \pi^{new}$
  2. Генерируем траектории, используя текущую политику $\pi^{old}$ .
  3. Считаем $A^{\pi^{old}}(s_i, a_i) \quad \forall (s_i, a_i)$
  4. Цикл по эпохам:
    1. Цикл по батчам:
      1. Считаем градиент суррогатной функции: $g = \nabla L = \nabla \frac{1}{N} \sum^{N}_{i=0} \frac{\pi^{new}}{\pi^{old}} A^{\pi^{old}}(s_i, a_i)$
      2. Аналитически считаем матрицу вторых производных KL-дивергенции: $K = \nabla^2 \frac{1}{N} KL(\pi^{new}(s_i) || \pi^{old}(s_i)))$
      3. Методом сопряженных градиентов ищем направление обновления весов, решая систему методом сопряжённых градиентов: $d = \text{ConjugateGradient} (Kx = g) = K^{-1} \nabla L$
      4. Линейным поиском подбираем параметр $\alpha$ , в направлении с ограничением: $\frac{1}{N} \sum KL(\pi^{new}(s_i)| \pi^{old}(s_i)) < \alpha$ и одновременной проверкой значения: $\frac{1}{N} \sum^{N} \frac{\pi^{new}(s_i)}{\pi^{old}(s_i))} A^{\pi^{old}}(s_i, a_i) \rightarrow \max$
      5. Обновляем параметры политики $\theta_{t+1} = \theta_t + \alpha_t \cdot d$

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

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

3.3. Proximal Policy Optimization (PPO), 2017

Proximal Policy Optimization Algorithms

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

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

$\mathbb{E}_{s \sim d_{\pi^{old}}} \mathbb{E}_{a \sim \pi^{old}}\frac{\pi^{new}(a|s)}{\pi^{old}(a|s)}A^{\pi^{old}}(s, a){\color{red}{- C \cdot KL(\pi^{new}(a|s) \| \pi^{old}(a|s))}}\rightarrow \max$

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

$KL(\pi^{new}(a|s) \| \pi^{old}(a|s))$ не достаточно штрафует за разницу между политиками и выражение $\frac{\pi^{new}(a|s)}{\pi^{old}(a|s)}$ часто уходит либо в беcконечность, либо в ноль и обучение разваливается.

Ограничиваем изменение политики, "подрезая" (clipping) слишком большие значения importance ratio $\frac{\pi^{new}(a|s)}{\pi^{old}(a|s)}$ .

$\mathbb{E}_{s \sim d_{\pi^{old}}} \mathbb{E}_{a \sim \pi^{old}} \;{\color{red}{Clip \left[ \frac{\pi^{new}(a|s)}{\pi^{old}(a|s)}, 1 - \epsilon, 1 + \epsilon \right]}}\cdot A^{\pi^{old}}(s, a)\rightarrow \max$ $\epsilon \in [0.1, 0.3]$

Если $\frac{\pi^{new}(a|s)}{\pi^{old}(a|s)}$ выходит за границы $1 \pm \epsilon$ , то считаем его равным $1 \pm \epsilon$ .Таким образом, новая политика не сможет выйти за регион $1 \pm \epsilon$ .

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

$\begin{multline} \mathbb{E}_{s \sim d_{\pi^{old}}} \mathbb{E}_{a \sim \pi^{old}} \;{\color{red}{\min \Big(}} \frac{\pi^{new}(a|s)}{\pi^{old}(a|s)} \cdot A^{\pi^{old}}(s, a), \\Clip \left[ \frac{\pi^{new}(a|s)}{\pi^{old}(a|s)}, 1 - \epsilon, 1 + \epsilon \right] \cdot A^{\pi^{old}}(s, a){\color{red}{\Big)}} \rightarrow \max \end{multline}$

В статье PPO были представлены следующие графики. По вертикальной оси находятся значения функции ошибки, а по горизонтальной значение функции $r = \frac{\pi^{new}(a|s)}{\pi^{old}(a|s)}$ .
В красной точке $\pi^{new}$ совпадает с $\pi^{old}$ , соответсвенно .

Допустим , то тогда и должно увеличиваться, поскольку мы хотим увеличить вероятность этого действия. Однако, когда доходит до числа $1 + \epsilon$ , дальнейшее изменение параметров не имеет смысла с точки зрения функции потерь, даже если бы это было на пользу. В обратной ситуации аналогично, если , то должно уменьшаться. Но, при достижении границы $1 - \epsilon$ выгода с точки зрения функции потерь теряется.

На графике ниже представлена линейная интерполяция изменения весов. Слева расположена $\pi^{old}$ , справа последняя версия $\pi^{new}$ .

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

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

Инициализируем параметры политики $\pi^{new}$
Loop:
1. Сохраняем старую политику $\pi^{old} = \pi^{new}$
2. Генерируем траектории c $\pi^{old}$
3. Считаем advantage estimate
4. Цикл по эпохам:
  1. Цикл по батчам:
    1. Считаем importance ratio: $r(s,a) = \frac{\pi^{new}(a|s)}{\pi^{old}(a|s)}$
    2. Считаем clipped importance ratio: $r^{clip}(s,a) = Clip(r(s,a), 1-\epsilon, 1+\epsilon) \quad \epsilon \in [0.1, 0.3]$
    3. Считаем Сlipped objective: $L_{Clip} = \frac{1}{N} \sum_{t} \Big[ \min \Big( r(s,a) \cdot A_t, \quad r^{clip}(s,a) \cdot A_t\Big) \Big]$
    4. Обновляем параметры $\theta = \theta + \alpha \nabla_{\theta} L_{Clip}$

4. Generalized Advantage Estimation (GAE)

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

$\mathbb{E}_{s \sim d_{\pi^{old}}} \mathbb{E}_{a \sim \pi^{old}} \;\frac{\pi^{new}(a|s)}{\pi^{old}(a|s)} \cdot {\color{red}{A^{\pi^{old}}(s, a)}}\rightarrow \max$ $\text{где} \quad A(s,a) = Q(s,a) - V(s) \quad [4.1]$

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

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

Функцию ценности обычно аппроксимируют нейронной сетью $V(s) \rightarrow V_{\phi}(s)$ с функцией потерь

$L_{V_{\varphi}} = \frac{1}{2}\Big[V_{\phi}(s_t) - y_t \Big]^2 \quad [4.2]$

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

Monte-Carlo возврат $y_t = \sum_t^{\infty} \gamma^t r_t$ несмещённый, но имеет огромную дисперсию.
Поэтому на практике используют усечённую -шаговую оценку:

$y_t \approx r_t + \gamma r_{t+1} + \dots + \gamma^{n-1} r_{t+n-1} +\gamma^n \text{detach}\Big(V_{\phi}(s_{t+n})\Big) \approx Q(s,a)$

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

$y_t = A(s_t, a_t) + \text{detach}\Big( V_{\phi}(s_t)\Big) \quad [4.3]$

Тогда:

$\frac{dL}{dV_{\phi}(s_t)} = V_{\phi}(s_t) - y_t$

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

$\begin{multline}V_{\phi}(s_t) \leftarrow V_{\phi}(s_t) - \alpha \Big( V_{\phi}(s_t) - y_t \Big) = \\ = V_{\phi}(s_t) - \alpha \Big( V_{\phi}(s_t) - A(s_t, a_t) - \text{detach}\Big( V_{\phi}(s_t)\Big) \Big) =V_{\phi}(s_t) + \alpha A(s_t, a_t)\end{multline}$ $V_{\phi}(s_t) \leftarrow V_{\phi}(s_t) + \alpha A_1(s_t, a_t) \quad [4.4]$

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

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

$A(s_t,a_t) = {\color{green}{Q(s_t,a_t)}} - V_{\phi}(s_t)$

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

$\Psi(s_t,a_t) = {\color{green}{\sum_{t'=0}^{\infty} \gamma^{t'} r_{t'}}} - V_{\phi}(s_t)$

Можно -функцию оценить одношаговой оценкой из знаний $V_{\varphi}$ :

$\Psi_1(s_t,a_t) = {\color{green}{r_t + \gamma V(s_{t+1})}} - V_{\phi}(s_t)$

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

Возникает задача умного подбора шага при оценивании функции $\Psi$ :

$\Psi_{N}(s_t,a_t) = {\color{green}{r_t + \gamma r_{t+1} + \gamma^2 r_{t+2} + \dots + \gamma^{N} V_{\phi}(s_{t+N})}} - V_{\phi}(s_t)$ $\Psi_{N}(s_t,a_t) = \sum_{k=0}^{N-1} \Big[ \gamma^k r_{t+k} \Big] + \gamma^{N} V_{\phi}(s_{t+N}) - V_{\phi}(s_t)$

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

$L_{TRPO} = \mathbb{E}_{s \sim d_{\pi^{old}}} \mathbb{E}_{a \sim \pi^{old}} \;\frac{\pi^{new}(a|s)}{\pi^{old}(a|s)} \cdot {\color{red}{\Psi^{\pi^{old}}(s, a)}}\rightarrow \max$

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

$L_{V_{\varphi}} = \frac{1}{2}\Big[V_{\phi}(s_t) - {\color{red}{y_t}} \Big]^2 \quad \text{, где} \quad y_t = {\color{red}{\Psi(s_t, a_t)}} + \text{detach}\Big( V_{\phi}(s_t)\Big)$

4.3. Generalized Advantage Estimation (GAE)

Вместо того чтобы подбирать количество шагов , можно использовать все -шаговые оценки Advantage с затухающими весами:

Estimation	$\Psi_{1}(s,a)$	$\Psi_{2}(s,a)$	$\Psi_3(s,a)$	$\dots$
Weight	$1-\lambda$	$(1-\lambda)\lambda$	$(1-\lambda)\lambda^2$	$\dots$

$\begin{multline}\Psi_{GAE}(s,a) = (1-\lambda) \cdot \Psi_{1}(s,a) + (1-\lambda)\lambda \cdot \Psi_{2}(s,a) + (1-\lambda)\lambda^2 \cdot \Psi_{3}(s,a) + \dots \\ \dots + (1-\lambda)\lambda^{N-1} \cdot \Psi_{N}(s,a)\end{multline}$ $\Psi_{GAE}(s,a) =(1-\lambda) \sum_{N=1}^{\infty} \lambda^{N-1} \Psi_{N}(s,a) \quad [4.5]$ $\text{где} \quad \Psi_{N}(s_t, a_t) = \sum_{k=1}^{N-1} \gamma^k r_{t+k} + \gamma^N V(s_{t + N}) - V(s_t) \quad [4.6]$

Однако, пересчитывать все $\Psi_N \quad \forall N \in (1,\infty)$ слишком долго и не оптимально.

Выпишем двушаговый $\Psi_2(s,a)$ через $\Psi_1(s,a)$ .

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

$\begin{align}\Psi_1(s_t,a_t) = r_t + \gamma V(s_{t+1}) - V(s_t) \\\Psi_2(s_t,a_t) = r_t + \gamma r_{t+1} + \gamma^2 V(s_{t+2}) - V(s_{t})\end{align}$

Выпишем одношаговый $\Psi_1$ , но для следующего состояния $s_{t+1}$ и следующего действия $a_{t+1}$ :

$\Psi_1(s_{t+1},a_{t+1}) = r_{t+1} + \gamma V(s_{t+2}) - V(s_{t+1})$

Умножим на $\gamma$ :

$\gamma \Psi_1(s_{t+1},a_{t+1}) = \gamma r_{t+1} + \gamma^2 V(s_{t+2}) - \gamma V(s_{t+1})$

Добавим и вычтем $\gamma V(s_{t+1})$ из $\Psi_2(s_t,a_t)$ :

$\begin{multline} = \Big[ r_t + \gamma V(s_{t+1}) - V(s_{t}) \Big] + \Big[ \gamma r_{t+1} + \gamma^2 V(s_{t+2}) - \gamma V(s_{t+1}) \Big] = \\ = \Psi_1(s_t, a_t) + \gamma \Psi_1(s_{t+1}, a_{t+1}) \end{multline}$ $\Psi_2(s_t,a_t) = \Psi_1(s_t, a_t) + \gamma \Psi_1(s_{t+1}, a_{t+1}) \quad [4.7]$

Если обобщить на шагов, то получится общая формула -шаговой ошибки через одношаговые:

$\Psi_N(s,a) = \sum_{k=0}^{N-1} \gamma^k \Psi_1(s_{t}, a_{t}) \quad [4.8]$

Подставив получившуюся формулу в можно доказать, что:

$\Psi_{GAE}(s,a) =(1-\lambda) \sum_{N=1}^{\infty} \lambda^{N-1} \Psi_{N}(s,a) = \sum_{N=1}^{\infty} (\gamma \lambda)^t \Psi_1(s_{t}, a_{t})$ $\Psi_{GAE}(s,a) =\sum_{N=1}^{\infty} (\gamma \lambda)^t \Psi_1(s_{t}, a_{t})\quad [4.9]$

Однако $\Psi_{GAE}(s,a)$ считаем не до конца траектории, а только по сделанным шагам, потому ограничим до :

$\Psi_{GAE}(s,a) =\sum_{N=1}^{T} (\gamma \lambda)^t \Psi_1(s_{t}, a_{t})\quad [4.10]$

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

$\Psi_{GAE}(s_t,a_t) = \Psi_1(s_t,a_t) +\lambda \gamma \Psi_{GAE}(s_{t+1}, a_{t+1}) \quad \text{if not } done_{t+1} \quad [4.11]$

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

$V_{\phi}(s_t) = V_{\phi}(s_t) + \alpha \Psi_{GAE}(s_t,a_t)$

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

Вместо выбора одного , $\lambda$ строит смесь всех -шаговых возвратов. Таким образором, выбирается не конкретный горизонт, а скорость затухания памяти. При $\lambda = 0$ алгоритм имеет короткую память и оценка близка к(низкая дисперсия, но высокое смещение). При $\lambda = 1$ – долгая память, оценка стремится к полному Монте-Карло возврату (низкое смещение, но высокая дисперсия). Эмпирически для многих задач оптимальный баланс достигается в диапазоне $\lambda = 0.9 – 0.97$ .

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

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

Инициализируем параметры политики $\pi^{new}$ и параметры критика $V_{\phi}$ .
Loop:
1. Сохраняем политику $\pi^{old} = \pi^{new}$
2. Генерируем траектории (rollouts) из политики $\pi^{old}$ , сохраняя логиты $\pi^{old}(a|s)$ для подсчета importance ratio.
3. Сохраняем выходы критика $V^{old}(s_t) = V_{\phi}(s_t)$
4. Считаем одношаговые оценки для всех состояний:
  if not $done_{t+1}$ :
  $\Psi_{(1)}(s_t, a_t) = r_t + \gamma V_{\phi}(s_{t+1}) - V_{\phi}(s_t)$
  else:
  $\Psi_{(1)}(s_t, a_t) = r_t - V_{\phi}(s_t)$
5. Считаем рекурсивно по формуле GAE оценки:
  $\Psi_{GAE}(s_{N-1}, a_{N-1}) = \Psi_{(1)}(s_{N-1}, a_{N-1})$
  for t in (N-2, 0, -1):
  if not $done_{t+1}$ :
  $\Psi_{GAE}(s_t,a_t) = \Psi_1(s_t,a_t) +\lambda \gamma \Psi_{GAE}(s_{t+1}, a_{t+1})$
  else:
  $\Psi_{GAE}(s_t,a_t) = \Psi_1(s_t,a_t)$
6. Считаем целевые значения критика: $y(s_t) = \Psi_{GAE}(s_t,a_t) + V_{\phi}(s_t)$
7. for epoch in n_epochs:
  1. for batch in batches:
    1. Считаем importance ratio: $\rho(s,a) = \frac{\pi^{new}(a|s)}{\pi^{old}(a|s)}$
    2. Считаем clipping importance ratio $\rho^{clip}(s,a) = Clip(\rho(s,a), 1-\epsilon, 1+\epsilon) \quad \epsilon \in [0.1, 0.3]$
    3. Считаем clipped objective (лосс актора): $L_{Clip} = \frac{1}{B} \sum \min \Big( \rho(s,a) \cdot \Psi_{GAE}(s, a), \quad \rho^{clip}(s,a) \cdot \Psi_{GAE}(s, a) \Big) \rightarrow \max$
    4. Обновляем параметры: $\theta = \theta + \alpha \nabla_{\theta} L_{Clip}$
    5. Считаем лосс критика: $L_{MSE} = \frac{1}{B} \sum (y(s) - V_{\phi}(s))^2$
    6. Обновляем параметры: $\phi = \phi - \alpha \nabla_{\phi} L_{MSE}$