Policy gradient a PPO: učení přímo strategie

Nechť (S, A, P, R, γ) je MDP (značení z dílu Markovovy rozhodovací procesy) s konečnou množinou akcí A a počátečním rozdělením p(s₀) na S, a nechť θ ∈ ℝ^m jsou parametry. Parametrizovaná (stochastická) strategie je systém strategií (π_θ)_{θ ∈ ℝ^m}, kde pro každé θ a každé s ∈ S je π_θ(· | s) rozdělení pravděpodobnosti na A a π_θ(a | s) je diferencovatelná funkce θ. Pro konečnou A se bere síť f_θ : S → ℝ^|A| a softmax jejích výstupů:

π_θ(a | s) = e^f_θ(s)_a/∑_{b ∈ A} e^f_θ(s)_b;

pro spojité akce síť vrací třeba střed a rozptyl normálního rozdělení. Cíl je očekávaný výnos jako funkce parametrů

J : ℝ^m → ℝ,   J(θ) = E_πθ[G₀],

a hledá se θ s co největším J(θ) gradientním výstupem θ ← θ + α ∇J(θ) s rychlostí učení α > 0 (∇ značí gradient podle θ).

Odměna je +1 za každý krok a diskont γ = 0,99. Zbývá-li do konce epizody T − t kroků, je výnos geometrická řada:

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

Na začátku epizody, která vydrží celých 500 kroků: G₀ = 1 − 0,99⁵⁰⁰/0,01 = 99,34. Dvacet kroků před pádem: G = 1 − 0,99²⁰/0,01 = 18,21. Krok těsně před pádem: G = 1. Výnos je tedy vždy kladný a nikdy nepřesáhne 1/1 − γ = 100 — to se bude hodit v oddílu o základní úrovni.

Věta o gradientu strategie. Nechť (S, A, P, R, γ) je MDP s konečnými S a A a γ ∈ [0, 1], ve kterém každá epizoda s₀, a₀, r₁, s₁, …, s_T skončí nejpozději po T_max krocích (T ≤ T_max), a nechť π_θ je parametrizovaná strategie s π_θ(a | s) > 0 pro všechna s, a, θ. Pak

∇J(θ) = E_πθ[ ∑_t=0^T−1 γ^t ∇ log π_θ(a_t | s_t) · G_t ],   G_t = ∑_k=t^T−1 γ^k−t r_k+1.

Pro γ = 1 činitel γ^t zmizí a dostaneme tvar v rozboru níže. Praktické implementace (i naše) ho vynechávají i pro γ < 1; pak výsledek už není nestranný odhad ∇J(θ) (obecně dokonce není gradientem žádné funkce; Chris Nota a Philip Thomas, 2020), protože pozdější kroky epizody dostanou větší váhu, než jim podle věty patří — v praxi to přesto funguje dobře.

REINFORCE (Ronald Williams, 1992) je algoritmus:

1. zvol počáteční θ ∈ ℝ^m (náhodné váhy sítě) a rychlost učení α > 0;
2. odehraj strategií π_θ epizody (u nás dohromady aspoň 2 000 kroků);
3. v každé epizodě spočítej odzadu G_T = 0 a G_t = r_t+1 + γ G_t+1 pro t = T − 1, …, 0;
4. spočítej odhad gradientu ĝ = průměr hodnot ∇ log π_θ(a_t | s_t) · G_t přes všechny odehrané kroky;
5. θ ← θ + α ĝ, data zahoď a pokračuj krokem 2, dokud nevyčerpáš rozpočet kroků v prostředí.

Pro softmax π = softmax(z) s výstupy sítě z ∈ ℝ^|A| je gradient podle z roven ∇_z log π_a = 1_a − π, kde 1_a je vektor s jedničkou na místě akce a a nulami jinde.

Nechť platí předpoklady věty o gradientu strategie. Epizodu (trajektorii) τ = (s₀, a₀, s₁, a₁, …, s_T) má strategie π_θ s pravděpodobností

