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))))