MaskedCategorical¶
- class torchrl.modules.MaskedCategorical(logits: Optional[Tensor] = None, probs: Optional[Tensor] = None, *, mask: Optional[Tensor] = None, indices: Optional[Tensor] = None, neg_inf: float = - inf, padding_value: Optional[int] = None)[source]¶
MaskedCategorical distribution.
Reference: https://www.tensorflow.org/agents/api_docs/python/tf_agents/distributions/masked/MaskedCategorical
- Parameters:
logits (torch.Tensor) – event log probabilities (unnormalized)
probs (torch.Tensor) – event probabilities. If provided, the probabilities corresponding to masked items will be zeroed and the probability re-normalized along its last dimension.
- Keyword Arguments:
mask (torch.Tensor) – A boolean mask of the same shape as
logits
/probs
whereFalse
entries are the ones to be masked. Alternatively, ifsparse_mask
is True, it represents the list of valid indices in the distribution. Exclusive withindices
.indices (torch.Tensor) – A dense index tensor representing which actions must be taken into account. Exclusive with
mask
.neg_inf (
float
, optional) – The log-probability value allocated to invalid (out-of-mask) indices. Defaults to -inf.padding_value – The padding value in the mask tensor. When sparse_mask == True, the padding_value will be ignored.
torch.manual_seed (>>>) –
torch.randn (>>> logits =) –
torch.tensor (>>> mask =) –
MaskedCategorical (>>> dist =) –
dist.sample (>>> sample =) –
print (>>>) –
tensor ([-inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf]) –
print –
-1.0831, (tensor([-1.1203, -1.0928, -1.0831, -1.1203, -1.1203, -1.0831, -1.1203,) – -1.1203, -1.1203])
print –
tensor –
probabilities (>>> # with) –
torch.ones (>>> prob =) –
prob.sum() (>>> prob = prob /) –
torch.tensor –
MaskedCategorical –
print –
-2.1972, (tensor([ -inf, -2.1972, -2.1972, -2.1972, -2.1972, -2.1972, -2.1972,) – -2.1972, -2.1972])