policy gradient

The goal of reinforcement learning is to find an optimal behavior strategy for the agent to obtain optimal rewards. Policy gradient methods target at modeling and optimizing the policy directly (as opposed to value-based methods that model a value function and derive the policy from it). The policy is usually modeled with a parameterized function respect to $\theta$, ${\pi }_{\theta }(s|a)$ The value of the reward (objective) function depends on this policy and then various algorithms can be applied to optimize for the best reward.
The reward function is defined as:

$$J(\theta )=\sum _{s\in \mathcal{S}}{d}^{\pi }(s)V^{\pi }(s)=\sum _{s\in \mathcal{S}}{d}^{\pi }(s)\sum _{a\in \mathcal{A}}{\pi }_{\theta }(a|s)Q^{\pi }(s,a)$$

where $d^{\pi }(s)={\lim }_{t\to \infty }P(s_t=s\mid s_0,{\pi }_0)$ is the stationary distribution of the markov chain for ${\pi }_{\theta }$ (on-policy state distribution under $\pi$).
Using gradient ascent, we can move $\theta$ toward the direction suggested by the gradient ${\nabla }_{\theta }J(\theta )$ to find the $\theta$ that maximizes the return.

Via the policy gradient theorem, we can drop $d^{\pi }$ from the gradient (i.e. don’t need to know / differentiate the environment).
We end up with the vanilla policy gradient:

$${\nabla }_{\theta }J(\theta )={\mathbb{E}}_{\pi }[Q^{\pi }(s,a){\nabla }_{\theta }\ln {\pi }_{\theta }(a\mid s)]$$

See vanilla policy gradient for a more detailed explanation of the most basic policy gradient algorithm / how we derive the gradient of this objective. Below is a more general formulation of different policy gradients, based on the High-Dimensional Continuous Control Using Generalized Advantage Estimation paper.

Policy gradient methods

Let $s_0$ be sampled from initial distribution ${\rho }_0$. A trajectory $(s_0,a_0,s_1,a_1,\dots )$ is generated by:

• Sampling actions: $a_t\sim \pi (a_t\mid s_t)$
• Sampling states: $s_{t+1}\sim P(s_{t+1}\mid s_t,a_t)$

until reaching a terminal state. At each timestep, a reward $r_t=r(s_t,a_t,s_{t+1})$ is received. The goal is to maximize the expected total reward ${\sum }_{t=0}^{\infty }{r}_t$ (assumed to be finite).

Policy gradient methods maximize the expected total reward by estimating the gradient $g$:

$$g:={\nabla }_{\theta }{\mathbb{E}}_{{\pi }_{\theta }}\left[\sum _{t=0}^{\infty }{r}_t\right]\approx \mathbb{E}\left[\sum _{t=0}^{\infty }{\Psi }_t{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)\right]$$

where ${\Psi }_t$ can take various forms, to estimate the contribution of each action to the total reward:

${\sum }_{t=0}^{\infty }{r}_t$ … Total trajectory reward, which is the original REINFORCE weight. Unbiased but highest variance: every action is credited with the whole trajectory’s reward, including rewards collected before $a_t$ that it couldn’t have caused.
${\sum }_{t^{\prime }=t}^{\infty }{r}_{t^{\prime }}$ … Future reward after action $a_t$ (reward-to-go), drops the past rewards $a_t$ can’t influence. It’s variance is strictly lower, but the expectation stays the same. This is the standard REINFORCE form in practice.
${\sum }_{t^{\prime }=t}^{\infty }{r}_{t^{\prime }}-b(s_t)$ … Baselined reward-to-go; $b$ is supposed to remove variance of the state, and focus the reward on the reward difference caused by actions in a specific state. $b$ does not bias the gradient, because it depends only on the state and not the action, so it cancels in expectation over actions (proof below).

The above options pass a sampled return to ${\Psi }_t$: this is the last unbiased option.
The below instead passes a learned value function, trading bias (approximation error) for lower variance.

