def test_in_local_sampler(policy, envs):
true_workers = WorkerFactory(seed=100,
n_workers=N_TRAJ,
max_path_length=MAX_PATH_LENGTH)
true_sampler = LocalSampler.from_worker_factory(true_workers, policy, envs)
vec_workers = WorkerFactory(seed=100,
n_workers=1,
worker_class=VecWorker,
worker_args=dict(n_envs=N_TRAJ),
max_path_length=MAX_PATH_LENGTH)
vec_sampler = LocalSampler.from_worker_factory(vec_workers, policy, [envs])
n_samples = 100
true_trajs = true_sampler.obtain_samples(0, n_samples, None)
vec_trajs = vec_sampler.obtain_samples(0, n_samples, None)
assert vec_trajs.lengths.sum() >= n_samples
assert_trajs_eq(true_trajs, vec_trajs)
# Test start_rollout optimization
true_trajs = true_sampler.obtain_samples(0, n_samples, None)
vec_trajs = vec_sampler.obtain_samples(0, n_samples, None)
assert vec_trajs.lengths.sum() >= n_samples
assert_trajs_eq(true_trajs, vec_trajs)
true_sampler.shutdown_worker()
vec_sampler.shutdown_worker()
def test_local_batch_sampler(self):
workers = WorkerFactory(seed=100,
max_path_length=self.algo.max_path_length)
sampler1 = LocalSampler.from_worker_factory(workers, self.policy,
self.env)
sampler2 = OnPolicyVectorizedSampler(self.algo, self.env)
sampler2.start_worker()
trajs1 = sampler1.obtain_samples(
0, 1000, tuple(self.algo.policy.get_param_values()))
trajs2 = sampler2.obtain_samples(0, 1000)
# pylint: disable=superfluous-parens
assert trajs1.observations.shape[0] >= 1000
assert trajs1.actions.shape[0] >= 1000
assert (sum(trajs1.rewards[:trajs1.lengths[0]]) == sum(
trajs2[0]['rewards']) == 1)
true_obs = np.array([0, 1, 2, 6, 10, 14])
true_actions = np.array([2, 2, 1, 1, 1, 2])
true_rewards = np.array([0, 0, 0, 0, 0, 1])
start = 0
for length in trajs1.lengths:
observations = trajs1.observations[start:start + length]
actions = trajs1.actions[start:start + length]
rewards = trajs1.rewards[start:start + length]
assert np.array_equal(observations, true_obs)
assert np.array_equal(actions, true_actions)
assert np.array_equal(rewards, true_rewards)
start += length
sampler1.shutdown_worker()
sampler2.shutdown_worker()
def test_update_envs_env_update():
max_path_length = 16
env = TfEnv(PointEnv())
policy = FixedPolicy(env.spec,
scripted_actions=[
env.action_space.sample()
for _ in range(max_path_length)
])
tasks = SetTaskSampler(PointEnv)
n_workers = 8
workers = WorkerFactory(seed=100,
max_path_length=max_path_length,
n_workers=n_workers)
sampler = LocalSampler.from_worker_factory(workers, policy, env)
rollouts = sampler.obtain_samples(0,
161,
np.asarray(policy.get_param_values()),
env_update=tasks.sample(n_workers))
mean_rewards = []
goals = []
for rollout in rollouts.split():
mean_rewards.append(rollout.rewards.mean())
goals.append(rollout.env_infos['task'][0]['goal'])
assert np.var(mean_rewards) > 0
assert np.var(goals) > 0
with pytest.raises(ValueError):
sampler.obtain_samples(0,
10,
np.asarray(policy.get_param_values()),
env_update=tasks.sample(n_workers + 1))
def test_pickle():
max_path_length = 16
env = MetaRLEnv(PointEnv())
policy = FixedPolicy(env.spec,
scripted_actions=[
env.action_space.sample()
for _ in range(max_path_length)
])
tasks = SetTaskSampler(PointEnv)
n_workers = 8
workers = WorkerFactory(seed=100,
max_path_length=max_path_length,
n_workers=n_workers)
sampler = MultiprocessingSampler.from_worker_factory(workers, policy, env)
sampler_pickled = pickle.dumps(sampler)
sampler.shutdown_worker()
sampler2 = pickle.loads(sampler_pickled)
rollouts = sampler2.obtain_samples(0,
161,
np.asarray(policy.get_param_values()),
env_update=tasks.sample(n_workers))
mean_rewards = []
goals = []
for rollout in rollouts.split():
mean_rewards.append(rollout.rewards.mean())
goals.append(rollout.env_infos['task'][0]['goal'])
assert np.var(mean_rewards) > 0
assert np.var(goals) > 0
sampler2.shutdown_worker()
env.close()
def test_obtain_exact_trajectories():
max_path_length = 15
n_workers = 8
env = TfEnv(PointEnv())
per_worker_actions = [env.action_space.sample() for _ in range(n_workers)]
policies = [
FixedPolicy(env.spec, [action] * max_path_length)
for action in per_worker_actions
]
workers = WorkerFactory(seed=100,
max_path_length=max_path_length,
n_workers=n_workers)
sampler = LocalSampler.from_worker_factory(workers, policies, envs=env)
n_traj_per_worker = 3
rollouts = sampler.obtain_exact_trajectories(n_traj_per_worker,
agent_update=policies)
# At least one action per trajectory.
assert sum(rollouts.lengths) >= n_workers * n_traj_per_worker
# All of the trajectories.
assert len(rollouts.lengths) == n_workers * n_traj_per_worker
worker = -1
for count, rollout in enumerate(rollouts.split()):
if count % n_traj_per_worker == 0:
worker += 1
assert (rollout.actions == per_worker_actions[worker]).all()
def test_init_with_crashed_worker():
max_path_length = 16
env = MetaRLEnv(PointEnv())
policy = FixedPolicy(env.spec,
scripted_actions=[
env.action_space.sample()
for _ in range(max_path_length)
])
tasks = SetTaskSampler(lambda: MetaRLEnv(PointEnv()))
n_workers = 2
workers = WorkerFactory(seed=100,
max_path_length=max_path_length,
n_workers=n_workers)
class CrashingPolicy:
def reset(self, **kwargs):
raise Exception('Intentional subprocess crash')
bad_policy = CrashingPolicy()
# This causes worker 2 to crash.
sampler = MultiprocessingSampler.from_worker_factory(
workers, [policy, bad_policy], envs=tasks.sample(n_workers))
rollouts = sampler.obtain_samples(0, 160, None)
assert sum(rollouts.lengths) >= 160
sampler.shutdown_worker()
env.close()
def test_ray_batch_sampler(self, ray_local_session_fixture):
del ray_local_session_fixture
assert ray.is_initialized()
workers = WorkerFactory(seed=100,
max_path_length=self.algo.max_path_length)
sampler1 = RaySampler(workers, self.policy, self.env)
sampler1.start_worker()
sampler1.shutdown_worker()
def test_init_with_env_updates(policy, envs):
task_sampler = EnvPoolSampler(envs)
envs = task_sampler.sample(N_TRAJ)
true_workers = WorkerFactory(seed=100,
n_workers=N_TRAJ,
max_path_length=MAX_PATH_LENGTH)
true_sampler = LocalSampler.from_worker_factory(true_workers, policy, envs)
vec_workers = WorkerFactory(seed=100,
n_workers=1,
worker_class=VecWorker,
worker_args=dict(n_envs=N_TRAJ),
max_path_length=MAX_PATH_LENGTH)
vec_sampler = LocalSampler.from_worker_factory(vec_workers, [policy],
[envs])
n_samples = 100
true_trajs = true_sampler.obtain_samples(0, n_samples, None)
vec_trajs = vec_sampler.obtain_samples(0, n_samples, None)
assert vec_trajs.lengths.sum() >= n_samples
assert_trajs_eq(true_trajs, vec_trajs)
true_sampler.shutdown_worker()
vec_sampler.shutdown_worker()
def test_reset_optimization(policy, envs, other_envs):
true_workers = WorkerFactory(seed=100,
n_workers=N_TRAJ,
max_path_length=MAX_PATH_LENGTH)
true_sampler = LocalSampler.from_worker_factory(true_workers, policy, envs)
vec_workers = WorkerFactory(seed=100,
n_workers=1,
worker_class=VecWorker,
worker_args=dict(n_envs=N_TRAJ),
max_path_length=MAX_PATH_LENGTH)
vec_sampler = LocalSampler.from_worker_factory(vec_workers, [policy],
[envs])
n_samples = 4 * MAX_PATH_LENGTH
true_sampler.obtain_samples(0, n_samples, None)
true_sampler.obtain_samples(0, n_samples, None)
true_trajs = true_sampler.obtain_samples(0, n_samples, None, other_envs)
vec_trajs = vec_sampler.obtain_samples(0, n_samples, None, [other_envs])
assert vec_trajs.lengths.sum() >= n_samples
assert_trajs_eq(true_trajs, vec_trajs)
true_sampler.shutdown_worker()
vec_sampler.shutdown_worker()
def evaluate(self, algo, test_rollouts_per_task=None):
"""Evaluate the Meta-RL algorithm on the test tasks.
Args:
algo (metarl.np.algos.MetaRLAlgorithm): The algorithm to evaluate.
test_rollouts_per_task (int or None): Number of rollouts per task.
"""
if test_rollouts_per_task is None:
test_rollouts_per_task = self._n_test_rollouts
adapted_trajectories = []
logger.log('Sampling for adapation and meta-testing...')
if self._test_sampler is None:
self._test_sampler = LocalSampler.from_worker_factory(
WorkerFactory(seed=get_seed(),
max_path_length=self._max_path_length,
n_workers=1,
worker_class=self._worker_class,
worker_args=self._worker_args),
agents=algo.get_exploration_policy(),
envs=self._test_task_sampler.sample(1))
for env_up in self._test_task_sampler.sample(self._n_test_tasks):
policy = algo.get_exploration_policy()
traj = TrajectoryBatch.concatenate(*[
self._test_sampler.obtain_samples(self._eval_itr, 1, policy,
env_up)
for _ in range(self._n_exploration_traj)
])
adapted_policy = algo.adapt_policy(policy, traj)
adapted_traj = self._test_sampler.obtain_samples(
self._eval_itr, test_rollouts_per_task * self._max_path_length,
adapted_policy)
adapted_trajectories.append(adapted_traj)
logger.log('Finished meta-testing...')
if self._test_task_names is not None:
name_map = dict(enumerate(self._test_task_names))
else:
name_map = None
with tabular.prefix(self._prefix + '/' if self._prefix else ''):
log_multitask_performance(
self._eval_itr,
TrajectoryBatch.concatenate(*adapted_trajectories),
getattr(algo, 'discount', 1.0),
name_map=name_map)
self._eval_itr += 1
def test_init_with_env_updates():
max_path_length = 16
env = TfEnv(PointEnv())
policy = FixedPolicy(env.spec,
scripted_actions=[
env.action_space.sample()
for _ in range(max_path_length)
])
tasks = SetTaskSampler(lambda: TfEnv(PointEnv()))
n_workers = 8
workers = WorkerFactory(seed=100,
max_path_length=max_path_length,
n_workers=n_workers)
sampler = LocalSampler.from_worker_factory(workers,
policy,
envs=tasks.sample(n_workers))
rollouts = sampler.obtain_samples(0, 160, policy)
assert sum(rollouts.lengths) >= 160
def test_init_with_env_updates(ray_local_session_fixture):
del ray_local_session_fixture
assert ray.is_initialized()
max_path_length = 16
env = MetaRLEnv(PointEnv())
policy = FixedPolicy(env.spec,
scripted_actions=[
env.action_space.sample()
for _ in range(max_path_length)
])
tasks = SetTaskSampler(lambda: MetaRLEnv(PointEnv()))
n_workers = 8
workers = WorkerFactory(seed=100,
max_path_length=max_path_length,
n_workers=n_workers)
sampler = RaySampler.from_worker_factory(workers,
policy,
envs=tasks.sample(n_workers))
rollouts = sampler.obtain_samples(0, 160, policy)
assert sum(rollouts.lengths) >= 160
def test_adapt_policy(self, set_params, test_params):
"""Test if policy can adapt to samples."""
worker = WorkerFactory(seed=100, max_path_length=100)
sampler = LocalSampler.from_worker_factory(worker, self.policy,
self.env)
self.policy.apply(partial(set_params, 0.1))
adapt_policy = self.algo.get_exploration_policy()
trajs = sampler.obtain_samples(0, 100, adapt_policy)
self.algo.adapt_policy(adapt_policy, trajs)
# Old policy should remain untouched
self.policy.apply(partial(test_params, 0.1))
# Adapted policy should not be identical to old policy
for v1, v2 in zip(adapt_policy.parameters(), self.policy.parameters()):
if v1.data.ne(v2.data).sum() > 0:
break
else:
pytest.fail("Parameters of adapted policy should not be "
"identical to the old policy.")
def test_obtain_samples():
env = MetaRLEnv(GridWorldEnv(desc='4x4'))
policy = ScriptedPolicy(
scripted_actions=[2, 2, 1, 0, 3, 1, 1, 1, 2, 2, 1, 1, 1, 2, 2, 1])
algo = Mock(env_spec=env.spec, policy=policy, max_path_length=16)
workers = WorkerFactory(seed=100,
max_path_length=algo.max_path_length,
n_workers=8)
sampler1 = MultiprocessingSampler.from_worker_factory(workers, policy, env)
sampler2 = LocalSampler.from_worker_factory(workers, policy, env)
trajs1 = sampler1.obtain_samples(0, 1000,
tuple(algo.policy.get_param_values()))
trajs2 = sampler2.obtain_samples(0, 1000,
tuple(algo.policy.get_param_values()))
# pylint: disable=superfluous-parens
assert trajs1.observations.shape[0] >= 1000
assert trajs1.actions.shape[0] >= 1000
assert (sum(trajs1.rewards[:trajs1.lengths[0]]) == sum(
trajs2.rewards[:trajs2.lengths[0]]) == 1)
true_obs = np.array([0, 1, 2, 6, 10, 14])
true_actions = np.array([2, 2, 1, 1, 1, 2])
true_rewards = np.array([0, 0, 0, 0, 0, 1])
start = 0
for length in trajs1.lengths:
observations = trajs1.observations[start:start + length]
actions = trajs1.actions[start:start + length]
rewards = trajs1.rewards[start:start + length]
assert np.array_equal(observations, true_obs)
assert np.array_equal(actions, true_actions)
assert np.array_equal(rewards, true_rewards)
start += length
sampler1.shutdown_worker()
sampler2.shutdown_worker()
env.close()
def test_ray_batch_sampler(ray_local_session_fixture):
del ray_local_session_fixture
env = MetaRLEnv(GridWorldEnv(desc='4x4'))
policy = ScriptedPolicy(
scripted_actions=[2, 2, 1, 0, 3, 1, 1, 1, 2, 2, 1, 1, 1, 2, 2, 1])
algo = Mock(env_spec=env.spec, policy=policy, max_path_length=16)
assert ray.is_initialized()
workers = WorkerFactory(seed=100, max_path_length=algo.max_path_length)
sampler1 = RaySampler(workers, policy, env)
sampler1.start_worker()
sampler2 = OnPolicyVectorizedSampler(algo, env)
sampler2.start_worker()
trajs1 = sampler1.obtain_samples(0, 1000,
tuple(algo.policy.get_param_values()))
trajs2 = sampler2.obtain_samples(0, 1000)
# pylint: disable=superfluous-parens
assert trajs1.observations.shape[0] >= 1000
assert trajs1.actions.shape[0] >= 1000
assert (sum(trajs1.rewards[:trajs1.lengths[0]]) == sum(
trajs2[0]['rewards']) == 1)
true_obs = np.array([0, 1, 2, 6, 10, 14])
true_actions = np.array([2, 2, 1, 1, 1, 2])
true_rewards = np.array([0, 0, 0, 0, 0, 1])
start = 0
for length in trajs1.lengths:
observations = trajs1.observations[start:start + length]
actions = trajs1.actions[start:start + length]
rewards = trajs1.rewards[start:start + length]
assert np.array_equal(observations, true_obs)
assert np.array_equal(actions, true_actions)
assert np.array_equal(rewards, true_rewards)
start += length
sampler1.shutdown_worker()
sampler2.shutdown_worker()
env.close()
def test_ray_batch_sampler(self):
workers = WorkerFactory(seed=100,
max_path_length=self.algo.max_path_length)
sampler1 = RaySampler(workers, self.policy, self.env)
sampler1.start_worker()
sampler1.shutdown_worker()