def unflatten_n(space, obs):
if isinstance(space, gym.spaces.Box):
obs = np.asarray(obs)
return obs.reshape((obs.shape[0], ) + space.shape)
elif isinstance(space, gym.spaces.Discrete):
return special.from_onehot_n(obs)
elif isinstance(space, gym.spaces.Tuple):
dims = [flat_dim(c) for c in space.spaces]
flat_xs = np.split(obs, np.cumsum(dims)[:-1], axis=-1)
unflat_xs = [
unflatten_n(c, xi) for c, xi in zip(space.spaces, flat_xs)
]
unflat_xs_grouped = list(zip(*unflat_xs))
return unflat_xs_grouped
else:
raise NotImplementedError
def unflatten_n(self, x):
return special.from_onehot_n(x)