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

Probability distribution for Box spaces.

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

Generate a transformed probability distribution.

Construct a base distribution, apply a sequence of transformations, and return the resulting transformed distribution.

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.

Returns:

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

Return type:

th.distributions.Distribution

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)

Mathematical Details

The BoxDist distribution is constructed by transforming a base Normal distribution on \(\mathbb{R}\) to the bounded Box interval \([\mathrm{low},\ \mathrm{high}]\) using three steps:

  1. Normal Distribution (Base): The base distribution is a Normal (Gaussian) distribution on \(\mathbb{R}\). The user must provide the parameters: - loc: the mean of the Normal distribution - scale: the standard deviation of the Normal distribution

    Note

    Here, \(z\) is sampled from \(\mathcal{N}(\text{loc},\ \text{scale})\), i.e., \(z \sim \mathcal{N}(\text{loc},\ \text{scale})\).

  2. Sigmoid Transform: Maps \(z \in \mathbb{R}\) to \((0, 1)\) via the sigmoid function:

    \[x = \sigma(z) = \frac{1}{1 + e^{-z}}\]

  3. Affine Transform: Maps \(x \in (0, 1)\) to \([\mathrm{low},\ \mathrm{high}]\):

    \[y = \mathrm{low} + (\mathrm{high} - \mathrm{low}) \cdot x\]

So, a sample \(z\) from the base distribution is transformed as:

\[y = \mathrm{low} + (\mathrm{high} - \mathrm{low}) \cdot \sigma(z)\]

The probability density is adjusted using the change-of-variables formula, ensuring the resulting distribution is properly normalized over the Box bounds.

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)