mirror of
https://github.com/Farama-Foundation/Gymnasium.git
synced 2025-08-02 06:16:32 +00:00
137 lines
3.7 KiB
Python
137 lines
3.7 KiB
Python
from itertools import zip_longest
|
|
from typing import Optional
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import gymnasium
|
|
from gymnasium.spaces import Box, Graph, utils
|
|
from gymnasium.utils.env_checker import data_equivalence
|
|
from tests.spaces.utils import TESTING_SPACES, TESTING_SPACES_IDS
|
|
|
|
TESTING_SPACES_EXPECTED_FLATDIMS = [
|
|
# Discrete
|
|
3,
|
|
3,
|
|
# Box
|
|
1,
|
|
4,
|
|
2,
|
|
2,
|
|
2,
|
|
# Multi-discrete
|
|
4,
|
|
10,
|
|
# Multi-binary
|
|
8,
|
|
6,
|
|
# Text
|
|
6,
|
|
6,
|
|
6,
|
|
# Tuple
|
|
9,
|
|
7,
|
|
10,
|
|
6,
|
|
None,
|
|
# Dict
|
|
7,
|
|
8,
|
|
17,
|
|
None,
|
|
# Graph
|
|
None,
|
|
None,
|
|
None,
|
|
# Sequence
|
|
None,
|
|
None,
|
|
None,
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
["space", "flatdim"],
|
|
zip_longest(TESTING_SPACES, TESTING_SPACES_EXPECTED_FLATDIMS),
|
|
ids=TESTING_SPACES_IDS,
|
|
)
|
|
def test_flatdim(space: gymnasium.spaces.Space, flatdim: Optional[int]):
|
|
"""Checks that the flattened dims of the space is equal to an expected value."""
|
|
if space.is_np_flattenable:
|
|
dim = utils.flatdim(space)
|
|
assert dim == flatdim, f"Expected {dim} to equal {flatdim}"
|
|
else:
|
|
with pytest.raises(
|
|
ValueError,
|
|
):
|
|
utils.flatdim(space)
|
|
|
|
|
|
@pytest.mark.parametrize("space", TESTING_SPACES, ids=TESTING_SPACES_IDS)
|
|
def test_flatten_space(space):
|
|
"""Test that the flattened spaces are a box and have the `flatdim` shape."""
|
|
flat_space = utils.flatten_space(space)
|
|
|
|
if space.is_np_flattenable:
|
|
assert isinstance(flat_space, Box)
|
|
(single_dim,) = flat_space.shape
|
|
flatdim = utils.flatdim(space)
|
|
|
|
assert single_dim == flatdim
|
|
elif isinstance(flat_space, Graph):
|
|
assert isinstance(space, Graph)
|
|
|
|
(node_single_dim,) = flat_space.node_space.shape
|
|
node_flatdim = utils.flatdim(space.node_space)
|
|
assert node_single_dim == node_flatdim
|
|
|
|
if flat_space.edge_space is not None:
|
|
(edge_single_dim,) = flat_space.edge_space.shape
|
|
edge_flatdim = utils.flatdim(space.edge_space)
|
|
assert edge_single_dim == edge_flatdim
|
|
else:
|
|
assert isinstance(
|
|
space,
|
|
(gymnasium.spaces.Tuple, gymnasium.spaces.Dict, gymnasium.spaces.Sequence),
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("space", TESTING_SPACES, ids=TESTING_SPACES_IDS)
|
|
def test_flatten(space):
|
|
"""Test that a flattened sample have the `flatdim` shape."""
|
|
flattened_sample = utils.flatten(space, space.sample())
|
|
|
|
if space.is_np_flattenable:
|
|
assert isinstance(flattened_sample, np.ndarray)
|
|
(single_dim,) = flattened_sample.shape
|
|
flatdim = utils.flatdim(space)
|
|
|
|
assert single_dim == flatdim
|
|
else:
|
|
assert isinstance(flattened_sample, (tuple, dict, Graph))
|
|
|
|
|
|
@pytest.mark.parametrize("space", TESTING_SPACES, ids=TESTING_SPACES_IDS)
|
|
def test_flat_space_contains_flat_points(space):
|
|
"""Test that the flattened samples are contained within the flattened space."""
|
|
flattened_samples = [utils.flatten(space, space.sample()) for _ in range(10)]
|
|
flat_space = utils.flatten_space(space)
|
|
|
|
for flat_sample in flattened_samples:
|
|
assert flat_sample in flat_space
|
|
|
|
|
|
@pytest.mark.parametrize("space", TESTING_SPACES, ids=TESTING_SPACES_IDS)
|
|
def test_flatten_roundtripping(space):
|
|
"""Tests roundtripping with flattening and unflattening are equal to the original sample."""
|
|
samples = [space.sample() for _ in range(10)]
|
|
|
|
flattened_samples = [utils.flatten(space, sample) for sample in samples]
|
|
unflattened_samples = [
|
|
utils.unflatten(space, sample) for sample in flattened_samples
|
|
]
|
|
|
|
for original, roundtripped in zip(samples, unflattened_samples):
|
|
assert data_equivalence(original, roundtripped)
|