Strange behavior of log_prob when calculating gradient

cijerezg: <p>I have the following model:</p>
<pre><code class="lang-auto">from torch.distributions import Normal

class Policy(nn.Module):
    def __init__(self, num_inp, num_out, hidden_dim):
        super().__init__()
        
        self.linear1 = nn.Linear(num_inp, hidden_dim)
        self.linear2 = nn.Linear(hidden_dim, hidden_dim)
        
        self.mean_linear = nn.Linear(hidden_dim, num_out)
        self.log_std_linear = nn.Linear(hidden_dim, num_out)

    def forward(self, obs):
        x = F.relu(self.linear1(obs))
        x = F.relu(self.linear2(x))
        mean = self.mean_linear(x)
        log_std = torch.clamp(self.log_std_linear(x), min=-2, max=20)
        std = log_std.exp()
        normal = Normal(mean, std)
        x_t = normal.rsample()
        log_prob = normal.log_prob(x_t)
        log_prob = log_prob.sum()
        return x_t, log_prob
</code></pre>
<p>the gradient with respect to the parameters can be calculated as:</p>
<pre><code class="lang-auto">import torch.autograd as autograd

model = Policy(17, 6, 50)
inp = torch.rand(1, 17)
x_t, log_prob = model(inp)

grads = autograd.grad(log_prob, model.parameters())
</code></pre>
<p>However, the gradient with respect to <code>mean_linear</code> (weight and bias) is <code>0</code>; all other gradients seem correct (linear1, linear2, log_std_linear). I have absolutely no idea why this is the case. When checking the source code of <code>log_prob`` I can see it calls both, </code>self.loc<code>(the mean) and</code>self.scale<code>(the std), therefore the gradient should propagate, but somehow the output is always</code>0```. Why does this happen? and how to fix it?</p>

KFrank: <p>Hi cijerezg!</p>
<aside class="quote no-group quote-modified" data-username="cijerezg" data-post="1" data-topic="156270" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/cijerezg/48/14945_2.png" class="avatar"> cijerezg:</div>
<blockquote>
<pre><code class="lang-auto">        normal = Normal(mean, std)
        x_t = normal.rsample()
        log_prob = normal.log_prob(x_t)
</code></pre>
<p>However, the gradient with respect to <code>mean_linear</code> (weight and bias) is <code>0</code>;</p>
</blockquote>
</aside>
<p>When you call <code>.log_prob (x_t)</code> using the same <code>Normal</code> distribution from<br>
which you <code>.rsample()</code>ed <code>x_t</code>, the <code>mean</code> drops out, so you are correctly<br>
getting zero for your gradient.</p>
<p>What’s going on is that when you use “reparameterization sampling,” the<br>
<code>Distribution</code> generates some random variate (I assume a uniform variate,<br>
but I don’t really know.) and transforms it into a variate from your desired<br>
distribution.  This transformation is differentiable with respect to the<br>
parameters that describe your distribution, so you can backpropagate<br>
through the <code>rsample()</code>ed value.</p>
<p>In the case of <code>Normal</code>, that random variate tells you how improbable<br>
your sample should be – how many standard deviations it should be<br>
from the mean of your <code>Normal</code> distribution.  <code>log_prob()</code> then tells you<br>
how improbable your sample is – that is, how many standard deviations<br>
it is <em>from that same mean.</em>  So <code>log_prob()</code> is independent of <code>mean</code>, and<br>
your gradient is zero.</p>
<p>Consider:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'1.12.0'
&gt;&gt;&gt; normal10 = torch.distributions.Normal (10.0, 1.234)
&gt;&gt;&gt; normal100 = torch.distributions.Normal (100.0, 1.234)
&gt;&gt;&gt; _ = torch.manual_seed (2022)
&gt;&gt;&gt; samp10 = normal10.rsample ([5])
&gt;&gt;&gt; _ = torch.manual_seed (2022)
&gt;&gt;&gt; samp100 = normal100.rsample ([5])
&gt;&gt;&gt; samp10
tensor([10.2363, 10.4080, 10.2846, 11.1027,  9.7478])
&gt;&gt;&gt; samp100
tensor([100.2363, 100.4080, 100.2846, 101.1027,  99.7478])
&gt;&gt;&gt; normal10.log_prob (samp10)
tensor([-1.1475, -1.1839, -1.1558, -1.5285, -1.1501])
&gt;&gt;&gt; normal100.log_prob (samp100)
tensor([-1.1475, -1.1839, -1.1558, -1.5285, -1.1501])
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

cijerezg: <p>Thanks for the insightful answer. This gave me another question, in case you know. When using RL one differentiates the log of the policy with respect to the parameters, e.g., the mean and the variance. Arguably, having the correct mean is more important than having the correct variance, e.g., some RL algorithms assume a variance and learn the mean.</p>
<p>How can one go about getting the gradient of the mean? Does it have to be done manually?</p>