Lecture: Policy Gradient

For simplicity, let us first consider a one-step decision-making process (also known as a contextual bandit).

Problem Setup

Let \(s \in \set S\) be the state vector of the environment, and let \(a \in \set A\) be the action taken by the agent. The reward function is denoted by \(r(a,s)\). The agent’s behavior is characterized by a policy, which is conveniently represented as a conditional probability distribution, \(\pi(a \mid s)\), specifying the probability of selecting each action \(a\) given a state \(s\). The objective is to find the optimal policy \(\pi^*\) that maximizes the expected reward: \[%R(\pi \mid s) R(\pi) = \mathbb{E}_{s\sim p_0}\mathbb{E}_{a\sim \pi(\cdot | s)} [r(a,s)] = \sum_{s,a} p_\pi(a, s) r(a,s),\] where \(p_0\) denotes the distribution of states, and we have \(p_{\pi}(a,s)\) denotes the joint distribution of \((a,s)\) when policy \(\pi\) is used: \[p_\pi(a,s) = \pi(a\mid s) p_0(s).\] Typically, \(p_0\) is unknown to us, but observed through a collection of data \(\{s^{(i)}\}\) drawn from \(p_0\).

Policy Gradient

In practice, the policy \(\pi(a \mid s) = \pi_\theta(a \mid s)\) is parameterized by a function with parameter \(\theta\). Assume the state-action space \(\mathcal S \times \mathcal A\) is finite or accountable, the gradient of the expected reward with respect to \(\theta\) is \[\begin{align} \nabla_\theta R(\textcolor{blue}{\pi_\theta}) & = \nabla_\theta \left ( \sum_{s,a} p_0(s) \textcolor{blue}{\pi_\theta(a \mid s)} r(a, s) \right ) \notag \\ & = \sum_{s,a} p_0(s) \textcolor{blue}{\nabla_\theta \pi_\theta(a \mid s)} r(a, s) \notag \\ & = \sum_{s,a} p_0(s) \textcolor{blue}{\pi_\theta(a \mid s)}\textcolor{blue}{ \frac{ \nabla_\theta \pi_\theta(a \mid s) }{\pi_\theta(a \mid s)} } r(a, s) \notag \\ & = \mathbb{E}_{\textcolor{blue}{p_{\pi_\theta}}} \left[ r(a, s) \textcolor{blue}{\nabla_\theta \log \pi_\theta(a \mid s)} \right], \notag \end{align}\] where the parts depending on \(\theta\) are highlighted, and we use the log-derivative trick \(\nabla_\theta \log \pi = \nabla_\theta \pi_\theta / \pi_\theta\). This allows us to express the derivative of an expectation with respect to a distribution as the expectation of the reward weighted by \(\nabla_\theta \log \pi_\theta(a \mid s)\).

Reward Baseline

For any policy \(\pi^\theta\), it is easy to prove the following identity: \[\mathbb{E}_{a\sim \pi_\theta(\cdot|s)}[\nabla_\theta \log \pi_\theta(a\mid s)] =0,~~~~~ \forall s,~\theta.\]

Hence, we can generalize the gradient formula above to \[\begin{align} \label{equ:pcgrad2} \nabla_\theta R(\pi_\theta ) & = \mathbb{E}_{p_{\pi_\theta}} \left [ (r(a,s) – \textcolor{magenta}{v(s)}) \textcolor{blue}{\nabla_\theta \log \pi_\theta (a\mid s)} \right ], \end{align}\] where \(v(s)\) is any function that does not depend on the action \(a\). The choice of \(v\) does not alter the expectation of the formula, but a proper choice of \(v\) can reduce the variance of the mean estimation given empirical observations.

One common choice of baseline is \(v(s) = \mathbb{E}_{a\sim \pi_\theta(\cdot|s)}[r(a,s)]\), in which case the difference \(r(a,s) – v(s)\) is known as the advantage function: \[A(a, s) = r(a,s) – v(s).\] A positive value \(A(a,s) > 0\) (respectively, negative \(<0\)) indicates that the action \(a\) performs above (respectively, below) the average.

