def test_custom_space_sync_vector_env():
env_fns = [make_custom_space_env(i) for i in range(4)]
try:
env = SyncVectorEnv(env_fns)
reset_observations = env.reset()
actions = ('action-2', 'action-3', 'action-5', 'action-7')
step_observations, rewards, dones, _ = env.step(actions)
finally:
env.close()
assert isinstance(env.single_observation_space, CustomSpace)
assert isinstance(env.observation_space, Tuple)
assert isinstance(reset_observations, tuple)
assert reset_observations == ('reset', 'reset', 'reset', 'reset')
assert isinstance(step_observations, tuple)
assert step_observations == ('step(action-2)', 'step(action-3)',
'step(action-5)', 'step(action-7)')
def test_custom_space_async_vector_env():
env_fns = [make_custom_space_env(i) for i in range(4)]
try:
env = AsyncVectorEnv(env_fns, shared_memory=False)
reset_observations = env.reset()
actions = ("action-2", "action-3", "action-5", "action-7")
step_observations, rewards, dones, _ = env.step(actions)
finally:
env.close()
assert isinstance(env.single_observation_space, CustomSpace)
assert isinstance(env.observation_space, Tuple)
assert isinstance(reset_observations, tuple)
assert reset_observations == ("reset", "reset", "reset", "reset")
assert isinstance(step_observations, tuple)
assert step_observations == (
"step(action-2)",
"step(action-3)",
"step(action-5)",
"step(action-7)",
)
def test_custom_space_async_vector_env_shared_memory():
env_fns = [make_custom_space_env(i) for i in range(4)]
with pytest.raises(ValueError):
env = AsyncVectorEnv(env_fns, shared_memory=True)
env.close(terminate=True)