wehub-resource-sync
252 lines
8.6 KiB
Python
252 lines
8.6 KiB
Python
import unittest
|
|
from collections import OrderedDict
|
|
|
|
import gymnasium as gym
|
|
import numpy as np
|
|
|
|
from ray.rllib.utils.serialization import (
|
|
convert_numpy_to_python_primitives,
|
|
gym_space_from_dict,
|
|
gym_space_to_dict,
|
|
space_from_dict,
|
|
space_to_dict,
|
|
)
|
|
from ray.rllib.utils.spaces.flexdict import FlexDict
|
|
from ray.rllib.utils.spaces.repeated import Repeated
|
|
from ray.rllib.utils.spaces.simplex import Simplex
|
|
|
|
|
|
def _assert_array_equal(eq, a1, a2, margin=None):
|
|
for a in zip(a1, a2):
|
|
eq(a[0], a[1], margin)
|
|
|
|
|
|
class TestGymCheckEnv(unittest.TestCase):
|
|
def test_box_space(self):
|
|
env = gym.make("CartPole-v1")
|
|
d = gym_space_to_dict(env.observation_space)
|
|
sp = gym_space_from_dict(d)
|
|
|
|
obs_space = env.observation_space
|
|
_assert_array_equal(
|
|
self.assertAlmostEqual, sp.low.tolist(), obs_space.low.tolist(), 0.001
|
|
)
|
|
_assert_array_equal(
|
|
self.assertAlmostEqual, sp.high.tolist(), obs_space.high.tolist(), 0.001
|
|
)
|
|
_assert_array_equal(self.assertEqual, sp._shape, obs_space._shape)
|
|
self.assertEqual(sp.dtype, obs_space.dtype)
|
|
|
|
def test_discrete_space(self):
|
|
env = gym.make("CartPole-v1")
|
|
d = gym_space_to_dict(env.action_space)
|
|
sp = gym_space_from_dict(d)
|
|
|
|
action_space = env.action_space
|
|
self.assertEqual(sp.n, action_space.n)
|
|
|
|
def test_multi_binary_space(self):
|
|
mb = gym.spaces.MultiBinary((2, 3))
|
|
d = space_to_dict(mb)
|
|
sp = space_from_dict(d)
|
|
|
|
self.assertEqual(sp.n, mb.n)
|
|
|
|
def test_multi_discrete_space(self):
|
|
md_space = gym.spaces.MultiDiscrete(nvec=np.array([3, 4, 5]))
|
|
d = gym_space_to_dict(md_space)
|
|
sp = gym_space_from_dict(d)
|
|
|
|
_assert_array_equal(self.assertAlmostEqual, sp.nvec, md_space.nvec, 0.001)
|
|
self.assertEqual(md_space.dtype, sp.dtype)
|
|
|
|
def test_tuple_space(self):
|
|
env = gym.make("CartPole-v1")
|
|
space = gym.spaces.Tuple(spaces=[env.observation_space, env.action_space])
|
|
d = gym_space_to_dict(space)
|
|
sp = gym_space_from_dict(d)
|
|
|
|
_assert_array_equal(
|
|
self.assertAlmostEqual,
|
|
sp.spaces[0].low.tolist(),
|
|
space.spaces[0].low.tolist(),
|
|
0.001,
|
|
)
|
|
_assert_array_equal(
|
|
self.assertAlmostEqual,
|
|
sp.spaces[0].high.tolist(),
|
|
space.spaces[0].high.tolist(),
|
|
0.001,
|
|
)
|
|
_assert_array_equal(
|
|
self.assertEqual, sp.spaces[0]._shape, space.spaces[0]._shape
|
|
)
|
|
self.assertEqual(sp.dtype, space.dtype)
|
|
|
|
self.assertEqual(sp.spaces[1].n, space.spaces[1].n)
|
|
|
|
def test_dict_space(self):
|
|
env = gym.make("CartPole-v1")
|
|
space = gym.spaces.Dict(
|
|
spaces={"obs": env.observation_space, "action": env.action_space}
|
|
)
|
|
d = gym_space_to_dict(space)
|
|
sp = gym_space_from_dict(d)
|
|
|
|
_assert_array_equal(
|
|
self.assertAlmostEqual,
|
|
sp.spaces["obs"].low.tolist(),
|
|
space.spaces["obs"].low.tolist(),
|
|
0.001,
|
|
)
|
|
_assert_array_equal(
|
|
self.assertAlmostEqual,
|
|
sp.spaces["obs"].high.tolist(),
|
|
space.spaces["obs"].high.tolist(),
|
|
0.001,
|
|
)
|
|
_assert_array_equal(
|
|
self.assertEqual, sp.spaces["obs"]._shape, space.spaces["obs"]._shape
|
|
)
|
|
self.assertEqual(sp.dtype, space.dtype)
|
|
|
|
self.assertEqual(sp.spaces["action"].n, space.spaces["action"].n)
|
|
|
|
def test_dict_space_with_ordered_dict(self):
|
|
"""Tests whether correct dict order is restored based on the original order."""
|
|
# User provides an OrderedDict -> gymnasium should take it and not further
|
|
# sort the keys. The same (user-provided) order must be restored.
|
|
input_space = gym.spaces.Dict(
|
|
OrderedDict(
|
|
{
|
|
"b_key": gym.spaces.Box(low=np.array([-1.0]), high=np.array([1.0])),
|
|
"a_key": gym.spaces.Discrete(n=3),
|
|
}
|
|
)
|
|
)
|
|
serialized_dict = space_to_dict(input_space)
|
|
deserialized_space = space_from_dict(serialized_dict)
|
|
self.assertTrue(input_space == deserialized_space)
|
|
|
|
# User provides a simple dict -> gymnasium automatically sorts all keys
|
|
# alphabetically. The same (alphabetical) order must be restored.
|
|
input_space = gym.spaces.Dict(
|
|
{
|
|
"b_key": gym.spaces.Box(low=np.array([-1.0]), high=np.array([1.0])),
|
|
"a_key": gym.spaces.Discrete(n=3),
|
|
}
|
|
)
|
|
serialized_dict = space_to_dict(input_space)
|
|
deserialized_space = space_from_dict(serialized_dict)
|
|
self.assertTrue(input_space == deserialized_space)
|
|
|
|
def test_simplex_space(self):
|
|
space = Simplex(shape=(3, 4), concentration=np.array((1, 2, 1, 2)))
|
|
|
|
d = gym_space_to_dict(space)
|
|
sp = gym_space_from_dict(d)
|
|
|
|
_assert_array_equal(self.assertEqual, space.shape, sp.shape)
|
|
_assert_array_equal(
|
|
self.assertAlmostEqual, space.concentration, sp.concentration
|
|
)
|
|
self.assertEqual(space.dtype, sp.dtype)
|
|
|
|
def test_repeated(self):
|
|
space = Repeated(gym.spaces.Box(low=-1, high=1, shape=(1, 200)), max_len=8)
|
|
|
|
d = gym_space_to_dict(space)
|
|
sp = gym_space_from_dict(d)
|
|
|
|
self.assertTrue(isinstance(sp.child_space, gym.spaces.Box))
|
|
self.assertEqual(space.max_len, sp.max_len)
|
|
self.assertEqual(space.dtype, sp.dtype)
|
|
|
|
def test_flex_dict(self):
|
|
space = FlexDict({})
|
|
space["box"] = gym.spaces.Box(low=-1, high=1, shape=(1, 200))
|
|
space["discrete"] = gym.spaces.Discrete(2)
|
|
space["tuple"] = gym.spaces.Tuple(
|
|
(gym.spaces.Box(low=-1, high=1, shape=(1, 200)), gym.spaces.Discrete(2))
|
|
)
|
|
|
|
d = gym_space_to_dict(space)
|
|
sp = gym_space_from_dict(d)
|
|
|
|
self.assertTrue(isinstance(sp["box"], gym.spaces.Box))
|
|
self.assertTrue(isinstance(sp["discrete"], gym.spaces.Discrete))
|
|
self.assertTrue(isinstance(sp["tuple"], gym.spaces.Tuple))
|
|
|
|
def test_text(self):
|
|
expected_space = gym.spaces.Text(min_length=3, max_length=10, charset="abc")
|
|
d = gym_space_to_dict(expected_space)
|
|
sp = gym_space_from_dict(d)
|
|
|
|
self.assertEqual(expected_space.max_length, sp.max_length)
|
|
self.assertEqual(expected_space.min_length, sp.min_length)
|
|
|
|
charset = getattr(expected_space, "character_set", None)
|
|
if charset is not None:
|
|
self.assertEqual(expected_space.character_set, sp.character_set)
|
|
else:
|
|
charset = getattr(expected_space, "charset", None)
|
|
if charset is None:
|
|
raise ValueError(
|
|
"Text space does not have charset or character_set attribute."
|
|
)
|
|
self.assertEqual(expected_space.charset, sp.charset)
|
|
|
|
def test_original_space(self):
|
|
space = gym.spaces.Box(low=0.0, high=1.0, shape=(10,))
|
|
space.original_space = gym.spaces.Dict(
|
|
{
|
|
"obs1": gym.spaces.Box(low=0.0, high=1.0, shape=(3,)),
|
|
"obs2": gym.spaces.Box(low=0.0, high=1.0, shape=(7,)),
|
|
}
|
|
)
|
|
|
|
d = space_to_dict(space)
|
|
sp = space_from_dict(d)
|
|
|
|
self.assertTrue(isinstance(sp, gym.spaces.Box))
|
|
self.assertTrue(isinstance(sp.original_space, gym.spaces.Dict))
|
|
self.assertTrue(isinstance(sp.original_space["obs1"], gym.spaces.Box))
|
|
self.assertTrue(isinstance(sp.original_space["obs2"], gym.spaces.Box))
|
|
|
|
def test_unknown_space_type_error_message(self):
|
|
with self.assertRaisesRegex(
|
|
ValueError, "Unknown space type for serialization:"
|
|
):
|
|
gym_space_to_dict(object())
|
|
|
|
def test_unknown_serialized_space_type_error_message(self):
|
|
with self.assertRaisesRegex(
|
|
ValueError, "Unknown space type for de-serialization: made-up-space"
|
|
):
|
|
gym_space_from_dict({"space": "made-up-space"})
|
|
|
|
|
|
class TestConvertNumpyToPythonPrimitives(unittest.TestCase):
|
|
def test_convert_numpy_to_python_primitives(self):
|
|
# test utility for converting numpy types to python primitives
|
|
test_cases = [
|
|
[1, 2, 3],
|
|
[1.0, 2.0, 3.0],
|
|
["abc", "def", "ghi"],
|
|
[True, False, True],
|
|
]
|
|
for test_case in test_cases:
|
|
_assert_array_equal(
|
|
self.assertEqual,
|
|
convert_numpy_to_python_primitives(np.array(test_case)),
|
|
test_case,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|