Spaces

Space Semantics

Envelope has two basic spaces: Discrete and Continuous. Both of them may have any shape, where Discrete naturally expresses a multi-discrete space if it's shape is any bigger than (). For convenience, both basic spaces have a from_shape method that creates a space with a full array of the given shape.

The two basic spaces can be nested in any PyTree, composing them into a PyTreeSpace. The PyTreeSpaces methods and properties are mapped onto the PyTree of subspaces, treating them as the leaves. For example

space = PyTreeSpace({
  "foo": Discrete([3, 5]),
  "bar": Continuous(low=-1.0, high=1.0)
})
space.shape  # {'foo': (2,), 'bar': ()}
space.dtype  # {'foo': dtype('int32'), 'bar': dtype('float32')}

Spaces can be batched using BatchedSpace, which returns a view that prepends a batch dimension and vectorizes sample and contains.

The constraints on the PyTreeSpace's members imply a strict ordering on the construction of spaces:

Discrete / Continuous → PyTreeSpace → BatchedSpace

API Reference

envelope.spaces.Space

Bases: ABC, FrozenPyTreeNode

Base class for all spaces. Spaces are immutable and hashable.

Attributes

dtype abstractmethod property

The dtype of the space.

shape abstractmethod property

The shape of the space.

Functions

contains(x) abstractmethod

Check if self contains a sample x.

sample(key) abstractmethod

Sample a random element from the space.

envelope.spaces.Discrete

Bases: Space

A discrete space with a given number of elements. n can be a scalar or an array. The shape and dtype of the space are inferred from n.

Parameters:

Name	Type	Description	Default
n	int | Array	The number of elements in the space.	required

Functions

from_shape(n, shape) classmethod

Create a Discrete space from a shape and a number of elements. This is a shorthand for Discrete with n being expanded to the shape.

envelope.spaces.Continuous

Bases: Space

A continuous space with a given lower and upper bound. low and high can be scalars or arrays. The shape and dtype of the space are inferred from low and high.

Parameters:

Name	Type	Description	Default
low	float | Array	The lower bound of the space.	required
high	float | Array	The upper bound of the space.	required

Functions

from_shape(low, high, shape) classmethod

Create a Continuous space from a shape and a lower and upper bound. This is a shorthand for Continuous with low and high being expanded to the shape.

envelope.spaces.PyTreeSpace

Bases: Space

A Space defined by a PyTree structure of other Spaces. While PyTreeSpace.tree might be an aribrarily nested PyTree, it's leaves must only be Discrete or Continuous, and not contain PyTreeSpace or BatchedSpace as leaves. The shape and dtype of the PyTreeSpace are PyTrees of the same structure, containing the shape and dtype of the leaves.

Parameters:

Name	Type	Description	Default
tree	PyTree	A PyTree with Discrete or Continuous leaves.	required

Attributes

dtype property

A PyTree of the same structure as tree, containing each leaf's dtype.

shape property

A PyTree of the same structure as tree, containing each leaf's shape.

envelope.spaces.BatchedSpace

Bases: Space

A view that adds a leading batch dimension to a base Space without materializing or broadcasting its parameters.

Parameters:

Name	Type	Description	Default
space	Space	The underlying base space.	required
batch_size	int	The leading batch dimension.	required

Attributes

dtype property

The dtype of the base space (batch dimensions don't affect dtype).

shape cached property

The shape of the BatchedSpace is the leading batch dimension prepend the shape of the wrapped Space. If the wrapped Space is a PyTreeSpace, the shape is a PyTree of the same structure, with the leading batch dimension prepended to the shape of each leaf Space.

Functions

contains(x)

BatchedSpace.contains checks if each entry of x along the leading dimension is contained in the base (unbatched) Space.

sample(key)

Sample a batch of samples from the wrapped Space. You may pass a single key or a batch of keys shaped (batch_size, ...).