diff --git a/.circleci/config.yml b/.circleci/config.yml index 6bd40c6fde..fa26256036 100644 --- a/.circleci/config.yml +++ b/.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: diff --git a/configs/tasks/roomnav_toy_data.yaml b/configs/tasks/roomnav_toy_data.yaml new file mode 100644 index 0000000000..10517108e0 --- /dev/null +++ b/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 diff --git a/habitat/config/default.py b/habitat/config/default.py index a433e997cc..0ae058a726 100644 --- a/habitat/config/default.py +++ b/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() diff --git a/habitat/datasets/registration.py b/habitat/datasets/registration.py index 265c43ef6e..cfd4d65360 100644 --- a/habitat/datasets/registration.py +++ b/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): diff --git a/habitat/datasets/roomnav/__init__.py b/habitat/datasets/roomnav/__init__.py new file mode 100644 index 0000000000..240697e324 --- /dev/null +++ b/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. diff --git a/habitat/datasets/roomnav/roomnav_dataset.py b/habitat/datasets/roomnav/roomnav_dataset.py new file mode 100644 index 0000000000..c67e15caa9 --- /dev/null +++ b/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) diff --git a/habitat/tasks/nav/nav_task.py b/habitat/tasks/nav/nav_task.py index 27b1451522..d90243b162 100644 --- a/habitat/tasks/nav/nav_task.py +++ b/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): diff --git a/habitat_baselines/common/environments.py b/habitat_baselines/common/environments.py index a6707ef826..3eb0ea4fc5 100644 --- a/habitat_baselines/common/environments.py +++ b/habitat_baselines/common/environments.py @@ -12,6 +12,8 @@ from typing import Optional, Type +import numpy as np + import habitat from habitat import Config, Dataset, SimulatorActions from habitat_baselines.common.baseline_registry import baseline_registry @@ -97,3 +99,83 @@ def get_done(self, observations): def get_info(self, observations): return self.habitat_env.get_metrics() + + +@baseline_registry.register_env(name="RoomNavRLEnv") +class RoomNavRLEnv(habitat.RLEnv): + def __init__(self, config: Config, dataset: Optional[Dataset] = None): + self._rl_config = config.RL + self._core_env_config = config.TASK_CONFIG + + self._previous_target_distance = None + self._previous_action = None + self._episode_distance_covered = None + self._success_distance = self._core_env_config.TASK.SUCCESS_DISTANCE + super().__init__(self._core_env_config, dataset) + + def reset(self): + self._previous_action = None + + observations = super().reset() + + self._previous_target_distance = self.habitat_env.current_episode.info[ + "geodesic_distance" + ] + return observations + + def step(self, action): + self._previous_action = action + return super().step(action) + + def get_reward_range(self): + return ( + self._rl_config.SLACK_REWARD - 1.0, + self._rl_config.SUCCESS_REWARD + 1.0, + ) + + def get_reward(self, observations): + reward = self._rl_config.SLACK_REWARD + + current_target_distance = self._distance_target() + reward += self._previous_target_distance - current_target_distance + self._previous_target_distance = current_target_distance + + if self._episode_success(): + reward += ( + self._rl_config.SUCCESS_REWARD + * self.get_info(observations)["roomnavmetric"] + ) + + return reward + + @staticmethod + def in_room(position, room_aabb): + return ( + room_aabb[0] + 0.20 < position[0] < room_aabb[2] - 0.20 + and room_aabb[1] + 0.20 < position[2] < room_aabb[3] - 0.20 + ) + + def _distance_target(self): + current_position = self._env.sim.get_agent_state().position.tolist() + target_position = self._env.current_episode.goals[0].position + distance = self._env.sim.geodesic_distance( + current_position, target_position + ) + return distance + + def _episode_success(self): + if self._previous_action == SimulatorActions.STOP and self.in_room( + self._env.sim.get_agent_state().position.tolist(), + self._env.current_episode.goals[0].room_aabb, + ): + return True + return False + + def get_done(self, observations): + done = False + if self._env.episode_over or self._episode_success(): + done = True + return done + + def get_info(self, observations): + return self.habitat_env.get_metrics() diff --git a/habitat_baselines/common/utils.py b/habitat_baselines/common/utils.py index feb24bf04a..5e650903b4 100644 --- a/habitat_baselines/common/utils.py +++ b/habitat_baselines/common/utils.py @@ -120,7 +120,8 @@ def generate_video( images: List[np.ndarray], episode_id: int, checkpoint_idx: int, - spl: float, + metric_name: str, + metric_value: float, tb_writer: TensorboardWriter, fps: int = 10, ) -> None: @@ -132,7 +133,8 @@ def generate_video( images: list of images to be converted to video. episode_id: episode id for video naming. checkpoint_idx: checkpoint index for video naming. - spl: SPL for this episode for video naming. + metric_name: name of the performance metric, e.g. "spl". + metric_value: value of metric. tb_writer: tensorboard writer object for uploading video. fps: fps for generated video. Returns: @@ -141,7 +143,7 @@ def generate_video( if len(images) < 1: return - video_name = f"episode{episode_id}_ckpt{checkpoint_idx}_spl{spl:.2f}" + video_name = f"episode{episode_id}_ckpt{checkpoint_idx}_{metric_name}{metric_value:.2f}" if "disk" in video_option: assert video_dir is not None images_to_video(images, video_dir, video_name) diff --git a/habitat_baselines/config/pointnav/ppo_pointnav.yaml b/habitat_baselines/config/pointnav/ppo_pointnav.yaml new file mode 100644 index 0000000000..c2d8a47ced --- /dev/null +++ b/habitat_baselines/config/pointnav/ppo_pointnav.yaml @@ -0,0 +1,36 @@ +BASE_TASK_CONFIG_PATH: "configs/tasks/pointnav.yaml" +TRAINER_NAME: "ppo" +ENV_NAME: "NavRLEnv" +SIMULATOR_GPU_ID: 0 +TORCH_GPU_ID: 0 +VIDEO_OPTION: ["disk", "tensorboard"] +TENSORBOARD_DIR: "tb" +VIDEO_DIR: "video_dir" +TEST_EPISODE_COUNT: 2 +EVAL_CKPT_PATH_DIR: "data/new_checkpoints" +NUM_PROCESSES: 1 +SENSORS: ["RGB_SENSOR", "DEPTH_SENSOR"] +CHECKPOINT_FOLDER: "data/new_checkpoints" +NUM_UPDATES: 10000 +LOG_INTERVAL: 10 +CHECKPOINT_INTERVAL: 50 + +RL: + PPO: + # ppo params + clip_param: 0.1 + ppo_epoch: 4 + num_mini_batch: 1 + value_loss_coef: 0.5 + entropy_coef: 0.01 + lr: 2.5e-4 + eps: 1e-5 + max_grad_norm: 0.5 + num_steps: 128 + hidden_size: 512 + use_gae: True + gamma: 0.99 + tau: 0.95 + use_linear_clip_decay: True + use_linear_lr_decay: True + reward_window_size: 50 diff --git a/habitat_baselines/config/roomnav/ppo_roomnav.yaml b/habitat_baselines/config/roomnav/ppo_roomnav.yaml new file mode 100644 index 0000000000..19d2bffcc4 --- /dev/null +++ b/habitat_baselines/config/roomnav/ppo_roomnav.yaml @@ -0,0 +1,35 @@ +BASE_TASK_CONFIG_PATH: "configs/tasks/roomnav_toy_data.yaml" +TRAINER_NAME: "ppo" +ENV_NAME: "RoomNavRLEnv" +SIMULATOR_GPU_ID: 0 +TORCH_GPU_ID: 0 +VIDEO_OPTION: [] +TENSORBOARD_DIR: "" +VIDEO_DIR: "" +TEST_EPISODE_COUNT: 2 +EVAL_CKPT_PATH_DIR: "data/checkpoints/roomnav/ckpt.0.pth" +NUM_PROCESSES: 1 +SENSORS: ["RGB_SENSOR", "DEPTH_SENSOR"] +CHECKPOINT_FOLDER: "data/checkpoints/roomnav" +NUM_UPDATES: 10000 +LOG_INTERVAL: 10 +CHECKPOINT_INTERVAL: 50 +RL: + PPO: + # ppo params + clip_param: 0.1 + ppo_epoch: 4 + num_mini_batch: 1 + value_loss_coef: 0.5 + entropy_coef: 0.01 + lr: 2.5e-4 + eps: 1e-5 + max_grad_norm: 0.5 + num_steps: 128 + hidden_size: 512 + use_gae: True + gamma: 0.99 + tau: 0.95 + use_linear_clip_decay: True + use_linear_lr_decay: True + reward_window_size: 50 diff --git a/habitat_baselines/config/pointnav/ppo_train_test.yaml b/habitat_baselines/config/test/ppo_pointnav_test.yaml similarity index 75% rename from habitat_baselines/config/pointnav/ppo_train_test.yaml rename to habitat_baselines/config/test/ppo_pointnav_test.yaml index 38dd96b808..cb3d16891c 100644 --- a/habitat_baselines/config/pointnav/ppo_train_test.yaml +++ b/habitat_baselines/config/test/ppo_pointnav_test.yaml @@ -5,12 +5,12 @@ SIMULATOR_GPU_ID: 0 TORCH_GPU_ID: 0 VIDEO_OPTION: [] TENSORBOARD_DIR: "" -EVAL_CKPT_PATH_DIR: "data/new_checkpoints" +EVAL_CKPT_PATH_DIR: "data/test_checkpoints/ppo/pointnav/ckpt.0.pth" NUM_PROCESSES: 1 -CHECKPOINT_FOLDER: "data/new_checkpoints" -NUM_UPDATES: 10 -LOG_INTERVAL: 1 -CHECKPOINT_INTERVAL: 1 +CHECKPOINT_FOLDER: "data/test_checkpoints/ppo/pointnav/" +NUM_UPDATES: 4 +LOG_INTERVAL: 100 +CHECKPOINT_INTERVAL: 2 RL: PPO: diff --git a/habitat_baselines/config/test/ppo_roomnav_test.yaml b/habitat_baselines/config/test/ppo_roomnav_test.yaml new file mode 100644 index 0000000000..c0249589a0 --- /dev/null +++ b/habitat_baselines/config/test/ppo_roomnav_test.yaml @@ -0,0 +1,36 @@ +BASE_TASK_CONFIG_PATH: "configs/tasks/roomnav_toy_data.yaml" +TRAINER_NAME: "ppo" +ENV_NAME: "RoomNavRLEnv" +SIMULATOR_GPU_ID: 0 +TORCH_GPU_ID: 0 +VIDEO_OPTION: [] +TENSORBOARD_DIR: "" +VIDEO_DIR: "" +TEST_EPISODE_COUNT: 2 +EVAL_CKPT_PATH_DIR: "data/checkpoints/roomnav/ckpt.0.pth" +NUM_PROCESSES: 1 +SENSORS: ["RGB_SENSOR", "DEPTH_SENSOR"] +CHECKPOINT_FOLDER: "data/checkpoints/roomnav" +NUM_UPDATES: 4 +LOG_INTERVAL: 100 +CHECKPOINT_INTERVAL: 2 + +RL: + PPO: + # ppo params + clip_param: 0.1 + ppo_epoch: 4 + num_mini_batch: 1 + value_loss_coef: 0.5 + entropy_coef: 0.01 + lr: 2.5e-4 + eps: 1e-5 + max_grad_norm: 0.5 + num_steps: 128 + hidden_size: 512 + use_gae: True + gamma: 0.99 + tau: 0.95 + use_linear_clip_decay: True + use_linear_lr_decay: True + reward_window_size: 50 diff --git a/habitat_baselines/rl/ppo/ppo_trainer.py b/habitat_baselines/rl/ppo/ppo_trainer.py index e508cf0d3d..5eba415e9e 100644 --- a/habitat_baselines/rl/ppo/ppo_trainer.py +++ b/habitat_baselines/rl/ppo/ppo_trainer.py @@ -379,6 +379,15 @@ def _eval_checkpoint( self.agent.load_state_dict(ckpt_dict["state_dict"]) self.actor_critic = self.agent.actor_critic + # get name of performance metric, e.g. "spl" + metric_name = self.config.TASK_CONFIG.TASK.MEASUREMENTS[0] + metric_cfg = getattr(self.config.TASK_CONFIG.TASK, metric_name) + measure_type = baseline_registry.get_measure(metric_cfg.TYPE) + assert measure_type is not None, "invalid measurement type {}".format( + metric_cfg.TYPE + ) + self.metric_uuid = measure_type(None, None)._get_uuid() + observations = self.envs.reset() batch = batch_obs(observations) for sensor in batch: @@ -457,8 +466,12 @@ def _eval_checkpoint( # episode ended if not_done_masks[i].item() == 0: episode_stats = dict() - episode_stats["spl"] = infos[i]["spl"] - episode_stats["success"] = int(infos[i]["spl"] > 0) + episode_stats[self.metric_uuid] = infos[i][ + self.metric_uuid + ] + episode_stats["success"] = int( + infos[i][self.metric_uuid] > 0 + ) episode_stats["reward"] = current_episode_reward[i].item() current_episode_reward[i] = 0 # use scene_id + episode_id as unique id for storing stats @@ -476,7 +489,8 @@ def _eval_checkpoint( images=rgb_frames[i], episode_id=current_episodes[i].episode_id, checkpoint_idx=checkpoint_index, - spl=infos[i]["spl"], + metric_name=self.metric_uuid, + metric_value=infos[i][self.metric_uuid], tb_writer=writer, ) @@ -516,12 +530,14 @@ def _eval_checkpoint( num_episodes = len(stats_episodes) episode_reward_mean = aggregated_stats["reward"] / num_episodes - episode_spl_mean = aggregated_stats["spl"] / num_episodes + episode_metric_mean = aggregated_stats[self.metric_uuid] / num_episodes episode_success_mean = aggregated_stats["success"] / num_episodes logger.info(f"Average episode reward: {episode_reward_mean:.6f}") logger.info(f"Average episode success: {episode_success_mean:.6f}") - logger.info(f"Average episode SPL: {episode_spl_mean:.6f}") + logger.info( + f"Average episode {self.metric_uuid}: {episode_metric_mean:.6f}" + ) writer.add_scalars( "eval_reward", @@ -529,7 +545,9 @@ def _eval_checkpoint( checkpoint_index, ) writer.add_scalars( - "eval_SPL", {"average SPL": episode_spl_mean}, checkpoint_index + f"eval_{self.metric_uuid}", + {f"average {self.metric_uuid}": episode_metric_mean}, + checkpoint_index, ) writer.add_scalars( "eval_success", diff --git a/habitat_baselines/run.py b/habitat_baselines/run.py index e8e1cb1d32..b44f299ca8 100644 --- a/habitat_baselines/run.py +++ b/habitat_baselines/run.py @@ -33,18 +33,33 @@ def main(): nargs=argparse.REMAINDER, help="Modify config options from command line", ) + args = parser.parse_args() - config = get_config(args.exp_config, args.opts) + run_exp(**vars(args)) + + +def run_exp(exp_config: str, run_type: str, opts=None) -> None: + r"""Runs experiment given mode and config + + Args: + exp_config: path to config file. + run_type: "train" or "eval. + opts: list of strings of additional config options. + + Returns: + None. + """ + config = get_config(exp_config, opts) + random.seed(config.TASK_CONFIG.SEED) np.random.seed(config.TASK_CONFIG.SEED) - trainer_init = baseline_registry.get_trainer(config.TRAINER_NAME) assert trainer_init is not None, f"{config.TRAINER_NAME} is not supported" trainer = trainer_init(config) - if args.run_type == "train": + if run_type == "train": trainer.train() - elif args.run_type == "eval": + elif run_type == "eval": trainer.eval() diff --git a/test/test_baseline_trainers.py b/test/test_baseline_trainers.py new file mode 100644 index 0000000000..19dccdefa2 --- /dev/null +++ b/test/test_baseline_trainers.py @@ -0,0 +1,66 @@ +#!/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 os +from glob import glob + +import pytest + +try: + from habitat_baselines.run import run_exp + from habitat_baselines.common.base_trainer import BaseRLTrainer + from habitat_baselines.config.default import get_config + + baseline_installed = True +except ImportError: + baseline_installed = False + + +@pytest.mark.skipif( + not baseline_installed, reason="baseline sub-module not installed" +) +@pytest.mark.parametrize( + "test_cfg_path,mode", + [ + (cfg, mode) + for mode in ("train", "eval") + for cfg in glob("habitat_baselines/config/test/*") + ], +) +def test_trainers(test_cfg_path, mode): + config = get_config(test_cfg_path).TASK_CONFIG + dataset_path = config.DATASET.DATA_PATH.format(split=config.DATASET.SPLIT) + if not os.path.exists(dataset_path): + pytest.skip( + f'No dataset "{config.DATASET.TYPE}",' + f'task "{config.TASK.TYPE}" skipped' + ) + run_exp(test_cfg_path, mode) + + +@pytest.mark.skipif( + not baseline_installed, reason="baseline sub-module not installed" +) +def test_eval_config(): + ckpt_opts = ["VIDEO_OPTION", "[]"] + eval_opts = ["VIDEO_OPTION", "['disk']"] + + ckpt_cfg = get_config(None, ckpt_opts) + assert ckpt_cfg.VIDEO_OPTION == [] + assert ckpt_cfg.CMD_TRAILING_OPTS == ["VIDEO_OPTION", "[]"] + + eval_cfg = get_config(None, eval_opts) + assert eval_cfg.VIDEO_OPTION == ["disk"] + assert eval_cfg.CMD_TRAILING_OPTS == ["VIDEO_OPTION", "['disk']"] + + trainer = BaseRLTrainer(get_config()) + assert trainer.config.VIDEO_OPTION == ["disk", "tensorboard"] + returned_config = trainer._setup_eval_config(checkpoint_config=ckpt_cfg) + assert returned_config.VIDEO_OPTION == [] + + trainer = BaseRLTrainer(eval_cfg) + returned_config = trainer._setup_eval_config(ckpt_cfg) + assert returned_config.VIDEO_OPTION == ["disk"]