models/research/cognitive_planning/envs/active_vision_dataset_env.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.
# ==============================================================================

"""Gym environment for the ActiveVision Dataset.

   The dataset is captured with a robot moving around and taking picture in
   multiple directions. The actions are moving in four directions, and rotate
   clockwise or counter clockwise. The observations are the output of vision
   pipelines such as object detectors. The goal is to find objects of interest
   in each environment. For more details, refer:
   http://cs.unc.edu/~ammirato/active_vision_dataset_website/.
"""
import tensorflow as tf
import collections
import copy
import json
import os
from StringIO import StringIO
import time
import gym
from gym.envs.registration import register
import gym.spaces
import networkx as nx
import numpy as np
import scipy.io as sio
from absl import logging
import gin
import cv2
import label_map_util
import visualization_utils as vis_util
from envs import task_env


register(
    id='active-vision-env-v0',
    entry_point=
    'cognitive_planning.envs.active_vision_dataset_env:ActiveVisionDatasetEnv',  # pylint: disable=line-too-long
)

_MAX_DEPTH_VALUE = 12102

SUPPORTED_ACTIONS = [
    'right', 'rotate_cw', 'rotate_ccw', 'forward', 'left', 'backward', 'stop'
]
SUPPORTED_MODALITIES = [
    task_env.ModalityTypes.SEMANTIC_SEGMENTATION,
    task_env.ModalityTypes.DEPTH,
    task_env.ModalityTypes.OBJECT_DETECTION,
    task_env.ModalityTypes.IMAGE,
    task_env.ModalityTypes.GOAL,
    task_env.ModalityTypes.PREV_ACTION,
    task_env.ModalityTypes.DISTANCE,
]

# Data structure for storing the information related to the graph of the world.
_Graph = collections.namedtuple('_Graph', [
    'graph', 'id_to_index', 'index_to_id', 'target_indexes', 'distance_to_goal'
])


def _init_category_index(label_map_path):
  """Creates category index from class indexes to name of the classes.

  Args:
    label_map_path: path to the mapping.
  Returns:
    A map for mapping int keys to string categories.
  """

  label_map = label_map_util.load_labelmap(label_map_path)
  num_classes = np.max(x.id for x in label_map.item)
  categories = label_map_util.convert_label_map_to_categories(
      label_map, max_num_classes=num_classes, use_display_name=True)
  category_index = label_map_util.create_category_index(categories)
  return category_index


def _draw_detections(image_np, detections, category_index):
  """Draws detections on to the image.

  Args:
    image_np: Image in the form of uint8 numpy array.
    detections: a dictionary that contains the detection outputs.
    category_index: contains the mapping between indexes and the category names.

  Returns:
    Does not return anything but draws the boxes on the
  """
  vis_util.visualize_boxes_and_labels_on_image_array(
      image_np,
      detections['detection_boxes'],
      detections['detection_classes'],
      detections['detection_scores'],
      category_index,
      use_normalized_coordinates=True,
      max_boxes_to_draw=1000,
      min_score_thresh=.0,
      agnostic_mode=False)


def generate_detection_image(detections,
                             image_size,
                             category_map,
                             num_classes,
                             is_binary=True):
  """Generates one_hot vector of the image using the detection boxes.

  Args:
    detections: 2D object detections from the image. It's a dictionary that
      contains detection_boxes, detection_classes, and detection_scores with
      dimensions of nx4, nx1, nx1 where n is the number of detections.
    image_size: The resolution of the output image.
    category_map: dictionary that maps label names to index.
    num_classes: Number of classes.
    is_binary: If true, it sets the corresponding channels to 0 and 1.
      Otherwise, sets the score in the corresponding channel.
  Returns:
    Returns image_size x image_size x num_classes image for the detection boxes.
  """
  res = np.zeros((image_size, image_size, num_classes), dtype=np.float32)
  boxes = detections['detection_boxes']
  labels = detections['detection_classes']
  scores = detections['detection_scores']
  for box, label, score in zip(boxes, labels, scores):
    transformed_boxes = [int(round(t)) for t in box * image_size]
    y1, x1, y2, x2 = transformed_boxes
    # Detector returns fixed number of detections. Boxes with area of zero
    # are equivalent of boxes that don't correspond to any detection box.
    # So, we need to skip the boxes with area 0.
    if (y2 - y1) * (x2 - x1) == 0:
      continue
    assert category_map[label] < num_classes, 'label = {}'.format(label)
    value = score
    if is_binary:
      value = 1
    res[y1:y2, x1:x2, category_map[label]] = value
  return res


def _get_detection_path(root, detection_folder_name, world):
  return os.path.join(root, 'Meta', detection_folder_name, world + '.npy')


def _get_image_folder(root, world):
  return os.path.join(root, world, 'jpg_rgb')