$Q^{\pi }(s_t,a_t)$ … state-action value replaces the sampled return with its expectation. This is the actor critic move: a learned critic estimates $Q$, cutting variance a lot, but adding bias whenever the critic is imperfect.
$A^{\pi }(s_t,a_t)$ … advantage function — $Q$ with $V$ as the baseline. The sweet spot: near-lowest variance because it measures better/worse than the state’s average action directly (see below). De-facto target of modern methods (A2C/A3C, TRPO, PPO); never known exactly, so it’s estimated, usually via GAE.
$r_t+V^{\pi }(s_{t+1})-V^{\pi }(s_t)$ … TD residual — a one-step, bootstrapped estimate of the advantage ($\mathbb{E}[{\delta }_t]=A^{\pi }$ under the true $V$). Cheapest advantage: needs only $V$, no separate $Q$-net. Most biased / most reliant on the value estimate; GAE dials between this (low variance, high bias) and the Monte-Carlo advantage (high variance, low bias) via $\lambda$.

where
$V^{\pi }(s_t):={\mathbb{E}}_{s_{t+1:\infty },a_{t:\infty }}\left[{\sum }_{l=0}^{\infty }{r}_{t+l}\right]$ … state-value; expected total reward from state $s_t$
$Q^{\pi }(s_t,a_t):={\mathbb{E}}_{s_{t+1:\infty },a_{t+1:\infty }}\left[{\sum }_{l=0}^{\infty }{r}_{t+l}\right]$ … state-action value; expected total reward from state $s_t$ and action $a_t$
$A^{\pi }(s_t,a_t):=Q^{\pi }(s_t,a_t)-V^{\pi }(s_t)$

Choosing the advantage function for ${\Psi }_t$ yields almost the lowest possible variance, though in practice, the advantage function is not known and must be estimated. This statement can be intuitively justified by the following interpretation of the policy gradient: that a step in the policy gradient direction should increase the probability of better-than-average actions and decrease the probability of worse-than-average actions. The advantage function, by it’s definition measures whether or not the action is better or worse than the policy’s default behavior. Hence, we should choose ${\Psi }_t$ to be the advantage function, so that the gradient term ${\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)$ points in the direction of increased ${\pi }_{\theta }(a_t\mid s_t)$ if and only if $A^{\pi }(s_t,a_t)>0$.

Why a state-dependent baseline doesn't bias the gradient

Subtracting $b(s_t)$ adds one extra term to the gradient: $\mathbb{E}[{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)b(s_t)]$. Because $b$ depends only on the state, pull it out of the expectation over actions and it collapses to zero:

$$b(s){\mathbb{E}}_{a\sim \pi }[{\nabla }_{\theta }\log {\pi }_{\theta }(a\mid s)]=b(s)\sum _a{\pi }_{\theta }(a\mid s){\nabla }_{\theta }\log {\pi }_{\theta }(a\mid s)=b(s)\sum _a{\nabla }_{\theta }{\pi }_{\theta }(a\mid s)=b(s){\nabla }_{\theta }\underset{=1}{\underbrace{{\sum }_a{\pi }_{\theta }(a\mid s)}}=0.$$

The middle step is the log-derivative trick in reverse ($\pi \nabla \log \pi =\nabla \pi$); the last step uses that probabilities sum to $1$, whose gradient is $0$. So the score ${\nabla }_{\theta }\log {\pi }_{\theta }$ has zero mean under ${\pi }_{\theta }$, and subtracting any action-independent $b(s)$ shifts the gradient by nothing in expectation — it only changes the estimator’s variance. A natural, near-optimal choice is $b(s)=V^{\pi }(s)$, which turns the weight into the advantage $A^{\pi }(s,a)=Q^{\pi }(s,a)-V^{\pi }(s)$ (see above). The requirement is strict: a baseline that depended on $a$ would not cancel and would bias the gradient.

References

$\theta$, but also inputs like $(s,a)$ are often ommitted in notation.

PPO

https://lilianweng.github.io/posts/2018-04-08-policy-gradient/