def test_unflatten_n(self):
disc = Discrete(3)
obs = np.asarray([0, 1, 2])
flat_arr = disc.flatten_n(obs)
base = np.asarray([0, 1, 2])
unflat_arr = disc.unflatten_n(flat_arr)
assert np.array_equal(unflat_arr, base)
def test_flatten_n(self):
disc = Discrete(3)
tup = Tuple((Discrete(2), disc))
obs = disc.flatten_n(np.asarray([0, 1, 0, 1, 2]))
unflat_ret = tup.unflatten_n(obs)
flat_ret = tup.flatten_n(unflat_ret)
base = np.asarray([[1., 0., 1., 0., 0.]])
assert np.array_equal(flat_ret, base)
def test_flatten_n(self):
disc = Discrete(3)
obs = np.asarray([0, 1, 2])
arr = disc.flatten_n(obs)
base = np.asarray([[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]])
assert np.array_equal(arr, base)
def test_unflatten_n(self):
disc = Discrete(3)
tup = Tuple((Discrete(2), disc))
obs = disc.flatten_n(np.asarray([0, 1, 0, 1, 2]))
ret = tup.unflatten_n(obs)
assert ret == [(0, 0)]
