Spinning Up: алгоритмы VPG и PPO

Где в курсе то же самое проще

Вывод градиента политики, reward-to-go, базовая линия, GAEобобщённая оценка преимущества (GAE)EN generalized advantage estimation (GAE)Schulman и др., 2015. Смешивает оценки преимущества разной длины с коэффициентом λ. и клиппинг PPO с числовыми примерами — в уроке 6.0.3 «Actor-critic, GAE и PPO». Как тот же PPO работает с языковой моделью (действие — токен, награда — от reward model) — в уроке 6.3.2 «RLHF — reward model и PPO», а как из него убирают критика — в 6.4.1 «GRPO». Эта страница — первоисточник в обозначениях, которые используют статьи.

Как устроен перевод

Переведены две страницы документации Spinning Up целиком: «Vanilla Policy Gradient» и «Proximal Policy Optimization». В оригинале раздел «Документация» вставляет docstring функций для двух версий кода — PyTorch и TensorFlow 1; они почти совпадают, поэтому мы перевели PyTorch-версию, а для TensorFlow оставили только то, что отличается. Абзацы «Словами», числовые примеры и раздел «Ключевые строки кода PPO» добавлены переводчиком. Псевдокод оригинала свёрстан пакетом algorithmic из LaTeX; здесь он записан нумерованным списком, формулы те же.

Vanilla Policy Gradient (простой градиент политики)

Основы

(Раньше: Введение в RL, часть 3)

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

Коротко

• VPG — алгоритм с обучением по текущей политикеon-policy (обучение на данных текущей политики)EN on-policyПротивоположность — off-policy, обучение на данных другой политики. (on-policy).
• VPG можно применять в средах как с дискретным, так и с непрерывным пространством действий.
• Реализация VPG в Spinning Up поддерживает распараллеливание через MPI.

Ключевые формулы

Пусть $\pi_{\theta}$ — политика с параметрами $\theta$, а $J(\pi_{\theta})$ — ожидаемая недисконтированная отдача политики на конечном горизонте. Градиент $J(\pi_{\theta})$ равен

$$\nabla_{\theta} J(\pi_{\theta}) = \underset{\tau \sim \pi_{\theta}}{\mathrm{E}}\left[ \sum_{t=0}^{T} \nabla_{\theta} \log \pi_{\theta}(a_t|s_t) A^{\pi_{\theta}}(s_t,a_t) \right],$$

где $\tau$ — траектория, а $A^{\pi_{\theta}}$ — функция преимуществапреимущество (advantage)EN advantageA = r − b или A(s, a) = Q(s, a) − V(s): насколько действие лучше среднего в этой ситуации. для текущей политики.

Словами: проходим по всем шагам траектории; на каждом шаге берём направление, в котором растёт логарифм вероятности выбранного действия, и умножаем его на преимущество — насколько это действие оказалось лучше среднего. Хорошие действия тянут параметры в свою сторону, плохие (с отрицательным $A$) — от себя.

Алгоритм градиента политики обновляет параметры политики стохастическим градиентным подъёмом по качеству политики:

$$\theta_{k+1} = \theta_k + \alpha \nabla_{\theta} J(\pi_{\theta_k})$$

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

Исследование и использование

VPG обучает стохастическую политику on-policy. Это значит, что она исследует среду, выбирая действия по самой свежей версии своей стохастической политики. Насколько случаен выбор действий, зависит и от начальных условий, и от процедуры обучения. По ходу обучения политика обычно становится всё менее случайной, потому что правило обновления подталкивает её использовать награды, которые она уже нашла. Из-за этого политика может застрять в локальном оптимуме.

Псевдокод

Алгоритм 1. Vanilla Policy Gradient

