def test_step_timeout_async_vector_env(shared_memory): env_fns = [make_slow_env(0.0, i) for i in range(4)] with pytest.raises(TimeoutError): try: env = AsyncVectorEnv(env_fns, shared_memory=shared_memory) env.reset() env.step_async([0.1, 0.1, 0.3, 0.1]) observations, rewards, dones, _ = env.step_wait(timeout=0.1) finally: env.close(terminate=True)
def test_step_async_vector_env(shared_memory, use_single_action_space): env_fns = [make_env("CubeCrash-v0", i) for i in range(8)] try: env = AsyncVectorEnv(env_fns, shared_memory=shared_memory) observations = env.reset() if use_single_action_space: actions = [env.single_action_space.sample() for _ in range(8)] else: actions = env.action_space.sample() observations, rewards, dones, _ = env.step(actions) finally: env.close() assert isinstance(env.observation_space, Box) assert isinstance(observations, np.ndarray) assert observations.dtype == env.observation_space.dtype assert observations.shape == (8, ) + env.single_observation_space.shape assert observations.shape == env.observation_space.shape assert isinstance(rewards, np.ndarray) assert isinstance(rewards[0], (float, np.floating)) assert rewards.ndim == 1 assert rewards.size == 8 assert isinstance(dones, np.ndarray) assert dones.dtype == np.bool_ assert dones.ndim == 1 assert dones.size == 8
def test_step_out_of_order_async_vector_env(shared_memory): env_fns = [make_env("CubeCrash-v0", i) for i in range(4)] with pytest.raises(NoAsyncCallError): try: env = AsyncVectorEnv(env_fns, shared_memory=shared_memory) actions = env.action_space.sample() observations = env.reset() observations, rewards, dones, infos = env.step_wait() except AlreadyPendingCallError as exception: assert exception.name == "step" raise finally: env.close(terminate=True) with pytest.raises(AlreadyPendingCallError): try: env = AsyncVectorEnv(env_fns, shared_memory=shared_memory) actions = env.action_space.sample() env.reset_async() env.step_async(actions) except AlreadyPendingCallError as exception: assert exception.name == "reset" raise finally: env.close(terminate=True)
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() assert isinstance(env.single_action_space, CustomSpace) assert isinstance(env.action_space, Tuple) 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_no_copy_async_vector_env(shared_memory): env_fns = [make_env("CartPole-v1", i) for i in range(8)] try: env = AsyncVectorEnv(env_fns, shared_memory=shared_memory, copy=False) observations = env.reset() observations[0] = 0 finally: env.close()
def test_no_copy_async_vector_env(shared_memory): env_fns = [make_env("CubeCrash-v0", i) for i in range(8)] try: env = AsyncVectorEnv(env_fns, shared_memory=shared_memory, copy=False) observations = env.reset() observations[0] = 128 assert np.all(env.observations[0] == 128) finally: env.close()
def main(): env_id = "Ant-v3" num_envs = 5 vec_env = AsyncVectorEnv([make_env(env_id) for i in range(num_envs)]) state = vec_env.reset() for i in range(5000): action = vec_env.action_space.sample() state, reward, done, _ = vec_env.step(action) if any(done): done_idx = [i for i, e in enumerate(done) if e] print(f"{done_idx}")
def test_reset_async_vector_env(shared_memory): env_fns = [make_env("CubeCrash-v0", i) for i in range(8)] try: env = AsyncVectorEnv(env_fns, shared_memory=shared_memory) observations = env.reset() finally: env.close() assert isinstance(env.observation_space, Box) assert isinstance(observations, np.ndarray) assert observations.dtype == env.observation_space.dtype assert observations.shape == (8, ) + env.single_observation_space.shape assert observations.shape == env.observation_space.shape
def test_reset_async_vector_env(shared_memory): env_fns = [make_env("CartPole-v1", i) for i in range(8)] try: env = AsyncVectorEnv(env_fns, shared_memory=shared_memory) observations = env.reset() finally: env.close() assert isinstance(env.observation_space, Box) assert isinstance(observations, np.ndarray) assert observations.dtype == env.observation_space.dtype assert observations.shape == (8, ) + env.single_observation_space.shape assert observations.shape == env.observation_space.shape try: env = AsyncVectorEnv(env_fns, shared_memory=shared_memory) observations = env.reset(return_info=False) finally: env.close() assert isinstance(env.observation_space, Box) assert isinstance(observations, np.ndarray) assert observations.dtype == env.observation_space.dtype assert observations.shape == (8, ) + env.single_observation_space.shape assert observations.shape == env.observation_space.shape try: env = AsyncVectorEnv(env_fns, shared_memory=shared_memory) observations, infos = env.reset(return_info=True) finally: env.close() assert isinstance(env.observation_space, Box) assert isinstance(observations, np.ndarray) assert observations.dtype == env.observation_space.dtype assert observations.shape == (8, ) + env.single_observation_space.shape assert observations.shape == env.observation_space.shape assert isinstance(infos, list) assert all([isinstance(info, dict) for info in infos])
def test_vector_env_equal(shared_memory): env_fns = [make_env("CubeCrash-v0", i) for i in range(4)] num_steps = 100 try: async_env = AsyncVectorEnv(env_fns, shared_memory=shared_memory) sync_env = SyncVectorEnv(env_fns) async_env.seed(0) sync_env.seed(0) assert async_env.num_envs == sync_env.num_envs assert async_env.observation_space == sync_env.observation_space assert async_env.single_observation_space == sync_env.single_observation_space assert async_env.action_space == sync_env.action_space assert async_env.single_action_space == sync_env.single_action_space async_observations = async_env.reset() sync_observations = sync_env.reset() assert np.all(async_observations == sync_observations) for _ in range(num_steps): actions = async_env.action_space.sample() assert actions in sync_env.action_space # fmt: off async_observations, async_rewards, async_dones, async_infos = async_env.step( actions) sync_observations, sync_rewards, sync_dones, sync_infos = sync_env.step( actions) # fmt: on for idx in range(len(sync_dones)): if sync_dones[idx]: assert "terminal_observation" in async_infos[idx] assert "terminal_observation" in sync_infos[idx] assert sync_dones[idx] assert np.all(async_observations == sync_observations) assert np.all(async_rewards == sync_rewards) assert np.all(async_dones == sync_dones) finally: async_env.close() sync_env.close()
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_call_async_vector_env(shared_memory): env_fns = [make_env("CartPole-v1", i) for i in range(4)] try: env = AsyncVectorEnv(env_fns, shared_memory=shared_memory) _ = env.reset() images = env.call("render", mode="rgb_array") gravity = env.call("gravity") finally: env.close() assert isinstance(images, tuple) assert len(images) == 4 for i in range(4): assert isinstance(images[i], np.ndarray) assert isinstance(gravity, tuple) assert len(gravity) == 4 for i in range(4): assert isinstance(gravity[i], float) assert gravity[i] == 9.8
def test_vector_env_equal(shared_memory): env_fns = [make_env("CartPole-v1", i) for i in range(4)] num_steps = 100 try: async_env = AsyncVectorEnv(env_fns, shared_memory=shared_memory) sync_env = SyncVectorEnv(env_fns) assert async_env.num_envs == sync_env.num_envs assert async_env.observation_space == sync_env.observation_space assert async_env.single_observation_space == sync_env.single_observation_space assert async_env.action_space == sync_env.action_space assert async_env.single_action_space == sync_env.single_action_space async_observations = async_env.reset(seed=0) sync_observations = sync_env.reset(seed=0) assert np.all(async_observations == sync_observations) for _ in range(num_steps): actions = async_env.action_space.sample() assert actions in sync_env.action_space # fmt: off async_observations, async_rewards, async_dones, async_infos = async_env.step( actions) sync_observations, sync_rewards, sync_dones, sync_infos = sync_env.step( actions) # fmt: on if any(sync_dones): assert "terminal_observation" in async_infos assert "_terminal_observation" in async_infos assert "terminal_observation" in sync_infos assert "_terminal_observation" in sync_infos assert np.all(async_observations == sync_observations) assert np.all(async_rewards == sync_rewards) assert np.all(async_dones == sync_dones) finally: async_env.close() sync_env.close()
def test_vector_env_equal(shared_memory): env_fns = [make_env('CubeCrash-v0', i) for i in range(4)] num_steps = 100 try: async_env = AsyncVectorEnv(env_fns, shared_memory=shared_memory) sync_env = SyncVectorEnv(env_fns) async_env.seed(0) sync_env.seed(0) assert async_env.num_envs == sync_env.num_envs assert async_env.observation_space == sync_env.observation_space assert async_env.single_observation_space == sync_env.single_observation_space assert async_env.action_space == sync_env.action_space assert async_env.single_action_space == sync_env.single_action_space async_observations = async_env.reset() sync_observations = sync_env.reset() assert np.all(async_observations == sync_observations) for _ in range(num_steps): actions = async_env.action_space.sample() assert actions in sync_env.action_space async_observations, async_rewards, async_dones, _ = async_env.step( actions) sync_observations, sync_rewards, sync_dones, _ = sync_env.step( actions) assert np.all(async_observations == sync_observations) assert np.all(async_rewards == sync_rewards) assert np.all(async_dones == sync_dones) finally: async_env.close() sync_env.close()
def test_already_closed_async_vector_env(shared_memory): env_fns = [make_env("CubeCrash-v0", i) for i in range(4)] with pytest.raises(ClosedEnvironmentError): env = AsyncVectorEnv(env_fns, shared_memory=shared_memory) env.close() observations = env.reset()
def test_already_closed_async_vector_env(shared_memory): env_fns = [make_env("CartPole-v1", i) for i in range(4)] with pytest.raises(ClosedEnvironmentError): env = AsyncVectorEnv(env_fns, shared_memory=shared_memory) env.close() env.reset()