Open youkaichao opened 4 years ago
+1
+1
Please add this feature
We are missing this feature!
I noticed that torch.distributions
supports batched probs/logits. Since my MultiDiscrete has the same length for each action dimension (it is actually the factorized, discretized policy from Sampled MuZero), it worked for me to use the current version of Categorical with
dist = torch.distributions.Categorical(probs=policy_probs)
action_bins = dist.sample()
, where policy_probs
is of size (batch_size, action_dims, num_bins)
, actions_bins
is of size (batch_size, action_dims)
.
For those who want a Multi-Categorical because you have a different length in each dimension, I guess you can try padding them to the same length, probably with torch.nn.utils.rnn.pad_sequence.
From my perspective, the @desaixie's answer provides an elegant solution to the raised problem. As for me, the issue can be closed.
🚀 Feature
Support for Multi-Categorical in torch.distributions
Motivation
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
Alternatives
Possible implementation
The implementation is very simple and straightforward. Here is just an example:
cc @vincentqb @fritzo @neerajprad @alicanb @vishwakftw