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...')
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...')
with tabular.prefix(self._prefix + '/' if self._prefix else ''):
log_multitask_performance(
self._eval_itr,
TrajectoryBatch.concatenate(*adapted_trajectories),
getattr(algo, 'discount', 1.0),
task_names=self._test_task_names)
self._eval_itr += 1
def _obtain_evaluation_samples(self,
env,
num_trajs=100,
max_path_length=1000):
"""Sample the policy for 10 trajectories and return average values.
Args:
env (metarl.envs.MetaRLEnv): The environement used to obtain
trajectories.
num_trajs (int): Number of trajectories.
max_path_length (int): Number of maximum steps in one batch.
Returns:
TrajectoryBatch: Evaluation trajectories, representing the best
current performance of the algorithm.
"""
paths = []
for _ in range(num_trajs):
path = rollout(env,
self.policy,
max_path_length=max_path_length,
deterministic=True)
paths.append(path)
return TrajectoryBatch.from_trajectory_list(self.env_spec, paths)
def collect_rollout(self):
"""Collect the current rollout, clearing the internal buffer.
Returns:
metarl.TrajectoryBatch: A batch of the trajectories completed since
the last call to collect_rollout().
"""
observations = self._observations
self._observations = []
last_observations = self._last_observations
self._last_observations = []
actions = self._actions
self._actions = []
rewards = self._rewards
self._rewards = []
terminals = self._terminals
self._terminals = []
env_infos = self._env_infos
self._env_infos = defaultdict(list)
agent_infos = self._agent_infos
self._agent_infos = defaultdict(list)
for k, v in agent_infos.items():
agent_infos[k] = np.asarray(v)
for k, v in env_infos.items():
env_infos[k] = np.asarray(v)
lengths = self._lengths
self._lengths = []
return TrajectoryBatch(self.env.spec, np.asarray(observations),
np.asarray(last_observations),
np.asarray(actions), np.asarray(rewards),
np.asarray(terminals), dict(env_infos),
dict(agent_infos), np.asarray(lengths,
dtype='i'))
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_act_box_env_spec_mismatch_traj(traj_data):
with pytest.raises(ValueError, match='actions should have'):
traj_data['env_spec'].action_space = akro.Box(low=1,
high=np.inf,
shape=(4, 3, 2),
dtype=np.float32)
t = TrajectoryBatch(**traj_data)
del t
def test_agent_infos_batch_mismatch_traj(traj_data):
with pytest.raises(
ValueError,
match='entry in agent_infos must have a batch dimension'):
traj_data['agent_infos']['hidden'] = traj_data['agent_infos'][
'hidden'][:-1]
t = TrajectoryBatch(**traj_data)
del t
def test_to_trajectory_list(traj_data):
t = TrajectoryBatch(**traj_data)
t_list = t.to_trajectory_list()
assert len(t_list) == len(traj_data['lengths'])
start = 0
for length, last_obs, s in zip(traj_data['lengths'],
traj_data['last_observations'], t_list):
stop = start + length
assert (
s['observations'] == traj_data['observations'][start:stop]).all()
assert (s['next_observations'] == np.concatenate(
(traj_data['observations'][start + 1:stop], [last_obs]))).all()
assert (s['actions'] == traj_data['actions'][start:stop]).all()
assert (s['rewards'] == traj_data['rewards'][start:stop]).all()
assert (s['dones'] == traj_data['terminals'][start:stop]).all()
start = stop
assert start == len(traj_data['rewards'])
def train_once(self, itr, paths):
"""Perform one step of policy optimization given one batch of samples.
Args:
itr (int): Iteration number.
paths (list[dict]): A list of collected paths.
Returns:
float: The average return in last epoch cycle.
"""
# -- Stage: Calculate baseline
if hasattr(self._baseline, 'predict_n'):
baseline_predictions = self._baseline.predict_n(paths)
else:
baseline_predictions = [
self._baseline.predict(path) for path in paths
]
# -- Stage: Pre-process samples based on collected paths
samples_data = paths_to_tensors(paths, self.max_path_length,
baseline_predictions, self._discount)
# -- Stage: Run and calculate performance of the algorithm
undiscounted_returns = log_performance(
itr,
TrajectoryBatch.from_trajectory_list(self._env_spec, paths),
discount=self._discount)
self._episode_reward_mean.extend(undiscounted_returns)
tabular.record('Extras/EpisodeRewardMean',
np.mean(self._episode_reward_mean))
samples_data['average_return'] = np.mean(undiscounted_returns)
epoch = itr // self._n_samples
i_sample = itr - epoch * self._n_samples
tabular.record('Epoch', epoch)
tabular.record('# Sample', i_sample)
rtn = samples_data['average_return']
self._all_returns.append(samples_data['average_return'])
if (itr + 1) % self._n_samples == 0:
avg_rtns = np.array(self._all_returns)
self._es.tell(self._all_params, -avg_rtns)
self.policy.set_param_values(self._es.best.get()[0])
# Clear for next epoch
rtn = max(self._all_returns)
self._all_returns.clear()
self._all_params = self._sample_params()
self._cur_params = self._all_params[(i_sample + 1) % self._n_samples]
self.policy.set_param_values(self._cur_params)
logger.log(tabular)
return rtn
def test_new_traj(traj_data):
t = TrajectoryBatch(**traj_data)
assert t.env_spec is traj_data['env_spec']
assert t.observations is traj_data['observations']
assert t.last_observations is traj_data['last_observations']
assert t.actions is traj_data['actions']
assert t.rewards is traj_data['rewards']
assert t.terminals is traj_data['terminals']
assert t.env_infos is traj_data['env_infos']
assert t.agent_infos is traj_data['agent_infos']
assert t.lengths is traj_data['lengths']
def test_log_multitask_performance_task_id():
lengths = np.array([10, 5, 1, 1])
batch = TrajectoryBatch(
EnvSpec(akro.Box(np.array([0., 0., 0.]), np.array([1., 1., 1.])),
akro.Box(np.array([-1., -1.]), np.array([0., 0.]))),
observations=np.ones((sum(lengths), 3), dtype=np.float32),
last_observations=np.ones((len(lengths), 3), dtype=np.float32),
actions=np.zeros((sum(lengths), 2), dtype=np.float32),
rewards=np.array([
0.34026529, 0.58263177, 0.84307509, 0.97651095, 0.81723901,
0.22631398, 0.03421301, 0.97515046, 0.64311832, 0.65068933,
0.17657714, 0.04783857, 0.73904013, 0.41364329, 0.52235551,
0.24203526, 0.43328910
]),
terminals=np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 1],
dtype=bool),
env_infos={
'success':
np.array([0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 1, 0, 1],
dtype=bool),
'task_id':
np.array([1] * 10 + [3] * 5 + [1] + [4])
},
agent_infos={},
lengths=lengths)
log_file = tempfile.NamedTemporaryFile()
csv_output = dowel.CsvOutput(log_file.name)
logger.add_output(csv_output)
log_multitask_performance(7, batch, 0.8, {
1: 'env1',
3: 'env2',
4: 'env3',
5: 'env4'
})
logger.log(tabular)
logger.dump_output_type(dowel.CsvOutput)
with open(log_file.name, 'r') as file:
rows = list(csv.DictReader(file))
res = {k: float(r) for (k, r) in rows[0].items()}
assert res['env1/Iteration'] == 7
assert res['env2/Iteration'] == 7
assert res['env3/Iteration'] == 7
assert res['env4/Iteration'] == 7
assert res['env1/NumTrajs'] == 2
assert res['env2/NumTrajs'] == 1
assert res['env3/NumTrajs'] == 1
assert res['env4/NumTrajs'] == 0
assert math.isclose(res['env1/SuccessRate'], 0.5)
assert math.isclose(res['env2/SuccessRate'], 1.0)
assert math.isclose(res['env3/SuccessRate'], 1.0)
assert math.isnan(res['env4/SuccessRate'])
assert math.isnan(res['env4/AverageReturn'])
def train_once(self, itr, paths):
"""Perform one step of policy optimization given one batch of samples.
Args:
itr (int): Iteration number.
paths (list[dict]): A list of collected paths.
Returns:
numpy.float64: Average return.
"""
# -- Stage: Calculate baseline
paths = [
dict(
observations=self._env_spec.observation_space.flatten_n(
path['observations'])
if self._flatten_input else path['observations'],
actions=(
self._env_spec.action_space.flatten_n( # noqa: E126
path['actions'])),
rewards=path['rewards'],
env_infos=path['env_infos'],
agent_infos=path['agent_infos'],
dones=path['dones']) for path in paths
]
if hasattr(self._baseline, 'predict_n'):
baseline_predictions = self._baseline.predict_n(paths)
else:
baseline_predictions = [
self._baseline.predict(path) for path in paths
]
# -- Stage: Pre-process samples based on collected paths
samples_data = paths_to_tensors(paths, self.max_path_length,
baseline_predictions, self._discount,
self._gae_lambda)
# -- Stage: Run and calculate performance of the algorithm
undiscounted_returns = log_performance(
itr,
TrajectoryBatch.from_trajectory_list(self._env_spec, paths),
discount=self._discount)
self._episode_reward_mean.extend(undiscounted_returns)
tabular.record('Extras/EpisodeRewardMean',
np.mean(self._episode_reward_mean))
samples_data['average_return'] = np.mean(undiscounted_returns)
self.log_diagnostics(samples_data)
logger.log('Optimizing policy...')
self.optimize_policy(samples_data)
return samples_data['average_return']
def collect_rollout(self):
"""Collect the current rollout, clearing the internal buffer.
One-hot task id is saved in env_infos['task_onehot']. Latent is saved
in agent_infos['latent']. Latent infos are saved in
agent_infos['latent_info_name'], where info_name is the original latent
info name.
Returns:
metarl.TrajectoryBatch: A batch of the trajectories completed since
the last call to collect_rollout().
"""
observations = self._observations
self._observations = []
last_observations = self._last_observations
self._last_observations = []
actions = self._actions
self._actions = []
rewards = self._rewards
self._rewards = []
terminals = self._terminals
self._terminals = []
latents = self._latents
self._latents = []
tasks = self._tasks
self._tasks = []
env_infos = self._env_infos
self._env_infos = defaultdict(list)
agent_infos = self._agent_infos
self._agent_infos = defaultdict(list)
latent_infos = self._latent_infos
self._latent_infos = defaultdict(list)
for k, v in latent_infos.items():
latent_infos[k] = np.asarray(v)
for k, v in agent_infos.items():
agent_infos[k] = np.asarray(v)
for k, v in env_infos.items():
env_infos[k] = np.asarray(v)
env_infos['task_onehot'] = np.asarray(tasks)
agent_infos['latent'] = np.asarray(latents)
for k, v in latent_infos.items():
agent_infos['latent_{}'.format(k)] = v
lengths = self._lengths
self._lengths = []
return TrajectoryBatch(self.env.spec, np.asarray(observations),
np.asarray(last_observations),
np.asarray(actions), np.asarray(rewards),
np.asarray(terminals), dict(env_infos),
dict(agent_infos), np.asarray(lengths,
dtype='i'))
def collect_rollout(self):
"""Collect all completed rollouts.
Returns:
metarl.TrajectoryBatch: A batch of the trajectories completed since
the last call to collect_rollout().
"""
if len(self._completed_rollouts) == 1:
result = self._completed_rollouts[0]
else:
result = TrajectoryBatch.concatenate(*self._completed_rollouts)
self._completed_rollouts = []
return result
def log_performance(self, itr, all_samples, loss_before, loss_after,
kl_before, kl, policy_entropy):
"""Evaluate performance of this batch.
Args:
itr (int): Iteration number.
all_samples (list[list[MAMLTrajectoryBatch]]): Two
dimensional list of MAMLTrajectoryBatch of size
[meta_batch_size * (num_grad_updates + 1)]
loss_before (float): Loss before optimization step.
loss_after (float): Loss after optimization step.
kl_before (float): KL divergence before optimization step.
kl (float): KL divergence after optimization step.
policy_entropy (float): Policy entropy.
Returns:
float: The average return in last epoch cycle.
"""
tabular.record('Iteration', itr)
rtns = []
for i in range(self._num_grad_updates + 1):
rtn = log_one_update_performance(
itr,
TrajectoryBatch.from_trajectory_list(
env_spec=self._env.spec,
paths=[
path for task_paths in all_samples
for path in task_paths[i].paths
]),
discount=self._inner_algo.discount,
task_names=self._env.task_names)
if i == self._num_grad_updates:
rtns.append(rtn)
with tabular.prefix(self._policy.name + '/'):
tabular.record('LossBefore', loss_before)
tabular.record('LossAfter', loss_after)
tabular.record('dLoss', loss_before - loss_after)
tabular.record('KLBefore', kl_before)
tabular.record('KLAfter', kl)
tabular.record('Entropy', policy_entropy)
return np.mean(rtns)
def _gather_rollout(self, rollout_number, last_observation):
assert 0 < self._path_lengths[rollout_number] <= self._max_path_length
traj = TrajectoryBatch(
self._envs[rollout_number].spec,
np.asarray(self._observations[rollout_number]),
np.asarray([last_observation]),
np.asarray(self._actions[rollout_number]),
np.asarray(self._rewards[rollout_number]),
np.asarray(self._terminals[rollout_number]),
self._env_infos[rollout_number], self._agent_infos[rollout_number],
np.asarray([self._path_lengths[rollout_number]], dtype='l'))
self._completed_rollouts.append(traj)
self._observations[rollout_number] = []
self._actions[rollout_number] = []
self._rewards[rollout_number] = []
self._terminals[rollout_number] = []
self._path_lengths[rollout_number] = 0
self._prev_obs[rollout_number] = self._envs[rollout_number].reset()
def test_log_performance():
lengths = np.array([10, 5, 1, 1])
batch = TrajectoryBatch(
EnvSpec(akro.Box(np.array([0., 0., 0.]), np.array([1., 1., 1.])),
akro.Box(np.array([-1., -1.]), np.array([0., 0.]))),
observations=np.ones((sum(lengths), 3), dtype=np.float32),
last_observations=np.ones((len(lengths), 3), dtype=np.float32),
actions=np.zeros((sum(lengths), 2), dtype=np.float32),
rewards=np.array([
0.34026529, 0.58263177, 0.84307509, 0.97651095, 0.81723901,
0.22631398, 0.03421301, 0.97515046, 0.64311832, 0.65068933,
0.17657714, 0.04783857, 0.73904013, 0.41364329, 0.52235551,
0.24203526, 0.43328910
]),
terminals=np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 1],
dtype=bool),
env_infos={
'success':
np.array([0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 1, 0, 1],
dtype=bool)
},
agent_infos={},
lengths=lengths)
log_file = tempfile.NamedTemporaryFile()
csv_output = dowel.CsvOutput(log_file.name)
logger.add_output(csv_output)
log_performance(7, batch, 0.8, prefix='test_log_performance')
logger.log(tabular)
logger.dump_output_type(dowel.CsvOutput)
with open(log_file.name, 'r') as file:
rows = list(csv.DictReader(file))
res = {k: float(r) for (k, r) in rows[0].items()}
assert res['test_log_performance/Iteration'] == 7
assert res['test_log_performance/NumTrajs'] == 4
assert math.isclose(res['test_log_performance/SuccessRate'], 0.75)
assert math.isclose(res['test_log_performance/CompletionRate'], 0.5)
assert math.isclose(res['test_log_performance/AverageDiscountedReturn'],
1.1131040640673113)
assert math.isclose(res['test_log_performance/AverageReturn'],
2.1659965525)
assert math.isclose(res['test_log_performance/StdReturn'],
2.354067152038576)
def train_once(self, itr, paths):
"""Train the algorithm once.
Args:
itr (int): Iteration number.
paths (list[dict]): A list of collected paths
Returns:
dict: Processed sample data, with key
* average_return: (float)
"""
batch_size = (self._training_batch_size if self._training_batch_size
else len(paths))
samples = self.process_samples(itr, paths)
for _ in range(self._training_epochs):
minibatch_ids_list = torch.randperm(len(paths)).split(batch_size)
for minibatch_ids in minibatch_ids_list:
obs, actions, rewards, valids, baselines = self._get_minibatch(samples, minibatch_ids)
loss = self._compute_loss(itr, obs, actions, rewards, valids,
baselines)
self._old_policy.load_state_dict(self.policy.state_dict())
self._optimizer.zero_grad()
loss.backward()
self._optimize(itr, obs, actions, rewards, valids, baselines)
self.baseline.fit(paths)
average_returns = log_performance(itr,
TrajectoryBatch.from_trajectory_list(
self.env_spec, paths),
discount=self.discount)
return np.mean(average_returns)
def _evaluate_policy(self, epoch):
"""Evaluate the performance of the policy via deterministic rollouts.
Statistics such as (average) discounted return and success rate are
recorded.
Args:
epoch (int): The current training epoch.
Returns:
float: The average return across self._num_evaluation_trajectories
trajectories
"""
eval_trajs = []
for _ in range(self._num_tasks):
eval_trajs.append(
self._obtain_evaluation_samples(
self._eval_env,
num_trajs=self._num_evaluation_trajectories))
eval_trajs = TrajectoryBatch.concatenate(*eval_trajs)
last_return = log_multitask_performance(epoch, eval_trajs,
self.discount)
return last_return
def test_lengths_shape_mismatch_traj(traj_data):
with pytest.raises(ValueError,
match='Lengths tensor must be a tensor of shape'):
traj_data['lengths'] = traj_data['lengths'].reshape((4, -1))
t = TrajectoryBatch(**traj_data)
del t
def test_act_env_spec_mismatch_traj(traj_data):
with pytest.raises(ValueError, match='actions must conform'):
traj_data['actions'] = traj_data['actions'][:, 0]
t = TrajectoryBatch(**traj_data)
del t
def train_once(self, itr, paths):
"""Train the algorithm once.
Args:
itr (int): Iteration number.
paths (list[dict]): A list of collected paths.
Returns:
numpy.float64: Calculated mean value of undiscounted returns.
"""
obs, actions, rewards, returns, valids, baselines = \
self.process_samples(paths)
if self._maximum_entropy:
policy_entropies = self._compute_policy_entropy(obs)
rewards += self._policy_ent_coeff * policy_entropies
obs_flat = torch.cat(filter_valids(obs, valids))
actions_flat = torch.cat(filter_valids(actions, valids))
rewards_flat = torch.cat(filter_valids(rewards, valids))
returns_flat = torch.cat(filter_valids(returns, valids))
advs_flat = self._compute_advantage(rewards, valids, baselines)
with torch.no_grad():
policy_loss_before = self._compute_loss_with_adv(
obs_flat, actions_flat, rewards_flat, advs_flat)
vf_loss_before = self._value_function.compute_loss(
obs_flat, returns_flat)
kl_before = self._compute_kl_constraint(obs)
self._train(obs_flat, actions_flat, rewards_flat, returns_flat,
advs_flat)
with torch.no_grad():
policy_loss_after = self._compute_loss_with_adv(
obs_flat, actions_flat, rewards_flat, advs_flat)
vf_loss_after = self._value_function.compute_loss(
obs_flat, returns_flat)
kl_after = self._compute_kl_constraint(obs)
policy_entropy = self._compute_policy_entropy(obs)
with tabular.prefix(self.policy.name):
tabular.record('/LossBefore', policy_loss_before.item())
tabular.record('/LossAfter', policy_loss_after.item())
tabular.record('/dLoss',
(policy_loss_before - policy_loss_after).item())
tabular.record('/KLBefore', kl_before.item())
tabular.record('/KL', kl_after.item())
tabular.record('/Entropy', policy_entropy.mean().item())
with tabular.prefix(self._value_function.name):
tabular.record('/LossBefore', vf_loss_before.item())
tabular.record('/LossAfter', vf_loss_after.item())
tabular.record('/dLoss',
vf_loss_before.item() - vf_loss_after.item())
self._old_policy.load_state_dict(self.policy.state_dict())
undiscounted_returns = log_performance(
itr,
TrajectoryBatch.from_trajectory_list(self._env_spec, paths),
discount=self.discount)
return np.mean(undiscounted_returns)
def test_lengths_dtype_mismatch_traj(traj_data):
with pytest.raises(ValueError,
match='Lengths tensor must have an integer dtype'):
traj_data['lengths'] = traj_data['lengths'].astype(np.float32)
t = TrajectoryBatch(**traj_data)
del t
def test_agent_infos_not_ndarray_traj(traj_data):
with pytest.raises(ValueError,
match='entry in agent_infos must be a numpy array'):
traj_data['agent_infos']['bar'] = list()
t = TrajectoryBatch(**traj_data)
del t
def test_terminals_dtype_mismatch_traj(traj_data):
with pytest.raises(ValueError, match='terminals tensor must be dtype'):
traj_data['terminals'] = traj_data['terminals'].astype(np.float32)
t = TrajectoryBatch(**traj_data)
del t
def obtain_exact_trajectories(self,
n_traj_per_worker,
agent_update,
env_update=None):
"""Sample an exact number of trajectories per worker.
Args:
n_traj_per_worker (int): Exact number of trajectories to gather for
each worker.
agent_update(object): Value which will be passed into the
`agent_update_fn` before doing rollouts. If a list is passed
in, it must have length exactly `factory.n_workers`, and will
be spread across the workers.
env_update(object): Value which will be passed into the
`env_update_fn` before doing rollouts. If a list is passed in,
it must have length exactly `factory.n_workers`, and will be
spread across the workers.
Returns:
TrajectoryBatch: Batch of gathered trajectories. Always in worker
order. In other words, first all trajectories from worker 0,
then all trajectories from worker 1, etc.
Raises:
AssertionError: On internal errors.
"""
self._agent_version += 1
updated_workers = set()
agent_ups = self._factory.prepare_worker_messages(
agent_update, cloudpickle.dumps)
env_ups = self._factory.prepare_worker_messages(env_update)
trajectories = defaultdict(list)
with click.progressbar(length=self._factory.n_workers,
label='Sampling') as pbar:
while any(
len(trajectories[i]) < n_traj_per_worker
for i in range(self._factory.n_workers)):
self._push_updates(updated_workers, agent_ups, env_ups)
tag, contents = self._to_sampler.get()
if tag == 'trajectory':
batch, version, worker_n = contents
if version == self._agent_version:
if len(trajectories[worker_n]) < n_traj_per_worker:
trajectories[worker_n].append(batch)
if len(trajectories[worker_n]) == n_traj_per_worker:
pbar.update(1)
try:
self._to_worker[worker_n].put_nowait(
('stop', ()))
except queue.Full:
pass
else:
raise AssertionError(
'Unknown tag {} with contents {}'.format(
tag, contents))
for q in self._to_worker:
try:
q.put_nowait(('stop', ()))
except queue.Full:
pass
ordered_trajectories = list(
itertools.chain(
*[trajectories[i] for i in range(self._factory.n_workers)]))
return TrajectoryBatch.concatenate(*ordered_trajectories)
def test_last_obs_env_spec_mismatch_traj(traj_data):
with pytest.raises(ValueError, match='last_observations must conform'):
traj_data['last_observations'] = \
traj_data['last_observations'][:, :, :, :1]
t = TrajectoryBatch(**traj_data)
del t
def test_obs_batch_mismatch_traj(traj_data):
with pytest.raises(ValueError, match='batch dimension of observations'):
traj_data['observations'] = traj_data['observations'][:-1]
t = TrajectoryBatch(**traj_data)
del t
def test_rewards_shape_mismatch_traj(traj_data):
with pytest.raises(ValueError, match='Rewards tensor'):
traj_data['rewards'] = traj_data['rewards'].reshape((2, -1))
t = TrajectoryBatch(**traj_data)
del t
def test_terminals_shape_mismatch_traj(traj_data):
with pytest.raises(ValueError, match='terminals tensor must have shape'):
traj_data['terminals'] = traj_data['terminals'].reshape((2, -1))
t = TrajectoryBatch(**traj_data)
del t
def obtain_samples(self, itr, num_samples, agent_update, env_update=None):
"""Collect at least a given number transitions (timesteps).
Args:
itr(int): The current iteration number. Using this argument is
deprecated.
num_samples(int): Minimum number of transitions / timesteps to
sample.
agent_update(object): Value which will be passed into the
`agent_update_fn` before doing rollouts. If a list is passed
in, it must have length exactly `factory.n_workers`, and will
be spread across the workers.
env_update(object): Value which will be passed into the
`env_update_fn` before doing rollouts. If a list is passed in,
it must have length exactly `factory.n_workers`, and will be
spread across the workers.
Returns:
metarl.TrajectoryBatch: The batch of collected trajectories.
Raises:
AssertionError: On internal errors.
"""
del itr
batches = []
completed_samples = 0
self._agent_version += 1
updated_workers = set()
agent_ups = self._factory.prepare_worker_messages(
agent_update, cloudpickle.dumps)
env_ups = self._factory.prepare_worker_messages(env_update)
with click.progressbar(length=num_samples, label='Sampling') as pbar:
while completed_samples < num_samples:
self._push_updates(updated_workers, agent_ups, env_ups)
for _ in range(self._factory.n_workers):
try:
tag, contents = self._to_sampler.get_nowait()
if tag == 'trajectory':
batch, version, worker_n = contents
del worker_n
if version == self._agent_version:
batches.append(batch)
num_returned_samples = batch.lengths.sum()
completed_samples += num_returned_samples
pbar.update(num_returned_samples)
else:
# Receiving paths from previous iterations is
# normal. Potentially, we could gather them
# here, if an off-policy method wants them.
pass
else:
raise AssertionError(
'Unknown tag {} with contents {}'.format(
tag, contents))
except queue.Empty:
pass
for q in self._to_worker:
try:
q.put_nowait(('stop', ()))
except queue.Full:
pass
return TrajectoryBatch.concatenate(*batches)