Deep RL Course, юнит 8: PPO

Это перевод теоретической части юнита 8 курса Hugging Face Deep RL Course. Практика — PPO с нуля по образцу CleanRL на LunarLander и асинхронный PPO в Sample Factory для Doom — переведена на странице «Юнит 8, практика». Карта курса — в путеводителе.

Как читать вместе с курсом ChinaAI

Прочитайте юнит перед разделом «PPO целиком» урока 6.0.3. Курс Hugging Face лучше всего объясняет клиппинг: таблица из шести случаев ниже — самое наглядное место. Урок добавляет то, чего здесь нет: GAE, несколько эпох на одном батче, мониторинг KL и clipfrac и полный код на PyTorch. Дальше та же формула встретится в RLHF (урок 6.3.2), где отношение вероятностей считается для каждого токена, а в GRPO (урок 6.4.1) из неё уберут критика.

О переводе

Перевод сделан по исходникам курса на GitHub (лицензия Apache-2.0). Формулы, которые в оригинале нарисованы на картинках, записаны формулами; таблица шести случаев из работы Дэниела Бика, которая в оригинале дана картинкой, набрана заново; остальные рисунки описаны словами или заменены схемами. В оригинале у юнита нет теста — в конце страницы вопросы от переводчика.

Введение

В юните 6 мы изучили Advantage Actor Critic (A2C) — гибридную архитектуру, сочетающую методы на основе ценности и на основе политики. Она стабилизирует обучение, уменьшая дисперсию, с помощью:

• актора, который управляет тем, как ведёт себя агент (метод на основе политики);
• критика, который оценивает, насколько хорошее действие совершено (метод на основе ценности).

Сегодня мы изучим Proximal Policy Optimization (PPO)проксимальная оптимизация политики (PPO)EN proximal policy optimization (PPO)Schulman и др., 2017. Ограничивает шаг обновления «клиппингом» отношения вероятностей. — архитектуру, которая повышает стабильность обучения агента, не допуская слишком больших обновлений политики. Для этого используется отношение, показывающее разницу между текущей и старой политикой, и это отношение обрезается (clip) до диапазона $[1 - \epsilon, 1 + \epsilon]$.

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

Юнит состоит из двух частей.

• В первой части вы изучите теорию PPO и напишете агента PPO с нуля по образцу реализации CleanRL. Проверять его надёжность будем на LunarLander-v2. LunarLander-v2 — первая среда, с которой вы начинали этот курс. Тогда вы не знали, как работает PPO, а теперь можете написать его с нуля и обучить. Разве не потрясающе 🤩.
• Во второй части мы глубже займёмся оптимизацией PPO с помощью Sample-Factory и обучим агента играть в ViZDoom (открытую версию Doom).

В оригинале здесь картинка: «Это среды, в которых вы будете обучать агентов: среды ViZDoom».

Звучит увлекательно? Начнём! 🚀

Интуиция PPO

Идея Proximal Policy Optimization (PPO) в том, чтобы повысить стабильность обучения политики, ограничивая, насколько сильно меняется политика на каждой эпохе обучения: мы хотим избежать слишком больших обновлений политики.

По двум причинам:

• Опыт показывает, что небольшие обновления политики во время обучения с большей вероятностью сходятся к оптимальному решению.
• Слишком большой шаг при обновлении политики может привести к «падению с обрыва» (получению плохой политики) и долгому восстановлению — или к тому, что восстановиться не удастся вовсе.

flowchart LR
A["Хорошая политика"] -->|"маленький шаг"| B["Чуть лучше"]
B -->|"маленький шаг"| C["Ещё лучше"]
A -->|"слишком большой шаг"| D["Обрыв: плохая политика"]
D -.->|"собирает плохие данные"| D

Маленькие шаги обновления политики делают обучение стабильнее (в оригинале — рисунок с обрывом по мотивам статьи Джонатана Хуэя)

