Spaces

Utilities to operate on Gymnasium spaces.



Overview

The utilities described in this section supports the following Gymnasium spaces:

Type

Supported spaces

Fundamental

Box, Discrete, and MultiDiscrete

Composite

Dict and Tuple

The following table provides a snapshot of the space sample conversion functions:

Input

Function

Output

Space (NumPy / int)

tensorize_space()

Space (PyTorch / JAX)

Space (PyTorch / JAX)

untensorize_space()

Space (NumPy / int)

Space (PyTorch / JAX)

flatten_tensorized_space()

PyTorch tensor / JAX array

PyTorch tensor / JAX array

unflatten_tensorized_space()

Space (PyTorch / JAX)


API (PyTorch)

skrl.utils.spaces.torch.compute_space_size(space: gymnasium.spaces.Space | Sequence[int] | int, occupied_size: bool = False) int

Get the size (number of elements) of a space.

Parameters:
  • space – Gymnasium space.

  • occupied_size – Whether the number of elements occupied by the space is returned (default: False). It only affects Discrete (occupied space is 1), and MultiDiscrete (occupied space is the number of discrete spaces).

Returns:

Size of the space (number of elements).

skrl.utils.spaces.torch.convert_gym_space(space: gym.Space, squeeze_batch_dimension: bool = False) gymnasium.Space

Converts a gym space to a gymnasium space.

Parameters:
  • space – Gym space to convert to.

  • squeeze_batch_dimension – Whether to remove fundamental spaces’ first dimension. It currently affects Box space only.

Raises:

ValueError – The given space is not supported.

Returns:

Converted space.

skrl.utils.spaces.torch.flatten_tensorized_space(x: Any) torch.Tensor

Flatten a tensorized space.

Parameters:

x – Tensorized space sample/value.

Raises:

ValueError – The given sample/value type is not supported.

Returns:

A tensor. The returned tensor will have shape (batch, space size).

skrl.utils.spaces.torch.sample_space(space: gymnasium.spaces.Space, batch_size: int = 1, backend: str = typing.Literal['numpy', 'torch'], device=None) Any

Generates a random sample from the specified space.

Parameters:
  • space – Gymnasium space.

  • batch_size – Size of the sampled batch (default: 1).

  • backend – Whether backend will be used to construct the fundamental spaces (default: "numpy").

  • device – Device on which a tensor/array is or will be allocated (default: None). This parameter is used when the backend is "torch".

Raises:

ValueError – The given space or backend is not supported.

Returns:

Sample of the space

skrl.utils.spaces.torch.tensorize_space(space: gymnasium.spaces.Space, x: Any, device: str | torch.device | None = None) Any

Convert the sample/value items of a given gymnasium space to PyTorch tensors.

Parameters:
  • space – Gymnasium space.

  • x – Sample/value of the given space to tensorize to.

  • device – Device on which a tensor/array is or will be allocated (default: None). This parameter is used when the space value is not a PyTorch tensor (e.g.: NumPy array, number).

Raises:

ValueError – The given space or the sample/value type is not supported.

Returns:

Sample/value space with items converted to tensors.

skrl.utils.spaces.torch.unflatten_tensorized_space(space: gymnasium.spaces.Space | Sequence[int] | int, x: torch.Tensor) Any

Unflatten a tensor to create a tensorized space.

Parameters:
  • space – Gymnasium space.

  • x – A tensor with shape (batch, space size).

Raises:

ValueError – The given space is not supported.

Returns:

Tensorized space value.

skrl.utils.spaces.torch.untensorize_space(space: gymnasium.spaces.Space, x: Any, squeeze_batch_dimension: bool = True) Any

Convert a tensorized space to a gymnasium space with expected sample/value item types.

Parameters:
  • space – Gymnasium space.

  • x – Tensorized space (Sample/value space where items are tensors).

  • squeeze_batch_dimension – Whether to remove the batch dimension. If True, only the sample/value with a batch dimension of size 1 will be affected

Raises:

ValueError – The given space or the sample/value type is not supported.

Returns:

Sample/value space with expected item types.


API (JAX)

skrl.utils.spaces.jax.compute_space_size(space: gymnasium.spaces.Space | Sequence[int] | int, occupied_size: bool = False) int

Get the size (number of elements) of a space.

Parameters:
  • space – Gymnasium space.

  • occupied_size – Whether the number of elements occupied by the space is returned (default: False). It only affects Discrete (occupied space is 1), and MultiDiscrete (occupied space is the number of discrete spaces).

Returns:

Size of the space (number of elements).

skrl.utils.spaces.jax.convert_gym_space(space: gym.Space, squeeze_batch_dimension: bool = False) gymnasium.Space

Converts a gym space to a gymnasium space.

Parameters:
  • space – Gym space to convert to.

  • squeeze_batch_dimension – Whether to remove fundamental spaces’ first dimension. It currently affects Box space only.

Raises:

ValueError – The given space is not supported.

Returns:

Converted space.

skrl.utils.spaces.jax.flatten_tensorized_space(x: Any, _jax: bool = True) jax.Array

Flatten a tensorized space.

Parameters:
  • x – Tensorized space sample/value.

  • _jax – Whether the space should be handled using JAX operations. It only affects composite spaces.

Raises:

ValueError – The given sample/value type is not supported.

Returns:

A tensor. The returned tensor will have shape (batch, space size).

skrl.utils.spaces.jax.sample_space(space: gymnasium.spaces.Space, batch_size: int = 1, backend: str = typing.Literal['numpy', 'jax'], device=None) Any

Generates a random sample from the specified space.

Parameters:
  • space – Gymnasium space.

  • batch_size – Size of the sampled batch (default: 1).

  • backend – Whether backend will be used to construct the fundamental spaces (default: "numpy").

  • device – Device on which a tensor/array is or will be allocated (default: None). This parameter is used when the backend is "jax".

Raises:

ValueError – The given space or backend is not supported.

Returns:

Sample of the space

skrl.utils.spaces.jax.tensorize_space(space: gymnasium.spaces.Space, x: Any, device: str | jax.Device | None = None, _jax: bool = True) Any

Convert the sample/value items of a given gymnasium space to JAX array.

Parameters:
  • space – Gymnasium space.

  • x – Sample/value of the given space to tensorize to.

  • device – Device on which a tensor/array is or will be allocated (default: None). This parameter is used when the space value is not a JAX array (e.g.: NumPy array, number).

  • _jax – Whether the converted value should be a JAX array. It only affects NumPy space values.

Raises:

ValueError – The given space or the sample/value type is not supported.

Returns:

Sample/value space with items converted to tensors.

skrl.utils.spaces.jax.unflatten_tensorized_space(space: gymnasium.spaces.Space | Sequence[int] | int, x: jax.Array) Any

Unflatten a tensor to create a tensorized space.

Parameters:
  • space – Gymnasium space.

  • x – A tensor with shape (batch, space size).

Raises:

ValueError – The given space is not supported.

Returns:

Tensorized space value.

skrl.utils.spaces.jax.untensorize_space(space: gymnasium.spaces.Space, x: Any, squeeze_batch_dimension: bool = True) Any

Convert a tensorized space to a gymnasium space with expected sample/value item types.

Parameters:
  • space – Gymnasium space.

  • x – Tensorized space (Sample/value space where items are tensors).

  • squeeze_batch_dimension – Whether to remove the batch dimension. If True, only the sample/value with a batch dimension of size 1 will be affected

Raises:

ValueError – The given space or the sample/value type is not supported.

Returns:

Sample/value space with expected item types.