Box Space

The BoxDist class extends the Gymnasium Box space to create continuous probability distributions.

Overview

BoxDist allows you to create continuous probability distributions for action spaces with bounded values. By default, it uses a transformed Normal distribution to ensure samples are within the specified bounds.

API Reference

class prob_spaces.box.BoxDist(low: ~typing.SupportsFloat | ~numpy.ndarray[tuple[int, ...], ~numpy.dtype[~typing.Any]], high: ~typing.SupportsFloat | ~numpy.ndarray[tuple[int, ...], ~numpy.dtype[~typing.Any]], shape: ~typing.Sequence[int] | None = None, dtype: type[~numpy.floating[~typing.Any]] | type[~numpy.integer[~typing.Any]] = <class 'numpy.float32'>, seed: int | ~numpy.random._generator.Generator | None = None, dist: None | ~typing.Type[~torch.distributions.distribution.Distribution] = None)[source]

Bases: Box

__call__(loc: Tensor, scale: Tensor) → Distribution[source]

Generates a transformed probability distribution based on the input location and scale parameters. The method constructs a base distribution, applies a sequence of transformations to it, and returns the resulting transformed distribution. This allows for creating flexible and expressive probability distributions.

Parameters:

  • loc – A tensor specifying the location parameters for the base distribution.

  • scale – A tensor specifying the scale parameters for the base distribution.

Returns:

A transformed distribution object derived from the specified base distribution and transformations.

Key Features

  • Support for arbitrary continuous ranges with lower and upper bounds

  • Built-in transformation to enforce bounds using sigmoid and affine transforms

  • Customizable base distribution (defaults to Normal)

Usage Examples

Basic usage:

import numpy as np
import torch as th
from prob_spaces.box import BoxDist

# Create a 2D box space with values between -1 and 1
space = BoxDist(low=-1, high=1, shape=(2,))

# Create parameters for the distribution
loc = th.zeros(2)    # Mean
scale = th.ones(2)   # Standard deviation

# Create a distribution
dist = space(loc, scale)

# 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.box import BoxDist

# Create a gymnasium space
gym_space = gym.spaces.Box(low=-1, high=1, shape=(3,))

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