Поэтому PPO обновляет политику осторожно. Для этого нужно измерить, насколько текущая политика изменилась по сравнению с прежней, — через отношение текущей политики к прежней. И мы обрезаем это отношение до диапазона $[1 - \epsilon, 1 + \epsilon]$, то есть убираем стимул для текущей политики уходить слишком далеко от старой (отсюда и слово proximal — «ближний»).

Примечание переводчика: почему обрыв особенно опасен в RL

В обучении с учителем плохой шаг градиента портит веса, но данные остаются прежними, и следующий шаг всё исправит. В RL данные собирает сама политика: испорченная политика начинает собирать плохие траектории, учится на них и может так и не выбраться. Поэтому ограничение шага в RL важнее, чем в обычном обучении.

Clipped Surrogate Objective — обрезанная суррогатная целевая функция

Напоминание: целевая функция политики

Вспомним, какую целевую функцию мы оптимизировали в REINFORCE:

$$L^{PG}(\theta) = \mathbb{E}_t\big[\log \pi_\theta(a_t \mid s_t) \cdot A_t\big]$$

Здесь $\log \pi_\theta(a_t \mid s_t)$ — логарифм вероятности совершить это действие в этом состоянии, а $A_t$ — преимущество: если $A > 0$, это действие лучше других возможных в этом состоянии.

Идея была в том, что, делая шаг градиентного подъёма по этой функции (что равносильно шагу градиентного спуска по ней со знаком минус), мы подталкиваем агента совершать действия, ведущие к большей награде, и избегать вредных действий.

Но проблема в размере шага:

• слишком маленький — обучение идёт слишком медленно;
• слишком большой — обучение слишком нестабильно.

Идея PPO — ограничить обновление политики новой целевой функцией, обрезанной суррогатной целевой функцией (clipped surrogate objective function), которая с помощью обрезки (clip) удерживает изменение политики в небольшом диапазоне.

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

$$L^{CLIP}(\theta) = \hat{\mathbb{E}}_t\Big[\min\big(r_t(\theta)\hat{A}_t,\ \operatorname{clip}\big(r_t(\theta), 1 - \epsilon, 1 + \epsilon\big)\hat{A}_t\big)\Big]$$

Разберём каждую часть, чтобы понять, как она работает.

Функция отношения

$$r_t(\theta) = \frac{\pi_\theta(a_t \mid s_t)}{\pi_{\theta_{old}}(a_t \mid s_t)}$$

Это вероятность совершить действие $a_t$ в состоянии $s_t$ при текущей политике, делённая на ту же вероятность при прежней политике.

Как видно, $r_t(\theta)$ — отношение вероятностей текущей и старой политики:

• если $r_t(\theta) > 1$, действие $a_t$ в состоянии $s_t$ при текущей политике вероятнее, чем при старой;
• если $r_t(\theta)$ между 0 и 1, действие при текущей политике менее вероятно, чем при старой.

Значит, это отношение вероятностей — простой способ оценить расхождение между старой и текущей политикой.

Необрезанная часть

Это отношение может заменить логарифм вероятности, который мы используем в целевой функции политики. Так получается левая часть новой целевой функции: отношение, умноженное на преимущество:

$$L^{CPI}(\theta) = \hat{\mathbb{E}}_t\left[\frac{\pi_\theta(a_t \mid s_t)}{\pi_{\theta_{old}}(a_t \mid s_t)}\hat{A}_t\right] = \hat{\mathbb{E}}_t\big[r_t(\theta)\hat{A}_t\big]$$

(Формула из статьи Proximal Policy Optimization Algorithms.)

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

Обрезанная часть

$$\operatorname{clip}\big(r_t(\theta), 1 - \epsilon, 1 + \epsilon\big)\hat{A}_t$$

Поэтому эту целевую функцию нужно ограничить, наказывая изменения, при которых отношение уходит далеко от 1 (в статье отношение может меняться только от 0,8 до 1,2).

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

Для этого есть два решения:

