def compute_policy_entropy(available_actions, policy, actions):
"""Compute total policy entropy.
Args: (same as compute_policy_log_probs)
Returns:
entropy: a scalar float tensor.
"""
_,arg_ids = actions
fn_pi, arg_pis = policy
fn_pi = mask_unavailable_actions(available_actions, fn_pi)
entropy = tf.reduce_mean(compute_entropy(fn_pi))
tf.summary.scalar('entropy/fn', entropy)
for arg_type in arg_ids.keys():
arg_id = arg_ids[arg_type]
arg_pi = arg_pis[arg_type]
batch_mask = tf.to_float(tf.not_equal(arg_id, -1))
arg_entropy = safe_div(
tf.reduce_sum(compute_entropy(arg_pi) * batch_mask),
tf.reduce_sum(batch_mask))
entropy += arg_entropy
tf.summary.scalar('used/arg/%s' % arg_type.name,
tf.reduce_mean(batch_mask))
tf.summary.scalar('entropy/arg/%s' % arg_type.name, arg_entropy)
return entropy
def sample_action(available_actions, policy):
def sample(probs):
dist = Categorical(probs=probs, allow_nan_stats=False)
return dist.sample()
fn_pi, arg_pis = policy
fn_pi = mask_unavailable_actions(available_actions, fn_pi)
fn_samples = sample(fn_pi)
arg_samples = dict()
for arg_type, arg_pi in arg_pis.items():
arg_samples[arg_type] = sample(arg_pi)
return fn_samples, arg_samples
def sample_actions(available_actions, policy):
"""Sample function ids and arguments from a predicted policy."""
def sample(probs):
dist = Categorical(probs=probs)
return dist.sample()
fn_pi, arg_pis = policy
fn_pi = mask_unavailable_actions(available_actions, fn_pi)
fn_samples = sample(fn_pi)
arg_samples = dict()
for arg_type, arg_pi in arg_pis.items():
arg_samples[arg_type] = sample(arg_pi)
return fn_samples, arg_samples
def compute_policy_log_probs(available_actions, policy, actions):
"""Compute action log probabilities given predicted policies and selected
actions.
Args:
available_actions: one-hot (in last dimenson) tensor of shape
[num_batch, NUM_FUNCTIONS].
policy: [fn_pi, {arg_0: arg_0_pi, ..., arg_n: arg_n_pi}]], where
each value is a tensor of shape [num_batch, num_params] representing
probability distributions over the function ids or over discrete
argument values.
actions: [fn_ids, {arg_0: arg_0_ids, ..., arg_n: arg_n_ids}], where
each value is a tensor of shape [num_batch] representing the selected
argument or actions ids. The argument id will be -1 if the argument is
not available for a specific (state, action) pair.
Returns:
log_prob: a tensor of shape [num_batch]
"""
def compute_log_probs(probs, labels):
# Select arbitrary element for unused arguments (log probs will be masked)
labels = tf.maximum(labels, 0)
indices = tf.stack([tf.range(tf.shape(labels)[0]), labels], axis=1)
# TODO tf.log should suffice
return safe_log(tf.gather_nd(probs, indices))
fn_id, arg_ids = actions
fn_pi, arg_pis = policy
# TODO: this should be unneccessary
fn_pi = mask_unavailable_actions(available_actions, fn_pi)
fn_log_prob = compute_log_probs(fn_pi, fn_id)
tf.summary.scalar('log_prob/fn', tf.reduce_mean(fn_log_prob))
log_prob = fn_log_prob
for arg_type in arg_ids.keys():
arg_id = arg_ids[arg_type]
arg_pi = arg_pis[arg_type]
arg_log_prob = compute_log_probs(arg_pi, arg_id)
arg_log_prob *= tf.to_float(tf.not_equal(arg_id, -1))
log_prob += arg_log_prob
tf.summary.scalar('log_prob/arg/%s' % arg_type.name,
tf.reduce_mean(arg_log_prob))
return log_prob