research/cognitive_planning/policies.py
# Copyright 2018 The TensorFlow Authors All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Interface for the policy of the agents use for navigation."""
import abc
import tensorflow as tf
from absl import logging
import embedders
from envs import task_env
slim = tf.contrib.slim
def _print_debug_ios(history, goal, output):
"""Prints sizes of history, goal and outputs."""
if history is not None:
shape = history.get_shape().as_list()
# logging.info('history embedding shape ')
# logging.info(shape)
if len(shape) != 3:
raise ValueError('history Tensor must have rank=3')
if goal is not None:
logging.info('goal embedding shape ')
logging.info(goal.get_shape().as_list())
if output is not None:
logging.info('targets shape ')
logging.info(output.get_shape().as_list())
class Policy(object):
"""Represents the policy of the agent for navigation tasks.
Instantiates a policy that takes embedders for each modality and builds a
model to infer the actions.
"""
__metaclass__ = abc.ABCMeta
def __init__(self, embedders_dict, action_size):
"""Instantiates the policy.
Args:
embedders_dict: Dictionary of embedders for different modalities. Keys
should be identical to keys of observation modality.
action_size: Number of possible actions.
"""
self._embedders = embedders_dict
self._action_size = action_size
@abc.abstractmethod
def build(self, observations, prev_state):
"""Builds the model that represents the policy of the agent.
Args:
observations: Dictionary of observations from different modalities. Keys
are the name of the modalities.
prev_state: The tensor of the previous state of the model. Should be set
to None if the policy is stateless
Returns:
Tuple of (action, state) where action is the action logits and state is
the state of the model after taking new observation.
"""
raise NotImplementedError(
'Needs implementation as part of Policy interface')
class LSTMPolicy(Policy):
"""Represents the implementation of the LSTM based policy.
The architecture of the model is as follows. It embeds all the observations
using the embedders, concatenates the embeddings of all the modalities. Feed
them through two fully connected layers. The lstm takes the features from
fully connected layer and the previous action and success of previous action
and feed them to LSTM. The value for each action is predicted afterwards.
Although the class name has the word LSTM in it, it also supports a mode that
builds the network without LSTM just for comparison purposes.
"""
def __init__(self,
modality_names,
embedders_dict,
action_size,
params,
max_episode_length,
feedforward_mode=False):
"""Instantiates the LSTM policy.
Args:
modality_names: List of modality names. Makes sure the ordering in
concatenation remains the same as modality_names list. Each modality
needs to be in the embedders_dict.
embedders_dict: Dictionary of embedders for different modalities. Keys
should be identical to keys of observation modality. Values should be
instance of Embedder class. All the observations except PREV_ACTION
requires embedder.
action_size: Number of possible actions.
params: is instance of tf.hparams and contains the hyperparameters for the
policy network.
max_episode_length: integer, specifying the maximum length of each
episode.
feedforward_mode: If True, it does not add LSTM to the model. It should
only be set True for comparison between LSTM and feedforward models.
"""
super(LSTMPolicy, self).__init__(embedders_dict, action_size)
self._modality_names = modality_names
self._lstm_state_size = params.lstm_state_size
self._fc_channels = params.fc_channels
self._weight_decay = params.weight_decay
self._target_embedding_size = params.target_embedding_size
self._max_episode_length = max_episode_length
self._feedforward_mode = feedforward_mode
def _build_lstm(self, encoded_inputs, prev_state, episode_length,
prev_action=None):
"""Builds an LSTM on top of the encoded inputs.
If prev_action is not None then it concatenates them to the input of LSTM.
Args:
encoded_inputs: The embedding of the observations and goal.
prev_state: previous state of LSTM.
episode_length: The tensor that contains the length of the sequence for
each element of the batch.
prev_action: tensor to previous chosen action and additional bit for
indicating whether the previous action was successful or not.
Returns:
a tuple of (lstm output, lstm state).
"""
# Adding prev action and success in addition to the embeddings of the
# modalities.
if prev_action is not None:
encoded_inputs = tf.concat([encoded_inputs, prev_action], axis=-1)
with tf.variable_scope('LSTM'):
lstm_cell = tf.nn.rnn_cell.BasicLSTMCell(self._lstm_state_size)
if prev_state is None:
# If prev state is set to None, a state of all zeros will be
# passed as a previous value for the cell. Should be used for the
# first step of each episode.
tf_prev_state = lstm_cell.zero_state(
encoded_inputs.get_shape().as_list()[0], dtype=tf.float32)
else:
tf_prev_state = tf.nn.rnn_cell.LSTMStateTuple(prev_state[0],
prev_state[1])
lstm_outputs, lstm_state = tf.nn.dynamic_rnn(
cell=lstm_cell,
inputs=encoded_inputs,
sequence_length=episode_length,
initial_state=tf_prev_state,
dtype=tf.float32,
)
lstm_outputs = tf.reshape(lstm_outputs, [-1, lstm_cell.output_size])
return lstm_outputs, lstm_state
def build(
self,
observations,
prev_state,
):
"""Builds the model that represents the policy of the agent.
Args:
observations: Dictionary of observations from different modalities. Keys
are the name of the modalities. Observation should have the following
key-values.
observations['goal']: One-hot tensor that indicates the semantic
category of the goal. The shape should be
(batch_size x max_sequence_length x goals).
observations[task_env.ModalityTypes.PREV_ACTION]: has action_size + 1
elements where the first action_size numbers are the one hot vector
of the previous action and the last element indicates whether the
previous action was successful or not. If
task_env.ModalityTypes.PREV_ACTION is not in the observation, it
will not be used in the policy.
prev_state: Previous state of the model. It should be a tuple of (c,h)
where c and h are the previous cell value and hidden state of the lstm.
Each element of tuple has shape of (batch_size x lstm_cell_size).
If it is set to None, then it initializes the state of the lstm with all
zeros.
Returns:
Tuple of (action, state) where action is the action logits and state is
the state of the model after taking new observation.
Raises:
ValueError: If any of the modality names is not in observations or
embedders_dict.
ValueError: If 'goal' is not in the observations.
"""
for modality_name in self._modality_names:
if modality_name not in observations:
raise ValueError('modality name does not exist in observations: {} not '
'in {}'.format(modality_name, observations.keys()))
if modality_name not in self._embedders:
if modality_name == task_env.ModalityTypes.PREV_ACTION:
continue
raise ValueError('modality name does not have corresponding embedder'
' {} not in {}'.format(modality_name,
self._embedders.keys()))
if task_env.ModalityTypes.GOAL not in observations:
raise ValueError('goal should be provided in the observations')
goal = observations[task_env.ModalityTypes.GOAL]
prev_action = None
if task_env.ModalityTypes.PREV_ACTION in observations:
prev_action = observations[task_env.ModalityTypes.PREV_ACTION]
with tf.variable_scope('policy'):
with slim.arg_scope(
[slim.fully_connected],
activation_fn=tf.nn.relu,
weights_initializer=tf.truncated_normal_initializer(stddev=0.01),
weights_regularizer=slim.l2_regularizer(self._weight_decay)):
all_inputs = []
# Concatenating the embedding of each modality by applying the embedders
# to corresponding observations.
def embed(name):
with tf.variable_scope('embed_{}'.format(name)):
# logging.info('Policy uses embedding %s', name)
return self._embedders[name].build(observations[name])
all_inputs = map(embed, [
x for x in self._modality_names
if x != task_env.ModalityTypes.PREV_ACTION
])
# Computing goal embedding.
shape = goal.get_shape().as_list()
with tf.variable_scope('embed_goal'):
encoded_goal = tf.reshape(goal, [shape[0] * shape[1], -1])
encoded_goal = slim.fully_connected(encoded_goal,
self._target_embedding_size)
encoded_goal = tf.reshape(encoded_goal, [shape[0], shape[1], -1])
all_inputs.append(encoded_goal)
# Concatenating all the modalities and goal.
all_inputs = tf.concat(all_inputs, axis=-1, name='concat_embeddings')
shape = all_inputs.get_shape().as_list()
all_inputs = tf.reshape(all_inputs, [shape[0] * shape[1], shape[2]])
# Applying fully connected layers.
encoded_inputs = slim.fully_connected(all_inputs, self._fc_channels)
encoded_inputs = slim.fully_connected(encoded_inputs, self._fc_channels)
if not self._feedforward_mode:
encoded_inputs = tf.reshape(encoded_inputs,
[shape[0], shape[1], self._fc_channels])
lstm_outputs, lstm_state = self._build_lstm(
encoded_inputs=encoded_inputs,
prev_state=prev_state,
episode_length=tf.ones((shape[0],), dtype=tf.float32) *
self._max_episode_length,
prev_action=prev_action,
)
else:
# If feedforward_mode=True, directly compute bypass the whole LSTM
# computations.
lstm_outputs = encoded_inputs
lstm_outputs = slim.fully_connected(lstm_outputs, self._fc_channels)
action_values = slim.fully_connected(
lstm_outputs, self._action_size, activation_fn=None)
action_values = tf.reshape(action_values, [shape[0], shape[1], -1])
if not self._feedforward_mode:
return action_values, lstm_state
else:
return action_values, None
class TaskPolicy(Policy):
"""A covenience abstract class providing functionality to deal with Tasks."""
def __init__(self,
task_config,
model_hparams=None,
embedder_hparams=None,
train_hparams=None):
"""Constructs a policy which knows how to work with tasks (see tasks.py).
It allows to read task history, goal and outputs in consistency with the
task config.
Args:
task_config: an object of type tasks.TaskIOConfig (see tasks.py)
model_hparams: a tf.HParams object containing parameter pertaining to
model (these are implementation specific)
embedder_hparams: a tf.HParams object containing parameter pertaining to
history, goal embedders (these are implementation specific)
train_hparams: a tf.HParams object containing parameter pertaining to
trainin (these are implementation specific)`
"""
super(TaskPolicy, self).__init__(None, None)
self._model_hparams = model_hparams
self._embedder_hparams = embedder_hparams
self._train_hparams = train_hparams
self._task_config = task_config
self._extra_train_ops = []
@property
def extra_train_ops(self):
"""Training ops in addition to the loss, e.g. batch norm updates.
Returns:
A list of tf ops.
"""
return self._extra_train_ops
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
@abc.abstractmethod
def build(self, feeds, prev_state):
pass
class ReactivePolicy(TaskPolicy):
"""A policy which ignores history.
It processes only the current observation (last element in history) and the
goal to output a prediction.
"""
def __init__(self, *args, **kwargs):
super(ReactivePolicy, self).__init__(*args, **kwargs)
# The current implementation ignores the prev_state as it is purely reactive.
# It returns None for the current 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
class RNNPolicy(TaskPolicy):
"""A policy which takes into account the full history via RNN.
The implementation might and will change.
The history, together with the goal, is processed using a stacked LSTM. The
output of the last LSTM step is used to produce a prediction. Currently, only
a single step output is supported.
"""
def __init__(self, lstm_hparams, *args, **kwargs):
super(RNNPolicy, self).__init__(*args, **kwargs)
self._lstm_hparams = lstm_hparams
# The prev_state is ignored as for now the full history is specified as first
# element of the feeds. It might turn out to be beneficial to keep the state
# as part of the policy object.
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