1. Вход: начальные параметры политики $\theta_0$, начальные параметры функции ценности $\phi_0$.
2. Для $k = 0, 1, 2, \dots$:
   1. Собрать набор траекторий $\mathcal{D}_k = \{\tau_i\}$, запуская в среде политику $\pi_k = \pi(\theta_k)$.
   2. Посчитать reward-to-go $\hat{R}_t$.
   3. Посчитать оценки преимущества $\hat{A}_t$ (любым методом оценки преимущества) на основе текущей функции ценности $V_{\phi_k}$.
   4. Оценить градиент политики: $$\hat{g}_k = \frac{1}{|\mathcal{D}_k|} \sum_{\tau \in \mathcal{D}_k} \sum_{t=0}^T \left. \nabla_{\theta} \log\pi_{\theta}(a_t|s_t)\right|_{\theta_k} \hat{A}_t.$$
   5. Обновить политику — либо обычным градиентным подъёмом, $$\theta_{k+1} = \theta_k + \alpha_k \hat{g}_k,$$ либо другим алгоритмом градиентного подъёма, например Adam.
   6. Подогнать функцию ценности регрессией по среднеквадратичной ошибке: $$\phi_{k+1} = \arg \min_{\phi} \frac{1}{|\mathcal{D}_k| T} \sum_{\tau \in \mathcal{D}_k} \sum_{t=0}^T\left( V_{\phi} (s_t) - \hat{R}_t \right)^2,$$ обычно каким-нибудь алгоритмом градиентного спуска.

Документация

Полезно знать

Ниже — документация реализаций VPG в Spinning Up на PyTorch и TensorFlow. Вызовы функций и docstring у них почти одинаковые, кроме деталей построения модели.

Документация: версия на PyTorch

spinup.vpg_pytorch(env_fn, actor_critic=<MLPActorCritic>, ac_kwargs={}, seed=0,
                   steps_per_epoch=4000, epochs=50, gamma=0.99, pi_lr=0.0003,
                   vf_lr=0.001, train_v_iters=80, lam=0.97, max_ep_len=1000,
                   logger_kwargs={}, save_freq=10)

Vanilla Policy Gradient (с GAE-Lambda для оценки преимущества).

Параметры:

• env_fn — функция, которая создаёт копию среды. Среда должна удовлетворять API OpenAI Gym.

• actor_critic — конструктор модуля PyTorch с методами step и act, модулем pi и модулем v. Метод step должен принимать батч наблюдений и возвращать:

  Символ	Форма	Описание
  a	(batch, act_dim)	массив NumPy с действиями для каждого наблюдения
  v	(batch,)	массив NumPy с оценками ценности для поданных наблюдений
  logp_a	(batch,)	массив NumPy с логарифмами вероятностей действий из a

  Метод act ведёт себя так же, как step, но возвращает только a.

  Прямой проход модуля pi должен принимать батч наблюдений и, необязательно, батч действий, и возвращать:

  Символ	Форма	Описание
  pi	—	объект Distribution из PyTorch с батчем распределений, описывающих политику для поданных наблюдений
  logp_a	(batch,)	необязательно (возвращается, только если подан батч действий): тензор с логарифмами вероятностей поданных действий по политике. Если действия не поданы, здесь будет None

  Прямой проход модуля v должен принимать батч наблюдений и возвращать:

  Символ	Форма	Описание
  v	(batch,)	тензор с оценками ценности для поданных наблюдений. (Критично: не забудьте сделать его плоским!)

• ac_kwargs (dict) — любые именованные аргументы для объекта ActorCritic, который вы передали в VPG.

• seed (int) — сид генераторов случайных чисел.

• steps_per_epoch (int) — число шагов взаимодействия (пар «состояние — действие») агента со средой за эпоху.

• epochs (int) — число эпох взаимодействия (равно числу обновлений политики).

• gamma (float) — коэффициент дисконтирования. (Всегда между 0 и 1.)

• pi_lr (float) — скорость обучения оптимизатора политики.

• vf_lr (float) — скорость обучения оптимизатора функции ценности.

• train_v_iters (int) — число шагов градиентного спуска по функции ценности за эпоху.

• lam (float) — лямбда для GAE-Lambda. (Всегда между 0 и 1, близко к 1.)

• max_ep_len (int) — максимальная длина траектории / эпизода / роллаута.

• logger_kwargs (dict) — именованные аргументы для EpochLogger.

• save_freq (int) — как часто (через сколько эпох) сохранять текущие политику и функцию ценности.

Содержимое сохранённой модели: версия на PyTorch

Сохранённую модель PyTorch можно загрузить командой ac = torch.load('path/to/model.pt') — получится объект актор-критика (ac) со свойствами, описанными в docstring функции vpg_pytorch.

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

