notes on policy gradient method

The Mathematical Derivation

The transformation relies on a simple identity from calculus: $\nabla_x \log f(x) = \frac{\nabla_x f(x)}{f(x)}$. Rearranging this gives us the key identity, log-derivative trick:

$$\nabla_x f(x) = f(x) \nabla_x \log f(x)$$

Now, let’s apply this to the policy gradient derivation.

1. Start with the objective function in its integral form (monte carlo estimation of expectation)

   $$J(\theta) = \mathbb{E}_{\tau\sim \pi_\theta}[R(\tau)]=\int_{\tau} p(\tau|\theta) R(\tau) d\tau$$

2. Apply the gradient operator.

   $$\nabla_{\theta} J(\theta) = \nabla_{\theta} \int_{\tau} p(\tau|\theta) R(\tau) d\tau$$

3. Move the gradient inside the integral. (This is valid under mild conditions).

   $$\nabla_{\theta} J(\theta) = \int_{\tau} \nabla_{\theta} p(\tau|\theta) R(\tau) d\tau$$

4. Apply the log-derivative trick. We substitute $\nabla_{\theta} p(\tau|\theta)$ with $p(\tau|\theta) \nabla_{\theta} \log p(\tau|\theta)$.

   $$\nabla_{\theta} J(\theta) = \int_{\tau} p(\tau|\theta) \left( \nabla_{\theta} \log p(\tau|\theta) \right) R(\tau) d\tau$$

5. Recognize the expression as an expectation. The integral is now an expectation of the term $\left( \nabla_{\theta} \log p(\tau|\theta) \right) R(\tau)$ over the distribution $p(\tau|\theta)$.

   $$\nabla_{\theta} J(\theta) = \mathbb{E}_{\tau \sim p(\tau|\theta)} \left[ \left( \nabla_{\theta} \log p(\tau|\theta) \right) R(\tau) \right]$$

   This completes the derivation.

Why exactly do we need the log-derivative trick?

There are two reasons to use the log-derivative trick:

• Fundamental reason: not in Monte Carlo form

  This is the more fundamental issue from a machine learning perspective. Our goal is to approximate the integral using samples. The standard Monte Carlo estimation principle is used to approximate an expectation of a function f(x) under a distribution p(x):

  $$\mathbb{E}_{x \sim p(x)}[f(x)] = \int p(x) f(x) dx \approx \frac{1}{N} \sum_{i=1}^{N} f(x_i), \quad \text{where } x_i \sim p(x)$$

  The recipe is simple: sample from $p(x)$ and average the value of $f(x)$ .

  Now, let’s look at the problematic form from Step 3:

  $\int_{\tau} \underbrace{\left[ \nabla_{\theta} p(\tau|\theta) \right] R(\tau)}_{\text{The whole integrand}} d\tau$

  This integral is not in the standard Monte Carlo form $\int p(\tau|\theta) f(\tau) d\tau$. The probability distribution we can easily sample from, $p(\tau|\theta)$, is part of the term being differentiated. We don’t have a clean separation between the “distribution to sample from” and the “function to evaluate”.

  monte carlo estimation require a $p$ and a $f$，step3 only has $\nabla p$, thus can’t write in MC estimation form to $\mathbb{E}[...]$

  To estimate this integral, you would need to:

  1. Sample trajectories $\tau_i$ from some base distribution (e.g., a uniform distribution over all possible trajectories, which is impossible).
  2. For each $\tau_i$, evaluate the entire integrand $\left[ \nabla_{\theta} p(\tau_i|\theta) \right] R(\tau_i)$. This is not a viable path.

  How the log-trick solves this: The trick rearranges the integral into the perfect Monte Carlo form

  Now we have a clear recipe that matches the Monte Carlo template:

  $$\nabla_{\theta} J(\theta) = \int_{\tau} \underbrace{p(\tau|\theta)}_{\text{Distribution } p(x)} \underbrace{\left[ \left( \nabla_{\theta} \log p(\tau|\theta) \right) R(\tau) \right]}_{\text{Function } f(x)} d\tau$$
  1. Sample a trajectory $\tau_i$ from the distribution $p(\tau|\theta)$. This is easy: just run the current policy $\pi_{\theta}$ in the environment.
  2. Evaluate the function $f(\tau_i) = \left( \nabla_{\theta} \log p(\tau_i|\theta) \right) R(\tau_i)$ for that trajectory.
  3. Average the results over many trajectories to get an estimate of the gradient.

