Policy Gradients

Recall that the goal of Reinforcement Learning is:

$$\begin{array}{rl}{\theta }^{\ast }& =\arg \underset{\theta }{\max }\sum _{t=0}^T{\mathbb{E}}_{(s_t,a_t)\sim p_{\theta }(s_t,a_t)}[r(s_t,a_t)]\\ & =\arg \underset{\theta }{\max }\underset{\text{This term can be denoted by }J(\theta )}{\underbrace{{\mathbb{E}}_{\tau \sim p_{\theta }(\tau )}\left[\sum _{t=0}^{T}r(s_t,a_t)\right]}}\end{array}$$

$J(\theta )$ can be evaluated from samples as follows:

$$J(\theta )={\mathbb{E}}_{\tau \sim p_{\theta }(\tau )}\left[\sum _{t=0}^{T}r(s_t,a_t)\right]\approx \frac{1}{N}\sum _{i=1}^N\sum _{t=0}^{T}r(s_{i,t},a_{i,t})$$

i.e. we collect $N$ trajectories by running policy ${\pi }_{\theta }$ and estimate the value of the policy $J(\theta )$. This is essentially a Monte-Carlo estimate. Letting $r(\tau )=\sum _{t=0}^{T}r(s_t,a_t)$, we can find the ${\nabla }_{\theta }J(\theta )$:

$$\begin{array}{rl}{\nabla }_{\theta }J(\theta )& ={\nabla }_{\theta }{\mathbb{E}}_{\tau \sim p_{\theta }(\tau )}[r(\tau )]\\ & ={\nabla }_{\theta }\int p_{\theta }(\tau )\cdot r(\tau )\cdot d\tau \\ & =\int ({\nabla }_{\theta }{p}_{\theta }(\tau ))\cdot r(\tau )\cdot d\tau \end{array}$$

We can use the identity $p_{\theta }(\tau ){\nabla }_{\theta }\log p_{\theta }(\tau )=p_{\theta }(\tau )\frac{{\nabla }_{\theta }{p}_{\theta }}{p_{\theta }}={\nabla }_{\theta }{p}_{\theta }(\tau )$ to simplify the above equation as follows:

$$\begin{array}{rl}{\nabla }_{\theta }J(\theta )& =\int (p_{\theta }(\tau )\cdot {\nabla }_{\theta }\log p_{\theta }(\tau ))\cdot r(\tau )\cdot d\tau \\ & ={\mathbb{E}}_{p_{\theta }(\tau )}\left[{\nabla }_{\theta }\log p_{\theta }(\tau )\cdot r(\tau )\right]...[1]\end{array}$$

We already know that:

$$\begin{array}{rl}{p}_{\theta }(\tau )& =p_{\theta }(s_0,a_0,s_1,a_1,...,s_T,a_T,s_{T+1})\\ & =\rho (s_0)\prod _{t=0}^T{\pi }_{\theta }(a_t|s_t)\cdot p(s_{t+1}|s_t,a_t)\\ ⟹\log p_{\theta }(\tau )& =\log \left(\rho (s_0)\prod _{t=0}^T{\pi }_{\theta }(a_t|s_t)\cdot p(s_{t+1}|s_t,a_t)\right)\\ & =\underset{\text{independent of }\theta }{\underbrace{\log \rho (s_0)}}+\sum _{t=0}^T\log {\pi }_{\theta }(a_t|s_t)+\underset{\text{independent of }\theta }{\underbrace{\log p(s_{t+1}|s_t,a_t)}}\\ ⟹{\nabla }_{\theta }\log p_{\theta }(\tau )& ={\nabla }_{\theta }\sum _{t=0}^T\log {\pi }_{\theta }(a_t|s_t)=\sum _{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t|s_t)\end{array}$$

Thus equation $[1]$ becomes

