[Feature Request] Add MaskedCategorical Distribution

Motivation

In RL use cases, it is very common that we have an action mask to only choose from valid actions. Then in such cases, MaskedCategorical distribution can be helpful to sample and compute log_prob and entropy.

Solution

Add a MaskedCategorical like TFAgents
https://www.tensorflow.org/agents/api_docs/python/tf_agents/distributions/masked/MaskedCategorical

Here we also need to add a sparse_mask argument to support passing mask as a sparse_vector. This is helpful especially in RLHF cases. In Language Models, the action space is vocab_size which is pretty large, but the number of candidate actions can be much smaller, then it can help to improve both speed and space with sparse_mask. An example of sparse_mask is shown below.

Alternatives

In RL with action masks cases, we can manually mask the logits to implement this.

Additional context

Add any other context or screenshots about the feature request here.

Checklist

• I have checked that there is no similar issue in the repo (required)