lit_sb3.distributions.categorical
source module lit_sb3.distributions.categorical
Categorical distribution for discrete action spaces.
Classes
-
CategoricalDistribution — Categorical distribution for discrete actions.
source class CategoricalDistribution(action_dim: int)
Bases : Distribution
Categorical distribution for discrete actions.
Parameters
-
action_dim : int — Number of discrete actions
Methods
-
proba_distribution_net — Create the layer that represents the distribution.
source method CategoricalDistribution.proba_distribution_net(latent_dim: int) → nn.Module
Create the layer that represents the distribution.
it will be the logits of the Categorical distribution. You can then get probabilities using a softmax.
Parameters
-
latent_dim : int — Dimension of the last layer of the policy network (before the action layer)
source method CategoricalDistribution.proba_distribution(action_logits: torch.Tensor) → SelfCategoricalDistribution
source method CategoricalDistribution.log_prob(actions: torch.Tensor) → torch.Tensor
source method CategoricalDistribution.entropy() → torch.Tensor
source method CategoricalDistribution.sample() → torch.Tensor
source method CategoricalDistribution.mode() → torch.Tensor
source method CategoricalDistribution.actions_from_params(action_logits: torch.Tensor, deterministic: bool = False) → torch.Tensor
source method CategoricalDistribution.log_prob_from_params(action_logits: torch.Tensor) → tuple[torch.Tensor, torch.Tensor]