• Transform a product gradient into a sum

  Let’s expand what $p(\tau|\theta)$ actually is for a trajectory of length $T$:

  $$p(\tau|\theta) = p(s_0) \prod_{t=0}^{T-1} \pi_{\theta}(a_t|s_t) p(s_{t+1}|s_t, a_t)$$

  Now, let’s try to compute its gradient, $\nabla_{\theta} p(\tau|\theta)$. Since the environment dynamics $p(s_{t+1}|s_t, a_t)$ and initial state $p(s_0)$ do not depend on $\theta$, we are taking the gradient of a long product of policy terms:

  $$\nabla_{\theta} p(\tau|\theta) \propto \nabla_{\theta} \left[ \prod_{t=0}^{T-1} \pi_{\theta}(a_t|s_t) \right]$$

  To differentiate a product, you must use the product rule. For just three terms, it’s already messy:

  $$\nabla(f_1 f_2 f_3) = (\nabla f_1)f_2 f_3 + f_1(\nabla f_2)f_3 + f_1 f_2(\nabla f_3)$$

  For a typical trajectory with hundreds of steps ($T \gg 1$), this results in an explosion of terms. It’s computationally very expensive and numerically unstable, as you are multiplying many small probability values, leading to potential underflow (vanishing gradients).

  $\log p(\tau|\theta) = \text{const} + \sum_{t=0}^{T-1} \log \pi_{\theta}(a_t|s_t)$

How the log-trick solves this: The logarithm turns products into sums:

$$
\nabla_{\theta} \log p(\tau|\theta) = \sum_{t=0}^{T-1} \nabla_{\theta} \log \pi_{\theta}(a_t|s_t)
$$

The gradient of a sum is the sum of gradients:

$$
\nabla_{\theta} \log p(\tau|\theta) = \sum_{t=0}^{T-1} \nabla_{\theta} \log \pi_{\theta}(a_t|s_t)
$$

This is vastly simpler to compute and more numerically stable.

$$
\nabla_{\theta} \log p(\tau|\theta) = \sum_{t=0}^{T-1} \nabla_{\theta} \log \pi_{\theta}(a_t|s_t)
$$

Surrogate losses in practice

Why does the policy gradient loss look different in code for on-policy and off-policy cases?

When deriving policy gradient methods, we usually start from the objective

$$J(\theta) = \mathbb{E}_{\pi_\theta}[R(\tau)]$$

where $\tau$ denotes a trajectory and $R(\tau)$ its total return.

If we consider an off-policy setting and apply importance sampling, the objective becomes

$$J(\theta) = \mathbb{E}_{\pi_{\text{old}}}\!\left[\frac{\pi_\theta(\tau)}{\pi_{\text{old}}(\tau)} R(\tau)\right].$$

However, in practical implementations, we often observe that:

• in the on-policy case, the loss is implemented as

  loss = log_probs * advantage

• in the off-policy case, the loss becomes

  loss = (pi_theta / pi_old) * advantage

Why do these implementations look different from the original objective $J(\theta)$?

---

Short answer

The loss used in code is not a direct implementation of the objective function $J(\theta)$.

Instead, it is a surrogate loss, constructed such that its gradient with respect to $\theta$ is exactly the gradient of the original objective we want to optimize.

Since we usually perform gradient descent in code, maximizing $J(\theta)$ via gradient ascent is equivalent to minimizing $-J(\theta)$. The implemented loss therefore includes a negative sign and is designed purely to produce the correct gradient.

Below, we derive this step by step.

---

Step 1: Objective and gradient

Our goal is to maximize

$$J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta}[R(\tau)].$$

To optimize it using gradient-based methods, we need

$$\nabla_\theta J(\theta) = \nabla_\theta \mathbb{E}_{\tau \sim \pi_\theta}[R(\tau)].$$

The main difficulty is that the expectation is taken with respect to $\pi_\theta$, which itself depends on $\theta$

---

Step 2: Log-derivative trick

The key mathematical tool behind policy gradient methods is the log-derivative trick, which allows us to move the gradient inside the expectation:

$$\nabla_\theta J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta} \big[\nabla_\theta \log \pi_\theta(\tau) \cdot R(\tau)\big].$$

Here,

$$\log \pi_\theta(\tau) = \log \prod_{t=0}^T \pi_\theta(a_t \mid s_t) = \sum_{t=0}^T \log \pi_\theta(a_t \mid s_t),$$

where we ignore environment transition probabilities since they do not depend on $\theta$.

Substituting this back, we obtain the standard form of the policy gradient theorem:

$$\nabla_\theta J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta} \left[\sum_{t=0}^T \nabla_\theta \log \pi_\theta(a_t \mid s_t)\, R(\tau)\right].$$

