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)