Abstract

Till now we have considered only learning the Value or Q function and estimating the policy from those. In the next few posts, we are going to look into directly learning the policy. Why directly learn the policy? First, Q learning has a lot of issues involving the Deadly Triad; second, if we have continuous actions we cannot really use it; and lastly, Q learning always learns a deterministic policy, and in cases of partially observed stochastic environments (which is nearly always what we have), having a stochastic policy is proven to be better.

Policy Gradient

I teased a bit, yes we are going to use (stochastic) Gradient Descent to optimize (this is also known as policy search) the following loss:

$J(\pi )\triangleq E_{\pi }[{\sum }_{t=0}^{\infty }{\gamma }^t{R}_{t+1}]$ $={\sum }_{t=0}^{\infty }{\gamma }^t{\sum }_s\left({\sum }_{s_0}{p}_0(s_0)p^{\pi }(s_0\to s,t)\right){\sum }_a\pi (a|s)R(s,a)$ $={\sum }_s\left({\sum }_{s_0}{\sum }_{t=0}^{\infty }{\gamma }^t{p}_0(s_0)p^{\pi }(s_0\to s,t)\right){\sum }_a\pi (a|s)R(s,a)$ $={\sum }_s{\rho }_{\pi }^{\gamma }(s){\sum }_a\pi (a|s)R(s,a)$

• ${\rho }_{\pi }^{\gamma }(s)\triangleq {\sum }_{t=0}^{\infty }{\gamma }^t{\sum }_{s_0}{p}_0(s_0)p^{\pi }(s_0\to s,t)$ this measures the time spent in non-terminal states
  • it is NOT a probability measure, since it is not normalised but can be normalized by exploiting ${\sum }_{i=0}^{\infty }{\gamma }^t=\frac{1}{1-\gamma }$ if $\gamma <1$ this yields
    • $p_{\pi }^{\gamma }(s)=(1-\gamma ){\rho }_{\pi }^{\gamma }(s)=(1-\gamma ){\sum }_{t=0}^{\infty }{\gamma }^t{p}_t(s)$
• $p_t^{\pi }(s)={\sum }_{s_0}{p}_0(s_0)p^{\pi }(s_0\to s,t)$ this is the marginal probability of being in state $s$ at time $t$
• $p^{\pi }(s_0\to s,t)$ is the probability of going from $s_0$ to $s$ in $t$ steps

We abuse the notation a bit by treating $\rho$ as a probability measure we get: $E_{{\rho }_{\pi }^{\gamma }(s)}[f(s)]={\sum }_s{\rho }_{\pi }^{\gamma }(s)f(s)$

And our final loss is: $J(\pi )=E_{{\rho }_{\pi }^{\gamma }(s),\pi (a|s)}[R(s,a)]$

Theorem

Now we differentiate the loss to get:

${\nabla }_{\theta }J(\theta )={\sum }_s{\rho }_{\pi }^{\gamma }(s){\sum }_a{Q}^{\pi }(s,a){\nabla }_{\theta }{\pi }_{\theta }(a|s)$ $={\sum }_s{\rho }_{\pi }^{\gamma }(s){\sum }_a{Q}^{{\pi }_{\theta }}(s,a){\pi }_{\theta }(a|s){\nabla }_{\theta }\log {\pi }_{\theta }(a|s)$ $=E_{{\rho }_{\pi }^{\gamma }(s){\pi }_{\theta }(a|s)}[Q^{{\pi }_{\theta }}(s,a){\nabla }_{\theta }\log {\pi }_{\theta }(a|s)]$

• ${\nabla }_{\theta }\log {\pi }_{\theta }(a|s)]$ this is also known as the score function and it is totally unrelated to the score function in Denoising Diffusion, which is the gradient with respect to a log probability: ${\nabla }_{\theta }\log {\pi }_{\theta }(a|s)$

Since it is a gradient, we can follow it to regions with higher reward! Now you can also see that the equation contains the Q function! Even in policy-based methods, we frequently use the Value, Q, or Advantage function (difference between the Value and Q function) since they stabilize the loss.

Reinforce

So naive policy gradient is said to have high variance! This is because we do a Monte Carlo rollout of the policy, which has low bias but, as said, high variance. To reduce the variance we need to introduce a baseline function, let’s look at the equations:

${\nabla }_{\theta }J({\pi }_{\theta })={\sum }_{t=0}^{\infty }{\gamma }^t{E}_{p_t(s){\pi }_{\theta }(a_t|s_t)}[{\nabla }_{\theta }\log {\pi }_{\theta }(a_t|s_t)Q_{{\pi }_{\theta }}(s_t,a_t)]$

$\approx {\sum }_{t=0}^{T-1}{\gamma }^t{G}_t{\nabla }_{\theta }\log {\pi }_{\theta }(a_t|s_t)$

• $G_t\triangleq r_t+\gamma r_{t+1}+{\gamma }^2{r}_{t+2}+\cdots +{\gamma }^{T-t-1}{r}_{T-1}={\sum }_{k=0}^{T-t-1}{\gamma }^k{r}_{t+k}={\sum }_{j=t}^{T-1}{\gamma }^{j-t}{r}_j$ this is the reward-to-go and is estimated using Markov Chain rollout of the policy, hence the high variance part

As mentioned before, we introduce a baseline function $b(s)$:

${\nabla }_{\theta }J({\pi }_{\theta })=E_{{\rho }_{\theta }(s){\pi }_{\theta }(a|s)}[{\nabla }_{\theta }\log {\pi }_{\theta }(a|s)(Q_{{\pi }_{\theta }}(s,a)-b(s))]$

Baseline Functions

It can be anything, it has just one requirement:

$E[{\nabla }_{\theta }b(s)]=0$

A common choice is either the Value Function $b(s)=V_{{\pi }_{\theta }}(s)$ or Advantage Function $b(s)=Q_{{\pi }_{\theta }}(a,s)$

Estimator

Thus our update for the parameters $\theta$ becomes:

$\theta \leftarrow \theta +\eta {\sum }_{t=0}^{T-1}{\gamma }^t(G_t-b(s_t)){\nabla }_{\theta }\log {\pi }_{\theta }(a_t|s_t)$

The update has an intuitive explanation:

We compute the sum of discounted future reward induced by a trajectory, compared to a baseline and if it is positive we increase $\theta$ to make this trajectory more likely otherwise we decrease $\theta$, or in simpler terms we reinforce good behavior and penalize bad.

Algorithm

Final Remarks

This is the basis of Policy-Based Reinforcement Learning. In the next posts, we will look into an alternative approach called Actor-Critic methods, where instead of MC policy rollout we use Temporal Difference to estimate $G_t$. However, as it turns out, neither Reinforce nor Actor-Critic methods guarantee monotonic improvement in the learned policy, and we turn to Policy Improvement methods which contain the currently popular Proximal Policy Optimization (PPO) method known from Large Language Models!