def build_q_perfect(self, cfg, goal_embedding):
"""Build the Q-function for perfect action space.
Args:
cfg: configuration of the experiments
goal_embedding: embedding tensor of the instructions
Returns:
q_out: output q values
input_select: object with maximum q value
action_select: actions with maximum q value
"""
per_input_ac_dim = cfg.per_input_ac_dim
des_len, inner_len = cfg.descriptor_length, cfg.inner_product_length
num_object = tf.shape(self.inputs)[1]
tp_concat = vector_tensor_product(self.inputs, self.inputs)
conv_layer_cfg = [[des_len * 8, 1, 1], [des_len * 4, 1, 1], [des_len, 1, 1]]
# [B, ?, ?, des_len]
tp_concat = stack_conv_layer(tp_concat, conv_layer_cfg)
# similarity with goal
goal_key = stack_dense_layer(
goal_embedding, [inner_len*2, inner_len]) # [B, d_inner]
goal_key = tf.expand_dims(goal_key, 1) # [B, 1, d_inner]
# [B, ?, ?, d_inner]
obs_query = tf.layers.conv2d(tp_concat, inner_len, 1, padding='same')
# [B, ?*?, d_inner]
obs_query = tf.reshape(obs_query, [-1, num_object**2, inner_len])
obs_query_t = tf.transpose(obs_query, perm=(0, 2, 1)) # [B, d_inner, ?*?]
inner = tf.matmul(goal_key, obs_query_t) # [B, 1, ?*?]
weight = tf.nn.softmax(inner, axis=-1) # [B, 1, ?*?]
prod = tf.matmul(
weight,
tf.reshape(tp_concat, [-1, num_object**2, des_len])) # [B, 1, des_len]
goal_embedding_ = tf.expand_dims(goal_embedding, 1) # [B, 1, dg]
# [B, ?, dg]
goal_embedding_ = tf.tile(goal_embedding_, multiples=[1, num_object, 1])
# [B, ?, des_len]
pair_wise_summary = tf.tile(prod, multiples=[1, num_object, 1])
# [B, ?, des_len+di+dg]
augemented_inputs = tf.concat(
[self.inputs, pair_wise_summary, goal_embedding_], axis=-1)
# [B, ?, 1, des_len+di+dg]
augemented_inputs = tf.expand_dims(augemented_inputs, axis=2)
conv_layer_cfg = [
[per_input_ac_dim*64, 1, 1],
[per_input_ac_dim*64, 1, 1],
[per_input_ac_dim, 1, 1]
]
# [B, ?, per_input_ac_dim]
q_out = tf.squeeze(
stack_conv_layer(augemented_inputs, conv_layer_cfg), axis=2)
input_max_q = tf.reduce_max(q_out, axis=2)
input_select = tf.argmax(input_max_q, axis=1)
action_max_q = tf.reduce_max(q_out, axis=1)
action_select = tf.argmax(action_max_q, axis=1)
return q_out, input_select, action_select
def build_q_discrete(self, cfg, goal_embedding):
"""Returns the q function for discrete action space.
Args:
cfg: configuration of the experiments
goal_embedding: embedding tensor of the instructions
Returns:
output q values of all actions
"""
ac_dim = cfg.ac_dim
des_len, inner_len = cfg.descriptor_length, cfg.inner_product_length
num_object = tf.shape(self.inputs)[1]
tp_concat = vector_tensor_product(self.inputs, self.inputs)
conv_layer_cfg = [[des_len * 8, 1, 1], [des_len * 4, 1, 1],
[des_len, 1, 1]]
# [B, ?, ?, des_len]
tp_concat = stack_conv_layer(tp_concat, conv_layer_cfg)
# similarity with goal
goal_key = stack_dense_layer(
goal_embedding, [inner_len * 2, inner_len]) # [B, d_inner]
goal_key = tf.expand_dims(goal_key, 1) # [B, 1, d_inner]
# [B, ?, ?, d_inner]
obs_query = tf.layers.conv2d(tp_concat, inner_len, 1, padding='same')
# [B, ?*?, d_inner]
obs_query = tf.reshape(obs_query, [-1, num_object**2, inner_len])
obs_query_t = tf.transpose(obs_query,
perm=(0, 2, 1)) # [B, d_inner, ?*?]
inner = tf.matmul(goal_key, obs_query_t) # [B, 1, ?*?]
weight = tf.nn.softmax(inner, axis=-1) # [B, 1, ?*?]
prod = tf.matmul(weight,
tf.reshape(
tp_concat,
[-1, num_object**2, des_len])) # [B, 1, des_len]
goal_embedding_ = tf.expand_dims(goal_embedding, 1) # [B, 1, dg]
# [B, ?, dg]
goal_embedding_ = tf.tile(goal_embedding_,
multiples=[1, num_object, 1])
# [B, ?, des_len]
pair_wise_summary = tf.tile(prod, multiples=[1, num_object, 1])
# [B, ?, des_len+di+dg]
augemented_inputs = tf.concat(
[self.inputs, pair_wise_summary, goal_embedding_], axis=-1)
# [B, ?, 1, des_len+di+dg]
augemented_inputs = tf.expand_dims(augemented_inputs, axis=2)
cfg = [[ac_dim // 8, 1, 1], [ac_dim // 8, 1, 1]]
heads = []
for _ in range(8):
# [B, ?, 1, ac_dim//8]
head_out = stack_conv_layer(augemented_inputs, cfg)
weights = tf.layers.conv2d(head_out, 1, 1) # [B, ?, 1, 1]
softmax_weights = tf.nn.softmax(weights, axis=1) # [B, ?, 1, 1]
heads.append(tf.reduce_sum(softmax_weights * head_out,
axis=(1, 2)))
# heads = 8 X [B, ac_dim//8]
out = tf.concat(heads, axis=1) # [B, ac_dim]
return tf.layers.dense(out, ac_dim)
