Understanding log_prob for Normal distribution in pytorch

username: sabeaussan
cooked: <p>I’m currently trying to solve Pendulum-v0 from the openAi gym environment which has a continuous action space. As a result, I need to use a Normal Distribution to sample my actions. What I don’t understand is the dimension of the log_prob when using it :</p>
<pre><code class="lang-auto">import torch
from torch.distributions import Normal 

means = torch.tensor([[0.0538],
        [0.0651]])
stds = torch.tensor([[0.7865],
        [0.7792]])

dist = Normal(means, stds)
a = torch.tensor([1.2,3.4])
d = dist.log_prob(a)
print(d.size())
</code></pre>
<p>I was expecting a tensor of size 2 (one log_prob for each actions) but it output a tensor of size(2,2).</p>
<p>However, when using a Categorical distribution for discrete environment the log_prob has the expected size:</p>
<pre><code class="lang-auto">logits = torch.tensor([[-0.0657, -0.0949],
        [-0.0586, -0.1007]])

dist = Categorical(logits = logits)
a = torch.tensor([1, 1])
print(dist.log_prob(a).size())
</code></pre>
<p>give me a tensor a size(2).</p>
<p>Why is the log_prob for Normal distribution of a different size ?</p>

username: KFrank
cooked: <p>Hello sabeaussan!</p>
<aside class="quote no-group quote-modified" data-username="sabeaussan" data-post="1" data-topic="73809" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/dec6dc/48.png" class="avatar"> sabeaussan:</div>
<blockquote>
<p>What I don’t understand is the dimension of the log_prob when using it:</p>
<pre><code class="lang-auto">means = torch.tensor([[0.0538],
        [0.0651]])
stds = torch.tensor([[0.7865],
        [0.7792]])

dist = Normal(means, stds)
</code></pre>
</blockquote>
</aside>
<p>You have created <code>dist</code> as a “batch” of two gaussian distributions.<br>
(The first has <code>mean = 0.0538</code> and <code>std = 0.7865</code>, and the second<br>
has <code>mean = 0.0651</code> and <code>std = 0.7792</code>.)</p>
<aside class="quote no-group">
<blockquote>
<p>I was expecting a tensor of size 2 (one log_prob for each actions) but it output a tensor of size(2,2).</p>
</blockquote>
</aside>
<p>You get two <code>log_prob</code>s (one for each gaussian in your “batch”) for<br>
each of the two elements of the tensor <code>a</code>, for a total of four <code>log_prob</code>s.</p>
<p>This is illustrated by some additions I made to the code you posted:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
from torch.distributions import Normal
torch.__version__

torch.manual_seed (2020)

means = torch.FloatTensor([[0.0538],
        [0.0651]])
stds = torch.FloatTensor([[0.7865],
        [0.7792]])

dist = Normal(means, stds)
a = torch.FloatTensor([1.2,3.4])
d = dist.log_prob(a)
print(d.size())

disB = Normal (torch.FloatTensor ([0.0]), torch.FloatTensor ([1.0]))

dist.log_prob (a)
disB.log_prob (a)

dist.sample()
disB.sample()
dist.sample_n (3)
disB.sample_n (3)

b = torch.FloatTensor([5.6])
dist.log_prob (b)
disB.log_prob (b)
</code></pre>
<p>And here are the results:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; from torch.distributions import Normal
&gt;&gt;&gt; torch.__version__
'0.3.0b0+591e73e'
&gt;&gt;&gt;
&gt;&gt;&gt; torch.manual_seed (2020)
&lt;torch._C.Generator object at 0x000002A2CFE06630&gt;
&gt;&gt;&gt;
&gt;&gt;&gt; means = torch.FloatTensor([[0.0538],
...         [0.0651]])
&gt;&gt;&gt; stds = torch.FloatTensor([[0.7865],
...         [0.7792]])
&gt;&gt;&gt;
&gt;&gt;&gt; dist = Normal(means, stds)
&gt;&gt;&gt; a = torch.FloatTensor([1.2,3.4])
&gt;&gt;&gt; d = dist.log_prob(a)
&gt;&gt;&gt; print(d.size())
torch.Size([2, 2])
&gt;&gt;&gt;
&gt;&gt;&gt; disB = Normal (torch.FloatTensor ([0.0]), torch.FloatTensor ([1.0]))
&gt;&gt;&gt;
&gt;&gt;&gt; dist.log_prob (a)

-1.7407 -9.7294
-1.7301 -9.8282
[torch.FloatTensor of size 2x2]

&gt;&gt;&gt; disB.log_prob (a)

-1.6389
-6.6989
[torch.FloatTensor of size 2]

&gt;&gt;&gt;
&gt;&gt;&gt; dist.sample()

 1.0269
-0.6832
[torch.FloatTensor of size 2x1]

&gt;&gt;&gt; disB.sample()

 1.5415
[torch.FloatTensor of size 1]

&gt;&gt;&gt; dist.sample_n (3)

(0 ,.,.) =
 -0.2670
  0.7512

(1 ,.,.) =
  0.0954
  0.1236

(2 ,.,.) =
  0.4295
 -0.4615
[torch.FloatTensor of size 3x2x1]

&gt;&gt;&gt; disB.sample_n (3)

-2.1489
-1.1463
-0.2720
[torch.FloatTensor of size 3x1]

&gt;&gt;&gt;
&gt;&gt;&gt; b = torch.FloatTensor([5.6])
&gt;&gt;&gt; dist.log_prob (b)

-25.5424
-25.8980
[torch.FloatTensor of size 2x1]

&gt;&gt;&gt; disB.log_prob (b)

-16.5989
[torch.FloatTensor of size 1]
</code></pre>
<p>Best.</p>
<p>K. Frank</p>