Gradiente de la política y PPO: aprender directamente la política

Una política estocástica π_θ(a | s) es una red con parámetros θ que, para un estado s, devuelve una distribución de probabilidad sobre las acciones: para un número finito de acciones, el softmax de las salidas de la red; para acciones continuas, por ejemplo la media y la varianza de una distribución normal. El objetivo es maximizar el retorno esperado J(θ) = E_πθ[G₀] mediante ascenso de gradiente θ ← θ + α ∇J(θ).

La recompensa es +1 por cada paso y el descuento γ = 0,99. Si quedan T − t pasos hasta el final del episodio, el retorno es una serie geométrica:

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

Al comienzo de un episodio que aguanta los 500 pasos completos: G₀ = (1 − 0,99⁵⁰⁰) / 0,01 = 99,34. Veinte pasos antes de la caída: G = (1 − 0,99²⁰) / 0,01 = 18,21. El paso justo antes de la caída: G = 1. El retorno es, por tanto, siempre positivo y nunca supera 1 / (1 − γ) = 100; esto será útil en la sección sobre la línea base.

El gradiente del retorno esperado se puede escribir como una esperanza, que se puede estimar a partir de episodios jugados:

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

REINFORCE (Ronald Williams, 1992) lo hace al pie de la letra: juega episodios, calcula para cada paso el retorno G_t desde ese paso en adelante y desplaza θ en la dirección ∇ log π(a_t | s_t) · G_t. Para el softmax, ∇ log π(a) respecto a las salidas de la red es simplemente (1_a − π).

Un episodio τ = (s₀, a₀, s₁, a₁, …) tiene, con la política π_θ, la probabilidad P_θ(τ) = p(s₀) ∏_t π_θ(a_t | s_t) P(s_t+1 | s_t, a_t) y el retorno G(τ). Entonces

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

porque ∇P = P · ∇ log P. Demostración (esbozo): el logaritmo convierte el producto en una suma, log P_θ(τ) = log p(s₀) + ∑_t log π_θ(a_t | s_t) + ∑_t log P(s_t+1 | s_t, a_t). De θ depende solo la suma del medio, así que no necesitamos conocer la física del entorno P: ∇ log P_θ(τ) = ∑_t ∇ log π_θ(a_t | s_t). Por último, para cada a_t descartamos las recompensas anteriores al paso t: no dependen de a_t, y por eso su contribución tiene esperanza nula (véase el teorema de la sección siguiente). Queda G_t. ∎ (El enunciado exacto para el retorno descontado es algo más sutil; aquí, simplificado.)

Dos acciones, salidas de la red z = (0; 0), así que π = softmax(z) = (0,50; 0,50). El agente jugó la acción 1 («a la derecha») y el retorno desde ese paso fue G = 10. Para el softmax, ∇_z log π(a) = 1_a − π, es decir

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

Un paso de ascenso con α = 0,1: z ← (0; 0) + 0,1 · (−5; 5) = (−0,5; 0,5). La nueva política π = (e^−0,5; e^0,5) / (e^−0,5 + e^0,5) = (0,269; 0,731). La acción que salió bien es ahora más probable. (En una red real cambian los pesos θ, y el gradiente llega a z por retropropagación; la dirección es la misma.)

A G_t se le puede restar cualquier función del estado b(s_t) sin que cambie la esperanza de la estimación del gradiente, porque E[∇ log π(a | s)] = 0. En cambio, la varianza de la estimación puede reducirse mucho. La mejor elección habitual es el valor del estado V(s), aprendido por una segunda red; la diferencia A(s, a) = Q(s, a) − V(s) se llama ventaja (advantage). Los métodos con una política aprendida («actor») y un valor aprendido («crítico») se llaman actor-crítico.

Para todo estado s se cumple E_{a ∼ πθ}[∇ log π_θ(a | s)] = 0, y por tanto E[∇ log π_θ(a | s) · b(s)] = 0 para cualquier función b que no dependa de a.

Demostración: ∑_a π_θ(a | s) ∇ log π_θ(a | s) = ∑_a ∇π_θ(a | s) = ∇ ∑_a π_θ(a | s) = ∇1 = 0. El término b(s) es una constante para s fijo, así que puede sacarse fuera de la suma. ∎

Un estado, dos acciones con π = (0,5; 0,5). La acción 1 da un retorno de 12; la acción 0, de 10: ambos positivos, como en el carro. La componente de la estimación del gradiente respecto a la salida z₁ es, a partir de una sola muestra, (1_a=1 − 0,5) · (G − b).

Sin línea base (b = 0): si sale la acción 1, la estimación es 0,5 · 12 = 6; si sale la acción 0, es −0,5 · 10 = −5. Media (6 − 5) / 2 = 0,5, varianza ((6 − 0,5)² + (−5 − 0,5)²) / 2 = 30,25.

