lit_sb3.distributions

source package lit_sb3.distributions

Probability distributions for reinforcement learning.

Classes

• Distribution — Abstract base class for distributions.

• DiagGaussianDistribution — Gaussian distribution with diagonal covariance matrix, for continuous actions.

• SquashedDiagGaussianDistribution — Gaussian distribution with diagonal covariance matrix, followed by a squashing function (tanh) to ensure bounds.

• CategoricalDistribution — Categorical distribution for discrete actions.

• MultiCategoricalDistribution — MultiCategorical distribution for multi discrete actions.

• BernoulliDistribution — Bernoulli distribution for MultiBinary action spaces.

• StateDependentNoiseDistribution — Distribution class for using generalized State Dependent Exploration (gSDE).

• TanhBijector — Bijective transformation of a probability distribution.

Functions

• sum_independent_dims — Continuous actions are usually considered to be independent,.

source class Distribution()

Bases : ABC

Abstract base class for distributions.

Methods

• proba_distribution_net — Create the layers and parameters that represent the distribution.

• proba_distribution — Set parameters of the distribution.

• log_prob — Returns the log likelihood.

• entropy — Returns Shannon's entropy of the probability.

• sample — Returns a sample from the probability distribution.

• mode — Returns the most likely action (deterministic output).

• get_actions — Return actions according to the probability distribution.

• actions_from_params — Returns samples from the probability distribution.

• log_prob_from_params — Returns samples and the associated log probabilities.

source method Distribution.proba_distribution_net(*args, **kwargs) → nn.Module | tuple[nn.Module, nn.Parameter]

Create the layers and parameters that represent the distribution.

Subclasses must define this, but the arguments and return type vary between concrete classes.

source method Distribution.proba_distribution(*args, **kwargs) → SelfDistribution

Set parameters of the distribution.

Returns

• SelfDistribution — self

source method Distribution.log_prob(x: torch.Tensor) → torch.Tensor

Returns the log likelihood.

Parameters

• x : torch.Tensor — the taken action

Returns

• torch.Tensor — The log likelihood of the distribution

source method Distribution.entropy() → torch.Tensor | None

Returns Shannon's entropy of the probability.

Returns

• torch.Tensor | None — the entropy, or None if no analytical form is known

source method Distribution.sample() → torch.Tensor

Returns a sample from the probability distribution.

Returns

• torch.Tensor — the stochastic action

source method Distribution.mode() → torch.Tensor

Returns the most likely action (deterministic output).

from the probability distribution.

Returns

• torch.Tensor — the stochastic action

source method Distribution.get_actions(deterministic: bool = False) → torch.Tensor

Return actions according to the probability distribution.

Parameters

• deterministic : bool — Whether to use deterministic actions.

source method Distribution.actions_from_params(*args, **kwargs) → torch.Tensor

Returns samples from the probability distribution.

given its parameters.

Returns

• torch.Tensor — actions

source method Distribution.log_prob_from_params(*args, **kwargs) → tuple[torch.Tensor, torch.Tensor]

Returns samples and the associated log probabilities.

from the probability distribution given its parameters.

Returns

• tuple[torch.Tensor, torch.Tensor] — actions and log prob

source sum_independent_dims(tensor: torch.Tensor) → torch.Tensor

Continuous actions are usually considered to be independent,.

so we can sum components of the log_prob or the entropy.

Parameters

• tensor : torch.Tensor — shape: (n_batch, n_actions) or (n_batch,)

Returns

• shape — (n_batch,) for (n_batch, n_actions) input, scalar for (n_batch,) input

source class DiagGaussianDistribution(action_dim: int)

Bases : Distribution

Gaussian distribution with diagonal covariance matrix, for continuous actions.

Parameters

• action_dim : int — Dimension of the action space.

Methods

• proba_distribution_net — Create the layers and parameter that represent the distribution.

• proba_distribution — Create the distribution given its parameters (mean, std).

• log_prob — Get the log probabilities of actions according to the distribution.

• entropy

• sample

• mode

• actions_from_params

• log_prob_from_params — Compute the log probability of taking an action.

source method DiagGaussianDistribution.proba_distribution_net(latent_dim: int, log_std_init: float = 0.0) → tuple[nn.Module, nn.Parameter]

Create the layers and parameter that represent the distribution.