actions = ac.act(torch.as_tensor(obs, dtype=torch.float32))

Документация: версия на TensorFlow

Функция spinup.vpg_tf1 принимает те же параметры, только вместо модуля PyTorch actor_critic — функция, которая строит граф вычислений актор-критика в TensorFlow 1.

Содержимое сохранённой модели: версия на TensorFlow

Граф вычислений, который сохраняет логгер, содержит:

Ключ	Значение
x	placeholder TensorFlow для входного состояния
pi	выбирает (сэмплирует) действие агента при состояниях из x
v	даёт оценку ценности для состояний из x

К сохранённой модели можно обратиться

• запустив обученную политику инструментом test_policy.py,
• или загрузив весь сохранённый граф в программу функцией restore_tf_graph.

Ссылки

Статьи по теме

• Policy Gradient Methods for Reinforcement Learning with Function Approximation, Sutton et al. 2000
• Optimizing Expectations: From Deep Reinforcement Learning to Stochastic Computation Graphs, Schulman 2016(a)
• Benchmarking Deep Reinforcement Learning for Continuous Control, Duan et al. 2016
• High Dimensional Continuous Control Using Generalized Advantage Estimation, Schulman et al. 2016(b)

Почему эти статьи?

Sutton 2000 включена, потому что это вневременная классика теории обучения с подкреплением, и в ней есть ссылки на более ранние работы, которые привели к современным градиентам политики. Schulman 2016(a) включена, потому что её глава 2 — ясное введение в теорию алгоритмов градиента политики, с псевдокодом. Duan 2016 — понятная недавняя статья-бенчмарк, которая показывает, как простой градиент политики в глубоком RL (то есть с нейросетевыми политиками и оптимизатором Adam) выглядит на фоне других алгоритмов глубокого RL. Schulman 2016(b) включена, потому что наша реализация VPG использует Generalized Advantage Estimation для вычисления градиента политики.

Другие открытые реализации

• rllab
• rllib (Ray)

---

Proximal Policy Optimization (проксимальная оптимизация политики)

Основы

(Раньше: основы TRPO)

PPOпроксимальная оптимизация политики (PPO)EN proximal policy optimization (PPO)Schulman и др., 2017. Ограничивает шаг обновления «клиппингом» отношения вероятностей. отвечает на тот же вопрос, что и TRPOTRPO (оптимизация политики в доверительной области)EN trust region policy optimization (TRPO)Schulman и др., 2015 — шаг политики с ограничением на KL к старой политике; предшественник PPO.: как сделать самый большой возможный шаг улучшения политики на тех данных, что у нас есть, и не шагнуть так далеко, чтобы случайно обрушить качество? TRPO решает эту задачу сложным методом второго порядка, а PPO — это семейство методов первого порядка, которые с помощью нескольких других приёмов держат новую политику рядом со старой. Методы PPO заметно проще реализовать, а на практике они, похоже, работают не хуже TRPO.

Есть два основных варианта PPO: PPO-Penalty и PPO-Clip.

PPO-Penalty приближённо решает задачу обновления с ограничением на KL, как TRPO, но вместо жёсткого ограничения штрафует KL-дивергенциюKL-дивергенция (дивергенция Кульбака — Лейблера)EN Kullback–Leibler divergenceKL(p ‖ q) = Σ p log(p / q) = H(p, q) − H(p) ≥ 0. Несимметрична, поэтому не настоящее расстояние. в целевой функции и по ходу обучения автоматически подстраивает коэффициент штрафа, чтобы он был нужного масштаба.

PPO-Clip не содержит в целевой функции слагаемого с KL-дивергенцией и вообще не имеет ограничения. Вместо этого он полагается на особую обрезку (clipping) в целевой функции, которая убирает у новой политики стимул уходить далеко от старой.

Здесь мы сосредоточимся только на PPO-Clip (основном варианте, который используют в OpenAI).

Коротко

• PPO — алгоритм с обучением по текущей политике (on-policy).
• PPO можно применять в средах как с дискретным, так и с непрерывным пространством действий.
• Реализация PPO в Spinning Up поддерживает распараллеливание через MPI.

Ключевые формулы

PPO-Clip обновляет политику так:

$$\theta_{k+1} = \arg \max_{\theta} \underset{s,a \sim \pi_{\theta_k}}{{\mathrm E}}\left[ L(s,a,\theta_k, \theta)\right],$$

обычно делая несколько шагов (как правило, минибатчевого) SGD для максимизации цели. Здесь $L$ задаётся так:

$$L(s,a,\theta_k,\theta) = \min\left( \frac{\pi_{\theta}(a|s)}{\pi_{\theta_k}(a|s)} A^{\pi_{\theta_k}}(s,a), \;\; \text{clip}\left(\frac{\pi_{\theta}(a|s)}{\pi_{\theta_k}(a|s)}, 1 - \epsilon, 1+\epsilon \right) A^{\pi_{\theta_k}}(s,a) \right),$$

где $\epsilon$ — (маленький) гиперпараметр, который грубо говоря задаёт, насколько далеко новой политике разрешено уйти от старой.

Словами: $\frac{\pi_{\theta}(a|s)}{\pi_{\theta_k}(a|s)}$ — отношение вероятностей: во сколько раз новая политика выбирает действие $a$ чаще старой (при $\theta = \theta_k$ оно равно 1). Первое слагаемое под минимумом — обычная «суррогатная» цель, второе — то же самое, но с отношением, зажатым в коридор $[1-\epsilon,\, 1+\epsilon]$. Минимум из двух выбирает более осторожную, пессимистичную оценку.

Выражение довольно сложное, и с первого взгляда трудно понять, что оно делает и как помогает держать новую политику рядом со старой. Оказывается, у этой цели есть заметно упрощённая версия ¹, с которой легче разобраться (и именно её мы реализуем в коде):

$$L(s,a,\theta_k,\theta) = \min\left( \frac{\pi_{\theta}(a|s)}{\pi_{\theta_k}(a|s)} A^{\pi_{\theta_k}}(s,a), \;\; g(\epsilon, A^{\pi_{\theta_k}}(s,a)) \right),$$

где

