question on Implementing the Simplest Policy Gradient from spinningup

Question

In Part 3: Intro to Policy Optimization from spinningup documentation, there is a formula to compute the estimate of the policy gradient:

''' This is an expectation, which means that we can estimate it with a sample mean. If we collect a set of trajectories $ D = \{\tau_i\}_{i=1,...,N} $ where each trajectory is obtained by letting the agent act in the environment using the policy $\pi_{\theta}$, the policy gradient can be estimated with

$$ \hat{g} = \frac{1}{|D|} \sum_{\tau \in D} \sum_{t=0}^{T} \nabla_{\theta} \log \pi_{\theta}(a_t |s_t) R(\tau), $$

where $|D|$ is the number of trajectories in $D$ (here, $N$). '''

Also in the code on the same page to compute loss:

# make loss function whose gradient, for the right data, is policy gradient 
def compute_loss(obs, act, weights):
    logp = get_policy(obs).log_prob(act)
    return -(logp * weights).mean()

Here the mean() is calculated over all the (state, action) pairs. However, the formula above calculates the mean over trajectories. There seems to be a disconnect. Are these two forms equivalant?

Answer 1

Note that the documentation states that the weights parameter is set to $R(\tau)$. Therefore, the $\nabla_\theta \log \pi_\theta (a_t|s_t)R(\tau)$ terms are the same whether calculating the policy gradient over trajectories (the equation) or over state-action pairs (the pseudocode). Also, note that each state-action pair is associated with exactly one trajectory; therefore,
$$\hat{h} = \sum_{\tau \in D} \sum_{t=0}^T \nabla_\theta \log \pi_\theta (a_t|s_t)R(\tau)$$ calculates both the unnormalized policy gradient over the trajectories and also the unnormalized policy gradient over the state-action pairs. The policy gradient over the trajectories is given as $\hat{g} = \frac{1}{|D|}\hat{h}$. Assuming there are $|K|$ state-action pairs, the policy gradient instead averaged over the state-action pairs is $\frac{1}{|K|}\hat{h} = \frac{|D|}{|K|}\left(\frac{1}{|D|}\hat{h}\right) = \frac{|D|}{|K|}\hat{g}$. Thus, the policy gradient averaged over state-action pairs is proportional by factor $\frac{|D|}{|K|}$ to the policy gradient over trajectories. Since multiplying a gradient by a constant (generally) does not change the underlying optimization problem, these two forms are practically equivalent.