def _get_json_path(root, world):
  return os.path.join(root, world, 'annotations.json')


def _get_image_path(root, world, image_id):
  return os.path.join(_get_image_folder(root, world), image_id + '.jpg')


def _get_image_list(path, worlds):
  """Builds a dictionary for all the worlds.

  Args:
    path: the path to the dataset on cns.
    worlds: list of the worlds.

  Returns:
    dictionary where the key is the world names and the values
    are the image_ids of that world.
  """
  world_id_dict = {}
  for loc in worlds:
    files = [t[:-4] for t in tf.gfile.ListDir(_get_image_folder(path, loc))]
    world_id_dict[loc] = files
  return world_id_dict


def read_all_poses(dataset_root, world):
  """Reads all the poses for each world.

  Args:
    dataset_root: the path to the root of the dataset.
    world: string, name of the world.

  Returns:
    Dictionary of poses for all the images in each world. The key is the image
    id of each view and the values are tuple of (x, z, R, scale). Where x and z
    are the first and third coordinate of translation. R is the 3x3 rotation
    matrix and scale is a float scalar that indicates the scale that needs to
    be multipled to x and z in order to get the real world coordinates.

  Raises:
    ValueError: if the number of images do not match the number of poses read.
  """
  path = os.path.join(dataset_root, world, 'image_structs.mat')
  with tf.gfile.Open(path) as f:
    data = sio.loadmat(f)
  xyz = data['image_structs']['world_pos']
  image_names = data['image_structs']['image_name'][0]
  rot = data['image_structs']['R'][0]
  scale = data['scale'][0][0]
  n = xyz.shape[1]
  x = [xyz[0][i][0][0] for i in range(n)]
  z = [xyz[0][i][2][0] for i in range(n)]
  names = [name[0][:-4] for name in image_names]
  if len(names) != len(x):
    raise ValueError('number of image names are not equal to the number of '
                     'poses {} != {}'.format(len(names), len(x)))
  output = {}
  for i in range(n):
    if rot[i].shape[0] != 0:
      assert rot[i].shape[0] == 3
      assert rot[i].shape[1] == 3
      output[names[i]] = (x[i], z[i], rot[i], scale)
    else:
      output[names[i]] = (x[i], z[i], None, scale)

  return output


def read_cached_data(should_load_images, dataset_root, segmentation_file_name,
                     targets_file_name, output_size):
  """Reads all the necessary cached data.

  Args:
    should_load_images: whether to load the images or not.
    dataset_root: path to the root of the dataset.
    segmentation_file_name: The name of the file that contains semantic
      segmentation annotations.
    targets_file_name: The name of the file the contains targets annotated for
      each world.
    output_size: Size of the output images. This is used for pre-processing the
      loaded images.
  Returns:
    Dictionary of all the cached data.
  """

  load_start = time.time()
  result_data = {}

  annotated_target_path = os.path.join(dataset_root, 'Meta',
                                       targets_file_name + '.npy')

  logging.info('loading targets: %s', annotated_target_path)
  with tf.gfile.Open(annotated_target_path) as f:
    result_data['targets'] = np.load(f).item()

  depth_image_path = os.path.join(dataset_root, 'Meta/depth_imgs.npy')
  logging.info('loading depth: %s', depth_image_path)
  with tf.gfile.Open(depth_image_path) as f:
    depth_data = np.load(f).item()

  logging.info('processing depth')
  for home_id in depth_data:
    images = depth_data[home_id]
    for image_id in images:
      depth = images[image_id]
      depth = cv2.resize(
          depth / _MAX_DEPTH_VALUE, (output_size, output_size),
          interpolation=cv2.INTER_NEAREST)
      depth_mask = (depth > 0).astype(np.float32)
      depth = np.dstack((depth, depth_mask))
      images[image_id] = depth
  result_data[task_env.ModalityTypes.DEPTH] = depth_data

  sseg_path = os.path.join(dataset_root, 'Meta',
                           segmentation_file_name + '.npy')
  logging.info('loading sseg: %s', sseg_path)
  with tf.gfile.Open(sseg_path) as f:
    sseg_data = np.load(f).item()

  logging.info('processing sseg')
  for home_id in sseg_data:
    images = sseg_data[home_id]
    for image_id in images:
      sseg = images[image_id]
      sseg = cv2.resize(
          sseg, (output_size, output_size), interpolation=cv2.INTER_NEAREST)
      images[image_id] = np.expand_dims(sseg, axis=-1).astype(np.float32)
  result_data[task_env.ModalityTypes.SEMANTIC_SEGMENTATION] = sseg_data

  if should_load_images:
    image_path = os.path.join(dataset_root, 'Meta/imgs.npy')
    logging.info('loading imgs: %s', image_path)
    with tf.gfile.Open(image_path) as f:
      image_data = np.load(f).item()

    result_data[task_env.ModalityTypes.IMAGE] = image_data

  with tf.gfile.Open(os.path.join(dataset_root, 'Meta/world_id_dict.npy')) as f:
    result_data['world_id_dict'] = np.load(f).item()

  logging.info('logging done in %f seconds', time.time() - load_start)
  return result_data


