def build(self, feeds, state): history, goal, _ = self._embed_task_ios(feeds) _print_debug_ios(history, goal, None) params = self._lstm_hparams cell = lambda: tf.contrib.rnn.BasicLSTMCell(params.cell_size) stacked_lstm = tf.contrib.rnn.MultiRNNCell( [cell() for _ in range(params.num_layers)]) # history is of shape batch_size x seq_len x embedding_dimension batch_size, seq_len, _ = tuple(history.get_shape().as_list()) if state is None: state = stacked_lstm.zero_state(batch_size, tf.float32) for t in range(seq_len): if params.concat_goal_everywhere: lstm_input = tf.concat([tf.squeeze(history[:, t, :]), goal], axis=1) else: lstm_input = tf.squeeze(history[:, t, :]) output, state = stacked_lstm(lstm_input, state) with tf.variable_scope('output_decoder'): oconfig = self._task_config.output.shape assert len(oconfig) == 1 features = tf.concat([output, goal], axis=1) assert len(output.get_shape().as_list()) == 2 assert len(goal.get_shape().as_list()) == 2 decoder = embedders.MLPEmbedder( layers=self._embedder_hparams.predictions.layer_sizes + oconfig) # Prediction is done off the last step lstm output and the goal. predictions = decoder.build(features) return predictions, state
def build(self, feeds, prev_state): history, goal, _ = self._embed_task_ios(feeds) _print_debug_ios(history, goal, None) with tf.variable_scope('output_decoder'): # Concatenate the embeddings of the current observation and the goal. reactive_input = tf.concat([tf.squeeze(history[:, -1, :]), goal], axis=1) oconfig = self._task_config.output.shape assert len(oconfig) == 1 decoder = embedders.MLPEmbedder( layers=self._embedder_hparams.predictions.layer_sizes + oconfig) predictions = decoder.build(reactive_input) return predictions, None
def _embed_task_ios(self, streams): """Embeds a list of heterogenous streams. These streams correspond to task history, goal and output. The number of streams is equal to the total number of history, plus one for the goal if present, plus one for the output. If the number of history is k, then the first k streams are the history. The used embedders depend on the input (or goal) types. If an input is an image, then a ResNet embedder is used, otherwise MLPEmbedder (see embedders.py). Args: streams: a list of Tensors. Returns: Three float Tensors history, goal, output. If there are no history, or no goal, then the corresponding returned values are None. The shape of the embedded history is batch_size x sequence_length x sum of all embedding dimensions for all history. The shape of the goal is embedding dimension. """ # EMBED history. index = 0 inps = [] scopes = [] for c in self._task_config.inputs: if c == task_env.ModalityTypes.IMAGE: scope_name = 'image_embedder/image' reuse = scope_name in scopes scopes.append(scope_name) with tf.variable_scope(scope_name, reuse=reuse): resnet_embedder = embedders.ResNet( self._embedder_hparams.image) image_embeddings = resnet_embedder.build(streams[index]) # Uncover batch norm ops. if self._embedder_hparams.image.is_train: self._extra_train_ops += resnet_embedder.extra_train_ops inps.append(image_embeddings) index += 1 else: scope_name = 'input_embedder/vector' reuse = scope_name in scopes scopes.append(scope_name) with tf.variable_scope(scope_name, reuse=reuse): input_vector_embedder = embedders.MLPEmbedder( layers=self._embedder_hparams.vector) vector_embedder = input_vector_embedder.build( streams[index]) inps.append(vector_embedder) index += 1 history = tf.concat(inps, axis=2) if inps else None # EMBED goal. goal = None if self._task_config.query is not None: scope_name = 'image_embedder/query' reuse = scope_name in scopes scopes.append(scope_name) with tf.variable_scope(scope_name, reuse=reuse): resnet_goal_embedder = embedders.ResNet( self._embedder_hparams.goal) goal = resnet_goal_embedder.build(streams[index]) if self._embedder_hparams.goal.is_train: self._extra_train_ops += resnet_goal_embedder.extra_train_ops index += 1 # Embed true targets if needed (tbd). true_target = streams[index] return history, goal, true_target