Spaces¶
Utilities to operate on Gymnasium spaces.
Overview¶
The utilities described in this section supports the following Gymnasium spaces:
Type |
Supported spaces |
---|---|
Fundamental |
|
Composite |
The following table provides a snapshot of the space sample conversion functions:
Input |
Function |
Output |
---|---|---|
Space (NumPy / int) |
Space (PyTorch / JAX) |
|
Space (PyTorch / JAX) |
Space (NumPy / int) |
|
Space (PyTorch / JAX) |
PyTorch tensor / JAX array |
|
PyTorch tensor / JAX array |
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 affectsDiscrete
(occupied space is 1), andMultiDiscrete
(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 affectsDiscrete
(occupied space is 1), andMultiDiscrete
(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.