Policy Gradient Theorem

If we want to train a reinforcement learning agent, one of the most direct ways is to optimize the policy parameters directly. Instead of learning a value function and selecting actions greedily, we want to parameterize the policy and use gradient ascent to maximize the expected return. This is where the policy gradient theorem comes in, giving us a mathematically clean way to compute the gradient of our objective function without needing to know the environment transition dynamics.

We want to find policy parameters $\theta$ that maximize expected return:

$$J(\theta)=\mathbb{E}_{\pi_\theta}\Big[\sum_{t=0}^{\infty}\gamma^t r_t\Big]$$

This expectation is taken over all possible trajectories:

$$\tau = (s_0,a_0,r_0,s_1,a_1,\dots)$$

generated by the policy:

$$\pi_\theta(a|s)$$

So, what we want to do is take trajectories and also take into account the return along those trajectories. We take into account the probability of taking that trajectory and the total return along that trajectory:

$$J(\theta)=\sum_{\tau} P_\theta(\tau)\,R(\tau)$$

where:

• $P_\theta(\tau)$ is the probability of the trajectory
• $R(\tau)$ is the total discounted return along that trajectory

We want the gradient:

$$\nabla_\theta J(\theta)=\nabla_\theta \sum_{\tau}P_\theta(\tau)R(\tau)$$

This gradient tells us to move in the direction that maximizes the expected return of our policy. We do this because we want to account for the probability of each trajectory and its expected return:

• If a trajectory has very high reward but extremely tiny probability, its influence on the gradient is small.
• If a trajectory is common and moderately good, it influences learning more.

We can rewrite the gradient using the log derivative trick. Recall the derivative of a logarithm using the chain rule:

$$\frac{d}{dx}\log f(x)=\frac{f'(x)}{f(x)}$$

If we rearrange this, we get:

$$f'(x)=f(x) \frac{d}{dx}\log f(x)$$

So, our original gradient can be written as:

$$\nabla_\theta J(\theta) = \sum_{\tau} P_\theta(\tau)\, \nabla_\theta \log P_\theta(\tau)\, R(\tau)$$

This matches the definition of an expectation:

$$\mathbb{E}_{x \sim P}[f(x)] = \sum_{x} P(x) f(x)$$

So, we can rewrite the gradient as an expectation:

$$\nabla_\theta J(\theta) = \mathbb{E}_{\tau\sim P_\theta} \Big[ \nabla_\theta \log P_\theta(\tau)\,R(\tau) \Big]$$

By converting it to an expectation, we can use Monte Carlo sampling and avoid summing over every possible trajectory the agent could take, since that would be infinite and impossible. We simply compute $\nabla_{\theta}\log P_{\theta}(\tau)R(\tau)$ for our samples and average them to get our gradient update.

Wait, how does $\log P_\theta(\tau)$ turn into policy terms? If we write out the probability of a trajectory:

$$P_\theta(\tau) = P(s_0) \prod_t \pi_\theta(a_t|s_t) P(s_{t+1}|s_t, a_t)$$

When we take the log and then the gradient with respect to $\theta$, the transition dynamics and initial state distribution drop out since they do not depend on $\theta$ at all:

$$\nabla_\theta \log P_\theta(\tau) = \sum_t \nabla_\theta \log \pi_\theta(a_t|s_t)$$

This is clean because it means we do not need to know the transition dynamics of the environment to compute the gradient.

So the final formula becomes:

$$\nabla_\theta J(\theta) = \mathbb{E}_{\pi_\theta} \Big[ \sum_t \nabla_\theta \log \pi_\theta(a_t|s_t) \;R(\tau) \Big]$$

Now, instead of using the full return of the entire trajectory $R(\tau)$ at every step, we can use the future return from that step onward:

$$G_t = \sum_{k=t}^{\infty}\gamma^{k-t}r_k$$

We can do this because of causality: actions taken at time step $t$ cannot affect rewards that were already received in the past (before $t$). Mathematically, the expectation of the gradient terms multiplied by past rewards is zero, so dropping them does not change the expected gradient. This trick helps reduce variance.

So we get:

$$\nabla_\theta J(\theta) = \mathbb{E}_{\pi_\theta} \Big[ \sum_t \nabla_\theta \log \pi_\theta(a_t|s_t) \;G_t \Big]$$

Reinforce Algorithm

Reinforce is a specific Monte Carlo implementation of the policy gradient theorem. Instead of using the action value function $Q^\pi(s,a)$, it uses the discounted return from time $t$ onward:

$$G_t = \sum_{k=t}^{\infty}\gamma^{k-t}r_k$$

The update becomes:

$$\nabla_\theta J(\theta) \approx \sum_t \nabla_\theta \log \pi_\theta(a_t|s_t)\,G_t$$

This is the Reinforce update.

How it works:

1. The agent interacts with the environment for a fixed number of steps or until an episode is complete using the current policy: $$\pi_\theta(a|s)$$
2. For each time step $t$, observe state $s_t$, select action $a_t \sim \pi_\theta(\cdot|s_t)$, execute it, and receive reward $r_t$. Store this transition. This produces a trajectory: $$\tau = (s_0,a_0,r_0,s_1,a_1,r_1,\dots)$$
3. We compute the future returns from time $t$: $$G_t = r_t + \gamma r_{t+1} + \gamma^2 r_{t+2} + \dots$$
4. For each time step, we compute the gradient of the log probability of the action: $$g_t = \nabla_\theta \log \pi_\theta(a_t|s_t)$$
5. We scale this gradient by the return: $g_t G_t$.
6. We sum or average the gradients across the trajectory: $$\nabla_\theta J(\theta) \approx \sum_t g_t G_t$$
7. Apply gradient ascent to update the policy parameters: $$\theta \leftarrow \theta + \alpha \sum_t \nabla_\theta \log \pi_\theta(a_t|s_t)\,G_t$$

And that is the core of REINFORCE. It is simple to implement but has a reputation for high variance since a single bad rollout can wildly skew the gradient estimate. To fix that, we usually introduce a baseline to reduce variance, leading to actor critic methods which make training way more stable.