Dict Space

The DictDist class extends the Gymnasium Dict space to create nested distributions.

Overview

DictDist allows you to create probability distributions for complex action spaces that consist of multiple sub-spaces organized in a dictionary structure.

API Reference

Bases: Dict

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

Create a dict of distributions based on input probabilities.

Parameters:

prob – Dictionary of probability tensors for each space
mask – Optional dictionary of masks for each space

Returns:

Dictionary of distribution objects

Usage Examples

Basic usage:

import torch as th
import gymnasium as gym
from prob_spaces.dict import DictDist
from prob_spaces.discrete import DiscreteDist
from prob_spaces.box import BoxDist

# Create a dictionary space with a discrete and a box space
space = DictDist({
    'discrete': DiscreteDist(n=3),
    'continuous': BoxDist(low=-1, high=1, shape=(2,))
})

# Create probabilities for each space
probs = {
    'discrete': th.tensor([0.3, 0.4, 0.3]),
    'continuous': (th.zeros(2), th.ones(2))  # (loc, scale) for box space
}

# Create distributions
dists = space(probs)

# Sample from each distribution
samples = {k: dist.sample() for k, dist in dists.items()}

# Compute log probabilities
log_probs = {k: dist.log_prob(samples[k]) for k, dist in dists.items()}