Skip to content

lit_sb3.distributions.maskable.maskable_categorical

source module lit_sb3.distributions.maskable.maskable_categorical

Modified PyTorch Categorical distribution with action masking support.

Classes

  • MaskableCategorical — Modified PyTorch Categorical distribution with support for invalid action masking.

source class MaskableCategorical(probs: torch.Tensor | None = None, logits: torch.Tensor | None = None, validate_args: bool | None = None, masks: MaybeMasks = None)

Bases : Categorical

Modified PyTorch Categorical distribution with support for invalid action masking.

To instantiate, must provide either probs or logits, but not both.

Parameters

  • probs : torch.Tensor | None — Tensor containing finite non-negative values, which will be renormalized to sum to 1 along the last dimension.

  • logits : torch.Tensor | None — Tensor of unnormalized log probabilities.

  • validate_args : bool | None — Whether or not to validate that arguments to methods like lob_prob() and icdf() match the distribution's shape, support, etc.

  • masks : MaybeMasks — An optional boolean ndarray of compatible shape with the distribution. If True, the corresponding choice's logit value is preserved. If False, it is set to a large negative value, resulting in near 0 probability.

Attributes

  • batch_shape : torch.Size — Returns the shape over which parameters are batched.

  • event_shape : torch.Size — Returns the shape of a single sample (without batching).

  • stddev : Tensor — Returns the standard deviation of the distribution.

Methods

  • apply_masking — Eliminate ("mask out") chosen categorical outcomes by setting their probability to 0.

  • entropy

source method MaskableCategorical.apply_masking(masks: MaybeMasks) → None

Eliminate ("mask out") chosen categorical outcomes by setting their probability to 0.

Parameters

  • masks : MaybeMasks — An optional boolean ndarray of compatible shape with the distribution. If True, the corresponding choice's logit value is preserved. If False, it is set to a large negative value, resulting in near 0 probability. If masks is None, any previously applied masking is removed, and the original logits are restored.

source method MaskableCategorical.entropy() → torch.Tensor