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 learn_cartpole():
"""Train an agent."""
env = gym.make('CartPole-v0')
try:
agent = ActorCritic(gym_space_distribution(env.action_space),
gym_space_vectorizer(env.observation_space))
with tf.Session() as sess:
a2c = A2C(sess, agent, target_kl=0.03)
roller = BasicRoller(env, agent, min_episodes=8, min_steps=1024)
while True:
with agent.frozen():
rollouts = roller.rollouts()
print('mean=%f' % (mean_total_reward(rollouts), ))
agent.actor.extend(
a2c.policy_update(rollouts,
STEP_SIZE,
NUM_STEPS,
min_leaf=30))
agent.critic.extend(
a2c.value_update(rollouts,
VAL_STEP,
NUM_STEPS,
min_leaf=30))
finally:
env.close()
def test_truncation(stateful, state_tuple):
"""
Test sequence truncation for TruncatedRoller with a
batch of one environment.
"""
def env_fn():
return SimpleEnv(7, (5, 3), 'uint8')
env = env_fn()
model = SimpleModel(env.action_space.low.shape,
stateful=stateful,
state_tuple=state_tuple)
basic_roller = BasicRoller(env, model, min_episodes=5)
expected = basic_roller.rollouts()
total_timesteps = sum([x.num_steps for x in expected])
batched_env = batched_gym_env([env_fn], sync=True)
trunc_roller = TruncatedRoller(batched_env, model,
total_timesteps // 2 + 1)
actual1 = trunc_roller.rollouts()
assert actual1[-1].trunc_end
actual2 = trunc_roller.rollouts()
expected1, expected2 = _artificial_truncation(expected,
len(actual1) - 1,
actual1[-1].num_steps)
assert len(actual2) == len(expected2) + 1
actual2 = actual2[:-1]
_compare_rollout_batch(actual1, expected1)
_compare_rollout_batch(actual2, expected2)
def run_ppo():
"""
Run a training worker.
"""
env = gym.make('CartPole-v0')
action_dist = gym_space_distribution(env.action_space)
obs_vectorizer = gym_space_vectorizer(env.observation_space)
with tf.Session() as sess:
model = MLP(sess, action_dist, obs_vectorizer, layer_sizes=[32])
# Deal with CartPole-v0 reward scale.
model.scale_outputs(20)
roller = BasicRoller(env, model, min_episodes=30)
ppo = PPO(model)
optimizer = MPIOptimizer(tf.train.AdamOptimizer(learning_rate=1e-3),
-ppo.objective)
sess.run(tf.global_variables_initializer())
optimizer.sync_from_root(sess)
for i in range(50):
rollouts = roller.rollouts()
# pylint: disable=E1101
print('batch %d: rank=%d mean=%f' %
(i, MPI.COMM_WORLD.Get_rank(), mean_total_reward(rollouts)))
mpi_ppo(ppo, optimizer, rollouts, log_fn=print)
def _test_truncation_case(self, stateful, state_tuple):
"""
Test rollout truncation and continuation for a
specific set of model parameters.
"""
env_fn = lambda: SimpleEnv(7, (5, 3), 'uint8')
env = env_fn()
model = SimpleModel(env.action_space.low.shape,
stateful=stateful,
state_tuple=state_tuple)
basic_roller = BasicRoller(env, model, min_episodes=5)
expected = basic_roller.rollouts()
total_timesteps = sum([x.num_steps for x in expected])
batched_env = batched_gym_env([env_fn], sync=True)
trunc_roller = TruncatedRoller(batched_env, model,
total_timesteps // 2 + 1)
actual1 = trunc_roller.rollouts()
self.assertTrue(actual1[-1].trunc_end)
actual2 = trunc_roller.rollouts()
expected1, expected2 = _artificial_truncation(expected,
len(actual1) - 1,
actual1[-1].num_steps)
self.assertEqual(len(actual2), len(expected2) + 1)
actual2 = actual2[:-1]
_compare_rollout_batch(self, actual1, expected1)
_compare_rollout_batch(self, actual2, expected2)
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):
"""
Test BasicRoller equivalence for a specific set of
model settings.
"""
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, min_episodes=5)
expected = basic_roller.rollouts()
total_timesteps = sum([x.num_steps for x in expected])
batched_env = batched_gym_env([env_fn], sync=True)
trunc_roller = TruncatedRoller(batched_env, model, total_timesteps)
actual = trunc_roller.rollouts()
_compare_rollout_batch(self, 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_batches_consistency(self, batch_size, trunc_start):
"""
Make sure that batches() produces the same outputs
that we got with step().
"""
env = TupleCartPole()
try:
roller = BasicRoller(env, self.model, min_episodes=7)
rollouts = roller.rollouts()
if trunc_start:
rollouts = self._truncate_first(rollouts)
num_batches = 10
for batch in self.model.batches(rollouts, batch_size=batch_size):
num_batches -= 1
if num_batches == 0:
break
self._test_batch(rollouts, batch)
finally:
env.close()
def test_trunc_basic_equivalence(stateful, state_tuple):
"""
Test that TruncatedRoller is equivalent to BasicRoller
for batches of one environment when the episodes end
cleanly.
"""
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, min_episodes=5)
expected = basic_roller.rollouts()
total_timesteps = sum([x.num_steps for x in expected])
batched_env = batched_gym_env([env_fn], sync=True)
trunc_roller = TruncatedRoller(batched_env, model, total_timesteps)
actual = trunc_roller.rollouts()
_compare_rollout_batch(actual, expected)
def run_algorithm(algo_name):
"""
Run the specified training algorithm.
"""
env = gym.make('CartPole-v0')
action_dist = gym_space_distribution(env.action_space)
obs_vectorizer = gym_space_vectorizer(env.observation_space)
with tf.Session() as sess:
model = MLP(sess, action_dist, obs_vectorizer, layer_sizes=[32])
# Deal with CartPole-v0 reward scale.
model.scale_outputs(20)
roller = BasicRoller(env, model, min_episodes=30)
inner_loop = algorithm_inner_loop(algo_name, model)
sess.run(tf.global_variables_initializer())
print('running algorithm:', algo_name)
for i in range(50):
rollouts = roller.rollouts()
print('batch %d: mean=%f' % (i, mean_total_reward(rollouts)))
inner_loop(rollouts)
if __name__ == '__main__':
monitor = "results/unbiased_random/1"
action_repeat = True
single_life = True
render = None
env = retro.make("SuperMarioBros-Nes")
env = MarioDiscretizer(env)
if single_life:
env = SingleLifeEnv(env)
if monitor is not None:
env = Monitor(env, monitor, video_callable=lambda i: False)
if render is not None:
env = AutoRenderer(env, auto_render_period=render)
if action_repeat:
env = FrameStack(env, 4)
# model = WeightedRandomAgent()
model = RandomAgent(lambda: env.action_space.sample())
player = BasicRoller(env, model, min_episodes=1)
# total_rollouts = [player.rollouts() for rollout_i in trange(40)]
# flat_rollouts = reduce(list.__add__, total_rollouts)
# total_rewards = map(lambda r: r.total_reward, flat_rollouts)
# [filename for path in dirs for filename in os.listdir(path)]
total_rewards = []
for i in tqdm(range(150)):
rollouts = player.rollouts()
total_rewards += [roll.total_reward for roll in rollouts]
print(total_rewards)
rewards_numbers = list(zip(count(), total_rewards))
sorted_reward_numbers = sorted(rewards_numbers, key=lambda t: t[1])
print(sorted_reward_numbers[-5:])
