lit_sb3.distributions.maskable.multi_categorical

source module lit_sb3.distributions.maskable.multi_categorical

Maskable multi-categorical distribution for multi-discrete action spaces.

Classes

• MaskableMultiCategoricalDistribution — MultiCategorical distribution for multi discrete actions. Supports invalid action masking.

source class MaskableMultiCategoricalDistribution(action_dims: list[int])

Bases : MaskableDistribution

MultiCategorical distribution for multi discrete actions. Supports invalid action masking.

Parameters

• action_dims : list[int] — List of sizes of discrete action spaces

Methods

• proba_distribution_net — Create the layer that represents the distribution.

• proba_distribution

• log_prob

• entropy

• sample

• mode

• actions_from_params

• log_prob_from_params

• apply_masking

source method MaskableMultiCategoricalDistribution.proba_distribution_net(latent_dim: int) → nn.Module

Create the layer that represents the distribution.

it will be the logits (flattened) of the MultiCategorical distribution. You can then get probabilities using a softmax on each sub-space.

Parameters

• latent_dim : int — Dimension of the last layer of the policy network (before the action layer)

source method MaskableMultiCategoricalDistribution.proba_distribution(action_logits: torch.Tensor) → SelfMaskableMultiCategoricalDistribution

source method MaskableMultiCategoricalDistribution.log_prob(actions: torch.Tensor) → torch.Tensor

source method MaskableMultiCategoricalDistribution.entropy() → torch.Tensor

source method MaskableMultiCategoricalDistribution.sample() → torch.Tensor

source method MaskableMultiCategoricalDistribution.mode() → torch.Tensor

source method MaskableMultiCategoricalDistribution.actions_from_params(action_logits: torch.Tensor, deterministic: bool = False) → torch.Tensor

source method MaskableMultiCategoricalDistribution.log_prob_from_params(action_logits: torch.Tensor) → tuple[torch.Tensor, torch.Tensor]

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