• TRPOTRPO (оптимизация политики в доверительной области)EN trust region policy optimization (TRPO)Schulman и др., 2015 — шаг политики с ограничением на KL к старой политике; предшественник PPO. (Trust Region Policy Optimization, оптимизация политики в доверительной области) ограничивает обновление политики условием на KL-дивергенциюKL-дивергенция (дивергенция Кульбака — Лейблера)EN Kullback–Leibler divergenceKL(p ‖ q) = Σ p log(p / q) = H(p, q) − H(p) ≥ 0. Несимметрична, поэтому не настоящее расстояние., вынесенным за пределы целевой функции. Но этот метод сложно реализовать, и он дольше считается.
• PPO обрезает отношение вероятностей прямо в целевой функции — своей обрезанной суррогатной целевой функцией.

Обрезанная часть — это вариант, в котором $r_t(\theta)$ обрезано до $[1 - \epsilon, 1 + \epsilon]$.

В обрезанной суррогатной целевой функции два отношения вероятностей: необрезанное и обрезанное до диапазона $[1 - \epsilon, 1 + \epsilon]$. Эпсилон — гиперпараметр, задающий этот диапазон обрезки (в статье $\epsilon = 0{,}2$).

Затем мы берём минимум из обрезанной и необрезанной целевой функции, поэтому итоговая цель — нижняя (пессимистическая) оценка необрезанной цели.

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

Визуализация обрезанной суррогатной целевой функции

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

В оригинале здесь картинка — таблица из работы Дэниела Бика «Towards Delivering a Coherent Self-Contained Explanation of Proximal Policy Optimization» (PDF). Вот она (в работе отношение обозначено $p_t(\theta)$, здесь — $r_t(\theta)$, как в тексте курса):

Случай	Отношение $r_t(\theta)$	Знак $A_t$	Значение минимума	Цель обрезана?	Знак цели	Градиент
1	в $[1 - \epsilon, 1 + \epsilon]$	+	$r_t(\theta) A_t$	нет	+	есть
2	в $[1 - \epsilon, 1 + \epsilon]$	−	$r_t(\theta) A_t$	нет	−	есть
3	$< 1 - \epsilon$	+	$r_t(\theta) A_t$	нет	+	есть
4	$< 1 - \epsilon$	−	$(1 - \epsilon) A_t$	да	−	0
5	$> 1 + \epsilon$	+	$(1 + \epsilon) A_t$	да	+	0
6	$> 1 + \epsilon$	−	$r_t(\theta) A_t$	нет	−	есть

Всего шесть разных ситуаций. Сначала вспомните, что мы берём минимум из обрезанной и необрезанной цели.

Случаи 1 и 2: отношение внутри диапазона

В ситуациях 1 и 2 обрезка не срабатывает, потому что отношение внутри диапазона $[1 - \epsilon, 1 + \epsilon]$.

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

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

В ситуации 2 преимущество отрицательное: действие хуже среднего по всем действиям в этом состоянии. Значит, надо отучать текущую политику от этого действия в этом состоянии.

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

Случаи 3 и 4: отношение ниже диапазона

Если отношение вероятностей меньше $1 - \epsilon$, вероятность выбрать это действие в этом состоянии намного ниже, чем у старой политики.

Если, как в ситуации 3, оценка преимущества положительная ($A > 0$), вы хотите увеличить вероятность выбрать это действие в этом состоянии.

Но если, как в ситуации 4, оценка преимущества отрицательная, мы не хотим уменьшать ещё сильнее вероятность выбрать это действие в этом состоянии. Поэтому градиент равен 0 (мы на горизонтальном участке), и веса не обновляются.

Случаи 5 и 6: отношение выше диапазона

Если отношение вероятностей больше $1 + \epsilon$, вероятность выбрать это действие в этом состоянии при текущей политике намного выше, чем при прежней.

Если, как в ситуации 5, преимущество положительное, мы не хотим жадничать. Вероятность выбрать это действие в этом состоянии и так уже выше, чем у прежней политики. Поэтому градиент равен 0 (мы на горизонтальном участке), и веса не обновляются.

Если, как в ситуации 6, преимущество отрицательное, мы хотим уменьшить вероятность выбрать это действие в этом состоянии.