On-Policy Estimation

In practice, given observations \(\{s^{(i)}\}\) drawn from \(p_0\) and \(a^{(i)} \sim \pi^\theta(\cdot|s^{(i)})\) drawn from the ongoing policy \(\pi^\theta\), the gradient can be estimated by \[\begin{align} \label{equ:onpolicyestimate} \nabla_\theta R(\pi_\theta ) \approx \frac{1}{n}\sum_{i=1}^n A^{(i)} \nabla_\theta \log \pi_\theta (a^{(i)} \mid s^{(i)}),~~~~~~ A^{(i)} = A(a^{(i)}, s^{(i)}). \end{align}\] The gradient ascent update, also known as REINFORCE, \[\theta\gets \theta + \epsilon \nabla_\theta R(\pi_\theta ).\] Intuitively, the gradient \(\nabla_\theta \log \pi_\theta (a^{(i)} \mid s^{(i)})\) for data points \((a^{(i)}, s^{(i)})\) with positive advantages \(A^{(i)} > 0\) is positively weighted, increasing the probability \(\pi_\theta(a \mid s)\) during gradient ascent. Conversely, for data points with negative advantages, the gradient is negatively weighted, decreasing the probability.

Connection to Weighted MLE

Assuming the data \(\{a^{(i)}, s^{(i)}\}\) are fixed, the policy gradient \(\nabla_\theta R(\pi_\theta)\) coincides with the gradient of the \(A^{(i)}\)-weighted log-likelihood function: \[\ell(\pi_\theta) = \frac{1}{n} \sum_{i=1}^n A^{(i)} \log \pi_\theta(a^{(i)} \mid s^{(i)}).\] This aligns with the intuition of maximizing \(\log \pi(a \mid s)\) for data points with large (and positive) \(A^{(i)}\), while minimizing the log-likelihood for negative \(A^{(i)}\).

The key difference, however, is that the data \(\{a^{(i)}, s^{(i)}\}\) are drawn from the policy \(\pi_\theta\) and thus depend on the current parameter \(\theta\). As \(\theta\) is updated during policy optimization, the data must either be regenerated or reweighted to reflect changes in \(\pi_\theta\). In contrast, maximum likelihood estimation assumes the data are fixed.

Hence, policy optimization can be viewed as an adaptively weighted maximum likelihood estimation (MLE), where the weights (\(A^{(i)}\)) are adjusted across iterations using newly generated data.

Off-Policy Estimation via Importance Sampling

The gradient estimation in [equ:onpolicyestimate] is on-policy, because the actions \(a\) must be drawn from the current policy \(\pi_\theta(\cdot | s)\). This does not yield efficient use of data as new actions must be rolled out once the policy is updated, and data from different policies cannot be used. In comparison, off-policy methods allow us to leverage actions from different policies.

One common approach to off-policy estimation is importance sampling. Assume that the data \(\{a^{(i)}, s^{(i)}\}\) are drawn from a behavior policy \(\pi^{\mathtt{data}}\). Using importance sampling, we have

\[\begin{align*} \nabla_\theta R(\pi_\theta) & = \mathbb{E}_{p_{\pi^{\texttt{data}}}} \left[\frac{\pi_\theta(a|s)}{\pi^{\mathtt{data}}(a|s)} A(a,s) \nabla_\theta \log \pi_\theta (a \mid s) \right] \\ & \approx \frac{1}{z} \sum_{i=1}^n w^{(i)} A^{(i)} \nabla_\theta \log \pi_\theta (a^{(i)} \mid s^{(i)}), ~~~~~ \end{align*}\] where \(w^{(i)}\) is the importance weight: \[w^{(i)} = \frac{\pi_\theta(a^{(i)}|s^{(i)})}{\pi^{\mathtt{data}}(a^{(i)}|s^{(i)})},\] and the normalization constant \(z\) is often taken as \(z = \sum_i w^{(i)}\).