Policy gradient a PPO: učenie priamo stratégie

Nech (S, A, P, R, γ) je MDP (značenie z dielu Markovove rozhodovacie procesy) s konečnou množinou akcií A a počiatočným rozdelením p(s₀) na S, a nech θ ∈ ℝ^m sú parametre. Parametrizovaná (stochastická) stratégia je systém stratégií (π_θ)_{θ ∈ ℝ^m}, kde pre každé θ a každé s ∈ S je π_θ(· | s) rozdelenie pravdepodobnosti na A a π_θ(a | s) je diferencovateľná funkcia θ. Pre konečnú A sa berie sieť f_θ : S → ℝ^|A| a softmax jej výstupov:

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

pre spojité akcie sieť vracia napríklad stred a rozptyl normálneho rozdelenia. Cieľ je očakávaný výnos ako funkcia parametrov

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

a hľadá sa θ s čo najväčším J(θ) gradientným vzostupom θ ← θ + α ∇J(θ) s rýchlosťou učenia α > 0 (∇ značí gradient podľa θ).

Odmena je +1 za každý krok a diskont γ = 0,99. Ak do konca epizódy zostáva T − t krokov, je výnos geometrický rad:

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

Na začiatku epizódy, ktorá vydrží celých 500 krokov: G₀ = 1 − 0,99⁵⁰⁰/0,01 = 99,34. Dvadsať krokov pred pádom: G = 1 − 0,99²⁰/0,01 = 18,21. Krok tesne pred pádom: G = 1. Výnos je teda vždy kladný a nikdy neprekročí 1/1 − γ = 100 — to sa bude hodiť v oddiele o základnej úrovni.

Veta o gradiente stratégie. Nech (S, A, P, R, γ) je MDP s konečnými S a A a γ ∈ [0, 1], v ktorom každá epizóda s₀, a₀, r₁, s₁, …, s_T skončí najneskôr po T_max krokoch (T ≤ T_max), a nech π_θ je parametrizovaná stratégia s π_θ(a | s) > 0 pre všetky s, a, θ. Potom

∇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.

Pre γ = 1 činiteľ γ^t zmizne a dostaneme tvar v rozbore nižšie. Praktické implementácie (aj naša) ho vynechávajú aj pre γ < 1; potom výsledok už nie je nestranný odhad ∇J(θ) (vo všeobecnosti dokonca nie je gradientom žiadnej funkcie; Chris Nota a Philip Thomas, 2020), pretože neskoršie kroky epizódy dostanú väčšiu váhu, než im podľa vety patrí — v praxi to napriek tomu funguje dobre.

REINFORCE (Ronald Williams, 1992) je algoritmus:

1. zvoľ počiatočné θ ∈ ℝ^m (náhodné váhy siete) a rýchlosť učenia α > 0;
2. odohraj stratégiou π_θ epizódy (u nás dohromady aspoň 2 000 krokov);
3. v každej epizóde vypočítaj odzadu G_T = 0 a G_t = r_t+1 + γ G_t+1 pre t = T − 1, …, 0;
4. vypočítaj odhad gradientu ĝ = priemer hodnôt ∇ log π_θ(a_t | s_t) · G_t cez všetky odohrané kroky;
5. θ ← θ + α ĝ, dáta zahoď a pokračuj krokom 2, kým nevyčerpáš rozpočet krokov v prostredí.

Pre softmax π = softmax(z) s výstupmi siete z ∈ ℝ^|A| je gradient podľa z rovný ∇_z log π_a = 1_a − π, kde 1_a je vektor s jednotkou na mieste akcie a a nulami inde.

Nech platia predpoklady vety o gradiente stratégie. Epizódu (trajektóriu) τ = (s₀, a₀, s₁, a₁, …, s_T) má stratégia π_θ s pravdepodobnosťou

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; trajektórií je konečne veľa, takže J(θ) = ∑_τ P_θ(τ) G(τ). Potom

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

pretože ∇P = P · ∇ log P. Dôkaz (náčrt): logaritmus zmení súčin na súčet, log P_θ(τ) = log p(s₀) + ∑_t log π_θ(a_t | s_t) + ∑_t log P(s_t+1 | s_t, a_t). Od θ závisí len prostredný súčet, takže fyziku prostredia P poznať nemusíme: ∇ log P_θ(τ) = ∑_t ∇ log π_θ(a_t | s_t). Nakoniec pri každom a_t zahodíme odmeny pred krokom t: od a_t nezávisia, a ich príspevok má preto nulovú strednú hodnotu (pozri vetu v ďalšom oddiele). Zo súčtu ∑_k γ^k r_k+1 tak pri kroku t zostane ∑_{k ≥ t} γ^k r_k+1 = γ^t G_t, čo je tvrdenie vety o gradiente stratégie. ∎

