What happens with policy gradient methods if rewards are differentiable?

Question

I would like some help with understanding why there is no explicit flow of information from the reward gradient to the parameters of the policy in policy gradient methods.

What I mean is the following, there are 2 scenarios:

1st - deterministic framework with given initial state $s_0$, actions $a_t = \mu_\theta(s_t)$, rewards $r(s_t, a_t, s_{t+1})$, and transitions $s_{t+1}=f(s_t, a_t)$. Assume all of these things are differentiable (maybe all is continuous). By drawing the computational graph I found I can compute the gradient of $J(\mu_\theta) = \sum_{t} r(s_t, a_t, s_{t+1})$ with respect to $\theta$. I could optimize for cumulative reward by doing gradient ascent on this objective.

2nd - framework in https://spinningup.openai.com/en/latest/spinningup/rl_intro3.html#id16 , which seems general, reading $\nabla_\theta J(\pi_\theta) = \nabla_\theta E_{\tau \sim \pi_\theta}[R(\tau)] = E_{\tau \sim \pi_\theta}[R(\tau)\sum_{t=0}^T \nabla_\theta \log{\pi_\theta (a_t | s_t)}]$.

What I don't understand (and I certainly feel confused about in an ignorant way) is why the derivation assumes that $\nabla_\theta R(\tau) = 0$ when this scenario (as it is stochastic) should include the 1st as a particular case, which does have a derivative!

It makes sense that $R(\tau)$ is only affected by $\theta$ through the change in probability of the trajectories $\tau$, but it still feels strange that $\tau = (s_0, a_0, \dots)$ is indeed one sample of $\tau = (s_0\sim \rho(s_0), a_0 \sim \pi_\theta(s_0), \dots)$ which depends on $\theta$. The deterministic reduction is obvious, but one could also think about reparametrization tricks in order to show the same point.

In other words, if the reward function were differentiable, i.e. fully differentiable known environment, how could I use this information?

Answer 1

That's exactly the point of the Policy Gradient Theorem. Let's go through he proof of this theorem - it relies on our ability to "loop" the reward expression $J(\theta,s)$, expressing it through $J(\theta,s')$ in the next state:

$$J(\theta,s) = \mathbb{E}_{\tau\sim\pi_\theta}\left[ R(\tau) | s_0 = s\right]$$ $$\begin{align} J(\theta,s) & = \sum_a\pi_\theta(a|s)\sum_{s',r}P(s',r|s,a)(r + J(\theta, s'))\\ & = \sum_aJ(\theta,a,s) \end{align} $$

Here, I've introduced the notation $J(\theta,a,s)$ to keep further derivation sane. Taking the gradient:

$$\begin{align} \nabla_\theta J(\theta,s) = \sum_{a}\left(\vphantom{\sum_a}\right.&\nabla_\theta \pi_\theta(a|s) \left(\sum_{s',r}P(s',r|s,a)(r + J(s',\theta))\right) + \\ &+\pi_\theta(a|s) \nabla_\theta \sum_{s',r}P(s',r|s,a)(r + J(s',\theta)\left.\vphantom{\sum_a}\right)\end{align} $$

With the first term in the sum we do the log trick: $$ \frac{\nabla_\theta \pi_\theta(a|s)}{\pi_\theta(a|s)} \left(\pi_\theta(a|s)\sum_{s',r}P(s',r|s,a)(r + J(s',\theta))\right) = J(\theta,a,s)\nabla_\theta\log\pi_\theta(a|s) $$ And the second term simplifies to:

$$\pi_\theta(a|s) \nabla_\theta \sum_{s',r}P(s',r|s,a)(r + J(s',\theta)) = \pi_\theta(a|s) \sum_{s'}P(s'|s,a) \nabla_\theta J(s',\theta) $$ The last equality is exactly where we loose the rewards $r$ - the gradient of the constant $r$ is zero, and we can sum over $r$ because of the transition probability normalization $\sum_{r}P(s',r|s,a) = P(s'|s,a)$.

So we've finally get the expression for the gradient:

$$\nabla_\theta J(\theta,s) = \sum_{a}\left(J(\theta,a,s)\nabla_\theta\log\pi_\theta(a|s) + \pi_\theta(a|s) \sum_{s'}P(s'|s,a) \nabla_\theta J(s',\theta) \right)$$

The Policy Gradient Theorem proof goes on though a couple more steps: unroll the last expression, rewrite through the distribution over states, then convert back from summation to the expectation over trajectories -- going through all this in detail here would too long. While we've already passed the crucial point for your question - we've got rid of the gradient of the rewards.

Edit: In response to your comment, let me just stress again that this expression is a statement of the Policy Gradient Theorem.

$$\nabla_\theta J(\theta) = \mathbb{E}_{\tau\sim\pi}\left[R(\tau)\sum\nabla_\theta\log \pi_\theta(a_t|s_t)\right]$$

The proof of Policy Gradient Theorem is pretty involved and certainly is not as simple as saying that $\nabla_\theta R(\tau) = 0$. Some authors abuse notation or cut corners when going through the proof of the theorem - for a reasonably strict exposition of the PG theorem I recommend Sutton and Barto, Chapter 13.

Another potential point of confusion is the notation $r(s_t,a_t,s_{t+1})$ that you've been using. In the most general MDP formulation rewards are random variables $r$. The probability of getting a reward $r$ and ending up in state $s_{t+1}$ is encoded in the transition probability $P(s_{t+1},r|s_t,a_t)$. (This covers the case of deterministic rewards by having a deterministic distribution over $r$ ). The derivation above uses this, most general, formulation of the MDP.