180 lines
7.1 KiB
Python
180 lines
7.1 KiB
Python
"""Implementation of a space that represents finite-length sequences."""
|
|
from __future__ import annotations
|
|
|
|
import typing
|
|
from typing import Any, Union
|
|
|
|
import numpy as np
|
|
from numpy.typing import NDArray
|
|
|
|
import gymnasium as gym
|
|
from gymnasium.spaces.space import Space
|
|
|
|
|
|
class Sequence(Space[Union[typing.Tuple[Any, ...], Any]]):
|
|
r"""This space represent sets of finite-length sequences.
|
|
|
|
This space represents the set of tuples of the form :math:`(a_0, \dots, a_n)` where the :math:`a_i` belong
|
|
to some space that is specified during initialization and the integer :math:`n` is not fixed
|
|
|
|
Example:
|
|
>>> from gymnasium.spaces import Sequence, Box
|
|
>>> observation_space = Sequence(Box(0, 1), seed=2)
|
|
>>> observation_space.sample()
|
|
(array([0.26161215], dtype=float32),)
|
|
>>> observation_space = Sequence(Box(0, 1), seed=0)
|
|
>>> observation_space.sample()
|
|
(array([0.6369617], dtype=float32), array([0.26978672], dtype=float32), array([0.04097353], dtype=float32))
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
space: Space[Any],
|
|
seed: int | np.random.Generator | None = None,
|
|
stack: bool = False,
|
|
):
|
|
"""Constructor of the :class:`Sequence` space.
|
|
|
|
Args:
|
|
space: Elements in the sequences this space represent must belong to this space.
|
|
seed: Optionally, you can use this argument to seed the RNG that is used to sample from the space.
|
|
stack: If `True` then the resulting samples would be stacked.
|
|
"""
|
|
assert isinstance(
|
|
space, Space
|
|
), f"Expects the feature space to be instance of a gym Space, actual type: {type(space)}"
|
|
self.feature_space = space
|
|
self.stack = stack
|
|
if self.stack:
|
|
self.stacked_feature_space: Space = gym.vector.utils.batch_space(
|
|
self.feature_space, 1
|
|
)
|
|
|
|
# None for shape and dtype, since it'll require special handling
|
|
super().__init__(None, None, seed)
|
|
|
|
def seed(self, seed: int | None = None) -> list[int]:
|
|
"""Seed the PRNG of this space and the feature space."""
|
|
seeds = super().seed(seed)
|
|
seeds += self.feature_space.seed(seed)
|
|
return seeds
|
|
|
|
@property
|
|
def is_np_flattenable(self):
|
|
"""Checks whether this space can be flattened to a :class:`spaces.Box`."""
|
|
return False
|
|
|
|
def sample(
|
|
self,
|
|
mask: None
|
|
| (
|
|
tuple[
|
|
None | np.integer | NDArray[np.integer],
|
|
Any,
|
|
]
|
|
) = None,
|
|
) -> tuple[Any] | Any:
|
|
"""Generates a single random sample from this space.
|
|
|
|
Args:
|
|
mask: An optional mask for (optionally) the length of the sequence and (optionally) the values in the sequence.
|
|
If you specify `mask`, it is expected to be a tuple of the form `(length_mask, sample_mask)` where `length_mask`
|
|
is
|
|
|
|
* ``None`` The length will be randomly drawn from a geometric distribution
|
|
* ``np.ndarray`` of integers, in which case the length of the sampled sequence is randomly drawn from this array.
|
|
* ``int`` for a fixed length sample
|
|
|
|
The second element of the mask tuple `sample` mask specifies a mask that is applied when
|
|
sampling elements from the base space. The mask is applied for each feature space sample.
|
|
|
|
Returns:
|
|
A tuple of random length with random samples of elements from the :attr:`feature_space`.
|
|
"""
|
|
if mask is not None:
|
|
length_mask, feature_mask = mask
|
|
else:
|
|
length_mask, feature_mask = None, None
|
|
|
|
if length_mask is not None:
|
|
if np.issubdtype(type(length_mask), np.integer):
|
|
assert (
|
|
0 <= length_mask
|
|
), f"Expects the length mask to be greater than or equal to zero, actual value: {length_mask}"
|
|
length = length_mask
|
|
elif isinstance(length_mask, np.ndarray):
|
|
assert (
|
|
len(length_mask.shape) == 1
|
|
), f"Expects the shape of the length mask to be 1-dimensional, actual shape: {length_mask.shape}"
|
|
assert np.all(
|
|
0 <= length_mask
|
|
), f"Expects all values in the length_mask to be greater than or equal to zero, actual values: {length_mask}"
|
|
assert np.issubdtype(
|
|
length_mask.dtype, np.integer
|
|
), f"Expects the length mask array to have dtype to be an numpy integer, actual type: {length_mask.dtype}"
|
|
length = self.np_random.choice(length_mask)
|
|
else:
|
|
raise TypeError(
|
|
f"Expects the type of length_mask to an integer or a np.ndarray, actual type: {type(length_mask)}"
|
|
)
|
|
else:
|
|
# The choice of 0.25 is arbitrary
|
|
length = self.np_random.geometric(0.25)
|
|
|
|
# Generate sample values from feature_space.
|
|
sampled_values = tuple(
|
|
self.feature_space.sample(mask=feature_mask) for _ in range(length)
|
|
)
|
|
|
|
if self.stack:
|
|
# Concatenate values if stacked.
|
|
out = gym.vector.utils.create_empty_array(
|
|
self.feature_space, len(sampled_values)
|
|
)
|
|
return gym.vector.utils.concatenate(self.feature_space, sampled_values, out)
|
|
|
|
return sampled_values
|
|
|
|
def contains(self, x: Any) -> bool:
|
|
"""Return boolean specifying if x is a valid member of this space."""
|
|
# by definition, any sequence is an iterable
|
|
if self.stack:
|
|
return all(
|
|
item in self.feature_space
|
|
for item in gym.vector.utils.iterate(self.stacked_feature_space, x)
|
|
)
|
|
else:
|
|
return isinstance(x, tuple) and all(
|
|
self.feature_space.contains(item) for item in x
|
|
)
|
|
|
|
def __repr__(self) -> str:
|
|
"""Gives a string representation of this space."""
|
|
return f"Sequence({self.feature_space}, stack={self.stack})"
|
|
|
|
def to_jsonable(
|
|
self, sample_n: typing.Sequence[tuple[Any, ...] | Any]
|
|
) -> list[list[Any]]:
|
|
"""Convert a batch of samples from this space to a JSONable data type."""
|
|
if self.stack:
|
|
return self.stacked_feature_space.to_jsonable(sample_n)
|
|
else:
|
|
return [self.feature_space.to_jsonable(sample) for sample in sample_n]
|
|
|
|
def from_jsonable(self, sample_n: list[list[Any]]) -> list[tuple[Any, ...] | Any]:
|
|
"""Convert a JSONable data type to a batch of samples from this space."""
|
|
if self.stack:
|
|
return self.stacked_feature_space.from_jsonable(sample_n)
|
|
else:
|
|
return [
|
|
tuple(self.feature_space.from_jsonable(sample)) for sample in sample_n
|
|
]
|
|
|
|
def __eq__(self, other: Any) -> bool:
|
|
"""Check whether ``other`` is equivalent to this instance."""
|
|
return (
|
|
isinstance(other, Sequence)
|
|
and self.feature_space == other.feature_space
|
|
and self.stack == other.stack
|
|
)
|