P_θ(τ) = p(s₀) ∏_t=0^T−1 π_θ(a_t | s_t) P(s_t+1 | s_t, a_t)

a nechť G(τ) = G₀ je její výnos; trajektorií je konečně mnoho, takže J(θ) = ∑_τ P_θ(τ) G(τ). Pak

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

protože ∇P = P · ∇ log P. Důkaz (náčrt): logaritmus změní součin na součet, log P_θ(τ) = log p(s₀) + ∑_t log π_θ(a_t | s_t) + ∑_t log P(s_t+1 | s_t, a_t). Na θ závisí jen prostřední součet, takže fyziku prostředí P znát nemusíme: ∇ log P_θ(τ) = ∑_t ∇ log π_θ(a_t | s_t). Nakonec u každého a_t zahodíme odměny před krokem t: na a_t nezávisejí, a jejich příspěvek má proto nulovou střední hodnotu (viz věta v dalším oddílu). Ze součtu ∑_k γ^k r_k+1 tak u kroku t zbude ∑_{k ≥ t} γ^k r_k+1 = γ^t G_t, což je tvrzení věty o gradientu strategie. ∎

Dvě akce, výstupy sítě z = (0; 0), takže π = softmax(z) = (0,50; 0,50). Agent zahrál akci 1 („doprava“) a výnos od toho kroku byl G = 10. Pro softmax je ∇_z log π(a) = 1_a − π, tedy

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

(barvy podle rozboru vzorce výše; výsledek je odhad ∇J z jediného kroku).

Krok výstupem s α = 0,1: z ← (0; 0) + 0,1 · (−5; 5) = (−0,5; 0,5). Nová strategie π = (e^−0,5; e^0,5)/e^−0,5 + e^0,5 = (0,269; 0,731). Akce, která vyšla, je teď pravděpodobnější. (Ve skutečné síti se mění váhy θ a k z se gradient dostane zpětným šířením; směr je týž.)

Nechť (S, A, P, R, γ) je MDP a π_θ parametrizovaná strategie.

1. Základní úroveň (baseline) je libovolná funkce b : S → ℝ, která závisí jen na stavu, ne na akci. Odhad gradientu se základní úrovní je ĝ = ∑_t ∇ log π_θ(a_t | s_t) · (G_t − b(s_t)); podle věty níže má stejnou střední hodnotu jako odhad bez b, jen se tím může výrazně zmenšit rozptyl.
2. Výhoda (advantage) strategie π je funkce A^π : S × A → ℝ, A^π(s, a) = Q^π(s, a) − V^π(s), kde Q^π a V^π jsou hodnotové funkce z dílu o MDP (písmeno A s argumenty tu značí výhodu, ne množinu akcí). Nejlepší běžná volba základní úrovně je b = V^π_θ, v praxi odhadnutá druhou sítí; pak G_t − b(s_t) je odhad výhody.
3. Actor-critic je metoda, která současně učí parametrizovanou strategii π_θ („herec“) a odhad hodnotové funkce V ≈ V^π_θ („kritik“) a odhad kritika používá v odhadu gradientu strategie.

Nechť A je konečná množina akcí, π_θ parametrizovaná strategie s π_θ(a | s) > 0, s ∈ S stav a b : S → ℝ funkce. Pak

1. E_{a ∼ πθ(· | s)}[∇ log π_θ(a | s)] = ∑_{a ∈ A} π_θ(a | s) ∇ log π_θ(a | s) = 0;
2. E_{a ∼ πθ(· | s)}[∇ log π_θ(a | s) · b(s)] = 0.

Odečtení b(s_t) od G_t tedy střední hodnotu odhadu gradientu nezmění (stačí podmínit stavem s_t).

Důkaz: ∑_a π_θ(a | s) ∇ log π_θ(a | s) = ∑_a ∇π_θ(a | s) = ∇ ∑_a π_θ(a | s) = ∇1 = 0. Člen b(s) je pro pevné s konstanta, takže se dá vytknout před součet. ∎