one output will be the mean of the Gaussian, the other parameter will be the standard deviation (log std in fact to allow negative values).

Parameters

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

• log_std_init : float — Initial value for the log standard deviation

source method DiagGaussianDistribution.proba_distribution(mean_actions: torch.Tensor, log_std: torch.Tensor) → SelfDiagGaussianDistribution

Create the distribution given its parameters (mean, std).

Parameters

• mean_actions : torch.Tensor — Mean of the action distribution.

• log_std : torch.Tensor — Log standard deviation of the action distribution.

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

Get the log probabilities of actions according to the distribution.

Note that you must first call the proba_distribution() method.

Parameters

• actions : torch.Tensor — Actions taken.

source method DiagGaussianDistribution.entropy() → torch.Tensor | None

source method DiagGaussianDistribution.sample() → torch.Tensor

source method DiagGaussianDistribution.mode() → torch.Tensor

source method DiagGaussianDistribution.actions_from_params(mean_actions: torch.Tensor, log_std: torch.Tensor, deterministic: bool = False) → torch.Tensor

source method DiagGaussianDistribution.log_prob_from_params(mean_actions: torch.Tensor, log_std: torch.Tensor) → tuple[torch.Tensor, torch.Tensor]

Compute the log probability of taking an action.

given the distribution parameters.

Parameters

• mean_actions : torch.Tensor — Mean of the action distribution.

• log_std : torch.Tensor — Log standard deviation of the action distribution.

source class SquashedDiagGaussianDistribution(action_dim: int, epsilon: float = 1e-06)

Bases : DiagGaussianDistribution

Gaussian distribution with diagonal covariance matrix, followed by a squashing function (tanh) to ensure bounds.

Parameters

• action_dim : int — Dimension of the action space.

• epsilon : float — small value to avoid NaN due to numerical imprecision.

Methods

• proba_distribution

• log_prob

• entropy

• sample

• mode

• log_prob_from_params

source method SquashedDiagGaussianDistribution.proba_distribution(mean_actions: torch.Tensor, log_std: torch.Tensor) → SelfSquashedDiagGaussianDistribution

source method SquashedDiagGaussianDistribution.log_prob(actions: torch.Tensor, gaussian_actions: torch.Tensor | None = None) → torch.Tensor

source method SquashedDiagGaussianDistribution.entropy() → torch.Tensor | None

source method SquashedDiagGaussianDistribution.sample() → torch.Tensor

source method SquashedDiagGaussianDistribution.mode() → torch.Tensor

source method SquashedDiagGaussianDistribution.log_prob_from_params(mean_actions: torch.Tensor, log_std: torch.Tensor) → tuple[torch.Tensor, torch.Tensor]

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.

• proba_distribution

• log_prob

• entropy

• sample

• mode

• actions_from_params

• log_prob_from_params

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]

source class MultiCategoricalDistribution(action_dims: list[int])

Bases : Distribution

MultiCategorical distribution for multi discrete actions.

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

source method MultiCategoricalDistribution.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 MultiCategoricalDistribution.proba_distribution(action_logits: torch.Tensor) → SelfMultiCategoricalDistribution

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

source method MultiCategoricalDistribution.entropy() → torch.Tensor

source method MultiCategoricalDistribution.sample() → torch.Tensor

source method MultiCategoricalDistribution.mode() → torch.Tensor

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

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

source class BernoulliDistribution(action_dims: int)

Bases : Distribution

Bernoulli distribution for MultiBinary action spaces.

Parameters

• action_dims : int — Number of binary actions

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

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

Create the layer that represents the distribution.

it will be the logits of the Bernoulli distribution.

Parameters

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

source method BernoulliDistribution.proba_distribution(action_logits: torch.Tensor) → SelfBernoulliDistribution

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

source method BernoulliDistribution.entropy() → torch.Tensor

source method BernoulliDistribution.sample() → torch.Tensor

source method BernoulliDistribution.mode() → torch.Tensor

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

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

source class StateDependentNoiseDistribution(action_dim: int, full_std: bool = True, use_expln: bool = False, squash_output: bool = False, learn_features: bool = False, epsilon: float = 1e-06)

Bases : Distribution

Distribution class for using generalized State Dependent Exploration (gSDE).

Paper: https://arxiv.org/abs/2005.05719.

It is used to create the noise exploration matrix and compute the log probability of an action with that noise.

