Policy gradient і PPO: навчання безпосередньо політики

Стохастична політика π_θ(a | s) — це мережа з параметрами θ, яка для стану s повертає розподіл імовірностей на діях: для скінченної кількості дій — softmax виходів мережі, для неперервних дій — скажімо, середнє й дисперсію нормального розподілу. Мета — максимізувати очікувану віддачу J(θ) = E_πθ[G₀] градієнтним підйомом (gradient ascent) θ ← θ + α ∇J(θ).

Винагорода +1 за кожен крок, дисконт γ = 0,99. Якщо до кінця епізоду лишається T − t кроків, віддача — це геометричний ряд:

G_t = 1 + γ + γ² + … + γ^T−t−1 = (1 − γ^T−t) / (1 − γ).

На початку епізоду, який протримається всі 500 кроків: G₀ = (1 − 0,99⁵⁰⁰) / 0,01 = 99,34. За двадцять кроків до падіння: G = (1 − 0,99²⁰) / 0,01 = 18,21. Крок безпосередньо перед падінням: G = 1. Отже, віддача завжди додатна й ніколи не перевищить 1 / (1 − γ) = 100 — це знадобиться в розділі про базовий рівень.

Градієнт очікуваної віддачі можна записати як математичне сподівання, яке вдається оцінити з зіграних епізодів:

∇J(θ) = E_πθ[ ∑_t ∇ log π_θ(a_t | s_t) · G_t ].

REINFORCE (Рональд Вільямс, 1992) робить це буквально: зіграй епізоди, для кожного кроку обчисли віддачу G_t від цього кроку далі й зсунь θ у напрямку ∇ log π(a_t | s_t) · G_t. Для softmax ∇ log π(a) за виходами мережі — це просто (1_a − π).

Епізод τ = (s₀, a₀, s₁, a₁, …) за політики π_θ має ймовірність P_θ(τ) = p(s₀) ∏_t π_θ(a_t | s_t) P(s_t+1 | s_t, a_t) і віддачу G(τ). Тоді

∇J(θ) = ∇ ∑_τ P_θ(τ) G(τ) = ∑_τ P_θ(τ) · ∇ log P_θ(τ) · G(τ) = E_{τ ∼ πθ}[ ∇ log P_θ(τ) · G(τ) ],

бо ∇P = P · ∇ log P. Доведення (ескіз): логарифм перетворює добуток на суму, log P_θ(τ) = log p(s₀) + ∑_t log π_θ(a_t | s_t) + ∑_t log P(s_t+1 | s_t, a_t). Від θ залежить лише середня сума, тож фізику середовища P знати не потрібно: ∇ log P_θ(τ) = ∑_t ∇ log π_θ(a_t | s_t). Нарешті для кожного a_t відкидаємо винагороди до кроку t: вони не залежать від a_t, тому їхній внесок має нульове математичне сподівання (див. теорему в наступному розділі). Лишається G_t. ∎ (Точне формулювання для дисконтованої віддачі дещо тонше; тут спрощено.)

Дві дії, виходи мережі z = (0; 0), тож π = softmax(z) = (0,50; 0,50). Агент зіграв дію 1 («праворуч»), і віддача від цього кроку склала G = 10. Для softmax ∇_z log π(a) = 1_a − π, отже

∇_z log π(1) · G = ((0; 1) − (0,50; 0,50)) · 10 = (−5; 5).

Крок підйому з α = 0,1: z ← (0; 0) + 0,1 · (−5; 5) = (−0,5; 0,5). Нова політика π = (e^−0,5; e^0,5) / (e^−0,5 + e^0,5) = (0,269; 0,731). Дія, яка вдалася, тепер імовірніша. (У справжній мережі змінюються ваги θ, а до z градієнт доходить зворотним поширенням; напрямок той самий.)

Від G_t можна відняти довільну функцію стану b(s_t), не змінивши математичного сподівання оцінки градієнта, бо E[∇ log π(a | s)] = 0. Зате дисперсія оцінки може помітно зменшитися. Найкращий поширений вибір — цінність стану V(s), яку вивчає друга мережа; різницю A(s, a) = Q(s, a) − V(s) називають перевагою (advantage). Методи з вивченою політикою («актор») і цінністю («критик») називають actor-critic (актор–критик).

Для кожного стану s виконується E_{a ∼ πθ}[∇ log π_θ(a | s)] = 0, а отже E[∇ log π_θ(a | s) · b(s)] = 0 для довільної функції b, яка не залежить від a.

Доведення: ∑_a π_θ(a | s) ∇ log π_θ(a | s) = ∑_a ∇π_θ(a | s) = ∇ ∑_a π_θ(a | s) = ∇1 = 0. Член b(s) для фіксованого s — стала, тож її можна винести за знак суми. ∎

Один стан, дві дії з π = (0,5; 0,5). Дія 1 дає віддачу 12, дія 0 — віддачу 10; обидві додатні, як на візку. Компонента оцінки градієнта за виходом z₁ з однієї вибірки — це (1_a=1 − 0,5) · (G − b).

Без базового рівня (b = 0): якщо випаде дія 1, оцінка 0,5 · 12 = 6; якщо дія 0, то −0,5 · 10 = −5. Середнє (6 − 5) / 2 = 0,5, дисперсія ((6 − 0,5)² + (−5 − 0,5)²) / 2 = 30,25.

