def subproc_mpi():
# this is necessary for the MPI test to fail
from mpi4py import MPI
env = SubprocEnv(env_fn=create_test_env_mpi,
env_kwargs=dict(episode_len=2))
rew, obs, first = env.observe()
def test_subproc_cloudpickle():
def make_env():
env = IdentityEnv()
env.square = lambda x: x**2
return env
env = SubprocEnv(env_fn=make_env)
assert env.callmethod("square", 2) == 4
def test_subproc():
env = SubprocEnv(env_fn=create_test_env, env_kwargs=dict(episode_len=2))
rew, obs, first = env.observe()
assert first
assert rew == 0.0
env.act(obs)
env.act(obs)
env.act(obs)
infos = env.get_info()
assert infos[0]["step"] == 3
def test_subproc_exception():
with pytest.raises(SubprocError):
env = SubprocEnv(env_fn=create_test_env_exception_on_create)
env = SubprocEnv(env_fn=create_test_env_exception_on_act)
rew, obs, first = env.observe()
env.act(obs)
with pytest.raises(SubprocError):
# because act does not get the return value, the error actually appears during the
# next call to observe()
env.observe()
with pytest.raises(SubprocError):
env = SubprocEnv(env_fn=create_test_env_exit_on_create)
env = SubprocEnv(env_fn=create_test_env_exit_on_act)
rew, obs, first = env.observe()
env.act(obs)
with pytest.raises(SubprocError):
env.observe()
def test_subproc_retro():
env = retro.make("Airstriker-Genesis")
with pytest.raises(RuntimeError):
env = retro.make("Airstriker-Genesis")
envs = [
SubprocEnv(env_fn=make_retro_env,
env_kwargs=dict(game="Airstriker-Genesis"))
for _ in range(2)
]
env = ConcatEnv(envs)
rew, ob, first = env.observe()
assert np.array_equal(ob[0], ob[1])
def vectorize_gym(num,
env_fn=None,
env_kwargs=None,
use_subproc=True,
render_mode=None,
seed=None):
"""
Given a function that creates a gym environment and a number of environments to create,
create the environments in subprocesses and combine them into a single gym3 Env.
This is meant as a replacement for baselines' SubprocVecEnv and DummyVecEnv
If you want to use this for a registered gym env, the default is to use gym.make as the function
to call:
env = vectorize_gym(num=2, env_kwargs={"id": "Pendulum-v0"})
:param num: number of gym environments to create
:param env_fn: function to call to create the gym environment, defaults to `gym.make`
:param env_kwargs: keyword arguments to pass to env_fn
:param use_subproc: if set to False, create the environment in the current process
:param render_mode: if set, this will be passed to the `FromGymEnv` adapter,
see the documentation for `FromGymEnv` for more information
"""
if env_fn is None:
import gym
env_fn = gym.make
if env_kwargs is None:
env_kwargs = {}
if use_subproc:
envs = [
SubprocEnv(
env_fn=_make_gym_env,
env_kwargs=dict(
env_fn=env_fn,
env_kwargs=env_kwargs,
render_mode=render_mode,
seed=seed,
),
) for _ in range(num)
]
else:
envs = [
_make_gym_env(env_fn=env_fn,
env_kwargs=env_kwargs,
render_mode=render_mode,
seed=seed) for _ in range(num)
]
return ConcatEnv(envs)
def test_subproc_callmethod():
env = SubprocEnv(env_fn=create_callmethod_test_env)
assert env.callmethod("square", 2) == 4