def __init__(self, env):
super().__init__(env)
self.env = env
self.action_space = spaces.Tuple((
spaces.Discrete(2),
spaces.Box(low=-1, high=1, shape=(2, ), dtype=np.float32),
spaces.Box(low=-1, high=1, shape=(2, ), dtype=np.float32),
))
def __init__(self):
super(NanAndInfEnv, self).__init__()
self.action_space = spaces.Box(low=-np.inf,
high=np.inf,
shape=(1, ),
dtype=np.float64)
self.observation_space = spaces.Box(low=-np.inf,
high=np.inf,
shape=(1, ),
dtype=np.float64)
def __init__(self):
super(DummyDictEnv, self).__init__()
self.observation_space = spaces.Dict({
"observation":
spaces.Box(low=-20.0, high=20.0, shape=(4, ), dtype=np.float32),
"achieved_goal":
spaces.Box(low=-20.0, high=20.0, shape=(4, ), dtype=np.float32),
"desired_goal":
spaces.Box(low=-20.0, high=20.0, shape=(4, ), dtype=np.float32),
})
self.action_space = spaces.Box(low=-1,
high=1,
shape=(3, ),
dtype=np.float32)
def make_transposed_image_env():
return CustomGymEnv(
spaces.Box(
low=np.zeros(transposed_image_space_shape),
high=np.ones(transposed_image_space_shape) * 255,
dtype=np.uint8,
))
def __init__(self, venv: VecEnv):
super(ObsDictWrapper, self).__init__(venv, venv.observation_space, venv.action_space)
self.venv = venv
self.spaces = list(venv.observation_space.spaces.values())
# get dimensions of observation and goal
if isinstance(self.spaces[0], spaces.Discrete):
self.obs_dim = 1
self.goal_dim = 1
else:
self.obs_dim = venv.observation_space.spaces["observation"].shape[0]
self.goal_dim = venv.observation_space.spaces["achieved_goal"].shape[0]
# new observation space with concatenated observation and (desired) goal
# for the different types of spaces
if isinstance(self.spaces[0], spaces.Box):
low_values = np.concatenate(
[venv.observation_space.spaces["observation"].low, venv.observation_space.spaces["desired_goal"].low]
)
high_values = np.concatenate(
[venv.observation_space.spaces["observation"].high, venv.observation_space.spaces["desired_goal"].high]
)
self.observation_space = spaces.Box(low_values, high_values, dtype=np.float32)
elif isinstance(self.spaces[0], spaces.MultiBinary):
total_dim = self.obs_dim + self.goal_dim
self.observation_space = spaces.MultiBinary(total_dim)
elif isinstance(self.spaces[0], spaces.Discrete):
dimensions = [venv.observation_space.spaces["observation"].n, venv.observation_space.spaces["desired_goal"].n]
self.observation_space = spaces.MultiDiscrete(dimensions)
else:
raise NotImplementedError(f"{type(self.spaces[0])} space is not supported")
def __init__(self, env: gym.Env, width: int = 84, height: int = 84):
gym.ObservationWrapper.__init__(self, env)
self.width = width
self.height = height
self.observation_space = spaces.Box(low=0,
high=255,
shape=(self.height, self.width, 1),
dtype=env.observation_space.dtype)
def __init__(self, max_steps):
"""Gym environment for testing that terminal observation is inserted
correctly."""
self.action_space = spaces.Discrete(2)
self.observation_space = spaces.Box(np.array([0]),
np.array([999]),
dtype="int")
self.max_steps = max_steps
self.current_step = 0
def __init__(self):
super(DummyEnv, self).__init__()
# First entry in teh spaces.Tuple to be the selection of the skill
# the remaining entries are the paramters for the specific skills
# So if you have 2 skills the tuple will need 3 entries in total:
# The first will be a a discrete space with n = 2
# The second one will be the parameters for skill1
# The third will be the parameters for skill2
self.action_space = spaces.Tuple(
(
spaces.Discrete(2),
spaces.Box(low=-1, high=1, shape=(4,), dtype=np.float32),
spaces.Box(low=-1, high=1, shape=(4,), dtype=np.float32),
)
)
# self.action_space = spaces.Box(low=-1, high=1, shape=(1, ))
# Example for using image as input:
self.observation_space = spaces.Box(low=0, high=255, shape=(2,), dtype=np.float32)
def test_image_space_checks():
not_image_space = spaces.Box(0, 1, shape=(10, ))
assert not is_image_space(not_image_space)
# Not uint8
not_image_space = spaces.Box(0, 255, shape=(10, 10, 3))
assert not is_image_space(not_image_space)
# Not correct shape
not_image_space = spaces.Box(0, 255, shape=(10, 10), dtype=np.uint8)
assert not is_image_space(not_image_space)
# Not correct low/high
not_image_space = spaces.Box(0, 10, shape=(10, 10, 3), dtype=np.uint8)
assert not is_image_space(not_image_space)
# Not correct space
not_image_space = spaces.Discrete(n=10)
assert not is_image_space(not_image_space)
an_image_space = spaces.Box(0, 255, shape=(10, 10, 3), dtype=np.uint8)
assert is_image_space(an_image_space)
an_image_space_with_odd_channels = spaces.Box(0,
255,
shape=(10, 10, 5),
dtype=np.uint8)
assert is_image_space(an_image_space_with_odd_channels)
# Should not pass if we check if channels are valid for an image
assert not is_image_space(an_image_space_with_odd_channels,
check_channels=True)
# Test if channel-check works
channel_first_space = spaces.Box(0, 255, shape=(3, 10, 10), dtype=np.uint8)
assert is_image_space_channels_first(channel_first_space)
channel_last_space = spaces.Box(0, 255, shape=(10, 10, 3), dtype=np.uint8)
assert not is_image_space_channels_first(channel_last_space)
channel_mid_space = spaces.Box(0, 255, shape=(10, 3, 10), dtype=np.uint8)
# Should raise a warning
with pytest.warns(Warning):
assert not is_image_space_channels_first(channel_mid_space)
def transpose_space(observation_space: spaces.Box) -> spaces.Box:
"""
Transpose an observation space (re-order channels).
:param observation_space:
:return:
"""
assert is_image_space(
observation_space), "The observation space must be an image"
width, height, channels = observation_space.shape
new_shape = (channels, width, height)
return spaces.Box(low=0,
high=255,
shape=new_shape,
dtype=observation_space.dtype)
def test_high_dimension_action_space():
"""
Test for continuous action space
with more than one action.
"""
env = FakeImageEnv()
# Patch the action space
env.action_space = spaces.Box(low=-1,
high=1,
shape=(20, ),
dtype=np.float32)
# Patch to avoid error
def patched_step(_action):
return env.observation_space.sample(), 0.0, False, {}
env.step = patched_step
check_env(env)
def __init__(self,
venv: VecEnv,
n_stack: int,
channels_order: Optional[str] = None):
self.venv = venv
self.n_stack = n_stack
wrapped_obs_space = venv.observation_space
assert isinstance(
wrapped_obs_space, spaces.Box
), "VecFrameStack only work with hmlf.spaces.Box observation space"
if channels_order is None:
# Detect channel location automatically for images
if is_image_space(wrapped_obs_space):
self.channels_first = is_image_space_channels_first(
wrapped_obs_space)
else:
# Default behavior for non-image space, stack on the last axis
self.channels_first = False
else:
assert channels_order in {
"last", "first"
}, "`channels_order` must be one of following: 'last', 'first'"
self.channels_first = channels_order == "first"
# This includes the vec-env dimension (first)
self.stack_dimension = 1 if self.channels_first else -1
repeat_axis = 0 if self.channels_first else -1
low = np.repeat(wrapped_obs_space.low, self.n_stack, axis=repeat_axis)
high = np.repeat(wrapped_obs_space.high,
self.n_stack,
axis=repeat_axis)
self.stackedobs = np.zeros((venv.num_envs, ) + low.shape, low.dtype)
observation_space = spaces.Box(low=low,
high=high,
dtype=venv.observation_space.dtype)
VecEnvWrapper.__init__(self, venv, observation_space=observation_space)
def make_non_image_env():
return CustomGymEnv(
spaces.Box(low=np.zeros((2, )), high=np.ones((2, ))))
def __init__(self, nvec):
super(DummyMultiDiscreteSpace, self).__init__()
self.observation_space = spaces.MultiDiscrete(nvec)
self.action_space = spaces.Box(low=-1, high=1, shape=(2,), dtype=np.float32)
def __init__(self, n):
super(DummyMultiBinary, self).__init__()
self.observation_space = spaces.MultiBinary(n)
self.action_space = spaces.Box(low=-1, high=1, shape=(2,), dtype=np.float32)
if not isinstance(vec_env, VecNormalize):
# more precise tests that we can't do with VecNormalize
# (which changes observation values)
assert np.all(prev_obs + 1 == terminal_obs)
assert np.all(obs == 0)
prev_obs_b = obs_b
vec_env.close()
SPACES = collections.OrderedDict([
("discrete", spaces.Discrete(2)),
("multidiscrete", spaces.MultiDiscrete([2, 3])),
("multibinary", spaces.MultiBinary(3)),
("continuous", spaces.Box(low=np.zeros(2), high=np.ones(2))),
])
def check_vecenv_spaces(vec_env_class, space, obs_assert):
"""Helper method to check observation spaces in vectorized environments."""
def make_env():
return CustomGymEnv(space)
vec_env = vec_env_class([make_env for _ in range(N_ENVS)])
obs = vec_env.reset()
obs_assert(obs)
dones = [False] * N_ENVS
while not any(dones):
actions = [vec_env.action_space.sample() for _ in range(N_ENVS)]
shape=(20, ),
dtype=np.float32)
# Patch to avoid error
def patched_step(_action):
return env.observation_space.sample(), 0.0, False, {}
env.step = patched_step
check_env(env)
@pytest.mark.parametrize(
"new_obs_space",
[
# Small image
spaces.Box(low=0, high=255, shape=(32, 32, 3), dtype=np.uint8),
# Range not in [0, 255]
spaces.Box(low=0, high=1, shape=(64, 64, 3), dtype=np.uint8),
# Wrong dtype
spaces.Box(low=0, high=255, shape=(64, 64, 3), dtype=np.float32),
# Not an image, it should be a 1D vector
spaces.Box(low=-1, high=1, shape=(64, 3), dtype=np.float32),
# Tuple space is not supported by SB
spaces.Tuple([spaces.Discrete(5),
spaces.Discrete(10)]),
# Dict space is not supported by SB when env is not a GoalEnv
spaces.Dict({"position": spaces.Discrete(5)}),
],
)
def test_non_default_spaces(new_obs_space):
env = FakeImageEnv()
def make_env():
return CustomGymEnv(spaces.Box(low=np.zeros(2), high=np.ones(2)))
def make_monitored_env():
return Monitor(
CustomGymEnv(spaces.Box(low=np.zeros(2), high=np.ones(2))))