$$\begin{array}{rl}{\nabla }_{\theta }J(\theta )& ={\mathbb{E}}_{p_{\theta }(\tau )}\left[\sum _{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t|s_t)\cdot \sum _{t=0}^{T}r(s_t,a_t)\right]\\ & \approx \underset{\text{Generate Samples}}{\underbrace{\frac{1}{N}\sum _{i=1}^N}}\left(\sum _{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t^i|s_t^i)\right)\underset{\text{Estimate Rewards}}{\underbrace{\left(\sum _{t=0}^{T}r(s_t^i,a_t^i)\right)}}\end{array}$$

Looking back at the anatomy of RL, The first part of the equation corresponds to the trajectory collection phase and the last part corresponds to the value estimation part. The policy update part is essentially:

$$\theta \leftarrow \theta +\alpha {\nabla }_{\theta }J(\theta )$$

where $\alpha$ is the learning rate.

---

REINFORCE Algorithm

Goal: Obtain the optimal policy ${\pi }^{\ast }$
While $|V^{{\pi }_{\theta }}-V^{{\pi }^{\ast }}|<ϵ$:
1. Collect $N$ trajectories $\{{\tau }^i{\}}_{i=1}^N$ by executing policy ${\pi }_{\theta }(a_t|s_t)$ in the environment
2. Estimate ${\nabla }_{\theta }J(\theta )\approx \frac{1}{N}\sum _{i=1}^N\left(\sum _{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t^i|s_t^i)\right)\left(\sum _{t=0}^{T}r(s_t^i,a_t^i)\right)$
3. Update $\theta \leftarrow \theta +\alpha {\nabla }_{\theta }J(\theta )$, i.e. take step towards the positive gradient

Comments:

1. Recall that ${\nabla }_{\theta }{J}_{MLE}(\theta )\approx \frac{1}{N}\sum _{i=1}^N\left(\sum _{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t^i|s_t^i)\right)$, i.e. finding the maximum likelihood estimate (MLE) using ERM in behavior cloning (BC). Comparing with the PG algorithm, we can see that there is an extra term $r(\tau )$ in PG which is “weighing” the action log likelihoods given states. Thus PG can be thought of as a “trial-and-error” algorithm where the actions which lead to more rewarding trajectories are given more importance during gradient update by reinforcing the behavior to take more rewarding actions conditional on states.
2. PG algorithms have very high variance. For instance, we can take a running example: consider the reward system +1 for each time-step (fixed horizon length of 100), and only if the agent reaches the goal, then there is an extra reward of +1. Now consider an alternate situation where at each step there is 0 reward instead for survival. PG algorithms work better in the latter situation, since we can interpret it as giving 100 weightage to all bad trajectories and 101 to good trajectories.
3. To reduce variance, we can simply introduce the concept of baselines:${\nabla }_{\theta }J(\theta )\approx \frac{1}{N}\sum _{i=1}^N\left(\sum _{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t^i|s_t^i)\right)\left(\sum _{t=0}^{T}r(s_t^i,a_t^i)-b\right)$which is essentially just subtracting a baseline value to give more importance to good trajectories, i.e. say if we subtract $b=100$, then good trajectories get 1 weightage and everything else gets 0 weightage. We can simply set $b=\frac{1}{N}\sum _{i=1}^{N}r(\tau )$ since:$\mathbb{E}[{\nabla }_{\theta }\log p_{\theta }(\tau )b]=\int p_{\theta }(\tau ){\nabla }_{\theta }\log p_{\theta }(\tau )b\cdot d\tau =b\int {\nabla }_{\theta }{p}_{\theta }(\tau )d\tau =b{\nabla }_{\theta }\int p_{\theta }(\tau )d\tau =b{\nabla }_{\theta }1=0$which essentially means that subtracting a baseline is unbiased in expectation. The baseline can also be learned during training (using back-propagation).
4. One other way to reduce variance is by noticing the causality during calculating the rewards, i.e. a policy at time $t^{\prime }$ cannot affect the rewards obtained at time $t$ where $t<t^{\prime }$. Thus gradient calculation equation can be modified as:
   ${\nabla }_{\theta }J(\theta )\approx \frac{1}{N}\sum _{i=1}^N\sum _{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t^i|s_t^i)\underset{\hat{Q}(s_{t^{\prime }}^i,a_{t^{\prime }}^i)}{\underbrace{\left(\sum _{t^{\prime }=t}^{T}r(s_{t^{\prime }}^i,a_{t^{\prime }}^i)\right)}}$
   where $\hat{Q}(s_{t^{\prime }}^i,a_{t^{\prime }}^i)$ is essentially the MC estimate of the state-action value function. Now notice that in this situation, we cannot introduce a constant baseline $b$, since for each state $s_t$ the baseline will be different. For our running example, say at $s_0$, the baseline value should be 100, but at $s_{50}$ the baseline value should be 50, and finally at $s_{100}$ the baseline value should be 0. Therefore there is a need to learn state dependent baselines:

$${\nabla }_{\theta }J(\theta )\approx \frac{1}{N}\sum _{i=1}^N\sum _{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t^i|s_t^i)\left(\hat{Q}(s_{t^{\prime }}^i,a_{t^{\prime }}^i)-b(s_t)\right)$$

ANOVA (for point 3)

From $[1]$ and including baselines we get ${\nabla }_{\theta }J(\theta )={\mathbb{E}}_{\tau \sim p_{\theta }(\tau )}\left[{\nabla }_{\theta }\log p_{\theta }(\tau )(r(\tau )-b)\right]$. Now the variance is:

$$\begin{array}{rl}Var& ={\mathbb{E}}_{\tau \sim p_{\theta }(\tau )}\left[({\nabla }_{\theta }\log p_{\theta }(\tau )(r(\tau )-b){)}^2\right]-{\mathbb{E}}_{\tau \sim p_{\theta }(\tau )}{\left[{\nabla }_{\theta }\log p_{\theta }(\tau )(r(\tau )-b)\right]}^2\\ & ={\mathbb{E}}_{\tau \sim p_{\theta }(\tau )}\left[({\nabla }_{\theta }\log p_{\theta }(\tau )(r(\tau )-b){)}^2\right]-{\mathbb{E}}_{\tau \sim p_{\theta }(\tau )}{\left[{\nabla }_{\theta }\log p_{\theta }(\tau )(r(\tau ))\right]}^2\end{array}$$

since baselines are unbiased in expectation

$$\begin{array}{rl}\frac{dVar}{db}& =\frac{d}{db}\mathbb{E}[g(\tau {)}^2(r(\tau )-b{)}^2]+0\\ & =\frac{d}{db}(\mathbb{E}[g(\tau {)}^{2}r(\tau {)}^2]-2\mathbb{E}[g(\tau {)}^{2}r(\tau )b+b^2\mathbb{E}[g(\tau {)}^2]])\\ & =0-2\mathbb{E}[g(\tau {)}^{2}r(\tau )]+2b\mathbb{E}[g(\tau {)}^2]=0\\ ⟹b& =\frac{\mathbb{E}[g(\tau {)}^{2}r(\tau )]}{\mathbb{E}[g(\tau {)}^2]}\end{array}$$

which is essentially just the expected rewards, weighted by gradient magnitudes

---

Off-policy Policy Gradients

The problem with the above algorithm is that it is completely on-policy, i.e. once the update happens, all the collected trajectories are discarded and new trajectories must to collected to perform the next policy gradient update. It would be more efficient if we did not have to discard all the previously collected trajectories since it took effort to collect the same. Hence, we look at ways to adapt the PG algorithm for the off-policy case.

We start with Importance Sampling, say we have an expectation of some function of $x$ w.r.t. distribution $p(x)$:

$$\begin{array}{rl}{\mathbb{E}}_{x\sim p(x)}[f(x)]& =\underset{x}{\int }p(x)\cdot f(x)\cdot dx\\ & =\underset{x}{\int }\frac{q(x)}{q(x)}\cdot p(x)\cdot f(x)\cdot dx\\ & =\underset{x}{\int }q(x)\cdot \frac{p(x)}{q(x)}\cdot f(x)\cdot dx\\ & ={\mathbb{E}}_{x\sim q(x)}\left[\frac{p(x)}{q(x)}\cdot f(x)\right]\end{array}$$

