def test_ep_batches(stateful, state_tuple, limits):
"""
Test that EpisodeRoller is equivalent to a
BasicRoller when run on a batch of envs.
"""
def env_fn():
return SimpleEnv(3, (4, 5), 'uint8')
model = SimpleModel((4, 5), stateful=stateful, state_tuple=state_tuple)
batched_env = batched_gym_env([env_fn] * 21, num_sub_batches=7, sync=True)
ep_roller = EpisodeRoller(batched_env, model, **limits)
actual = ep_roller.rollouts()
total_steps = sum([r.num_steps for r in actual])
assert len(actual) >= ep_roller.min_episodes
assert total_steps >= ep_roller.min_steps
if 'min_steps' not in limits:
num_eps = ep_roller.min_episodes + batched_env.num_envs - 1
assert len(actual) == num_eps
basic_roller = BasicRoller(env_fn(), model, min_episodes=len(actual))
expected = basic_roller.rollouts()
_compare_rollout_batch(actual, expected)
def _test_batch_equivalence_case(self, stateful, state_tuple,
**roller_kwargs):
"""
Test BasicRoller equivalence when using a batch of
environments.
"""
env_fn = lambda: SimpleEnv(3, (4, 5), 'uint8')
model = SimpleModel((4, 5), stateful=stateful, state_tuple=state_tuple)
batched_env = batched_gym_env([env_fn] * 21,
num_sub_batches=7,
sync=True)
ep_roller = EpisodeRoller(batched_env, model, **roller_kwargs)
actual = ep_roller.rollouts()
total_steps = sum([r.num_steps for r in actual])
self.assertTrue(len(actual) >= ep_roller.min_episodes)
self.assertTrue(total_steps >= ep_roller.min_steps)
if 'min_steps' not in roller_kwargs:
num_eps = ep_roller.min_episodes + batched_env.num_envs - 1
self.assertTrue(len(actual) == num_eps)
basic_roller = BasicRoller(env_fn(), model, min_episodes=len(actual))
expected = basic_roller.rollouts()
_compare_rollout_batch(self, actual, expected)
def test_ep_basic_equivalence(stateful, state_tuple, limits):
"""
Test that EpisodeRoller is equivalent to a
BasicRoller when run on a single environment.
"""
env_fn = lambda: SimpleEnv(3, (4, 5), 'uint8')
env = env_fn()
model = SimpleModel(env.action_space.low.shape,
stateful=stateful,
state_tuple=state_tuple)
basic_roller = BasicRoller(env, model, **limits)
expected = basic_roller.rollouts()
batched_env = batched_gym_env([env_fn], sync=True)
ep_roller = EpisodeRoller(batched_env, model, **limits)
actual = ep_roller.rollouts()
_compare_rollout_batch(actual, expected)
def _test_basic_equivalence_case(self, stateful, state_tuple,
**roller_kwargs):
"""
Test BasicRoller equivalence for a single env in a
specific case.
"""
env_fn = lambda: SimpleEnv(3, (4, 5), 'uint8')
env = env_fn()
model = SimpleModel(env.action_space.low.shape,
stateful=stateful,
state_tuple=state_tuple)
basic_roller = BasicRoller(env, model, **roller_kwargs)
expected = basic_roller.rollouts()
batched_env = batched_gym_env([env_fn], sync=True)
ep_roller = EpisodeRoller(batched_env, model, **roller_kwargs)
actual = ep_roller.rollouts()
_compare_rollout_batch(self, actual, expected)
def test_multiple_batches(self):
"""
Make sure calling rollouts multiple times works.
"""
env_fn = lambda: SimpleEnv(3, (4, 5), 'uint8')
env = env_fn()
try:
model = SimpleModel(env.action_space.low.shape)
finally:
env.close()
batched_env = batched_gym_env([env_fn], sync=True)
try:
ep_roller = EpisodeRoller(batched_env,
model,
min_episodes=5,
min_steps=7)
first = ep_roller.rollouts()
for _ in range(3):
_compare_rollout_batch(self, first, ep_roller.rollouts())
finally:
batched_env.close()