API Reference

This section provides detailed information about the prob-spaces API.

Probability Spaces

Probability distribution classes for various Gymnasium spaces.

class prob_spaces.DictDist(spaces: None | dict[str, Space] | Sequence[tuple[str, Space]] = None, seed: dict | int | Generator | None = None, **spaces_kwargs: Space)[source]: Probability distribution for Dict spaces.

class prob_spaces.DiscreteDist(n: int | integer[Any], seed: int | Generator | None = None, start: int | integer[Any] = 0)[source]

Probability distribution for Discrete spaces.

classmethod from_space(space: Discrete) → DiscreteDist[source]

Create a DiscreteDist from a gymnasium Discrete space.

Returns:: An instance of DiscreteDist created from the given gymnasium Discrete space.
Return type:: DiscreteDist

class prob_spaces.MultiDiscreteDist(nvec: ~numpy.ndarray[tuple[int, ...], ~numpy.dtype[~numpy.integer[~typing.Any]]] | list[int], dtype: str | type[~numpy.integer[~typing.Any]] = <class 'numpy.int64'>, seed: int | ~numpy.random._generator.Generator | None = None, start: ~numpy.ndarray[tuple[int, ...], ~numpy.dtype[~numpy.integer[~typing.Any]]] | list[int] | None = None)[source]

Probability distribution for MultiDiscrete spaces.

classmethod from_space(space: MultiDiscrete) → MultiDiscreteDist[source]

Create a MultiDiscreteDist from a gymnasium MultiDiscrete space.

Returns:: An instance of MultiDiscreteDist created from the given gymnasium MultiDiscrete space.
Return type:: MultiDiscreteDist

property prob_last_dim: int: Return the last dimension size for probability tensors.

Convert an action space into its corresponding probability distribution space.

This function supports different types of action spaces and creates an appropriate distribution space for each one. Supported action spaces include MultiDiscrete, Discrete, Box, and Dict. For Dict action spaces, the function recursively converts each subspace into its probability distribution space.

Parameters:: action_space (Spaces) – The input action space to be converted. This can be an instance of gym.spaces.MultiDiscrete, gym.spaces.Discrete, gym.spaces.Box, or gym.spaces.Dict.
Raises:: NotImplementedError – If the input action space type is not supported.
Returns:: The corresponding probability distribution space created based on the input action space type.
Returns:: The corresponding probability distribution space created based on the input action space type.
Return type:: DistSpaces

Converter

Module for converting gymnasium action spaces to probability distribution spaces.

Convert an action space into its corresponding probability distribution space.

This function supports different types of action spaces and creates an appropriate distribution space for each one. Supported action spaces include MultiDiscrete, Discrete, Box, and Dict. For Dict action spaces, the function recursively converts each subspace into its probability distribution space.

Parameters:: action_space (Spaces) – The input action space to be converted. This can be an instance of gym.spaces.MultiDiscrete, gym.spaces.Discrete, gym.spaces.Box, or gym.spaces.Dict.
Raises:: NotImplementedError – If the input action space type is not supported.
Returns:: The corresponding probability distribution space created based on the input action space type.
Returns:: The corresponding probability distribution space created based on the input action space type.
Return type:: DistSpaces

Distributions

Probability distribution package for prob_spaces.dists.

Module for categorical probability distributions with masking support.

class prob_spaces.dists.categorical.CategoricalDist(logits: Tensor | None = None, probs: Tensor | None = None, *, mask: Tensor = None, indices: Tensor = None, neg_inf: float = -inf, padding_value: int | None = None, start: int | integer[Any] | ndarray[tuple[int, ...], dtype[integer[Any]]] | list[int] = 0)[source]

Categorical distribution with masking and start offset support.

log_prob(value: Tensor) → Tensor[source]

Compute the log probability of a value, accounting for start offset.

Returns:: The log probability tensor, accounting for the start offset.
Return type:: torch.Tensor

sample(sample_shape: Size | Sequence[int] | None = None) → Tensor[source]

Sample from the categorical distribution with start offset.

Returns:: A tensor containing samples from the categorical distribution, adjusted by the start offset.
Return type:: th.Tensor

property th_start: Tensor: Return the start offset as a tensor.