What is the gradient of the objective function in the Soft Actor-Critic paper?

Question

In the paper Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor, they define the loss function for the policy network as

$$ J_\pi(\phi)=\mathbb E_{s_t\sim \mathcal D}\left[D_{KL}\left(\pi_\phi(\cdot|s_t)\Big\Vert {\exp(Q_\theta(s_t,\cdot)\over Z_\theta(s_t)}\right)\right] $$

Applying the reparameterization trick, let $a_t=f_\phi(\epsilon_t;s_t)$, then the objective could be rewritten as

$$ J_\pi(\phi)=\mathbb E_{s_t\sim \mathcal D, \epsilon \sim\mathcal N}[\log \pi_\phi(f_\phi(\epsilon_;s_t)|s_t)-Q_\theta(s_t,f_\phi(\epsilon_t;s_t))] $$

They compute the gradient of the above objective as follows

$$ \nabla_\phi J_\pi(\phi)=\nabla_\phi\log\pi_\phi(a_t|s_t)+(\nabla_{a_t}\log\pi_\phi(a_t|s_t)-\nabla_{a_t}Q(s_t,a_t))\nabla_\phi f_\phi(\epsilon_t;s_t) $$

The thing confuses me is the first term in the gradient, where does it come from? To my best knowledge, the second large term is already the gradient we need, why do they add the first term?

Answer 1

I'll give it a go here and try to answer your question, I'm not sure if this is entirely correct, so if someone thinks that it isn't please correct me.
I'll disregard expectation here to make things simpler. First, note that policy $\pi$ depends on parameter vector $\phi$ and function $f_\phi(\epsilon_t;s_t)$, and value function $Q$ depends on parameter vector $\theta$ and same function $f_\phi(\epsilon_t;s_t)$. Also, one important thing that authors mention in the paper and you didn't mention is that this solution is approximate gradient not the true gradient.
Our goal is to calculate gradient of objective function $J_\pi$ with respect to $\phi$, so disregarding the expectation we have:

$\nabla_\phi J_\pi (\phi) = \nabla_\phi \log\pi(\phi,f_\phi (\epsilon_t;s_t)) - \nabla_\phi Q(s_t,\theta,f_\phi (\epsilon_t;s_t))$

Let's see the gradient of first term on right hand side. To get the full gradient we need to calculate derivative w.r.t to both variables, $\phi$ and $f_\phi (\epsilon_t;s_t)$, so we have:

$\nabla_\phi \log\pi(\phi,f_\phi (\epsilon_t;s_t)) = \frac {\partial \log\pi(\phi,f_\phi (\epsilon_t;s_t))}{\partial \phi} + \frac{\partial \log\pi(\phi,f_\phi (\epsilon_t;s_t))}{\partial f_\phi(\epsilon_t;s_t)} \frac{\partial f_\phi(\epsilon_t;s_t)}{\partial \phi}$

This is where approximation comes, they replace $f_\phi (\epsilon_t;s_t)$ with $a_t$ in some places and we have:

$\nabla_\phi \log\pi(\phi,f_\phi (\epsilon_t;s_t)) \approx \frac {\partial \log\pi(\phi,a_t)}{\partial \phi} + \frac{\partial \log\pi(\phi,a_t)}{\partial a_t} \frac{\partial f_\phi(\epsilon_t;s_t)}{\partial \phi}$
$\nabla_\phi \log\pi(\phi,f_\phi (\epsilon_t;s_t)) \approx \nabla_\phi \log\pi(\phi,a_t) + \nabla_{a_t} \log\pi(\phi,a_t) \nabla_\phi f_\phi (\epsilon_t;s_t)$

For the second term in first expression on right hand side we have:

$\nabla_\phi Q(s_t,\theta,f_\phi (\epsilon_t;s_t)) = \frac {\partial Q(s_t,\theta,f_\phi (\epsilon_t;s_t))}{\partial \phi} + \frac{\partial Q(s_t,\theta,f_\phi (\epsilon_t;s_t))}{\partial f_\phi(\epsilon_t;s_t)} \frac{\partial f_\phi(\epsilon_t;s_t)}{\partial \phi}$
$\nabla_\phi Q(s_t,\theta,f_\phi (\epsilon_t;s_t)) \approx \frac {\partial Q(s_t,\theta,a_t)}{\partial \phi} + \frac{\partial Q(s_t,\theta,a_t)}{\partial a_t} \frac{\partial f_\phi(\epsilon_t;s_t)}{\partial \phi}$

Fist term on right hand side is 0 because $Q$ does not depend on $\phi$ so we have:

$\nabla_\phi Q(s_t,\theta,f_\phi (\epsilon_t;s_t)) \approx \nabla_{a_t}Q(s_t, \theta,a_t)\nabla_\phi f_\phi(\epsilon_t;s_t)$

Now you add up things and you get the final result.

Answer 2

This is more meant like a comment to the previous answer. I also originally thought that

$$ \nabla_{\theta}\log \pi_{\theta}(f_{\theta}(\varepsilon, s)\mid s) = \nabla_{a}\log\pi_{\theta}(a\mid s)\vert_{a=f_{\theta}(\varepsilon,s)}\nabla_{\theta}f_{\theta}(\varepsilon, s), $$ instead of

$$ \nabla_{\theta}\log \pi_{\theta}(f_{\theta}(\varepsilon, s)\mid s) = \nabla_{a}\log\pi_{\theta}(a\mid s)\vert_{a=f_{\theta}(\varepsilon,s)}\nabla_{\theta}f_{\theta}(\varepsilon, s) + \nabla_{\theta}\log\pi_{\theta}(a\mid s)\vert_{a=f_{\theta}(\varepsilon, s)}. $$ The following is certainly less elegant, but I hope that it gives some additional intuition why we need to take the gradient with respect to $a$ and $\theta$. For simplicity, I will assume that $a$ is one-dimensional, but the same argument would apply for higher dimensions. In the SAC paper, they assume that $\pi_{\theta}$ is a Gaussian distribution $\mathcal{N}(\mu_{\theta}(s), \sigma_{\theta}(s))$. Therefore: $$ \log\pi_{\theta}(a\mid s)=-\frac{1}{2}\log(2\pi) - \log\sigma_{\theta}(s)-\frac{(a-\mu_{\theta}(s))^2}{2\sigma_{\theta}(s)^2}. $$ Then the gradient becomes: \begin{align} \nabla_{\theta}\log \pi_{\theta}(f_{\theta}(\varepsilon, s)\mid s)&=-\frac{\nabla_{\theta}\sigma_{\theta}(s)}{\sigma_{\theta}(s)}-\frac{(f_{\theta}(\varepsilon, s)-\mu_{\theta}(s))(\nabla_{\theta}f_{\theta}(\varepsilon,s)-\nabla_{\theta}\mu_{\theta}(s))}{\sigma_{\theta}(s)^2} \\&+\frac{(f_{\theta}(\varepsilon, s)-\mu_{\theta}(s))^2\nabla_{\theta}\sigma_{\theta}(s)}{\sigma_{\theta}(s)^3}.\end{align}

Let us calculate now the terms on the right hand side:

$$ \nabla_{a}\log\pi_{\theta}(a\mid s)\vert_{a=f_{\theta}(\varepsilon,s)}=-\frac{f_{\theta}(\varepsilon,s)-\mu_{\theta}(s)}{\sigma_{\theta}(s)^2} $$

and \begin{align}\nabla_{\theta}\log\pi_{\theta}(a\mid s)\vert_{a=f_{\theta}(\varepsilon, s)}&=-\frac{\nabla_{\theta}\sigma_{\theta}(s)}{\sigma_{\theta}(s)}+\frac{(f_{\theta}(\varepsilon, s)-\mu_{\theta}(s))\nabla_{\theta}\mu_{\theta}(s)}{\sigma_{\theta}(s)^2}\\ &+\frac{(f_{\theta}(\varepsilon, s)-\mu_{\theta}(s))^2\nabla_{\theta}\sigma_{\theta}(s)}{\sigma_{\theta}(s)^3},\end{align} which proves the equality. For the Q-function, we apply the chain rule as usual as $Q$ does not depend on $\theta$.