This expression is now computable. In code,

log_probs corresponds to $\log \pi_\theta(a_t \mid s_t)$.

---

Step 3: Variance reduction (from return to advantage)

Although the above estimator is unbiased, it has very high variance. Two standard refinements are applied.

1. Causality (reward-to-go)

An action at time $t$ should not be influenced by rewards received before $t$. Therefore, the total return $R(\tau)$ can be replaced by the reward-to-go $G_t$:

$$\nabla_\theta J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta} \left[\sum_{t=0}^T \nabla_\theta \log \pi_\theta(a_t \mid s_t)\, G_t\right].$$

2. Baseline

We may subtract a baseline $b(s_t)$ that depends only on the state, without changing the expected gradient. A common choice is the value function $V(s_t)$.

This leads to the advantage function

$$A(s_t, a_t) = G_t - V(s_t)$$

and the gradient becomes

$$\nabla_\theta J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta} \left[\sum_{t=0}^T \nabla_\theta \log \pi_\theta(a_t \mid s_t)\, A(s_t, a_t)\right].$$

---

Step 4: Constructing the surrogate loss (connection to code)

In PyTorch , we define a scalar loss and call loss.backward() to compute gradients.

However, the reward or advantage is an external scalar with no gradient. Using it directly as a loss would not allow gradients to flow to $\theta$.

Therefore, we construct a surrogate objective that:

1. has the same gradient w.r.t. $\theta$ as the original objective;
2. explicitly depends on $\theta$, enabling backpropagation.

If we define the surrogate loss

$$L_{\text{PG}}(\theta) = - \mathbb{E}_{\tau \sim \pi_\theta} \left[\sum_{t=0}^T \log \pi_\theta(a_t \mid s_t)\, A(s_t, a_t)\right],$$

then

$$\nabla_\theta L_{\text{PG}}(\theta) = - \nabla_\theta J(\theta).$$

In practice, the expectation is approximated by a minibatch, and

• log_probs corresponds to $\log \pi_\theta(a_t \mid s_t)$
• advantage corresponds to $A(s_t, a_t)$

Thus, the implementation

loss = - (log_probs * advantage).mean()

is exactly minimizing $L_{\text{PG}}$, which is equivalent to maximizing $J(\theta)$.

Note that the advantage is treated as a constant and is usually detached from the computation graph.

---

On importance sampling and off-policy objectives

In some implementations (e.g. in ROLL), we encounter a loss of the form

$$\text{loss} = \frac{\pi_\theta(a_t \mid s_t)}{\pi_{\text{old}}(a_t \mid s_t)} A(s_t, a_t).$$

In certain cases, $\pi_\theta$ and $\pi_{\text{old}}$ are identical, so the ratio equals 1. Writing it in this form is still valid because $\pi_{\text{old}}$ is treated as detached, while gradients flow through the numerator $\pi_\theta$.

Many RL frameworks use a unified loss

$L = r(\theta) A$

for both on-policy and off-policy training.

In the on-policy case, although $r(\theta)=1$ in value,

the denominator is detached, so gradients still flow through $\pi_\theta$

More importantly, this objective is a first-order equivalent surrogate of the standard policy gradient loss around $\theta_{\text{old}}$.

---

Why does PPO-style off-policy loss not multiply by log-probabilities?

Naively extending the on-policy surrogate would suggest a loss like

$$- \exp(\log \pi_\theta - \log \pi_{\text{old}}) \cdot \log \pi_\theta \cdot A.$$

However, in practice we use

$$L_{\text{PPO}}(\theta) = \mathbb{E}_\tau \big[\min(r_t(\theta) A_t,\; \text{clip}(r_t(\theta), 1-\epsilon, 1+\epsilon) A_t)\big],$$

where

$$r_t(\theta) = \frac{\pi_\theta(a_t \mid s_t)}{\pi_{\text{old}}(a_t \mid s_t)}.$$

The key point is that we are not patching the on-policy loss to make it off-policy.

Instead, we start from first principles and define an off-policy objective:

$$J(\theta) = \mathbb{E}_{\tau \sim \pi_{\text{old}}} \left[\frac{\pi_\theta(a \mid s)}{\pi_{\text{old}}(a \mid s)} A(s, a)\right].$$

Here, the importance ratio already contains $\pi_\theta$, so it naturally enables gradient flow.

Thus, there is no need to multiply by $\log \pi_\theta$ again.

In other words, the importance ratio serves a dual role: it corrects the distribution mismatch and simultaneously provides the necessary dependence on $\theta$ for gradient propagation.