З базовим рівнем b = 11 (середня віддача): 0,5 · (12 − 11) = 0,5 і −0,5 · (10 − 11) = 0,5. Середнє й далі 0,5, дисперсія 0.

Справжній градієнт — π₁(1 − π₁) · (12 − 10) = 0,25 · 2 = 0,5. Обидві оцінки в середньому цілять правильно, але без планки одна вибірка промахується на ±5,5, а з планкою — зовсім ні. Насправді дисперсія до нуля не падає (віддачі випадкові, а станів багато), але принцип той самий.

З K = 40 пакетів маємо оцінки градієнта ĝ⁽¹⁾, …, ĝ^(K) (вектори всіх ваг мережі). Для кожної компоненти j обчислюємо дисперсію за пакетами й підсумовуємо:

∑_j (1/K) ∑_k (ĝ_j^(k) − ḡ_j)²,   ḡ_j = (1/K) ∑_k ĝ_j^(k).

Це слід коваріаційної матриці оцінки — одне число, яке показує, наскільки оцінка тремтить. Спад: 1 − 0,131 / 0,177 = 0,26, тобто 26 %.

PPO (Джон Шульман та ін., OpenAI, 2017) збирає пакет переходів старою політикою π_old, обчислює переваги Â_t, а потім кілька епох максимізує обрізану ціль

L(θ) = E_t[ min( r_t(θ) Â_t, clip(r_t(θ), 1 − ε, 1 + ε) Â_t ) ],   r_t(θ) = π_θ(a_t | s_t) / π_old(a_t | s_t),

зазвичай з ε = 0,2. Щойно ймовірність дії зміниться більш ніж на ε в напрямку, який покращує ціль, градієнт від цієї вибірки зникає. Переваги зазвичай оцінюють методом GAE (generalized advantage estimation) — зваженим середнім TD-похибок із параметром λ.

Стара політика дала зіграній дії ймовірність 0,5. Після кількох кроків навчання це 0,65, тож r = 0,65 / 0,5 = 1,3 і clip(1,3; 0,8; 1,2) = 1,2.

Перевага +2: min(1,3 · 2; 1,2 · 2) = min(2,6; 2,4) = 2,4. Переміг обрізаний член, який від θ не залежить, — градієнт від цієї вибірки нульовий. Дія вже на 30 % частіша, і більше не буде.

Перевага −2 (те саме відношення): min(−2,6; −2,4) = −2,6. Переміг необрізаний член: політика помилково зробила погану дію ймовірнішою, тож градієнт і далі тягне її назад.

Перевага +2, але ймовірність упала до 0,35: r = 0,7, min(1,4; 0,8 · 2 = 1,6) = 1,4. Знову необрізаний член — добра дія, яку політика придушила, повертається. Обрізання гальмує лише ті кроки, які покращили б ціль «занадто».

KL-дивергенція з частини про теорію інформації, тут із натуральним логарифмом (одиниця — нат; 1 нат = 1 / ln 2 ≈ 1,44 біта), усереднена за станами з пакета:

D_KL(π_old ‖ π_θ) = E_{s, a ∼ πold}[ log π_old(a | s) − log π_θ(a | s) ] ≈ (1/N) ∑_i=1^N [ log π_old(a_i | s_i) − log π_θ(a_i | s_i) ].

Дії a_i зіграла стара політика, тож середнє за пакетом — незміщена оцінка. Це середній логарифм відношення, а отже e^−D — середнє геометричне відношення π_θ / π_old за зібраними діями.

Що це означає в числах: у PPO щонайбільше 0,06 ната, e^−0,06 = 0,94 — дії зі старого пакета зазвичай на кілька відсотків менш імовірні, дані й далі щось кажуть про нову політику. Без обрізання аж 24,9 ната: e^−24,9 ≈ 1,5 · 10⁻¹¹, нова політика майже ніколи не зіграла б дій, з яких щойно навчалася.

Для запиту x і відповіді y, згенерованої моделлю π_θ, яку донавчають, PPO отримує винагороду

R(x, y) = r_φ(x, y) − β · log [ π_θ(y | x) / π_ref(y | x) ],

де r_φ — модель винагороди, π_ref — заморожена модель після донавчання на прикладах (SFT), а β — сила штрафу (так це записує InstructGPT, Лонг Оуян та ін., 2022). Математичне сподівання другого члена за відповідями — це точно β · D_KL(π_θ ‖ π_ref). Епізод — одна відповідь, дії — токени, винагорода r_φ надходить лише після останнього токена; решта — це PPO саме так, як описано вище.

Модель відповідає на задачу чотири рази, правильно лише вперше: винагороди (1; 0; 0; 0). Середнє 0,25, стандартне відхилення √((0,75² + 3 · 0,25²) / 4) = √0,1875 = 0,433.

Переваги: (1 − 0,25) / 0,433 = 1,732 для правильної відповіді і (0 − 0,25) / 0,433 = −0,577 для кожної з трьох неправильних. Сума дорівнює нулю: правильна підсилюється, неправильні придушуються. Якби всі чотири відповіді були правильні (або всі неправильні), усі переваги нульові й запитання нічого не навчить — модель учиться лише із запитань, на які вона іноді влучає, а іноді ні. (Ділимо на G; деякі реалізації ділять на G − 1, і тоді числа трохи відрізняються.)