Add RoomNav support

.circleci/config.yml

@@ -196,7 +196,6 @@ jobs:
               . activate habitat; cd habitat-api
               python setup.py develop --all
               python setup.py test
-              python -u habitat_baselines/run.py --exp-config habitat_baselines/config/pointnav/ppo_train_test.yaml --run-type train
 
 
 workflows:

configs/tasks/roomnav_toy_data.yaml

@@ -0,0 +1,27 @@
+ENVIRONMENT:
+  MAX_EPISODE_STEPS: 2000
+SIMULATOR:
+  AGENT_0:
+    SENSORS: ['DEPTH_SENSOR']
+  HABITAT_SIM_V0:
+    GPU_DEVICE_ID: 0
+  RGB_SENSOR:
+    WIDTH: 256
+    HEIGHT: 256
+  DEPTH_SENSOR:
+    WIDTH: 256
+    HEIGHT: 256
+TASK:
+  TYPE: Nav-v0
+  GOAL_SENSOR_UUID: "roomgoal"
+  SENSORS: ['ROOMGOAL_SENSOR']
+  ROOMGOAL_SENSOR:
+    TYPE: RoomGoalSensor
+  MEASUREMENTS: ['ROOMNAVMETRIC']
+  ROOMNAVMETRIC:
+    TYPE: RoomNavMetric
+
+DATASET:
+  TYPE: RoomNav-v1
+  SPLIT: train
+  DATA_PATH: data/datasets/roomnav/roomnav_mp3d_tiny/v1/{split}/{split}.json.gz

habitat/config/default.py

@@ -44,6 +44,11 @@
 _C.TASK.POINTGOAL_SENSOR.TYPE = "PointGoalSensor"
 _C.TASK.POINTGOAL_SENSOR.GOAL_FORMAT = "POLAR"
 # -----------------------------------------------------------------------------
+# # ROOMGOAL SENSOR
+# -----------------------------------------------------------------------------
+_C.TASK.ROOMGOAL_SENSOR = CN()
+_C.TASK.ROOMGOAL_SENSOR.TYPE = "RoomGoalSensor"
+# -----------------------------------------------------------------------------
 # # STATIC POINTGOAL SENSOR
 # -----------------------------------------------------------------------------
 _C.TASK.STATIC_POINTGOAL_SENSOR = CN()
@@ -67,6 +72,11 @@
 _C.TASK.SPL.TYPE = "SPL"
 _C.TASK.SPL.SUCCESS_DISTANCE = 0.2
 # -----------------------------------------------------------------------------
+# # ROOMNAVMETRIC MEASUREMENT
+# -----------------------------------------------------------------------------
+_C.TASK.ROOMNAVMETRIC = CN()
+_C.TASK.ROOMNAVMETRIC.TYPE = "RoomNavMetric"
+# -----------------------------------------------------------------------------
 # # TopDownMap MEASUREMENT
 # -----------------------------------------------------------------------------
 _C.TASK.TOP_DOWN_MAP = CN()

habitat/datasets/registration.py

@@ -7,6 +7,7 @@
 from habitat.core.registry import registry
 from habitat.datasets.eqa.mp3d_eqa_dataset import Matterport3dDatasetV1
 from habitat.datasets.pointnav.pointnav_dataset import PointNavDatasetV1
+from habitat.datasets.roomnav.roomnav_dataset import RoomNavDatasetV1
 
 
 def make_dataset(id_dataset, **kwargs):

habitat/datasets/roomnav/__init__.py

@@ -0,0 +1,5 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.

habitat/datasets/roomnav/roomnav_dataset.py

