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
