mirror of
https://github.com/vale981/ray
synced 2025-03-12 14:16:39 -04:00
487 lines
14 KiB
Python
487 lines
14 KiB
Python
from gym.spaces import Box, Discrete, MultiDiscrete
|
|
import numpy as np
|
|
import tree # pip install dm_tree
|
|
import unittest
|
|
|
|
import ray
|
|
from ray.rllib.utils.framework import try_import_tf, try_import_torch
|
|
from ray.rllib.utils.numpy import flatten_inputs_to_1d_tensor as flatten_np
|
|
from ray.rllib.utils.test_utils import check
|
|
from ray.rllib.utils.tf_utils import flatten_inputs_to_1d_tensor as flatten_tf
|
|
from ray.rllib.utils.torch_utils import flatten_inputs_to_1d_tensor as flatten_torch
|
|
|
|
tf1, tf, tfv = try_import_tf()
|
|
torch, _ = try_import_torch()
|
|
|
|
|
|
class TestUtils(unittest.TestCase):
|
|
# Nested struct of data with B=3.
|
|
struct = {
|
|
"a": np.array([1, 3, 2]),
|
|
"b": (
|
|
np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]),
|
|
np.array(
|
|
[[[8.0], [7.0], [6.0]], [[5.0], [4.0], [3.0]], [[2.0], [1.0], [0.0]]]
|
|
),
|
|
),
|
|
"c": {
|
|
"ca": np.array([[1, 2], [3, 5], [0, 1]]),
|
|
"cb": np.array([1.0, 2.0, 3.0]),
|
|
},
|
|
}
|
|
# Nested struct of data with B=2 and T=1.
|
|
struct_w_time_axis = {
|
|
"a": np.array([[1], [3]]),
|
|
"b": (
|
|
np.array([[[1.0, 2.0, 3.0]], [[4.0, 5.0, 6.0]]]),
|
|
np.array([[[[8.0], [7.0], [6.0]]], [[[5.0], [4.0], [3.0]]]]),
|
|
),
|
|
"c": {"ca": np.array([[[1, 2]], [[3, 5]]]), "cb": np.array([[1.0], [2.0]])},
|
|
}
|
|
# Corresponding space struct.
|
|
spaces = dict(
|
|
{
|
|
"a": Discrete(4),
|
|
"b": (Box(-1.0, 10.0, (3,)), Box(-1.0, 1.0, (3, 1))),
|
|
"c": dict(
|
|
{
|
|
"ca": MultiDiscrete([4, 6]),
|
|
"cb": Box(-1.0, 1.0, ()),
|
|
}
|
|
),
|
|
}
|
|
)
|
|
|
|
@classmethod
|
|
def setUpClass(cls) -> None:
|
|
ray.init()
|
|
|
|
@classmethod
|
|
def tearDownClass(cls) -> None:
|
|
ray.shutdown()
|
|
|
|
def test_flatten_inputs_to_1d_tensor(self):
|
|
# B=3; no time axis.
|
|
check(
|
|
flatten_np(self.struct, spaces_struct=self.spaces),
|
|
np.array(
|
|
[
|
|
[
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
2.0,
|
|
3.0,
|
|
8.0,
|
|
7.0,
|
|
6.0,
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
],
|
|
[
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
4.0,
|
|
5.0,
|
|
6.0,
|
|
5.0,
|
|
4.0,
|
|
3.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
2.0,
|
|
],
|
|
[
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
7.0,
|
|
8.0,
|
|
9.0,
|
|
2.0,
|
|
1.0,
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
3.0,
|
|
],
|
|
]
|
|
),
|
|
)
|
|
|
|
struct_tf = tree.map_structure(lambda s: tf.convert_to_tensor(s), self.struct)
|
|
check(
|
|
flatten_tf(struct_tf, spaces_struct=self.spaces),
|
|
np.array(
|
|
[
|
|
[
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
2.0,
|
|
3.0,
|
|
8.0,
|
|
7.0,
|
|
6.0,
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
],
|
|
[
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
4.0,
|
|
5.0,
|
|
6.0,
|
|
5.0,
|
|
4.0,
|
|
3.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
2.0,
|
|
],
|
|
[
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
7.0,
|
|
8.0,
|
|
9.0,
|
|
2.0,
|
|
1.0,
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
3.0,
|
|
],
|
|
]
|
|
),
|
|
)
|
|
|
|
struct_torch = tree.map_structure(lambda s: torch.from_numpy(s), self.struct)
|
|
check(
|
|
flatten_torch(struct_torch, spaces_struct=self.spaces),
|
|
np.array(
|
|
[
|
|
[
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
2.0,
|
|
3.0,
|
|
8.0,
|
|
7.0,
|
|
6.0,
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
],
|
|
[
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
4.0,
|
|
5.0,
|
|
6.0,
|
|
5.0,
|
|
4.0,
|
|
3.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
2.0,
|
|
],
|
|
[
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
7.0,
|
|
8.0,
|
|
9.0,
|
|
2.0,
|
|
1.0,
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
3.0,
|
|
],
|
|
]
|
|
),
|
|
)
|
|
|
|
def test_flatten_inputs_to_1d_tensor_w_time_axis(self):
|
|
# B=2; T=1
|
|
check(
|
|
flatten_np(
|
|
self.struct_w_time_axis, spaces_struct=self.spaces, time_axis=True
|
|
),
|
|
np.array(
|
|
[
|
|
[
|
|
[
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
2.0,
|
|
3.0,
|
|
8.0,
|
|
7.0,
|
|
6.0,
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
]
|
|
],
|
|
[
|
|
[
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
4.0,
|
|
5.0,
|
|
6.0,
|
|
5.0,
|
|
4.0,
|
|
3.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
2.0,
|
|
]
|
|
],
|
|
]
|
|
),
|
|
)
|
|
|
|
struct_tf = tree.map_structure(
|
|
lambda s: tf.convert_to_tensor(s), self.struct_w_time_axis
|
|
)
|
|
check(
|
|
flatten_tf(struct_tf, spaces_struct=self.spaces, time_axis=True),
|
|
np.array(
|
|
[
|
|
[
|
|
[
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
2.0,
|
|
3.0,
|
|
8.0,
|
|
7.0,
|
|
6.0,
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
]
|
|
],
|
|
[
|
|
[
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
4.0,
|
|
5.0,
|
|
6.0,
|
|
5.0,
|
|
4.0,
|
|
3.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
2.0,
|
|
]
|
|
],
|
|
]
|
|
),
|
|
)
|
|
|
|
struct_torch = tree.map_structure(
|
|
lambda s: torch.from_numpy(s), self.struct_w_time_axis
|
|
)
|
|
check(
|
|
flatten_torch(struct_torch, spaces_struct=self.spaces, time_axis=True),
|
|
np.array(
|
|
[
|
|
[
|
|
[
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
2.0,
|
|
3.0,
|
|
8.0,
|
|
7.0,
|
|
6.0,
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
]
|
|
],
|
|
[
|
|
[
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
4.0,
|
|
5.0,
|
|
6.0,
|
|
5.0,
|
|
4.0,
|
|
3.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
1.0,
|
|
2.0,
|
|
]
|
|
],
|
|
]
|
|
),
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import pytest
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|