Con línea base b = 11 (retorno medio): 0,5 · (12 − 11) = 0,5 y −0,5 · (10 − 11) = 0,5. La media sigue siendo 0,5, la varianza 0.

El gradiente verdadero es π₁(1 − π₁) · (12 − 10) = 0,25 · 2 = 0,5. Las dos estimaciones apuntan bien en promedio, pero sin listón una sola muestra se desvía ±5,5, y con listón nada. En la realidad la varianza no baja a cero (los retornos son aleatorios y hay muchos estados), pero el principio es el mismo.

De K = 40 lotes tenemos estimaciones del gradiente ĝ⁽¹⁾, …, ĝ^(K) (vectores de todos los pesos de la red). Para cada componente j calculamos la varianza entre lotes y sumamos:

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

Es la traza de la matriz de covarianza de la estimación: un único número que dice cuánto tiembla la estimación. Reducción: 1 − 0,131 / 0,177 = 0,26, es decir, un 26 %.

PPO (John Schulman et al., OpenAI, 2017) recoge un lote de transiciones con la política antigua π_old, calcula las ventajas Â_t y luego, durante varias épocas, maximiza el objetivo recortado

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

normalmente con ε = 0,2. En cuanto la probabilidad de una acción cambia más de ε en la dirección que mejora el objetivo, el gradiente de esa muestra desaparece. Las ventajas suelen estimarse con GAE (generalized advantage estimation), una media ponderada de los errores TD con el parámetro λ.

La política antigua daba a la acción jugada una probabilidad de 0,5. Tras unos cuantos pasos de aprendizaje es 0,65, así que r = 0,65 / 0,5 = 1,3 y clip(1,3; 0,8; 1,2) = 1,2.

Ventaja +2: min(1,3 · 2; 1,2 · 2) = min(2,6; 2,4) = 2,4. Gana el término recortado, que no depende de θ: el gradiente de esta muestra es cero. La acción ya es un 30 % más frecuente, y no más.

Ventaja −2 (misma razón): min(−2,6; −2,4) = −2,6. Gana el término sin recortar: la política hizo por error más probable una acción mala, y por eso el gradiente sigue tirando de ella hacia atrás.

Ventaja +2, pero la probabilidad bajó a 0,35: r = 0,7, min(1,4; 0,8 · 2 = 1,6) = 1,4. Otra vez el término sin recortar: una buena acción que la política había suprimido vuelve. El recorte solo frena los pasos que mejorarían el objetivo «demasiado».

La divergencia KL de la parte sobre teoría de la información, aquí con el logaritmo natural (unidad: nat; 1 nat = 1 / ln 2 ≈ 1,44 bits), promediada sobre los estados del lote:

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) ].

Las acciones a_i las jugó la política antigua, así que la media sobre el lote es una estimación insesgada. Es el logaritmo medio de la razón, y por tanto e^−D es la media geométrica de la razón π_θ / π_old sobre las acciones recogidas.

Qué significa en números: en PPO, como mucho 0,06 nats, e^−0,06 = 0,94: las acciones del lote antiguo son típicamente unos pocos puntos porcentuales menos probables, y los datos siguen diciendo algo sobre la política nueva. Sin recorte, hasta 24,9 nats: e^−24,9 ≈ 1,5 · 10⁻¹¹; la política nueva casi nunca jugaría las acciones de las que acaba de aprender.

Para una consulta x y una respuesta y generada por el modelo que se está ajustando, π_θ, PPO recibe la recompensa

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

donde r_φ es el modelo de recompensa, π_ref el modelo congelado tras el ajuste fino con ejemplos (SFT) y β la fuerza de la penalización (así lo escribe InstructGPT, Long Ouyang et al., 2022). La esperanza del segundo término sobre las respuestas es exactamente β · D_KL(π_θ ‖ π_ref). Un episodio es una respuesta, las acciones son tokens, y la recompensa r_φ llega solo después del último token; el resto es PPO tal como se ha descrito arriba.

El modelo responde cuatro veces a un problema, y solo acierta la primera: recompensas (1; 0; 0; 0). Media 0,25, desviación típica √((0,75² + 3 · 0,25²) / 4) = √0,1875 = 0,433.

Ventajas: (1 − 0,25) / 0,433 = 1,732 para la respuesta correcta y (0 − 0,25) / 0,433 = −0,577 para cada una de las tres incorrectas. La suma es cero: la correcta se refuerza y las incorrectas se suprimen. Si las cuatro respuestas fueran correctas (o todas incorrectas), todas las ventajas serían nulas y la pregunta no enseñaría nada: solo se aprende de las preguntas en las que el modelo unas veces acierta y otras no. (Dividimos entre G; algunas implementaciones dividen entre G − 1, y entonces los números difieren ligeramente.)