Discrete Space

The DiscreteDist class extends the Gymnasium Discrete space to create categorical distributions.

Overview

DiscreteDist allows you to create a categorical probability distribution from a discrete action space. It is useful for environments with a fixed number of distinct actions.

API Reference

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

Bases: Discrete

__call__(prob: Tensor, mask: Tensor = None) → MaskedCategorical[source]

Compute and return a masked categorical distribution based on the given probability tensor and an optional mask. The distribution incorporates specific probabilities and constraints defined by the provided input.

Parameters:

prob – A tensor representing the probabilities for each category.
mask – A tensor specifying a mask to limit the valid categories. Defaults to a tensor of ones if not provided.

Returns:

A MaskedCategorical distribution constructed with given probabilities, mask, and starting values.

Usage Examples

Basic usage:

import torch as th
from prob_spaces.discrete import DiscreteDist

# Create a discrete space with 5 possible actions
space = DiscreteDist(n=5)

# Create logits/probabilities for each action
probs = th.tensor([0.1, 0.2, 0.3, 0.3, 0.1])

# Create a distribution
dist = space(probs)

# Sample an action
action = dist.sample()

# Compute log probability
log_prob = dist.log_prob(action)

Converting from gymnasium:

import gymnasium as gym
from prob_spaces.discrete import DiscreteDist

# Create a gymnasium space
gym_space = gym.spaces.Discrete(n=10)

# Convert to a DiscreteDist
space_dist = DiscreteDist.from_space(gym_space)