Content deleted Content added
→Motivation: KL |
I fixed the formula so it would express "the total reward from time $ |
||
(45 intermediate revisions by 8 users not shown) | |||
Line 6:
== Overview ==
In policy-based RL, the actor is a parameterized policy function <math>\pi_\theta</math>, where <math>\theta</math> are the parameters of the actor. The actor takes as argument the state of the environment <math>s</math> and produces a [[probability distribution]] <math>\pi_\theta(\cdot
If the action space is discrete, then <math>\sum_{a} \pi_\theta(a
The goal of policy optimization is to find some <math>\theta</math> that maximizes the expected episodic reward <math>J(\theta)</math>:<math display="block">J(\theta) = \mathbb{E}_{\pi_\theta}\left[\sum_{t\in 0:T} \gamma^t R_t \Big| S_0 = s_0 \right]</math>where <math>
\gamma
</math> is the [[discount factor]], <math>
Line 31 ⟶ 29:
=== Policy gradient ===
The '''REINFORCE algorithm''' was the first policy gradient method.<ref>{{Cite journal |last=Williams |first=Ronald J. |date=May 1992 |title=Simple statistical gradient-following algorithms for connectionist reinforcement learning |url=http://link.springer.com/10.1007/BF00992696 |journal=Machine Learning |language=en |volume=8 |issue=3–4 |pages=229–256 |doi=10.1007/BF00992696 |issn=0885-6125}}</ref> It is based on the identity for the policy gradient<math display="block">\nabla_\theta J(\theta)= \mathbb{E}_{\pi_\theta}\left[
\sum_{
\Big|S_0 = s_0
\right]</math> which can be improved via the "causality trick"<ref>{{Cite journal |last1=Sutton |first1=Richard S |last2=McAllester |first2=David |last3=Singh |first3=Satinder |last4=Mansour |first4=Yishay |date=1999 |title=Policy Gradient Methods for Reinforcement Learning with Function Approximation |url=https://proceedings.neurips.cc/paper_files/paper/1999/hash/464d828b85b0bed98e80ade0a5c43b0f-Abstract.html |journal=Advances in Neural Information Processing Systems |publisher=MIT Press |volume=12}}</ref><math display="block">
\nabla_\theta J(\theta)= \mathbb{E}_{\pi_\theta}\left[\sum_{
\Big|S_0 = s_0 \right]
</math>
Line 49 ⟶ 47:
}}
{{hidden begin|style=width:100%|ta1=center|border=1px #aaa solid|title=
{{Math proof|title=Proof of the lemma|proof=
Use the [[reparameterization trick#REINFORCE estimator|reparameterization trick]].
Line 84 ⟶ 82:
\end{aligned}
</math>
}}
{{Math proof|title=Proof of the two identities|proof=
Applying the [[reparameterization trick#REINFORCE estimator|reparameterization trick]],
Line 106 ⟶ 102:
}}
{{hidden end}}Thus, we have an [[unbiased estimator]] of the policy gradient:<math display="block">
\nabla_\theta J(\theta) \approx \frac 1N \sum_{
</math>where the index <math>
The [[Score (statistics)|score function]] <math>\nabla_\theta \ln \pi_\theta (A_t
=== Algorithm ===
Line 115 ⟶ 111:
# Rollout <math>N</math> trajectories in the environment, using <math>\pi_{\theta_t}</math> as the policy function.
# Compute the policy gradient estimation: <math>
# Update the policy by gradient ascent: <math>\theta_{
Here, <math>\
== Variance reduction ==
REINFORCE is an '''on-policy''' algorithm, meaning that the trajectories used for the update must be sampled from the current policy <math>\pi_\theta</math>. This can lead to high variance in the updates, as the returns <math>R(\tau)</math> can vary significantly between trajectories. Many variants of REINFORCE
=== REINFORCE with baseline ===
A common way for reducing variance is the '''REINFORCE with baseline''' algorithm, based on the following identity:<math display="block">\nabla_\theta J(\theta)= \mathbb{E}_{\pi_\theta}\left[\sum_{
\Big|S_0 = s_0 \right]</math>for any function <math>b: \text{States} \to \R</math>. This can be proven by applying the previous lemma.
The algorithm uses the modified gradient estimator<math display="block">
\frac 1N \sum_{
=== Actor-critic methods ===
{{Main|Actor-critic algorithm}}
If <math display="inline">b_i</math> is chosen well, such that <math display="inline">b_i(S_t) \approx \sum_{\tau \in t:T} (\gamma^\tau R_\tau) = \gamma^t V^{\pi_{\theta_i}}(S_t)</math>, this could significantly decrease variance in the gradient estimation. That is, the baseline should be as close to the '''value function''' <math>V^{\pi_{\theta_i}}(S_t)</math> as possible, approaching the ideal of:<math display="block">\nabla_\theta J(\theta)= \mathbb{E}_{\pi_\theta}\left[\sum_{t\in 0:T} \nabla_\theta\ln\pi_\theta(A_t| S_t)\left(\sum_{\tau \in t:T} (\gamma^\tau R_\tau) - \gamma^t V^{\pi_\theta}(S_t)\right)
\Big|S_0 = s_0 \right]</math>Note that, as the policy <math>\pi_{\theta_t}</math> updates, the value function <math>V^{\pi_{\theta_i}}(S_t)</math> updates as well, so the baseline should also be updated. One common approach is to train a separate function that estimates the value function, and use that as the baseline. This is one of the [[actor-critic method]]s, where the policy function is the actor and the value function is the critic.
The '''Q-function''' <math>Q^\pi</math> can also be used as the critic, since<math display="block">\nabla_\theta J(\theta)= E_{\pi_\theta}\left[\sum_{0\leq
\cdot Q^{\pi_\theta}(
\Big|S_0 = s_0 \right]</math> by a similar argument using the tower law.
Subtracting the value function as a baseline, we find that the '''advantage function''' <math>A^{\pi}(S,A) = Q^{\pi}(S,A) - V^{\pi}(S)</math> can be used as the critic as well:<math display="block">\nabla_\theta J(\theta)= E_{\pi_\theta}\left[\sum_{0\leq
\cdot A^{\pi_\theta}(
\Big|S_0 = s_0 \right]</math>In summary, there are many unbiased estimators for <math display="inline">\nabla_\theta J_\theta</math>, all in the form of: <math display="block">\nabla_\theta J(\theta) = E_{\pi_\theta}\left[\sum_{0\leq
\cdot \
\Big|S_0 = s_0 \right]</math> where <math display="inline">\
* <math display="inline">\sum_{0 \leq
* <math display="inline">\gamma^
* <math display="inline">\gamma^
* <math display="inline">\gamma^
* <math display="inline">\gamma^
* <math display="inline">\gamma^
Some more possible <math display="inline">\
* <math display="inline">\gamma^
* <math display="inline">\gamma^
* <math display="inline">\gamma^
== Natural policy gradient ==
{{Anchor|NPG}}
The natural policy gradient method is a variant of the policy gradient method, proposed by [[Sham Kakade]] in 2001.<ref>{{Cite journal |last=Kakade |first=Sham M |date=2001 |title=A Natural Policy Gradient |url=https://proceedings.neurips.cc/paper_files/paper/2001/hash/4b86abe48d358ecf194c56c69108433e-Abstract.html |journal=Advances in Neural Information Processing Systems |publisher=MIT Press |volume=14}}</ref> Unlike standard policy gradient methods, which depend on the choice of parameters <math>\theta</math> (making updates coordinate-dependent), the natural policy gradient aims to provide a [[coordinate-free]] update, which is geometrically "natural".
=== Motivation ===
Standard policy gradient updates <math>\theta_{
\begin{cases}
\max_{\theta_{
\|\theta_{
\end{cases}
</math>
While the objective (linearized improvement) is geometrically meaningful, the Euclidean constraint <math>\|\theta_{
\max_{\theta_{i+1}} J(\theta_i) + (\theta_{i+1} - \theta_i)^T \nabla_\theta J(\theta_i)\\
\
\end{cases}</math>where the KL divergence between two policies is '''averaged''' over the state distribution under policy <math>\pi_{\theta_i}</math>. That is,<math display="block">\bar{D}_{KL}(\pi_{\theta_{i+1}} \| \pi_{\theta_{i}}) := \mathbb E_{s \sim \pi_{\theta_i}}[D_{KL}( \pi_{\theta_{i+1}}(\cdot | s) \| \pi_{\theta_{i}}(\cdot | s) )]</math> This ensures updates are invariant to invertible affine parameter transformations.
=== Fisher information approximation ===
For small <math>\epsilon</math>, the KL divergence is approximated by the [[Fisher information metric]]:<math display="block">
</math>where <math>F(\theta)</math> is the [[Fisher information matrix]] of the policy, defined as:<math display="block">
F(\theta) = \mathbb{E}_{s, a \sim \pi_\theta}\left[ \nabla_\theta \ln \pi_\theta(a|s) \left(\nabla_\theta \ln \pi_\theta(a|s)\right)^T \right]
</math> This transforms the problem into a problem in [[quadratic programming]], yielding the natural policy gradient update:<math display="block">
\theta_{
</math>The step size <math>\alpha</math> is typically adjusted to maintain the KL constraint, with <math display="inline">\alpha \
Inverting <math>F(\theta)</math> is computationally intensive, especially for high-dimensional parameters (e.g., neural networks). Practical implementations often use approximations.
== Trust Region Policy Optimization (TRPO) ==
{{Anchor|TRPO}}
'''Trust Region Policy Optimization''' (TRPO) is a policy gradient method that extends the natural policy gradient approach by enforcing a [[trust region]] constraint on policy updates.<ref name=":3">{{Cite journal |last1=Schulman |first1=John |last2=Levine |first2=Sergey |last3=Moritz |first3=Philipp |last4=Jordan |first4=Michael |last5=Abbeel |first5=Pieter |date=2015-07-06 |title=Trust region policy optimization |url=https://dl.acm.org/doi/10.5555/3045118.3045319 |journal=Proceedings of the 32nd International Conference on International Conference on Machine Learning - Volume 37 |series=ICML'15 |___location=Lille, France |publisher=JMLR.org |pages=1889–1897}}</ref> Developed by Schulman et al. in 2015, TRPO improves upon the natural policy gradient method.
The natural gradient descent is theoretically optimal, ''if'' the objective is truly a quadratic function, but this is only an approximation. TRPO's line search and KL constraint attempts to restrict the solution to within a "trust region" in which this approximation does not break down. This makes TRPO more robust in practice.
=== Formulation ===
Like natural policy gradient, TRPO iteratively updates the policy parameters <math>\theta</math> by solving a constrained optimization problem specified coordinate-free:<math display="block">
\begin{cases}
\max_{\theta} L(\theta, \theta_i)\\
\bar{D}_{KL}(\pi_{\theta} \| \pi_{\theta_{i}}) \leq \epsilon
\end{cases}
</math>where
* <math>L(\theta, \theta_i) = \mathbb{E}_{s, a \sim \pi_{\theta_i}}\left[ \frac{\pi_\theta(a|s)}{\pi_{\theta_i}(a|s)} A^{\pi_{\theta_i}}(s, a) \right]</math> is the '''surrogate advantage''', measuring the performance of <math>\pi_\theta</math> relative to the old policy <math>\pi_{\theta_i}</math>.
* <math>\epsilon</math> is the trust region radius.
Note that in general, other surrogate advantages are possible:<math display="block">L(\theta, \theta_i) = \mathbb{E}_{s, a \sim \pi_{\theta_i}}\left[ \frac{\pi_\theta(a|s)}{\pi_{\theta_i}(a|s)}\Psi^{\pi_{\theta_i}}(s, a) \right]</math>where <math>\Psi</math> is any linear sum of the previously mentioned type. Indeed, OpenAI recommended using the Generalized Advantage Estimate, instead of the plain advantage <math>A^{\pi_\theta}</math>.
The surrogate advantage <math>L(\theta, \theta_t)
</math> is designed to align with the policy gradient <math>\nabla_\theta J(\theta)</math>. Specifically, when <math>\theta = \theta_t</math>, '''<math>
\nabla_\theta L(\theta, \theta_t)
</math>''' equals the policy gradient derived from the advantage function:
<math display="block">\nabla_\theta J(\theta) = \mathbb{E}_{(s, a) \sim \pi_\theta}\left[\nabla_\theta \ln \pi_\theta(a | s) \cdot A^{\pi_\theta}(s, a) \right] = \nabla_\theta L(\theta, \theta_t)</math>However, when <math>\theta \neq \theta_i</math>, this is not necessarily true. Thus it is a "surrogate" of the real objective.
As with natural policy gradient, for small policy updates, TRPO approximates the surrogate advantage and KL divergence using Taylor expansions around <math>\theta_t</math>:<math display="block">
\begin{aligned}
L(\theta, \theta_i) &\approx g^T (\theta - \theta_i), \\
\bar{D}_{\text{KL}}(\pi_{\theta} \| \pi_{\theta_i}) &\approx \frac{1}{2} (\theta - \theta_i)^T H (\theta - \theta_i),
\end{aligned}
</math>
where:
* <math>g = \nabla_\theta L(\theta, \theta_i) \big|_{\theta = \theta_i}</math> is the policy gradient.
* <math>F = \nabla_\theta^2 \bar{D}_{\text{KL}}(\pi_{\theta} \| \pi_{\theta_i}) \big|_{\theta = \theta_i}</math> is the Fisher information matrix.
This reduces the problem to a quadratic optimization, yielding the natural policy gradient update:
<math display="block">
\theta_{i+1} = \theta_i + \sqrt{\frac{2\epsilon}{g^T F^{-1} g}} F^{-1} g.
</math>So far, this is essentially the same as natural gradient method. However, TRPO improves upon it by two modifications:
* Use [[conjugate gradient method]] to solve for <math>
x
</math> in <math>Fx = g</math> iteratively without explicit matrix inversion.
* Use [[backtracking line search]] to ensure the trust-region constraint is satisfied. Specifically, it backtracks the step size to ensure the KL constraint and policy improvement. That is, it tests each of the following test-solutions<math display="block">
\theta_{i+1} = \theta_i + \sqrt{\frac{2\epsilon}{x^T F x}} x, \; \theta_i + \alpha \sqrt{\frac{2\epsilon}{x^T F x}} x, \; \theta_i + \alpha^2 \sqrt{\frac{2\epsilon}{x^T F x}} x, \; \dots
</math> until it finds one that both satisfies the KL constraint <math>\bar{D}_{KL}(\pi_{\theta_{i+1}} \| \pi_{\theta_{i}}) \leq \epsilon </math> and results in a higher <math>
L(\theta_{i+1}, \theta_i) \geq L(\theta_i, \theta_i)
</math>. Here, <math>\alpha \in (0,1)</math> is the backtracking coefficient.
== Proximal Policy Optimization (PPO) ==
{{Anchor|PPO}}
A further improvement is [[proximal policy optimization]] (PPO), which avoids even computing <math>F(\theta)</math> and <math>F(\theta)^{-1}</math> via a first-order approximation using clipped probability ratios.<ref name=":0">{{Citation |last1=Schulman |first1=John |title=Proximal Policy Optimization Algorithms |date=2017-08-28 |arxiv=1707.06347 |last2=Wolski |first2=Filip |last3=Dhariwal |first3=Prafulla |last4=Radford |first4=Alec |last5=Klimov |first5=Oleg}}</ref>
Specifically, instead of maximizing the surrogate advantage<math display="block">
\max_\theta L(\theta, \theta_t) = \mathbb{E}_{s, a \sim \pi_{\theta_t}}\left[ \frac{\pi_\theta(a|s)}{\pi_{\theta_t}(a|s)} A^{\pi_{\theta_t}}(s, a) \right]
</math> under a KL divergence constraint, it directly inserts the constraint into the surrogate advantage:<math display="block">
\max_\theta \mathbb{E}_{s, a \sim \pi_{\theta_t}}\left[
\begin{cases}
\min \left(\frac{\pi_\theta(a|s)}{\pi_{\theta_t}(a|s)}, 1 + \epsilon \right) A^{\pi_{\theta_t}}(s, a) & \text{ if } A^{\pi_{\theta_t}}(s, a) > 0
\\
\max \left(\frac{\pi_\theta(a|s)}{\pi_{\theta_t}(a|s)}, 1 - \epsilon \right) A^{\pi_{\theta_t}}(s, a) & \text{ if } A^{\pi_{\theta_t}}(s, a) < 0
\end{cases}
\right]
</math> and PPO maximizes the surrogate advantage by stochastic gradient descent, as usual.
In words, gradient-ascending the new surrogate advantage function means that, at some state <math>
s, a
</math>, if the advantage is positive: <math>
A^{\pi_{\theta_t}} (s, a) > 0
</math>, then the gradient should direct <math>
\theta
</math> towards the direction that increases the probability of performing action <math>
a
</math> under the state <math>
s
</math>. However, as soon as <math>
\theta
</math> has changed so much that <math>
\pi_\theta(a | s) \geq (1 + \epsilon) \pi_{\theta_t}(a | s)
</math>, then the gradient should stop pointing it in that direction. And similarly if <math>
A^{\pi_{\theta_t}} (s, a) < 0
</math>. Thus, PPO avoids pushing the parameter update too hard, and avoids changing the policy too much.
To be more precise, to update <math>
\theta_t
</math> to <math>
\theta_{t+1}
</math> requires multiple update steps on the same batch of data. It would initialize <math>
\theta = \theta_t
</math>, then repeatedly apply gradient descent (such as the [[Adam optimizer]]) to update <math>
\theta
</math> until the surrogate advantage has stabilized. It would then assign <math>
\theta_{t+1}
</math> to <math>
\theta
</math>, and do it again.
During this inner-loop, the first update to <math>
\theta
</math> would not hit the <math>
1 - \epsilon, 1 + \epsilon
</math> bounds, but as <math>
\theta
</math> is updated further and further away from <math>
\theta_t
</math>, it eventually starts hitting the bounds. For each such bound hit, the corresponding gradient becomes zero, and thus PPO avoid updating <math>
\theta
</math> too far away from <math>
\theta_t
</math>.
This is important, because the surrogate loss assumes that the state-action pair <math>
s, a
</math> is sampled from what the agent would see if the agent runs the policy <math>
\pi_{\theta_t}
</math>, but policy gradient should be on-policy. So, as <math>
\theta
</math> changes, the surrogate loss becomes more and more ''off''-policy. This is why keeping <math>
\theta
</math> ''proximal'' to <math>
\theta_t
</math> is necessary.
If there is a reference policy <math>
\pi_{\text{ref}}
</math> that the trained policy should not diverge too far from, then additional KL divergence penalty can be added:<math display="block">-\beta \mathbb{E}_{s, a \sim \pi_{\theta_t}}\left[\log\left(\frac{\pi_{\theta}(a|s)}{\pi_{\text{ref}}(a|s)}\right) \right]</math>where <math>
\beta
</math> adjusts the strength of the penalty. This has been used in training [[Reasoning language model|reasoning language models]] with [[reinforcement learning from human feedback]].<ref name="summarizationpaper">{{cite journal |author=Nisan Stiennon |author2=Long Ouyang |author3=Jeffrey Wu |author4=Daniel Ziegler |author5=Ryan Lowe |author6=Chelsea Voss |author7=Alec Radford |author8=Dario Amodei |author9=Paul F. Christiano |date=2020 |title=Learning to summarize with human feedback |url=https://proceedings.neurips.cc/paper/2020/hash/1f89885d556929e98d3ef9b86448f951-Abstract.html |journal=Advances in Neural Information Processing Systems |language=en |volume=33}}</ref> The KL divergence penalty term can be estimated with lower variance using the equivalent form (see [[f-divergence]] for details):<ref name=":1" /><math display="block">-\beta \mathbb{E}_{s, a \sim \pi_{\theta_t}}\left[
\log\left(\frac{\pi_{\theta}(a|s)}{\pi_{\text{ref}}(a|s)}\right)
+ \frac{\pi_{\text{ref}}(a|s)} {\pi_{\theta}(a|s)}
-1
\right]</math>
=== Group Relative Policy Optimization (GRPO) ===
{{Anchor|GRPO}}
The Group Relative Policy Optimization (GRPO) is a minor variant of PPO that omits the value function estimator <math>V</math>. Instead, for each state <math>s </math>, it samples multiple actions <math>a_1, \dots, a_G </math> from the policy <math>\pi_{\theta_t} </math>, then calculate the group-relative advantage<ref name=":1">{{Citation |last1=Shao |first1=Zhihong |title=DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models |date=2024-04-27 |arxiv=2402.03300 |last2=Wang |first2=Peiyi |last3=Zhu |first3=Qihao |last4=Xu |first4=Runxin |last5=Song |first5=Junxiao |last6=Bi |first6=Xiao |last7=Zhang |first7=Haowei |last8=Zhang |first8=Mingchuan |last9=Li |first9=Y. K.}}</ref><math display="block">A^{\pi_{\theta_t}}(s, a_{j}) = \frac{r(s, a_{j}) - \mu}{\sigma} </math> where <math>\mu, \sigma </math> are the mean and standard deviation of <math>r(s, a_1), \dots, r(s, a_G) </math>. That is, it is the [[standard score]] of the rewards.
Then, it maximizes the PPO objective, averaged over all actions:<math display="block">
\max_\theta \frac{1}{G} \sum_{i=1}^G \mathbb{E}_{(s, a_1, \dots, a_G) \sim \pi_{\theta_t}}\left[
\begin{cases}
\min \left(\frac{\pi_\theta(a_i|s)}{\pi_{\theta_t}(a_i|s)}, 1 + \epsilon \right) A^{\pi_{\theta_t}}(s, a_i) & \text{ if } A^{\pi_{\theta_t}}(s, a_i) > 0
\\
\max \left(\frac{\pi_\theta(a_i|s)}{\pi_{\theta_t}(a_i|s)}, 1 - \epsilon \right) A^{\pi_{\theta_t}}(s, a_i) & \text{ if } A^{\pi_{\theta_t}}(s, a_i) < 0
\end{cases}
\right]
</math>Intuitively, each policy update step in GRPO makes the policy more likely to respond to each state with an action that performed relatively better than other actions tried at that state, and less likely to respond with one that performed relatively worse.
As before, the KL penalty term can be applied to encourage the trained policy to stay close to a reference policy. GRPO was first proposed in the context of training [[reasoning language model|reasoning language models]] by researchers at [[DeepSeek]].<ref name=":1" />
== See also ==
Line 211 ⟶ 349:
* {{Cite web |last=Weng |first=Lilian |date=2018-04-08 |title=Policy Gradient Algorithms |url=https://lilianweng.github.io/posts/2018-04-08-policy-gradient/ |access-date=2025-01-25 |website=lilianweng.github.io |language=en}}
* {{Cite web |title=Vanilla Policy Gradient — Spinning Up documentation |url=https://spinningup.openai.com/en/latest/algorithms/vpg.html |access-date=2025-01-25 |website=spinningup.openai.com}}
{{Artificial intelligence navbox}}
[[Category:Reinforcement learning]]
[[Category:Machine learning algorithms]]
|