Policy Gradients

The core idea of Reinforcement learning is to learn some kind of behavior through optimizing for rewards. The behavior learned by an agent i.e. the schema it follows while going through this process is the learned policy that it uses to decide which action to take and thus, the transition from one state to another. One way to close the loop for the agent to learn is by evaluating the states and actions through value functions and thus, our way to measure the learned policy is seen through these value functions, approximated by lookup tables, linear combinations, Neural Networks e.t.c. Policy Gradient methods take a different approach where they bypass the need for a value function by parameterizing the policy directly. While the agent can still use a value function to learn, it need not use it for selecting actions. The advantages that policy gradient methods offer are 3 fold:

1. Approximating the policy might be simpler than approximating action values and a policy-based method might typically learn faster and yield a superior asymptotic policy. One very good example that illustrates this is the work by Simsek et. al on the game of Tetris where they showed that it is possible to choose amongst actions without really evaluating them.
2. Policy gradient methods can handle stochastic policies. The case of card games with imperfect information, like poker, is a direct example where the optimal play might be to do 2 different things with certain probabilities. If we are maximizing the actions based on value approximations, we don’t really have a natural way of finding stochastic policies. Policy Gradient methods can do this.
3. Policy gradient methods offer stronger convergence guarantees since with continuous policies the action probabilities change smoothly. This is not the case with the fixed $ϵ$ -greedy evaluation since there is always a probability to do something random.
4. The choice of policy parameterization is sometimes a good way of injecting prior knowledge about a desired form of the policy into the system. This is especially helpful when we look at introducing Meta-Learning strategies into Reinforcement Learning.

In the following sections, I first use the theoretical treatment done in Sutton and Barto’s book since I was better able to understand policy gradients’ essence through this. However, I again do the derivation by looking at the whole thing from the viewpoint of trajectories, since I find it more intuitive.

Policy Gradient Theorem

The issue with the parameterization of the policy is that the policy affects both the action selections and the distribution of states in which those selections are made. While going from state to action is straightforward, going the other way round involves the environment and thus, the parameterization is typically unknown. Thus, with this unknown effect of policy changes on the state distributions, the issue is evaluating the gradients of the performance. This is where the policy gradient theorem comes into the picture, as it shows that the gradient of the policy w.r.t its parameters does not involve the derivative of the state distribution. For episodic tasks, if we assume that every episode starts in some particular state $s_0$, then we can write a performance measure as the value of the start state of the episode

$$J(\theta )=v_{{\pi }_{\theta }}(s_0)$$

For simplicity, we remove the $\theta$ from the subscript of $\pi$. Now, to get a derivative of this measure, we start by differentiating this value function w.r.t $\theta$:

$$\begin{array}{rl}{\mathrm{\nabla }}_{\theta }J(\theta )& =\mathrm{\nabla }{v}_{\pi }(s)\\ & =\mathrm{\nabla }[\sum _{a\in \mathcal{A}}\pi (a|s)q_{\pi }(s,a)]\\ & =\sum _a[\mathrm{\nabla }\pi (a|s)q_{\pi }(s,a)+\pi (a|s)\mathrm{\nabla }{q}_{\pi }(s,a)]\\ & =\sum _a[\mathrm{\nabla }\pi (a|s)q_{\pi }(s,a)+\pi (a|s)\mathrm{\nabla }\sum _{s^{\prime }\in \mathcal{S},r\in \mathcal{R}}p(s^{\prime },r|s,a)(r+v_{\pi }(s^{\prime }))]\end{array}$$

we now extend $q_{\pi }(s,a)$ in the second term on the right to the rollout for the new state $s^{\prime }$.

$$\mathrm{\nabla }{v}_{\pi }(s)=\sum _a[\mathrm{\nabla }\pi (a|s)q_{\pi }(s,a)+\pi (a|s)\mathrm{\nabla }\sum _{s^{\prime }\in \mathcal{S},r\in \mathcal{R}}p(s^{\prime },r|s,a)(r+v_{\pi }(s^{\prime }))]$$

The reward is independent of the parameters $\theta$, so we can set that derivative inside the sum to 0, and so, we get:

$$\mathrm{\nabla }{v}_{\pi }(s)=\sum _a[\mathrm{\nabla }\pi (a|s)q_{\pi }(s,a)+\pi (a|s)\mathrm{\nabla }\sum _{s^{\prime }\in \mathcal{S}}p(s^{\prime }|s,a)v_{\pi }(s^{\prime })]$$

