d.MultiCategorical Distribution

rsarpongstreetor: <p>I am creating a ProbabilisticActor for my training network. I used an action spec with  a Composite MultiCategorical spec . I need  to implement the d.MultiCategorical in the <strong>ProbabilisticActor</strong> constructor but cannot find <em><strong>MultiCategorical</strong></em> in <strong>torch.distributions</strong> nor  can I find it in  <strong>torchrl.modules.distributions</strong>. Please can you  give me the solution to fix my this situation.</p>

KFrank: <p>Hi Richard!</p>
<aside class="quote no-group" data-username="rsarpongstreetor" data-post="1" data-topic="222761" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/rsarpongstreetor/48/68703_2.png" class="avatar"> rsarpongstreetor:</div>
<blockquote>
<p>I need  to implement the d.MultiCategorical in the <strong>ProbabilisticActor</strong> constructor but cannot find <em><strong>MultiCategorical</strong></em> in <strong>torch.distributions</strong></p>
</blockquote>
</aside>
<p>This appears to be a feature that has been requested but not (yet?) implemented.</p>
<p>See this github issue (that contains some “Possible Implementation” code that<br>
might work for your use case):</p>
<aside class="onebox githubissue" data-onebox-src="https://github.com/pytorch/pytorch/issues/43250">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/issues/43250" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/pytorch</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue" data-github-private-repo="false">
	  <svg width="60" height="60" class="github-icon" viewBox="0 0 14 16" aria-hidden="true"><path fill-rule="evenodd" d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/43250" target="_blank" rel="noopener nofollow ugc">Support for Multi-Categorical in torch.distributions</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2020-08-19" data-time="03:16:27" data-timezone="UTC">03:16AM - 19 Aug 20 UTC</span>
      </div>


      <div class="user">
        <a href="https://github.com/youkaichao" target="_blank" rel="noopener nofollow ugc">
          <img alt="" src="https://avatars.githubusercontent.com/u/23236638?v=4" class="onebox-avatar-inline" width="20" height="20">
          youkaichao
        </a>
      </div>
    </div>

    <div class="labels">
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          module: distributions
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          feature
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          triaged
        </span>
    </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">## 🚀 Feature


Support for Multi-Categorical in torch.distributions

## Mot<span class="show-more-container"><a href="" rel="noopener" class="show-more">…</a></span><span class="excerpt hidden">ivation



As openai gym supports ``MultiDiscrete`` space, it would be nice if pytorch can support the corresponding distribution, too. This is also a common scenario, where the action consists of multiple sub-actions, each sub-action has different number of choices.

## Pitch


```python
# I have three sub-actions, each have 3, 5, 4 choices, respectively
dist = multi_categorical_maker([3, 5, 4])
# give it batch of logits with 5 items. 
# the first 3 logits corresponds to the first sub-action, 
# the next 5 for the second, 
# and the last 4 for the last sub-action
d = dist(logits=torch.rand(5, 12)) 

# I get 5 actions, the first sub-action ranges in [0, 3), the second in [0, 5), the last in [0, 4)
d.sample()
tensor([[0, 4, 3],
        [0, 4, 0],
        [0, 3, 1],
        [0, 0, 3],
        [0, 0, 3]])

# return the log_prob of actions, which are the sum of log_prob of sub-actions
d.log_prob(d.sample())
tensor([-3.8733, -4.3669, -4.6643, -3.5386, -4.1125])
```

## Alternatives



## Possible implementation

The implementation is very simple and straightforward. Here is just an example:
```python
import torch
from torch.distributions import Categorical, Distribution
from typing import List


class MultiCategorical(Distribution):

    def __init__(self, dists: List[Categorical]):
        super().__init__()
        self.dists = dists

    def log_prob(self, value):
        ans = []
        for d, v in zip(self.dists, torch.split(value, 1, dim=-1)):
            ans.append(d.log_prob(v.squeeze(-1)))
        return torch.stack(ans, dim=-1).sum(dim=-1)

    def entropy(self):
        return torch.stack([d.entropy() for d in self.dists], dim=-1).sum(dim=-1)

    def sample(self, sample_shape=torch.Size()):
        return torch.stack([d.sample(sample_shape) for d in self.dists], dim=-1)


def multi_categorical_maker(nvec):
    def get_multi_categorical(logits):
        start = 0
        ans = []
        for n in nvec:
            ans.append(Categorical(logits=logits[:, start: start + n]))
            start += n
        return MultiCategorical(ans)
    return get_multi_categorical

```

cc @vincentqb @fritzo @neerajprad @alicanb @vishwakftw</span></p>
  </div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>Best.</p>
<p>K. Frank</p>