Why does (not) the distribution of states depend on the policy parameters that induce it?

Question

I came across the following proof of what's commonly referred to as the log-derivative trick in policy-gradient algorithms, and I have a question -

While transitioning from the first line to the second, the gradient with respect to policy parameters $\theta$ was pushed into the summation. What bothers me is how it skipped over $\mu (s)$, the distribution of states - which (the way I understand it), is induced by the policy $\pi_\theta$ itself! Why then does it not depend on $\theta$?

Let me know what's going wrong! Thank you!

Answer 1

The proof you are given in the above post is not wrong. It's just they skip some of the steps and directly written the final answer. Let me go through those steps:

I will simplify some of the things to avoid complication but the generosity remains the same. Like I will think of the reward as only dependent on the current state, $s$, and current action, $a$. So, $r = r(s,a)$

First, we will define the average reward as: $$r(\pi) = \sum_s \mu(s)\sum_a \pi(a|s)\sum_{s^{\prime}} P_{ss'}^{a} r $$ We can further simplify average reward as: $$r(\pi) = \sum_s \mu(s)\sum_a \pi(a|s)r(s,a) $$ My notation may be slightly different than the aforementioned slides since I'm only following Sutton's book on RL. Our objective function is: $$ J(\theta) = r(\pi) $$ We want to prove that: $$ \nabla_{\theta} J(\theta) = \nabla_{\theta}r(\pi) = \sum_s \mu(s) \sum_a \nabla_{\theta}\pi(a|s) Q(s,a)$$

Now let's start the proof: $$\nabla_{\theta}V(s) = \nabla_{\theta} \sum_{a} \pi(a|s) Q(s,a)$$ $$\nabla_{\theta}V(s) = \sum_{a} [Q(s,a) \nabla_{\theta} \pi(a|s) + \pi(a|s) \nabla_{\theta}Q(s,a)]$$ $$\nabla_{\theta}V(s) = \sum_{a} [Q(s,a) \nabla_{\theta} \pi(a|s) + \pi(a|s) \nabla_{\theta}[R(s,a) - r(\pi) + \sum_{s^{\prime}}P_{ss^{\prime}}^{a}V(s^{\prime})]]$$ $$\nabla_{\theta}V(s) = \sum_{a} [Q(s,a) \nabla_{\theta} \pi(a|s) + \pi(a|s) [- \nabla_{\theta}r(\pi) + \sum_{s^{\prime}}P_{ss^{\prime}}^{a}\nabla_{\theta}V(s^{\prime})]]$$ $$\nabla_{\theta}V(s) = \sum_{a} [Q(s,a) \nabla_{\theta} \pi(a|s) + \pi(a|s) \sum_{s^{\prime}}P_{ss^{\prime}}^{a}\nabla_{\theta}V(s^{\prime})] - \nabla_{\theta}r(\pi)\sum_{a}\pi(a|s)$$ Now we will rearrange this: $$\nabla_{\theta}r(\pi) = \sum_{a} [Q(s,a) \nabla_{\theta} \pi(a|s) + \pi(a|s) \sum_{s^{\prime}}P_{ss^{\prime}}^{a}\nabla_{\theta}V(s^{\prime})] - \nabla_{\theta}V(s)$$ Multiplying both sides by $\mu(s)$ and summing over $s$: $$\nabla_{\theta}r(\pi) \sum_{s}\mu(s)= \sum_{s}\mu(s) \sum_{a} Q(s,a) \nabla_{\theta} \pi(a|s) + \sum_{s}\mu(s) \sum_a \pi(a|s) \sum_{s^{\prime}}P_{ss^{\prime}}^{a}\nabla_{\theta}V(s^{\prime}) - \sum_{s}\mu(s) \nabla_{\theta}V(s)$$ $$\nabla_{\theta}r(\pi) = \sum_{s}\mu(s) \sum_{a} Q(s,a) \nabla_{\theta} \pi(a|s) + \sum_{s^{\prime}}\mu(s^{\prime})\nabla_{\theta}V(s^{\prime}) - \sum_{s}\mu(s) \nabla_{\theta}V(s)$$ Now we are there: $$\nabla_{\theta}r(\pi) = \sum_{s}\mu(s) \sum_{a} Q(s,a) \nabla_{\theta} \pi(a|s)$$ This is the policy gradient theoram for average reward formulation (ref. Policy gradient).

Answer 2

The reason you are confused is because this is not the full derivation of the Policy Gradient Theorem. You are correct in thinking that $\mu(s)$ depends on the policy $\pi$ which in turn depends on the policy parameters $\theta$, and so there should be a derivative of $\mu$ wrt $\theta$, however the Policy Gradient Theorem doesn't require you to take this derivative.

In fact, the great thing about the Policy Gradient Theorem is that the final result does not require you to take a derivative of the state distribution with respect to the policy parameters. I would encourage you to read and go through the derivation of the Policy Gradient Theorem from e.g. Sutton and Barto to see why you don't need to take the derivative.

Above is an image of the Policy Gradient Theorem proof from the Sutton and Barto book. If you carefully go through this line by line you will see that you are not required to take a derivative of the state distribution anywhere in the proof.