We can choose any distribution $q(x)$ (it’s under our control) as long as the support of the distribution is same as that of the $p(x)$ and $q(x)>0$ for all such $x$ where $p(x)>0$.

Assuming we have trajectories from some other distribution $\bar{p}$ instead of $p_{\theta }$, we can modify the PG algorithm as follows:

$$\begin{array}{rl}J(\theta )& ={\mathbb{E}}_{\tau \sim \bar{p}(\tau )}\left[\frac{p_{\theta }(\tau )}{\bar{p}(\tau )}\cdot r(\tau )\right]\end{array}$$

Now,

$$\begin{array}{r}\frac{p_{\theta }(\tau )}{\bar{p}(\tau )}=\frac{\rho (s_0)\prod _{t=0}^T{\pi }_{\theta }(a_t|s_t)\cdot p(s_{t+1}|a_{t+1})}{\rho (s_0)\prod _{t=0}^T{\bar{\pi }}_{\theta }(a_t|s_t)\cdot p(s_{t+1}|a_{t+1})}=\frac{\prod _{t=0}^T{\pi }_{\theta }(a_t|s_t)}{\prod _{t=0}^T{\bar{\pi }}_{\theta }(a_t|s_t)}...[2]\end{array}$$

Let’s assume that our current policy (after the $t^{\text{th}}$) policy update is parameterized by ${\theta }^t$ (that is used for collecting new transitions) whereas the policies that collected the previous samples (which are stored in a buffer) can be parameterized by ${\theta }^{t-k}$. In doing so we are overloading some notations, therefore we will assume that the previously collected dataset was collected by some policy ${\pi }^{\theta }$ and the current policy is ${\pi }^{{\theta }^{\prime }}$ (note that by convention objects in the past are denoted by a letter and for corresponding future objects a prime is added, e.g. if current state is $s$, next state is often denoted as $s^{\prime }$). Therefore we would like to update ${\pi }_{{\theta }^{\prime }}$ using samples from ${\pi }_{\theta }:$

$$\begin{array}{rl}J({\theta }^{\prime })& ={\mathbb{E}}_{\tau \sim p_{\theta }(\tau )}\left[\frac{p_{{\theta }^{\prime }}(\tau )}{p_{\theta }(\tau )}\cdot r(\tau )\right]\\ {\nabla }_{thet{a}^{\prime }}J({\theta }^{\prime })& ={\mathbb{E}}_{\tau \sim p_{\theta }(\tau )}\left[\frac{{\nabla }_{{\theta }^{\prime }}{p}_{{\theta }^{\prime }}(\tau )}{p_{\theta }(\tau )}\cdot r(\tau )\right]\\ & ={\mathbb{E}}_{\tau \sim p_{\theta }(\tau )}\left[\frac{p_{{\theta }^{\prime }}(\tau )}{p_{\theta }(\tau )}\cdot {\nabla }_{{\theta }^{\prime }}\log p_{{\theta }^{\prime }}(\tau )\cdot r(\tau )\right]\end{array}$$

Replacing with $[2]$ and expanding the remaining terms we get:

$$\begin{array}{rl}{\nabla }_{{\theta }^{\prime }}J({\theta }^{\prime })& ={\mathbb{E}}_{\tau \sim p_{\theta }(\tau )}\left[\left(\prod _{t=0}^T\frac{{\pi }_{{\theta }^{\prime }}(a_t|s_t)}{{\pi }_{\theta }(a_t|s_t)}\right)\left(\sum _{t=0}^T{\nabla }_{{\theta }^{\prime }}\log {\pi }_{{\theta }^{\prime }}(a_t|s_t)\right)\left(\sum _{t=0}^{T}r(s_t,a_t)\right)\right]\end{array}$$

Now in order to decrease the variance in estimates, we can modify the above by introducing the notion of causality. We have already seen how the current reward is independent of the trajectory history. In this case notice further that the importance weights also have a notion of causality involved, i.e. current importance weights are not affected by future actions:

$$\begin{array}{rl}{\nabla }_{{\theta }^{\prime }}J({\theta }^{\prime })& ={\mathbb{E}}_{\tau \sim p_{\theta }(\tau )}\left[\sum _{t=0}^T{\nabla }_{{\theta }^{\prime }}\log {\pi }_{{\theta }^{\prime }}(a_t|s_t)\left(\prod _{t^{\prime }=0}^t\frac{{\pi }_{{\theta }^{\prime }}(a_{t^{\prime }}|s_{t^{\prime }})}{{\pi }_{\theta }(a_{t^{\prime }}|s_{t^{\prime }})}\right)\left(\sum _{t^{\prime }=t}^{T}r(s_{t^{\prime }},a_{t^{\prime }})\left(\prod _{t^{\prime \prime }=t}^{t^{\prime }}\frac{{\pi }_{\theta }(a_{t^{\prime \prime }}|s_{t^{\prime \prime }})}{{\pi }_{\theta }(a_{t^{\prime \prime }}|s_{t^{\prime \prime }})}\right)\right)\right]\end{array}$$

If we ignore the final term, we can recover an algorithm called policy iteration (cite) which has some nice convergence guarantees. However even in the reduced expression

$$\begin{array}{rl}{\nabla }_{{\theta }^{\prime }}J({\theta }^{\prime })& ={\mathbb{E}}_{\tau \sim p_{\theta }(\tau )}\left[\sum _{t=0}^T{\nabla }_{{\theta }^{\prime }}\log {\pi }_{{\theta }^{\prime }}(a_t|s_t)\left(\prod _{t^{\prime }=0}^t\frac{{\pi }_{{\theta }^{\prime }}(a_{t^{\prime }}|s_{t^{\prime }})}{{\pi }_{\theta }(a_{t^{\prime }}|s_{t^{\prime }})}\right)\left(\sum _{t^{\prime }=t}^{T}r(s_{t^{\prime }},a_{t^{\prime }})\right)\right]\end{array}$$

notice that the second term is actually exponentially hard to compute and doesn’t work very well for long horizon tasks. Hence we can look into something like a “first-order” approximation of the same, which is essentially just modifying the on-policy algorithm by adding importance weights:

$$\begin{array}{rl}{\nabla }_{\theta }{J}_{\text{On Policy}}(\theta )& \approx \frac{1}{N}\sum _{i=1}^N\sum _{t=1}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t^i|s_t^i){\hat{Q}}_t^i\\ {\nabla }_{{\theta }^{\prime }}{J}_{\text{Off Policy}}({\theta }^{\prime })& \approx \frac{1}{N}\sum _{i=1}^N\sum _{t=1}^T\frac{{\pi }_{{\theta }^{\prime }}(s_t^i,a_t^i)}{{\pi }_{\theta }(s_t^i,a_t^i)}{\nabla }_{{\theta }^{\prime }}\log {\pi }_{{\theta }^{\prime }}(a_t^i|s_t^i){\hat{Q}}_t^i\\ & \approx \frac{1}{N}\sum _{i=1}^N\sum _{t=1}^T\frac{{\pi }_{{\theta }^{\prime }}(s_t^i){\pi }_{{\theta }^{\prime }}(a_t^i|s_t^i)}{{\pi }_{\theta }(s_t^i){\pi }_{\theta }(a_t^i|s_t^i)}{\nabla }_{{\theta }^{\prime }}\log {\pi }_{{\theta }^{\prime }}(a_t^i|s_t^i){\hat{Q}}_t^i\\ & \approx \frac{1}{N}\sum _{i=1}^N\sum _{t=1}^T\frac{{\pi }_{{\theta }^{\prime }}(a_t^i|s_t^i)}{{\pi }_{\theta }(a_t^i|s_t^i)}{\nabla }_{{\theta }^{\prime }}\log {\pi }_{{\theta }^{\prime }}(a_t^i|s_t^i){\hat{Q}}_t^i\end{array}$$

It is reasonable to assume that the marginals are similar (cite) and can be ignored.

---

Straight Through Gumble Softmax (STGS) for PG Algorithms