Need for reparameterization trick in RL (and others)?

Question

This is a multi-fold question that has a number of closely related questions; that is why I will pose them all here, instead of separate questions.

In RL you have a parameterized policy that dictates the probability with which you pick action $a$ given parameters $\theta$ and state $s_t$. That is typically denoted as $\pi(a|\theta,s)=P(A_t=a|S_t=s,\theta_t=\theta$). When one is dealing with continuous actions, then one cannot talk about the probability of choosing an action. Instead, one chooses a density, for example Gaussian. Then the policy is given by:

$$\pi(a|s,\theta)=\dfrac{1}{\sigma(s,\theta)\sqrt{2\pi}}\exp\left(-\dfrac{(a-\mu(s,\theta)^2}{2\sigma(s,\theta)^2}\right)$$

One is interested in computing $\nabla_\theta \log \pi$. I'll do a preliminary calculation for mere illustrative purposes where I calculate the derivative with respect to $\mu_\theta$.

$$\dfrac{\partial \log\pi}{d\mu_\theta}=\dfrac{\partial \log \pi}{\partial \mu_\theta}=\dfrac{a-\mu_\theta}{\sigma}$$

1) Since an action $a$ has been sampled, can't that action $a$ be put back in $\nabla_\theta\log\pi$ an evaluate the gradient without any reparameterization trick? For instance, just plugging in $a,s,\theta$ in the above equation would result in the gradient to update the mean.

Moreover, for a Gaussian, the reparemeterization trick is as follows: $$x=\mu(s,\theta)+\sigma(s,\theta) \varepsilon$$

where $\varepsilon \sim \mathcal{N}(0, 1)$. Then $x$ is a Gaussian sample. If we want to compute the likelihood (not probability) of sample $x$, then that is: $$\pi(x|s,\theta)=\dfrac{1}{\sigma\sqrt{2\pi}}\exp\left(-\dfrac{(\mu+\sigma\varepsilon-\mu)^2}{2\sigma^2}\right)=\dfrac{1}{\sigma\sqrt{2\pi}}\exp\left(-\dfrac{(\sigma\varepsilon)^2}{2\sigma^2}\right)=\dfrac{1}{\sigma\sqrt{2\pi}}\exp\left(-\dfrac{\varepsilon}{2}\right)$$.

so, there is no dependence on $\mu$ anymore. This is corroborated by pytorch. In particular, this answer discusses that. Furthermore, when calculating $\nabla_\theta \log \pi$, the dependence on the sample is completely gone as $\nabla_\theta \log \exp(-\varepsilon/2)=0$

2) How to interpret this? how is it possible there is no dependance on $\mu$? or that $\nabla \log \pi$ does not depend at all on the value of $x$? This result seems contradictory to what I derived before ($d\log\pi/ d\phi$)

I know the idea of the reparameterization trick is that we can differentiate through it, but I don't see how it agrees with the derivation I had in the second equation. And more importantly, I don't really see why it is necessary in RL.

Answer 1

I did some more digging (which I should have done before posting this, sorry). I found the answer to my question. I would say it is very summarized on this reddit thread. The short answer is that both are possible: using the log of the distribution as well as using the reparemeterization trick.

The longer answer is (please feel free to correct me if I say something incorrect):

The gradient of the RL objective is of the form:

$$\nabla_\theta\mathbb{E}_{\tau\sim\pi(\theta)}[r(\tau)]=\nabla_\theta\int_\tau r(\tau)\pi(\theta)d\tau=\int_\tau r(\tau)\nabla_\theta\pi(\theta)d\tau \tag{1}$$

But note that this is a problem because $\int_\tau r(\tau)\nabla_\theta\pi(\theta)d\tau\neq \mathbb{E}_{\tau\sim\pi(\theta}[r(\tau)\nabla_{\theta}\pi(\theta)]$. Therefore, something needs to be done about the expression in equation (1).

One way to go about it is with the log trick, which goes as follows:

$$\int_\tau r(\tau)\pi(\theta)\dfrac{\nabla_\theta \pi(\theta)}{\pi(\theta)}d\tau=\int_{\tau} \pi(\theta)r(\tau)\nabla_\theta\log\pi(\theta)d\tau=\mathbb{E}_{\tau\sim \pi(\theta)}[r(\tau)\nabla_\theta\pi(\theta)]$$

which is a quantity that can be sampled.

The other way to go about this is with the reparameterization trick. If $\pi$ is a Gaussian, and recall that $\tau\sim \pi$. Then, the Gaussian can be sampled as $\tau=\mu(\theta)+\sigma(\theta) \varepsilon$, where $\varepsilon \sim \mathcal{N}(0,1)$

Then, (somehow that I still don't understand or I am doing wrong), we get that:

$$\int_\tau r(\tau)\nabla_\theta\pi(\theta)d\tau=\int_{\varepsilon}p_\varepsilon r(\tau)\nabla_\theta\pi(\theta)d\varepsilon=\mathbb{E}_{\tau\sim \varepsilon}[r(\tau)\nabla_\theta\pi(\theta)]$$

Intuitively, the reparameterization trick allows to separate the parameters with which the gradient is taken of from the parameters of the sampling process (the expectation).

There are surely mistakes in my derivation, especially in the reparameterization trick. I'll try to address them myself, but would appreciate if someone can help me.