# What happens with policy gradient methods if rewards are differentiable?

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?

• When you say "reward gradient", what specific function are you considering the gradient of, and with respect to which parameters? I don't think that $R(\tau)$ is a reward function (although it may be composed of a sum of many reward functions), and you may have trouble with different notation and emphasis used in the two sources. May 21, 2021 at 14:34
• I agree that it's probably a notation confusion. $R(\tau)$ is indeed not a reward function, but probably it's better defined as $\sum_t r_e(s_t, a_t, s_{t+1})$ where $r_e(\cdot) = E[r(\cdot)]$ being $r(\cdot)$ a possibly non-deterministic funcion. The gradients are with respect to the parameters $\theta$ of the policy $\pi_\theta$. May 22, 2021 at 10:21
• You are asking about the gradients of the "reward function". The policy clearly has a gradient with respect to $\theta$ since it is a direct function of $\theta$. The reward function $r(s,a,s')$ clearly has 0 gradient with respect to $\theta$, since no change to $\theta$ changes it. The situation is less clear when it comees to $R(\tau)$, because it depends on context - e.g. whether you consider $\tau$ to be a random variable, or a specific trajectory - but are you certain that any of your sources claim that has a zero gradient? May 22, 2021 at 12:17
• I don't agree with the reward function having 0 gradient, and I think the cumulative reward should behave accordingly. It's just enough to write $r(s, \mu_\theta(s), s')$ and the dependecy with $\theta$ is revealed. May 22, 2021 at 18:35
• The difference is between random variables and observations. There is no dependency on $\theta$ for any given observation, but there is when dealing with expected results. That is why is important to notice when the equations are referencing $\tau$ as an observation (where probabilties from the policy are already resolved) or as a random variable. May 22, 2021 at 19:22

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.

• Another great answer. May 19, 2021 at 22:23
• Thank you very much for the detailed answer. Sadly, I fail to understand why this is different from the derivation in the link in the answer. I know that the gradient of a constant is 0, the question is why is $r$ a constant. Could you try to answer the question and address the cases and contradiction there posed? May 21, 2021 at 13:42
• Hey! thanks for the update! I will go a little further by commenting another question. First, I agree that $r(s_t, a_t, s_{t+1})$ is possibly stochastic. I would ask the following, why is $P(s', r|s, a)$ treated as a constant when I could write $P(s', r | s, \pi^s_theta(s))$ (where $\pi^s$ is after sampling)? Moreover, for deterministic everything and known environment dynamics, $J(\theta, s) = E_{\tau \sim \pi_\theta}[R(\tau) | s_0] = R(\tau = \text{rollout}(\mu_\theta(s_0))) = \sum_t r_e(s_t, \mu_\theta(s_t), s_{t+1})$. Taking gradient at RHS would involve gradient of $r_e$. Is this right? May 22, 2021 at 10:30
• PS I have read the reference, but the same issue with the desaparition of the function to get a constant is there. I assume it's something like $E[R(tau)] = \int_\theta R(\tau)P(\tau)$ which for the deterministic case $=R(\tau(\theta))$ because $P$ is a delta. Differentiating this yields $\nabla_\theta R(\tau(\theta))$. I have the feeling it's a tricky math issue I'm not noticing :( can you see my point? May 22, 2021 at 10:39
• PS2 reading your answer again I think you have a point. I am messing up notaiton. Could you remind me of some math? Can I differentiate a random variable $r$ coming from a distribution parametrized by $\theta$? if not, why not, if it depends on $\theta$? May 22, 2021 at 10:48