Jeden stav, dvě akce s π = (0,5; 0,5). Akce 1 vynese výnos 12, akce 0 výnos 10 — obojí kladné jako na vozíku. Složka odhadu gradientu podle výstupu z₁ je z jednoho vzorku (1_a=1 − 0,5) · (G − b) (barvy podle vzorce výše).

Bez základní úrovně (b = 0): padne-li akce 1, odhad je 0,5 · 12 = 6; padne-li akce 0, je −0,5 · 10 = −5. Průměr 6 − 5/2 = 0,5, rozptyl (6 − 0,5)² + (−5 − 0,5)²/2 = 30,25.

Se základní úrovní b = 11 (průměrný výnos): 0,5 · (12 − 11) = 0,5 a −0,5 · (10 − 11) = 0,5. Průměr je pořád 0,5, rozptyl 0.

Skutečný gradient je π₁(1 − π₁) · (12 − 10) = 0,25 · 2 = 0,5. Oba odhady míří v průměru správně, ale bez laťky se jeden vzorek trefí vedle o ±5,5, s laťkou vůbec. Ve skutečnosti rozptyl na nulu neklesne (výnosy jsou náhodné a stavů je mnoho), ale princip je týž.

Nechť K ≥ 1 (u nás K = 40) a ĝ⁽¹⁾, …, ĝ^(K) ∈ ℝ^m jsou odhady gradientu z K nezávislých dávek (vektory přes všech m vah sítě). Pro j = 1, …, m označme ḡ_j = 1/K ∑_k=1^K ĝ_j^(k) průměr j-té složky. Součet rozptylů složek je číslo

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

Je to stopa (výběrové) kovarianční matice odhadu — jedno číslo pro to, jak moc se odhad třese. Pokles: 1 − 0,131/0,177 = 0,26, tedy 26 %.

Nechť π_θ je parametrizovaná strategie, V : S → ℝ naučený odhad hodnotové funkce (kritik), ε > 0 (typicky ε = 0,2), γ ∈ [0, 1] a λ ∈ [0, 1]. Pro x ∈ ℝ označme clip(x, 1 − ε, 1 + ε) = min(max(x, 1 − ε), 1 + ε) číslo x sevřené do intervalu [1 − ε, 1 + ε]. PPO (Proximal Policy Optimization; John Schulman a kol., OpenAI, 2017) opakuje tyto kroky:

1. Polož θ_old = θ, π_old = π_θold a nasbírej touto strategií dávku kroků (s_t, a_t, r_t+1, s_t+1).
2. Pro každý krok spočítej TD chybu δ_t = r_t+1 + γ V(s_t+1) − V(s_t) (po pádu V(s_t+1) = 0) a odhad výhody metodou GAE (generalized advantage estimation) Â_t = ∑_l≥0 (γλ)^l δ_t+l, kde součet končí s koncem epizody.
3. Několik epoch (u nás 10) maximalizuj gradientními kroky vzhledem k θ ořezaný cíl L(θ) = E_t[ min( r_t(θ) Â_t, clip(r_t(θ), 1 − ε, 1 + ε) Â_t ) ],   r_t(θ) = π_θ(a_t | s_t)/π_old(a_t | s_t), kde E_t je průměr přes kroky dávky (minidávky) a r_t(θ) je poměr pravděpodobností (nezaměňovat s odměnou r_t+1); zároveň uč kritika V zmenšováním průměru (V(s_t) − y_t)², kde y_t = Â_t + V_old(s_t) je cíl spočtený v kroku 2 z hodnot kritika před učením.
4. Dávku zahoď a pokračuj krokem 1, dokud nevyčerpáš rozpočet kroků v prostředí.

Jakmile se pravděpodobnost akce změní o víc než faktor 1 ± ε ve směru, který cíl zlepšuje, gradient z toho vzorku zmizí.

Stará strategie dala zahrané akci pravděpodobnost 0,5. Po několika krocích učení je to 0,65, takže r = 0,65/0,5 = 1,3 a clip(1,3; 0,8; 1,2) = 1,2 (barvy podle rozboru ořezaného cíle výše).

