prob-spaces: Probability Distributions from Gymnasium Spaces

prob-spaces is a Python package that allows you to create probability distributions from Gymnasium spaces. It provides a simple and intuitive interface for working with various probability spaces in reinforcement learning environments.

Key Features:

• Create probability distributions directly from Gymnasium spaces
• Support for common space types: Discrete, MultiDiscrete, Box, and Dict
• Seamless integration with PyTorch for sampling and computing log probabilities
• Support for masking operations to constrain valid actions

Installation

From PyPI

To install prob-spaces from PyPI:

pip install prob-spaces

GPU Support

prob-spaces uses PyTorch, which can be installed with CUDA support for GPU acceleration. The package configuration includes a PyTorch CUDA 12.4 index. To use a different CUDA version, you may need to modify the PyTorch installation separately.

Example Usage

Here's a simple example of how to use prob-spaces:

import gymnasium as gym
import torch as th
from prob_spaces.converter import convert_to_prob_space

# Create a Gymnasium space
action_space = gym.spaces.Discrete(5)

# Convert to a probability space
prob_space = convert_to_prob_space(action_space)

# Create a probability distribution
probs = th.ones(5)  # Uniform distribution
dist = prob_space(probs)

# Sample from the distribution
action = dist.sample()

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

Documentation

Documentation is available online and provides detailed information on how to use the package, including examples and API references.