Categorical distributions and LogSoftmax

ipaulo: <p>The question concerns the <code>torch.distributions</code> implementation. This is the canonical example from the <a href="https://github.com/pytorch/pytorch/releases" rel="nofollow noopener">relase page</a>,</p>
<pre><code>probs = policy_network(state)
# NOTE: categorical is equivalent to what used to be called multinomial
m = torch.distributions.Categorical(probs)
action = m.sample()
next_state, reward = env.step(action)
loss = -m.log_prob(action) * reward
loss.backward()
</code></pre>
<p>Usually, the probabilities are obtained from <code>policy_network</code> as a result of a softmax operation.<br>
However, if we subsequently take a logarithm of the softmax probabilities, are not we loosing the numerical precision?<br>
And if so, why  the <code>Categorical</code> distribution is not implemented in the way that allows object construction with only logits provided?</p>

smth: <blockquote>
<p>Usually, the probabilities are obtained from policy_network as a result of a softmax operation.</p>
</blockquote>
<p>you should put a <code>LogSoftmax</code> in the policy network’s last layer, not <code>Softmax</code>, precisely for the reason you described.</p>

ipaulo: <blockquote>
<p>you should put a LogSoftmax in the policy network’s last layer, not Softmax</p>
</blockquote>
<p>If I do so, I will be  required to  subsequently call <code>exp(log_probs)</code> to match a signature of the <code>torch.distributions.Categorical</code>.<br>
Also, the subsequent call to <code>torch.distributions.Categorica(probs).log_prob(0)</code> will also take the <code>log</code>, which, again, will cause all the numerical issues.<br>
So,  either I don’t understand the purpose of the current implementation of <code>torch.distributions</code> or the example on the github release page is misleading.</p>
<p>Also, probably something is weird with sampling from <code>Categorical</code>. After executing</p>
<pre><code>probs = F.softmax(torch.autograd.Variable(torch.Tensor([.25, .6])))
dist = torch.distributions.Categorical(probs)
</code></pre>
<p>this code works find</p>
<pre><code>torch.multinomial(dist.probs, 2, True)
</code></pre>
<p>but this one does not</p>
<pre><code>dist.sample_n(10)
</code></pre>
<p>yielding the following stack trace</p>
<pre><code>---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
&lt;ipython-input-8-792c4214dc6f&gt; in &lt;module&gt;()
----&gt; 1 dist.sample_n(10)

~/anaconda3/lib/python3.6/site-packages/torch/distributions.py in sample_n(self, n)
    143             return self.sample().expand(1, 1)
    144         else:
--&gt; 145             return torch.multinomial(self.probs, n, True).t()
    146 
    147     def log_prob(self, value):

RuntimeError: invalid argument 2: out of range at /Users/soumith/minicondabuild3/conda-bld/pytorch_1512381214802/work/torch/lib/TH/generic/THTensor.c:479
</code></pre>
<p>This happens because of the last <code>.t()</code>.</p>

Kaixhin: <p><a class="mention" href="/u/fritzo">@fritzo</a> does it make sense for us to have a keyword argument <code>logits</code> for <code>Categorical</code> like <a href="https://www.tensorflow.org/api_docs/python/tf/distributions/Categorical" rel="nofollow noopener">TensorFlow</a>?</p>

fritzo: <p>Yes, it makes sense, and we already do that in Pyro: <a href="https://github.com/uber/pyro/blob/dev/pyro/distributions/categorical.py#L34" rel="nofollow noopener">https://github.com/uber/pyro/blob/dev/pyro/distributions/categorical.py#L34</a></p>

Mauricio_Maroto: <p>Hi <a class="mention" href="/u/ipaulo">@ipaulo</a> Did you find the answer to your question? I am interested too.</p>

XHZhu: <p>Hi, I think the <code>log_prob</code> comes from the policy optimization algorithm (i.e., convert multiply to sum as <a href="https://spinningup.openai.com/en/latest/spinningup/rl_intro3.html" rel="noopener nofollow ugc">here</a>). The output of the policy network should still be a distribution over your action space.</p>