@@ -0,0 +1,129 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import gzip
+import json
+import os
+from typing import List, Optional
+
+from habitat.config import Config
+from habitat.core.dataset import Dataset
+from habitat.core.registry import registry
+from habitat.tasks.nav.nav_task import (
+    RoomGoal,
+    RoomNavigationEpisode,
+    ShortestPathPoint,
+)
+
+ALL_SCENES_MASK = "*"
+CONTENT_SCENES_PATH_FIELD = "content_scenes_path"
+DEFAULT_SCENE_PATH_PREFIX = "data/scene_datasets/"
+
+
+@registry.register_dataset(name="RoomNav-v1")
+class RoomNavDatasetV1(Dataset):
+    r"""Class inherited from Dataset that loads Room Navigation dataset.
+    """
+
+    episodes: List[RoomNavigationEpisode]
+    content_scenes_path: str = "{data_path}/content/{scene}.json.gz"
+
+    @staticmethod
+    def check_config_paths_exist(config: Config) -> bool:
+        return os.path.exists(
+            config.DATA_PATH.format(split=config.SPLIT)
+        ) and os.path.exists(config.SCENES_DIR)
+
+    @staticmethod
+    def get_scenes_to_load(config: Config) -> List[str]:
+        r"""Return list of scene ids for which dataset has separate files with
+        episodes.
+        """
+        assert RoomNavDatasetV1.check_config_paths_exist(config)
+        dataset_dir = os.path.dirname(
+            config.DATA_PATH.format(split=config.SPLIT)
+        )
+
+        cfg = config.clone()
+        cfg.defrost()
+        cfg.CONTENT_SCENES = []
+        dataset = RoomNavDatasetV1(cfg)
+        return RoomNavDatasetV1._get_scenes_from_folder(
+            content_scenes_path=dataset.content_scenes_path,
+            dataset_dir=dataset_dir,
+        )
+
+    @staticmethod
+    def _get_scenes_from_folder(content_scenes_path, dataset_dir):
+        scenes = []
+        content_dir = content_scenes_path.split("{scene}")[0]
+        scene_dataset_ext = content_scenes_path.split("{scene}")[1]
+        content_dir = content_dir.format(data_path=dataset_dir)
+        if not os.path.exists(content_dir):
+            print("Path doesnt exist: ", content_dir)
+            return scenes
+
+        for filename in os.listdir(content_dir):
+            if filename.endswith(scene_dataset_ext):
+                scene = filename[: -len(scene_dataset_ext)]
+                scenes.append(scene)
+        scenes.sort()
+        return scenes
+
+    def __init__(self, config: Optional[Config] = None) -> None:
+        self.episodes = []
+
+        if config is None:
+            return
+
+        datasetfile_path = config.DATA_PATH.format(split=config.SPLIT)
+
+        with gzip.open(datasetfile_path, "rt") as f:
+            self.from_json(f.read(), scenes_dir=config.SCENES_DIR)
+
+        # Read separate file for each scene
+        dataset_dir = os.path.dirname(datasetfile_path)
+        scenes = config.CONTENT_SCENES
+        if ALL_SCENES_MASK in scenes:
+            scenes = RoomNavDatasetV1._get_scenes_from_folder(
+                content_scenes_path=self.content_scenes_path,
+                dataset_dir=dataset_dir,
+            )
+
+        for scene in scenes:
+            scene_filename = self.content_scenes_path.format(
+                data_path=dataset_dir, scene=scene
+            )
+            with gzip.open(scene_filename, "rt") as f:
+                self.from_json(f.read(), scenes_dir=config.SCENES_DIR)
+
+    def from_json(
+        self, json_str: str, scenes_dir: Optional[str] = None
+    ) -> None:
+        deserialized = json.loads(json_str)
+        if CONTENT_SCENES_PATH_FIELD in deserialized:
+            self.content_scenes_path = deserialized[CONTENT_SCENES_PATH_FIELD]
+
+        for episode in deserialized["episodes"]:
+            episode = RoomNavigationEpisode(**episode)
+
+            scene_id_start = episode.scene_id.find("data/scene_datasets")
+            episode.scene_id = episode.scene_id[scene_id_start:]
+            if scenes_dir is not None:
+                if episode.scene_id.startswith(DEFAULT_SCENE_PATH_PREFIX):
+                    episode.scene_id = episode.scene_id[
+                        len(DEFAULT_SCENE_PATH_PREFIX) :
+                    ]
+
+                episode.scene_id = os.path.join(scenes_dir, episode.scene_id)
+
+            for g_index, goal in enumerate(episode.goals):
+                episode.goals[g_index] = RoomGoal(**goal)
+            if episode.shortest_paths is not None:
+                for path in episode.shortest_paths:
+                    for p_index, point in enumerate(path):
+                        path[p_index] = ShortestPathPoint(**point)
+            self.episodes.append(episode)

habitat/tasks/nav/nav_task.py

@@ -4,7 +4,7 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-from typing import Any, List, Optional, Type
+from typing import Any, List, Optional, Tuple, Type
 
 import attr
 import cv2
@@ -75,9 +75,11 @@ class ObjectGoal(NavigationGoal):
 class RoomGoal(NavigationGoal):
     r"""Room goal that can be specified by room_id or position with radius.
     """
-
-    room_id: str = attr.ib(default=None, validator=not_none_validator)
-    room_name: Optional[str] = None
+    room_aabb: Tuple[float] = attr.ib(
+        default=None, validator=not_none_validator
+    )
+    # room_id: str = attr.ib(default=None, validator=not_none_validator)
+    room_name: str = attr.ib(default=None, validator=not_none_validator)
 
 
 @attr.s(auto_attribs=True, kw_only=True)
@@ -105,6 +107,26 @@ class NavigationEpisode(Episode):
     shortest_paths: Optional[List[ShortestPathPoint]] = None
 
 