Dve akcie, výstupy siete z = (0; 0), takže π = softmax(z) = (0,50; 0,50). Agent zahral akciu 1 („doprava“) a výnos od toho kroku bol G = 10. Pre softmax je ∇_z log π(a) = 1_a − π, teda

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

(farby podľa rozboru vzorca vyššie; výsledok je odhad ∇J z jediného kroku).

Krok vzostupom s α = 0,1: z ← (0; 0) + 0,1 · (−5; 5) = (−0,5; 0,5). Nová stratégia π = (e^−0,5; e^0,5)/e^−0,5 + e^0,5 = (0,269; 0,731). Akcia, ktorá vyšla, je teraz pravdepodobnejšia. (V skutočnej sieti sa menia váhy θ a k z sa gradient dostane spätným šírením; smer je ten istý.)

Nech (S, A, P, R, γ) je MDP a π_θ parametrizovaná stratégia.

1. Základná úroveň (baseline) je ľubovoľná funkcia b : S → ℝ, ktorá závisí len od stavu, nie od akcie. Odhad gradientu so základnou úrovňou je ĝ = ∑_t ∇ log π_θ(a_t | s_t) · (G_t − b(s_t)); podľa vety nižšie má rovnakú strednú hodnotu ako odhad bez b, len sa tým môže výrazne zmenšiť rozptyl.
2. Výhoda (advantage) stratégie π je funkcia A^π : S × A → ℝ, A^π(s, a) = Q^π(s, a) − V^π(s), kde Q^π a V^π sú hodnotové funkcie z dielu o MDP (písmeno A s argumentmi tu značí výhodu, nie množinu akcií). Najlepšia bežná voľba základnej úrovne je b = V^π_θ, v praxi odhadnutá druhou sieťou; potom G_t − b(s_t) je odhad výhody.
3. Actor-critic je metóda, ktorá súčasne učí parametrizovanú stratégiu π_θ („herec“) a odhad hodnotovej funkcie V ≈ V^π_θ („kritik“) a odhad kritika používa v odhade gradientu stratégie.

Nech A je konečná množina akcií, π_θ parametrizovaná stratégia s π_θ(a | s) > 0, s ∈ S stav a b : S → ℝ funkcia. Potom

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.

Odčítanie b(s_t) od G_t teda strednú hodnotu odhadu gradientu nezmení (stačí podmieniť stavom s_t).

Dôkaz: ∑_a π_θ(a | s) ∇ log π_θ(a | s) = ∑_a ∇π_θ(a | s) = ∇ ∑_a π_θ(a | s) = ∇1 = 0. Člen b(s) je pre pevné s konštanta, takže sa dá vybrať pred súčet. ∎

Jeden stav, dve akcie s π = (0,5; 0,5). Akcia 1 prinesie výnos 12, akcia 0 výnos 10 — oboje kladné ako na vozíku. Zložka odhadu gradientu podľa výstupu z₁ je z jednej vzorky (1_a=1 − 0,5) · (G − b) (farby podľa vzorca vyššie).

Bez základnej úrovne (b = 0): ak padne akcia 1, odhad je 0,5 · 12 = 6; ak padne akcia 0, je −0,5 · 10 = −5. Priemer 6 − 5/2 = 0,5, rozptyl (6 − 0,5)² + (−5 − 0,5)²/2 = 30,25.

So základnou úrovňou b = 11 (priemerný výnos): 0,5 · (12 − 11) = 0,5 a −0,5 · (10 − 11) = 0,5. Priemer je stále 0,5, rozptyl 0.

Skutočný gradient je π₁(1 − π₁) · (12 − 10) = 0,25 · 2 = 0,5. Oba odhady mieria v priemere správne, ale bez latky sa jedna vzorka trafí vedľa o ±5,5, s latkou vôbec. V skutočnosti rozptyl na nulu neklesne (výnosy sú náhodné a stavov je veľa), ale princíp je ten istý.

Nech K ≥ 1 (u nás K = 40) a ĝ⁽¹⁾, …, ĝ^(K) ∈ ℝ^m sú odhady gradientu z K nezávislých dávok (vektory cez všetkých m váh siete). Pre j = 1, …, m označme ḡ_j = 1/K ∑_k=1^K ĝ_j^(k) priemer j-tej zložky. Súčet rozptylov zložiek je číslo

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

Je to stopa (výberovej) kovariančnej matice odhadu — jedno číslo pre to, ako veľmi sa odhad trasie. Pokles: 1 − 0,131/0,177 = 0,26, teda 26 %.

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