@gin.configurable
def get_spec_dtype_map():
  return {gym.spaces.Box: np.float32}


@gin.configurable
class ActiveVisionDatasetEnv(task_env.TaskEnv):
  """Simulates the environment from ActiveVisionDataset."""
  cached_data = None

  def __init__(
      self,
      episode_length,
      modality_types,
      confidence_threshold,
      output_size,
      worlds,
      targets,
      compute_distance,
      should_draw_detections,
      dataset_root,
      labelmap_path,
      reward_collision,
      reward_goal_range,
      num_detection_classes,
      segmentation_file_name,
      detection_folder_name,
      actions,
      targets_file_name,
      eval_init_points_file_name=None,
      shaped_reward=False,
  ):
    """Instantiates the environment for ActiveVision Dataset.

    Args:
      episode_length: the length of each episode.
      modality_types: a list of the strings where each entry indicates the name
        of the modalities to be loaded. Valid entries are "sseg", "det",
        "depth", "image", "distance", and "prev_action". "distance" should be
        used for computing metrics in tf agents.
      confidence_threshold: Consider detections more than confidence_threshold
        for potential targets.
      output_size: Resolution of the output image.
      worlds: List of the name of the worlds.
      targets: List of the target names. Each entry is a string label of the
        target category (e.g. 'fridge', 'microwave', so on).
      compute_distance: If True, outputs the distance of the view to the goal.
      should_draw_detections (bool): If True, the image returned for the
        observation will contains the bounding boxes.
      dataset_root: the path to the root folder of the dataset.
      labelmap_path: path to the dictionary that converts label strings to
        indexes.
      reward_collision: the reward the agents get after hitting an obstacle.
        It should be a non-positive number.
      reward_goal_range: the number of steps from goal, such that the agent is
        considered to have reached the goal. If the agent's distance is less
        than the specified goal range, the episode is also finishes by setting
        done = True.
      num_detection_classes: number of classes that detector outputs.
      segmentation_file_name: the name of the file that contains the semantic
        information. The file should be in the dataset_root/Meta/ folder.
      detection_folder_name: Name of the folder that contains the detections
        for each world. The folder should be under dataset_root/Meta/ folder.
      actions: The list of the action names. Valid entries are listed in
        SUPPORTED_ACTIONS.
      targets_file_name: the name of the file that contains the annotated
        targets. The file should be in the dataset_root/Meta/Folder
      eval_init_points_file_name: The name of the file that contains the initial
        points for evaluating the performance of the agent. If set to None,
        episodes start at random locations. Should be only set for evaluation.
      shaped_reward: Whether to add delta goal distance to the reward each step.

    Raises:
      ValueError: If one of the targets are not available in the annotated
        targets or the modality names are not from the domain specified above.
      ValueError: If one of the actions is not in SUPPORTED_ACTIONS.
      ValueError: If the reward_collision is a positive number.
      ValueError: If there is no action other than stop provided.
    """
    if reward_collision > 0:
      raise ValueError('"reward" for collision should be non positive')

    if reward_goal_range < 0:
      logging.warning('environment does not terminate the episode if the agent '
                      'is too close to the environment')

    if not modality_types:
      raise ValueError('modality names can not be empty')

    for name in modality_types:
      if name not in SUPPORTED_MODALITIES:
        raise ValueError('invalid modality type: {}'.format(name))

    actions_other_than_stop_found = False
    for a in actions:
      if a != 'stop':
        actions_other_than_stop_found = True
      if a not in SUPPORTED_ACTIONS:
        raise ValueError('invalid action %s', a)

    if not actions_other_than_stop_found:
      raise ValueError('environment needs to have actions other than stop.')

    super(ActiveVisionDatasetEnv, self).__init__()

    self._episode_length = episode_length
    self._modality_types = set(modality_types)
    self._confidence_threshold = confidence_threshold
    self._output_size = output_size
    self._dataset_root = dataset_root
    self._worlds = worlds
    self._targets = targets
    self._all_graph = {}
    for world in self._worlds:
      with tf.gfile.Open(_get_json_path(self._dataset_root, world), 'r') as f:
        file_content = f.read()
        file_content = file_content.replace('.jpg', '')
        io = StringIO(file_content)
        self._all_graph[world] = json.load(io)

    self._cur_world = ''
    self._cur_image_id = ''
    self._cur_graph = None  # Loaded by _update_graph
    self._steps_taken = 0
    self._last_action_success = True
    self._category_index = _init_category_index(labelmap_path)
    self._category_map = dict(
        [(c, i) for i, c in enumerate(self._category_index)])
    self._detection_cache = {}
    if not ActiveVisionDatasetEnv.cached_data:
      ActiveVisionDatasetEnv.cached_data = read_cached_data(
          True, self._dataset_root, segmentation_file_name, targets_file_name,
          self._output_size)
    cached_data = ActiveVisionDatasetEnv.cached_data

    self._world_id_dict = cached_data['world_id_dict']
    self._depth_images = cached_data[task_env.ModalityTypes.DEPTH]
    self._semantic_segmentations = cached_data[
        task_env.ModalityTypes.SEMANTIC_SEGMENTATION]
    self._annotated_targets = cached_data['targets']
    self._cached_imgs = cached_data[task_env.ModalityTypes.IMAGE]
    self._graph_cache = {}
    self._compute_distance = compute_distance
    self._should_draw_detections = should_draw_detections
    self._reward_collision = reward_collision
    self._reward_goal_range = reward_goal_range
    self._num_detection_classes = num_detection_classes
    self._actions = actions
    self._detection_folder_name = detection_folder_name
    self._shaped_reward = shaped_reward

    self._eval_init_points = None
    if eval_init_points_file_name is not None:
      self._eval_init_index = 0
      init_points_path = os.path.join(self._dataset_root, 'Meta',
                                      eval_init_points_file_name + '.npy')
      with tf.gfile.Open(init_points_path) as points_file:
        data = np.load(points_file).item()
      self._eval_init_points = []
      for world in self._worlds:
        for goal in self._targets:
          if world in self._annotated_targets[goal]:
            for image_id in data[world]:
              self._eval_init_points.append((world, image_id[0], goal))
        logging.info('loaded %d eval init points', len(self._eval_init_points))

    self.action_space = gym.spaces.Discrete(len(self._actions))

    obs_shapes = {}
    if task_env.ModalityTypes.SEMANTIC_SEGMENTATION in self._modality_types:
      obs_shapes[task_env.ModalityTypes.SEMANTIC_SEGMENTATION] = gym.spaces.Box(
          low=0, high=255, shape=(self._output_size, self._output_size, 1))
    if task_env.ModalityTypes.OBJECT_DETECTION in self._modality_types:
      obs_shapes[task_env.ModalityTypes.OBJECT_DETECTION] = gym.spaces.Box(
          low=0,
          high=255,
          shape=(self._output_size, self._output_size,
                 self._num_detection_classes))
    if task_env.ModalityTypes.DEPTH in self._modality_types:
      obs_shapes[task_env.ModalityTypes.DEPTH] = gym.spaces.Box(
          low=0,
          high=_MAX_DEPTH_VALUE,
          shape=(self._output_size, self._output_size, 2))
    if task_env.ModalityTypes.IMAGE in self._modality_types:
      obs_shapes[task_env.ModalityTypes.IMAGE] = gym.spaces.Box(
          low=0, high=255, shape=(self._output_size, self._output_size, 3))
    if task_env.ModalityTypes.GOAL in self._modality_types:
      obs_shapes[task_env.ModalityTypes.GOAL] = gym.spaces.Box(
          low=0, high=1., shape=(len(self._targets),))
    if task_env.ModalityTypes.PREV_ACTION in self._modality_types:
      obs_shapes[task_env.ModalityTypes.PREV_ACTION] = gym.spaces.Box(
          low=0, high=1., shape=(len(self._actions) + 1,))
    if task_env.ModalityTypes.DISTANCE in self._modality_types:
      obs_shapes[task_env.ModalityTypes.DISTANCE] = gym.spaces.Box(
          low=0, high=255, shape=(1,))
    self.observation_space = gym.spaces.Dict(obs_shapes)

    self._prev_action = np.zeros((len(self._actions) + 1), dtype=np.float32)

    # Loading all the poses.
    all_poses = {}
    for world in self._worlds:
      all_poses[world] = read_all_poses(self._dataset_root, world)
    self._cached_poses = all_poses
    self._vertex_to_pose = {}
    self._pose_to_vertex = {}

  @property
  def actions(self):
    """Returns list of actions for the env."""
    return self._actions

  def _next_image(self, image_id, action):
    """Given the action, returns the name of the image that agent ends up in.

    Args:
      image_id: The image id of the current view.
      action: valid actions are ['right', 'rotate_cw', 'rotate_ccw',
      'forward', 'left']. Each rotation is 30 degrees.

    Returns:
      The image name for the next location of the agent. If the action results
      in collision or it is not possible for the agent to execute that action,
      returns empty string.
    """
    assert action in self._actions, 'invalid action : {}'.format(action)
    assert self._cur_world in self._all_graph, 'invalid world {}'.format(
        self._cur_world)
    assert image_id in self._all_graph[
        self._cur_world], 'image_id {} is not in {}'.format(
            image_id, self._cur_world)
    return self._all_graph[self._cur_world][image_id][action]

  def _largest_detection_for_image(self, image_id, detections_dict):
    """Assigns area of the largest box for the view with given image id.

    Args:
      image_id: Image id of the view.
      detections_dict: Detections for the view.
    """
    for cls, box, score in zip(detections_dict['detection_classes'],
                               detections_dict['detection_boxes'],
                               detections_dict['detection_scores']):
      if cls not in self._targets:
        continue
      if score < self._confidence_threshold:
        continue
      ymin, xmin, ymax, xmax = box
      area = (ymax - ymin) * (xmax - xmin)
      if abs(area) < 1e-5:
        continue
      if image_id not in self._detection_area:
        self._detection_area[image_id] = area
      else:
        self._detection_area[image_id] = max(self._detection_area[image_id],
                                             area)

  def _compute_goal_indexes(self):
    """Computes the goal indexes for the environment.

    Returns:
      The indexes of the goals that are closest to target categories. A vertex
      is goal vertice if the desired objects are detected in the image and the
      target categories are not seen by moving forward from that vertice.
    """
    for image_id in self._world_id_dict[self._cur_world]:
      detections_dict = self._detection_table[image_id]
      self._largest_detection_for_image(image_id, detections_dict)
    goal_indexes = []
    for image_id in self._world_id_dict[self._cur_world]:
      if image_id not in self._detection_area:
        continue
      # Detection box is large enough.
      if self._detection_area[image_id] < 0.01:
        continue
      ok = True
      next_image_id = self._next_image(image_id, 'forward')
      if next_image_id:
        if next_image_id in self._detection_area:
          ok = False
      if ok:
        goal_indexes.append(self._cur_graph.id_to_index[image_id])
    return goal_indexes

  def to_image_id(self, vid):
    """Converts vertex id to the image id.

    Args:
      vid: vertex id of the view.
    Returns:
      image id of the input vertex id.
    """
    return self._cur_graph.index_to_id[vid]

  def to_vertex(self, image_id):
    return self._cur_graph.id_to_index[image_id]

  def observation(self, view_pose):
    """Returns the observation at the given the vertex.

    Args:
      view_pose: pose of the view of interest.

    Returns:
      Observation at the given view point.

    Raises:
      ValueError: if the given view pose is not similar to any of the poses in
        the current world.
    """
    vertex = self.pose_to_vertex(view_pose)
    if vertex is None:
      raise ValueError('The given found is not close enough to any of the poses'
                       ' in the environment.')
    image_id = self._cur_graph.index_to_id[vertex]
    output = collections.OrderedDict()

    if task_env.ModalityTypes.SEMANTIC_SEGMENTATION in self._modality_types:
      output[task_env.ModalityTypes.
             SEMANTIC_SEGMENTATION] = self._semantic_segmentations[
                 self._cur_world][image_id]

    detection = None
    need_det = (
        task_env.ModalityTypes.OBJECT_DETECTION in self._modality_types or
        (task_env.ModalityTypes.IMAGE in self._modality_types and
         self._should_draw_detections))
    if need_det:
      detection = self._detection_table[image_id]
      detection_image = generate_detection_image(
          detection,
          self._output_size,
          self._category_map,
          num_classes=self._num_detection_classes)

    if task_env.ModalityTypes.OBJECT_DETECTION in self._modality_types:
      output[task_env.ModalityTypes.OBJECT_DETECTION] = detection_image

    if task_env.ModalityTypes.DEPTH in self._modality_types:
      output[task_env.ModalityTypes.DEPTH] = self._depth_images[
          self._cur_world][image_id]

    if task_env.ModalityTypes.IMAGE in self._modality_types:
      output_img = self._cached_imgs[self._cur_world][image_id]
      if self._should_draw_detections:
        output_img = output_img.copy()
        _draw_detections(output_img, detection, self._category_index)
      output[task_env.ModalityTypes.IMAGE] = output_img

    if task_env.ModalityTypes.GOAL in self._modality_types:
      goal = np.zeros((len(self._targets),), dtype=np.float32)
      goal[self._targets.index(self._cur_goal)] = 1.
      output[task_env.ModalityTypes.GOAL] = goal

    if task_env.ModalityTypes.PREV_ACTION in self._modality_types:
      output[task_env.ModalityTypes.PREV_ACTION] = self._prev_action

    if task_env.ModalityTypes.DISTANCE in self._modality_types:
      output[task_env.ModalityTypes.DISTANCE] = np.asarray(
          [self.gt_value(self._cur_goal, vertex)], dtype=np.float32)

    return output

  def _step_no_reward(self, action):
    """Performs a step in the environment with given action.

    Args:
      action: Action that is used to step in the environment. Action can be
        string or integer. If the type is integer then it uses the ith element
        from self._actions list. Otherwise, uses the string value as the action.

    Returns:
      observation, done, info
      observation: dictonary that contains all the observations specified in
        modality_types.
        observation[task_env.ModalityTypes.OBJECT_DETECTION]: contains the
        detection of the current view.
        observation[task_env.ModalityTypes.IMAGE]: contains the
          image of the current view. Note that if using the images for training,
          should_load_images should be set to false.
        observation[task_env.ModalityTypes.SEMANTIC_SEGMENTATION]: contains the
          semantic segmentation of the current view.
        observation[task_env.ModalityTypes.DEPTH]: If selected, returns the
          depth map for the current view.
        observation[task_env.ModalityTypes.PREV_ACTION]: If selected, returns
          a numpy of (action_size + 1,). The first action_size elements indicate
          the action and the last element indicates whether the previous action
          was successful or not.
      done: True after episode_length steps have been taken, False otherwise.
      info: Empty dictionary.

    Raises:
      ValueError: for invalid actions.
    """
    # Primarily used for gym interface.
    if not isinstance(action, str):
      if not self.action_space.contains(action):
        raise ValueError('Not a valid actions: %d', action)

      action = self._actions[action]

    if action not in self._actions:
      raise ValueError('Not a valid action: %s', action)

    action_index = self._actions.index(action)

    if action == 'stop':
      next_image_id = self._cur_image_id
      done = True
      success = True
    else:
      next_image_id = self._next_image(self._cur_image_id, action)
      self._steps_taken += 1
      done = False
      success = True
    if not next_image_id:
      success = False
    else:
      self._cur_image_id = next_image_id

    if self._steps_taken >= self._episode_length:
      done = True

    cur_vertex = self._cur_graph.id_to_index[self._cur_image_id]
    observation = self.observation(self.vertex_to_pose(cur_vertex))

    # Concatenation of one-hot prev action + a binary number for success of
    # previous actions.
    self._prev_action = np.zeros((len(self._actions) + 1,), dtype=np.float32)
    self._prev_action[action_index] = 1.
    self._prev_action[-1] = float(success)

    distance_to_goal = self.gt_value(self._cur_goal, cur_vertex)
    if success:
      if distance_to_goal <= self._reward_goal_range:
        done = True

    return observation, done, {'success': success}

  @property
  def graph(self):
    return self._cur_graph.graph

  def state(self):
    return self.vertex_to_pose(self.to_vertex(self._cur_image_id))

  def gt_value(self, goal, v):
    """Computes the distance to the goal from vertex v.

    Args:
      goal: name of the goal.
      v: vertex id.

    Returns:
      Minimmum number of steps to the given goal.
    """
    assert goal in self._cur_graph.distance_to_goal, 'goal: {}'.format(goal)
    assert v in self._cur_graph.distance_to_goal[goal]
    res = self._cur_graph.distance_to_goal[goal][v]
    return res

  def _update_graph(self):
    """Creates the graph for each environment and updates the _cur_graph."""
    if self._cur_world not in self._graph_cache:
      graph = nx.DiGraph()
      id_to_index = {}
      index_to_id = {}
      image_list = self._world_id_dict[self._cur_world]
      for i, image_id in enumerate(image_list):
        id_to_index[image_id] = i
        index_to_id[i] = image_id
        graph.add_node(i)

      for image_id in image_list:
        for action in self._actions:
          if action == 'stop':
            continue
          next_image = self._all_graph[self._cur_world][image_id][action]
          if next_image:
            graph.add_edge(
                id_to_index[image_id], id_to_index[next_image], action=action)
      target_indexes = {}
      number_of_nodes_without_targets = graph.number_of_nodes()
      distance_to_goal = {}
      for goal in self._targets:
        if self._cur_world not in self._annotated_targets[goal]:
          continue
        goal_indexes = [
            id_to_index[i]
            for i in self._annotated_targets[goal][self._cur_world]
            if i
        ]
        super_source_index = graph.number_of_nodes()
        target_indexes[goal] = super_source_index
        graph.add_node(super_source_index)
        index_to_id[super_source_index] = goal
        id_to_index[goal] = super_source_index
        for v in goal_indexes:
          graph.add_edge(v, super_source_index, action='stop')
          graph.add_edge(super_source_index, v, action='stop')
        distance_to_goal[goal] = {}
        for v in range(number_of_nodes_without_targets):
          distance_to_goal[goal][v] = len(
              nx.shortest_path(graph, v, super_source_index)) - 2

      self._graph_cache[self._cur_world] = _Graph(
          graph, id_to_index, index_to_id, target_indexes, distance_to_goal)
    self._cur_graph = self._graph_cache[self._cur_world]

  def reset_for_eval(self, new_world, new_goal, new_image_id):
    """Resets to the given goal and image_id."""
    return self._reset_env(new_world=new_world, new_goal=new_goal, new_image_id=new_image_id)

  def get_init_config(self, path):
    """Exposes the initial state of the agent for the given path.

    Args:
      path: sequences of the vertexes that the agent moves.

    Returns:
      image_id of the first view, world, and the goal.
    """
    return self._cur_graph.index_to_id[path[0]], self._cur_world, self._cur_goal

  def _reset_env(
      self,
      new_world=None,
      new_goal=None,
      new_image_id=None,
  ):
    """Resets the agent in a random world and random id.

    Args:
      new_world: If not None, sets the new world to new_world.
      new_goal: If not None, sets the new goal to new_goal.
      new_image_id: If not None, sets the first image id to new_image_id.

    Returns:
      observation: dictionary of the observations. Content of the observation
      is similar to that of the step function.
    Raises:
      ValueError: if it can't find a world and annotated goal.
    """
    self._steps_taken = 0
    # The first prev_action is special all zero vector + success=1.
    self._prev_action = np.zeros((len(self._actions) + 1,), dtype=np.float32)
    self._prev_action[len(self._actions)] = 1.
    if self._eval_init_points is not None:
      if self._eval_init_index >= len(self._eval_init_points):
        self._eval_init_index = 0
      a = self._eval_init_points[self._eval_init_index]
      self._cur_world, self._cur_image_id, self._cur_goal = a
      self._eval_init_index += 1
    elif not new_world:
      attempts = 100
      found = False
      while attempts >= 0:
        attempts -= 1
        self._cur_goal = np.random.choice(self._targets)
        available_worlds = list(
            set(self._annotated_targets[self._cur_goal].keys()).intersection(
                set(self._worlds)))
        if available_worlds:
          found = True
          break
      if not found:
        raise ValueError('could not find a world that has a target annotated')
      self._cur_world = np.random.choice(available_worlds)
    else:
      self._cur_world = new_world
      self._cur_goal = new_goal
      if new_world not in self._annotated_targets[new_goal]:
        return None

    self._cur_goal_index = self._targets.index(self._cur_goal)
    if new_image_id:
      self._cur_image_id = new_image_id
    else:
      self._cur_image_id = np.random.choice(
          self._world_id_dict[self._cur_world])
    if self._cur_world not in self._detection_cache:
      with tf.gfile.Open(
          _get_detection_path(self._dataset_root, self._detection_folder_name,
                              self._cur_world)) as f:
        # Each file contains a dictionary with image ids as keys and detection
        # dicts as values.
        self._detection_cache[self._cur_world] = np.load(f).item()
    self._detection_table = self._detection_cache[self._cur_world]
    self._detection_area = {}
    self._update_graph()
    if self._cur_world not in self._vertex_to_pose:
      # adding fake pose for the super nodes of each target categories.
      self._vertex_to_pose[self._cur_world] = {
          index: (-index,) for index in self._cur_graph.target_indexes.values()
      }
      # Calling vetex_to_pose for each vertex results in filling out the
      # dictionaries that contain pose related data.
      for image_id in self._world_id_dict[self._cur_world]:
        self.vertex_to_pose(self.to_vertex(image_id))

      # Filling out pose_to_vertex from vertex_to_pose.
      self._pose_to_vertex[self._cur_world] = {
          tuple(v): k
          for k, v in self._vertex_to_pose[self._cur_world].iteritems()
      }

    cur_vertex = self._cur_graph.id_to_index[self._cur_image_id]
    observation = self.observation(self.vertex_to_pose(cur_vertex))
    return observation

  def cur_vertex(self):
    return self._cur_graph.id_to_index[self._cur_image_id]

  def cur_image_id(self):
    return self._cur_image_id

  def path_to_goal(self, image_id=None):
    """Returns the path from image_id to the self._cur_goal.

    Args:
      image_id: If set to None, computes the path from the current view.
        Otherwise, sets the current view to the given image_id.
    Returns:
      The path to the goal.
    Raises:
      Exception if there's no path from the view to the goal.
    """
    if image_id is None:
      image_id = self._cur_image_id
    super_source = self._cur_graph.target_indexes[self._cur_goal]
    try:
      path = nx.shortest_path(self._cur_graph.graph,
                              self._cur_graph.id_to_index[image_id],
                              super_source)
    except:
      print 'path not found, image_id = ', self._cur_world, self._cur_image_id
      raise
    return path[:-1]

  def targets(self):
    return [self.vertex_to_pose(self._cur_graph.target_indexes[self._cur_goal])]

  def vertex_to_pose(self, v):
    """Returns pose of the view for a given vertex.

    Args:
      v: integer, vertex index.

    Returns:
      (x, z, dir_x, dir_z) where x and z are the tranlation and dir_x, dir_z are
        a vector giving direction of the view.
    """
    if v in self._vertex_to_pose[self._cur_world]:
      return np.copy(self._vertex_to_pose[self._cur_world][v])

    x, z, rot, scale = self._cached_poses[self._cur_world][self.to_image_id(
        v)]
    if rot is None:  # if rotation is not provided for the given vertex.
      self._vertex_to_pose[self._cur_world][v] = np.asarray(
          [x * scale, z * scale, v])
      return np.copy(self._vertex_to_pose[self._cur_world][v])
    # Multiply rotation matrix by [0,0,1] to get a vector of length 1 in the
    # direction of the ray.
    direction = np.zeros((3, 1), dtype=np.float32)
    direction[2][0] = 1
    direction = np.matmul(np.transpose(rot), direction)
    direction = [direction[0][0], direction[2][0]]
    self._vertex_to_pose[self._cur_world][v] = np.asarray(
        [x * scale, z * scale, direction[0], direction[1]])
    return np.copy(self._vertex_to_pose[self._cur_world][v])

  def pose_to_vertex(self, pose):
    """Returns the vertex id for the given pose."""
    if tuple(pose) not in self._pose_to_vertex[self._cur_world]:
      raise ValueError(
          'The given pose is not present in the dictionary: {}'.format(
              tuple(pose)))

    return self._pose_to_vertex[self._cur_world][tuple(pose)]

  def check_scene_graph(self, world, goal):
    """Checks the connectivity of the scene graph.

    Goes over all the views. computes the shortest path to the goal. If it
    crashes it means that it's not connected. Otherwise, the env graph is fine.

    Args:
      world: the string name of the world.
      goal: the string label for the goal.
    Returns:
      Nothing.
    """
    obs = self._reset_env(new_world=world, new_goal=goal)
    if not obs:
      print '{} is not availble in {}'.format(goal, world)
      return True
    for image_id in self._world_id_dict[self._cur_world]:
      print 'check image_id = {}'.format(image_id)
      self._cur_image_id = image_id
      path = self.path_to_goal()
      actions = []
      for i in range(len(path) - 2):
        actions.append(self.action(path[i], path[i + 1]))
      actions.append('stop')

  @property
  def goal_one_hot(self):
    res = np.zeros((len(self._targets),), dtype=np.float32)
    res[self._cur_goal_index] = 1.
    return res

  @property
  def goal_index(self):
    return self._cur_goal_index

  @property
  def goal_string(self):
    return self._cur_goal

  @property
  def worlds(self):
    return self._worlds

  @property
  def possible_targets(self):
    return self._targets

  def action(self, from_pose, to_pose):
    """Returns the action that takes source vertex to destination vertex.

    Args:
      from_pose: pose of the source.
      to_pose: pose of the destination.
    Returns:
      Returns the index of the action.
    Raises:
      ValueError: If it is not possible to go from the first vertice to second
      vertice with one action, it raises value error.
    """
    from_index = self.pose_to_vertex(from_pose)
    to_index = self.pose_to_vertex(to_pose)
    if to_index not in self.graph[from_index]:
      from_image_id = self.to_image_id(from_index)
      to_image_id = self.to_image_id(to_index)
      raise ValueError('{},{} is not connected to {},{}'.format(
          from_index, from_image_id, to_index, to_image_id))
    return self._actions.index(self.graph[from_index][to_index]['action'])

  def random_step_sequence(self, min_len=None, max_len=None):
    """Generates random step sequence that takes agent to the goal.

    Args:
      min_len: integer, minimum length of a step sequence. Not yet implemented.
      max_len: integer, should be set to an integer and it is the maximum number
        of observations and path length to be max_len.
    Returns:
      Tuple of (path, actions, states, step_outputs).
        path: a random path from a random starting point and random environment.
        actions: actions of the returned path.
        states: viewpoints of all the states in between.
        step_outputs: list of step() return tuples.
    Raises:
      ValueError: if first_n is not greater than zero; if min_len is different
        from None.
    """
    if max_len is None:
      raise ValueError('max_len can not be set as None')
    if max_len < 1:
      raise ValueError('first_n must be greater or equal to 1.')
    if min_len is not None:
      raise ValueError('min_len is not yet implemented.')

    path = []
    actions = []
    states = []
    step_outputs = []
    obs = self.reset()
    last_obs_tuple = [obs, 0, False, {}]
    for _ in xrange(max_len):
      action = np.random.choice(self._actions)
      # We don't want to sample stop action because stop does not add new
      # information.
      while action == 'stop':
        action = np.random.choice(self._actions)
      path.append(self.to_vertex(self._cur_image_id))
      onehot = np.zeros((len(self._actions),), dtype=np.float32)
      onehot[self._actions.index(action)] = 1.
      actions.append(onehot)
      states.append(self.vertex_to_pose(path[-1]))
      step_outputs.append(copy.deepcopy(last_obs_tuple))
      last_obs_tuple = self.step(action)

    return path, actions, states, step_outputs