def test_asynchronous_exception():
env = IdentityEnv()
env = ExceptionOnAct(env)
env = AsynchronousWrapper(env)
env.observe()
env.act(0)
with pytest.raises(CustomException):
env.observe()
def test_asynchronous_ordering():
num_steps = 100
env = IdentityEnv()
env = RecordActs(env)
env = AsynchronousWrapper(env)
for i in range(num_steps):
env.act(i)
env.observe()
assert env.env.acts == list(range(num_steps))
def test_identity_env():
env = IdentityEnv(space=types.discrete_scalar(1024), episode_len=2)
rew, ob, first = env.observe()
assert ob == ob
assert first
assert rew == 0
env.act(ob)
rew, ob, first = env.observe()
assert not first
assert rew == 1
env.act(ob)
rew, ob, first = env.observe()
assert first
assert rew == 1
def test_recorder():
with tempfile.TemporaryDirectory() as tmpdir:
ep_len1 = 3
ep_len2 = 4
env1 = IdentityEnv(
space=types.TensorType(eltype=types.Discrete(256),
shape=(3, 3, 3)),
episode_len=ep_len1,
)
env2 = IdentityEnv(
space=types.TensorType(eltype=types.Discrete(256),
shape=(3, 3, 3)),
episode_len=ep_len2,
seed=1,
)
env = ConcatEnv([env1, env2])
env = TrajectoryRecorderWrapper(env=env, directory=tmpdir)
_, obs, _ = env.observe()
action = types_np.zeros(env.ac_space, bshape=(env.num, ))
action[1] = 1
num_acs = 10
for _ in range(num_acs):
env.act(action)
files = sorted(glob(os.path.join(tmpdir, "*.pickle")))
print(files)
assert len(files) == (num_acs // ep_len1) + (num_acs // ep_len2)
with open(files[0], "rb") as f:
loaded_traj = pickle.load(f)
assert len(loaded_traj["ob"]) == ep_len1
assert np.allclose(loaded_traj["ob"][0], obs[0])
assert np.allclose(loaded_traj["act"][0], action[0])
with open(files[1], "rb") as f:
loaded_traj = pickle.load(f)
assert len(loaded_traj["ob"]) == ep_len2
assert np.allclose(loaded_traj["ob"][0], obs[1])
assert np.allclose(loaded_traj["act"][0], action[1])
def test_identity_env_delay():
delay = 2
for space in [types.discrete_scalar(1024)]:
env = IdentityEnv(space=space,
episode_len=delay * 2,
delay_steps=delay)
_rew, obs, _first = env.observe()
obs_queue = [obs]
for i in range(delay):
env.act(0)
_rew, obs, _first = env.observe()
obs_queue.append(obs)
first = False
for i in range(delay):
env.act(obs_queue.pop(0))
rew, obs, first = env.observe()
obs_queue.append(obs)
assert rew == 1
assert first
def test_concat(space_type):
if space_type == "binary":
space = types.discrete_scalar(2)
elif space_type == "dict":
space = types.DictType(degrees=types.discrete_scalar(360))
elif space_type == "real_dict":
space = types.DictType(
degrees=types.TensorType(shape=(), eltype=types.Real()))
else:
raise Exception(f"invalid space_type {space_type}")
base_env1 = IdentityEnv(space=space, episode_len=1, seed=0)
base_env1.f = lambda x: [x[0]**2]
base_env2 = IdentityEnv(space=space, episode_len=2, seed=1)
base_env2.f = lambda x: [2 * x[0]]
env1 = AssertSpacesWrapper(base_env1)
env1 = AddInfo(env=env1, id=1)
env2 = AssertSpacesWrapper(base_env2)
env2 = AddInfo(env=env2, id=2)
env = AssertSpacesWrapper(ConcatEnv([env1, env2]))
rew, ob, first = env.observe()
assert np.array_equal(rew, np.array([0, 0]))
if isinstance(space, types.DictType):
ob = ob["degrees"]
assert ob.shape == (2, )
assert ob[0] != ob[1]
assert np.array_equal(first, np.array([1, 1]))
act = np.array([ob[0], ob[0]])
if isinstance(space, types.DictType):
act = dict(degrees=act)
env.act(act)
rew, _ob, first = env.observe()
if space_type == "real_dict":
assert rew[0] == 0
assert rew[1] < 0
else:
assert np.array_equal(rew, np.array([1, 0]))
assert np.array_equal(first, np.array([1, 0]))
assert env.get_info() == [{"id": 1}, {"id": 2}]
with pytest.raises(AssertionError):
env.callmethod("f", [2, 3, 4])
assert env.callmethod("f", [2, 3]) == [2**2, 2 * 3]
def test_recorder():
with tempfile.TemporaryDirectory() as tmpdir:
env = IdentityEnv(
space=types.TensorType(eltype=types.Discrete(256), shape=(64, 64, 3))
)
writer_kwargs = {
"codec": "libx264rgb",
"pixelformat": "bgr24",
"output_params": ["-crf", "0"],
}
env = VideoRecorderWrapper(
env=env, directory=tmpdir, env_index=0, writer_kwargs=writer_kwargs
)
_, obs, _ = env.observe()
for _ in range(2):
env.act(types_np.zeros(env.ac_space, bshape=(env.num,)))
video_files = sorted(glob(os.path.join(tmpdir, "*.mp4")))
assert len(video_files) > 0
with imageio.get_reader(video_files[0]) as r:
for im in r:
assert np.allclose(im, obs[0])
break
def create_test_env_exit_on_act():
env = IdentityEnv()
env = ExitOnAct(env)
return env
def create_test_env_exception_on_act():
env = IdentityEnv()
env = ExceptionOnAct(env)
return env
def create_callmethod_test_env():
env = IdentityEnv()
env.square = lambda x: x**2
return env
def create_test_env(episode_len):
env = IdentityEnv(episode_len=episode_len)
env = StepInfo(env)
return env
def make_env():
env = IdentityEnv()
env.square = lambda x: x**2
return env
def create_test_env_mpi(episode_len):
# this is necessary for the MPI test to fail
from mpi4py import MPI
env = IdentityEnv(episode_len=episode_len)
return env