Thus, we now have a recursive formulation of $\mathrm{\nabla }{v}_{\pi }(s)$ in terms of $\mathrm{\nabla }{v}_{\pi }(s^{\prime })$. To calculate this derivative in the infinite horizon episodic case we just need to unroll this infinitely many times, which can be written as

$$\sum _{x\in \mathcal{s}}\sum _{k=0}^{\mathrm{\infty }}P(s\to x,k,\pi )\sum _{a\in \mathcal{A}}\mathrm{\nabla }\pi (a|s)q_{\pi }(x,a)$$

Here, $P(s\to x,k,\pi )$ is the probability of transitioning from state $s$ to state $x$ in $k$ steps under policy $\pi$. To estimate this probability we use something called the stationary distribution of the Markov chains. This term comes from the Fundamental Theorem of Markov Chains which intuitively says that in very long random walks the probability of ending up at some state is independent of where you started. When we club all these probabilities into a distribution over the states, then we have a stationary distribution, denoted by $\mu (s)$. In on-policy training, we usually estimate this distribution by the fraction of time spent in a state. In our case of episodic tasks, if we let $\eta (s)$ denote the total time spent in a state in an episode, then we can calculate $\mu (s)$ as

$$\mu (s)=\frac{\eta (s)}{\sum _s\eta (s)}$$

In our derivation, $P(s\to x,k,\pi )$ for very long walks can be estimated by the total time spent in the state $s$. Thus, we can inject $\eta (s)$ into our equation as follows:

$$\begin{array}{rl}{\mathrm{\nabla }}_{\theta }J(\theta )& =\sum _{s\in \mathcal{S}}\eta (s)\sum _{a\in \mathcal{A}}\mathrm{\nabla }\pi (a|s)q_{\pi }(s,a)\\ & =\sum _{s\in \mathcal{S}}\eta (s)\sum _{s\in \mathcal{S}}\frac{\eta (s)}{\sum _{s\in \mathcal{S}}\eta (s)}\sum _{a\in \mathcal{A}}\mathrm{\nabla }\pi (a|s)q_{\pi }(s,a)\\ & \propto \sum _{s\in \mathcal{S}}\mu (s)\sum _{a\in \mathcal{A}}\mathrm{\nabla }\pi (a|s)q_{\pi }(s,a)\end{array}$$

Thus, we get the form of the theorem as

$${\mathrm{\nabla }}_{\theta }J(\theta )\propto \sum _{s\in \mathcal{S}}\mu (s)\sum _{a\in \mathcal{A}}{q}_{\pi }(s,a)\mathrm{\nabla }\pi (a|s)$$

The proportionality is the average length of an episode. In the case of continuous tasks, this is 1. Thus, now we see that we have a gradient over a parameterized policy, which allows us to move in the direction of maximizing this gradient i.e gradient ascent. We can estimate this gradient through different means.

REINFORCE: Monte-Carlo Sampling

For Monte-Carlo sampling of the policy, our essential requirement is sampling from a distribution that allows us to get an estimate of the policy. From the equation of the policy gradient theorem, we can write this again as an expectation over a sample of states $S_t\sim s$ in the direction of the policy gradient

$${\mathrm{\nabla }}_{\theta }J(\theta )={\mathbb{E}}_{\pi }[\sum _{a\in \mathcal{A}}{q}_{\pi }(S_t,a){\mathrm{\nabla }}_{\theta }\pi (a|S_t,\theta )]$$

The expectation above would be an expectation over the actions if we were to include the probability of selecting the actions as the weight. Thus, we can do that to remove the sum over actions too:

$$\begin{array}{rl}{\mathrm{\nabla }}_{\theta }J(\theta )& ={\mathbb{E}}_{\pi }[\sum _{a\in \mathcal{A}}{q}_{\pi }(S_t,a){\mathrm{\nabla }}_{\theta }\pi (a|S_t,\theta )\frac{\pi (a|S_t,\theta )}{\pi (a|S_t,\theta )}]\\ & ={\mathbb{E}}_{\pi }[\sum _{a\in \mathcal{A}}\pi (a|S_t,\theta )q_{\pi }(S_t,a)\frac{{\mathrm{\nabla }}_{\theta }\pi (a|S_t,\theta )}{\pi (a|S_t,\theta )}]\\ & ={\mathbb{E}}_{\pi }[q_{\pi }(S_t,A_t)\frac{{\mathrm{\nabla }}_{\theta }\pi (A_t|S_t,\theta )}{\pi (A_t|S_t,\theta )}]\end{array}$$

The expectation over $q_{\pi }(S_t,A_t)$ is essentially the return $G_t$. Thus, we can replace that in the above equation to get:

$${\mathrm{\nabla }}_{\theta }J(\theta )={\mathbb{E}}_{\pi }[G_t\frac{{\mathrm{\nabla }}_{\theta }\pi (A_t|S_t,\theta )}{\pi (A_t|S_t,\theta )}]$$

We now have a full sampling of the states and actions conditioned on our parameters in the gradients. This can be considered a sample from the policy and we can update our parameters using this quantity to get our update rule as:

$$\begin{array}{rl}{\theta }_{t+1}& ={\theta }_t+\alpha {\mathrm{\nabla }}_{\theta }J(\theta )\\ & ={\theta }_t+\alpha (G_t\frac{{\mathrm{\nabla }}_{\theta }\pi (A_t|S_t,\theta )}{\pi (A_t|S_t,\theta )})\end{array}$$

This is the REINFORCE Algorithm! We have each update which is simply the learning rate $\alpha$ multiplied by a quantity that is proportional to the return and a vector of gradients of the probability of taking a certain action in a state. From the gradient ascent logic, we can see that this vector is the direction of maximizing the probability of taking action $A_t$ again, whenever we visit $S_t$. Moreover, The update is increasing the parameter vector in this direction proportional to the return, and inversely proportional to the action probability. Since the return is evaluated till the end of the episode, this is a Monte-Carlo Algorithm. We can further refine this using the identity of log differentiation and adding the discount factor to get the update as:

$${\theta }_{t+1}=\theta +\alpha {\gamma }^t{G}_t{\mathrm{\nabla }}_{\theta }\ln \pi (A_t|S_t,\theta )$$

Looking at Trajectories

Another way to look at the above formulation is through sampled trajectories, as done in Sergey Levine’s slides. While theoretically, this treatment is similar to the one done above, I just find the notation more intuitive. Recall, a trajectory is a sequence of states and actions over time, and the rewards accumulated in this sequence qualify this trajectory. Thus, we can say that the utility objective $J(\theta )$ is the sum of the accumulated rewards of some number of trajectories sampled from a policy ${\pi }_{\theta }$

$$J(\theta )={\mathbb{E}}_{\tau \sim {\pi }_{\theta }(\tau )}[\sum _{t}r(s_t,a_t)]\approx \frac{1}{N}\sum _i\sum _{t}r(s_{i:t},a_{i:t})$$

Let this sum of reward be denoted by $G(\tau )$ for a trajectory $\tau$. Thus, we can re-write the above equation as

$$J(\theta )={\mathbb{E}}_{\tau \sim {\pi }_{\theta }(\tau )}[G(\tau )]$$

This expectation is essentially sampling a trajectory from a policy and weighing it with the accumulated rewards. Thus, we can write it as

$$J(\theta )=\int {\pi }_{\theta }(\tau )G(\tau )d\tau$$

Now, we just differentiate this objective, and add a convenient policy term to make it an expectation:

$$\begin{array}{rl}{\mathrm{\nabla }}_{\theta }J(\theta )& ={\mathrm{\nabla }}_{\theta }\int {\pi }_{\theta }(\tau )G(\tau )d\tau \\ & =\int {\mathrm{\nabla }}_{\theta }{\pi }_{\theta }(\tau )G(\tau )d\tau \\ & =\int {\pi }_{\theta }(\tau )\frac{{\mathrm{\nabla }}_{\theta }{\pi }_{\theta }(\tau )}{{\pi }_{\theta }(\tau )}G(\tau )d\tau \end{array}$$

Now, we just use an identity $\frac{dx}{x}=d\log x$ and get