Итак, политика обновляется только по необрезанной части цели. Когда минимум — обрезанная часть, веса политики не обновляются: градиент будет равен 0.

Значит, политика обновляется, только если:

• отношение внутри диапазона $[1 - \epsilon, 1 + \epsilon]$;
• отношение вне диапазона, но преимущество тянет его обратно к диапазону:
  • отношение ниже диапазона, но преимущество $A_t > 0$;
  • отношение выше диапазона, но преимущество $A_t < 0$.

Вы можете спросить, почему градиент равен 0, когда минимум — обрезанное отношение. Когда отношение обрезано, производная берётся уже не от $r_t(\theta) \cdot A_t$, а от $(1 - \epsilon) \cdot A_t$ или $(1 + \epsilon) \cdot A_t$, и обе равны 0.

Подведём итог: благодаря обрезанной суррогатной цели мы ограничиваем, насколько текущая политика может отойти от старой. Мы убираем стимул выводить отношение вероятностей за пределы интервала, потому что обрезка обнуляет градиент. Если отношение $> 1 + \epsilon$ или $< 1 - \epsilon$, градиент равен 0.

flowchart TB
Q1{"Знак преимущества A_t"} -->|"A_t > 0: действие лучше среднего"| P["Тянем r_t вверх, но не выше 1+ε"]
Q1 -->|"A_t отрицательно: действие хуже среднего"| N["Тянем r_t вниз, но не ниже 1−ε"]
P --> P2["r_t выше 1+ε: градиент 0 (случай 5)"]
N --> N2["r_t ниже 1−ε: градиент 0 (случай 4)"]

Шпаргалка к таблице: клиппинг останавливает только «слишком рьяное» движение в выгодную сторону

Примечание переводчика: шесть случаев на числах

Пусть $\epsilon = 0{,}2$ и $A_t = 2$ (действие хорошее). При $r_t = 1{,}1$: $\min(1{,}1 \cdot 2;\ 1{,}1 \cdot 2) = 2{,}2$ — обрезки нет, градиент есть (случай 1). При $r_t = 1{,}5$: $\min(1{,}5 \cdot 2;\ 1{,}2 \cdot 2) = \min(3;\ 2{,}4) = 2{,}4$ — константа, градиент 0 (случай 5): политика уже достаточно подняла вероятность. Теперь $A_t = -2$ (действие плохое) и $r_t = 1{,}5$: $\min(1{,}5 \cdot (-2);\ 1{,}2 \cdot (-2)) = \min(-3;\ -2{,}4) = -3$ — выбрана необрезанная часть, градиент есть (случай 6): ошибку «слишком часто выбираем плохое действие» исправлять можно всегда.

Итоговая функция потерь PPO в стиле actor-critic — сочетание обрезанной суррогатной целевой функции, функции потерь ценности и энтропийного бонуса:

$$L_t^{CLIP+VF+S}(\theta) = \hat{\mathbb{E}}_t\Big[L_t^{CLIP}(\theta) - c_1 L_t^{VF}(\theta) + c_2\, S\big[\pi_\theta\big](s_t)\Big]$$

Здесь $c_1$ и $c_2$ — коэффициенты, $L_t^{VF} = \big(V_\theta(s_t) - V_t^{targ}\big)^2$ — квадратичная ошибка критика, а $S$ — энтропийный бонусэнтропийный бонусEN entropy bonusДобавка c·H[π(·|s)] к цели — награда за «неуверенность» политики, чтобы она не схлопнулась раньше времени и продолжала исследовать.: он поощряет исследование, не давая политике слишком рано стать детерминированной.

Это было довольно сложно. Не спешите разобраться в этих ситуациях, глядя на таблицу и график. Вы должны понять, почему это осмысленно. Если хотите копнуть глубже, лучший источник — работа «Towards Delivering a Coherent Self-Contained Explanation of Proximal Policy Optimization» Дэниела Бика, особенно раздел 3.4.

Примечание переводчика: от формулы к коду