+@attr.s(auto_attribs=True, kw_only=True)
+class RoomNavigationEpisode(NavigationEpisode):
+    r"""Class for episode specification that includes initial position and
+    rotation of agent, scene name, goal and optional shortest paths. An
+    episode is a description of one task instance for the agent.
+    Args:
+        episode_id: id of episode in the dataset, usually episode number
+        scene_id: id of scene in scene dataset
+        start_position: numpy ndarray containing 3 entries for (x, y, z)
+        start_rotation: numpy ndarray with 4 entries for (x, y, z, w)
+            elements of unit quaternion (versor) representing agent 3D
+            orientation. ref: https://en.wikipedia.org/wiki/Versor
+        goals: list of goals specifications
+        start_room: room id
+        shortest_paths: list containing shortest paths to goals
+    """
+
+    goals: List[RoomGoal] = attr.ib(default=None, validator=not_none_validator)
+
+
 @registry.register_sensor
 class PointGoalSensor(Sensor):
     r"""Sensor for PointGoal observations which are used in the PointNav task.
@@ -243,6 +265,52 @@ def get_observation(self, observations, episode):
         return self._initial_vector
 
 
+@registry.register_sensor
+class RoomGoalSensor(Sensor):
+    r"""Sensor for RoomGoal observations which are used in the RoomNav task.
+    For the agent in simulator the forward direction is along negative-z.
+    In polar coordinate format the angle returned is azimuth to the goal.
+    Args:
+        sim: reference to the simulator for calculating task observations.
+        config: config for the RoomGoal sensor. Can contain field for
+            GOAL_FORMAT which can be used to specify the format in which
+            the roomgoal is specified. Current options for goal format are
+            cartesian and polar.
+    Attributes:
+        _goal_format: format for specifying the goal which can be done
+            in cartesian or polar coordinates.
+    """
+
+    def __init__(self, sim: Simulator, config: Config):
+        self._sim = sim
+        self.room_name_to_id = {
+            "bathroom": 0,
+            "bedroom": 1,
+            "dining room": 2,
+            "kitchen": 3,
+            "living room": 4,
+        }
+
+        super().__init__(config=config)
+
+    def _get_uuid(self, *args: Any, **kwargs: Any):
+        return "roomgoal"
+
+    def _get_sensor_type(self, *args: Any, **kwargs: Any):
+        return SensorTypes.PATH
+
+    def _get_observation_space(self, *args: Any, **kwargs: Any):
+        return spaces.Box(
+            low=min(self.room_name_to_id.values()),
+            high=max(self.room_name_to_id.values()),
+            shape=(1,),
+            dtype=np.int64,
+        )
+
+    def get_observation(self, observations, episode):
+        return np.array([self.room_name_to_id[episode.goals[0].room_name]])
+
+
 @registry.register_sensor
 class HeadingSensor(Sensor):
     r"""Sensor for observing the agent's heading in the global coordinate
@@ -377,6 +445,81 @@ def update_metric(self, episode, action):
         )
 
 
+@registry.register_measure
+class RoomNavMetric(Measure):
+    r"""RoomNavMetric - SPL but for RoomNav
+    """
+
+    def __init__(self, sim: Simulator, config: Config):
+        self._previous_position = None
+        self._start_end_episode_distance = None
+        self._agent_episode_distance = None
+        self._sim = sim
+        self._config = config
+
+        super().__init__()
+
+    def _get_uuid(self, *args: Any, **kwargs: Any):
+        return "roomnavmetric"
+
+    def nearest_point_in_room(self, start_position, room_aabb):
+        x_axes = np.arange(room_aabb[0], room_aabb[2], 0.1)
+        y_axes = np.arange(room_aabb[1], room_aabb[3], 0.1)
+
+        shortest_distance = 100000.0
+        for i in x_axes:
+            for j in y_axes:
+                if self._sim.is_navigable([i, start_position[1], j]):
+                    dist = self._sim.geodesic_distance(
+                        start_position, [i, start_position[1], j]
+                    )
+                    shortest_distance = min(dist, shortest_distance)
+
+        return shortest_distance
+
+    def reset_metric(self, episode):
+        self._previous_position = self._sim.get_agent_state().position.tolist()
+        self._start_end_episode_distance = self.nearest_point_in_room(
+            episode.start_position, episode.goals[0].room_aabb
+        )
+        self._agent_episode_distance = 0.0
+        self._metric = None
+
+    def _euclidean_distance(self, position_a, position_b):
+        return np.linalg.norm(
+            np.array(position_b) - np.array(position_a), ord=2
+        )
+
+    @staticmethod
+    def in_room(position, room_aabb):
+        return (
+            room_aabb[0] + 0.5 < position[0] < room_aabb[2] - 0.5
+            and room_aabb[1] + 0.5 < position[2] < room_aabb[3] - 0.5
+        )
+
+    def update_metric(self, episode, action):
+        ep_success = 0
+        current_position = self._sim.get_agent_state().position.tolist()
+
+        if action == self._sim.index_stop_action and self.in_room(
+            current_position, episode.goals[0].room_aabb
+        ):
+            ep_success = 1
+
+        self._agent_episode_distance += self._euclidean_distance(
+            current_position, self._previous_position
+        )
+
+        self._previous_position = current_position
+
+        self._metric = ep_success * (
+            self._start_end_episode_distance
+            / max(
+                self._start_end_episode_distance, self._agent_episode_distance
+            )
+        )
+
+
 @registry.register_measure
 class Collisions(Measure):
     def __init__(self, sim, config):