1. Polož θ_old = θ, π_old = π_θold a nazbieraj touto stratégiou dávku krokov (s_t, a_t, r_t+1, s_t+1).
2. Pre každý krok vypočítaj TD chybu δ_t = r_t+1 + γ V(s_t+1) − V(s_t) (po páde V(s_t+1) = 0) a odhad výhody metódou GAE (generalized advantage estimation) Â_t = ∑_l≥0 (γλ)^l δ_t+l, kde súčet končí s koncom epizódy.
3. Niekoľko epoch (u nás 10) maximalizuj gradientnými krokmi vzhľadom na θ orezaný cieľ 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 priemer cez kroky dávky (minidávky) a r_t(θ) je pomer pravdepodobností (nezamieňať s odmenou r_t+1); zároveň uč kritika V zmenšovaním priemeru (V(s_t) − y_t)², kde y_t = Â_t + V_old(s_t) je cieľ vypočítaný v kroku 2 z hodnôt kritika pred učením.
4. Dávku zahoď a pokračuj krokom 1, kým nevyčerpáš rozpočet krokov v prostredí.

Hneď ako sa pravdepodobnosť akcie zmení o viac než faktor 1 ± ε v smere, ktorý cieľ zlepšuje, gradient z tej vzorky zmizne.

Stará stratégia dala zahranej akcii pravdepodobnosť 0,5. Po niekoľkých krokoch učenia je to 0,65, takže r = 0,65/0,5 = 1,3 a clip(1,3; 0,8; 1,2) = 1,2 (farby podľa rozboru orezaného cieľa vyššie).

Výhoda +2: min(1,3 · 2; 1,2 · 2) = min(2,6; 2,4) = 2,4. Vyhral orezaný člen, ktorý od θ nezávisí — gradient z tejto vzorky je nula. Akcia je už o 30 % častejšia, viac nie.

Výhoda −2 (ten istý pomer): min(1,3 · (−2); 1,2 · (−2)) = min(−2,6; −2,4) = −2,6. Vyhral neorezaný člen: stratégia omylom spravdepodobnila zlú akciu, a gradient ju preto ďalej sťahuje späť.

Výhoda +2, ale pravdepodobnosť 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 neorezaný člen — dobrá akcia, ktorú stratégia potlačila, sa vracia. Orezanie brzdí len kroky, ktoré by cieľ zlepšili „príliš“.

Nech π_old a π_θ sú stratégie na konečnej množine akcií A s π_θ(a | s) > 0 a nech (s_i, a_i), i = 1, …, N, sú kroky dávky nazbierané stratégiou π_old. KL divergencia z dielu o teórii informácie, tu s prirodzeným logaritmom (jednotka nat; 1 nat = 1/ln 2 ≈ 1,44 bitu), spriemerovaná cez stavy, ktoré stará stratégia navštevuje, 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) ].

Akcie a_i zahrala stará stratégia, takže priemer cez dávku je nestranný odhad. Je to priemerný logaritmus pomeru, a teda e^−D̂ je geometrický priemer pomeru π_θ/π_old cez nazbierané akcie.

Čo to znamená v číslach: pri PPO najviac 0,06 natu, e^−0,06 = 0,94 — akcie zo starej dávky sú typicky o pár percent menej pravdepodobné, dáta o novej stratégii stále vypovedajú. Bez orezania až 24,9 natu: e^−24,9 ≈ 1,5 · 10⁻¹¹, nová stratégia by akcie, z ktorých sa práve učila, takmer nikdy nezahrala.

Nech x je otázka a y = (y₁, …, y_n) odpoveď (postupnosti tokenov) vygenerovaná ladeným jazykovým modelom π_θ, kde π_θ(y | x) = ∏_k=1ⁿ π_θ(y_k | x, y₁, …, y_k−1). Nech r_φ(x, y) ∈ ℝ je výstup modelu odmeny, π_ref zmrazený model po doladení na príkladoch (SFT) a β > 0 sila trestu. PPO potom dostane odmenu

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

(tak to zapisuje InstructGPT, Long Ouyang a kol., 2022). Stredná hodnota druhého člena cez odpovede je presne β · D_KL(π_θ ‖ π_ref). Epizóda je jedna odpoveď, akcie sú tokeny, odmena r_φ príde až po poslednom tokene; zvyšok je PPO tak, ako je opísané vyššie.

Model odpovie na úlohu štyrikrát, správne len prvýkrát: odmeny (1; 0; 0; 0). Priemer 0,25, smerodajná odchýlka √0,75² + 3 · 0,25²/4 = √0,1875 = 0,433.

Výhody (farby podľa vzorca vyššie): 1 − 0,25/0,433 = 1,732 pre správnu odpoveď a 0 − 0,25/0,433 = −0,577 pre každú z troch zlých. Súčet je nula: správna sa posilní, zlé potlačia. Keby boli všetky štyri odpovede správne (alebo všetky zlé), sú všetky výhody nulové a otázka nič nenaučí — učí sa len z otázok, kde sa model niekedy trafí a niekedy nie. (Delíme G; niektoré implementácie delia G − 1, čísla sa potom mierne líšia.)