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]
class prob_spaces.DiscreteDist(n: int | integer[Any], seed: int | Generator | None = None, start: int | integer[Any] = 0)[source]
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]
classmethod from_space(space: MultiDiscrete) MultiDiscreteDist[source]

Convert a gymnasium space to a MultiDiscreteDist.

prob_spaces.convert_to_prob_space(action_space: Box | Discrete | MultiDiscrete) BoxDist | DiscreteDist | MultiDiscreteDist | None[source]

Converts 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.

Converter

prob_spaces.converter.convert_to_prob_space(action_space: Box | Discrete | MultiDiscrete) BoxDist | DiscreteDist | MultiDiscreteDist | None[source]

Converts 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.

Distributions

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]
log_prob(value: Tensor) Tensor[source]

Returns the log of the probability density/mass function evaluated at value.

Parameters:

value (Tensor)

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

Generates a sample_shape shaped reparameterized sample or sample_shape shaped batch of reparameterized samples if the distribution parameters are batched.

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

Generates a sample_shape shaped sample or sample_shape shaped batch of samples if the distribution parameters are batched.