def clean_dict(tensor_dict):
"""
Detach tensor values in nested dict.
Args:
tensor_dict (dict): Dict containing torch tensor.
Returns:
dict: Dict containing numpy arrays.
"""
# Un-nest.
param = define_by_run_flatten(tensor_dict)
ret = {}
# Detach tensor values.
for key, value in param.items():
if isinstance(value, torch.Tensor):
ret[key] = value.detach().numpy()
# Pack again.
return define_by_run_unflatten(ret)
def force_torch_tensors(params, requires_grad=False):
"""
Converts input params to torch tensors
Args:
params (list): Input args.
requires_grad (bool): If gradients need to be computed from these arguments.
Returns:
list: List of Torch tensors.
"""
if get_backend() == "pytorch":
tensor_params = []
for param in params:
if isinstance(param, dict):
# Flatten dict.
param = define_by_run_flatten(param)
ret = {}
for key, value in param.items():
ret[key] = convert_param(value, requires_grad)
tensor_params.append(ret)
else:
tensor_params.append(convert_param(param, requires_grad))
return tensor_params
def _graph_fn_update_from_external_batch(root, preprocessed_states,
actions, rewards, terminals,
sequence_indices,
apply_postprocessing):
"""
Calls iterative optimization by repeatedly sub-sampling.
"""
multi_gpu_sync_optimizer = root.sub_components.get(
"multi-gpu-synchronizer")
# Return values.
loss, loss_per_item, vf_loss, vf_loss_per_item = None, None, None, None
policy = root.get_sub_component_by_name(agent.policy.scope)
value_function = root.get_sub_component_by_name(
agent.value_function.scope)
optimizer = root.get_sub_component_by_name(agent.optimizer.scope)
loss_function = root.get_sub_component_by_name(
agent.loss_function.scope)
value_function_optimizer = root.get_sub_component_by_name(
agent.value_function_optimizer.scope)
vars_merger = root.get_sub_component_by_name(
agent.vars_merger.scope)
gae_function = root.get_sub_component_by_name(
agent.gae_function.scope)
prev_log_probs = policy.get_log_likelihood(
preprocessed_states, actions)["log_likelihood"]
if get_backend() == "tf":
# Log probs before update.
prev_log_probs = tf.stop_gradient(prev_log_probs)
batch_size = tf.shape(preprocessed_states)[0]
prior_baseline_values = tf.stop_gradient(
value_function.value_output(preprocessed_states))
# Advantages are based on prior baseline values.
advantages = tf.cond(
pred=apply_postprocessing,
true_fn=lambda: gae_function.calc_gae_values(
prior_baseline_values, rewards, terminals,
sequence_indices),
false_fn=lambda: rewards)
if self.standardize_advantages:
mean, std = tf.nn.moments(x=advantages, axes=[0])
advantages = (advantages - mean) / std
def opt_body(index_, loss_, loss_per_item_, vf_loss_,
vf_loss_per_item_):
start = tf.random_uniform(shape=(),
minval=0,
maxval=batch_size - 1,
dtype=tf.int32)
indices = tf.range(
start=start,
limit=start + agent.sample_size) % batch_size
sample_states = tf.gather(params=preprocessed_states,
indices=indices)
if isinstance(actions, ContainerDataOp):
sample_actions = FlattenedDataOp()
for name, action in flatten_op(actions).items():
sample_actions[name] = tf.gather(params=action,
indices=indices)
sample_actions = unflatten_op(sample_actions)
else:
sample_actions = tf.gather(params=actions,
indices=indices)
sample_prior_log_probs = tf.gather(params=prev_log_probs,
indices=indices)
sample_rewards = tf.gather(params=rewards, indices=indices)
sample_terminals = tf.gather(params=terminals,
indices=indices)
sample_sequence_indices = tf.gather(
params=sequence_indices, indices=indices)
sample_advantages = tf.gather(params=advantages,
indices=indices)
sample_advantages.set_shape((self.sample_size, ))
sample_baseline_values = value_function.value_output(
sample_states)
sample_prior_baseline_values = tf.gather(
params=prior_baseline_values, indices=indices)
# If we are a multi-GPU root:
# Simply feeds everything into the multi-GPU sync optimizer's method and return.
if multi_gpu_sync_optimizer is not None:
main_policy_vars = agent.policy.variables()
main_vf_vars = agent.value_function.variables()
all_vars = agent.vars_merger.merge(
main_policy_vars, main_vf_vars)
# grads_and_vars, loss, loss_per_item, vf_loss, vf_loss_per_item = \
out = multi_gpu_sync_optimizer.calculate_update_from_external_batch(
all_vars,
sample_states,
sample_actions,
sample_rewards,
sample_terminals,
sample_sequence_indices,
apply_postprocessing=apply_postprocessing)
avg_grads_and_vars_policy, avg_grads_and_vars_vf = agent.vars_splitter.call(
out["avg_grads_and_vars_by_component"])
policy_step_op = agent.optimizer.apply_gradients(
avg_grads_and_vars_policy)
vf_step_op = agent.value_function_optimizer.apply_gradients(
avg_grads_and_vars_vf)
step_op = root._graph_fn_group(policy_step_op,
vf_step_op)
step_and_sync_op = multi_gpu_sync_optimizer.sync_variables_to_towers(
step_op, all_vars)
loss_vf, loss_per_item_vf = out[
"additional_return_0"], out["additional_return_1"]
# Have to set all shapes here due to strict loop-var shape requirements.
out["loss"].set_shape(())
loss_vf.set_shape(())
loss_per_item_vf.set_shape((agent.sample_size, ))
out["loss_per_item"].set_shape((agent.sample_size, ))
with tf.control_dependencies([step_and_sync_op]):
if index_ == 0:
# Increase the global training step counter.
out["loss"] = root._graph_fn_training_step(
out["loss"])
return index_ + 1, out["loss"], out[
"loss_per_item"], loss_vf, loss_per_item_vf
policy_probs = policy.get_log_likelihood(
sample_states, sample_actions)["log_likelihood"]
baseline_values = value_function.value_output(
tf.stop_gradient(sample_states))
sample_rewards = tf.cond(
pred=apply_postprocessing,
true_fn=lambda: gae_function.calc_gae_values(
baseline_values, sample_rewards, sample_terminals,
sample_sequence_indices),
false_fn=lambda: sample_rewards)
sample_rewards.set_shape((agent.sample_size, ))
entropy = policy.get_entropy(sample_states)["entropy"]
loss, loss_per_item, vf_loss, vf_loss_per_item = \
loss_function.loss(
policy_probs, sample_prior_log_probs,
sample_baseline_values, sample_prior_baseline_values, sample_advantages, entropy
)
if hasattr(root, "is_multi_gpu_tower"
) and root.is_multi_gpu_tower is True:
policy_grads_and_vars = optimizer.calculate_gradients(
policy.variables(), loss)
vf_grads_and_vars = value_function_optimizer.calculate_gradients(
value_function.variables(), vf_loss)
grads_and_vars_by_component = vars_merger.merge(
policy_grads_and_vars, vf_grads_and_vars)
return grads_and_vars_by_component, loss, loss_per_item, vf_loss, vf_loss_per_item
else:
step_op, loss, loss_per_item = optimizer.step(
policy.variables(), loss, loss_per_item)
loss.set_shape(())
loss_per_item.set_shape((agent.sample_size, ))
vf_step_op, vf_loss, vf_loss_per_item = value_function_optimizer.step(
value_function.variables(), vf_loss,
vf_loss_per_item)
vf_loss.set_shape(())
vf_loss_per_item.set_shape((agent.sample_size, ))
with tf.control_dependencies([step_op, vf_step_op]):
return index_ + 1, loss, loss_per_item, vf_loss, vf_loss_per_item
def cond(index_, loss_, loss_per_item_, v_loss_,
v_loss_per_item_):
return index_ < agent.iterations
init_loop_vars = [
0,
tf.zeros(shape=(), dtype=tf.float32),
tf.zeros(shape=(agent.sample_size, )),
tf.zeros(shape=(), dtype=tf.float32),
tf.zeros(shape=(agent.sample_size, ))
]
if hasattr(root, "is_multi_gpu_tower"
) and root.is_multi_gpu_tower is True:
return opt_body(*init_loop_vars)
else:
index, loss, loss_per_item, vf_loss, vf_loss_per_item = tf.while_loop(
cond=cond,
body=opt_body,
loop_vars=init_loop_vars,
parallel_iterations=1)
# Increase the global training step counter.
loss = root._graph_fn_training_step(loss)
return loss, loss_per_item, vf_loss, vf_loss_per_item
elif get_backend() == "pytorch":
if isinstance(prev_log_probs, dict):
for name in actions.keys():
prev_log_probs[name] = prev_log_probs[name].detach()
else:
prev_log_probs = prev_log_probs.detach()
batch_size = preprocessed_states.shape[0]
sample_size = min(batch_size, agent.sample_size)
prior_baseline_values = value_function.value_output(
preprocessed_states).detach()
if apply_postprocessing:
advantages = gae_function.calc_gae_values(
prior_baseline_values, rewards, terminals,
sequence_indices)
else:
advantages = rewards
if self.standardize_advantages:
advantages = (advantages - torch.mean(advantages)
) / torch.std(advantages)
for _ in range(agent.iterations):
start = int(torch.rand(1) * (batch_size - 1))
indices = torch.arange(start=start,
end=start + sample_size,
dtype=torch.long) % batch_size
sample_states = torch.index_select(preprocessed_states, 0,
indices)
if isinstance(actions, dict):
sample_actions = DataOpDict()
sample_prior_log_probs = DataOpDict()
for name, action in define_by_run_flatten(
actions,
scope_separator_at_start=False).items():
sample_actions[name] = torch.index_select(
action, 0, indices)
sample_prior_log_probs[name] = torch.index_select(
prev_log_probs[name], 0, indices)
else:
sample_actions = torch.index_select(
actions, 0, indices)
sample_prior_log_probs = torch.index_select(
prev_log_probs, 0, indices)
sample_advantages = torch.index_select(
advantages, 0, indices)
sample_prior_baseline_values = torch.index_select(
prior_baseline_values, 0, indices)
policy_probs = policy.get_log_likelihood(
sample_states, sample_actions)["log_likelihood"]
sample_baseline_values = value_function.value_output(
sample_states)
entropy = policy.get_entropy(sample_states)["entropy"]
loss, loss_per_item, vf_loss, vf_loss_per_item = loss_function.loss(
policy_probs, sample_prior_log_probs,
sample_baseline_values, sample_prior_baseline_values,
sample_advantages, entropy)
# Do not need step op.
_, loss, loss_per_item = optimizer.step(
policy.variables(), loss, loss_per_item)
_, vf_loss, vf_loss_per_item = \
value_function_optimizer.step(value_function.variables(), vf_loss, vf_loss_per_item)
return loss, loss_per_item, vf_loss, vf_loss_per_item