$${\mathrm{\nabla }}_{\theta }J(\theta )=\int {\pi }_{\theta }(\tau ){\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(\tau )G(\tau )d\tau$$

Thus, we can write this as:

$${\mathrm{\nabla }}_{\theta }J(\theta )={\mathbb{E}}_{\tau \sim {\pi }_{\theta }(\tau )}[{\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(\tau )G(\tau )]$$

Hence, we get the same final result as before: the gradient depends on the gradient of the log policy weighted by the rewards accumulated over the trajectory. Now, we can translate this to states and actions over the trajectory by simply considering what the policy of the trajectory represents:

$$\begin{array}{rl}& {\pi }_{\theta }(s_1,a_1,...,s_T,a_T)=\mu (s)\prod _{t=1}^T{\pi }_{\theta }(a_t|s_t)p(s_{t+1}|s_t,a_t)\\ ⟹& \log {\pi }_{\theta }(s_1,a_1,...,s_T,a_T)=\log \mu (s)+\sum _{t=1}^T\log {\pi }_{\theta }(a_t|s_t)+\log p(s_{t+1}|s_t,a_t)\end{array}$$ $${\mathrm{\nabla }}_{\theta }J(\theta )={\mathbb{E}}_{\tau \sim {\pi }_{\theta }(\tau )}[\sum _{t=1}^T{\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(a_t|s_t)\sum _{t=1}^{T}r(s_t,a_t)]$$

Then we take average of the samples as the expectation, we get

$${\mathrm{\nabla }}_{\theta }J(\theta )\approx \frac{1}{N}\sum _{i=1}^N[\sum _{t=1}^T{\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(a_t|s_t)\sum _{t=1}^{T}r(s_t,a_t)]$$

And this is the key formula behind REINFORCE again and the update can thus, be written as

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

This is where I find it more intuitive to just use $\tau$ for trajectories since now we can just write our REINFORCE algorithm as :

1. Sample a set of trajectories $\{{\tau }^i\}$ from the policy $$\pi_{\mathbf{\theta}}(a_t

   s_t)$$

2. Estimate ${\mathrm{\nabla }}_{\theta }J(\theta )$
3. Update the parameters $\theta \leftarrow \theta +\alpha {\mathrm{\nabla }}_{\theta }J(\theta )$

Reducing Variance

The idea of parameterizing the policy and working directly with the sampled trajectories has an intuitive appeal due to its clarity. However, this approach suffers from high variance. This is because when we compute the expectation over trajectories, we are essentially sampling $N$ different trajectories and then taking the average of the accumulated rewards. If we take this sampling to be uniform, we can easily imagine scenarios, where the trajectories sampled, have wildly different accumulated rewards. Thus, the chances of getting a high variance in the values that we are averaging over increase. If we were to scale $N$ to $\mathrm{\infty }$ then our average becomes closer to the true expectation. However, this is computationally expensive. There are multiple ways to reduce this variance.

Rewards to go

In the trajectory formulation, we are accumulating all of the rewards from $t=0$ to $t=N$. However, one way to make this online would be to just consider the rewards from the timestep at which we take the policy log value till the end of the horizon. This can be written as

$${\mathrm{\nabla }}_{\theta }J(\theta )\approx \frac{1}{N}\sum _{i=1}^N[\sum _{t=1}^T{\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(a_t|s_t)\sum _{t^{\prime }=t}^{T}r(s_{t^{\prime }},a_{t^{\prime }})]$$

Thus, by reducing the number of rewards we consider for each policy evaluation, we are essentially better able to control the variance up to a certain extent

Baselining

Another way to control the variance is to realize that the actions in a state are the quantities that create a variance for each state in trajectory. However, we see that our return $G(\tau )$ is dependent on both states and actions. Thus, if we could compare this value to a baseline value $b(s)$, we eliminate the variance resulting from fixed state selection. In the policy gradient theorem, this would look like

$${\mathrm{\nabla }}_{\theta }J(\theta )\propto \sum _{s\in \mathcal{S}}\mu (s)\sum _{a\in \mathcal{A}}(q_{\pi }(s,a)-b(s))\mathrm{\nabla }\pi (a|s)$$

since $b(s)$ does not vary with actiosn, it would not have an effect on our summation since its sum over all actions would be 0:

$$\begin{array}{rl}\sum _{a}b(s){\mathrm{\nabla }}_{\theta }\pi (a|s,\theta )& =b(s){\mathrm{\nabla }}_{\theta }\sum _a\pi (a|s,\theta )\\ & =b(s){\mathrm{\nabla }}_{\theta }1\\ & =0\end{array}$$

Thus, we can effectively update our REINFORCE update with this baseline to get

$${\theta }_{t+1}=={\theta }_t+\alpha (G_t-b(s))\frac{{\mathrm{\nabla }}_{\theta }\pi (A_t|S_t,\theta )}{\pi (A_t|S_t,\theta )}$$

Or, in the trajectory formulation

$${\theta }_{t+1}=={\theta }_t+\alpha {\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(\tau )(G(\tau )-b(s))$$

One good function for $b(s)$ could be the estimate of the state value $\hat{v}(s_t,\mathbf{w})$ . Doing something like this may seem like going to the realm of actor-critic methods, where we are parameterizing the policy and using the value function, but this is not the case here since we are not using the value function to bootstrap. We are stabilizing the variance by using the estimate of the value function as a baseline. Baselines are not just limited to this. We can inject all kinds of things into the baseline to try scaling up our policy gradient. For example, an interesting paper published recently talks about using functions that take into account causal dependency as a Baseline. There are many other extensions like Deterministic Policy Gradients, Deep Deterministic Policy Gradients, Proximal Optimization e.t.c that look deeper into this problem.