Výhoda +2: min(1,3 · 2; 1,2 · 2) = min(2,6; 2,4) = 2,4. Vyhrál ořezaný člen, který na θ nezávisí — gradient z tohoto vzorku je nula. Akce už je o 30 % častější, víc ne.

Výhoda −2 (týž poměr): min(1,3 · (−2); 1,2 · (−2)) = min(−2,6; −2,4) = −2,6. Vyhrál neořezaný člen: strategie omylem zpravděpodobnila špatnou akci, a gradient ji proto dál stahuje zpátky.

Výhoda +2, ale pravděpodobnost klesla na 0,35: r = 0,35/0,5 = 0,7, min(0,7 · 2 = 1,4; 0,8 · 2 = 1,6) = 1,4. Zase neořezaný člen — dobrá akce, kterou strategie potlačila, se vrací. Ořez brzdí jen kroky, které by cíl zlepšily „příliš“.

Nechť π_old a π_θ jsou strategie na konečné množině akcí A s π_θ(a | s) > 0 a nechť (s_i, a_i), i = 1, …, N, jsou kroky dávky nasbírané strategií π_old. KL divergence z dílu o teorii informace, tady s přirozeným logaritmem (jednotka nat; 1 nat = 1/ln 2 ≈ 1,44 bitu), průměrovaná přes stavy, které stará strategie navštěvuje, je

D_KL(π_old ‖ π_θ) = E_s[ ∑_{a ∈ A} π_old(a | s) log π_old(a | s)/π_θ(a | s) ] = E_{s, a ∼ πold}[ log π_old(a | s) − log π_θ(a | s) ]

a její odhad z dávky je

D̂ = 1/N ∑_i=1^N [ log π_old(a_i | s_i) − log π_θ(a_i | s_i) ].

Akce a_i zahrála stará strategie, takže průměr přes dávku je nestranný odhad. Je to průměrný logaritmus poměru, a tedy e^−D̂ je geometrický průměr poměru π_θ/π_old přes nasbírané akce.

Co to znamená v číslech: u PPO nejvýš 0,06 natu, e^−0,06 = 0,94 — akce ze staré dávky jsou typicky o pár procent méně pravděpodobné, data o nové strategii pořád vypovídají. Bez ořezu až 24,9 natu: e^−24,9 ≈ 1,5 · 10⁻¹¹, nová strategie by akce, ze kterých se právě učila, téměř nikdy nezahrála.

Nechť x je dotaz a y = (y₁, …, y_n) odpověď (posloupnosti tokenů) vygenerovaná laděným jazykovým modelem π_θ, kde π_θ(y | x) = ∏_k=1ⁿ π_θ(y_k | x, y₁, …, y_k−1). Nechť r_φ(x, y) ∈ ℝ je výstup modelu odměny, π_ref zmrazený model po doladění na příkladech (SFT) a β > 0 síla trestu. PPO pak dostane odměnu

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

(tak to zapisuje InstructGPT, Long Ouyang a kol., 2022). Střední hodnota druhého členu přes odpovědi je přesně β · D_KL(π_θ ‖ π_ref). Epizoda je jedna odpověď, akce jsou tokeny, odměna r_φ přijde až po posledním tokenu; zbytek je PPO tak, jak je popsané výše.

Model odpoví na úlohu čtyřikrát, správně jen poprvé: odměny (1; 0; 0; 0). Průměr 0,25, směrodatná odchylka √0,75² + 3 · 0,25²/4 = √0,1875 = 0,433.

Výhody (barvy podle vzorce výše): 1 − 0,25/0,433 = 1,732 pro správnou odpověď a 0 − 0,25/0,433 = −0,577 pro každou ze tří špatných. Součet je nula: správná se posílí, špatné potlačí. Kdyby byly všechny čtyři odpovědi správně (nebo všechny špatně), jsou všechny výhody nulové a otázka nic nenaučí — učí se jen z otázek, kde se model někdy trefí a někdy ne. (Dělíme G; některé implementace dělí G − 1, čísla se pak mírně liší.)