Parameters

• action_dim : int — Dimension of the action space.

• full_std : bool — Whether to use (n_features x n_actions) parameters for the std instead of only (n_features,)

• use_expln : bool — Use expln() function instead of exp() to ensure a positive standard deviation (cf paper). It allows to keep variance above zero and prevent it from growing too fast. In practice, exp() is usually enough.

• squash_output : bool — Whether to squash the output using a tanh function, this ensures bounds are satisfied.

• learn_features : bool — Whether to learn features for gSDE or not. This will enable gradients to be backpropagated through the features latent_sde in the code.

• epsilon : float — small value to avoid NaN due to numerical imprecision.

Methods

• get_std — Get the standard deviation from the learned parameter.

• sample_weights — Sample weights for the noise exploration matrix,.

• proba_distribution_net — Create the layers and parameter that represent the distribution.

• proba_distribution — Create the distribution given its parameters (mean, std).

• log_prob

• entropy

• sample

• mode

• get_noise

• actions_from_params

• log_prob_from_params

source method StateDependentNoiseDistribution.get_std(log_std: torch.Tensor) → torch.Tensor

Get the standard deviation from the learned parameter.

(log of it by default). This ensures that the std is positive.

Parameters

• log_std : torch.Tensor — Log standard deviation of the action distribution.

source method StateDependentNoiseDistribution.sample_weights(log_std: torch.Tensor, batch_size: int = 1) → None

Sample weights for the noise exploration matrix,.

using a centered Gaussian distribution.

Parameters

• log_std : torch.Tensor — Log standard deviation of the action distribution.

• batch_size : int — Number of elements to sample.

source method StateDependentNoiseDistribution.proba_distribution_net(latent_dim: int, log_std_init: float = -2.0, latent_sde_dim: int | None = None) → tuple[nn.Module, nn.Parameter]

Create the layers and parameter that represent the distribution.

one output will be the deterministic action, the other parameter will be the standard deviation of the distribution that control the weights of the noise matrix.

Parameters

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

• log_std_init : float — Initial value for the log standard deviation

• latent_sde_dim : int | None — Dimension of the last layer of the features extractor for gSDE. By default, it is shared with the policy network.

source method StateDependentNoiseDistribution.proba_distribution(mean_actions: torch.Tensor, log_std: torch.Tensor, latent_sde: torch.Tensor) → SelfStateDependentNoiseDistribution

Create the distribution given its parameters (mean, std).

Parameters

• mean_actions : torch.Tensor — Mean of the action distribution.

• log_std : torch.Tensor — Log standard deviation of the action distribution.

• latent_sde : torch.Tensor — Latent variable for the SDE exploration noise.

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

source method StateDependentNoiseDistribution.entropy() → torch.Tensor | None

source method StateDependentNoiseDistribution.sample() → torch.Tensor

source method StateDependentNoiseDistribution.mode() → torch.Tensor

source method StateDependentNoiseDistribution.get_noise(latent_sde: torch.Tensor) → torch.Tensor

source method StateDependentNoiseDistribution.actions_from_params(mean_actions: torch.Tensor, log_std: torch.Tensor, latent_sde: torch.Tensor, deterministic: bool = False) → torch.Tensor

source method StateDependentNoiseDistribution.log_prob_from_params(mean_actions: torch.Tensor, log_std: torch.Tensor, latent_sde: torch.Tensor) → tuple[torch.Tensor, torch.Tensor]

source class TanhBijector(epsilon: float = 1e-06)

Bijective transformation of a probability distribution.

using a squashing function (tanh).

Parameters

• epsilon : float — small value to avoid NaN due to numerical imprecision.

Methods

• forward

• atanh — Inverse of Tanh.

• inverse — Inverse tanh.

• log_prob_correction

source staticmethod TanhBijector.forward(x: torch.Tensor) → torch.Tensor

source staticmethod TanhBijector.atanh(x: torch.Tensor) → torch.Tensor

Inverse of Tanh.

Taken from Pyro: https://github.com/pyro-ppl/pyro 0.5 * torch.log((1 + x ) / (1 - x))

source staticmethod TanhBijector.inverse(y: torch.Tensor) → torch.Tensor

Inverse tanh.

Parameters

• y : torch.Tensor — The output value.

source method TanhBijector.log_prob_correction(x: torch.Tensor) → torch.Tensor