В коде (см. практику) целевую функцию максимизируют, минимизируя её со знаком минус. Поэтому в CleanRL вместо $\min$ стоит $\max$ от отрицательных величин: pg_loss = torch.max(-A * ratio, -A * torch.clamp(ratio, 1 - eps, 1 + eps)).mean(), а полный loss — pg_loss - ent_coef * entropy + vf_coef * v_loss. Отношение считают через логарифмы: ratio = exp(newlogprob - oldlogprob).

Заключение

На сегодня всё. Поздравляем с окончанием юнита и практического занятия!

Лучший способ учиться — практиковаться и пробовать. Почему бы не доработать реализацию, чтобы она принимала на вход кадры?

Увидимся во второй части юнита 🔥 Продолжайте учиться и оставайтесь крутыми 🤗

Дополнительное чтение

Это необязательные материалы для тех, кто хочет копнуть глубже.

Объяснения PPO

• Towards Delivering a Coherent Self-Contained Explanation of Proximal Policy Optimization — Дэниел Бик.
• What is the way to understand Proximal Policy Optimization Algorithm in RL? (Stack Overflow).
• Foundations of Deep RL Series, L4: TRPO and PPO — Питер Аббил (видео).
• Пост OpenAI о PPO.
• PPO в Spinning Up.
• Статья Proximal Policy Optimization Algorithms.

Детали реализации PPO

• The 37 Implementation Details of Proximal Policy Optimization.
• Part 1 of 3 — Proximal Policy Optimization Implementation: 11 Core Implementation Details (видео).

Выборка по значимости

• Importance Sampling Explained (видео).

Примечание переводчика: что читать по-русски и по-китайски

PPO целиком (GAE, несколько эпох на батче, value loss, энтропия, KL и clipfrac) и код на 70 строк — урок 6.0.3. PPO для языковых моделей — урок 6.3.2; что меняет GRPO — урок 6.4.1. Отношение $r_t(\theta)$ — это выборка по значимости (importance sampling): данные собраны старой политикой, а оптимизируем новую. По-китайски: «Обучение с подкреплением своими руками» (Hands-on RL), глава 12 «Алгоритм PPO», и «Грибная книга» (easy-rl), глава 5 «Алгоритм PPO».

Вопросы от переводчика

Проверьте себя по юниту 84 вопроса·

1 Что показывает отношение $r_t(\theta) = 1{,}3$?

Действие при текущей политике на 30% вероятнее, чем при старойПреимущество действия равно 1,3Политика ухудшилась на 30%

2 При $\epsilon = 0{,}2$, $A_t = 1$ и $r_t = 1{,}5$ чему равна обрезанная цель $\min(r_t A_t,\ \mathrm{clip}(r_t, 0{,}8, 1{,}2) A_t)$?

3 В каких случаях таблицы градиент равен нулю?Можно выбрать несколько ответов

Случай 1: отношение внутри диапазона, $A > 0$Случай 4: отношение ниже $1 - \epsilon$, $A < 0$Случай 5: отношение выше $1 + \epsilon$, $A > 0$Случай 6: отношение выше $1 + \epsilon$, $A < 0$

4 Чем PPO проще TRPO?

PPO не использует преимуществоPPO ограничивает шаг обрезкой прямо в целевой функции, а TRPO — отдельным условием на KL-дивергенциюPPO не нужна нейросеть

Проверить ответы

Словарик

English	Русский
Proximal Policy Optimization (PPO)	проксимальная оптимизация политики
policy update	обновление политики
probability ratio $r_t(\theta)$	отношение вероятностей
clip, clip range $\epsilon$	обрезка (клиппинг), диапазон обрезки
clipped surrogate objective	обрезанная суррогатная целевая функция
unclipped / clipped part	необрезанная / обрезанная часть
lower (pessimistic) bound	нижняя (пессимистическая) оценка
Trust Region Policy Optimization (TRPO)	оптимизация политики в доверительной области
KL divergence constraint	ограничение на KL-дивергенцию
value loss	функция потерь ценности (критика)
entropy bonus	энтропийный бонус
importance sampling	выборка по значимости