$$g(\epsilon, A) = \left\{ \begin{array}{ll} (1 + \epsilon) A & A \geq 0 \\ (1 - \epsilon) A & A < 0. \end{array} \right.$$

Чтобы понять, какую интуицию отсюда вынести, рассмотрим одну пару «состояние — действие» $(s,a)$ и разберём случаи.

Преимущество положительное. Пусть преимущество для этой пары положительно; тогда её вклад в цель сводится к

$$L(s,a,\theta_k,\theta) = \min\left( \frac{\pi_{\theta}(a|s)}{\pi_{\theta_k}(a|s)}, (1 + \epsilon) \right) A^{\pi_{\theta_k}}(s,a).$$

Раз преимущество положительно, цель растёт, если действие становится вероятнее — то есть если $\pi_{\theta}(a|s)$ увеличивается. Но минимум в этом слагаемом ограничивает, насколько цель может вырасти. Как только $\pi_{\theta}(a|s) > (1+\epsilon) \pi_{\theta_k}(a|s)$, срабатывает минимум, и слагаемое упирается в потолок $(1+\epsilon) A^{\pi_{\theta_k}}(s,a)$. Итак: новой политике невыгодно уходить далеко от старой.

Преимущество отрицательное. Пусть преимущество для этой пары отрицательно; тогда её вклад в цель сводится к

$$L(s,a,\theta_k,\theta) = \max\left( \frac{\pi_{\theta}(a|s)}{\pi_{\theta_k}(a|s)}, (1 - \epsilon) \right) A^{\pi_{\theta_k}}(s,a).$$

Раз преимущество отрицательно, цель растёт, если действие становится менее вероятным — то есть если $\pi_{\theta}(a|s)$ уменьшается. Но максимум в этом слагаемом ограничивает, насколько цель может вырасти. Как только $\pi_{\theta}(a|s) < (1-\epsilon) \pi_{\theta_k}(a|s)$, срабатывает максимум, и слагаемое упирается в потолок $(1-\epsilon) A^{\pi_{\theta_k}}(s,a)$. И снова: новой политике невыгодно уходить далеко от старой.

Числовой пример (добавлен переводчиком)

Пусть $\epsilon = 0{,}2$.

• $A = 2$, новая политика выбирает действие в 1,5 раза чаще старой (отношение $r = 1{,}5$). Без обрезки вклад был бы $1{,}5 \cdot 2 = 3$, с обрезкой — $\min(1{,}5;\ 1{,}2) \cdot 2 = 2{,}4$. Дальнейший рост $r$ цель не увеличивает, и градиент по этому примеру равен нулю.
• $A = -1$, новая политика выбирает действие вдвое реже ($r = 0{,}5$). Вклад $\max(0{,}5;\ 0{,}8) \cdot (-1) = -0{,}8$: уменьшать вероятность ниже 80 % от старой уже бесполезно.
• $A = 2$, но $r = 0{,}9$ (действие стало чуть реже). Вклад $\min(0{,}9;\ 1{,}2) \cdot 2 = 1{,}8$ — обрезка не мешает исправлять ошибку и снова повышать вероятность хорошего действия.

Итак, мы увидели, что обрезка работает как регуляризатор: она убирает у политики стимул резко меняться, а гиперпараметр $\epsilon$ соответствует тому, насколько далеко новая политика может уйти от старой, всё ещё выигрывая по цели.

Полезно знать

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

Когда освоитесь с базовой математикой и деталями реализации, стоит посмотреть другие реализации и то, как они решают эту проблему!

Исследование и использование

PPO обучает стохастическую политику on-policy. Это значит, что он исследует среду, выбирая действия по самой свежей версии своей стохастической политики. Насколько случаен выбор действий, зависит и от начальных условий, и от процедуры обучения. По ходу обучения политика обычно становится всё менее случайной, потому что правило обновления подталкивает её использовать награды, которые она уже нашла. Из-за этого политика может застрять в локальном оптимуме.

Псевдокод

Алгоритм 1. PPO-Clip

1. Вход: начальные параметры политики $\theta_0$, начальные параметры функции ценности $\phi_0$.
2. Для $k = 0, 1, 2, \dots$:
   1. Собрать набор траекторий $\mathcal{D}_k = \{\tau_i\}$, запуская в среде политику $\pi_k = \pi(\theta_k)$.
   2. Посчитать reward-to-go $\hat{R}_t$.
   3. Посчитать оценки преимущества $\hat{A}_t$ (любым методом оценки преимущества) на основе текущей функции ценности $V_{\phi_k}$.
   4. Обновить политику, максимизируя цель PPO-Clip: $$\theta_{k+1} = \arg \max_{\theta} \frac{1}{|\mathcal{D}_k| T} \sum_{\tau \in \mathcal{D}_k} \sum_{t=0}^T \min\left( \frac{\pi_{\theta}(a_t|s_t)}{\pi_{\theta_k}(a_t|s_t)} A^{\pi_{\theta_k}}(s_t,a_t), \;\; g(\epsilon, A^{\pi_{\theta_k}}(s_t,a_t)) \right),$$ обычно стохастическим градиентным подъёмом с Adam.
   5. Подогнать функцию ценности регрессией по среднеквадратичной ошибке: $$\phi_{k+1} = \arg \min_{\phi} \frac{1}{|\mathcal{D}_k| T} \sum_{\tau \in \mathcal{D}_k} \sum_{t=0}^T\left( V_{\phi} (s_t) - \hat{R}_t \right)^2,$$ обычно каким-нибудь алгоритмом градиентного спуска.

Документация

Полезно знать

Ниже — документация реализаций PPO в Spinning Up на PyTorch и TensorFlow. Вызовы функций и docstring у них почти одинаковые, кроме деталей построения модели.

Документация: версия на PyTorch

spinup.ppo_pytorch(env_fn, actor_critic=<MLPActorCritic>, ac_kwargs={}, seed=0,
                   steps_per_epoch=4000, epochs=50, gamma=0.99, clip_ratio=0.2,
                   pi_lr=0.0003, vf_lr=0.001, train_pi_iters=80, train_v_iters=80,
                   lam=0.97, max_ep_len=1000, target_kl=0.01, logger_kwargs={},
                   save_freq=10)

Proximal Policy Optimization (с обрезкой), с ранней остановкой по приближённой KL.

Параметры:

• env_fn — функция, которая создаёт копию среды. Среда должна удовлетворять API OpenAI Gym.
• actor_critic — конструктор модуля PyTorch с методами step и act, модулем pi и модулем v. Требования к step, act, pi и v те же, что у VPG (см. таблицы выше): step возвращает a, v, logp_a; pi возвращает распределение и, если поданы действия, их логарифмы вероятностей; v возвращает плоский тензор оценок ценности.
• ac_kwargs (dict) — любые именованные аргументы для объекта ActorCritic, который вы передали в PPO.
• seed (int) — сид генераторов случайных чисел.
• steps_per_epoch (int) — число шагов взаимодействия (пар «состояние — действие») агента со средой за эпоху.
• epochs (int) — число эпох взаимодействия (равно числу обновлений политики).
• gamma (float) — коэффициент дисконтирования. (Всегда между 0 и 1.)
• clip_ratio (float) — гиперпараметр обрезки в цели политики. Грубо говоря: насколько далеко новая политика может уйти от старой, всё ещё выигрывая (улучшая целевую функцию)? Новая политика может уйти и дальше, чем говорит clip_ratio, но цели это больше не помогает. (Обычно маленький, от 0,1 до 0,3.) Обычно обозначается $\epsilon$.
• pi_lr (float) — скорость обучения оптимизатора политики.
• vf_lr (float) — скорость обучения оптимизатора функции ценности.
• train_pi_iters (int) — максимальное число шагов градиентного спуска по функции потерь политики за эпоху. (Из-за ранней остановки оптимизатор может сделать меньше.)
• train_v_iters (int) — число шагов градиентного спуска по функции ценности за эпоху.
• lam (float) — лямбда для GAE-Lambda. (Всегда между 0 и 1, близко к 1.)
• max_ep_len (int) — максимальная длина траектории / эпизода / роллаута.
• target_kl (float) — грубо говоря, какую KL-дивергенцию между новой и старой политикой после обновления мы считаем допустимой. Используется для ранней остановки. (Обычно маленькая, 0,01 или 0,05.)
• logger_kwargs (dict) — именованные аргументы для EpochLogger.
• save_freq (int) — как часто (через сколько эпох) сохранять текущие политику и функцию ценности.

Содержимое сохранённой модели: версия на PyTorch

Сохранённую модель PyTorch можно загрузить командой ac = torch.load('path/to/model.pt') — получится объект актор-критика (ac) со свойствами, описанными в docstring функции ppo_pytorch.

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

actions = ac.act(torch.as_tensor(obs, dtype=torch.float32))

Документация и сохранённая модель: версия на TensorFlow

Функция spinup.ppo_tf1 принимает те же параметры (с функцией построения графа вместо модуля PyTorch). Граф, который сохраняет логгер, содержит те же ключи, что у VPG: x (placeholder для состояния), pi (сэмплирует действие), v (оценка ценности). Запускать его можно через test_policy.py или загружать через restore_tf_graph.

Ссылки

Статьи по теме

• Proximal Policy Optimization Algorithms, Schulman et al. 2017
• High Dimensional Continuous Control Using Generalized Advantage Estimation, Schulman et al. 2016
• Emergence of Locomotion Behaviours in Rich Environments, Heess et al. 2017

Почему эти статьи?

Schulman 2017 включена, потому что это исходная статья с описанием PPO. Schulman 2016 включена, потому что наша реализация PPO использует Generalized Advantage Estimation для вычисления градиента политики. Heess 2017 включена, потому что в ней есть крупный эмпирический анализ поведения, которому учатся агенты PPO в сложных средах (хотя там используется PPO-Penalty, а не PPO-Clip).

Другие открытые реализации

• Baselines
• ModularRL (Осторожно: здесь реализован PPO-Penalty, а не PPO-Clip.)
• rllab (Осторожно: здесь реализован PPO-Penalty, а не PPO-Clip.)
• rllib (Ray)

---

Ключевые строки кода PPO (добавлено переводчиком)

Чтобы связать формулы с кодом, ниже — три фрагмента из файла spinup/algos/pytorch/ppo/ppo.py (MIT). Код приведён как в оригинале, русские комментарии добавлены переводчиком. Целиком файл запускать не нужно: он рассчитан на старый Gym и MPI. Его полезно читать.

1. GAE и reward-to-go в конце траектории

def finish_path(self, last_val=0):
    # срез буфера с текущей траекторией
    path_slice = slice(self.path_start_idx, self.ptr)
    # last_val = 0, если эпизод закончился «смертью», иначе V(s_T) —
    # так мы достраиваем отдачу за пределами обрезанного эпизода (бутстрэп)
    rews = np.append(self.rew_buf[path_slice], last_val)
    vals = np.append(self.val_buf[path_slice], last_val)

    # следующие две строки — оценка преимущества GAE-Lambda:
    # δ_t = r_t + γ·V(s_{t+1}) − V(s_t), затем A_t = Σ (γλ)^l · δ_{t+l}
    deltas = rews[:-1] + self.gamma * vals[1:] - vals[:-1]
    self.adv_buf[path_slice] = core.discount_cumsum(deltas, self.gamma * self.lam)

    # reward-to-go R̂_t = Σ γ^l · r_{t+l} — цель для регрессии функции ценности
    self.ret_buf[path_slice] = core.discount_cumsum(rews, self.gamma)[:-1]

    self.path_start_idx = self.ptr

Перед обновлением преимущества ещё и нормируются: из них вычитают среднее и делят на стандартное отклонение («трюк нормализации преимущества» в методе get).

2. Функция потерь политики PPO-Clip

def compute_loss_pi(data):
    obs, act, adv, logp_old = data['obs'], data['act'], data['adv'], data['logp']

    # Потери политики
    pi, logp = ac.pi(obs, act)
    ratio = torch.exp(logp - logp_old)          # π_θ(a|s) / π_θk(a|s) через разность логарифмов
    clip_adv = torch.clamp(ratio, 1-clip_ratio, 1+clip_ratio) * adv
    loss_pi = -(torch.min(ratio * adv, clip_adv)).mean()   # минус: оптимизатор минимизирует

    # Полезная дополнительная информация
    approx_kl = (logp_old - logp).mean().item()  # грубая оценка KL(старая || новая)
    ent = pi.entropy().mean().item()             # энтропия политики
    clipped = ratio.gt(1+clip_ratio) | ratio.lt(1-clip_ratio)
    clipfrac = torch.as_tensor(clipped, dtype=torch.float32).mean().item()  # доля обрезанных примеров
    pi_info = dict(kl=approx_kl, ent=ent, cf=clipfrac)

    return loss_pi, pi_info

Обратите внимание: здесь используется полная форма с clamp, а не функция $g(\epsilon, A)$. Обе дают одно и то же значение; упрощённая форма нужна, чтобы проще рассуждать.

3. Обновление: много шагов по одним и тем же данным и ранняя остановка

# Обучаем политику несколькими шагами градиентного спуска
for i in range(train_pi_iters):
    pi_optimizer.zero_grad()
    loss_pi, pi_info = compute_loss_pi(data)
    kl = mpi_avg(pi_info['kl'])
    if kl > 1.5 * target_kl:                 # политика ушла слишком далеко — стоп
        logger.log('Early stopping at step %d due to reaching max kl.'%i)
        break
    loss_pi.backward()
    mpi_avg_grads(ac.pi)    # усредняем градиенты между процессами MPI
    pi_optimizer.step()

# Обучение функции ценности: регрессия V(s) на reward-to-go
for i in range(train_v_iters):
    vf_optimizer.zero_grad()
    loss_v = compute_loss_v(data)            # ((ac.v(obs) - ret)**2).mean()
    loss_v.backward()
    mpi_avg_grads(ac.v)
    vf_optimizer.step()

Что изменится в RLHF и GRPO

В RLHF (урок 6.3.2) те же три фрагмента применяются к токенам: наблюдение — промпт и уже сгенерированная часть ответа, действие — следующий токен, logp — лог-вероятность токена, к награде добавляется KL-штраф к референсной модели. В GRPO (урок 6.4.1) исчезает второй цикл — функции ценности нет, а преимущество считают нормировкой наград внутри группы ответов на один промпт; функция потерь с clamp остаётся почти без изменений.

Что дальше

• TRPO, DDPG, TD3, SAC — остальные четыре алгоритма Spinning Up.
• Упражнения: упражнение 1.2 — написать гауссову политику для этого PPO.
• Ключевые статьи — откуда что взялось.

Footnotes

1. Вывод упрощённой формы цели PPO-Clip — в этой заметке. ↩