import os
import argparse
import json
import jsonlines
import numpy as np
import gzip
import collections

def eval_result_file(
    result_file, force_stop=False,  max_steps=500,
):
    scores = collections.defaultdict(list)
    oracle_scores = collections.defaultdict(list)
    episode_ids = set()
    scene_ids = set()
    action_frac = collections.Counter()
    nwrong = 0
    with jsonlines.open(result_file, 'r') as f:
        for x in f:
            if np.isinf(x['infos'][0]['distance_to_goal']):
                nwrong += 1
                continue

            if x['episode_id'] in episode_ids:
                print('repeat', x['episode_id'])
                continue

            episode_ids.add(x['episode_id'])
            scene_ids.add(x['episode_id'].split('_')[0])
            scores['num_steps'].append(len(x['actions']))
            scores['collisions'].append(sum([info['collisions']['is_collision'] for info in x['infos']]))
            scores['start_dists'].append(x['infos'][0]['distance_to_goal'])
            action_frac.update(x['actions'])
            
            if max_steps < 500:
                x['infos'] = x['infos'][:min(max_steps, len(x['infos']))]
            if x['infos'][-1]['spl'] == 1:
                x['infos'][-1]['softspl'] = 1
            for k, v in x['infos'][-1].items():
                if force_stop and k == 'success':
                    v = x['infos'][-1]['distance_to_goal'] < 0.1
                if k != 'collisions':
                    scores[k].append(v)

            for t, info in enumerate(x['infos']):
                if info['distance_to_goal'] < 0.1:
                    oracle_scores['success'].append(1)
                    oracle_scores['num_steps'].append(t+1)
                    break
            else:
                oracle_scores['num_steps'].append(len(x['infos']))
                oracle_scores['success'].append(0)

    action_sum = np.sum(list(action_frac.values()))
    for a, c in action_frac.most_common():
        print(a, '%.2f'%(c/action_sum*100))
                
    print(sorted(list(scene_ids)))
    print('#episodes: %d'%(len(scores['num_steps'])))
    if nwrong > 0:
        print('\t#wrong start episodes', nwrong)
    print('min steps', np.min(scores['num_steps']))
    for k, v in scores.items():
        v = np.array(v)
        num_nans = np.sum(np.isnan(v))
        if num_nans > 0:
            print('\tNaN', num_nans)
            # print(v)
        print(k, np.mean(v[~np.isnan(v)]))

    print('\noracle scores')
    for k, v in oracle_scores.items():
        print(k, np.mean(v))
    print()

    out_keys = ['num_steps', 'collisions', 'distance_to_goal', 'success', 'spl', 'softspl', 'goal_vis_pixels', 'oracle_success']
    out_values = []
    for key in out_keys[:-1]:
        value = np.array(scores[key])
        value = np.mean(value[~np.isnan(value)])
        if key in ['success', 'spl', 'softspl', 'goal_vis_pixels']:
            value *= 100
        out_values.append(value)
    out_values.append(np.mean(oracle_scores['success']) * 100)
    print(','.join(out_keys))
    print(','.join(['%.2f'%v for v in out_values]))

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('result_file')
    parser.add_argument('--force_stop', action='store_true', default=False, help='force the agent to stop at the end of the episode')
    parser.add_argument('--max_steps', default=500, type=int)
    args = parser.parse_args()

    eval_result_file(
        args.result_file, force_stop=args.force_stop, 
        max_steps=args.max_steps,
    )

if __name__ == '__main__':
    main()


evaluate_results.py
cshizhe/onav_rim
import argparse
import habitat

from habitat.utils.visualizations.utils import observations_to_image, images_to_video, append_text_to_image

config = habitat.get_config("configs/tasks/objectnav_mp3d_il.yaml")


def make_videos(observations_list, output_prefix, ep_id):
    prefix = output_prefix + "_{}".format(ep_id)
    images_to_video(observations_list[0], output_dir="demos", video_name=prefix)


def run_reference_replay(
    cfg, num_episodes=None, output_prefix=None
):
    possible_actions = cfg.TASK.POSSIBLE_ACTIONS  
    with habitat.Env(cfg) as env:
        total_success = 0
        spl = 0

        num_episodes = min(num_episodes, len(env.episodes))
        print("Replaying {}/{} episodes".format(num_episodes, len(env.episodes)))
        for ep_id in range(num_episodes):
            observation_list = []
            env.reset()

            step_index = 1
            total_reward = 0.0
            episode = env.current_episode

            for data in env.current_episode.reference_replay[step_index:]:
                action = possible_actions.index(data.action)
                action_name = env.task.get_action_name(
                    action
                )

                observations = env.step(action=action)

                info = env.get_metrics()
                frame = observations_to_image({"rgb": observations["rgb"]}, info)
                frame = append_text_to_image(frame, "Find and go to {}".format(episode.object_category))

                observation_list.append(frame)
                if action_name == "STOP":
                    break
            make_videos([observation_list], output_prefix, ep_id)
            print("Total reward for trajectory: {}".format(total_reward))

            if len(episode.reference_replay) <= 500:
                total_success += info["success"]
                spl += info["spl"]

        print("SPL: {}, {}, {}".format(spl/num_episodes, spl, num_episodes))
        print("Success: {}, {}, {}".format(total_success/num_episodes, total_success, num_episodes))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--path", type=str, default="replays/demo_1.json.gz"
    )
    parser.add_argument(
        "--output-prefix", type=str, default="demo"
    )
    parser.add_argument(
        "--num-episodes", type=int, default=10
    )
    args = parser.parse_args()
    cfg = config
    cfg.defrost()
    cfg.DATASET.DATA_PATH = args.path

    # Set to a high limit to allow replaying episodes with 
    # number of steps greater than ObjectNav episode step
    # limit.
    cfg.ENVIRONMENT.MAX_EPISODE_STEPS = 5000

    # Set to a high limit to allow loading episodes with 
    # number of steps greater than ObjectNav episode step
    # limit in the replay buffer.
    cfg.DATASET.MAX_REPLAY_STEPS = 5000
    cfg.freeze()

    run_reference_replay(
        cfg,
        num_episodes=args.num_episodes,
        output_prefix=args.output_prefix
    )


if __name__ == "__main__":
    main()

examples/objectnav_replay.py
cshizhe/onav_rim
import argparse
import attr
import habitat
import numpy as np

from habitat import logger
from habitat.tasks.nav.shortest_path_follower import ShortestPathFollower
from habitat.sims.habitat_simulator.actions import HabitatSimActions

from habitat_sim.utils import viz_utils as vut
from habitat_sim.utils.common import quat_to_coeffs, quat_from_magnum

from scripts.utils.utils import write_json, write_gzip, load_dataset

config = habitat.get_config("configs/tasks/objectnav_mp3d.yaml")


def get_action(action):
    if action == HabitatSimActions.TURN_RIGHT:
        return "TURN_RIGHT"
    elif action == HabitatSimActions.TURN_LEFT:
        return "TURN_LEFT"
    elif action == HabitatSimActions.MOVE_FORWARD:
        return "MOVE_FORWARD"
    elif action == HabitatSimActions.LOOK_UP:
        return "LOOK_UP"
    elif action == HabitatSimActions.LOOK_DOWN:
        return "LOOK_DOWN"
    return "STOP"


def get_agent_pose(sim):
    agent_translation = sim._default_agent.body.object.translation
    agent_rotation = sim._default_agent.body.object.rotation
    sensor_data = {}
    for sensor_key, v in sim._default_agent._sensors.items():
        rotation = quat_from_magnum(v.object.rotation)
        rotation = quat_to_coeffs(rotation).tolist()
        translation = v.object.translation
        sensor_data[sensor_key] = {
            "rotation": rotation,
            "translation": np.array(translation).tolist()
        }
    
    return {
        "position": np.array(agent_translation).tolist(),
        "rotation": quat_to_coeffs(quat_from_magnum(agent_rotation)).tolist(),
        "sensor_data": sensor_data
    }


def get_action_data(action, sim):
    data = {}
    data["action"] = get_action(action)
    data["agent_state"] = get_agent_pose(sim)
    return data


def get_episode_json(episode, reference_replay):
    episode.reference_replay = reference_replay
    episode.scene_id = episode.scene_id
    ep_json = attr.asdict(episode)
    del ep_json["_shortest_path_cache"]
    return ep_json


def get_closest_goal(episode, sim):
    min_dist = 1000.0
    goal_location = None
    agent_position = sim.get_agent_state().position
    for goal in episode.goals:
        for view_point in goal.view_points:
            position = view_point.agent_state.position
            
            dist = sim.geodesic_distance(
                agent_position, position
            )
            if min_dist > dist:
                min_dist = dist
                goal_location = position
    return goal_location


def generate_trajectories(cfg, episode_path, output_path=""):
    with habitat.Env(cfg) as env:
        goal_radius = 0.1
        spl = 0
        total_success = 0.0
        total_episodes = 0.0

        dataset = load_dataset(episode_path)

        dataset["episodes"] = []

        logger.info("Total episodes: {}".format(len(env.episodes)))
        for _ in range(len(env.episodes)):
            follower = ShortestPathFollower(
                env._sim, goal_radius, False
            )
            env.reset()
            success = 0
            reference_replay = []
            episode = env.current_episode
            goal_position = get_closest_goal(episode, env.sim)
            
            info = {}
            reference_replay.append(get_action_data(HabitatSimActions.STOP, env._sim))
            if goal_position is None:
                continue
            while not env.episode_over:
                best_action = follower.get_next_action(
                    goal_position
                )

                if "distance_to_goal" in info.keys() and info["distance_to_goal"] < 0.1 and best_action != HabitatSimActions.STOP:
                    best_action = HabitatSimActions.STOP

                env.step(best_action)

                info = env.get_metrics()
                success = info["success"]

                action_data = get_action_data(best_action, env._sim)
                reference_replay.append(action_data)

            ep_data = get_episode_json(env.current_episode, reference_replay)
            total_success += success
            spl += info["spl"]
            total_episodes += 1

            dataset["episodes"].append(ep_data)
        
        print("Total episodes: {}".format(total_episodes))

        print("\n\nEpisode success: {}".format(total_success / total_episodes))
        print("SPL: {}, {}, {}".format(spl/total_episodes, spl, total_episodes))
        print("Success: {}, {}, {}".format(total_success/total_episodes, total_success, total_episodes))
        write_json(dataset, output_path)
        write_gzip(output_path, output_path)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--episodes", type=str, default="data/episodes/sampled.json.gz"
    )
    parser.add_argument(
        "--output-path", type=str, default="data/episodes/shortest_paths.json"
    )
    args = parser.parse_args()
    cfg = config
    cfg.defrost()
    cfg.DATASET.DATA_PATH = args.episodes
    cfg.freeze()

    generate_trajectories(cfg, args.episodes, output_path=args.output_path)

if __name__ == "__main__":
    main()

examples/objectnav_shortest_path_generator.py
cshizhe/onav_rim
import argparse
import habitat

from habitat.utils.visualizations.utils import observations_to_image, images_to_video

config = habitat.get_config("configs/tasks/pickplace_mp3d.yaml")


def make_videos(observations_list, output_prefix, ep_id):
    prefix = output_prefix + "_{}".format(ep_id)
    images_to_video(observations_list[0], output_dir="demos", video_name=prefix)


def run_reference_replay(cfg, restore_state=False, step_env=False, num_episodes=None, output_prefix=None):
    possible_actions = cfg.TASK.POSSIBLE_ACTIONS
    with habitat.Env(cfg) as env:
        total_success = 0
        total_spl = 0

        num_episodes = min(num_episodes, len(env.episodes))

        print("Replaying {}/{} episodes".format(num_episodes, len(env.episodes)))
        for ep_id in range(num_episodes):
            observation_list = []
            env.reset()
            i = 0
            step_index = 1

            for data in env.current_episode.reference_replay[step_index:]:
                action = possible_actions.index(data.action)            

                if step_env:
                    observations = env.step(action=action)
                elif not restore_state:
                    observations = env.step(action=action, replay_data=data)
                else:
                    agent_state = data.agent_state
                    sensor_states = data.agent_state.sensor_data
                    object_states = data.object_states
                    observations = env._sim.get_observations_at(agent_state.position, agent_state.rotation, sensor_states, object_states)

                info = env.get_metrics()
                frame = observations_to_image({"rgb": observations["rgb"]}, {})

                observation_list.append(frame)
                i+=1
            
            total_success += info["success"]
            total_spl += info["spl"]

            make_videos([observation_list], output_prefix, ep_id)

        print("Average success: {} - {} - {}".format(total_success / num_episodes, total_success, num_episodes))
        print("Average SPL: {} - {} - {}".format(total_spl / num_episodes, total_spl, num_episodes))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--path", type=str, default="replays/demo_1.json.gz"
    )
    parser.add_argument(
        "--output-prefix", type=str, default="demo"
    )
    parser.add_argument(
        "--restore-state", dest='restore_state', action='store_true'
    )
    parser.add_argument(
        "--step-env", dest='step_env', action='store_true'
    )
    parser.add_argument(
        "--num-episodes", type=int, default=10
    )
    args = parser.parse_args()
    cfg = config
    cfg.defrost()
    cfg.DATASET.TYPE = "PickPlaceDataset-v1"
    cfg.DATASET.DATA_PATH = args.path
    cfg.freeze()

    run_reference_replay(
        cfg,
        args.restore_state,
        args.step_env,
        num_episodes=args.num_episodes,
        output_prefix=args.output_prefix
    )

if __name__ == "__main__":
    main()

examples/pickplace_replay.py
cshizhe/onav_rim
#!/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
import argparse
from tarfile import is_tarfile

import numpy as np

import habitat
from habitat.core.utils import try_cv2_import
from habitat.tasks.nav.shortest_path_follower import ShortestPathFollower
from habitat.utils.visualizations import maps
from habitat.utils.visualizations.utils import images_to_video

from habitat_baselines.config.default import get_config
from habitat_baselines.common.environments import NavRLEnv

cv2 = try_cv2_import()


def draw_top_down_map(info, output_size):
    return maps.colorize_draw_agent_and_fit_to_height(
        info["top_down_map"], output_size
    )


def viz_human_demos(scene_id, output_dir, num_episodes=10, shortest_path=False):
    config = get_config(
        config_paths="habitat_baselines/config/objectnav/il_objectnav.yaml"
    )
    print(config)
    config.defrost()
    config.TASK_CONFIG.TASK.MEASUREMENTS.append("TOP_DOWN_MAP")
    config.TASK_CONFIG.DATASET.CONTENT_SCENES = [scene_id]
    config.freeze()        

    with NavRLEnv(config=config) as env:
        print("Environment creation successful")
        for i in range(num_episodes):
            env.reset()
            current_episode = env.current_episode

            dirname = os.path.join(output_dir, scene_id)
            os.makedirs(dirname, exist_ok=True)
            print("Agent stepping around inside environment.")

            if shortest_path:
                if current_episode.is_thda:
                    goal_radius = 1.0
                    follower = ShortestPathFollower(
                        env.habitat_env.sim, goal_radius, False
                    )
                else:
                    actions = [x.action if x.action is not None else 0 for x in current_episode.shortest_paths[0]]

            else:
                actions = [x.action for x in current_episode.reference_replay]
                assert actions[0] == 'STOP'
                actions = actions[1:]

            images = []
            t = 0
            while True:
                if shortest_path and current_episode.is_thda:
                    action = follower.get_next_action(
                        env.habitat_env.current_episode.goals[0].position
                    )   # int
                else:
                    action = actions[t]

                if action == 'STOP' or action == 0:
                    break
                observations, reward, done, info = env.step(action=action)
                im = observations["rgb"]
                top_down_map = draw_top_down_map(info, im.shape[0])
                output_im = np.concatenate((im, top_down_map), axis=1)
                images.append(output_im)
                t += 1

            images_to_video(
                images, dirname, 
                f'{current_episode.object_category}-{current_episode.episode_id}'
            )
            print(f"{i} Episode finished")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--output_dir', required=True)
    parser.add_argument('--num_episodes', type=int, default=10)
    parser.add_argument('--shortest_path', action='store_true', default=False)
    args = parser.parse_args()

    scene_ids = [
        x.split('.')[0] for x in \
            os.listdir('data/datasets/objectnav/objectnav_mp3d_70k/train/content')
    ]
    for scene_id in scene_ids[1:10]:
        viz_human_demos(
            scene_id, args.output_dir, 
            num_episodes=args.num_episodes,
            shortest_path=args.shortest_path,
        )


if __name__ == "__main__":
    main()

examples/visualize_human_demonstrations.py
cshizhe/onav_rim
#!/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
import argparse
from tarfile import is_tarfile

import numpy as np

import habitat
from habitat.core.utils import try_cv2_import
from habitat.tasks.nav.shortest_path_follower import ShortestPathFollower
from habitat.utils.visualizations import maps
from habitat.utils.visualizations.utils import images_to_video

from habitat_baselines.config.default import get_config
from habitat_baselines.common.environments import NavRLEnv

cv2 = try_cv2_import()


def draw_top_down_map(info, output_size):
    return maps.colorize_draw_agent_and_fit_to_height(
        info["top_down_map"], output_size
    )


def viz_human_demos(scene_id, output_dir, num_episodes=10, shortest_path=False):
    config = get_config(
        config_paths="habitat_baselines/config/objectnav/il_objectnav.yaml"
    )
    print(config)
    config.defrost()
    config.TASK_CONFIG.TASK.MEASUREMENTS.append("TOP_DOWN_MAP")
    config.TASK_CONFIG.DATASET.CONTENT_SCENES = [scene_id]
    config.freeze()        

    with NavRLEnv(config=config) as env:
        print("Environment creation successful")
        for i in range(num_episodes):
            env.reset()
            current_episode = env.current_episode

            dirname = os.path.join(output_dir, scene_id)
            os.makedirs(dirname, exist_ok=True)
            print("Agent stepping around inside environment.")

            if shortest_path:
                if current_episode.is_thda:
                    goal_radius = 1.0
                    follower = ShortestPathFollower(
                        env.habitat_env.sim, goal_radius, False
                    )
                else:
                    actions = [x.action if x.action is not None else 0 for x in current_episode.shortest_paths[0]]

            else:
                actions = [x.action for x in current_episode.reference_replay]
                assert actions[0] == 'STOP'
                actions = actions[1:]

            images = []
            t = 0
            while True:
                if shortest_path and current_episode.is_thda:
                    action = follower.get_next_action(
                        env.habitat_env.current_episode.goals[0].position
                    )   # int
                else:
                    action = actions[t]

                if action == 'STOP' or action == 0:
                    break
                observations, reward, done, info = env.step(action=action)
                im = observations["rgb"]
                top_down_map = draw_top_down_map(info, im.shape[0])
                output_im = np.concatenate((im, top_down_map), axis=1)
                images.append(output_im)
                t += 1

            images_to_video(
                images, dirname, 
                f'{current_episode.object_category}-{current_episode.episode_id}'
            )
            print(f"{i} Episode finished")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--output_dir', required=True)
    parser.add_argument('--num_episodes', type=int, default=10)
    parser.add_argument('--shortest_path', action='store_true', default=False)
    args = parser.parse_args()

    scene_ids = [
        x.split('.')[0] for x in \
            os.listdir('data/datasets/objectnav/objectnav_mp3d_70k/train/content')
    ]
    for scene_id in scene_ids[1:10]:
        viz_human_demos(
            scene_id, args.output_dir, 
            num_episodes=args.num_episodes,
            shortest_path=args.shortest_path,
        )


if __name__ == "__main__":
    main()

examples/visualize_pred_results.py
cshizhe/onav_rim
#!/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.

from typing import List, Optional, Union

import yacs.config


# Default Habitat config node
class Config(yacs.config.CfgNode):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs, new_allowed=True)


CN = Config

DEFAULT_CONFIG_DIR = "configs/"
CONFIG_FILE_SEPARATOR = ","

# -----------------------------------------------------------------------------
# Config definition
# -----------------------------------------------------------------------------
_C = CN()
_C.SEED = 100
# -----------------------------------------------------------------------------
# ENVIRONMENT
# -----------------------------------------------------------------------------
_C.ENVIRONMENT = CN()
_C.ENVIRONMENT.MAX_EPISODE_STEPS = 1000
_C.ENVIRONMENT.MAX_EPISODE_SECONDS = 10000000
_C.ENVIRONMENT.ITERATOR_OPTIONS = CN()
_C.ENVIRONMENT.ITERATOR_OPTIONS.CYCLE = True
_C.ENVIRONMENT.ITERATOR_OPTIONS.SHUFFLE = True
_C.ENVIRONMENT.ITERATOR_OPTIONS.GROUP_BY_SCENE = True
_C.ENVIRONMENT.ITERATOR_OPTIONS.NUM_EPISODE_SAMPLE = -1
_C.ENVIRONMENT.ITERATOR_OPTIONS.MAX_SCENE_REPEAT_EPISODES = -1
_C.ENVIRONMENT.ITERATOR_OPTIONS.MAX_SCENE_REPEAT_STEPS = int(1e4)
_C.ENVIRONMENT.ITERATOR_OPTIONS.STEP_REPETITION_RANGE = 0.2
# -----------------------------------------------------------------------------
# TASK
# -----------------------------------------------------------------------------
# -----------------------------------------------------------------------------
# # NAVIGATION TASK
# -----------------------------------------------------------------------------
_C.TASK = CN()
_C.TASK.TYPE = "Nav-v0"
_C.TASK.SUCCESS_DISTANCE = 0.2
_C.TASK.SENSORS = []
_C.TASK.MEASUREMENTS = []
_C.TASK.GOAL_SENSOR_UUID = "pointgoal"
_C.TASK.POSSIBLE_ACTIONS = ["STOP", "MOVE_FORWARD", "TURN_LEFT", "TURN_RIGHT"]
# -----------------------------------------------------------------------------
# # ACTIONS
# -----------------------------------------------------------------------------
ACTIONS = CN()
ACTIONS.STOP = CN()
ACTIONS.STOP.TYPE = "StopAction"
# -----------------------------------------------------------------------------
# # NAVIGATION ACTIONS
# -----------------------------------------------------------------------------
ACTIONS.MOVE_FORWARD = CN()
ACTIONS.MOVE_FORWARD.TYPE = "MoveForwardAction"
ACTIONS.TURN_LEFT = CN()
ACTIONS.TURN_LEFT.TYPE = "TurnLeftAction"
ACTIONS.TURN_RIGHT = CN()
ACTIONS.TURN_RIGHT.TYPE = "TurnRightAction"
ACTIONS.LOOK_UP = CN()
ACTIONS.LOOK_UP.TYPE = "LookUpAction"
ACTIONS.LOOK_DOWN = CN()
ACTIONS.LOOK_DOWN.TYPE = "LookDownAction"
ACTIONS.TELEPORT = CN()
ACTIONS.TELEPORT.TYPE = "TeleportAction"

_C.TASK.ACTIONS = ACTIONS
# -----------------------------------------------------------------------------
# # TASK SENSORS
# -----------------------------------------------------------------------------
# -----------------------------------------------------------------------------
# POINTGOAL SENSOR
# -----------------------------------------------------------------------------
_C.TASK.POINTGOAL_SENSOR = CN()
_C.TASK.POINTGOAL_SENSOR.TYPE = "PointGoalSensor"
_C.TASK.POINTGOAL_SENSOR.GOAL_FORMAT = "POLAR"
_C.TASK.POINTGOAL_SENSOR.DIMENSIONALITY = 2
# -----------------------------------------------------------------------------
# POINTGOAL WITH GPS+COMPASS SENSOR
# -----------------------------------------------------------------------------
_C.TASK.POINTGOAL_WITH_GPS_COMPASS_SENSOR = _C.TASK.POINTGOAL_SENSOR.clone()
_C.TASK.POINTGOAL_WITH_GPS_COMPASS_SENSOR.TYPE = (
    "PointGoalWithGPSCompassSensor"
)
# -----------------------------------------------------------------------------
# OBJECTGOAL SENSOR
# -----------------------------------------------------------------------------
_C.TASK.OBJECTGOAL_SENSOR = CN()
_C.TASK.OBJECTGOAL_SENSOR.TYPE = "ObjectGoalSensor"
_C.TASK.OBJECTGOAL_SENSOR.GOAL_SPEC = "TASK_CATEGORY_ID"
_C.TASK.OBJECTGOAL_SENSOR.GOAL_SPEC_MAX_VAL = 50
# -----------------------------------------------------------------------------
# IMAGEGOAL SENSOR
# -----------------------------------------------------------------------------
_C.TASK.IMAGEGOAL_SENSOR = CN()
_C.TASK.IMAGEGOAL_SENSOR.TYPE = "ImageGoalSensor"
# -----------------------------------------------------------------------------
# HEADING SENSOR
# -----------------------------------------------------------------------------
_C.TASK.HEADING_SENSOR = CN()
_C.TASK.HEADING_SENSOR.TYPE = "HeadingSensor"
# -----------------------------------------------------------------------------
# COMPASS SENSOR
# -----------------------------------------------------------------------------
_C.TASK.COMPASS_SENSOR = CN()
_C.TASK.COMPASS_SENSOR.TYPE = "CompassSensor"
# -----------------------------------------------------------------------------
# GPS SENSOR
# -----------------------------------------------------------------------------
_C.TASK.GPS_SENSOR = CN()
_C.TASK.GPS_SENSOR.TYPE = "GPSSensor"
_C.TASK.GPS_SENSOR.DIMENSIONALITY = 2
# -----------------------------------------------------------------------------
# ALL OBJECT POSITION SENSOR
# -----------------------------------------------------------------------------
_C.TASK.ALL_OBJECT_POSITIONS = CN()
_C.TASK.ALL_OBJECT_POSITIONS.TYPE = "AllObjectPositions"
# -----------------------------------------------------------------------------
# PROXIMITY SENSOR
# -----------------------------------------------------------------------------
_C.TASK.PROXIMITY_SENSOR = CN()
_C.TASK.PROXIMITY_SENSOR.TYPE = "ProximitySensor"
_C.TASK.PROXIMITY_SENSOR.MAX_DETECTION_RADIUS = 2.0
# -----------------------------------------------------------------------------
# SUCCESS MEASUREMENT
# -----------------------------------------------------------------------------
_C.TASK.SUCCESS = CN()
_C.TASK.SUCCESS.TYPE = "Success"
_C.TASK.SUCCESS.SUCCESS_DISTANCE = 0.2
# -----------------------------------------------------------------------------
# SPL MEASUREMENT
# -----------------------------------------------------------------------------
_C.TASK.SPL = CN()
_C.TASK.SPL.TYPE = "SPL"
# -----------------------------------------------------------------------------
# SOFT-SPL MEASUREMENT
# -----------------------------------------------------------------------------
_C.TASK.SOFT_SPL = CN()
_C.TASK.SOFT_SPL.TYPE = "SoftSPL"
# -----------------------------------------------------------------------------
# TopDownMap MEASUREMENT
# -----------------------------------------------------------------------------
_C.TASK.TOP_DOWN_MAP = CN()
_C.TASK.TOP_DOWN_MAP.TYPE = "TopDownMap"
_C.TASK.TOP_DOWN_MAP.MAX_EPISODE_STEPS = _C.ENVIRONMENT.MAX_EPISODE_STEPS
_C.TASK.TOP_DOWN_MAP.MAP_PADDING = 3
_C.TASK.TOP_DOWN_MAP.MAP_RESOLUTION = 256
_C.TASK.TOP_DOWN_MAP.DRAW_SOURCE = True
_C.TASK.TOP_DOWN_MAP.DRAW_BORDER = True
_C.TASK.TOP_DOWN_MAP.DRAW_SHORTEST_PATH = True
_C.TASK.TOP_DOWN_MAP.FOG_OF_WAR = CN()
_C.TASK.TOP_DOWN_MAP.FOG_OF_WAR.DRAW = True
_C.TASK.TOP_DOWN_MAP.FOG_OF_WAR.VISIBILITY_DIST = 5.0
_C.TASK.TOP_DOWN_MAP.FOG_OF_WAR.FOV = 90
_C.TASK.TOP_DOWN_MAP.DRAW_VIEW_POINTS = True
_C.TASK.TOP_DOWN_MAP.DRAW_GOAL_POSITIONS = True
# Axes aligned bounding boxes
_C.TASK.TOP_DOWN_MAP.DRAW_GOAL_AABBS = True

# -----------------------------------------------------------------------------
# # COVERAGE MEASUREMENT
# -----------------------------------------------------------------------------
_C.TASK.COVERAGE = CN()
_C.TASK.COVERAGE.TYPE = "Coverage"
_C.TASK.COVERAGE.GRID_DELTA = 0.25
_C.TASK.COVERAGE.EGOCENTRIC = False # Agent will spawn centered on a tile, as opposed to global


_C.TASK.GOAL_OBJECT_VISIBLE = CN() # nb: this is fraction of view, not pixels
_C.TASK.GOAL_OBJECT_VISIBLE.TYPE = "GoalObjectVisible"

# -----------------------------------------------------------------------------
# ETN REWARD MEASUREMENT
# -----------------------------------------------------------------------------
# The next reward (OBJNAV_REWARD) was ported from master, this is my own simpler writeup.
_C.TASK.ETN_REWARD = CN()
_C.TASK.ETN_REWARD.TYPE = "ExploreThenNavReward"
_C.TASK.ETN_REWARD.EXPLORE_GOAL_SEEN_THRESHOLD = 0.02
# Duplicted from coverage since measure only receives its own config
_C.TASK.ETN_REWARD.REWARD = 0.25
_C.TASK.ETN_REWARD.ATTENUATION = 0.995
_C.TASK.ETN_REWARD.VISIT_EXP = 1

# -----------------------------------------------------------------------------
# COLLISIONS MEASUREMENT
# -----------------------------------------------------------------------------
_C.TASK.COLLISIONS = CN()
_C.TASK.COLLISIONS.TYPE = "Collisions"
# -----------------------------------------------------------------------------
# -----------------------------------------------------------------------------
# # EQA TASK
# -----------------------------------------------------------------------------
_C.TASK.ACTIONS.ANSWER = CN()
_C.TASK.ACTIONS.ANSWER.TYPE = "AnswerAction"
# # EQA TASK QUESTION SENSOR
# -----------------------------------------------------------------------------
_C.TASK.QUESTION_SENSOR = CN()
_C.TASK.QUESTION_SENSOR.TYPE = "QuestionSensor"
# -----------------------------------------------------------------------------
# # EQA TASK CORRECT_ANSWER measure for training
# -----------------------------------------------------------------------------
_C.TASK.CORRECT_ANSWER = CN()
_C.TASK.CORRECT_ANSWER.TYPE = "CorrectAnswer"
# -----------------------------------------------------------------------------
# # EQA TASK ANSWER SENSOR
# -----------------------------------------------------------------------------
_C.TASK.EPISODE_INFO = CN()
_C.TASK.EPISODE_INFO.TYPE = "EpisodeInfo"
# -----------------------------------------------------------------------------
# # VLN TASK INSTRUCTION SENSOR
# -----------------------------------------------------------------------------
_C.TASK.INSTRUCTION_SENSOR = CN()
_C.TASK.INSTRUCTION_SENSOR.TYPE = "InstructionSensor"
_C.TASK.INSTRUCTION_SENSOR_UUID = "instruction"
# -----------------------------------------------------------------------------
# # DISTANCE_TO_GOAL MEASUREMENT
# -----------------------------------------------------------------------------
_C.TASK.DISTANCE_TO_GOAL = CN()
_C.TASK.DISTANCE_TO_GOAL.TYPE = "DistanceToGoal"
_C.TASK.DISTANCE_TO_GOAL.DISTANCE_TO = "POINT"
# -----------------------------------------------------------------------------
# # ANSWER_ACCURACY MEASUREMENT
# -----------------------------------------------------------------------------
_C.TASK.ANSWER_ACCURACY = CN()
_C.TASK.ANSWER_ACCURACY.TYPE = "AnswerAccuracy"

_C.TASK.OBJ_GEN_NEAR_DIST = 1.0
_C.TASK.OBJ_GEN_FAR_DIST = 5.0
_C.TASK.INCREASE_DIST = False
_C.TASK.NUM_UNIQ_SELECTED_OBJECTS = 5
_C.TASK.NUM_OBJECT_COPIES = 3
# -----------------------------------------------------------------------------
# SIMULATOR
# -----------------------------------------------------------------------------
_C.SIMULATOR = CN()
_C.SIMULATOR.TYPE = "Sim-v0"
_C.SIMULATOR.ACTION_SPACE_CONFIG = "v0"
_C.SIMULATOR.FORWARD_STEP_SIZE = 0.25  # in metres
_C.SIMULATOR.SCENE = (
    "data/scene_datasets/habitat-test-scenes/van-gogh-room.glb"
)
_C.SIMULATOR.SEED = _C.SEED
_C.SIMULATOR.TURN_ANGLE = 10  # angle to rotate left or right in degrees
_C.SIMULATOR.TILT_ANGLE = 15  # angle to tilt the camera up or down in degrees
_C.SIMULATOR.DEFAULT_AGENT_ID = 0
# -----------------------------------------------------------------------------
# SIMULATOR SENSORS
# -----------------------------------------------------------------------------
SIMULATOR_SENSOR = CN()
SIMULATOR_SENSOR.HEIGHT = 480
SIMULATOR_SENSOR.WIDTH = 640
SIMULATOR_SENSOR.HFOV = 90  # horizontal field of view in degrees
SIMULATOR_SENSOR.POSITION = [0, 1.25, 0]
SIMULATOR_SENSOR.ORIENTATION = [0.0, 0.0, 0.0]  # Euler's angles
# -----------------------------------------------------------------------------
# RGB SENSOR
# -----------------------------------------------------------------------------
_C.SIMULATOR.RGB_SENSOR = SIMULATOR_SENSOR.clone()
_C.SIMULATOR.RGB_SENSOR.TYPE = "HabitatSimRGBSensor"
# -----------------------------------------------------------------------------
# DEPTH SENSOR
# -----------------------------------------------------------------------------
_C.SIMULATOR.DEPTH_SENSOR = SIMULATOR_SENSOR.clone()
_C.SIMULATOR.DEPTH_SENSOR.TYPE = "HabitatSimDepthSensor"
_C.SIMULATOR.DEPTH_SENSOR.MIN_DEPTH = 0.0
_C.SIMULATOR.DEPTH_SENSOR.MAX_DEPTH = 10.0
_C.SIMULATOR.DEPTH_SENSOR.NORMALIZE_DEPTH = True
# -----------------------------------------------------------------------------
# SEMANTIC SENSOR
# -----------------------------------------------------------------------------
_C.SIMULATOR.SEMANTIC_SENSOR = SIMULATOR_SENSOR.clone()
_C.SIMULATOR.SEMANTIC_SENSOR.TYPE = "HabitatSimSemanticSensor"
# -----------------------------------------------------------------------------
# AGENT
# -----------------------------------------------------------------------------
_C.SIMULATOR.AGENT_0 = CN()
_C.SIMULATOR.AGENT_0.HEIGHT = 1.5
_C.SIMULATOR.AGENT_0.RADIUS = 0.1
_C.SIMULATOR.AGENT_0.MASS = 32.0
_C.SIMULATOR.AGENT_0.LINEAR_ACCELERATION = 20.0
_C.SIMULATOR.AGENT_0.ANGULAR_ACCELERATION = 4 * 3.14
_C.SIMULATOR.AGENT_0.LINEAR_FRICTION = 0.5
_C.SIMULATOR.AGENT_0.ANGULAR_FRICTION = 1.0
_C.SIMULATOR.AGENT_0.COEFFICIENT_OF_RESTITUTION = 0.0
_C.SIMULATOR.AGENT_0.SENSORS = ["RGB_SENSOR"]
_C.SIMULATOR.AGENT_0.IS_SET_START_STATE = False
_C.SIMULATOR.AGENT_0.START_POSITION = [0, 0, 0]
_C.SIMULATOR.AGENT_0.START_ROTATION = [0, 0, 0, 1]
_C.SIMULATOR.AGENTS = ["AGENT_0"]
# -----------------------------------------------------------------------------
# SIMULATOR HABITAT_SIM_V0
# -----------------------------------------------------------------------------
_C.SIMULATOR.HABITAT_SIM_V0 = CN()
_C.SIMULATOR.HABITAT_SIM_V0.GPU_DEVICE_ID = 0
# Use Habitat-Sim's GPU->GPU copy mode to return rendering results
# in PyTorch tensors.  Requires Habitat-Sim to be built
# with --with-cuda
# This will generally imply sharing CUDA tensors between processes.
# Read here: https://pytorch.org/docs/stable/multiprocessing.html#sharing-cuda-tensors
# for the caveats that results in
_C.SIMULATOR.HABITAT_SIM_V0.GPU_GPU = False
# Whether or not the agent slides on collisions
_C.SIMULATOR.HABITAT_SIM_V0.ALLOW_SLIDING = True
_C.SIMULATOR.HABITAT_SIM_V0.ENABLE_PHYSICS = False
_C.SIMULATOR.HABITAT_SIM_V0.PHYSICS_CONFIG_FILE = (
    "./data/default.physics_config.json"
)
# -----------------------------------------------------------------------------
# PYROBOT
# -----------------------------------------------------------------------------
_C.PYROBOT = CN()
_C.PYROBOT.ROBOTS = ["locobot"]  # types of robots supported
_C.PYROBOT.ROBOT = "locobot"
_C.PYROBOT.SENSORS = ["RGB_SENSOR", "DEPTH_SENSOR", "BUMP_SENSOR"]
_C.PYROBOT.BASE_CONTROLLER = "proportional"
_C.PYROBOT.BASE_PLANNER = "none"
# -----------------------------------------------------------------------------
# SENSORS
# -----------------------------------------------------------------------------
PYROBOT_VISUAL_SENSOR = CN()
PYROBOT_VISUAL_SENSOR.HEIGHT = 480
PYROBOT_VISUAL_SENSOR.WIDTH = 640
# -----------------------------------------------------------------------------
# RGB SENSOR
# -----------------------------------------------------------------------------
_C.PYROBOT.RGB_SENSOR = PYROBOT_VISUAL_SENSOR.clone()
_C.PYROBOT.RGB_SENSOR.TYPE = "PyRobotRGBSensor"
_C.PYROBOT.RGB_SENSOR.CENTER_CROP = False
# -----------------------------------------------------------------------------
# DEPTH SENSOR
# -----------------------------------------------------------------------------
_C.PYROBOT.DEPTH_SENSOR = PYROBOT_VISUAL_SENSOR.clone()
_C.PYROBOT.DEPTH_SENSOR.TYPE = "PyRobotDepthSensor"
_C.PYROBOT.DEPTH_SENSOR.MIN_DEPTH = 0.0
_C.PYROBOT.DEPTH_SENSOR.MAX_DEPTH = 5.0
_C.PYROBOT.DEPTH_SENSOR.NORMALIZE_DEPTH = True
_C.PYROBOT.DEPTH_SENSOR.CENTER_CROP = False
# -----------------------------------------------------------------------------
# BUMP SENSOR
# -----------------------------------------------------------------------------
_C.PYROBOT.BUMP_SENSOR = CN()
_C.PYROBOT.BUMP_SENSOR.TYPE = "PyRobotBumpSensor"
# -----------------------------------------------------------------------------
# ACTIONS LOCOBOT
# -----------------------------------------------------------------------------
_C.PYROBOT.LOCOBOT = CN()
_C.PYROBOT.LOCOBOT.ACTIONS = ["BASE_ACTIONS", "CAMERA_ACTIONS"]
_C.PYROBOT.LOCOBOT.BASE_ACTIONS = ["go_to_relative", "go_to_absolute"]
_C.PYROBOT.LOCOBOT.CAMERA_ACTIONS = ["set_pan", "set_tilt", "set_pan_tilt"]
# TODO(akadian): add support for Arm actions
# -----------------------------------------------------------------------------
# DATASET
# -----------------------------------------------------------------------------
_C.DATASET = CN()
_C.DATASET.TYPE = "PointNav-v1"
_C.DATASET.SPLIT = "train"
_C.DATASET.SCENES_DIR = "data/scene_datasets"
_C.DATASET.CONTENT_SCENES = ["*"]
_C.DATASET.DATA_PATH = (
    "data/datasets/pointnav/habitat-test-scenes/v1/{split}/{split}.json.gz"
)

# -----------------------------------------------------------------------------


def get_config(
    config_paths: Optional[Union[List[str], str]] = None,
    opts: Optional[list] = None,
) -> CN:
    r"""Create a unified config with default values overwritten by values from
    :p:`config_paths` and overwritten by options from :p:`opts`.

    :param config_paths: List of config paths or string that contains comma
        separated list of config paths.
    :param opts: Config options (keys, values) in a list (e.g., passed from
        command line into the config. For example,
        :py:`opts = ['FOO.BAR', 0.5]`. Argument can be used for parameter
        sweeping or quick tests.
    """
    config = _C.clone()
    if config_paths:
        if isinstance(config_paths, str):
            if CONFIG_FILE_SEPARATOR in config_paths:
                config_paths = config_paths.split(CONFIG_FILE_SEPARATOR)
            else:
                config_paths = [config_paths]

        for config_path in config_paths:
            config.merge_from_file(config_path)

    if opts:
        config.merge_from_list(opts)

    config.freeze()
    return config

habitat/config/default.py
cshizhe/onav_rim
#!/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.
r"""Implements dataset functionality to be used ``habitat.EmbodiedTask``.
``habitat.core.dataset`` abstracts over a collection of
``habitat.core.Episode``. Each episode consists of a single instantiation
of a ``habitat.Agent`` inside ``habitat.Env``.
"""
import copy
import json
import os
import random
from itertools import groupby
from typing import (
    Any,
    Callable,
    Dict,
    Generic,
    Iterator,
    List,
    Optional,
    Sequence,
    TypeVar,
    Union,
)

import attr
import numpy as np
from numpy import ndarray

from habitat.config import Config
from habitat.core.utils import not_none_validator

ALL_SCENES_MASK = "*"


@attr.s(auto_attribs=True, kw_only=True)
class ObjectInScene:
    object_id: int = attr.ib(default=None, validator=not_none_validator)
    semantic_category_id: int = attr.ib(default=None)
    object_template: str = attr.ib(default=None, validator=not_none_validator)
    scale: float = attr.ib(default=None)
    position: List[float] = attr.ib(default=None)
    rotation: List[float] = attr.ib(default=None)


@attr.s(auto_attribs=True, kw_only=True)
class SceneState:
    objects: List[ObjectInScene] = attr.ib(default=None)


@attr.s(auto_attribs=True, kw_only=True)
class Episode:
    r"""Base class for episode specification that includes initial position and
    rotation of agent, scene id, episode.

    :property episode_id: id of episode in the dataset, usually episode number.
    :property scene_id: id of scene in dataset.
    :property start_position: list of length 3 for cartesian coordinates
        :py:`(x, y, z)`.
    :property start_rotation: list of length 4 for (x, y, z, w) elements
        of unit quaternion (versor) representing 3D agent orientation
        (https://en.wikipedia.org/wiki/Versor). The rotation specifying the
        agent's orientation is relative to the world coordinate axes.

    This information is provided by a :ref:`Dataset` instance.
    """

    episode_id: str = attr.ib(default=None, validator=not_none_validator)
    scene_id: str = attr.ib(default=None, validator=not_none_validator)
    start_position: List[float] = attr.ib(
        default=None, validator=not_none_validator
    )
    start_rotation: List[float] = attr.ib(
        default=None, validator=not_none_validator
    )
    info: Optional[Dict[str, Any]] = None
    _shortest_path_cache: Any = attr.ib(init=False, default=None)

    def __getstate__(self):
        return {
            k: v
            for k, v in self.__dict__.items()
            if k not in {"_shortest_path_cache"}
        }

    def __setstate__(self, state):
        self.__dict__.update(state)
        self.__dict__["_shortest_path_cache"] = None


T = TypeVar("T", bound=Episode)


class Dataset(Generic[T]):
    r"""Base class for dataset specification."""
    episodes: List[T]

    @staticmethod
    def scene_from_scene_path(scene_path: str) -> str:
        r"""Helper method to get the scene name from an episode.

        :param scene_path: The path to the scene, assumes this is formatted
                            ``/path/to/<scene_name>.<ext>``

        :return: <scene_name> from the path
        """
        return os.path.splitext(os.path.basename(scene_path))[0]

    @classmethod
    def get_scenes_to_load(cls, config: Config) -> List[str]:
        r"""Returns a list of scene names that would be loaded with this dataset.

        Useful for determing what scenes to split up among different workers.

        :param config: The config for the dataset

        :return: A list of scene names that would be loaded with the dataset
        """
        assert cls.check_config_paths_exist(config)  # type: ignore[attr-defined]
        dataset = cls(config)  # type: ignore[call-arg]
        return list(map(cls.scene_from_scene_path, dataset.scene_ids))

    @classmethod
    def build_content_scenes_filter(cls, config) -> Callable[[T], bool]:
        r"""Returns a filter function that takes an episode and returns True if that
        episode is valid under the CONTENT_SCENES feild of the provided config
        """
        scenes_to_load = set(config.CONTENT_SCENES)

        def _filter(ep: T) -> bool:
            return (
                ALL_SCENES_MASK in scenes_to_load
                or cls.scene_from_scene_path(ep.scene_id) in scenes_to_load
            )

        return _filter

    @property
    def num_episodes(self) -> int:
        r"""number of episodes in the dataset"""
        return len(self.episodes)

    @property
    def scene_ids(self) -> List[str]:
        r"""unique scene ids present in the dataset."""
        return sorted({episode.scene_id for episode in self.episodes})

    def get_scene_episodes(self, scene_id: str) -> List[T]:
        r"""..

        :param scene_id: id of scene in scene dataset.
        :return: list of episodes for the :p:`scene_id`.
        """
        return list(
            filter(lambda x: x.scene_id == scene_id, iter(self.episodes))
        )

    def get_episodes(self, indexes: List[int]) -> List[T]:
        r"""..

        :param indexes: episode indices in dataset.
        :return: list of episodes corresponding to indexes.
        """
        return [self.episodes[episode_id] for episode_id in indexes]

    def get_episode_iterator(self, *args: Any, **kwargs: Any) -> Iterator:
        r"""Gets episode iterator with options. Options are specified in
        :ref:`EpisodeIterator` documentation.

        :param args: positional args for iterator constructor
        :param kwargs: keyword args for iterator constructor
        :return: episode iterator with specified behavior

        To further customize iterator behavior for your :ref:`Dataset`
        subclass, create a customized iterator class like
        :ref:`EpisodeIterator` and override this method.
        """
        return EpisodeIterator(self.episodes, *args, **kwargs)

    def to_json(self) -> str:
        class DatasetJSONEncoder(json.JSONEncoder):
            def default(self, obj):
                if isinstance(obj, np.ndarray):
                    return obj.tolist()

                return (
                    obj.__getstate__()
                    if hasattr(obj, "__getstate__")
                    else obj.__dict__
                )

        result = DatasetJSONEncoder().encode(self)
        return result

    def from_json(
        self, json_str: str, scenes_dir: Optional[str] = None
    ) -> None:
        r"""Creates dataset from :p:`json_str`.

        :param json_str: JSON string containing episodes information.
        :param scenes_dir: directory containing graphical assets relevant
            for episodes present in :p:`json_str`.

        Directory containing relevant graphical assets of scenes is passed
        through :p:`scenes_dir`.
        """
        raise NotImplementedError

    def filter_episodes(self, filter_fn: Callable[[T], bool]) -> "Dataset":
        r"""Returns a new dataset with only the filtered episodes from the
        original dataset.

        :param filter_fn: function used to filter the episodes.
        :return: the new dataset.
        """
        new_episodes = []
        for episode in self.episodes:
            if filter_fn(episode):
                new_episodes.append(episode)
        new_dataset = copy.copy(self)
        new_dataset.episodes = new_episodes
        return new_dataset

    def get_splits(
        self,
        num_splits: int,
        episodes_per_split: Optional[int] = None,
        remove_unused_episodes: bool = False,
        collate_scene_ids: bool = True,
        sort_by_episode_id: bool = False,
        allow_uneven_splits: bool = False,
    ) -> List["Dataset"]:
        r"""Returns a list of new datasets, each with a subset of the original
        episodes.

        :param num_splits: the number of splits to create.
        :param episodes_per_split: if provided, each split will have up to this
            many episodes. If it is not provided, each dataset will have
            :py:`len(original_dataset.episodes) // num_splits` episodes. If
            max_episodes_per_split is provided and is larger than this value,
            it will be capped to this value.
        :param remove_unused_episodes: once the splits are created, the extra
            episodes will be destroyed from the original dataset. This saves
            memory for large datasets.
        :param collate_scene_ids: if true, episodes with the same scene id are
            next to each other. This saves on overhead of switching between
            scenes, but means multiple sequential episodes will be related to
            each other because they will be in the same scene.
        :param sort_by_episode_id: if true, sequences are sorted by their
            episode ID in the returned splits.
        :param allow_uneven_splits: if true, the last splits can be shorter
            than the others. This is especially useful for splitting over
            validation/test datasets in order to make sure that all episodes
            are copied but none are duplicated.
        :return: a list of new datasets, each with their own subset of
            episodes.

        All splits will have the same number of episodes, but no episodes will
        be duplicated.
        """
        if self.num_episodes < num_splits:
            raise ValueError(
                "Not enough episodes to create those many splits."
            )

        if episodes_per_split is not None:
            if allow_uneven_splits:
                raise ValueError(
                    "You probably don't want to specify allow_uneven_splits"
                    " and episodes_per_split."
                )

            if num_splits * episodes_per_split > self.num_episodes:
                raise ValueError(
                    "Not enough episodes to create those many splits."
                )

        new_datasets = []

        if episodes_per_split is not None:
            stride = episodes_per_split
        else:
            stride = self.num_episodes // num_splits
        split_lengths = [stride] * num_splits

        if allow_uneven_splits:
            episodes_left = self.num_episodes - stride * num_splits
            split_lengths[:episodes_left] = [stride + 1] * episodes_left
            assert sum(split_lengths) == self.num_episodes

        num_episodes = sum(split_lengths)

        rand_items = np.random.choice(
            self.num_episodes, num_episodes, replace=False
        )
        if collate_scene_ids:
            scene_ids: Dict[str, List[int]] = {}
            for rand_ind in rand_items:
                scene = self.episodes[rand_ind].scene_id
                if scene not in scene_ids:
                    scene_ids[scene] = []
                scene_ids[scene].append(rand_ind)
            rand_items = []
            list(map(rand_items.extend, scene_ids.values()))
        ep_ind = 0
        new_episodes = []
        for nn in range(num_splits):
            new_dataset = copy.copy(self)  # Creates a shallow copy
            new_dataset.episodes = []
            new_datasets.append(new_dataset)
            for _ii in range(split_lengths[nn]):
                new_dataset.episodes.append(self.episodes[rand_items[ep_ind]])
                ep_ind += 1
            if sort_by_episode_id:
                new_dataset.episodes.sort(key=lambda ep: ep.episode_id)
            new_episodes.extend(new_dataset.episodes)
        if remove_unused_episodes:
            self.episodes = new_episodes
        return new_datasets


class EpisodeIterator(Iterator):
    r"""Episode Iterator class that gives options for how a list of episodes
    should be iterated.

    Some of those options are desirable for the internal simulator to get
    higher performance. More context: simulator suffers overhead when switching
    between scenes, therefore episodes of the same scene should be loaded
    consecutively. However, if too many consecutive episodes from same scene
    are feed into RL model, the model will risk to overfit that scene.
    Therefore it's better to load same scene consecutively and switch once a
    number threshold is reached.

    Currently supports the following features:

    Cycling:
        when all episodes are iterated, cycle back to start instead of throwing
        StopIteration.
    Cycling with shuffle:
        when cycling back, shuffle episodes groups grouped by scene.
    Group by scene:
        episodes of same scene will be grouped and loaded consecutively.
    Set max scene repeat:
        set a number threshold on how many episodes from the same scene can be
        loaded consecutively.
    Sample episodes:
        sample the specified number of episodes.
    """

    def __init__(
        self,
        episodes: Sequence[T],
        cycle: bool = True,
        shuffle: bool = False,
        group_by_scene: bool = True,
        max_scene_repeat_episodes: int = -1,
        max_scene_repeat_steps: int = -1,
        num_episode_sample: int = -1,
        step_repetition_range: float = 0.2,
        seed: int = None,
    ) -> None:
        r"""..

        :param episodes: list of episodes.
        :param cycle: if :py:`True`, cycle back to first episodes when
            StopIteration.
        :param shuffle: if :py:`True`, shuffle scene groups when cycle. No
            effect if cycle is set to :py:`False`. Will shuffle grouped scenes
            if :p:`group_by_scene` is :py:`True`.
        :param group_by_scene: if :py:`True`, group episodes from same scene.
        :param max_scene_repeat_episodes: threshold of how many episodes from the same
            scene can be loaded consecutively. :py:`-1` for no limit
        :param max_scene_repeat_steps: threshold of how many steps from the same
            scene can be taken consecutively. :py:`-1` for no limit
        :param num_episode_sample: number of episodes to be sampled. :py:`-1`
            for no sampling.
        :param step_repetition_range: The maximum number of steps within each scene is
            uniformly drawn from
            [1 - step_repeat_range, 1 + step_repeat_range] * max_scene_repeat_steps
            on each scene switch.  This stops all workers from swapping scenes at
            the same time
        """
        if seed:
            random.seed(seed)
            np.random.seed(seed)

        # sample episodes
        if num_episode_sample >= 0:
            episodes = np.random.choice(
                episodes, num_episode_sample, replace=False
            )

        if not isinstance(episodes, list):
            episodes = list(episodes)

        self.episodes = episodes
        self.cycle = cycle
        self.group_by_scene = group_by_scene
        self.shuffle = shuffle

        if shuffle:
            random.shuffle(self.episodes)

        if group_by_scene:
            self.episodes = self._group_scenes(self.episodes)

        self.max_scene_repetition_episodes = max_scene_repeat_episodes
        self.max_scene_repetition_steps = max_scene_repeat_steps

        self._rep_count = -1  # 0 corresponds to first episode already returned
        self._step_count = 0
        self._prev_scene_id: Optional[str] = None

        self._iterator = iter(self.episodes)

        self.step_repetition_range = step_repetition_range
        self._set_shuffle_intervals()

    def __iter__(self) -> "EpisodeIterator":
        return self

    def __next__(self) -> Episode:
        r"""The main logic for handling how episodes will be iterated.

        :return: next episode.
        """
        self._forced_scene_switch_if()

        next_episode = next(self._iterator, None)
        if next_episode is None:
            if not self.cycle:
                raise StopIteration

            self._iterator = iter(self.episodes)

            if self.shuffle:
                self._shuffle()

            next_episode = next(self._iterator)

        if (
            self._prev_scene_id != next_episode.scene_id
            and self._prev_scene_id is not None
        ):
            self._rep_count = 0
            self._step_count = 0

        self._prev_scene_id = next_episode.scene_id
        return next_episode

    def _forced_scene_switch(self) -> None:
        r"""Internal method to switch the scene. Moves remaining episodes
        from current scene to the end and switch to next scene episodes.
        """
        grouped_episodes = [
            list(g)
            for k, g in groupby(self._iterator, key=lambda x: x.scene_id)
        ]

        if len(grouped_episodes) > 1:
            # Ensure we swap by moving the current group to the end
            grouped_episodes = grouped_episodes[1:] + grouped_episodes[0:1]

        self._iterator = iter(sum(grouped_episodes, []))

    def _shuffle(self) -> None:
        r"""Internal method that shuffles the remaining episodes.
        If self.group_by_scene is true, then shuffle groups of scenes.
        """
        assert self.shuffle
        episodes = list(self._iterator)

        random.shuffle(episodes)

        if self.group_by_scene:
            episodes = self._group_scenes(episodes)

        self._iterator = iter(episodes)

    def _group_scenes(
        self, episodes: Union[Sequence[Episode], List[Episode], ndarray]
    ) -> List[T]:
        r"""Internal method that groups episodes by scene
        Groups will be ordered by the order the first episode of a given
        scene is in the list of episodes

        So if the episodes list shuffled before calling this method,
        the scenes will be in a random order
        """
        assert self.group_by_scene

        scene_sort_keys: Dict[str, int] = {}
        for e in episodes:
            if e.scene_id not in scene_sort_keys:
                scene_sort_keys[e.scene_id] = len(scene_sort_keys)

        return sorted(episodes, key=lambda e: scene_sort_keys[e.scene_id])  # type: ignore[arg-type]

    def step_taken(self) -> None:
        self._step_count += 1

    @staticmethod
    def _randomize_value(value: int, value_range: float) -> int:
        return random.randint(
            int(value * (1 - value_range)), int(value * (1 + value_range))
        )

    def _set_shuffle_intervals(self) -> None:
        if self.max_scene_repetition_episodes > 0:
            self._max_rep_episode = self.max_scene_repetition_episodes
        else:
            self._max_rep_episode = None

        if self.max_scene_repetition_steps > 0:
            self._max_rep_step = self._randomize_value(
                self.max_scene_repetition_steps, self.step_repetition_range
            )
        else:
            self._max_rep_step = None

    def _forced_scene_switch_if(self) -> None:
        do_switch = False
        self._rep_count += 1

        # Shuffle if a scene has been selected more than _max_rep_episode times in a row
        if (
            self._max_rep_episode is not None
            and self._rep_count >= self._max_rep_episode
        ):
            do_switch = True

        # Shuffle if a scene has been used for more than _max_rep_step steps in a row
        if (
            self._max_rep_step is not None
            and self._step_count >= self._max_rep_step
        ):
            do_switch = True

        if do_switch:
            self._forced_scene_switch()
            self._set_shuffle_intervals()

habitat/core/dataset.py
cshizhe/onav_rim
#!/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.
r"""Implements tasks and measurements needed for training and benchmarking of
``habitat.Agent`` inside ``habitat.Env``.
"""

from collections import OrderedDict
from typing import Any, Dict, Iterable, List, Optional, Union

import numpy as np

from habitat.core.logging import logger
from habitat.config import Config
from habitat.core.dataset import Dataset, Episode
from habitat.core.simulator import Observations, SensorSuite, Simulator
from habitat.core.spaces import ActionSpace, EmptySpace, Space


class Action:
    r"""
    An action that can be performed by an agent solving a task in environment.
    For example for navigation task action classes will be:
    ``MoveForwardAction, TurnLeftAction, TurnRightAction``. The action can
    use ``Task`` members to pass a state to another action, as well as keep
    own state and reset when new episode starts.
    """

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        return

    def reset(self, *args: Any, **kwargs: Any) -> None:
        r"""Reset method is called from ``Env`` on each reset for each new
        episode. Goal of the method is to reset ``Action``'s state for each
        episode.
        """
        raise NotImplementedError

    def step(self, *args: Any, **kwargs: Any) -> Observations:
        r"""Step method is called from ``Env`` on each ``step``. Can call
        simulator or task method, change task's state.

        :param kwargs: optional parameters for the action, like distance/force.
        :return: observations after taking action in the task, including ones
            coming from a simulator.
        """
        raise NotImplementedError

    @property
    def action_space(self) -> Space:
        r"""a current Action's action space."""
        raise NotImplementedError


class SimulatorTaskAction(Action):
    r"""
    An ``EmbodiedTask`` action that is wrapping simulator action.
    """

    def __init__(
        self, *args: Any, config: Config, sim: Simulator, **kwargs: Any
    ) -> None:
        self._config = config
        self._sim = sim

    @property
    def action_space(self):
        return EmptySpace()

    def reset(self, *args: Any, **kwargs: Any) -> None:
        return None

    def step(self, *args: Any, **kwargs: Any) -> Observations:
        r"""Step method is called from ``Env`` on each ``step``."""
        raise NotImplementedError


class Measure:
    r"""Represents a measure that provides measurement on top of environment
    and task.

    :data uuid: universally unique id.
    :data _metric: metric for the :ref:`Measure`, this has to be updated with
        each :ref:`step() <env.Env.step()>` call on :ref:`env.Env`.

    This can be used for tracking statistics when running experiments. The
    user of this class needs to implement the :ref:`reset_metric()` and
    :ref:`update_metric()` method and the user is also required to set the
    :ref:`uuid <Measure.uuid>` and :ref:`_metric` attributes.

    .. (uuid is a builtin Python module, so just :ref:`uuid` would link there)
    """

    _metric: Any
    uuid: str

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        self.uuid = self._get_uuid(*args, **kwargs)
        self._metric = None

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        raise NotImplementedError

    def reset_metric(self, *args: Any, **kwargs: Any) -> None:
        r"""Reset :ref:`_metric`, this method is called from :ref:`env.Env` on
        each reset.
        """
        raise NotImplementedError

    def update_metric(self, *args: Any, **kwargs: Any) -> None:
        r"""Update :ref:`_metric`, this method is called from :ref:`env.Env`
        on each :ref:`step() <env.Env.step()>`
        """
        raise NotImplementedError

    def get_metric(self):
        r"""..

        :return: the current metric for :ref:`Measure`.
        """
        return self._metric


class Metrics(dict):
    r"""Dictionary containing measurements."""

    def __init__(self, measures: Dict[str, Measure]) -> None:
        """Constructor

        :param measures: list of :ref:`Measure` whose metrics are fetched and
            packaged.
        """
        data = [
            (uuid, measure.get_metric()) for uuid, measure in measures.items()
        ]
        super().__init__(data)


class Measurements:
    r"""Represents a set of Measures, with each :ref:`Measure` being
    identified through a unique id.
    """

    measures: Dict[str, Measure]

    def __init__(self, measures: Iterable[Measure]) -> None:
        """Constructor

        :param measures: list containing :ref:`Measure`, uuid of each
            :ref:`Measure` must be unique.
        """
        self.measures = OrderedDict()
        for measure in measures:
            assert (
                measure.uuid not in self.measures
            ), "'{}' is duplicated measure uuid".format(measure.uuid)
            self.measures[measure.uuid] = measure

    def reset_measures(self, *args: Any, **kwargs: Any) -> None:
        for measure in self.measures.values():
            measure.reset_metric(*args, **kwargs)

    def update_measures(self, *args: Any, **kwargs: Any) -> None:
        for measure in self.measures.values():
            measure.update_metric(*args, **kwargs)

    def get_metrics(self) -> Metrics:
        r"""Collects measurement from all :ref:`Measure`\ s and returns it
        packaged inside :ref:`Metrics`.
        """
        return Metrics(self.measures)

    def _get_measure_index(self, measure_name):
        return list(self.measures.keys()).index(measure_name)

    def check_measure_dependencies(
        self, measure_name: str, dependencies: List[str]
    ):
        r"""Checks if dependencies measures are enabled and calculatethat the measure
        :param measure_name: a name of the measure for which has dependencies.
        :param dependencies: a list of a measure names that are required by
        the measure.
        :return:
        """
        measure_index = self._get_measure_index(measure_name)
        for dependency_measure in dependencies:
            assert (
                dependency_measure in self.measures
            ), f"""{measure_name} measure requires {dependency_measure}
                listed in tje measures list in the config."""

        for dependency_measure in dependencies:
            assert measure_index > self._get_measure_index(
                dependency_measure
            ), f"""{measure_name} measure requires be listed after {dependency_measure}
                in tje measures list in the config."""


class EmbodiedTask:
    r"""Base class for embodied tasks. ``EmbodiedTask`` holds definition of
    a task that agent needs to solve: action space, observation space,
    measures, simulator usage. ``EmbodiedTask`` has :ref:`reset` and
    :ref:`step` methods that are called by ``Env``. ``EmbodiedTask`` is the
    one of main dimensions for the framework extension. Once new embodied task
    is introduced implementation of ``EmbodiedTask`` is a formal definition of
    the task that opens opportunity for others to propose solutions and
    include it into benchmark results.

    Args:
        config: config for the task.
        sim: reference to the simulator for calculating task observations.
        dataset: reference to dataset for task instance level information.

    :data measurements: set of task measures.
    :data sensor_suite: suite of task sensors.
    """

    _config: Any
    _sim: Optional[Simulator]
    _dataset: Optional[Dataset]
    _is_episode_active: bool
    _is_resetting: bool
    measurements: Measurements
    sensor_suite: SensorSuite

    def __init__(
        self, config: Config, sim: Simulator, dataset: Optional[Dataset] = None
    ) -> None:
        from habitat.core.registry import registry

        self._config = config
        self._sim = sim
        self._dataset = dataset

        self.measurements = Measurements(
            self._init_entities(
                entity_names=config.MEASUREMENTS,
                register_func=registry.get_measure,
                entities_config=config,
            ).values()
        )

        self.sensor_suite = SensorSuite(
            self._init_entities(
                entity_names=config.SENSORS,
                register_func=registry.get_sensor,
                entities_config=config,
            ).values()
        )

        self.actions = self._init_entities(
            entity_names=config.POSSIBLE_ACTIONS,
            register_func=registry.get_task_action,
            entities_config=self._config.ACTIONS,
        )
        self._action_keys = list(self.actions.keys())

    def _init_entities(
        self, entity_names, register_func, entities_config=None
    ) -> OrderedDict:
        if entities_config is None:
            entities_config = self._config

        entities = OrderedDict()
        for entity_name in entity_names:
            entity_cfg = getattr(entities_config, entity_name)
            entity_type = register_func(entity_cfg.TYPE)
            assert (
                entity_type is not None
            ), f"invalid {entity_name} type {entity_cfg.TYPE}"
            entities[entity_name] = entity_type(
                sim=self._sim,
                config=entity_cfg,
                dataset=self._dataset,
                task=self,
            )
        return entities

    def reset(self, episode: Episode):
        self._is_resetting = True
        observations = self._sim.reset()
        observations.update(
            self.sensor_suite.get_observations(
                observations=observations, episode=episode, task=self
            )
        )

        for action_instance in self.actions.values():
            action_instance.reset(episode=episode, task=self)
        self._is_resetting = False

        return observations

    def step(self, action: Dict[str, Any], episode: Episode):
        if "action_args" not in action or action["action_args"] is None:
            action["action_args"] = {}
        action_name = action["action"]
        if isinstance(action_name, (int, np.integer)):
            action_name = self.get_action_name(action_name)
        assert (
            action_name in self.actions
        ), f"Can't find '{action_name}' action in {self.actions.keys()}."

        task_action = self.actions[action_name]

        observations = task_action.step(**action["action_args"], task=self)
        observations.update(
            self.sensor_suite.get_observations(
                observations=observations,
                episode=episode,
                action=action,
                task=self,
            )
        )

        self._is_episode_active = self._check_episode_is_active(
            observations=observations, action=action, episode=episode
        )

        return observations

    def get_action_name(self, action_index: int):
        if action_index >= len(self.actions):
            raise ValueError(f"Action index '{action_index}' is out of range.")
        return self._action_keys[action_index]

    @property
    def action_space(self) -> Space:
        return ActionSpace(
            {
                action_name: action_instance.action_space
                for action_name, action_instance in self.actions.items()
            }
        )

    def overwrite_sim_config(
        self, sim_config: Config, episode: Episode
    ) -> Config:
        r"""Update config merging information from :p:`sim_config` and
        :p:`episode`.

        :param sim_config: config for simulator.
        :param episode: current episode.
        """
        raise NotImplementedError

    def _check_episode_is_active(
        self,
        *args: Any,
        action: Union[int, Dict[str, Any]],
        episode: Episode,
        **kwargs: Any,
    ) -> bool:
        raise NotImplementedError

    @property
    def is_episode_active(self):
        return self._is_episode_active
    
    @property
    def is_resetting(self):
        return self._is_resetting

    def seed(self, seed: int) -> None:
        return

habitat/core/embodied_task.py
cshizhe/onav_rim
#!/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.

from multiprocessing import current_process
import random
import time
from typing import Any, Dict, Iterator, List, Optional, Tuple, Union, cast

import gym
import numba
import numpy as np
from gym import spaces

from habitat.config import Config
from habitat.core.dataset import Dataset, Episode, EpisodeIterator
from habitat.core.embodied_task import EmbodiedTask, Metrics
from habitat.core.simulator import Observations, Simulator
from habitat.datasets import make_dataset
from habitat.sims import make_sim
from habitat.tasks import make_task
from habitat.utils import profiling_wrapper
from habitat.tasks.nav.object_nav_task import mapping_mpcat40_to_goal21


class Env:
    r"""Fundamental environment class for :ref:`habitat`.

    :data observation_space: ``SpaceDict`` object corresponding to sensor in
        sim and task.
    :data action_space: ``gym.space`` object corresponding to valid actions.

    All the information  needed for working on embodied tasks with simulator
    is abstracted inside :ref:`Env`. Acts as a base for other derived
    environment classes. :ref:`Env` consists of three major components:
    ``dataset`` (`episodes`), ``simulator`` (:ref:`sim`) and :ref:`task` and
    connects all the three components together.
    """

    observation_space: spaces.Dict
    action_space: spaces.Dict
    _config: Config
    _dataset: Optional[Dataset]
    number_of_episodes: Optional[int]
    _episodes: List[Episode]
    _current_episode_index: Optional[int]
    _current_episode: Optional[Episode]
    _episode_iterator: Optional[Iterator]
    _sim: Simulator
    _task: EmbodiedTask
    _max_episode_seconds: int
    _max_episode_steps: int
    _elapsed_steps: int
    _episode_start_time: Optional[float]
    _episode_over: bool

    def __init__(
        self, config: Config, dataset: Optional[Dataset] = None
    ) -> None:
        """Constructor

        :param config: config for the environment. Should contain id for
            simulator and ``task_name`` which are passed into ``make_sim`` and
            ``make_task``.
        :param dataset: reference to dataset for task instance level
            information. Can be defined as :py:`None` in which case
            ``_episodes`` should be populated from outside.
        """

        assert config.is_frozen(), (
            "Freeze the config before creating the "
            "environment, use config.freeze()."
        )
        self._config = config
        self._dataset = dataset
        self._current_episode_index = None
        if self._dataset is None and config.DATASET.TYPE:
            self._dataset = make_dataset(
                id_dataset=config.DATASET.TYPE, config=config.DATASET
            )
        self._episodes = (
            self._dataset.episodes
            if self._dataset
            else cast(List[Episode], [])
        )
        self._current_episode = None
        iter_option_dict = {
            k.lower(): v
            for k, v in config.ENVIRONMENT.ITERATOR_OPTIONS.items()
        }
        iter_option_dict["seed"] = config.SEED
        self._episode_iterator = self._dataset.get_episode_iterator(
            **iter_option_dict
        )

        # load the first scene if dataset is present
        if self._dataset:
            assert (
                len(self._dataset.episodes) > 0
            ), "dataset should have non-empty episodes list"
            self._config.defrost()
            self._config.SIMULATOR.SCENE = self._dataset.episodes[0].scene_id
            self._config.freeze()

            self.number_of_episodes = len(self._dataset.episodes)
        else:
            self.number_of_episodes = None

        self._sim = make_sim(
            id_sim=self._config.SIMULATOR.TYPE, config=self._config.SIMULATOR
        )
        self._task = make_task(
            self._config.TASK.TYPE,
            config=self._config.TASK,
            sim=self._sim,
            dataset=self._dataset,
        )
        self.observation_space = spaces.Dict(
            {
                **self._sim.sensor_suite.observation_spaces.spaces,
                **self._task.sensor_suite.observation_spaces.spaces,
            }
        )
        self.action_space = self._task.action_space
        self._max_episode_seconds = (
            self._config.ENVIRONMENT.MAX_EPISODE_SECONDS
        )
        self._max_episode_steps = self._config.ENVIRONMENT.MAX_EPISODE_STEPS
        self._elapsed_steps = 0
        self._episode_start_time: Optional[float] = None
        self._episode_over = False

    @property
    def current_episode(self) -> Episode:
        assert self._current_episode is not None
        return self._current_episode

    @current_episode.setter
    def current_episode(self, episode: Episode) -> None:
        self._current_episode = episode

    @property
    def episode_iterator(self) -> Iterator:
        return self._episode_iterator

    @episode_iterator.setter
    def episode_iterator(self, new_iter: Iterator) -> None:
        self._episode_iterator = new_iter

    @property
    def episodes(self) -> List[Episode]:
        return self._episodes

    @episodes.setter
    def episodes(self, episodes: List[Episode]) -> None:
        assert (
            len(episodes) > 0
        ), "Environment doesn't accept empty episodes list."
        self._episodes = episodes

    @property
    def sim(self) -> Simulator:
        return self._sim

    @property
    def episode_start_time(self) -> Optional[float]:
        return self._episode_start_time

    @property
    def episode_over(self) -> bool:
        return self._episode_over

    @property
    def task(self) -> EmbodiedTask:
        return self._task

    @property
    def _elapsed_seconds(self) -> float:
        assert (
            self._episode_start_time
        ), "Elapsed seconds requested before episode was started."
        return time.time() - self._episode_start_time

    def get_metrics(self) -> Metrics:
        return self._task.measurements.get_metrics()

    def _past_limit(self) -> bool:
        if (
            self._max_episode_steps != 0
            and self._max_episode_steps <= self._elapsed_steps
        ):
            return True
        elif (
            self._max_episode_seconds != 0
            and self._max_episode_seconds <= self._elapsed_seconds
        ):
            return True
        return False

    def _reset_stats(self) -> None:
        self._episode_start_time = time.time()
        self._elapsed_steps = 0
        self._episode_over = False

    def reset(self, current_episode=None) -> Observations:
        r"""Resets the environments and returns the initial observations.

        :return: initial observations from the environment.
        """
        self._reset_stats()

        assert len(self.episodes) > 0, "Episodes list is empty"
        if current_episode is None:
            # Delete the shortest path cache of the current episode
            # Caching it for the next time we see this episode isn't really worth
            # it
            if self._current_episode is not None:
                self._current_episode._shortest_path_cache = None

            self._current_episode = next(self._episode_iterator)
        else:
            self._current_episode = current_episode

        self.reconfigure(self._config)
 
        observations = self.task.reset(episode=self.current_episode)
        # obtain mapping from instance id to semantic label id
        scene = self.sim.semantic_annotations()

        # TODO: for goal_vis_pixels sensor
        instance_id_to_label_id = {int(obj.id.split("_")[-1]): obj.category.index() for obj in scene.objects}
        # instance_id_to_label_id = {}
        self.mapping = np.array([ instance_id_to_label_id[i] for i in range(len(instance_id_to_label_id)) ])
       
        is_rearrangement = "Rearrangement" in self._config.DATASET.TYPE
        # ! MP3D object id to category ID mapping
        if not is_rearrangement and 'semantic' in observations: # If no map, then we're using semantics on scene without
            # Should raise a warning, but just driving ahead for now
            if self.mapping.size > 0:
                if len(self.sim.obj_sem_id_to_sem_category_mapping.keys()) > 0:
                    instance_id_to_label_id.update(self.sim.obj_sem_id_to_sem_category_mapping)
                    # print('env instance_id_to_label_id', instance_id_to_label_id)
                    self.mapping = np.zeros(
                        max(max(instance_id_to_label_id.keys()) + 1, 51), dtype=np.int8
                    )  # debug hack
                    for key, value in instance_id_to_label_id.items():
                        self.mapping[key] = value # mp3d category id
                        # # remap everything void/unlabeled/misc/etc .. to misc class 40
                        # # (void:0,  unlabeled: 41, misc=40)
                        # self.mapping[key] = 40 if value <= 0 or value == 41 else value
                        # # Map MP3D to THDA SemSeg Categories
                        # if value not in mapping_mpcat40_to_goal21.keys():
                        #     # Map every unknown category to 28 
                        #     self.mapping[key] = 28
                        # else:
                        #     self.mapping[key] = mapping_mpcat40_to_goal21[self.mapping[key]]
                        #     print(key, self.mapping[key])
                    observations["semantic"] = self.mapping[observations["semantic"]]
                else:
                    observations['semantic'] = np.take(self.mapping, observations['semantic'])
            else:
                observations['semantic'] = observations['semantic'].astype(int)

        self._task.measurements.reset_measures(
            episode=self.current_episode,
            task=self.task, observations=observations
        )

        return observations

    def _update_step_stats(self) -> None:
        self._elapsed_steps += 1
        self._episode_over = not self._task.is_episode_active
        if self._past_limit():
            self._episode_over = True

        if self.episode_iterator is not None and isinstance(
            self.episode_iterator, EpisodeIterator
        ):
            self.episode_iterator.step_taken()

    def step(
        self, action: Union[int, str, Dict[str, Any]], **kwargs
    ) -> Observations:
        r"""Perform an action in the environment and return observations.

        :param action: action (belonging to :ref:`action_space`) to be
            performed inside the environment. Action is a name or index of
            allowed task's action and action arguments (belonging to action's
            :ref:`action_space`) to support parametrized and continuous
            actions.
        :return: observations after taking action in environment.
        """

        assert (
            self._episode_start_time is not None
        ), "Cannot call step before calling reset"
        assert (
            self._episode_over is False
        ), "Episode over, call reset before calling step"

        # Support simpler interface as well
        if isinstance(action, (str, int, np.integer)):
            action = {"action": action}
            if "replay_data" in kwargs.keys():
                action["action_args"] = {}
                action["action_args"]["replay_data"] = kwargs["replay_data"]

        observations = self.task.step(
            action=action, episode=self.current_episode
        )

        is_rearrangement = "Rearrangement" in self._config.DATASET.TYPE
        # Instance to sem seg
        if not is_rearrangement and 'semantic' in observations:
            if self.mapping.size > 0:
                observations['semantic'] = np.take(self.mapping, observations['semantic'])
            else:
                observations['semantic'] = observations['semantic'].astype(int)

        self._task.measurements.update_measures(
            episode=self.current_episode, action=action,
            task=self.task, observations=observations
        )

        self._update_step_stats()

        return observations

    @staticmethod
    @numba.njit
    def _seed_numba(seed: int):
        random.seed(seed)
        np.random.seed(seed)

    def seed(self, seed: int) -> None:
        random.seed(seed)
        np.random.seed(seed)
        self._seed_numba(seed)
        self._sim.seed(seed)
        self._task.seed(seed)

    def reconfigure(self, config: Config) -> None:
        self._config = config

        self._config.defrost()
        self._config.SIMULATOR = self._task.overwrite_sim_config(
            self._config.SIMULATOR, self.current_episode
        )
        self._config.freeze()

        self._sim.reconfigure(self._config.SIMULATOR)

    def render(self, mode="rgb") -> np.ndarray:
        return self._sim.render(mode)

    def close(self) -> None:
        self._sim.close()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()


class RLEnv(gym.Env):
    r"""Reinforcement Learning (RL) environment class which subclasses ``gym.Env``.

    This is a wrapper over :ref:`Env` for RL users. To create custom RL
    environments users should subclass `RLEnv` and define the following
    methods: :ref:`get_reward_range()`, :ref:`get_reward()`,
    :ref:`get_done()`, :ref:`get_info()`.

    As this is a subclass of ``gym.Env``, it implements `reset()` and
    `step()`.
    """

    _env: Env

    def __init__(
        self, config: Config, dataset: Optional[Dataset] = None
    ) -> None:
        """Constructor

        :param config: config to construct :ref:`Env`
        :param dataset: dataset to construct :ref:`Env`.
        """

        self._env = Env(config, dataset)
        self.observation_space = self._env.observation_space
        self.action_space = self._env.action_space
        self.number_of_episodes = self._env.number_of_episodes
        self.reward_range = self.get_reward_range()

    @property
    def habitat_env(self) -> Env:
        return self._env

    @property
    def episodes(self) -> List[Episode]:
        return self._env.episodes

    @episodes.setter
    def episodes(self, episodes: List[Episode]) -> None:
        self._env.episodes = episodes

    @property
    def current_episode(self) -> Episode:
        return self._env.current_episode

    @profiling_wrapper.RangeContext("RLEnv.reset")
    def reset(self, **kwargs) -> Observations:
        return self._env.reset(**kwargs)

    def get_reward_range(self):
        r"""Get min, max range of reward.

        :return: :py:`[min, max]` range of reward.
        """
        raise NotImplementedError

    def get_reward(self, observations: Observations) -> Any:
        r"""Returns reward after action has been performed.

        :param observations: observations from simulator and task.
        :return: reward after performing the last action.

        This method is called inside the :ref:`step()` method.
        """
        raise NotImplementedError

    def get_done(self, observations: Observations) -> bool:
        r"""Returns boolean indicating whether episode is done after performing
        the last action.

        :param observations: observations from simulator and task.
        :return: done boolean after performing the last action.

        This method is called inside the step method.
        """
        raise NotImplementedError

    def get_info(self, observations) -> Dict[Any, Any]:
        r"""..

        :param observations: observations from simulator and task.
        :return: info after performing the last action.
        """
        raise NotImplementedError

    @profiling_wrapper.RangeContext("RLEnv.step")
    def step(self, *args, **kwargs) -> Tuple[Observations, Any, bool, dict]:
        r"""Perform an action in the environment.

        :return: :py:`(observations, reward, done, info)`
        """

        observations = self._env.step(*args, **kwargs)
        reward = self.get_reward(observations)
        done = self.get_done(observations)
        info = self.get_info(observations)

        return observations, reward, done, info

    def seed(self, seed: Optional[int] = None) -> None:
        self._env.seed(seed)

    def render(self, mode: str = "rgb") -> np.ndarray:
        return self._env.render(mode)

    def close(self) -> None:
        self._env.close()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

habitat/core/env.py
cshizhe/onav_rim
#!/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 signal
import warnings
from multiprocessing.connection import Connection
from multiprocessing.context import BaseContext
from queue import Queue
from threading import Thread
from typing import (
    Any,
    Callable,
    Dict,
    Iterator,
    List,
    Optional,
    Sequence,
    Set,
    Tuple,
    Union,
    cast,
)

import gym
import numpy as np
from gym import spaces

import habitat
from habitat.config import Config
from habitat.core.env import Env, RLEnv
from habitat.core.logging import logger
from habitat.core.utils import tile_images
from habitat.utils import profiling_wrapper

try:
    # Use torch.multiprocessing if we can.
    # We have yet to find a reason to not use it and
    # you are required to use it when sending a torch.Tensor
    # between processes
    import torch
    from torch import multiprocessing as mp  # type:ignore
except ImportError:
    torch = None
    import multiprocessing as mp  # type:ignore

STEP_COMMAND = "step"
RESET_COMMAND = "reset"
RENDER_COMMAND = "render"
CLOSE_COMMAND = "close"
OBSERVATION_SPACE_COMMAND = "observation_space"
ACTION_SPACE_COMMAND = "action_space"
NUMBER_OF_EPISODES_COMMAND = "number_of_episodes"
CALL_COMMAND = "call"
EPISODE_COMMAND = "current_episode"
COUNT_EPISODES_COMMAND = "count_episodes"
EPISODE_OVER = "episode_over"
GET_METRICS = "get_metrics"
GET_AGENT_POSE = "get_agent_pose"
GET_OBJECT_STATE = "get_current_object_states"
SCENES_SPLITS = "scene_splits"
DATASETS = "datasets"
SCENE_COMMAND = "current_scene"


def _make_env_fn(
    config: Config, dataset: Optional[habitat.Dataset] = None, rank: int = 0
) -> Env:
    """Constructor for default habitat :ref:`env.Env`.

    :param config: configuration for environment.
    :param dataset: dataset for environment.
    :param rank: rank for setting seed of environment
    :return: :ref:`env.Env` / :ref:`env.RLEnv` object
    """
    habitat_env = Env(config=config, dataset=dataset)
    habitat_env.seed(config.SEED + rank)
    return habitat_env


class VectorEnv:
    r"""Vectorized environment which creates multiple processes where each
    process runs its own environment. Main class for parallelization of
    training and evaluation.


    All the environments are synchronized on step and reset methods.
    """

    observation_spaces: List[spaces.Dict]
    number_of_episodes: List[Optional[int]]
    action_spaces: List[spaces.Dict]
    _workers: List[Union[mp.Process, Thread]]
    _is_waiting: bool
    _num_envs: int
    _auto_reset_done: bool
    _mp_ctx: BaseContext
    _connection_read_fns: List[Callable[[], Any]]
    _connection_write_fns: List[Callable[[Any], None]]

    def __init__(
        self,
        make_env_fn: Callable[..., Union[Env, RLEnv]] = _make_env_fn,
        env_fn_args: Sequence[Tuple] = None,
        auto_reset_done: bool = True,
        multiprocessing_start_method: str = "forkserver",
        workers_ignore_signals: bool = False,
    ) -> None:
        """..

        :param make_env_fn: function which creates a single environment. An
            environment can be of type :ref:`env.Env` or :ref:`env.RLEnv`
        :param env_fn_args: tuple of tuple of args to pass to the
            :ref:`_make_env_fn`.
        :param auto_reset_done: automatically reset the environment when
            done. This functionality is provided for seamless training
            of vectorized environments.
        :param multiprocessing_start_method: the multiprocessing method used to
            spawn worker processes. Valid methods are
            :py:`{'spawn', 'forkserver', 'fork'}`; :py:`'forkserver'` is the
            recommended method as it works well with CUDA. If :py:`'fork'` is
            used, the subproccess  must be started before any other GPU useage.
        :param workers_ignore_signals: Whether or not workers will ignore SIGINT and SIGTERM
            and instead will only exit when :ref:`close` is called
        """
        self._is_waiting = False
        self._is_closed = True

        assert (
            env_fn_args is not None and len(env_fn_args) > 0
        ), "number of environments to be created should be greater than 0"

        self._num_envs = len(env_fn_args)

        assert multiprocessing_start_method in self._valid_start_methods, (
            "multiprocessing_start_method must be one of {}. Got '{}'"
        ).format(self._valid_start_methods, multiprocessing_start_method)
        self._auto_reset_done = auto_reset_done
        self._mp_ctx = mp.get_context(multiprocessing_start_method)
        self._workers = []
        (
            self._connection_read_fns,
            self._connection_write_fns,
        ) = self._spawn_workers(  # noqa
            env_fn_args,
            make_env_fn,
            workers_ignore_signals=workers_ignore_signals,
        )

        self._is_closed = False

        for write_fn in self._connection_write_fns:
            write_fn((OBSERVATION_SPACE_COMMAND, None))
        self.observation_spaces = [
            read_fn() for read_fn in self._connection_read_fns
        ]
        for write_fn in self._connection_write_fns:
            write_fn((ACTION_SPACE_COMMAND, None))
        self.action_spaces = [
            read_fn() for read_fn in self._connection_read_fns
        ]
        for write_fn in self._connection_write_fns:
            write_fn((NUMBER_OF_EPISODES_COMMAND, None))
        self.number_of_episodes = [
            read_fn() for read_fn in self._connection_read_fns
        ]
        self._paused: List[Tuple] = []

    @property
    def num_envs(self):
        r"""number of individual environments."""
        return self._num_envs - len(self._paused)

    @staticmethod
    @profiling_wrapper.RangeContext("_worker_env")
    def _worker_env(
        connection_read_fn: Callable,
        connection_write_fn: Callable,
        env_fn: Callable,
        env_fn_args: Tuple[Any],
        auto_reset_done: bool,
        mask_signals: bool = False,
        child_pipe: Optional[Connection] = None,
        parent_pipe: Optional[Connection] = None,
    ) -> None:
        r"""process worker for creating and interacting with the environment."""
        if mask_signals:
            signal.signal(signal.SIGINT, signal.SIG_IGN)
            signal.signal(signal.SIGTERM, signal.SIG_IGN)

            signal.signal(signal.SIGUSR1, signal.SIG_IGN)
            signal.signal(signal.SIGUSR2, signal.SIG_IGN)

        env = env_fn(*env_fn_args)
        if parent_pipe is not None:
            parent_pipe.close()
        try:
            command, data = connection_read_fn()
            while command != CLOSE_COMMAND:
                if command == STEP_COMMAND:
                    # different step methods for habitat.RLEnv and habitat.Env
                    if isinstance(env, (habitat.RLEnv, gym.Env)):
                        # habitat.RLEnv
                        observations, reward, done, info = env.step(**data)
                        if auto_reset_done and done:
                            observations = env.reset()
                        with profiling_wrapper.RangeContext(
                            "worker write after step"
                        ):
                            connection_write_fn(
                                (observations, reward, done, info)
                            )
                    elif isinstance(env, habitat.Env):  # type: ignore
                        # habitat.Env
                        observations = env.step(**data)
                        if auto_reset_done and env.episode_over:
                            observations = env.reset()
                        connection_write_fn(observations)
                    else:
                        raise NotImplementedError

                elif command == RESET_COMMAND:
                    observations = env.reset()
                    connection_write_fn(observations)

                elif command == RENDER_COMMAND:
                    connection_write_fn(env.render(*data[0], **data[1]))

                elif command in {
                    OBSERVATION_SPACE_COMMAND,
                    ACTION_SPACE_COMMAND,
                    NUMBER_OF_EPISODES_COMMAND,
                }:
                    connection_write_fn(getattr(env, command))
                elif command == CALL_COMMAND:
                    function_name, function_args = data
                    if function_args is None or len(function_args) == 0:
                        result = getattr(env, function_name)()
                    else:
                        result = getattr(env, function_name)(**function_args)
                    connection_write_fn(result)

                # TODO: update CALL_COMMAND for getting attribute like this
                elif command == EPISODE_COMMAND:
                    connection_write_fn(env.current_episode)

                elif command == COUNT_EPISODES_COMMAND:
                    connection_write_fn(len(env.episodes))

                elif command == EPISODE_OVER:
                    connection_write_fn(env.episode_over)

                elif command == GET_METRICS:
                    result = env.get_metrics()
                    connection_write_fn(result)

                elif command == GET_AGENT_POSE:
                    result = env._env._sim.get_agent_pose()
                    connection_write_fn(result)

                elif command == GET_OBJECT_STATE:
                    result = env._env._sim.get_current_object_states()
                    connection_write_fn(result)

                elif command == SCENES_SPLITS:
                    result = env._env._config.DATASET.CONTENT_SCENES
                    connection_write_fn(result)

                elif command == DATASETS:
                    result = env._env._dataset
                    connection_write_fn(result)

                elif command == SCENE_COMMAND:
                    result = env.current_episode.scene_id
                    connection_write_fn(result)

                else:
                    raise NotImplementedError

                with profiling_wrapper.RangeContext("worker wait for command"):
                    command, data = connection_read_fn()

            if child_pipe is not None:
                child_pipe.close()
        except KeyboardInterrupt:
            logger.info("Worker KeyboardInterrupt")
        finally:
            env.close()

    def _spawn_workers(
        self,
        env_fn_args: Sequence[Tuple],
        make_env_fn: Callable[..., Union[Env, RLEnv]] = _make_env_fn,
        workers_ignore_signals: bool = False,
    ) -> Tuple[List[Callable[[], Any]], List[Callable[[Any], None]]]:
        parent_connections, worker_connections = zip(
            *[self._mp_ctx.Pipe(duplex=True) for _ in range(self._num_envs)]
        )
        self._workers = []
        for worker_conn, parent_conn, env_args in zip(
            worker_connections, parent_connections, env_fn_args
        ):
            ps = self._mp_ctx.Process(
                target=self._worker_env,
                args=(
                    worker_conn.recv,
                    worker_conn.send,
                    make_env_fn,
                    env_args,
                    self._auto_reset_done,
                    workers_ignore_signals,
                    worker_conn,
                    parent_conn,
                ),
            )
            self._workers.append(cast(mp.Process, ps))
            ps.daemon = True
            ps.start()
            worker_conn.close()
        return (
            [p.recv for p in parent_connections],
            [p.send for p in parent_connections],
        )

    def current_episodes(self):
        self._is_waiting = True
        for write_fn in self._connection_write_fns:
            write_fn((EPISODE_COMMAND, None))
        results = []
        for read_fn in self._connection_read_fns:
            results.append(read_fn())
        self._is_waiting = False
        return results

    def current_scenes(self):
        self._is_waiting = True
        for write_fn in self._connection_write_fns:
            write_fn((SCENE_COMMAND, None))
        results = []
        for read_fn in self._connection_read_fns:
            results.append(read_fn())
        self._is_waiting = False
        return results

    def count_episodes(self):
        self._is_waiting = True
        for write_fn in self._connection_write_fns:
            write_fn((COUNT_EPISODES_COMMAND, None))
        results = []
        for read_fn in self._connection_read_fns:
            results.append(read_fn())
        self._is_waiting = False
        return results

    def get_agent_pose(self):
        self._is_waiting = True
        for write_fn in self._connection_write_fns:
            write_fn((GET_AGENT_POSE, None))
        results = []
        for read_fn in self._connection_read_fns:
            results.append(read_fn())
        self._is_waiting = False
        return results

    def scene_splits(self):
        self._is_waiting = True
        for write_fn in self._connection_write_fns:
            write_fn((SCENES_SPLITS, None))
        results = []
        for read_fn in self._connection_read_fns:
            results.append(read_fn())
        self._is_waiting = False
        return results

    def datasets(self):
        self._is_waiting = True
        for write_fn in self._connection_write_fns:
            write_fn((DATASETS, None))
        results = []
        for read_fn in self._connection_read_fns:
            results.append(read_fn())
        self._is_waiting = False
        return results

    def get_current_object_states(self):
        self._is_waiting = True
        for write_fn in self._connection_write_fns:
            write_fn((GET_OBJECT_STATE, None))
        results = []
        for read_fn in self._connection_read_fns:
            results.append(read_fn())
        self._is_waiting = False
        return results

    def episode_over(self):
        self._is_waiting = True
        for write_fn in self._connection_write_fns:
            write_fn((EPISODE_OVER, None))
        results = []
        for read_fn in self._connection_read_fns:
            results.append(read_fn())
        self._is_waiting = False
        return results

    def get_metrics(self):
        self._is_waiting = True
        for write_fn in self._connection_write_fns:
            write_fn((GET_METRICS, None))
        results = []
        for read_fn in self._connection_read_fns:
            results.append(read_fn())
        self._is_waiting = False
        return results

    def reset(self):
        r"""Reset all the vectorized environments

        :return: list of outputs from the reset method of envs.
        """
        self._is_waiting = True
        for write_fn in self._connection_write_fns:
            write_fn((RESET_COMMAND, None))
        results = []
        for read_fn in self._connection_read_fns:
            results.append(read_fn())
        self._is_waiting = False
        return results

    def reset_at(self, index_env: int):
        r"""Reset in the index_env environment in the vector.

        :param index_env: index of the environment to be reset
        :return: list containing the output of reset method of indexed env.
        """
        self._is_waiting = True
        self._connection_write_fns[index_env]((RESET_COMMAND, None))
        results = [self._connection_read_fns[index_env]()]
        self._is_waiting = False
        return results

    def step_at(self, index_env: int, action: Dict[str, Any]):
        r"""Step in the index_env environment in the vector.

        :param index_env: index of the environment to be stepped into
        :param action: action to be taken
        :return: list containing the output of step method of indexed env.
        """
        self._is_waiting = True
        self._warn_cuda_tensors(action)
        self._connection_write_fns[index_env]((STEP_COMMAND, action))
        results = [self._connection_read_fns[index_env]()]
        self._is_waiting = False
        return results

    def async_step(self, data: List[Union[int, str, Dict[str, Any]]]) -> None:
        r"""Asynchronously step in the environments.

        :param data: list of size _num_envs containing keyword arguments to
            pass to :ref:`step` method for each Environment. For example,
            :py:`[{"action": "TURN_LEFT", "action_args": {...}}, ...]`.
        """
        # Backward compatibility
        if isinstance(data[0], (int, np.integer, str)):
            actions = [{"action": {"action": action}} for action in data]
        else:
            actions = cast(List[Dict[str, Any]], data)

        self._is_waiting = True
        for write_fn, action in zip(self._connection_write_fns, actions):
            self._warn_cuda_tensors(action)
            write_fn((STEP_COMMAND, action))

    @profiling_wrapper.RangeContext("wait_step")
    def wait_step(self) -> List[Any]:
        r"""Wait until all the asynchronized environments have synchronized."""
        observations = []
        for read_fn in self._connection_read_fns:
            observations.append(read_fn())
        self._is_waiting = False
        return observations

    def step(self, data: List[Union[int, str, Dict[str, Any]]]) -> List[Any]:
        r"""Perform actions in the vectorized environments.

        :param data: list of size _num_envs containing keyword arguments to
            pass to :ref:`step` method for each Environment. For example,
            :py:`[{"action": "TURN_LEFT", "action_args": {...}}, ...]`.
        :return: list of outputs from the step method of envs.
        """
        self.async_step(data)
        return self.wait_step()

    def close(self) -> None:
        if self._is_closed:
            return

        if self._is_waiting:
            for read_fn in self._connection_read_fns:
                read_fn()

        for write_fn in self._connection_write_fns:
            write_fn((CLOSE_COMMAND, None))

        for _, _, write_fn, _ in self._paused:
            write_fn((CLOSE_COMMAND, None))

        for process in self._workers:
            process.join()

        for _, _, _, process in self._paused:
            process.join()

        self._is_closed = True

    def pause_at(self, index: int) -> None:
        r"""Pauses computation on this env without destroying the env.

        :param index: which env to pause. All indexes after this one will be
            shifted down by one.

        This is useful for not needing to call steps on all environments when
        only some are active (for example during the last episodes of running
        eval episodes).
        """
        if self._is_waiting:
            for read_fn in self._connection_read_fns:
                read_fn()
        read_fn = self._connection_read_fns.pop(index)
        write_fn = self._connection_write_fns.pop(index)
        worker = self._workers.pop(index)
        self._paused.append((index, read_fn, write_fn, worker))

    def resume_all(self) -> None:
        r"""Resumes any paused envs."""
        for index, read_fn, write_fn, worker in reversed(self._paused):
            self._connection_read_fns.insert(index, read_fn)
            self._connection_write_fns.insert(index, write_fn)
            self._workers.insert(index, worker)
        self._paused = []

    def call_at(
        self,
        index: int,
        function_name: str,
        function_args: Optional[Dict[str, Any]] = None,
    ) -> Any:
        r"""Calls a function (which is passed by name) on the selected env and
        returns the result.

        :param index: which env to call the function on.
        :param function_name: the name of the function to call on the env.
        :param function_args: optional function args.
        :return: result of calling the function.
        """
        self._is_waiting = True
        self._connection_write_fns[index](
            (CALL_COMMAND, (function_name, function_args))
        )
        result = self._connection_read_fns[index]()
        self._is_waiting = False
        return result

    def call(
        self,
        function_names: List[str],
        function_args_list: Optional[List[Any]] = None,
    ) -> List[Any]:
        r"""Calls a list of functions (which are passed by name) on the
        corresponding env (by index).

        :param function_names: the name of the functions to call on the envs.
        :param function_args_list: list of function args for each function. If
            provided, :py:`len(function_args_list)` should be as long as
            :py:`len(function_names)`.
        :return: result of calling the function.
        """
        self._is_waiting = True
        if function_args_list is None:
            function_args_list = [None] * len(function_names)
        assert len(function_names) == len(function_args_list)
        func_args = zip(function_names, function_args_list)
        for write_fn, func_args_on in zip(
            self._connection_write_fns, func_args
        ):
            write_fn((CALL_COMMAND, func_args_on))
        results = []
        for read_fn in self._connection_read_fns:
            results.append(read_fn())
        self._is_waiting = False
        return results

    def render(
        self, mode: str = "human", *args, **kwargs
    ) -> Union[np.ndarray, None]:
        r"""Render observations from all environments in a tiled image."""
        for write_fn in self._connection_write_fns:
            write_fn((RENDER_COMMAND, (args, {"mode": "rgb", **kwargs})))
        images = [read_fn() for read_fn in self._connection_read_fns]
        tile = tile_images(images)
        if mode == "human":
            from habitat.core.utils import try_cv2_import

            cv2 = try_cv2_import()

            cv2.imshow("vecenv", tile[:, :, ::-1])
            cv2.waitKey(1)
            return None
        elif mode == "rgb_array":
            return tile
        else:
            raise NotImplementedError

    @property
    def _valid_start_methods(self) -> Set[str]:
        return {"forkserver", "spawn", "fork"}

    def _warn_cuda_tensors(
        self, action: Dict[str, Any], prefix: Optional[str] = None
    ):
        if torch is None:
            return

        for k, v in action.items():
            if isinstance(v, dict):
                subk = f"{prefix}.{k}" if prefix is not None else k
                self._warn_cuda_tensors(v, prefix=subk)
            elif torch.is_tensor(v) and v.device.type == "cuda":
                subk = f"{prefix}.{k}" if prefix is not None else k
                warnings.warn(
                    "Action with key {} is a CUDA tensor."
                    "  This will result in a CUDA context in the subproccess worker."
                    "  Using CPU tensors instead is recommended.".format(subk)
                )

    def __del__(self):
        self.close()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()


class ThreadedVectorEnv(VectorEnv):
    r"""Provides same functionality as :ref:`VectorEnv`, the only difference
    is it runs in a multi-thread setup inside a single process.

    The :ref:`VectorEnv` runs in a multi-proc setup. This makes it much easier
    to debug when using :ref:`VectorEnv` because you can actually put break
    points in the environment methods. It should not be used for best
    performance.
    """

    def _spawn_workers(
        self,
        env_fn_args: Sequence[Tuple],
        make_env_fn: Callable[..., Env] = _make_env_fn,
        workers_ignore_signals: bool = False,
    ) -> Tuple[List[Callable[[], Any]], List[Callable[[Any], None]]]:
        queues: Iterator[Tuple[Any, ...]] = zip(
            *[(Queue(), Queue()) for _ in range(self._num_envs)]
        )
        parent_read_queues, parent_write_queues = queues
        self._workers = []
        for parent_read_queue, parent_write_queue, env_args in zip(
            parent_read_queues, parent_write_queues, env_fn_args
        ):
            thread = Thread(
                target=self._worker_env,
                args=(
                    parent_write_queue.get,
                    parent_read_queue.put,
                    make_env_fn,
                    env_args,
                    self._auto_reset_done,
                ),
            )
            self._workers.append(thread)
            thread.daemon = True
            thread.start()
        return (
            [q.get for q in parent_read_queues],
            [q.put for q in parent_write_queues],
        )

habitat/core/vector_env.py
cshizhe/onav_rim
import glob
import gzip
import json
import multiprocessing
import os
import os.path as osp

import numpy as np
import pydash
import tqdm

import habitat
import habitat_sim
from habitat.datasets.object_nav.create.object_nav_generator import (
    build_goal,
    generate_objectnav_episode,
)
from habitat_sim.utils.common import quat_from_two_vectors, quat_rotate_vector

num_episodes_per_scene = int(5e4)

whitelist_classes = [
    {"children": [], "id": 1, "name": "backpack", "parents": []},
    {"children": [], "id": 3, "name": "basket", "parents": []},
    {"children": [], "id": 4, "name": "bathtub", "parents": []},
    {"children": [], "id": 6, "name": "beanbag", "parents": []},
    {"children": [], "id": 7, "name": "bed", "parents": []},
    {"children": [], "id": 8, "name": "bench", "parents": []},
    {"children": [], "id": 9, "name": "bike", "parents": []},
    {"children": [], "id": 10, "name": "bin", "parents": []},
    {"children": [], "id": 11, "name": "blanket", "parents": []},
    {"children": [], "id": 13, "name": "book", "parents": []},
    {"children": [], "id": 14, "name": "bottle", "parents": []},
    {"children": [], "id": 15, "name": "box", "parents": []},
    {"children": [], "id": 16, "name": "bowl", "parents": []},
    {"children": [], "id": 17, "name": "camera", "parents": []},
    {"children": [], "id": 18, "name": "cabinet", "parents": []},
    {"children": [], "id": 19, "name": "candle", "parents": []},
    {"children": [], "id": 20, "name": "chair", "parents": []},
    {"children": [], "id": 21, "name": "chopping-board", "parents": []},
    {"children": [], "id": 22, "name": "clock", "parents": []},
    {"children": [], "id": 23, "name": "cloth", "parents": []},
    {"children": [], "id": 24, "name": "clothing", "parents": []},
    {"children": [], "id": 25, "name": "coaster", "parents": []},
    {"children": [], "id": 26, "name": "comforter", "parents": []},
    {"children": [], "id": 27, "name": "computer-keyboard", "parents": []},
    {"children": [], "id": 28, "name": "cup", "parents": []},
    {"children": [], "id": 29, "name": "cushion", "parents": []},
    {"children": [], "id": 30, "name": "curtain", "parents": []},
    {"children": [], "id": 31, "name": "ceiling", "parents": []},
    {"children": [], "id": 32, "name": "cooktop", "parents": []},
    {"children": [], "id": 33, "name": "countertop", "parents": []},
    {"children": [], "id": 34, "name": "desk", "parents": []},
    {"children": [], "id": 35, "name": "desk-organizer", "parents": []},
    {"children": [], "id": 36, "name": "desktop-computer", "parents": []},
    {"children": [], "id": 37, "name": "door", "parents": []},
    {"children": [], "id": 38, "name": "exercise-ball", "parents": []},
    {"children": [], "id": 39, "name": "faucet", "parents": []},
    {"children": [], "id": 40, "name": "floor", "parents": []},
    {"children": [], "id": 41, "name": "handbag", "parents": []},
    {"children": [], "id": 42, "name": "hair-dryer", "parents": []},
    {"children": [], "id": 44, "name": "indoor-plant", "parents": []},
    {"children": [], "id": 45, "name": "knife-block", "parents": []},
    {"children": [], "id": 47, "name": "lamp", "parents": []},
    {"children": [], "id": 48, "name": "laptop", "parents": []},
    {"children": [], "id": 50, "name": "mat", "parents": []},
    {"children": [], "id": 51, "name": "microwave", "parents": []},
    {"children": [], "id": 52, "name": "monitor", "parents": []},
    {"children": [], "id": 53, "name": "mouse", "parents": []},
    {"children": [], "id": 54, "name": "nightstand", "parents": []},
    {"children": [], "id": 55, "name": "pan", "parents": []},
    {"children": [], "id": 57, "name": "paper-towel", "parents": []},
    {"children": [], "id": 58, "name": "phone", "parents": []},
    {"children": [], "id": 59, "name": "picture", "parents": []},
    {"children": [], "id": 60, "name": "pillar", "parents": []},
    {"children": [], "id": 61, "name": "pillow", "parents": []},
    {"children": [], "id": 62, "name": "pipe", "parents": []},
    {"children": [], "id": 64, "name": "plate", "parents": []},
    {"children": [], "id": 65, "name": "pot", "parents": []},
    {"children": [], "id": 66, "name": "rack", "parents": []},
    {"children": [], "id": 67, "name": "refrigerator", "parents": []},
    {"children": [], "id": 68, "name": "remote-control", "parents": []},
    {"children": [], "id": 69, "name": "scarf", "parents": []},
    {"children": [], "id": 70, "name": "sculpture", "parents": []},
    {"children": [], "id": 73, "name": "shower-stall", "parents": []},
    {"children": [], "id": 74, "name": "sink", "parents": []},
    {"children": [], "id": 76, "name": "sofa", "parents": []},
    {"children": [], "id": 77, "name": "stair", "parents": []},
    {"children": [], "id": 78, "name": "stool", "parents": []},
    {"children": [], "id": 79, "name": "switch", "parents": []},
    {"children": [], "id": 80, "name": "table", "parents": []},
    {"children": [], "id": 83, "name": "tissue-paper", "parents": []},
    {"children": [], "id": 84, "name": "toilet", "parents": []},
    {"children": [], "id": 85, "name": "toothbrush", "parents": []},
    {"children": [], "id": 86, "name": "towel", "parents": []},
    {"children": [], "id": 87, "name": "tv-screen", "parents": []},
    {"children": [], "id": 89, "name": "umbrella", "parents": []},
    {"children": [], "id": 90, "name": "utensil-holder", "parents": []},
    {"children": [], "id": 91, "name": "vase", "parents": []},
    {"children": [], "id": 96, "name": "wardrobe", "parents": []},
    {"children": [], "id": 98, "name": "rug", "parents": []},
    {"children": [], "id": 100, "name": "bag", "parents": []},
]

name_to_id = (
    pydash.chain()
    .group_by("name")
    .map_values(lambda v: v[0])
    .map_values("id")(whitelist_classes)
)
id_to_name = {v: k for k, v in name_to_id.items()}


def _generate_fn(scene):
    cfg = habitat.get_config()
    cfg.defrost()
    cfg.SIMULATOR.SCENE = scene
    cfg.SIMULATOR.AGENT_0.SENSORS = ["SEMANTIC_SENSOR"]
    cfg.SIMULATOR.SEMANTIC_SENSOR.HEIGHT = 256
    cfg.SIMULATOR.SEMANTIC_SENSOR.WIDTH = 256
    cfg.freeze()

    with open(osp.join(osp.dirname(scene), "info_semantic.json"), "r") as f:
        semantics = json.load(f)

    objects = semantics["objects"]
    object_by_class = pydash.group_by(objects, "class_id")
    object_by_class = {
        k: v
        for k, v in object_by_class.items()
        if k in pydash.map_(whitelist_classes, "id")
    }

    if False:
        for to_find in [
            "desk",
            "refrigerator",
            "toilet",
            "shower-stall",
            "bed",
            "sculpture",
        ]:
            if name_to_id[to_find] in object_by_class.keys():
                print("found {}".format(to_find))

    sim = habitat.sims.make_sim("Sim-v0", config=cfg.SIMULATOR)

    goals_by_class = {}
    for k, v in tqdm.tqdm(object_by_class.items()):
        goals = []
        for o in tqdm.tqdm(v, leave=False):
            goals.append(
                build_goal(
                    sim,
                    o["id"],
                    o["class_id"],
                    quat_rotate_vector(
                        quat_from_two_vectors(
                            np.array([0, 0, -1]), habitat_sim.geo.GRAVITY
                        ),
                        np.array(o["oriented_bbox"]["abb"]["center"]),
                    ),
                )
            )
        goals = [g for g in goals if g is not None]

        if len(goals) == 0:
            continue

        goals_by_class[k] = goals

    if len(goals_by_class) == 0:
        return None

    dset = habitat.datasets.make_dataset("ObjectNav-v1")
    eps_per_obj = 500
    with tqdm.tqdm(total=len(goals_by_class) * eps_per_obj) as pbar:
        for goals in goals_by_class.values():
            try:
                for ep in generate_objectnav_episode(
                    sim, goals, num_episodes=eps_per_obj
                ):
                    dset.episodes.append(ep)
                    pbar.update()
            except RuntimeError:
                pbar.update()

    for ep in dset.episodes:
        ep.scene_id = ep.scene_id[len("./data/scene_datasets/") :]

    scene_key = osp.basename(osp.dirname(osp.dirname(scene)))
    fname = f"./data/datasets/objectnav/replica/v2/train/content/{scene_key}.json.gz"
    os.makedirs(osp.dirname(fname), exist_ok=True)
    with gzip.open(fname, "wt") as f:
        f.write(dset.to_json())

    sim.close()


if __name__ == "__main__":
    mp_ctx = multiprocessing.get_context("forkserver")
    scenes = glob.glob(
        "./data/scene_datasets/replica_dataset/*/habitat/mesh_semantic.ply"
    )
    blacklist = {"room_1", "room_2", "hotel_0", "office_0", "office_1"}
    blacklist = blacklist.union({"apartment_0", "hotel_0", "aparment_2"})
    blacklist = {}
    scenes = [s for s in scenes if not any(k in s for k in blacklist)]
    with mp_ctx.Pool(20) as pool, tqdm.tqdm(total=len(scenes)) as pbar:
        for _ in pool.imap_unordered(_generate_fn, scenes):
            pbar.update()

habitat/datasets/object_nav/create/create_objectnav_dataset.py
cshizhe/onav_rim
#!/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 csv
import glob
import gzip
import itertools
import json
import lzma
import multiprocessing
import os
import os.path as osp
import random
import sys
import traceback
from collections import defaultdict

import GPUtil
import numpy as np
import pydash
import tqdm

import habitat
import habitat_sim
from habitat.config.default import get_config
from habitat.core.dataset import ObjectInScene, SceneState
from habitat.core.embodied_task import Episode
from habitat.core.logging import logger
from habitat.core.simulator import AgentState
from habitat.datasets import make_dataset

# from habitat.datasets.object_nav import object_nav_dataset
from habitat.datasets.object_nav.create.object_nav_generator import (
    generate_objectnav_episode,
)

# from habitat.datasets.object_nav.create.object_nav_generator_old import (
#     generate_new_obj_nav_episode,
# )
from habitat.datasets.pointnav.create.utils import (
    get_gibson_scenes,
    get_habitat_gibson_scenes,
    get_mp3d_scenes,
)

CFG_TEST = "configs/test/habitat_mp3d_eqa_test.yaml"
CLOSE_STEP_THRESHOLD = 0.028

COMPRESSION = ".gz"
OUTPUT_OBJ_FOLDER = (
    "./data/datasets/objectnav/mp3d/v1_allobjs_quaternion_fixed_2"
)

scenes_source = "mp3d"
OUTPUT_JSON_FOLDER = (
    f"./data/datasets/objectnav/{scenes_source}/objnav_ycb_objs_200k_single/"
)
# "./data/datasets/objectnav/mp3d/objnav_challenge2020_final_q_fixed_ud_50k_train_split_small_1m_local_planner_min/"
# "./data/datasets/objectnav/mp3d/v1_floatencoded_cats_90_normalized_round_obb_fixed_6"
NUM_GPUS = len(GPUtil.getAvailable(limit=256))
TASKS_PER_GPU = 12
tiny_gibson_scenes = [
    "Allensville",
    "Beechwood",
    "Benevolence",
    "Coffeen",
    "Collierville",
    "Corozal",
    "Cosmos",
    "Darden",
    "Forkland",
    "Hanson",
    "Hiteman",
    "Ihlen",
    "Klickitat",
    "Lakeville",
    "Leonardo",
    "Lindenwood",
    "Markleeville",
    "Marstons",
    "McDade",
    "Merom",
    "Mifflinburg",
    "Muleshoe",
    "Newfields",
    "Noxapater",
    "Onaga",
    "Pinesdale",
    "Pomaria",
    "Ranchester",
    "Shelbyville",
    "Stockman",
    "Tolstoy",
    "Uvalda",
    "Wainscott",
    "Wiconisco",
    "Woodbine",
]

LARGE_MP3D_SCENE = "1pXnuDYAj8r"

mp3dcat40_map = defaultdict(lambda: 41)
mp3dcat40_mapinv = dict()
mp3dcat_nametosynset = dict()
with open("examples/mpcat40.tsv", "r") as f:
    reader = csv.reader(f, delimiter="\t")
    next(reader)
    for line in reader:
        id_ = int(line[0].strip())
        synsets = line[3].strip().split(",")
        if len(synsets) == 0:
            continue
        mp3dcat40_mapinv[id_] = line[1].strip()
        for synset in synsets:
            mp3dcat40_map[synset] = id_
        mp3dcat_nametosynset[line[1].strip()] = synset

# print(mp3dcat_nametosynset)


def read_ycb_google_objects(
    path="habitat/datasets/object_nav/create/ycb_google_16k_meta.csv",
    ycb_objects_dir="data/assets/objects/ycb_google_16k_v2/configs_gltf/",
):
    with open(path, "r") as f:
        csv_r = csv.reader(f)
        headers = next(csv_r)
        filename_ind = headers.index("filename")
        object_name_ind = headers.index("object_name")
        synset_ind = headers.index("synset")
        object_category_ind = headers.index("category")
        ycb_objects_lib = [
            {
                "filename": ycb_objects_dir + line[filename_ind],
                "object_name": line[object_name_ind],
                "synset": line[synset_ind],
                "object_category": line[object_category_ind],
            }
            for line in csv_r
        ]
        ycb_objects_categories = {
            object["object_category"] for object in ycb_objects_lib
        }
        return ycb_objects_lib, sorted(list(ycb_objects_categories))


ycb_objects_lib, ycb_objects_categories = read_ycb_google_objects()
ycb_objects_categories = [
    "foodstuff",
    "stationery",
    "fruit",
    "plaything",
    "hand_tool",
    "game_equipment",
    "kitchenware",
]


def sample_objects_for_episode(
    objects_lib, num_uniq_selected_objects=5, num_copies=3
):
    # Up to 5 unique objects in the scene
    uniq_selected_objects = np.random.choice(
        objects_lib,
        size=np.random.randint(low=1, high=num_uniq_selected_objects),
    )
    if num_copies > 1:
        # Up to 2 instances of the same object
        selected_objects = [
            object.copy()
            for object in uniq_selected_objects
            for _ in range(np.random.randint(low=1, high=num_copies))
        ]
        return selected_objects
    return uniq_selected_objects


def sample_scene_state(selected_objects):
    max_semantic_object_id = 1 << 16  #  shift to avoid using first 2 bytes
    # If scene is loaded and sim is available.
    # sem_scene = sim.semantic_annotations()
    # if sem_scene and len(sem_scene.objects) > 0:
    #     max_semantic_object_id = max(
    #         [int(obj.id.split("_")[-1]) for obj in sem_scene.objects])

    for object_id, object in enumerate(selected_objects):
        object["object_id"] = str(max_semantic_object_id + object_id)

    return SceneState(
        objects=[
            ObjectInScene(
                object_id=str(max_semantic_object_id + object_id),
                object_template=f"{object['filename'].replace('.glb', '')}",
                scale=5,
                semantic_category_id=category_to_mp3d_category_id[
                    object["object_category"]
                ],
            )
            for object_id, object in enumerate(selected_objects)
        ]
    )


# with open("examples/coco_to_synset_edited.json", "r") as f:
#     coco_map = json.load(f)

mp3d_valid_ids = np.array(
    [3]
    + list(range(5, 9))
    + list(range(10, 12))
    + list(range(13, 16))
    + list(range(18, 21))
    + list(range(22, 24))
    + list(range(25, 28))
    + list(range(33, 35))
    + [38]
)

# np.array(list(set(mp3dcat40_map[synset] for synset in coco_map.values())))
mp3d_valid_ids = np.sort(mp3d_valid_ids)
mp3d_wordlist = [mp3dcat40_mapinv[a] for a in mp3d_valid_ids]
category_to_task_category_id = {k: int(v) for v, k in enumerate(mp3d_wordlist)}
category_to_mp3d_category_id = {
    k: int(v) for k, v in zip(mp3d_wordlist, mp3d_valid_ids)
}

# Appending YCB categories
max_task_cat = max(category_to_task_category_id.values())
max_mp3d_cat = 42  # Max MP3D cat value
for ycb_category_id, category in enumerate(ycb_objects_categories):
    category_to_task_category_id[category] = max_task_cat + ycb_category_id + 1
    category_to_mp3d_category_id[category] = max_mp3d_cat + ycb_category_id + 1

# print(ycb_objects_categories)
# print(category_to_task_category_id)


def get_objnav_config(i, scene):

    CFG_EQA = "configs/tasks/challenge_objectnav2021.local.rgbd.yaml"
    objnav_config = get_config(CFG_EQA).clone()
    objnav_config.defrost()
    # print(objnav_config.SIMULATOR.AGENT_0.RADIUS)
    # objnav_config.SIMULATOR_GPU_ID = i % 2
    objnav_config.TASK.SENSORS = []
    objnav_config.SIMULATOR.AGENT_0.SENSORS = ["SEMANTIC_SENSOR"]
    # objnav_config.SIMULATOR.AGENT_0.RADIUS = 0.2
    # HEIGHT = objnav_config.SIMULATOR.AGENT_0.HEIGHT
    # objnav_config.SIMULATOR.AGENT_0.HEIGHT = HEIGHT
    # objnav_config.SIMULATOR.SEMANTIC_SENSOR.HEIGHT = 64
    # objnav_config.SIMULATOR.SEMANTIC_SENSOR.WIDTH = 64
    # objnav_config.SIMULATOR.RGB_SENSOR.HEIGHT = 64
    # objnav_config.SIMULATOR.RGB_SENSOR.WIDTH = 64
    FOV = 90
    objnav_config.SIMULATOR.RGB_SENSOR.HFOV = FOV
    # objnav_config.SIMULATOR.RGB_SENSOR.VFOV = FOV
    # objnav_config.SIMULATOR.RGB_SENSOR.POSITION = [0, HEIGHT, 0]
    objnav_config.SIMULATOR.DEPTH_SENSOR.HFOV = FOV
    # objnav_config.SIMULATOR.DEPTH_SENSOR.VFOV = FOV
    # objnav_config.SIMULATOR.DEPTH_SENSOR.POSITION = [0, HEIGHT, 0]
    objnav_config.SIMULATOR.SEMANTIC_SENSOR.HFOV = FOV
    # objnav_config.SIMULATOR.SEMANTIC_SENSOR.VFOV = FOV
    # objnav_config.SIMULATOR.SEMANTIC_SENSOR.POSITION = [0, HEIGHT, 0]

    # TODO lower resolution of semantic sensor
    objnav_config.TASK.MEASUREMENTS = []
    # objnav_config.TASK.POSSIBLE_ACTIONS = ['MOVE_FORWARD', 'TURN_LEFT', 'TURN_RIGHT', 'STOP']
    deviceIds = GPUtil.getAvailable(
        order="memory", limit=1, maxLoad=1.0, maxMemory=1.0
    )
    if i < NUM_GPUS * TASKS_PER_GPU or len(deviceIds) == 0:
        deviceId = i % NUM_GPUS
    else:
        deviceId = deviceIds[0]
    objnav_config.SIMULATOR.HABITAT_SIM_V0.GPU_DEVICE_ID = (
        deviceId  # i % NUM_GPUS
    )
    objnav_config.DATASET.DATA_PATH = (
        "./data/datasets/pointnav/mp3d/v1/val/val.json.gz"
    )
    objnav_config.DATASET.SCENES_DIR = "./data/scene_datasets/"
    objnav_config.DATASET.SPLIT = "train"
    # objnav_config.TASK.MEASUREMENTS.append("TOP_DOWN_MAP")
    # objnav_config.TASK.SENSORS.append("HEADING_SENSOR")
    # new_test_scene = "/private/home/maksymets/data/gibson/gibson_tiny_manual_verified/" + tiny_gibson_scenes[0] + '.glb'
    # new_test_scene = scene#'data/scene_datasets/mp3d/17DRP5sb8fy/17DRP5sb8fy.glb'
    objnav_config.SIMULATOR.SCENE = scene  # new_test_scene
    objnav_config.freeze()
    return objnav_config


def get_simulator(objnav_config):
    sim = habitat.sims.make_sim("Sim-v0", config=objnav_config.SIMULATOR)
    navmesh_settings = habitat_sim.NavMeshSettings()
    navmesh_settings.set_defaults()
    navmesh_settings.agent_radius = objnav_config.SIMULATOR.AGENT_0.RADIUS
    navmesh_settings.agent_height = objnav_config.SIMULATOR.AGENT_0.HEIGHT
    sim.recompute_navmesh(sim.pathfinder, navmesh_settings)
    return sim


def get_gravity_mobb(object_obb: habitat_sim.geo.OBB):
    bounding_area = [
        (object_obb.local_to_world @ np.array([x, y, z, 1]))[:-1]
        for x, y, z in itertools.product(*([[-1, 1]] * 3))
    ]
    bounding_area = np.array(bounding_area, dtype=np.float32)
    # print('Bounding Area: %s' % bounding_area)
    # TODO Maybe Cache this
    return habitat_sim.geo.compute_gravity_aligned_MOBB(
        habitat_sim.geo.GRAVITY, bounding_area
    )


def get_goals_view_points(sim, goal_objects, agent_state):
    from habitat.tasks.nav.object_nav_task import (
        ObjectGoal,
        ObjectGoalNavEpisode,
        ObjectViewLocation,
    )

    goals = [
        ObjectGoal(
            object_id=object["object_id"],
            object_category=object["object_category"],
            position=object["position"],
            view_points=[
                ObjectViewLocation(
                    agent_state=AgentState(
                        position=np.array(
                            sim.pathfinder.snap_point(object["position"])
                        )
                    ),
                    iou=0.5 # TODO add some legit value
                )
            ],
        )
        for object in goal_objects
    ]

    goal_targets = (
        [vp.agent_state.position for vp in goal.view_points] for goal in goals
    )

    closest_goal_targets = (
        sim.geodesic_distance(agent_state, vps) for vps in goal_targets
    )
    # closest_goal_targets, goals_sorted = zip(
    #     *sorted(
    #         zip(closest_goal_targets, goals),
    #         key=lambda x: x[0],
    #     )
    # )
    return goals, list(closest_goal_targets)


def reset_goal_objects(sim, goal_objects, agent_state, closest_goal_targets):
    retries = 0
    from habitat_sim.physics import MotionType

    while np.inf in closest_goal_targets and retries < 20:
        retries += 1

        for object_order, object in enumerate(goal_objects):
            if closest_goal_targets[object_order] != np.inf and retries < 10:
                continue
            sim_object_id = sim.objid_to_sim_object_mapping[
                object["object_id"]
            ]
            sim.set_object_motion_type(MotionType.KINEMATIC, sim_object_id)
            # sim.set_translation(sim.pathfinder.get_random_navigable_point(), sim_object_id), sim.get_translation(sim_object_id)
            while not sim.sample_object_state(sim_object_id):
                pass
            sim.set_object_motion_type(MotionType.STATIC, sim_object_id)
            object["position"] = list(sim.get_translation(sim_object_id))
            object["rotation"] = list(sim.get_rotation_vec(sim_object_id))
            object["obb"] = sim.get_object_scene_node(
                sim_object_id
            ).cumulative_bb
            sim.recompute_navmesh(sim.pathfinder, sim.navmesh_settings, True)
        if np.all(closest_goal_targets == np.inf):
            s = sim.sample_navigable_point()
        goals, closest_goal_targets = get_goals_view_points(
            sim, goal_objects=goal_objects, agent_state=agent_state
        )
        # logger.info(f"closest_goal_targets: {closest_goal_targets} rt {retries}")


def fill_episode(
    episode,
    agent_state,
    episode_counter,
    goal_category,
    goals,
    selected_objects,
    sim,
):
    episode.start_position = agent_state
    angle = np.random.uniform(0, 2 * np.pi)
    episode.start_rotation = [0, np.sin(angle / 2), 0, np.cos(angle / 2)]
    sim.set_agent_state(agent_state, episode.start_rotation)
    episode.goals = goals
    episode_counter += 1
    episode.episode_id = str(episode_counter)
    episode.object_category = goal_category
    from habitat.core.dataset import ObjectInScene, SceneState

    episode.scene_state = SceneState(
        objects=[
            ObjectInScene(
                object_id=object["object_id"],
                object_template=f"{object['filename'].replace('.glb', '')}",
                rotation=object["rotation"],
                position=object["position"],
                scale=5,
                semantic_category_id=category_to_mp3d_category_id[
                    object["object_category"]
                ],
            )
            for object in selected_objects
        ]
    )


def reset_episode(
    episode, sim, selected_objects, near_dist=1, far_dist=5, episode_counter=0
):

    for object in selected_objects:
        sim_object_id = sim.objid_to_sim_object_mapping[object["object_id"]]
        object["position"] = list(sim.get_translation(sim_object_id))
        object["rotation"] = list(sim.get_rotation_vec(sim_object_id))
        object["obb"] = sim.get_object_scene_node(sim_object_id).cumulative_bb

    category_to_object = {}
    for object in selected_objects:
        category_to_object[object["object_category"]] = category_to_object.get(
            object["object_category"], []
        ) + [object]
    goal_category = np.random.choice(list(category_to_object.keys()))

    from habitat.datasets.pointnav.pointnav_generator import (
        ISLAND_RADIUS_LIMIT,
    )

    s = sim.sample_navigable_point()
    while sim.island_radius(s) < ISLAND_RADIUS_LIMIT:
        s = sim.sample_navigable_point()

    goals, closest_goal_targets = get_goals_view_points(
        sim, goal_objects=category_to_object[goal_category], agent_state=s
    )

    reset_goal_objects(
        sim,
        goal_objects=category_to_object[goal_category],
        agent_state=s,
        closest_goal_targets=closest_goal_targets,
    )

    retries = 0
    while (
        np.inf in closest_goal_targets
        or not near_dist <= min(closest_goal_targets) <= far_dist
    ):
        retries += 1
        if retries > 200 and np.inf in closest_goal_targets:
            logger.info(f"Falling back to previous episode as dist is inf.")
            return None
        s = sim.sample_navigable_point()
        goals, closest_goal_targets = get_goals_view_points(
            sim, goal_objects=category_to_object[goal_category], agent_state=s
        )

    fill_episode(
        episode,
        s,
        episode_counter,
        goal_category,
        goals,
        selected_objects,
        sim,
    )

    return episode


def generate_scene(args):
    i, scene, split = args
    objnav_config = get_objnav_config(i, scene)

    sim = get_simulator(objnav_config)
    # Check there exists a navigable point
    test_point = sim.sample_navigable_point()
    if not sim.is_navigable(np.array(test_point)):
        print("Scene is not navigable: %s" % scene)
        sim.close()
        return scene, 0, defaultdict(list), None

    # goals_by_class = defaultdict(list)
    # cell_size = objnav_config.SIMULATOR.AGENT_0.RADIUS / 2.0
    # for obj in tqdm.tqdm(selected_objects, desc="Objects for %s:" % scene):
    #     print("Object id: %d" % obj["id"])
    #     print(obj["category_name"])
    #
    #     goal = build_goal(
    #         sim,
    #         object_id=obj["object_id"],
    #         object_name_id=obj["object_name"],
    #         object_category_name=obj["category_name"],
    #         object_category_id=1, #obj["category_id"],
    #         object_position=obj["position"],
    #         object_aabb=obj["aabb"],
    #         object_obb=obj["obb"],
    #         object_gmobb=obj["gravity_mobb"],
    #         cell_size=cell_size,
    #         grid_radius=3.0,
    #     )
    #     if goal == None:  # or len(goal.view_points) == 0:
    #         continue
    #     goals_by_class[obj["category_id"]].append(goal)
    #     # break
    # os.makedirs(osp.dirname(fname_obj), exist_ok=True)
    # total_objects_by_cat = {k: len(v) for k, v in goals_by_class.items()}
    # with open(fname_obj, "wb") as f:
    #     pickle.dump(goals_by_class, f)

    scene_key = habitat.Dataset.scene_from_scene_path(scene)
    fname = (
        f"{OUTPUT_JSON_FOLDER}/{split}/content/{scene_key}.json{COMPRESSION}"
    )

    if os.path.exists(fname):
        print("Scene already generated. Skipping")
        sim.close()
        return scene, 0, 0, None

    total_episodes = 2000 if split == "train" else 10
    # if split == 'val':
    #    total_obs_json = 100
    # if split == 'val_mini':
    #    total_episodes = 30
    dset = habitat.datasets.make_dataset("ObjectNav-v1")
    dset.category_to_task_category_id = category_to_task_category_id
    dset.category_to_mp3d_category_id = category_to_mp3d_category_id
    eps_generated = 0
    with tqdm.tqdm(total=total_episodes, desc=scene) as pbar:

        try:
            while eps_generated < total_episodes:
                sim.habitat_config.defrost()
                selected_objects = sample_objects_for_episode(ycb_objects_lib)
                sim.habitat_config.scene_state = [
                    sample_scene_state(selected_objects).__dict__
                ]
                sim.habitat_config.freeze()
                sim._initialize_scene()

                for object in selected_objects:
                    sim_object_id = sim.objid_to_sim_object_mapping[
                        object["object_id"]
                    ]
                    object["position"] = list(
                        sim.get_translation(sim_object_id)
                    )
                    object["rotation"] = list(
                        sim.get_rotation_vec(sim_object_id)
                    )
                    object["obb"] = sim.get_object_scene_node(
                        sim_object_id
                    ).cumulative_bb

                category_to_object = {}
                for object in selected_objects:
                    category_to_object[
                        object["object_category"]
                    ] = category_to_object.get(
                        object["object_category"], []
                    ) + [
                        object
                    ]

                from habitat.core.simulator import AgentState
                from habitat.tasks.nav.object_nav_task import (
                    ObjectGoal,
                    ObjectGoalNavEpisode,
                    ObjectViewLocation,
                )

                # all_view_points = [np.array(
                #                             sim.pathfinder.snap_point(
                #                                 object['position']))
                #             for object in selected_objects]
                #
                # source_position = sim.sample_navigable_point()
                # while sim.geodesic_distance(source_position, all_view_points) is not math.inf:
                #     source_position = sim.sample_navigable_point()
                # angle = np.random.uniform(0, 2 * np.pi)
                # source_rotation = [
                #     0,
                #     np.sin(angle / 2),
                #     0,
                #     np.cos(angle / 2),
                # ]  # Pick random starting rotation

                for goal_category in category_to_object.keys():
                    # Generate episode for each category
                    # goal_category = \
                    # np.random.choice(list(category_to_object.keys()), size=1)[
                    #     0]
                    try:
                        for ep in generate_objectnav_episode(
                            sim,
                            goals=[
                                ObjectGoal(
                                    object_id=object["object_id"],
                                    object_category=object["object_category"],
                                    position=object["position"],
                                    view_points=[
                                        ObjectViewLocation(
                                            agent_state=AgentState(
                                                position=np.array(
                                                    sim.pathfinder.snap_point(
                                                        object["position"]
                                                    )
                                                )
                                            ),
                                        )
                                    ],
                                )
                                for object in category_to_object[goal_category]
                            ],
                            scene_state=SceneState(
                                objects=[
                                    ObjectInScene(
                                        object_id=object["object_id"],
                                        object_template=f"{object['filename'].replace('.glb', '')}",
                                        rotation=object["rotation"],
                                        position=object["position"],
                                        scale=5,
                                        semantic_category_id=category_to_mp3d_category_id[
                                            object["object_category"]
                                        ],
                                    )
                                    for object in selected_objects
                                ]
                            ),
                            num_episodes=1,
                        ):
                            ep.episode_id = eps_generated
                            dset.episodes.append(ep)
                            pbar.update()
                            eps_generated += 1
                    except RuntimeError:
                        pbar.update()

                    # episode = ObjectGoalNavEpisode(
                    #     episode_id=str(eps_generated),
                    #     scene_id=sim.habitat_config.SCENE,
                    #     scene_state=SceneState(
                    #         objects=[ObjectInScene(
                    #             object_id=str(object_id),
                    #             object_template=f"{object['filename'].replace('.glb', '')}",
                    #             rotation=object['rotation'],
                    #             position=object['position'],
                    #             scale=5,
                    #             semantic_id=category_to_mp3d_category_id[
                    #                 object["object_category"]]
                    #         ) for object_id, object in
                    #             enumerate(selected_objects)]
                    #     ),
                    #     start_position=source_position,
                    #     start_rotation=source_rotation,
                    #     object_category=goal_category,
                    #     goals=[
                    #         ObjectGoal(
                    #             object_id=str(object['object_id']),
                    #             object_category=object['object_category'],
                    #             position=object['position'],
                    #             view_points=[ObjectViewLocation(
                    #                 agent_state=AgentState(
                    #                     position=np.array(
                    #                         sim.pathfinder.snap_point(
                    #                             object['position']))
                    #                 )
                    #             )]
                    #         )
                    #         for object in category_to_object[goal_category]
                    #     ],
                    # )
                # print("\n +1 iterations")
        except RuntimeError:
            traceback.print_exc()
            print("Skipping category")

    for ep in dset.episodes:
        # STRIP OUT PATH
        ep.scene_id = ep.scene_id[len("./data/scene_datasets/") :]

    os.makedirs(osp.dirname(fname), exist_ok=True)
    save_dataset(dset, fname)
    print(f"Saved dataset {fname}.")
    # dset2 = habitat.datasets.make_dataset('ObjectNav-v1')
    # with gzip.open(fname, 'rt') as f:
    #     #dset2f.read()
    #     dset2.from_json(f.read())
    # assert dset2.category_vocab_dict_dense.word_list == dset.category_vocab_dict_dense.word_list
    # assert dset2.category_vocab_dict_sparse == dset.category_vocab_dict_sparse
    sim.close()
    return scene, eps_generated, 0, fname


SPLIT = None


def get_file_opener(fname):
    ext = os.path.splitext(fname)[-1]

    if ext == ".gz":
        file_opener = gzip.open
    elif ext == ".xz":
        file_opener = lzma.open
    else:
        print(ext)
        assert False
        # ile_opener = open
    return file_opener


def save_dataset(dset: habitat.Dataset, fname: str):
    file_opener = get_file_opener(fname)
    # compression = gzip if format == 'gzip' else lzma
    # dset = dset.from_json(dset.to_json())
    # ddset = dset.from_json(json.loads(json.dumps(dset.to_json())))
    if (
        os.path.basename(os.path.dirname(fname)) == "content"
        and len(dset.episodes) == 0
    ):
        print("WARNING UNEXPECTED EMPTY EPISODES: %s" % fname)
        return
    with file_opener(fname, "wt") as f:
        if len(dset.episodes) == 0:
            print("WARNING EMPTY EPISODES: %s" % fname)
            f.write(
                json.dumps(
                    {
                        "episodes": [],
                        "category_to_task_category_id": dset.category_to_task_category_id,
                        "category_to_mp3d_category_id": dset.category_to_mp3d_category_id,
                    }
                )
            )
        else:
            # dset = object_nav_dataset.ObjectNavDatasetV1.dedup_goals_dset(dset)
            f.write(dset.to_json())


def read_dset(json_fname):
    print(json_fname)
    dset2 = habitat.datasets.make_dataset("ObjectNav-v1")
    file_opener = get_file_opener(json_fname)

    # compression = gzip if os.path.splitext(json_fname)[-1] == 'gz' else lzma
    with file_opener(json_fname, "rt") as f:
        dset2.from_json(f.read())
    return dset2


if __name__ == "__main__":
    mp_ctx = multiprocessing.get_context("fork")

    new_test_scene = (
        "/private/home/maksymets/data/gibson/gibson_tiny_manual_verified/"
        + tiny_gibson_scenes[0]
        + ".glb"
    )

    a = dict()
    split = sys.argv[1]
    mode = sys.argv[2] if len(sys.argv) > 2 else "run"
    # read_dset('./data/datasets/objectnav/mp3d/objnav_challenge2020_v2/val/content/2azQ1b91cZZ.json.gz')
    # Just for testing
    def prepare_args(split):
        print(split)
        if split in ["test_standard", "test_challenge"]:
            filename = "examples/mp3d_" + split + ".txt"
            print(filename)
            with open(filename, "r") as f:
                # print(f.read())
                scenes = [line.strip() for line in f]
                print(scenes)
        else:
            if scenes_source == "gibson":
                scenes = get_gibson_scenes(split)  #  get_gibson_scenes(split)
            else:
                scenes = get_mp3d_scenes(split)
                # scenes = [LARGE_MP3D_SCENE]  # debug remove

        random.shuffle(scenes)
        if scenes_source == "gibson":
            scenes = [
                f"./data/scene_datasets/gibson/{scene}.glb" for scene in scenes
            ]

        else:
            scenes = [
                f"./data/scene_datasets/mp3d/{scene}/{scene}.glb"
                for scene in scenes
            ]

        blacklist = ["JmbYfDe2QKZ"]
        # if split != 'val' assert
        scenes = [
            scene
            for scene in scenes
            if os.path.basename(os.path.dirname(scene)) not in blacklist
        ]
        if split == "val_mini":
            scenes = [scenes[0]]
        return [(i, scene, split) for i, scene in enumerate(scenes)]

    np.random.seed(1234)
    if split == "*":
        args = []
        for split in ["train", "test", "val"]:
            args += prepare_args(split)
    else:
        args = prepare_args(split)

    ############### GENERATE EPISODES ########################
    GPU_THREADS = NUM_GPUS * TASKS_PER_GPU  # 1  #debug
    print(GPU_THREADS)
    print("*" * 1000)
    CPU_THREADS = int(multiprocessing.cpu_count() / 4) if mode == "run" else 1
    print(f"CPU_THREADS: {CPU_THREADS}")
    with mp_ctx.Pool(CPU_THREADS, maxtasksperchild=1) as pool, tqdm.tqdm(
        total=len(args)
    ) as pbar, open("train_subtotals.json", "w") as f:
        total_all = 0
        subtotals = []
        for scene, subtotal, subtotal_by_cat, fname in pool.imap_unordered(
            generate_scene, args
        ):
            a[scene] = (subtotal, subtotal_by_cat)
            # print("*" * 10)
            pbar.update()
            # print(a)
            total_all += subtotal
            subtotals.append(subtotal_by_cat)
        print(total_all)
        print(subtotals)

        json.dump({"total_objects:": total_all, "subtotal": subtotals}, f)
    # sys.exit(0)

    if split == "*":
        splits = ["val", "test", "train"]
    else:
        splits = [split]

    ############### GENERATE EMPTY root dataset and mini_val ########################
    for split in splits:
        dset = habitat.datasets.make_dataset("ObjectNav-v1")
        dset.category_to_task_category_id = category_to_task_category_id
        dset.category_to_mp3d_category_id = category_to_mp3d_category_id
        global_dset = f"{OUTPUT_JSON_FOLDER}/{split}/{split}.json{COMPRESSION}"
        if os.path.exists(global_dset):
            os.remove(global_dset)
        if not os.path.exists(os.path.dirname(global_dset)):
            os.mkdir(os.path.dirname(global_dset))
        # if split != "train":
        jsons_gz = glob.glob(
            f"{OUTPUT_JSON_FOLDER}/{split}/content/*.json{COMPRESSION}"
        )

        with mp_ctx.Pool(CPU_THREADS) as pool, tqdm.tqdm(
            total=len(jsons_gz), desc=f"Loading jsons for {split}"
        ) as pbar:
            for i, dset2 in enumerate(pool.map(read_dset, jsons_gz)):
                if i == 0 and split == "val":
                    MINI_SAMPLE = 10  # 30
                    print("Sampling %d episodes for minival" % MINI_SAMPLE)
                    min_dset_fname = f"{OUTPUT_JSON_FOLDER}/{split}_mini/{split}_mini.json{COMPRESSION}"
                    os.makedirs(os.path.dirname(min_dset_fname), exist_ok=True)

                    min_dset = habitat.datasets.make_dataset("ObjectNav-v1")
                    min_dset.category_to_task_category_id = (
                        category_to_task_category_id
                    )
                    min_dset.category_to_mp3d_category_id = (
                        category_to_mp3d_category_id
                    )

                    min_dset.episodes = list(
                        dset2.get_episode_iterator(
                            num_episode_sample=MINI_SAMPLE, cycle=False
                        )
                    )
                    save_dataset(min_dset, min_dset_fname)
                    # with gzip.open(min_dset_fname, "wt") as f:
                    #    f.write(min_dset.to_json())

                # dset.episodes.extend(dset2.episodes)
                pbar.update()
        save_dataset(dset, global_dset)

habitat/datasets/object_nav/create/create_objectnav_dataset_with_added_objects.py
cshizhe/onav_rim
import os

import imageio
import numpy as np

IMAGE_DIR = "data/images/objnav_dataset_gen/debug"
MAX_DIST = [0, 0, 200]  # Blue
NON_NAVIGABLE = [150, 150, 150]  # Grey
POINT_COLOR = [150, 150, 150]  # Grey
VIEW_POINT_COLOR = [0, 200, 0]  # Green
CENTER_POINT_COLOR = [200, 0, 0]  # Red


def plot_area(points, points_iou, object_center, object_name_id):
    max_coord = 1000  # int(np.max(points) * 100)
    image = np.zeros((max_coord, max_coord, 3), dtype=np.uint8)

    def mark_points(points, color):
        int_points = [
            (int(p[0] * 10 + max_coord / 2), int(p[2] * 10 + max_coord / 2))
            for p in points
        ]
        for p in int_points:
            image[p[0], p[1]] = color

    def iou_points(points, color):
        int_points = [
            (
                int(p[0] * 10 + max_coord / 2),
                int(p[2] * 10 + max_coord / 2),
                iou,
            )
            for iou, p, _ in points
        ]
        for p in int_points:
            if p[2] == -1:
                image[p[0], p[1]] = NON_NAVIGABLE
            elif p[2] == -0.5:
                image[p[0], p[1]] = MAX_DIST
            else:
                color = int((p[2] + 1.0) * 255)
                image[p[0], p[1]] = [color, color, color]

    iou_points(points, POINT_COLOR)
    mark_points(points_iou, VIEW_POINT_COLOR)
    mark_points(object_center, CENTER_POINT_COLOR)

    imageio.imsave(
        os.path.join(IMAGE_DIR, f"objnav_image_{object_name_id}.png"), image
    )

habitat/datasets/object_nav/create/debug_utils.py
cshizhe/onav_rim
import gzip
import json
import os
import random
from typing import Any, Dict, List

import numpy as np
from tqdm import tqdm
from yacs.config import CfgNode

import habitat
import habitat.datasets.pointnav.create.utils as utils
from habitat.config.default import get_config
from habitat.datasets.pointnav.create.pointnav_dataset_creator import (
    PointNavDatasetCreator,
)
from habitat.datasets.semantic_nav import generate_new_obj_nav_episode

# from habitat.datasets.pointnav.pointnav_dataset import PointNavDatasetV1
# from habitat.tasks.nav.nav_task import NavigationEpisode, NavigationGoal

try:
    from _json import encode_basestring_ascii
except ImportError:
    encode_basestring_ascii = None
try:
    from _json import encode_basestring
except ImportError:
    encode_basestring = None


class DatasetFloatJSONEncoder(json.JSONEncoder):
    """
        JSON Encoder that set float precision for space saving purpose.
    """

    # Version of JSON library that encoder is compatible with.
    __version__ = "2.0.9"

    def default(self, object):
        return object.__dict__

    # Overriding method to inject own `_repr` function for floats with needed
    # precision.
    def iterencode(self, o, _one_shot=False):

        if self.check_circular:
            markers = {}
        else:
            markers = None
        if self.ensure_ascii:
            _encoder = encode_basestring_ascii
        else:
            _encoder = encode_basestring

        def floatstr(
            o,
            allow_nan=self.allow_nan,
            _repr=lambda x: format(x, ".5f"),
            _inf=float("inf"),
            _neginf=-float("inf"),
        ):
            if o != o:
                text = "NaN"
            elif o == _inf:
                text = "Infinity"
            elif o == _neginf:
                text = "-Infinity"
            else:
                return _repr(o)

            if not allow_nan:
                raise ValueError(
                    "Out of range float values are not JSON compliant: "
                    + repr(o)
                )

            return text

        _iterencode = json.encoder._make_iterencode(
            markers,
            self.default,
            _encoder,
            self.indent,
            floatstr,
            self.key_separator,
            self.item_separator,
            self.sort_keys,
            self.skipkeys,
            _one_shot,
        )
        return _iterencode(o, 0)


class ObjectGoalNavDatasetCreatorV1(PointNavDatasetCreator):
    r"""Point Navigation dataset.
    """

    def generate_dataset(
        self,
        config: Any = None,
        scenes: List[str] = None,
        query_data: Dict = None,
    ) -> None:
        self.episodes = []

        if config is None:
            return

        dataset = EQDDatasetV1()

        if not scenes:
            scenes = self.__class__.get_scenes(config.split)
            # val_mini_scenes = self.__class__.get_scenes("val_mini",
            #                                            config.scene_path)
            # scenes = [scene for scene in scenes if scene not in val_mini_scenes]

        if config.num_houses > 0:
            scenes = scenes[: config.num_houses]
        progress_bar = tqdm(total=config.num_episodes)
        np.random.seed(config.seed)

        env_cfg = get_config("configs/tasks/eqd_create.yaml")
        env_cfg.defrost()
        env_cfg.gpu_device_id = config.gpu_id
        env_cfg.task_name = "Nav-v0"
        env_cfg.freeze()

        num_queries = sum([len(query_data[scene].items()) for scene in scenes])
        num_episodes_per_query = round(config.num_episodes / num_queries)

        for scene in scenes:
            if scene not in query_data:
                print(
                    "{} skipped as not presented in query data.".format(scene)
                )
                continue
            print(f"\nQueries from: {scene}")
            env_cfg.defrost()
            env_cfg.SIMULATOR.SCENE = config.scene_path.format(scene=scene)
            env_cfg.freeze()
            env = habitat.Env(config=env_cfg)
            env.seed(config.seed)
            env.episodes = [env.episodes[0]]
            env.episodes[0].scene_id = config.scene_path.format(scene=scene)
            env.reset()

            from time import time

            for query_id, query in enumerate(query_data[scene].items()):
                query[1]["type"] = "objects:room:level"
                print(f"Query id: {query_id}, {query[0]}, {query[1]}")
                t_gen_episode = time()
                for episode in generate_new_obj_nav_episode(
                    env, query, num_episodes_per_query
                ):
                    print(f"t_gen_episode: {time() - t_gen_episode }")
                    t_gen_episode = time()
                    episode.episode_id = len(dataset.episodes)
                    dataset.episodes.append(episode)

                    progress_bar.update(1)
                if len(dataset.episodes) > 1:
                    break
            env.close()
            del env

        print("scene_count: ", len(scenes))
        print("episodes_count: ", len(dataset.episodes))
        return dataset

    @staticmethod
    def get_scenes(split, scene_path="{scene}"):
        return utils.get_mp3d_scenes(split, scene_path)


def get_default_mp3d_v1_config(split: str = "train"):
    # "test_standard" test_challenge
    config = CfgNode()
    config.seed = 704184
    config.num_episodes = 30
    config.num_houses = 0
    config.split = split  # data/datasets/pointnav/gibson/v1
    config.output_path = "data/datasets/eqd/mp3d/v1/{}_30_eps/" "".format(
        config.split
    )

    config.IS_SINGLE_FILE = False
    config.OUTPUT_PATH_SINGLE_SCENE = "{}content/{{scene}}.json.gz".format(
        config.output_path
    )

    config.gpu_id = 0
    config.scene_path = "data/scene_datasets/mp3d/{scene}/{scene}.glb"
    return config


def main():
    config = get_default_mp3d_v1_config()
    np.random.seed(config.seed)
    random.seed(config.seed)

    if not os.path.exists(config.output_path):
        os.mkdir(config.output_path)

    if config.IS_SINGLE_FILE:
        dataset = ObjectGoalNavDatasetCreatorV1(config)
        json_str = str(dataset.to_json())

        main_dataset_file = "{}{}.json.gz".format(
            config.output_path, config.split
        )
        with gzip.GzipFile(main_dataset_file, "wb") as f:
            f.write(json_str.encode("utf-8"))

        print("Dataset file: {}".format(main_dataset_file))

    else:
        cur_config = config.clone()
        dataset = ObjectGoalNavDatasetCreatorV1(cur_config)
        json_str = str(dataset.to_json())

        with gzip.GzipFile(
            "{}{}.json.gz".format(config.output_path, config.split), "wb"
        ) as f:
            f.write(json_str.encode("utf-8"))

        print("Dataset file: {}".format(config.output_path))

        scenes = ObjectGoalNavDatasetCreatorV1.get_scenes(config.split)

        if config.num_houses > 0:
            scenes = scenes[: config.num_houses]
        progress_bar = tqdm(total=config.num_episodes)
        difficulty_counts = {}

        finished_episodes = 0
        for scene in scenes:
            scene_dataset_path = config.OUTPUT_PATH_SINGLE_SCENE.format(
                scene=scene
            )
            scene_path = config.scene_path.format(scene=scene)

            if os.path.exists(scene_dataset_path):
                print("{} exists.".format(scene_dataset_path))
                finished_episodes += round(config.num_episodes / len(scenes))
                continue

            cur_config.num_episodes = (
                round(config.num_episodes / len(scenes))
                if scene != scenes[-1]
                else config.num_episodes - finished_episodes
            )
            dataset = ObjectGoalNavDatasetCreatorV1(
                config=cur_config, scenes=[scene_path]
            )
            finished_episodes += len(dataset.episodes)
            json_str = str(dataset.to_json())

            if not os.path.exists(os.path.dirname(scene_dataset_path)):
                os.makedirs(os.path.dirname(scene_dataset_path))

            with gzip.GzipFile(scene_dataset_path, "wb") as f:
                f.write(json_str.encode("utf-8"))
            print("Dataset file: {}.".format(scene_dataset_path))
            progress_bar.update(len(dataset.episodes))


if __name__ == "__main__":
    main()

habitat/datasets/object_nav/create/obj_nav_dataset_creator.py
cshizhe/onav_rim
#!/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 itertools
import random
from typing import List, Optional

import numpy as np
import pydash

import habitat_sim
from habitat.core.dataset import SceneState
from habitat.core.simulator import AgentState, ShortestPathPoint
from habitat.datasets.pointnav.pointnav_generator import (
    ISLAND_RADIUS_LIMIT,
    _ratio_sample_rate,
)
from habitat.datasets.utils import get_action_shortest_path
from habitat.sims.habitat_simulator.actions import HabitatSimActions
from habitat.sims.habitat_simulator.habitat_simulator import HabitatSim
from habitat.tasks.nav.nav import NavigationEpisode
from habitat.tasks.nav.object_nav_task import (
    ObjectGoal,
    ObjectGoalNavEpisode,
    ObjectViewLocation,
)
from habitat.tasks.utils import compute_pixel_coverage
from habitat.utils.geometry_utils import (
    angle_between_quaternions,
    quaternion_from_two_vectors,
)

# from habitat.utils.geometry_utils import direction_to_quaternion
from habitat_sim.errors import GreedyFollowerError
from habitat_sim.utils.common import (
    quat_from_angle_axis,
    quat_to_angle_axis,
    quat_to_coeffs,
)


def _direction_to_quaternion(direction_vector: np.array):
    origin_vector = np.array([0, 0, -1])
    return quaternion_from_two_vectors(origin_vector, direction_vector)


def _get_multipath(sim: HabitatSim, start, ends):
    multi_goal = habitat_sim.MultiGoalShortestPath()
    multi_goal.requested_start = start
    multi_goal.requested_ends = ends
    sim.pathfinder.find_path(multi_goal)
    return multi_goal


def _get_action_shortest_path(
    sim: HabitatSim, start_pos, start_rot, goal_pos, goal_radius=0.05
):
    sim.set_agent_state(start_pos, start_rot, reset_sensors=True)
    greedy_follower = sim.make_greedy_follower()
    return greedy_follower.find_path(goal_pos)


def is_compatible_episode(
    s,
    t,
    sim: HabitatSim,
    goals: List[ObjectGoal],
    near_dist,
    far_dist,
    geodesic_to_euclid_ratio,
):
    FAIL_TUPLE = False, 0, 0, [], [], [], []
    if sim.island_radius(s) < ISLAND_RADIUS_LIMIT:
        return FAIL_TUPLE
    s = np.array(s)
    # distance = np.array(list(map(lambda pos: sim.geodesic_distance(s, pos), t)))

    # assert(len(goals_repeated) == len(t))
    # TWO TARGETS MAY BE BETWEEN TWO GOALS
    goal_targets = (
        [vp.agent_state.position for vp in goal.view_points] for goal in goals
    )

    # old_school
    # closest_goal_targets = (min((sim.geodesic_distance(s, vp), vp) for vp in vps) for vps in goal_targets)
    # closest_goal_targets, goals_sorted = zip(*sorted(zip(closest_goal_targets, goals)))
    # d_separation, closest_target = closest_goal_targets[0]
    # closest_goal_targets = [_get_multipath(sim ,s, goal_targets)]
    closest_goal_targets = (
        sim.geodesic_distance(s, vps) for vps in goal_targets
    )
    closest_goal_targets, goals_sorted = zip(
        *sorted(zip(closest_goal_targets, goals), key=lambda x: x[0])
    )
    d_separation = closest_goal_targets[0]

    if (
        np.inf in closest_goal_targets
        or not near_dist <= d_separation <= far_dist
    ):
        return FAIL_TUPLE

    # shortest_path = sim.get_straight_shortest_path_points(s, closest_target)
    shortest_path = None
    # shortest_path = closest_goal_targets[0].points
    shortest_path_pos = shortest_path
    euclid_dist = np.linalg.norm(s - goals_sorted[0].position)
    distances_ratio = d_separation / euclid_dist
    if distances_ratio < geodesic_to_euclid_ratio and (
        np.random.rand()
        > _ratio_sample_rate(distances_ratio, geodesic_to_euclid_ratio)
    ):
        return FAIL_TUPLE

    # geodesic_distances, _ = zip(*closest_goal_targets)
    geodesic_distances = closest_goal_targets

    # euclid_distances_target = np.linalg.norm(np.array([goal.position for goal in goals]) - shortest_path[-1:], axis=1)
    # _, goals_sorted = zip(*sorted(zip(euclid_distances_target, goals)))
    # goal_index = np.argmin(euclid_distances_target)
    # goals_sorted = [goals[goal_index]] + goals[:goal_index] + goals[goal_index+1:]

    angle = np.random.uniform(0, 2 * np.pi)
    source_rotation = [
        0,
        np.sin(angle / 2),
        0,
        np.cos(angle / 2),
    ]  # Pick random starting rotation

    # try:
    #     action_shortest_path = _get_action_shortest_path(
    #         sim, s, source_rotation, shortest_path[-1]
    #     )
    #     if action_shortest_path == None:
    #         return FAIL_TUPLE
    #     shortest_path = (
    #         action_shortest_path
    #     )
    #     # [ShortestPathPoint(point, [0,0,0], action) for point, action in zip(shortest_path, action_shortest_path)]
    # except GreedyFollowerError:
    #     print(
    #         "Could not find path between %s and %s"
    #         % (str(s), str(shortest_path[-1]))
    #     )
    #     return FAIL_TUPLE
    ending_state = None
    # ending_state = closest_goal_targets[0].points[-1]  # sim.get_agent_state()
    # ending_position = ending_state.position
    #
    # e_q = ending_state.rotation  # We presume the agent is upright
    # # print(shortest_path)
    # goal_direction = goals_sorted[0].position - ending_position
    # goal_direction[1] = 0
    # a_q = _direction_to_quaternion(goal_direction)
    # quat_delta = e_q - a_q
    # # The angle between the two quaternions should be how much to turn
    # theta, _ = quat_to_angle_axis(quat_delta)
    # if theta < 0 or theta > np.pi:
    #     turn = HabitatSimActions.TURN_LEFT
    #     if theta > np.pi:
    #         theta = 2 * np.pi - theta
    # else:
    #     turn = HabitatSimActions.TURN_RIGHT
    # turn_angle = np.deg2rad(sim.config.TURN_ANGLE)
    # num_of_turns = int(theta / turn_angle + 0.5)
    # shortest_path += [turn] * num_of_turns
    # angle = angle_between_quaternions(a_q, e_q)

    # if len(shortest_path) > 750:
    #     print("ERROR SHORTEST PATH IS TOO LONG")
    #     return FAIL_TUPLE

    # shortest_path = get_action_shortest_path(sim,
    #     s,
    #     source_rotation,
    #     goals_sorted[0].position,
    #     max_episode_steps=750)
    # #CANNOT CATCH ERRORS
    # if shortest_path == None:
    #     return FAIL_TUPLE
    # Make it so it doesn't have to see the object initially
    # TODO Consider turning this check back on later?
    # num_rotation_attempts = 20
    # for attempt in range(num_rotation_attempts):
    #     angle = np.random.uniform(0, 2 * np.pi)
    #     source_rotation = [0, np.sin(angle / 2), 0, np.cos(angle / 2)]

    #     obs = sim.get_observations_at(s, source_rotation)

    #     if all(
    #         (compute_pixel_coverage(obs["semantic"], g.object_id) / g.best_iou) < 0.01
    #         for g in goals_sorted
    #     ):
    #         break

    # if attempt == num_rotation_attempts:
    #     return FAIL_TUPLE

    return (
        True,
        d_separation,
        euclid_dist,
        source_rotation,
        geodesic_distances,
        goals_sorted,
        shortest_path,
        ending_state,
    )


def build_goal(
    sim: HabitatSim,
    object_id: int,
    object_name_id: int,
    object_category_name: str,
    object_category_id: int,
    object_position,
    object_aabb: habitat_sim.geo.BBox,
    object_obb: habitat_sim.geo.OBB,
    object_gmobb: habitat_sim.geo.OBB,
    cell_size: float = 1.0,
    grid_radius: float = 10.0,
    turn_radians: float = np.pi / 9,
    max_distance: float = 1.0,
):
    object_position = object_aabb.center
    eps = 1e-5

    x_len, y_len, z_len = object_aabb.sizes / 2.0 + max_distance
    x_bxp = np.arange(-x_len, x_len + eps, step=cell_size) + object_position[0]
    # y_bxp = np.arange(-y_len, y_len + eps, step=1) + object_position[1]
    # y = np.arange(-y_len, y_len, step=cell_size)
    z_bxp = np.arange(-z_len, z_len + eps, step=cell_size) + object_position[2]
    # print('Original obb %s' % object_obb.center )
    # #print('Original calculated: %f %f %f' % (x_len, y_len, z_len))
    # print('Obb calculated: %s ' % obb.center)
    # print('Obb calculated: %s' % obb.half_extents)
    # print('Obb calculated sizes: %s' % (obb.sizes / 2.0))
    # print('AABB: %s %s' % (object_aabb.center, object_aabb.sizes))
    candiatate_poses = [
        np.array([x, object_position[1], z])
        for x, z in itertools.product(x_bxp, z_bxp)
    ]

    # x_len, y_len, z_len = obb.half_extents + MAX_DISTANCE

    # x_arr = np.arange(-x_len, x_len + eps, step=cell_size, dtype=np.float32) / obb.half_extents[0]
    # z_arr = np.arange(-z_len, z_len + eps, step=cell_size, dtype=np.float32) / obb.half_extents[2]

    # candiatate_poses = [(obb.local_to_world @ np.array([x, 0, z, 1]))[:-1] for x, z in itertools.product(x_arr, z_arr)]
    # print('New Bounding Area: %s' % np.array([(obb.local_to_world @ np.array([x, y, z, 1]))[:-1] for x,y,z in
    #     itertools.product(*([[-1, 1]] * 3))]))
    # #print('Candiate Poses %s' % str(candiatate_poses))
    # print('Candiate Poses variance: %s ' % str(habitat_sim.geo.compute_gravity_aligned_MOBB(habitat_sim.geo.GRAVITY, candiatate_poses).sizes))
    # #candiatate_poses = [np.array([p[0], obb.center[1], p[2]]) for p in candiatate_poses] #if obb.distance(np.array([p[0], obb.center[1], p[2]])) <= MAX_DISTANCE]
    # print(candiatate_poses)
    # obb = habitat_sim.geo.OBB(object_aabb)
    # thetas = np.arange(0, 2 * np.pi, step=turn_radians)
    def down_is_navigable(pt):
        pf = sim.pathfinder

        delta_y = 0.05
        max_steps = int(2 / delta_y)
        step = 0
        is_navigable = pf.is_navigable(pt, 2)
        while not is_navigable:
            pt[1] -= delta_y
            is_navigable = pf.is_navigable(pt)
            step += 1
            if step == max_steps:
                return False
        return True

    def _get_iou(x, y, z):
        pt = np.array([x, y, z])

        if not (
            object_gmobb.distance(pt) <= max_distance
            and habitat_sim.geo.OBB(object_aabb).distance(pt) <= max_distance
        ):
            return -0.5, pt, None

        if not down_is_navigable(pt):
            return -1.0, pt, None
        # if not pf.is_navigable(pt, 2) and not pf.is_navigable(np.array(pf.snap_point(pt)), 2):
        #     # other_points = [(pf.is_navigable(p), p) for p in itertools.product([-x,x], [-y,y], [-z,z])]
        #     # if any(t[0] for t in other_points):
        #     #     print(pt)
        #     #     print(other_points)
        #     return -1.0, pt, None
        pf = sim.pathfinder
        pt = np.array(pf.snap_point(pt))

        goal_direction = object_position - pt
        # print(goal_direction)
        goal_direction[1] = 0

        q = _direction_to_quaternion(goal_direction)
        # POSITION MAY BE UNREACHABLE
        # goal_angle, _ = quat_to_angle_axis(goal_rot)
        # theta = goal

        # print(pt)
        # print(pt)
        # print('*' * 100)

        # q = quat_from_angle_axis(theta, habitat_sim.geo.UP)
        cov = 0
        agent = sim.get_agent(0)
        for act in [
            HabitatSimActions.LOOK_DOWN,
            HabitatSimActions.LOOK_UP,
            HabitatSimActions.LOOK_UP,
        ]:
            # if act is not None:
            agent.act(act)
            for v in agent._sensors.values():
                v.set_transformation_from_spec()
            obs = sim.get_observations_at(pt, q, keep_agent_at_new_pose=True)
            cov += compute_pixel_coverage(obs["semantic"], object_id)
        from habitat.utils.visualizations.utils import observations_to_image
        import imageio
        import os

        obs = sim.get_observations_at(pt, q, keep_agent_at_new_pose=True)
        imageio.imsave(
            os.path.join(
                "data/images/objnav_dataset_gen",
                f"{object_name_id}_{object_id}_{x}_{z}_.png",
            ),
            observations_to_image(obs, info={}),
        )

        return cov, pt, q

    # candiatate_poses = [
    #     object_position + np.array([x, 0, z])
    #     #(obb.local_to_world @ np.array([x, y_len, z, 1]))[:-1]
    #     for x, z in itertools.product(x, z)
    # ]
    # PROPER
    # END
    # candiatate_poses = (object_obb.local_to_world @ np.array([x, 0, z, 1]) for x, y, z in itertools.product(x, z))
    # candiatate_poses = [p[:3] for p in candiatate_poses]
    # candiatate_poses = [(x, y, z) for x, z in itertools.product(x, z)]
    candiatate_poses_ious = list(_get_iou(*pos) for pos in candiatate_poses)

    # print(np.mean([a[1] - object_position for a in candiatate_poses_ious if np.any(a[1] != None) and a[0] != 0], axis=0))
    best_iou = (
        max(v[0] for v in candiatate_poses_ious)
        if len(candiatate_poses_ious) != 0
        else 0
    )
    if best_iou <= 0.0:
        print(
            f"No view points found for {object_name_id}_{object_id}: {best_iou}"
        )
        return None

    keep_thresh = 0  # 1/(256**2)#0.75 * best_iou

    view_locations = [
        ObjectViewLocation(
            AgentState(pt.tolist(), quat_to_coeffs(q).tolist()), iou
        )
        for iou, pt, q in candiatate_poses_ious
        if iou is not None and iou > keep_thresh
    ]
    import habitat.datasets.object_nav.create.debug_utils as debug_utils

    debug_utils.plot_area(
        candiatate_poses_ious,
        [v.agent_state.position for v in view_locations],
        [object_position],
        object_category_name + object_name_id,
    )

    view_locations = sorted(view_locations, reverse=True, key=lambda v: v.iou)
    if len(view_locations) == 0:
        print(
            f"No valid views found for {object_name_id}_{object_id}: {best_iou}"
        )
        return None
    # for view in view_locations:
    #     obs = sim.get_observations_at(
    #         view.agent_state.position, view.agent_state.rotation
    #     )

    #     from habitat.utils.visualizations.utils import (
    #        observations_to_image,
    #        get_image_with_obj_overlay,
    #     )
    #     import imageio
    #     import os

    #     obs["rgb"] = get_image_with_obj_overlay(obs, [object_id])
    #     imageio.imsave(
    #        os.path.join(
    #            "data/images/objnav_dataset_gen/",
    #            f"{object_name_id}_{object_id}_{view.iou}_{view.agent_state.position}.png",
    #        ),
    #        observations_to_image(obs, info={}).astype(np.uint8),
    #     )
    goal = ObjectGoal(
        position=object_position.tolist(),
        view_points=view_locations,
        object_id=object_id,
        object_category=object_category_name,
        object_name=object_name_id,
    )

    return goal


def _create_episode(
    episode_id,
    scene_id,
    start_position,
    start_rotation,
    goals,
    shortest_paths=None,
    scene_state=None,
    info=None,
):
    return ObjectGoalNavEpisode(
        episode_id=str(episode_id),
        goals=goals,
        scene_id=scene_id,
        object_category=goals[0].object_category,
        # scene_state=scene_state,
        start_position=start_position,
        start_rotation=start_rotation,
        shortest_paths=shortest_paths,
        info=info,
    )


def generate_objectnav_episode(
    sim: HabitatSim,
    goals: List[ObjectGoal],
    scene_state: SceneState = None,
    num_episodes: int = -1,
    closest_dist_limit: float = 0.2,
    furthest_dist_limit: float = 30,
    geodesic_to_euclid_min_ratio: float = 1.05,
    number_retries_per_target: int = 1000,
):
    r"""Generator function that generates PointGoal navigation episodes.
    An episode is trivial if there is an obstacle-free, straight line between
    the start and goal positions. A good measure of the navigation
    complexity of an episode is the ratio of
    geodesic shortest path position to Euclidean distance between start and
    goal positions to the corresponding Euclidean distance.
    If the ratio is nearly 1, it indicates there are few obstacles, and the
    episode is easy; if the ratio is larger than 1, the
    episode is difficult because strategic navigation is required.
    To keep the navigation complexity of the precomputed episodes reasonably
    high, we perform aggressive rejection sampling for episodes with the above
    ratio falling in the range [1, 1.1].
    Following this, there is a significant decrease in the number of
    straight-line episodes.
    :param sim: simulator with loaded scene for generation.
    :param num_episodes: number of episodes needed to generate
    :param is_gen_shortest_path: option to generate shortest paths
    :param shortest_path_success_distance: success distance when agent should
    stop during shortest path generation
    :param shortest_path_max_steps maximum number of steps shortest path
    expected to be
    :param closest_dist_limit episode geodesic distance lowest limit
    :param furthest_dist_limit episode geodesic distance highest limit
    :param geodesic_to_euclid_min_ratio geodesic shortest path to Euclid
    distance ratio upper limit till aggressive sampling is applied.
    :return: navigation episode that satisfy specified distribution for
    currently loaded into simulator scene.
    """
    episode_count = 0
    while episode_count < num_episodes or num_episodes < 0:
        # target_positions = (
        #     pydash.chain()
        #     .map(lambda g: g.view_points)
        #     .flatten().map(lambda v: v.agent_state.position)(goals)
        # )
        # Cache this transformation
        target_positions = np.array(
            list(
                itertools.chain(
                    *(
                        (
                            view_point.agent_state.position
                            for view_point in g.view_points
                        )
                        for g in goals
                    )
                )
            )
        )

        for retry in range(number_retries_per_target):
            source_position = sim.sample_navigable_point()
            if (
                source_position is None
                or np.any(np.isnan(source_position))
                or not sim.is_navigable(source_position)
            ):
                raise RuntimeError("Unable to find valid starting location")
            if sim.island_radius(source_position) < ISLAND_RADIUS_LIMIT:
                continue
            compat_outputs = is_compatible_episode(
                source_position,
                target_positions,
                sim,
                goals,
                near_dist=closest_dist_limit,
                far_dist=furthest_dist_limit,
                geodesic_to_euclid_ratio=geodesic_to_euclid_min_ratio,
            )
            is_compatible = compat_outputs[0]

            if is_compatible:
                is_compatible, dist, euclid_dist, source_rotation, geodesic_distances, goals_sorted, shortest_path, ending_state = (
                    compat_outputs
                )
                shortest_paths = None #[shortest_path]

                episode = _create_episode(
                    episode_id=episode_count,
                    scene_id=sim.habitat_config.SCENE,
                    start_position=source_position,
                    start_rotation=source_rotation,
                    shortest_paths=shortest_paths,
                    scene_state=scene_state,
                    info={
                        "geodesic_distance": dist,
                        "euclidean_distance": euclid_dist,
                        "closest_goal_object_id": goals_sorted[0].object_id,
                        # "navigation_bounds": sim.pathfinder.get_bounds(),
                        # "best_viewpoint_position": ending_state.position,
                    },
                    goals=goals_sorted,
                )

                episode_count += 1
                yield episode
                break

        if retry == number_retries_per_target:
            raise RuntimeError("Unable to find valid starting location")


def generate_objectnav_episode_with_added_objects(
    sim: HabitatSim,
    objects: List,
    goal_category: str,
    goals: List[ObjectGoal],
    scene_state: SceneState = None,
    num_episodes: int = -1,
    closest_dist_limit: float = 0.2,
    furthest_dist_limit: float = 30,
    geodesic_to_euclid_min_ratio: float = 1.05,
    number_retries_per_target: int = 1000,
):
    r"""Generator function that generates PointGoal navigation episodes.
    An episode is trivial if there is an obstacle-free, straight line between
    the start and goal positions. A good measure of the navigation
    complexity of an episode is the ratio of
    geodesic shortest path position to Euclidean distance between start and
    goal positions to the corresponding Euclidean distance.
    If the ratio is nearly 1, it indicates there are few obstacles, and the
    episode is easy; if the ratio is larger than 1, the
    episode is difficult because strategic navigation is required.
    To keep the navigation complexity of the precomputed episodes reasonably
    high, we perform aggressive rejection sampling for episodes with the above
    ratio falling in the range [1, 1.1].
    Following this, there is a significant decrease in the number of
    straight-line episodes.
    :param sim: simulator with loaded scene for generation.
    :param num_episodes: number of episodes needed to generate
    :param is_gen_shortest_path: option to generate shortest paths
    :param shortest_path_success_distance: success distance when agent should
    stop during shortest path generation
    :param shortest_path_max_steps maximum number of steps shortest path
    expected to be
    :param closest_dist_limit episode geodesic distance lowest limit
    :param furthest_dist_limit episode geodesic distance highest limit
    :param geodesic_to_euclid_min_ratio geodesic shortest path to Euclid
    distance ratio upper limit till aggressive sampling is applied.
    :return: navigation episode that satisfy specified distribution for
    currently loaded into simulator scene.
    """
    episode_count = 0
    while episode_count < num_episodes or num_episodes < 0:
        # target_positions = (
        #     pydash.chain()
        #     .map(lambda g: g.view_points)
        #     .flatten().map(lambda v: v.agent_state.position)(goals)
        # )
        # Cache this transformation
        target_positions = np.array(
            list(
                itertools.chain(
                    *(
                        (
                            view_point.agent_state.position
                            for view_point in g.view_points
                        )
                        for g in goals
                    )
                )
            )
        )

        for retry in range(number_retries_per_target):
            source_position = sim.sample_navigable_point()
            if (
                source_position is None
                or np.any(np.isnan(source_position))
                or not sim.is_navigable(source_position)
            ):
                raise RuntimeError("Unable to find valid starting location")
            if sim.island_radius(source_position) < ISLAND_RADIUS_LIMIT:
                continue
            compat_outputs = is_compatible_episode(
                source_position,
                target_positions,
                sim,
                goals,
                near_dist=closest_dist_limit,
                far_dist=furthest_dist_limit,
                geodesic_to_euclid_ratio=geodesic_to_euclid_min_ratio,
            )
            is_compatible = compat_outputs[0]

            if is_compatible:
                is_compatible, dist, euclid_dist, source_rotation, geodesic_distances, goals_sorted, shortest_path, ending_state = (
                    compat_outputs
                )
                shortest_paths = [shortest_path]

                episode = _create_episode(
                    episode_id=episode_count,
                    scene_id=sim.habitat_config.SCENE,
                    start_position=source_position,
                    start_rotation=source_rotation,
                    shortest_paths=shortest_paths,
                    scene_state=scene_state,
                    info={
                        "geodesic_distance": dist,
                        "euclidean_distance": euclid_dist,
                        "closest_goal_object_id": goals_sorted[0].object_id,
                        # "navigation_bounds": sim.pathfinder.get_bounds(),
                        # "best_viewpoint_position": ending_state.position,
                    },
                    goals=goals_sorted,
                )

                episode_count += 1
                yield episode
                break

        if retry == number_retries_per_target:
            raise RuntimeError("Unable to find valid starting location")

habitat/datasets/object_nav/create/object_nav_generator.py
cshizhe/onav_rim
import random
from time import time
from typing import List

import numpy as np
from pydash import py_

import habitat
from habitat.core.simulator import AgentState, ShortestPathPoint
from habitat.datasets.pointnav.pointnav_generator import is_compatible_episode
from habitat.tasks.nav.nav import NavigationEpisode
from habitat.tasks.nav.object_nav_task import ObjectGoal
from habitat.tasks.nav.shortest_path_follower import (  # mp3d_direction_to_hsim_quaternion,
    ShortestPathFollower,
)
from habitat.utils.geometry_utils import (
    direction_to_quaternion,
    quaternion_to_list,
)

GEODESIC_TO_EUCLID_RATIO_THRESHOLD = 1.1
NUMBER_RETRIES_PER_TARGET = 10000
NEAR_DIST_LIMIT = 1
FAR_DIST_LIMIT = 30
ISLAND_RADIUS_LIMIT = 1.5

DIFFICULTIES_PROP = {
    "easy": {"start": NEAR_DIST_LIMIT, "end": 6},
    "medium": {"start": 6, "end": 8},
    "hard": {"start": 8, "end": FAR_DIST_LIMIT},
}

EPISODE_COUNT = 0

TEMPLATES = {
    "find the {OBJ} in the {ROOM}{LEVEL}",
    "show me the {OBJ} in the {ROOM}{LEVEL}",
    "go to the {ROOM}{LEVEL} and find all the {OBJ} there",
    "check out the {OBJ} in the {ROOM}{LEVEL}",
}


def create_query(obj_name, obj_count, room, level=None):
    """
    :param obj_name: object category
    :param obj_count: the number of objects (not used for now)
    :param room: room name
    :param level: level id
    :return: query string
    """
    level_txt = " on level {}".format(level) if level else ""
    template = random.choice(TEMPLATES)
    txt = template.format(OBJ=obj_name, ROOM=room, LEVEL=level_txt)

    return txt


def _ratio_sample_rate(ratio):
    """
     :param ratio: geodesic distance ratio to Euclid distance
    :return: value between 0.008 and 0.144 for ration 1 and 1.1
    """
    return 20 * (ratio - 0.98) ** 2


def draw_top_down_map(info, heading, output_size):
    from habitat.utils.visualizations import maps
    import cv2

    top_down_map = maps.colorize_topdown_map(info["top_down_map"]["map"])
    original_map_size = top_down_map.shape[:2]
    map_scale = np.array(
        (1, original_map_size[1] * 1.0 / original_map_size[0])
    )
    new_map_size = np.round(output_size * map_scale).astype(np.int32)
    # OpenCV expects w, h but map size is in h, w
    top_down_map = cv2.resize(top_down_map, (new_map_size[1], new_map_size[0]))

    map_agent_pos = info["top_down_map"]["agent_map_coord"]
    map_agent_pos = np.round(
        map_agent_pos * new_map_size / original_map_size
    ).astype(np.int32)
    top_down_map = maps.draw_agent(
        top_down_map,
        map_agent_pos,
        heading - np.pi / 2,
        agent_radius_px=top_down_map.shape[0] / 40,
    )
    return top_down_map


def create_episode(
    episode, objects, start_position, start_rotation, shortest_paths, info
) -> NavigationEpisode:
    goals = [
        ObjectGoal(
            object_id=obj["id"],
            object_name=obj["instance_id"],
            object_category=obj["category_name"],
            view_points=[
                AgentState(position=view[0][0], rotation=view[0][1])
                for view in obj["view_points"]
            ],
            position=obj["center"],
        )
        for obj in objects
    ]
    new_episode = NavigationEpisode(
        episode_id=episode.episode_id,
        goals=goals,
        scene_id=episode.scene_id,
        start_position=start_position,
        start_rotation=start_rotation,
        start_room=episode.start_room,
        shortest_paths=shortest_paths,
        info=info,
    )
    return new_episode


def get_action_shortest_path(
    env, goal_position, goal_rotation, radius=0.05, record_video=False
) -> List[ShortestPathPoint]:
    import os
    from habitat.utils.visualizations.utils import images_to_video

    # print("BEFORE:", env.sim.get_agent_state())
    # pos_bef = env.sim.get_agent_state().position.tolist()
    # rot_bef = env.sim.get_agent_state().rotation.tolist()
    observations = env.reset()
    if record_video:
        images = []
        im = observations["rgb"]
        top_down_map = draw_top_down_map(
            env.get_metrics(), observations["heading"], im.shape[0]
        )
        output_im = np.concatenate((im, top_down_map), axis=1)
        images.append(output_im)

        print("Environment creation successful")
        IMAGE_DIR = os.path.join("data", "videos")
        dirname = os.path.join(IMAGE_DIR, "eqd_v1")
    # print("EPISODE:", env.current_episode.start_position,
    #       type(env.current_episode.start_position),
    #       env.current_episode.start_rotation,
    #       type(env.current_episode.start_rotation)
    #       )
    # # env.sim.set_agent_state(env.current_episode.start_position, env.current_episode.start_rotation)
    # print("AFTER:", env.sim.get_agent_state())
    # env._task.measurements.reset_measures(episode=env.current_episode)

    mode = "greedy"

    goal_radius = radius  # env.episodes[0].goals[0].radius
    if goal_radius is None or goal_radius == 0:
        goal_radius = env._config.SIMULATOR.FORWARD_STEP_SIZE
    # scene = env.sim.semantic_annotations()
    # object_id = env.current_episode.goals[0].object_id
    # goal_center = scene.objects[object_id].obb.center.tolist()
    # env.current_episode.goals[0].position = goal_center

    follower = ShortestPathFollower(env.sim, goal_radius, False)
    follower.mode = mode

    # print("Agent stepping around inside environment.")

    shortest_path = []

    while not env.episode_over:
        action = follower.get_next_action(
            goal_position,  # env.current_episode.goals[0].position
            # goal_rotation,
            env.current_episode.goals[0].position,
        )
        state = env.sim.get_agent_state()
        # print("state.position: ", state.position, ", state.rotation: ",
        #       state.rotation)

        shortest_path.append(
            ShortestPathPoint(
                state.position.tolist(),
                quaternion_to_list(state.rotation),
                action.value,
            )
        )

        observations = env.step(action.value)
        if record_video:
            im = observations["rgb"]
            top_down_map = draw_top_down_map(
                env.get_metrics(), observations["heading"], im.shape[0]
            )
            output_im = np.concatenate((im, top_down_map), axis=1)
            images.append(output_im)

    state = env.sim.get_agent_state()
    # print("state.position: ", state.position, ", state.rotation: ",
    #       state.rotation)

    shortest_path.append(
        ShortestPathPoint(
            state.position.tolist(),
            quaternion_to_list(state.rotation),
            action.value,
        )
    )

    if record_video:
        print(
            "goal_position: ",
            goal_position,
            ", goal_rotation: ",
            goal_rotation,
            "goal_rot_orig: ",
            goal_rotation,
        )
        print(
            "start_position: ",
            env.current_episode.start_position,
            ", " "start_rotation: ",
            env.current_episode.start_rotation,
        )
        video_filename = "{}_{}_{}".format(
            env.current_episode.episode_id,
            env.current_episode.goals[0].name,
            env.current_episode.goals[0].room_name,
        )
        video_filename = video_filename.replace(" ", "_")
        # print( "Heights difference: ", env.sim.get_agent_state().position[1] -
        #        goal_position[1])
        images_to_video(images, dirname, video_filename)
        print("Episode finished: {}".format(dirname + video_filename))
        # print("Shortest path length: {}".format(len(shortest_path)))

    if len(shortest_path) == env._config.ENVIRONMENT.MAX_EPISODE_STEPS:
        shortest_path = []
        print(
            "BAD SHORTEST PATH of the EPISODE: {}".format(
                env.current_episode.episode_id
            )
        )
    return [shortest_path]


def get_image_with_obj_overlay(observations, objects):
    rgb = observations["rgb"]
    sem = observations["semantic"] == -1
    for obj in objects:
        sem = np.logical_or(
            sem, observations["semantic"] == obj["instance_id"]
        )
    color_sem = np.array([(0, 0, 0), (255, 0, 0)])[np.asanyarray(sem, np.int)]
    return rgb + 0.5 * color_sem


def get_action_reference_path(
    env, objects, record_video=False
) -> List[ShortestPathPoint]:
    import os
    from habitat.utils.visualizations.utils import images_to_video
    from habitat.sims.habitat_simulator import SimulatorActions

    radius = 0.1

    # print("BEFORE:", env.sim.get_agent_state())
    # pos_bef = env.sim.get_agent_state().position.tolist()
    # rot_bef = env.sim.get_agent_state().rotation.tolist()
    observations = env.reset()
    if record_video:
        images = []
        observations = env.reset()
        im = observations["rgb"]
        top_down_map = draw_top_down_map(
            env.get_metrics(), observations["heading"], im.shape[0]
        )
        output_im = np.concatenate((im, top_down_map), axis=1)
        images.append(output_im)
        IMAGE_DIR = os.path.join("data", "videos")
        dirname = os.path.join(IMAGE_DIR, "eqd_v1")

    # print("EPISODE:", env.current_episode.start_position,
    #       type(env.current_episode.start_position),
    #       env.current_episode.start_rotation,
    #       type(env.current_episode.start_rotation)
    #       )
    # # env.sim.set_agent_state(env.current_episode.start_position, env.current_episode.start_rotation)
    # print("AFTER:", env.sim.get_agent_state())
    # env._task.measurements.reset_measures(episode=env.current_episode)

    mode = "greedy"
    goal_radius = radius  # env.episodes[0].goals[0].radius
    if goal_radius is None or goal_radius == 0:
        goal_radius = env._config.SIMULATOR.FORWARD_STEP_SIZE
    # scene = env.sim.semantic_annotations()
    # object_id = env.current_episode.goals[0].object_id
    # goal_center = scene.objects[object_id].obb.center.tolist()
    # env.current_episode.goals[0].position = goal_center

    follower = ShortestPathFollower(env.sim, goal_radius, False)
    follower.mode = mode

    shortest_path = []
    geo_dist_time = 0
    objects_to_visit = objects.copy()
    while len(objects_to_visit) > 0:

        cur_position = env.sim.get_agent_state().position
        view_points = [
            (obj_id, view)
            for obj_id, obj in enumerate(objects_to_visit)
            for view in obj["view_points"]
        ]

        # view_dist = np.array(list(map(lambda o: [env.sim.geodesic_distance(
        #     cur_position, o[1][0][0]), o[1][0]], view_points)))

        view_dist = np.asarray(
            list(
                map(
                    lambda o: env.sim.geodesic_distance(
                        cur_position, o[1][0][0]
                    ),
                    view_points,
                )
            )
        )
        next_view = view_points[np.argmin(view_dist, axis=0)]
        goal_position = next_view[1][0][0]
        goal_rotation = next_view[1][0][1]
        del objects_to_visit[next_view[0]]
        action = None

        while (
            action != SimulatorActions.STOP
            and len(shortest_path) <= env._config.ENVIRONMENT.MAX_EPISODE_STEPS
        ):
            t_geo_dist_time = time()
            action = follower.get_next_action(
                goal_position,  # env.current_episode.goals[0].position
                goal_rotation,
                # env.current_episode.goals[0].position
            )
            geo_dist_time += time() - t_geo_dist_time
            state = env.sim.get_agent_state()
            # print("state.position: ", state.position, ", state.rotation: ",
            #       state.rotation)

            shortest_path.append(
                ShortestPathPoint(
                    state.position.tolist(),
                    quaternion_to_list(state.rotation),
                    action.value,
                )
            )

            if action == SimulatorActions.STOP:
                # print("Target #{} found.".format(len(objects_to_visit) + 1))
                break
            observations = env.step(action.value)
            env.sim._is_episode_active = True
            env._episode_over = False

            if record_video:
                im = get_image_with_obj_overlay(observations, objects)
                top_down_map = draw_top_down_map(
                    env.get_metrics(), observations["heading"], im.shape[0]
                )
                output_im = np.concatenate((im, top_down_map), axis=1)
                images.append(output_im)

    state = env.sim.get_agent_state()
    # print("state.position: ", state.position, ", state.rotation: ",
    #       state.rotation)
    if action is not None:
        shortest_path.append(
            ShortestPathPoint(
                state.position.tolist(),
                quaternion_to_list(state.rotation),
                action.value,
            )
        )

    # print("goal_position: ", goal_position, ", goal_rotation: ",
    #       goal_rotation, "goal_rot_orig: ", goal_rotation)
    # print("start_position: ", env.current_episode.start_position, ", "
    #                                                               "start_rotation: ",
    #       env.current_episode.start_rotation)
    print(f"shortest_follower_time: {geo_dist_time}")
    if record_video:
        video_filename = "{}_{}_{}".format(
            env.current_episode.episode_id,
            env.current_episode.goals[0].name,
            env.current_episode.goals[0].room_name,
        )
        video_filename = video_filename.replace(" ", "_")
        # print( "Heights difference: ", env.sim.get_agent_state().position[1] -
        #        goal_position[1])
        images_to_video(images, dirname, video_filename)
        print("Episode finished: {}".format(dirname + video_filename))

    if len(shortest_path) >= env._config.ENVIRONMENT.MAX_EPISODE_STEPS:
        shortest_path = []
        print(
            "BAD SHORTEST PATH of the EPISODE: {}".format(
                env.current_episode.episode_id
            )
        )
    return [shortest_path]


# ShortestPathPoint(position, rotation, action)
#                      for position, rotation, action in
#                      zip(action_shortest_path.points,
#                          action_shortest_path.rotations,
#                          actions)


def generate_new_start_eqa_episode(env, num_episodes=-1):
    episode = None
    episode_count = 0
    # while True: #not episode:

    target_position = env.current_episode.shortest_paths[0][-1].position
    target_rotation = habitat.utils.geometry_utils.quaternion_xyzw_to_wxyz(
        env.current_episode.shortest_paths[0][-1].rotation
    )
    if (
        env._sim._sim.pathfinder.island_radius(target_position)
        < ISLAND_RADIUS_LIMIT
    ):
        print("ISLAND_RADIUS_LIMIT")
        yield None

    if not env._sim.is_navigable(target_position):
        print("NOT NAVIGABLE END")
        yield None
    while episode_count < num_episodes or num_episodes < 0:
        # target_position =
        # env._sim._sim.pathfinder.get_random_navigable_point()
        # habitat.utils.geometry_utils.quaternion_xyzw_to_wxyz(env.current_episode.shortest_paths[0][-1].rotation)
        # quaternion_wxyz_to_xyzw
        # print("TARGET radius: ", env._sim._sim.pathfinder.island_radius(
        #    target_position), target_position[1])
        if (
            env._sim._sim.pathfinder.island_radius(target_position)
            < ISLAND_RADIUS_LIMIT
        ):
            # print("ISLAND_RADIUS_LIMIT")
            continue

        for retry in range(NUMBER_RETRIES_PER_TARGET):
            source_position = (
                env._sim._sim.pathfinder.get_random_navigable_point()
            )
            # print("source radius: ", env._sim._sim.pathfinder.island_radius(
            #    source_position), source_position[1])
            is_compatible, dist = is_compatible_episode(
                source_position,
                target_position,
                env,
                near_dist=DIFFICULTIES_PROP["easy"]["start"],
                far_dist=DIFFICULTIES_PROP["hard"]["end"],
            )
            if is_compatible:
                angle = np.random.uniform(0, 2 * np.pi)
                source_rotation = [0, np.sin(angle / 2), 0, np.cos(angle / 2)]
                env.sim.set_agent_state(source_position, source_rotation)
                env.current_episode.goals[0].position = target_position
                env.current_episode.start_position = source_position.tolist()
                env.current_episode.start_rotation = source_rotation

                # scene = env.sim.semantic_annotations()
                # object_id = env.current_episode.goals[0].object_id
                # target_position = scene.objects[object_id].obb.center

                shortest_paths = get_action_shortest_path(
                    env,
                    goal_position=target_position,
                    goal_rotation=target_rotation,
                    # env.current_episode.goals[0].position #target_position
                )
                # logger.info("Geo dist: {}".format(dist))
                # print("SUCCESS TARGET radius: ",
                #       env._sim._sim.pathfinder.island_radius(
                #           target_position), target_position[1])
                # print("SUCCESS source radius: ",
                #       env._sim._sim.pathfinder.island_radius(
                #           source_position), source_position[1])

                episode = create_episode(
                    episode=env.current_episode,
                    start_position=source_position.tolist(),
                    start_rotation=source_rotation,
                    # target_position=target_position,#.tolist(),
                    shortest_paths=shortest_paths,
                    info={"geodesic_distance": dist},
                )

                # make_greedy_path_video(episode, env,
                #                 "data/videos/mp3d/debug/" if "mp3d" in
                #                                             env._sim.config.SCENE else "data/videos/gibson/debug/"
                #                 )
                episode_count += 1
                # print("is compatible")
                yield episode
                break
    # return episode


def compute_iou(cand_mask, ref_mask=None):
    """
    Given (h, w) cand_mask, we wanna our ref_mask to be in the center of image,
    with [0.25h:0.75h, 0.25w:0.75w] occupied.
    """
    if ref_mask is None:
        h, w = cand_mask.shape[0], cand_mask.shape[1]
        ref_mask = np.zeros((h, w), np.int8)
        ref_mask[
            int(0.25 * h) : int(0.85 * h), int(0.25 * w) : int(0.75 * w)
        ] = 1

    inter = (cand_mask > 0) & (ref_mask > 0)
    union = (cand_mask > 0) | (ref_mask > 0)
    iou = inter.sum() / (union.sum())
    # print("compute_iou: {}".format(iou))
    return iou


def get_view_state_object(sim, target_centroid):
    step_size = habitat.get_config().SIMULATOR.FORWARD_STEP_SIZE
    width = 2.0
    xv, yv, zv = np.meshgrid(
        np.linspace(-width, width, num=int(2.0 * width / step_size))
        + target_centroid[0],
        np.linspace(-1, 0, num=int(1.0 * width / step_size))
        + target_centroid[1],
        np.linspace(-width, width, num=int(2.0 * width / step_size))
        + target_centroid[2],
    )
    for x, y, z in zip(xv.flatten(), yv.flatten(), zv.flatten()):
        candidate_point = [x, y, z]  # target_centroid[1] - 0.6,
        if sim.is_navigable(candidate_point):
            return candidate_point
    return None


def get_view_states(sim: habitat.Simulator, obj, target_centroid):
    step_size = habitat.get_config().SIMULATOR.FORWARD_STEP_SIZE
    turn_angle = 2 * np.pi * habitat.get_config().SIMULATOR.TURN_ANGLE / 360

    width = 2.0
    xv, yv, zv = np.meshgrid(
        np.linspace(-width, width, num=int(2.0 * width / step_size))
        + target_centroid[0],
        np.linspace(-1, 0, num=20) + target_centroid[1],
        np.linspace(-width, width, num=int(2.0 * width / step_size))
        + target_centroid[2],
    )

    thetas = np.linspace(
        0.0, 2.0 * np.pi - turn_angle, num=int(2.0 * np.pi / turn_angle)
    )
    headings = [[0, np.sin(t / 2), 0, np.cos(t / 2)] for t in thetas]

    def count_pixels(pos):
        goal_direction = np.array(target_centroid) - np.array(pos)
        goal_direction[1] = 0
        # goal_rot = mp3d_direction_to_hsim_quaternion(goal_direction)
        goal_rot = direction_to_quaternion(goal_direction)

        def _iou(heading):
            sim.set_agent_state(pos, heading)
            sim_obs = sim._sim.get_sensor_observations()
            obs = sim._sensor_suite.get_observations(sim_obs)

            result = (
                (pos, heading),
                compute_iou((obs["semantic"] == obj["instance_id"])),
            )
            # if (result[1] > 0):
            #     imageio.imsave(
            #         os.path.join(
            #             "/private/home/maksymets/hapi_eqd/data/images/",
            #             "{}_{}_target_image.png".format(obj[
            #                                                 'category_name'],
            #                                                result[1]
            #             )),
            #         obs["rgb"]
            #     )
            return result

        # object_id = int(query[1]["objects"][0]["id"].split("_")[-1])
        #     retun (py_().map(compute_iou)(objs))

        return (
            py_()
            .map(_iou)
            .max_by(1)([goal_rot])
            #                (headings[np.random.randint(0, 1)::4])
        )

    def sample_positions(positions):
        inds = np.arange(len(positions))
        max_samples = int(4e3)
        if len(positions) > max_samples:
            inds = np.random.choice(inds, size=max_samples, replace=False)

        return [positions[i] for i in inds]

    # py_().filter(lambda pos: sim.is_navigable(pos)).thru(sample_positions).map(
    #     count_pixels).filter(lambda v: v[1] > 0.075).sort_by(1,
    #                                                          reverse=True).take_while(
    #     lambda v, idx, arr: (v[1] / (arr[0][1] + 1e-10)) > 0.85).map(0).map(
    #     lambda hp: (hp[0], hp[1]))(
    #     zip(xv.flatten(), yv.flatten(), zv.flatten()))

    return (
        py_()
        # .map(lambda xy: np.array([xy[0], xy[2], target_centroid[1]]))
        .map(lambda pos: np.array(pos))
        .filter(lambda pos: sim.is_navigable(pos))
        .filter(
            lambda pos: sim._sim.pathfinder.island_radius(pos)
            > ISLAND_RADIUS_LIMIT
        )
        # .thru(sample_positions)
        .map(count_pixels)
        .filter(lambda v: v[1] > 0.075)
        .take_while(lambda v, idx, arr: (v[1] / (arr[0][1] + 1e-10)) > 0.85)
        .sort_by(1, reverse=True)
        #            .map(0)
        #            .map(lambda hp: (hp[0].tolist(), hp[1].tolist(), hp[2].tolist()))
        (zip(xv.flatten(), yv.flatten(), zv.flatten()))
    )


def generate_new_obj_nav_episode(
    env,
    # sim: Simulator,
    objects,
    num_episodes: int = -1,
    is_gen_shortest_path: bool = True,
    shortest_path_success_distance: float = 0.2,
    shortest_path_max_steps: int = 500,
    closest_dist_limit: float = 1,
    furthest_dist_limit: float = 30,
    geodesic_to_euclid_min_ratio: float = 1.1,
    number_retries_per_target: int = 10,
) -> NavigationEpisode:
    env.current_episode.episode_id = 0
    print("Number of objects: ", len(objects))
    for obj in objects:
        if isinstance(obj["id"], str):
            obj["instance_id"] = int(obj["id"].split("_")[-1])
        else:
            obj["instance_id"] = obj["id"]
        # Single navigable point around.
        # target_position = get_view_state_object(env.sim, objects[0]["center"])
        view_points = get_view_states(env.sim, obj, obj["center"])
        if not view_points:
            print("Failed find position for {}".format(obj["category_name"]))
            return
            yield

        print(
            "Number of views: ",
            len(view_points),
            " for ",
            obj["category_name"],
        )
        # target_position = view_points[0][0][0]
        # target_rotation = view_points[0][0][1]
        obj["view_points"] = view_points
        # target_position = env.current_episode.shortest_paths[0][-1].position

    shortest_paths = []

    env.current_episode.goals[0].position = objects[0]["center"]

    episode_count = 0
    gen_shortest_path = 0
    while episode_count < num_episodes or num_episodes < 0:
        for retry in range(number_retries_per_target):
            source_position = env.sim.sample_navigable_point()

            comp_results = list(
                map(
                    lambda obj: is_compatible_episode(
                        source_position,
                        obj["view_points"][0][0][0],
                        env.sim,
                        near_dist=closest_dist_limit,
                        far_dist=furthest_dist_limit,
                        geodesic_to_euclid_ratio=geodesic_to_euclid_min_ratio,
                    ),
                    objects,
                )
            )
            from functools import reduce

            all_compatible = reduce(
                lambda acum, el: acum and el[0], comp_results
            )
            # closest_obj = np.argmin(np.array(list(map(lambda x:x[1], comp_results))))

        if all_compatible:
            angle = np.random.uniform(0, 2 * np.pi)
            source_rotation = [0, np.sin(angle / 2), 0, np.cos(angle / 2)]
            env.current_episode.start_position = source_position
            env.current_episode.start_rotation = source_rotation

            env.current_episode.goals[0].name = objects[0]["category_name"]
            t_gen_shortest_path = time()
            cur_shortest_paths = get_action_reference_path(env, objects)
            print(f"t_gen_shortest_path: {time() - t_gen_shortest_path}")

            if cur_shortest_paths and len(cur_shortest_paths) > 0:
                shortest_paths.append(cur_shortest_paths[0])

            episode = create_episode(
                episode=env.current_episode,
                query=query,
                start_position=env.current_episode.start_position,
                start_rotation=env.current_episode.start_rotation,
                shortest_paths=shortest_paths,
                info={
                    # "geodesic_distance": env._sim.geodesic_distance(env.current_episode.start_position, target_position),
                },
            )

            episode_count += 1
            yield episode

habitat/datasets/object_nav/create/object_nav_generator_old.py
cshizhe/onav_rim
import csv
import itertools

import nltk
from nltk.corpus import wordnet as wn


def ycb_google_16k_meta_filename_to_word(filename: str):
    words = filename.replace(".glb", "")[4:].replace("_", " ").split()

    if len(words) == 1:
        return words[0], " ".join(words)

    if {"box", "lego", "toy", "pitcher"}.intersection(words):
        return words[-2], " ".join(words)

    return words[-1], " ".join(words)


ycb_synsets = {
    "container.n.01",
    "bread.n.01",
    "food",
    "fruit.n.01",
    "cooking_utensil.n.01",
    "tableware.n.01",
    "cutlery.n.02",
    "tool.n.01",
    "lever.n.01'",
    "lock.n.01",
    "hand_tool.n.01",
    "device.n.01",
    "equipment.n.01",
    "cube.n.05",
    "game.n.01",
    "artifact.n.01",
}


def ycb_google_16k_meta_get_synset(word: str, object_name: str):
    gen_synsets = {synset for synset in wn.synsets(word)}
    gen_synsets = gen_synsets.union(
        {hyper for synset in gen_synsets for hyper in synset.hypernyms()}
    )
    gen_synsets = gen_synsets.union(
        {hyper for synset in gen_synsets for hyper in synset.hypernyms()}
    )
    gen_synsets = {synset.name() for synset in gen_synsets}
    # two_level_synsets = {synset.name() for synset in wn.synsets(word)}.union(
    #     {hyper.name() for synset in wn.synsets(word) for hyper in synset.hypernyms()})
    key_synsets = ycb_synsets.intersection(gen_synsets)
    if key_synsets:
        return key_synsets.pop()
    return None


with open(
    "habitat/datasets/object_nav/create/ycb_google_16k_meta.csv", "r"
) as f:
    csv_r = csv.reader(f)
    headers = next(csv_r)
    print(headers)
    key_ind = headers.index("filename")
    rows = [
        (line[key_ind], ycb_google_16k_meta_filename_to_word(line[key_ind]))
        for line in csv_r
        if line[key_ind] != ""
    ]
    rows_synset = [
        (filename, word[1], word[0], wn.synsets(word[0]))
        for filename, word in rows
    ]
    synsets_source = [
        (
            filename,
            word[1],
            word[0],
            ycb_google_16k_meta_get_synset(word[0], word[1]),
            # wn.synsets(word[0])[0].definition(), wn.synsets(word[0])[0].hypernyms(),
            # wn.synsets(word[0])[0].hypernyms()[0].hypernyms()
        )
        for filename, word in rows
    ]
    # [(line[0], wn.synset(line[key_ind])) for line in csv_r if line[key_ind] != '']
    for row in synsets_source:
        print(f"{row}")

    with open(
        "habitat/datasets/object_nav/create/ycb_google_16k_meta_output.csv",
        "w",
    ) as of:
        writer = csv.writer(of)
        writer.writerow(["filename"])
        for row in synsets_source:
            writer.writerow(row)  # [row[0], row[1], row[0]])

# container.n.01
# bread.n.01
# food
# fruit.n.01
# cooking_utensil.n.01
# tableware.n.01
# cutlery.n.02
# tool.n.01
# lever.n.01'
# lock.n.01
# hand_tool.n.01
# device.n.01
# equipment.n.01
# cube.n.05
# game.n.01
# artifact.n.01

with open("category_mapping.tsv", "r") as f:
    csv_r = csv.reader(f, delimiter="\t")
    headers = next(csv_r)
    key_ind = headers.index("wnsynsetkey")
    print(headers)
    synsets_source = [
        (line[1], wn.synset(line[key_ind]))
        for line in csv_r
        if line[key_ind] != ""
    ]

with open("dictionaries.csv", "r") as f:
    csv_r = csv.reader(f)
    next(csv_r)
    next(csv_r)
    synsets_target = [
        [(line[1], sn) for sn in wn.synsets(line[1])] for line in csv_r
    ]
    synsets_target = list(itertools.chain(*synsets_target))

best_matches = []
for word1, sn1 in synsets_source:
    best = max(
        [
            [
                sn1.path_similarity(sn2, simulate_root=False),
                word1,
                word2,
                sn1,
                sn2,
            ]
            for word2, sn2 in synsets_target
            if sn1.path_similarity(sn2, simulate_root=False) is not None
        ]
    )
    best = best + [best[-2].lowest_common_hypernyms(best[-1])]
    best_matches.append(best)

print("Similarity,word1,word2,synset1,synset2,lcasynset")
for best_match in sorted(best_matches, key=lambda x: x[1]):
    # Only difference between the results posted is this if statement
    if best_match[-3] in best_match[-1] or best_match[-2] in best_match[-1]:
        print(best_match)

habitat/datasets/object_nav/create/synset_unify_py.py
cshizhe/onav_rim
#!/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 json
import os
import random
from typing import Any, Dict, List, Optional, Sequence

from habitat.config import Config
from habitat.core.dataset import ObjectInScene, SceneState
from habitat.core.registry import registry
from habitat.core.simulator import AgentState, ShortestPathPoint
from habitat.core.utils import DatasetFloatJSONEncoder
from habitat.datasets.pointnav.pointnav_dataset import (
    CONTENT_SCENES_PATH_FIELD,
    DEFAULT_SCENE_PATH_PREFIX,
    PointNavDatasetV1,
)
from habitat.tasks.nav.object_nav_task import (
    ObjectGoal,
    ObjectGoalNavEpisode,
    ObjectViewLocation,
)
from habitat.tasks.nav.object_nav_task import ReplayActionSpec, AgentStateSpec


@registry.register_dataset(name="ObjectNav-v1")
class ObjectNavDatasetV1(PointNavDatasetV1):
    r"""Class inherited from PointNavDataset that loads Object Navigation dataset."""
    category_to_task_category_id: Dict[str, int]
    category_to_scene_annotation_category_id: Dict[str, int]
    episodes: List[ObjectGoalNavEpisode] = []  # type: ignore
    content_scenes_path: str = "{data_path}/content/{scene}.json.gz"
    goals_by_category: Dict[str, Sequence[ObjectGoal]]
    count = 0

    @staticmethod
    def dedup_goals(dataset: Dict[str, Any]) -> Dict[str, Any]:
        if len(dataset["episodes"]) == 0:
            return dataset

        goals_by_category = {}
        for i, ep in enumerate(dataset["episodes"]):
            dataset["episodes"][i]["object_category"] = ep["goals"][0][
                "object_category"
            ]
            ep = ObjectGoalNavEpisode(**ep)

            goals_key = ep.goals_key
            if goals_key not in goals_by_category:
                goals_by_category[goals_key] = ep.goals

            dataset["episodes"][i]["goals"] = []

        dataset["goals_by_category"] = goals_by_category

        return dataset
    
    @staticmethod
    def dedup_goals_dset(dataset):
        if len(dataset.episodes) == 0:
            return dataset
        goals_by_category = dict()
        for i, ep in enumerate(dataset.episodes):
            dataset.episodes[i].object_category = ep.goals[0].object_category
            # ep = ObjectGoalNavEpisode(**ep)
            goals_key = ep.goals_key
            if goals_key not in goals_by_category:
                goals_by_category[goals_key] = ep.goals
            dataset.episodes[i].goals = []
        dataset.goals_by_category = goals_by_category
        return dataset

    def to_json(self) -> str:
        for i in range(len(self.episodes)):
            self.episodes[i].goals = []

        result = DatasetFloatJSONEncoder().encode(self)

        for i in range(len(self.episodes)):
            goals = self.goals_by_category[self.episodes[i].goals_key]
            if not isinstance(goals, list):
                goals = list(goals)
            self.episodes[i].goals = goals

        return result

    def __init__(self, config: Optional[Config] = None) -> None:
        self.goals_by_category = {}
        super().__init__(config)
        self.episodes = list(self.episodes)

    @staticmethod
    def __deserialize_goal(serialized_goal: Dict[str, Any]) -> ObjectGoal:
        g = ObjectGoal(**serialized_goal)

        for vidx, view in enumerate(g.view_points):
            view_location = ObjectViewLocation(**view)  # type: ignore
            view_location.agent_state = AgentState(**view_location.agent_state)  # type: ignore
            g.view_points[vidx] = view_location

        return g

    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]

        if "category_to_task_category_id" in deserialized:
            self.category_to_task_category_id = deserialized[
                "category_to_task_category_id"
            ]

        if "category_to_scene_annotation_category_id" in deserialized:
            self.category_to_scene_annotation_category_id = deserialized[
                "category_to_scene_annotation_category_id"
            ]

        if "category_to_mp3d_category_id" in deserialized:
            self.category_to_scene_annotation_category_id = deserialized[
                "category_to_mp3d_category_id"
            ]

        assert len(self.category_to_task_category_id) == len(
            self.category_to_scene_annotation_category_id
        )

        assert set(self.category_to_task_category_id.keys()) == set(
            self.category_to_scene_annotation_category_id.keys()
        ), "category_to_task and category_to_mp3d must have the same keys"

        if len(deserialized["episodes"]) == 0:
            return

        if "goals_by_category" not in deserialized:
            deserialized = self.dedup_goals(deserialized)

        for k, v in deserialized["goals_by_category"].items():
            self.goals_by_category[k] = [self.__deserialize_goal(g) for g in v]

        for i, episode in enumerate(deserialized["episodes"]):
            episode = ObjectGoalNavEpisode(**episode)
            episode.episode_id = "{}_{}".format(episode.scene_id.split("/")[-1].split(".")[0], str(i))

            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)

            episode.goals = self.goals_by_category[episode.goals_key]

            if episode.shortest_paths is not None:
                for path in episode.shortest_paths:
                    for p_index, point in enumerate(path):
                        if point is None or isinstance(point, (int, str)):
                            point = {
                                "action": point,
                                "rotation": None,
                                "position": None,
                            }

                        path[p_index] = ShortestPathPoint(**point)

            self.episodes.append(episode)  # type: ignore [attr-defined]
        # random.shuffle(self.episodes)


@registry.register_dataset(name="ObjectNav-v2")
class ObjectNavDatasetV2(PointNavDatasetV1):
    r"""Class inherited from PointNavDataset that loads Object Navigation dataset."""
    category_to_task_category_id: Dict[str, int]
    category_to_scene_annotation_category_id: Dict[str, int]
    episodes: List[ObjectGoalNavEpisode] = []  # type: ignore
    content_scenes_path: str = "{data_path}/content/{scene}.json.gz"
    goals_by_category: Dict[str, Sequence[ObjectGoal]]
    gibson_to_mp3d_category_map: Dict[str, str] = {'couch': 'sofa', 'toilet': 'toilet', 'bed': 'bed', 'tv': 'tv_monitor', 'potted plant': 'plant', 'chair': 'chair'}

    @staticmethod
    def dedup_goals(dataset: Dict[str, Any]) -> Dict[str, Any]:
        if len(dataset["episodes"]) == 0:
            return dataset

        goals_by_category = {}
        for i, ep in enumerate(dataset["episodes"]):
            dataset["episodes"][i]["object_category"] = ep["goals"][0][
                "object_category"
            ]
            ep = ObjectGoalNavEpisode(**ep)

            goals_key = ep.goals_key
            if goals_key not in goals_by_category:
                goals_by_category[goals_key] = ep.goals

            dataset["episodes"][i]["goals"] = []

        dataset["goals_by_category"] = goals_by_category

        return dataset

    def to_json(self) -> str:
        for i in range(len(self.episodes)):
            self.episodes[i].goals = []

        result = DatasetFloatJSONEncoder().encode(self)

        for i in range(len(self.episodes)):
            goals = self.goals_by_category[self.episodes[i].goals_key]
            if not isinstance(goals, list):
                goals = list(goals)
            self.episodes[i].goals = goals

        return result

    def __init__(self, config: Optional[Config] = None) -> None:
        self.goals_by_category = {}
        self.max_replay_steps = config.MAX_REPLAY_STEPS
        super().__init__(config)
        self.episodes = list(self.episodes)

    @staticmethod
    def __deserialize_goal(serialized_goal: Dict[str, Any]) -> ObjectGoal:
        g = ObjectGoal(**serialized_goal)

        for vidx, view in enumerate(g.view_points):
            view_location = ObjectViewLocation(**view)  # type: ignore
            view_location.agent_state = AgentState(**view_location.agent_state)  # type: ignore
            g.view_points[vidx] = view_location

        return g

    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]

        if "category_to_task_category_id" in deserialized:
            self.category_to_task_category_id = deserialized[
                "category_to_task_category_id"
            ]

        if "category_to_scene_annotation_category_id" in deserialized:
            self.category_to_scene_annotation_category_id = deserialized[
                "category_to_scene_annotation_category_id"
            ]

        if "category_to_mp3d_category_id" in deserialized:
            self.category_to_scene_annotation_category_id = deserialized[
                "category_to_mp3d_category_id"
            ]

        assert len(self.category_to_task_category_id) == len(
            self.category_to_scene_annotation_category_id
        )

        assert set(self.category_to_task_category_id.keys()) == set(
            self.category_to_scene_annotation_category_id.keys()
        ), "category_to_task and category_to_mp3d must have the same keys"

        if len(deserialized["episodes"]) == 0:
            return

        if "goals_by_category" not in deserialized:
            deserialized = self.dedup_goals(deserialized)

        for k, v in deserialized["goals_by_category"].items():
            self.goals_by_category[k] = [self.__deserialize_goal(g) for g in v]

        for i, episode in enumerate(deserialized["episodes"]):
            if "_shortest_path_cache" in episode:
                del episode["_shortest_path_cache"]
            
            if "gibson" in episode["scene_id"]:
                episode["scene_id"] = "gibson_semantic/{}".format(episode["scene_id"].split("/")[-1])

            episode = ObjectGoalNavEpisode(**episode)
            # episode.episode_id = str(i)

            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)

            if not episode.is_thda:
                episode.goals = self.goals_by_category[episode.goals_key]
                if episode.scene_dataset == "gibson":
                    episode.object_category = self.gibson_to_mp3d_category_map[episode.object_category]
            else:
                goals = []
                for g in episode.goals:
                    g = ObjectGoal(**g)
                    for vidx, view in enumerate(g.view_points):
                        view_location = ObjectViewLocation(**view)  # type: ignore
                        view_location.agent_state = AgentState(**view_location.agent_state)  # type: ignore
                        g.view_points[vidx] = view_location
                    goals.append(g)
                episode.goals = goals

                objects = [ObjectInScene(**o) for o in episode.scene_state["objects"]]
                scene_state = [SceneState(objects=objects).__dict__]
                episode.scene_state = scene_state

            if episode.reference_replay is not None:
                for i, replay_step in enumerate(episode.reference_replay):
                    # replay_step["agent_state"] = AgentStateSpec(**replay_step["agent_state"])
                    replay_step["agent_state"] = None
                    episode.reference_replay[i] = ReplayActionSpec(**replay_step)

            if episode.shortest_paths is not None:
                for path in episode.shortest_paths:
                    for p_index, point in enumerate(path):
                        if point is None or isinstance(point, (int, str)):
                            point = {
                                "action": point,
                                "rotation": None,
                                "position": None,
                            }

                        path[p_index] = ShortestPathPoint(**point)
            
            if len(episode.reference_replay) > self.max_replay_steps:
                continue

            self.episodes.append(episode)  # type: ignore [attr-defined]


habitat/datasets/object_nav/object_nav_dataset.py
cshizhe/onav_rim
#!/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.

from habitat.core.dataset import Dataset
from habitat.core.registry import registry


def _try_register_pickplace_dataset():
    try:
        from habitat.datasets.pickplace.pickplace_dataset import (  # noqa: F401 isort:skip
            PickPlaceDatasetV1,
        )
    except ImportError as e:
        rearrangement_dataset_import_error = e

        @registry.register_dataset(name="RearrangementDS-v1")
        class PickPlaceDatasetImportError(Dataset):
            def __init__(self, *args, **kwargs):
                raise rearrangement_dataset_import_error

habitat/datasets/pickplace/__init__.py
cshizhe/onav_rim
#!/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
import sys
from typing import List, Optional

from habitat.config import Config
from habitat.core.dataset import Dataset
from habitat.core.registry import registry
from habitat.datasets.utils import VocabFromText
from habitat.tasks.pickplace.pickplace import (
    InstructionData,
    RearrangementEpisode,
    RearrangementSpec,
    RearrangementObjectSpec,
    ReplayActionSpec,
    GrabReleaseActionSpec,
    AgentStateSpec,
    ObjectStateSpec
)

DEFAULT_SCENE_PATH_PREFIX = "data/scene_datasets/"


@registry.register_dataset(name="PickPlaceDataset-v0")
class PickPlaceDatasetV1(Dataset):
    r"""Class inherited from Dataset that loads a Vision and Language
    Navigation dataset.
    """

    episodes: List[RearrangementEpisode]
    instruction_vocab: VocabFromText

    @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)

    def __init__(self, config: Optional[Config] = None) -> None:
        self.episodes = []

        if config is None:
            return

        dataset_filename = config.DATA_PATH.format(split=config.SPLIT)
        with gzip.open(dataset_filename, "rt") as f:
            self.from_json(f.read(), scenes_dir=config.SCENES_DIR)

        self.episodes = list(
            filter(self.build_content_scenes_filter(config), self.episodes)
        )

    def from_json(
        self, json_str: str, scenes_dir: Optional[str] = None
    ) -> None:

        deserialized = json.loads(json_str)
        max_semantic_object_id = 1 << 16
        self.instruction_vocab = VocabFromText(
           sentences=deserialized["instruction_vocab"]["sentences"]
        )

        for i, episode in enumerate(deserialized["episodes"]):
            episode["reference_replay"] = []
            episode = RearrangementEpisode(**episode)

            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)

            episode.instruction = InstructionData(**episode.instruction)
            for i, goal in enumerate(episode.goals):
                if goal["info"]["is_receptacle"]:
                    episode.goals[1] = RearrangementSpec(**goal)
                else:
                    episode.goals[0] = RearrangementSpec(**goal)

            for i, obj in enumerate(episode.objects):
                episode.objects[i]["semantic_object_id"] = (episode.objects[i]["object_id"] + (1<<16))
                episode.objects[i] = RearrangementObjectSpec(**obj)
            if len(episode.reference_replay) > 1500:
                continue
            self.episodes.append(episode)


@registry.register_dataset(name="PickPlaceDataset-v1")
class PickPlaceDatasetV2(Dataset):
    r"""Class inherited from Dataset that loads a Vision and Language
    Navigation dataset.
    """

    episodes: List[RearrangementEpisode]
    instruction_vocab: VocabFromText

    @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)

    def __init__(self, config: Optional[Config] = None) -> None:
        self.episodes = []

        if config is None:
            return

        dataset_filename = config.DATA_PATH.format(split=config.SPLIT)
        with gzip.open(dataset_filename, "rt") as f:
            self.from_json(f.read(), scenes_dir=config.SCENES_DIR)

        self.episodes = list(
            filter(self.build_content_scenes_filter(config), self.episodes)
        )

    def from_json(
        self, json_str: str, scenes_dir: Optional[str] = None
    ) -> None:

        deserialized = json.loads(json_str)
        max_semantic_object_id = 1 << 16
        self.instruction_vocab = VocabFromText(
           sentences=deserialized["instruction_vocab"]["sentences"]
        )

        for i, episode in enumerate(deserialized["episodes"]):
            episode = RearrangementEpisode(**episode)
            episode.episode_id = str(i)

            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)

            episode.instruction = InstructionData(**episode.instruction)
            for i, replay_step in enumerate(episode.reference_replay):
                if "action_data" in replay_step.keys():
                    replay_step["action_data"] = GrabReleaseActionSpec(**replay_step["action_data"])
                if "agent_state" in replay_step.keys():
                    replay_step["agent_state"] = AgentStateSpec(**replay_step["agent_state"])
                if "object_states" in replay_step.keys():
                    for j in range(len(replay_step["object_states"])):
                        replay_step["object_states"][j] = ObjectStateSpec(**replay_step["object_states"][j])
                episode.reference_replay[i] = ReplayActionSpec(**replay_step)
            for i, goal in enumerate(episode.goals):
                if goal["info"]["is_receptacle"]:
                    episode.goals[1] = RearrangementSpec(**goal)
                else:
                    episode.goals[0] = RearrangementSpec(**goal)

            for i, obj in enumerate(episode.objects):
                episode.objects[i]["semantic_object_id"] = (episode.objects[i]["object_id"] + (1<<16))
                episode.objects[i] = RearrangementObjectSpec(**obj)
            
            self.episodes.append(episode)

habitat/datasets/pickplace/pickplace_dataset.py
cshizhe/onav_rim
import os
from typing import Any, List, Optional

import numpy as np
from tqdm import tqdm

import habitat
from habitat.config.default import get_config
from habitat.core.dataset import Dataset
from habitat.datasets.pointnav.generator import generate_pointnav_episode
from habitat.tasks.nav.nav import NavigationEpisode, NavigationGoal


class PointNavDatasetCreator(Dataset):
    r"""Point Navigation dataset.
    """
    episodes: List[NavigationEpisode]
    content_scenes_path: str = "{data_path}/content/{scene}.json.gz"

    @staticmethod
    def _create_episode(
        episode_id,
        scene_id,
        start_position,
        start_rotation,
        target_position,
        shortest_path,
        info=None,
    ) -> Optional[NavigationEpisode]:
        episode_id = str(episode_id)
        goals = [NavigationGoal(position=target_position)]

        episode = NavigationEpisode(
            episode_id=episode_id,
            goals=goals,
            scene_id=scene_id,
            start_position=start_position,
            start_rotation=start_rotation,
            shortest_paths=[shortest_path],
            info=info,
        )
        return episode

    @staticmethod
    def get_scenes(split, scene_path):
        raise NotImplementedError

    @staticmethod
    def check_config_paths_exist(config: Any) -> bool:
        return os.path.exists(
            config.data_path.format(split=config.split)
        ) and os.path.exists(config.scenes_path)

    def __init__(self, config: Any = None, scenes=None) -> None:
        self.episodes = []

        if config is None:
            return

        if not scenes:
            scenes = self.__class__.get_scenes(config.split, config.scene_path)
            # val_mini_scenes = self.__class__.get_scenes("val_mini",
            #                                            config.scene_path)
            # scenes = [scene for scene in scenes if scene not in val_mini_scenes]

        if config.num_houses > 0:
            scenes = scenes[: config.num_houses]
        progress_bar = tqdm(total=config.num_episodes)
        np.random.seed(config.seed)
        difficulty_counts = {}

        env_cfg = get_config()
        env_cfg.defrost()
        env_cfg.gpu_device_id = config.gpu_id
        env_cfg.task_name = "Nav-v0"
        env_cfg.freeze()

        for scene_path in scenes:
            env_cfg.defrost()
            env_cfg.SIMULATOR.SCENE = scene_path
            env_cfg.freeze()
            env = habitat.Env(config=env_cfg)
            env.seed(config.seed)

            num_episodes = round(config.num_episodes / len(scenes))
            # if
            # scene_path!=scenes[-1] else config.num_episodes - len(self.episodes)

            for episode in generate_pointnav_episode(env, num_episodes):
                episode.episode_id = len(self.episodes)
                self.episodes.append(episode)
                # dist_ratio = dataset.episodes[-1].info["distance_ratio"]
                # if dist_ratio < GEODESTIC_TO_EUCLID_RATIO:
                #     simple_episodes += 1

                progress_bar.update(1)
                # break
            env.close()
            del env

        print("scene_count: ", len(scenes))
        print("episodes_count: ", len(self.episodes))

    def to_json(self) -> str:
        result = habitat.core.utils.DatasetFloatJSONEncoder().encode(self)
        return result

habitat/datasets/pointnav/create/pointnav_dataset_creator.py
cshizhe/onav_rim
import gzip
import json
from collections import OrderedDict

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns

import habitat

config_dir_path = "./habitat-api/configs"


def load_dataset(
    config_path="configs/datasets/pointnav/gibson.yaml",
    splits=["val", "train"],
):  # "train",
    config = habitat.get_config(config_path)

    episodes = []

    for split in splits:
        config.defrost()
        config.DATASET.SPLIT = split
        dataset = habitat.make_dataset(
            id_dataset=config.DATASET.TYPE, config=config.DATASET
        )
        print(config.DATASET.SPLIT + ": Len episodes: ", len(dataset.episodes))
        episodes += (
            dataset.episodes
        )  # json.loads(dataset.to_json())["episodes"]

    print("TOTAL Len episodes: ", len(episodes))
    d_episodes = list(
        map(
            lambda x: {
                "episode_id": int(x.episode_id),
                "scene_id": x.scene_id.split("/")[-1],
                "geodesic_distance": x.info["geodesic_distance"],
                "start_position": x.start_position,
                "goal_position": x.goals[0].position,
                # "euclid_distance": x["info"]["euclid_distance"],
                # "distance_ratio": x["info"]["distance_ratio"],
                "euclid_distance": x.info["euclidean_distance"]
                if "euclidean_distance" in x.info
                else min(
                    [
                        np.power(
                            np.power(
                                np.array(
                                    np.array(x.start_position) - goal.position
                                ),
                                2,
                            ).sum(0),
                            0.5,
                        )
                        for goal in x.goals
                    ]
                ),
            },
            episodes,
        )
    )
    # print(d_episodes)
    ep_df = pd.DataFrame(
        d_episodes,
        columns=[
            "episode_id",
            "scene_id",
            "geodesic_distance",
            "difficulty",
            "start_position",
            "goal_position",
            "euclid_distance",
        ],
    )
    ep_df["ratio"] = ep_df["geodesic_distance"] / ep_df["euclid_distance"]
    return ep_df


def generate_dataset_stats():
    dataset_configs = OrderedDict(
        {
            "Gibson v1": "configs/datasets/pointnav/gibson.yaml",
            "Gibson v2": "configs/datasets/pointnav/gibson_v2.yaml",
            "MP3D objectnav v1": "configs/datasets/objectnav/mp3d.yaml",
            # "MP3D": "configs/pointnav/mp3d.yaml",
            # "SUNCG": "configs/pointnav/suncg.yaml",
        }
    )
    # datasets = {name: load_dataset(config) for name, config in
    #             dataset_configs.items()}
    datasets = OrderedDict()
    for split in [
        "train",
        "val_mini",
        "val",
    ]:  # , "test_challenge" "test_standard", ,"train" "test", "val_mini", "val", "test_challenge", "test_standard" "train" ,
        datasets.update(
            {
                name + " " + split: load_dataset(config, [split])
                for name, config in dataset_configs.items()
            }
        )
    # datasets["Gibson test"] = load_dataset(dataset_configs["Gibson"]) # ["7LK038WEKH"]

    sns.set(style="white", font_scale=3)
    # f, ax = plt.subplots(figsize=(20, 20))

    # import matplotlib.gridspec as gridspec
    #
    # fig = plt.figure(figsize=(11 * len(datasets.keys()), 36))
    # outer = gridspec.GridSpec(1, 3, wspace=0.2, hspace=0.2)
    #
    # for i in range(3):
    #     inner = gridspec.GridSpecFromSubplotSpec(2, 1,
    #                                              subplot_spec=outer[i],
    #                                              wspace=0.1, hspace=0.1)
    #
    #     for j in range(2):
    #         ax = plt.Subplot(fig, inner[j])
    #         t = ax.text(0.5, 0.5, 'outer=%d, inner=%d' % (i, j))
    #         t.set_ha('center')
    #         ax.set_xticks([])
    #         ax.set_yticks([])
    #         fig.add_subplot(ax)
    n_plots = len(datasets.keys())
    n_datasets = len(dataset_configs.keys())
    n_splits_per_ds = len(datasets.keys()) // n_datasets

    fig, axs = plt.subplots(
        nrows=max(n_datasets, 2),
        ncols=3 * n_splits_per_ds,
        figsize=(20 * len(datasets.keys()), 30),
        sharex="col",  # sharey='row'
        #                            sharey=True) # , sharex=True
    )

    #     fig, axs = plt.subplots(nrows=n_datasets, ncols=3 * len(datasets.keys()) // n_datasets,
    #                             figsize=(11 * len(datasets.keys()), 20),
    #                             sharex='col', #sharey='row'
    # #                            sharey=True) # , sharex=True
    #                             )
    # ylabel='Number of episodes'
    # axs.set(ylabel='Number of episodes')
    # for ax in axs:

    for dataset_id, dataset_pair in enumerate(datasets.items()):
        column = (dataset_id // n_datasets) * 3
        axs[dataset_id % n_datasets][column].set_ylabel(
            "{} split\nNumber of episodes".format(dataset_pair[0])
        )
        sns.distplot(
            dataset_pair[1]["euclid_distance"],
            # axlabel="Euclidean Distance",
            axlabel="",
            ax=axs[dataset_id % n_datasets][column],
            kde=False,
            norm_hist=False,
        )
        sns.distplot(
            dataset_pair[1]["geodesic_distance"],
            # axlabel="Geodesic Distance",
            axlabel="",
            ax=axs[dataset_id % n_datasets][column + 1],
            kde=False,
            norm_hist=False,
        )
        sns.distplot(
            dataset_pair[1][dataset_pair[1]["ratio"] < 2]["ratio"],
            #    bins=30,
            #    hist=True,
            # axlabel="Geodesic-Euclidean Distance Ratio",
            #    label="WHAT",
            #    norm_hist=True,
            axlabel="",
            ax=axs[dataset_id % n_datasets][column + 2],
            kde=False,
            norm_hist=False,
        )

        axs[n_datasets - 1][column].set_xlabel("Euclidean " "Distance")
        axs[n_datasets - 1][column + 1].set_xlabel("Geodesic " "Distance")
        axs[n_datasets - 1][column + 2].set_xlabel(
            "Geodesic-Euclidean Distance " "Ratio"
        )
    fig.tight_layout()
    #    plt.show()

    plt.savefig("datasets/pointnav/analysis/output" "/splits_stats.png")


generate_dataset_stats()

# for dataset_id, dataset_pair in enumerate(datasets.items()):
#     axs[dataset_id][0].set_ylabel('{} split\nNumber of '
#                                   'episodes'.format(dataset_pair[0]))
#
#     sns.distplot(
#         dataset_pair[1]["euclid_distance"],
#         # axlabel="Euclidean Distance",
#         axlabel="",
#         ax=axs[dataset_id][0],
#         kde=False,
#         norm_hist=False,
#
#     )
#     sns.distplot(
#         dataset_pair[1]["geodesic_distance"],
#         # axlabel="Geodesic Distance",
#         axlabel="",
#         ax=axs[dataset_id][1],
#         kde=False,
#         norm_hist=False,
#     )
#     sns.distplot(
#         dataset_pair[1][dataset_pair[1]["ratio"] < 2]["ratio"],
#         #    bins=30,
#         #    hist=True,
#         # axlabel="Geodesic-Euclidean Distance Ratio",
#         #    label="WHAT",
#         #    norm_hist=True,
#         axlabel="",
#         ax=axs[dataset_id][2],
#         kde=False,
#         norm_hist=False,
#     )
#
# axs[n_plots - 1][0].set_xlabel("Euclidean Distance")
# axs[n_plots - 1][1].set_xlabel("Geodesic Distance")
# axs[n_plots - 1][2].set_xlabel("Geodesic-Euclidean Distance Ratio")

habitat/datasets/pointnav/create/stats.py
cshizhe/onav_rim
import csv
import gzip
import os
import random
from typing import List

import habitat

MP3D_SCENES_SPLITS_PATH = "habitat/datasets/scenes/scenes_mp3d.csv"
HABITAT_MP3D_SCENES_SPLITS_PATH = (
    "habitat/datasets/scenes/scenes_habitat_mp3d.csv"
)
GIBSON_SCENES_SPLITS_PATH = (
    "habitat/datasets/scenes/scenes_gibson_fullplus.csv"
)

GIBSON_HABITAT_SCENES_SPLITS_PATH = (
    "habitat/datasets/scenes/scenes_gibson_habtiat.csv"
)


def get_mp3d_scenes(
    split: str = "train", scene_template: str = "{scene}"
) -> List[str]:
    scenes = []
    with open(MP3D_SCENES_SPLITS_PATH, newline="") as csvfile:
        spamreader = csv.DictReader(csvfile, delimiter=",", quotechar="|")
        for row in spamreader:
            if split in row["set"].split() or split == "*":
                scenes.append(scene_template.format(scene=row["id"]))
    return scenes


def get_gibson_scenes(
    split: str = "train", scene_template: str = "{scene}"
) -> List[str]:
    scenes = []
    with open(GIBSON_SCENES_SPLITS_PATH, newline="") as csvfile:
        spamreader = csv.DictReader(csvfile, delimiter=",", quotechar="|")
        for row in spamreader:
            if row[split] == "1" or split == "*":
                scenes.append(scene_template.format(scene=row["id"]))
    return scenes


def get_habitat_mp3d_scenes(
    split: str = "train", scene_template: str = "{scene}"
) -> List[str]:
    scenes = []
    with open(HABITAT_MP3D_SCENES_SPLITS_PATH, newline="") as csvfile:
        spamreader = csv.DictReader(csvfile, delimiter="\t", quotechar="|")
        for row in spamreader:
            if split == "*" or split in row["set"].split():
                scenes.append(
                    scene_template.format(
                        scene=f'{row["folder"]}/{row["glb"]}'
                    )
                )
    return scenes


def get_habitat_gibson_scenes(
    split: str = "train", scene_template: str = "{scene}"
) -> List[str]:
    scenes = []
    with open(GIBSON_HABITAT_SCENES_SPLITS_PATH, newline="") as csvfile:
        spamreader = csv.DictReader(csvfile, delimiter=",", quotechar="|")
        for row in spamreader:
            if row[split] == "1" or split == "*":
                scenes.append(scene_template.format(scene=row["id"]))
    return scenes


def get_avg_geo_dist(dataset):
    sum_geo_dist = 0
    for episode in dataset.episodes:
        sum_geo_dist += episode.info["geodesic_distance"]

    return sum_geo_dist / len(dataset.episodes)


def generate_sampled_train(
    config_path="datasets/pointnav/gibson.yaml", num_episodes=1000
):
    config = habitat.get_config(config_path, config_dir="habitat-api/configs")
    config.defrost()
    config.DATASET.SPLIT = "train"
    config.freeze()
    print("Dataset is loading.")
    dataset = habitat.make_dataset(
        id_dataset=config.DATASET.TYPE, config=config.DATASET
    )
    print(config.DATASET.SPLIT + ": Len episodes: ", len(dataset.episodes))
    dataset.episodes = random.sample(dataset.episodes, num_episodes)
    print("Average geo distance: ", get_avg_geo_dist(dataset))

    json_str = str(dataset.to_json())

    output_dir = "data/datasets/pointnav/gibson/v1/{}_small_2/".format(
        config.DATASET.SPLIT
    )
    if not os.path.exists(output_dir):
        os.mkdir(output_dir)
    main_dataset_file = "{}/{}_small_2.json.gz".format(
        output_dir, config.DATASET.SPLIT
    )
    with gzip.GzipFile(main_dataset_file, "wb") as f:
        f.write(json_str.encode("utf-8"))

    print("Dataset file: {}".format(main_dataset_file))


def generate_mini_train_splits():
    generate_sampled_train("datasets/pointnav/gibson.yaml")
    # generate_sampled_train("datasets/pointnav/mp3d.yaml")

habitat/datasets/pointnav/create/utils.py
cshizhe/onav_rim
#!/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.

from habitat.core.logging import logger
from habitat.core.registry import registry
from habitat.datasets.eqa import _try_register_mp3d_eqa_dataset
from habitat.datasets.object_nav import _try_register_objectnavdatasetv1
from habitat.datasets.pointnav import _try_register_pointnavdatasetv1
from habitat.datasets.vln import _try_register_r2r_vln_dataset
from habitat.datasets.pickplace import _try_register_pickplace_dataset


def make_dataset(id_dataset, **kwargs):
    logger.info("Initializing dataset {}".format(id_dataset))
    _dataset = registry.get_dataset(id_dataset)
    assert _dataset is not None, "Could not find dataset {}".format(id_dataset)

    return _dataset(**kwargs)  # type: ignore


_try_register_objectnavdatasetv1()
_try_register_mp3d_eqa_dataset()
_try_register_pointnavdatasetv1()
_try_register_r2r_vln_dataset()
_try_register_pickplace_dataset()

habitat/datasets/registration.py
cshizhe/onav_rim
from habitat.core.registry import registry
from habitat.core.simulator import Simulator

# from habitat.sims.habitat_simulator.actions import (
#     HabitatSimV1ActionSpaceConfiguration,
# )


def _try_register_habitat_sim():
    try:
        import habitat_sim  # noqa: F401

        has_habitat_sim = True
    except ImportError as e:
        has_habitat_sim = False
        habitat_sim_import_error = e

    if has_habitat_sim:
        from habitat.sims.pickplace.simulator import PickPlaceSim
        from habitat.sims.habitat_simulator.actions import (  # noqa: F401
            HabitatSimV1ActionSpaceConfiguration,
        )
    else:

        @registry.register_simulator(name="Sim-v0")
        class HabitatSimImportError(Simulator):
            def __init__(self, *args, **kwargs):
                raise habitat_sim_import_error

habitat/sims/habitat_simulator/__init__.py
cshizhe/onav_rim
#!/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.

from enum import Enum
from typing import Dict, List, Any

import attr

import habitat_sim
from habitat.core.registry import registry
from habitat.core.simulator import ActionSpaceConfiguration
from habitat.core.utils import Singleton

from habitat.core.embodied_task import SimulatorTaskAction
from habitat_sim.agent.controls.controls import ActuationSpec


class _DefaultHabitatSimActions(Enum):
    STOP = 0
    MOVE_FORWARD = 1
    TURN_LEFT = 2
    TURN_RIGHT = 3
    LOOK_UP = 4
    LOOK_DOWN = 5


@attr.s(auto_attribs=True, slots=True)
class HabitatSimActionsSingleton(metaclass=Singleton):
    r"""Implements an extendable Enum for the mapping of action names
    to their integer values.

    This means that new action names can be added, but old action names cannot
    be removed nor can their mapping be altered. This also ensures that all
    actions are always contigously mapped in :py:`[0, len(HabitatSimActions) - 1]`

    This accesible as the global singleton :ref:`HabitatSimActions`
    """

    _known_actions: Dict[str, int] = attr.ib(init=False, factory=dict)

    def __attrs_post_init__(self):
        for action in _DefaultHabitatSimActions:
            self._known_actions[action.name] = action.value

    def extend_action_space(self, name: str) -> int:
        r"""Extends the action space to accomodate a new action with
        the name :p:`name`

        :param name: The name of the new action
        :return: The number the action is registered on

        Usage:

        .. code:: py

            from habitat.sims.habitat_simulator.actions import HabitatSimActions
            HabitatSimActions.extend_action_space("MY_ACTION")
            print(HabitatSimActions.MY_ACTION)
        """
        assert (
            name not in self._known_actions
        ), "Cannot register an action name twice"
        self._known_actions[name] = len(self._known_actions)

        return self._known_actions[name]

    def has_action(self, name: str) -> bool:
        r"""Checks to see if action :p:`name` is already register

        :param name: The name to check
        :return: Whether or not :p:`name` already exists
        """

        return name in self._known_actions

    def __getattr__(self, name):
        return self._known_actions[name]

    def __getitem__(self, name):
        return self._known_actions[name]

    def __len__(self):
        return len(self._known_actions)

    def __iter__(self):
        return iter(self._known_actions)


HabitatSimActions: HabitatSimActionsSingleton = HabitatSimActionsSingleton()


@registry.register_action_space_configuration(name="v0")
class HabitatSimV0ActionSpaceConfiguration(ActionSpaceConfiguration):
    def get(self):
        return {
            HabitatSimActions.STOP: habitat_sim.ActionSpec("stop"),
            HabitatSimActions.MOVE_FORWARD: habitat_sim.ActionSpec(
                "move_forward",
                habitat_sim.ActuationSpec(
                    amount=self.config.FORWARD_STEP_SIZE
                ),
            ),
            HabitatSimActions.TURN_LEFT: habitat_sim.ActionSpec(
                "turn_left",
                habitat_sim.ActuationSpec(amount=self.config.TURN_ANGLE),
            ),
            HabitatSimActions.TURN_RIGHT: habitat_sim.ActionSpec(
                "turn_right",
                habitat_sim.ActuationSpec(amount=self.config.TURN_ANGLE),
            ),
        }


@registry.register_action_space_configuration(name="v1")
class HabitatSimV1ActionSpaceConfiguration(
    HabitatSimV0ActionSpaceConfiguration
):
    def get(self):
        config = super().get()
        new_config = {
            HabitatSimActions.LOOK_UP: habitat_sim.ActionSpec(
                "look_up",
                habitat_sim.ActuationSpec(amount=self.config.TILT_ANGLE),
            ),
            HabitatSimActions.LOOK_DOWN: habitat_sim.ActionSpec(
                "look_down",
                habitat_sim.ActuationSpec(amount=self.config.TILT_ANGLE),
            ),
        }

        config.update(new_config)

        return config


@registry.register_action_space_configuration(name="pyrobotnoisy")
class HabitatSimPyRobotActionSpaceConfiguration(ActionSpaceConfiguration):
    def get(self):
        return {
            HabitatSimActions.STOP: habitat_sim.ActionSpec("stop"),
            HabitatSimActions.MOVE_FORWARD: habitat_sim.ActionSpec(
                "pyrobot_noisy_move_forward",
                habitat_sim.PyRobotNoisyActuationSpec(
                    amount=self.config.FORWARD_STEP_SIZE,
                    robot=self.config.NOISE_MODEL.ROBOT,
                    controller=self.config.NOISE_MODEL.CONTROLLER,
                    noise_multiplier=self.config.NOISE_MODEL.NOISE_MULTIPLIER,
                ),
            ),
            HabitatSimActions.TURN_LEFT: habitat_sim.ActionSpec(
                "pyrobot_noisy_turn_left",
                habitat_sim.PyRobotNoisyActuationSpec(
                    amount=self.config.TURN_ANGLE,
                    robot=self.config.NOISE_MODEL.ROBOT,
                    controller=self.config.NOISE_MODEL.CONTROLLER,
                    noise_multiplier=self.config.NOISE_MODEL.NOISE_MULTIPLIER,
                ),
            ),
            HabitatSimActions.TURN_RIGHT: habitat_sim.ActionSpec(
                "pyrobot_noisy_turn_right",
                habitat_sim.PyRobotNoisyActuationSpec(
                    amount=self.config.TURN_ANGLE,
                    robot=self.config.NOISE_MODEL.ROBOT,
                    controller=self.config.NOISE_MODEL.CONTROLLER,
                    noise_multiplier=self.config.NOISE_MODEL.NOISE_MULTIPLIER,
                ),
            ),
            HabitatSimActions.LOOK_UP: habitat_sim.ActionSpec(
                "look_up",
                habitat_sim.ActuationSpec(amount=self.config.TILT_ANGLE),
            ),
            HabitatSimActions.LOOK_DOWN: habitat_sim.ActionSpec(
                "look_down",
                habitat_sim.ActuationSpec(amount=self.config.TILT_ANGLE),
            ),
            # The perfect actions are needed for the oracle planner
            "_forward": habitat_sim.ActionSpec(
                "move_forward",
                habitat_sim.ActuationSpec(
                    amount=self.config.FORWARD_STEP_SIZE
                ),
            ),
            "_left": habitat_sim.ActionSpec(
                "turn_left",
                habitat_sim.ActuationSpec(amount=self.config.TURN_ANGLE),
            ),
            "_right": habitat_sim.ActionSpec(
                "turn_right",
                habitat_sim.ActuationSpec(amount=self.config.TURN_ANGLE),
            ),
        }

habitat/sims/habitat_simulator/actions.py
cshizhe/onav_rim
#!/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.

from logging import Logger
from typing import (
    TYPE_CHECKING,
    Any,
    Dict,
    List,
    Optional,
    Sequence,
    Union,
    cast,
)

import numpy as np
import math
import magnum as mn
import random
from gym import spaces
from gym.spaces.box import Box
from numpy import ndarray

if TYPE_CHECKING:
    from torch import Tensor

import habitat_sim
from habitat.core.dataset import Episode
from habitat.core.registry import registry
from habitat.core.simulator import (
    AgentState,
    Config,
    DepthSensor,
    Observations,
    RGBSensor,
    SemanticSensor,
    Sensor,
    SensorSuite,
    ShortestPathPoint,
    Simulator,
    VisualObservation,
)
from habitat.core.spaces import Space
from habitat_sim.nav import NavMeshSettings
from habitat_sim.utils.common import quat_from_coeffs, quat_to_magnum
from habitat_sim.physics import MotionType

RGBSENSOR_DIMENSION = 3


def overwrite_config(config_from: Config, config_to: Any) -> None:
    r"""Takes Habitat Lab config and Habitat-Sim config structures. Overwrites
    Habitat-Sim config with Habitat Lab values, where a field name is present
    in lowercase. Mostly used to avoid :ref:`sim_cfg.field = hapi_cfg.FIELD`
    code.

    Args:
        config_from: Habitat Lab config node.
        config_to: Habitat-Sim config structure.
    """

    def if_config_to_lower(config):
        if isinstance(config, Config):
            return {key.lower(): val for key, val in config.items()}
        else:
            return config

    for attr, value in config_from.items():
        if hasattr(config_to, attr.lower()):
            setattr(config_to, attr.lower(), if_config_to_lower(value))


@registry.register_sensor
class HabitatSimRGBSensor(RGBSensor):
    sim_sensor_type: habitat_sim.SensorType

    def __init__(self, config: Config) -> None:
        self.sim_sensor_type = habitat_sim.SensorType.COLOR
        super().__init__(config=config)

    def _get_observation_space(self, *args: Any, **kwargs: Any) -> Box:
        return spaces.Box(
            low=0,
            high=255,
            shape=(self.config.HEIGHT, self.config.WIDTH, RGBSENSOR_DIMENSION),
            dtype=np.uint8,
        )

    def get_observation(
        self, sim_obs: Dict[str, Union[ndarray, bool, "Tensor"]]
    ) -> VisualObservation:
        obs = cast(Optional[VisualObservation], sim_obs.get(self.uuid, None))
        check_sim_obs(obs, self)

        # remove alpha channel
        obs = obs[:, :, :RGBSENSOR_DIMENSION]  # type: ignore[index]
        return obs


@registry.register_sensor
class HabitatSimDepthSensor(DepthSensor):
    sim_sensor_type: habitat_sim.SensorType
    min_depth_value: float
    max_depth_value: float

    def __init__(self, config: Config) -> None:
        self.sim_sensor_type = habitat_sim.SensorType.DEPTH

        if config.NORMALIZE_DEPTH:
            self.min_depth_value = 0
            self.max_depth_value = 1
        else:
            self.min_depth_value = config.MIN_DEPTH
            self.max_depth_value = config.MAX_DEPTH

        super().__init__(config=config)

    def _get_observation_space(self, *args: Any, **kwargs: Any) -> Box:
        return spaces.Box(
            low=self.min_depth_value,
            high=self.max_depth_value,
            shape=(self.config.HEIGHT, self.config.WIDTH, 1),
            dtype=np.float32,
        )

    def get_observation(
        self, sim_obs: Dict[str, Union[ndarray, bool, "Tensor"]]
    ) -> VisualObservation:
        obs = cast(Optional[VisualObservation], sim_obs.get(self.uuid, None))
        check_sim_obs(obs, self)
        if isinstance(obs, np.ndarray):
            obs = np.clip(obs, self.config.MIN_DEPTH, self.config.MAX_DEPTH)

            obs = np.expand_dims(
                obs, axis=2
            )  # make depth observation a 3D array
        else:
            obs = obs.clamp(self.config.MIN_DEPTH, self.config.MAX_DEPTH)  # type: ignore[attr-defined]

            obs = obs.unsqueeze(-1)  # type: ignore[attr-defined]

        if self.config.NORMALIZE_DEPTH:
            # normalize depth observation to [0, 1]
            obs = (obs - self.config.MIN_DEPTH) / (
                self.config.MAX_DEPTH - self.config.MIN_DEPTH
            )

        return obs


@registry.register_sensor
class HabitatSimSemanticSensor(SemanticSensor):
    sim_sensor_type: habitat_sim.SensorType

    def __init__(self, config):
        self.sim_sensor_type = habitat_sim.SensorType.SEMANTIC
        super().__init__(config=config)

    def _get_observation_space(self, *args: Any, **kwargs: Any):
        return spaces.Box(
            low=np.iinfo(np.uint32).min,
            high=np.iinfo(np.uint32).max,
            shape=(self.config.HEIGHT, self.config.WIDTH),
            dtype=np.uint32,
        )

    def get_observation(
        self, sim_obs: Dict[str, Union[ndarray, bool, "Tensor"]]
    ) -> VisualObservation:
        obs = cast(Optional[VisualObservation], sim_obs.get(self.uuid, None))
        check_sim_obs(obs, self)
        return obs


def check_sim_obs(obs: ndarray, sensor: Sensor) -> None:
    assert obs is not None, (
        "Observation corresponding to {} not present in "
        "simulator's observations".format(sensor.uuid)
    )


HabitatSimVizSensors = Union[
    HabitatSimRGBSensor, HabitatSimDepthSensor, HabitatSimSemanticSensor
]


@registry.register_simulator(name="Sim-v0")
class HabitatSim(habitat_sim.Simulator, Simulator):
    r"""Simulator wrapper over habitat-sim

    habitat-sim repo: https://github.com/facebookresearch/habitat-sim

    Args:
        config: configuration for initializing the simulator.
    """

    def __init__(self, config: Config) -> None:
        self.habitat_config = config
        agent_config = self._get_agent_config()

        sim_sensors = []
        for sensor_name in agent_config.SENSORS:
            sensor_cfg = getattr(self.habitat_config, sensor_name)
            sensor_type = registry.get_sensor(sensor_cfg.TYPE)

            assert sensor_type is not None, "invalid sensor type {}".format(
                sensor_cfg.TYPE
            )
            sim_sensors.append(sensor_type(sensor_cfg))

        self._sensor_suite = SensorSuite(sim_sensors)
        self.sim_config = self.create_sim_config(self._sensor_suite)
        self._current_scene = self.sim_config.sim_cfg.scene_id
        super().__init__(self.sim_config)
        self._action_space = spaces.Discrete(
            len(self.sim_config.agents[0].action_space)
        )
        self._prev_sim_obs: Optional[Observations] = None

        self.navmesh_settings = habitat_sim.NavMeshSettings()
        self.navmesh_settings.set_defaults()
        self.navmesh_settings.agent_radius = agent_config.RADIUS
        self.navmesh_settings.agent_height = agent_config.HEIGHT

    def create_sim_config(
        self, _sensor_suite: SensorSuite
    ) -> habitat_sim.Configuration:
        sim_config = habitat_sim.SimulatorConfiguration()
        # Check if Habitat-Sim is post Scene Config Update
        if not hasattr(sim_config, "scene_id"):
            raise RuntimeError(
                "Incompatible version of Habitat-Sim detected, please upgrade habitat_sim"
            )
        overwrite_config(
            config_from=self.habitat_config.HABITAT_SIM_V0,
            config_to=sim_config,
        )
        sim_config.scene_id = self.habitat_config.SCENE
        agent_config = habitat_sim.AgentConfiguration()
        overwrite_config(
            config_from=self._get_agent_config(), config_to=agent_config
        )

        sensor_specifications = []
        for sensor in _sensor_suite.sensors.values():
            sim_sensor_cfg = habitat_sim.SensorSpec()
            overwrite_config(
                config_from=sensor.config, config_to=sim_sensor_cfg
            )
            sim_sensor_cfg.uuid = sensor.uuid
            sim_sensor_cfg.resolution = list(
                sensor.observation_space.shape[:2]
            )
            sim_sensor_cfg.parameters["hfov"] = str(sensor.config.HFOV)

            # TODO(maksymets): Add configure method to Sensor API to avoid
            # accessing child attributes through parent interface
            # We know that the Sensor has to be one of these Sensors
            sensor = cast(HabitatSimVizSensors, sensor)
            sim_sensor_cfg.sensor_type = sensor.sim_sensor_type
            sim_sensor_cfg.gpu2gpu_transfer = (
                self.habitat_config.HABITAT_SIM_V0.GPU_GPU
            )
            sensor_specifications.append(sim_sensor_cfg)

        agent_config.sensor_specifications = sensor_specifications
        agent_config.action_space = registry.get_action_space_configuration(
            self.habitat_config.ACTION_SPACE_CONFIG
        )(self.habitat_config).get()

        return habitat_sim.Configuration(sim_config, [agent_config])

    @property
    def sensor_suite(self) -> SensorSuite:
        return self._sensor_suite

    @property
    def action_space(self) -> Space:
        return self._action_space

    def _update_agents_state(self) -> bool:
        is_updated = False
        for agent_id, _ in enumerate(self.habitat_config.AGENTS):
            agent_cfg = self._get_agent_config(agent_id)
            if agent_cfg.IS_SET_START_STATE:
                self.set_agent_state(
                    agent_cfg.START_POSITION,
                    agent_cfg.START_ROTATION,
                    agent_id,
                )
                is_updated = True

        return is_updated

    def reset(self) -> Observations:
        sim_obs = super().reset()
        if self._update_agents_state():
            sim_obs = self.get_sensor_observations()

        self._prev_sim_obs = sim_obs
        return self._sensor_suite.get_observations(sim_obs)

    def step(self, action: Union[str, int]) -> Observations:
        sim_obs = super().step(action)
        self._prev_sim_obs = sim_obs
        observations = self._sensor_suite.get_observations(sim_obs)
        return observations

    def render(self, mode: str = "rgb") -> Any:
        r"""
        Args:
            mode: sensor whose observation is used for returning the frame,
                eg: "rgb", "depth", "semantic"

        Returns:
            rendered frame according to the mode
        """
        sim_obs = self.get_sensor_observations()
        observations = self._sensor_suite.get_observations(sim_obs)

        output = observations.get(mode)
        assert output is not None, "mode {} sensor is not active".format(mode)
        if not isinstance(output, np.ndarray):
            # If it is not a numpy array, it is a torch tensor
            # The function expects the result to be a numpy array
            output = output.to("cpu").numpy()

        return output

    def reconfigure(self, habitat_config: Config) -> None:
        # TODO(maksymets): Switch to Habitat-Sim more efficient caching
        is_same_scene = habitat_config.SCENE == self._current_scene
        self.habitat_config = habitat_config
        self.sim_config = self.create_sim_config(self._sensor_suite)
        if not is_same_scene:
            self._current_scene = habitat_config.SCENE
            self.close()
            super().reconfigure(self.sim_config)

        self._update_agents_state()

        self.sim_object_to_objid_mapping = {}
        self.objid_to_sim_object_mapping = {}
        self.obj_sem_id_to_sem_category_mapping = {}

        self._initialize_THDA_scene()

    def remove_all_objects(self):
        for id in self.get_existing_object_ids():
            self.remove_object(id)

    def get_rotation_vec(self, object_id):
        quat = self.get_rotation(object_id)
        return np.array(quat.vector).tolist() + [quat.scalar]

    # THDA Initialization
    def _initialize_THDA_scene(self):
        self.remove_all_objects()
        self.recompute_navmesh(self.pathfinder, self.navmesh_settings, True)
        # Retrieving scene state from config
        if (
            not hasattr(self.habitat_config, "scene_state")
            or self.habitat_config.scene_state is None
        ):
            return
        objects = self.habitat_config.scene_state[0]["objects"]
        obj_templates_mgr = self.get_object_template_manager()

        # self.remove_all_objects()
        # self.recompute_navmesh(self.pathfinder, self.navmesh_settings, True)

        self.sim_object_to_objid_mapping = {}
        self.objid_to_sim_object_mapping = {}
        self.obj_sem_id_to_sem_category_mapping = {}

        if objects is not None:
            for object in objects:
                object_template = f"{object.object_template}"
                object_pos = object.position
                object_rot = object.rotation
                object_template_handles = (
                    obj_templates_mgr.get_file_template_handles(
                        object_template
                    )
                )

                if len(object_template_handles) >= 1:
                    object_template_handle = object_template_handles[0]
                    # logger.info(f"Reusing template handle.{object_template_handle} | {obj_templates_mgr.get_num_templates()}")
                else:
                    # if obj_templates_mgr.get_num_templates() < 20:
                    # logger.info(f"Creating new template handle. {object_template} | {obj_templates_mgr.get_num_templates()}")
                    _object_template_ids = (
                        obj_templates_mgr.load_object_configs(object_template)
                    )
                    assert (
                        len(_object_template_ids) > 0
                    ), f"Object template '{object_template}' wasn't found."
                    object_template_id = _object_template_ids[0]
                    object_attr = obj_templates_mgr.get_template_by_ID(
                        object_template_id
                    )
                    if object.scale:
                        object_attr.scale = mn.Vector3(object.scale)
                        # print(f"object.scale: {object.scale}")
                    else:
                        object_attr.scale = mn.Vector3(5)
                    obj_templates_mgr.register_template(object_attr)

                    # logger.error(f"SCENE: {self.habitat_config.SCENE}")
                    object_template_handle = object_attr.handle
                    # else:
                    # object_template_handle = obj_templates_mgr.get_random_template_handle()
                    # logger.info(f"Random object {object_template_handle} instead of {object_template}")

                object_id = self.add_object_by_handle(object_template_handle)

                self.sim_object_to_objid_mapping[object_id] = object.object_id
                self.objid_to_sim_object_mapping[object.object_id] = object_id

                # Return back
                self.set_object_semantic_id(int(object.object_id), object_id)
                # [self.get_object_scene_node(id).semantic_id for id in self.get_existing_object_ids()]
                self.obj_sem_id_to_sem_category_mapping[
                    int(object.object_id)
                ] = object.semantic_category_id

                if object_pos is not None:
                    self.set_translation(object_pos, object_id)
                    if isinstance(object_rot, list):
                        object_rot = quat_from_coeffs(object_rot)
                    if object_rot is not None:
                        object_rot = quat_to_magnum(object_rot)
                        self.set_rotation(object_rot, object_id)
                else:
                    self.sample_object_state(object_id)
                    self.recompute_navmesh(
                        self.pathfinder, self.navmesh_settings, True
                    )

                # self.sample_object_state(object_id)

                self.set_object_motion_type(MotionType.STATIC, object_id)
                # print(f"rotation: {self.get_rotation_vec(object_id)}")
                # print(f"position: {np.array(self.get_translation(object_id)).tolist()}")
                # print(f"{self.sample_navigable_point()}")

            # Recompute the navmesh after placing all the objects.
            self.recompute_navmesh(
                self.pathfinder, self.navmesh_settings, True
            )
            # logger.info("Inserting objects and recompute navmesh")
    
    def sample_object_state(
        sim,
        object_id,
        from_navmesh=True,
        maintain_object_up=True,
        max_tries=100,
        bb=None,
    ):
        # check that the object is not STATIC
        if (
            sim.get_object_motion_type(object_id)
            is habitat_sim.physics.MotionType.STATIC
        ):
            print("sample_object_state : Object is STATIC, aborting.")
        if from_navmesh:
            if not sim.pathfinder.is_loaded:
                print("sample_object_state : No pathfinder, aborting.")
                return False
        elif not bb:
            print(
                "sample_object_state : from_navmesh not specified and no bounding box provided, aborting."
            )
            return False
        tries = 0
        valid_placement = False
        # Note: following assumes sim was not reconfigured without close
        scene_collision_margin = 0.04
        while not valid_placement and tries < max_tries:
            tries += 1
            # initialize sample location to random point in scene bounding box
            sample_location = np.array([0, 0, 0])
            if from_navmesh:
                # query random navigable point
                sample_location = sim.pathfinder.get_random_navigable_point()
            else:
                sample_location = np.random.uniform(bb.min, bb.max)
            # set the test state
            # to debug sim.set_translation(sim.pathfinder.get_random_navigable_point(), object_id), sim.get_translation(object_id)
            sim.set_translation(sample_location, object_id)
            if maintain_object_up:
                # random rotation only on the Y axis
                y_rotation = mn.Quaternion.rotation(
                    mn.Rad(random.random() * 2 * math.pi),
                    mn.Vector3(0, 1.0, 0),
                )
                sim.set_rotation(
                    y_rotation * sim.get_rotation(object_id), object_id
                )
            else:
                # unconstrained random rotation
                sim.set_rotation(ut.random_quaternion(), object_id)

            # raise object such that lowest bounding box corner is above the navmesh sample point.
            if from_navmesh:
                obj_node = sim.get_object_scene_node(object_id)
                xform_bb = habitat_sim.geo.get_transformed_bb(
                    obj_node.cumulative_bb, obj_node.transformation
                )
                # also account for collision margin of the scene
                y_translation = mn.Vector3(
                    0, xform_bb.size_y() / 2.0 + scene_collision_margin, 0
                )
                sim.set_translation(
                    y_translation + sim.get_translation(object_id), object_id
                )

            # test for penetration with the environment
            if not sim.contact_test(object_id):
                valid_placement = True

        if not valid_placement:
            return False
        return True

    def geodesic_distance(
        self,
        position_a: Union[Sequence[float], ndarray],
        position_b: Union[Sequence[float], Sequence[Sequence[float]]],
        episode: Optional[Episode] = None,
    ) -> float:
        if episode is None or episode._shortest_path_cache is None:
            path = habitat_sim.MultiGoalShortestPath()
            if isinstance(position_b[0], (Sequence, np.ndarray)):
                path.requested_ends = np.array(position_b, dtype=np.float32)
            else:
                path.requested_ends = np.array(
                    [np.array(position_b, dtype=np.float32)]
                )
        else:
            path = episode._shortest_path_cache

        path.requested_start = np.array(position_a, dtype=np.float32)

        self.pathfinder.find_path(path)

        if episode is not None:
            episode._shortest_path_cache = path

        return path.geodesic_distance

    def action_space_shortest_path(
        self,
        source: AgentState,
        targets: Sequence[AgentState],
        agent_id: int = 0,
    ) -> List[ShortestPathPoint]:
        r"""
        Returns:
            List of agent states and actions along the shortest path from
            source to the nearest target (both included). If one of the
            target(s) is identical to the source, a list containing only
            one node with the identical agent state is returned. Returns
            an empty list in case none of the targets are reachable from
            the source. For the last item in the returned list the action
            will be None.
        """
        raise NotImplementedError(
            "This function is no longer implemented. Please use the greedy "
            "follower instead"
        )

    @property
    def up_vector(self) -> np.ndarray:
        return np.array([0.0, 1.0, 0.0])

    @property
    def forward_vector(self) -> np.ndarray:
        return -np.array([0.0, 0.0, 1.0])

    def get_straight_shortest_path_points(self, position_a, position_b):
        path = habitat_sim.ShortestPath()
        path.requested_start = position_a
        path.requested_end = position_b
        self.pathfinder.find_path(path)
        return path.points

    def sample_navigable_point(self) -> List[float]:
        return self.pathfinder.get_random_navigable_point().tolist()

    def is_navigable(self, point: List[float]) -> bool:
        return self.pathfinder.is_navigable(point)

    def semantic_annotations(self):
        r"""
        Returns:
            SemanticScene which is a three level hierarchy of semantic
            annotations for the current scene. Specifically this method
            returns a SemanticScene which contains a list of SemanticLevel's
            where each SemanticLevel contains a list of SemanticRegion's where
            each SemanticRegion contains a list of SemanticObject's.

            SemanticScene has attributes: aabb(axis-aligned bounding box) which
            has attributes aabb.center and aabb.sizes which are 3d vectors,
            categories, levels, objects, regions.

            SemanticLevel has attributes: id, aabb, objects and regions.

            SemanticRegion has attributes: id, level, aabb, category (to get
            name of category use category.name()) and objects.

            SemanticObject has attributes: id, region, aabb, obb (oriented
            bounding box) and category.

            SemanticScene contains List[SemanticLevels]
            SemanticLevel contains List[SemanticRegion]
            SemanticRegion contains List[SemanticObject]

            Example to loop through in a hierarchical fashion:
            for level in semantic_scene.levels:
                for region in level.regions:
                    for obj in region.objects:
        """
        return self.semantic_scene

    def _get_agent_config(self, agent_id: Optional[int] = None) -> Any:
        if agent_id is None:
            agent_id = self.habitat_config.DEFAULT_AGENT_ID
        agent_name = self.habitat_config.AGENTS[agent_id]
        agent_config = getattr(self.habitat_config, agent_name)
        return agent_config

    def get_agent_state(self, agent_id: int = 0) -> habitat_sim.AgentState:
        assert agent_id == 0, "No support of multi agent in {} yet.".format(
            self.__class__.__name__
        )
        return self.get_agent(agent_id).get_state()

    def set_agent_state(
        self,
        position: List[float],
        rotation: List[float],
        agent_id: int = 0,
        reset_sensors: bool = True,
    ) -> bool:
        r"""Sets agent state similar to initialize_agent, but without agents
        creation. On failure to place the agent in the proper position, it is
        moved back to its previous pose.

        Args:
            position: list containing 3 entries for (x, y, z).
            rotation: list with 4 entries for (x, y, z, w) elements of unit
                quaternion (versor) representing agent 3D orientation,
                (https://en.wikipedia.org/wiki/Versor)
            agent_id: int identification of agent from multiagent setup.
            reset_sensors: bool for if sensor changes (e.g. tilt) should be
                reset).

        Returns:
            True if the set was successful else moves the agent back to its
            original pose and returns false.
        """
        agent = self.get_agent(agent_id)
        new_state = self.get_agent_state(agent_id)
        new_state.position = position
        new_state.rotation = rotation

        # NB: The agent state also contains the sensor states in _absolute_
        # coordinates. In order to set the agent's body to a specific
        # location and have the sensors follow, we must not provide any
        # state for the sensors. This will cause them to follow the agent's
        # body
        new_state.sensor_states = {}
        agent.set_state(new_state, reset_sensors)
        return True

    def get_observations_at(
        self,
        position: Optional[List[float]] = None,
        rotation: Optional[List[float]] = None,
        keep_agent_at_new_pose: bool = False,
    ) -> Optional[Observations]:
        current_state = self.get_agent_state()
        if position is None or rotation is None:
            success = True
        else:
            success = self.set_agent_state(
                position, rotation, reset_sensors=False
            )

        if success:
            sim_obs = self.get_sensor_observations()

            self._prev_sim_obs = sim_obs

            observations = self._sensor_suite.get_observations(sim_obs)
            if not keep_agent_at_new_pose:
                self.set_agent_state(
                    current_state.position,
                    current_state.rotation,
                    reset_sensors=False,
                )
            return observations
        else:
            return None

    def distance_to_closest_obstacle(
        self, position: ndarray, max_search_radius: float = 2.0
    ) -> float:
        return self.pathfinder.distance_to_closest_obstacle(
            position, max_search_radius
        )

    def island_radius(self, position: Sequence[float]) -> float:
        return self.pathfinder.island_radius(position)

    @property
    def previous_step_collided(self):
        r"""Whether or not the previous step resulted in a collision

        Returns:
            bool: True if the previous step resulted in a collision, false otherwise

        Warning:
            This feild is only updated when :meth:`step`, :meth:`reset`, or :meth:`get_observations_at` are
            called.  It does not update when the agent is moved to a new loction.  Furthermore, it
            will _always_ be false after :meth:`reset` or :meth:`get_observations_at` as neither of those
            result in an action (step) being taken.
        """
        return self._prev_sim_obs.get("collided", False)

habitat/sims/habitat_simulator/habitat_simulator.py
cshizhe/onav_rim
from habitat.core.registry import registry
from habitat.core.simulator import Simulator


def _try_register_pickplace_sim():
    try:
        import habitat_sim  # noqa: F401

        has_habitat_sim = True
    except ImportError as e:
        has_habitat_sim = False
        rearrangement_sim_import_error = e

    if has_habitat_sim:
        from habitat.sims.pickplace.actions import (  # noqa: F401
            PickPlaceSimV0ActionSpaceConfiguration,
        )
    else:

        @registry.register_simulator(name="PickPlaceSim-v0")
        class PickPlaceSimImportError(Simulator):
            def __init__(self, *args, **kwargs):
                raise rearrangement_sim_import_error

habitat/sims/pickplace/__init__.py
cshizhe/onav_rim

#!/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.

from enum import Enum
from typing import Dict, List, Any

import attr

import habitat_sim
from habitat.core.embodied_task import EmbodiedTask
from habitat.core.registry import registry
from habitat.core.simulator import ActionSpaceConfiguration
from habitat.core.utils import Singleton

from habitat.core.embodied_task import SimulatorTaskAction
from habitat_sim.agent.controls.controls import ActuationSpec
from habitat.sims.habitat_simulator.actions import (
    HabitatSimActions,
    HabitatSimV1ActionSpaceConfiguration
)


@attr.s(auto_attribs=True, slots=True)
class GrabReleaseActuationSpec(ActuationSpec):
    visual_sensor_name: str = "rgb"
    crosshair_pos: List[int] = [128, 128]
    amount: float = 2.0


@registry.register_action_space_configuration(name="PickPlaceActions-v0")
class PickPlaceSimV0ActionSpaceConfiguration(ActionSpaceConfiguration):
    def __init__(self, config):
        super().__init__(config)
        # self.config = config
        if not HabitatSimActions.has_action("GRAB_RELEASE"):
            HabitatSimActions.extend_action_space("GRAB_RELEASE")
        if not HabitatSimActions.has_action("MOVE_BACKWARD"):
            HabitatSimActions.extend_action_space("MOVE_BACKWARD")
        if not HabitatSimActions.has_action("NO_OP"):
            HabitatSimActions.extend_action_space("NO_OP")
        if not HabitatSimActions.has_action("STOP"):
            HabitatSimActions.extend_action_space("STOP")

    def get(self):
        #config = super().get()
        new_config = {
            HabitatSimActions.STOP: habitat_sim.ActionSpec("stop"),
            HabitatSimActions.MOVE_FORWARD: habitat_sim.ActionSpec(
                "move_forward",
                habitat_sim.ActuationSpec(amount=self.config.FORWARD_STEP_SIZE),
            ),
            HabitatSimActions.TURN_LEFT: habitat_sim.ActionSpec(
                "turn_left",
                habitat_sim.ActuationSpec(amount=self.config.TURN_ANGLE),
            ),
            HabitatSimActions.TURN_RIGHT: habitat_sim.ActionSpec(
                "turn_right",
                habitat_sim.ActuationSpec(amount=self.config.TURN_ANGLE),
            ),
            HabitatSimActions.LOOK_UP: habitat_sim.ActionSpec(
                "look_up",
                habitat_sim.ActuationSpec(amount=self.config.TILT_ANGLE),
            ),
            HabitatSimActions.LOOK_DOWN: habitat_sim.ActionSpec(
                "look_down",
                habitat_sim.ActuationSpec(amount=self.config.TILT_ANGLE),
            ),
            HabitatSimActions.MOVE_BACKWARD: habitat_sim.ActionSpec(
                "move_backward",
                habitat_sim.ActuationSpec(amount=self.config.FORWARD_STEP_SIZE),
            ),
            HabitatSimActions.NO_OP: habitat_sim.ActionSpec(
                "no_op",
                habitat_sim.ActuationSpec(amount=0.05),
            ),
            HabitatSimActions.GRAB_RELEASE: habitat_sim.ActionSpec(
                "grab_or_release_object_under_crosshair",
                GrabReleaseActuationSpec(
                    crosshair_pos=self.config.CROSSHAIR_POS,
                    amount=self.config.GRAB_DISTANCE,
                ),
            ),
        }

        return new_config


@registry.register_task_action
class NoOpAction(SimulatorTaskAction):
    name: str = "NO_OP"

    def step(self, *args: Any, **kwargs: Any):
        r"""Update ``_metric``, this method is called from ``Env`` on each
        ``step``.
        """
        if "replay_data" in kwargs.keys():
            return self._sim.step_from_replay(
                HabitatSimActions.NO_OP,
                replay_data=kwargs["replay_data"]
            )
        return self._sim.step(HabitatSimActions.NO_OP)


@registry.register_task_action
class GrabOrReleaseAction(SimulatorTaskAction):

    def step(self, task: EmbodiedTask, *args: Any, **kwargs: Any):
        r"""This method is called from ``Env`` on each ``step``."""
        if "replay_data" in kwargs.keys():
            return self._sim.step_from_replay(
                HabitatSimActions.GRAB_RELEASE,
                replay_data=kwargs["replay_data"]
            )
        return self._sim.step(HabitatSimActions.GRAB_RELEASE)


@registry.register_task_action
class MoveBackwardAction(SimulatorTaskAction):
    name: str = "MOVE_BACKWARD"

    def step(self, *args: Any, **kwargs: Any):
        r"""Update ``_metric``, this method is called from ``Env`` on each
        ``step``.
        """
        if "replay_data" in kwargs.keys():
            return self._sim.step_from_replay(
                HabitatSimActions.MOVE_BACKWARD,
                replay_data=kwargs["replay_data"]
            )
        return self._sim.step(HabitatSimActions.MOVE_BACKWARD)


habitat/sims/pickplace/actions.py
cshizhe/onav_rim
#!/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.

from typing import Dict, List, Optional

import numpy as np
import math
import magnum as mn

from collections import defaultdict

from habitat_sim.nav import NavMeshSettings
from habitat_sim.utils.common import quat_from_coeffs, quat_to_magnum, quat_to_coeffs, quat_from_magnum
from habitat_sim.physics import MotionType

from habitat.core.registry import registry
from habitat.core.simulator import (
    Config,
    Observations
)
from habitat.sims.habitat_simulator.habitat_simulator import HabitatSim
from habitat.tasks.pickplace.pickplace import ObjectStateSpec


@registry.register_simulator(name="PickPlaceSim-v0")
class PickPlaceSim(HabitatSim):
    r"""Simulator wrapper over habitat-sim with
    object rearrangement functionalities.
    """

    def __init__(self, config: Config) -> None:
        self.did_reset = False
        super().__init__(config=config)
        self.nearest_object_id = -1
        self.gripped_object_id = -1
        self.gripped_object_transformation = None
        self.max_distance = 2.0
        self.agent_object_handle = "cylinderSolid_rings_1_segments_12_halfLen_1_useTexCoords_false_useTangents_false_capEnds_true"

        agent_id = self.habitat_config.DEFAULT_AGENT_ID
        agent_config = self._get_agent_config(agent_id)

        self.navmesh_settings = NavMeshSettings()
        self.navmesh_settings.set_defaults()
        self.navmesh_settings.agent_radius = agent_config.RADIUS
        self.navmesh_settings.agent_height = agent_config.HEIGHT
        self.default_agent_id = agent_id
        self.ISLAND_RADIUS_LIMIT = 2.0

    def reconfigure(self, config: Config) -> None:
        self.remove_existing_objects()
        super().reconfigure(config)
        self._initialize_objects(config)

    def reset(self):
        sim_obs = super().reset()
        if self._update_agents_state():
            sim_obs = self.get_sensor_observations(agent_ids=self.default_agent_id)

        self._prev_sim_obs = sim_obs
        self._prev_sim_obs["gripped_object_id"] = -1
        self._prev_sim_obs["object_under_cross_hair"] = -1
        self._prev_step_data = defaultdict(int)
        self.did_reset = True
        self.gripped_object_id = -1
        self.nearest_object_id = -1
        return self._sensor_suite.get_observations(sim_obs)

    def remove_existing_objects(self):
        # first remove all existing objects
        existing_object_ids = self.get_existing_object_ids()

        for obj_id in existing_object_ids:
            self.remove_object(obj_id)
            object_ = self.get_object_from_scene(obj_id)
            if object_ is not None:
                self.remove_contact_test_object(object_.object_handle)
        self.remove_contact_test_object(self.agent_object_handle)
        self.clear_recycled_object_ids()
        self.clear_scene_objects()

    def _initialize_objects(self, sim_config):
        # Add contact test object for agent
        self.add_contact_test_object(self.agent_object_handle)
        objects = sim_config.objects
        obj_attr_mgr = self.get_object_template_manager()

        self.obj_id_to_semantic_obj_id_map = defaultdict(int)

        if objects is not None:
            # Sort objects by object id
            object_map = {}
            for object_ in objects:
                object_map[object_.object_id] = object_

            for key in sorted(object_map.keys()):
                object_ = object_map[key]
                object_handle = object_.object_template.split('/')[-1]
                object_template = "data/test_assets/objects/{}".format(object_handle)
                object_pos = object_.position
                rotation = quat_from_coeffs(object_.rotation)
                object_rotation = quat_to_magnum(rotation)

                object_template_id = obj_attr_mgr.load_object_configs(
                    object_template
                )[0]
                object_attr = obj_attr_mgr.get_template_by_ID(object_template_id)
                obj_attr_mgr.register_template(object_attr)

                object_id = self.add_object_by_handle(object_attr.handle)
                self.add_contact_test_object(object_attr.handle)

                self.set_translation(object_pos, object_id)
                self.set_rotation(object_rotation, object_id)

                object_.object_handle = "data/test_assets/objects/{}".format(object_handle)
                self.add_object_in_scene(object_id, object_)

                self.set_object_motion_type(MotionType.DYNAMIC, object_id)
                self.obj_id_to_semantic_obj_id_map[object_id] = object_.semantic_object_id
                self.set_object_semantic_id(object_.semantic_object_id, object_id)


    def get_resolution(self):
        resolution = self._default_agent.agent_config.sensor_specifications[
            0
        ].resolution
        return mn.Vector2(list(map(int, resolution)))

    def clear_scene_objects(self):
        self._scene_objects = []

    def add_object_in_scene(self, objectId, data):
        data.object_id = objectId
        self._scene_objects.append(data)

    def update_object_in_scene(self, prevObjectId, newObjectId):
        for index in range(len(self._scene_objects)):
            if self._scene_objects[index].object_id == prevObjectId:
                self._scene_objects[index].object_id = newObjectId

    def get_object_from_scene(self, objectId):
        for index in range(len(self._scene_objects)):
            if self._scene_objects[index].object_id == objectId:
                return self._scene_objects[index]
        return None
    
    def check_object_exists_in_scene(self, object_id):
        exists = object_id in self.get_existing_object_ids()
        return exists
    
    def is_collision(self, handle, translation, is_navigation_test = False):
        collision_filter = 1
        collision_mask = -1
        if is_navigation_test:
            collision_mask = 1
        return self.pre_add_contact_test(
            handle, translation, is_navigation_test, collision_filter, collision_mask
        )

    def is_agent_colliding(self, action, agentTransform):
        stepSize = 0.15
        if action == "move_forward":
            position = agentTransform.backward * (-1 * stepSize)
            newPosition = agentTransform.translation + position
            filteredPoint = self.pathfinder.try_step(
                agentTransform.translation,
                newPosition
            )
            filterDiff = filteredPoint - newPosition
            # adding buffer of 0.05 y to avoid collision with navmesh
            offsets = [0, 0.05, -0.05, 0.1, -0.1, -0.2]
            is_colliding = False
            for offset in offsets:
                finalPosition = newPosition + filterDiff + mn.Vector3(0.0, offset, 0.0)
                collision = self.is_collision(self.agent_object_handle, finalPosition, True)
                if collision:
                    is_colliding = True
            return {
                "collision": is_colliding,
            }
        elif action == "move_backward":
            position = agentTransform.backward * stepSize
            newPosition = agentTransform.translation + position
            filteredPoint = self.pathfinder.try_step(
                agentTransform.translation,
                newPosition
            )
            filterDiff = filteredPoint - newPosition
            # adding buffer of 0.05 y to avoid collision with navmesh
            offsets = [0, 0.05, -0.05, 0.1, -0.1, -0.2]
            is_colliding = False
            for offset in offsets:
                finalPosition = newPosition + filterDiff + mn.Vector3(0.0, offset, 0.0)
                collision = self.is_collision(self.agent_object_handle, finalPosition, True)
                if collision:
                    is_colliding = True
            return {
                "collision": is_colliding,
            }
        return {
            "collision": False
        }

    def draw_bb_around_nearest_object(self, object_id):
        if object_id == -1:
            if self.nearest_object_id != -1 and self.gripped_object_id != self.nearest_object_id:
                self.set_object_bb_draw(False, self.nearest_object_id)
                self.nearest_object_id = object_id
        else:
            if self.nearest_object_id != -1 and self.gripped_object_id != self.nearest_object_id:
                self.set_object_bb_draw(False, self.nearest_object_id)
                self.nearest_object_id = -1
            object_ = self.get_object_from_scene(object_id)
            if object_.is_receptacle == True:
                return
            if self.nearest_object_id != object_id:
                self.nearest_object_id = object_id
                self.set_object_bb_draw(True, self.nearest_object_id, 0)

    def step(self, action: int):
        dt = 1.0 / 20.0
        self._num_total_frames += 1
        collided = False

        agent_config = self._default_agent.agent_config
        action_spec = agent_config.action_space[action]

        crosshair_pos = np.array(self.get_resolution()) // 2
        crosshair_pos = list(map(int, crosshair_pos))

        ray = self.unproject(crosshair_pos)
        cross_hair_point = ray.direction
        ref_point = self._default_agent.body.object.absolute_translation

        nearest_object_id = self.find_nearest_object_under_crosshair(
            cross_hair_point, ref_point, self.get_resolution(), self.max_distance
        )
        object_under_cross_hair = nearest_object_id
        self._prev_step_data["nearest_object_id"] = nearest_object_id
        self._prev_step_data["gripped_object_id"] = self.gripped_object_id

        if action_spec.name == "grab_or_release_object_under_crosshair":
            self._prev_step_data["action"] = action_spec.name
            self._prev_step_data["is_grab_action"] = self.gripped_object_id == -1
            self._prev_step_data["is_release_action"] = self.gripped_object_id != -1
            # already gripped an object
            if self.gripped_object_id != -1:
                ref_transform = self._default_agent.body.object.transformation
                ray_hit_info = self.find_floor_position_under_crosshair(
                    cross_hair_point, ref_transform,
                    self.get_resolution(), action_spec.actuation.amount
                )

                floor_position = ray_hit_info.point
                if floor_position is None:
                    return True

                new_object_position = mn.Vector3(
                    floor_position.x, floor_position.y, floor_position.z
                )
                scene_object = self.get_object_from_scene(self.gripped_object_id)

                # find no collision point
                count = 0
                contact = self.is_collision(scene_object.object_handle, new_object_position)
                while contact and count < 4:
                    new_object_position = mn.Vector3(
                        new_object_position.x,
                        new_object_position.y + 0.25,
                        new_object_position.z,
                    )
                    contact = self.is_collision(scene_object.object_handle, new_object_position)
                    count += 1

                snapped_point = self.pathfinder.snap_point(new_object_position)
                snapped_point = np.array(new_object_position)
                is_point_on_separate_island = self.island_radius(snapped_point) < self.ISLAND_RADIUS_LIMIT

                if not (contact or is_point_on_separate_island):
                    new_object_id = self.add_object_by_handle(
                        scene_object.object_handle
                    )
                    self.set_translation(new_object_position, new_object_id)
                    self.set_object_semantic_id(scene_object.semantic_object_id, new_object_id)

                    self.update_object_in_scene(self.gripped_object_id, new_object_id)
                    self._prev_step_data["new_object_translation"] = np.array(new_object_position).tolist()
                    self._prev_step_data["new_object_id"] = new_object_id
                    self._prev_step_data["object_handle"] = scene_object.object_handle
                    self._prev_step_data["gripped_object_id"] = self.gripped_object_id
                    self.gripped_object_id = -1
            elif nearest_object_id != -1:
                self.gripped_object_transformation = self.get_transformation(
                    nearest_object_id
                )
                object_ = self.get_object_from_scene(nearest_object_id)
                if object_.is_receptacle != True:
                    self.remove_object(nearest_object_id)
                    self.gripped_object_id = nearest_object_id
                self._prev_step_data["gripped_object_id"] = self.gripped_object_id
        elif action_spec.name == "no_op":
            pass
        else:
            agent_transform = self._default_agent.body.object.transformation
            data = self.is_agent_colliding(action_spec.name, agent_transform)
            if not data["collision"]:
                self._default_agent.act(action)
                collided = data["collision"]
                self._last_state = self._default_agent.get_state()

        # step world physics
        super().step_world(dt)

        # self.draw_bb_around_nearest_object(object_under_cross_hair)

        # obtain observations
        self._prev_sim_obs = self.get_sensor_observations(agent_ids=self.default_agent_id, draw_crosshair=False)
        self._prev_sim_obs["collided"] = collided
        self._prev_sim_obs["gripped_object_id"] = self.gripped_object_id
        self._prev_sim_obs["object_under_cross_hair"] = object_under_cross_hair
        self._prev_step_data["collided"] = collided

        observations = self._sensor_suite.get_observations(self._prev_sim_obs)
        return observations

    def is_object_out_of_bounds(self):
        object_ids = self.get_existing_object_ids()
        for object_id in object_ids:
            position = self.get_translation(object_id)
            snapped_position = self.pathfinder.snap_point(position)

            if np.isnan(snapped_position.x) or math.isnan(snapped_position.x):
                return True
        return False

    def restore_object_states(self, object_states: Dict = {}):
        object_ids = []
        for object_state in object_states:
            object_id = object_state.object_id
            translation = object_state.translation
            object_rotation = object_state.rotation

            object_translation = mn.Vector3(translation)
            if isinstance(object_rotation, list):
                object_rotation = quat_from_coeffs(object_rotation)

            object_rotation = quat_to_magnum(object_rotation)
            self.set_translation(object_translation, object_id)
            self.set_rotation(object_rotation, object_id)
            object_ids.append(object_id)
        return object_ids
    
    def get_current_object_states(self):
        existing_object_ids = self.get_existing_object_ids()
        object_states = []
        for object_id in existing_object_ids:
            translation = self.get_translation(object_id)
            rotation = self.get_rotation(object_id)
            rotation = quat_from_magnum(rotation)
            scene_object = self.get_object_from_scene(object_id)

            object_state = ObjectStateSpec(
                object_id=object_id,
                translation=np.array(translation).tolist(),
                rotation=quat_to_coeffs(rotation).tolist(),
                object_handle=scene_object.object_handle,
            )
            object_states.append(object_state)
        return object_states
    
    def get_agent_pose(self):
        agent_translation = self._default_agent.body.object.translation
        agent_rotation = self._default_agent.body.object.rotation
        sensor_data = {}
        for sensor_key, v in self._default_agent._sensors.items():
            rotation = quat_from_magnum(v.object.rotation)
            rotation = quat_to_coeffs(rotation).tolist()
            translation = v.object.translation
            sensor_data[sensor_key] = {
                "rotation": rotation,
                "translation": np.array(translation).tolist()
            }
        
        return {
            "position": np.array(agent_translation).tolist(),
            "rotation": quat_to_coeffs(quat_from_magnum(agent_rotation)).tolist(),
            "sensor_data": sensor_data
        }
    
    def restore_sensor_states(self, sensor_data: Dict = {}):
        for sensor_key, v in self._default_agent._sensors.items():
            rotation = None
            if sensor_key in sensor_data.keys():
                rotation = sensor_data[sensor_key]["rotation"]
            else:
                rotation = sensor_data["rgb"]["rotation"]
            rotation = quat_from_coeffs(rotation)
            agent_rotation = quat_to_magnum(rotation)
            v.object.rotation = agent_rotation
    
    def get_sensor_states(self):
        sensor_states = {}
        for sensor_key, v in self._default_agent._sensors.items():
            rotation = v.object.rotation
            rotation = quat_from_magnum(rotation)
            rotation = quat_to_coeffs(rotation)
            sensor_states[sensor_key] = {
                "rotation": rotation
            }
        return sensor_states

    def add_objects_by_handle(self, objects):
        for object_ in objects:
            object_handle = object_.object_handle
            object_id = self.add_object_by_handle(object_handle)
            self.set_object_semantic_id(object_.semantic_object_id, object_id)

    def step_from_replay(self, action: int, replay_data: Dict = {}):
        dt = 1.0 / 20.0
        self._num_total_frames += 1
        collided = False

        agent_config = self._default_agent.agent_config
        action_spec = agent_config.action_space[action]

        if action_spec.name == "grab_or_release_object_under_crosshair":
            action_data = replay_data.action_data
            if action_data is not None and action_data.released_object_id != -1:
                if replay_data.is_release_action:
                    # Fetch object handle and drop point from replay
                    new_object_position = mn.Vector3(action_data.released_object_position)
                    scene_object = self.get_object_from_scene(action_data.released_object_id)
                    new_object_id = self.add_object_by_handle(
                        scene_object.object_handle
                    )
                    self.set_translation(new_object_position, new_object_id)
                    self.set_object_semantic_id(scene_object.semantic_object_id, new_object_id)

                    self.update_object_in_scene(new_object_id, action_data.released_object_id)
                    self.gripped_object_id = replay_data.gripped_object_id
                elif replay_data.is_grab_action:
                    self.gripped_object_transformation = self.get_transformation(
                        action_data.grab_object_id
                    )
                    self.remove_object(action_data.grab_object_id)
                    self.gripped_object_id = replay_data.gripped_object_id
        elif action_spec.name == "no_op":
            self.restore_object_states(replay_data.object_states)
        else:
            if replay_data.agent_state is not None:
                if action_spec.name == "look_up" or action_spec.name == "look_down":
                    sensor_data = replay_data.agent_state.sensor_data
                    self.restore_sensor_states(sensor_data)
                else:
                    position = np.array(replay_data.agent_state.position)
                    rotation = np.array(replay_data.agent_state.rotation)
                    success = self.set_agent_state(
                        position, rotation, reset_sensors=False
                    )
                collided = replay_data.collision
                self._last_state = self._default_agent.get_state()
            else:
                collided = replay_data.collision
                if not collided:
                    self._default_agent.act(action)

        # self.draw_bb_around_nearest_object(replay_data.object_under_cross_hair)

        # obtain observations
        self._prev_sim_obs = self.get_sensor_observations(agent_ids=self.default_agent_id, draw_crosshair=False)
        self._prev_sim_obs["collided"] = collided
        self._prev_sim_obs["gripped_object_id"] = self.gripped_object_id

        observations = self._sensor_suite.get_observations(self._prev_sim_obs)
        return observations

    def get_observations_at(
        self,
        position: Optional[List[float]] = None,
        rotation: Optional[List[float]] = None,
        sensor_states: Optional[List[Dict]] = None,
        object_states: Optional[List[Dict]] = None,
        keep_agent_at_new_pose: bool = False,
    ) -> Optional[Observations]:
        current_state = self.get_agent_state()
        current_object_states = self.get_current_object_states()
        current_sensor_states = self.get_sensor_states()
        if position is None or rotation is None:
            success = True
        else:
            success = self.set_agent_state(
                np.array(position), np.array(rotation), reset_sensors=False
            )
        if sensor_states is not None:
            self.restore_sensor_states(sensor_states)

        object_to_re_add = []
        if object_states is not None:
            current_state_objects = self.restore_object_states(object_states)

            for object_id in self.get_existing_object_ids():
                if object_id not in current_state_objects:
                    self.remove_object(object_id)
                    object_to_re_add.append(self.get_object_from_scene(object_id))

        if success:
            sim_obs = self.get_sensor_observations(agent_ids=self.default_agent_id)

            self._prev_sim_obs = sim_obs

            observations = self._sensor_suite.get_observations(sim_obs)
            if not keep_agent_at_new_pose:
                self.set_agent_state(
                    current_state.position,
                    current_state.rotation,
                    reset_sensors=False,
                )
                self.add_objects_by_handle(object_to_re_add)
                self.restore_object_states(current_object_states)
                self.restore_sensor_states(current_sensor_states)
                
            return observations
        else:
            return None

habitat/sims/pickplace/simulator.py
cshizhe/onav_rim
#!/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.

from habitat.core.logging import logger
from habitat.core.registry import registry
from habitat.sims.habitat_simulator import _try_register_habitat_sim
from habitat.sims.pyrobot import _try_register_pyrobot
from habitat.sims.pickplace import _try_register_pickplace_sim


def make_sim(id_sim, **kwargs):
    logger.info("initializing sim {}".format(id_sim))
    _sim = registry.get_simulator(id_sim)
    assert _sim is not None, "Could not find simulator with name {}".format(
        id_sim
    )
    return _sim(**kwargs)


_try_register_habitat_sim()
_try_register_pyrobot()
_try_register_pickplace_sim()

habitat/sims/registration.py
cshizhe/onav_rim
#!/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.

# TODO, lots of typing errors in here

from typing import Any, List, Optional, Tuple

import attr
import numpy as np
from gym import spaces

from habitat.config import Config
from habitat.core.dataset import Dataset, Episode
from habitat.core.embodied_task import (
    EmbodiedTask,
    Measure,
    SimulatorTaskAction,
)
from habitat.core.logging import logger
from habitat.core.registry import registry
from habitat.core.simulator import (
    AgentState,
    RGBSensor,
    Sensor,
    SensorTypes,
    ShortestPathPoint,
    Simulator,
)
from habitat.core.utils import not_none_validator, try_cv2_import
from habitat.sims.habitat_simulator.actions import HabitatSimActions
from habitat.tasks.utils import cartesian_to_polar
from habitat.utils.geometry_utils import (
    quaternion_from_coeff,
    quaternion_rotate_vector,
)
from habitat.utils.visualizations import fog_of_war, maps

try:
    from habitat.sims.habitat_simulator.habitat_simulator import HabitatSim
except ImportError:
    pass
cv2 = try_cv2_import()


MAP_THICKNESS_SCALAR: int = 128


def merge_sim_episode_config(sim_config: Config, episode: Episode) -> Any:
    sim_config.defrost()
    sim_config.SCENE = episode.scene_id
    sim_config.freeze()
    if (
        episode.start_position is not None
        and episode.start_rotation is not None
    ):
        agent_name = sim_config.AGENTS[sim_config.DEFAULT_AGENT_ID]
        agent_cfg = getattr(sim_config, agent_name)
        agent_cfg.defrost()
        agent_cfg.START_POSITION = episode.start_position
        agent_cfg.START_ROTATION = episode.start_rotation
        agent_cfg.IS_SET_START_STATE = True
        agent_cfg.freeze()
    return sim_config


@attr.s(auto_attribs=True, kw_only=True)
class NavigationGoal:
    r"""Base class for a goal specification hierarchy."""

    position: List[float] = attr.ib(default=None, validator=not_none_validator)
    radius: Optional[float] = None


@attr.s(auto_attribs=True, kw_only=True)
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


@attr.s(auto_attribs=True, kw_only=True)
class NavigationEpisode(Episode):
    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[NavigationGoal] = attr.ib(
        default=None, validator=not_none_validator
    )
    start_room: Optional[str] = None
    shortest_paths: Optional[List[List[ShortestPathPoint]]] = None


@registry.register_sensor
class PointGoalSensor(Sensor):
    r"""Sensor for PointGoal observations which are used in PointGoal Navigation.

    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 PointGoal sensor. Can contain field for
            GOAL_FORMAT which can be used to specify the format in which
            the pointgoal is specified. Current options for goal format are
            cartesian and polar.

            Also contains a DIMENSIONALITY field which specifes the number
            of dimensions ued to specify the goal, must be in [2, 3]

    Attributes:
        _goal_format: format for specifying the goal which can be done
            in cartesian or polar coordinates.
        _dimensionality: number of dimensions used to specify the goal
    """
    cls_uuid: str = "pointgoal"

    def __init__(
        self, sim: Simulator, config: Config, *args: Any, **kwargs: Any
    ):
        self._sim = sim

        self._goal_format = getattr(config, "GOAL_FORMAT", "CARTESIAN")
        assert self._goal_format in ["CARTESIAN", "POLAR"]

        self._dimensionality = getattr(config, "DIMENSIONALITY", 2)
        assert self._dimensionality in [2, 3]

        super().__init__(config=config)

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        return self.cls_uuid

    def _get_sensor_type(self, *args: Any, **kwargs: Any):
        return SensorTypes.PATH

    def _get_observation_space(self, *args: Any, **kwargs: Any):
        sensor_shape = (self._dimensionality,)

        return spaces.Box(
            low=np.finfo(np.float32).min,
            high=np.finfo(np.float32).max,
            shape=sensor_shape,
            dtype=np.float32,
        )

    def _compute_pointgoal(
        self, source_position, source_rotation, goal_position
    ):
        direction_vector = goal_position - source_position
        direction_vector_agent = quaternion_rotate_vector(
            source_rotation.inverse(), direction_vector
        )

        if self._goal_format == "POLAR":
            if self._dimensionality == 2:
                rho, phi = cartesian_to_polar(
                    -direction_vector_agent[2], direction_vector_agent[0]
                )
                return np.array([rho, -phi], dtype=np.float32)
            else:
                _, phi = cartesian_to_polar(
                    -direction_vector_agent[2], direction_vector_agent[0]
                )
                theta = np.arccos(
                    direction_vector_agent[1]
                    / np.linalg.norm(direction_vector_agent)
                )
                rho = np.linalg.norm(direction_vector_agent)

                return np.array([rho, -phi, theta], dtype=np.float32)
        else:
            if self._dimensionality == 2:
                return np.array(
                    [-direction_vector_agent[2], direction_vector_agent[0]],
                    dtype=np.float32,
                )
            else:
                return direction_vector_agent

    def get_observation(
        self,
        observations,
        episode: NavigationEpisode,
        *args: Any,
        **kwargs: Any,
    ):
        source_position = np.array(episode.start_position, dtype=np.float32)
        rotation_world_start = quaternion_from_coeff(episode.start_rotation)
        goal_position = np.array(episode.goals[0].position, dtype=np.float32)

        return self._compute_pointgoal(
            source_position, rotation_world_start, goal_position
        )


@registry.register_sensor
class ImageGoalSensor(Sensor):
    r"""Sensor for ImageGoal observations which are used in ImageGoal Navigation.

    RGBSensor needs to be one of the Simulator sensors.
    This sensor return the rgb image taken from the goal position to reach with
    random rotation.

    Args:
        sim: reference to the simulator for calculating task observations.
        config: config for the ImageGoal sensor.
    """
    cls_uuid: str = "imagegoal"

    def __init__(
        self, *args: Any, sim: Simulator, config: Config, **kwargs: Any
    ):
        self._sim = sim
        sensors = self._sim.sensor_suite.sensors
        rgb_sensor_uuids = [
            uuid
            for uuid, sensor in sensors.items()
            if isinstance(sensor, RGBSensor)
        ]
        if len(rgb_sensor_uuids) != 1:
            raise ValueError(
                f"ImageGoalNav requires one RGB sensor, {len(rgb_sensor_uuids)} detected"
            )

        (self._rgb_sensor_uuid,) = rgb_sensor_uuids
        self._current_episode_id: Optional[str] = None
        self._current_image_goal = None
        super().__init__(config=config)

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        return self.cls_uuid

    def _get_sensor_type(self, *args: Any, **kwargs: Any):
        return SensorTypes.PATH

    def _get_observation_space(self, *args: Any, **kwargs: Any):
        return self._sim.sensor_suite.observation_spaces.spaces[
            self._rgb_sensor_uuid
        ]

    def _get_pointnav_episode_image_goal(self, episode: NavigationEpisode):
        goal_position = np.array(episode.goals[0].position, dtype=np.float32)
        # to be sure that the rotation is the same for the same episode_id
        # since the task is currently using pointnav Dataset.
        seed = abs(hash(episode.episode_id)) % (2 ** 32)
        rng = np.random.RandomState(seed)
        angle = rng.uniform(0, 2 * np.pi)
        source_rotation = [0, np.sin(angle / 2), 0, np.cos(angle / 2)]
        goal_observation = self._sim.get_observations_at(
            position=goal_position.tolist(), rotation=source_rotation
        )
        return goal_observation[self._rgb_sensor_uuid]

    def get_observation(
        self,
        *args: Any,
        observations,
        episode: NavigationEpisode,
        **kwargs: Any,
    ):
        episode_uniq_id = f"{episode.scene_id} {episode.episode_id}"
        if episode_uniq_id == self._current_episode_id:
            return self._current_image_goal

        self._current_image_goal = self._get_pointnav_episode_image_goal(
            episode
        )
        self._current_episode_id = episode_uniq_id

        return self._current_image_goal


@registry.register_sensor(name="PointGoalWithGPSCompassSensor")
class IntegratedPointGoalGPSAndCompassSensor(PointGoalSensor):
    r"""Sensor that integrates PointGoals observations (which are used PointGoal Navigation) and GPS+Compass.

    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 PointGoal sensor. Can contain field for
            GOAL_FORMAT which can be used to specify the format in which
            the pointgoal is specified. Current options for goal format are
            cartesian and polar.

            Also contains a DIMENSIONALITY field which specifes the number
            of dimensions ued to specify the goal, must be in [2, 3]

    Attributes:
        _goal_format: format for specifying the goal which can be done
            in cartesian or polar coordinates.
        _dimensionality: number of dimensions used to specify the goal
    """
    cls_uuid: str = "pointgoal_with_gps_compass"

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        return self.cls_uuid

    def get_observation(
        self, observations, episode, *args: Any, **kwargs: Any
    ):
        agent_state = self._sim.get_agent_state()
        agent_position = agent_state.position
        rotation_world_agent = agent_state.rotation
        goal_position = np.array(episode.goals[0].position, dtype=np.float32)

        return self._compute_pointgoal(
            agent_position, rotation_world_agent, goal_position
        )


@registry.register_sensor
class HeadingSensor(Sensor):
    r"""Sensor for observing the agent's heading in the global coordinate
    frame.

    Args:
        sim: reference to the simulator for calculating task observations.
        config: config for the sensor.
    """
    cls_uuid: str = "heading"

    def __init__(
        self, sim: Simulator, config: Config, *args: Any, **kwargs: Any
    ):
        self._sim = sim
        super().__init__(config=config)

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        return self.cls_uuid

    def _get_sensor_type(self, *args: Any, **kwargs: Any):
        return SensorTypes.HEADING

    def _get_observation_space(self, *args: Any, **kwargs: Any):
        return spaces.Box(low=-np.pi, high=np.pi, shape=(1,), dtype=np.float)

    def _quat_to_xy_heading(self, quat):
        direction_vector = np.array([0, 0, -1])

        heading_vector = quaternion_rotate_vector(quat, direction_vector)

        phi = cartesian_to_polar(-heading_vector[2], heading_vector[0])[1]
        return np.array([phi], dtype=np.float32)

    def get_observation(
        self, observations, episode, *args: Any, **kwargs: Any
    ):
        agent_state = self._sim.get_agent_state()
        rotation_world_agent = agent_state.rotation

        return self._quat_to_xy_heading(rotation_world_agent.inverse())


@registry.register_sensor(name="CompassSensor")
class EpisodicCompassSensor(HeadingSensor):
    r"""The agents heading in the coordinate frame defined by the epiosde,
    theta=0 is defined by the agents state at t=0
    """
    cls_uuid: str = "compass"

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        return self.cls_uuid

    def get_observation(
        self, observations, episode, *args: Any, **kwargs: Any
    ):
        agent_state = self._sim.get_agent_state()
        rotation_world_agent = agent_state.rotation
        rotation_world_start = quaternion_from_coeff(episode.start_rotation)

        return self._quat_to_xy_heading(
            rotation_world_agent.inverse() * rotation_world_start
        )


@registry.register_sensor(name="GPSSensor")
class EpisodicGPSSensor(Sensor):
    r"""The agents current location in the coordinate frame defined by the episode,
    i.e. the axis it faces along and the origin is defined by its state at t=0

    Args:
        sim: reference to the simulator for calculating task observations.
        config: Contains the DIMENSIONALITY field for the number of dimensions to express the agents position
    Attributes:
        _dimensionality: number of dimensions used to specify the agents position
    """
    cls_uuid: str = "gps"

    def __init__(
        self, sim: Simulator, config: Config, *args: Any, **kwargs: Any
    ):
        self._sim = sim

        self._dimensionality = getattr(config, "DIMENSIONALITY", 2)
        assert self._dimensionality in [2, 3]
        super().__init__(config=config)

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        return self.cls_uuid

    def _get_sensor_type(self, *args: Any, **kwargs: Any):
        return SensorTypes.POSITION

    def _get_observation_space(self, *args: Any, **kwargs: Any):
        sensor_shape = (self._dimensionality,)
        return spaces.Box(
            low=np.finfo(np.float32).min,
            high=np.finfo(np.float32).max,
            shape=sensor_shape,
            dtype=np.float32,
        )

    def get_observation(
        self, observations, episode, *args: Any, **kwargs: Any
    ):
        agent_state = self._sim.get_agent_state()

        origin = np.array(episode.start_position, dtype=np.float32)
        rotation_world_start = quaternion_from_coeff(episode.start_rotation)

        agent_position = agent_state.position

        agent_position = quaternion_rotate_vector(
            rotation_world_start.inverse(), agent_position - origin
        )
        if self._dimensionality == 2:
            return np.array(
                [-agent_position[2], agent_position[0]], dtype=np.float32
            )
        else:
            return agent_position.astype(np.float32)


@registry.register_sensor
class ProximitySensor(Sensor):
    r"""Sensor for observing the distance to the closest obstacle

    Args:
        sim: reference to the simulator for calculating task observations.
        config: config for the sensor.
    """
    cls_uuid: str = "proximity"

    def __init__(self, sim, config, *args: Any, **kwargs: Any):
        self._sim = sim
        self._max_detection_radius = getattr(
            config, "MAX_DETECTION_RADIUS", 2.0
        )
        super().__init__(config=config)

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        return self.cls_uuid

    def _get_sensor_type(self, *args: Any, **kwargs: Any):
        return SensorTypes.TACTILE

    def _get_observation_space(self, *args: Any, **kwargs: Any):
        return spaces.Box(
            low=0.0,
            high=self._max_detection_radius,
            shape=(1,),
            dtype=np.float32,
        )

    def get_observation(
        self, observations, *args: Any, episode, **kwargs: Any
    ):
        current_position = self._sim.get_agent_state().position

        return np.array(
            [
                self._sim.distance_to_closest_obstacle(
                    current_position, self._max_detection_radius
                )
            ],
            dtype=np.float32,
        )


@registry.register_measure
class Success(Measure):
    r"""Whether or not the agent succeeded at its task

    This measure depends on DistanceToGoal measure.
    """

    cls_uuid: str = "success"

    def __init__(
        self, sim: Simulator, config: Config, *args: Any, **kwargs: Any
    ):
        self._sim = sim
        self._config = config

        super().__init__()

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        return self.cls_uuid

    def reset_metric(self, episode, task, *args: Any, **kwargs: Any):
        task.measurements.check_measure_dependencies(
            self.uuid, [DistanceToGoal.cls_uuid]
        )
        self.update_metric(episode=episode, task=task, *args, **kwargs)  # type: ignore

    def update_metric(
        self, episode, task: EmbodiedTask, *args: Any, **kwargs: Any
    ):
        distance_to_target = task.measurements.measures[
            DistanceToGoal.cls_uuid
        ].get_metric()
        success_distance = self._config.SUCCESS_DISTANCE

        if hasattr(episode, "is_thda") and episode.is_thda:
            success_distance = self._config.THDA_SUCCESS_DISTANCE

        if (
            hasattr(task, "is_stop_called")
            and task.is_stop_called  # type: ignore
            and distance_to_target < success_distance
        ):
            self._metric = 1.0
        else:
            self._metric = 0.0


@registry.register_measure
class SPL(Measure):
    r"""SPL (Success weighted by Path Length)

    ref: On Evaluation of Embodied Agents - Anderson et. al
    https://arxiv.org/pdf/1807.06757.pdf
    The measure depends on Distance to Goal measure and Success measure
    to improve computational
    performance for sophisticated goal areas.
    """

    def __init__(
        self, sim: Simulator, config: Config, *args: Any, **kwargs: Any
    ):
        self._previous_position = None
        self._start_end_episode_distance = None
        self._agent_episode_distance: Optional[float] = None
        self._episode_view_points = None
        self._sim = sim
        self._config = config

        super().__init__()

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        return "spl"

    def reset_metric(self, episode, task, *args: Any, **kwargs: Any):
        task.measurements.check_measure_dependencies(
            self.uuid, [DistanceToGoal.cls_uuid, Success.cls_uuid]
        )

        self._previous_position = self._sim.get_agent_state().position
        self._agent_episode_distance = 0.0
        self._start_end_episode_distance = task.measurements.measures[
            DistanceToGoal.cls_uuid
        ].get_metric()
        self.update_metric(  # type:ignore
            episode=episode, task=task, *args, **kwargs
        )

    def _euclidean_distance(self, position_a, position_b):
        return np.linalg.norm(position_b - position_a, ord=2)

    def update_metric(
        self, episode, task: EmbodiedTask, *args: Any, **kwargs: Any
    ):
        ep_success = task.measurements.measures[Success.cls_uuid].get_metric()

        current_position = self._sim.get_agent_state().position
        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
            )
        )
        # print(self._start_end_episode_distance, self._agent_episode_distance, self._metric)


@registry.register_measure
class SoftSPL(SPL):
    r"""Soft SPL

    Similar to SPL with a relaxed soft-success criteria. Instead of a boolean
    success is now calculated as 1 - (ratio of distance covered to target).
    """

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        return "softspl"

    def reset_metric(self, episode, task, *args: Any, **kwargs: Any):
        task.measurements.check_measure_dependencies(
            self.uuid, [DistanceToGoal.cls_uuid]
        )

        self._previous_position = self._sim.get_agent_state().position
        self._agent_episode_distance = 0.0
        self._start_end_episode_distance = task.measurements.measures[
            DistanceToGoal.cls_uuid
        ].get_metric()
        self.update_metric(episode=episode, task=task, *args, **kwargs)  # type: ignore

    def update_metric(self, episode, task, *args: Any, **kwargs: Any):
        current_position = self._sim.get_agent_state().position
        distance_to_target = task.measurements.measures[
            DistanceToGoal.cls_uuid
        ].get_metric()

        ep_soft_success = max(
            0, (1 - distance_to_target / self._start_end_episode_distance)
        )

        self._agent_episode_distance += self._euclidean_distance(
            current_position, self._previous_position
        )

        self._previous_position = current_position

        self._metric = ep_soft_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, *args: Any, **kwargs: Any):
        self._sim = sim
        self._config = config
        self._metric = None
        super().__init__()

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        return "collisions"

    def reset_metric(self, episode, *args: Any, **kwargs: Any):
        self._metric = None

    def update_metric(self, episode, action, *args: Any, **kwargs: Any):
        if self._metric is None:
            self._metric = {"count": 0, "is_collision": False}
        self._metric["is_collision"] = False
        if self._sim.previous_step_collided:
            self._metric["count"] += 1
            self._metric["is_collision"] = True


@registry.register_measure
class TopDownMap(Measure):
    r"""Top Down Map measure"""
    cls_uuid = "top_down_map"

    def __init__(
        self, sim: "HabitatSim", config: Config, *args: Any, **kwargs: Any
    ):
        self._sim = sim
        self._config = config
        self._grid_delta = config.MAP_PADDING
        self._step_count: Optional[int] = None
        self._map_resolution = config.MAP_RESOLUTION
        self._ind_x_min: Optional[int] = None
        self._ind_x_max: Optional[int] = None
        self._ind_y_min: Optional[int] = None
        self._ind_y_max: Optional[int] = None
        self._previous_xy_location: Optional[Tuple[int, int]] = None
        self._top_down_map: Optional[np.ndarray] = None
        self._shortest_path_points: Optional[List[Tuple[int, int]]] = None
        self.line_thickness = int(
            np.round(self._map_resolution * 2 / MAP_THICKNESS_SCALAR)
        )
        self.point_padding = 2 * int(
            np.ceil(self._map_resolution / MAP_THICKNESS_SCALAR)
        )
        super().__init__()

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        return "top_down_map"

    def get_original_map(self):
        top_down_map = maps.get_topdown_map_from_sim(
            self._sim,
            map_resolution=self._map_resolution,
            draw_border=self._config.DRAW_BORDER,
        )

        if self._config.FOG_OF_WAR.DRAW:
            self._fog_of_war_mask = np.zeros_like(top_down_map)
        else:
            self._fog_of_war_mask = None

        return top_down_map

    def _draw_point(self, position, point_type):
        t_x, t_y = maps.to_grid(
            position[2],
            position[0],
            self._top_down_map.shape[0:2],
            sim=self._sim,
        )
        self._top_down_map[
            t_x - self.point_padding : t_x + self.point_padding + 1,
            t_y - self.point_padding : t_y + self.point_padding + 1,
        ] = point_type

    def _draw_goals_view_points(self, episode):
        if self._config.DRAW_VIEW_POINTS:
            agent_position = self._sim.get_agent_state().position
            for goal in episode.goals:
                try:
                    if goal.view_points is not None:
                        for view_point in goal.view_points:
                            # if abs(agent_position[1] - view_point.agent_state.position[1]) > 0.3:
                            #     continue
                            self._draw_point(
                                view_point.agent_state.position,
                                maps.MAP_VIEW_POINT_INDICATOR,
                            )
                except AttributeError:
                    pass

    def _draw_goals_positions(self, episode):
        if self._config.DRAW_GOAL_POSITIONS:
            for goal in episode.goals:
                agent_position = self._sim.get_agent_state().position
                # if abs(agent_position[1] - goal.position[1]) > 0.3:
                #     continue
                try:
                    if len(episode.goals) == 2 and goal.info["is_receptacle"] == True:
                        self._draw_point(
                            goal.position, maps.MAP_RECEPTACLE_COLOR
                        )
                    else:
                        self._draw_point(
                            goal.position, maps.MAP_TARGET_POINT_INDICATOR
                        )
                except AttributeError:
                    pass

    def _draw_goals_aabb(self, episode):
        if self._config.DRAW_GOAL_AABBS:
            for goal in episode.goals:
                try:
                    sem_scene = self._sim.semantic_annotations()
                    object_id = goal.object_id
                    assert int(
                        sem_scene.objects[object_id].id.split("_")[-1]
                    ) == int(
                        goal.object_id
                    ), f"Object_id doesn't correspond to id in semantic scene objects dictionary for episode: {episode}"

                    center = sem_scene.objects[object_id].aabb.center
                    x_len, _, z_len = (
                        sem_scene.objects[object_id].aabb.sizes / 2.0
                    )
                    # Nodes to draw rectangle
                    corners = [
                        center + np.array([x, 0, z])
                        for x, z in [
                            (-x_len, -z_len),
                            (-x_len, z_len),
                            (x_len, z_len),
                            (x_len, -z_len),
                            (-x_len, -z_len),
                        ]
                    ]

                    map_corners = [
                        maps.to_grid(
                            p[2],
                            p[0],
                            self._top_down_map.shape[0:2],
                            sim=self._sim,
                        )
                        for p in corners
                    ]

                    maps.draw_path(
                        self._top_down_map,
                        map_corners,
                        maps.MAP_TARGET_BOUNDING_BOX,
                        self.line_thickness,
                    )
                except AttributeError:
                    pass

    def _draw_shortest_path(
        self, episode: NavigationEpisode, agent_position: AgentState
    ):
        if self._config.DRAW_SHORTEST_PATH:
            _shortest_path_points = (
                self._sim.get_straight_shortest_path_points(
                    agent_position, episode.goals[0].position
                )
            )
            self._shortest_path_points = [
                maps.to_grid(
                    p[2], p[0], self._top_down_map.shape[0:2], sim=self._sim
                )
                for p in _shortest_path_points
            ]
            maps.draw_path(
                self._top_down_map,
                self._shortest_path_points,
                maps.MAP_SHORTEST_PATH_COLOR,
                self.line_thickness,
            )

    def reset_metric(self, episode, *args: Any, **kwargs: Any):
        self._step_count = 0
        self._metric = None
        self._top_down_map = self.get_original_map()
        agent_position = self._sim.get_agent_state().position
        a_x, a_y = maps.to_grid(
            agent_position[2],
            agent_position[0],
            self._top_down_map.shape[0:2],
            sim=self._sim,
        )
        self._previous_xy_location = (a_y, a_x)

        self.update_fog_of_war_mask(np.array([a_x, a_y]))

        # draw source and target parts last to avoid overlap
        self._draw_goals_view_points(episode)
        self._draw_goals_aabb(episode)
        self._draw_goals_positions(episode)

        self._draw_shortest_path(episode, agent_position)

        if self._config.DRAW_SOURCE:
            self._draw_point(
                episode.start_position, maps.MAP_SOURCE_POINT_INDICATOR
            )

    def update_metric(self, episode, action, *args: Any, **kwargs: Any):
        self._step_count += 1
        house_map, map_agent_x, map_agent_y = self.update_map(
            self._sim.get_agent_state().position
        )

        self._metric = {
            "map": house_map,
            "fog_of_war_mask": self._fog_of_war_mask,
            "agent_map_coord": (map_agent_x, map_agent_y),
            "agent_angle": self.get_polar_angle(),
        }

    def get_polar_angle(self):
        agent_state = self._sim.get_agent_state()
        # quaternion is in x, y, z, w format
        ref_rotation = agent_state.rotation

        heading_vector = quaternion_rotate_vector(
            ref_rotation.inverse(), np.array([0, 0, -1])
        )

        phi = cartesian_to_polar(-heading_vector[2], heading_vector[0])[1]
        z_neg_z_flip = np.pi
        return np.array(phi) + z_neg_z_flip

    def update_map(self, agent_position):
        a_x, a_y = maps.to_grid(
            agent_position[2],
            agent_position[0],
            self._top_down_map.shape[0:2],
            sim=self._sim,
        )
        # Don't draw over the source point
        if self._top_down_map[a_x, a_y] != maps.MAP_SOURCE_POINT_INDICATOR:
            color = 10 + min(
                self._step_count * 245 // self._config.MAX_EPISODE_STEPS, 245
            )

            thickness = self.line_thickness
            cv2.line(
                self._top_down_map,
                self._previous_xy_location,
                (a_y, a_x),
                color,
                thickness=thickness,
            )

        self.update_fog_of_war_mask(np.array([a_x, a_y]))

        self._previous_xy_location = (a_y, a_x)
        return self._top_down_map, a_x, a_y

    def update_fog_of_war_mask(self, agent_position):
        if self._config.FOG_OF_WAR.DRAW:
            self._fog_of_war_mask = fog_of_war.reveal_fog_of_war(
                self._top_down_map,
                self._fog_of_war_mask,
                agent_position,
                self.get_polar_angle(),
                fov=self._config.FOG_OF_WAR.FOV,
                max_line_len=self._config.FOG_OF_WAR.VISIBILITY_DIST
                / maps.calculate_meters_per_pixel(
                    self._map_resolution, sim=self._sim
                ),
            )


@registry.register_measure
class DistanceToGoal(Measure):
    """The measure calculates a distance towards the goal."""

    cls_uuid: str = "distance_to_goal"

    def __init__(
        self, sim: Simulator, config: Config, *args: Any, **kwargs: Any
    ):
        self._previous_position: Optional[Tuple[float, float, float]] = None
        self._sim = sim
        self._config = config
        self._episode_view_points: Optional[
            List[Tuple[float, float, float]]
        ] = None

        super().__init__(**kwargs)

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        return self.cls_uuid

    def reset_metric(self, episode, *args: Any, **kwargs: Any):
        self._previous_position = None
        self._metric = None
        if self._config.DISTANCE_TO == "VIEW_POINTS":
            self._episode_view_points = [
                view_point.agent_state.position
                for goal in episode.goals
                for view_point in goal.view_points
            ]
        self.update_metric(episode=episode, *args, **kwargs)  # type: ignore

    def _euclidean_distance(self, position_a, position_b):
        return np.linalg.norm(position_b - position_a, ord=2)

    def update_metric(
        self, episode: NavigationEpisode, *args: Any, **kwargs: Any
    ):
        current_position = self._sim.get_agent_state().position

        if self._previous_position is None or not np.allclose(
            self._previous_position, current_position, atol=1e-4
        ):
            if self._config.DISTANCE_TO == "POINT":
                if hasattr(episode, "is_thda") and episode.is_thda:
                    distance_to_target = np.inf
                    for goal in episode.goals:
                        distance_to_goal = self._euclidean_distance(
                            current_position, goal.position
                        )
                        distance_to_target = min(distance_to_target, distance_to_goal)
                else:
                    distance_to_target = self._sim.geodesic_distance(
                        current_position,
                        [goal.position for goal in episode.goals],
                        episode,
                    )
            elif self._config.DISTANCE_TO == "VIEW_POINTS":
                if hasattr(episode, "is_thda") and episode.is_thda:
                    distance_to_target = np.inf
                    for goal in episode.goals:
                        distance_to_goal = self._euclidean_distance(
                            current_position, goal.position
                        )
                        distance_to_target = min(distance_to_target, distance_to_goal)
                else:
                    distance_to_target = self._sim.geodesic_distance(
                        current_position, self._episode_view_points, episode
                    )
            else:
                logger.error(
                    f"Non valid DISTANCE_TO parameter was provided: {self._config.DISTANCE_TO}"
                )

            self._previous_position = current_position
            self._metric = distance_to_target


@registry.register_task_action
class MoveForwardAction(SimulatorTaskAction):
    name: str = "MOVE_FORWARD"

    def step(self, *args: Any, **kwargs: Any):
        r"""Update ``_metric``, this method is called from ``Env`` on each
        ``step``.
        """
        if "replay_data" in kwargs.keys():
            return self._sim.step_from_replay(
                HabitatSimActions.MOVE_FORWARD,
                replay_data=kwargs["replay_data"]
            )
        return self._sim.step(HabitatSimActions.MOVE_FORWARD)


@registry.register_task_action
class TurnLeftAction(SimulatorTaskAction):
    def step(self, *args: Any, **kwargs: Any):
        r"""Update ``_metric``, this method is called from ``Env`` on each
        ``step``.
        """
        if "replay_data" in kwargs.keys():
            return self._sim.step_from_replay(
                HabitatSimActions.TURN_LEFT,
                replay_data=kwargs["replay_data"]
            )
        return self._sim.step(HabitatSimActions.TURN_LEFT)


@registry.register_task_action
class TurnRightAction(SimulatorTaskAction):
    def step(self, *args: Any, **kwargs: Any):
        r"""Update ``_metric``, this method is called from ``Env`` on each
        ``step``.
        """
        if "replay_data" in kwargs.keys():
            return self._sim.step_from_replay(
                HabitatSimActions.TURN_RIGHT,
                replay_data=kwargs["replay_data"]
            )
        return self._sim.step(HabitatSimActions.TURN_RIGHT)


@registry.register_task_action
class StopAction(SimulatorTaskAction):
    name: str = "STOP"

    def reset(self, task: EmbodiedTask, *args: Any, **kwargs: Any):
        task.is_stop_called = False  # type: ignore

    def step(self, task: EmbodiedTask, *args: Any, **kwargs: Any):
        r"""Update ``_metric``, this method is called from ``Env`` on each
        ``step``.
        """
        task.is_stop_called = True  # type: ignore
        return self._sim.get_observations_at()  # type: ignore


@registry.register_task_action
class LookUpAction(SimulatorTaskAction):
    def step(self, *args: Any, **kwargs: Any):
        r"""Update ``_metric``, this method is called from ``Env`` on each
        ``step``.
        """
        if "replay_data" in kwargs.keys():
            return self._sim.step_from_replay(
                HabitatSimActions.LOOK_UP,
                replay_data=kwargs["replay_data"]
            )
        return self._sim.step(HabitatSimActions.LOOK_UP)


@registry.register_task_action
class LookDownAction(SimulatorTaskAction):
    def step(self, *args: Any, **kwargs: Any):
        r"""Update ``_metric``, this method is called from ``Env`` on each
        ``step``.
        """
        if "replay_data" in kwargs.keys():
            return self._sim.step_from_replay(
                HabitatSimActions.LOOK_DOWN,
                replay_data=kwargs["replay_data"]
            )
        return self._sim.step(HabitatSimActions.LOOK_DOWN)


@registry.register_task_action
class TeleportAction(SimulatorTaskAction):
    # TODO @maksymets: Propagate through Simulator class
    COORDINATE_EPSILON = 1e-6
    COORDINATE_MIN = -62.3241 - COORDINATE_EPSILON
    COORDINATE_MAX = 90.0399 + COORDINATE_EPSILON

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        return "TELEPORT"

    def step(
        self,
        *args: Any,
        position: List[float],
        rotation: List[float],
        **kwargs: Any,
    ):
        r"""Update ``_metric``, this method is called from ``Env`` on each
        ``step``.
        """

        if not isinstance(rotation, list):
            rotation = list(rotation)

        if not self._sim.is_navigable(position):
            return self._sim.get_observations_at()  # type: ignore

        return self._sim.get_observations_at(
            position=position, rotation=rotation, keep_agent_at_new_pose=True
        )

    @property
    def action_space(self) -> spaces.Dict:
        return spaces.Dict(
            {
                "position": spaces.Box(
                    low=np.array([self.COORDINATE_MIN] * 3),
                    high=np.array([self.COORDINATE_MAX] * 3),
                    dtype=np.float32,
                ),
                "rotation": spaces.Box(
                    low=np.array([-1.0, -1.0, -1.0, -1.0]),
                    high=np.array([1.0, 1.0, 1.0, 1.0]),
                    dtype=np.float32,
                ),
            }
        )


@registry.register_task(name="Nav-v0")
class NavigationTask(EmbodiedTask):
    def __init__(
        self, config: Config, sim: Simulator, dataset: Optional[Dataset] = None
    ) -> None:
        super().__init__(config=config, sim=sim, dataset=dataset)

    def overwrite_sim_config(self, sim_config: Any, episode: Episode) -> Any:
        return merge_sim_episode_config(sim_config, episode)

    def _check_episode_is_active(self, *args: Any, **kwargs: Any) -> bool:
        return not getattr(self, "is_stop_called", False)

habitat/tasks/nav/nav.py
cshizhe/onav_rim
# 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 typing import Any, List, Optional

import attr
from cv2 import log
import numpy as np
from gym import spaces

from habitat.config import Config
from habitat.core.dataset import SceneState
from habitat.core.logging import logger
from habitat.core.registry import registry
from habitat.core.simulator import AgentState, Sensor, SensorTypes
from habitat.core.utils import not_none_validator
from habitat.tasks.nav.nav import (
    NavigationEpisode,
    NavigationGoal,
    NavigationTask
)

try:
    from habitat.datasets.object_nav.object_nav_dataset import (
        ObjectNavDatasetV1,
    )
except ImportError:
    pass


task_cat2mpcat40 = [
    3,  # ('chair', 2, 0)
    5,  # ('table', 4, 1)
    6,  # ('picture', 5, 2)
    7,  # ('cabinet', 6, 3)
    8,  # ('cushion', 7, 4)
    10,  # ('sofa', 9, 5),
    11,  # ('bed', 10, 6)
    13,  # ('chest_of_drawers', 12, 7),
    14,  # ('plant', 13, 8)
    15,  # ('sink', 14, 9)
    18,  # ('toilet', 17, 10),
    19,  # ('stool', 18, 11),
    20,  # ('towel', 19, 12)
    22,  # ('tv_monitor', 21, 13)
    23,  # ('shower', 22, 14)
    25,  # ('bathtub', 24, 15)
    26,  # ('counter', 25, 16),
    27,  # ('fireplace', 26, 17),
    33,  # ('gym_equipment', 32, 18),
    34,  # ('seating', 33, 19),
    38,  # ('clothes', 37, 20),
    43,  # ('foodstuff', 42, 21),
    44,  # ('stationery', 43, 22),
    45,  # ('fruit', 44, 23),
    46,  # ('plaything', 45, 24),
    47,  # ('hand_tool', 46, 25),
    48,  # ('game_equipment', 47, 26),
    49,  # ('kitchenware', 48, 27)
]

mapping_mpcat40_to_goal21 = {
    3: 1,
    5: 2,
    6: 3,
    7: 4,
    8: 5,
    10: 6,
    11: 7,
    13: 8,
    14: 9,
    15: 10,
    18: 11,
    19: 12,
    20: 13,
    22: 14,
    23: 15,
    25: 16,
    26: 17,
    27: 18,
    33: 19,
    34: 20,
    38: 21,
    43: 22,  #  ('foodstuff', 42, task_cat: 21)
    44: 28,  #  ('stationery', 43, task_cat: 22)
    45: 26,  #  ('fruit', 44, task_cat: 23)
    46: 25,  #  ('plaything', 45, task_cat: 24)
    47: 24,  # ('hand_tool', 46, task_cat: 25)
    48: 23,  # ('game_equipment', 47, task_cat: 26)
    49: 27,  # ('kitchenware', 48, task_cat: 27)
}

mapping_mpcat40_to_goal20 = {
    3: 0,
    5: 1,
    6: 2,
    7: 3,
    8: 4,
    10: 5,
    11: 6,
    13: 7,
    14: 8,
    15: 9,
    18: 10,
    19: 11,
    20: 12,
    22: 13,
    23: 14,
    25: 15,
    26: 16,
    27: 17,
    33: 18,
    34: 19,
    38: 20,
    43: 21,  #  ('foodstuff', 42, task_cat: 21)
    44: 22,  #  ('stationery', 43, task_cat: 22)
    45: 23,  #  ('fruit', 44, task_cat: 23)
    46: 24,  #  ('plaything', 45, task_cat: 24)
    47: 25,  # ('hand_tool', 46, task_cat: 25)
    48: 26,  # ('game_equipment', 47, task_cat: 26)
    49: 27,  # ('kitchenware', 48, task_cat: 27)
}


@attr.s(auto_attribs=True, kw_only=True)
class AgentStateSpec:
    r"""Agent data specifications that capture states of agent and sensor in replay state.
    """
    position: Optional[List[float]] = attr.ib(default=None)
    rotation: Optional[List[float]] = attr.ib(default=None)
    sensor_data: Optional[dict] = attr.ib(default=None)


@attr.s(auto_attribs=True, kw_only=True)
class ReplayActionSpec:
    r"""Replay specifications that capture metadata associated with action.
    """
    action: str = attr.ib(default=None, validator=not_none_validator)
    agent_state: Optional[AgentStateSpec] = attr.ib(default=None)


@attr.s(auto_attribs=True, kw_only=True)
class ObjectGoalNavEpisode(NavigationEpisode):
    r"""ObjectGoal Navigation Episode

    :param object_category: Category of the obect
    """
    object_category: Optional[str] = None
    reference_replay: Optional[List[ReplayActionSpec]] = None
    scene_state: Optional[List[SceneState]] = None
    is_thda: Optional[bool] = False
    scene_dataset: Optional[str] = "mp3d"

    @property
    def goals_key(self) -> str:
        r"""The key to retrieve the goals"""
        return f"{os.path.basename(self.scene_id)}_{self.object_category}"


@attr.s(auto_attribs=True)
class ObjectViewLocation:
    r"""ObjectViewLocation provides information about a position around an object goal
    usually that is navigable and the object is visible with specific agent
    configuration that episode's dataset was created.
     that is target for
    navigation. That can be specify object_id, position and object
    category. An important part for metrics calculation are view points that
     describe success area for the navigation.

    Args:
        agent_state: navigable AgentState with a position and a rotation where
        the object is visible.
        iou: an intersection of a union of the object and a rectangle in the
        center of view. This metric is used to evaluate how good is the object
        view form current position. Higher iou means better view, iou equals
        1.0 if whole object is inside of the rectangle and no pixel inside
        the rectangle belongs to anything except the object.
    """
    agent_state: AgentState
    iou: Optional[float]


@attr.s(auto_attribs=True, kw_only=True)
class ObjectGoal(NavigationGoal):
    r"""Object goal provides information about an object that is target for
    navigation. That can be specify object_id, position and object
    category. An important part for metrics calculation are view points that
     describe success area for the navigation.

    Args:
        object_id: id that can be used to retrieve object from the semantic
        scene annotation
        object_name: name of the object
        object_category: object category name usually similar to scene semantic
        categories
        room_id: id of a room where object is located, can be used to retrieve
        room from the semantic scene annotation
        room_name: name of the room, where object is located
        view_points: navigable positions around the object with specified
        proximity of the object surface used for navigation metrics calculation.
        The object is visible from these positions.
    """

    object_id: str = attr.ib(default=None, validator=not_none_validator)
    object_name: Optional[str] = None
    object_name_id: Optional[int] = None
    object_category: Optional[str] = None
    room_id: Optional[str] = None
    room_name: Optional[str] = None
    view_points: Optional[List[ObjectViewLocation]] = None


@registry.register_sensor
class ObjectGoalSensor(Sensor):
    r"""A sensor for Object Goal specification as observations which is used in
    ObjectGoal Navigation. The goal is expected to be specified by object_id or
    semantic category id.
    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: a reference to the simulator for calculating task observations.
        config: a config for the ObjectGoalSensor sensor. Can contain field
            GOAL_SPEC that specifies which id use for goal specification,
            GOAL_SPEC_MAX_VAL the maximum object_id possible used for
            observation space definition.
        dataset: a Object Goal navigation dataset that contains dictionaries
        of categories id to text mapping.
    """
    cls_uuid: str = "objectgoal"

    def __init__(
        self,
        sim,
        config: Config,
        dataset: "ObjectNavDatasetV1",
        *args: Any,
        **kwargs: Any,
    ):
        self._sim = sim
        self._dataset = dataset
        super().__init__(config=config)

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        return self.cls_uuid

    def _get_sensor_type(self, *args: Any, **kwargs: Any):
        return SensorTypes.SEMANTIC

    def _get_observation_space(self, *args: Any, **kwargs: Any):
        sensor_shape = (1,)
        max_value = self.config.GOAL_SPEC_MAX_VAL - 1
        if self.config.GOAL_SPEC == "TASK_CATEGORY_ID":
            max_value = max(
                self._dataset.category_to_task_category_id.values()
            )
            logger.info("max object cat: {}".format(max_value))
            logger.info("cats: {}".format(self._dataset.category_to_task_category_id.values()))

        return spaces.Box(
            low=0, high=max_value, shape=sensor_shape, dtype=np.int64
        )

    def get_observation(
        self,
        observations,
        *args: Any,
        episode: ObjectGoalNavEpisode,
        **kwargs: Any,
    ) -> Optional[int]:

        if len(episode.goals) == 0:
            logger.error(
                f"No goal specified for episode {episode.episode_id}."
            )
            return None
        if not isinstance(episode.goals[0], ObjectGoal):
            logger.error(
                f"First goal should be ObjectGoal, episode {episode.episode_id}."
            )
            return None
        category_name = episode.object_category
        if self.config.GOAL_SPEC == "TASK_CATEGORY_ID":
            return np.array(
                [self._dataset.category_to_task_category_id[category_name]],
                dtype=np.int64,
            )
        elif self.config.GOAL_SPEC == "OBJECT_ID":
            obj_goal = episode.goals[0]
            assert isinstance(obj_goal, ObjectGoal)  # for type checking
            return np.array([obj_goal.object_name_id], dtype=np.int64)
        else:
            raise RuntimeError(
                "Wrong GOAL_SPEC specified for ObjectGoalSensor."
            )


@registry.register_task(name="ObjectNav-v1")
class ObjectNavigationTask(NavigationTask):
    r"""An Object Navigation Task class for a task specific methods.
    Used to explicitly state a type of the task in config.
    """
    _is_episode_active: bool
    _prev_action: int

    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)
        self._is_episode_active = False

    def overwrite_sim_config(self, sim_config, episode):
        super().overwrite_sim_config(sim_config, episode)

        sim_config.defrost()
        sim_config.scene_state = getattr(episode, 'scene_state', None)
        sim_config.freeze()
        
        return sim_config

    def _check_episode_is_active(self,  action, *args: Any, **kwargs: Any) -> bool:
        return not getattr(self, "is_stop_called", False)

habitat/tasks/nav/object_nav_task.py
cshizhe/onav_rim
#!/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 warnings
from typing import Optional, Union

import numpy as np

import habitat_sim
from habitat.sims.habitat_simulator.actions import HabitatSimActions
from habitat.sims.habitat_simulator.habitat_simulator import HabitatSim


def action_to_one_hot(action: int) -> np.array:
    one_hot = np.zeros(len(HabitatSimActions), dtype=np.float32)
    one_hot[action] = 1
    return one_hot


class ShortestPathFollower:
    r"""Utility class for extracting the action on the shortest path to the
        goal.

    :param sim: HabitatSim instance.
    :param goal_radius: Distance between the agent and the goal for it to be
            considered successful.
    :param return_one_hot: If true, returns a one-hot encoding of the action
            (useful for training ML agents). If false, returns the
            SimulatorAction.
    :param stop_on_error: Return stop if the follower is unable to determine a
                          suitable action to take next.  If false, will raise
                          a habitat_sim.errors.GreedyFollowerError instead
    """

    def __init__(
        self,
        sim: HabitatSim,
        goal_radius: float,
        return_one_hot: bool = True,
        stop_on_error: bool = True,
    ):

        self._return_one_hot = return_one_hot
        self._sim = sim
        self._goal_radius = goal_radius
        self._follower: Optional[habitat_sim.GreedyGeodesicFollower] = None
        self._current_scene = None
        self._stop_on_error = stop_on_error

    def _build_follower(self):
        if self._current_scene != self._sim.habitat_config.SCENE:
            self._follower = self._sim.make_greedy_follower(
                0,
                self._goal_radius,
                stop_key=HabitatSimActions.STOP,
                forward_key=HabitatSimActions.MOVE_FORWARD,
                left_key=HabitatSimActions.TURN_LEFT,
                right_key=HabitatSimActions.TURN_RIGHT,
            )
            self._current_scene = self._sim.habitat_config.SCENE

    def _get_return_value(self, action) -> Union[int, np.array]:
        if self._return_one_hot:
            return action_to_one_hot(action)
        else:
            return action

    def get_next_action(
        self, goal_pos: np.array
    ) -> Optional[Union[int, np.array]]:
        """Returns the next action along the shortest path."""
        self._build_follower()
        assert self._follower is not None
        try:
            next_action = self._follower.next_action_along(goal_pos)
            # print('follower: {}'.format(next_action))
        except habitat_sim.errors.GreedyFollowerError as e:
            if self._stop_on_error:
                next_action = HabitatSimActions.STOP
            else:
                raise e

        return self._get_return_value(next_action)

    def find_path(
        self, goal_pos: np.ndarray
    ):
        self._build_follower()
        assert self._follower is not None
        try:
            actions = self._follower.find_path(goal_pos)
        except habitat_sim.errors.GreedyFollowerError as e:
            if self._stop_on_error:
                actions = None
            else:
                raise e
        return actions

    @property
    def mode(self):
        warnings.warn(".mode is depricated", DeprecationWarning)
        return ""

    @mode.setter
    def mode(self, new_mode: str):
        warnings.warn(".mode is depricated", DeprecationWarning)

habitat/tasks/nav/shortest_path_follower.py
cshizhe/onav_rim
#!/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.

from habitat.core.embodied_task import EmbodiedTask
from habitat.core.registry import registry


def _try_register_pickplace_task():
    try:
        from habitat.tasks.pickplace.pickplace import PickPlaceTask  # noqa
    except ImportError as e:
        rearrangement_task_import_error = e

        @registry.register_task(name="PickPlaceTask-v0")
        class PickPlaceTaskImportError(EmbodiedTask):
            def __init__(self, *args, **kwargs):
                raise rearrangement_task_import_error

habitat/tasks/pickplace/__init__.py
cshizhe/onav_rim
#!/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.

from typing import Any, Dict, List, Optional, Type

import attr
import habitat_sim
import math
import numpy as np
from gym import spaces

from collections import defaultdict
from habitat.core.logging import logger
from habitat.config import Config
from habitat.core.embodied_task import EmbodiedTask
from habitat.core.registry import registry
from habitat.core.simulator import Observations, Sensor, Simulator, SensorTypes
from habitat.core.utils import not_none_validator
from habitat.tasks.nav.nav import (
    merge_sim_episode_config,
    DistanceToGoal,
    TopDownMap,
    EpisodicGPSSensor,
    PointGoalSensor
)
from habitat.core.dataset import Dataset, Episode
from habitat.core.embodied_task import SimulatorTaskAction, Measure
from habitat.sims.habitat_simulator.actions import (
    HabitatSimActions,
    HabitatSimV1ActionSpaceConfiguration,
)
from habitat.tasks.utils import get_habitat_sim_action, get_habitat_sim_action_str
from habitat.sims.habitat_simulator.habitat_simulator import HabitatSim


task_cat2mpcat40 = [
    3,  # ('chair', 2, 0)
    5,  # ('table', 4, 1)
    6,  # ('picture', 5, 2)
    7,  # ('cabinet', 6, 3)
    8,  # ('cushion', 7, 4)
    10,  # ('sofa', 9, 5),
    11,  # ('bed', 10, 6)
    13,  # ('chest_of_drawers', 12, 7),
    14,  # ('plant', 13, 8)
    15,  # ('sink', 14, 9)
    18,  # ('toilet', 17, 10),
    19,  # ('stool', 18, 11),
    20,  # ('towel', 19, 12)
    22,  # ('tv_monitor', 21, 13)
    23,  # ('shower', 22, 14)
    25,  # ('bathtub', 24, 15)
    26,  # ('counter', 25, 16),
    27,  # ('fireplace', 26, 17),
    33,  # ('gym_equipment', 32, 18),
    34,  # ('seating', 33, 19),
    38,  # ('clothes', 37, 20),
    43,  # ('foodstuff', 42, 21),
    44,  # ('stationery', 43, 22),
    45,  # ('fruit', 44, 23),
    46,  # ('plaything', 45, 24),
    47,  # ('hand_tool', 46, 25),
    48,  # ('game_equipment', 47, 26),
    49,  # ('kitchenware', 48, 27)
]

@attr.s(auto_attribs=True)
class InstructionData:
    instruction_text: str
    instruction_tokens: List[int]


@attr.s(auto_attribs=True, kw_only=True)
class RearrangementSpec:
    r"""Specifications that capture a particular position of final position
    or initial position of the object.
    """

    position: List[float] = attr.ib(default=None, validator=not_none_validator)
    rotation: List[float] = attr.ib(default=None, validator=not_none_validator)
    info: Optional[Dict[str, str]] = attr.ib(default=None)


@attr.s(auto_attribs=True, kw_only=True)
class RearrangementObjectSpec(RearrangementSpec):
    r"""Object specifications that capture position of each object in the scene,
    the associated object template.
    """
    object_id: str = attr.ib(default=None, validator=not_none_validator)
    semantic_object_id: str = attr.ib(default=None, validator=not_none_validator)
    object_handle: Optional[str] = attr.ib(
        default="", validator=not_none_validator
    )
    object_template: Optional[str] = attr.ib(
        default="", validator=not_none_validator
    )
    object_icon: Optional[str] = attr.ib(
        default="", validator=not_none_validator
    )
    motion_type: Optional[str] = attr.ib(default=None)
    is_receptacle: Optional[bool] = attr.ib(default=None)


@attr.s(auto_attribs=True, kw_only=True)
class GrabReleaseActionSpec:
    r"""Grab/Release action reaply data specifications that capture states
     of each grab/release action.
    """
    released_object_position: Optional[List[float]] = attr.ib(default=None)
    released_object_id: Optional[int] = attr.ib(default=None)
    released_object_handle: Optional[str] = attr.ib(default=None)
    grab_object_id: Optional[int] = attr.ib(default=None)


@attr.s(auto_attribs=True, kw_only=True)
class ObjectStateSpec:
    r"""Object data specifications that capture states of each object in replay state.
    """
    object_id: Optional[int] = attr.ib(default=None)
    translation: Optional[List[float]] = attr.ib(default=None)
    rotation: Optional[List[float]] = attr.ib(default=None)
    motion_type: Optional[str] = attr.ib(default=None)
    object_handle: Optional[str] = attr.ib(default=None)


@attr.s(auto_attribs=True, kw_only=True)
class AgentStateSpec:
    r"""Agent data specifications that capture states of agent and sensor in replay state.
    """
    position: Optional[List[float]] = attr.ib(default=None)
    rotation: Optional[List[float]] = attr.ib(default=None)
    sensor_data: Optional[dict] = attr.ib(default=None)


@attr.s(auto_attribs=True, kw_only=True)
class ReplayActionSpec:
    r"""Replay specifications that capture metadata associated with action.
    """
    action: str = attr.ib(default=None, validator=not_none_validator)
    object_under_cross_hair: Optional[int] = attr.ib(default=None)
    action_data: Optional[GrabReleaseActionSpec] = attr.ib(default=None)
    is_grab_action: Optional[bool] = attr.ib(default=None)
    is_release_action: Optional[bool] = attr.ib(default=None)
    object_states: Optional[List[ObjectStateSpec]] = attr.ib(default=None)
    agent_state: Optional[AgentStateSpec] = attr.ib(default=None)
    collision: Optional[dict] = attr.ib(default=None)
    gripped_object_id: Optional[int] = attr.ib(default=None)


@attr.s(auto_attribs=True, kw_only=True)
class RearrangementEpisode(Episode):
    r"""Specification of episode that includes initial position and rotation
    of agent, goal specifications, instruction specifications, reference path,
    and optional shortest paths.

    Args:
        episode_id: id of episode in the dataset
        scene_id: id of scene inside the simulator.
        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.
        instruction: single natural language instruction for the task.
        reference_replay: List of keypresses which gives the reference
            actions to the goal that aligns with the instruction.
    """
    goals: List[RearrangementSpec] = attr.ib(
        default=None, validator=not_none_validator
    )
    reference_replay: List[Dict] = attr.ib(
        default=None, validator=not_none_validator
    )
    instruction: InstructionData = attr.ib(
        default=None, validator=not_none_validator
    )
    objects: List[RearrangementObjectSpec] = attr.ib(
        default=None, validator=not_none_validator
    )


@registry.register_sensor(name="InstructionSensor")
class InstructionSensor(Sensor):
    def __init__(self, **kwargs):
        self.uuid = "instruction"
        super().__init__()

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        return self.uuid

    def _get_observation_space(self, *args: Any, **kwargs: Any):
        return spaces.Box(
            low=0,
            high=66,
            shape=(11,),
            dtype=np.int64,
        )

    def _get_sensor_type(self, *args:Any, **kwargs: Any):
        return SensorTypes.TOKEN_IDS

    def _get_observation(
        self,
        observations: Dict[str, Observations],
        episode: RearrangementEpisode,
        **kwargs
    ):
        return episode.instruction.instruction_tokens

    def get_observation(self, **kwargs):
        return self._get_observation(**kwargs)


@registry.register_sensor(name="DemonstrationSensor")
class DemonstrationSensor(Sensor):
    def __init__(self, **kwargs):
        self.uuid = "demonstration"
        self.observation_space = spaces.Discrete(1)
        self.timestep = 0
        self.prev_action = 0

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        return self.uuid

    def _get_observation(
        self,
        observations: Dict[str, Observations],
        episode,
        task: EmbodiedTask,
        **kwargs
    ):
        # Fetch next action as observation
        if task.is_resetting:  # reset
            self.timestep = 1
        
        if self.timestep < len(episode.reference_replay):
            action_name = episode.reference_replay[self.timestep].action
            action = get_habitat_sim_action(action_name)
        else:
            action = 0

        self.timestep += 1
        return action

    def get_observation(self, **kwargs):
        return self._get_observation(**kwargs)


@registry.register_sensor(name="InflectionWeightSensor")
class InflectionWeightSensor(Sensor):
    def __init__(self, config: Config, **kwargs):
        self.uuid = "inflection_weight"
        self.observation_space = spaces.Discrete(1)
        self._config = config
        self.timestep = 0

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        return self.uuid

    def _get_observation(
        self,
        observations: Dict[str, Observations],
        episode,
        task: EmbodiedTask,
        **kwargs
    ):
        if task.is_resetting:  # reset
            self.timestep = 0
        
        inflection_weight = 1.0
        if self.timestep == 0:
            inflection_weight = 1.0
        elif self.timestep >= len(episode.reference_replay):
            inflection_weight = 1.0 
        elif episode.reference_replay[self.timestep - 1].action != episode.reference_replay[self.timestep].action:
            inflection_weight = self._config.INFLECTION_COEF
        self.timestep += 1
        return inflection_weight

    def get_observation(self, **kwargs):
        return self._get_observation(**kwargs)


@registry.register_sensor(name="GrippedObjectSensor")
class GrippedObjectSensor(Sensor):
    def __init__(self, *args, sim: HabitatSim, config: Config, **kwargs):
        self._sim = sim
        self.uuid = "gripped_object_id"
        super().__init__(config=config)

    def _get_observation_space(self, *args: Any, **kwargs: Any):
        return spaces.Discrete(
            len(self._sim.get_existing_object_ids())
        )

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        return self.uuid

    def _get_sensor_type(self, *args:Any, **kwargs: Any):
        return SensorTypes.MEASUREMENT

    def _get_observation(
        self,
        observations: Dict[str, Observations],
        episode: RearrangementEpisode,
        *args: Any,
        **kwargs
    ):
        return self._sim.gripped_object_id

    def get_observation(self, **kwargs):
        return self._get_observation(**kwargs)


@registry.register_measure
class ObjectToReceptacleDistance(Measure):
    """The measure calculates distance of object towards the goal."""

    cls_uuid: str = "object_receptacle_distance"

    def __init__(
        self, sim: Simulator, config: Config, *args: Any, **kwargs: Any
    ):
        self._sim = sim
        self._config = config

        super().__init__(**kwargs)

    @staticmethod
    def _get_uuid(*args: Any, **kwargs: Any):
        return ObjectToReceptacleDistance.cls_uuid

    def reset_metric(self, episode, *args: Any, **kwargs: Any):
        self.update_metric(*args, episode=episode, **kwargs)

    def _geo_dist(self, src_pos, goal_pos: np.array) -> float:
        return self._sim.geodesic_distance(src_pos, [goal_pos])

    def _euclidean_distance(self, position_a, position_b):
        return np.linalg.norm(
            np.array(position_b) - np.array(position_a), ord=2
        )

    def get_navigable_position(self, position):
        if not self._sim.pathfinder.is_navigable(position):
            position = self._sim.pathfinder.snap_point(position)
            position = np.array(position).tolist()
        return position

    def update_metric(self, episode, *args: Any, **kwargs: Any):
        object_ids = self._sim.get_existing_object_ids()
        obj_id = -1
        receptacle_id = -1
        for object_id in object_ids:
            scene_object = self._sim.get_object_from_scene(object_id)
            if scene_object.is_receptacle == False:
                obj_id = scene_object.object_id
            else:
                receptacle_id = scene_object.object_id
        was_nan = False
        object_position = np.array([0, 0, 0])
        if receptacle_id == -1:
            self._metric = 0.0
        elif obj_id != -1:
            object_position = np.array(
                self._sim.get_translation(obj_id)
            ).tolist()

            receptacle_position = np.array(
                self._sim.get_translation(receptacle_id)
            ).tolist()

            object_position = self.get_navigable_position(object_position)
            receptacle_position = self.get_navigable_position(receptacle_position)

            self._metric = self._geo_dist(
                object_position, receptacle_position
            )
            if self._metric == np.inf or self._metric == np.nan:
                was_nan = True
                self._metric = 2.0
        else:
            receptacle_position = np.array(
                self._sim.get_translation(receptacle_id)
            ).tolist()
            receptacle_position = self.get_navigable_position(receptacle_position)

            agent_state = self._sim.get_agent_state()
            agent_position = agent_state.position

            self._metric = self._geo_dist(
                agent_position, receptacle_position
            )
            if self._metric == np.inf or self._metric == np.nan:
                was_nan = True
                self._metric = 2.0
        self._metric = {
            "was_nan": was_nan,
            "metric": self._metric
        }
        # if was_nan:
        #     logger.error("object receptacle distance is nan: {} -- {} - {} -- {}".format(receptacle_id, obj_id, object_position, episode.episode_id))


@registry.register_measure
class AgentToObjectDistance(Measure):
    """The measure calculates the distance of objects from the agent"""

    cls_uuid: str = "agent_object_distance"

    def __init__(
        self, sim: Simulator, config: Config, *args: Any, **kwargs: Any
    ):
        self._sim = sim
        self._config = config

        super().__init__(**kwargs)

    @staticmethod
    def _get_uuid(*args: Any, **kwargs: Any):
        return AgentToObjectDistance.cls_uuid

    def reset_metric(self, episode, *args: Any, **kwargs: Any):
        self.update_metric(*args, episode=episode, **kwargs)

    def _euclidean_distance(self, position_a, position_b):
        return np.linalg.norm(
            np.array(position_b) - np.array(position_a), ord=2
        )

    def _geo_dist(self, src_pos, object_pos: np.array) -> float:
        return self._sim.geodesic_distance(src_pos, [object_pos])

    def get_navigable_position(self, position):
        if not self._sim.pathfinder.is_navigable(position):
            position = self._sim.pathfinder.snap_point(position)
            position = np.array(position).tolist()
        return position

    def update_metric(self, episode, *args: Any, **kwargs: Any):
        object_ids = self._sim.get_existing_object_ids()

        sim_obj_id = -1
        for object_id in object_ids:
            scene_object = self._sim.get_object_from_scene(object_id)
            if scene_object.is_receptacle == False:
                sim_obj_id = scene_object.object_id
        was_nan = False
        object_position = np.array([0, 0, 0])
        if sim_obj_id != -1:
            object_position = np.array(
                self._sim.get_translation(sim_obj_id)
            ).tolist()

            object_position = self.get_navigable_position(object_position)

            agent_state = self._sim.get_agent_state()
            agent_position = agent_state.position

            self._metric = self._geo_dist(
                agent_position, object_position
            )
            if self._metric == np.inf or self._metric == np.nan:
                was_nan = True
                self._metric = 2.0
        else:
            self._metric = 0.0
        self._metric = {
            "was_nan": was_nan,
            "metric": self._metric
        }
        if was_nan:
            logger.error("Agent object distance is inf: {} -- {} -- {}".format(sim_obj_id, object_position, episode.episode_id))


@registry.register_measure
class AgentToReceptacleDistance(Measure):
    """The measure calculates the distance of receptacle from the agent"""

    cls_uuid: str = "agent_receptacle_distance"

    def __init__(
        self, sim: Simulator, config: Config, *args: Any, **kwargs: Any
    ):
        self._sim = sim
        self._config = config

        super().__init__(**kwargs)

    @staticmethod
    def _get_uuid(*args: Any, **kwargs: Any):
        return AgentToReceptacleDistance.cls_uuid

    def reset_metric(self, episode, *args: Any, **kwargs: Any):
        self.update_metric(*args, episode=episode, **kwargs)

    def _euclidean_distance(self, position_a, position_b):
        return np.linalg.norm(
            np.array(position_b) - np.array(position_a), ord=2
        )

    def _geo_dist(self, src_pos, object_pos: np.array) -> float:
        return self._sim.geodesic_distance(src_pos, [object_pos])

    def get_navigable_position(self, position):
        if not self._sim.pathfinder.is_navigable(position):
            position = self._sim.pathfinder.snap_point(position)
            position = np.array(position).tolist()
        return position

    def update_metric(self, episode, *args: Any, **kwargs: Any):
        object_ids = self._sim.get_existing_object_ids()

        sim_obj_id = -1
        for object_id in object_ids:
            scene_object = self._sim.get_object_from_scene(object_id)
            if scene_object.is_receptacle == True:
                sim_obj_id = scene_object.object_id
        was_nan = False
        if sim_obj_id != -1:
            receptacle_position = np.array(
                self._sim.get_translation(sim_obj_id)
            ).tolist()

            receptacle_position = self.get_navigable_position(receptacle_position)

            agent_state = self._sim.get_agent_state()
            agent_position = agent_state.position

            self._metric = self._geo_dist(
                agent_position, receptacle_position 
            )
            if self._metric == np.inf or self._metric == np.nan:
                was_nan = True
                self._metric = 2.0
        else:
            self._metric = 0.0
        self._metric = {
            "was_nan": was_nan,
            "metric": self._metric
        }


@registry.register_measure
class GoalObjectVisible(Measure):
    r"""GoalObjectVisible"""
    cls_uuid = "goal_vis_pixels"

    def __init__(self, sim: Simulator, config: Config, *args: Any, **kwargs: Any):
        self._sim = sim
        self._config = config
        self.task_cat2mpcat40 = task_cat2mpcat40

        super().__init__()

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        return self.cls_uuid

    def reset_metric(self, episode, task, *args: Any, **kwargs: Any):
        self.update_metric(
            episode=episode, task=task, *args, **kwargs
        )

    def update_metric(
        self,
        episode,
        task: EmbodiedTask,
        observations,
        *args: Any,
        **kwargs: Any,
    ):
        self._metric = 0
        if "semantic" in observations:
            semantic_obs = observations["semantic"]

            if self._config.INSERTED_OBJECTS:
                object_ids = self._sim.get_existing_object_ids()
                goal_visible_pixels = 0
                semantic_object_id = self._sim.obj_id_to_semantic_obj_id_map[0]
                # If object is gripped caclulate visible pixels for receptacle
                if observations["gripped_object_id"] != -1:
                    semantic_object_id = self._sim.obj_id_to_semantic_obj_id_map[1]
                
                goal_visible_pixels += (semantic_obs == semantic_object_id).sum() # Sum over all since we're not batched
                goal_visible_area = goal_visible_pixels / semantic_obs.size
                self._metric = goal_visible_area
            else:
                # permute tensor to dimension [CHANNEL x HEIGHT X WIDTH]
                idx = self.task_cat2mpcat40[
                    observations["objectgoal"][0]
                ]  # task._dataset.category_to_task_category_id[episode.object_category], task._dataset.category_to_scene_annotation_category_id[episode.object_category], observations["objectgoal"][0]

                goal_visible_pixels = (semantic_obs == idx).sum() # Sum over all since we're not batched
                goal_visible_area = goal_visible_pixels / semantic_obs.size
                self._metric = goal_visible_area


@registry.register_measure
class Coverage(Measure):
    """Coverage
    Number of visited squares in a gridded environment
    - this is not exactly what we want to reward, but a decent starting point
    - the semantic coverage is agent-based i.e. we should reward new visuals/objects discovered
    - Note, internal grid has origin at bottom left
    EGOCENTRIC -- center + align the coverage reward.
    """
    cls_uuid: str = "coverage"

    def __init__(self, sim, config, **kwargs: Any):
        self._sim = sim
        self._config = config

        self._visited = None  # number of visits
        self._mini_visited = None
        self._step = None
        self._reached_count = None
        self._mini_reached = None
        self._mini_delta = 0.5
        self._grid_delta = config.GRID_DELTA
        super().__init__()

    def _get_uuid(self, *args: Any, **kwargs: Any):
        return self.cls_uuid

    def _to_grid(self, delta, sim_x, sim_y, sim_z=0):
        # ! Note we actually get sim_x, sim_z in 2D case but that doesn't affect code
        grid_x = int((sim_x) / delta)
        grid_y = int((sim_y) / delta)
        grid_z = int((sim_z) / delta)
        return grid_x, grid_y, grid_z

    def reset_metric(self, episode, task, observations, *args: Any, **kwargs: Any):
        self._visited = {}
        self._mini_visited = {}
        self._reached_count = 0
        # Used for coverage prediction (not elegant, I know)
        self._mini_reached = 0
        self._step = 0  # Tracking episode length
        current_visit = self._visit(task, observations)
        self._metric = {
            "reached": self._reached_count,
            "mini_reached": self._mini_reached,
            "visit_count": current_visit,
            "step": self._step
        }

    def _visit(self, task, observations):
        """ Returns number of times visited current square """
        self._step += 1
        if self._config.EGOCENTRIC:
            global_loc = observations[EpisodicGPSSensor.cls_uuid]
        else:
            global_loc = self._sim.get_agent_state().position.tolist()

        mini_loc = self._to_grid(self._mini_delta, *global_loc)
        if mini_loc in self._mini_visited:
            self._mini_visited[mini_loc] += 1
        else:
            self._mini_visited[mini_loc] = 1
            self._mini_reached += 1

        grid_loc = self._to_grid(self._grid_delta, *global_loc)
        if grid_loc in self._visited:
            self._visited[grid_loc] += 1
            return self._visited[grid_loc]
        self._visited[grid_loc] = 1
        self._reached_count += 1
        return self._visited[grid_loc]

    def update_metric(
        self, episode, action, task: EmbodiedTask, observations, *args: Any, **kwargs: Any
    ):
        current_visit = self._visit(task, observations)
        self._metric = {
            "reached": self._reached_count,
            "mini_reached": self._mini_reached,
            "visit_count": current_visit,
            "step": self._step
        }


@registry.register_measure
class CoverageExplorationReward(Measure):
    # Parallels ExploreRLEnv in `environments.py`

    cls_uuid: str = "coverage_explore_reward"

    def __init__(self, sim, config, **kwargs: Any):
        self._sim = sim
        self._config = config
        self._attenuation_penalty = 1.0
        super().__init__()

    def _get_uuid(self, *args: Any, **kwargs: Any):
        return self.cls_uuid

    def reset_metric(self, episode, task, *args: Any, **kwargs: Any):
        task.measurements.check_measure_dependencies(
            self.cls_uuid,
            [
                Coverage.cls_uuid,
            ],
        )
        self._attenuation_penalty = 1.0
        self._metric = 0
        self.update_metric(episode=episode, task=task, *args, **kwargs)

    def update_metric(
        self, episode, task: EmbodiedTask, *args: Any, **kwargs: Any
    ):
        self._attenuation_penalty *= self._config.ATTENUATION
        visit = task.measurements.measures[
            Coverage.cls_uuid
        ].get_metric()["visit_count"]
        self._metric = self._attenuation_penalty * self._config.REWARD / (visit ** self._config.VISIT_EXP)


@registry.register_measure
class ExploreThenNavReward(CoverageExplorationReward):
    # Stop count-based coverage once goal is seen

    cls_uuid: str = "explore_then_nav_reward"
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._goal_was_seen = False
        self._previous_measure = None

    def reset_metric(self, episode, task, *args, **kwargs):
        task.measurements.check_measure_dependencies(
            self.cls_uuid,
            [
                Coverage.cls_uuid,
                GoalObjectVisible.cls_uuid,
                DistanceToGoal.cls_uuid,
            ]
        )
        self._goal_was_seen = False
        self._previous_measure = None
        super().reset_metric(episode, task, *args, **kwargs)

    def update_metric(
        self, episode, task: EmbodiedTask, *args: Any, **kwargs: Any
    ):
        if self._goal_was_seen:
            measure = task.measurements.measures[DistanceToGoal.cls_uuid].get_metric()
            self._metric = self._previous_measure - measure
            return
        goal_vis = task.measurements.measures[GoalObjectVisible.cls_uuid].get_metric()
        if goal_vis > self._config.EXPLORE_GOAL_SEEN_THRESHOLD:
            self._goal_was_seen = True
            self._previous_measure = task.measurements.measures[DistanceToGoal.cls_uuid].get_metric()
        super().update_metric(*args, episode, task, **kwargs)


@registry.register_measure
class ReleaseFailed(Measure):
    r"""Grab Success - whether an object was grabbed during episode or not
    """

    def __init__(
        self, sim: Simulator, config: Config, *args: Any, **kwargs: Any
    ):
        self._sim = sim
        self._config = config
        self.prev_gripped_object_id = -1

        super().__init__()

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        return "release_failed"

    def reset_metric(self, episode, task, observations, *args: Any, **kwargs: Any):
        self._metric = 0
        self.prev_gripped_object_id = -1
        self.update_metric(episode=episode, task=task, observations=observations, action={"action": 0}, *args, **kwargs)  # type: ignore

    def update_metric(
        self, episode, task: EmbodiedTask, observations, action, *args: Any, **kwargs: Any
    ):
        action_name = task.get_action_name(action["action"])
        gripped_object_id = observations["gripped_object_id"]
        if action_name == "GRAB_RELEASE" and gripped_object_id != -1 and self.prev_gripped_object_id != -1 and gripped_object_id == self.prev_gripped_object_id:
            self._metric += 1
        self.prev_gripped_object_id = gripped_object_id


@registry.register_sensor
class AllObjectPositions(PointGoalSensor):
    cls_uuid = "all_object_positions"
    
    def _get_observation_space(self, *args: Any, **kwargs: Any):
        sensor_shape = (2, self._dimensionality,)

        return spaces.Box(
            low=np.finfo(np.float32).min,
            high=np.finfo(np.float32).max,
            shape=sensor_shape,
            dtype=np.float32,
        )

    def get_observation(
        self, observations, episode, *args: Any, **kwargs: Any
    ):  
        agent_state = self._sim.get_agent_state()
        agent_position = agent_state.position
        rotation_world_agent = agent_state.rotation
        sensor_data = np.zeros((2, self._dimensionality))
        objects = self._sim._scene_objects
        existing_object_ids = self._sim.get_existing_object_ids()

        for obj in objects:
            object_id = obj.object_id

            if object_id not in existing_object_ids:
                sensor_data[object_id] = self._compute_pointgoal(
                    agent_position, rotation_world_agent, agent_position
                )
            else:
                object_position = self._sim.get_translation(object_id)
                sensor_data[object_id] = self._compute_pointgoal(
                    agent_position, rotation_world_agent, object_position
                )
        return sensor_data


@registry.register_measure
class RearrangementReward(Measure):
    r"""RearrangementReward"""

    def __init__(self, sim, config: Config, *args: Any, **kwargs: Any):
        self._sim = sim
        self._config = config
        self.step_count = 0
        self._gripped_object_count = defaultdict(int)
        super().__init__()

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        return "rearrangement_reward"

    def reset_metric(self, episode, task, observations, *args: Any, **kwargs: Any):
        task.measurements.check_measure_dependencies(
            self.uuid,
            [
                AgentToObjectDistance.cls_uuid,
                AgentToReceptacleDistance.cls_uuid,
                GoalObjectVisible.cls_uuid,
                RearrangementSuccess.cls_uuid,
                Coverage.cls_uuid,
            ],
        )
        self._metric = 0
        self.step_count = 0
        self._object_was_seen = False
        self._recetpacle_was_seen = False
        self._attenuation_penalty = 1.0
        self._gripped_object_count = defaultdict(int)
        self._previous_agent_object_distance = task.measurements.measures[AgentToObjectDistance.cls_uuid].get_metric()["metric"]
        self._previous_agent_receptacle_distance = task.measurements.measures[AgentToReceptacleDistance.cls_uuid].get_metric()["metric"]
        self._previous_gripped_object_id = -1
        self._previous_exploration_area = None
        self._previous_placement_success = 0
        self._placed_success = False
        self.update_metric(episode=episode, task=task, action={"action": 0}, observations=observations, *args, **kwargs)
    
    def reward_coverage(
        self,
        episode,
        task: EmbodiedTask,
        observations,
        *args: Any,
        **kwargs: Any,
    ):
        self._attenuation_penalty *= self._config.ATTENUATION
        visit = task.measurements.measures[
            Coverage.cls_uuid
        ].get_metric()["visit_count"]
        coverage_reward = self._attenuation_penalty * self._config.COVERAGE_REWARD / (visit ** self._config.VISIT_EXP)
        return coverage_reward

    def reward_grab_success(self, episode, task, action, observations, *args: Any, **kwargs: Any):
        reward = 0.0
        action_name = task.get_action_name(action["action"])
        current_agent_receptacle_distance = task.measurements.measures[AgentToReceptacleDistance.cls_uuid].get_metric()["metric"]
        current_object_receptacle_distance = task.measurements.measures[ObjectToReceptacleDistance.cls_uuid].get_metric()["metric"]

        if action_name != "GRAB_RELEASE":
            return reward

        if observations["gripped_object_id"] != -1 and observations["gripped_object_id"] != self._previous_gripped_object_id:
            obj_id = observations["gripped_object_id"]
            self._gripped_object_count[obj_id] += 1
            # Reward only on first grab action
            if self._gripped_object_count[obj_id] <= 1:
                reward += (
                    self._config.GRAB_SUCCESS_REWARD /
                    self._gripped_object_count[obj_id]
                )
        # Add penalty if agent drops object too far
        if self._config.ENABLE_PENALTY:
            drop_penalty_dist_threshold = (
                self._config.DROP_PENALTY_DIST_THRESHOLD
            )

            # If agent drops object too far from receptacle add penalty
            if self._previous_gripped_object_id != -1 and current_agent_receptacle_distance > drop_penalty_dist_threshold:
                reward += self._config.DROP_PENALTY
        
        # If the agent drops an object at it's successful position, 
        # give a positive reward, if it removes from a successful position,
        # give a negative reward
        if observations['gripped_object_id'] == -1:
            placement_threshold = self._config.SUCCESS_DISTANCE
            is_success = int(current_object_receptacle_distance <= placement_threshold)
            # New success will be in {-1, 0, 1}
            new_success = (is_success - self._previous_placement_success)

            reward += self._config.GRAB_SUCCESS_REWARD * new_success
            self._previous_placement_success = is_success
        
        return reward

    def reward_distance_to_object(
        self,
        episode,
        task: EmbodiedTask,
        action,
        observations,
        *args: Any,
        **kwargs: Any,
    ):
        reward = 0
        current_agent_object_distance = task.measurements.measures[AgentToObjectDistance.cls_uuid].get_metric()["metric"]
        current_agent_receptacle_distance = task.measurements.measures[AgentToReceptacleDistance.cls_uuid].get_metric()["metric"]
        action_name = task.get_action_name(action["action"])

        if action_name != "GRAB_RELEASE" and self._previous_gripped_object_id == -1:
            agent_object_dist_reward = self._previous_agent_object_distance - current_agent_object_distance
            reward += agent_object_dist_reward

        if action_name != "GRAB_RELEASE" and self._previous_gripped_object_id != -1:
            agent_receptacle_dist_reward = self._previous_agent_receptacle_distance - current_agent_receptacle_distance
            reward += agent_receptacle_dist_reward
        return reward

    def reward_distance_to_goal(
        self,
        episode,
        task: EmbodiedTask,
        action,
        observations,
        *args: Any,
        **kwargs: Any,
    ):
        reward = 0
        # Grab success reward
        reward += self.reward_grab_success(
            episode=episode,
            task=task,
            action=action,
            observations=observations,
            *args,
            **kwargs,
        )

        reward += self.reward_distance_to_object(
            episode=episode,
            task=task,
            action=action,
            observations=observations,
            *args,
            **kwargs,
        )
        return reward

    def reward_object_seen(
        self,
        task,
        observations
    ):
        reward = 0
        goal_object_visible = task.measurements.measures[
            GoalObjectVisible.cls_uuid
        ].get_metric()
        # Object seen threshold is small as we have a lot of small objects
        object_seen_reward_threshold = 0.001
        receptacle_seen_reward_threshold = 0.001
        if hasattr(self._config, "OBJECT_SEEN_REWARD_THRESHOLD"):
            object_seen_reward_threshold = (
                self._config.OBJECT_SEEN_REWARD_THRESHOLD
            )
        if hasattr(self._config, "RECEPTACLE_SEEN_REWARD_THRESHOLD"):
            receptacle_seen_reward_threshold = (
                self._config.RECEPTACLE_SEEN_REWARD_THRESHOLD
            )

        if (
            goal_object_visible > object_seen_reward_threshold
            and not self._object_was_seen
        ):
            self._object_was_seen = True
            reward += self._config.OBJECT_SEEN_REWARD

        if (
            goal_object_visible > receptacle_seen_reward_threshold
            and observations["gripped_object_id"] != -1
            and not self._recetpacle_was_seen
        ):
            self._recetpacle_was_seen = True
            reward += self._config.OBJECT_SEEN_REWARD
        return reward

    def reward_distance_to_object_plus_visible(
        self,
        episode,
        task: EmbodiedTask,
        action,
        observations,
        *args: Any,
        **kwargs: Any,
    ):
        reward = 0
        # Object seen reward
        reward += self.reward_object_seen(task=task, observations=observations)

        # Grab success reward
        reward += self.reward_grab_success(
            episode=episode,
            task=task,
            action=action,
            observations=observations,
            *args,
            **kwargs,
        )

        reward += self.reward_distance_to_object(
            episode=episode,
            task=task,
            action=action,
            observations=observations,
            *args,
            **kwargs,
        )

        return reward

    def reward_explore_dt_when_object_visible(
        self,
        episode,
        task: EmbodiedTask,
        action,
        observations,
        *args: Any,
        **kwargs: Any,
    ):
        reward = 0
        # Object seen reward
        reward += self.reward_object_seen(task=task, observations=observations)

        # Grab success reward
        reward += self.reward_grab_success(
            episode=episode,
            task=task,
            action=action,
            observations=observations,
            *args,
            **kwargs,
        )

        if not self._object_was_seen and observations["gripped_object_id"] == -1:
            reward += self.reward_coverage(
                episode=episode,
                task=task,
                observations=observations,
                *args,
                **kwargs,
            )
        else:
            reward += self.reward_distance_to_object(
                episode=episode,
                task=task,
                action=action,
                observations=observations,
                *args,
                **kwargs,
            )
        return reward

    def reward_explore_dt_when_object_or_receptacle_visible(
        self,
        episode,
        task: EmbodiedTask,
        action,
        observations,
        *args: Any,
        **kwargs: Any,
    ):
        reward = 0
        # Object seen reward
        reward += self.reward_object_seen(task=task, observations=observations)

        # Grab success reward
        reward += self.reward_grab_success(
            episode=episode,
            task=task,
            action=action,
            observations=observations,
            *args,
            **kwargs,
        )

        # Explore until you find the object
        if not self._object_was_seen and observations["gripped_object_id"] == -1:
            # print("object not seen coverage")
            reward += self.reward_coverage(
                episode=episode,
                task=task,
                observations=observations,
                *args,
                **kwargs,
            )
        # Explore after you grab the object until you find the receptacle
        elif not self._recetpacle_was_seen and observations["gripped_object_id"] != -1:
            # print("receptacle not seen coverage")
            reward += self.reward_coverage(
                episode=episode,
                task=task,
                observations=observations,
                *args,
                **kwargs,
            )
        # Use distance reward if you saw object and receptacle when solving the task
        else:
            # print("Distance when obj/recept seen")
            reward += self.reward_distance_to_object(
                episode=episode,
                task=task,
                action=action,
                observations=observations,
                *args,
                **kwargs,
            )
        return reward

    def update_metric(
        self,
        episode,
        task: EmbodiedTask,
        action,
        observations,
        *args: Any,
        **kwargs: Any,
    ):
        self.step_count += 1
        # reward = self._config.SLACK_REWARD if self.step_count > 20 else 0
        reward = 0

        current_agent_object_distance = task.measurements.measures[AgentToObjectDistance.cls_uuid].get_metric()["metric"]
        current_agent_receptacle_distance = task.measurements.measures[AgentToReceptacleDistance.cls_uuid].get_metric()["metric"]

        if self._config.MODE == "DISTANCE_TO_GOAL_PLUS_VISIBLE":
            reward += self.reward_distance_to_object_plus_visible(
                episode=episode,
                task=task,
                action=action,
                observations=observations,
                *args,
                **kwargs,
            )
        elif self._config.MODE == "DISTANCE_TO_GOAL_WHEN_VISIBLE_ELSE_EXPLORE":
            reward += self.reward_explore_dt_when_object_visible(
                episode=episode,
                task=task,
                action=action,
                observations=observations,
                *args,
                **kwargs,
            )
        elif self._config.MODE == "DISTANCE_TO_OBJ_THEN_RECPT_VISIBLE_ELSE_EXPLORE":
            reward += self.reward_explore_dt_when_object_or_receptacle_visible(
                episode=episode,
                task=task,
                action=action,
                observations=observations,
                *args,
                **kwargs,
            )
        elif self._config.MODE == "DISTANCE_TO_GOAL":
            reward += self.reward_distance_to_goal(
                episode=episode,
                task=task,
                action=action,
                observations=observations,
                *args,
                **kwargs,
            )

        if task.measurements.measures[RearrangementSuccess.cls_uuid].get_metric():
            reward += self._config.SUCCESS_REWARD

        self._previous_agent_object_distance = current_agent_object_distance
        self._previous_agent_receptacle_distance = current_agent_receptacle_distance
        self._previous_gripped_object_id = observations["gripped_object_id"]
        self._metric = reward


@registry.register_measure
class RearrangementSuccess(Measure):
    r"""Whether or not the agent succeeded at its task

    This measure depends on DistanceToGoal measure.
    """

    cls_uuid: str = "success"

    def __init__(
        self, sim: Simulator, config: Config, *args: Any, **kwargs: Any
    ):
        self._sim = sim
        self._config = config

        super().__init__()

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        return self.cls_uuid

    def reset_metric(self, episode, task, *args: Any, **kwargs: Any):
        task.measurements.check_measure_dependencies(
            self.uuid, [ObjectToReceptacleDistance.cls_uuid]
        )
        self.update_metric(episode=episode, task=task, *args, **kwargs)  # type: ignore

    def update_metric(
        self, episode, task: EmbodiedTask, *args: Any, **kwargs: Any
    ):
        distance_to_target = task.measurements.measures[
            ObjectToReceptacleDistance.cls_uuid
        ].get_metric()["metric"]
        object_ids = self._sim.get_existing_object_ids()

        obj_id = -1
        receptacle_id = -1
        for object_id in object_ids:
            scene_object = self._sim.get_object_from_scene(object_id)
            if scene_object.is_receptacle == False:
                obj_id = scene_object.object_id
            else:
                receptacle_id = scene_object.object_id

        is_object_stacked = False
        if obj_id != -1 and receptacle_id != -1:
            object_position = self._sim.get_translation(obj_id)
            receptacle_position = self._sim.get_translation(receptacle_id)

            object_y = object_position.y
            receptacle_y = receptacle_position.y + self._sim.get_object_bb_y_coord(receptacle_id)
            is_object_stacked = (object_y > receptacle_y)
        gripped_object_id = self._sim.gripped_object_id

        if (
            hasattr(task, "is_stop_called")
            and task.is_stop_called # type: ignore
            and distance_to_target <= self._config.SUCCESS_DISTANCE
            and is_object_stacked
            and gripped_object_id == -1
        ):
            self._metric = 1.0
        else:
            self._metric = 0.0


@registry.register_measure
class RearrangementSPL(Measure):
    r"""SPL (Success weighted by Path Length)

    ref: On Evaluation of Embodied Agents - Anderson et. al
    https://arxiv.org/pdf/1807.06757.pdf
    The measure depends on Distance to Goal measure and Success measure
    to improve computational
    performance for sophisticated goal areas.
    """

    def __init__(
        self, sim: Simulator, config: Config, *args: Any, **kwargs: Any
    ):
        self._previous_position = None
        self._start_end_episode_distance = None
        self._agent_episode_distance: Optional[float] = None
        self._episode_view_points = None
        self._sim = sim
        self._config = config

        super().__init__()

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        return "spl"

    def reset_metric(self, episode, task, *args: Any, **kwargs: Any):
        task.measurements.check_measure_dependencies(
            self.uuid, [DistanceToGoal.cls_uuid, RearrangementSuccess.cls_uuid]
        )

        self._previous_position = self._sim.get_agent_state().position
        self._agent_episode_distance = 0.0
        self._start_end_episode_distance = task.measurements.measures[
            AgentToObjectDistance.cls_uuid
        ].get_metric()["metric"] + task.measurements.measures[
            ObjectToReceptacleDistance.cls_uuid
        ].get_metric()["metric"] 
        self.update_metric(  # type:ignore
            episode=episode, task=task, *args, **kwargs
        )

    def _euclidean_distance(self, position_a, position_b):
        return np.linalg.norm(position_b - position_a, ord=2)

    def update_metric(
        self, episode, task: EmbodiedTask, *args: Any, **kwargs: Any
    ):
        ep_success = task.measurements.measures[RearrangementSuccess.cls_uuid].get_metric()

        current_position = self._sim.get_agent_state().position
        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 GrabSuccess(Measure):
    r"""Grab Success - whether an object was grabbed during episode or not
    """

    def __init__(
        self, sim: Simulator, config: Config, *args: Any, **kwargs: Any
    ):
        self._sim = sim
        self._config = config
        self.prev_gripped_object_id = -1

        super().__init__()

    def _get_uuid(self, *args: Any, **kwargs: Any) -> str:
        return "grab_success"

    def reset_metric(self, episode, task, *args: Any, **kwargs: Any):
        self._metric = 0
        self.update_metric(episode=episode, task=task, *args, **kwargs)  # type: ignore

    def update_metric(
        self, episode, task: EmbodiedTask, *args: Any, **kwargs: Any
    ):
        gripped_object_id = self._sim.gripped_object_id
        if gripped_object_id != -1 and gripped_object_id != self.prev_gripped_object_id:
            self._metric += 1
        self.prev_gripped_object_id = gripped_object_id


def merge_sim_episode_with_object_config(
    sim_config: Config, episode: Type[Episode]
) -> Any:
    sim_config = merge_sim_episode_config(sim_config, episode)
    sim_config.defrost()
    sim_config.objects = episode.objects
    sim_config.freeze()

    return sim_config


@registry.register_task(name="PickPlaceTask-v0")
class PickPlaceTask(EmbodiedTask):
    r"""Language based Pick Place Task
    Goal: An agent must rearrange objects in a 3D environment
        specified by a natural language instruction.
    Usage example:
        examples/object_rearrangement_example.py
    """

    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)
        self._is_episode_active = False
    
    def reset(self, **kwargs):
        observations = super().reset(**kwargs)
        return observations

    def _check_episode_is_active(self, *args: Any, **kwargs: Any) -> bool:
        is_object_out_of_bounds = self._sim.is_object_out_of_bounds()
        if is_object_out_of_bounds:
            logger.info("Object is OOB terminating episodes")
        return not getattr(self, "is_stop_called", False) and not is_object_out_of_bounds

    def overwrite_sim_config(self, sim_config, episode):
        return merge_sim_episode_with_object_config(sim_config, episode)

habitat/tasks/pickplace/pickplace.py
cshizhe/onav_rim
#!/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.

from habitat.core.logging import logger
from habitat.core.registry import registry
from habitat.tasks.eqa import _try_register_eqa_task
from habitat.tasks.nav import _try_register_nav_task
from habitat.tasks.vln import _try_register_vln_task
from habitat.tasks.pickplace import _try_register_pickplace_task


def make_task(id_task, **kwargs):
    logger.info("Initializing task {}".format(id_task))
    _task = registry.get_task(id_task)
    assert _task is not None, "Could not find task with name {}".format(
        id_task
    )

    return _task(**kwargs)


_try_register_eqa_task()
_try_register_nav_task()
_try_register_vln_task()
_try_register_pickplace_task()

habitat/tasks/registration.py
cshizhe/onav_rim
#!/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 numpy as np
import quaternion  # noqa # pylint: disable=unused-import

from habitat.sims.habitat_simulator.actions import HabitatSimActions


def quaternion_to_rotation(q_r, q_i, q_j, q_k):
    r"""
    ref: https://en.wikipedia.org/wiki/Quaternions_and_spatial_rotation
    """
    s = 1  # unit quaternion
    rotation_mat = np.array(
        [
            [
                1 - 2 * s * (q_j ** 2 + q_k ** 2),
                2 * s * (q_i * q_j - q_k * q_r),
                2 * s * (q_i * q_k + q_j * q_r),
            ],
            [
                2 * s * (q_i * q_j + q_k * q_r),
                1 - 2 * s * (q_i ** 2 + q_k ** 2),
                2 * s * (q_j * q_k - q_i * q_r),
            ],
            [
                2 * s * (q_i * q_k - q_j * q_r),
                2 * s * (q_j * q_k + q_i * q_r),
                1 - 2 * s * (q_i ** 2 + q_j ** 2),
            ],
        ],
        dtype=np.float32,
    )
    return rotation_mat


def cartesian_to_polar(x, y):
    rho = np.sqrt(x ** 2 + y ** 2)
    phi = np.arctan2(y, x)
    return rho, phi


def compute_pixel_coverage(instance_seg, object_id):
    cand_mask = instance_seg == object_id
    score = cand_mask.sum().astype(np.float64) / cand_mask.size
    return score


def get_habitat_sim_action(action):
    if action == "TURN_RIGHT":
        return HabitatSimActions.TURN_RIGHT
    elif action == "TURN_LEFT":
        return HabitatSimActions.TURN_LEFT
    elif action == "MOVE_FORWARD":
        return HabitatSimActions.MOVE_FORWARD
    elif action == "MOVE_BACKWARD":
        return HabitatSimActions.MOVE_BACKWARD
    elif action == "LOOK_UP":
        return HabitatSimActions.LOOK_UP
    elif action == "LOOK_DOWN":
        return HabitatSimActions.LOOK_DOWN
    return HabitatSimActions.STOP


def get_habitat_sim_action_str(action):
    if action == HabitatSimActions.TURN_RIGHT:
        return "turnRight"
    elif action == HabitatSimActions.TURN_LEFT:
        return "turnLeft"
    elif action == HabitatSimActions.MOVE_FORWARD:
        return "moveForward"
    elif action == HabitatSimActions.MOVE_BACKWARD:
        return "moveBackward"
    elif action == HabitatSimActions.LOOK_UP:
        return "lookUp"
    elif action == HabitatSimActions.LOOK_DOWN:
        return "lookDown"
    elif action == HabitatSimActions.GRAB_RELEASE:
        return "grabReleaseObject"
    elif action == HabitatSimActions.NO_OP:
        return "stepPhysics"
    elif action == HabitatSimActions.START:
        return "start"
    return "stop"

habitat/tasks/utils.py
cshizhe/onav_rim
# Copyright 2016 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.
# ==============================================================================

"""Utilities for processing depth images.
"""
from argparse import Namespace

import itertools
import numpy as np
import torch

import habitat.utils.rotation_utils as ru


def get_camera_matrix(width, height, fov):
    """Returns a camera matrix from image size and fov."""
    xc = (width - 1.) / 2.
    zc = (height - 1.) / 2.
    f = (width / 2.) / np.tan(np.deg2rad(fov / 2.))
    camera_matrix = {'xc': xc, 'zc': zc, 'f': f}
    camera_matrix = Namespace(**camera_matrix)
    return camera_matrix


def get_point_cloud_from_z(Y, camera_matrix, scale=1):
    """Projects the depth image Y into a 3D point cloud.
    Inputs:
        Y is ...xHxW
        camera_matrix
    Outputs:
        X is positive going right
        Y is positive into the image
        Z is positive up in the image
        XYZ is ...xHxWx3
    """
    x, z = np.meshgrid(np.arange(Y.shape[-1]),
                       np.arange(Y.shape[-2] - 1, -1, -1))
    for _ in range(Y.ndim - 2):
        x = np.expand_dims(x, axis=0)
        z = np.expand_dims(z, axis=0)
    X = (x[::scale, ::scale] - camera_matrix.xc) * \
        Y[::scale, ::scale] / camera_matrix.f
    Z = (z[::scale, ::scale] - camera_matrix.zc) * \
        Y[::scale, ::scale] / camera_matrix.f
    XYZ = np.concatenate((X[..., np.newaxis],
                          Y[::scale, ::scale][..., np.newaxis],
                          Z[..., np.newaxis]), axis=X.ndim)
    return XYZ


def transform_camera_view(XYZ, sensor_height, camera_elevation_degree):
    """
    Transforms the point cloud into geocentric frame to account for
    camera elevation and angle
    Input:
        XYZ                     : ...x3
        sensor_height           : height of the sensor
        camera_elevation_degree : camera elevation to rectify.
    Output:
        XYZ : ...x3
    """
    R = ru.get_r_matrix(
        [1., 0., 0.], angle=np.deg2rad(camera_elevation_degree))
    XYZ = np.matmul(XYZ.reshape(-1, 3), R.T).reshape(XYZ.shape)
    XYZ[..., 2] = XYZ[..., 2] + sensor_height
    return XYZ


def transform_pose(XYZ, current_pose):
    """
    Transforms the point cloud into geocentric frame to account for
    camera position
    Input:
        XYZ                     : ...x3
        current_pose            : camera position (x, y, theta (radians))
    Output:
        XYZ : ...x3
    """
    R = ru.get_r_matrix([0., 0., 1.], angle=current_pose[2] - np.pi / 2.)
    XYZ = np.matmul(XYZ.reshape(-1, 3), R.T).reshape(XYZ.shape)
    XYZ[:, :, 0] = XYZ[:, :, 0] + current_pose[0]
    XYZ[:, :, 1] = XYZ[:, :, 1] + current_pose[1]
    return XYZ


def bin_points(XYZ_cms, map_size, z_bins, xy_resolution):
    """Bins points into xy-z bins
    XYZ_cms is ... x H x W x3
    Outputs is ... x map_size x map_size x (len(z_bins)+1)
    """
    sh = XYZ_cms.shape
    XYZ_cms = XYZ_cms.reshape([-1, sh[-3], sh[-2], sh[-1]])
    n_z_bins = len(z_bins) + 1
    counts = []
    for XYZ_cm in XYZ_cms:
        isnotnan = np.logical_not(np.isnan(XYZ_cm[:, :, 0]))
        X_bin = np.round(XYZ_cm[:, :, 0] / xy_resolution).astype(np.int32)
        Y_bin = np.round(XYZ_cm[:, :, 1] / xy_resolution).astype(np.int32)
        Z_bin = np.digitize(XYZ_cm[:, :, 2], bins=z_bins).astype(np.int32)

        isvalid = np.array([X_bin >= 0, X_bin < map_size, Y_bin >= 0,
                            Y_bin < map_size,
                            Z_bin >= 0, Z_bin < n_z_bins, isnotnan])
        isvalid = np.all(isvalid, axis=0)

        ind = (Y_bin * map_size + X_bin) * n_z_bins + Z_bin
        ind[np.logical_not(isvalid)] = 0
        count = np.bincount(ind.ravel(), isvalid.ravel().astype(np.int32),
                            minlength=map_size * map_size * n_z_bins)
        counts = np.reshape(count, [map_size, map_size, n_z_bins])

    counts = counts.reshape(list(sh[:-3]) + [map_size, map_size, n_z_bins])

    return counts


def get_point_cloud_from_z_t(Y_t, camera_matrix, device, scale=1):
    """Projects the depth image Y into a 3D point cloud.
    Inputs:
        Y is ...xHxW
        camera_matrix
    Outputs:
        X is positive going right
        Y is positive into the image
        Z is positive up in the image
        XYZ is ...xHxWx3
    """
    grid_x, grid_z = torch.meshgrid(torch.arange(Y_t.shape[-1]),
                                    torch.arange(Y_t.shape[-2] - 1, -1, -1))
    grid_x = grid_x.transpose(1, 0).to(device)
    grid_z = grid_z.transpose(1, 0).to(device)
    grid_x = grid_x.unsqueeze(0).expand(Y_t.size())
    grid_z = grid_z.unsqueeze(0).expand(Y_t.size())

    X_t = (grid_x[:, ::scale, ::scale] - camera_matrix.xc) * \
        Y_t[:, ::scale, ::scale] / camera_matrix.f
    Z_t = (grid_z[:, ::scale, ::scale] - camera_matrix.zc) * \
        Y_t[:, ::scale, ::scale] / camera_matrix.f

    XYZ = torch.stack(
        (X_t, Y_t[:, ::scale, ::scale], Z_t), dim=len(Y_t.size()))

    return XYZ


def transform_camera_view_t(
        XYZ, sensor_height, camera_elevation_degree, device):
    """
    Transforms the point cloud into geocentric frame to account for
    camera elevation and angle
    Input:
        XYZ                     : ...x3
        sensor_height           : height of the sensor
        camera_elevation_degree : camera elevation to rectify.
    Output:
        XYZ : ...x3
    """
    R = ru.get_r_matrix(
        [1., 0., 0.], angle=np.deg2rad(camera_elevation_degree))
    XYZ = torch.matmul(XYZ.reshape(-1, 3),
                       torch.from_numpy(R).float().transpose(1, 0).to(device)
                       ).reshape(XYZ.shape)
    XYZ[..., 2] = XYZ[..., 2] + sensor_height
    return XYZ


def transform_pose_t(XYZ, current_pose, device):
    """
    Transforms the point cloud into geocentric frame to account for
    camera position
    Input:
        XYZ                     : ...x3
        current_pose            : camera position (x, y, theta (radians))
    Output:
        XYZ : ...x3
    """
    # R = ru.get_r_matrix([0., 0., 1.], angle=current_pose[2] - np.pi / 2.) # facing east as 0
    R = ru.get_r_matrix([0., 0., 1.], angle=current_pose[2])
    XYZ = torch.matmul(XYZ.reshape(-1, 3),
                       torch.from_numpy(R).float().transpose(1, 0).to(device)
                       ).reshape(XYZ.shape)
    XYZ[..., 0] += current_pose[0]
    XYZ[..., 1] += current_pose[1]
    return XYZ


def splat_feat_nd(init_grid, feat, coords):
    """
    Args:
        init_grid: B X nF X W X H X D X ..
        feat: B X nF X nPt
        coords: B X nDims X nPt in [-1, 1]
    Returns:
        grid: B X nF X W X H X D X ..
    """
    wts_dim = []
    pos_dim = []
    grid_dims = init_grid.shape[2:]

    B = init_grid.shape[0]
    F = init_grid.shape[1]

    n_dims = len(grid_dims)

    grid_flat = init_grid.view(B, F, -1)

    for d in range(n_dims):
        pos = coords[:, [d], :] * grid_dims[d] / 2 + grid_dims[d] / 2
        pos_d = []
        wts_d = []

        for ix in [0, 1]:
            pos_ix = torch.floor(pos) + ix
            safe_ix = (pos_ix > 0) & (pos_ix < grid_dims[d])
            safe_ix = safe_ix.type(pos.dtype)

            # put the position with decimals into the neighboring 2 indexes
            wts_ix = 1 - torch.abs(pos - pos_ix)

            wts_ix = wts_ix * safe_ix
            pos_ix = pos_ix * safe_ix

            pos_d.append(pos_ix)    # (B, nPts), indexes in the grid dimension
            wts_d.append(wts_ix)    # (B, nPts), weights in the index of current indexes

        pos_dim.append(pos_d)
        wts_dim.append(wts_d)

    l_ix = [[0, 1] for d in range(n_dims)]

    for ix_d in itertools.product(*l_ix):
        # ix_d: (0, 0, 0), (0, 0, 1), ..., (1, 1, 1)
        wts = torch.ones_like(wts_dim[0][0])    # (B, nPts)
        index = torch.zeros_like(wts_dim[0][0])
        for d in range(n_dims):
            index = index * grid_dims[d] + pos_dim[d][ix_d[d]]
            wts = wts * wts_dim[d][ix_d[d]]

        index = index.long()
        grid_flat.scatter_add_(2, index.expand(-1, F, -1), feat * wts)
        grid_flat = torch.round(grid_flat)

    return grid_flat.view(init_grid.shape)

habitat/utils/depth_utils.py
cshizhe/onav_rim
# Copyright 2016 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.
# ==============================================================================

"""Utilities for generating and applying rotation matrices.
"""
import numpy as np

ANGLE_EPS = 0.001


def normalize(v):
    return v / np.linalg.norm(v)


def get_r_matrix(ax_, angle):
    # Rodrigues' rotation formula
    # https://www.cnblogs.com/wtyuan/p/12324495.html
    ax = normalize(ax_)
    if np.abs(angle) > ANGLE_EPS:
        S_hat = np.array(
            [[0.0, -ax[2], ax[1]], [ax[2], 0.0, -ax[0]], [-ax[1], ax[0], 0.0]],
            dtype=np.float32)
        R = np.eye(3) + np.sin(angle) * S_hat + \
            (1 - np.cos(angle)) * (np.linalg.matrix_power(S_hat, 2))
    else:
        R = np.eye(3)
    return R


def r_between(v_from_, v_to_):
    v_from = normalize(v_from_)
    v_to = normalize(v_to_)
    ax = normalize(np.cross(v_from, v_to))
    angle = np.arccos(np.dot(v_from, v_to))
    return get_r_matrix(ax, angle)


def rotate_camera_to_point_at(up_from, lookat_from, up_to, lookat_to):
    inputs = [up_from, lookat_from, up_to, lookat_to]
    for i in range(4):
        inputs[i] = normalize(np.array(inputs[i]).reshape((-1,)))
    up_from, lookat_from, up_to, lookat_to = inputs
    r1 = r_between(lookat_from, lookat_to)

    new_x = np.dot(r1, np.array([1, 0, 0]).reshape((-1, 1))).reshape((-1))
    to_x = normalize(np.cross(lookat_to, up_to))
    angle = np.arccos(np.dot(new_x, to_x))
    if angle > ANGLE_EPS:
        if angle < np.pi - ANGLE_EPS:
            ax = normalize(np.cross(new_x, to_x))
            flip = np.dot(lookat_to, ax)
            if flip > 0:
                r2 = get_r_matrix(lookat_to, angle)
            elif flip < 0:
                r2 = get_r_matrix(lookat_to, -1. * angle)
        else:
            # Angle of rotation is too close to 180 degrees, direction of
            # rotation does not matter.
            r2 = get_r_matrix(lookat_to, angle)
    else:
        r2 = np.eye(3)
    return np.dot(r2, r1)

habitat/utils/rotation_utils.py
cshizhe/onav_rim
#!/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 typing import Any, Dict, List, Optional, Sequence, Tuple

import imageio
import numpy as np
import scipy.ndimage

from habitat.core.utils import try_cv2_import
from habitat.utils.visualizations import utils

try:
    from habitat.sims.habitat_simulator.habitat_simulator import HabitatSim
except ImportError:
    pass

cv2 = try_cv2_import()


AGENT_SPRITE = imageio.imread(
    os.path.join(
        os.path.dirname(__file__),
        "assets",
        "maps_topdown_agent_sprite",
        "100x100.png",
    )
)
AGENT_SPRITE = np.ascontiguousarray(np.flipud(AGENT_SPRITE))

MAP_INVALID_POINT = 0
MAP_VALID_POINT = 1
MAP_BORDER_INDICATOR = 2
MAP_SOURCE_POINT_INDICATOR = 4
MAP_TARGET_POINT_INDICATOR = 6
MAP_SHORTEST_PATH_COLOR = 7
MAP_VIEW_POINT_INDICATOR = 8
MAP_TARGET_BOUNDING_BOX = 9
MAP_RECEPTACLE_COLOR = 10
TOP_DOWN_MAP_COLORS = np.full((256, 3), 150, dtype=np.uint8)
TOP_DOWN_MAP_COLORS[10:] = cv2.applyColorMap(
    np.arange(246, dtype=np.uint8), cv2.COLORMAP_JET
).squeeze(1)[:, ::-1]
TOP_DOWN_MAP_COLORS[MAP_INVALID_POINT] = [255, 255, 255]  # White
TOP_DOWN_MAP_COLORS[MAP_VALID_POINT] = [150, 150, 150]  # Light Grey
TOP_DOWN_MAP_COLORS[MAP_BORDER_INDICATOR] = [50, 50, 50]  # Grey
TOP_DOWN_MAP_COLORS[MAP_SOURCE_POINT_INDICATOR] = [0, 0, 200]  # Blue
TOP_DOWN_MAP_COLORS[MAP_TARGET_POINT_INDICATOR] = [200, 0, 0]  # Red
TOP_DOWN_MAP_COLORS[MAP_SHORTEST_PATH_COLOR] = [0, 200, 0]  # Green
TOP_DOWN_MAP_COLORS[MAP_VIEW_POINT_INDICATOR] = [245, 150, 150]  # Light Red
TOP_DOWN_MAP_COLORS[MAP_TARGET_BOUNDING_BOX] = [0, 175, 0]  # Green
TOP_DOWN_MAP_COLORS[MAP_RECEPTACLE_COLOR] = [255,215,0] # Yellow


def draw_agent(
    image: np.ndarray,
    agent_center_coord: Tuple[int, int],
    agent_rotation: float,
    agent_radius_px: int = 5,
) -> np.ndarray:
    r"""Return an image with the agent image composited onto it.
    Args:
        image: the image onto which to put the agent.
        agent_center_coord: the image coordinates where to paste the agent.
        agent_rotation: the agent's current rotation in radians.
        agent_radius_px: 1/2 number of pixels the agent will be resized to.
    Returns:
        The modified background image. This operation is in place.
    """

    # Rotate before resize to keep good resolution.
    rotated_agent = scipy.ndimage.interpolation.rotate(
        AGENT_SPRITE, agent_rotation * 180 / np.pi
    )
    # Rescale because rotation may result in larger image than original, but
    # the agent sprite size should stay the same.
    initial_agent_size = AGENT_SPRITE.shape[0]
    new_size = rotated_agent.shape[0]
    agent_size_px = max(
        1, int(agent_radius_px * 2 * new_size / initial_agent_size)
    )
    resized_agent = cv2.resize(
        rotated_agent,
        (agent_size_px, agent_size_px),
        interpolation=cv2.INTER_LINEAR,
    )
    utils.paste_overlapping_image(image, resized_agent, agent_center_coord)
    return image


def pointnav_draw_target_birdseye_view(
    agent_position: np.ndarray,
    agent_heading: float,
    goal_position: np.ndarray,
    resolution_px: int = 800,
    goal_radius: float = 0.2,
    agent_radius_px: int = 20,
    target_band_radii: Optional[List[float]] = None,
    target_band_colors: Optional[List[Tuple[int, int, int]]] = None,
) -> np.ndarray:
    r"""Return an image of agent w.r.t. centered target location for pointnav
    tasks.

    Args:
        agent_position: the agent's current position.
        agent_heading: the agent's current rotation in radians. This can be
            found using the HeadingSensor.
        goal_position: the pointnav task goal position.
        resolution_px: number of pixels for the output image width and height.
        goal_radius: how near the agent needs to be to be successful for the
            pointnav task.
        agent_radius_px: 1/2 number of pixels the agent will be resized to.
        target_band_radii: distance in meters to the outer-radius of each band
            in the target image.
        target_band_colors: colors in RGB 0-255 for the bands in the target.
    Returns:
        Image centered on the goal with the agent's current relative position
        and rotation represented by an arrow. To make the rotations align
        visually with habitat, positive-z is up, positive-x is left and a
        rotation of 0 points upwards in the output image and rotates clockwise.
    """
    if target_band_radii is None:
        target_band_radii = [20, 10, 5, 2.5, 1]
    if target_band_colors is None:
        target_band_colors = [
            (47, 19, 122),
            (22, 99, 170),
            (92, 177, 0),
            (226, 169, 0),
            (226, 12, 29),
        ]

    assert len(target_band_radii) == len(
        target_band_colors
    ), "There must be an equal number of scales and colors."

    goal_agent_dist = np.linalg.norm(agent_position - goal_position, 2)

    goal_distance_padding = max(
        2, 2 ** np.ceil(np.log(max(1e-6, goal_agent_dist)) / np.log(2))
    )
    movement_scale = 1.0 / goal_distance_padding
    half_res = resolution_px // 2
    im_position = np.full(
        (resolution_px, resolution_px, 3), 255, dtype=np.uint8
    )

    # Draw bands:
    for scale, color in zip(target_band_radii, target_band_colors):
        if goal_distance_padding * 4 > scale:
            cv2.circle(
                im_position,
                (half_res, half_res),
                max(2, int(half_res * scale * movement_scale)),
                color,
                thickness=-1,
            )

    # Draw such that the agent being inside the radius is the circles
    # overlapping.
    cv2.circle(
        im_position,
        (half_res, half_res),
        max(2, int(half_res * goal_radius * movement_scale)),
        (127, 0, 0),
        thickness=-1,
    )

    relative_position = agent_position - goal_position
    # swap x and z, remove y for (x,y,z) -> image coordinates.
    relative_position = relative_position[[2, 0]]
    relative_position *= half_res * movement_scale
    relative_position += half_res
    relative_position = np.round(relative_position).astype(np.int32)

    # Draw the agent
    draw_agent(im_position, relative_position, agent_heading, agent_radius_px)

    # Rotate twice to fix coordinate system to upwards being positive-z.
    # Rotate instead of flip to keep agent rotations in sync with egocentric
    # view.
    im_position = np.rot90(im_position, 2)
    return im_position


def to_grid(
    realworld_x: float,
    realworld_y: float,
    grid_resolution: Tuple[int, int],
    sim: Optional["HabitatSim"] = None,
    pathfinder=None,
) -> Tuple[int, int]:
    r"""Return gridworld index of realworld coordinates assuming top-left corner
    is the origin. The real world coordinates of lower left corner are
    (coordinate_min, coordinate_min) and of top right corner are
    (coordinate_max, coordinate_max)
    """
    if sim is None and pathfinder is None:
        raise RuntimeError(
            "Must provide either a simulator or pathfinder instance"
        )

    if pathfinder is None:
        pathfinder = sim.pathfinder

    lower_bound, upper_bound = pathfinder.get_bounds()

    grid_size = (
        abs(upper_bound[2] - lower_bound[2]) / grid_resolution[0],
        abs(upper_bound[0] - lower_bound[0]) / grid_resolution[1],
    )
    grid_x = int((realworld_x - lower_bound[2]) / grid_size[0])
    grid_y = int((realworld_y - lower_bound[0]) / grid_size[1])
    return grid_x, grid_y


def from_grid(
    grid_x: int,
    grid_y: int,
    grid_resolution: Tuple[int, int],
    sim: Optional["HabitatSim"] = None,
    pathfinder=None,
) -> Tuple[float, float]:
    r"""Inverse of _to_grid function. Return real world coordinate from
    gridworld assuming top-left corner is the origin. The real world
    coordinates of lower left corner are (coordinate_min, coordinate_min) and
    of top right corner are (coordinate_max, coordinate_max)
    """

    if sim is None and pathfinder is None:
        raise RuntimeError(
            "Must provide either a simulator or pathfinder instance"
        )

    if pathfinder is None:
        pathfinder = sim.pathfinder

    lower_bound, upper_bound = pathfinder.get_bounds()

    grid_size = (
        abs(upper_bound[2] - lower_bound[2]) / grid_resolution[0],
        abs(upper_bound[0] - lower_bound[0]) / grid_resolution[1],
    )
    realworld_x = lower_bound[2] + grid_x * grid_size[0]
    realworld_y = lower_bound[0] + grid_y * grid_size[1]
    return realworld_x, realworld_y


def _outline_border(top_down_map):
    left_right_block_nav = (top_down_map[:, :-1] == 1) & (
        top_down_map[:, :-1] != top_down_map[:, 1:]
    )
    left_right_nav_block = (top_down_map[:, 1:] == 1) & (
        top_down_map[:, :-1] != top_down_map[:, 1:]
    )

    up_down_block_nav = (top_down_map[:-1] == 1) & (
        top_down_map[:-1] != top_down_map[1:]
    )
    up_down_nav_block = (top_down_map[1:] == 1) & (
        top_down_map[:-1] != top_down_map[1:]
    )

    top_down_map[:, :-1][left_right_block_nav] = MAP_BORDER_INDICATOR
    top_down_map[:, 1:][left_right_nav_block] = MAP_BORDER_INDICATOR

    top_down_map[:-1][up_down_block_nav] = MAP_BORDER_INDICATOR
    top_down_map[1:][up_down_nav_block] = MAP_BORDER_INDICATOR


def calculate_meters_per_pixel(
    map_resolution: int, sim: Optional["HabitatSim"] = None, pathfinder=None
):
    r"""Calculate the meters_per_pixel for a given map resolution"""
    if sim is None and pathfinder is None:
        raise RuntimeError(
            "Must provide either a simulator or pathfinder instance"
        )

    if pathfinder is None:
        pathfinder = sim.pathfinder

    lower_bound, upper_bound = pathfinder.get_bounds()
    return min(
        abs(upper_bound[coord] - lower_bound[coord]) / map_resolution
        for coord in [0, 2]
    )


def get_topdown_map(
    pathfinder,
    height: float,
    map_resolution: int = 1024,
    draw_border: bool = True,
    meters_per_pixel: Optional[float] = None,
) -> np.ndarray:
    r"""Return a top-down occupancy map for a sim. Note, this only returns valid
    values for whatever floor the agent is currently on.

    :param pathfinder: A habitat-sim pathfinder instances to get the map from
    :param height: The height in the environment to make the topdown map
    :param map_resolution: Length of the longest side of the map.  Used to calculate :p:`meters_per_pixel`
    :param draw_border: Whether or not to draw a border
    :param meters_per_pixel: Overrides map_resolution an

    :return: Image containing 0 if occupied, 1 if unoccupied, and 2 if border (if
        the flag is set).
    """

    if meters_per_pixel is None:
        meters_per_pixel = calculate_meters_per_pixel(
            map_resolution, pathfinder=pathfinder
        )

    top_down_map = pathfinder.get_topdown_view(
        meters_per_pixel=meters_per_pixel, height=height
    ).astype(np.uint8)

    # Draw border if necessary
    if draw_border:
        _outline_border(top_down_map)

    return np.ascontiguousarray(top_down_map)


def get_topdown_map_from_sim(
    sim: "HabitatSim",
    map_resolution: int = 1024,
    draw_border: bool = True,
    meters_per_pixel: Optional[float] = None,
    agent_id: int = 0,
) -> np.ndarray:
    r"""Wrapper around :py:`get_topdown_map` that retrieves that pathfinder and heigh from the current simulator

    :param sim: Simulator instance.
    :param agent_id: The agent ID
    """
    return get_topdown_map(
        sim.pathfinder,
        sim.get_agent(agent_id).state.position[1],
        map_resolution,
        draw_border,
        meters_per_pixel,
    )


def colorize_topdown_map(
    top_down_map: np.ndarray,
    fog_of_war_mask: Optional[np.ndarray] = None,
    fog_of_war_desat_amount: float = 0.5,
) -> np.ndarray:
    r"""Convert the top down map to RGB based on the indicator values.
    Args:
        top_down_map: A non-colored version of the map.
        fog_of_war_mask: A mask used to determine which parts of the
            top_down_map are visible
            Non-visible parts will be desaturated
        fog_of_war_desat_amount: Amount to desaturate the color of unexplored areas
            Decreasing this value will make unexplored areas darker
            Default: 0.5
    Returns:
        A colored version of the top-down map.
    """
    _map = TOP_DOWN_MAP_COLORS[top_down_map]

    if fog_of_war_mask is not None:
        fog_of_war_desat_values = np.array([[fog_of_war_desat_amount], [1.0]])
        # Only desaturate things that are valid points as only valid points get revealed
        desat_mask = top_down_map != MAP_INVALID_POINT

        _map[desat_mask] = (
            _map * fog_of_war_desat_values[fog_of_war_mask]
        ).astype(np.uint8)[desat_mask]

    return _map


def draw_path(
    top_down_map: np.ndarray,
    path_points: Sequence[Tuple],
    color: int = 10,
    thickness: int = 2,
) -> None:
    r"""Draw path on top_down_map (in place) with specified color.
    Args:
        top_down_map: A colored version of the map.
        color: color code of the path, from TOP_DOWN_MAP_COLORS.
        path_points: list of points that specify the path to be drawn
        thickness: thickness of the path.
    """
    for prev_pt, next_pt in zip(path_points[:-1], path_points[1:]):
        # Swapping x y
        cv2.line(
            top_down_map,
            prev_pt[::-1],
            next_pt[::-1],
            color,
            thickness=thickness,
        )


def colorize_draw_agent_and_fit_to_height(
    topdown_map_info: Dict[str, Any], output_height: int
):
    r"""Given the output of the TopDownMap measure, colorizes the map, draws the agent,
    and fits to a desired output height

    :param topdown_map_info: The output of the TopDownMap measure
    :param output_height: The desired output height
    """
    top_down_map = topdown_map_info["map"]
    top_down_map = colorize_topdown_map(
        top_down_map, topdown_map_info["fog_of_war_mask"]
    )
    map_agent_pos = topdown_map_info["agent_map_coord"]
    top_down_map = draw_agent(
        image=top_down_map,
        agent_center_coord=map_agent_pos,
        agent_rotation=topdown_map_info["agent_angle"],
        agent_radius_px=min(top_down_map.shape[0:2]) // 32,
    )

    if top_down_map.shape[0] > top_down_map.shape[1]:
        top_down_map = np.rot90(top_down_map, 1)

    # scale top down map to align with rgb view
    old_h, old_w, _ = top_down_map.shape
    top_down_height = output_height
    top_down_width = int(float(top_down_height) / old_h * old_w)
    # cv2 resize (dsize is width first)
    top_down_map = cv2.resize(
        top_down_map,
        (top_down_width, top_down_height),
        interpolation=cv2.INTER_CUBIC,
    )

    return top_down_map

habitat/utils/visualizations/maps.py
cshizhe/onav_rim
#!/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 cv2
import os
import textwrap
from typing import Dict, List, Optional, Tuple

import colorsys
import imageio
import numpy as np
import tqdm

from habitat.core.logging import logger
from habitat.core.utils import try_cv2_import
from habitat.utils.visualizations import maps
from habitat_sim.utils import viz_utils as vut

cv2 = try_cv2_import()

def make_rgb_palette(n=40):
    HSV_tuples = [(x*1.0/n, 0.8, 0.8) for x in range(n)]
    RGB_map = np.array(list(map(lambda x: colorsys.hsv_to_rgb(*x), HSV_tuples)))
    return RGB_map


def paste_overlapping_image(
    background: np.ndarray,
    foreground: np.ndarray,
    location: Tuple[int, int],
    mask: Optional[np.ndarray] = None,
):
    r"""Composites the foreground onto the background dealing with edge
    boundaries.
    Args:
        background: the background image to paste on.
        foreground: the image to paste. Can be RGB or RGBA. If using alpha
            blending, values for foreground and background should both be
            between 0 and 255. Otherwise behavior is undefined.
        location: the image coordinates to paste the foreground.
        mask: If not None, a mask for deciding what part of the foreground to
            use. Must be the same size as the foreground if provided.
    Returns:
        The modified background image. This operation is in place.
    """
    assert mask is None or mask.shape[:2] == foreground.shape[:2]
    foreground_size = foreground.shape[:2]
    min_pad = (
        max(0, foreground_size[0] // 2 - location[0]),
        max(0, foreground_size[1] // 2 - location[1]),
    )

    max_pad = (
        max(
            0,
            (location[0] + (foreground_size[0] - foreground_size[0] // 2))
            - background.shape[0],
        ),
        max(
            0,
            (location[1] + (foreground_size[1] - foreground_size[1] // 2))
            - background.shape[1],
        ),
    )

    background_patch = background[
        (location[0] - foreground_size[0] // 2 + min_pad[0]) : (
            location[0]
            + (foreground_size[0] - foreground_size[0] // 2)
            - max_pad[0]
        ),
        (location[1] - foreground_size[1] // 2 + min_pad[1]) : (
            location[1]
            + (foreground_size[1] - foreground_size[1] // 2)
            - max_pad[1]
        ),
    ]
    foreground = foreground[
        min_pad[0] : foreground.shape[0] - max_pad[0],
        min_pad[1] : foreground.shape[1] - max_pad[1],
    ]
    if foreground.size == 0 or background_patch.size == 0:
        # Nothing to do, no overlap.
        return background

    if mask is not None:
        mask = mask[
            min_pad[0] : foreground.shape[0] - max_pad[0],
            min_pad[1] : foreground.shape[1] - max_pad[1],
        ]

    if foreground.shape[2] == 4:
        # Alpha blending
        foreground = (
            background_patch.astype(np.int32) * (255 - foreground[:, :, [3]])
            + foreground[:, :, :3].astype(np.int32) * foreground[:, :, [3]]
        ) // 255
    if mask is not None:
        background_patch[mask] = foreground[mask]
    else:
        background_patch[:] = foreground
    return background


def images_to_video(
    images: List[np.ndarray],
    output_dir: str,
    video_name: str,
    fps: int = 15,
    quality: Optional[float] = 5,
    **kwargs,
):
    r"""Calls imageio to run FFMPEG on a list of images. For more info on
    parameters, see https://imageio.readthedocs.io/en/stable/format_ffmpeg.html
    Args:
        images: The list of images. Images should be HxWx3 in RGB order.
        output_dir: The folder to put the video in.
        video_name: The name for the video.
        fps: Frames per second for the video. Not all values work with FFMPEG,
            use at your own risk.
        quality: Default is 5. Uses variable bit rate. Highest quality is 10,
            lowest is 0.  Set to None to prevent variable bitrate flags to
            FFMPEG so you can manually specify them using output_params
            instead. Specifying a fixed bitrate using ‘bitrate’ disables
            this parameter.
    """
    assert 0 <= quality <= 10
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    video_name = video_name.replace(" ", "_").replace("\n", "_") + ".mp4"
    writer = imageio.get_writer(
        os.path.join(output_dir, video_name),
        fps=fps,
        quality=quality,
        **kwargs,
    )
    logger.info(f"Video created: {os.path.join(output_dir, video_name)}")
    for im in tqdm.tqdm(images):
        writer.append_data(im)
    writer.close()


def draw_collision(view: np.ndarray, alpha: float = 0.4) -> np.ndarray:
    r"""Draw translucent red strips on the border of input view to indicate
    a collision has taken place.
    Args:
        view: input view of size HxWx3 in RGB order.
        alpha: Opacity of red collision strip. 1 is completely non-transparent.
    Returns:
        A view with collision effect drawn.
    """
    strip_width = view.shape[0] // 20
    mask = np.ones(view.shape)
    mask[strip_width:-strip_width, strip_width:-strip_width] = 0
    mask = mask == 1
    view[mask] = (alpha * np.array([255, 0, 0]) + (1.0 - alpha) * view)[mask]
    return view


def observations_to_image(observation: Dict, info: Dict, top_down_map_only=False) -> np.ndarray:
    r"""Generate image of single frame from observation and info
    returned from a single environment step().

    Args:
        observation: observation returned from an environment step().
        info: info returned from an environment step().

    Returns:
        generated image of a single frame.
    """
    egocentric_view_l: List[np.ndarray] = []
    if "rgb" in observation:
        rgb = observation["rgb"]
        if not isinstance(rgb, np.ndarray):
            rgb = rgb.cpu().numpy()

        egocentric_view_l.append(rgb)

    # draw depth map if observation has depth info
    if "depth" in observation:
        depth_map = observation["depth"].squeeze() * 255.0
        if not isinstance(depth_map, np.ndarray):
            depth_map = depth_map.cpu().numpy()

        depth_map = depth_map.astype(np.uint8)
        depth_map = np.stack([depth_map for _ in range(3)], axis=2)
        egocentric_view_l.append(depth_map)

    if "semantic" in observation:
        semantic_map = observation["semantic"].squeeze()
        if not isinstance(semantic_map, np.ndarray):
            semantic_map = semantic_map.cpu().numpy()
        colors = make_rgb_palette(45)
        semantic_colors = colors[semantic_map % 45] * 255
        semantic_colors = semantic_colors.astype(np.uint8)
        egocentric_view_l.append(semantic_colors)

        if "gt_semantic" in observation:
            semantic_map = observation["gt_semantic"].squeeze()
            if not isinstance(semantic_map, np.ndarray):
                semantic_map = semantic_map.cpu().numpy()
            colors = make_rgb_palette(45)
            semantic_colors = colors[semantic_map % 45] * 255
            semantic_colors = semantic_colors.astype(np.uint8)
            egocentric_view_l.append(semantic_colors)

    # add image goal if observation has image_goal info
    if "imagegoal" in observation:
        rgb = observation["imagegoal"]
        if not isinstance(rgb, np.ndarray):
            rgb = rgb.cpu().numpy()

        egocentric_view_l.append(rgb)

    assert (
        len(egocentric_view_l) > 0
    ), "Expected at least one visual sensor enabled."
    egocentric_view = np.concatenate(egocentric_view_l, axis=1)

    # draw collision
    if "collisions" in info and info["collisions"]["is_collision"]:
        egocentric_view = draw_collision(egocentric_view)

    frame = egocentric_view

    if "top_down_map" in info:
        top_down_map = maps.colorize_draw_agent_and_fit_to_height(
            info["top_down_map"], egocentric_view.shape[0]
        )
        frame = np.concatenate((egocentric_view, top_down_map), axis=1)
        if top_down_map_only:
            top_down_map = maps.colorize_draw_agent_and_fit_to_height(
                info["top_down_map"], 2048
            )
            frame = top_down_map
    return frame


def append_text_to_image(image: np.ndarray, text: str):
    r"""Appends text underneath an image of size (height, width, channels).
    The returned image has white text on a black background. Uses textwrap to
    split long text into multiple lines.
    Args:
        image: the image to put text underneath
        text: a string to display
    Returns:
        A new image with text inserted underneath the input image
    """
    h, w, c = image.shape
    font_size = 2.5
    font_thickness = 1
    font = cv2.FONT_HERSHEY_SIMPLEX
    blank_image = np.zeros(image.shape, dtype=np.uint8)

    char_size = cv2.getTextSize(" ", font, font_size, font_thickness)[0]
    wrapped_text = textwrap.wrap(text, width=int(w / char_size[0]))

    y = 0
    for line in wrapped_text:
        textsize = cv2.getTextSize(line, font, font_size, font_thickness)[0]
        y += textsize[1] + 15
        x = 10
        cv2.putText(
            blank_image,
            line,
            (x, y),
            font,
            font_size,
            (255, 255, 255),
            font_thickness,
            lineType=cv2.LINE_AA,
        )
    text_image = blank_image[0 : y + 10, 0:w]
    final = np.concatenate((image, text_image), axis=0)
    return final


def make_video_cv2(
    observations, cross_hair=None, prefix="dummy", open_vid=True, fps=15, output_path="./demos/"
):
    sensor_keys = list(observations[0])
    videodims = observations[0][sensor_keys[0]].shape
    videodims = (videodims[1], videodims[0])  # flip to w,h order
    print(videodims, sensor_keys[0])
    video_file = output_path + prefix + ".mp4"
    print("Encoding the video: %s " % video_file)
    writer = vut.get_fast_video_writer(video_file, fps=fps)
    for ob in observations:
        # If in RGB/RGBA format, remove the alpha channel
        rgb_im_1st_person = cv2.cvtColor(ob["rgb"], cv2.COLOR_RGBA2RGB)
        if cross_hair is not None:
            rgb_im_1st_person[
                cross_hair[0] - 2 : cross_hair[0] + 2,
                cross_hair[1] - 2 : cross_hair[1] + 2,
            ] = [255, 0, 0]

        if rgb_im_1st_person.shape[:2] != videodims:
            rgb_im_1st_person = cv2.resize(
                rgb_im_1st_person, videodims, interpolation=cv2.INTER_AREA
            )
        # write the 1st person observation to video
        writer.append_data(rgb_im_1st_person)
    writer.close()

    if open_vid:
        print("Displaying video")
        vut.display_video(video_file)

habitat/utils/visualizations/utils.py
cshizhe/onav_rim
#!/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.

from habitat_baselines.common.base_il_trainer import BaseILTrainer
from habitat_baselines.common.base_trainer import BaseRLTrainer, BaseTrainer
from habitat_baselines.rl.ddppo import DDPPOTrainer  # noqa: F401
from habitat_baselines.rl.ppo.ppo_trainer import PPOTrainer, RolloutStorage
from habitat_baselines.il.disk_based.il_trainer import RearrangementBCTrainer
from habitat_baselines.il.disk_based.il_ddp_trainer import RearrangementBCDistribTrainer
from habitat_baselines.il.env_based.il_trainer import ILEnvTrainer
from habitat_baselines.il.env_based.il_ddp_trainer import ILEnvDDPTrainer

__all__ = [
    "BaseTrainer",
    "BaseRLTrainer",
    "BaseILTrainer",
    "PPOTrainer",
    "RolloutStorage",
    "ILEnvTrainer",
    "ILEnvDDPTrainer",
    "RearrangementBCTrainer",
    "RearrangementBCDistribTrainer",
]

habitat_baselines/__init__.py
cshizhe/onav_rim
#!/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 argparse
import random
from typing import Dict, Optional

import numpy as np
import torch
from gym.spaces import Box
from gym.spaces import Dict as SpaceDict
from gym.spaces import Discrete

import habitat
from habitat.config import Config
from habitat.core.agent import Agent
from habitat.core.simulator import Observations
from habitat_baselines.rl.ppo import PointNavBaselinePolicy
from habitat_baselines.utils.common import batch_obs


def get_default_config() -> Config:
    c = Config()
    c.INPUT_TYPE = "blind"
    c.MODEL_PATH = "data/checkpoints/blind.pth"
    c.RESOLUTION = 256
    c.HIDDEN_SIZE = 512
    c.RANDOM_SEED = 7
    c.PTH_GPU_ID = 0
    c.GOAL_SENSOR_UUID = "pointgoal_with_gps_compass"
    return c


class PPOAgent(Agent):
    def __init__(self, config: Config) -> None:
        spaces = {
            get_default_config().GOAL_SENSOR_UUID: Box(
                low=np.finfo(np.float32).min,
                high=np.finfo(np.float32).max,
                shape=(2,),
                dtype=np.float32,
            )
        }

        if config.INPUT_TYPE in ["depth", "rgbd"]:
            spaces["depth"] = Box(
                low=0,
                high=1,
                shape=(config.RESOLUTION, config.RESOLUTION, 1),
                dtype=np.float32,
            )

        if config.INPUT_TYPE in ["rgb", "rgbd"]:
            spaces["rgb"] = Box(
                low=0,
                high=255,
                shape=(config.RESOLUTION, config.RESOLUTION, 3),
                dtype=np.uint8,
            )
        observation_spaces = SpaceDict(spaces)

        action_spaces = Discrete(4)

        self.device = (
            torch.device("cuda:{}".format(config.PTH_GPU_ID))
            if torch.cuda.is_available()
            else torch.device("cpu")
        )
        self.hidden_size = config.HIDDEN_SIZE

        random.seed(config.RANDOM_SEED)
        torch.random.manual_seed(config.RANDOM_SEED)
        if torch.cuda.is_available():
            torch.backends.cudnn.deterministic = True  # type: ignore

        self.actor_critic = PointNavBaselinePolicy(
            observation_space=observation_spaces,
            action_space=action_spaces,
            hidden_size=self.hidden_size,
        )
        self.actor_critic.to(self.device)

        if config.MODEL_PATH:
            ckpt = torch.load(config.MODEL_PATH, map_location='cpu') #self.device)
            #  Filter only actor_critic weights
            self.actor_critic.load_state_dict(
                {
                    k[len("actor_critic.") :]: v
                    for k, v in ckpt["state_dict"].items()
                    if "actor_critic" in k
                }
            )

        else:
            habitat.logger.error(
                "Model checkpoint wasn't loaded, evaluating " "a random model."
            )

        self.test_recurrent_hidden_states: Optional[torch.Tensor] = None
        self.not_done_masks: Optional[torch.Tensor] = None
        self.prev_actions: Optional[torch.Tensor] = None

    def reset(self) -> None:
        self.test_recurrent_hidden_states = torch.zeros(
            self.actor_critic.net.num_recurrent_layers,
            1,
            self.hidden_size,
            device=self.device,
        )
        self.not_done_masks = torch.zeros(1, 1, device=self.device)
        self.prev_actions = torch.zeros(
            1, 1, dtype=torch.long, device=self.device
        )

    def act(self, observations: Observations) -> Dict[str, int]:
        batch = batch_obs([observations], device=self.device)
        with torch.no_grad():
            (
                _,
                actions,
                _,
                self.test_recurrent_hidden_states,
            ) = self.actor_critic.act(
                batch,
                self.test_recurrent_hidden_states,
                self.prev_actions,
                self.not_done_masks,
                deterministic=False,
            )
            #  Make masks not done till reset (end of episode) will be called
            self.not_done_masks = torch.ones(1, 1, device=self.device)
            self.prev_actions.copy_(actions)  # type: ignore

        return {"action": actions[0][0].item()}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input-type",
        default="blind",
        choices=["blind", "rgb", "depth", "rgbd"],
    )
    parser.add_argument("--model-path", default="", type=str)
    parser.add_argument(
        "--task-config", type=str, default="configs/tasks/pointnav.yaml"
    )
    args = parser.parse_args()

    agent_config = get_default_config()
    agent_config.INPUT_TYPE = args.input_type
    agent_config.MODEL_PATH = args.model_path

    agent = PPOAgent(agent_config)
    benchmark = habitat.Benchmark(config_paths=args.task_config)
    metrics = benchmark.evaluate(agent)

    for k, v in metrics.items():
        habitat.logger.info("{}: {:.3f}".format(k, v))


if __name__ == "__main__":
    main()

habitat_baselines/agents/ppo_agents.py
cshizhe/onav_rim
#!/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
import time
from typing import Any, ClassVar, Dict, List, Tuple, Union

import torch
from numpy import ndarray
from torch import Tensor

from habitat import Config, logger
from habitat.core.env import Env, RLEnv
from habitat.core.vector_env import VectorEnv
from habitat_baselines.common.tensorboard_utils import TensorboardWriter
from habitat_baselines.utils.common import (
    get_checkpoint_id,
    poll_checkpoint_folder,
)


class BaseTrainer:
    r"""Generic trainer class that serves as a base template for more
    specific trainer classes like RL trainer, SLAM or imitation learner.
    Includes only the most basic functionality.
    """

    supported_tasks: ClassVar[List[str]]

    def train(self) -> None:
        raise NotImplementedError

    def _setup_eval_config(self, checkpoint_config: Config) -> Config:
        r"""Sets up and returns a merged config for evaluation. Config
            object saved from checkpoint is merged into config file specified
            at evaluation time with the following overwrite priority:
                  eval_opts > ckpt_opts > eval_cfg > ckpt_cfg
            If the saved config is outdated, only the eval config is returned.

        Args:
            checkpoint_config: saved config from checkpoint.

        Returns:
            Config: merged config for eval.
        """

        config = self.config.clone()

        ckpt_cmd_opts = checkpoint_config.CMD_TRAILING_OPTS
        eval_cmd_opts = config.CMD_TRAILING_OPTS

        try:
            config.merge_from_other_cfg(checkpoint_config)
            config.merge_from_other_cfg(self.config)
            config.merge_from_list(ckpt_cmd_opts)
            config.merge_from_list(eval_cmd_opts)
        except KeyError:
            logger.info("Saved config is outdated, using solely eval config")
            config = self.config.clone()
            config.merge_from_list(eval_cmd_opts)
        config.defrost()
        if config.TASK_CONFIG.DATASET.SPLIT == "train":
            config.TASK_CONFIG.DATASET.SPLIT = "val"
        config.TASK_CONFIG.SIMULATOR.AGENT_0.SENSORS = self.config.SENSORS
        config.freeze()

        return config

    def eval(self) -> None:
        r"""Main method of trainer evaluation. Calls _eval_checkpoint() that
        is specified in Trainer class that inherits from BaseRLTrainer
        or BaseILTrainer

        Returns:
            None
        """
        self.device = (
            torch.device("cuda", self.config.TORCH_GPU_ID)
            if torch.cuda.is_available()
            else torch.device("cpu")
        )

        if "tensorboard" in self.config.VIDEO_OPTION:
            assert (
                len(self.config.TENSORBOARD_DIR) > 0
            ), "Must specify a tensorboard directory for video display"
            os.makedirs(self.config.TENSORBOARD_DIR, exist_ok=True)
        if "disk" in self.config.VIDEO_OPTION:
            assert (
                len(self.config.VIDEO_DIR) > 0
            ), "Must specify a directory for storing videos on disk"

        with TensorboardWriter(
            self.config.TENSORBOARD_DIR, flush_secs=self.flush_secs
        ) as writer:
            if os.path.isfile(self.config.EVAL_CKPT_PATH_DIR):
                # evaluate singe checkpoint
                proposed_index = get_checkpoint_id(
                    self.config.EVAL_CKPT_PATH_DIR
                )
                if proposed_index is not None:
                    ckpt_idx = proposed_index
                else:
                    ckpt_idx = 0
                self._eval_checkpoint(
                    self.config.EVAL_CKPT_PATH_DIR,
                    writer,
                    checkpoint_index=ckpt_idx,
                )
            else:
                # evaluate multiple checkpoints in order
                prev_ckpt_ind = -1
                while True:
                    current_ckpt = None
                    while current_ckpt is None:
                        current_ckpt = poll_checkpoint_folder(
                            self.config.EVAL_CKPT_PATH_DIR, prev_ckpt_ind
                        )
                        time.sleep(2)  # sleep for 2 secs before polling again
                    logger.info(f"=======current_ckpt: {current_ckpt}=======")
                    prev_ckpt_ind += 1
                    self._eval_checkpoint(
                        checkpoint_path=current_ckpt,
                        writer=writer,
                        checkpoint_index=prev_ckpt_ind,
                    )
                    print("eval deone!!!")
                print("eval complete!!!")

    def _eval_checkpoint(
        self,
        checkpoint_path: str,
        writer: TensorboardWriter,
        checkpoint_index: int = 0,
    ) -> None:
        raise NotImplementedError

    def save_checkpoint(self, file_name) -> None:
        raise NotImplementedError

    def load_checkpoint(self, checkpoint_path, *args, **kwargs) -> Dict:
        raise NotImplementedError


class BaseRLTrainer(BaseTrainer):
    r"""Base trainer class for RL trainers. Future RL-specific
    methods should be hosted here.
    """
    device: torch.device  # type: ignore
    config: Config
    video_option: List[str]
    _flush_secs: int

    def __init__(self, config: Config) -> None:
        super().__init__()
        assert config is not None, "needs config file to initialize trainer"
        self.config = config
        self._flush_secs = 30

    @property
    def flush_secs(self):
        return self._flush_secs

    @flush_secs.setter
    def flush_secs(self, value: int):
        self._flush_secs = value

    def train(self) -> None:
        raise NotImplementedError

    def _eval_checkpoint(
        self,
        checkpoint_path: str,
        writer: TensorboardWriter,
        checkpoint_index: int = 0,
    ) -> None:
        r"""Evaluates a single checkpoint. Trainer algorithms should
        implement this.

        Args:
            checkpoint_path: path of checkpoint
            writer: tensorboard writer object for logging to tensorboard
            checkpoint_index: index of cur checkpoint for logging

        Returns:
            None
        """
        raise NotImplementedError

    def save_checkpoint(self, file_name) -> None:
        raise NotImplementedError

    def load_checkpoint(self, checkpoint_path, *args, **kwargs) -> Dict:
        raise NotImplementedError

    @staticmethod
    def _pause_envs(
        envs_to_pause: List[int],
        envs: Union[VectorEnv, RLEnv, Env],
        test_recurrent_hidden_states: Tensor,
        not_done_masks: Tensor,
        current_episode_reward: Tensor,
        prev_actions: Tensor,
        batch: Dict[str, Tensor],
        rgb_frames: Union[List[List[Any]], List[List[ndarray]]],
    ) -> Tuple[
        Union[VectorEnv, RLEnv, Env],
        Tensor,
        Tensor,
        Tensor,
        Tensor,
        Dict[str, Tensor],
        List[List[Any]],
    ]:
        # pausing self.envs with no new episode
        if len(envs_to_pause) > 0:
            state_index = list(range(envs.num_envs))
            for idx in reversed(envs_to_pause):
                state_index.pop(idx)
                envs.pause_at(idx)

            # indexing along the batch dimensions
            test_recurrent_hidden_states = test_recurrent_hidden_states[
                :, state_index
            ]
            not_done_masks = not_done_masks[state_index]
            current_episode_reward = current_episode_reward[state_index]
            prev_actions = prev_actions[state_index]

            for k, v in batch.items():
                batch[k] = v[state_index]

            rgb_frames = [rgb_frames[i] for i in state_index]

        return (
            envs,
            test_recurrent_hidden_states,
            not_done_masks,
            current_episode_reward,
            prev_actions,
            batch,
            rgb_frames,
        )


habitat_baselines/common/base_trainer.py
cshizhe/onav_rim
#!/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.
r"""
This file hosts task-specific or trainer-specific environments for trainers.
All environments here should be a (direct or indirect ) subclass of Env class
in habitat. Customized environments should be registered using
``@baseline_registry.register_env(name="myEnv")` for reusability
"""
import numpy as np
from collections import defaultdict
from logging import Logger, log
from typing import Optional, Type
import skimage.morphology

import habitat
from habitat import Config, Dataset, logger
from habitat_baselines.common.baseline_registry import baseline_registry
from habitat.tasks.nav.nav import (
    TopDownMap
)
from habitat.tasks.pickplace.pickplace import (
    GoalObjectVisible,
    AgentToReceptacleDistance,
    ObjectToReceptacleDistance,
    AgentToObjectDistance
)

from habitat.tasks.utils import cartesian_to_polar
from habitat.utils.geometry_utils import (
    quaternion_from_coeff,
    quaternion_rotate_vector,
)
import habitat.utils.pose as pu


def get_env_class(env_name: str) -> Type[habitat.RLEnv]:
    r"""Return environment class based on name.

    Args:
        env_name: name of the environment.

    Returns:
        Type[habitat.RLEnv]: env class.
    """
    return baseline_registry.get_env(env_name)


@baseline_registry.register_env(name="NavRLEnv")
class NavRLEnv(habitat.RLEnv):
    def __init__(self, config: Config, dataset: Optional[Dataset] = None):
        self._rl_config = config.RL
        self._core_env_config = config.TASK_CONFIG
        self._reward_measure_name = self._rl_config.REWARD_MEASURE
        self._success_measure_name = self._rl_config.SUCCESS_MEASURE

        self._previous_measure = None
        self._previous_action = None
        super().__init__(self._core_env_config, dataset)

    def reset(self, **kwargs):
        self._previous_action = None
        observations = super().reset(**kwargs)
        self._previous_measure = self._env.get_metrics()[
            self._reward_measure_name
        ]
        return observations

    def step(self, *args, **kwargs):
        self._previous_action = kwargs["action"]
        return super().step(*args, **kwargs)

    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_measure = self._env.get_metrics()[self._reward_measure_name]

        reward += self._previous_measure - current_measure
        self._previous_measure = current_measure

        if self._episode_success():
            reward += self._rl_config.SUCCESS_REWARD

        return reward

    def _episode_success(self):
        return self._env.get_metrics()[self._success_measure_name]

    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()


@baseline_registry.register_env(name="RearrangementRLEnv")
class RearrangementRLEnv(habitat.RLEnv):
    def __init__(self, config: Config, dataset: Optional[Dataset] = None):
        self._rl_config = config.RL
        self._core_env_config = config.TASK_CONFIG
        self._reward_measure_name = self._rl_config.REWARD_MEASURE
        self._success_measure_name = self._rl_config.SUCCESS_MEASURE
        super().__init__(self._core_env_config, dataset)

    def reset(self):
        observations = super().reset()
        return observations

    def step(self, *args, **kwargs):
        return super().step(*args, **kwargs)

    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
        # Get reward from RearrangementReward measure
        current_measure = self._env.get_metrics()[self._reward_measure_name]
        reward += current_measure
        return reward

    def _episode_success(self):
        return self._env.get_metrics()[self._success_measure_name]
    

    def _end_early(self):
        end_early = 0
        end_early += self._env.get_metrics()[AgentToObjectDistance.cls_uuid]["was_nan"]
        end_early += self._env.get_metrics()[AgentToReceptacleDistance.cls_uuid]["was_nan"]
        end_early += self._env.get_metrics()[ObjectToReceptacleDistance.cls_uuid]["was_nan"]
        return end_early

    def get_done(self, observations):
        done = False
        if self._end_early():
            logger.info("End early")

        if self._env.episode_over or self._episode_success() or self._end_early():
            done = True
        return done

    def get_info(self, observations):
        return self.habitat_env.get_metrics()


@baseline_registry.register_env(name="ExploreThenNavRLEnv")
class ExploreThenNavRLEnv(NavRLEnv):
    r"""
        We want to train an agent that overfits less. We provide an exploration reward and a delayed gratification success reward.
        This is to avoid weird shaping loss.
        We provide quickly attenuating coverage reward, and a gradually increasing success reward.
    """
    def __init__(self, config, dataset=None): # add coverage to the metrics
        super().__init__(config, dataset)
        self.step_penalty = 1
        self._goal_was_seen = False # Switch for turning off coverage, turning on success shaping
        self.visit_bonus = 0
        self._previous_view = 0

    def step(self, *args, **kwargs):
        self.step_penalty *= self._rl_config.COVERAGE_ATTENUATION
        return super().step(*args, **kwargs)

    def reset(self):
        self.step_penalty = 1
        self._goal_was_seen = False
        self._previous_view = 0
        # self.visit_bonus = 0
        return super().reset()

    def get_reward_range(self):
        old_low, old_hi = super().get_reward_range()
        return old_low, old_hi + self._rl_config.COVERAGE_REWARD

    def get_reward(self, observations):
        # Distance reward after agent sees the object
        if self._goal_was_seen:
            return super().get_reward(observations)
            # return reward
            # if self._episode_success():
            #     return reward + self.visit_bonus * self._rl_config.COVERAGE_BONUS_SCALE

        # ! If there's no GoalObjectVisible measure (i.e. eval), automatically use shaping
        if GoalObjectVisible.cls_uuid not in self.habitat_env.get_metrics() or self.habitat_env.get_metrics()[GoalObjectVisible.cls_uuid] > self._rl_config.EXPLORE_GOAL_SEEN_THRESHOLD:
            self._goal_was_seen = True
            super().get_reward(observations) # ! Hack -- clear the shaping reward from exploration phase

        # distance = self._env.get_metrics()["distance_to_goal"]
        # if distance < self._rl_config.COVERAGE_FALLOFF_RADIUS or self.step_penalty < 0.01: # ~150 steps have passed
            # self.visit_bonus = math.sqrt(self.habitat_env.get_metrics()["coverage"]["reached"])

        if self._rl_config.COVERAGE_TYPE == "VISIT":
            visit = self.habitat_env.get_metrics()["coverage"]["visit_count"]
            # logger.info("Coverage reward {}".format(self.step_penalty * self._rl_config.COVERAGE_REWARD / (visit ** self._rl_config.COVERAGE_VISIT_EXP)))
            # No reward till goal is not seen
            return 0
            # return self.step_penalty * self._rl_config.COVERAGE_REWARD / (visit ** self._rl_config.COVERAGE_VISIT_EXP)
            # return self.step_penalty * self._rl_config.COVERAGE_REWARD / (visit ** self._rl_config.COVERAGE_VISIT_EXP)
        else: # VIEW
            # print(self.habitat_env.get_metrics().keys())
            # Here we provide a bonus for the new view
            reward = 0
            map_measures = self.habitat_env.get_metrics()[TopDownMap.cls_uuid]
            # print(map_measures.keys())
            if map_measures:
                # print(map_measures["fog_of_war_mask"].size)
                # 18K / 1.8M, 40K  / 2.3M, 6K / 2M, 9K / 1M -> let's just divide by 50K
                # print(map_measures["fog_of_war_mask"].sum())
                explore_view = np.log(1 +
                    map_measures["fog_of_war_mask"].sum() / 50000.0
                    # / map_measures["fog_of_war_mask"].size
                ) / np.log(2)
                # import pdb
                # pdb.set_trace()
            else:
                explore_view = 0
            if self._previous_view > 0:
                reward += (
                    explore_view - self._previous_view
                ) * self._rl_config.COVERAGE_REWARD * self.step_penalty
            # print(reward)
            self._previous_view = explore_view
            return reward
habitat_baselines/common/environments.py
cshizhe/onav_rim
import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F

import habitat.utils.depth_utils as du


def get_local_map_boundaries(
    agent_loc, local_sizes, full_sizes, global_downscaling=2,
):
    loc_r, loc_c = agent_loc
    local_w, local_h = local_sizes
    full_w, full_h = full_sizes
    
    if global_downscaling > 1:
        gx1, gy1 = loc_r - local_w // 2, loc_c - local_h // 2
        gx2, gy2 = gx1 + local_w, gy1 + local_h
        if gx1 < 0:
            gx1, gx2 = 0, local_w
        if gx2 > full_w:
            gx1, gx2 = full_w - local_w, full_w

        if gy1 < 0:
            gy1, gy2 = 0, local_h
        if gy2 > full_h:
            gy1, gy2 = full_h - local_h, full_h
    else:
        gx1, gx2, gy1, gy2 = 0, full_w, 0, full_h

    return [gx1, gx2, gy1, gy2]

    
def init_map_and_pose(map_cfg, device, num_scenes=1, remove_batch_dim=False):
    # Initialize map variables:
    # Full map consists of multiple channels containing the following:
    # 1. Obstacle Map
    # 2. Exploread Area
    nc = 2
    visited_locs_in_full_map = [[] for _ in range(num_scenes)]

    # Calculating full and local map sizes
    map_size = map_cfg.map_size_cm // map_cfg.map_resolution
    full_w, full_h = map_size, map_size
    local_w = int(full_w / map_cfg.global_downscaling)
    local_h = int(full_h / map_cfg.global_downscaling)

    # Initializing full and local map
    full_map = torch.zeros(num_scenes, nc, full_w, full_h).float().to(device)
    local_map = torch.zeros(num_scenes, nc, local_w, local_h).float().to(device)

    # Initial full and local pose: (x, y, o) x,y: meters, o: degree
    full_pose = torch.zeros(num_scenes, 3).float().to(device)
    local_pose = torch.zeros(num_scenes, 3).float().to(device)

    # Origin of local map: (column, row, 0) meters (upper left point)
    origins = np.zeros((num_scenes, 3))

    # Local Map Boundaries: (x1, x2, y1, y2) map unit
    lmb = np.zeros((num_scenes, 4)).astype(int)
        
    for e in range(num_scenes):
        full_map[e].fill_(0.)
        full_pose[e].fill_(0.)
        # an agent is initialized at the center of the map
        full_pose[e, :2] = map_cfg.map_size_cm / 100.0 / 2.0

        locs = full_pose[e].cpu().numpy()
    
        r, c = locs[0], locs[1]
        loc_r, loc_c = [int(r * 100.0 / map_cfg.map_resolution),
                        int(c * 100.0 / map_cfg.map_resolution)]

        # initialize explored area and current agent location
        visited_locs_in_full_map[e].append((loc_r, loc_c))

        lmb[e] = get_local_map_boundaries(
            (loc_r, loc_c), (local_w, local_h), (full_w, full_h),
            map_cfg.global_downscaling
        )

        origins[e] = [lmb[e][0] * map_cfg.map_resolution / 100.0,
                      lmb[e][2] * map_cfg.map_resolution / 100.0, 0.]

        local_map[e] = full_map[e, :,
                                lmb[e, 0]:lmb[e, 1],
                                lmb[e, 2]:lmb[e, 3]]
        local_pose[e] = full_pose[e] - \
            torch.from_numpy(origins[e]).to(device).float()
    
    if num_scenes == 1 and remove_batch_dim:
        return full_map[0], full_pose[0], local_map[0], local_pose[0], \
               lmb[0], origins[0], visited_locs_in_full_map[0]
    return full_map, full_pose, local_map, local_pose, lmb, origins, \
           visited_locs_in_full_map
           
  
def get_grid(pose, grid_size, device):
    """
    Input:
        `pose` FloatTensor(bs, 3)
        `grid_size` 4-tuple (bs, _, grid_h, grid_w)
        `device` torch.device (cpu or gpu)
    Output:
        `rot_grid` FloatTensor(bs, grid_h, grid_w, 2)
        `trans_grid` FloatTensor(bs, grid_h, grid_w, 2)

    """
    pose = pose.float()
    x = pose[:, 0]
    y = pose[:, 1]
    t = pose[:, 2]

    bs = x.size(0)
    t = t * np.pi / 180.
    cos_t = t.cos()
    sin_t = t.sin()

    theta11 = torch.stack([cos_t, -sin_t,
                           torch.zeros(cos_t.shape).float().to(device)], 1)
    theta12 = torch.stack([sin_t, cos_t,
                           torch.zeros(cos_t.shape).float().to(device)], 1)
    theta1 = torch.stack([theta11, theta12], 1)

    theta21 = torch.stack([torch.ones(x.shape).to(device),
                           -torch.zeros(x.shape).to(device), x], 1)
    theta22 = torch.stack([torch.zeros(x.shape).to(device),
                           torch.ones(x.shape).to(device), y], 1)
    theta2 = torch.stack([theta21, theta22], 1)

    rot_grid = F.affine_grid(theta1, torch.Size(grid_size))
    trans_grid = F.affine_grid(theta2, torch.Size(grid_size))

    return rot_grid, trans_grid


  
class NavMapping(nn.Module):

    def __init__(self, config, device):
        super().__init__()
        
        self.num_processes = 1

        self.device = device
        self.screen_h = config.TASK_CONFIG.SIMULATOR.DEPTH_SENSOR.HEIGHT
        self.screen_w = config.TASK_CONFIG.SIMULATOR.DEPTH_SENSOR.WIDTH
        self.fov = config.TASK_CONFIG.SIMULATOR.DEPTH_SENSOR.HFOV
        
        self.resolution = config.MAPPING.map_resolution
        self.z_resolution = config.MAPPING.map_resolution
        self.map_size_cm = config.MAPPING.map_size_cm // config.MAPPING.global_downscaling
        self.vision_range = config.MAPPING.vision_range
        self.du_scale = config.MAPPING.du_scale
        self.exp_pred_threshold = config.MAPPING.exp_pred_threshold
        self.map_pred_threshold = config.MAPPING.map_pred_threshold

        self.max_height = int(360 / self.z_resolution)
        self.min_height = int(-40 / self.z_resolution)
        self.agent_height = config.TASK_CONFIG.SIMULATOR.AGENT_0.HEIGHT * 100.
        self.shift_loc = [self.vision_range * self.resolution // 2, 0, 0]
        self.camera_matrix = du.get_camera_matrix(
            self.screen_w, self.screen_h, self.fov
        )

        vr = self.vision_range

        self.init_grid = torch.zeros(
            self.num_processes, 1, vr, vr,
            self.max_height - self.min_height
        ).float().to(self.device)

        self.feat = torch.ones(
            self.num_processes, 1,
            self.screen_h // self.du_scale * self.screen_w // self.du_scale
        ).float().to(self.device)

    def forward(self, obs, pose_obs, maps_last, poses_last, camera_elevation_degree=0):
        depth = torch.from_numpy(obs['processed_depth']).to(self.device).unsqueeze(0)
                
        point_cloud_t = du.get_point_cloud_from_z_t(
            depth, self.camera_matrix, self.device, scale=self.du_scale
        )

        # rotate along x-axis, change elevation
        agent_view_t = du.transform_camera_view_t(
            point_cloud_t, self.agent_height, camera_elevation_degree, self.device
        )

        # do not rotate along z-axis: shift_loc[-1]=0 (rotate the map later)
        agent_view_centered_t = du.transform_pose_t(
            agent_view_t, self.shift_loc, self.device
        )

        max_h = self.max_height
        min_h = self.min_height
        xy_resolution = self.resolution
        z_resolution = self.z_resolution
        vision_range = self.vision_range
        XYZ_cm_std = agent_view_centered_t.float()
        XYZ_cm_std[..., :2] = (XYZ_cm_std[..., :2] / xy_resolution)
        XYZ_cm_std[..., :2] = (XYZ_cm_std[..., :2] -
                               vision_range // 2.) / vision_range * 2.
        XYZ_cm_std[..., 2] = XYZ_cm_std[..., 2] / z_resolution
        XYZ_cm_std[..., 2] = (XYZ_cm_std[..., 2] -
                              (max_h + min_h) // 2.) / (max_h - min_h) * 2.
        
        XYZ_cm_std = XYZ_cm_std.permute(0, 3, 1, 2)
        XYZ_cm_std = XYZ_cm_std.view(XYZ_cm_std.shape[0],
                                     XYZ_cm_std.shape[1],
                                     XYZ_cm_std.shape[2] * XYZ_cm_std.shape[3])

        voxels = du.splat_feat_nd(
            self.init_grid * 0., self.feat, XYZ_cm_std
        ).transpose(2, 3) # tranpose because x-axis should denote columns instead of rows
        voxels = torch.flip(voxels, (2, ))  # facing north at the front
        # (num_scenes, channel, hx, hy, hz)

        min_z = int(25 / z_resolution - min_h)
        max_z = int((self.agent_height + 1) / z_resolution - min_h)

        agent_height_proj = voxels[..., min_z:max_z].sum(4)
        all_height_proj = voxels.sum(4)

        fp_map_pred = agent_height_proj[:, 0:1, :, :]
        fp_exp_pred = all_height_proj[:, 0:1, :, :]
        fp_map_pred = fp_map_pred / self.map_pred_threshold
        fp_exp_pred = fp_exp_pred / self.exp_pred_threshold
        fp_map_pred = torch.clamp(fp_map_pred, min=0.0, max=1.0)
        fp_exp_pred = torch.clamp(fp_exp_pred, min=0.0, max=1.0)

        pose_pred = poses_last

        agent_view = torch.zeros(self.num_processes, 2,
                                 self.map_size_cm // self.resolution,
                                 self.map_size_cm // self.resolution
                                 ).to(self.device)

        x1 = self.map_size_cm // (self.resolution * 2) - self.vision_range // 2
        x2 = x1 + self.vision_range
        # y1 = self.map_size_cm // (self.resolution * 2)
        y1 = self.map_size_cm // (self.resolution * 2) - self.vision_range # facing north at the front
        y2 = y1 + self.vision_range
        agent_view[:, 0:1, y1:y2, x1:x2] = fp_map_pred
        agent_view[:, 1:2, y1:y2, x1:x2] = fp_exp_pred

        corrected_pose = pose_obs

        def get_new_pose_batch(pose, rel_pose_change):

            pose[:, 1] += rel_pose_change[:, 0] * \
                torch.sin(pose[:, 2] / 57.29577951308232) \
                + rel_pose_change[:, 1] * \
                torch.cos(pose[:, 2] / 57.29577951308232)
            pose[:, 0] += rel_pose_change[:, 0] * \
                torch.cos(pose[:, 2] / 57.29577951308232) \
                - rel_pose_change[:, 1] * \
                torch.sin(pose[:, 2] / 57.29577951308232)
            pose[:, 2] += rel_pose_change[:, 2] * 57.29577951308232

            pose[:, 2] = torch.fmod(pose[:, 2] - 180.0, 360.0) + 180.0
            pose[:, 2] = torch.fmod(pose[:, 2] + 180.0, 360.0) - 180.0

            return pose

        current_poses = get_new_pose_batch(poses_last, corrected_pose)
        st_pose = current_poses.clone().detach()

        st_pose[:, :2] = - (st_pose[:, :2] * 100.0 / self.resolution
                            - self.map_size_cm // (self.resolution * 2)) /\
            (self.map_size_cm // (self.resolution * 2))
        # st_pose[:, 2] = 90. - (st_pose[:, 2])
        st_pose[:, :2] = torch.flip(st_pose[:, :2], (1, )) # transpose x and y axis
        st_pose[:, 2] = st_pose[:, 2] - 180

        rot_mat, trans_mat = get_grid(st_pose, agent_view.size(),
                                      self.device)

        rotated = F.grid_sample(agent_view, rot_mat, align_corners=True)
        translated = F.grid_sample(rotated, trans_mat, align_corners=True)

        maps2 = torch.cat((maps_last.unsqueeze(1), translated.unsqueeze(1)), 1)

        map_pred, _ = torch.max(maps2, 1)

        return fp_map_pred, map_pred, pose_pred, current_poses

habitat_baselines/common/nav_mapping.py
cshizhe/onav_rim
#!/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.

r"""This module defines various ObservationTransformers that can be used
to transform the output of the simulator before they are fed into the
policy of the neural network. This can include various useful preprocessing
including faking a semantic sensor using RGB input and MaskRCNN or faking
a depth sensor using RGB input. You can also stich together multiple sensors.
This code runs on the batched of inputs to these networks efficiently.
ObservationTransformer all run as nn.modules and can be used for encoders or
any other neural networks preprocessing steps.
Assumes the input is on CUDA.

They also implement a function that transforms that observation space so help
fake or modify sensor input from the simulator.

This module API is experimental and likely to change
"""
import abc
import copy
from logging import log
import numbers
from enum import Enum
from typing import Dict, Iterable, List, Optional, Tuple, Union

import numpy as np
import torch
from gym import spaces
from torch import nn

from habitat.config import Config
from habitat.core.logging import logger
from habitat_baselines.common.baseline_registry import baseline_registry
from habitat_baselines.utils.common import (
    center_crop,
    get_image_height_width,
    image_resize_shortest_edge,
    overwrite_gym_box_shape,
)


class ObservationTransformer(nn.Module, metaclass=abc.ABCMeta):
    """This is the base ObservationTransformer class that all other observation
    Transformers should extend. from_config must be implemented by the transformer.
    transform_observation_space is only needed if the observation_space ie.
    (resolution, range, or num of channels change)."""

    def transform_observation_space(
        self, observation_space: spaces.Dict, **kwargs
    ):
        return observation_space

    @classmethod
    @abc.abstractmethod
    def from_config(cls, config: Config):
        pass

    def forward(
        self, observations: Dict[str, torch.Tensor]
    ) -> Dict[str, torch.Tensor]:
        return observations


@baseline_registry.register_obs_transformer()
class ResizeShortestEdge(ObservationTransformer):
    r"""An nn module the resizes your the shortest edge of the input while maintaining aspect ratio.
    This module assumes that all images in the batch are of the same size.
    """

    def __init__(
        self,
        size: int,
        channels_last: bool = True,
        trans_keys: Tuple[str] = ("rgb", "depth", "semantic"),
    ):
        """Args:
        size: The size you want to resize the shortest edge to
        channels_last: indicates if channels is the last dimension
        """
        super(ResizeShortestEdge, self).__init__()
        self._size: int = size
        self.channels_last: bool = channels_last
        self.trans_keys: Tuple[str] = trans_keys

    def transform_observation_space(
        self,
        observation_space: spaces.Dict,
    ):
        size = self._size
        observation_space = copy.deepcopy(observation_space)
        if size:
            for key in observation_space.spaces:
                if key in self.trans_keys:
                    # In the observation space dict, the channels are always last
                    h, w = get_image_height_width(
                        observation_space.spaces[key], channels_last=True
                    )
                    if size == min(h, w):
                        continue
                    scale = size / min(h, w)
                    new_h = int(h * scale)
                    new_w = int(w * scale)
                    new_size = (new_h, new_w)
                    logger.info(
                        "Resizing observation of %s: from %s to %s"
                        % (key, (h, w), new_size)
                    )
                    observation_space.spaces[key] = overwrite_gym_box_shape(
                        observation_space.spaces[key], new_size
                    )
        return observation_space

    def _transform_obs(self, obs: torch.Tensor) -> torch.Tensor:
        return image_resize_shortest_edge(
            obs, self._size, channels_last=self.channels_last
        )

    @torch.no_grad()
    def forward(
        self, observations: Dict[str, torch.Tensor]
    ) -> Dict[str, torch.Tensor]:
        if self._size is not None:
            observations.update(
                {
                    sensor: self._transform_obs(observations[sensor])
                    for sensor in self.trans_keys
                    if sensor in observations
                }
            )
        return observations

    @classmethod
    def from_config(cls, config: Config):
        return cls(config.RL.POLICY.OBS_TRANSFORMS.RESIZE_SHORTEST_EDGE.SIZE)


@baseline_registry.register_obs_transformer()
class CenterCropper(ObservationTransformer):
    """An observation transformer is a simple nn module that center crops your input."""

    def __init__(
        self,
        size: Union[int, Tuple[int, int]],
        channels_last: bool = True,
        trans_keys: Tuple[str] = ("rgb", "depth", "semantic"),
    ):
        """Args:
        size: A sequence (h, w) or int of the size you wish to resize/center_crop.
                If int, assumes square crop
        channels_list: indicates if channels is the last dimension
        trans_keys: The list of sensors it will try to centercrop.
        """
        super().__init__()
        if isinstance(size, numbers.Number):
            size = (int(size), int(size))
        assert len(size) == 2, "forced input size must be len of 2 (h, w)"
        self._size = size
        self.channels_last = channels_last
        self.trans_keys = trans_keys  # TODO: Add to from_config constructor

    def transform_observation_space(
        self,
        observation_space: spaces.Dict,
    ):
        size = self._size
        observation_space = copy.deepcopy(observation_space)
        if size:
            for key in observation_space.spaces:
                if (
                    key in self.trans_keys
                    and observation_space.spaces[key].shape[-3:-1] != size
                ):
                    h, w = get_image_height_width(
                        observation_space.spaces[key], channels_last=True
                    )
                    logger.info(
                        "Center cropping observation size of %s from %s to %s"
                        % (key, (h, w), size)
                    )

                    observation_space.spaces[key] = overwrite_gym_box_shape(
                        observation_space.spaces[key], size
                    )
        return observation_space

    def _transform_obs(self, obs: torch.Tensor) -> torch.Tensor:
        return center_crop(
            obs,
            self._size,
            channels_last=self.channels_last,
        )

    @torch.no_grad()
    def forward(
        self, observations: Dict[str, torch.Tensor]
    ) -> Dict[str, torch.Tensor]:
        if self._size is not None:
            observations.update(
                {
                    sensor: self._transform_obs(observations[sensor])
                    for sensor in self.trans_keys
                    if sensor in observations
                }
            )
        return observations

    @classmethod
    def from_config(cls, config: Config):
        cc_config = config.RL.POLICY.OBS_TRANSFORMS.CENTER_CROPPER
        return cls(
            (
                cc_config.HEIGHT,
                cc_config.WIDTH,
            )
        )

@baseline_registry.register_obs_transformer(name="RGBRandomNoise")
class RGBRandomNoise(ObservationTransformer):
    """An observation transformer is a simple nn module that adds random noise to your input."""

    def __init__(
        self,
        mean: float,
        std: float,
        trans_keys: Tuple[str] = ("rgb"),
    ):
        """Args:
        size: A sequence (h, w) or int of the size you wish to resize/center_crop.
                If int, assumes square crop
        channels_list: indicates if channels is the last dimension
        trans_keys: The list of sensors it will try to centercrop.
        """
        super().__init__()
        self.mean = mean
        self.std = std
        self.trans_keys = trans_keys  # TODO: Add to from_config constructor

    def transform_observation_space(
        self,
        observation_space: spaces.Dict,
    ):
        size = self._size
        observation_space = copy.deepcopy(observation_space)
        return observation_space

    def _transform_obs(self, obs: torch.Tensor) -> torch.Tensor:
        obs_shape = obs.size()
        mean_vector = torch.ones(obs_shape) * self.mean
        std_vector = torch.ones(obs_shape) * self.std
        random_noise = torch.normal(mean=mean_vector, std=std_vector)
        transformed_obs = random_noise + obs      
        return transformed_obs

    @torch.no_grad()
    def forward(
        self, observations: Dict[str, torch.Tensor]
    ) -> Dict[str, torch.Tensor]:
        if self.mean is not None:
            observations.update(
                {
                    sensor: self._transform_obs(observations[sensor])
                    for sensor in self.trans_keys
                    if sensor in observations
                }
            )
        return observations

    @classmethod
    def from_config(cls, config: Config):
        cc_config = config.IL.OBS_TRANSFORMS.RGBRandomNoise
        return cls(
            cc_config.MEAN,
            cc_config.STD,
        )


class _DepthFrom(Enum):
    Z_VAL = 0
    OPTI_CENTER = 1


class CameraProjection(metaclass=abc.ABCMeta):
    """This is the base CameraProjection class that converts
    projection model of images into different one. This can be used for
    conversion between cubemap, equirect, fisheye images, etc.
    projection that project 3D points onto the image plane and
    unprojection that project image points onto unit sphere
    must be implemented."""

    def __init__(
        self,
        img_h: int,
        img_w: int,
        R: Optional[torch.Tensor] = None,
        depth_from: _DepthFrom = _DepthFrom.OPTI_CENTER,
    ):
        """Args:
        img_h: (int) the height of camera image
        img_w: (int) the width of camera image
        R: (torch.Tensor) 3x3 rotation matrix of camera
        depth_from: (_DepthFrom) the depth from z value or optical center
        """
        self.img_h = img_h
        self.img_w = img_w
        self.depth_from = depth_from

        # Camera rotation: points in world coord = R @ points in camera coord
        if R is not None:
            self.R = R.float()
        else:
            self.R = None

    @abc.abstractmethod
    def projection(
        self, world_pts: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Project points in world coord onto image planes.
        Args:
            world_pts: 3D points in world coord
        Returns:
            proj_pts: Projected points for grid_sample, -1 <= proj_pts <= 1
            valid_mask: True if the point is valid (inside FoV)
        """

    @abc.abstractmethod
    def unprojection(
        self, with_rotation: bool = True
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Unproject 2D image points onto unit sphere.
        Args:
            with_rotation: If True, unprojected points is in world coord.
                           If False, unprojected points is in camera coord.
        Returns:
            unproj_pts: Unprojected 3D points on unit sphere
            valid_mask: True if the point is valid (inside FoV)
        """

    @property
    def rotation(self):
        """Camera rotation: points in world coord = R @ points in camera coord"""
        if self.R is None:
            return torch.eye(3, dtype=torch.float32)
        else:
            return self.R

    @property
    def shape(self):
        """Camera image shape: (img_h, img_w)"""
        return (self.img_h, self.img_w)

    def size(self):
        """Camera image shape: (img_h, img_w)"""
        return self.shape

    def camcoord2worldcoord(self, pts: torch.Tensor):
        """Convert points in camera coords into points in world coords.
        Args:
            pts: 3D points in camera coords
        Returns:
            rotated_pts: 3D points in world coords
        """
        if self.R is None:
            return pts
        else:
            # Rotate points according to camera rotation
            _h, _w, _ = pts.shape
            # points in world coord = R @ points in camera coord
            rotated_pts = torch.matmul(pts.view((-1, 3)), self.R.T)
            return rotated_pts.view(_h, _w, 3)

    def worldcoord2camcoord(self, pts: torch.Tensor):
        """Convert points in world coords into points in camera coords.
        Args:
            pts: 3D points in world coords
        Returns:
            rotated_pts: 3D points in camera coords
        """
        if self.R is None:
            return pts
        else:
            # Rotate points according to camera rotation
            _h, _w, _ = pts.shape
            # points in camera coord = R.T @ points in world coord
            rotated_pts = torch.matmul(pts.view((-1, 3)), self.R)
            return rotated_pts.view(_h, _w, 3)


class PerspectiveProjection(CameraProjection):
    """This is the perspective camera projection class."""

    def __init__(
        self,
        img_h: int,
        img_w: int,
        f: Optional[float] = None,
        R: Optional[torch.Tensor] = None,
    ):
        """Args:
        img_h: (int) the height of camera image
        img_w: (int) the width of camera image
        f: (float) the focal length of camera
        R: (torch.Tensor) 3x3 rotation matrix of camera
        """
        super(PerspectiveProjection, self).__init__(
            img_h, img_w, R, _DepthFrom.Z_VAL
        )
        if f is None:
            self.f = max(img_h, img_w) / 2
        else:
            self.f = f

    def projection(
        self, world_pts: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        # Rotate world points according to camera rotation
        world_pts = self.worldcoord2camcoord(world_pts)

        # Project points onto image plane
        img_pts = self.f * world_pts / torch.abs(world_pts[..., 2:3])
        cx = self.img_w / 2
        cy = self.img_h / 2
        u = img_pts[..., 0] + cx
        v = img_pts[..., 1] + cy

        # For grid_sample, -1 <= proj_pts <= 1
        mapx = 2 * u / self.img_w - 1.0
        mapy = 2 * v / self.img_h - 1.0
        proj_pts = torch.stack([mapx, mapy], dim=-1)

        # Valid mask
        valid_mask = torch.abs(proj_pts).max(-1)[0] <= 1  # -1 <= grid.xy <= 1
        valid_mask *= img_pts[..., 2] > 0
        return proj_pts, valid_mask

    def unprojection(
        self, with_rotation: bool = True
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        v, u = torch.meshgrid(
            torch.arange(self.img_h), torch.arange(self.img_w)
        )
        x = (u + 0.5) - self.img_w / 2
        y = (v + 0.5) - self.img_h / 2
        z = torch.full_like(x, self.f, dtype=torch.float)
        unproj_pts = torch.stack([x, y, z], dim=-1)
        # Project on unit shpere
        unproj_pts /= torch.norm(unproj_pts, dim=-1, keepdim=True)
        # All points in image are valid
        valid_mask = torch.full(unproj_pts.shape[:2], True, dtype=torch.bool)

        # Rotate unproj_pts points according to camera rotation
        if with_rotation:
            unproj_pts = self.camcoord2worldcoord(unproj_pts)

        return unproj_pts, valid_mask


class EquirectProjection(CameraProjection):
    """This is the equirectanglar camera projection class."""

    def __init__(
        self, img_h: int, img_w: int, R: Optional[torch.Tensor] = None
    ):
        """Args:
        img_h: (int) the height of equirectanglar camera image
        img_w: (int) the width of equirectanglar camera image
        R: (torch.Tensor) 3x3 rotation matrix of camera
        """
        super(EquirectProjection, self).__init__(img_h, img_w, R)

    def projection(
        self, world_pts: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        # Rotate world points according to camera rotation
        world_pts = self.worldcoord2camcoord(world_pts)

        x, y, z = world_pts[..., 0], world_pts[..., 1], world_pts[..., 2]
        # x,y,z to theta, phi
        theta = torch.atan2(x, z)
        c = torch.sqrt(x * x + z * z)
        phi = torch.atan2(y, c)

        # For grid_sample, -1 <= proj_pts <= 1
        mapx = theta / np.pi
        mapy = phi / (np.pi / 2)
        proj_pts = torch.stack([mapx, mapy], dim=-1)

        # All points in image are valid
        valid_mask = torch.full(proj_pts.shape[:2], True, dtype=torch.bool)
        return proj_pts, valid_mask

    def unprojection(
        self, with_rotation: bool = True
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        theta_map, phi_map = self.get_theta_phi_map(self.img_h, self.img_w)
        unproj_pts = self.angle2sphere(theta_map, phi_map)
        # All points in image are valid
        valid_mask = torch.full(unproj_pts.shape[:2], True, dtype=torch.bool)
        # Rotate unproj_pts points according to camera rotation
        if with_rotation:
            unproj_pts = self.camcoord2worldcoord(unproj_pts)
        return unproj_pts, valid_mask

    def get_theta_phi_map(
        self, img_h: int, img_w: int
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Get theta and phi map for equirectangular image.
        PI < theta_map < PI,  PI/2 < phi_map < PI/2
        """
        phi, theta = torch.meshgrid(torch.arange(img_h), torch.arange(img_w))
        theta_map = (theta + 0.5) * 2 * np.pi / img_w - np.pi
        phi_map = (phi + 0.5) * np.pi / img_h - np.pi / 2
        return theta_map, phi_map

    def angle2sphere(
        self, theta_map: torch.Tensor, phi_map: torch.Tensor
    ) -> torch.Tensor:
        """Project points on unit sphere based on theta and phi map."""
        sin_theta = torch.sin(theta_map)
        cos_theta = torch.cos(theta_map)
        sin_phi = torch.sin(phi_map)
        cos_phi = torch.cos(phi_map)
        return torch.stack(
            [cos_phi * sin_theta, sin_phi, cos_phi * cos_theta], dim=-1
        )


class FisheyeProjection(CameraProjection):
    r"""This is the fisheye camera projection class.
    The camera model is based on the Double Sphere Camera Model (Usenko et. al.;3DV 2018).
    Paper: https://arxiv.org/abs/1807.08957
    Implementation: https://github.com/matsuren/dscamera
    """

    def __init__(
        self,
        img_h: int,
        img_w: int,
        fish_fov: float,
        cx: float,
        cy: float,
        fx: float,
        fy: float,
        xi: float,
        alpha: float,
        R: Optional[torch.Tensor] = None,
    ):
        """Args:
        img_h: (int) the height of fisheye camera image
        img_w: (int) the width of fisheye camera image
        fish_fov: (float) the fov of fisheye camera in degrees
        cx, cy: (float) the optical center of the fisheye camera
        fx, fy, xi, alpha: (float) the fisheye camera model parameters
        R: (torch.Tensor) 3x3 rotation matrix of camera
        """
        super(FisheyeProjection, self).__init__(img_h, img_w, R)

        self.fish_fov = fish_fov  # FoV in degrees
        fov_rad = self.fish_fov / 180 * np.pi  # FoV in radians
        self.fov_cos = np.cos(fov_rad / 2)
        self.fish_param = [cx, cy, fx, fy, xi, alpha]

    def projection(
        self, world_pts: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        # Rotate world points according to camera rotation
        world_pts = self.worldcoord2camcoord(world_pts)

        # Unpack parameters
        cx, cy, fx, fy, xi, alpha = self.fish_param
        # Unpack 3D world points
        x, y, z = world_pts[..., 0], world_pts[..., 1], world_pts[..., 2]

        # Calculate fov
        world_pts_fov_cos = z  # point3D @ z_axis
        fov_mask = world_pts_fov_cos >= self.fov_cos

        # Calculate projection
        x2 = x * x
        y2 = y * y
        z2 = z * z
        d1 = torch.sqrt(x2 + y2 + z2)
        zxi = xi * d1 + z
        d2 = torch.sqrt(x2 + y2 + zxi * zxi)

        div = alpha * d2 + (1 - alpha) * zxi
        u = fx * x / div + cx
        v = fy * y / div + cy

        # Projected points on image plane
        # For grid_sample, -1 <= proj_pts <= 1
        mapx = 2 * u / self.img_w - 1.0
        mapy = 2 * v / self.img_h - 1.0
        proj_pts = torch.stack([mapx, mapy], dim=-1)

        # Check valid area
        if alpha <= 0.5:
            w1 = alpha / (1 - alpha)
        else:
            w1 = (1 - alpha) / alpha
        w2 = w1 + xi / np.sqrt(2 * w1 * xi + xi * xi + 1)
        valid_mask = z > -w2 * d1
        valid_mask *= fov_mask

        return proj_pts, valid_mask

    def unprojection(
        self, with_rotation: bool = True
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        # Unpack parameters
        cx, cy, fx, fy, xi, alpha = self.fish_param

        # Calculate unprojection
        v, u = torch.meshgrid(
            [torch.arange(self.img_h), torch.arange(self.img_w)]
        )
        mx = (u - cx) / fx
        my = (v - cy) / fy
        r2 = mx * mx + my * my
        mz = (1 - alpha * alpha * r2) / (
            alpha * torch.sqrt(1 - (2 * alpha - 1) * r2) + 1 - alpha
        )
        mz2 = mz * mz

        k1 = mz * xi + torch.sqrt(mz2 + (1 - xi * xi) * r2)
        k2 = mz2 + r2
        k = k1 / k2

        # Unprojected unit vectors
        unproj_pts = k.unsqueeze(-1) * torch.stack([mx, my, mz], dim=-1)
        unproj_pts[..., 2] -= xi

        # Calculate fov
        unproj_fov_cos = unproj_pts[..., 2]  # unproj_pts @ z_axis
        fov_mask = unproj_fov_cos >= self.fov_cos
        if alpha > 0.5:
            fov_mask *= r2 <= (1 / (2 * alpha - 1))

        # Rotate unproj_pts points according to camera rotation
        if with_rotation:
            unproj_pts = self.camcoord2worldcoord(unproj_pts)

        return unproj_pts, fov_mask


class ProjectionConverter(nn.Module):
    r"""This is the implementation to convert {cubemap, equirect, fisheye} images
    into {perspective, equirect, fisheye} images.
    """

    def __init__(
        self,
        input_projections: Union[List[CameraProjection], CameraProjection],
        output_projections: Union[List[CameraProjection], CameraProjection],
    ):
        """Args:
        input_projections: input images of projection models
        output_projections: generated image of projection models
        """
        super(ProjectionConverter, self).__init__()
        # Convert to list
        if not isinstance(input_projections, list):
            input_projections = [input_projections]
        if not isinstance(output_projections, list):
            output_projections = [output_projections]

        self.input_models = input_projections
        self.output_models = output_projections
        self.input_len = len(self.input_models)
        self.output_len = len(self.output_models)

        # Check image size
        input_size = self.input_models[0].size()
        for it in self.input_models:
            assert (
                input_size == it.size()
            ), "All input models must have the same image size"

        output_size = self.output_models[0].size()
        for it in self.output_models:
            assert (
                output_size == it.size()
            ), "All output models must have the same image size"

        # Check if depth conversion is required
        # If depth is in z value in input, conversion is required
        self.input_zfactor = self.calculate_zfactor(self.input_models)
        # If depth is in z value in output, inverse conversion is required
        self.output_zfactor = self.calculate_zfactor(
            self.output_models, inverse=True
        )

        # grids shape: (output_len, input_len, output_img_h, output_img_w, 2)
        self.grids = self.generate_grid()
        # _grids_cache shape: (batch_size*output_len*input_len, output_img_h, output_img_w, 2)
        self._grids_cache = None

    def _generate_grid_one_output(
        self, output_model: CameraProjection
    ) -> torch.Tensor:
        # Obtain points on unit sphere
        world_pts, not_assigned_mask = output_model.unprojection()
        # Generate grid
        grids = []
        for input_model in self.input_models:
            grid, input_mask = input_model.projection(world_pts)
            # Make sure each point is only assigned to single input
            input_mask *= not_assigned_mask
            # Values bigger than one will be ignored by grid_sample
            grid[~input_mask] = 2
            # Update not_assigned_mask
            not_assigned_mask *= ~input_mask
            grids.append(grid)
        grids = torch.stack(grids, dim=0)
        return grids

    def generate_grid(self) -> torch.Tensor:
        multi_output_grids = []
        for output_model in self.output_models:
            grids = self._generate_grid_one_output(output_model)
            multi_output_grids.append(grids.unsqueeze(1))
        multi_output_grids = torch.cat(multi_output_grids, dim=1)
        return multi_output_grids  # input_len, output_len, output_img_h, output_img_w, 2

    def _convert(self, batch: torch.Tensor) -> torch.Tensor:
        """Takes a batch of images stacked in proper order and converts thems,
        reduces batch size by input_len."""
        batch_size, ch, _H, _W = batch.shape
        out_h, out_w = self.output_models[0].size()
        if batch_size == 0 or batch_size % self.input_len != 0:
            raise ValueError(f"Batch size should be {self.input_len}x")
        output = torch.nn.functional.grid_sample(
            batch,
            self._grids_cache,
            align_corners=True,
            padding_mode="zeros",
        )
        output = output.view(
            batch_size // self.input_len,
            self.input_len,
            ch,
            out_h,
            out_w,
        ).sum(dim=1)
        return output  # output_len * batch_size, ch, output_model.img_h, output_model.img_w

    def to_converted_tensor(self, batch: torch.Tensor) -> torch.Tensor:
        """Convert tensors based on projection models. If there are two
        batches from two envs (R_1st, G_1st, B_1st) and (R_2nd, G_2nd, B_2nd),
        the input order is [R_1st, G_1st, B_1st, R_2nd, G_2nd, B_2nd]
        """
        # batch tensor order should be NCHW
        batch_size, ch, in_h, in_w = batch.size()

        out_h, out_w = self.output_models[0].size()

        # Check whether batch size is len(self.input_models) x
        if batch_size == 0 or batch_size % self.input_len != 0:
            raise ValueError(f"Batch size should be {self.input_len}x")

        # How many sets of input.
        num_input_set = batch_size // self.input_len

        # to(device) is a NOOP after the first call
        self.grids = self.grids.to(batch.device)

        # Adjust batch for multiple outputs
        # batch must be [1st batch * output_len, 2nd batch * output_len, ...]
        # not that [1st batch, 2nd batch, ...] * output_len
        multi_out_batch = (
            batch.view(num_input_set, self.input_len, ch, in_h, in_w)
            .repeat(1, self.output_len, 1, 1, 1)
            .view(self.output_len * batch_size, ch, in_h, in_w)
        )

        # Cache the repeated grids for subsequent batches
        if (
            self._grids_cache is None
            or self._grids_cache.size()[0] != multi_out_batch.size()[0]
        ):
            # batch size is more than one
            self._grids_cache = self.grids.repeat(
                num_input_set, 1, 1, 1, 1
            ).view(batch_size * self.output_len, out_h, out_w, 2)
        self._grids_cache = self._grids_cache.to(batch.device)

        return self._convert(multi_out_batch)

    def calculate_zfactor(
        self, projections: List[CameraProjection], inverse: bool = False
    ) -> Optional[torch.Tensor]:
        """Calculate z factor based on camera projection models. z_factor is
        used for converting depth in z value to depth from optical center
        (for input_models) or conversion of depth from optical center to depth
        in z value (inverse = True, for output_models). Whether the conversion
        is required or not is decided based on depth_from property of
        CameraProjection class.
        Args:
            projections: input or output projection models
            inverse: True to convert depth from optical center to z value
                     False to convert z value to depth from optical center
        Returns:
            z_factors: z factor. Return None if conversion is not required.
        """
        z_factors = []
        for cam in projections:
            if cam.depth_from == _DepthFrom.Z_VAL:
                pts_on_sphere, _ = cam.unprojection(with_rotation=False)
                zval_to_optcenter = 1 / pts_on_sphere[..., 2]
                z_factors.append(zval_to_optcenter.unsqueeze(0))
            else:
                all_one = torch.full(
                    (1, cam.img_h, cam.img_w), 1.0, dtype=torch.float
                )
                z_factors.append(all_one)
        z_factors = torch.stack(z_factors)

        if (z_factors == 1.0).all():
            # All input cameras have depth from optical center
            return None
        else:
            if not inverse:
                # for input_models
                return z_factors
            else:
                # for output_models
                return 1 / z_factors

    def forward(
        self, batch: torch.Tensor, is_depth: bool = False
    ) -> torch.Tensor:

        # Depth conversion for input tensors
        if is_depth and self.input_zfactor is not None:
            input_b = batch.size()[0] // self.input_len
            self.input_zfactor = self.input_zfactor.to(batch.device)
            batch = batch * self.input_zfactor.repeat(input_b, 1, 1, 1)

        # Common operator to convert projection models
        out = self.to_converted_tensor(batch)

        # Depth conversion for output tensors
        if is_depth and self.output_zfactor is not None:
            output_b = out.size()[0] // self.output_len
            self.output_zfactor = self.output_zfactor.to(batch.device)
            out = out * self.output_zfactor.repeat(output_b, 1, 1, 1)

        return out


def get_cubemap_projections(
    img_h: int = 256, img_w: int = 256
) -> List[CameraProjection]:
    """Get cubemap camera projections that consist of six PerspectiveCameras.
    The orders are 'BACK', 'DOWN', 'FRONT', 'LEFT', 'RIGHT', 'UP'.
    Args:
        img_h: (int) the height of camera image
        img_w: (int) the width of camera image

    The rotation matrices are equivalent to
    .. code-block:: python
        from scipy.spatial.transform import Rotation
        rotations = [
            Rotation.from_euler("y", 180, degrees=True),  # Back
            Rotation.from_euler("x", -90, degrees=True),  # Down
            Rotation.from_euler("x", 0, degrees=True),  # Front
            Rotation.from_euler("y", -90, degrees=True),  # Left
            Rotation.from_euler("y", 90, degrees=True),  # Right
            Rotation.from_euler("x", 90, degrees=True)  # Up
        ]
    """
    rotations = [
        torch.tensor([[-1, 0, 0], [0, 1, 0], [0, 0, -1]]),  # Back
        torch.tensor([[1, 0, 0], [0, 0, 1], [0, -1, 0]]),  # Down
        torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, 1]]),  # Front
        torch.tensor([[0, 0, -1], [0, 1, 0], [1, 0, 0]]),  # Left
        torch.tensor([[0, 0, 1], [0, 1, 0], [-1, 0, 0]]),  # Right
        torch.tensor([[1, 0, 0], [0, 0, -1], [0, 1, 0]]),  # Up
    ]

    projections = []
    for rot in rotations:
        cam = PerspectiveProjection(img_h, img_w, R=rot)
        projections.append(cam)
    return projections


class Cube2Equirect(ProjectionConverter):
    """This is the backend Cube2Equirect nn.module that does the stiching.
    Inspired from https://github.com/fuenwang/PanoramaUtility and
    optimized for modern PyTorch."""

    def __init__(self, equ_h: int, equ_w: int):
        """Args:
        equ_h: (int) the height of the generated equirect
        equ_w: (int) the width of the generated equirect
        """

        # Cubemap input
        input_projections = get_cubemap_projections()

        # Equirectangular output
        output_projection = EquirectProjection(equ_h, equ_w)
        super(Cube2Equirect, self).__init__(
            input_projections, output_projection
        )


class ProjectionTransformer(ObservationTransformer):
    r"""
    ProjectionTransformer base class. It can be used to  convert {cubemap, equirect, fisheye} images
    into {perspective, equirect, fisheye} images in ObservationTransformer.
    """

    def __init__(
        self,
        converter: ProjectionConverter,
        sensor_uuids: List[str],
        image_shape: Tuple[int, int],
        channels_last: bool = False,
        target_uuids: Optional[List[str]] = None,
        depth_key: str = "depth",
    ):
        r""":param converter: ProjectionConverter class
        :param sensor_uuids: List of sensor_uuids
        :param image_shape: The shape of the output image (height, width)
        :param channels_last: Are the channels last in the input
        :param target_uuids: Optional List of which of the sensor_uuids to overwrite
        :param depth_key: If sensor_uuids has depth_key substring, they are processed as depth
        """
        super(ProjectionTransformer, self).__init__()
        num_sensors = len(sensor_uuids)
        assert (
            num_sensors % converter.input_len == 0 and num_sensors != 0
        ), f"{len(sensor_uuids)}: length of sensors is not a multiple of {converter.input_len}"
        # TODO verify attributes of the sensors in the config if possible. Think about API design
        assert (
            len(image_shape) == 2
        ), f"image_shape must be a tuple of (height, width), given: {image_shape}"
        self.sensor_uuids: List[str] = sensor_uuids
        self.img_shape: Tuple[int, int] = image_shape
        self.channels_last: bool = channels_last
        self.converter = converter
        if target_uuids is None:
            self.target_uuids: List[str] = self.sensor_uuids[::6]
        else:
            self.target_uuids: List[str] = target_uuids
        self.depth_key = depth_key

    def transform_observation_space(
        self,
        observation_space: spaces.Dict,
    ):
        r"""Transforms the target UUID's sensor obs_space so it matches the new shape (H, W)"""
        # Transforms the observation space to of the target UUID
        for i, key in enumerate(self.target_uuids):
            assert (
                key in observation_space.spaces
            ), f"{key} not found in observation space: {observation_space.spaces}"
            h, w = get_image_height_width(
                observation_space.spaces[key], channels_last=True
            )
            in_len = self.converter.input_len
            logger.info(
                f"Overwrite sensor: {key} from size of ({h}, {w}) to image of"
                f" {self.img_shape} from sensors: {self.sensor_uuids[i*in_len:(i+1)*in_len]}"
            )
            if (h, w) != self.img_shape:
                observation_space.spaces[key] = overwrite_gym_box_shape(
                    observation_space.spaces[key], self.img_shape
                )
        return observation_space

    @torch.no_grad()
    def forward(
        self, observations: Dict[str, torch.Tensor]
    ) -> Dict[str, torch.Tensor]:

        for i, target_sensor_uuid in enumerate(self.target_uuids):
            # number of input and input sensor uuids
            in_len = self.converter.input_len
            in_sensor_uuids = self.sensor_uuids[i * in_len : (i + 1) * in_len]

            # If the sensor is depth
            is_depth = any(self.depth_key in s for s in in_sensor_uuids)

            # The UUID we are overwriting
            assert target_sensor_uuid in in_sensor_uuids
            sensor_obs = [observations[sensor] for sensor in in_sensor_uuids]
            target_obs = observations[target_sensor_uuid]
            sensor_dtype = target_obs.dtype
            # Stacking along axis makes the flattening go in the right order.
            imgs = torch.stack(sensor_obs, axis=1)
            imgs = torch.flatten(imgs, end_dim=1)
            if not self.channels_last:
                imgs = imgs.permute((0, 3, 1, 2))  # NHWC => NCHW
            imgs = imgs.float()  # NCHW
            # Here is where the projection conversion happens
            output = self.converter(imgs, is_depth=is_depth)

            # for debugging
            # torchvision.utils.save_image(output, f'sample_eqr_{target_sensor_uuid}.jpg', normalize=True, range=(0, 255) if 'rgb' in target_sensor_uuid else (0, 1))
            output = output.to(dtype=sensor_dtype)
            if not self.channels_last:
                output = output.permute((0, 2, 3, 1))  # NCHW => NHWC
            observations[target_sensor_uuid] = output
        return observations


@baseline_registry.register_obs_transformer()
class CubeMap2Equirect(ProjectionTransformer):
    r"""This is an experimental use of ObservationTransformer that converts a cubemap
    output to an equirectangular one through projection. This needs to be fed
    a list of 6 cameras at various orientations but will be able to stitch a
    360 sensor out of these inputs. The code below will generate a config that
    has the 6 sensors in the proper orientations. This code also assumes a 90
    FOV.

    Sensor order for cubemap stiching is Back, Down, Front, Left, Right, Up.
    The output will be writen the UUID of the first sensor.
    """

    def __init__(
        self,
        sensor_uuids: List[str],
        eq_shape: Tuple[int, int],
        channels_last: bool = False,
        target_uuids: Optional[List[str]] = None,
        depth_key: str = "depth",
    ):
        r""":param sensor_uuids: List of sensor_uuids: Back, Down, Front, Left, Right, Up.
        :param eq_shape: The shape of the equirectangular output (height, width)
        :param channels_last: Are the channels last in the input
        :param target_uuids: Optional List of which of the sensor_uuids to overwrite
        :param depth_key: If sensor_uuids has depth_key substring, they are processed as depth
        """

        converter = Cube2Equirect(eq_shape[0], eq_shape[1])
        super(CubeMap2Equirect, self).__init__(
            converter,
            sensor_uuids,
            eq_shape,
            channels_last,
            target_uuids,
            depth_key,
        )

    @classmethod
    def from_config(cls, config):
        cube2eq_config = config.RL.POLICY.OBS_TRANSFORMS.CUBE2EQ
        if hasattr(cube2eq_config, "TARGET_UUIDS"):
            # Optional Config Value to specify target UUID
            target_uuids = cube2eq_config.TARGET_UUIDS
        else:
            target_uuids = None
        return cls(
            cube2eq_config.SENSOR_UUIDS,
            eq_shape=(
                cube2eq_config.HEIGHT,
                cube2eq_config.WIDTH,
            ),
            target_uuids=target_uuids,
        )


class Cube2Fisheye(ProjectionConverter):
    r"""This is the implementation to generate fisheye images from cubemap images.
    The camera model is based on the Double Sphere Camera Model (Usenko et. al.;3DV 2018).
    Paper: https://arxiv.org/abs/1807.08957
    """

    def __init__(
        self,
        fish_h: int,
        fish_w: int,
        fish_fov: float,
        cx: float,
        cy: float,
        fx: float,
        fy: float,
        xi: float,
        alpha: float,
    ):
        """Args:
        fish_h: (int) the height of the generated fisheye
        fish_w: (int) the width of the generated fisheye
        fish_fov: (float) the fov of the generated fisheye in degrees
        cx, cy: (float) the optical center of the generated fisheye
        fx, fy, xi, alpha: (float) the fisheye camera model parameters
        """

        # Cubemap input
        input_projections = get_cubemap_projections()

        # Fisheye output
        output_projection = FisheyeProjection(
            fish_h, fish_w, fish_fov, cx, cy, fx, fy, xi, alpha
        )
        super(Cube2Fisheye, self).__init__(
            input_projections, output_projection
        )


@baseline_registry.register_obs_transformer()
class CubeMap2Fisheye(ProjectionTransformer):
    r"""This is an experimental use of ObservationTransformer that converts a cubemap
    output to a fisheye one through projection. This needs to be fed
    a list of 6 cameras at various orientations but will be able to stitch a
    fisheye image out of these inputs. The code below will generate a config that
    has the 6 sensors in the proper orientations. This code also assumes a 90
    FOV.

    Sensor order for cubemap stiching is Back, Down, Front, Left, Right, Up.
    The output will be writen the UUID of the first sensor.
    """

    def __init__(
        self,
        sensor_uuids: List[str],
        fish_shape: Tuple[int, int],
        fish_fov: float,
        fish_params: Tuple[float],
        channels_last: bool = False,
        target_uuids: Optional[List[str]] = None,
        depth_key: str = "depth",
    ):
        r""":param sensor_uuids: List of sensor_uuids: Back, Down, Front, Left, Right, Up.
        :param fish_shape: The shape of the fisheye output (height, width)
        :param fish_fov: The FoV of the fisheye output in degrees
        :param fish_params: The camera parameters of fisheye output (f, xi, alpha)
        :param channels_last: Are the channels last in the input
        :param target_uuids: Optional List of which of the sensor_uuids to overwrite
        :param depth_key: If sensor_uuids has depth_key substring, they are processed as depth
        """

        assert (
            len(fish_params) == 3
        ), "fish_params must have three parameters (f, xi, alpha)"
        # fisheye camera parameters
        fx = fish_params[0] * min(fish_shape)
        fy = fx
        cx = fish_shape[1] / 2
        cy = fish_shape[0] / 2
        xi = fish_params[1]
        alpha = fish_params[2]
        converter: ProjectionConverter = Cube2Fisheye(
            fish_shape[0], fish_shape[1], fish_fov, cx, cy, fx, fy, xi, alpha
        )

        super(CubeMap2Fisheye, self).__init__(
            converter,
            sensor_uuids,
            fish_shape,
            channels_last,
            target_uuids,
            depth_key,
        )

    @classmethod
    def from_config(cls, config):
        cube2fish_config = config.RL.POLICY.OBS_TRANSFORMS.CUBE2FISH
        if hasattr(cube2fish_config, "TARGET_UUIDS"):
            # Optional Config Value to specify target UUID
            target_uuids = cube2fish_config.TARGET_UUIDS
        else:
            target_uuids = None
        return cls(
            cube2fish_config.SENSOR_UUIDS,
            fish_shape=(
                cube2fish_config.HEIGHT,
                cube2fish_config.WIDTH,
            ),
            fish_fov=cube2fish_config.FOV,
            fish_params=cube2fish_config.PARAMS,
            target_uuids=target_uuids,
        )


class Equirect2Cube(ProjectionConverter):
    """This is the backend Equirect2CubeMap that converts equirectangular image
    to cubemap images."""

    def __init__(self, img_h: int, img_w: int):
        """Args:
        img_h: (int) the height of the generated cubemap
        img_w: (int) the width of the generated cubemap
        """

        # Equirectangular input
        input_projection = EquirectProjection(256, 512)

        #  Cubemap output
        output_projections = get_cubemap_projections(img_h, img_w)
        super(Equirect2Cube, self).__init__(
            input_projection, output_projections
        )


@baseline_registry.register_obs_transformer()
class Equirect2CubeMap(ProjectionTransformer):
    r"""This is an experimental use of ObservationTransformer that converts
    an equirectangular image to cubemap images.
    Cubemap order is Back, Down, Front, Left, Right, Up.
    The output will be writen the UUID of the first sensor.
    """

    def __init__(
        self,
        sensor_uuids: List[str],
        img_shape: Tuple[int, int],
        channels_last: bool = False,
        target_uuids: Optional[List[str]] = None,
        depth_key: str = "depth",
    ):
        r""":param sensor_uuids: List of sensor_uuids: Back, Down, Front, Left, Right, Up.
        :param img_shape: The shape of the equirectangular output (height, width)
        :param channels_last: Are the channels last in the input
        :param target_uuids: Optional List of which of the sensor_uuids to overwrite
        :param depth_key: If sensor_uuids has depth_key substring, they are processed as depth
        """

        converter = Equirect2Cube(img_shape[0], img_shape[1])
        super(Equirect2CubeMap, self).__init__(
            converter,
            sensor_uuids,
            img_shape,
            channels_last,
            target_uuids,
            depth_key,
        )

    @classmethod
    def from_config(cls, config):
        eq2cube_config = config.RL.POLICY.OBS_TRANSFORMS.EQ2CUBE

        if hasattr(eq2cube_config, "TARGET_UUIDS"):
            # Optional Config Value to specify target UUID
            target_uuids = eq2cube_config.TARGET_UUIDS
        else:
            target_uuids = None
        return cls(
            eq2cube_config.SENSOR_UUIDS,
            img_shape=(
                eq2cube_config.HEIGHT,
                eq2cube_config.WIDTH,
            ),
            target_uuids=target_uuids,
        )


def get_active_obs_transforms(config: Config) -> List[ObservationTransformer]:
    active_obs_transforms = []
    if hasattr(config.RL.POLICY, "OBS_TRANSFORMS"):
        obs_transform_names = (
            config.RL.POLICY.OBS_TRANSFORMS.ENABLED_TRANSFORMS
        )
        for obs_transform_name in obs_transform_names:
            obs_trans_cls = baseline_registry.get_obs_transformer(
                obs_transform_name
            )
            obs_transform = obs_trans_cls.from_config(config)
            active_obs_transforms.append(obs_transform)
    
    if hasattr(config, "IL") and hasattr(config.IL, "OBS_TRANSFORMS"):
        obs_transform_names = (
            config.IL.OBS_TRANSFORMS.ENABLED_TRANSFORMS, 
        )
        for obs_transform_name in obs_transform_names:
            obs_trans_cls = baseline_registry.get_obs_transformer(
                obs_transform_name
            )
            obs_transform = obs_trans_cls.from_config(config)
            active_obs_transforms.append(obs_transform)
    return active_obs_transforms


def apply_obs_transforms_batch(
    batch: Dict[str, torch.Tensor],
    obs_transforms: Iterable[ObservationTransformer],
) -> Dict[str, torch.Tensor]:
    for obs_transform in obs_transforms:
        batch = obs_transform(batch)
    return batch


def apply_obs_transforms_obs_space(
    obs_space: spaces.Dict, obs_transforms: Iterable[ObservationTransformer]
) -> spaces.Dict:
    for obs_transform in obs_transforms:
        obs_space = obs_transform.transform_observation_space(obs_space)
    return obs_space

habitat_baselines/common/obs_transformers.py
cshizhe/onav_rim
#!/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.

from typing import Any

import numpy as np
import torch
#from torch.utils.tensorboard import SummaryWriter
from tensorboardX import SummaryWriter


class TensorboardWriter:
    def __init__(self, log_dir: str, *args: Any, **kwargs: Any):
        r"""A Wrapper for tensorboard SummaryWriter. It creates a dummy writer
        when log_dir is empty string or None. It also has functionality that
        generates tb video directly from numpy images.

        Args:
            log_dir: Save directory location. Will not write to disk if
            log_dir is an empty string.
            *args: Additional positional args for SummaryWriter
            **kwargs: Additional keyword args for SummaryWriter
        """
        self.writer = None
        if log_dir is not None and len(log_dir) > 0:
            self.writer = SummaryWriter(log_dir, *args, **kwargs)

    def __getattr__(self, item):
        if self.writer:
            return self.writer.__getattribute__(item)
        else:
            return lambda *args, **kwargs: None

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.writer:
            self.writer.close()

    def add_video_from_np_images(
        self, video_name: str, step_idx: int, images: np.ndarray, fps: int = 10
    ) -> None:
        r"""Write video into tensorboard from images frames.

        Args:
            video_name: name of video string.
            step_idx: int of checkpoint index to be displayed.
            images: list of n frames. Each frame is a np.ndarray of shape.
            fps: frame per second for output video.

        Returns:
            None.
        """
        if not self.writer:
            return
        # initial shape of np.ndarray list: N * (H, W, 3)
        frame_tensors = [
            torch.from_numpy(np_arr).unsqueeze(0) for np_arr in images
        ]
        video_tensor = torch.cat(tuple(frame_tensors))
        video_tensor = video_tensor.permute(0, 3, 1, 2).unsqueeze(0)
        # final shape of video tensor: (1, n, 3, H, W)
        self.writer.add_video(
            video_name, video_tensor, fps=fps, global_step=step_idx
        )

habitat_baselines/common/tensorboard_utils.py
cshizhe/onav_rim
#!/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.

from typing import List, Optional, Union

import numpy as np
from copy import deepcopy

from habitat import get_config as get_task_config
from habitat.config import Config as CN

DEFAULT_CONFIG_DIR = "configs/"
CONFIG_FILE_SEPARATOR = ","
# -----------------------------------------------------------------------------
# EXPERIMENT CONFIG
# -----------------------------------------------------------------------------
_C = CN()
# task config can be a list of conifgs like "A.yaml,B.yaml"
_C.BASE_TASK_CONFIG_PATH = "configs/tasks/pointnav.yaml"
_C.TASK_CONFIG = CN()  # task_config will be stored as a config node
_C.CMD_TRAILING_OPTS = []  # store command line options as list of strings
_C.TRAINER_NAME = "ppo"
_C.ENV_NAME = "NavRLEnv"
_C.SIMULATOR_GPU_ID = 0
_C.TORCH_GPU_ID = 0
_C.VIDEO_OPTION = ["disk", "tensorboard"]
_C.TENSORBOARD_DIR = "tb"
_C.VIDEO_DIR = "video_dir"
_C.TEST_EPISODE_COUNT = -1
_C.EVAL_CKPT_PATH_DIR = "data/checkpoints"  # path to ckpt or path to ckpts dir
_C.NUM_PROCESSES = 16
_C.SENSORS = ["RGB_SENSOR", "DEPTH_SENSOR"]
_C.CHECKPOINT_FOLDER = "data/checkpoints"
_C.NUM_UPDATES = 10000
_C.LOG_INTERVAL = 10
_C.LOG_FILE = "train.log"
_C.CHECKPOINT_INTERVAL = 50
_C.FORCE_BLIND_POLICY = False
# -----------------------------------------------------------------------------
# EVAL CONFIG
# -----------------------------------------------------------------------------
_C.EVAL = CN()
# The split to evaluate on
_C.EVAL.SPLIT = "val"
_C.EVAL.USE_CKPT_CONFIG = True
# -----------------------------------------------------------------------------
# REINFORCEMENT LEARNING (RL) ENVIRONMENT CONFIG
# -----------------------------------------------------------------------------
_C.RL = CN()
_C.RL.REWARD_MEASURE = "distance_to_goal"
_C.RL.SUCCESS_MEASURE = "spl"
_C.RL.SUCCESS_REWARD = 2.5
_C.RL.SLACK_REWARD = -0.01
_C.RL.COVERAGE_REWARD = 0.25
_C.RL.COVERAGE_ATTENUATION = 0.99
_C.RL.COVERAGE_VISIT_EXP = 1
_C.RL.COVERAGE_BONUS_SCALE = 0.5 # We see coverages of around 5-10
_C.RL.COVERAGE_FALLOFF_RADIUS = 2.0
_C.RL.COVERAGE_TYPE = "VISIT" # Or "FOG_OF_WAR"
_C.RL.EXPLORE_GOAL_SEEN_THRESHOLD = 0.05 # Transition for ExploreThenNav
# -----------------------------------------------------------------------------
# POLICY CONFIG
# -----------------------------------------------------------------------------
_C.RL.POLICY = CN()
_C.RL.POLICY.name = "PointNavBaselinePolicy"
# -----------------------------------------------------------------------------
# OBS_TRANSFORMS CONFIG
# -----------------------------------------------------------------------------
_C.RL.POLICY.OBS_TRANSFORMS = CN()
_C.RL.POLICY.OBS_TRANSFORMS.ENABLED_TRANSFORMS = tuple()
_C.RL.POLICY.OBS_TRANSFORMS.CENTER_CROPPER = CN()
_C.RL.POLICY.OBS_TRANSFORMS.CENTER_CROPPER.HEIGHT = 256
_C.RL.POLICY.OBS_TRANSFORMS.CENTER_CROPPER.WIDTH = 256
_C.RL.POLICY.OBS_TRANSFORMS.RESIZE_SHORTEST_EDGE = CN()
_C.RL.POLICY.OBS_TRANSFORMS.RESIZE_SHORTEST_EDGE.SIZE = 256
_C.RL.POLICY.OBS_TRANSFORMS.CUBE2EQ = CN()
_C.RL.POLICY.OBS_TRANSFORMS.CUBE2EQ.HEIGHT = 256
_C.RL.POLICY.OBS_TRANSFORMS.CUBE2EQ.WIDTH = 512
_C.RL.POLICY.OBS_TRANSFORMS.CUBE2EQ.SENSOR_UUIDS = list()
_C.RL.POLICY.OBS_TRANSFORMS.CUBE2FISH = CN()
_C.RL.POLICY.OBS_TRANSFORMS.CUBE2FISH.HEIGHT = 256
_C.RL.POLICY.OBS_TRANSFORMS.CUBE2FISH.WIDTH = 256
_C.RL.POLICY.OBS_TRANSFORMS.CUBE2FISH.FOV = 180
_C.RL.POLICY.OBS_TRANSFORMS.CUBE2FISH.PARAMS = (0.2, 0.2, 0.2)
_C.RL.POLICY.OBS_TRANSFORMS.CUBE2FISH.SENSOR_UUIDS = list()
_C.RL.POLICY.OBS_TRANSFORMS.EQ2CUBE = CN()
_C.RL.POLICY.OBS_TRANSFORMS.EQ2CUBE.HEIGHT = 256
_C.RL.POLICY.OBS_TRANSFORMS.EQ2CUBE.WIDTH = 256
_C.RL.POLICY.OBS_TRANSFORMS.EQ2CUBE.SENSOR_UUIDS = list()
# -----------------------------------------------------------------------------
# PROXIMAL POLICY OPTIMIZATION (PPO)
# -----------------------------------------------------------------------------
_C.RL.PPO = CN()
_C.RL.PPO.clip_param = 0.2
_C.RL.PPO.ppo_epoch = 4
_C.RL.PPO.num_mini_batch = 16
_C.RL.PPO.value_loss_coef = 0.5
_C.RL.PPO.entropy_coef = 0.01
_C.RL.PPO.lr = 7e-4
_C.RL.PPO.eps = 1e-5
_C.RL.PPO.max_grad_norm = 0.5
_C.RL.PPO.num_steps = 5
_C.RL.PPO.use_gae = True
_C.RL.PPO.use_linear_lr_decay = False
_C.RL.PPO.use_linear_clip_decay = False
_C.RL.PPO.gamma = 0.99
_C.RL.PPO.tau = 0.95
_C.RL.PPO.reward_window_size = 50
_C.RL.PPO.use_normalized_advantage = True
_C.RL.PPO.hidden_size = 512
# -----------------------------------------------------------------------------
# DECENTRALIZED DISTRIBUTED PROXIMAL POLICY OPTIMIZATION (DD-PPO)
# -----------------------------------------------------------------------------
_C.RL.DDPPO = CN()
_C.RL.DDPPO.sync_frac = 0.6
_C.RL.DDPPO.distrib_backend = "GLOO"
_C.RL.DDPPO.rnn_type = "LSTM"
_C.RL.DDPPO.num_recurrent_layers = 2
_C.RL.DDPPO.backbone = "resnet50"
_C.RL.DDPPO.pretrained_weights = "data/ddppo-models/gibson-2plus-resnet50.pth"
# Loads pretrained weights
_C.RL.DDPPO.pretrained = False
# Loads just the visual encoder backbone weights
_C.RL.DDPPO.pretrained_encoder = False
# Whether or not the visual encoder backbone will be trained
_C.RL.DDPPO.train_encoder = True
# Whether or not to reset the critic linear layer
_C.RL.DDPPO.reset_critic = True
# -----------------------------------------------------------------------------
# ORBSLAM2 BASELINE
# -----------------------------------------------------------------------------
_C.ORBSLAM2 = CN()
_C.ORBSLAM2.SLAM_VOCAB_PATH = "habitat_baselines/slambased/data/ORBvoc.txt"
_C.ORBSLAM2.SLAM_SETTINGS_PATH = (
    "habitat_baselines/slambased/data/mp3d3_small1k.yaml"
)
_C.ORBSLAM2.MAP_CELL_SIZE = 0.1
_C.ORBSLAM2.MAP_SIZE = 40
_C.ORBSLAM2.CAMERA_HEIGHT = get_task_config().SIMULATOR.DEPTH_SENSOR.POSITION[
    1
]
_C.ORBSLAM2.BETA = 100
_C.ORBSLAM2.H_OBSTACLE_MIN = 0.3 * _C.ORBSLAM2.CAMERA_HEIGHT
_C.ORBSLAM2.H_OBSTACLE_MAX = 1.0 * _C.ORBSLAM2.CAMERA_HEIGHT
_C.ORBSLAM2.D_OBSTACLE_MIN = 0.1
_C.ORBSLAM2.D_OBSTACLE_MAX = 4.0
_C.ORBSLAM2.PREPROCESS_MAP = True
_C.ORBSLAM2.MIN_PTS_IN_OBSTACLE = (
    get_task_config().SIMULATOR.DEPTH_SENSOR.WIDTH / 2.0
)
_C.ORBSLAM2.ANGLE_TH = float(np.deg2rad(15))
_C.ORBSLAM2.DIST_REACHED_TH = 0.15
_C.ORBSLAM2.NEXT_WAYPOINT_TH = 0.5
_C.ORBSLAM2.NUM_ACTIONS = 3
_C.ORBSLAM2.DIST_TO_STOP = 0.05
_C.ORBSLAM2.PLANNER_MAX_STEPS = 500
_C.ORBSLAM2.DEPTH_DENORM = get_task_config().SIMULATOR.DEPTH_SENSOR.MAX_DEPTH
# -----------------------------------------------------------------------------
# PROFILING
# -----------------------------------------------------------------------------
_C.PROFILING = CN()
_C.PROFILING.CAPTURE_START_STEP = -1
_C.PROFILING.NUM_STEPS_TO_CAPTURE = -1
# -----------------------------------------------------------------------------
# MODEL
# -----------------------------------------------------------------------------
_C.MODEL = CN()
_C.MODEL.num_ft_views = None
# -----------------------------------------------------------------------------
# MAPPING
# -----------------------------------------------------------------------------
_C.MAPPING = CN()
_C.MAPPING.map_size_cm = 2400       # cm
_C.MAPPING.map_resolution = 5       # cm
_C.MAPPING.global_downscaling = 2   # full map -> local map
_C.MAPPING.vision_range = 100       # grids (100x100) in local map updated per step
_C.MAPPING.du_scale = 1             # subsampling depth points in the image
_C.MAPPING.map_pred_threshold = 1.0
_C.MAPPING.exp_pred_threshold = 1.0
_C.MAPPING.collision_threshold = 0.2


def add_pano_sensors_to_config(config):
    """Dynamically adds RGB and Depth cameras to config.TASK_CONFIG, forming
    an N-frame panorama. The PanoRGB and PanoDepth observation transformers
    can be used to stack these frames together.
    """
    # https://github.com/facebookresearch/habitat-lab/issues/531
    # “orientation: numpy.ndarray[numpy.float32[3, 1]]”. 
    # The value of the orientation as "[x,y,z]" mean that the pitch, yaw, and roll of the sensor are x, y, z, respectively 
    # And it means to look more toward the ceiling if x is positive
    # In addition, the value of x needs to be set as radian (i.e., 30/180*pi) rather than degree (30 degrees).

    num_ft_views = config.MODEL.num_ft_views

    num_cameras = 12 # horizontal rotation: 30 degrees
    if num_ft_views == 1:
        return config
    elif num_ft_views == 2:
        orient = [(0, 0, 0), (- np.pi / 6, 0, 0)]
    elif num_ft_views == '2_up':
        orient = [(0, 0, 0), (np.pi / 6, 0, 0)]
    elif num_ft_views == 3:
        orient = [(0, 0, 0), (np.pi / 6, 0, 0), (- np.pi / 6, 0, 0)]
    elif num_ft_views == 12:
        orient = [(0, np.pi * 2 / num_cameras * i, 0) for i in range(num_cameras)]
    elif num_ft_views == 14:
        orient = [(0, np.pi * 2 / num_cameras * i, 0) for i in range(num_cameras)] + \
                 [(np.pi / 6, 0, 0), (- np.pi / 6, 0, 0)]
    else:
        raise NotImplementedError(f'unspport #cameras: {num_ft_views}')
    
    config.defrost()

    # RGB_SENSOR
    config.TASK_CONFIG.SIMULATOR.RGB_SENSOR.ORIENTATION = orient[0]
    for camera_id in range(1, len(orient)):
        camera_template = f"RGB_{camera_id}"
        camera_config = deepcopy(config.TASK_CONFIG.SIMULATOR.RGB_SENSOR)
        camera_config.ORIENTATION = orient[camera_id]

        camera_config.UUID = camera_template.lower()
        setattr(
            config.TASK_CONFIG.SIMULATOR, camera_template, camera_config
        )
        config.TASK_CONFIG.SIMULATOR.AGENT_0.SENSORS.append(
            camera_template
        )

    # DEPTH_SENSOR
    config.TASK_CONFIG.SIMULATOR.DEPTH_SENSOR.ORIENTATION = orient[0]
    for camera_id in range(1, len(orient)):
        camera_template = f"DEPTH_{camera_id}"
        camera_config = deepcopy(config.TASK_CONFIG.SIMULATOR.DEPTH_SENSOR)
        camera_config.ORIENTATION = orient[camera_id]
        camera_config.UUID = camera_template.lower()

        setattr(
            config.TASK_CONFIG.SIMULATOR, camera_template, camera_config
        )
        config.TASK_CONFIG.SIMULATOR.AGENT_0.SENSORS.append(
            camera_template
        )

    config.SENSORS = config.TASK_CONFIG.SIMULATOR.AGENT_0.SENSORS
    config.freeze()
    return config


def get_config(
    config_paths: Optional[Union[List[str], str]] = None,
    opts: Optional[list] = None,
) -> CN:
    r"""Create a unified config with default values overwritten by values from
    :ref:`config_paths` and overwritten by options from :ref:`opts`.

    Args:
        config_paths: List of config paths or string that contains comma
        separated list of config paths.
        opts: Config options (keys, values) in a list (e.g., passed from
        command line into the config. For example, ``opts = ['FOO.BAR',
        0.5]``. Argument can be used for parameter sweeping or quick tests.
    """
    config = _C.clone()
    if config_paths:
        if isinstance(config_paths, str):
            if CONFIG_FILE_SEPARATOR in config_paths:
                config_paths = config_paths.split(CONFIG_FILE_SEPARATOR)
            else:
                config_paths = [config_paths]

        for config_path in config_paths:
            config.merge_from_file(config_path)

    if opts:
        for k, v in zip(opts[0::2], opts[1::2]):
            if k == "BASE_TASK_CONFIG_PATH":
                config.BASE_TASK_CONFIG_PATH = v

    config.TASK_CONFIG = get_task_config(config.BASE_TASK_CONFIG_PATH)
    if opts:
        config.CMD_TRAILING_OPTS = config.CMD_TRAILING_OPTS + opts
        config.merge_from_list(config.CMD_TRAILING_OPTS)

    if config.MODEL.num_ft_views is None:
        config.MODEL.num_ft_views = 1
    config = add_pano_sensors_to_config(config)

    config.freeze()
    return config

habitat_baselines/config/default.py
cshizhe/onav_rim
import gzip
import json

import torch
import torch.nn as nn
from habitat import Config, logger


class InstructionEncoder(nn.Module):
    def __init__(self, config: Config):
        r"""An encoder that uses RNN to encode an instruction. Returns
        the final hidden state after processing the instruction sequence.
        Args:
            config: must have
                vocab_size: number of words in the vocabulary
                embedding_size: The dimension of each embedding vector
                use_pretrained_embeddings:
                embedding_file:
                fine_tune_embeddings:
                dataset_vocab:
                hidden_size: The hidden (output) size
                rnn_type: The RNN cell type.  Must be GRU or LSTM
                final_state_only: Whether or not to return just the final state
        """
        super().__init__()

        self.config = config

        if self.config.use_pretrained_embeddings:
            self.embedding_layer = nn.Embedding.from_pretrained(
                embeddings=self._load_embeddings(),
                freeze=not self.config.fine_tune_embeddings,
            )
        else:  # each embedding initialized to sampled Gaussian
            self.embedding_layer = nn.Embedding(
                num_embeddings=config.vocab_size,
                embedding_dim=config.embedding_size,
                padding_idx=0,
            )

        rnn = nn.GRU if self.config.rnn_type == "GRU" else nn.LSTM
        self.bidir = config.bidirectional
        self.encoder_rnn = rnn(
            input_size=config.embedding_size,
            hidden_size=config.hidden_size,
            bidirectional=self.bidir,
        )
        self.final_state_only = config.final_state_only

    @property
    def output_size(self):
        return self.config.hidden_size * (2 if self.bidir else 1)

    def _load_embeddings(self):
        """ Loads word embeddings from a pretrained embeddings file.
        PAD: index 0. [0.0, ... 0.0]
        UNK: index 1. mean of all R2R word embeddings: [mean_0, ..., mean_n]
        why UNK is averaged:
            https://groups.google.com/forum/#!searchin/globalvectors/unk|sort:date/globalvectors/9w8ZADXJclA/hRdn4prm-XUJ
        Returns:
            embeddings tensor of size [num_words x embedding_dim]
        """
        with gzip.open(self.config.embedding_file, "rt") as f:
            embeddings = torch.tensor(json.load(f))
        return embeddings

    def forward(self, observations):
        """
        Tensor sizes after computation:
            instruction: [batch_size x seq_length]
            lengths: [batch_size]
            hidden_state: [batch_size x hidden_size]
        """
        instruction = observations["instruction"].long()

        if len(instruction.size()) == 3:
            instruction = instruction.contiguous().view(
                -1, instruction.size(2)
            )

        lengths = (instruction != 0.0).long().sum(dim=1).cpu()
        embedded = self.embedding_layer(instruction)

        packed_seq = nn.utils.rnn.pack_padded_sequence(
            embedded, lengths, batch_first=True, enforce_sorted=False
        )

        output, final_state = self.encoder_rnn(packed_seq)

        if self.config.rnn_type == "LSTM":
            final_state = final_state[0]

        if self.final_state_only:
            return final_state.squeeze(0)
        else:
            return nn.utils.rnn.pad_packed_sequence(output, batch_first=True)[
                0
            ].permute(0, 2, 1)

habitat_baselines/il/common/encoders/instruction.py
cshizhe/onav_rim
from typing import Dict, Tuple

import numpy as np
from PIL import Image

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as models
from torchvision import transforms as T
from torchvision.transforms import functional as TF

from gym import spaces
from habitat import logger
from habitat_baselines.utils.common import Flatten
from habitat_baselines.rl.ddppo.policy import resnet
from habitat_baselines.rl.ddppo.policy.resnet_policy import ResNetEncoder

import clip

class VlnResnetDepthEncoder(nn.Module):
    def __init__(
        self,
        observation_space,
        output_size=128,
        checkpoint="NONE",
        backbone="resnet50",
        resnet_baseplanes=32,
        normalize_visual_inputs=False,
        trainable=False,
        spatial_output: bool = False,
    ):
        super().__init__()
        self.visual_encoder = ResNetEncoder(
            spaces.Dict({"depth": observation_space.spaces["depth"]}),
            baseplanes=resnet_baseplanes,
            ngroups=resnet_baseplanes // 2,
            make_backbone=getattr(resnet, backbone),
            normalize_visual_inputs=normalize_visual_inputs,
        )

        for param in self.visual_encoder.parameters():
            param.requires_grad_(trainable)

        if checkpoint != "NONE":
            ddppo_weights = torch.load(checkpoint, map_location=torch.device('cpu'))

            weights_dict = {}
            for k, v in ddppo_weights["state_dict"].items():
                split_layer_name = k.split(".")[2:]
                if split_layer_name[0] != "visual_encoder":
                    continue

                layer_name = ".".join(split_layer_name[1:])
                weights_dict[layer_name] = v

            del ddppo_weights
            self.visual_encoder.load_state_dict(weights_dict, strict=True)

        self.spatial_output = spatial_output

        if not self.spatial_output:
            self.output_shape = (output_size,)
            self.visual_fc = nn.Sequential(
                # Flatten(),
                nn.Linear(np.prod(self.visual_encoder.output_shape), output_size),
                nn.ReLU(True),
            )
        else:
            self.spatial_embeddings = nn.Embedding(
                self.visual_encoder.output_shape[1]
                * self.visual_encoder.output_shape[2],
                64,
            )

            self.output_shape = list(self.visual_encoder.output_shape)
            self.output_shape[0] += self.spatial_embeddings.embedding_dim
            self.output_shape = tuple(self.output_shape)

    def forward(self, observations, no_fc_layer=False):
        """
        Args:
            observations: [BATCH, HEIGHT, WIDTH, CHANNEL]
        Returns:
            [BATCH, OUTPUT_SIZE]
        """
        if "depth_features" in observations:
            x = observations["depth_features"]
        else:
            obs_depth = observations["depth"]
            if len(obs_depth.size()) == 5:
                observations["depth"] = obs_depth.contiguous().view(
                    -1, obs_depth.size(2), obs_depth.size(3), obs_depth.size(4)
                )
            x = self.visual_encoder(observations)
            x = x.flatten(start_dim=1)

        x = x.detach()
        if no_fc_layer:
            return x

        if self.spatial_output:
            b, c, h, w = x.size()

            spatial_features = (
                self.spatial_embeddings(
                    torch.arange(
                        0,
                        self.spatial_embeddings.num_embeddings,
                        device=x.device,
                        dtype=torch.long,
                    )
                )
                .view(1, -1, h, w)
                .expand(b, self.spatial_embeddings.embedding_dim, h, w)
            )

            return torch.cat([x, spatial_features], dim=1)
        else:
            return self.visual_fc(x)

class ResnetRGBEncoder(nn.Module):
    def __init__(
        self,
        observation_space,
        output_size=256,
        backbone="resnet50",
        resnet_baseplanes=32,
        normalize_visual_inputs=False,
        trainable=False,
        spatial_output: bool = False,
    ):
        super().__init__()

        backbone_split = backbone.split("_")
        logger.info("backbone: {}".format(backbone_split))
        make_backbone = getattr(resnet, backbone_split[0])

        self.visual_encoder = ResNetEncoder(
            spaces.Dict({"rgb": observation_space.spaces["rgb"]}),
            baseplanes=resnet_baseplanes,
            ngroups=resnet_baseplanes // 2,
            make_backbone=make_backbone,
            normalize_visual_inputs=normalize_visual_inputs,
        )

        for param in self.visual_encoder.parameters():
            param.requires_grad_(trainable)

        self.spatial_output = spatial_output

        if not self.spatial_output:
            self.output_shape = (output_size,)
            self.visual_fc = nn.Sequential(
                Flatten(),
                nn.Linear(np.prod(self.visual_encoder.output_shape), output_size),
                nn.ReLU(True),
            )
        else:
            self.spatial_embeddings = nn.Embedding(
                self.visual_encoder.output_shape[1]
                * self.visual_encoder.output_shape[2],
                64,
            )

            self.output_shape = list(self.visual_encoder.output_shape)
            self.output_shape[0] += self.spatial_embeddings.embedding_dim
            self.output_shape = tuple(self.output_shape)
    
    @property
    def is_blind(self):
        return self._n_input_rgb == 0

    def forward(self, observations):
        """
        Args:
            observations: [BATCH, HEIGHT, WIDTH, CHANNEL]
        Returns:
            [BATCH, OUTPUT_SIZE]
        """
        obs_rgb = observations["rgb"]
        if len(obs_rgb.size()) == 5:
            observations["rgb"] = obs_rgb.contiguous().view(
                -1, obs_rgb.size(2), obs_rgb.size(3), obs_rgb.size(4)
            )

        if "rgb_features" in observations:
            x = observations["rgb_features"]
        else:
            x = self.visual_encoder(observations)

        if self.spatial_output:
            b, c, h, w = x.size()

            spatial_features = (
                self.spatial_embeddings(
                    torch.arange(
                        0,
                        self.spatial_embeddings.num_embeddings,
                        device=x.device,
                        dtype=torch.long,
                    )
                )
                .view(1, -1, h, w)
                .expand(b, self.spatial_embeddings.embedding_dim, h, w)
            )

            return torch.cat([x, spatial_features], dim=1)
        else:
            return self.visual_fc(x)

class ResNetCLIPEncoder(nn.Module):
    def __init__(
        self,
        observation_space: spaces.Dict,
        output_size=256,
        pooling='attnpool',
        device: torch.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    ):
        super().__init__()

        model, preprocess = clip.load("RN50", device=device)

        # expected input: C x H x W (np.uint8 in [0-255])
        self.preprocess = T.Compose([
            T.Resize(224),
            T.CenterCrop(size=(224, 224)),
            T.ToTensor(),
            # preprocess.transforms[0],  
            # preprocess.transforms[1],
            # # already tensor, but want float
            # T.ConvertImageDtype(torch.float), # if the input is uint8, will convert [0, 255] -> [0, 1]
            # T.Normalize(mean=(0., 0., 0.), std=(255., 255., 255.)),   # [0, 255] -> convert to [0, 1]
            # # normalize with CLIP mean, std
            preprocess.transforms[4],
        ])
        # expected output: C x H x W (np.float32)

        self.backbone = model.visual

        if pooling == 'none':
            self.backbone.attnpool = nn.Identity()
            self.output_shape = (2048, 7, 7)
        elif pooling == 'avgpool':
            self.backbone.attnpool = nn.Sequential(
                nn.AdaptiveAvgPool2d(output_size=(1,1)),
                nn.Flatten()
            )
            self.output_shape = (2048,)
        else:
            self.output_shape = (1024,)

        for param in self.backbone.parameters():
            param.requires_grad = False
        for module in self.backbone.modules():
            if "BatchNorm" in type(module).__name__:
                module.momentum = 0.0
        self.backbone.eval()

        self.visual_fc = nn.Sequential(
            nn.Linear(np.prod(self.output_shape), output_size),
            nn.ReLU(True),
        )
        self.output_shape = (output_size, )

    def forward(self, observations: Dict[str, torch.Tensor], no_fc_layer=False) -> torch.Tensor:  # type: ignore
        for module in self.backbone.modules():
            if "BatchNorm" in type(module).__name__:
                module.eval()

        if "rgb_features" in observations:
            x = observations["rgb_features"]
        else:
            device = observations["rgb"].device
            rgb_observations = observations["rgb"]  # (batch, height, width, channel)

            rgb_observations = rgb_observations.data.cpu().numpy().astype(np.uint8)
            rgb_observations = [self.preprocess(Image.fromarray(rgb_image)) for rgb_image in rgb_observations]
            rgb_observations = torch.stack(rgb_observations, 0).to(device)

            # rgb_observations = rgb_observations.permute(0, 3, 1, 2) # BATCH x CHANNEL x HEIGHT X WIDTH
            # rgb_observations = torch.stack(
            #     [self.preprocess(rgb_image) for rgb_image in rgb_observations]
            # )  # [BATCH x CHANNEL x HEIGHT X WIDTH] in torch.float32
            x = self.backbone(rgb_observations).float()

        x = x.detach()

        if no_fc_layer:
            return x

        x = self.visual_fc(x)
        return x

class ResNetImageNetEncoder(nn.Module):
    def __init__(self, observation_space: spaces.Dict, output_size=256, backbone='resnet50'):
        super().__init__()

        self.preprocess = T.Compose(
            [
                T.Resize(224),
                T.CenterCrop(size=(224, 224)),
                # T.ConvertImageDtype(torch.float),
                T.ToTensor(),
                T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            ]
        )
        self.backbone = getattr(models, backbone)(pretrained=True)
        self.backbone.fc = nn.Identity()

        for param in self.backbone.parameters():
            param.requires_grad = False
        for module in self.backbone.modules():
            if "BatchNorm" in type(module).__name__:
                module.momentum = 0.0
        self.backbone.eval()

        self.output_shape = (2048,) if backbone != 'resnet18' else 512

        self.visual_fc = nn.Sequential(
            nn.Linear(np.prod(self.output_shape), output_size),
            nn.ReLU(True),
        )
        self.output_shape = (output_size, )

    def forward(self, observations: Dict[str, torch.Tensor]) -> torch.Tensor:  # type: ignore
        for module in self.backbone.modules():
            if "BatchNorm" in type(module).__name__:
                module.eval()

        if "rgb_features" in observations:
            x = observations["rgb_features"]
        else:
            device = observations["rgb"].device
            rgb_observations = observations["rgb"]  # (batch, height, width, channel)

            rgb_observations = rgb_observations.data.cpu().numpy().astype(np.uint8)
            rgb_observations = [self.preprocess(Image.fromarray(rgb_image)) for rgb_image in rgb_observations]
            rgb_observations = torch.stack(rgb_observations, 0).to(device)

            # rgb_observations = rgb_observations.permute(0, 3, 1, 2) # BATCH x CHANNEL x HEIGHT X WIDTH
            # rgb_observations = torch.stack(
            #     [self.preprocess(rgb_image) for rgb_image in rgb_observations]
            # )  # [BATCH x CHANNEL x HEIGHT X WIDTH] in torch.float32

            x = self.backbone(rgb_observations).float()            

        x = self.visual_fc(x)
        return x


class ResnetSemSeqEncoder(nn.Module):
    def __init__(
        self,
        observation_space,
        output_size=256,
        backbone="resnet18",
        resnet_baseplanes=32,
        normalize_visual_inputs=False,
        trainable=False,
        spatial_output: bool = False,
        semantic_embedding_size=4,
        use_pred_semantics=False,
        use_goal_seg=False,
        is_thda=False,
    ):
        super().__init__()
        if not use_goal_seg:
            sem_input_size = 40 + 2
            self.semantic_embedder = nn.Embedding(sem_input_size, semantic_embedding_size)

        self.visual_encoder = ResNetEncoder(
            spaces.Dict({"semantic": observation_space.spaces["semantic"]}),
            baseplanes=resnet_baseplanes,
            ngroups=resnet_baseplanes // 2,
            make_backbone=getattr(resnet, backbone),
            normalize_visual_inputs=normalize_visual_inputs,
            sem_embedding_size=semantic_embedding_size,
        )

        for param in self.visual_encoder.parameters():
            param.requires_grad_(trainable)

        self.spatial_output = spatial_output
        self.use_goal_seg = use_goal_seg

        if not self.spatial_output:
            self.output_shape = (output_size,)
            self.visual_fc = nn.Sequential(
                Flatten(),
                nn.Linear(np.prod(self.visual_encoder.output_shape), output_size),
                nn.ReLU(True),
            )
        else:
            self.spatial_embeddings = nn.Embedding(
                self.visual_encoder.output_shape[1]
                * self.visual_encoder.output_shape[2],
                64,
            )

            self.output_shape = list(self.visual_encoder.output_shape)
            self.output_shape[0] += self.spatial_embeddings.embedding_dim
            self.output_shape = tuple(self.output_shape)
    
    @property
    def is_blind(self):
        return self._n_input_rgb == 0

    def forward(self, observations):
        """
        Args:
            observations: [BATCH, HEIGHT, WIDTH, CHANNEL]
        Returns:
            [BATCH, OUTPUT_SIZE]
        """
        obs_semantic = observations["semantic"]
        if len(obs_semantic.size()) == 5:
            observations["semantic"] = obs_semantic.contiguous().view(
                -1, obs_semantic.size(2), obs_semantic.size(3), obs_semantic.size(4)
            )

        if "semantic_features" in observations:
            x = observations["semantic_features"]
        else:
            # Embed input when using all object categories
            if not self.use_goal_seg:
                categories = observations["semantic"].long() + 1
                observations["semantic"] = self.semantic_embedder(categories)
            x = self.visual_encoder(observations)

        if self.spatial_output:
            b, c, h, w = x.size()

            spatial_features = (
                self.spatial_embeddings(
                    torch.arange(
                        0,
                        self.spatial_embeddings.num_embeddings,
                        device=x.device,
                        dtype=torch.long,
                    )
                )
                .view(1, -1, h, w)
                .expand(b, self.spatial_embeddings.embedding_dim, h, w)
            )

            return torch.cat([x, spatial_features], dim=1)
        else:
            return self.visual_fc(x)


class ResnetEncoder(nn.Module):
    def __init__(
        self,
        observation_space,
        output_size=256,
        checkpoint="NONE",
        backbone="resnet50",
        resnet_baseplanes=32,
        normalize_visual_inputs=False,
        trainable=True,
        spatial_output: bool = False,
        sem_embedding_size=4,
    ):
        super().__init__()
        self.visual_encoder = ResNetEncoder(
            observation_space,
            baseplanes=resnet_baseplanes,
            ngroups=resnet_baseplanes // 2,
            make_backbone=getattr(resnet, backbone),
            normalize_visual_inputs=normalize_visual_inputs,
            sem_embedding_size=sem_embedding_size,
        )

        for param in self.visual_encoder.parameters():
            param.requires_grad_(trainable)

        self.spatial_output = spatial_output

        if not self.spatial_output:
            self.output_shape = (output_size,)
            self.visual_fc = nn.Sequential(
                Flatten(),
                nn.Linear(np.prod(self.visual_encoder.output_shape), output_size),
                nn.ReLU(True),
            )
        else:
            self.spatial_embeddings = nn.Embedding(
                self.visual_encoder.output_shape[1]
                * self.visual_encoder.output_shape[2],
                64,
            )

            self.output_shape = list(self.visual_encoder.output_shape)
            self.output_shape[0] += self.spatial_embeddings.embedding_dim
            self.output_shape = tuple(self.output_shape)

    def forward(self, observations):
        """
        Args:
            observations: [BATCH, HEIGHT, WIDTH, CHANNEL]
        Returns:
            [BATCH, OUTPUT_SIZE]
        """
        obs_rgb = observations["rgb"]
        if len(obs_rgb.size()) == 5:
            observations["rgb"] = obs_rgb.contiguous().view(
                -1, obs_rgb.size(2), obs_rgb.size(3), obs_rgb.size(4)
            )
        obs_depth = observations["depth"]
        if len(obs_rgb.size()) == 5:
            observations["depth"] = obs_depth.contiguous().view(
                -1, obs_depth.size(2), obs_depth.size(3), obs_depth.size(4)
            )
        obs_semantic = observations["semantic"]
        if len(obs_rgb.size()) == 5:
            observations["semantic"] = obs_semantic.contiguous().view(
                -1, obs_semantic.size(2), obs_semantic.size(3), obs_semantic.size(4)
            )

        if "rgb_features" in observations:
            x = observations["rgb_features"]
        else:
            x = self.visual_encoder(observations)

        if self.spatial_output:
            b, c, h, w = x.size()

            spatial_features = (
                self.spatial_embeddings(
                    torch.arange(
                        0,
                        self.spatial_embeddings.num_embeddings,
                        device=x.device,
                        dtype=torch.long,
                    )
                )
                .view(1, -1, h, w)
                .expand(b, self.spatial_embeddings.embedding_dim, h, w)
            )

            return torch.cat([x, spatial_features], dim=1)
        else:
            return self.visual_fc(x)

habitat_baselines/il/common/encoders/resnet_encoders.py
cshizhe/onav_rim
import sys
import numpy as np
import torch
import cv2

sys.path.append('/home/shichen/codes/objgoalnav/ViT-Adapter/segmentation//')

import mmcv_custom   # noqa: F401,F403
import mmseg_custom   # noqa: F401,F403
from mmseg.apis import inference_segmentor, init_segmentor
from mmseg.core.evaluation import get_palette
from mmcv.runner import load_checkpoint
from mmseg.core import get_classes

MP3D_CATEGORIES = [
    'chair', 'table', 'picture', 'cabinet', 'cushion', 'sofa', 'bed', 'chest_of_drawers',
    'plant', 'sink', 'toilet', 'stool', 'towel', 'tv_monitor', 'shower', 'bathtub',
    'counter', 'fireplace', 'gym_equipment', 'seating', 'clothes', 
    # 'foodstuff', 'stationery', 'fruit', 'plaything', 
    # 'hand_tool', 'game_equipment', 'kitchenware',
]

mp3d_to_ade20k = {
    'chair': ['chair', 'armchair', 'swivel chair'], 
    'table': ['table', 'desk', 'pool table', 'coffee table'],
    'picture': ['painting'], 
    'cabinet': ['cabinet'], 
    'cushion': ['cushion', 'pillow'], 
    'sofa': ['sofa'], 
    'bed': ['bed '], 
    'chest_of_drawers': ['wardrobe', 'chest of drawers'],
    'plant': ['plant', 'flower'], 
    'sink': ['sink'], 
    'toilet': ['toilet'], 
    'stool': ['stool'], 
    'towel': ['towel'], 
    'tv_monitor': ['screen', 'crt screen', 'monitor'], 
    'shower': ['shower'], 
    'bathtub': ['bathtub'],
    'counter': ['counter', 'countertop'], 
    'fireplace': ['fireplace'], 
    'gym_equipment': ['minibike', 'bicycle', ], 
    'seating': ['seat', 'bench'], 
    'clothes': ['apparel'],
}
ade20k_to_mp3d = {}
for k, vs in mp3d_to_ade20k.items():
    for v in vs:
        ade20k_to_mp3d[v] = k

class SegmentationModel(object):
    def __init__(self, device) -> None:
        self.load_pretrained_semantic_models(device)

        self.category_mapping = np.zeros((len(self.classes), ), dtype=np.int64) - 1
        for i, label in enumerate(self.classes):
            if label in ade20k_to_mp3d:
                self.category_mapping[i] = MP3D_CATEGORIES.index(
                    ade20k_to_mp3d[label]
                )
        
    def load_pretrained_semantic_models(self, device='cuda:0'):
        config_file = '/home/shichen/codes/objgoalnav/ViT-Adapter/segmentation/configs/ade20k/mask2former_beit_adapter_large_640_160k_ade20k_ss.py'
        checkpoint_file = '/home/shichen/codes/objgoalnav/ViT-Adapter/segmentation/released/mask2former_beit_adapter_large_640_160k_ade20k.pth.tar'

        palette = 'ade20k'

        model = init_segmentor(config_file, checkpoint=None, device=device)
        checkpoint = load_checkpoint(model, checkpoint_file, map_location='cpu')
        if 'CLASSES' in checkpoint.get('meta', {}):
            model.CLASSES = checkpoint['meta']['CLASSES']
        else:
            model.CLASSES = get_classes(palette)

        self.model = model
        self.classes = model.CLASSES

    def predict_target(self, img, depth_img, target_obj, outfile=None):
        # should convert image from RGB to BGR
        img = torch.flip(img, (2, ))
        result = inference_segmentor(self.model, img.data.cpu().numpy())
    
        seg_masks = self.category_mapping[result[0]]
        seg_tgt_masks = seg_masks == target_obj.item()

        area, med_dist = 0, np.inf
        
        # print(np.unique(seg_masks))
        if np.sum(seg_tgt_masks) > 0:
            area = np.sum(seg_tgt_masks) / np.prod(seg_masks.shape)
            dists = depth_img[seg_tgt_masks]
            # print('dists', dists)
            med_dist = torch.median(dists).item()
            med_dist = med_dist * 4.5 + 0.5

        if outfile is not None:
            cv2.imwrite(outfile, result[0])

        return area, med_dist
habitat_baselines/il/common/encoders/semantic_predictor.py
cshizhe/onav_rim
import msgpack_numpy
import os
import torch
from collections import defaultdict
from typing import List

import lmdb
import magnum as mn
import numpy as np
from torch.utils.data import Dataset
from tqdm import tqdm

import habitat
from habitat import logger
from habitat.datasets.utils import VocabDict
from habitat.tasks.pickplace.pickplace import RearrangementEpisode


class ObservationsDict(dict):
    def pin_memory(self):
        for k, v in self.items():
            self[k] = v.pin_memory()

        return self


def collate_fn(batch):
    """Each sample in batch: (
            obs,
            prev_actions,
            oracle_actions,
            inflec_weight,
        )
    """

    def _pad_helper(t, max_len, fill_val=0):
        pad_amount = max_len - t.size(0)
        if pad_amount == 0:
            return t

        pad = torch.full_like(t[0:1], fill_val).expand(pad_amount, *t.size()[1:])
        return torch.cat([t, pad], dim=0)

    transposed = list(zip(*batch))
    
    observations_batch = list(transposed[1])
    next_actions_batch = list(transposed[2])
    prev_actions_batch = list(transposed[3])
    weights_batch = list(transposed[4])
    B = len(prev_actions_batch)

    new_observations_batch = defaultdict(list)
    for sensor in observations_batch[0]:
        for bid in range(B):
            new_observations_batch[sensor].append(observations_batch[bid][sensor])

    observations_batch = new_observations_batch

    max_traj_len = max(ele.size(0) for ele in prev_actions_batch)
    for bid in range(B):
        for sensor in observations_batch:
            observations_batch[sensor][bid] = _pad_helper(
                observations_batch[sensor][bid], max_traj_len, fill_val=1.0
            )
        next_actions_batch[bid] = _pad_helper(next_actions_batch[bid], max_traj_len)
        prev_actions_batch[bid] = _pad_helper(prev_actions_batch[bid], max_traj_len)
        weights_batch[bid] = _pad_helper(weights_batch[bid], max_traj_len)

    for sensor in observations_batch:
        observations_batch[sensor] = torch.stack(observations_batch[sensor], dim=1)
    
    next_actions_batch = torch.stack(next_actions_batch, dim=1)
    prev_actions_batch = torch.stack(prev_actions_batch, dim=1)
    weights_batch = torch.stack(weights_batch, dim=1)
    not_done_masks = torch.ones_like(next_actions_batch, dtype=torch.float)
    not_done_masks[0] = 0

    observations_batch = ObservationsDict(observations_batch)

    return (
        observations_batch,
        prev_actions_batch,
        not_done_masks,
        next_actions_batch,
        weights_batch,
    )


class PickPlaceDataset(Dataset):
    """Pytorch dataset for object rearrangement task for each episode"""

    def __init__(self, config, content_scenes=["*"], mode="train", use_iw=False, inflection_weight_coef=1.0):
        """
        Args:
            env (habitat.Env): Habitat environment
            config: Config
            mode: 'train'/'val'
        """
        scene_split_name = "train"
        if content_scenes[0] != "*":
            scene_split_name = "_".join(content_scenes)

        self.config = config.TASK_CONFIG
        self.dataset_path = config.DATASET_PATH.format(split=mode, scene_split=scene_split_name)

        self.config.defrost()
        self.config.DATASET.CONTENT_SCENES = content_scenes
        self.config.freeze()

        self.resolution = [self.config.SIMULATOR.RGB_SENSOR.WIDTH, self.config.SIMULATOR.RGB_SENSOR.HEIGHT]
        self.possible_actions = config.TASK_CONFIG.TASK.POSSIBLE_ACTIONS

        self.total_actions = 0
        self.inflections = 0
        self.inflection_weight_coef = inflection_weight_coef

        if use_iw:
            self.inflec_weight = torch.tensor([1.0, inflection_weight_coef])
        else:
            self.inflec_weight = torch.tensor([1.0, 1.0])

        if not self.cache_exists():
            """
            for each scene > load scene in memory > save frames for each
            episode corresponding to that scene
            """
            self.env = habitat.Env(config=self.config)
            self.episodes = self.env._dataset.episodes
            self.instruction_vocab = self.env._dataset.instruction_vocab

            logger.info(
                "Dataset cache not found. Saving rgb, seg, depth scene images"
            )
            logger.info(
                "Number of {} episodes: {}".format(mode, len(self.episodes))
            )

            self.scene_ids = []
            self.scene_episode_dict = {}

            # dict for storing list of episodes for each scene
            for episode in self.episodes:
                if episode.scene_id not in self.scene_ids:
                    self.scene_ids.append(episode.scene_id)
                    self.scene_episode_dict[episode.scene_id] = [episode]
                else:
                    self.scene_episode_dict[episode.scene_id].append(episode)

            self.lmdb_env = lmdb.open(
                self.dataset_path,
                map_size=int(2e12),
                writemap=True,
            )

            self.count = 0
            for scene in tqdm(list(self.scene_episode_dict.keys())):
                for episode in tqdm(self.scene_episode_dict[scene]):
                    self.load_scene(scene, episode)
                    state_index_queue = []
                    try:
                        # TODO: Consider alternative for shortest_paths
                        state_index_queue.extend(range(0, len(episode.reference_replay) - 1))
                    except AttributeError as e:
                        logger.error(e)
                    self.save_frames(state_index_queue, episode)
            print("Inflection weight coef: {}, N: {}, nI: {}".format(self.total_actions / self.inflections, self.total_actions, self.inflections))
            logger.info("Rearrangement database ready!")
            self.env.close()
        else:
            logger.info("Dataset cache found.")
            self.lmdb_env = lmdb.open(
                self.dataset_path,
                readonly=True,
                lock=False,
            )

        self.dataset_length = int(self.lmdb_env.begin().stat()["entries"] / 4)
        self.lmdb_env.close()
        self.lmdb_env = None

    def save_frames(
        self, state_index_queue: List[int], episode: RearrangementEpisode
    ) -> None:
        r"""
        Writes rgb, seg, depth frames to LMDB.
        """
        next_actions = []
        prev_actions = []
        observations = {
            "rgb": [],
            "depth": [],
            "instruction": [],
        }
        reference_replay = episode.reference_replay
        instruction = episode.instruction
        print("Replay len: {}".format(len(reference_replay)))
        for state_index in state_index_queue:
            instruction_tokens = np.array(instruction.instruction_tokens)

            state = reference_replay[state_index]
            position = state.agent_state.position
            rotation = state.agent_state.rotation
            object_states = state.object_states
            sensor_states = state.agent_state.sensor_data

            observation = self.env.sim.get_observations_at(
                position, rotation, sensor_states, object_states
            )

            next_state = reference_replay[state_index + 1]
            next_action = self.possible_actions.index(next_state.action)

            prev_state = reference_replay[state_index]
            prev_action = self.possible_actions.index(prev_state.action)

            observations["depth"].append(observation["depth"])
            observations["rgb"].append(observation["rgb"])
            observations["instruction"].append(instruction_tokens)
            next_actions.append(next_action)
            prev_actions.append(prev_action)
        
        oracle_actions = np.array(next_actions)
        inflection_weights = np.concatenate(([1], oracle_actions[1:] != oracle_actions[:-1]))
        self.total_actions += inflection_weights.shape[0]
        self.inflections += np.sum(inflection_weights)
        inflection_weights = self.inflec_weight[torch.from_numpy(inflection_weights)].numpy()

        sample_key = "{0:0=6d}".format(self.count)
        with self.lmdb_env.begin(write=True) as txn:
            txn.put((sample_key + "_obs").encode(), msgpack_numpy.packb(observations, use_bin_type=True))
            txn.put((sample_key + "_next_action").encode(), np.array(next_actions).tobytes())
            txn.put((sample_key + "_prev_action").encode(), np.array(prev_actions).tobytes())
            txn.put((sample_key + "_weights").encode(), inflection_weights.tobytes())
        
        self.count += 1
        # images_to_video(images=obs_list, output_dir="demos", video_name="dummy_{}".format(self.count))

    def cache_exists(self) -> bool:
        if os.path.exists(self.dataset_path):
            if os.listdir(self.dataset_path):
                return True
        else:
            os.makedirs(self.dataset_path)
        return False
    
    def get_vocab_dict(self) -> VocabDict:
        r"""Returns Instruction VocabDicts"""
        return self.instruction_vocab

    def load_scene(self, scene, episode) -> None:
        self.config.defrost()
        self.config.SIMULATOR.SCENE = scene
        self.config.SIMULATOR.objects = episode.objects
        self.config.freeze()
        self.env.sim.reconfigure(self.config.SIMULATOR)

    def __len__(self) -> int:
        return self.dataset_length

    def __getitem__(self, idx: int):
        r"""Returns batches to trainer.

        batch: (rgb, depth, seg)

        """
        if self.lmdb_env is None:
            self.lmdb_env = lmdb.open(
                self.dataset_path,
                map_size=int(2e12),
                writemap=True,
            )
            self.lmdb_txn = self.lmdb_env.begin()
            self.lmdb_cursor = self.lmdb_txn.cursor()
        
        height, width = int(self.resolution[0]), int(self.resolution[1])

        obs_idx = "{0:0=6d}_obs".format(idx)
        observations_binary = self.lmdb_cursor.get(obs_idx.encode())
        observations = msgpack_numpy.unpackb(observations_binary, raw=False)
        for k, v in observations.items():
            obs = np.array(observations[k])
            observations[k] = torch.from_numpy(obs)

        next_action_idx = "{0:0=6d}_next_action".format(idx)
        next_action_binary = self.lmdb_cursor.get(next_action_idx.encode())
        next_action = np.frombuffer(next_action_binary, dtype="int")
        next_action = torch.from_numpy(np.copy(next_action))

        prev_action_idx = "{0:0=6d}_prev_action".format(idx)
        prev_action_binary = self.lmdb_cursor.get(prev_action_idx.encode())
        prev_action = np.frombuffer(prev_action_binary, dtype="int")
        prev_action = torch.from_numpy(np.copy(prev_action))

        weight_idx = "{0:0=6d}_weights".format(idx)
        weight_binary = self.lmdb_cursor.get(weight_idx.encode())
        weight = np.frombuffer(weight_binary, dtype="float32")
        weight = torch.from_numpy(np.copy(weight))
        weight = torch.where(weight != 1.0, self.inflection_weight_coef, 1.0)

        return idx, observations, next_action, prev_action, weight

habitat_baselines/il/disk_based/dataset/dataset.py
cshizhe/onav_rim
#!/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 contextlib
import numpy as np
import os
import random
import time
import torch
import tqdm

from collections import defaultdict
from typing import Any, Dict, List, Tuple, Union
from torch import Tensor
from torch import distributed as distrib
from numpy import ndarray
from torch.utils.data import DataLoader, ConcatDataset

from habitat import logger
from habitat.core.env import Env, RLEnv
from habitat.core.vector_env import VectorEnv
from habitat.utils import profiling_wrapper
from habitat.utils.visualizations.utils import observations_to_image
from habitat_baselines.common.base_il_trainer import BaseILTrainer
from habitat_baselines.common.baseline_registry import baseline_registry
from habitat_baselines.common.environments import get_env_class
from habitat_baselines.common.tensorboard_utils import TensorboardWriter
from habitat_baselines.il.disk_based.dataset.dataset import PickPlaceDataset, collate_fn
from habitat_baselines.il.disk_based.policy.resnet_policy import PickPlacePolicy
from habitat_baselines.utils.env_utils import construct_envs
from habitat_baselines.utils.common import (
    batch_obs,
    generate_video,
)
from habitat_baselines.rl.ddppo.algo.ddp_utils import (
    EXIT,
    REQUEUE,
    add_signal_handlers,
    init_distrib_slurm,
    load_interrupted_state,
    requeue_job,
    save_interrupted_state,
)


@baseline_registry.register_trainer(name="rearrangement-behavior-cloning-distrib")
class RearrangementBCDistribTrainer(BaseILTrainer):
    r"""Trainer class for PPO algorithm
    Paper: https://arxiv.org/abs/1707.06347.
    """
    supported_tasks = ["PickPlaceTask-v0"]

    def __init__(self, config=None):
        super().__init__(config)

        self.obs_transforms = []

        if config is not None:
            logger.info(f"config: {config}")

        self.device = (
            torch.device("cuda", self.config.TORCH_GPU_ID)
            if torch.cuda.is_available()
            else torch.device("cpu")
        )

    def _make_results_dir(self, split="val"):
        r"""Makes directory for saving eqa-cnn-pretrain eval results."""
        for s_type in ["rgb", "seg", "depth", "top_down_map"]:
            dir_name = self.config.RESULTS_DIR.format(split=split, type=s_type)
            if not os.path.isdir(dir_name):
                os.makedirs(dir_name)

    METRICS_BLACKLIST = {"top_down_map", "collisions.is_collision", "goal_vis_pixels", "rearrangement_reward", "coverage", "collisions.count", "release_failed", "object_receptacle_distance", "exploration_metrics", "room_visitation_map"}

    @classmethod
    def _extract_scalars_from_info(
        cls, info: Dict[str, Any]
    ) -> Dict[str, float]:
        result = {}
        for k, v in info.items():
            if k in cls.METRICS_BLACKLIST:
                continue

            if isinstance(v, dict):
                result.update(
                    {
                        k + "." + subk: subv
                        for subk, subv in cls._extract_scalars_from_info(
                            v
                        ).items()
                        if (k + "." + subk) not in cls.METRICS_BLACKLIST
                    }
                )
            # Things that are scalar-like will have an np.size of 1.
            # Strings also have an np.size of 1, so explicitly ban those
            elif np.size(v) == 1 and not isinstance(v, str):
                result[k] = float(v)

        return result

    @classmethod
    def _extract_scalars_from_infos(
        cls, infos: List[Dict[str, Any]]
    ) -> Dict[str, List[float]]:

        results = defaultdict(list)
        for i in range(len(infos)):
            for k, v in cls._extract_scalars_from_info(infos[i]).items():
                results[k].append(v)

        return results
    
    @staticmethod
    def _pause_envs(
        envs_to_pause: List[int],
        envs: Union[VectorEnv, RLEnv, Env],
        test_recurrent_hidden_states: Tensor,
        not_done_masks: Tensor,
        current_episode_reward: Tensor,
        prev_actions: Tensor,
        batch: Dict[str, Tensor],
        rgb_frames: Union[List[List[Any]], List[List[ndarray]]],
    ) -> Tuple[
        Union[VectorEnv, RLEnv, Env],
        Tensor,
        Tensor,
        Tensor,
        Tensor,
        Dict[str, Tensor],
        List[List[Any]],
    ]:
        # pausing self.envs with no new episode
        if len(envs_to_pause) > 0:
            state_index = list(range(envs.num_envs))
            for idx in reversed(envs_to_pause):
                state_index.pop(idx)
                envs.pause_at(idx)
            
            # indexing along the batch dimensions
            test_recurrent_hidden_states = test_recurrent_hidden_states[
                :, state_index
            ]

            not_done_masks = not_done_masks[state_index]
            current_episode_reward = current_episode_reward[state_index]
            prev_actions = prev_actions[state_index]

            for k, v in batch.items():
                batch[k] = v[state_index]

            rgb_frames = [rgb_frames[i] for i in state_index]

        return (
            envs,
            test_recurrent_hidden_states,
            not_done_masks,
            current_episode_reward,
            prev_actions,
            batch,
            rgb_frames,
        )

    def _setup_model(self, observation_space, action_space, model_config):
        model_config.defrost()
        model_config.TORCH_GPU_ID = self.config.TORCH_GPU_ID
        model_config.freeze()

        model = PickPlacePolicy(observation_space, action_space, model_config)
        return model

    def _setup_dataset(self):
        config = self.config

        content_scenes = self.envs.scene_splits()[0]
        logger.info("Scene splits: {}".format(content_scenes))
        datasets = []
        for scene in ["q9vSo1VnCiC"]:
            dataset = PickPlaceDataset(
                config,
                content_scenes=[scene],
                use_iw=config.IL.USE_IW,
                inflection_weight_coef=config.MODEL.inflection_weight_coef
            )
            datasets.append(dataset)

        concat_dataset = ConcatDataset(datasets)
        return concat_dataset


    @profiling_wrapper.RangeContext("train")
    def train(self) -> None:
        r"""Main method for training PPO.

        Returns:
            None
        """

        config = self.config

        self.local_rank, tcp_store = init_distrib_slurm(
            self.config.IL.distrib_backend
        )
        add_signal_handlers()

        self.world_rank = distrib.get_rank()
        self.world_size = distrib.get_world_size()

        self.config.defrost()
        self.config.TORCH_GPU_ID = self.local_rank
        self.config.SIMULATOR_GPU_ID = self.local_rank
        # Multiply by the number of simulators to make sure they also get unique seeds
        self.config.TASK_CONFIG.SEED += (
            self.world_rank * self.config.NUM_PROCESSES
        )
        self.config.freeze()

        if torch.cuda.is_available():
            self.device = torch.device("cuda", self.local_rank)
            torch.cuda.set_device(self.device)
        else:
            self.device = torch.device("cpu")

        random.seed(self.config.TASK_CONFIG.SEED)
        np.random.seed(self.config.TASK_CONFIG.SEED)
        torch.manual_seed(self.config.TASK_CONFIG.SEED)

        self.envs = construct_envs(
            config,
            get_env_class(config.ENV_NAME),
            workers_ignore_signals=True
        )

        rearrangement_dataset = self._setup_dataset()
        batch_size = config.IL.BehaviorCloning.batch_size

        train_sampler = torch.utils.data.distributed.DistributedSampler(
            rearrangement_dataset,
            num_replicas=self.world_size,
            rank=self.world_rank,
            shuffle=True,
            drop_last=True,
        )
        logger.info("Setup dataloader")

        train_loader = DataLoader(
            rearrangement_dataset,
            collate_fn=collate_fn,
            batch_size=batch_size,
            shuffle=False,
            sampler=train_sampler,
            num_workers=0,
        )
        logger.info("Dataloader setup")

        logger.info(
            "[ train_loader has {} samples ]".format(
                len(rearrangement_dataset)
            )
        )
        logger.info("Setting up distributed model...")

        action_space = self.envs.action_spaces[0]

        self.model = self._setup_model(
            self.envs.observation_spaces[0],
            action_space,
            config.MODEL
        )

        self.model.to(self.device)
        # Distributed data parallel setup
        if torch.cuda.is_available():
            self.model = torch.nn.parallel.DistributedDataParallel(
                self.model,
                device_ids=[self.device],
                output_device=self.device
            )
        else:
            self.model = torch.nn.parallel.DistributedDataParallel(self.model)

        optim = torch.optim.Adam(
            filter(lambda p: p.requires_grad, self.model.parameters()),
            lr=float(config.IL.BehaviorCloning.lr),
        )

        scheduler = torch.optim.lr_scheduler.OneCycleLR(
            optim, max_lr=config.IL.BehaviorCloning.lr,
            steps_per_epoch=len(train_loader), epochs=config.IL.BehaviorCloning.max_epochs
        )
        cross_entropy_loss = torch.nn.CrossEntropyLoss(reduction="none")

        epoch, t = 1, 0
        softmax = torch.nn.Softmax(dim=1)

        interrupted_state = load_interrupted_state()
        if interrupted_state is not None:
            self.model.load_state_dict(interrupted_state["state_dict"])
            optim.load_state_dict(
                interrupted_state["optim_state"]
            )
            scheduler.load_state_dict(interrupted_state["lr_sched_state"])

            requeue_stats = interrupted_state["requeue_stats"]
            epoch = requeue_stats["epoch"]
            t = requeue_stats["t"]

        logger.info("Starting training")
        with (
            TensorboardWriter(
                self.config.TENSORBOARD_DIR, flush_secs=self.flush_secs
            )
            if self.world_rank == 0
            else contextlib.suppress()
        ) as writer:
            while epoch <= config.IL.BehaviorCloning.max_epochs:
                logger.info("Epoch: {}".format(epoch))
                train_loader.sampler.set_epoch(epoch)
                start_time = time.time()
                avg_loss = 0.0
                avg_load_time = 0.0
                avg_slice_time = 0.0
                avg_train_time = 0.0

                batch_start_time = time.time()

                if EXIT.is_set():
                    profiling_wrapper.range_pop()  # train update

                    self.envs.close()

                    if REQUEUE.is_set() and self.world_rank == 0:
                        requeue_stats = dict(
                            epoch=epoch,
                            t=t,
                        )
                        save_interrupted_state(
                            dict(
                                state_dict=self.model.state_dict(),
                                optim_state=optim.state_dict(),
                                lr_sched_state=scheduler.state_dict(),
                                config=self.config,
                                requeue_stats=requeue_stats,
                            ),
                            config.INTERRUPTED_STATE_FILE_PATH
                        )

                    requeue_job()
                    return

                for batch in train_loader:
                    torch.cuda.empty_cache()
                    t += 1

                    (
                     observations_batch,
                     gt_prev_action,
                     episode_not_done,
                     gt_next_action,
                     inflec_weights
                    ) = batch

                    avg_load_time += ((time.time() - batch_start_time) / 60)

                    rnn_hidden_states = torch.zeros(
                        config.MODEL.STATE_ENCODER.num_recurrent_layers,
                        gt_prev_action.shape[1],
                        config.MODEL.STATE_ENCODER.hidden_size,
                        device=self.device,
                    )

                    optim.zero_grad()

                    num_samples = gt_prev_action.shape[0]
                    timestep_batch_size = config.IL.BehaviorCloning.timestep_batch_size
                    num_steps = num_samples // timestep_batch_size + (num_samples % timestep_batch_size != 0)
                    batch_loss = 0
                    for i in range(num_steps):
                        slice_start_time = time.time()
                        start_idx = i * timestep_batch_size
                        end_idx = start_idx + timestep_batch_size
                        observations_batch_sample = {
                            k: v[start_idx:end_idx].to(device=self.device)
                            for k, v in observations_batch.items()
                        }

                        gt_next_action_sample = gt_next_action[start_idx:end_idx].long().to(self.device)
                        gt_prev_action_sample = gt_prev_action[start_idx:end_idx].long().to(self.device)
                        episode_not_dones_sample = episode_not_done[start_idx:end_idx].long().to(self.device)
                        inflec_weights_sample = inflec_weights[start_idx:end_idx].long().to(self.device)

                        avg_slice_time += ((time.time() - slice_start_time) / 60)

                        train_time = time.time()

                        if i != num_steps - 1 :
                            with self.model.no_sync():
                                logits, rnn_hidden_states = self.model(
                                    observations_batch_sample,
                                    rnn_hidden_states,
                                    gt_prev_action_sample,
                                    episode_not_dones_sample
                                )

                                T, N = gt_next_action_sample.shape
                                logits = logits.view(T, N, -1)

                                action_loss = cross_entropy_loss(logits.permute(0, 2, 1), gt_next_action_sample)
                                denom = inflec_weights_sample.sum(0)
                                denom[denom == 0.0] = 1
                                action_loss = ((inflec_weights_sample * action_loss).sum(0) / denom).mean()
                                loss = (action_loss / num_steps)
                                loss.backward()
                        else:
                            logits, rnn_hidden_states = self.model(
                                observations_batch_sample,
                                rnn_hidden_states,
                                gt_prev_action_sample,
                                episode_not_dones_sample
                            )

                            T, N = gt_next_action_sample.shape
                            logits = logits.view(T, N, -1)

                            action_loss = cross_entropy_loss(logits.permute(0, 2, 1), gt_next_action_sample)
                            denom = inflec_weights_sample.sum(0)
                            denom[denom == 0.0] = 1
                            action_loss = ((inflec_weights_sample * action_loss).sum(0) / denom).mean()
                            loss = (action_loss / num_steps)
                            loss.backward()
                        batch_loss += loss.item()
                        avg_train_time += ((time.time() - train_time) / 60)
                        rnn_hidden_states = rnn_hidden_states.detach()

                    # Sync loss
                    stats = torch.tensor(
                        [batch_loss],
                        device=self.device,
                    )
                    distrib.all_reduce(stats)
                    batch_loss = stats[0].item()

                    if t % config.LOG_INTERVAL == 0:
                        logger.info(
                            "[ Epoch: {}; iter: {}; loss: {:.3f}; load time: {:.3f}; train time: {:.3f};]".format(
                                epoch, t, batch_loss / self.world_size, avg_load_time / t, avg_train_time / t,
                            )
                        )

                    optim.step()
                    scheduler.step()
                    batch_start_time = time.time()
                    avg_loss += batch_loss

                end_time = time.time()
                time_taken = "{:.1f}".format((end_time - start_time) / 60)
                avg_loss = avg_loss / len(train_loader)
                avg_train_time = avg_train_time / len(train_loader)
                avg_load_time = avg_load_time / len(train_loader)

                if self.world_rank == 0:
                    logger.info(
                        "[ Epoch {} completed. Time taken: {} minutes. Load time: {} mins. Train time: {} mins]".format(
                            epoch, time_taken, avg_load_time, avg_train_time
                        )
                    )
                    logger.info("[ Average loss: {:.3f} ]".format(avg_loss / self.world_size))
                    writer.add_scalar("avg_train_loss", avg_loss / self.world_size, epoch)

                    print("-----------------------------------------")

                    if epoch % config.CHECKPOINT_INTERVAL == 0:
                        self.save_checkpoint(
                            self.model.state_dict(), "model_{}.ckpt".format(epoch)
                        )

                epoch += 1
        logger.info("Epochs ended")
        self.envs.close()
        logger.info("Closing envs")

    def _eval_checkpoint(
        self,
        checkpoint_path: str,
        writer: TensorboardWriter,
        checkpoint_index: int = 0,
    ) -> None:
        r"""Evaluates a single checkpoint.

        Args:
            checkpoint_path: path of checkpoint
            writer: tensorboard writer object for logging to tensorboard
            checkpoint_index: index of cur checkpoint for logging

        Returns:
            None
        """
        config = self.config

        config.defrost()
        config.TASK_CONFIG.DATASET.SPLIT = config.EVAL.SPLIT
        config.freeze()

        if len(self.config.VIDEO_OPTION) > 0:
            config.defrost()
            config.TASK_CONFIG.TASK.MEASUREMENTS.append("COLLISIONS")
            if config.SHOW_TOP_DOWN_MAP:
                config.TASK_CONFIG.TASK.MEASUREMENTS.append("TOP_DOWN_MAP")
            config.freeze()
        
        config.defrost()
        if hasattr(config.EVAL, "semantic_metrics") and config.EVAL.semantic_metrics:
            config.TASK_CONFIG.TASK.MEASUREMENTS = config.TASK_CONFIG.TASK.MEASUREMENTS + ["ROOM_VISITATION_MAP", "EXPLORATION_METRICS"]
            logger.info("Setting up semantic exploration metrics")
        config.freeze()

        self.envs = construct_envs(config, get_env_class(config.ENV_NAME))
        batch_size = config.IL.BehaviorCloning.batch_size

        logger.info(
            "[ val_loader has {} samples ]".format(
                self.envs.count_episodes()
            )
        )

        action_space = self.envs.action_spaces[0]

        self.model = self._setup_model(
            self.envs.observation_spaces[0],
            action_space,
            config.MODEL
        )
        
        # Map location CPU is almost always better than mapping to a CUDA device.
        ckpt_dict = torch.load(checkpoint_path, map_location='cpu') #self.device)
        self.model.load_state_dict(
            {
                k.replace("module.", ""): v
                for k, v in ckpt_dict.items()
                if "module" in k
            }, strict=True)
        self.model.to(self.device)
        self.model.eval()

        observations = self.envs.reset()
        batch = batch_obs(observations, device=self.device)

        current_episode_reward = torch.zeros(
            self.envs.num_envs, 1, device=self.device
        )

        prev_actions = torch.zeros(
            self.envs.num_envs, 1, device=self.device, dtype=torch.long
        )
        not_done_masks = torch.zeros(
            self.envs.num_envs, 1, device=self.device
        )
        rnn_hidden_states = torch.zeros(
            config.MODEL.STATE_ENCODER.num_recurrent_layers,
            self.envs.num_envs,
            config.MODEL.STATE_ENCODER.hidden_size,
            device=self.device,
        )
        stats_episodes: Dict[
            Any, Any
        ] = {}  # dict of dicts that stores stats per episode

        rgb_frames = [
            [] for _ in range(self.config.NUM_PROCESSES)
        ]  # type: List[List[np.ndarray]]
        if len(self.config.VIDEO_OPTION) > 0:
            os.makedirs(self.config.VIDEO_DIR, exist_ok=True)

        number_of_eval_episodes = self.config.TEST_EPISODE_COUNT
        if number_of_eval_episodes == -1:
            number_of_eval_episodes = sum(self.envs.number_of_episodes)
        else:
            total_num_eps = sum(self.envs.number_of_episodes)
            if total_num_eps < number_of_eval_episodes:
                logger.warn(
                    f"Config specified {number_of_eval_episodes} eval episodes"
                    ", dataset only has {total_num_eps}."
                )
                logger.warn(f"Evaluating with {total_num_eps} instead.")
                number_of_eval_episodes = total_num_eps

        pbar = tqdm.tqdm(total=number_of_eval_episodes)
        possible_actions = self.config.TASK_CONFIG.TASK.POSSIBLE_ACTIONS
        episode_count = 0

        while (
            len(stats_episodes) < number_of_eval_episodes
            and self.envs.num_envs > 0
        ):
            current_episodes = self.envs.current_episodes()

            with torch.no_grad():
                (
                    logits,
                    rnn_hidden_states
                ) = self.model(
                    batch,
                    rnn_hidden_states,
                    prev_actions,
                    not_done_masks
                )

                actions = torch.argmax(logits, dim=1)
                prev_actions.copy_(actions.unsqueeze(1))

            action_names = [possible_actions[a.item()] for a in actions.to(device="cpu")]
            # NB: Move actions to CPU.  If CUDA tensors are
            # sent in to env.step(), that will create CUDA contexts
            # in the subprocesses.
            # For backwards compatibility, we also call .item() to convert to
            # an int
            step_data = [a.item() for a in actions.to(device="cpu")]
            outputs = self.envs.step(step_data)

            observations, rewards_l, dones, infos = [
                list(x) for x in zip(*outputs)
            ]
            batch = batch_obs(observations, device=self.device)

            not_done_masks = torch.tensor(
                [[0.0] if done else [1.0] for done in dones],
                dtype=torch.float,
                device=self.device,
            )

            rewards = torch.tensor(
                rewards_l, dtype=torch.float, device=self.device
            ).unsqueeze(1)
            current_episode_reward += rewards
            next_episodes = self.envs.current_episodes()
            envs_to_pause = []
            n_envs = self.envs.num_envs
            for i in range(n_envs):
                if (
                    next_episodes[i].scene_id,
                    next_episodes[i].episode_id,
                ) in stats_episodes:
                    envs_to_pause.append(i)

                # episode ended
                if not_done_masks[i].item() == 0:
                    pbar.update()
                    episode_stats = {}
                    episode_stats["reward"] = current_episode_reward[i].item()
                    episode_stats.update(
                        self._extract_scalars_from_info(infos[i])
                    )
                    current_episode_reward[i] = 0
                    # use scene_id + episode_id as unique id for storing stats
                    stats_episodes[
                        (
                            current_episodes[i].scene_id,
                            current_episodes[i].episode_id,
                        )
                    ] = episode_stats
                    next_episodes = self.envs.current_episodes()
                    episode_count += 1

                    if len(self.config.VIDEO_OPTION) > 0:
                        generate_video(
                            video_option=self.config.VIDEO_OPTION,
                            video_dir=self.config.VIDEO_DIR,
                            images=rgb_frames[i],
                            episode_id=episode_count, #current_episodes[i].episode_id,
                            checkpoint_idx=checkpoint_index,
                            metrics=self._extract_scalars_from_info(infos[i]),
                            tb_writer=writer,
                        )

                        rgb_frames[i] = []
                # episode continues
                elif len(self.config.VIDEO_OPTION) > 0:
                    # TODO move normalization / channel changing out of the policy and undo it here
                    frame = observations_to_image(
                        {k: v[i] for k, v in batch.items()}, infos[i]
                    )
                    rgb_frames[i].append(frame)


            (
                self.envs,
                rnn_hidden_states,
                not_done_masks,
                current_episode_reward,
                prev_actions,
                batch,
                rgb_frames,
            ) = self._pause_envs(
                envs_to_pause,
                self.envs,
                rnn_hidden_states,
                not_done_masks,
                current_episode_reward,
                prev_actions,
                batch,
                rgb_frames,
            )

        num_episodes = len(stats_episodes)
        aggregated_stats = {}
        for stat_key in next(iter(stats_episodes.values())).keys():
            aggregated_stats[stat_key] = (
                sum(v[stat_key] for v in stats_episodes.values())
                / num_episodes
            )

        for k, v in aggregated_stats.items():
            logger.info(f"Average episode {k}: {v:.4f}")

        step_id = checkpoint_index
        if "extra_state" in ckpt_dict and "step" in ckpt_dict["extra_state"]:
            step_id = ckpt_dict["extra_state"]["step"]

        writer.add_scalars(
            "eval_reward",
            {"average reward": aggregated_stats["reward"]},
            step_id,
        )

        metrics = {k: v for k, v in aggregated_stats.items() if k != "reward"}
        if len(metrics) > 0:
            writer.add_scalars("eval_metrics", metrics, step_id)

        self.envs.close()

        print ("environments closed")

habitat_baselines/il/disk_based/il_ddp_trainer.py
cshizhe/onav_rim
#!/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
import numpy as np
import random
import time
import torch
import tqdm

from collections import defaultdict
from typing import Any, Dict, List, Tuple, Union

from torch import Tensor
from numpy import ndarray
from torch.utils.data import DataLoader, ConcatDataset

from habitat import logger, make_dataset
from habitat.core.env import Env, RLEnv
from habitat.core.vector_env import VectorEnv
from habitat.utils import profiling_wrapper
from habitat.utils.visualizations.utils import observations_to_image, append_text_to_image
from habitat_baselines.common.base_il_trainer import BaseILTrainer
from habitat_baselines.common.baseline_registry import baseline_registry
from habitat_baselines.common.environments import get_env_class
from habitat_baselines.common.tensorboard_utils import TensorboardWriter
from habitat_baselines.il.disk_based.dataset.dataset import PickPlaceDataset, collate_fn
from habitat_baselines.il.disk_based.policy.resnet_policy import PickPlacePolicy
from habitat_baselines.utils.env_utils import construct_envs
from habitat_baselines.utils.common import (
    batch_obs,
    generate_video,
)


@baseline_registry.register_trainer(name="rearrangement-behavior-cloning")
class RearrangementBCTrainer(BaseILTrainer):
    r"""Trainer class for PPO algorithm
    Paper: https://arxiv.org/abs/1707.06347.
    """
    supported_tasks = ["PickPlaceTask-v0"]

    def __init__(self, config=None):
        super().__init__(config)

        self.obs_transforms = []

        if config is not None:
            logger.info(f"config: {config}")

        self.device = (
            torch.device("cuda", self.config.TORCH_GPU_ID)
            if torch.cuda.is_available()
            else torch.device("cpu")
        )

    def _make_results_dir(self, split="val"):
        r"""Makes directory for saving eqa-cnn-pretrain eval results."""
        for s_type in ["rgb", "seg", "depth", "top_down_map"]:
            dir_name = self.config.RESULTS_DIR.format(split=split, type=s_type)
            if not os.path.isdir(dir_name):
                os.makedirs(dir_name)

    METRICS_BLACKLIST = {"top_down_map", "collisions.is_collision", "goal_vis_pixels", "rearrangement_reward", "coverage"}

    @classmethod
    def _extract_scalars_from_info(
        cls, info: Dict[str, Any]
    ) -> Dict[str, float]:
        result = {}
        for k, v in info.items():
            if k in cls.METRICS_BLACKLIST:
                continue

            if isinstance(v, dict):
                result.update(
                    {
                        k + "." + subk: subv
                        for subk, subv in cls._extract_scalars_from_info(
                            v
                        ).items()
                        if (k + "." + subk) not in cls.METRICS_BLACKLIST
                    }
                )
            # Things that are scalar-like will have an np.size of 1.
            # Strings also have an np.size of 1, so explicitly ban those
            elif np.size(v) == 1 and not isinstance(v, str):
                result[k] = float(v)

        return result

    @classmethod
    def _extract_scalars_from_infos(
        cls, infos: List[Dict[str, Any]]
    ) -> Dict[str, List[float]]:

        results = defaultdict(list)
        for i in range(len(infos)):
            for k, v in cls._extract_scalars_from_info(infos[i]).items():
                results[k].append(v)

        return results
    
    @staticmethod
    def _pause_envs(
        envs_to_pause: List[int],
        envs: Union[VectorEnv, RLEnv, Env],
        test_recurrent_hidden_states: Tensor,
        not_done_masks: Tensor,
        current_episode_reward: Tensor,
        prev_actions: Tensor,
        batch: Dict[str, Tensor],
        rgb_frames: Union[List[List[Any]], List[List[ndarray]]],
    ) -> Tuple[
        Union[VectorEnv, RLEnv, Env],
        Tensor,
        Tensor,
        Tensor,
        Tensor,
        Dict[str, Tensor],
        List[List[Any]],
    ]:
        # pausing self.envs with no new episode
        if len(envs_to_pause) > 0:
            state_index = list(range(envs.num_envs))
            for idx in reversed(envs_to_pause):
                state_index.pop(idx)
                envs.pause_at(idx)
            
            # indexing along the batch dimensions
            test_recurrent_hidden_states = test_recurrent_hidden_states[
                :, state_index
            ]

            not_done_masks = not_done_masks[state_index]
            current_episode_reward = current_episode_reward[state_index]
            prev_actions = prev_actions[state_index]

            for k, v in batch.items():
                batch[k] = v[state_index]

            rgb_frames = [rgb_frames[i] for i in state_index]

        return (
            envs,
            test_recurrent_hidden_states,
            not_done_masks,
            current_episode_reward,
            prev_actions,
            batch,
            rgb_frames,
        )

    def _setup_actor_critic_agent(self, observation_space, action_space, model_config):
        model_config.defrost()
        model_config.TORCH_GPU_ID = self.config.TORCH_GPU_ID
        model_config.freeze()

        policy = baseline_registry.get_policy(self.config.IL.POLICY.name)
        policy = policy.from_config(
            self.config, observation_space, action_space
        )
        return policy

    def _setup_dataset(self):
        config = self.config

        content_scenes = config.TASK_CONFIG.DATASET.CONTENT_SCENES
        if "*" in content_scenes:
            dataset = make_dataset(
                config.TASK_CONFIG.DATASET.TYPE, config=config.TASK_CONFIG.DATASET
            )
            content_scenes = dataset.get_scenes_to_load(config.TASK_CONFIG.DATASET)
        datasets = []
        print(content_scenes)
        for scene in ["q9vSo1VnCiC"]:
            dataset = PickPlaceDataset(
                config,
                content_scenes=[scene],
                use_iw=config.IL.USE_IW,
                inflection_weight_coef=config.MODEL.inflection_weight_coef
            )
            datasets.append(dataset)

        concat_dataset = ConcatDataset(datasets)
        return concat_dataset


    @profiling_wrapper.RangeContext("train")
    def train(self) -> None:
        r"""Main method for training PPO.

        Returns:
            None
        """

        config = self.config

        random.seed(self.config.TASK_CONFIG.SEED)
        np.random.seed(self.config.TASK_CONFIG.SEED)
        torch.manual_seed(self.config.TASK_CONFIG.SEED)

        self.envs = construct_envs(
            config,
            get_env_class(config.ENV_NAME),
            workers_ignore_signals=True
        )

        rearrangement_dataset = self._setup_dataset()
        batch_size = config.IL.BehaviorCloning.batch_size

        logger.info("Setup dataloader")

        train_loader = DataLoader(
            rearrangement_dataset,
            collate_fn=collate_fn,
            batch_size=batch_size,
            shuffle=True,
            num_workers=8,
            drop_last=True,
        )
        logger.info("Dataloader setup")

        logger.info(
            "[ train_loader has {} samples ]".format(
                len(rearrangement_dataset)
            )
        )

        action_space = self.envs.action_spaces[0]

        self.policy = self._setup_actor_critic_agent(
            self.envs.observation_spaces[0],
            action_space,
            config.MODEL
        )
        self.policy = torch.nn.DataParallel(self.policy, dim=1)
        self.policy.to(self.device)

        optim = torch.optim.Adam(
            filter(lambda p: p.requires_grad, self.policy.parameters()),
            lr=float(config.IL.BehaviorCloning.lr),
        )

        scheduler = torch.optim.lr_scheduler.OneCycleLR(
            optim, max_lr=config.IL.BehaviorCloning.lr,
            steps_per_epoch=len(train_loader), epochs=config.IL.BehaviorCloning.max_epochs
        )

        cross_entropy_loss = torch.nn.CrossEntropyLoss(reduction="none")

        epoch, t = 1, 0
        logger.info("Starting training: {}".format(self.device))
        with (
            TensorboardWriter(
                self.config.TENSORBOARD_DIR, flush_secs=self.flush_secs
            )
        ) as writer:
            while epoch <= config.IL.BehaviorCloning.max_epochs:
                logger.info("Epoch: {}".format(epoch))
                start_time = time.time()
                avg_loss = 0.0
                avg_load_time = 0.0
                avg_slice_time = 0.0
                avg_train_time = 0.0

                batch_start_time = time.time()

                for batch in train_loader:
                    torch.cuda.empty_cache()
                    t += 1

                    (
                     observations_batch,
                     gt_prev_action,
                     episode_not_done,
                     gt_next_action,
                     inflec_weights
                    ) = batch

                    avg_load_time += ((time.time() - batch_start_time) / 60)

                    rnn_hidden_states = torch.zeros(
                        config.MODEL.STATE_ENCODER.num_recurrent_layers,
                        gt_prev_action.shape[1],
                        config.MODEL.STATE_ENCODER.hidden_size,
                        device=self.device,
                    )

                    optim.zero_grad()

                    num_samples = gt_prev_action.shape[0]
                    timestep_batch_size = config.IL.BehaviorCloning.timestep_batch_size
                    num_steps = num_samples // timestep_batch_size + (num_samples % timestep_batch_size != 0)
                    batch_loss = 0
                    for i in range(num_steps):
                        slice_start_time = time.time()
                        start_idx = i * timestep_batch_size
                        end_idx = start_idx + timestep_batch_size
                        observations_batch_sample = {
                            k: v[start_idx:end_idx].to(device=self.device)
                            for k, v in observations_batch.items()
                        }

                        gt_next_action_sample = gt_next_action[start_idx:end_idx].long().to(self.device)
                        gt_prev_action_sample = gt_prev_action[start_idx:end_idx].long().to(self.device)
                        episode_not_dones_sample = episode_not_done[start_idx:end_idx].long().to(self.device)
                        inflec_weights_sample = inflec_weights[start_idx:end_idx].long().to(self.device)

                        avg_slice_time += ((time.time() - slice_start_time) / 60)

                        train_time = time.time()

                        logits, rnn_hidden_states = self.policy(
                            observations_batch_sample,
                            rnn_hidden_states,
                            gt_prev_action_sample,
                            episode_not_dones_sample
                        )

                        T, N = gt_next_action_sample.shape
                        logits = logits.view(T, N, -1)

                        action_loss = cross_entropy_loss(logits.permute(0, 2, 1), gt_next_action_sample)
                        denom = inflec_weights_sample.sum(0)
                        denom[denom == 0.0] = 1
                        action_loss = ((inflec_weights_sample * action_loss).sum(0) / denom).mean()
                        loss = (action_loss / num_steps)
                        loss.backward()
                        batch_loss += loss.item()
                        rnn_hidden_states = rnn_hidden_states.detach()

                        avg_train_time += ((time.time() - train_time) / 60)

                    if t % config.LOG_INTERVAL == 0:
                        logger.info(
                            "[ Epoch: {}; iter: {}; loss: {:.3f}; load time: {:.3f}; train time: {:.3f};]".format(
                                epoch, t, batch_loss, avg_load_time / t, avg_train_time / t,
                            )
                        )

                    optim.step()
                    scheduler.step()
                    batch_start_time = time.time()
                    avg_loss += batch_loss

                end_time = time.time()
                time_taken = "{:.1f}".format((end_time - start_time) / 60)
                avg_loss = avg_loss / len(train_loader)
                avg_train_time = avg_train_time / len(train_loader)
                avg_load_time = avg_load_time / len(train_loader)

                logger.info(
                    "[ Epoch {} completed. Time taken: {} minutes. Load time: {} mins. Train time: {} mins]".format(
                        epoch, time_taken, avg_load_time, avg_train_time
                    )
                )
                logger.info("[ Average loss: {:.3f} ]".format(avg_loss))
                writer.add_scalar("avg_train_loss", avg_loss, epoch)

                print("-----------------------------------------")

                if epoch % config.CHECKPOINT_INTERVAL == 0:
                    self.save_checkpoint(
                        self.policy.module.state_dict(), "model_{}.ckpt".format(epoch)
                    )

                epoch += 1
        logger.info("Epochs ended")
        self.envs.close()
        logger.info("Closing envs")

    def _eval_checkpoint(
        self,
        checkpoint_path: str,
        writer: TensorboardWriter,
        checkpoint_index: int = 0,
    ) -> None:
        r"""Evaluates a single checkpoint.
        Args:
            checkpoint_path: path of checkpoint
            writer: tensorboard writer object for logging to tensorboard
            checkpoint_index: index of cur checkpoint for logging
        Returns:
            None
        """
        config = self.config

        config.defrost()
        config.TASK_CONFIG.DATASET.SPLIT = config.EVAL.SPLIT
        config.freeze()

        if len(self.config.VIDEO_OPTION) > 0:
            config.defrost()
            config.TASK_CONFIG.TASK.MEASUREMENTS.append("COLLISIONS")
            if config.SHOW_TOP_DOWN_MAP:
                config.TASK_CONFIG.TASK.MEASUREMENTS.append("TOP_DOWN_MAP")
            config.freeze()

        self.envs = construct_envs(config, get_env_class(config.ENV_NAME))
        batch_size = config.IL.BehaviorCloning.batch_size

        logger.info(
            "[ val_loader has {} samples ]".format(
                self.envs.count_episodes()
            )
        )

        action_space = self.envs.action_spaces[0]

        self.policy = self._setup_actor_critic_agent(
            self.envs.observation_spaces[0],
            action_space,
            config.MODEL
        )

        # Map location CPU is almost always better than mapping to a CUDA device.
        ckpt_dict = torch.load(checkpoint_path, map_location='cpu') #self.device)
        self.policy.load_state_dict(ckpt_dict, strict=True)
        self.policy.to(self.device)
        self.policy.eval()

        observations = self.envs.reset()
        batch = batch_obs(observations, device=self.device)

        current_episode_reward = torch.zeros(
            self.envs.num_envs, 1, device=self.device
        )

        prev_actions = torch.zeros(
            self.envs.num_envs, 1, device=self.device, dtype=torch.long
        )
        not_done_masks = torch.zeros(
            self.envs.num_envs, 1, device=self.device
        )
        rnn_hidden_states = torch.zeros(
            config.MODEL.STATE_ENCODER.num_recurrent_layers,
            self.envs.num_envs,
            config.MODEL.STATE_ENCODER.hidden_size,
            device=self.device,
        )
        stats_episodes: Dict[
            Any, Any
        ] = {}  # dict of dicts that stores stats per episode

        rgb_frames = [
            [] for _ in range(self.config.NUM_PROCESSES)
        ]  # type: List[List[np.ndarray]]
        if len(self.config.VIDEO_OPTION) > 0:
            os.makedirs(self.config.VIDEO_DIR, exist_ok=True)

        number_of_eval_episodes = self.config.TEST_EPISODE_COUNT
        if number_of_eval_episodes == -1:
            number_of_eval_episodes = sum(self.envs.number_of_episodes)
        else:
            total_num_eps = sum(self.envs.number_of_episodes)
            if total_num_eps < number_of_eval_episodes:
                logger.warn(
                    f"Config specified {number_of_eval_episodes} eval episodes"
                    ", dataset only has {total_num_eps}."
                )
                logger.warn(f"Evaluating with {total_num_eps} instead.")
                number_of_eval_episodes = total_num_eps

        pbar = tqdm.tqdm(total=number_of_eval_episodes)
        action_indices = [1] * self.envs.num_envs
        action_match_count = [(0, 0)] * self.envs.num_envs
        possible_actions = self.config.TASK_CONFIG.TASK.POSSIBLE_ACTIONS
        episode_count = 0
        while (
            len(stats_episodes) < number_of_eval_episodes
            and self.envs.num_envs > 0
        ):
            current_episodes = self.envs.current_episodes()

            with torch.no_grad():
                (
                    logits,
                    rnn_hidden_states
                ) = self.policy(
                    batch,
                    rnn_hidden_states,
                    prev_actions,
                    not_done_masks
                )

                actions = torch.argmax(logits, dim=1)
                prev_actions.copy_(actions.unsqueeze(1))

            action_names = [possible_actions[a.item()] for a in actions.to(device="cpu")]
            # NB: Move actions to CPU.  If CUDA tensors are
            # sent in to env.step(), that will create CUDA contexts
            # in the subprocesses.
            # For backwards compatibility, we also call .item() to convert to
            # an int
            step_data = [a.item() for a in actions.to(device="cpu")]
            outputs = self.envs.step(step_data)

            observations, rewards_l, dones, infos = [
                list(x) for x in zip(*outputs)
            ]
            batch = batch_obs(observations, device=self.device)

            not_done_masks = torch.tensor(
                [[0.0] if done else [1.0] for done in dones],
                dtype=torch.float,
                device=self.device,
            )

            rewards = torch.tensor(
                rewards_l, dtype=torch.float, device=self.device
            ).unsqueeze(1)
            current_episode_reward += rewards
            next_episodes = self.envs.current_episodes()
            envs_to_pause = []
            n_envs = self.envs.num_envs
            for i in range(n_envs):
                if (
                    next_episodes[i].scene_id,
                    next_episodes[i].episode_id,
                ) in stats_episodes:
                    envs_to_pause.append(i)

                # episode ended
                if not_done_masks[i].item() == 0:
                    pbar.update()
                    episode_stats = {}
                    episode_stats["reward"] = current_episode_reward[i].item()
                    episode_stats.update(
                        self._extract_scalars_from_info(infos[i])
                    )
                    current_episode_reward[i] = 0
                    # use scene_id + episode_id as unique id for storing stats
                    stats_episodes[
                        (
                            current_episodes[i].scene_id,
                            current_episodes[i].episode_id,
                        )
                    ] = episode_stats
                    next_episodes = self.envs.current_episodes()
                    episode_count += 1

                    if len(self.config.VIDEO_OPTION) > 0:
                        generate_video(
                            video_option=self.config.VIDEO_OPTION,
                            video_dir=self.config.VIDEO_DIR,
                            images=rgb_frames[i],
                            episode_id=episode_count,
                            checkpoint_idx=checkpoint_index,
                            metrics=self._extract_scalars_from_info(infos[i]),
                            tb_writer=writer,
                        )

                        rgb_frames[i] = []
                # episode continues
                elif len(self.config.VIDEO_OPTION) > 0:
                    # TODO move normalization / channel changing out of the policy and undo it here
                    frame = observations_to_image(
                        {k: v[i] for k, v in batch.items()}, infos[i]
                    )
                    frame = append_text_to_image(frame, "Instruction: {}".format(next_episodes[i].instruction.instruction_text))
                    rgb_frames[i].append(frame)


            (
                self.envs,
                rnn_hidden_states,
                not_done_masks,
                current_episode_reward,
                prev_actions,
                batch,
                rgb_frames,
            ) = self._pause_envs(
                envs_to_pause,
                self.envs,
                rnn_hidden_states,
                not_done_masks,
                current_episode_reward,
                prev_actions,
                batch,
                rgb_frames,
            )

        num_episodes = len(stats_episodes)
        aggregated_stats = {}
        for stat_key in next(iter(stats_episodes.values())).keys():
            aggregated_stats[stat_key] = (
                sum(v[stat_key] for v in stats_episodes.values())
                / num_episodes
            )

        for k, v in aggregated_stats.items():
            logger.info(f"Average episode {k}: {v:.4f}")

        step_id = checkpoint_index
        if "extra_state" in ckpt_dict and "step" in ckpt_dict["extra_state"]:
            step_id = ckpt_dict["extra_state"]["step"]

        writer.add_scalars(
            "eval_reward",
            {"average reward": aggregated_stats["reward"]},
            step_id,
        )

        metrics = {k: v for k, v in aggregated_stats.items() if k != "reward"}
        if len(metrics) > 0:
            writer.add_scalars("eval_metrics", metrics, step_id)

        self.envs.close()

        print ("environments closed")

habitat_baselines/il/disk_based/il_trainer.py
cshizhe/onav_rim
import math
import sys
from typing import Dict, Iterable, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F

from gym import Space
from habitat import Config, logger
from habitat.tasks.nav.nav import (
    EpisodicCompassSensor,
    EpisodicGPSSensor,
)
from habitat_baselines.il.common.encoders.instruction import InstructionEncoder
from habitat_baselines.il.common.encoders.resnet_encoders import (
    VlnResnetDepthEncoder,
    ResnetRGBEncoder,
)
from habitat_baselines.rl.models.rnn_state_encoder import RNNStateEncoder
from habitat_baselines.common.baseline_registry import baseline_registry
from habitat_baselines.rl.ppo import Net, Policy


class PickPlaceNet(Net):
    r"""A baseline sequence to sequence network that concatenates instruction,
    RGB, and depth encodings before decoding an action distribution with an RNN.
    Modules:
        Instruction encoder
        Depth encoder
        RGB encoder
        RNN state encoder
    """

    def __init__(self, observation_space: Space, model_config: Config, num_actions):
        super().__init__()
        self.model_config = model_config

        # Init the instruction encoder
        self.instruction_encoder = InstructionEncoder(model_config.INSTRUCTION_ENCODER)

        # Init the depth encoder
        assert model_config.DEPTH_ENCODER.cnn_type in [
            "VlnResnetDepthEncoder",
            "Noe"
        ], "DEPTH_ENCODER.cnn_type must be SimpleDepthCNN or VlnResnetDepthEncoder"
        if model_config.DEPTH_ENCODER.cnn_type == "VlnResnetDepthEncoder":
            self.depth_encoder = VlnResnetDepthEncoder(
                observation_space,
                output_size=model_config.DEPTH_ENCODER.output_size,
                checkpoint=model_config.DEPTH_ENCODER.ddppo_checkpoint,
                backbone=model_config.DEPTH_ENCODER.backbone,
                trainable=model_config.DEPTH_ENCODER.trainable,
            )

        # Init the RGB visual encoder
        assert model_config.RGB_ENCODER.cnn_type in [
            "ResnetRGBEncoder",
        ], "RGB_ENCODER.cnn_type must be either 'SimpleRGBCNN' or 'TorchVisionResNet50'."

        if model_config.RGB_ENCODER.cnn_type == "ResnetRGBEncoder":
            self.rgb_encoder = ResnetRGBEncoder(
                observation_space,
                output_size=model_config.RGB_ENCODER.output_size,
                backbone=model_config.RGB_ENCODER.backbone,
                trainable=model_config.RGB_ENCODER.train_encoder,
            )

        # Init the RNN state decoder
        rnn_input_size = (
            self.instruction_encoder.output_size
            + model_config.DEPTH_ENCODER.output_size
            + model_config.RGB_ENCODER.output_size
        )

        if EpisodicGPSSensor.cls_uuid in observation_space.spaces:
            input_gps_dim = observation_space.spaces[
                EpisodicGPSSensor.cls_uuid
            ].shape[0]
            self.gps_embedding = nn.Linear(input_gps_dim, 32)
            rnn_input_size += 32
            logger.info("\n\nSetting up GPS sensor")
        
        if EpisodicCompassSensor.cls_uuid in observation_space.spaces:
            assert (
                observation_space.spaces[EpisodicCompassSensor.cls_uuid].shape[
                    0
                ]
                == 1
            ), "Expected compass with 2D rotation."
            input_compass_dim = 2  # cos and sin of the angle
            self.compass_embedding = nn.Linear(input_compass_dim, 32)
            rnn_input_size += 32
            logger.info("\n\nSetting up Compass sensor")

        if model_config.SEQ2SEQ.use_prev_action:
            self.prev_action_embedding = nn.Embedding(num_actions + 1, 32)
            rnn_input_size += self.prev_action_embedding.embedding_dim

        self.state_encoder = RNNStateEncoder(
            input_size=rnn_input_size,
            hidden_size=model_config.STATE_ENCODER.hidden_size,
            num_layers=model_config.STATE_ENCODER.num_recurrent_layers,
            rnn_type=model_config.STATE_ENCODER.rnn_type,
        )

        self.train()

    @property
    def output_size(self):
        return self.model_config.STATE_ENCODER.hidden_size

    @property
    def is_blind(self):
        return self.rgb_encoder.is_blind or self.depth_encoder.is_blind

    @property
    def num_recurrent_layers(self):
        return self.state_encoder.num_recurrent_layers

    def forward(self, observations, rnn_hidden_states, prev_actions, masks):
        r"""
        instruction_embedding: [batch_size x INSTRUCTION_ENCODER.output_size]
        depth_embedding: [batch_size x DEPTH_ENCODER.output_size]
        rgb_embedding: [batch_size x RGB_ENCODER.output_size]
        """
        instruction_embedding = self.instruction_encoder(observations)
        depth_embedding = self.depth_encoder(observations)
        rgb_embedding = self.rgb_encoder(observations)

        x = [instruction_embedding, depth_embedding, rgb_embedding]

        if EpisodicGPSSensor.cls_uuid in observations:
            x.append(
                self.gps_embedding(observations[EpisodicGPSSensor.cls_uuid])
            )
        
        if EpisodicCompassSensor.cls_uuid in observations:
            compass_observations = torch.stack(
                [
                    torch.cos(observations[EpisodicCompassSensor.cls_uuid]),
                    torch.sin(observations[EpisodicCompassSensor.cls_uuid]),
                ],
                -1,
            )
            x.append(
                self.compass_embedding(compass_observations.squeeze(dim=1))
            )

        if self.model_config.SEQ2SEQ.use_prev_action:
            prev_actions_embedding = self.prev_action_embedding(
                ((prev_actions.float() + 1) * masks).long().view(-1)
            )
            x.append(prev_actions_embedding)

        x = torch.cat(x, dim=1)
        x, rnn_hidden_states = self.state_encoder(x, rnn_hidden_states, masks)

        return x, rnn_hidden_states

@baseline_registry.register_policy
class PickPlacePolicy(Policy):
    def __init__(
        self, observation_space: Space, action_space: Space, model_config: Config
    ):
        super().__init__(
            PickPlaceNet(
                observation_space=observation_space,
                model_config=model_config,
                num_actions=action_space.n,
            ),
            action_space.n,
            no_critic=True
        )

    @classmethod
    def from_config(
        cls, config: Config, observation_space, action_space
    ):
        return cls(
            observation_space=observation_space,
            action_space=action_space,
            model_config=config.MODEL,
        )

habitat_baselines/il/disk_based/policy/resnet_policy.py
cshizhe/onav_rim
#!/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.

from habitat_baselines.il.env_based.policy.resnet_policy import ObjectNavILPolicy

__all__ = ["ObjectNavILPolicy"]

habitat_baselines/il/env_based/__init__.py
cshizhe/onav_rim
#!/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.

from typing import Optional, Tuple

import torch
from torch import Tensor
from torch import nn as nn
from torch import optim as optim

from habitat import logger
from habitat.utils import profiling_wrapper


class ILAgent(nn.Module):
    def __init__(
        self,
        model: nn.Module,
        num_envs: int,
        num_mini_batch: int,
        lr: Optional[float] = None,
        eps: Optional[float] = None,
        max_grad_norm: Optional[float] = None,
    ) -> None:

        super().__init__()

        self.model = model

        self.num_mini_batch = num_mini_batch

        self.max_grad_norm = max_grad_norm
        self.num_envs = num_envs

        self.optimizer = optim.Adam(
            list(filter(lambda p: p.requires_grad, model.parameters())),
            lr=lr,
            eps=eps,
        )
        self.device = next(model.parameters()).device

    def forward(self, *x):
        raise NotImplementedError

    def update(self, rollouts) -> Tuple[float, float, float]:
        total_loss_epoch = 0.0

        profiling_wrapper.range_push("BC.update epoch")
        data_generator = rollouts.recurrent_generator(
            self.num_mini_batch
        )
        cross_entropy_loss = torch.nn.CrossEntropyLoss(reduction="none")
        hidden_states = []

        for sample in data_generator:
            (
                obs_batch,
                recurrent_hidden_states_batch,
                actions_batch,
                prev_actions_batch,
                masks_batch,
                idx
            ) = sample
            # for k, v in obs_batch.items():
            #     print(k, v.size()) # (num_steps, num_mini_batch, xxx)

            # Reshape to do in a single forward pass for all steps
            (
                logits,
                rnn_hidden_states,
            ) = self.model(
                obs_batch,
                recurrent_hidden_states_batch,
                prev_actions_batch,
                masks_batch,
            )

            T, N, _ = actions_batch.shape
            logits = logits.view(T, N, -1)

            action_loss = cross_entropy_loss(logits.permute(0, 2, 1), actions_batch.squeeze(-1))

            self.optimizer.zero_grad()
            inflections_batch = obs_batch["inflection_weight"]

            total_loss = ((inflections_batch * action_loss).sum(0) / inflections_batch.sum(0)).mean()

            self.before_backward(total_loss)
            total_loss.backward()
            self.after_backward(total_loss)

            self.before_step()
            self.optimizer.step()
            self.after_step()

            total_loss_epoch += total_loss.item()
            hidden_states.append(rnn_hidden_states)

        profiling_wrapper.range_pop()

        hidden_states = torch.cat(hidden_states, dim=1)

        total_loss_epoch /= self.num_mini_batch

        return total_loss_epoch, hidden_states

    def before_backward(self, loss: Tensor) -> None:
        pass

    def after_backward(self, loss: Tensor) -> None:
        pass

    def before_step(self) -> None:
        nn.utils.clip_grad_norm_(
            self.model.parameters(), self.max_grad_norm
        )

    def after_step(self) -> None:
        pass


EPS_PPO = 1e-5


class DecentralizedDistributedMixin:

    def init_distributed(self, find_unused_params: bool = True) -> None:
        r"""Initializes distributed training for the model

        1. Broadcasts the model weights from world_rank 0 to all other workers
        2. Adds gradient hooks to the model

        :param find_unused_params: Whether or not to filter out unused parameters
                                   before gradient reduction.  This *must* be True if
                                   there are any parameters in the model that where unused in the
                                   forward pass, otherwise the gradient reduction
                                   will not work correctly.
        """
        # NB: Used to hide the hooks from the nn.Module,
        # so they don't show up in the state_dict
        class Guard:
            def __init__(self, model, device):
                if torch.cuda.is_available():
                    self.ddp = torch.nn.parallel.DistributedDataParallel(
                        model, device_ids=[device], output_device=device
                    )
                else:
                    self.ddp = torch.nn.parallel.DistributedDataParallel(model)

        self._ddp_hooks = Guard(self.model, self.device)  # type: ignore
        # self.get_advantages = self._get_advantages_distributed

        self.reducer = self._ddp_hooks.ddp.reducer
        self.find_unused_params = find_unused_params

    def before_backward(self, loss: Tensor) -> None:
        super().before_backward(loss)  # type: ignore

        if self.find_unused_params:
            self.reducer.prepare_for_backward([loss])  # type: ignore
        else:
            self.reducer.prepare_for_backward([])  # type: ignore


class DDPILAgent(DecentralizedDistributedMixin, ILAgent):
    pass

habitat_baselines/il/env_based/algos/agent.py
cshizhe/onav_rim
#!/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 torch

from collections import defaultdict


class RolloutStorage:
    r"""Class for storing rollout information for RL trainers."""

    def __init__(
        self,
        num_steps,
        num_envs,
        observation_space,
        action_space,
        recurrent_hidden_state_size,
        num_recurrent_layers=1,
    ):
        self.observations = {}

        for sensor in observation_space.spaces:
            self.observations[sensor] = torch.zeros(
                num_steps + 1,
                num_envs,
                *observation_space.spaces[sensor].shape
            )

        self.recurrent_hidden_states = torch.zeros(
            1,
            num_recurrent_layers,
            num_envs,
            recurrent_hidden_state_size,
        )

        self.rewards = torch.zeros(num_steps, num_envs, 1)

        self.action_log_probs = torch.zeros(num_steps, num_envs, 1)
        if action_space.__class__.__name__ == "ActionSpace":
            action_shape = 1
        else:
            action_shape = action_space.shape[0]

        self.actions = torch.zeros(num_steps, num_envs, 1)
        self.prev_actions = torch.zeros(num_steps + 1, num_envs, 1)
        if action_space.__class__.__name__ == "ActionSpace":
            self.actions = self.actions.long()
            self.prev_actions = self.prev_actions.long()

        self.masks = torch.zeros(num_steps + 1, num_envs, 1)
        self.episode_step_index = [1] * (num_steps + 1)

        self.num_steps = num_steps
        self.num_envs = num_envs
        self.step = 0

    def to(self, device):
        for sensor in self.observations:
            self.observations[sensor] = self.observations[sensor].to(device)

        self.recurrent_hidden_states = self.recurrent_hidden_states.to(device)
        self.rewards = self.rewards.to(device)
        self.action_log_probs = self.action_log_probs.to(device)
        self.actions = self.actions.to(device)
        self.prev_actions = self.prev_actions.to(device)
        self.masks = self.masks.to(device)

    def insert(
        self,
        observations,
        actions,
        rewards,
        masks,
    ):
        for sensor in observations:
            self.observations[sensor][self.step + 1].copy_(
                observations[sensor]
            )
        self.actions[self.step].copy_(actions)
        self.prev_actions[self.step + 1].copy_(actions)
        self.rewards[self.step].copy_(rewards)
        self.masks[self.step + 1].copy_(masks)

        self.step = self.step + 1
    
    def update_running_episode_step(self, masks):
        for i in range(self.num_envs):
            self.episode_step_index[i] += 1
            if masks[self.step][i].item() == 0:
                self.episode_step_index[i] = 1

    def after_update(self, rnn_hidden_states):
        for sensor in self.observations:
            self.observations[sensor][0].copy_(
                self.observations[sensor][self.step]
            )

        self.recurrent_hidden_states[0].copy_(
            rnn_hidden_states.detach()
        )
        self.masks[0].copy_(self.masks[self.step])
        self.prev_actions[0].copy_(self.prev_actions[self.step])
        self.step = 0

    def get_next_actions(self):
        next_action_observations = self.observations["demonstration"][self.step]
        actions = next_action_observations.clone()
        return actions

    def recurrent_generator(self, num_mini_batch):
        num_processes = self.rewards.size(1)
        assert num_processes >= num_mini_batch, (
            "Trainer requires the number of processes ({}) "
            "to be greater than or equal to the number of "
            "trainer mini batches ({}).".format(num_processes, num_mini_batch)
        )
        num_envs_per_batch = num_processes // num_mini_batch
        perm = torch.randperm(num_processes)
        for start_ind in range(0, num_processes, num_envs_per_batch):
            observations_batch = defaultdict(list)

            recurrent_hidden_states_batch = []
            actions_batch = []
            prev_actions_batch = []
            masks_batch = []
            index_batch = []

            for offset in range(num_envs_per_batch):
                # ind = perm[start_ind + offset]
                ind = start_ind + offset
                # Ignore OOB index
                if ind >= num_processes:
                    continue

                for sensor in self.observations:
                    observations_batch[sensor].append(
                        self.observations[sensor][: self.step, ind]
                    )

                recurrent_hidden_states_batch.append(
                    self.recurrent_hidden_states[0, :, ind]
                )

                actions_batch.append(self.actions[: self.step, ind])
                prev_actions_batch.append(self.prev_actions[: self.step, ind])
                masks_batch.append(self.masks[: self.step, ind])
                index_batch.append(ind)

            T, N = self.step, num_envs_per_batch

            # These are all tensors of size (T, N, -1)
            for sensor in observations_batch:
                observations_batch[sensor] = torch.stack(
                    observations_batch[sensor], 1
                )

            actions_batch = torch.stack(actions_batch, 1)
            prev_actions_batch = torch.stack(prev_actions_batch, 1)
            masks_batch = torch.stack(masks_batch, 1)

            # States is just a (num_recurrent_layers, N, -1) tensor
            recurrent_hidden_states_batch = torch.stack(
                recurrent_hidden_states_batch, 1
            )

            yield (
                observations_batch,
                recurrent_hidden_states_batch,
                actions_batch,
                prev_actions_batch,
                masks_batch,
                index_batch
            )

    @staticmethod
    def _flatten_helper(t: int, n: int, tensor: torch.Tensor) -> torch.Tensor:
        r"""Given a tensor of size (t, n, ..), flatten it to size (t*n, ...).

        Args:
            t: first dimension of tensor.
            n: second dimension of tensor.
            tensor: target tensor to be flattened.

        Returns:
            flattened tensor of size (t*n, ...)
        """
        return tensor.view(t * n, *tensor.size()[2:])

habitat_baselines/il/env_based/common/rollout_storage.py
cshizhe/onav_rim
#!/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 contextlib
import os
import random
import time
import numpy as np
import torch

from collections import defaultdict, deque
from typing import DefaultDict, Optional

from torch import distributed as distrib
from torch import nn as nn
from torch.optim.lr_scheduler import LambdaLR

from habitat import Config, logger
from habitat.utils import profiling_wrapper
from habitat_baselines.common.baseline_registry import baseline_registry
from habitat_baselines.common.environments import get_env_class
from habitat_baselines.common.obs_transformers import (
    apply_obs_transforms_batch,
    apply_obs_transforms_obs_space,
    get_active_obs_transforms
)
from habitat_baselines.il.env_based.common.rollout_storage import RolloutStorage
from habitat_baselines.common.tensorboard_utils import TensorboardWriter
from habitat_baselines.rl.ddppo.algo.ddp_utils import (
    EXIT,
    REQUEUE,
    add_signal_handlers,
    init_distrib_slurm,
    load_interrupted_state,
    requeue_job,
    save_interrupted_state,
)
from habitat_baselines.utils.common import batch_obs, linear_decay
from habitat_baselines.utils.env_utils import construct_envs
from habitat_baselines.il.env_based.algos.agent import DDPILAgent
from habitat_baselines.il.env_based.il_trainer import ILEnvTrainer
from habitat_baselines.il.env_based.policy.rednet import load_rednet


@baseline_registry.register_trainer(name="ddp-il-trainer")
class ILEnvDDPTrainer(ILEnvTrainer):
    # DD-PPO cuts rollouts short to mitigate the straggler effect
    # This, in theory, can cause some rollouts to be very short.
    # All rollouts contributed equally to the loss/model-update,
    # thus very short rollouts can be problematic.  This threshold
    # limits the how short a short rollout can be as a fraction of the
    # max rollout length
    SHORT_ROLLOUT_THRESHOLD: float = 0.25

    def __init__(self, config: Optional[Config] = None) -> None:
        interrupted_state = load_interrupted_state()
        if interrupted_state is not None:
            config = interrupted_state["config"]

        super().__init__(config)  

    def _setup_actor_critic_agent(self, il_cfg: Config, model_config: Config) -> None:
        r"""Sets up actor critic and agent for PPO.

        Args:
            ppo_cfg: config node with relevant params

        Returns:
            None
        """
        logger.add_filehandler(self.config.LOG_FILE)

        observation_space = self.envs.observation_spaces[0]
        self.obs_transforms = get_active_obs_transforms(self.config)
        observation_space = apply_obs_transforms_obs_space(
            observation_space, self.obs_transforms
        )
        self.obs_space = observation_space

        model_config.defrost()
        model_config.TORCH_GPU_ID = self.config.TORCH_GPU_ID
        model_config.freeze()

        policy = baseline_registry.get_policy(self.config.IL.POLICY.name)
        self.policy = policy.from_config(
            self.config, observation_space, self.envs.action_spaces[0]
        )
        self.policy.to(self.device)

        self.semantic_predictor = None
        if model_config.USE_PRED_SEMANTICS:
            self.semantic_predictor = load_rednet(
                self.device,
                ckpt=model_config.SEMANTIC_ENCODER.rednet_ckpt,
                resize=True, # since we train on half-vision
                num_classes=model_config.SEMANTIC_ENCODER.num_classes
            )
            self.semantic_predictor.eval()

        self.agent = DDPILAgent(
            model=self.policy,
            num_envs=self.envs.num_envs,
            num_mini_batch=il_cfg.num_mini_batch,
            lr=il_cfg.lr,
            eps=il_cfg.eps,
            max_grad_norm=il_cfg.max_grad_norm,
        )

    @profiling_wrapper.RangeContext("train")
    def train(self) -> None:
        r"""Main method for DD-PPO.

        Returns:
            None
        """
        self.local_rank, tcp_store = init_distrib_slurm(
            self.config.IL.distrib_backend
        )
        add_signal_handlers()

        profiling_wrapper.configure(
            capture_start_step=self.config.PROFILING.CAPTURE_START_STEP,
            num_steps_to_capture=self.config.PROFILING.NUM_STEPS_TO_CAPTURE,
        )

        interrupted_state = load_interrupted_state()
        if interrupted_state is not None:
            logger.info("Overriding current config with interrupted state config")
            self.config = interrupted_state["config"]

        # Stores the number of workers that have finished their rollout
        num_rollouts_done_store = distrib.PrefixStore(
            "rollout_tracker", tcp_store
        )
        num_rollouts_done_store.set("num_done", "0")

        self.world_rank = distrib.get_rank()
        self.world_size = distrib.get_world_size()

        self.config.defrost()
        self.config.TORCH_GPU_ID = self.local_rank
        self.config.SIMULATOR_GPU_ID = self.local_rank
        # Multiply by the number of simulators to make sure they also get unique seeds
        self.config.TASK_CONFIG.SEED += (
            self.world_rank * self.config.NUM_PROCESSES
        )
        self.config.freeze()

        random.seed(self.config.TASK_CONFIG.SEED)
        np.random.seed(self.config.TASK_CONFIG.SEED)
        torch.manual_seed(self.config.TASK_CONFIG.SEED)

        if torch.cuda.is_available():
            self.device = torch.device("cuda", self.local_rank)
            torch.cuda.set_device(self.device)
        else:
            self.device = torch.device("cpu")

        self.envs = construct_envs(
            self.config,
            get_env_class(self.config.ENV_NAME),
            workers_ignore_signals=True,
        )

        logger.info(
            "[ train_loader has {} samples ]".format(
                self.envs.count_episodes()
            )
        )

        il_cfg = self.config.IL.BehaviorCloning
        if (
            not os.path.isdir(self.config.CHECKPOINT_FOLDER)
            and self.world_rank == 0
        ):
            os.makedirs(self.config.CHECKPOINT_FOLDER)

        self._setup_actor_critic_agent(il_cfg, self.config.MODEL)
        self.agent.init_distributed(find_unused_params=True)
        self.agent.train()

        if self.world_rank == 0:
            logger.info(
                "agent number of trainable parameters: {}".format(
                    sum(
                        param.numel()
                        for param in self.agent.parameters()
                        if param.requires_grad
                    )
                )
            )

        observations = self.envs.reset()
        batch = batch_obs(observations, device=self.device)
        batch = apply_obs_transforms_batch(batch, self.obs_transforms)

        obs_space = self.obs_space

        rollouts = RolloutStorage(
            il_cfg.num_steps,
            self.envs.num_envs,
            obs_space,
            self.envs.action_spaces[0],
            self.config.MODEL.STATE_ENCODER.hidden_size,
            num_recurrent_layers=self.config.MODEL.STATE_ENCODER.num_recurrent_layers,
        )
        rollouts.to(self.device)

        for sensor in rollouts.observations:
            rollouts.observations[sensor][0].copy_(batch[sensor])
            # Use first semantic observations from RedNet predictor as well
            if sensor == "semantic" and self.config.MODEL.USE_PRED_SEMANTICS:
                semantic_obs = self.semantic_predictor(batch["rgb"], batch["depth"])
                # Subtract 1 from class labels for THDA YCB categories
                if self.config.MODEL.SEMANTIC_ENCODER.is_thda:
                    semantic_obs = semantic_obs - 1
                rollouts.observations[sensor][0].copy_(semantic_obs)

        # batch and observations may contain shared PyTorch CUDA
        # tensors.  We must explicitly clear them here otherwise
        # they will be kept in memory for the entire duration of training!
        batch = None
        observations = None

        current_episode_reward = torch.zeros(
            self.envs.num_envs, 1, device=self.device
        )
        running_episode_stats = dict(
            count=torch.zeros(self.envs.num_envs, 1, device=self.device),
            reward=torch.zeros(self.envs.num_envs, 1, device=self.device),
        )
        window_episode_stats: DefaultDict[str, deque] = defaultdict(
            lambda: deque(maxlen=il_cfg.reward_window_size)
        )

        t_start = time.time()
        env_time = 0
        pth_time = 0
        count_steps: float = 0
        count_checkpoints = 0
        start_update = 0
        prev_time = 0

        lr_scheduler = LambdaLR(
            optimizer=self.agent.optimizer,
            lr_lambda=lambda x: linear_decay(x, self.config.NUM_UPDATES),  # type: ignore
        )
        if interrupted_state is not None:
            self.agent.load_state_dict(interrupted_state["state_dict"])
            self.agent.optimizer.load_state_dict(
                interrupted_state["optim_state"]
            )
            lr_scheduler.load_state_dict(interrupted_state["lr_sched_state"])

            requeue_stats = interrupted_state["requeue_stats"]
            env_time = requeue_stats["env_time"]
            pth_time = requeue_stats["pth_time"]
            count_steps = requeue_stats["count_steps"]
            count_checkpoints = requeue_stats["count_checkpoints"]
            start_update = requeue_stats["start_update"]
            prev_time = requeue_stats["prev_time"]

        with (
            TensorboardWriter(
                self.config.TENSORBOARD_DIR, flush_secs=self.flush_secs
            )
            if self.world_rank == 0
            else contextlib.suppress()
        ) as writer:
            for update in range(start_update, self.config.NUM_UPDATES):
                profiling_wrapper.on_start_step()
                profiling_wrapper.range_push("train update")
                self.current_update = update

                if update > 0 and il_cfg.use_linear_lr_decay:
                    lr_scheduler.step()  # type: ignore

                if update > 0 and  il_cfg.use_linear_clip_decay:
                    self.agent.clip_param = il_cfg.clip_param * linear_decay(
                        update, self.config.NUM_UPDATES
                    )

                if EXIT.is_set():
                    profiling_wrapper.range_pop()  # train update

                    self.envs.close()

                    if REQUEUE.is_set() and self.world_rank == 0:
                        requeue_stats = dict(
                            env_time=env_time,
                            pth_time=pth_time,
                            count_steps=count_steps,
                            count_checkpoints=count_checkpoints,
                            start_update=update,
                            prev_time=(time.time() - t_start) + prev_time,
                        )
                        save_interrupted_state(
                            dict(
                                state_dict=self.agent.state_dict(),
                                optim_state=self.agent.optimizer.state_dict(),
                                lr_sched_state=lr_scheduler.state_dict(),
                                config=self.config,
                                requeue_stats=requeue_stats,
                            )
                        )

                    requeue_job()
                    return

                count_steps_delta = 0
                self.agent.eval()
                profiling_wrapper.range_push("rollouts loop")
                for step in range(il_cfg.num_steps):
                    (
                        delta_pth_time,
                        delta_env_time,
                        delta_steps,
                    ) = self._collect_rollout_step(
                        rollouts, current_episode_reward, running_episode_stats
                    )
                    pth_time += delta_pth_time
                    env_time += delta_env_time
                    count_steps_delta += delta_steps

                    # This is where the preemption of workers happens.  If a
                    # worker detects it will be a straggler, it preempts itself!
                    if (
                        step
                        >= il_cfg.num_steps * self.SHORT_ROLLOUT_THRESHOLD
                    ) and int(num_rollouts_done_store.get("num_done")) > (
                        il_cfg.sync_frac * self.world_size
                    ):
                        break
                profiling_wrapper.range_pop()  # rollouts loop

                num_rollouts_done_store.add("num_done", 1)

                self.agent.train()
                (
                    delta_pth_time,
                    total_loss
                ) = self._update_agent(il_cfg, rollouts)
                pth_time += delta_pth_time

                stats_ordering = sorted(running_episode_stats.keys())
                stats = torch.stack(
                    [running_episode_stats[k] for k in stats_ordering], 0
                )
                distrib.all_reduce(stats)

                for i, k in enumerate(stats_ordering):
                    window_episode_stats[k].append(stats[i].clone())

                stats = torch.tensor(
                    [total_loss, count_steps_delta],
                    device=self.device,
                )
                distrib.all_reduce(stats)
                count_steps += stats[1].item()

                if self.world_rank == 0:
                    num_rollouts_done_store.set("num_done", "0")

                    losses = [
                        stats[0].item() / self.world_size,
                    ]
                    deltas = {
                        k: (
                            (v[-1] - v[0]).sum().item()
                            if len(v) > 1
                            else v[0].sum().item()
                        )
                        for k, v in window_episode_stats.items()
                    }
                    deltas["count"] = max(deltas["count"], 1.0)

                    writer.add_scalar(
                        "reward",
                        deltas["reward"] / deltas["count"],
                        count_steps,
                    )

                    # Check to see if there are any metrics
                    # that haven't been logged yet
                    metrics = {
                        k: v / deltas["count"]
                        for k, v in deltas.items()
                        if k not in {"reward", "count"}
                    }
                    if len(metrics) > 0:
                        writer.add_scalars("metrics", metrics, count_steps)

                    writer.add_scalars(
                        "losses",
                        {k: l for l, k in zip(losses, ["action"])},
                        count_steps,
                    )

                    # log stats
                    if update > 0 and update % self.config.LOG_INTERVAL == 0:
                        logger.info(
                            "update: {}\tfps: {:.3f}\tloss: {:.3f}".format(
                                update,
                                count_steps
                                / ((time.time() - t_start) + prev_time), losses[0]
                            )
                        )

                        logger.info(
                            "update: {}\tenv-time: {:.3f}s\tpth-time: {:.3f}s\t"
                            "frames: {}".format(
                                update, env_time, pth_time, count_steps
                            )
                        )
                        logger.info(
                            "Average window size: {}  {}".format(
                                len(window_episode_stats["count"]),
                                "  ".join(
                                    "{}: {:.3f}".format(k, v / deltas["count"])
                                    for k, v in deltas.items()
                                    if k != "count"
                                ),
                            )
                        )

                    if update == self.config.MODEL.SWITCH_TO_PRED_SEMANTICS_UPDATE - 1:
                        self.save_checkpoint(
                            f"ckpt_gt_best.{count_checkpoints}.pth",
                            dict(step=count_steps),
                        )

                    # checkpoint model
                    if update % self.config.CHECKPOINT_INTERVAL == 0:
                        self.save_checkpoint(
                            f"ckpt.{count_checkpoints}.pth",
                            dict(step=count_steps),
                        )
                        count_checkpoints += 1


                profiling_wrapper.range_pop()  # train update

            self.envs.close()

habitat_baselines/il/env_based/il_ddp_trainer.py
cshizhe/onav_rim
#!/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
import time
import copy
import jsonlines
import collections
from filelock import FileLock

from collections import defaultdict, deque
from typing import Any, DefaultDict, Dict, List, Optional

import numpy as np
import torch
import tqdm
from torch.optim.lr_scheduler import LambdaLR

from habitat import Config, logger
from habitat.utils import profiling_wrapper
from habitat.utils.visualizations.utils import observations_to_image
from habitat.sims.habitat_simulator.actions import HabitatSimActions

from habitat_baselines.common.base_trainer import BaseRLTrainer
from habitat_baselines.common.baseline_registry import baseline_registry
from habitat_baselines.common.environments import get_env_class
from habitat_baselines.common.obs_transformers import (
    apply_obs_transforms_batch,
    apply_obs_transforms_obs_space,
    get_active_obs_transforms,
)
from habitat_baselines.il.env_based.common.rollout_storage import RolloutStorage
from habitat_baselines.il.env_based.algos.agent import ILAgent
from habitat_baselines.common.tensorboard_utils import TensorboardWriter
from habitat_baselines.utils.common import (
    batch_obs,
    generate_video,
    linear_decay,
)
from habitat_baselines.utils.env_utils import make_env_fn, construct_envs
from habitat_baselines.il.env_based.policy.rednet import load_rednet

from habitat.tasks.nav.shortest_path_follower import ShortestPathFollower
from habitat.tasks.utils import cartesian_to_polar
from habitat.utils.geometry_utils import (
    quaternion_from_coeff,
    quaternion_rotate_vector,
    quaternion_to_list,
)

try:
    from habitat_baselines.il.common.encoders.semantic_predictor import SegmentationModel
except:
    pass

def write_to_jsonlines(filepath, item):
    with FileLock(filepath + '.lock'):
        with jsonlines.open(filepath, mode='a') as outf:
            outf.write(item)

def convert_offline_model(ckpt_dict, is_map_cmt=False):
    state_dict = {}
    for k, v in ckpt_dict.items():
        if k.split('.')[0] in ['action_distribution', 'fine_goto_action_fc']:
            if is_map_cmt:
                newk = 'model.' + k
            else:
                newk = k.replace('action_distribution', 'model.action_distribution.linear')
        else:
            if k.startswith('vis_pred_layer') or k.startswith('map_pred_layer') or k.startswith('sem_pred_layer'):
                continue
            newk = 'model.net.' + k
        state_dict[newk] = v
    ckpt_dict = {'state_dict': state_dict}
    return ckpt_dict

def process_batch_gps_compass(batch, gpscompass_noise_type):
    if gpscompass_noise_type == 'zero':
        batch['gps'][:] = 0
        batch['compass'][:] = 0
    elif gpscompass_noise_type == 'gaussian':
        device = batch['gps'].device
        batch['gps'] += torch.randn(batch['gps'].size(), device=device)
        batch['compass'] += torch.randn(batch['compass'].size(), device=device)
    return batch

@baseline_registry.register_trainer(name="il-trainer")
class ILEnvTrainer(BaseRLTrainer):
    r"""Trainer class for behavior cloning.
    """
    supported_tasks = ["ObjectNav-v1"]

    def __init__(self, config=None):
        super().__init__(config)
        self.policy = None
        self.agent = None
        self.envs = None
        self.obs_transforms = []
        if config is not None:
            logger.info(f"config: {config}")

    def _setup_actor_critic_agent(self, il_cfg: Config, model_config: Config, single_env=False) -> None:
        r"""Sets up actor critic and agent for PPO.

        Args:
            ppo_cfg: config node with relevant params

        Returns:
            None
        """
        logger.add_filehandler(self.config.LOG_FILE)

        if not single_env:
            observation_space = self.envs.observation_spaces[0]
            action_space = self.envs.action_spaces[0]
            num_envs = self.envs.num_envs
        else:
            observation_space = self.envs.observation_space
            action_space = self.envs.action_space
            num_envs = 1

        self.obs_transforms = get_active_obs_transforms(self.config)
        observation_space = apply_obs_transforms_obs_space(
            observation_space, self.obs_transforms
        )
        self.obs_space = observation_space

        model_config.defrost()
        model_config.TORCH_GPU_ID = self.config.TORCH_GPU_ID
        model_config.freeze()

        policy = baseline_registry.get_policy(self.config.IL.POLICY.name)
        self.policy = policy.from_config(
            self.config, observation_space, action_space
        )
        self.policy.to(self.device)

        self.semantic_predictor = None
        if model_config.USE_PRED_SEMANTICS:
            self.semantic_predictor = load_rednet(
                self.device,
                ckpt=model_config.SEMANTIC_ENCODER.rednet_ckpt,
                resize=True, # since we train on half-vision
                num_classes=model_config.SEMANTIC_ENCODER.num_classes
            )
            self.semantic_predictor.eval()

        if model_config.USE_PRETRAINED_SEGMENTATION:
            self.segmentation_model = SegmentationModel(self.device)

        self.agent = ILAgent(
            model=self.policy,
            num_envs=num_envs,
            num_mini_batch=il_cfg.num_mini_batch,
            lr=il_cfg.lr,
            eps=il_cfg.eps,
            max_grad_norm=il_cfg.max_grad_norm,
        )

    def _make_results_dir(self, split="val"):
        r"""Makes directory for saving eqa-cnn-pretrain eval results."""
        for s_type in ["rgb", "seg", "depth", "top_down_map"]:
            dir_name = self.config.RESULTS_DIR.format(split=split, type=s_type)
            os.makedirs(dir_name, exist_ok=True)

    @profiling_wrapper.RangeContext("save_checkpoint")
    def save_checkpoint(
        self, file_name: str, extra_state: Optional[Dict] = None
    ) -> None:
        r"""Save checkpoint with specified name.

        Args:
            file_name: file name for checkpoint

        Returns:
            None
        """
        checkpoint = {
            "state_dict": self.agent.state_dict(),
            "config": self.config,
        }
        if extra_state is not None:
            checkpoint["extra_state"] = extra_state

        torch.save(
            checkpoint, os.path.join(self.config.CHECKPOINT_FOLDER, file_name)
        )

    def load_checkpoint(self, checkpoint_path: str, *args, **kwargs) -> Dict:
        r"""Load checkpoint of specified path as a dict.

        Args:
            checkpoint_path: path of target checkpoint
            *args: additional positional args
            **kwargs: additional keyword args

        Returns:
            dict containing checkpoint info
        """
        return torch.load(checkpoint_path, *args, **kwargs)

    METRICS_BLACKLIST = {"top_down_map", "collisions.is_collision", "room_visitation_map", "exploration_metrics"}

    @classmethod
    def _extract_scalars_from_info(
        cls, info: Dict[str, Any]
    ) -> Dict[str, float]:
        result = {}
        for k, v in info.items():
            if k in cls.METRICS_BLACKLIST:
                continue

            if isinstance(v, dict):
                result.update(
                    {
                        k + "." + subk: subv
                        for subk, subv in cls._extract_scalars_from_info(
                            v
                        ).items()
                        if (k + "." + subk) not in cls.METRICS_BLACKLIST
                    }
                )
            # Things that are scalar-like will have an np.size of 1.
            # Strings also have an np.size of 1, so explicitly ban those
            elif np.size(v) == 1 and not isinstance(v, str):
                result[k] = float(v)

        return result

    @classmethod
    def _extract_scalars_from_infos(
        cls, infos: List[Dict[str, Any]]
    ) -> Dict[str, List[float]]:

        results = defaultdict(list)
        for i in range(len(infos)):
            for k, v in cls._extract_scalars_from_info(infos[i]).items():
                results[k].append(v)

        return results

    @profiling_wrapper.RangeContext("_collect_rollout_step")
    def _collect_rollout_step(
        self, rollouts, current_episode_reward, running_episode_stats
    ):
        pth_time = 0.0
        env_time = 0.0

        t_sample_action = time.time()

        # fetch actions and environment state from replay buffer
        next_actions = rollouts.get_next_actions()
        actions = next_actions.long().unsqueeze(-1)
        step_data = [a.item() for a in next_actions.long().to(device="cpu")]

        pth_time += time.time() - t_sample_action

        t_step_env = time.time()
        profiling_wrapper.range_pop()  # compute actions

        outputs = self.envs.step(step_data)
        observations, rewards_l, dones, infos = [
            list(x) for x in zip(*outputs)
        ]

        env_time += time.time() - t_step_env

        t_update_stats = time.time()
        batch = batch_obs(observations, device=self.device)
        if self.config.MODEL.USE_PRED_SEMANTICS and self.current_update >= self.config.MODEL.SWITCH_TO_PRED_SEMANTICS_UPDATE:
            batch["semantic"] = self.semantic_predictor(batch["rgb"], batch["depth"])
            # Subtract 1 from class labels for THDA YCB categories
            if self.config.MODEL.SEMANTIC_ENCODER.is_thda:
                batch["semantic"] = batch["semantic"] - 1
        batch = apply_obs_transforms_batch(batch, self.obs_transforms)

        rewards = torch.tensor(
            rewards_l, dtype=torch.float, device=current_episode_reward.device
        )
        rewards = rewards.unsqueeze(1)

        masks = torch.tensor(
            [[0.0] if done else [1.0] for done in dones],
            dtype=torch.float,
            device=current_episode_reward.device,
        )

        current_episode_reward += rewards
        running_episode_stats["reward"] += (1 - masks) * current_episode_reward  # type: ignore
        running_episode_stats["count"] += 1 - masks  # type: ignore
        for k, v_k in self._extract_scalars_from_infos(infos).items():
            v = torch.tensor(
                v_k, dtype=torch.float, device=current_episode_reward.device
            ).unsqueeze(1)
            if k not in running_episode_stats:
                running_episode_stats[k] = torch.zeros_like(
                    running_episode_stats["count"]
                )

            running_episode_stats[k] += (1 - masks) * v  # type: ignore

        current_episode_reward *= masks

        rollouts.insert(
            batch,
            actions,
            rewards,
            masks,
        )

        pth_time += time.time() - t_update_stats

        return pth_time, env_time, self.envs.num_envs

    @profiling_wrapper.RangeContext("_update_agent")
    def _update_agent(self, ppo_cfg, rollouts):
        t_update_model = time.time()

        total_loss, rnn_hidden_states = self.agent.update(rollouts)

        rollouts.after_update(rnn_hidden_states)

        return (
            time.time() - t_update_model,
            total_loss,
        )

    @profiling_wrapper.RangeContext("train")
    def train(self) -> None:
        r"""Main method for training PPO.

        Returns:
            None
        """

        profiling_wrapper.configure(
            capture_start_step=self.config.PROFILING.CAPTURE_START_STEP,
            num_steps_to_capture=self.config.PROFILING.NUM_STEPS_TO_CAPTURE,
        )

        self.envs = construct_envs(
            self.config, get_env_class(self.config.ENV_NAME)
        )

        il_cfg = self.config.IL.BehaviorCloning
        self.device = (
            torch.device("cuda", self.config.TORCH_GPU_ID)
            if torch.cuda.is_available()
            else torch.device("cpu")
        )
        if not os.path.isdir(self.config.CHECKPOINT_FOLDER):
            os.makedirs(self.config.CHECKPOINT_FOLDER)

        self._setup_actor_critic_agent(il_cfg, self.config.MODEL)
        logger.info(
            "agent number of parameters: {}".format(
                sum(param.numel() for param in self.agent.parameters())
            )
        )

        rollouts = RolloutStorage(
            il_cfg.num_steps,
            self.envs.num_envs,
            self.obs_space,
            self.envs.action_spaces[0],
            self.config.MODEL.STATE_ENCODER.hidden_size,
            self.config.MODEL.STATE_ENCODER.num_recurrent_layers,
        )
        rollouts.to(self.device)

        observations = self.envs.reset()
        batch = batch_obs(observations, device=self.device)
        batch = apply_obs_transforms_batch(batch, self.obs_transforms)

        for sensor in rollouts.observations:
            rollouts.observations[sensor][0].copy_(batch[sensor])
            # Use first semantic observations from RedNet predictor as well
            if sensor == "semantic" and self.config.MODEL.USE_PRED_SEMANTICS:
                semantic_obs = self.semantic_predictor(batch["rgb"], batch["depth"])
                # Subtract 1 from class labels for THDA YCB categories
                if self.config.MODEL.SEMANTIC_ENCODER.is_thda:
                    semantic_obs = semantic_obs - 1
                rollouts.observations[sensor][0].copy_(semantic_obs)

        # batch and observations may contain shared PyTorch CUDA
        # tensors.  We must explicitly clear them here otherwise
        # they will be kept in memory for the entire duration of training!
        batch = None
        observations = None

        current_episode_reward = torch.zeros(self.envs.num_envs, 1)
        running_episode_stats = dict(
            count=torch.zeros(self.envs.num_envs, 1),
            reward=torch.zeros(self.envs.num_envs, 1),
        )
        window_episode_stats: DefaultDict[str, deque] = defaultdict(
            lambda: deque(maxlen=il_cfg.reward_window_size)
        )

        t_start = time.time()
        env_time = 0
        pth_time = 0
        count_steps = 0
        count_checkpoints = 0

        lr_scheduler = LambdaLR(
            optimizer=self.agent.optimizer,
            lr_lambda=lambda x: linear_decay(x, self.config.NUM_UPDATES),  # type: ignore
        )
        self.possible_actions = self.config.TASK_CONFIG.TASK.POSSIBLE_ACTIONS

        with TensorboardWriter(
            self.config.TENSORBOARD_DIR, flush_secs=self.flush_secs
        ) as writer:
            for update in range(self.config.NUM_UPDATES):
                profiling_wrapper.on_start_step()
                profiling_wrapper.range_push("train update")
                
                self.current_update = update

                if il_cfg.use_linear_lr_decay and update > 0:
                    lr_scheduler.step()  # type: ignore

                if il_cfg.use_linear_clip_decay and update > 0:
                    self.agent.clip_param = il_cfg.clip_param * linear_decay(
                        update, self.config.NUM_UPDATES
                    )

                profiling_wrapper.range_push("rollouts loop")
                for _step in range(il_cfg.num_steps):
                    (
                        delta_pth_time,
                        delta_env_time,
                        delta_steps,
                    ) = self._collect_rollout_step(
                        rollouts, current_episode_reward, running_episode_stats
                    )
                    pth_time += delta_pth_time
                    env_time += delta_env_time
                    count_steps += delta_steps
                profiling_wrapper.range_pop()  # rollouts loop

                (
                    delta_pth_time,
                    total_loss
                ) = self._update_agent(il_cfg, rollouts)
                pth_time += delta_pth_time

                for k, v in running_episode_stats.items():
                    window_episode_stats[k].append(v.clone())

                deltas = {
                    k: (
                        (v[-1] - v[0]).sum().item()
                        if len(v) > 1
                        else v[0].sum().item()
                    )
                    for k, v in window_episode_stats.items()
                }
                deltas["count"] = max(deltas["count"], 1.0) # number of episodes

                writer.add_scalar(
                    "reward", deltas["reward"] / deltas["count"], count_steps
                )

                # Check to see if there are any metrics
                # that haven't been logged yet
                metrics = {
                    k: v / deltas["count"]
                    for k, v in deltas.items()
                    if k not in {"reward", "count"}
                }
                if len(metrics) > 0:
                    writer.add_scalars("metrics", metrics, count_steps)

                losses = [total_loss]
                writer.add_scalars(
                    "losses",
                    {k: l for l, k in zip(losses, ["action"])},
                    count_steps,
                )

                # log stats
                if update % self.config.LOG_INTERVAL == 0:
                    logger.info(
                        "update: {}\tfps: {:.3f}\tloss: {:.3f}".format(
                            update, count_steps / (time.time() - t_start), total_loss
                        )
                    )

                    logger.info(
                        "update: {}\tenv-time: {:.3f}s\tpth-time: {:.3f}s\t"
                        "frames: {}".format(
                            update, env_time, pth_time, count_steps
                        )
                    )

                    logger.info(
                        "Average window size: {}  {}".format(
                            len(window_episode_stats["count"]),
                            "  ".join(
                                "{}: {:.3f}".format(k, v / deltas["count"])
                                for k, v in deltas.items()
                                if k != "count"
                            ),
                        )
                    )

   
                if update == self.config.MODEL.SWITCH_TO_PRED_SEMANTICS_UPDATE - 1:
                    self.save_checkpoint(
                        f"ckpt_gt_best.{count_checkpoints}.pth",
                        dict(step=count_steps),
                    )

                # checkpoint model
                if update % self.config.CHECKPOINT_INTERVAL == 0:
                    self.save_checkpoint(
                        f"ckpt.{count_checkpoints}.pth", dict(step=count_steps)
                    )
                    count_checkpoints += 1

                profiling_wrapper.range_pop()  # train update

            self.envs.close()

    def _eval_checkpoint(
        self, checkpoint_path: str, writer: TensorboardWriter, checkpoint_index: int = 0
    ) -> None:
        r"""Evaluates a single checkpoint.

        Args:
            checkpoint_path: path of checkpoint
            writer: tensorboard writer object for logging to tensorboard
            checkpoint_index: index of cur checkpoint for logging

        Returns:
            None
        """
        torch.set_grad_enabled(False)

        # Map location CPU is almost always better than mapping to a CUDA device.
        ckpt_dict = self.load_checkpoint(checkpoint_path, map_location="cpu")
        if self.config.EVAL_CKPT_FROM_OFFLINEBC:
            ckpt_dict = convert_offline_model(ckpt_dict)

        self._make_results_dir(self.config.EVAL.SPLIT)

        if self.config.EVAL.USE_CKPT_CONFIG:
            config = self._setup_eval_config(ckpt_dict["config"])
        else:
            config = self.config.clone()

        il_cfg = config.IL.BehaviorCloning

        config.defrost()
        config.NUM_PROCESSES = 1 # one episode at a time
        config.TASK_CONFIG.DATASET.SPLIT = config.EVAL.SPLIT
        config.TASK_CONFIG.DATASET.TYPE = "ObjectNav-v1"
        config.TASK_CONFIG.ENVIRONMENT.MAX_EPISODE_STEPS = 500
        config.TASK_CONFIG.TASK.MEASUREMENTS.append("COLLISIONS")   # Just for self-test
        config.freeze()

        use_subgoal = getattr(config.MODEL, 'use_subgoal', False)
        episode_iterative = getattr(config, 'EPISODE_ITERATIVE', False)
        single_env = use_subgoal or episode_iterative

        if config.TORCH_GPU_ID == 0:
            logger.info(f"env config: {config}")
        if use_subgoal or episode_iterative:
            self.envs = make_env_fn(
                config, get_env_class(config.ENV_NAME)
            )
        else:
            self.envs = construct_envs(
                config, get_env_class(config.ENV_NAME), 
                auto_reset_done=False if config.MODEL.model_class != 'ObjectNavRNN' else True
            )
        
        self._setup_actor_critic_agent(il_cfg, config.MODEL, single_env=single_env)

        self.agent.load_state_dict(ckpt_dict["state_dict"], strict=True)    
        self.policy = self.agent.model
        self.policy.eval()

        number_of_eval_episodes = config.TEST_EPISODE_COUNT
        if number_of_eval_episodes == -1:
            if single_env:
                number_of_eval_episodes = self.envs.number_of_episodes
            else:
                number_of_eval_episodes = sum(self.envs.number_of_episodes)
        else:
            if single_env:
                total_num_eps = self.envs.number_of_episodes
            else:
                total_num_eps = sum(self.envs.number_of_episodes)
            if total_num_eps < number_of_eval_episodes:
                logger.warn(
                    f"Config specified {number_of_eval_episodes} eval episodes"
                    ", dataset only has {total_num_eps}."
                )
                logger.warn(f"Evaluating with {total_num_eps} instead.")
                number_of_eval_episodes = total_num_eps

        pred_outfile = os.path.join(config.EVAL_RESULTS_DIR, '%s_pred_trajectories.jsonl'%(config.EVAL.SPLIT))
        if len(config.VIDEO_OPTION) > 0:
            os.makedirs(config.VIDEO_DIR, exist_ok=True)

        
        model_class = config.MODEL.model_class
        if model_class == 'ObjectNavTransformer':
            if episode_iterative:
                stats_episodes = self._eval_checkpoint_transformer_iterative_setting(
                    config, number_of_eval_episodes, pred_outfile, writer, checkpoint_index
                )
            else:
                stats_episodes = self._eval_checkpoint_transformer(
                    config, number_of_eval_episodes, pred_outfile, writer, checkpoint_index
                )
        elif model_class == 'ObjectNavImapSingleTransformer':
            stats_episodes = self._eval_checkpoint_recursive_transformer(
                 config, number_of_eval_episodes, pred_outfile, writer, checkpoint_index
            )
        elif model_class == 'ObjectNavRNN':
            stats_episodes = self._eval_checkpoint_rnn(
                config, number_of_eval_episodes, pred_outfile, writer, checkpoint_index
            )
        else:
            raise NotImplementedError('unsupported model class: %s' % str(model_class))


        num_episodes = len(stats_episodes)
        aggregated_stats = {}
        for stat_key in next(iter(stats_episodes.values())).keys():
            aggregated_stats[stat_key] = (
                sum(v[stat_key] for v in stats_episodes.values())
                / num_episodes
            )

        for k, v in aggregated_stats.items():
            logger.info(f"Average episode {k}: {v:.4f}")

        step_id = checkpoint_index
        if "extra_state" in ckpt_dict and "step" in ckpt_dict["extra_state"]:
            step_id = ckpt_dict["extra_state"]["step"]

        writer.add_scalars(
            "eval_reward",
            {"average reward": aggregated_stats["reward"]},
            step_id,
        )

        metrics = {k: v for k, v in aggregated_stats.items() if k not in ["reward", "pred_reward"]}
        if len(metrics) > 0:
            writer.add_scalars("eval_metrics", metrics, step_id)

        self.envs.close()

    def _eval_checkpoint_rnn(
        self, config, number_of_eval_episodes, pred_outfile, writer, checkpoint_index
    ):
        stats_episodes: Dict[
            Any, Any
        ] = {}  # dict of dicts that stores stats per episode

        current_episode_reward = torch.zeros(self.envs.num_envs, 1, device=self.device)
        test_recurrent_hidden_states = torch.zeros(
            config.MODEL.STATE_ENCODER.num_recurrent_layers,
            config.NUM_PROCESSES,
            config.MODEL.hidden_size,
            device=self.device,
        )
        prev_actions = torch.zeros(
            config.NUM_PROCESSES, 1, device=self.device, dtype=torch.long
        ) - 1
        not_done_masks = torch.zeros(
            config.NUM_PROCESSES, 1, device=self.device
        )
        current_episode_steps = torch.zeros(
            self.envs.num_envs, 1, device=self.device
        )

        pred_trajectories = [
            {'actions': [], 'infos': []} for _ in range(config.NUM_PROCESSES)
        ]
        rgb_frames = [
            [] for _ in range(config.NUM_PROCESSES)
        ]  # type: List[List[np.ndarray]]

        observations = self.envs.reset()
        batch = batch_obs(observations, device=self.device)
        batch = apply_obs_transforms_batch(batch, self.obs_transforms)

        pbar = tqdm.tqdm(total=number_of_eval_episodes)
        action_masks = torch.zeros(self.envs.num_envs, 6).bool().to(self.device)
        last_compass = [obs['compass'] for obs in observations]
        last_gps = [obs['gps'] for obs in observations]
        is_collides = [False for _ in observations]

        if len(config.VIDEO_OPTION) > 0:
            os.makedirs(config.VIDEO_DIR, exist_ok=True)
            for i in range(config.NUM_PROCESSES):
                frame = observations_to_image(
                    {"rgb": batch["rgb"][i]}, {}
                )
                rgb_frames[i].append(frame)

        while (
            len(stats_episodes) < number_of_eval_episodes
            and self.envs.num_envs > 0
        ):
            current_episodes = self.envs.current_episodes()

            with torch.no_grad():
                if self.semantic_predictor is not None:
                    batch["semantic"] = self.semantic_predictor(batch["rgb"], batch["depth"])
                    if self.config.MODEL.SEMANTIC_ENCODER.is_thda:
                        batch["semantic"] = batch["semantic"] - 1
                if (not is_collides[0]) or (config.MODEL.enc_collide_steps):
                    (
                        logits,
                        test_recurrent_hidden_states,
                    ) = self.policy(
                        batch,
                        test_recurrent_hidden_states,
                        prev_actions,
                        not_done_masks,
                    )
                current_episode_steps += 1

                logits = logits.masked_fill(action_masks[:logits.size(0)].to(logits.device), -float('inf'))
                actions = torch.argmax(logits, dim=1)
                prev_actions.copy_(actions.unsqueeze(1))  # type: ignore

            # NB: Move actions to CPU.  If CUDA tensors are
            # sent in to env.step(), that will create CUDA contexts
            # in the subprocesses.
            # For backwards compatibility, we also call .item() to convert to
            # an int
            step_data = [a.item() for a in actions.to(device="cpu")]

            outputs = self.envs.step(step_data)

            observations, rewards_l, dones, infos = [
                list(x) for x in zip(*outputs)
            ]
            #print('action', step_data)
            #print('info', infos)
            for ib, aidx in enumerate(step_data):
                pred_trajectories[ib]['actions'].append(aidx)
                pred_trajectories[ib]['infos'].append(infos[ib])

                # heuristic rule: collide
                if all(np.isclose(last_compass[ib], observations[ib]['compass'])) and \
                   all(np.isclose(last_gps[ib], observations[ib]['gps'])) and aidx in [1, 2, 3]:
                    is_collides[ib] = True
                    action_masks[ib, aidx] = True
                else:
                    is_collides[ib] = False
                    action_masks[ib, :] = False
            last_compass = [obs['compass'] for obs in observations]
            last_gps = [obs['gps'] for obs in observations]
                
            batch = batch_obs(observations, device=self.device)
            batch = apply_obs_transforms_batch(batch, self.obs_transforms)

            not_done_masks = torch.tensor(
                [[0.0] if done else [1.0] for done in dones],
                dtype=torch.float,
                device=self.device,
            )
            prev_actions.masked_fill_(not_done_masks.bool().logical_not(), -1) # the init action is -1

            rewards = torch.tensor(
                rewards_l, dtype=torch.float, device=self.device
            ).unsqueeze(1)
            current_episode_reward += rewards
            next_episodes = self.envs.current_episodes()
            envs_to_pause = []
            n_envs = self.envs.num_envs
            for i in range(n_envs):
                if (
                    next_episodes[i].scene_id,
                    next_episodes[i].episode_id,
                ) in stats_episodes:
                    envs_to_pause.append(i)

                # episode ended
                if not_done_masks[i].item() == 0:
                    pbar.update()
                    episode_stats = {}
                    episode_stats["reward"] = current_episode_reward[i].item()
                    episode_stats.update(
                        self._extract_scalars_from_info(infos[i])
                    )
                    current_episode_reward[i] = 0
                    current_episode_steps[i] = 0

                    # use scene_id + episode_id as unique id for storing stats
                    stats_episodes[
                        (
                            current_episodes[i].scene_id,
                            current_episodes[i].episode_id,
                        )
                    ] = episode_stats

                    pred_trajectories[i].update({'episode_id': current_episodes[i].episode_id})
                    write_to_jsonlines(pred_outfile, pred_trajectories[i])
                    pred_trajectories[i] = {'actions': [], 'infos': []}

                    if len(self.config.VIDEO_OPTION) > 0:
                        generate_video(
                            video_option=self.config.VIDEO_OPTION,
                            video_dir=self.config.VIDEO_DIR,
                            images=rgb_frames[i],
                            episode_id=current_episodes[i].episode_id,
                            checkpoint_idx=checkpoint_index,
                            metrics=self._extract_scalars_from_info(infos[i]),
                            tb_writer=writer,
                        )

                        rgb_frames[i] = []
                        frame = observations_to_image(
                            {"rgb": batch["rgb"][i]}, infos[i]
                        )
                        rgb_frames[i].append(frame)

                # episode continues
                elif len(self.config.VIDEO_OPTION) > 0:
                    # TODO move normalization / channel changing out of the policy and undo it here
                    frame = observations_to_image(
                        {"rgb": batch["rgb"][i]}, infos[i]
                    )
                    rgb_frames[i].append(frame)

            (
                self.envs,
                test_recurrent_hidden_states,
                not_done_masks,
                current_episode_reward,
                prev_actions,
                batch,
                rgb_frames,
            ) = self._pause_envs(
                envs_to_pause,
                self.envs,
                test_recurrent_hidden_states,
                not_done_masks,
                current_episode_reward,
                prev_actions,
                batch,
                rgb_frames,
            )

        return stats_episodes
     
    def _eval_checkpoint_transformer(
        self, config, number_of_eval_episodes, 
        pred_outfile, writer, checkpoint_index
    ): 
        stats_episodes: Dict[
            Any, Any
        ] = {}  # dict of dicts that stores stats per episode

        gpscompass_noise_type = getattr(self.config.MODEL, 'gpscompass_noise_type', None)
        print(gpscompass_noise_type)
        for _ in tqdm.trange(number_of_eval_episodes):
            pred_trajectories = {
                'actions': [], 'infos': []
            }
            rgb_frames = []
            
            observations = self.envs.reset()
            current_episodes = self.envs.current_episodes()

            batch = batch_obs(observations, device=self.device)
            batch = process_batch_gps_compass(batch, gpscompass_noise_type)
            batch = apply_obs_transforms_batch(batch, self.obs_transforms)
            if self.semantic_predictor is not None:
                batch["semantic"] = self.semantic_predictor(batch["rgb"], batch["depth"])
                if config.MODEL.SEMANTIC_ENCODER.is_thda:
                    batch["semantic"] = batch["semantic"] - 1

            prev_actions = torch.zeros(
                self.envs.num_envs, 1, device=self.device, dtype=torch.long
            ) - 1
            # initialize at the begning of each episode
            history_embeds = None
            last_gps, last_compass = observations[0]['gps'], observations[0]['compass']
            action_masks = torch.zeros(self.envs.num_envs, 6).bool().to(self.device)
            current_episode_reward = 0

            done = False
            nav_step = 0
            is_collide = False
            while not done:
                if len(config.VIDEO_OPTION) > 0:
                    frame = observations_to_image(
                        {"rgb": batch["rgb"][0]}, {}
                    )
                    rgb_frames.append(frame)

                if config.MODEL.enc_collide_steps or (not is_collide):
                    logits, history_embeds = self.policy(batch, history_embeds, prev_actions, nav_step)
                # print('step', nav_step, history_embeds.size())
                logits = logits.masked_fill(action_masks, -float('inf'))
                # print(nav_step, history_embeds.size(), action_masks, is_collide)

                actions = torch.argmax(logits, dim=1)
                step_data = [a.item() for a in actions.to(device="cpu")]

                outputs = self.envs.step(step_data)
                observations, rewards_l, dones, infos = [
                    list(x) for x in zip(*outputs)
                ]
                done = dones[0]
                current_episode_reward += rewards_l[0]
                
                # if (all(observations[0]['gps'] == last_gps)) and \
                #    (all(observations[0]['compass'] == last_compass)) and \
                #    (step_data[0] in [1, 2, 3]):
                if (all(np.isclose(observations[0]['gps'], last_gps))) and \
                    (all(np.isclose(observations[0]['compass'], last_compass))) and \
                    (step_data[0] in [1, 2, 3]):
                # if infos[0]['collisions']['is_collision']:
                    is_collide = True
                else:
                    is_collide = False
                # print(nav_step, is_collide, infos[0]['collisions'])
                # print(last_gps, last_compass)
                # print(observations[0]['gps'], observations[0]['compass'])
                last_gps = observations[0]['gps']
                last_compass = observations[0]['compass']

                # action_masks = torch.zeros(len(step_data), 6).bool()
                if is_collide:
                    action_masks[0, step_data[0]] = True
                else:
                    action_masks[0, :] = False
                pred_trajectories['actions'].append(step_data[0])
                pred_trajectories['infos'].append(infos[0])

                if config.MODEL.enc_collide_steps or (not is_collide):
                    batch = batch_obs(observations, device=self.device)
                    batch = process_batch_gps_compass(batch, gpscompass_noise_type)
                    batch = apply_obs_transforms_batch(batch, self.obs_transforms)
                    if self.semantic_predictor is not None:
                        batch["semantic"] = self.semantic_predictor(batch["rgb"], batch["depth"])
                        if self.config.MODEL.SEMANTIC_ENCODER.is_thda:
                            batch["semantic"] = batch["semantic"] - 1

                    nav_step += 1
                    prev_actions.copy_(actions.unsqueeze(1))  # type: ignore
                
            # print(current_episodes[0].object_category)
            # print(pred_trajectories['actions'])
            # print(pred_trajectories['infos'][0])
            # print(infos[0])
            # print()
            episode_stats = {}
            episode_stats["reward"] = current_episode_reward
            episode_stats.update(
                self._extract_scalars_from_info(infos[0])
            )
            stats_episodes[
                (
                    current_episodes[0].scene_id,
                    current_episodes[0].episode_id,
                )
            ] = episode_stats

            pred_trajectories.update({'episode_id': current_episodes[0].episode_id})
            write_to_jsonlines(pred_outfile, pred_trajectories)
            if len(self.config.VIDEO_OPTION) > 0:
                generate_video(
                    video_option=self.config.VIDEO_OPTION,
                    video_dir=self.config.VIDEO_DIR,
                    images=rgb_frames,
                    episode_id=current_episodes[0].episode_id,
                    checkpoint_idx=checkpoint_index,
                    metrics=self._extract_scalars_from_info(infos[0]),
                    tb_writer=writer,
                )

        return stats_episodes

    def _eval_checkpoint_recursive_transformer(
        self, config, number_of_eval_episodes, pred_outfile, writer, checkpoint_index
    ):
        stats_episodes: Dict[
            Any, Any
        ] = {}  # dict of dicts that stores stats per episode

        model_config = self.config.MODEL
        gpscompass_noise_type = getattr(self.config.MODEL, 'gpscompass_noise_type', None)
        for _ in tqdm.trange(number_of_eval_episodes):
            pred_trajectories = {
                'actions': [], 'infos': []
            }
            rgb_frames = []
            
            observations = self.envs.reset()
            current_episodes = self.envs.current_episodes()

            batch = batch_obs(observations, device=self.device)
            batch = process_batch_gps_compass(batch, gpscompass_noise_type)
            batch = apply_obs_transforms_batch(batch, self.obs_transforms)
            if self.semantic_predictor is not None:
                batch["semantic"] = self.semantic_predictor(batch["rgb"], batch["depth"])
                if self.config.MODEL.SEMANTIC_ENCODER.is_thda:
                    batch["semantic"] = batch["semantic"] - 1

            if model_config.USE_PRETRAINED_SEGMENTATION:
                batch['obj_areas'], batch['obj_dists'] = [], []
                for ib in range(len(batch['rgb'])):
                    obj_area, obj_dist = self.segmentation_model.predict_target(
                        batch['rgb'][ib], batch['depth'][ib], 
                        batch['objectgoal'][ib]
                    )
                    batch['obj_areas'].append(obj_area)
                    batch['obj_dists'].append(obj_dist)

            recursive_states = None
            if config.INIT_MAP_EMBED_DIR is not None:
                init_map_embeds = np.load(os.path.join(config.INIT_MAP_EMBED_DIR, '%s.npy'%current_episodes[0].episode_id))
                recursive_states = (torch.from_numpy(init_map_embeds).to(self.device), None)
            prev_actions = torch.zeros(
                self.envs.num_envs, 1, device=self.device, dtype=torch.long
            ) - 1
            # initialize at the begning of each episode
            last_gps, last_compass = observations[0]['gps'], observations[0]['compass']
            action_masks = torch.zeros(self.envs.num_envs, 6).bool().to(self.device)
            current_episode_reward = 0

            done = False
            is_collide = False
            nav_step = 0
            while not done:
                if len(self.config.VIDEO_OPTION) > 0:
                    frame = observations_to_image(
                        {"rgb": batch["rgb"][0]}, {}
                    )
                    rgb_frames.append(frame)

                if config.MODEL.enc_collide_steps or (not is_collide):
                    logits, recursive_states = self.policy(batch, recursive_states, prev_actions, nav_step)
                logits = logits.masked_fill(action_masks, -float('inf'))

                actions = torch.argmax(logits, dim=1)
                step_data = [a.item() for a in actions.to(device="cpu")]

                if model_config.USE_PRETRAINED_SEGMENTATION:
                    step_data = [0 if batch['obj_areas'][ib] > 0.1 and batch['obj_dists'][ib] < 0.9 \
                                 else a for ib, a in enumerate(step_data)]

                outputs = self.envs.step(step_data)
                observations, rewards_l, dones, infos = [
                    list(x) for x in zip(*outputs)
                ]
                done = dones[0]
                current_episode_reward += rewards_l[0]

                if model_config.USE_PRETRAINED_SEGMENTATION:
                    infos[0]['obj_area'] = batch['obj_areas'][0]
                    infos[0]['obj_dist'] = batch['obj_dists'][0]
                
                if (all(np.isclose(observations[0]['gps'],  last_gps))) and \
                   (all(np.isclose(observations[0]['compass'], last_compass))) and \
                   (step_data[0] in [1, 2, 3]):
                    is_collide = True
                else:
                    is_collide = False
                # print(nav_step, is_collide, infos[0]['collisions'])
                # print(last_gps, last_compass)
                # print(observations[0]['gps'], observations[0]['compass'])
                last_gps = observations[0]['gps']
                last_compass = observations[0]['compass']

                if is_collide:
                    action_masks[0, step_data[0]] = True
                else:
                    action_masks[0, :] = False
                pred_trajectories['actions'].append(step_data[0])
                pred_trajectories['infos'].append(infos[0])

                if config.MODEL.enc_collide_steps or (not is_collide):
                    batch = batch_obs(observations, device=self.device)
                    batch = process_batch_gps_compass(batch, gpscompass_noise_type)
                    batch = apply_obs_transforms_batch(batch, self.obs_transforms)
                    if self.semantic_predictor is not None:
                        batch["semantic"] = self.semantic_predictor(batch["rgb"], batch["depth"])
                        if self.config.MODEL.SEMANTIC_ENCODER.is_thda:
                            batch["semantic"] = batch["semantic"] - 1
                    if model_config.USE_PRETRAINED_SEGMENTATION:
                        batch['obj_areas'], batch['obj_dists'] = [], []
                        for ib in range(len(batch['rgb'])):
                            obj_area, obj_dist = self.segmentation_model.predict_target(
                                batch['rgb'][ib], batch['depth'][ib], batch['objectgoal'][ib],
                                # outfile=os.path.join('../notebooks/segmasks', '%04d.png'%(nav_step+1))
                            )
                            # import cv2
                            # cv2.imwrite('../notebooks/segmasks/rgb_%04d.png'%(nav_step+1), batch['rgb'][ib].data.cpu().numpy())
                            # print(nav_step, batch['objectgoal'][ib], step_data[ib], obj_area, obj_dist)
                            batch['obj_areas'].append(obj_area)
                            batch['obj_dists'].append(obj_dist)

                    nav_step += 1
                    prev_actions.copy_(actions.unsqueeze(1))  # type: ignore

            # print(current_episodes[0].object_category)
            # print(pred_trajectories['actions'])
            # print(pred_trajectories['infos'][0])
            # print(infos[0])
            # print()
            if config.SAVE_RECURSIVE_STATE:
                state_outdir = os.path.join(config.EVAL_RESULTS_DIR, 'recursive_states')
                os.makedirs(state_outdir, exist_ok=True)
                np.save(os.path.join(state_outdir, '%s.npy'%(current_episodes[0].episode_id)), recursive_states[0].data.cpu().numpy())

            episode_stats = {}
            episode_stats["reward"] = current_episode_reward
            episode_stats.update(
                self._extract_scalars_from_info(infos[0])
            )
            stats_episodes[
                (
                    current_episodes[0].scene_id,
                    current_episodes[0].episode_id,
                )
            ] = episode_stats

            pred_trajectories.update({'episode_id': current_episodes[0].episode_id})
            write_to_jsonlines(pred_outfile, pred_trajectories)
            if len(self.config.VIDEO_OPTION) > 0:
                generate_video(
                    video_option=self.config.VIDEO_OPTION,
                    video_dir=self.config.VIDEO_DIR,
                    images=rgb_frames,
                    episode_id=current_episodes[0].episode_id,
                    checkpoint_idx=checkpoint_index,
                    metrics=self._extract_scalars_from_info(infos[0]),
                    tb_writer=writer,
                )

        return stats_episodes


def get_episodic_gps(agent_state, start_position, start_rotation):
    origin = np.array(start_position, dtype=np.float32)
    rotation_world_start = quaternion_from_coeff(start_rotation)

    agent_position = agent_state.position
    agent_position = quaternion_rotate_vector(
        rotation_world_start.inverse(), agent_position - origin
    )

    return np.array(
        [-agent_position[2], agent_position[0]], dtype=np.float32
    )

def get_episodic_compass(agent_state, start_rotation):
    rotation_world_agent = agent_state.rotation
    rotation_world_start = quaternion_from_coeff(start_rotation)

    quat = rotation_world_agent.inverse() * rotation_world_start
    direction_vector = np.array([0, 0, -1])
    heading_vector = quaternion_rotate_vector(quat, direction_vector)
    phi = cartesian_to_polar(-heading_vector[2], heading_vector[0])[1]
    return np.array([phi], dtype=np.float32)

habitat_baselines/il/env_based/il_trainer.py
cshizhe/onav_rim
import numpy as np
import torch
import torch.nn as nn

from gym import Space
from gym import spaces
from habitat import Config, logger

from habitat.tasks.nav.nav import (
    EpisodicCompassSensor,
    EpisodicGPSSensor,
)
from habitat.tasks.nav.object_nav_task import ObjectGoalSensor
from habitat.tasks.nav.object_nav_task import (
    task_cat2mpcat40,
    mapping_mpcat40_to_goal21
)
from habitat_baselines.il.common.encoders.resnet_encoders import (
    VlnResnetDepthEncoder,
    ResnetRGBEncoder,
    ResnetSemSeqEncoder,
    ResNetImageNetEncoder,
    ResNetCLIPEncoder,
)
from habitat_baselines.rl.ppo import Net


class ObjectNavBase(Net):
    r"""A baseline sequence to sequence network that concatenates instruction,
    RGB, and depth encodings before decoding an action distribution with an RNN.
    Modules:
        Instruction encoder
        Depth encoder
        RGB encoder
    """

    def __init__(self, observation_space: Space, model_config: Config, num_actions, device=None):
        super().__init__()
        self.model_config = model_config
        self.device = device
        
        if isinstance(model_config.num_ft_views, str):
            self.model_config.defrost()
            self.model_config.num_ft_views = int(self.model_config.num_ft_views.split('_')[0]) # 2_up
            self.model_config.freeze()

        self.step_input_sizes = {}
        self.build_model(observation_space, model_config, num_actions)

        self.train()

    def build_visual_encoders(self, observation_space, model_config):
        # Init the depth encoder
        assert model_config.DEPTH_ENCODER.cnn_type in [
            "VlnResnetDepthEncoder",
            "None", None
        ], "DEPTH_ENCODER.cnn_type %s must be VlnResnetDepthEncoder" % (model_config.DEPTH_ENCODER.cnn_type)
        if model_config.DEPTH_ENCODER.cnn_type == "VlnResnetDepthEncoder":
            self.depth_encoder = VlnResnetDepthEncoder(
                observation_space,
                output_size=model_config.DEPTH_ENCODER.output_size,
                checkpoint=model_config.DEPTH_ENCODER.ddppo_checkpoint,
                backbone=model_config.DEPTH_ENCODER.backbone,
                trainable=model_config.DEPTH_ENCODER.trainable,
            )
            self.step_input_sizes['depth'] = model_config.DEPTH_ENCODER.output_size * model_config.num_ft_views
        else:
            self.depth_encoder = None

        # Init the RGB visual encoder
        assert model_config.RGB_ENCODER.cnn_type in [
            "ResnetRGBEncoder",
            "resnet50_imagenet",
            "resnet18_imagenet",
            "resnet50_clip",
            "None",
        ], "RGB_ENCODER.cnn_type must be 'ResnetRGBEncoder'."

        if model_config.RGB_ENCODER.cnn_type == "ResnetRGBEncoder":
            self.rgb_encoder = ResnetRGBEncoder(
                observation_space,
                output_size=model_config.RGB_ENCODER.output_size,
                backbone=model_config.RGB_ENCODER.backbone,
                trainable=model_config.RGB_ENCODER.train_encoder,
                normalize_visual_inputs=model_config.normalize_visual_inputs,
            )
            self.step_input_sizes['rgb'] = model_config.RGB_ENCODER.output_size * model_config.num_ft_views
        elif model_config.RGB_ENCODER.cnn_type in ["resnet50_imagenet", 'resnet18_imagenet']:
            self.rgb_encoder = ResNetImageNetEncoder(
                observation_space,
                output_size=model_config.RGB_ENCODER.output_size,
                backbone=model_config.RGB_ENCODER.cnn_type.split('_')[0],
            )
            self.step_input_sizes['rgb'] = model_config.RGB_ENCODER.output_size * model_config.num_ft_views
        elif model_config.RGB_ENCODER.cnn_type.startswith("resnet50_clip"):
            self.rgb_encoder = ResNetCLIPEncoder(
                observation_space,
                output_size=model_config.RGB_ENCODER.output_size,
                pooling='avgpool',
                device=self.device
            )
            self.step_input_sizes['rgb'] = model_config.RGB_ENCODER.output_size * model_config.num_ft_views
        else:
            self.rgb_encoder = None
            logger.info("RGB encoder is none")

        sem_seg_output_size = 0
        self.semantic_predictor = None
        self.is_thda = False
        if model_config.USE_SEMANTICS:
            sem_embedding_size = model_config.SEMANTIC_ENCODER.embedding_size

            self.is_thda = model_config.SEMANTIC_ENCODER.is_thda
            sem_spaces = {
                "semantic": spaces.Box(
                    low=0,
                    high=255,
                    shape=(480, 640, sem_embedding_size),
                    dtype=np.uint8,
                ),
            }
            sem_obs_space = spaces.Dict(sem_spaces)
            self.sem_seg_encoder = ResnetSemSeqEncoder(
                sem_obs_space,
                output_size=model_config.SEMANTIC_ENCODER.output_size,
                backbone=model_config.SEMANTIC_ENCODER.backbone,
                trainable=model_config.SEMANTIC_ENCODER.train_encoder,
                semantic_embedding_size=sem_embedding_size,
                is_thda=self.is_thda
            )
            sem_seg_output_size = model_config.SEMANTIC_ENCODER.output_size
            logger.info("Setting up Sem Seg model")
            self.step_input_sizes['sem'] = sem_seg_output_size

            self.embed_sge = model_config.embed_sge
            if self.embed_sge:
                self.task_cat2mpcat40 = torch.tensor(task_cat2mpcat40)
                self.mapping_mpcat40_to_goal = np.zeros(
                    max(
                        max(mapping_mpcat40_to_goal21.keys()) + 1,
                        50,
                    ),
                    dtype=np.int8,
                )

                for key, value in mapping_mpcat40_to_goal21.items():
                    self.mapping_mpcat40_to_goal[key] = value
                self.mapping_mpcat40_to_goal = torch.tensor(self.mapping_mpcat40_to_goal)
                self.step_input_sizes['sem_sge'] = 1
        
    def build_object_category_encoder(self, model_config, hidden_size):
        self._n_object_categories = 28
        logger.info("Object categories: {}".format(self._n_object_categories))
        obj_embed_file = getattr(model_config, 'obj_embed_file', None)
        if obj_embed_file is None:
            self.obj_category_to_embeds = None
            self.obj_categories_embedding = nn.Embedding(
                self._n_object_categories, hidden_size
            )
        else:
            obj_embeds = np.load(obj_embed_file, allow_pickle=True).item()
            category_to_task_category_id = {'chair': 0, 'table': 1, 'picture': 2, 'cabinet': 3, 'cushion': 4, 'sofa': 5, 'bed': 6, 'chest_of_drawers': 7, 'plant': 8, 'sink': 9, 'toilet': 10, 'stool': 11, 'towel': 12, 'tv_monitor': 13, 'shower': 14, 'bathtub': 15, 'counter': 16, 'fireplace': 17, 'gym_equipment': 18, 'seating': 19, 'clothes': 20}

            self.obj_category_to_embeds = {
                category_to_task_category_id[k]: torch.from_numpy(v.astype(np.float32)) \
                    for k, v in obj_embeds.items() if k in category_to_task_category_id
            }
            self.obj_categories_embedding = nn.Linear(
                list(obj_embeds.values())[0].shape[0], hidden_size
            )
        self.step_input_sizes['objectgoal'] = hidden_size

    def build_encoders_concat(self, observation_space, model_config, num_actions):
        '''Habitat-web baseline: concatenate the embeded features
        '''
        self.build_visual_encoders(observation_space, model_config)

        if EpisodicGPSSensor.cls_uuid in observation_space.spaces:
            input_gps_dim = observation_space.spaces[
                EpisodicGPSSensor.cls_uuid
            ].shape[0]
            self.gps_embedding = nn.Linear(input_gps_dim, 32)
            self.step_input_sizes['gps'] = 32
            logger.info("\n\nSetting up GPS sensor")
        
        if EpisodicCompassSensor.cls_uuid in observation_space.spaces:
            assert (
                observation_space.spaces[EpisodicCompassSensor.cls_uuid].shape[
                    0
                ]
                == 1
            ), "Expected compass with 2D rotation."
            input_compass_dim = 2  # cos and sin of the angle
            self.compass_embedding_dim = 32
            self.compass_embedding = nn.Linear(input_compass_dim, self.compass_embedding_dim)
            self.step_input_sizes['compass'] = 32
            logger.info("\n\nSetting up Compass sensor")

        self.build_object_category_encoder(model_config, 32)

        if model_config.SEQ2SEQ.use_prev_action:
            self.prev_action_embedding = nn.Embedding(num_actions + 1, 32)
            self.step_input_sizes['prev_action'] = self.prev_action_embedding.embedding_dim

        logger.info('step input size: %d' % (np.sum(list(self.step_input_sizes.values()))))

        if model_config.encoder_type == 'concat_linear':
            step_input_size = np.sum(list(self.step_input_sizes.values()))
            self.ft_fusion_layer = nn.Linear(step_input_size, self.output_size)
        else:
            self.ft_fusion_layer = None

    def build_encoders_add(self, observation_space, model_config, num_actions):
        '''Typical approach in transformer: add all the embeded features
        '''
        hidden_size = model_config.hidden_size

        model_config.defrost()
        model_config.DEPTH_ENCODER.output_size = hidden_size
        model_config.RGB_ENCODER.output_size = hidden_size
        model_config.freeze()

        # RGB and depth encoder
        self.build_visual_encoders(observation_space, model_config)
        self.rgb_layer_norm = nn.LayerNorm(hidden_size)
        self.depth_layer_norm = nn.LayerNorm(hidden_size)

        if EpisodicGPSSensor.cls_uuid in observation_space.spaces:
            input_gps_dim = observation_space.spaces[
                EpisodicGPSSensor.cls_uuid
            ].shape[0]
            self.gps_embedding = nn.Sequential(
                nn.Linear(input_gps_dim, hidden_size),
                nn.LayerNorm(hidden_size)
            )
            self.step_input_sizes['gps'] = hidden_size
            logger.info("\n\nSetting up GPS sensor")
        
        if EpisodicCompassSensor.cls_uuid in observation_space.spaces:
            assert (
                observation_space.spaces[EpisodicCompassSensor.cls_uuid].shape[
                    0
                ]
                == 1
            ), "Expected compass with 2D rotation."
            input_compass_dim = 2  # cos and sin of the angle
            self.compass_embedding = nn.Sequential(
                nn.Linear(input_compass_dim, hidden_size),
                nn.LayerNorm(hidden_size),
            )
            self.step_input_sizes['compass'] = hidden_size
            logger.info("\n\nSetting up Compass sensor")

        self.build_object_category_encoder(model_config, hidden_size)

        if model_config.SEQ2SEQ.use_prev_action:
            self.prev_action_embedding = nn.Embedding(num_actions + 1, hidden_size)
            self.step_input_sizes['prev_action'] = hidden_size

        self.ft_fusion_layer = nn.Sequential(
            nn.LayerNorm(hidden_size),
            nn.Dropout(model_config.dropout_rate),
        )

    def build_encoders(self, observation_space, model_config, num_actions):
        # encoders
        if model_config.encoder_type.startswith('concat'):
            self.build_encoders_concat(observation_space, model_config, num_actions)
        elif model_config.encoder_type == 'add':
            self.build_encoders_add(observation_space, model_config, num_actions)

    def build_model(self, observation_space, model_config, num_actions):
        raise NotImplementedError

    @property
    def is_blind(self):
        return self.rgb_encoder.is_blind and self.depth_encoder.is_blind

    def _extract_sge(self, observations):
        # recalculating to keep this self-contained instead of depending on training infra
        if "semantic" in observations and "objectgoal" in observations:
            obj_semantic = observations["semantic"].contiguous().flatten(start_dim=1)
            
            if len(observations["objectgoal"].size()) == 3:
                observations["objectgoal"] = observations["objectgoal"].contiguous().view(
                    -1, observations["objectgoal"].size(2)
                )

            idx = self.task_cat2mpcat40[
                observations["objectgoal"].long()
            ]
            if self.is_thda:
                idx = self.mapping_mpcat40_to_goal[idx].long()
            idx = idx.to(obj_semantic.device)

            if len(idx.size()) == 3:
                idx = idx.squeeze(1)

            goal_visible_pixels = (obj_semantic == idx).sum(dim=1)
            goal_visible_area = torch.true_divide(goal_visible_pixels, obj_semantic.size(-1)).float()
            return goal_visible_area.unsqueeze(-1)

    def encode_visual_obs(self, observations):
        rgb_obs = observations.get("rgb", None)
        depth_obs = observations.get("depth", None)

        outs = {}

        if self.depth_encoder is not None:
            if depth_obs is not None and len(depth_obs.size()) == 5:  # (num_steps, num_envs, h, w, c)
                observations["depth"] = depth_obs.contiguous().view(
                    -1, depth_obs.size(2), depth_obs.size(3), depth_obs.size(4)
                )
            depth_embedding = self.depth_encoder(observations)
            if self.model_config.num_ft_views > 1: # multiple cameras
                depth_embedding = [depth_embedding]
                for k in range(1, self.model_config.num_ft_views):
                    depth_embedding.append(self.depth_encoder({'depth': observations[f'depth_{k}']}))
                depth_embedding = torch.cat(depth_embedding, -1)
            outs['depth'] = depth_embedding

        if self.rgb_encoder is not None:
            if rgb_obs is not None and len(rgb_obs.size()) == 5:    # (num_steps, num_envs, h, w, c)
                observations["rgb"] = rgb_obs.contiguous().view(
                    -1, rgb_obs.size(2), rgb_obs.size(3), rgb_obs.size(4)
                )
            rgb_embedding = self.rgb_encoder(observations)
            if self.model_config.num_ft_views > 1: # multiple cameras
                rgb_embedding = [rgb_embedding]
                for k in range(1, self.model_config.num_ft_views):
                    rgb_embedding.append(self.rgb_encoder({'rgb': observations[f'rgb_{k}']}))
                rgb_embedding = torch.cat(rgb_embedding, -1)
            outs['rgb'] = rgb_embedding

        if self.model_config.USE_SEMANTICS:
            semantic_obs = observations["semantic"] # (num_steps, num_envs, h, w)
            if len(semantic_obs.size()) == 4:
                observations["semantic"] = semantic_obs.contiguous().view(
                    -1, semantic_obs.size(2), semantic_obs.size(3)
                )
            if self.embed_sge:
                sge_embedding = self._extract_sge(observations)
                outs['sem_sge'] = sge_embedding

            sem_seg_embedding = self.sem_seg_encoder(observations)
            outs['sem'] = sem_seg_embedding

        return outs

    def encode_step_obs_concat(self, observations, prev_actions=None, step_embeddings=None, offline=False):
        outs = self.encode_visual_obs(observations)
        embeds = []
        for key in ['depth', 'rgb', 'sem', 'sem_sge']:
            if key in outs:
                embeds.append(outs[key])
        device = embeds[-1].device

        if EpisodicGPSSensor.cls_uuid in observations:
            obs_gps = observations[EpisodicGPSSensor.cls_uuid]  # (num_steps, num_envs, 2)
            if len(obs_gps.size()) == 3 and (not offline):
                obs_gps = obs_gps.contiguous().view(-1, obs_gps.size(2))
            outs['gps'] = self.gps_embedding(obs_gps)
            embeds.append(outs['gps'])
        
        if EpisodicCompassSensor.cls_uuid in observations:
            obs_compass = observations["compass"]  # (num_steps, num_envs, 1)
            if len(obs_compass.size()) == 3 and (not offline):
                obs_compass = obs_compass.contiguous().view(-1, obs_compass.size(2))
            if offline:
                compass_observations = torch.concat(
                    [torch.cos(obs_compass), torch.sin(obs_compass)], -1,
                )
                compass_embedding = self.compass_embedding(compass_observations)
            else:
                compass_observations = torch.stack(
                    [torch.cos(obs_compass), torch.sin(obs_compass)], -1,
                )
                compass_embedding = self.compass_embedding(compass_observations.squeeze(dim=1))
            outs['compass'] = compass_embedding
            embeds.append(outs['compass'])

        object_goal = observations[ObjectGoalSensor.cls_uuid].long()  # (num_steps, num_envs, 1)
        if len(object_goal.size()) == 3 and (not offline):
            object_goal = object_goal.contiguous().view(-1, object_goal.size(2))
        if self.obj_category_to_embeds is None:
            if offline:
                embeds.append(self.obj_categories_embedding(object_goal))
            else:
                embeds.append(self.obj_categories_embedding(object_goal).squeeze(dim=1))
        else:
            obj_embedding = torch.stack(
                [self.obj_category_to_embeds[k] for k in object_goal[..., 0].data.cpu().numpy()], 0
            )
            embeds.append(self.obj_categories_embedding(obj_embedding.to(device)))

        if self.model_config.SEQ2SEQ.use_prev_action:
            if offline:
                prev_actions_embedding = self.prev_action_embedding(
                    prev_actions + 1
                )   # (batch_size, num_steps)
            else:
                prev_actions_embedding = self.prev_action_embedding(
                    (prev_actions + 1).view(-1)
                )  # prev_actions: (num_steps, num_envs, 1)
            embeds.append(prev_actions_embedding)
        
        if offline:
            for i, x in enumerate(embeds):
                if len(x.size()) != 3: # (3, (4, 32))
                    print(i, x.size())
        embeds = torch.cat(embeds, dim=-1)
        if self.ft_fusion_layer is not None:
            embeds = self.ft_fusion_layer(embeds)

        if step_embeddings is not None:
            embeds = embeds + step_embeddings

        outs['fused_embeds'] = embeds
        
        return outs

    def encode_step_obs_add(self, observations, prev_actions=None, step_embeddings=None):
        outs = self.encode_visual_obs(observations)

        if self.depth_encoder is not None:
            outs['depth'] = self.depth_layer_norm(outs['depth'])

        if self.rgb_encoder is not None:
            outs['rgb'] = self.rgb_layer_norm(outs['rgb'])

        if EpisodicGPSSensor.cls_uuid in observations:
            obs_gps = observations[EpisodicGPSSensor.cls_uuid]  # (num_steps, num_envs, 2)
            if len(obs_gps.size()) == 3:
                obs_gps = obs_gps.contiguous().view(-1, obs_gps.size(2))
            outs['gps'] = self.gps_embedding(obs_gps)
        
        if EpisodicCompassSensor.cls_uuid in observations:
            obs_compass = observations["compass"]  # (num_steps, num_envs, 1)
            if len(obs_compass.size()) == 3:
                obs_compass = obs_compass.contiguous().view(-1, obs_compass.size(2))
            compass_observations = torch.stack(
                [torch.cos(obs_compass), torch.sin(obs_compass)], -1,
            )
            compass_embedding = self.compass_embedding(compass_observations.squeeze(dim=1))
            outs['compass'] = compass_embedding

        if self.model_config.SEQ2SEQ.use_prev_action:
            prev_actions_embedding = self.prev_action_embedding(
                (prev_actions + 1).view(-1)
            )  # prev_actions: (num_steps, num_envs, 1)
            outs['prev_action'] = prev_actions_embedding

        # (batch_size, hidden_size)
        fused_embeds = torch.sum(torch.stack(list(outs.values()), 0), 0)

        object_goal = observations[ObjectGoalSensor.cls_uuid].long()  # (num_steps, num_envs, 1)
        if len(object_goal.size()) == 3:
            object_goal = object_goal.contiguous().view(-1, object_goal.size(2))
        outs['objectgoal'] = self.obj_categories_embedding(object_goal).squeeze(dim=1)

        add_objgoal = self.model_config.get('encoder_add_objgoal', True)
        if add_objgoal:
            fused_embeds = fused_embeds + outs['objectgoal']

        if step_embeddings is not None:
            fused_embeds = fused_embeds + step_embeddings

        fused_embeds = self.ft_fusion_layer(fused_embeds)
        outs['fused_embeds'] = fused_embeds

        return outs

    def encode_step_obs(self, observations, **kwargs):
        if self.model_config.encoder_type.startswith('concat'):
            return self.encode_step_obs_concat(observations, **kwargs)
        elif self.model_config.encoder_type == 'add':
            return self.encode_step_obs_add(observations, **kwargs)

    def forward(self, observations, hidden_states, prev_actions, nav_step):
        raise NotImplementedError


habitat_baselines/il/env_based/policy/onav_base.py
cshizhe/onav_rim
import torch
import torch.nn as nn
import einops

from habitat_baselines.il.env_based.policy.onav_base import ObjectNavBase
from habitat_baselines.il.env_based.policy.onav_vis_models import ObjectNavRNN

from offline_bc.models.transformer import (
    PositionalEncoding, 
    TransformerEncoderLayer,
    TransformerEncoder
)
from offline_bc.models.onav_imap_models import (
    ImapEmbedding, GpsCompassEmbedding
)


class ObjectNavImapSingleTransformer(ObjectNavRNN):

    @property
    def output_size(self):
        return self.model_config.hidden_size

    def build_model(self, observation_space, model_config, num_actions):
        self.build_encoders(observation_space, model_config, num_actions)

        if model_config.STATE_ENCODER.learnable_step_embedding:
            self.step_embedding = nn.Embedding(
                model_config.STATE_ENCODER.max_steps,
                model_config.hidden_size,
            )
        else:
            self.step_embedding = PositionalEncoding(
                model_config.hidden_size,
                max_len=2000 #model_config.STATE_ENCODER.max_steps,
            )

        encoder_layer = TransformerEncoderLayer(
            model_config.hidden_size, 
            model_config.STATE_ENCODER.num_attention_heads, 
            dim_feedforward=model_config.STATE_ENCODER.intermediate_size, 
            dropout=model_config.STATE_ENCODER.dropout_prob,
            activation=model_config.STATE_ENCODER.hidden_act, 
            normalize_before=False
        )
        self.state_encoder = TransformerEncoder(
            encoder_layer, 
            model_config.STATE_ENCODER.num_hidden_layers, 
            norm=None, batch_first=True
        )
        if model_config.STATE_ENCODER.add_pos_attn:
            self.obs_pos_layer = GpsCompassEmbedding(model_config.hidden_size)

        self.imap_embedding = ImapEmbedding(model_config)

        subgoal_pred_type = getattr(model_config, 'subgoal_pred_type', None)
        if subgoal_pred_type == 'coord':
            self.subgoal_pred_layer = nn.Sequential(
                nn.Dropout(model_config.dropout_rate),
                nn.Linear(model_config.hidden_size, 3 + 22 + 2),
            )
        elif subgoal_pred_type == 'heatmap':
            x, y = torch.meshgrid(
                torch.arange(model_config.subgoal_goto_heatmap_size), 
                torch.arange(model_config.subgoal_goto_heatmap_size)
            )
            xy = torch.stack([x, y], dim=2)
            xy = (xy + 0.5 - model_config.subgoal_goto_heatmap_size / 2).float() # relative distance to the center
            xy = xy.view(-1, 2) # (heatmapsize**2, 2)
            self.subgoal_goto_heatmap_locs = xy * model_config.subgoal_goto_heatmap_grid_meter
            self.subgoal_pred_layer = nn.Sequential(
                nn.Dropout(model_config.dropout_rate),
                nn.Linear(model_config.hidden_size, 3 + 22 + model_config.subgoal_goto_heatmap_size**2),
            )    
        else:
            self.subgoal_pred_layer = None        

    def forward(self, batch, recursive_states, prev_actions, step_id, pred_subgoal=False):
        batch_size = batch['gps'].size(0)
        device = batch['gps'].device

        step_ids = torch.zeros(batch_size).long().to(device) + step_id
        stepid_embeds = self.step_embedding(step_ids)

        # (batch_size, hidden_size)
        t_inputs = self.encode_step_obs(
            batch, prev_actions=prev_actions, 
            step_embeddings=stepid_embeds
        )
        t_embeds = t_inputs['fused_embeds']

        if recursive_states is None:
            imap_embeds, imap_pos_embeds = self.imap_embedding(batch_size)
        elif recursive_states[1] is None:
            _, imap_pos_embeds = self.imap_embedding(batch_size)
            imap_embeds = recursive_states[0]
        else:
            imap_embeds, imap_pos_embeds = recursive_states
        
        if self.model_config.STATE_ENCODER.add_pos_attn:
            obs_pos_embeds = self.obs_pos_layer(batch['gps'], batch['compass'])
            t_pos_embeds = torch.cat(
                [imap_pos_embeds, obs_pos_embeds.unsqueeze(1)], dim=1
            )
        else:
            t_pos_embeds = None

        hiddens, layer_attn_weights = self.state_encoder(
            torch.cat([imap_embeds, t_embeds.unsqueeze(1)], dim=1), 
            pos=t_pos_embeds, return_attn_weights=True, 
        )

        imap_embeds = hiddens[:, :-1]
        obs_hiddens = hiddens[:, -1]
        add_objgoal = self.model_config.get('encoder_add_objgoal', True)
        if not add_objgoal:
            obs_hiddens = obs_hiddens + t_inputs['objectgoal']

        return obs_hiddens, (imap_embeds, imap_pos_embeds)

habitat_baselines/il/env_based/policy/onav_imap_models.py
cshizhe/onav_rim
import torch
import torch.nn as nn


from habitat_baselines.rl.models.rnn_state_encoder import RNNStateEncoder
from habitat_baselines.il.env_based.policy.onav_base import ObjectNavBase

from offline_bc.models.transformer import (
    triangular_mask, PositionalEncoding,
    TransformerEncoderLayer,
    TransformerEncoder,
)


class ObjectNavRNN(ObjectNavBase):
    r"""A baseline sequence to sequence network that concatenates instruction,
    RGB, and depth encodings before decoding an action distribution with an RNN.
    Modules:
        Instruction encoder
        Depth encoder
        RGB encoder
        RNN state encoder
    """
    @property
    def num_recurrent_layers(self):
        return self.state_encoder.num_recurrent_layers

    @property
    def output_size(self):
        return self.model_config.hidden_size

    def build_model(self, observation_space, model_config, num_actions):
        self.build_encoders(observation_space, model_config, num_actions)
        
        self.state_encoder = RNNStateEncoder(
            input_size=self.output_size,
            hidden_size=model_config.hidden_size,
            num_layers=model_config.STATE_ENCODER.num_recurrent_layers,
            rnn_type=model_config.STATE_ENCODER.rnn_type,
        )

    def forward(self, observations, rnn_hidden_states, prev_actions, masks):
        r"""
        instruction_embedding: [batch_size x INSTRUCTION_ENCODER.output_size]
        depth_embedding: [batch_size x DEPTH_ENCODER.output_size]
        rgb_embedding: [batch_size x RGB_ENCODER.output_size]
        """
        x = self.encode_step_obs(observations, prev_actions=prev_actions)['fused_embeds']
        x, rnn_hidden_states = self.state_encoder(x, rnn_hidden_states, masks)

        return x, rnn_hidden_states


class ObjectNavTransformer(ObjectNavBase):

    @property
    def output_size(self):
        return self.model_config.hidden_size

    def build_model(self, observation_space, model_config, num_actions):
        self.build_encoders(observation_space, model_config, num_actions)

        if model_config.STATE_ENCODER.learnable_step_embedding:
            self.step_embedding = nn.Embedding(
                model_config.STATE_ENCODER.max_steps,
                model_config.hidden_size,
            )
        else:
            self.step_embedding = PositionalEncoding(
                model_config.hidden_size,
                max_len=model_config.STATE_ENCODER.max_steps,
            )

        encoder_layer = TransformerEncoderLayer(
            model_config.hidden_size, 
            model_config.STATE_ENCODER.num_attention_heads, 
            dim_feedforward=model_config.STATE_ENCODER.intermediate_size, 
            dropout=model_config.STATE_ENCODER.dropout_prob,
            activation=model_config.STATE_ENCODER.hidden_act, 
            normalize_before=False
        )
        self.state_encoder = TransformerEncoder(
            encoder_layer, 
            model_config.STATE_ENCODER.num_hidden_layers, 
            norm=None, batch_first=True
        )

    def forward(self, batch, history_embeds, prev_actions, step_id):
        batch_size = batch['gps'].size(0)
        device = batch['gps'].device

        step_ids = torch.zeros(batch_size).long().to(device) + step_id
        step_embeddings = self.step_embedding(step_ids)

        # (batch_size, hidden_size)
        inputs = self.encode_step_obs(
            batch, prev_actions=prev_actions, 
            step_embeddings=step_embeddings
        )
        embeds = inputs['fused_embeds']

        if history_embeds is None:
            history_embeds = embeds.unsqueeze(1)
        else:
            history_embeds = torch.cat([history_embeds, embeds.unsqueeze(1)], dim=1)
        
        causal_attn_masks = triangular_mask(
            history_embeds.size(1), device, diagonal_shift=1
        )
        hiddens = self.state_encoder(
            history_embeds, mask=causal_attn_masks, 
        )

        hiddens = hiddens[:, -1]
        add_objgoal = self.model_config.get('encoder_add_objgoal', True)
        if not add_objgoal:
            hiddens = hiddens + inputs['objectgoal']

        return hiddens, history_embeds


habitat_baselines/il/env_based/policy/onav_vis_models.py
cshizhe/onav_rim
import torch
import torch.nn as nn
import einops

from gym import Space
from habitat import Config, logger

from habitat_baselines.common.baseline_registry import baseline_registry
from habitat_baselines.rl.ppo import Policy

from habitat_baselines.il.env_based.policy.onav_vis_models import (
    ObjectNavRNN,
    ObjectNavTransformer,
)
from habitat_baselines.il.env_based.policy.onav_imap_models import (
    ObjectNavImapSingleTransformer
)
from offline_bc.models.onav_base import ClsPrediction


model_factory = {
    'ObjectNavRNN': ObjectNavRNN,
    'ObjectNavTransformer': ObjectNavTransformer,
    'ObjectNavImapSingleTransformer': ObjectNavImapSingleTransformer,
}

@baseline_registry.register_policy
class ObjectNavILPolicy(Policy):
    def __init__(
        self, observation_space: Space, action_space: Space, model_config: Config,
        no_critic=True
    ):
        model_class = model_factory[model_config.model_class]
        model = model_class(
            observation_space=observation_space,
            model_config=model_config,
            num_actions=action_space.n,
        )

        if model_config.action_clf_class == 'ClsPrediction':
            clf_net = ClsPrediction(model.output_size, action_space.n)
        elif model_config.action_clf_class == 'linear':
            clf_net = nn.Sequential(
                nn.Dropout(model_config.dropout_rate),
                nn.Linear(model.output_size, action_space.n),
            )
        else:
            clf_net = None
        
        super().__init__(
            model,
            action_space.n,
            no_critic=no_critic,
            clf_net=clf_net
        )

    @classmethod
    def from_config(
        cls, config: Config, observation_space, action_space, no_critic=True
    ):
        return cls(
            observation_space=observation_space,
            action_space=action_space,
            model_config=config.MODEL,         
            no_critic=no_critic,   
        )


habitat_baselines/il/env_based/policy/resnet_policy.py
cshizhe/onav_rim
import functools
import os
import shlex
import signal
import subprocess
import threading
from os import path as osp
from typing import Any, Callable, Optional, Tuple, Union, overload

import ifcfg
import torch
from torch import distributed as distrib

from habitat import logger

EXIT = threading.Event()
EXIT.clear()
REQUEUE = threading.Event()
REQUEUE.clear()


# Default port to initialized the TCP store on
DEFAULT_PORT = 8738
DEFAULT_PORT_RANGE = 127
# Default address of world rank 0
DEFAULT_MASTER_ADDR = "127.0.0.1"

SLURM_JOBID = os.environ.get("SLURM_JOB_ID", None)
INTERRUPTED_STATE_FILE = osp.join(
    "/scratch/shichen/datasets/habitat_data/exprs/", "interrupted_states", f"{SLURM_JOBID}.pth"
)


def is_slurm_job() -> bool:
    return SLURM_JOBID is not None


def is_slurm_batch_job() -> bool:
    r"""Heuristic to determine if a slurm job is a batch job or not. Batch jobs
    will have a job name that is not a shell unless the user specifically set the job
    name to that of a shell. Interactive jobs have a shell name as their job name.
    """
    return is_slurm_job() and os.environ.get("SLURM_JOB_NAME", None) not in (
        None,
        "bash",
        "zsh",
        "fish",
        "tcsh",
        "sh",
    )


@overload
def rank0_only() -> bool:
    ...


@overload
def rank0_only(fn: Callable) -> Callable:
    ...


def rank0_only(fn: Optional[Callable] = None) -> Union[Callable, bool]:
    r"""Helper function to only execute code if a process is world rank 0
    Can be used both as a function in an if statement,
    .. code:: py
        if rank0_only():
            ...
    or as a decorator,
    .. code:: py
        @rank0_only
        def fn_for_r0_only(...):
            ...
    :param fn: Function to wrap and only execute if the process is rank 0.
        If a process is rank 0, the function will be run and it's return value
        will be returned.  If a process is not rank 0, then the function will not
        be ran and :py:`None` will be returned.
    :return: The wrapped function if :p:`fn` is not :py:`None`, otherwise
        whether or not this process is rank 0
    """
    if fn is None:
        return (
            not torch.distributed.is_initialized()
            or torch.distributed.get_rank() == 0
        )

    @functools.wraps(fn)
    def _wrapper(*args, **kwargs):
        if rank0_only():
            return fn(*args, **kwargs)
        return None

    return _wrapper


def _clean_exit_handler(signum, frame):
    EXIT.set()
    print("Exiting cleanly", flush=True)


def _requeue_handler(signal, frame):
    print("Got signal to requeue", flush=True)
    EXIT.set()
    REQUEUE.set()


def add_signal_handlers() -> None:
    signal.signal(signal.SIGINT, _clean_exit_handler)
    signal.signal(signal.SIGTERM, _clean_exit_handler)

    # SIGUSR2 can be sent to all processes to have them cleanup
    # and exit nicely.  This is nice to use with SLURM as scancel <job_id>
    # sets a 30 second timer for the job to exit, and it can take more than
    # 30 seconds for the job to cleanup and exit nicely.  When using NCCL,
    # forcing the job to exit without cleaning up can be bad.
    # scancel --signal SIGUSR2 <job_id> will set no such timer and will give
    # the job ample time to cleanup and exit.
    signal.signal(signal.SIGUSR2, _clean_exit_handler)

    signal.signal(signal.SIGUSR1, _requeue_handler)


@rank0_only
def save_interrupted_state(state: Any, filename: str = None):
    r"""Saves the interrupted job state to the specified filename.
        This is useful when working with preemptable job partitions.
    This method will do nothing if SLURM is not currently being used and the filename is the default
    :param state: The state to save
    :param filename: The filename.  Defaults to "${HOME}/.interrupted_states/${SLURM_JOBID}.pth"
    """
    if SLURM_JOBID is None and filename is None:
        logger.warn("SLURM_JOBID is none, not saving interrupted state")
        return

    if filename is None:
        filename = INTERRUPTED_STATE_FILE
        if not osp.exists(osp.dirname(INTERRUPTED_STATE_FILE)):
            raise RuntimeError(
                "Please create a .interrupted_states directory in your home directory for job preemption"
                "(This is intentionally not created automatically as it can get quite large)"
            )

    torch.save(state, filename)


def load_interrupted_state(filename: str = None) -> Optional[Any]:
    r"""Loads the saved interrupted state
    :param filename: The filename of the saved state.
        Defaults to "${HOME}/.interrupted_states/${SLURM_JOBID}.pth"
    :return: The saved state if the file exists, else none
    """
    if SLURM_JOBID is None and filename is None:
        return None

    if filename is None:
        filename = INTERRUPTED_STATE_FILE

    if not osp.exists(filename):
        return None

    return torch.load(filename, map_location="cpu")


def requeue_job():
    r"""Requeues the job by calling ``scontrol requeue ${SLURM_JOBID}``"""
    if SLURM_JOBID is None:
        return

    if not REQUEUE.is_set():
        return

    if distrib.is_initialized():
        distrib.barrier()

    if rank0_only():
        logger.info(f"Requeueing job {SLURM_JOBID}")
        subprocess.check_call(shlex.split(f"scontrol requeue {SLURM_JOBID}"))


def get_ifname() -> str:
    return ifcfg.default_interface()["device"]


def get_distrib_size() -> Tuple[int, int, int]:
    # Check to see if we should parse from torch.distributed.launch
    if os.environ.get("LOCAL_RANK", None) is not None:
        local_rank = int(os.environ["LOCAL_RANK"])
        world_rank = int(os.environ["RANK"])
        world_size = int(os.environ["WORLD_SIZE"])
    # Else parse from SLURM is using SLURM
    elif os.environ.get("SLURM_JOBID", None) is not None:
        local_rank = int(os.environ["SLURM_LOCALID"])
        world_rank = int(os.environ["SLURM_PROCID"])
        world_size = int(os.environ["SLURM_NTASKS"])
    # Otherwise setup for just 1 process, this is nice for testing
    else:
        local_rank = 0
        world_rank = 0
        world_size = 1

    return local_rank, world_rank, world_size


def init_distrib_slurm(
    backend: str = "nccl",
) -> Tuple[int, torch.distributed.TCPStore]:  # type: ignore
    r"""Initializes torch.distributed by parsing environment variables set
        by SLURM when ``srun`` is used or by parsing environment variables set
        by torch.distributed.launch
    :param backend: Which torch.distributed backend to use
    :returns: Tuple of the local_rank (aka which GPU to use for this process)
        and the TCPStore used for the rendezvous
    """
    assert (
        torch.distributed.is_available()
    ), "torch.distributed must be available"

    if "GLOO_SOCKET_IFNAME" not in os.environ:
        os.environ["GLOO_SOCKET_IFNAME"] = get_ifname()

    if "NCCL_SOCKET_IFNAME" not in os.environ:
        os.environ["NCCL_SOCKET_IFNAME"] = get_ifname()

    local_rank, world_rank, world_size = get_distrib_size()

    master_port = int(os.environ.get("MASTER_PORT", DEFAULT_PORT))
    if SLURM_JOBID is not None:
        master_port += int(SLURM_JOBID) % int(
            os.environ.get("MASTER_PORT_RANGE", DEFAULT_PORT_RANGE)
        )
    master_addr = os.environ.get("MASTER_ADDR", DEFAULT_MASTER_ADDR)

    tcp_store = distrib.TCPStore(  # type: ignore
        master_addr, master_port, world_size, world_rank == 0
    )
    distrib.init_process_group(
        backend, store=tcp_store, rank=world_rank, world_size=world_size
    )

    return local_rank, tcp_store
habitat_baselines/rl/ddppo/algo/ddp_utils.py
cshizhe/onav_rim
#!/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 contextlib
import os
import random
import time
from collections import defaultdict, deque
from typing import DefaultDict, Optional

import numpy as np
import torch
from gym import spaces
from torch import distributed as distrib
from torch import nn as nn
from torch.optim.lr_scheduler import LambdaLR

from habitat import Config, logger
from habitat.utils import profiling_wrapper
from habitat_baselines.common.baseline_registry import baseline_registry
from habitat_baselines.common.environments import get_env_class
from habitat_baselines.common.obs_transformers import (
    apply_obs_transforms_batch,
    apply_obs_transforms_obs_space,
    get_active_obs_transforms,
)
from habitat_baselines.common.rollout_storage import RolloutStorage
from habitat_baselines.common.tensorboard_utils import TensorboardWriter
from habitat_baselines.rl.ddppo.algo.ddp_utils import (
    EXIT,
    REQUEUE,
    add_signal_handlers,
    init_distrib_slurm,
    load_interrupted_state,
    requeue_job,
    save_interrupted_state,
)
from habitat_baselines.rl.ddppo.algo.ddppo import DDPPO
from habitat_baselines.rl.ddppo.policy.resnet_policy import (  # noqa: F401
    PointNavResNetPolicy,
)
from habitat_baselines.rl.ppo.ppo_trainer import PPOTrainer
from habitat_baselines.utils.common import batch_obs, linear_decay
from habitat_baselines.utils.env_utils import construct_envs


@baseline_registry.register_trainer(name="ddppo")
class DDPPOTrainer(PPOTrainer):
    # DD-PPO cuts rollouts short to mitigate the straggler effect
    # This, in theory, can cause some rollouts to be very short.
    # All rollouts contributed equally to the loss/model-update,
    # thus very short rollouts can be problematic.  This threshold
    # limits the how short a short rollout can be as a fraction of the
    # max rollout length
    SHORT_ROLLOUT_THRESHOLD: float = 0.25

    def __init__(self, config: Optional[Config] = None) -> None:
        interrupted_state = load_interrupted_state()
        if interrupted_state is not None:
            config = interrupted_state["config"]

        super().__init__(config)

    def _setup_actor_critic_agent(self, ppo_cfg: Config) -> None:
        r"""Sets up actor critic and agent for DD-PPO.

        Args:
            ppo_cfg: config node with relevant params

        Returns:
            None
        """
        logger.add_filehandler(self.config.LOG_FILE)

        policy = baseline_registry.get_policy(self.config.RL.POLICY.name)
        self.obs_transforms = get_active_obs_transforms(self.config)
        observation_space = self.envs.observation_spaces[0]
        observation_space = apply_obs_transforms_obs_space(
            observation_space, self.obs_transforms
        )
        self.actor_critic = policy.from_config(
            self.config, observation_space, self.envs.action_spaces[0]
        )
        self.obs_space = observation_space
        self.actor_critic.to(self.device)

        if (
            self.config.RL.DDPPO.pretrained_encoder
            or self.config.RL.DDPPO.pretrained
        ):
            pretrained_state = torch.load(
                self.config.RL.DDPPO.pretrained_weights, map_location="cpu"
            )

        if self.config.RL.DDPPO.pretrained:
            self.actor_critic.load_state_dict(
                {
                    k[len("actor_critic.") :]: v
                    for k, v in pretrained_state["state_dict"].items()
                }, strict=False
            )
        elif self.config.RL.DDPPO.pretrained_encoder:
            prefix = "actor_critic.net.visual_encoder."
            self.actor_critic.net.visual_encoder.load_state_dict(
                {
                    k[len(prefix) :]: v
                    for k, v in pretrained_state["state_dict"].items()
                    if k.startswith(prefix)
                }
            )

        if not self.config.RL.DDPPO.train_encoder:
            self._static_encoder = True
            for param in self.actor_critic.net.visual_encoder.parameters():
                param.requires_grad_(False)
        
        self._static_vis_encoder = False
        if hasattr(self.config, "MODEL") and self.config.MODEL.freeze_encoders:
            self._static_vis_encoder = True
            for param in self.actor_critic.net.depth_encoder.parameters():
                param.requires_grad_(False)
            for param in self.actor_critic.net.rgb_encoder.parameters():
                param.requires_grad_(False)
            for param in self.actor_critic.net.sem_seg_encoder.parameters():
                param.requires_grad_(False)

        if self.config.RL.DDPPO.reset_critic:
            nn.init.orthogonal_(self.actor_critic.critic.fc.weight)
            nn.init.constant_(self.actor_critic.critic.fc.bias, 0)

        self.agent = DDPPO(
            actor_critic=self.actor_critic,
            clip_param=ppo_cfg.clip_param,
            ppo_epoch=ppo_cfg.ppo_epoch,
            num_mini_batch=ppo_cfg.num_mini_batch,
            value_loss_coef=ppo_cfg.value_loss_coef,
            entropy_coef=ppo_cfg.entropy_coef,
            lr=ppo_cfg.lr,
            eps=ppo_cfg.eps,
            max_grad_norm=ppo_cfg.max_grad_norm,
            use_normalized_advantage=ppo_cfg.use_normalized_advantage,
        )

    @profiling_wrapper.RangeContext("train")
    def train(self) -> None:
        r"""Main method for DD-PPO.

        Returns:
            None
        """
        self.local_rank, tcp_store = init_distrib_slurm(
            self.config.RL.DDPPO.distrib_backend
        )
        add_signal_handlers()

        profiling_wrapper.configure(
            capture_start_step=self.config.PROFILING.CAPTURE_START_STEP,
            num_steps_to_capture=self.config.PROFILING.NUM_STEPS_TO_CAPTURE,
        )

        # Stores the number of workers that have finished their rollout
        num_rollouts_done_store = distrib.PrefixStore(
            "rollout_tracker", tcp_store
        )
        num_rollouts_done_store.set("num_done", "0")

        self.world_rank = distrib.get_rank()
        self.world_size = distrib.get_world_size()
        self.current_update = 0

        self.config.defrost()
        self.config.TORCH_GPU_ID = self.local_rank
        self.config.SIMULATOR_GPU_ID = self.local_rank
        # Multiply by the number of simulators to make sure they also get unique seeds
        self.config.TASK_CONFIG.SEED += (
            self.world_rank * self.config.NUM_PROCESSES
        )
        self.config.freeze()

        random.seed(self.config.TASK_CONFIG.SEED)
        np.random.seed(self.config.TASK_CONFIG.SEED)
        torch.manual_seed(self.config.TASK_CONFIG.SEED)

        if torch.cuda.is_available():
            self.device = torch.device("cuda", self.local_rank)
            torch.cuda.set_device(self.device)
        else:
            self.device = torch.device("cpu")

        self.envs = construct_envs(
            self.config,
            get_env_class(self.config.ENV_NAME),
            workers_ignore_signals=True,
        )

        ppo_cfg = self.config.RL.PPO
        if (
            not os.path.isdir(self.config.CHECKPOINT_FOLDER)
            and self.world_rank == 0
        ):
            os.makedirs(self.config.CHECKPOINT_FOLDER)

        self._setup_actor_critic_agent(ppo_cfg)
        self.agent.init_distributed(find_unused_params=True)

        if self.world_rank == 0:
            logger.info(
                "agent number of trainable parameters: {}".format(
                    sum(
                        param.numel()
                        for param in self.agent.parameters()
                        if param.requires_grad
                    )
                )
            )

        observations = self.envs.reset()
        batch = batch_obs(observations, device=self.device)

        if self.config.MODEL.USE_PRED_SEMANTICS and self.current_update >= self.config.MODEL.SWITCH_TO_PRED_SEMANTICS_UPDATE:
            batch["semantic"] = self.semantic_predictor(batch["rgb"], batch["depth"])
            # Subtract 1 from class labels for THDA YCB categories
            if self.config.MODEL.SEMANTIC_ENCODER.is_thda:
                batch["semantic"] = batch["semantic"] - 1

        batch = apply_obs_transforms_batch(batch, self.obs_transforms)

        obs_space = self.obs_space
        if self._static_encoder:
            self._encoder = self.actor_critic.net.visual_encoder
            obs_space = spaces.Dict(
                {
                    "visual_features": spaces.Box(
                        low=np.finfo(np.float32).min,
                        high=np.finfo(np.float32).max,
                        shape=self._encoder.output_shape,
                        dtype=np.float32,
                    ),
                    **obs_space.spaces,
                }
            )
            with torch.no_grad():
                batch["visual_features"] = self._encoder(batch)

        rollouts = RolloutStorage(
            ppo_cfg.num_steps,
            self.envs.num_envs,
            obs_space,
            self.envs.action_spaces[0],
            ppo_cfg.hidden_size,
            num_recurrent_layers=self.actor_critic.net.num_recurrent_layers,
        )
        rollouts.to(self.device)

        for sensor in rollouts.observations:
            rollouts.observations[sensor][0].copy_(batch[sensor])

        # batch and observations may contain shared PyTorch CUDA
        # tensors.  We must explicitly clear them here otherwise
        # they will be kept in memory for the entire duration of training!
        batch = None
        observations = None

        current_episode_reward = torch.zeros(
            self.envs.num_envs, 1, device=self.device
        )
        running_episode_stats = dict(
            count=torch.zeros(self.envs.num_envs, 1, device=self.device),
            reward=torch.zeros(self.envs.num_envs, 1, device=self.device),
        )
        window_episode_stats: DefaultDict[str, deque] = defaultdict(
            lambda: deque(maxlen=ppo_cfg.reward_window_size)
        )

        t_start = time.time()
        env_time = 0
        pth_time = 0
        count_steps: float = 0
        count_checkpoints = 0
        start_update = 0
        prev_time = 0

        lr_scheduler = LambdaLR(
            optimizer=self.agent.optimizer,
            lr_lambda=lambda x: linear_decay(x, self.config.NUM_UPDATES),  # type: ignore
        )

        interrupted_state = load_interrupted_state()
        if interrupted_state is not None:
            self.agent.load_state_dict(interrupted_state["state_dict"])
            self.agent.optimizer.load_state_dict(
                interrupted_state["optim_state"]
            )
            lr_scheduler.load_state_dict(interrupted_state["lr_sched_state"])

            requeue_stats = interrupted_state["requeue_stats"]
            env_time = requeue_stats["env_time"]
            pth_time = requeue_stats["pth_time"]
            count_steps = requeue_stats["count_steps"]
            count_checkpoints = requeue_stats["count_checkpoints"]
            start_update = requeue_stats["start_update"]
            prev_time = requeue_stats["prev_time"]

        with (
            TensorboardWriter(
                self.config.TENSORBOARD_DIR, flush_secs=self.flush_secs
            )
            if self.world_rank == 0
            else contextlib.suppress()
        ) as writer:
            for update in range(start_update, self.config.NUM_UPDATES):
                profiling_wrapper.on_start_step()
                profiling_wrapper.range_push("train update")
                self.current_update += 1

                if ppo_cfg.use_linear_lr_decay:
                    lr_scheduler.step()  # type: ignore

                if ppo_cfg.use_linear_clip_decay:
                    self.agent.clip_param = ppo_cfg.clip_param * linear_decay(
                        update, self.config.NUM_UPDATES
                    )

                if EXIT.is_set():
                    profiling_wrapper.range_pop()  # train update

                    self.envs.close()

                    if REQUEUE.is_set() and self.world_rank == 0:
                        requeue_stats = dict(
                            env_time=env_time,
                            pth_time=pth_time,
                            count_steps=count_steps,
                            count_checkpoints=count_checkpoints,
                            start_update=update,
                            prev_time=(time.time() - t_start) + prev_time,
                        )
                        save_interrupted_state(
                            dict(
                                state_dict=self.agent.state_dict(),
                                optim_state=self.agent.optimizer.state_dict(),
                                lr_sched_state=lr_scheduler.state_dict(),
                                config=self.config,
                                requeue_stats=requeue_stats,
                            )
                        )

                    requeue_job()
                    return

                count_steps_delta = 0
                self.agent.eval()
                profiling_wrapper.range_push("rollouts loop")
                for step in range(ppo_cfg.num_steps):

                    (
                        delta_pth_time,
                        delta_env_time,
                        delta_steps,
                    ) = self._collect_rollout_step(
                        rollouts, current_episode_reward, running_episode_stats
                    )
                    pth_time += delta_pth_time
                    env_time += delta_env_time
                    count_steps_delta += delta_steps

                    # This is where the preemption of workers happens.  If a
                    # worker detects it will be a straggler, it preempts itself!
                    if (
                        step
                        >= ppo_cfg.num_steps * self.SHORT_ROLLOUT_THRESHOLD
                    ) and int(num_rollouts_done_store.get("num_done")) > (
                        self.config.RL.DDPPO.sync_frac * self.world_size
                    ):
                        break
                profiling_wrapper.range_pop()  # rollouts loop

                num_rollouts_done_store.add("num_done", 1)

                self.agent.train()
                if self._static_encoder:
                    self._encoder.eval()
                
                if self._static_vis_encoder:
                    self.actor_critic.net.rgb_encoder.eval()
                    self.actor_critic.net.depth_encoder.eval()
                    self.actor_critic.net.sem_seg_encoder.eval()

                (
                    delta_pth_time,
                    value_loss,
                    action_loss,
                    dist_entropy,
                ) = self._update_agent(ppo_cfg, rollouts)
                pth_time += delta_pth_time

                stats_ordering = sorted(running_episode_stats.keys())
                stats = torch.stack(
                    [running_episode_stats[k] for k in stats_ordering], 0
                )
                distrib.all_reduce(stats)

                for i, k in enumerate(stats_ordering):
                    window_episode_stats[k].append(stats[i].clone())

                stats = torch.tensor(
                    [value_loss, action_loss, count_steps_delta],
                    device=self.device,
                )
                distrib.all_reduce(stats)
                count_steps += stats[2].item()

                if self.world_rank == 0:
                    num_rollouts_done_store.set("num_done", "0")

                    losses = [
                        stats[0].item() / self.world_size,
                        stats[1].item() / self.world_size,
                    ]
                    deltas = {
                        k: (
                            (v[-1] - v[0]).sum().item()
                            if len(v) > 1
                            else v[0].sum().item()
                        )
                        for k, v in window_episode_stats.items()
                    }
                    deltas["count"] = max(deltas["count"], 1.0)

                    writer.add_scalar(
                        "reward",
                        deltas["reward"] / deltas["count"],
                        count_steps,
                    )

                    # Check to see if there are any metrics
                    # that haven't been logged yet
                    metrics = {
                        k: v / deltas["count"]
                        for k, v in deltas.items()
                        if k not in {"reward", "count"}
                    }
                    if len(metrics) > 0:
                        writer.add_scalars("metrics", metrics, count_steps)

                    writer.add_scalars(
                        "losses",
                        {k: l for l, k in zip(losses, ["value", "policy"])},
                        count_steps,
                    )

                    # log stats
                    if update > 0 and update % self.config.LOG_INTERVAL == 0:
                        logger.info(
                            "update: {}\tfps: {:.3f}\t".format(
                                update,
                                count_steps
                                / ((time.time() - t_start) + prev_time),
                            )
                        )

                        logger.info(
                            "update: {}\tenv-time: {:.3f}s\tpth-time: {:.3f}s\t"
                            "frames: {}".format(
                                update, env_time, pth_time, count_steps
                            )
                        )
                        logger.info(
                            "Average window size: {}  {}".format(
                                len(window_episode_stats["count"]),
                                "  ".join(
                                    "{}: {:.3f}".format(k, v / deltas["count"])
                                    for k, v in deltas.items()
                                    if k != "count"
                                ),
                            )
                        )

                    # checkpoint model
                    if update % self.config.CHECKPOINT_INTERVAL == 0:
                        self.save_checkpoint(
                            f"ckpt.{count_checkpoints}.pth",
                            dict(step=count_steps),
                        )
                        count_checkpoints += 1

                profiling_wrapper.range_pop()  # train update

            self.envs.close()

habitat_baselines/rl/ddppo/algo/ddppo_trainer.py
cshizhe/onav_rim
import contextlib
import functools
import os
import signal
import socket
import subprocess
import threading
from os import path as osp
from typing import Any, Callable, Optional, Tuple, Union, overload

import ifcfg
import torch
from torch import distributed as distrib

from habitat import Config, logger

EXIT = threading.Event()
EXIT.clear()
REQUEUE = threading.Event()
REQUEUE.clear()
SAVE_STATE = threading.Event()
SAVE_STATE.clear()


# Default port to initialized the TCP store on
DEFAULT_PORT = 8738
DEFAULT_PORT_RANGE = 127
# Default address of world rank 0
DEFAULT_MAIN_ADDR = "127.0.0.1"

SLURM_JOBID = os.environ.get("SLURM_JOB_ID", None)
RESUME_STATE_BASE_NAME = ".habitat-resume-state"


def is_slurm_job() -> bool:
    return SLURM_JOBID is not None


def is_slurm_batch_job() -> bool:
    r"""Heuristic to determine if a slurm job is a batch job or not. Batch jobs
    will have a job name that is not a shell unless the user specifically set the job
    name to that of a shell. Interactive jobs have a shell name as their job name.
    """
    return is_slurm_job() and os.environ.get("SLURM_JOB_NAME", None) not in (
        None,
        "bash",
        "zsh",
        "fish",
        "tcsh",
        "sh",
        "interactive",
    )


def resume_state_filename(config: Config) -> str:
    fname = RESUME_STATE_BASE_NAME

    if is_slurm_job() and config.RL.preemption.append_slurm_job_id:
        fname += "-{}".format(SLURM_JOBID)

    return osp.join(config.CHECKPOINT_FOLDER, fname + ".pth")


@overload
def rank0_only() -> bool:
    ...


@overload
def rank0_only(fn: Callable) -> Callable:
    ...


def rank0_only(fn: Optional[Callable] = None) -> Union[Callable, bool]:
    r"""Helper function to only execute code if a process is world rank 0

    Can be used both as a function in an if statement,

    .. code:: py

        if rank0_only():
            ...

    or as a decorator,

    .. code:: py

        @rank0_only
        def fn_for_r0_only(...):
            ...

    :param fn: Function to wrap and only execute if the process is rank 0.
        If a process is rank 0, the function will be run and it's return value
        will be returned.  If a process is not rank 0, then the function will not
        be ran and :py:`None` will be returned.

    :return: The wrapped function if :p:`fn` is not :py:`None`, otherwise
        whether or not this process is rank 0
    """
    if fn is None:
        return (
            not torch.distributed.is_initialized()
            or torch.distributed.get_rank() == 0
        )

    @functools.wraps(fn)
    def _wrapper(*args, **kwargs):
        if rank0_only():
            return fn(*args, **kwargs)
        return None

    return _wrapper


def _ignore_handler(signum, frame):
    pass


def _clean_exit_handler(signum, frame):
    EXIT.set()
    print("Exiting cleanly", flush=True)


def _clean_exit_and_save_handler(signum, frame):
    EXIT.set()
    SAVE_STATE.set()
    print("Exiting cleanly and saving state", flush=True)


def _requeue_handler(signal, frame):
    REQUEUE.set()
    SAVE_STATE.set()
    EXIT.set()
    print("Got signal to requeue", flush=True)


def add_signal_handlers() -> None:
    signal.signal(signal.SIGCONT, _ignore_handler)
    signal.signal(signal.SIGINT, _clean_exit_handler)

    # SIGUSR2 can be sent to all processes to have them cleanup
    # and exit nicely.  This is nice to use with SLURM as scancel <job_id>
    # sets a 30 second timer for the job to exit, and it can take more than
    # 30 seconds for the job to cleanup and exit nicely.  When using NCCL,
    # forcing the job to exit without cleaning up can be bad.
    # scancel --signal SIGUSR2 <job_id> will set no such timer and will give
    # the job ample time to cleanup and exit.
    signal.signal(signal.SIGUSR2, _clean_exit_handler)

    # SLURM always sends SIGTERM so we can use this to save and exit
    signal.signal(signal.SIGTERM, _clean_exit_and_save_handler)

    signal.signal(signal.SIGUSR1, _requeue_handler)


@rank0_only
def save_resume_state(state: Any, filename_or_config: Union[Config, str]):
    r"""Saves the resume job state to the specified filename.
        This is useful when working with preemptable job partitions.

    :param state: The state to save
    :param filename_or_config: The filename of the saved state or the config to construct it.
    """
    if isinstance(filename_or_config, Config):
        filename = resume_state_filename(filename_or_config)
    else:
        filename = filename_or_config

    torch.save(state, filename)


def load_resume_state(filename_or_config: Union[Config, str]) -> Optional[Any]:
    r"""Loads the saved resume state

    :param filename_or_config: The filename of the saved state or the config to construct it.

    :return: The saved state if the file exists, else none
    """
    if isinstance(filename_or_config, Config):
        filename = resume_state_filename(filename_or_config)
    else:
        filename = filename_or_config

    if not osp.exists(filename):
        return None

    if rank0_only():
        logger.info(f"Loading resume state: {filename}")

    return torch.load(filename, map_location="cpu")


def requeue_job():
    r"""Requeues the job by calling ``scontrol requeue ${SLURM_JOBID}``"""
    if not is_slurm_batch_job():
        return

    if not REQUEUE.is_set():
        return

    if distrib.is_initialized():
        distrib.barrier()

    if rank0_only():
        logger.info(f"Requeueing job {SLURM_JOBID}")
        subprocess.check_call(["scontrol", "requeue", str(SLURM_JOBID)])


def get_ifname() -> str:
    return ifcfg.default_interface()["device"]


def get_distrib_size() -> Tuple[int, int, int]:
    # Check to see if we should parse from torch.distributed.launch
    if os.environ.get("LOCAL_RANK", None) is not None:
        local_rank = int(os.environ["LOCAL_RANK"])
        world_rank = int(os.environ["RANK"])
        world_size = int(os.environ["WORLD_SIZE"])
    # Else parse from SLURM is using SLURM
    elif os.environ.get("SLURM_JOBID", None) is not None:
        local_rank = int(os.environ["SLURM_LOCALID"])
        world_rank = int(os.environ["SLURM_PROCID"])
        world_size = int(os.environ["SLURM_NTASKS"])
    # Otherwise setup for just 1 process, this is nice for testing
    else:
        local_rank = 0
        world_rank = 0
        world_size = 1

    return local_rank, world_rank, world_size


def init_distrib_slurm(
    backend: str = "nccl",
) -> Tuple[int, torch.distributed.TCPStore]:  # type: ignore
    r"""Initializes torch.distributed by parsing environment variables set
        by SLURM when ``srun`` is used or by parsing environment variables set
        by torch.distributed.launch

    :param backend: Which torch.distributed backend to use

    :returns: Tuple of the local_rank (aka which GPU to use for this process)
        and the TCPStore used for the rendezvous
    """
    assert (
        torch.distributed.is_available()
    ), "torch.distributed must be available"

    if "GLOO_SOCKET_IFNAME" not in os.environ:
        os.environ["GLOO_SOCKET_IFNAME"] = get_ifname()

    if "NCCL_SOCKET_IFNAME" not in os.environ:
        os.environ["NCCL_SOCKET_IFNAME"] = get_ifname()

    local_rank, world_rank, world_size = get_distrib_size()

    main_port = int(os.environ.get("MAIN_PORT", DEFAULT_PORT))
    if SLURM_JOBID is not None:
        main_port += int(SLURM_JOBID) % int(
            os.environ.get("MAIN_PORT_RANGE", DEFAULT_PORT_RANGE)
        )
    main_addr = os.environ.get("MAIN_ADDR", DEFAULT_MAIN_ADDR)

    tcp_store = distrib.TCPStore(  # type: ignore
        main_addr, main_port, world_size, world_rank == 0
    )
    distrib.init_process_group(
        backend, store=tcp_store, rank=world_rank, world_size=world_size
    )

    return local_rank, tcp_store


def find_free_port() -> int:
    """
    Returns a free port on the system.
    Note that this can only be used to find a port for torch.distribted
    if it's called by a process on the node that will have
    world_rank == 0 and then all ranks are created. If you
    just called `find_free_port()` on each rank independently, every
    rank will have a different port!
    """
    with contextlib.closing(
        socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    ) as sock:
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        sock.bind(("localhost", 0))
        _, port = sock.getsockname()
        return port

habitat_baselines/rl/ddppo/ddp_utils.py
cshizhe/onav_rim
#!/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.

from typing import List, Optional, Type, Union, cast

from torch import Tensor
from torch import nn as nn
from torch.nn.modules.container import Sequential
from torch.nn.modules.conv import Conv2d


def conv3x3(
    in_planes: int, out_planes: int, stride: int = 1, groups: int = 1
) -> Conv2d:
    """3x3 convolution with padding"""
    return nn.Conv2d(
        in_planes,
        out_planes,
        kernel_size=3,
        stride=stride,
        padding=1,
        bias=False,
        groups=groups,
    )


def conv1x1(in_planes: int, out_planes: int, stride: int = 1) -> Conv2d:
    """1x1 convolution"""
    return nn.Conv2d(
        in_planes, out_planes, kernel_size=1, stride=stride, bias=False
    )


class BasicBlock(nn.Module):
    expansion = 1
    resneXt = False

    def __init__(
        self,
        inplanes,
        planes,
        ngroups,
        stride=1,
        downsample=None,
        cardinality=1,
    ):
        super(BasicBlock, self).__init__()
        self.convs = nn.Sequential(
            conv3x3(inplanes, planes, stride, groups=cardinality),
            nn.GroupNorm(ngroups, planes),
            nn.ReLU(True),
            conv3x3(planes, planes, groups=cardinality),
            nn.GroupNorm(ngroups, planes),
        )
        self.downsample = downsample
        self.relu = nn.ReLU(True)

    def forward(self, x):
        residual = x

        out = self.convs(x)

        if self.downsample is not None:
            residual = self.downsample(x)

        return self.relu(out + residual)


def _build_bottleneck_branch(
    inplanes: int,
    planes: int,
    ngroups: int,
    stride: int,
    expansion: int,
    groups: int = 1,
) -> Sequential:
    return nn.Sequential(
        conv1x1(inplanes, planes),
        nn.GroupNorm(ngroups, planes),
        nn.ReLU(True),
        conv3x3(planes, planes, stride, groups=groups),
        nn.GroupNorm(ngroups, planes),
        nn.ReLU(True),
        conv1x1(planes, planes * expansion),
        nn.GroupNorm(ngroups, planes * expansion),
    )


class SE(nn.Module):
    def __init__(self, planes, r=16):
        super().__init__()
        self.squeeze = nn.AdaptiveAvgPool2d(1)
        self.excite = nn.Sequential(
            nn.Linear(planes, int(planes / r)),
            nn.ReLU(True),
            nn.Linear(int(planes / r), planes),
            nn.Sigmoid(),
        )

    def forward(self, x):
        b, c, _, _ = x.size()
        x = self.squeeze(x)
        x = x.view(b, c)
        x = self.excite(x)

        return x.view(b, c, 1, 1)


def _build_se_branch(planes, r=16):
    return SE(planes, r)


class Bottleneck(nn.Module):
    expansion = 4
    resneXt = False

    def __init__(
        self,
        inplanes: int,
        planes: int,
        ngroups: int,
        stride: int = 1,
        downsample: Optional[Sequential] = None,
        cardinality: int = 1,
    ) -> None:
        super().__init__()
        self.convs = _build_bottleneck_branch(
            inplanes,
            planes,
            ngroups,
            stride,
            self.expansion,
            groups=cardinality,
        )
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample

    def _impl(self, x: Tensor) -> Tensor:
        identity = x

        out = self.convs(x)

        if self.downsample is not None:
            identity = self.downsample(x)

        return self.relu(out + identity)

    def forward(self, x: Tensor) -> Tensor:
        return self._impl(x)


class SEBottleneck(Bottleneck):
    def __init__(
        self,
        inplanes,
        planes,
        ngroups,
        stride=1,
        downsample=None,
        cardinality=1,
    ):
        super().__init__(
            inplanes, planes, ngroups, stride, downsample, cardinality
        )

        self.se = _build_se_branch(planes * self.expansion)

    def _impl(self, x):
        identity = x

        out = self.convs(x)
        out = self.se(out) * out

        if self.downsample is not None:
            identity = self.downsample(x)

        return self.relu(out + identity)


class SEResNeXtBottleneck(SEBottleneck):
    expansion = 2
    resneXt = True


class ResNeXtBottleneck(Bottleneck):
    expansion = 2
    resneXt = True


Block = Union[Type[Bottleneck], Type[BasicBlock]]


class ResNet(nn.Module):
    def __init__(
        self,
        in_channels: int,
        base_planes: int,
        ngroups: int,
        block: Block,
        layers: List[int],
        cardinality: int = 1,
    ) -> None:
        super(ResNet, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(
                in_channels,
                base_planes,
                kernel_size=7,
                stride=2,
                padding=3,
                bias=False,
            ),
            nn.GroupNorm(ngroups, base_planes),
            nn.ReLU(True),
        )
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.cardinality = cardinality

        self.inplanes = base_planes
        if block.resneXt:
            base_planes *= 2

        self.layer1 = self._make_layer(block, ngroups, base_planes, layers[0])
        self.layer2 = self._make_layer(
            block, ngroups, base_planes * 2, layers[1], stride=2
        )
        self.layer3 = self._make_layer(
            block, ngroups, base_planes * 2 * 2, layers[2], stride=2
        )
        self.layer4 = self._make_layer(
            block, ngroups, base_planes * 2 * 2 * 2, layers[3], stride=2
        )

        self.final_channels = self.inplanes
        self.final_spatial_compress = 1.0 / (2 ** 5)

    def _make_layer(
        self,
        block: Block,
        ngroups: int,
        planes: int,
        blocks: int,
        stride: int = 1,
    ) -> Sequential:
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                conv1x1(self.inplanes, planes * block.expansion, stride),
                nn.GroupNorm(ngroups, planes * block.expansion),
            )

        layers = []
        layers.append(
            block(
                self.inplanes,
                planes,
                ngroups,
                stride,
                downsample,
                cardinality=self.cardinality,
            )
        )
        self.inplanes = planes * block.expansion
        for _i in range(1, blocks):
            layers.append(block(self.inplanes, planes, ngroups))

        return nn.Sequential(*layers)

    def forward(self, x) -> Tensor:
        x = self.conv1(x)
        x = self.maxpool(x)
        x = cast(Tensor, x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        return x


def resnet18(in_channels, base_planes, ngroups, dropout_prob=0.0):
    model = ResNet(in_channels, base_planes, ngroups, BasicBlock, [2, 2, 2, 2])

    return model


def resnet50(in_channels: int, base_planes: int, ngroups: int, dropout_prob=0.0) -> ResNet:
    model = ResNet(in_channels, base_planes, ngroups, Bottleneck, [3, 4, 6, 3])

    return model


def resneXt50(in_channels, base_planes, ngroups):
    model = ResNet(
        in_channels,
        base_planes,
        ngroups,
        ResNeXtBottleneck,
        [3, 4, 6, 3],
        cardinality=int(base_planes / 2),
    )

    return model


def se_resnet50(in_channels, base_planes, ngroups, dropout_prob=0.0):
    model = ResNet(
        in_channels, base_planes, ngroups, SEBottleneck, [3, 4, 6, 3]
    )

    return model


def se_resneXt50(in_channels, base_planes, ngroups):
    model = ResNet(
        in_channels,
        base_planes,
        ngroups,
        SEResNeXtBottleneck,
        [3, 4, 6, 3],
        cardinality=int(base_planes / 2),
    )

    return model


def se_resneXt101(in_channels, base_planes, ngroups):
    model = ResNet(
        in_channels,
        base_planes,
        ngroups,
        SEResNeXtBottleneck,
        [3, 4, 23, 3],
        cardinality=int(base_planes / 2),
    )

    return model

habitat_baselines/rl/ddppo/policy/resnet.py
cshizhe/onav_rim
import torch.nn as nn
import math


__all__ = ['ResNet', 'resnet50', 'resnet101']


def conv3x3(in_planes, out_planes, stride=1):
    "3x3 convolution with padding"
    return nn.Conv2d(
        in_planes,
        out_planes,
        kernel_size=3,
        stride=stride,
        padding=1,
        bias=False
    )

class BasicBlock(nn.Module):
    expansion = 1

    def __init__(
        self,
        inplanes,
        planes,
        ngroups,
        stride=1,
        downsample=None,
        dropout_prob=0.0,
    ):
        super(BasicBlock, self).__init__()
        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = nn.GroupNorm(ngroups, planes)
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = nn.GroupNorm(ngroups, planes)
        self.downsample = downsample
        self.relu = nn.ReLU(True)

        gn_init(self.bn1)
        gn_init(self.bn2, zero_init=True)

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        return self.relu(out + residual)


class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, inplanes, planes, ngroups, stride=1, downsample=None, dropout_prob=0.0):
        super(Bottleneck, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
        self.bn1 = nn.GroupNorm(ngroups, planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
                               padding=1, bias=False)
        self.bn2 = nn.GroupNorm(ngroups, planes)
        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
        self.bn3 = nn.GroupNorm(ngroups, planes * 4)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride
        self.dropout = nn.Dropout(p=dropout_prob)

        gn_init(self.bn1)
        gn_init(self.bn2)
        gn_init(self.bn3, zero_init=True)

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.dropout(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)
        out = self.dropout(out)

        out = self.conv3(out)
        out = self.bn3(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.dropout(self.relu(out))

        return out


def conv2d_init(m):
    assert isinstance(m, nn.Conv2d)
    n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
    m.weight.data.normal_(0, math.sqrt(2. / n))

def gn_init(m, zero_init=False):
    assert isinstance(m, nn.GroupNorm)
    m.weight.data.fill_(0. if zero_init else 1.)
    m.bias.data.zero_()


class ResNet(nn.Module):

    def __init__(self, in_channels, base_planes, ngroups, block, layers, dropout_prob=0.0):
        self.inplanes = base_planes
        super(ResNet, self).__init__()
        self.dropout_prob = dropout_prob
        self.conv1 = nn.Conv2d(in_channels, base_planes, kernel_size=7, stride=2, padding=3,
                               bias=False)
        self.bn1 = nn.GroupNorm(ngroups, base_planes)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.layer1 = self._make_layer(block, ngroups, base_planes, layers[0])
        self.layer2 = self._make_layer(block, ngroups, base_planes*2, layers[1], stride=2)
        self.layer3 = self._make_layer(block, ngroups, base_planes*4, layers[2], stride=2)
        self.layer4 = self._make_layer(block, ngroups, base_planes*8, layers[3], stride=2)
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        # self.fc = nn.Linear(512 * block.expansion, num_classes)
        self.final_channels = self.inplanes
        self.final_spatial_compress = 1.0 / (2 ** 5)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                conv2d_init(m)
        gn_init(self.bn1)

    def _make_layer(self, block, ngroups, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=False),
                nn.GroupNorm(ngroups, planes * block.expansion),
            )
            m = downsample[1]
            assert isinstance(m, nn.GroupNorm)
            gn_init(m)

        layers = []
        layers.append(block(self.inplanes, planes, ngroups, stride, downsample, dropout_prob=self.dropout_prob))
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes, ngroups, dropout_prob=self.dropout_prob))

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        # x = self.avgpool(x)
        # x = x.view(x.size(0), -1)
        # x = self.fc(x)

        return x

def resnet18(in_channels, base_planes, ngroups, dropout_prob=0.0):
    model = ResNet(in_channels, base_planes, ngroups, BasicBlock, [2, 2, 2, 2], dropout_prob=dropout_prob)
    # model = ResNet(in_channels, base_planes, ngroups, Bottleneck, [2, 2, 2, 2])
    return model


def resnet50(in_channels, base_planes, ngroups, dropout_prob=0.0):
    model = ResNet(in_channels, base_planes, ngroups, Bottleneck, [3, 4, 6, 3], dropout_prob=dropout_prob)
    return model

def resnet101(in_channels, base_planes, ngroups, dropout_prob=0.0):
    model = ResNet(in_channels, base_planes, ngroups, Bottleneck, [3, 4, 23, 3], dropout_prob=dropout_prob)
    return model
habitat_baselines/rl/ddppo/policy/resnet_gn.py
cshizhe/onav_rim
#!/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.


from typing import Dict, Tuple

import math
import numpy as np
import torch
from gym import spaces
from torch import nn as nn
from torch.nn import functional as F

from habitat import logger
from habitat.config import Config
from habitat.tasks.nav.nav import (
    EpisodicCompassSensor,
    EpisodicGPSSensor,
    HeadingSensor,
    ImageGoalSensor,
    IntegratedPointGoalGPSAndCompassSensor,
    PointGoalSensor,
    ProximitySensor,
)
from habitat.tasks.nav.object_nav_task import ObjectGoalSensor
from habitat_baselines.common.baseline_registry import baseline_registry
from habitat_baselines.rl.ddppo.policy import resnet
from habitat_baselines.rl.ddppo.policy.running_mean_and_var import (
    RunningMeanAndVar,
)
from habitat_baselines.rl.models.rnn_state_encoder import RNNStateEncoder
from habitat_baselines.rl.ppo import Net, Policy
from habitat_baselines.utils.common import Flatten


@baseline_registry.register_policy
class PointNavResNetPolicy(Policy):
    def __init__(
        self,
        observation_space: spaces.Dict,
        action_space,
        hidden_size: int = 512,
        num_recurrent_layers: int = 2,
        rnn_type: str = "LSTM",
        resnet_baseplanes: int = 32,
        backbone: str = "resnet50",
        normalize_visual_inputs: bool = False,
        force_blind_policy: bool = False,
        **kwargs
    ):
        super().__init__(
            PointNavResNetNet(
                observation_space=observation_space,
                action_space=action_space,
                hidden_size=hidden_size,
                num_recurrent_layers=num_recurrent_layers,
                rnn_type=rnn_type,
                backbone=backbone,
                resnet_baseplanes=resnet_baseplanes,
                normalize_visual_inputs=normalize_visual_inputs,
                force_blind_policy=force_blind_policy,
            ),
            action_space.n,
        )

    @classmethod
    def from_config(
        cls, config: Config, observation_space: spaces.Dict, action_space
    ):
        return cls(
            observation_space=observation_space,
            action_space=action_space,
            model_config=config.MODEL,
        )


class ResNetEncoder(nn.Module):
    def __init__(
        self,
        observation_space: spaces.Dict,
        baseplanes: int = 32,
        ngroups: int = 32,
        spatial_size: int = 128,
        make_backbone=None,
        normalize_visual_inputs: bool = False,
        sem_embedding_size=4,
        dropout_prob: float = 0.0
    ):
        super().__init__()

        if "rgb" in observation_space.spaces:
            self._frame_size = tuple(observation_space.spaces["rgb"].shape[:2])
            self._n_input_rgb = observation_space.spaces["rgb"].shape[2]
            # spatial_size = observation_space.spaces["rgb"].shape[:2] // 2
            spatial_size = observation_space.spaces["rgb"].shape[:2]
        else:
            self._n_input_rgb = 0

        if "depth" in observation_space.spaces:
            self._frame_size = tuple(observation_space.spaces["depth"].shape[:2])
            self._n_input_depth = observation_space.spaces["depth"].shape[2]
            # spatial_size = observation_space.spaces["depth"].shape[:2] // 2
            spatial_size = observation_space.spaces["depth"].shape[:2]
        else:
            self._n_input_depth = 0
        
        if "semantic" in observation_space.spaces:
            self._frame_size = tuple(observation_space.spaces["semantic"].shape[:2])
            self._n_input_semantics = sem_embedding_size # observation_space.spaces["semantic"].shape[2]
        else:
            self._n_input_semantics = 0
        
        if self._frame_size == (256, 256):
            spatial_size = (128, 128)
        elif self._frame_size == (240, 320):
            spatial_size = (120, 108)
        elif self._frame_size == (480, 640):
            spatial_size = (120, 128)
        elif self._frame_size == (224, 224):
            spatial_size = (112, 112)

        if normalize_visual_inputs:
            self.running_mean_and_var: nn.Module = RunningMeanAndVar(
                self._n_input_depth + self._n_input_rgb
            )
        else:
            self.running_mean_and_var = nn.Sequential()

        if not self.is_blind:
            input_channels = self._n_input_depth + self._n_input_rgb + self._n_input_semantics
            self.backbone = make_backbone(input_channels, baseplanes, ngroups, dropout_prob=dropout_prob)

            final_spatial = np.array([math.ceil(
                d * self.backbone.final_spatial_compress
            ) for d in spatial_size])
            after_compression_flat_size = 2048
            num_compression_channels = int(
                round(after_compression_flat_size / np.prod(final_spatial))
            )
            self.compression = nn.Sequential(
                nn.Conv2d(
                    self.backbone.final_channels,
                    num_compression_channels,
                    kernel_size=3,
                    padding=1,
                    bias=False,
                ),
                nn.GroupNorm(1, num_compression_channels),
                nn.ReLU(True),
            )

            self.output_shape = (
                num_compression_channels,
                final_spatial[0],
                final_spatial[1],
            )

    @property
    def is_blind(self):
        return self._n_input_rgb + self._n_input_depth + self._n_input_semantics == 0

    def layer_init(self):
        for layer in self.modules():
            if isinstance(layer, (nn.Conv2d, nn.Linear)):
                nn.init.kaiming_normal_(
                    layer.weight, nn.init.calculate_gain("relu")
                )
                if layer.bias is not None:
                    nn.init.constant_(layer.bias, val=0)

    def forward(self, observations: Dict[str, torch.Tensor]) -> torch.Tensor:  # type: ignore
        if self.is_blind:
            return None

        cnn_input = []
        if self._n_input_rgb > 0:
            rgb_observations = observations["rgb"]
            # permute tensor to dimension [BATCH x CHANNEL x HEIGHT X WIDTH]
            rgb_observations = rgb_observations.permute(0, 3, 1, 2)
            rgb_observations = rgb_observations / 255.0  # normalize RGB
            cnn_input.append(rgb_observations)

        if self._n_input_depth > 0:
            depth_observations = observations["depth"]

            # permute tensor to dimension [BATCH x CHANNEL x HEIGHT X WIDTH]
            depth_observations = depth_observations.permute(0, 3, 1, 2)

            cnn_input.append(depth_observations)

        if self._n_input_semantics > 0:
            semantic_observations = observations["semantic"]

            # permute tensor to dimension [BATCH x CHANNEL x HEIGHT X WIDTH]
            semantic_observations = semantic_observations.permute(0, 3, 1, 2)

            cnn_input.append(semantic_observations)

        x = torch.cat(cnn_input, dim=1)
        if self._frame_size == (256, 256):
            x = F.avg_pool2d(x, 2)
        elif self._frame_size == (240, 320):
            x = F.avg_pool2d(x, (2, 3), padding=(0, 1)) # 240 x 324 -> 120 x 108
        elif self._frame_size == (480, 640):
            x = F.avg_pool2d(x, (4, 5))
        elif self._frame_size == (224, 224):
            x = F.avg_pool2d(x, 2)

        x = self.running_mean_and_var(x)
        x = self.backbone(x)
        x = self.compression(x)
        return x


class PointNavResNetNet(Net):
    """Network which passes the input image through CNN and concatenates
    goal vector with CNN's output and passes that through RNN.
    """

    def __init__(
        self,
        observation_space: spaces.Dict,
        action_space,
        hidden_size: int,
        num_recurrent_layers: int,
        rnn_type: str,
        backbone,
        resnet_baseplanes,
        normalize_visual_inputs: bool,
        force_blind_policy: bool = False,
    ):
        super().__init__()

        self.prev_action_embedding = nn.Embedding(action_space.n + 1, 32)
        self._n_prev_action = 32
        rnn_input_size = self._n_prev_action

        if (
            IntegratedPointGoalGPSAndCompassSensor.cls_uuid
            in observation_space.spaces
        ):
            n_input_goal = (
                observation_space.spaces[
                    IntegratedPointGoalGPSAndCompassSensor.cls_uuid
                ].shape[0]
                + 1
            )
            self.tgt_embeding = nn.Linear(n_input_goal, 32)
            rnn_input_size += 32

        if ObjectGoalSensor.cls_uuid in observation_space.spaces:
            self._n_object_categories = (
                int(
                    observation_space.spaces[ObjectGoalSensor.cls_uuid].high[0]
                )
                + 1
            )
            self.obj_categories_embedding = nn.Embedding(
                self._n_object_categories, 32
            )
            rnn_input_size += 32

        if EpisodicGPSSensor.cls_uuid in observation_space.spaces:
            input_gps_dim = observation_space.spaces[
                EpisodicGPSSensor.cls_uuid
            ].shape[0]
            self.gps_embedding = nn.Linear(input_gps_dim, 32)
            rnn_input_size += 32

        if PointGoalSensor.cls_uuid in observation_space.spaces:
            input_pointgoal_dim = observation_space.spaces[
                PointGoalSensor.cls_uuid
            ].shape[0]
            self.pointgoal_embedding = nn.Linear(input_pointgoal_dim, 32)
            rnn_input_size += 32

        if HeadingSensor.cls_uuid in observation_space.spaces:
            input_heading_dim = (
                observation_space.spaces[HeadingSensor.cls_uuid].shape[0] + 1
            )
            assert input_heading_dim == 2, "Expected heading with 2D rotation."
            self.heading_embedding = nn.Linear(input_heading_dim, 32)
            rnn_input_size += 32

        if ProximitySensor.cls_uuid in observation_space.spaces:
            input_proximity_dim = observation_space.spaces[
                ProximitySensor.cls_uuid
            ].shape[0]
            self.proximity_embedding = nn.Linear(input_proximity_dim, 32)
            rnn_input_size += 32

        if EpisodicCompassSensor.cls_uuid in observation_space.spaces:
            assert (
                observation_space.spaces[EpisodicCompassSensor.cls_uuid].shape[
                    0
                ]
                == 1
            ), "Expected compass with 2D rotation."
            input_compass_dim = 2  # cos and sin of the angle
            self.compass_embedding = nn.Linear(input_compass_dim, 32)
            rnn_input_size += 32

        if ImageGoalSensor.cls_uuid in observation_space.spaces:
            goal_observation_space = spaces.Dict(
                {"rgb": observation_space.spaces[ImageGoalSensor.cls_uuid]}
            )
            self.goal_visual_encoder = ResNetEncoder(
                goal_observation_space,
                baseplanes=resnet_baseplanes,
                ngroups=resnet_baseplanes // 2,
                make_backbone=getattr(resnet, backbone),
                normalize_visual_inputs=normalize_visual_inputs,
            )

            self.goal_visual_fc = nn.Sequential(
                Flatten(),
                nn.Linear(
                    np.prod(self.goal_visual_encoder.output_shape), hidden_size
                ),
                nn.ReLU(True),
            )

            rnn_input_size += hidden_size

        self._hidden_size = hidden_size

        self.visual_encoder = ResNetEncoder(
            observation_space if not force_blind_policy else spaces.Dict({}),
            baseplanes=resnet_baseplanes,
            ngroups=resnet_baseplanes // 2,
            make_backbone=getattr(resnet, backbone),
            normalize_visual_inputs=normalize_visual_inputs,
        )

        if not self.visual_encoder.is_blind:
            self.visual_fc = nn.Sequential(
                Flatten(),
                nn.Linear(
                    np.prod(self.visual_encoder.output_shape), hidden_size
                ),
                nn.ReLU(True),
            )

        self.state_encoder = RNNStateEncoder(
            (0 if self.is_blind else self._hidden_size) + rnn_input_size,
            self._hidden_size,
            rnn_type=rnn_type,
            num_layers=num_recurrent_layers,
        )

        self.train()

    @property
    def output_size(self):
        return self._hidden_size

    @property
    def is_blind(self):
        return self.visual_encoder.is_blind

    @property
    def num_recurrent_layers(self):
        return self.state_encoder.num_recurrent_layers

    def forward(
        self,
        observations: Dict[str, torch.Tensor],
        rnn_hidden_states,
        prev_actions,
        masks,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        x = []
        if not self.is_blind:
            if "visual_features" in observations:
                visual_feats = observations["visual_features"]
            else:
                visual_feats = self.visual_encoder(observations)

            visual_feats = self.visual_fc(visual_feats)
            x.append(visual_feats)

        if IntegratedPointGoalGPSAndCompassSensor.cls_uuid in observations:
            goal_observations = observations[
                IntegratedPointGoalGPSAndCompassSensor.cls_uuid
            ]
            goal_observations = torch.stack(
                [
                    goal_observations[:, 0],
                    torch.cos(-goal_observations[:, 1]),
                    torch.sin(-goal_observations[:, 1]),
                ],
                -1,
            )

            x.append(self.tgt_embeding(goal_observations))

        if PointGoalSensor.cls_uuid in observations:
            goal_observations = observations[PointGoalSensor.cls_uuid]
            x.append(self.pointgoal_embedding(goal_observations))

        if ProximitySensor.cls_uuid in observations:
            sensor_observations = observations[ProximitySensor.cls_uuid]
            x.append(self.proximity_embedding(sensor_observations))

        if HeadingSensor.cls_uuid in observations:
            sensor_observations = observations[HeadingSensor.cls_uuid]
            sensor_observations = torch.stack(
                [
                    torch.cos(sensor_observations[0]),
                    torch.sin(sensor_observations[0]),
                ],
                -1,
            )
            x.append(self.heading_embedding(sensor_observations))

        if ObjectGoalSensor.cls_uuid in observations:
            object_goal = observations[ObjectGoalSensor.cls_uuid].long()
            x.append(self.obj_categories_embedding(object_goal).squeeze(dim=1))

        if EpisodicCompassSensor.cls_uuid in observations:
            compass_observations = torch.stack(
                [
                    torch.cos(observations[EpisodicCompassSensor.cls_uuid]),
                    torch.sin(observations[EpisodicCompassSensor.cls_uuid]),
                ],
                -1,
            )
            x.append(
                self.compass_embedding(compass_observations.squeeze(dim=1))
            )

        if EpisodicGPSSensor.cls_uuid in observations:
            x.append(
                self.gps_embedding(observations[EpisodicGPSSensor.cls_uuid])
            )

        if ImageGoalSensor.cls_uuid in observations:
            goal_image = observations[ImageGoalSensor.cls_uuid]
            goal_output = self.goal_visual_encoder({"rgb": goal_image})
            x.append(self.goal_visual_fc(goal_output))

        prev_actions = self.prev_action_embedding(
            ((prev_actions.float() + 1) * masks).long().squeeze(dim=-1)
        )
        x.append(prev_actions)

        out = torch.cat(x, dim=1)
        out, rnn_hidden_states = self.state_encoder(
            out, rnn_hidden_states, masks
        )

        return out, rnn_hidden_states

habitat_baselines/rl/ddppo/policy/resnet_policy.py
cshizhe/onav_rim
import torch
from torch import nn as nn


class RNNStateEncoder(nn.Module):
    def __init__(
        self,
        input_size: int,
        hidden_size: int,
        num_layers: int = 1,
        rnn_type: str = "GRU",
    ):
        r"""An RNN for encoding the state in RL.

        Supports masking the hidden state during various timesteps in the forward lass

        Args:
            input_size: The input size of the RNN
            hidden_size: The hidden size
            num_layers: The number of recurrent layers
            rnn_type: The RNN cell type.  Must be GRU or LSTM
        """

        super().__init__()
        self._num_recurrent_layers = num_layers
        self._rnn_type = rnn_type

        self.rnn = getattr(nn, rnn_type)(
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
        )

        self.layer_init()

    def layer_init(self):
        for name, param in self.rnn.named_parameters():
            if "weight" in name:
                nn.init.orthogonal_(param)
            elif "bias" in name:
                nn.init.constant_(param, 0)

    @property
    def num_recurrent_layers(self):
        return self._num_recurrent_layers * (
            2 if "LSTM" in self._rnn_type else 1
        )

    def _pack_hidden(self, hidden_states):
        if "LSTM" in self._rnn_type:
            hidden_states = torch.cat(
                [hidden_states[0], hidden_states[1]], dim=0
            )

        return hidden_states

    def _unpack_hidden(self, hidden_states):
        if "LSTM" in self._rnn_type:
            hidden_states = (
                hidden_states[0 : self._num_recurrent_layers],
                hidden_states[self._num_recurrent_layers :],
            )

        return hidden_states

    def _mask_hidden(self, hidden_states, masks):
        if isinstance(hidden_states, tuple):
            hidden_states = tuple(v * masks for v in hidden_states)
        else:
            hidden_states = masks * hidden_states

        return hidden_states

    def single_forward(self, x, hidden_states, masks):
        r"""Forward for a non-sequence input"""
        hidden_states = self._unpack_hidden(hidden_states)
        x, hidden_states = self.rnn(
            x.unsqueeze(0),
            self._mask_hidden(hidden_states, masks.unsqueeze(0)),
        )
        x = x.squeeze(0)
        hidden_states = self._pack_hidden(hidden_states)
        return x, hidden_states

    def episode_forward(self, x, init_hidden_states):
        r'''Forward a mini-batch of episodes
        Args:
            x: (N, T, D)
            hidden_states: zero state
        '''
        outputs, hidden_states = self.rnn(x.permute(1, 0, 2), init_hidden_states)
        outputs = outputs.permute(1, 0, 2)
        return outputs, hidden_states

    def seq_forward(self, x, hidden_states, masks):
        r"""Forward for a sequence of length T

        Args:
            x: (T, N, -1) Tensor that has been flattened to (T * N, -1)
            hidden_states: The starting hidden state.
            masks: The masks to be applied to hidden state at every timestep.
                A (T, N) tensor flatten to (T * N)
        """
        # x is a (T, N, -1) tensor flattened to (T * N, -1)
        n = hidden_states.size(1)
        t = int(x.size(0) / n)

        # unflatten
        x = x.view(t, n, x.size(1))
        masks = masks.view(t, n)

        # steps in sequence which have zero for any agent. Assume t=0 has
        # a zero in it.
        # because multiple episodes can exist (0 is the start of a new episode)
        has_zeros = torch.nonzero((masks[1:] == 0.0).any(dim=-1)).squeeze().cpu()

        # +1 to correct the masks[1:]
        if has_zeros.dim() == 0:
            has_zeros = [has_zeros.item() + 1]  # handle scalar
        else:
            has_zeros = (has_zeros + 1).numpy().tolist()

        # add t=0 and t=T to the list
        has_zeros = [0] + has_zeros + [t]

        hidden_states = self._unpack_hidden(hidden_states)  # (nlayer, N, d)
        outputs = []
        for i in range(len(has_zeros) - 1):
            # process steps that don't have any zeros in masks together
            start_idx = has_zeros[i]
            end_idx = has_zeros[i + 1]

            rnn_scores, hidden_states = self.rnn(
                x[start_idx:end_idx],
                self._mask_hidden(
                    hidden_states, masks[start_idx].view(1, -1, 1)
                ),
            )

            outputs.append(rnn_scores)

        # x is a (T, N, -1) tensor
        x = torch.cat(outputs, dim=0)
        x = x.view(t * n, -1)  # flatten

        hidden_states = self._pack_hidden(hidden_states)
        return x, hidden_states

    def forward(self, x, hidden_states, masks=None, seq_single_episode=False):
        # if x.size(0) == hidden_states.size(1):
        #     return self.single_forward(x, hidden_states, masks)
        # else:
        if seq_single_episode:
            return self.episode_forward(x, hidden_states)
        else:
            return self.seq_forward(x, hidden_states, masks)

habitat_baselines/rl/models/rnn_state_encoder.py
cshizhe/onav_rim
#!/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 abc

import torch
from gym import spaces
from torch import nn as nn

from habitat.config import Config
from habitat.tasks.nav.nav import (
    ImageGoalSensor,
    IntegratedPointGoalGPSAndCompassSensor,
    PointGoalSensor,
)
from habitat_baselines.common.baseline_registry import baseline_registry
from habitat_baselines.rl.models.rnn_state_encoder import RNNStateEncoder
from habitat_baselines.rl.models.simple_cnn import SimpleCNN
from habitat_baselines.utils.common import CategoricalNet


class Policy(nn.Module, metaclass=abc.ABCMeta):
    def __init__(self, net, dim_actions, no_critic=False, clf_net=None):
        super().__init__()
        self.net = net
        self.dim_actions = dim_actions
        self.no_critic = no_critic

        self.action_distribution = CategoricalNet(
            self.net.output_size, self.dim_actions, clf_net=clf_net
        )
        if not self.no_critic:
            self.critic = CriticHead(self.net.output_size)

    def forward(self, *x, **kwargs):
        features, rnn_hidden_states = self.net(
            *x, **kwargs
        )
        distribution = self.action_distribution(features)

        return distribution.logits, rnn_hidden_states

    def act(
        self,
        observations,
        rnn_hidden_states,
        prev_actions,
        masks,
        deterministic=False,
    ):
        features, rnn_hidden_states = self.net(
            observations, rnn_hidden_states, prev_actions, masks
        )
        distribution = self.action_distribution(features)

        if deterministic:
            action = distribution.mode()
        else:
            action = distribution.sample()
        action_log_probs = distribution.log_probs(action)

        if self.no_critic:
            return action, rnn_hidden_states            

        value = self.critic(features)
        return value, action, action_log_probs, rnn_hidden_states

    def get_value(self, observations, rnn_hidden_states, prev_actions, masks):
        features, _ = self.net(
            observations, rnn_hidden_states, prev_actions, masks
        )
        return self.critic(features)

    def evaluate_actions(
        self, observations, rnn_hidden_states, prev_actions, masks, action
    ):
        features, rnn_hidden_states = self.net(
            observations, rnn_hidden_states, prev_actions, masks
        )
        distribution = self.action_distribution(features)
        value = self.critic(features)

        action_log_probs = distribution.log_probs(action)
        distribution_entropy = distribution.entropy().mean()

        return value, action_log_probs, distribution_entropy, rnn_hidden_states

    @classmethod
    @abc.abstractmethod
    def from_config(cls, config, observation_space, action_space):
        pass


class CriticHead(nn.Module):
    def __init__(self, input_size):
        super().__init__()
        self.fc = nn.Linear(input_size, 1)
        nn.init.orthogonal_(self.fc.weight)
        nn.init.constant_(self.fc.bias, 0)

    def forward(self, x):
        return self.fc(x)


@baseline_registry.register_policy
class PointNavBaselinePolicy(Policy):
    def __init__(
        self,
        observation_space: spaces.Dict,
        action_space,
        hidden_size: int = 512,
        **kwargs
    ):
        super().__init__(
            PointNavBaselineNet(  # type: ignore
                observation_space=observation_space,
                hidden_size=hidden_size,
                **kwargs,
            ),
            action_space.n,
        )

    @classmethod
    def from_config(
        cls, config: Config, observation_space: spaces.Dict, action_space
    ):
        return cls(
            observation_space=observation_space,
            action_space=action_space,
            hidden_size=config.RL.PPO.hidden_size,
        )


class Net(nn.Module, metaclass=abc.ABCMeta):
    @abc.abstractmethod
    def forward(self, observations, rnn_hidden_states, prev_actions, masks):
        pass

    @property
    @abc.abstractmethod
    def output_size(self):
        pass

    # @property
    # @abc.abstractmethod
    # def num_recurrent_layers(self):
    #     pass

    @property
    @abc.abstractmethod
    def is_blind(self):
        pass


class PointNavBaselineNet(Net):
    r"""Network which passes the input image through CNN and concatenates
    goal vector with CNN's output and passes that through RNN.
    """

    def __init__(
        self,
        observation_space: spaces.Dict,
        hidden_size: int,
    ):
        super().__init__()

        if (
            IntegratedPointGoalGPSAndCompassSensor.cls_uuid
            in observation_space.spaces
        ):
            self._n_input_goal = observation_space.spaces[
                IntegratedPointGoalGPSAndCompassSensor.cls_uuid
            ].shape[0]
        elif PointGoalSensor.cls_uuid in observation_space.spaces:
            self._n_input_goal = observation_space.spaces[
                PointGoalSensor.cls_uuid
            ].shape[0]
        elif ImageGoalSensor.cls_uuid in observation_space.spaces:
            goal_observation_space = spaces.Dict(
                {"rgb": observation_space.spaces[ImageGoalSensor.cls_uuid]}
            )
            self.goal_visual_encoder = SimpleCNN(
                goal_observation_space, hidden_size
            )
            self._n_input_goal = hidden_size

        self._hidden_size = hidden_size

        self.visual_encoder = SimpleCNN(observation_space, hidden_size)

        self.state_encoder = RNNStateEncoder(
            (0 if self.is_blind else self._hidden_size) + self._n_input_goal,
            self._hidden_size,
        )

        self.train()

    @property
    def output_size(self):
        return self._hidden_size

    @property
    def is_blind(self):
        return self.visual_encoder.is_blind

    @property
    def num_recurrent_layers(self):
        return self.state_encoder.num_recurrent_layers

    def forward(self, observations, rnn_hidden_states, prev_actions, masks):
        if IntegratedPointGoalGPSAndCompassSensor.cls_uuid in observations:
            target_encoding = observations[
                IntegratedPointGoalGPSAndCompassSensor.cls_uuid
            ]

        elif PointGoalSensor.cls_uuid in observations:
            target_encoding = observations[PointGoalSensor.cls_uuid]
        elif ImageGoalSensor.cls_uuid in observations:
            image_goal = observations[ImageGoalSensor.cls_uuid]
            target_encoding = self.goal_visual_encoder({"rgb": image_goal})

        x = [target_encoding]

        if not self.is_blind:
            perception_embed = self.visual_encoder(observations)
            x = [perception_embed] + x

        x_out = torch.cat(x, dim=1)
        x_out, rnn_hidden_states = self.state_encoder(
            x_out, rnn_hidden_states, masks
        )

        return x_out, rnn_hidden_states

habitat_baselines/rl/ppo/policy.py
cshizhe/onav_rim
#!/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
import time
from collections import defaultdict, deque
from typing import Any, DefaultDict, Dict, List, Optional

import numpy as np
import torch
import tqdm
from torch.optim.lr_scheduler import LambdaLR

from habitat import Config, logger
from habitat.utils import profiling_wrapper
from habitat.utils.visualizations.utils import observations_to_image
from habitat_baselines.common.base_trainer import BaseRLTrainer
from habitat_baselines.common.baseline_registry import baseline_registry
from habitat_baselines.common.environments import get_env_class
from habitat_baselines.common.obs_transformers import (
    apply_obs_transforms_batch,
    apply_obs_transforms_obs_space,
    get_active_obs_transforms,
)
from habitat_baselines.common.rollout_storage import RolloutStorage
from habitat_baselines.common.tensorboard_utils import TensorboardWriter
from habitat_baselines.rl.ppo import PPO
from habitat_baselines.utils.common import (
    batch_obs,
    generate_video,
    linear_decay,
)
from habitat_baselines.utils.env_utils import construct_envs


@baseline_registry.register_trainer(name="ppo")
class PPOTrainer(BaseRLTrainer):
    r"""Trainer class for PPO algorithm
    Paper: https://arxiv.org/abs/1707.06347.
    """
    supported_tasks = ["Nav-v0"]

    def __init__(self, config=None):
        super().__init__(config)
        self.actor_critic = None
        self.agent = None
        self.envs = None
        self.obs_transforms = []
        if config is not None:
            logger.info(f"config: {config}")

        self._static_encoder = False
        self._encoder = None

    def _setup_actor_critic_agent(self, ppo_cfg: Config) -> None:
        r"""Sets up actor critic and agent for PPO.

        Args:
            ppo_cfg: config node with relevant params

        Returns:
            None
        """
        logger.add_filehandler(self.config.LOG_FILE)

        policy = baseline_registry.get_policy(self.config.RL.POLICY.name)
        observation_space = self.envs.observation_spaces[0]
        self.obs_transforms = get_active_obs_transforms(self.config)
        observation_space = apply_obs_transforms_obs_space(
            observation_space, self.obs_transforms
        )
        self.obs_space = observation_space
        self.actor_critic = policy.from_config(
            self.config, observation_space, self.envs.action_spaces[0]
        )
        self.actor_critic.to(self.device)

        self.agent = PPO(
            actor_critic=self.actor_critic,
            clip_param=ppo_cfg.clip_param,
            ppo_epoch=ppo_cfg.ppo_epoch,
            num_mini_batch=ppo_cfg.num_mini_batch,
            value_loss_coef=ppo_cfg.value_loss_coef,
            entropy_coef=ppo_cfg.entropy_coef,
            lr=ppo_cfg.lr,
            eps=ppo_cfg.eps,
            max_grad_norm=ppo_cfg.max_grad_norm,
            use_normalized_advantage=ppo_cfg.use_normalized_advantage,
        )

    @profiling_wrapper.RangeContext("save_checkpoint")
    def save_checkpoint(
        self, file_name: str, extra_state: Optional[Dict] = None
    ) -> None:
        r"""Save checkpoint with specified name.

        Args:
            file_name: file name for checkpoint

        Returns:
            None
        """
        checkpoint = {
            "state_dict": self.agent.state_dict(),
            "config": self.config,
        }
        if extra_state is not None:
            checkpoint["extra_state"] = extra_state

        torch.save(
            checkpoint, os.path.join(self.config.CHECKPOINT_FOLDER, file_name)
        )

    def load_checkpoint(self, checkpoint_path: str, *args, **kwargs) -> Dict:
        r"""Load checkpoint of specified path as a dict.

        Args:
            checkpoint_path: path of target checkpoint
            *args: additional positional args
            **kwargs: additional keyword args

        Returns:
            dict containing checkpoint info
        """
        return torch.load(checkpoint_path, *args, **kwargs)

    METRICS_BLACKLIST = {"top_down_map", "collisions.is_collision"}

    @classmethod
    def _extract_scalars_from_info(
        cls, info: Dict[str, Any]
    ) -> Dict[str, float]:
        result = {}
        for k, v in info.items():
            if k in cls.METRICS_BLACKLIST:
                continue

            if isinstance(v, dict):
                result.update(
                    {
                        k + "." + subk: subv
                        for subk, subv in cls._extract_scalars_from_info(
                            v
                        ).items()
                        if (k + "." + subk) not in cls.METRICS_BLACKLIST
                    }
                )
            # Things that are scalar-like will have an np.size of 1.
            # Strings also have an np.size of 1, so explicitly ban those
            elif np.size(v) == 1 and not isinstance(v, str):
                result[k] = float(v)

        return result

    @classmethod
    def _extract_scalars_from_infos(
        cls, infos: List[Dict[str, Any]]
    ) -> Dict[str, List[float]]:

        results = defaultdict(list)
        for i in range(len(infos)):
            for k, v in cls._extract_scalars_from_info(infos[i]).items():
                results[k].append(v)

        return results

    @profiling_wrapper.RangeContext("_collect_rollout_step")
    def _collect_rollout_step(
        self, rollouts, current_episode_reward, running_episode_stats
    ):
        pth_time = 0.0
        env_time = 0.0

        t_sample_action = time.time()
        # sample actions
        with torch.no_grad():
            step_observation = {
                k: v[rollouts.step] for k, v in rollouts.observations.items()
            }

            profiling_wrapper.range_push("compute actions")
            (
                values,
                actions,
                actions_log_probs,
                recurrent_hidden_states,
            ) = self.actor_critic.act(
                step_observation,
                rollouts.recurrent_hidden_states[rollouts.step],
                rollouts.prev_actions[rollouts.step],
                rollouts.masks[rollouts.step],
            )

        pth_time += time.time() - t_sample_action

        t_step_env = time.time()
        # NB: Move actions to CPU.  If CUDA tensors are
        # sent in to env.step(), that will create CUDA contexts
        # in the subprocesses.
        # For backwards compatibility, we also call .item() to convert to
        # an int
        step_data = [a.item() for a in actions.to(device="cpu")]
        profiling_wrapper.range_pop()  # compute actions

        outputs = self.envs.step(step_data)
        observations, rewards_l, dones, infos = [
            list(x) for x in zip(*outputs)
        ]

        env_time += time.time() - t_step_env

        t_update_stats = time.time()
        batch = batch_obs(observations, device=self.device)
        batch = apply_obs_transforms_batch(batch, self.obs_transforms)

        rewards = torch.tensor(
            rewards_l, dtype=torch.float, device=current_episode_reward.device
        )
        rewards = rewards.unsqueeze(1)

        masks = torch.tensor(
            [[0.0] if done else [1.0] for done in dones],
            dtype=torch.float,
            device=current_episode_reward.device,
        )

        current_episode_reward += rewards
        running_episode_stats["reward"] += (1 - masks) * current_episode_reward  # type: ignore
        running_episode_stats["count"] += 1 - masks  # type: ignore
        for k, v_k in self._extract_scalars_from_infos(infos).items():
            v = torch.tensor(
                v_k, dtype=torch.float, device=current_episode_reward.device
            ).unsqueeze(1)
            if k not in running_episode_stats:
                running_episode_stats[k] = torch.zeros_like(
                    running_episode_stats["count"]
                )

            running_episode_stats[k] += (1 - masks) * v  # type: ignore

        current_episode_reward *= masks

        if self._static_encoder:
            with torch.no_grad():
                batch["visual_features"] = self._encoder(batch)

        rollouts.insert(
            batch,
            recurrent_hidden_states,
            actions,
            actions_log_probs,
            values,
            rewards,
            masks,
        )

        pth_time += time.time() - t_update_stats

        return pth_time, env_time, self.envs.num_envs

    @profiling_wrapper.RangeContext("_update_agent")
    def _update_agent(self, ppo_cfg, rollouts):
        t_update_model = time.time()
        with torch.no_grad():
            last_observation = {
                k: v[rollouts.step] for k, v in rollouts.observations.items()
            }
            next_value = self.actor_critic.get_value(
                last_observation,
                rollouts.recurrent_hidden_states[rollouts.step],
                rollouts.prev_actions[rollouts.step],
                rollouts.masks[rollouts.step],
            ).detach()

        rollouts.compute_returns(
            next_value, ppo_cfg.use_gae, ppo_cfg.gamma, ppo_cfg.tau
        )

        value_loss, action_loss, dist_entropy = self.agent.update(rollouts)

        rollouts.after_update()

        return (
            time.time() - t_update_model,
            value_loss,
            action_loss,
            dist_entropy,
        )

    @profiling_wrapper.RangeContext("train")
    def train(self) -> None:
        r"""Main method for training PPO.

        Returns:
            None
        """

        profiling_wrapper.configure(
            capture_start_step=self.config.PROFILING.CAPTURE_START_STEP,
            num_steps_to_capture=self.config.PROFILING.NUM_STEPS_TO_CAPTURE,
        )

        self.envs = construct_envs(
            self.config, get_env_class(self.config.ENV_NAME)
        )

        ppo_cfg = self.config.RL.PPO
        self.device = (
            torch.device("cuda", self.config.TORCH_GPU_ID)
            if torch.cuda.is_available()
            else torch.device("cpu")
        )
        if not os.path.isdir(self.config.CHECKPOINT_FOLDER):
            os.makedirs(self.config.CHECKPOINT_FOLDER)
        self._setup_actor_critic_agent(ppo_cfg)
        logger.info(
            "agent number of parameters: {}".format(
                sum(param.numel() for param in self.agent.parameters())
            )
        )
        self.current_update = 0

        rollouts = RolloutStorage(
            ppo_cfg.num_steps,
            self.envs.num_envs,
            self.obs_space,
            self.envs.action_spaces[0],
            ppo_cfg.hidden_size,
        )
        rollouts.to(self.device)

        observations = self.envs.reset()
        batch = batch_obs(observations, device=self.device)
        batch = apply_obs_transforms_batch(batch, self.obs_transforms)

        for sensor in rollouts.observations:
            rollouts.observations[sensor][0].copy_(batch[sensor])

        # batch and observations may contain shared PyTorch CUDA
        # tensors.  We must explicitly clear them here otherwise
        # they will be kept in memory for the entire duration of training!
        batch = None
        observations = None

        current_episode_reward = torch.zeros(self.envs.num_envs, 1)
        running_episode_stats = dict(
            count=torch.zeros(self.envs.num_envs, 1),
            reward=torch.zeros(self.envs.num_envs, 1),
        )
        window_episode_stats: DefaultDict[str, deque] = defaultdict(
            lambda: deque(maxlen=ppo_cfg.reward_window_size)
        )

        t_start = time.time()
        env_time = 0
        pth_time = 0
        count_steps = 0
        count_checkpoints = 0

        lr_scheduler = LambdaLR(
            optimizer=self.agent.optimizer,
            lr_lambda=lambda x: linear_decay(x, self.config.NUM_UPDATES),  # type: ignore
        )

        with TensorboardWriter(
            self.config.TENSORBOARD_DIR, flush_secs=self.flush_secs
        ) as writer:
            for update in range(self.config.NUM_UPDATES):
                profiling_wrapper.on_start_step()
                profiling_wrapper.range_push("train update")
                self.current_update += 1

                if ppo_cfg.use_linear_lr_decay:
                    lr_scheduler.step()  # type: ignore

                if ppo_cfg.use_linear_clip_decay:
                    self.agent.clip_param = ppo_cfg.clip_param * linear_decay(
                        update, self.config.NUM_UPDATES
                    )

                profiling_wrapper.range_push("rollouts loop")
                for _step in range(ppo_cfg.num_steps):
                    (
                        delta_pth_time,
                        delta_env_time,
                        delta_steps,
                    ) = self._collect_rollout_step(
                        rollouts, current_episode_reward, running_episode_stats
                    )
                    pth_time += delta_pth_time
                    env_time += delta_env_time
                    count_steps += delta_steps
                profiling_wrapper.range_pop()  # rollouts loop

                (
                    delta_pth_time,
                    value_loss,
                    action_loss,
                    dist_entropy,
                ) = self._update_agent(ppo_cfg, rollouts)
                pth_time += delta_pth_time

                for k, v in running_episode_stats.items():
                    window_episode_stats[k].append(v.clone())

                deltas = {
                    k: (
                        (v[-1] - v[0]).sum().item()
                        if len(v) > 1
                        else v[0].sum().item()
                    )
                    for k, v in window_episode_stats.items()
                }
                deltas["count"] = max(deltas["count"], 1.0)

                writer.add_scalar(
                    "reward", deltas["reward"] / deltas["count"], count_steps
                )

                # Check to see if there are any metrics
                # that haven't been logged yet
                metrics = {
                    k: v / deltas["count"]
                    for k, v in deltas.items()
                    if k not in {"reward", "count"}
                }
                if len(metrics) > 0:
                    writer.add_scalars("metrics", metrics, count_steps)

                losses = [value_loss, action_loss]
                writer.add_scalars(
                    "losses",
                    {k: l for l, k in zip(losses, ["value", "policy"])},
                    count_steps,
                )

                # log stats
                if update > 0 and update % self.config.LOG_INTERVAL == 0:
                    logger.info(
                        "update: {}\tfps: {:.3f}\t".format(
                            update, count_steps / (time.time() - t_start)
                        )
                    )

                    logger.info(
                        "update: {}\tenv-time: {:.3f}s\tpth-time: {:.3f}s\t"
                        "frames: {}".format(
                            update, env_time, pth_time, count_steps
                        )
                    )

                    logger.info(
                        "Average window size: {}  {}".format(
                            len(window_episode_stats["count"]),
                            "  ".join(
                                "{}: {:.3f}".format(k, v / deltas["count"])
                                for k, v in deltas.items()
                                if k != "count"
                            ),
                        )
                    )

                # checkpoint model
                if update % self.config.CHECKPOINT_INTERVAL == 0:
                    self.save_checkpoint(
                        f"ckpt.{count_checkpoints}.pth", dict(step=count_steps)
                    )
                    count_checkpoints += 1

                profiling_wrapper.range_pop()  # train update

            self.envs.close()

    def _eval_checkpoint(
        self,
        checkpoint_path: str,
        writer: TensorboardWriter,
        checkpoint_index: int = 0,
    ) -> None:
        r"""Evaluates a single checkpoint.

        Args:
            checkpoint_path: path of checkpoint
            writer: tensorboard writer object for logging to tensorboard
            checkpoint_index: index of cur checkpoint for logging

        Returns:
            None
        """
        # Map location CPU is almost always better than mapping to a CUDA device.
        ckpt_dict = self.load_checkpoint(checkpoint_path, map_location="cpu")

        if self.config.EVAL.USE_CKPT_CONFIG:
            config = self._setup_eval_config(ckpt_dict["config"])
        else:
            config = self.config.clone()

        ppo_cfg = config.RL.PPO

        config.defrost()
        config.TASK_CONFIG.DATASET.SPLIT = config.EVAL.SPLIT
        config.freeze()

        if len(self.config.VIDEO_OPTION) > 0:
            config.defrost()
            config.TASK_CONFIG.TASK.MEASUREMENTS.append("TOP_DOWN_MAP")
            config.TASK_CONFIG.TASK.MEASUREMENTS.append("COLLISIONS")
            config.freeze()

        logger.info(f"env config: {config}")
        self.envs = construct_envs(config, get_env_class(config.ENV_NAME))
        self._setup_actor_critic_agent(ppo_cfg)

        self.agent.load_state_dict(ckpt_dict["state_dict"])
        self.actor_critic = self.agent.actor_critic

        observations = self.envs.reset()
        batch = batch_obs(observations, device=self.device)
        batch = apply_obs_transforms_batch(batch, self.obs_transforms)

        current_episode_reward = torch.zeros(
            self.envs.num_envs, 1, device=self.device
        )

        test_recurrent_hidden_states = torch.zeros(
            self.actor_critic.net.num_recurrent_layers,
            self.config.NUM_PROCESSES,
            ppo_cfg.hidden_size,
            device=self.device,
        )
        prev_actions = torch.zeros(
            self.config.NUM_PROCESSES, 1, device=self.device, dtype=torch.long
        )
        not_done_masks = torch.zeros(
            self.config.NUM_PROCESSES, 1, device=self.device
        )
        stats_episodes: Dict[
            Any, Any
        ] = {}  # dict of dicts that stores stats per episode

        rgb_frames = [
            [] for _ in range(self.config.NUM_PROCESSES)
        ]  # type: List[List[np.ndarray]]
        if len(self.config.VIDEO_OPTION) > 0:
            os.makedirs(self.config.VIDEO_DIR, exist_ok=True)

        number_of_eval_episodes = self.config.TEST_EPISODE_COUNT
        if number_of_eval_episodes == -1:
            number_of_eval_episodes = sum(self.envs.number_of_episodes)
        else:
            total_num_eps = sum(self.envs.number_of_episodes)
            if total_num_eps < number_of_eval_episodes:
                logger.warn(
                    f"Config specified {number_of_eval_episodes} eval episodes"
                    ", dataset only has {total_num_eps}."
                )
                logger.warn(f"Evaluating with {total_num_eps} instead.")
                number_of_eval_episodes = total_num_eps

        pbar = tqdm.tqdm(total=number_of_eval_episodes)
        self.actor_critic.eval()
        while (
            len(stats_episodes) < number_of_eval_episodes
            and self.envs.num_envs > 0
        ):
            current_episodes = self.envs.current_episodes()

            with torch.no_grad():
                (
                    _,
                    actions,
                    _,
                    test_recurrent_hidden_states,
                ) = self.actor_critic.act(
                    batch,
                    test_recurrent_hidden_states,
                    prev_actions,
                    not_done_masks,
                    deterministic=False,
                )

                prev_actions.copy_(actions)  # type: ignore

            # NB: Move actions to CPU.  If CUDA tensors are
            # sent in to env.step(), that will create CUDA contexts
            # in the subprocesses.
            # For backwards compatibility, we also call .item() to convert to
            # an int
            step_data = [a.item() for a in actions.to(device="cpu")]

            outputs = self.envs.step(step_data)

            observations, rewards_l, dones, infos = [
                list(x) for x in zip(*outputs)
            ]
            batch = batch_obs(observations, device=self.device)
            batch = apply_obs_transforms_batch(batch, self.obs_transforms)

            not_done_masks = torch.tensor(
                [[0.0] if done else [1.0] for done in dones],
                dtype=torch.float,
                device=self.device,
            )

            rewards = torch.tensor(
                rewards_l, dtype=torch.float, device=self.device
            ).unsqueeze(1)
            current_episode_reward += rewards
            next_episodes = self.envs.current_episodes()
            envs_to_pause = []
            n_envs = self.envs.num_envs
            for i in range(n_envs):
                if (
                    next_episodes[i].scene_id,
                    next_episodes[i].episode_id,
                ) in stats_episodes:
                    envs_to_pause.append(i)

                # episode ended
                if not_done_masks[i].item() == 0:
                    pbar.update()
                    episode_stats = {}
                    episode_stats["reward"] = current_episode_reward[i].item()
                    episode_stats.update(
                        self._extract_scalars_from_info(infos[i])
                    )
                    current_episode_reward[i] = 0
                    # use scene_id + episode_id as unique id for storing stats
                    stats_episodes[
                        (
                            current_episodes[i].scene_id,
                            current_episodes[i].episode_id,
                        )
                    ] = episode_stats

                    if len(self.config.VIDEO_OPTION) > 0:
                        generate_video(
                            video_option=self.config.VIDEO_OPTION,
                            video_dir=self.config.VIDEO_DIR,
                            images=rgb_frames[i],
                            episode_id=current_episodes[i].episode_id,
                            checkpoint_idx=checkpoint_index,
                            metrics=self._extract_scalars_from_info(infos[i]),
                            tb_writer=writer,
                        )

                        rgb_frames[i] = []

                # episode continues
                elif len(self.config.VIDEO_OPTION) > 0:
                    # TODO move normalization / channel changing out of the policy and undo it here
                    frame = observations_to_image(
                        {k: v[i] for k, v in batch.items()}, infos[i]
                    )
                    rgb_frames[i].append(frame)

            (
                self.envs,
                test_recurrent_hidden_states,
                not_done_masks,
                current_episode_reward,
                prev_actions,
                batch,
                rgb_frames,
            ) = self._pause_envs(
                envs_to_pause,
                self.envs,
                test_recurrent_hidden_states,
                not_done_masks,
                current_episode_reward,
                prev_actions,
                batch,
                rgb_frames,
            )

        num_episodes = len(stats_episodes)
        aggregated_stats = {}
        for stat_key in next(iter(stats_episodes.values())).keys():
            aggregated_stats[stat_key] = (
                sum(v[stat_key] for v in stats_episodes.values())
                / num_episodes
            )

        for k, v in aggregated_stats.items():
            logger.info(f"Average episode {k}: {v:.4f}")

        step_id = checkpoint_index
        if "extra_state" in ckpt_dict and "step" in ckpt_dict["extra_state"]:
            step_id = ckpt_dict["extra_state"]["step"]

        writer.add_scalars(
            "eval_reward",
            {"average reward": aggregated_stats["reward"]},
            step_id,
        )

        metrics = {k: v for k, v in aggregated_stats.items() if k != "reward"}
        if len(metrics) > 0:
            writer.add_scalars("eval_metrics", metrics, step_id)

        self.envs.close()

habitat_baselines/rl/ppo/ppo_trainer.py
cshizhe/onav_rim
#!/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
import copy
import argparse
import random
import time
import numpy as np

import torch
import torch.multiprocessing as mp

from habitat.config import Config
from habitat_baselines.common.baseline_registry import baseline_registry
from habitat_baselines.config.default import get_config

def evaluation_worker(config):
    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)
    trainer.eval()

def multiprocess_evaluate(config, num_gpus):
    if config.TASK_CONFIG.DATASET.CONTENT_SCENES[0] == '*':
        content_dir = os.path.join(
            os.path.dirname(config.TASK_CONFIG.DATASET.DATA_PATH.format(split=config.EVAL.SPLIT)),
            'content'
        )
        if os.path.exists(content_dir):
            scenes = [x.split('.')[0] for x in os.listdir(content_dir) if x.endswith('.json.gz')]
        else:
            scenes = config.TASK_CONFIG.DATASET.CONTENT_SCENES
    else:
        scenes = config.TASK_CONFIG.DATASET.CONTENT_SCENES
    assert num_gpus <= len(scenes), '#scenes: %d should >= #process: %d' % (len(scenes), num_gpus)
    scenes.sort()

    os.makedirs(config.TENSORBOARD_DIR, exist_ok=True)

    def even_split_scenes(scenes, nparts):
        splited_scenes = [[] for _ in range(nparts)]
        for i, scene in enumerate(scenes):
            idx = i % nparts
            splited_scenes[idx].append(scene)
        return splited_scenes

    node_rank = int(os.environ.get('SLURM_NODEID', 0))
    num_gpus_per_node = int(os.environ.get('SLURM_NTASKS_PER_NODE', 1))
    cur_gpu_id = int(os.environ.get('SLURM_LOCALID', 0))
    scenes = even_split_scenes(scenes, num_gpus)[cur_gpu_id + node_rank * num_gpus_per_node]

    num_process_per_gpu = config.NUM_PROCESSES
    splited_scenes = even_split_scenes(scenes, num_process_per_gpu)

    processes = []
    for i in range(num_process_per_gpu):
        i_config = copy.deepcopy(config)
        i_config.defrost()
        i_config.TRAINER_NAME = 'il-trainer'
        i_config.TORCH_GPU_ID = cur_gpu_id
        i_config.SIMULATOR_GPU_ID = cur_gpu_id
        i_config.NUM_PROCESSES = 1
        i_config.TASK_CONFIG.DATASET.CONTENT_SCENES = splited_scenes[i]
        i_config.freeze()
        print('gpu_id %d/%d, proc %d/%d, scenes %d\n%s' %(
            cur_gpu_id+node_rank*num_gpus_per_node+1, num_gpus, i+1, num_process_per_gpu, len(splited_scenes[i]),
            ', '.join(splited_scenes[i])
        ))
        processes.append(
            mp.Process(target=evaluation_worker, args=(i_config, ))
        )
        processes[-1].start()
    
    for proc in processes:
        proc.join()

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--run-type",
        choices=["train", "eval"],
        required=True,
        help="run type of the experiment (train or eval)",
    )
    parser.add_argument(
        "--exp-config",
        type=str,
        required=True,
        help="path to config yaml containing info about experiment",
    )
    parser.add_argument(
        "opts",
        default=None,
        nargs=argparse.REMAINDER,
        help="Modify config options from command line",
    )

    args = parser.parse_args()
    run_exp(**vars(args))


def execute_exp(config: Config, run_type: str) -> None:
    r"""This function runs the specified config with the specified runtype
    Args:
    config: Habitat.config
    runtype: str {train or eval}
    """
    random.seed(config.TASK_CONFIG.SEED)
    np.random.seed(config.TASK_CONFIG.SEED)
    torch.manual_seed(config.TASK_CONFIG.SEED)

    if run_type == "train":
        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)
        trainer.train()

    elif run_type == "eval":
        result_file = os.path.join(config.EVAL_RESULTS_DIR, f'{config.EVAL.SPLIT}_pred_trajectories.jsonl')
        if os.path.exists(result_file):
            print(f'eval result dir {result_file} already exists')
            return
        mp.set_start_method('spawn')
        multiprocess_evaluate(config, int(os.environ.get('SLURM_NTASKS', 1)))

        # trainer.eval()


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)
    execute_exp(config, run_type)


if __name__ == "__main__":
    main()

habitat_baselines/run.py
cshizhe/onav_rim
#!/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.

r"""The code below is taken from
https://github.com/JunjH/Revisiting_Single_Depth_Estimation
Revisiting Single Image Depth Estimation: Toward Higher Resolution Maps With Accurate Object Boundaries
Junjie Hu and Mete Ozay and Yan Zhang and Takayuki Okatani
WACV 2019

ResNet code gently borrowed from
https://github.com/pytorch/vision/blob/master/torchvision/models/py
"""


import math

import numpy as np
import torch
import torch.nn.parallel
from PIL import Image
from torch import nn as nn
from torch.nn import functional as F
from torch.utils import model_zoo as model_zoo
from torchvision import transforms

accimage = None


__all__ = [
    "ResNet",
    "resnet18",
    "resnet34",
    "resnet50",
    "resnet101",
    "resnet152",
]


model_urls = {
    "resnet18": "https://download.pytorch.org/models/resnet18-5c106cde.pth",
    "resnet34": "https://download.pytorch.org/models/resnet34-333f7ec4.pth",
    "resnet50": "https://download.pytorch.org/models/resnet50-19c8e357.pth",
    "resnet101": "https://download.pytorch.org/models/resnet101-5d3b4d8f.pth",
    "resnet152": "https://download.pytorch.org/models/resnet152-b121ed2d.pth",
}


def conv3x3(in_planes, out_planes, stride=1):
    "3x3 convolution with padding"
    return nn.Conv2d(
        in_planes,
        out_planes,
        kernel_size=3,
        stride=stride,
        padding=1,
        bias=False,
    )


class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(BasicBlock, self).__init__()
        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = nn.BatchNorm2d(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = nn.BatchNorm2d(planes)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out


class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(Bottleneck, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(
            planes, planes, kernel_size=3, stride=stride, padding=1, bias=False
        )
        self.bn2 = nn.BatchNorm2d(planes)
        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(planes * 4)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out


class ResNet(nn.Module):
    def __init__(self, block, layers, num_classes=1000):
        self.inplanes = 64
        super(ResNet, self).__init__()
        self.conv1 = nn.Conv2d(
            3, 64, kernel_size=7, stride=2, padding=3, bias=False
        )
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
        self.avgpool = nn.AvgPool2d(7, stride=1)
        self.fc = nn.Linear(512 * block.expansion, num_classes)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2.0 / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

    def _make_layer(self, block, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(
                    self.inplanes,
                    planes * block.expansion,
                    kernel_size=1,
                    stride=stride,
                    bias=False,
                ),
                nn.BatchNorm2d(planes * block.expansion),
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes * block.expansion
        for _ in range(1, blocks):
            layers.append(block(self.inplanes, planes))

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = self.avgpool(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)

        return x


def resnet18(pretrained=False, **kwargs):
    r"""Constructs a ResNet-18 model.
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls["resnet18"]))
    return model


def resnet34(pretrained=False, **kwargs):
    r"""Constructs a ResNet-34 model.
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls["resnet34"]))
    return model


def resnet50(pretrained=False, **kwargs):
    r"""Constructs a ResNet-50 model.
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
    if pretrained:
        model.load_state_dict(
            model_zoo.load_url(
                model_urls["resnet50"], "pretrained_model/encoder"
            )
        )
    return model


def resnet101(pretrained=False, **kwargs):
    r"""Constructs a ResNet-101 model.
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls["resnet101"]))
    return model


def resnet152(pretrained=False, **kwargs):
    r"""Constructs a ResNet-152 model.
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls["resnet152"]))
    return model


class model(nn.Module):
    def __init__(self, Encoder, num_features, block_channel):

        super(model, self).__init__()

        self.E = Encoder
        self.D = D(num_features)
        self.MFF = MFF(block_channel)
        self.R = R(block_channel)

    def forward(self, x):
        x_block1, x_block2, x_block3, x_block4 = self.E(x)
        x_decoder = self.D(x_block1, x_block2, x_block3, x_block4)
        x_mff = self.MFF(
            x_block1,
            x_block2,
            x_block3,
            x_block4,
            [x_decoder.size(2), x_decoder.size(3)],
        )
        out = self.R(torch.cat((x_decoder, x_mff), 1))

        return out


class _UpProjection(nn.Sequential):
    def __init__(self, num_input_features, num_output_features):
        super(_UpProjection, self).__init__()

        self.conv1 = nn.Conv2d(
            num_input_features,
            num_output_features,
            kernel_size=5,
            stride=1,
            padding=2,
            bias=False,
        )
        self.bn1 = nn.BatchNorm2d(num_output_features)
        self.relu = nn.ReLU(inplace=True)
        self.conv1_2 = nn.Conv2d(
            num_output_features,
            num_output_features,
            kernel_size=3,
            stride=1,
            padding=1,
            bias=False,
        )
        self.bn1_2 = nn.BatchNorm2d(num_output_features)

        self.conv2 = nn.Conv2d(
            num_input_features,
            num_output_features,
            kernel_size=5,
            stride=1,
            padding=2,
            bias=False,
        )
        self.bn2 = nn.BatchNorm2d(num_output_features)

    def forward(self, x, size):
        x = F.upsample(x, size=size, mode="bilinear")
        x_conv1 = self.relu(self.bn1(self.conv1(x)))
        bran1 = self.bn1_2(self.conv1_2(x_conv1))
        bran2 = self.bn2(self.conv2(x))

        out = self.relu(bran1 + bran2)

        return out


class E_resnet(nn.Module):
    def __init__(self, original_model, num_features=2048):
        super(E_resnet, self).__init__()
        self.conv1 = original_model.conv1
        self.bn1 = original_model.bn1
        self.relu = original_model.relu
        self.maxpool = original_model.maxpool

        self.layer1 = original_model.layer1
        self.layer2 = original_model.layer2
        self.layer3 = original_model.layer3
        self.layer4 = original_model.layer4

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x_block1 = self.layer1(x)
        x_block2 = self.layer2(x_block1)
        x_block3 = self.layer3(x_block2)
        x_block4 = self.layer4(x_block3)

        return x_block1, x_block2, x_block3, x_block4


class D(nn.Module):
    def __init__(self, num_features=2048):
        super(D, self).__init__()
        self.conv = nn.Conv2d(
            num_features,
            num_features // 2,
            kernel_size=1,
            stride=1,
            bias=False,
        )
        num_features = num_features // 2
        self.bn = nn.BatchNorm2d(num_features)

        self.up1 = _UpProjection(
            num_input_features=num_features,
            num_output_features=num_features // 2,
        )
        num_features = num_features // 2

        self.up2 = _UpProjection(
            num_input_features=num_features,
            num_output_features=num_features // 2,
        )
        num_features = num_features // 2

        self.up3 = _UpProjection(
            num_input_features=num_features,
            num_output_features=num_features // 2,
        )
        num_features = num_features // 2

        self.up4 = _UpProjection(
            num_input_features=num_features,
            num_output_features=num_features // 2,
        )
        num_features = num_features // 2

    def forward(self, x_block1, x_block2, x_block3, x_block4):
        x_d0 = F.relu(self.bn(self.conv(x_block4)))
        x_d1 = self.up1(x_d0, [x_block3.size(2), x_block3.size(3)])
        x_d2 = self.up2(x_d1, [x_block2.size(2), x_block2.size(3)])
        x_d3 = self.up3(x_d2, [x_block1.size(2), x_block1.size(3)])
        x_d4 = self.up4(x_d3, [x_block1.size(2) * 2, x_block1.size(3) * 2])

        return x_d4


class MFF(nn.Module):
    def __init__(self, block_channel, num_features=64):

        super(MFF, self).__init__()

        self.up1 = _UpProjection(
            num_input_features=block_channel[0], num_output_features=16
        )

        self.up2 = _UpProjection(
            num_input_features=block_channel[1], num_output_features=16
        )

        self.up3 = _UpProjection(
            num_input_features=block_channel[2], num_output_features=16
        )

        self.up4 = _UpProjection(
            num_input_features=block_channel[3], num_output_features=16
        )

        self.conv = nn.Conv2d(
            num_features,
            num_features,
            kernel_size=5,
            stride=1,
            padding=2,
            bias=False,
        )
        self.bn = nn.BatchNorm2d(num_features)

    def forward(self, x_block1, x_block2, x_block3, x_block4, size):
        x_m1 = self.up1(x_block1, size)
        x_m2 = self.up2(x_block2, size)
        x_m3 = self.up3(x_block3, size)
        x_m4 = self.up4(x_block4, size)

        x = self.bn(self.conv(torch.cat((x_m1, x_m2, x_m3, x_m4), 1)))
        x = F.relu(x)

        return x


class R(nn.Module):
    def __init__(self, block_channel):

        super(R, self).__init__()

        num_features = 64 + block_channel[3] // 32
        self.conv0 = nn.Conv2d(
            num_features,
            num_features,
            kernel_size=5,
            stride=1,
            padding=2,
            bias=False,
        )
        self.bn0 = nn.BatchNorm2d(num_features)

        self.conv1 = nn.Conv2d(
            num_features,
            num_features,
            kernel_size=5,
            stride=1,
            padding=2,
            bias=False,
        )
        self.bn1 = nn.BatchNorm2d(num_features)

        self.conv2 = nn.Conv2d(
            num_features, 1, kernel_size=5, stride=1, padding=2, bias=True
        )

    def forward(self, x):
        x0 = self.conv0(x)
        x0 = self.bn0(x0)
        x0 = F.relu(x0)

        x1 = self.conv1(x0)
        x1 = self.bn1(x1)
        x1 = F.relu(x1)

        x2 = self.conv2(x1)

        return x2


def _is_pil_image(img):
    return isinstance(img, Image.Image)


def _is_numpy_image(img):
    return isinstance(img, np.ndarray) and (img.ndim in {2, 3})


class Scale(object):
    def __init__(self, size):
        self.size = size

    def __call__(self, image):
        image = self.changeScale(image, self.size)

        return image

    def changeScale(self, img, size, interpolation=Image.BILINEAR):
        ow, oh = size

        return img.resize((ow, oh), interpolation)


class CenterCrop(object):
    def __init__(self, size):
        self.size = size

    def __call__(self, image):
        image = self.centerCrop(image, self.size)

        return image

    def centerCrop(self, image, size):
        w1, h1 = image.size
        tw, th = size

        if w1 == tw and h1 == th:
            return image

        x1 = int(round((w1 - tw) / 2.0))
        y1 = int(round((h1 - th) / 2.0))

        image = image.crop((x1, y1, tw + x1, th + y1))

        return image


class ToTensor(object):
    r"""Convert a ``PIL.Image`` or ``numpy.ndarray`` to tensor.
    Converts a PIL.Image or numpy.ndarray (H x W x C) in the range
    [0, 255] to a torch.FloatTensor of shape (C x H x W) in the range [0.0, 1.0].
    """

    def __call__(self, image):
        image = self.to_tensor(image)

        return image

    def to_tensor(self, pic):
        if not (_is_pil_image(pic) or _is_numpy_image(pic)):
            raise TypeError(
                "pic should be PIL Image or ndarray. Got {}".format(type(pic))
            )

        if isinstance(pic, np.ndarray):

            img = torch.from_numpy(pic.transpose((2, 0, 1)))
            return img.float().div(255)

        if accimage is not None and isinstance(pic, accimage.Image):
            nppic = np.zeros(
                [pic.channels, pic.height, pic.width], dtype=np.float32
            )
            pic.copyto(nppic)
            return torch.from_numpy(nppic)

        # handle PIL Image
        if pic.mode == "I":
            img = torch.from_numpy(np.array(pic, np.int32, copy=False))
        elif pic.mode == "I;16":
            img = torch.from_numpy(np.array(pic, np.int16, copy=False))
        else:
            img = torch.ByteTensor(
                torch.ByteStorage.from_buffer(pic.tobytes())
            )
        # PIL image mode: 1, L, P, I, F, RGB, YCbCr, RGBA, CMYK
        if pic.mode == "YCbCr":
            nchannel = 3
        elif pic.mode == "I;16":
            nchannel = 1
        else:
            nchannel = len(pic.mode)
        img = img.view(pic.size[1], pic.size[0], nchannel)
        # put it from HWC to CHW format
        # yikes, this transpose takes 80% of the loading time/CPU
        img = img.transpose(0, 1).transpose(0, 2).contiguous()
        if isinstance(img, torch.ByteTensor):
            return img.float().div(255)
        else:
            return img


class Normalize(object):
    def __init__(self, mean, std):
        self.mean = mean
        self.std = std

    def __call__(self, image):
        image = self.normalize(image, self.mean, self.std)

        return image

    def normalize(self, tensor, mean, std):
        for t, m, s in zip(tensor, mean, std):
            t.sub_(m).div_(s)

        return tensor


def define_model(is_resnet, is_densenet, is_senet):
    if is_resnet:
        original_model = resnet50(pretrained=False)
        Encoder = E_resnet(original_model)
        model1 = model(
            Encoder, num_features=2048, block_channel=[256, 512, 1024, 2048]
        )
    if is_densenet:
        # original_model = dendensenet161(pretrained=False)
        # Encoder = E_densenet(original_model)
        # model1 = model(
        #    Encoder, num_features=2208, block_channel=[192, 384, 1056, 2208]
        # )
        raise NotImplementedError()
    if is_senet:
        # original_model = senet154(pretrained=False)
        # Encoder = E_senet(original_model)
        # model1 = model(
        #    Encoder, num_features=2048, block_channel=[256, 512, 1024, 2048]
        # )
        raise NotImplementedError()
    return model1


class MonoDepthEstimator:
    def __init__(self, checkpoint="./pretrained_model/model_resnet"):
        self.model = define_model(
            is_resnet=True, is_densenet=False, is_senet=False
        )
        self.model = torch.nn.DataParallel(self.model).cuda()
        cpt = torch.load(checkpoint, map_location='cpu')
        if "state_dict" in cpt.keys():
            cpt = cpt["state_dict"]
        self.model.load_state_dict(cpt)
        self.model.eval()
        self.init_preprocessor()

    def init_preprocessor(self):
        __imagenet_stats = {
            "mean": [0.485, 0.456, 0.406],
            "std": [0.229, 0.224, 0.225],
        }

        self.transform = transforms.Compose(
            [
                Scale([320, 240]),
                # CenterCrop([304, 228]),
                ToTensor(),
                Normalize(__imagenet_stats["mean"], __imagenet_stats["std"]),
            ]
        )

    def preprocess(self, image):
        image_torch = self.transform(image).unsqueeze(0)
        return image_torch.cuda()

    def compute_depth(self, image):
        # Input: image is a PIL image
        # Output: depth is a numpy array
        image_torch = self.preprocess(image)
        # print(image_torch.size())
        depth_torch = self.model(image_torch)
        depth = (
            depth_torch.view(depth_torch.size(2), depth_torch.size(3))
            .data.cpu()
            .numpy()
        )
        return depth

habitat_baselines/slambased/monodepth.py
cshizhe/onav_rim
#!/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 glob
import os
import re
import shutil
import tarfile
from collections import defaultdict
from io import BytesIO
from typing import (
    Any,
    DefaultDict,
    Dict,
    Iterable,
    List,
    Optional,
    Tuple,
    Union,
)

import numpy as np
import torch
from gym.spaces import Box
from numpy import ndarray
from PIL import Image
from torch import Size, Tensor
from torch import nn as nn

from habitat import logger
from habitat.core.dataset import Episode
from habitat.utils import profiling_wrapper
from habitat.utils.visualizations.utils import images_to_video
from habitat_baselines.common.tensorboard_utils import TensorboardWriter


class Flatten(nn.Module):
    def forward(self, x: Tensor) -> Tensor:
        return torch.flatten(x, start_dim=1)


class CustomFixedCategorical(torch.distributions.Categorical):  # type: ignore
    def sample(
        self, sample_shape: Size = torch.Size()  # noqa: B008
    ) -> Tensor:
        return super().sample(sample_shape).unsqueeze(-1)

    def log_probs(self, actions: Tensor) -> Tensor:
        return (
            super()
            .log_prob(actions.squeeze(-1))
            .view(actions.size(0), -1)
            .sum(-1)
            .unsqueeze(-1)
        )

    def mode(self):
        return self.probs.argmax(dim=-1, keepdim=True)


class CategoricalNet(nn.Module):
    def __init__(self, num_inputs: int, num_outputs: int, clf_net=None) -> None:
        super().__init__()

        if clf_net is None:
            self.linear = nn.Linear(num_inputs, num_outputs)
            nn.init.orthogonal_(self.linear.weight, gain=0.01)
            nn.init.constant_(self.linear.bias, 0)
        else:
            self.linear = clf_net

    def forward(self, x: Tensor, masks=None) -> CustomFixedCategorical:
        x = self.linear(x)
        if masks is not None:
            x = x.masked_fill(masks, -float('inf'))
        return CustomFixedCategorical(logits=x)


def linear_decay(epoch: int, total_num_updates: int) -> float:
    r"""Returns a multiplicative factor for linear value decay

    Args:
        epoch: current epoch number
        total_num_updates: total number of

    Returns:
        multiplicative factor that decreases param value linearly
    """
    return 1 - (epoch / float(total_num_updates))


def _to_tensor(v: Union[Tensor, ndarray]) -> torch.Tensor:
    if torch.is_tensor(v):
        return v
    elif isinstance(v, np.ndarray):
        if v.dtype == np.uint32:
            return torch.from_numpy(v.astype(int))
        else:
            return torch.from_numpy(v)
    else:
        return torch.tensor(v, dtype=torch.float)


@torch.no_grad()
@profiling_wrapper.RangeContext("batch_obs")
def batch_obs(
    observations: List[Dict],
    device: Optional[torch.device] = None,
) -> Dict[str, torch.Tensor]:
    r"""Transpose a batch of observation dicts to a dict of batched
    observations.

    Args:
        observations:  list of dicts of observations.
        device: The torch.device to put the resulting tensors on.
            Will not move the tensors if None

    Returns:
        transposed dict of torch.Tensor of observations.
    """
    batch: DefaultDict[str, List] = defaultdict(list)

    for obs in observations:
        for sensor in obs:
            batch[sensor].append(_to_tensor(obs[sensor]))

    batch_t: Dict[str, torch.Tensor] = {}

    for sensor in batch:
        batch_t[sensor] = torch.stack(batch[sensor], dim=0).to(device=device)

    return batch_t


def get_checkpoint_id(ckpt_path: str) -> Optional[int]:
    r"""Attempts to extract the ckpt_id from the filename of a checkpoint.
    Assumes structure of ckpt.ID.path .

    Args:
        ckpt_path: the path to the ckpt file

    Returns:
        returns an int if it is able to extract the ckpt_path else None
    """
    ckpt_path = os.path.basename(ckpt_path)
    nums: List[int] = [int(s) for s in ckpt_path.split(".") if s.isdigit()]
    if len(nums) > 0:
        return nums[-1]
    return None


def poll_checkpoint_folder(
    checkpoint_folder: str, previous_ckpt_ind: int
) -> Optional[str]:
    r"""Return (previous_ckpt_ind + 1)th checkpoint in checkpoint folder
    (sorted by time of last modification).

    Args:
        checkpoint_folder: directory to look for checkpoints.
        previous_ckpt_ind: index of checkpoint last returned.

    Returns:
        return checkpoint path if (previous_ckpt_ind + 1)th checkpoint is found
        else return None.
    """
    assert os.path.isdir(checkpoint_folder), (
        f"invalid checkpoint folder " f"path {checkpoint_folder}"
    )
    models_paths = list(
        filter(os.path.isfile, glob.glob(checkpoint_folder + "/*"))
    )
    models_paths.sort(key=os.path.getmtime)
    ind = previous_ckpt_ind + 1
    if ind < len(models_paths):
        return models_paths[ind]
    return None


def generate_video(
    video_option: List[str],
    video_dir: Optional[str],
    images: List[np.ndarray],
    episode_id: Union[int, str],
    checkpoint_idx: int,
    metrics: Dict[str, float],
    tb_writer: TensorboardWriter,
    fps: int = 15,
) -> None:
    r"""Generate video according to specified information.

    Args:
        video_option: string list of "tensorboard" or "disk" or both.
        video_dir: path to target video directory.
        images: list of images to be converted to video.
        episode_id: episode id for video naming.
        checkpoint_idx: checkpoint index 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:
        None
    """
    if len(images) < 1:
        return

    metric_strs = []
    for k, v in metrics.items():
        metric_strs.append(f"{k}={v:.2f}")

    video_name = f"episode={episode_id}-ckpt={checkpoint_idx}-" + "-".join(
        metric_strs
    )
    if "disk" in video_option:
        assert video_dir is not None
        images_to_video(images, video_dir, video_name)
    if "tensorboard" in video_option:
        tb_writer.add_video_from_np_images(
            f"episode{episode_id}", checkpoint_idx, images, fps=fps
        )


def tensor_to_depth_images(tensor: Union[torch.Tensor, List]) -> np.ndarray:
    r"""Converts tensor (or list) of n image tensors to list of n images.
    Args:
        tensor: tensor containing n image tensors
    Returns:
        list of images
    """
    images = []

    for img_tensor in tensor:
        image = img_tensor.permute(1, 2, 0).cpu().numpy() * 255
        images.append(image)

    return images


def tensor_to_bgr_images(
    tensor: Union[torch.Tensor, Iterable[torch.Tensor]]
) -> List[np.ndarray]:
    r"""Converts tensor of n image tensors to list of n BGR images.
    Args:
        tensor: tensor containing n image tensors
    Returns:
        list of images
    """
    import cv2

    images = []

    for img_tensor in tensor:
        img = img_tensor.permute(1, 2, 0).cpu().numpy() * 255
        img = img.astype(np.uint8)
        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        images.append(img)

    return images


def image_resize_shortest_edge(
    img: Tensor, size: int, channels_last: bool = False
) -> torch.Tensor:
    """Resizes an img so that the shortest side is length of size while
        preserving aspect ratio.

    Args:
        img: the array object that needs to be resized (HWC) or (NHWC)
        size: the size that you want the shortest edge to be resize to
        channels: a boolean that channel is the last dimension
    Returns:
        The resized array as a torch tensor.
    """
    img = _to_tensor(img)
    no_batch_dim = len(img.shape) == 3
    if len(img.shape) < 3 or len(img.shape) > 5:
        raise NotImplementedError()
    if no_batch_dim:
        img = img.unsqueeze(0)  # Adds a batch dimension
    h, w = get_image_height_width(img, channels_last=channels_last)
    if channels_last:
        if len(img.shape) == 4:
            # NHWC -> NCHW
            img = img.permute(0, 3, 1, 2)
        else:
            # NDHWC -> NDCHW
            img = img.permute(0, 1, 4, 2, 3)

    # Percentage resize
    scale = size / min(h, w)
    h = int(h * scale)
    w = int(w * scale)
    img = torch.nn.functional.interpolate(
        img.float(), size=(h, w), mode="area"
    ).to(dtype=img.dtype)
    if channels_last:
        if len(img.shape) == 4:
            # NCHW -> NHWC
            img = img.permute(0, 2, 3, 1)
        else:
            # NDCHW -> NDHWC
            img = img.permute(0, 1, 3, 4, 2)
    if no_batch_dim:
        img = img.squeeze(dim=0)  # Removes the batch dimension
    return img


def center_crop(
    img: Tensor, size: Union[int, Tuple[int, int]], channels_last: bool = False
) -> Tensor:
    """Performs a center crop on an image.

    Args:
        img: the array object that needs to be resized (either batched or unbatched)
        size: A sequence (h, w) or a python(int) that you want cropped
        channels_last: If the channels are the last dimension.
    Returns:
        the resized array
    """
    h, w = get_image_height_width(img, channels_last=channels_last)

    if isinstance(size, int):
        size_tuple: Tuple[int, int] = (int(size), int(size))
    else:
        size_tuple = size
    assert len(size_tuple) == 2, "size should be (h,w) you wish to resize to"
    cropy, cropx = size_tuple

    startx = w // 2 - (cropx // 2)
    starty = h // 2 - (cropy // 2)
    if channels_last:
        return img[..., starty : starty + cropy, startx : startx + cropx, :]
    else:
        return img[..., starty : starty + cropy, startx : startx + cropx]


def get_image_height_width(
    img: Union[Box, np.ndarray, torch.Tensor], channels_last: bool = False
) -> Tuple[int, int]:
    if img.shape is None or len(img.shape) < 3 or len(img.shape) > 5:
        raise NotImplementedError()
    if channels_last:
        # NHWC
        h, w = img.shape[-3:-1]
    else:
        # NCHW
        h, w = img.shape[-2:]
    return h, w


def overwrite_gym_box_shape(box: Box, shape) -> Box:
    if box.shape == shape:
        return box
    shape = list(shape) + list(box.shape[len(shape) :])
    low = box.low if np.isscalar(box.low) else np.min(box.low)
    high = box.high if np.isscalar(box.high) else np.max(box.high)
    return Box(low=low, high=high, shape=shape, dtype=box.dtype)


def get_scene_episode_dict(episodes: List[Episode]) -> Dict:
    scene_ids = []
    scene_episode_dict = {}

    for episode in episodes:
        if episode.scene_id not in scene_ids:
            scene_ids.append(episode.scene_id)
            scene_episode_dict[episode.scene_id] = [episode]
        else:
            scene_episode_dict[episode.scene_id].append(episode)

    return scene_episode_dict


def base_plus_ext(path: str) -> Union[Tuple[str, str], Tuple[None, None]]:
    """Helper method that splits off all extension.
    Returns base, allext.
    path: path with extensions
    returns: path with all extensions removed
    """
    match = re.match(r"^((?:.*/|)[^.]+)[.]([^/]*)$", path)
    if not match:
        return None, None
    return match.group(1), match.group(2)


def valid_sample(sample: Optional[Any]) -> bool:
    """Check whether a webdataset sample is valid.
    sample: sample to be checked
    """
    return (
        sample is not None
        and isinstance(sample, dict)
        and len(list(sample.keys())) > 0
        and not sample.get("__bad__", False)
    )


def img_bytes_2_np_array(
    x: Tuple[int, torch.Tensor, bytes]
) -> Tuple[int, torch.Tensor, bytes, np.ndarray]:
    """Mapper function to convert image bytes in webdataset sample to numpy
    arrays.
    Args:
        x: webdataset sample containing ep_id, question, answer and imgs
    Returns:
        Same sample with bytes turned into np arrays.
    """
    images = []
    img_bytes: bytes
    for img_bytes in x[3:]:
        bytes_obj = BytesIO()
        bytes_obj.write(img_bytes)
        image = np.array(Image.open(bytes_obj))
        img = image.transpose(2, 0, 1)
        img = img / 255.0
        images.append(img)
    return (*x[0:3], np.array(images, dtype=np.float32))


def create_tar_archive(archive_path: str, dataset_path: str) -> None:
    """Creates tar archive of dataset and returns status code.
    Used in VQA trainer's webdataset.
    """
    logger.info("[ Creating tar archive. This will take a few minutes. ]")

    with tarfile.open(archive_path, "w:gz") as tar:
        for file in sorted(os.listdir(dataset_path)):
            tar.add(os.path.join(dataset_path, file))


def delete_folder(path: str) -> None:
    shutil.rmtree(path)

habitat_baselines/utils/common.py
cshizhe/onav_rim
#!/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 random
from typing import List, Type, Union

import habitat
from habitat import Config, Env, RLEnv, VectorEnv, make_dataset, logger


def make_env_fn(
    config: Config, env_class: Union[Type[Env], Type[RLEnv]]
) -> Union[Env, RLEnv]:
    r"""Creates an env of type env_class with specified config and rank.
    This is to be passed in as an argument when creating VectorEnv.

    Args:
        config: root exp config that has core env config node as well as
            env-specific config node.
        env_class: class type of the env to be created.

    Returns:
        env object created according to specification.
    """
    dataset = make_dataset(
        config.TASK_CONFIG.DATASET.TYPE, config=config.TASK_CONFIG.DATASET
    )
    env = env_class(config=config, dataset=dataset)
    env.seed(config.TASK_CONFIG.SEED)
    return env


def construct_envs(
    config: Config,
    env_class: Union[Type[Env], Type[RLEnv]],
    workers_ignore_signals: bool = False,
    auto_reset_done: bool = True,
) -> VectorEnv:
    r"""Create VectorEnv object with specified config and env class type.
    To allow better performance, dataset are split into small ones for
    each individual env, grouped by scenes.

    :param config: configs that contain num_processes as well as information
    :param necessary to create individual environments.
    :param env_class: class type of the envs to be created.
    :param workers_ignore_signals: Passed to :ref:`habitat.VectorEnv`'s constructor

    :return: VectorEnv object created according to specification.
    """

    num_processes = config.NUM_PROCESSES
    configs = []
    env_classes = [env_class for _ in range(num_processes)]
    dataset = make_dataset(config.TASK_CONFIG.DATASET.TYPE, config=config.TASK_CONFIG.DATASET)
    scenes = config.TASK_CONFIG.DATASET.CONTENT_SCENES
    if "*" in config.TASK_CONFIG.DATASET.CONTENT_SCENES:
        scenes = dataset.get_scenes_to_load(config.TASK_CONFIG.DATASET)

    if num_processes > 1:
        if len(scenes) == 0:
            raise RuntimeError(
                "No scenes to load, multiple process logic relies on being able to split scenes uniquely between processes"
            )

        if len(scenes) < num_processes:
            raise RuntimeError(
                "reduce the number of processes as there "
                "aren't enough number of scenes"
            )

        random.shuffle(scenes)

    scene_splits: List[List[str]] = [[] for _ in range(num_processes)]
    for idx, scene in enumerate(scenes):
        scene_splits[idx % len(scene_splits)].append(scene)

    assert sum(map(len, scene_splits)) == len(scenes)

    for i in range(num_processes):
        proc_config = config.clone()
        proc_config.defrost()

        task_config = proc_config.TASK_CONFIG
        task_config.SEED = task_config.SEED + i
        if len(scenes) > 0:
            task_config.DATASET.CONTENT_SCENES = scene_splits[i]

        task_config.SIMULATOR.HABITAT_SIM_V0.GPU_DEVICE_ID = (
            config.SIMULATOR_GPU_ID
        )

        task_config.SIMULATOR.AGENT_0.SENSORS = config.SENSORS

        proc_config.freeze()
        configs.append(proc_config)

    envs = habitat.VectorEnv(
        make_env_fn=make_env_fn,
        env_fn_args=tuple(zip(configs, env_classes)),
        workers_ignore_signals=workers_ignore_signals,
        auto_reset_done=auto_reset_done
    )
    return envs


def construct_ddp_envs(
    config: Config,
    env_class: Union[Type[Env], Type[RLEnv]],
    workers_ignore_signals: bool = False,
    world_rank: int = 0,
    world_size: int = 0,
    auto_reset_done: bool = True,
) -> VectorEnv:
    r"""Create VectorEnv object with specified config and env class type.
    To allow better performance, dataset are split into small ones for
    each individual env, grouped by scenes.

    :param config: configs that contain num_processes as well as information
    :param necessary to create individual environments.
    :param env_class: class type of the envs to be created.
    :param workers_ignore_signals: Passed to :ref:`habitat.VectorEnv`'s constructor

    :return: VectorEnv object created according to specification.
    """

    num_processes = config.NUM_PROCESSES
    configs = []
    env_classes = [env_class for _ in range(num_processes)]
    dataset = make_dataset(config.TASK_CONFIG.DATASET.TYPE)
    scenes = config.TASK_CONFIG.DATASET.CONTENT_SCENES
    if "*" in config.TASK_CONFIG.DATASET.CONTENT_SCENES:
        scenes = dataset.get_scenes_to_load(config.TASK_CONFIG.DATASET)

    if num_processes > 1:
        if len(scenes) == 0:
            raise RuntimeError(
                "No scenes to load, multiple process logic relies on being able to split scenes uniquely between processes"
            )

        if len(scenes) < num_processes:
            raise RuntimeError(
                "reduce the number of processes as there "
                "aren't enough number of scenes"
            )

        # random.shuffle(scenes)
        num_scenes_per_node =  len(scenes) // world_size
        start_idx = world_rank * num_scenes_per_node
        end_idx = start_idx + num_scenes_per_node
        if world_rank == (world_size - 1):
            scenes = scenes[start_idx:]
        else:
            scenes = scenes[start_idx : end_idx]
        logger.info("world rank: {}, scenes: {} - {}".format(world_rank, len(scenes), scenes))


    scene_splits: List[List[str]] = [[] for _ in range(num_processes)]
    for idx, scene in enumerate(scenes):
        scene_splits[idx % len(scene_splits)].append(scene)

    assert sum(map(len, scene_splits)) == len(scenes)

    for i in range(num_processes):
        proc_config = config.clone()
        proc_config.defrost()

        task_config = proc_config.TASK_CONFIG
        task_config.SEED = task_config.SEED + i
        if len(scenes) > 0:
            task_config.DATASET.CONTENT_SCENES = scene_splits[i]

        task_config.SIMULATOR.HABITAT_SIM_V0.GPU_DEVICE_ID = (
            config.SIMULATOR_GPU_ID
        )

        task_config.SIMULATOR.AGENT_0.SENSORS = config.SENSORS

        proc_config.freeze()
        configs.append(proc_config)

    envs = habitat.VectorEnv(
        make_env_fn=make_env_fn,
        env_fn_args=tuple(zip(configs, env_classes)),
        workers_ignore_signals=workers_ignore_signals,
        auto_reset_done=auto_reset_done
    )
    return envs

habitat_baselines/utils/env_utils.py
cshizhe/onav_rim
#!/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 cv2
import numpy as np
import torch

from habitat_baselines.utils.common import (
    tensor_to_bgr_images,
    tensor_to_depth_images,
)

try:
    from habitat_sim.utils.common import d3_40_colors_rgb
except ImportError:
    d3_40_colors_rgb = None

from PIL import Image


def save_rgb_results(
    gt_rgb: torch.Tensor, pred_rgb: torch.Tensor, path: str
) -> None:
    r"""For saving RGB reconstruction results during EQA-CNN-Pretrain eval.

    Args:
        gt_rgb: RGB ground truth tensor
        pred_rgb: RGB reconstruction tensor
        path: to save images
    """
    path = path.format(split="val", type="rgb")
    gt_bgr_o, pred_bgr = tensor_to_bgr_images([gt_rgb, pred_rgb])
    cv2.imwrite(path + "_gt.jpg", gt_bgr_o)
    cv2.imwrite(path + "_pred.jpg", pred_bgr)


def save_seg_results(
    gt_seg: torch.Tensor, pred_seg: torch.Tensor, path: str
) -> None:
    r"""For saving predicted and ground truth seg maps during
    EQA-CNN-Pretrain eval.

    Args:
        gt_seg: ground truth segmentation tensor
        pred_seg: ouput segmentation tensor
        path: to save images
    """

    path = path.format(split="val", type="seg")

    gt_seg = gt_seg.cpu().numpy() % 40
    pred_seg = torch.argmax(pred_seg, 0).cpu().numpy() % 40

    gt_seg_colored = d3_40_colors_rgb[gt_seg]
    pred_seg_colored = d3_40_colors_rgb[pred_seg]

    cv2.imwrite(path + "_gt.jpg", gt_seg_colored)
    cv2.imwrite(path + "_pred.jpg", pred_seg_colored)


def save_depth_results(
    gt_depth: torch.Tensor, pred_depth: torch.Tensor, path: str
) -> None:
    r"""For saving predicted and ground truth depth maps during
    EQA-CNN-Pretrain eval.

    Args:
        gt_depth: ground truth depth tensor
        pred_depth: ouput depth tensor
        path: to save images
    """
    path = path.format(split="val", type="depth")

    gt_depth, pred_depth = tensor_to_depth_images([gt_depth, pred_depth])

    cv2.imwrite(path + "_gt.jpg", gt_depth)
    cv2.imwrite(path + "_pred.jpg", pred_depth)


def save_frame(
    frame, path: str
) -> None:
    r"""For saving RGB reconstruction results during EQA-CNN-Pretrain eval.

    Args:
        gt_rgb: RGB ground truth tensor
        pred_rgb: RGB reconstruction tensor
        path: to save images
    """
    im = Image.fromarray(frame)
    im.save(path + ".png")


def put_vqa_text_on_image(
    image: np.ndarray,
    question: str,
    prediction: str,
    ground_truth: str,
) -> np.ndarray:
    r"""For writing VQA question, prediction and ground truth answer
        on image.
    Args:
        image: image on which text has to be written
        question: input question to model
        prediction: model's answer prediction
        ground_truth: ground truth answer
    Returns:
        image with text
    """
    font = cv2.FONT_HERSHEY_SIMPLEX
    color = (0, 0, 0)
    scale = 0.4
    thickness = 1

    cv2.putText(
        image,
        "Question: " + question,
        (10, 15),
        font,
        scale,
        color,
        thickness,
    )
    cv2.putText(
        image,
        "Prediction: " + prediction,
        (10, 30),
        font,
        scale,
        color,
        thickness,
    )
    cv2.putText(
        image,
        "Ground truth: " + ground_truth,
        (10, 45),
        font,
        scale,
        color,
        thickness,
    )

    return image


def save_vqa_image_results(
    images_tensor: torch.Tensor,
    question: str,
    prediction: str,
    ground_truth: str,
    path: str,
) -> None:
    r"""For saving VQA input images with input question and predicted answer.
    Being used to save model predictions during eval.
    Args:
        images_tensor: images' tensor containing input frames
        question: input question to model
        prediction: model's answer prediction
        ground_truth: ground truth answer
        path: to save images
    Returns:
        None
    """

    images = tensor_to_bgr_images(images_tensor)

    collage_image = cv2.hconcat(images)
    collage_image = cv2.copyMakeBorder(
        collage_image,
        55,
        0,
        0,
        0,
        cv2.BORDER_CONSTANT,
        value=(255, 255, 255),
    )

    image = put_vqa_text_on_image(
        collage_image, question, prediction, ground_truth
    )

    cv2.imwrite(path, image)

habitat_baselines/utils/visualizations/utils.py
cshizhe/onav_rim
from typing import List, Optional, Union

import numpy as np

from habitat.config import Config as CN

CONFIG_FILE_SEPARATOR = ','

# -----------------------------------------------------------------------------
# EXPERIMENT CONFIG
# -----------------------------------------------------------------------------
_C = CN()
_C.SEED = 100
_C.CMD_TRAILING_OPTS = []  # store command line options as list of strings

# -----------------------------------------------------------------------------
# MODEL
# -----------------------------------------------------------------------------
_C.MODEL = CN()
_C.MODEL.MAX_EPISODE_STEPS = 1000

# -----------------------------------------------------------------------------
# DATASET
# -----------------------------------------------------------------------------
_C.DATASET = CN()
_C.DATASET.MAX_EPISODE_STEPS = 1000
# _C.DATASET.num_ft_views = 1

def get_config(
    config_paths: Optional[Union[List[str], str]] = None,
    opts: Optional[list] = None,
) -> CN:
    r"""Create a unified config with default values overwritten by values from
    :ref:`config_paths` and overwritten by options from :ref:`opts`.

    Args:
        config_paths: List of config paths or string that contains comma
        separated list of config paths.
        opts: Config options (keys, values) in a list (e.g., passed from
        command line into the config. For example, ``opts = ['FOO.BAR',
        0.5]``. Argument can be used for parameter sweeping or quick tests.
    """
    config = _C.clone()
    if config_paths:
        if isinstance(config_paths, str):
            if CONFIG_FILE_SEPARATOR in config_paths:
                config_paths = config_paths.split(CONFIG_FILE_SEPARATOR)
            else:
                config_paths = [config_paths]

        for config_path in config_paths:
            config.merge_from_file(config_path)
            
    if opts:
        config.CMD_TRAILING_OPTS = config.CMD_TRAILING_OPTS + opts
        config.merge_from_list(config.CMD_TRAILING_OPTS)

    config.freeze()
    return config

offline_bc/config/default.py
cshizhe/onav_rim
import os
import json
import jsonlines
import numpy as np
import cv2
import time
import math
import random
import collections

import lmdb
import msgpack
import msgpack_numpy
msgpack_numpy.patch()

import torch
import torch.nn.functional as F
from torch.utils.data import Dataset
from torch.nn.utils.rnn import pad_sequence

from offline_bc.utils.logger import LOGGER
from offline_bc.utils.ops import pad_tensors, gen_seq_masks
THDA_OBJECTS = set(['foodstuff', 'fruit', 'plaything', 'game_equipment', 'hand_tool', 'kitchenware', 'stationery'])

MAP_IMAGE_SIZE = 960

def crop_centered_local_map(
    agent_loc, global_map, local_sizes, 
    keep_centered_when_out_of_boundary=False,
    channel_before_size=False
):
    '''Crop a local map with local_sizes in the global_map centered at the agent's location.
    keep_centered_when_out_of_boundary: if True, we will pad zeros in the local map
    '''
    loc_r, loc_c = agent_loc
    local_w, local_h = local_sizes
    if channel_before_size:
        n_channel, full_w, full_h = global_map.shape
    else:
        full_w, full_h, n_channel = global_map.shape

    assert (local_w <= full_w) & (local_h <= full_h)

    gx1, gy1 = loc_r - local_w // 2, loc_c - local_h // 2
    gx2, gy2 = gx1 + local_w, gy1 + local_h
    
    if keep_centered_when_out_of_boundary:
        if isinstance(global_map, np.ndarray):
            if channel_before_size:
                local_map = np.zeros((n_channel, local_w, local_h), dtype=global_map.dtype)
            else:
                local_map = np.zeros((local_w, local_h, n_channel), dtype=global_map.dtype)
        elif isinstance(global_map, torch.Tensor):
            if channel_before_size:
                local_map = torch.zeros(
                    n_channel, local_w, local_h,
                    dtype=global_map.dtype, device=global_map.device
                )
            else:
                local_map = torch.zeros(
                    local_w, local_h, n_channel, 
                    dtype=global_map.dtype, device=global_map.device
                )
    
        lx1 = max(- gx1, 0)
        ly1 = max(- gy1, 0)
        if lx1 >= local_w or ly1 >= local_h:
            return local_map
        lx2 = lx1 + local_w - max(0, gx2 - full_w)
        ly2 = ly1 + local_h - max(0, gy2 - full_h)
        if lx2 <= 0 or ly2 <= 0:
            return local_map

        gx1, gy1 = max(0, gx1), max(0, gy1)
        gx2, gy2 = min(full_w, gx2), min(full_h, gy2)
       
        if channel_before_size:
            local_map[:, lx1: lx2, ly1: ly2] = global_map[:, gx1: gx2, gy1: gy2]
        else:
            local_map[lx1: lx2, ly1: ly2] = global_map[gx1: gx2, gy1: gy2]

    else:
        if gx1 < 0:
            gx1, gx2 = 0, local_w
        if gx2 > full_w:
            gx1, gx2 = full_w - local_w, full_w

        if gy1 < 0:
            gy1, gy2 = 0, local_h
        if gy2 > full_h:
            gy1, gy2 = full_h - local_h, full_h

        if channel_before_size:
            local_map = global_map[:, gx1: gx2, gy1: gy2]
        else:
            local_map = global_map[gx1: gx2, gy1: gy2]

    return local_map

class NavDemoDataset(Dataset):
    def __init__(
        self, trn_scene_ids, rgb_image_dir=None, rgb_ft_dir=None, 
        depth_ft_dir=None, sem_ft_dir=None, meta_dir=None, semseg_dir=None,
        rgb_image_size=224, max_steps=500, inflection_weight=None,
        num_ft_views=1, num_history_fts=None, use_thda=True, 
        val_scene_ids=[], validation=False, map_dir=None, 
        map_image_size=240, **kwargs
    ) -> None:
        if validation:
            self.scene_ids = val_scene_ids
        else:
            if trn_scene_ids == '*':
                self.scene_ids = [os.path.splitext(x)[0] for x in os.listdir(meta_dir)]
            else:
                self.scene_ids = trn_scene_ids
            # if len(val_scene_ids) > 0:
            #     self.scene_ids = [x for x in self.scene_ids if x not in val_scene_ids]

        self.rgb_image_dir = rgb_image_dir
        self.rgb_ft_dir = rgb_ft_dir
        self.depth_ft_dir = depth_ft_dir
        self.semseg_dir = semseg_dir
        self.sem_ft_dir = sem_ft_dir
        self.meta_dir = meta_dir
        self.rgb_image_size = rgb_image_size
        self.max_steps = max_steps
        self.num_history_fts = num_history_fts if num_history_fts is not None else max_steps
        self.inflection_weight = inflection_weight
        self.num_ft_views = num_ft_views
        self.use_thda = use_thda
        self.kwargs = kwargs

        self._all_lmdb_envs = []

        self.meta_info_envs, self.meta_info_txns = self.build_lmdb_envs(
            self.meta_dir, self.scene_ids
        )
        if rgb_image_dir is not None:
            self.rgb_image_envs, self.rgb_image_txns = self.build_lmdb_envs(rgb_image_dir, self.scene_ids)
        if rgb_ft_dir is not None:
            if isinstance(rgb_ft_dir, str):
                rgb_ft_dir = [rgb_ft_dir]
            self.rgb_ft_envs = collections.defaultdict(list)
            self.rgb_ft_txns = collections.defaultdict(list)
            for x in rgb_ft_dir:
                rgb_ft_envs, rgb_ft_txns = self.build_lmdb_envs(x, self.scene_ids)
                for k in rgb_ft_envs.keys():
                    self.rgb_ft_envs[k].append(rgb_ft_envs[k])
                    self.rgb_ft_txns[k].append(rgb_ft_txns[k])
        if depth_ft_dir is not None:
            if isinstance(depth_ft_dir, str):
                depth_ft_dir = [depth_ft_dir]
            self.depth_ft_envs = collections.defaultdict(list)
            self.depth_ft_txns = collections.defaultdict(list)
            for x in depth_ft_dir:
                depth_ft_envs, depth_ft_txns = self.build_lmdb_envs(x, self.scene_ids)
                for k in depth_ft_envs.keys():
                    self.depth_ft_envs[k].append(depth_ft_envs[k])
                    self.depth_ft_txns[k].append(depth_ft_txns[k])
        if semseg_dir is not None:
            self.semseg_envs, self.semseg_txns = self.build_lmdb_envs(semseg_dir, self.scene_ids)
        if sem_ft_dir is not None:
            self.sem_ft_envs, self.sem_ft_txns = self.build_lmdb_envs(sem_ft_dir, self.scene_ids)

        self.load_dataset_meta_infos()

        self.map_dir = map_dir
        self.map_image_size = map_image_size
        if self.map_dir is not None:
            self.data_idxs, self.num_episodes = [], 0
            self.visited_locs_in_full_map = {}
            for scene_id in self.scene_ids:
                self.visited_locs_in_full_map[scene_id] = {}
                with jsonlines.open(os.path.join(map_dir, 'meta_infos', f'{scene_id}.jsonl'), 'r') as f:
                    for item in f:
                        self.visited_locs_in_full_map[scene_id][item['episode_id']] = item['visited_locs_in_full_map']
                        self.data_idxs.append((scene_id, item['episode_id']))
                        self.num_episodes += 1
            resized_map_dir = os.path.join(self.map_dir, 'infer_map_%s_size_%d_pool_%d'%(
                '_'.join(self.kwargs['map_types']), self.map_image_size, 
                self.kwargs['map_image_max_pool'])
            )
            if os.path.exists(resized_map_dir):
                self.resized_map_dir = resized_map_dir
                map_env_dir = resized_map_dir
            else:
                self.resized_map_dir = None
                map_env_dir = os.path.join(map_dir, 'maps')
            self.map_envs, self.map_txns = self.build_lmdb_envs(map_env_dir, self.scene_ids)
            LOGGER.info(f'MAP: #scenes: {len(self.scene_ids)}, #episodes: {self.num_episodes}')

    def load_dataset_meta_infos(self):
        self.data_idxs, self.num_episodes = [], 0

        for scene_id in self.scene_ids:
            for episode_id, item in self.meta_info_txns[scene_id].cursor():
                episode_id = episode_id.decode('ascii')
                item = msgpack.unpackb(item)
                if (not self.use_thda) and (item['object_category'] in THDA_OBJECTS):
                    continue
                if len(item['demonstration']) > self.kwargs.get('max_episode_steps', float('inf')):
                    continue
                self.data_idxs.append((scene_id, episode_id))
                self.num_episodes += 1

        LOGGER.info(f'#scenes: {len(self.scene_ids)}, #episodes: {self.num_episodes}')

    def build_lmdb_envs(self, lmdb_dir, scene_ids):
        lmdb_envs, lmdb_txns = {}, {}
        for scene_id in scene_ids:
            lmdb_envs[scene_id] = lmdb.open(
                os.path.join(lmdb_dir, scene_id), 
                readonly=True, lock=False, readahead=False, meminit=False
            )
            lmdb_txns[scene_id] = lmdb_envs[scene_id].begin()
            self._all_lmdb_envs.append(lmdb_envs[scene_id])
        return lmdb_envs, lmdb_txns

    def load_map_at_step(self, txn, episode_id, step, map_names=['obs', 'exp']):
        imgs = []
        for map_name in map_names:
            key = f'{episode_id}_{step}_{map_name}'
            value = txn.get(key.encode('ascii'))
            if value is None:
                img = np.zeros((MAP_IMAGE_SIZE, MAP_IMAGE_SIZE), dtype=np.uint8)
                LOGGER.info(f'DATASET ERROR: {key} not exists')
            else:
                img = np.frombuffer(value, dtype=np.uint8)
                img = cv2.imdecode(img, cv2.IMREAD_GRAYSCALE)   # (MAP_IMAGE_SIZE, MAP_IMAGE_SIZE)
            imgs.append(img)
        imgs = np.stack(imgs, axis=2)
        return imgs
    
    def __del__(self):
        for lmdb_env in self._all_lmdb_envs:
            lmdb_env.close()

    def __len__(self):
        return self.num_episodes

    def load_episode_meta_info(self, txn, episode_id):
        key = episode_id.encode('ascii')
        value = txn.get(key)
        value = msgpack.unpackb(value)
        return value

    def load_episode_images(self, txn, episode_id, im_height, im_width, im_channel):
        key = episode_id.encode('ascii')
        value = txn.get(key)
        assert value is not None, episode_id
        images_flt = np.frombuffer(value, dtype=np.uint8)
        images_flt = cv2.imdecode(
            images_flt, 
            cv2.IMREAD_GRAYSCALE if im_channel == 1 else cv2.IMREAD_COLOR
        )
        images = images_flt.reshape(-1, im_height, im_width, im_channel)
        return images

    def load_episode_fts(self, txn, episode_id):
        key = episode_id.encode('ascii')
        value = txn.get(key)
        assert value is not None, episode_id
        fts = msgpack.unpackb(value)
        if len(fts.shape) == 3: # (nsteps, num_views, dim)
            if self.num_ft_views == 1:
                fts = fts[:, 0]
            elif self.num_ft_views == 12: # horizontal pano
                fts = fts[:, :12]
            elif self.num_ft_views == 3: # up and down
                fts = np.concatenate([fts[:, :1], fts[:, -2:]], axis=1)
            elif self.num_ft_views == 2: # down
                fts = np.concatenate([fts[:, :1], fts[:, -1:]], axis=1)
            elif self.num_ft_views == '2_up':
                fts = np.concatenate([fts[:, :1], fts[:, -2:-1]], axis=1)
        # BUG: The given NumPy array is not writeable if without np.array()
        fts = np.array(fts, dtype=np.float32)
        return fts

    def get_full_episode(self, scene_id, episode_id, meta):
        outs = {
            'episode_ids': episode_id,
            'scene_ids': scene_id,
            'objectgoal': meta['objectgoal'],
            'object_category': meta['object_category'],
            'compass': np.array(meta['compass'], dtype=np.float32), # (nsteps, 1)
            'gps': np.array(meta['gps'], dtype=np.float32), # (nsteps, 2)
            'demonstration': np.array(meta['demonstration'], dtype=np.int64), # (nsteps, )
        }

        if self.rgb_image_dir is not None:
            outs['rgb'] = self.load_episode_images(
                self.rgb_image_txns[scene_id], episode_id,
                self.rgb_image_size, self.rgb_image_size, 3
            )
        if self.rgb_ft_dir is not None:
            ridx = np.random.randint(len(self.rgb_ft_txns[scene_id]))
            outs['rgb_features'] = self.load_episode_fts(
                self.rgb_ft_txns[scene_id][ridx], episode_id
            )
        if self.depth_ft_dir is not None:
            ridx = np.random.randint(len(self.depth_ft_txns[scene_id]))
            outs['depth_features'] = self.load_episode_fts(
                self.depth_ft_txns[scene_id][ridx], episode_id
            )
        if self.sem_ft_dir is not None:
            sem_fts = self.load_episode_fts(
                self.sem_ft_txns[scene_id], episode_id
            )
            if meta['objectgoal'] < 21:
                sem_goal_fts = sem_fts[:, meta['objectgoal']][:, None]
            else:
                sem_goal_fts = np.zeros((sem_fts.shape[0], 1), dtype=np.float32)
            outs['sem_features'] = np.concatenate(
                [sem_fts, sem_goal_fts], 1
            )
        if self.semseg_dir is not None:
            outs['semantic'] = self.load_episode_images(
                self.semseg_txns[scene_id], episode_id,
                480, 640, 1
            ) - 1 # [1, 29] -> [0, 28]

        num_steps = len(meta['compass'])

        outs['inflection_weight'] = np.ones((num_steps, ), dtype=np.float32)
        if self.inflection_weight is not None:
            prev_actions = np.array([0] + meta['demonstration'][:-1])
            actions = np.array(meta['demonstration'])
            outs['inflection_weight'][prev_actions != actions] = self.inflection_weight
        stop_weight = self.kwargs.get('stop_weight', None)
        if stop_weight is not None:
            outs['inflection_weight'][-1] = stop_weight

        return outs

    def __getitem__(self, data_idx):
        scene_id, episode_id = self.data_idxs[data_idx]
        meta = self.load_episode_meta_info(self.meta_info_txns[scene_id], episode_id)

        outs = self.get_full_episode(scene_id, episode_id, meta)
        num_steps = len(meta['compass'])

        if num_steps > self.num_history_fts:
            even_start_idxs = np.arange(0, num_steps-self.num_history_fts+1, self.num_history_fts).tolist() + \
                              np.arange(num_steps-self.num_history_fts, 0, -self.num_history_fts).tolist()
            even_start_idxs = list(set(even_start_idxs))
            start_step = np.random.choice(even_start_idxs)
            if self.kwargs.get('no_random_start_step', False):
                start_step = 0
            end_step = start_step + self.num_history_fts
        else:
            start_step, end_step = 0, num_steps

        if self.kwargs.get('infer_visual_feature_task', False):
            max_future_step_size = self.kwargs.get('max_future_step_size', 100)
            infer_time_ids = np.array([
                np.random.randint(start_step, min(num_steps, t+max_future_step_size)) \
                    for t in range(start_step, end_step)
            ])
            outs['infer_gps'] = outs['gps'][infer_time_ids]
            outs['infer_compass'] = outs['compass'][infer_time_ids]
            outs['infer_visual_features'] = outs['rgb_features'][infer_time_ids]
            if self.kwargs.get('infer_depth_feature', False):
                outs['infer_visual_features'] = np.concatenate(
                    [outs['infer_visual_features'], outs['depth_features'][infer_time_ids]],
                    axis=-1
                )

        if self.kwargs.get('infer_local_map_task', False):
            map_types = self.kwargs.get('map_types', ['obs'])
            max_pool_size = self.kwargs.get('map_image_max_pool', 1)
            visited_locs_in_full_map = self.visited_locs_in_full_map[scene_id][episode_id]
            if self.resized_map_dir is None:
                local_maps = []
                for t in range(start_step, end_step):
                    local_maps.append(
                        crop_centered_local_map(
                            visited_locs_in_full_map[t], 
                            self.load_map_at_step(self.map_txns[scene_id], episode_id, t, map_types),
                            (self.map_image_size, self.map_image_size),
                            keep_centered_when_out_of_boundary=True
                        )
                    )
                local_maps = np.stack(local_maps, 0)
                local_maps = torch.from_numpy(local_maps) # (nsteps, height, width, nchannels)
                local_maps = F.max_pool2d(
                    local_maps.permute(0, 3, 1, 2).float(), max_pool_size, max_pool_size
                ).bool().view(end_step - start_step, len(map_types), -1)
            else:
                local_maps = msgpack.unpackb(self.map_txns[scene_id].get(episode_id.encode('ascii')))
                local_maps = local_maps[start_step: end_step]
                local_maps = torch.from_numpy(local_maps) # (nsteps, height, width, nchannels)
            outs['infer_local_maps'] = local_maps

        if num_steps > self.num_history_fts:
            for k, v in outs.items():
                if isinstance(v, np.ndarray) and (not k.startswith('infer')):
                    outs[k] = v[start_step: end_step]

        outs['step_ids'] = np.arange(start_step, end_step)
        if outs['step_ids'][-1] >= self.max_steps:
            outs['step_ids'] = outs['step_ids'] - start_step
        outs['ori_start_step'] = start_step
        outs['num_steps'] = min(self.num_history_fts, num_steps)

        gpscompass_noise_type = self.kwargs.get('gpscompass_noise_type', None)
        if gpscompass_noise_type == 'zero':
            outs['gps'][:] = 0
            outs['compass'][:] = 0
        elif gpscompass_noise_type == 'gaussian':
            outs['gps'] = outs['gps'] + np.random.normal(0, 1, size=outs['gps'].shape)
            outs['compass'] = outs['compass'] + np.random.normal(0, 1, size=outs['compass'].shape)

        return outs


def collate_fn(inputs):
    batch = {
        k: [x[k] for x in inputs] for k in inputs[0].keys()
    }

    max_steps = np.max(batch['num_steps'])
    # batch['num_steps'] = torch.LongTensor(batch['num_steps'])

    if 'rgb' in batch:
        # (batch x nsteps, 224, 224, 3)
        batch['rgb'] = np.concatenate(batch['rgb'], 0)
    if 'semantic' in batch:
        # (batch x nsteps, 480, 640, 1)
        batch['semantic'] = np.concatenate(batch['semantic'], 0)
    
    # for key in ['rgb_features', 'depth_features', 'compass', 'gps', \
    #             'infer_gps', 'infer_compass', 'infer_visual_features']:
    #     if key in batch:
    #         # (batch, max_steps, dim_ft)
    #         batch[key] = pad_tensors(
    #             [torch.FloatTensor(x) for x in batch[key]]
    #         )
    # if batch['infer_gps'].size(1) != batch['rgb_features'].size(1):
    #     print(batch['infer_gps'].size(), batch['rgb_features'].size())

    batch['objectgoal'] = torch.LongTensor(batch['objectgoal'])
    batch['demonstration'] = pad_sequence(
        [torch.LongTensor(x) for x in batch['demonstration']], 
        batch_first=True, padding_value=-100
    )
    batch['inflection_weight'] = pad_sequence(
        [torch.FloatTensor(x) for x in batch['inflection_weight']], 
        batch_first=True, padding_value=0
    )
    batch['step_ids'] = pad_sequence(
        [torch.from_numpy(x) for x in batch['step_ids']],
        batch_first=True, padding_value=0
    )
    return batch

offline_bc/data/dataset.py
cshizhe/onav_rim
"""
Copyright (c) Microsoft Corporation.
Licensed under the MIT license.

A prefetch loader to speedup data loading
Modified from Nvidia Deep Learning Examples
(https://github.com/NVIDIA/DeepLearningExamples/tree/master/PyTorch).
"""
import random
from typing import List, Dict, Tuple, Union, Iterator

import torch
from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
from torch.utils.data.distributed import DistributedSampler
import torch.distributed as dist


class MetaLoader:
    """wraps multiple data loaders"""

    def __init__(
        self, loaders, accum_steps: int = 1, distributed: bool = False, device=None
    ):
        assert isinstance(loaders, dict)
        self.name2loader = {}
        self.name2iter = {}
        self.name2pre_epoch = {}
        self.names: List[str] = []
        ratios: List[int] = []
        for n, l in loaders.items():
            if isinstance(l, tuple):
                l, r, p = l
            elif isinstance(l, DataLoader):
                r = 1
                p = lambda e: None
            else:
                raise ValueError()
            self.names.append(n)
            self.name2loader[n] = l
            self.name2iter[n] = iter(l)
            self.name2pre_epoch[n] = p
            ratios.append(r)

        self.accum_steps = accum_steps
        self.device = device
        self.sampling_ratios = torch.tensor(ratios).float().to(self.device)
        self.distributed = distributed
        self.step = 0

    def __iter__(self) -> Iterator[Tuple]:
        """this iterator will run indefinitely"""
        task_id = None
        epoch_id = 0
        while True:
            if self.step % self.accum_steps == 0:
                task_id = torch.multinomial(self.sampling_ratios, 1)
                if self.distributed:
                    # make sure all process is training same task
                    dist.broadcast(task_id, 0)
            self.step += 1
            task = self.names[task_id.cpu().item()]
            iter_ = self.name2iter[task]
            try:
                batch = next(iter_)
            except StopIteration:
                epoch_id += 1
                # In distributed mode, calling the set_epoch() method at the beginning of each epoch 
                # before creating the DataLoader iterator is necessary to make shuffling work properly 
                # across multiple epochs. Otherwise, the same ordering will be always used.
                self.name2pre_epoch[task](epoch_id)
                iter_ = iter(self.name2loader[task])
                batch = next(iter_)
                self.name2iter[task] = iter_

            yield task, batch


def move_to_cuda(batch: Union[List, Tuple, Dict, torch.Tensor], device: torch.device):
    if isinstance(batch, torch.Tensor):
        return batch.to(device, non_blocking=True)
    elif isinstance(batch, list):
        return [move_to_cuda(t, device) for t in batch]
    elif isinstance(batch, tuple):
        return tuple(move_to_cuda(t, device) for t in batch)
    elif isinstance(batch, dict):
        return {n: move_to_cuda(t, device) for n, t in batch.items()}
    return batch


class PrefetchLoader(object):
    """
    overlap compute and cuda data transfer
    """
    def __init__(self, loader, device: torch.device):
        self.loader = loader
        self.device = device

    def __iter__(self):
        loader_it = iter(self.loader)
        self.preload(loader_it)
        batch = self.next(loader_it)
        while batch is not None:
            yield batch
            batch = self.next(loader_it)

    def __len__(self):
        return len(self.loader)

    def preload(self, it):
        try:
            self.batch = next(it)
        except StopIteration:
            self.batch = None
            return
        self.batch = move_to_cuda(self.batch, self.device)

    def next(self, it):
        batch = self.batch
        self.preload(it)
        return batch

    def __getattr__(self, name):
        method = self.loader.__getattribute__(name)
        return method


def build_dataloader(dataset, collate_fn, is_train: bool, opts):

    batch_size = opts.train_batch_size if is_train else opts.val_batch_size

    if opts.local_rank == -1:
        if is_train:
            sampler: Union[
                RandomSampler, SequentialSampler, DistributedSampler
            ] = RandomSampler(dataset)
        else:
            sampler = SequentialSampler(dataset)

        size = torch.cuda.device_count() if torch.cuda.is_available() else 1
        pre_epoch = lambda e: None

        # DataParallel: scale the batch size by the number of GPUs
        if size > 1:
            batch_size *= size

    else:
        size = dist.get_world_size()
        sampler = DistributedSampler(
            dataset, num_replicas=size, rank=dist.get_rank(), 
            shuffle=is_train
        )
        pre_epoch = sampler.set_epoch

    loader = DataLoader(
        dataset,
        sampler=sampler,
        batch_size=batch_size,
        num_workers=opts.n_workers,
        pin_memory=opts.pin_mem,
        collate_fn=collate_fn,
        drop_last=False,
        prefetch_factor=2,
    )

    return loader, pre_epoch

offline_bc/data/loader.py
cshizhe/onav_rim
import os
import jsonlines
import numpy as np
import cv2
import time

import torch

from offline_bc.utils.logger import LOGGER
from offline_bc.data.dataset import (
    NavDemoDataset, collate_fn
)
from habitat_baselines.common.nav_mapping import get_local_map_boundaries

MAP_IMAGE_SIZE = 480


class NavMapDemoDataset(NavDemoDataset):
    def __init__(
        self, scene_ids, rgb_image_dir=None, rgb_ft_dir=None, 
        depth_ft_dir=None, meta_dir=None, semseg_dir=None,
        rgb_image_size=224, max_steps=500, inflection_weight=None,
        num_ft_views=1, use_thda=True, map_dir=None, map_image_size=240,
        **kwargs
    ) -> None:
        super().__init__(
            scene_ids, rgb_image_dir=rgb_image_dir, rgb_ft_dir=rgb_ft_dir,
            depth_ft_dir=depth_ft_dir, meta_dir=meta_dir, semseg_dir=semseg_dir,
            rgb_image_size=rgb_image_size, max_steps=max_steps, 
            inflection_weight=inflection_weight,
            num_ft_views=num_ft_views, use_thda=use_thda, **kwargs
        )
        self.map_dir = map_dir
        self.map_image_size = map_image_size

        self.visited_locs_in_full_map = {}
        for scene_id in self.scene_ids:
            self.visited_locs_in_full_map[scene_id] = {}
            with jsonlines.open(os.path.join(map_dir, 'meta_infos', f'{scene_id}.jsonl'), 'r') as f:
                for item in f:
                    self.visited_locs_in_full_map[scene_id][item['episode_id']] = item['visited_locs_in_full_map']

        self.map_envs, self.map_txns = self.build_lmdb_envs(
            os.path.join(map_dir, 'maps'), self.scene_ids
        ) 

    def load_map_at_step(self, txn, episode_id, step):
        imgs = []
        for map_name in ['obs', 'exp']:
            key = f'{episode_id}_{step}_{map_name}'
            value = txn.get(key.encode('ascii'))
            if value is None:
                img = np.zeros((MAP_IMAGE_SIZE, MAP_IMAGE_SIZE), dtype=np.uint8)
                LOGGER.info(f'DATASET ERROR: {key} not exists')
            else:
                img = np.frombuffer(value, dtype=np.uint8)
                img = cv2.imdecode(img, cv2.IMREAD_GRAYSCALE)   # (MAP_IMAGE_SIZE, MAP_IMAGE_SIZE)
            imgs.append(img)
        imgs = np.stack(imgs, axis=2)
        return imgs

    def plot_traj_image(self, traj_img, visited_locs_in_full_map, start_step, end_step=None):
        if end_step is None:
            end_step = start_step + 1
        for t in range(start_step, end_step):
            r, c = visited_locs_in_full_map[t]
            r = min(max(0, r), MAP_IMAGE_SIZE - 1)
            c = min(max(0, c), MAP_IMAGE_SIZE - 1)
            traj_img[max(0, r-1):r+2, max(0, c-1):c+2] = t + 1
        return traj_img

    def __getitem__(self, data_idx):
        outs = super().__getitem__(data_idx)

        scene_id, idx = self.data_idxs[data_idx]
        meta = self.meta_infos[scene_id][idx]
        episode_id = meta['episode_id']

        visited_locs_in_full_map = self.visited_locs_in_full_map[scene_id][episode_id]
        start_step, end_step = outs['start_step'], outs['end_step']

        # traj_img = np.zeros((MAP_IMAGE_SIZE, MAP_IMAGE_SIZE), dtype=np.float32) # memory costly
        # traj_img = self.plot_traj_image(traj_img, visited_locs_in_full_map, 0, start_step)

        outs['maps'] = []
        for step in range(start_step, end_step):
            map_imgs = self.load_map_at_step(
                self.map_txns[scene_id], episode_id, step, 
            )
            x1, x2, y1, y2 = get_local_map_boundaries(
                visited_locs_in_full_map[step], (self.map_image_size, self.map_image_size),
                (MAP_IMAGE_SIZE, MAP_IMAGE_SIZE), global_downscaling=2
            )

            # traj_img = self.plot_traj_image(traj_img, visited_locs_in_full_map, start_step)
            # normed_traj_img = traj_img / (step + 1)
            # img = np.concatenate([map_imgs, normed_traj_img[:, :, None]], axis=2)[x1:x2, y1:y2]

            img = map_imgs[x1:x2, y1:y2]
            outs['maps'].append(img)
        outs['maps'] = np.stack(outs['maps'], 0)
        outs['visited_locs_in_full_map'] = visited_locs_in_full_map[:end_step]

        return outs


def map_collate_fn(inputs):
    batch = collate_fn(inputs)
    
    # (batch x nsteps, 240, 240, 3)
    batch['maps'] = np.concatenate(batch['maps'], 0)
    
    return batch


if __name__ == '__main__':
    from torch.utils.data import DataLoader

    input_dir = 'data/datasets/objectnav/mp3d_70k_demos'
    map_dir = 'data/datasets/objectnav/mp3d_70k_demo_maps'
    scene_ids = ['vyrNrziPKCB']

    dataset = NavMapDemoDataset(
        scene_ids, 
        rgb_ft_dir=os.path.join(input_dir, 'rgb_fts', 'resnet50'), 
        depth_ft_dir=os.path.join(input_dir, 'depth_fts'), 
        meta_dir=os.path.join(input_dir, 'meta_infos'), 
        rgb_image_size=224, map_dir=map_dir, map_image_size=240,
    )

    # st = time.time()
    # for i in range(len(dataset)):
    #     outs = dataset[i]
    #     if i in range(3):
    #         for k, v in outs.items():
    #             if isinstance(v, np.ndarray):
    #                 print(k, v.dtype, v.shape)
    #                 print('\t', np.min(v), np.mean(v), np.max(v))
    #             else:
    #                 print(k, type(v), v)
    # print('cost time: %.2fmin' % ((time.time() - st) / 60.))    # ~6min for all


    # multi-workers BUG:
    # if the batch_size is large, there is:
    # struct.error: 'i' format requires -2147483648 <= number <= 2147483647
    dataloader = DataLoader(
        dataset, batch_size=16, shuffle=True, 
        num_workers=4, collate_fn=map_collate_fn, 
        pin_memory=True, drop_last=False
    )
    st = time.time()
    for n, batch in enumerate(dataloader):
        for k, v in batch.items():
            if isinstance(v, torch.Tensor):
                print(k, v.dtype, v.size())
            else:
                print(k, type(v), len(v), type(v[0]))
        # if n < 3:
        #     np.save('batch%02d.npy'%n, batch)
        print()
    print('cost time: %.2fmin' % ((time.time() - st) / 60.))    # ~2min for all with 4 workers


offline_bc/data/map_dataset.py
cshizhe/onav_rim
import os
import jsonlines
import numpy as np
import cv2
import time

import torch

from torch.nn.utils.rnn import pad_sequence
from offline_bc.utils.ops import pad_tensors

from offline_bc.utils.logger import LOGGER
from offline_bc.data.dataset import NavDemoDataset


def crop_centered_local_map(
    agent_loc, global_map, local_sizes, 
    keep_centered_when_out_of_boundary=False,
    channel_before_size=False
):
    '''Crop a local map with local_sizes in the global_map centered at the agent's location.
    keep_centered_when_out_of_boundary: if True, we will pad zeros in the local map
    '''
    loc_r, loc_c = agent_loc
    local_w, local_h = local_sizes
    if channel_before_size:
        n_channel, full_w, full_h = global_map.shape
    else:
        full_w, full_h, n_channel = global_map.shape

    assert (local_w <= full_w) & (local_h <= full_h)

    gx1, gy1 = loc_r - local_w // 2, loc_c - local_h // 2
    gx2, gy2 = gx1 + local_w, gy1 + local_h
    
    if keep_centered_when_out_of_boundary:
        if isinstance(global_map, np.ndarray):
            if channel_before_size:
                local_map = np.zeros((n_channel, local_w, local_h), dtype=global_map.dtype)
            else:
                local_map = np.zeros((local_w, local_h, n_channel), dtype=global_map.dtype)
        elif isinstance(global_map, torch.Tensor):
            if channel_before_size:
                local_map = torch.zeros(
                    n_channel, local_w, local_h,
                    dtype=global_map.dtype, device=global_map.device
                )
            else:
                local_map = torch.zeros(
                    local_w, local_h, n_channel, 
                    dtype=global_map.dtype, device=global_map.device
                )
    
        lx1 = max(- gx1, 0)
        ly1 = max(- gy1, 0)
        if lx1 >= local_w or ly1 >= local_h:
            return local_map
        lx2 = lx1 + local_w - max(0, gx2 - full_w)
        ly2 = ly1 + local_h - max(0, gy2 - full_h)
        if lx2 <= 0 or ly2 <= 0:
            return local_map

        gx1, gy1 = max(0, gx1), max(0, gy1)
        gx2, gy2 = min(full_w, gx2), min(full_h, gy2)
       
        if channel_before_size:
            local_map[:, lx1: lx2, ly1: ly2] = global_map[:, gx1: gx2, gy1: gy2]
        else:
            local_map[lx1: lx2, ly1: ly2] = global_map[gx1: gx2, gy1: gy2]

    else:
        if gx1 < 0:
            gx1, gx2 = 0, local_w
        if gx2 > full_w:
            gx1, gx2 = full_w - local_w, full_w

        if gy1 < 0:
            gy1, gy2 = 0, local_h
        if gy2 > full_h:
            gy1, gy2 = full_h - local_h, full_h

        if channel_before_size:
            local_map = global_map[:, gx1: gx2, gy1: gy2]
        else:
            local_map = global_map[gx1: gx2, gy1: gy2]

    return local_map

class NavMapStepDemoDataset(NavDemoDataset):
    def __init__(
        self, scene_ids, rgb_image_dir=None, rgb_ft_dir=None, 
        depth_ft_dir=None, meta_dir=None, semseg_dir=None,
        rgb_image_size=224, max_steps=500, inflection_weight=None,
        num_ft_views=1, use_thda=True, map_dir=None, map_image_size=240,
        num_history_fts=None, num_history_ft_gap=1, num_actions=6, **kwargs
    ) -> None:
        super().__init__(
            scene_ids, rgb_image_dir=rgb_image_dir, rgb_ft_dir=rgb_ft_dir,
            depth_ft_dir=depth_ft_dir, meta_dir=meta_dir, semseg_dir=semseg_dir,
            rgb_image_size=rgb_image_size, max_steps=max_steps, 
            inflection_weight=inflection_weight, num_history_fts=num_history_fts,
            num_ft_views=num_ft_views, use_thda=use_thda, save_step_idxs=True, 
            **kwargs
        )
        self.map_dir = map_dir
        self.map_image_size = map_image_size
        self.num_history_ft_gap = num_history_ft_gap
        self.num_actions = num_actions

        self.visited_locs_in_full_map = {}
        for scene_id in self.scene_ids:
            self.visited_locs_in_full_map[scene_id] = {}
            with jsonlines.open(os.path.join(map_dir, 'meta_infos', f'{scene_id}.jsonl'), 'r') as f:
                for item in f:
                    self.visited_locs_in_full_map[scene_id][item['episode_id']] = item['visited_locs_in_full_map']

        self.map_envs, self.map_txns = self.build_lmdb_envs(
            os.path.join(map_dir, 'maps'), self.scene_ids
        ) 

    def load_dataset_meta_infos(self):
        super().load_dataset_meta_infos()
        del self.data_idxs

    def __len__(self):
        return self.num_steps

    def load_map_at_step(self, txn, episode_id, step):
        imgs = []
        for map_name in ['obs', 'exp']:
            key = f'{episode_id}_{step}_{map_name}'
            value = txn.get(key.encode('ascii'))
            if value is None:
                img = np.zeros((960, 960), dtype=np.uint8)
                LOGGER.info(f'DATASET ERROR: {key} not exists')
            else:
                img = np.frombuffer(value, dtype=np.uint8)
                img = cv2.imdecode(img, cv2.IMREAD_GRAYSCALE)
            imgs.append(img)
        imgs = np.stack(imgs, axis=2)
        return imgs

    def plot_traj_image(self, traj_img, visited_locs_in_full_map, start_step, end_step=None):
        nrows, ncols = traj_img.shape
        if end_step is None:
            end_step = start_step + 1
        for t in range(start_step, end_step):
            r, c = visited_locs_in_full_map[t]
            r = min(max(0, r), nrows - 1)
            c = min(max(0, c), ncols - 1)
            traj_img[max(0, r-1):r+2, max(0, c-1):c+2] = 1 # t + 1
        return traj_img

    def _sample_steps(self, num_steps):
        step_ids = np.arange(num_steps, dtype=np.int32)
        # reverse the order to ensure the last step is selected
        step_ids = step_ids[::-self.num_history_ft_gap]
        step_ids = step_ids[:self.num_history_fts]
        step_ids = step_ids[::-1]
        return step_ids

    def __getitem__(self, item_idx):
        scene_id, data_idx, cur_step = self.step_idxs[item_idx]

        meta = self.load_episode_meta_info(self.meta_info_txns[scene_id], data_idx)
        episode_id = meta['episode_id']

        num_steps = len(meta['compass'])
        start_step = 0
        if num_steps > self.max_steps:
            start_step = num_steps - self.max_steps

        outs = {
            # 'episode_ids': episode_id,
            # 'scene_ids': scene_id,
            'objectgoal': meta['objectgoal'],
            'object_category': meta['object_category'],
        }

        step_ids = self._sample_steps(cur_step + 1)
         # for episode > max_steps, just use the last max_steps steps
        step_ids = step_ids[step_ids >= start_step]
        outs['step_ids'] = step_ids.copy() - start_step
        outs['num_history_fts'] = len(step_ids)

        # relative to the start orientation
        compass = np.array(meta['compass'], dtype=np.float32)[step_ids] # (nsteps, 1)
        # relative to the start position using the start orientation as coordinates
        # habitat-sim axis [-agent_position[2] -z, agent_position[0] x]
        gps = np.array(meta['gps'], dtype=np.float32)[step_ids] # (nsteps, 2)
        # add the <bos> token
        prev_actions = np.array([self.num_actions] + meta['demonstration'], dtype=np.int32)[step_ids] # (nsteps, )
        outs['prev_actions'] = prev_actions

        if self.rgb_ft_dir is not None:
            rgb_fts = self.load_episode_fts(
                self.rgb_ft_txns[scene_id], episode_id
            )
            outs['rgb_features'] = rgb_fts[step_ids]
        if self.depth_ft_dir is not None:
            depth_fts = self.load_episode_fts(
                self.depth_ft_txns[scene_id], episode_id
            )
            outs['depth_features'] = depth_fts[step_ids]

        visited_locs_in_full_map = self.visited_locs_in_full_map[scene_id][episode_id]
        map_imgs = self.load_map_at_step(
            self.map_txns[scene_id], episode_id, cur_step, 
        )
        traj_img = np.zeros(map_imgs.shape[:-1], dtype=np.uint8)
        traj_img = self.plot_traj_image(traj_img, visited_locs_in_full_map, 0, cur_step+1)
        map_imgs = np.concatenate([map_imgs, traj_img[:, :, None]], axis=2)
        outs['maps'] = crop_centered_local_map(
            visited_locs_in_full_map[cur_step], map_imgs, 
            (self.map_image_size, self.map_image_size), 
            keep_centered_when_out_of_boundary=True
        )

        outs['compass_features'] = np.concatenate([np.sin(compass), np.cos(compass)], axis=1)
        agent_loc = np.array(visited_locs_in_full_map[cur_step]) # (x, y)
        agent_loc = 0.05 * (agent_loc - 480) # each grid is 5cm
        outs['gps_features'] = np.stack([gps[:, 0], -gps[:, 1]], axis=1) # relative to the start position (meters)
        # print('agent_loc', agent_loc, 'gps', outs['gps_features'][-1]) # should be similar
        outs['gps_features'] = outs['gps_features'] - agent_loc # relative to the agent location

        outs['demonstration'] = meta['demonstration'][cur_step]

        outs['inflection_weight'] = 1
        if self.inflection_weight is not None:
            if cur_step == 0 or meta['demonstration'][cur_step] != meta['demonstration'][cur_step - 1]:
                outs['inflection_weight'] = self.inflection_weight
        if self.kwargs.get('stop_weight', None) is not None and outs['demonstration'] == 0:
            outs['inflection_weight'] = self.kwargs['stop_weight']
            
        return outs


def map_step_collate_fn(inputs):
    batch = {
        k: [x[k] for x in inputs] for k in inputs[0].keys()
    }

    # max_steps = np.max(batch['num_history_fts'])
    # batch['num_history_fts'] = torch.LongTensor(batch['num_history_fts'])
    
    # Do not pad the features here to save batch memory
    # for key in ['rgb_features', 'depth_features', 'compass_features', 'gps_features']:
    #     if key in batch:
    #         # (batch, max_steps, dim_ft)
    #         batch[key] = pad_tensors(
    #             [torch.FloatTensor(x) for x in batch[key]]
    #         )
    # batch['step_ids'] = pad_sequence(
    #     [torch.LongTensor(x) for x in batch['step_ids']], 
    #     batch_first=True, padding_value=0
    # )
    # batch['prev_actions]

    batch['maps'] = torch.from_numpy(np.stack(batch['maps'], 0)) # (batch, h, w, 3) uint8

    batch['objectgoal'] = torch.LongTensor(batch['objectgoal'])
    batch['demonstration'] = torch.LongTensor(batch['demonstration'])
    batch['inflection_weight'] = torch.FloatTensor(batch['inflection_weight'])
    return batch


if __name__ == '__main__':
    from torch.utils.data import DataLoader

    input_dir = 'data/datasets/objectnav/mp3d_70k_demos'
    map_dir = 'data/datasets/objectnav/mp3d_70k_demo_maps_48m'
    scene_ids = ['17DRP5sb8fy']

    dataset = NavMapStepDemoDataset(
        scene_ids, 
        rgb_ft_dir=os.path.join(input_dir, 'rgb_fts', 'clip'), 
        depth_ft_dir=os.path.join(input_dir, 'depth_fts'), 
        meta_dir=os.path.join(input_dir, 'meta_infos'), 
        map_dir=map_dir, map_image_size=240,
        num_history_fts=100, num_history_ft_gap=1, 
        inflection_weight=3,
    )

    # st = time.time()
    # for i in range(len(dataset)):
    #     outs = dataset[i]
    #     if i in range(3):
    #         for k, v in outs.items():
    #             if isinstance(v, np.ndarray):
    #                 print(k, v.dtype, v.shape)
    #                 print('\t', np.min(v), np.mean(v), np.max(v))
    #             else:
    #                 print(k, type(v), v)
    # print('cost time: %.2fmin' % ((time.time() - st) / 60.))  


    # multi-workers BUG:
    # if the batch_size is large, there is:
    # struct.error: 'i' format requires -2147483648 <= number <= 2147483647
    dataloader = DataLoader(
        dataset, batch_size=256, shuffle=True, 
        num_workers=4, collate_fn=map_step_collate_fn, 
        pin_memory=True, drop_last=False
    )
    print('num batches: %d' % (len(dataloader)))
    st = time.time()
    for n, batch in enumerate(dataloader):
        for k, v in batch.items():
            if isinstance(v, torch.Tensor):
                print(k, v.dtype, v.size())
            else:
                print(k, type(v), len(v), type(v[0]))
        if n < 1:
            np.save('batch%02d.npy'%n, batch)
        break
        print()
        
    print('cost time: %.2fmin' % ((time.time() - st) / 60.))    # ~10min for all with 4 workers


offline_bc/data/map_step_dataset.py
cshizhe/onav_rim
import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F
import einops

from gym import spaces

from habitat.tasks.nav.object_nav_task import (
    task_cat2mpcat40,
    mapping_mpcat40_to_goal21
)
from habitat_baselines.il.common.encoders.resnet_encoders import (
    VlnResnetDepthEncoder,
    ResnetRGBEncoder,
    ResnetSemSeqEncoder,
    ResNetImageNetEncoder,
    ResNetCLIPEncoder,
)

from offline_bc.utils.logger import LOGGER
from offline_bc.utils.ops import pad_tensors_wgrad
from offline_bc.utils.ops import pad_tensors, gen_seq_masks


class ClsPrediction(nn.Module):
    def __init__(self, hidden_size, output_size, input_size=None, dropout_rate=0):
        super().__init__()
        if input_size is None:
            input_size = hidden_size
        self.net = nn.Sequential(nn.Linear(input_size, hidden_size),
                                 nn.ReLU(),
                                 nn.LayerNorm(hidden_size, eps=1e-12),
                                 nn.Dropout(dropout_rate),
                                 nn.Linear(hidden_size, output_size))

    def forward(self, x):
        return self.net(x)


class NavILBaseModel(nn.Module):
    r"""Reimplementation of habitat_baselines.il.env_based.policy.resnet_policy.ObjectNavILPolicy
    """

    def __init__(self, model_config, device):
        super().__init__()
        self.model_config = model_config
        self.device = device

        self.step_input_sizes = {}
        self.build_model(model_config)

    def build_model(self, model_config):
        # encoders
        if model_config.encoder_type.startswith('concat'):
            self.build_encoders_concat(model_config)
        elif model_config.encoder_type == 'add':
            self.build_encoders_add(model_config)
        else:
            raise NotImplementedError('unsupported encoder type: %s' % (model_config.encoder_type))

        # action predictor
        if model_config.action_clf_class == 'ClsPrediction':
            self.action_distribution = ClsPrediction(
                self.output_size, model_config.num_actions, 
                dropout_rate=model_config.dropout_rate
            )
        else:
            self.action_distribution = nn.Sequential(
                nn.Dropout(model_config.dropout_rate),
                nn.Linear(self.output_size, model_config.num_actions),
            )

    def build_visual_encoders(self, model_config):

        observation_space = spaces.Dict({
            'depth': spaces.Box(low=0., high=1., shape=(480, 640, 1), dtype=np.float32),
            'rgb': spaces.Box(low=0, high=255, shape=(224, 224, 3), dtype=np.uint8),
            'semantic': spaces.Box(low=0, high=4294967295, shape=(480, 640), dtype=np.uint32),
        })

        # Init the depth encoder
        assert model_config.DEPTH_ENCODER.cnn_type in [
            "VlnResnetDepthEncoder",
            "None",
        ], "DEPTH_ENCODER.cnn_type must be VlnResnetDepthEncoder"
        if model_config.DEPTH_ENCODER.cnn_type == "VlnResnetDepthEncoder":
            self.depth_encoder = VlnResnetDepthEncoder(
                observation_space,
                output_size=model_config.DEPTH_ENCODER.output_size,
                checkpoint=model_config.DEPTH_ENCODER.ddppo_checkpoint,
                backbone=model_config.DEPTH_ENCODER.backbone,
                trainable=model_config.DEPTH_ENCODER.trainable,
            )
            self.step_input_sizes['depth'] = model_config.DEPTH_ENCODER.output_size * model_config.num_ft_views
        else:
            self.depth_encoder = None

        # Init the RGB visual encoder
        assert model_config.RGB_ENCODER.cnn_type in [
            "ResnetRGBEncoder",
            "resnet50_imagenet",
            "resnet18_imagenet",
            "resnet50_clip",
            "None",
        ], "RGB_ENCODER.cnn_type must be 'ResnetRGBEncoder'."

        if model_config.RGB_ENCODER.cnn_type == "ResnetRGBEncoder":
            self.rgb_encoder = ResnetRGBEncoder(
                observation_space,
                output_size=model_config.RGB_ENCODER.output_size,
                backbone=model_config.RGB_ENCODER.backbone,
                trainable=model_config.RGB_ENCODER.train_encoder,
                normalize_visual_inputs=model_config.RGB_ENCODER.normalize_visual_inputs,
            )
            self.step_input_sizes['rgb'] = model_config.RGB_ENCODER.output_size * model_config.num_ft_views
        elif model_config.RGB_ENCODER.cnn_type in ["resnet50_imagenet", 'resnet18_imagenet']:
            self.rgb_encoder = ResNetImageNetEncoder(
                observation_space,
                output_size=model_config.RGB_ENCODER.output_size,
                backbone=model_config.RGB_ENCODER.cnn_type.split('_')[0],
            )
            self.step_input_sizes['rgb'] = model_config.RGB_ENCODER.output_size * model_config.num_ft_views
        elif model_config.RGB_ENCODER.cnn_type.startswith("resnet50_clip"):
            self.rgb_encoder = ResNetCLIPEncoder(
                observation_space,
                output_size=model_config.RGB_ENCODER.output_size,
                pooling='avgpool',
                device=self.device
            )
            self.step_input_sizes['rgb'] = model_config.RGB_ENCODER.output_size * model_config.num_ft_views
        else:
            self.rgb_encoder = None
            LOGGER.info("RGB encoder is none")

        sem_seg_output_size = 0
        self.semantic_predictor = None
        self.is_thda = False
        if model_config.USE_SEMANTICS:
            sem_embedding_size = model_config.SEMANTIC_ENCODER.embedding_size

            self.is_thda = model_config.SEMANTIC_ENCODER.is_thda
            sem_spaces = {
                "semantic": spaces.Box(
                    low=0,
                    high=255,
                    shape=(480, 640, sem_embedding_size),
                    dtype=np.uint8,
                ),
            }
            sem_obs_space = spaces.Dict(sem_spaces)
            self.sem_seg_encoder = ResnetSemSeqEncoder(
                sem_obs_space,
                output_size=model_config.SEMANTIC_ENCODER.output_size,
                backbone=model_config.SEMANTIC_ENCODER.backbone,
                trainable=model_config.SEMANTIC_ENCODER.train_encoder,
                semantic_embedding_size=sem_embedding_size,
                is_thda=self.is_thda
            )
            sem_seg_output_size = model_config.SEMANTIC_ENCODER.output_size
            LOGGER.info("Setting up Sem Seg model")
            self.step_input_sizes['sem'] = sem_seg_output_size

            self.embed_sge = model_config.embed_sge
            if self.embed_sge:
                self.task_cat2mpcat40 = torch.tensor(task_cat2mpcat40)
                self.mapping_mpcat40_to_goal = np.zeros(
                    max(
                        max(mapping_mpcat40_to_goal21.keys()) + 1,
                        50,
                    ),
                    dtype=np.int8,
                )

                for key, value in mapping_mpcat40_to_goal21.items():
                    self.mapping_mpcat40_to_goal[key] = value
                self.mapping_mpcat40_to_goal = torch.tensor(self.mapping_mpcat40_to_goal)
                self.step_input_sizes['sem_sge'] = 1
        
    def build_encoders_concat(self, model_config):
        '''Habitat-web baseline: concatenate the embeded features
        '''
        self.build_visual_encoders(model_config)

        if hasattr(model_config, 'SEM_FT_ENCODER') and model_config.SEM_FT_ENCODER.type is not None:
            self.sem_ft_layer = nn.Linear(
                model_config.SEM_FT_ENCODER.input_size,
                model_config.SEM_FT_ENCODER.output_size
            )
            self.step_input_sizes['sem'] = model_config.SEM_FT_ENCODER.output_size
            LOGGER.info("\n\nSetting up Semantic sensor")
        else:
            self.sem_ft_layer = None

        if model_config.USE_GPS:
            self.gps_embedding = nn.Linear(2, 32)
            self.step_input_sizes['gps'] = 32
            LOGGER.info("\n\nSetting up GPS sensor")
        
        if model_config.USE_COMPASS:
            input_compass_dim = 2  # cos and sin of the angle
            self.compass_embedding_dim = 32
            self.compass_embedding = nn.Linear(input_compass_dim, self.compass_embedding_dim)
            self.step_input_sizes['compass'] = 32
            LOGGER.info("\n\nSetting up Compass sensor")

        self._n_object_categories = 28
        LOGGER.info("Object categories: {}".format(self._n_object_categories))
        if self.model_config.obj_embed_file is None:
            self.obj_category_to_embeds = None
            self.obj_categories_embedding = nn.Embedding(
                self._n_object_categories, 32
            )
        else:
            obj_embeds = np.load(self.model_config.obj_embed_file, allow_pickle=True).item()
            self.obj_category_to_embeds = {
                k: torch.from_numpy(v.astype(np.float32)).to(self.device) \
                    for k, v in obj_embeds.items()
            }
            self.obj_categories_embedding = nn.Linear(
                list(obj_embeds.values())[0].shape[0], 32
            )
        self.step_input_sizes['objectgoal'] = 32
        LOGGER.info("\n\nSetting up Object Goal sensor")

        if model_config.SEQ2SEQ.use_prev_action:
            self.prev_action_embedding = nn.Embedding(model_config.num_actions + 1, 32)
            self.step_input_sizes['prev_action'] = self.prev_action_embedding.embedding_dim

        LOGGER.info('step input size: %d' % (np.sum(list(self.step_input_sizes.values()))))

        if model_config.encoder_type == 'concat_linear':
            step_input_size = np.sum(list(self.step_input_sizes.values()))
            self.ft_fusion_layer = nn.Linear(step_input_size, self.output_size)
        else:
            self.ft_fusion_layer = None

    def build_encoders_add(self, model_config):
        '''Typical approach in transformer: add all the embeded features
        '''
        hidden_size = model_config.hidden_size

        model_config.defrost()
        model_config.DEPTH_ENCODER.output_size = hidden_size
        model_config.RGB_ENCODER.output_size = hidden_size
        model_config.freeze()

        # RGB and depth encoder
        self.build_visual_encoders(model_config)
        self.rgb_layer_norm = nn.LayerNorm(hidden_size)
        self.depth_layer_norm = nn.LayerNorm(hidden_size)

        if hasattr(model_config, 'SEM_FT_ENCODER') and model_config.SEM_FT_ENCODER.type is not None:
            self.sem_ft_layer = nn.Sequential(
                nn.Linear(
                    model_config.SEM_FT_ENCODER.input_size,
                    model_config.SEM_FT_ENCODER.hidden_size
                ),
                nn.LayerNorm(hidden_size)
            )
            self.step_input_sizes['sem'] = hidden_size
        else:
            self.sem_ft_layer = None

        if model_config.USE_GPS:
            self.gps_embedding = nn.Sequential(
                nn.Linear(2, hidden_size),
                nn.LayerNorm(hidden_size)
            )
            self.step_input_sizes['gps'] = hidden_size
        
        if model_config.USE_COMPASS:
            # cos and sin of the heading angle
            self.compass_embedding = nn.Sequential(
                nn.Linear(2, hidden_size),
                nn.LayerNorm(hidden_size),
            )
            self.step_input_sizes['compass'] = hidden_size

        self._n_object_categories = 28
        LOGGER.info("Object categories: {}".format(self._n_object_categories))
        if self.model_config.obj_embed_file is None:
            self.obj_category_to_embeds = None
            self.obj_categories_embedding = nn.Embedding(
                self._n_object_categories, hidden_size
            )
        else:
            obj_embeds = np.load(self.model_config.obj_embed_file, allow_pickle=True).item()
            self.obj_category_to_embeds = {
                k: torch.from_numpy(v.astype(np.float32)).to(self.device) \
                    for k, v in obj_embeds.items()
            }
            self.obj_categories_embedding = nn.Linear(
                iter(obj_embeds.values()).shape[0], hidden_size
            )
        self.step_input_sizes['objectgoal'] = hidden_size

        if model_config.SEQ2SEQ.use_prev_action:
            self.prev_action_embedding = nn.Embedding(model_config.num_actions + 1, hidden_size)
            self.step_input_sizes['prev_action'] = hidden_size

        self.ft_fusion_layer = nn.Sequential(
            nn.LayerNorm(hidden_size),
            nn.Dropout(model_config.dropout_rate),
        )

    @property
    def output_size(self):
        return self.model_config.hidden_size

    @property
    def num_parameters(self):
        nweights, nparams = 0, 0
        for k, v in self.named_parameters():
            nweights += np.prod(v.size())
            nparams += 1
        return nweights, nparams

    @property
    def num_trainable_parameters(self):
        nweights, nparams = 0, 0
        for k, v in self.named_parameters():
            if v.requires_grad:
                nweights += np.prod(v.size())
                nparams += 1
        return nweights, nparams

    def _extract_sge(self, batch):
        # recalculating to keep this self-contained instead of depending on training infra
        if "semantic" in batch and "objectgoal" in batch:
            # (batch x nsteps, height x width x 1)
            obj_semantic = batch["semantic"].contiguous().flatten(start_dim=1)
            
            object_goals = []
            for i, og in enumerate(batch['objectgoal'].data.cpu().numpy()):
                object_goals.extend([og] * batch['num_steps'][i])
            object_goals = torch.from_numpy(np.array(object_goals))

            idx = self.task_cat2mpcat40[object_goals]
            if self.is_thda:
                idx = self.mapping_mpcat40_to_goal[idx].long()
            idx = idx.to(obj_semantic.device).unsqueeze(1)

            goal_visible_pixels = (obj_semantic == idx).sum(dim=1)
            goal_visible_area = torch.true_divide(goal_visible_pixels, obj_semantic.size(-1)).float()
            return goal_visible_area.unsqueeze(-1)

    def prepare_batch(self, batch):
        for key in ['rgb_features', 'depth_features', 'sem_features', \
                    'compass', 'gps', 
                    'infer_gps', 'infer_compass', 'infer_visual_features']:
            if key in batch:
                # (batch, max_steps, dim_ft)
                batch[key] = pad_tensors(
                    [torch.FloatTensor(x) for x in batch[key]]
                )
        
        if 'infer_local_maps' in batch:
            batch['infer_local_maps'] = pad_tensors(
                [x.float() for x in batch['infer_local_maps']]
            )

        for k, v in batch.items():
            if k in ['rgb']:
                batch[k] = torch.from_numpy(v).to(self.device).float()
            elif k in ['semantic']:
                batch[k] = torch.from_numpy(v).to(self.device)
            else:
                if isinstance(v, torch.Tensor):
                    batch[k] = v.to(self.device)
        return batch

    def encode_step_obs_concat(self, batch, step_embeddings=None):
        batch_size, num_steps, _ = batch['gps'].size()

        x = []
        if self.depth_encoder is not None:
            depth_embedding = self.depth_encoder(batch)
            if self.model_config.num_ft_views > 1: # (batch, nsteps, nviews, dim)
                depth_embedding = depth_embedding.view(batch_size, num_steps, -1)
            x.append(depth_embedding)

        if self.rgb_encoder is not None:
            rgb_embedding = self.rgb_encoder(batch)
            if len(rgb_embedding.size()) == 2:  # (batch x nsteps, dim)
                rgb_embedding = torch.split(rgb_embedding, batch['num_steps'], dim=0)
                rgb_embedding = pad_tensors_wgrad(rgb_embedding)
            if self.model_config.num_ft_views > 1: # (batch, nsteps, nviews, dim)
                rgb_embedding = rgb_embedding.view(batch_size, num_steps, -1)
            x.append(rgb_embedding)

        if self.sem_ft_layer is not None:
            sem_embedding = self.sem_ft_layer(batch['sem_features'])
            x.append(sem_embedding)

        if self.model_config.USE_SEMANTICS:
            # batch["semantic"]: (num_steps x num_envs, h, w, 1)
            if self.embed_sge:
                sge_embedding = self._extract_sge(batch)
                if len(sge_embedding.size()) == 2:
                    sge_embedding = torch.split(sge_embedding, batch['num_steps'], dim=0)
                    sge_embedding = pad_tensors_wgrad(sge_embedding)
                x.append(sge_embedding)

            batch['semantic'] = batch['semantic'].squeeze(dim=3)
            sem_seg_embedding = self.sem_seg_encoder(batch)
            if len(sem_seg_embedding.size()) == 2:  # (batch x nsteps, dim)
                sem_seg_embedding = torch.split(sem_seg_embedding, batch['num_steps'], dim=0)
                sem_seg_embedding = pad_tensors_wgrad(sem_seg_embedding)
            x.append(sem_seg_embedding)

        if self.model_config.USE_GPS:
            x.append(self.gps_embedding(batch['gps']))
        
        if self.model_config.USE_COMPASS:
            compass_observations = torch.concat(
                [torch.cos(batch['compass']), torch.sin(batch['compass'])],
                -1,
            )
            compass_embedding = self.compass_embedding(compass_observations)
            x.append(compass_embedding)

        if self.obj_category_to_embeds is None:
            obj_embedding = self.obj_categories_embedding(batch['objectgoal'])
        else:
            obj_embedding = torch.stack(
                [self.obj_category_to_embeds[k] for k in batch['object_category']], 0
            )
            obj_embedding = self.obj_categories_embedding(obj_embedding)

        x.append(
            einops.repeat(
                obj_embedding, 'b d -> b t d', t=num_steps
            )
        )

        if self.model_config.SEQ2SEQ.use_prev_action:
            prev_actions = torch.zeros(batch_size, num_steps, dtype=torch.long).to(self.device)
            prev_actions[:, 0] = -1
            prev_actions[:, 1:] = batch['demonstration'][:, :-1]
            prev_actions[prev_actions == -100] = -1
            prev_actions_embedding = self.prev_action_embedding(
                prev_actions + 1
            )
            x.append(prev_actions_embedding)
        
        x = torch.cat(x, dim=2)
        if self.ft_fusion_layer is not None:
            x = self.ft_fusion_layer(x)

        if step_embeddings is not None:
            x = x + step_embeddings
        
        return x

    def encode_step_obs_add(self, batch, step_embeddings=None):
        x = {}

        if self.depth_encoder is not None:
            depth_embedding = self.depth_encoder(batch)
            depth_embedding = self.depth_layer_norm(depth_embedding)
            x['depth'] = depth_embedding

        if self.rgb_encoder is not None:
            rgb_embedding = self.rgb_encoder(batch)
            rgb_embedding = self.rgb_layer_norm(rgb_embedding)
            x['rgb'] = rgb_embedding

        if self.sem_ft_layer is not None:
            sem_embedding = self.sem_ft_layer(batch['sem_features'])
            x['sem'] = sem_embedding

        if self.model_config.USE_GPS:
            x['gps'] = self.gps_embedding(batch['gps'])
        
        if self.model_config.USE_COMPASS:
            compass_observations = torch.concat(
                [torch.cos(batch['compass']), torch.sin(batch['compass'])],
                -1,
            )
            compass_embedding = self.compass_embedding(compass_observations)
            x['compass'] = compass_embedding

        if self.model_config.SEQ2SEQ.use_prev_action:
            prev_actions = torch.zeros_like(batch['demonstration'])
            prev_actions[:, 0] = -1
            prev_actions[:, 1:] = batch['demonstration'][:, :-1]
            prev_actions[prev_actions == -100] = -1
            prev_actions_embedding = self.prev_action_embedding(
                prev_actions + 1
            )
            x['prev_action'] = prev_actions_embedding

        # (batch_size, nsteps, hidden_size)
        fused_embeds = torch.sum(torch.stack(list(x.values()), 0), 0)

        if self.obj_category_to_embeds is None:
            obj_embedding = self.obj_categories_embedding(batch['objectgoal'])
        else:
            obj_embedding = torch.stack(
                [self.obj_category_to_embeds[k] for k in batch['object_category']], 0
            )
            obj_embedding = self.obj_categories_embedding(obj_embedding)
        x['objectgoal'] = obj_embedding

        add_objgoal = self.model_config.get('encoder_add_objgoal', True)
        if add_objgoal:
            fused_embeds = fused_embeds + x['objectgoal'].unsqueeze(1)
        
        if step_embeddings is not None:
            fused_embeds = fused_embeds + step_embeddings
        fused_embeds = self.ft_fusion_layer(fused_embeds)
        x['fused_embeds'] = fused_embeds

        return x

    def encode_step_obs(self, batch, **kwargs):
        if self.model_config.encoder_type.startswith('concat'):
            return self.encode_step_obs_concat(batch, **kwargs)
        elif self.model_config.encoder_type == 'add':
            return self.encode_step_obs_add(batch, **kwargs)

    def compute_loss(self, logits, batch):
        losses = F.cross_entropy(
            logits.permute(0, 2, 1),
            batch['demonstration'],
            reduction='none', ignore_index=-100
        )   # (N, T)
        loss = torch.mean(
            torch.sum(losses * batch['inflection_weight'], 1) / \
                torch.sum(batch['inflection_weight'], 1)
        )
        return loss

    def forward(self, batch, compute_loss=False):
        raise NotImplementedError


offline_bc/models/onav_base.py
cshizhe/onav_rim
import torch
import torch.nn as nn
import torch.nn.functional as F
import einops

from offline_bc.models.transformer import (
    PositionalEncoding,
    TransformerEncoderLayer,
    TransformerEncoder,
)
from offline_bc.models.onav_base import NavILBaseModel
from offline_bc.utils.logger import LOGGER


class ImapEmbedding(nn.Module):
    def __init__(self, model_config) -> None:
        super().__init__()

        map_cfg = model_config.MAP_ENCODER
        if map_cfg.token_embed_type == 'single':
            self.imap_token_embedding = nn.Embedding(1, model_config.hidden_size)
        elif map_cfg.token_embed_type == 'multi':
            self.imap_token_embedding = nn.Embedding(map_cfg.imap_size**2, model_config.hidden_size)
        else:
            raise NotImplementedError(f'unsupported imap token embed {map_cfg.token_embed_type}')
        
        if map_cfg.encode_position:
            self.imap_pos_fts = self._create_imap_pos_features(map_cfg.imap_size)
            self.imap_pos_layer = nn.Sequential(
                nn.Linear(2, model_config.hidden_size), # (location: x, y)
                nn.LayerNorm(model_config.hidden_size)
            )
            
        self.ft_fusion_layer = nn.Sequential(
            nn.LayerNorm(model_config.hidden_size),
            nn.Dropout(model_config.dropout_rate),
        )

        self.imap_token_type = map_cfg.token_embed_type
        self.imap_num_tokens = map_cfg.imap_size**2
        self.imap_size = map_cfg.imap_size
        self.encode_position = map_cfg.encode_position

    def _create_imap_pos_features(self, imap_size):
        x, y = torch.meshgrid(torch.arange(imap_size), torch.arange(imap_size))
        xy = torch.stack([x, y], dim=2)
        xy = (xy + 0.5 - imap_size / 2).float() # relative distance to the center
        xy = xy.view(-1, 2)
        return xy

    def forward(self, batch_size):
        '''Get the initialized imap embedding
        '''
        device = self.imap_token_embedding.weight.device

        if self.imap_token_type == 'multi':
            token_types = torch.arange(self.imap_num_tokens, dtype=torch.long, device=device)
        else:
            token_types = torch.zeros(self.imap_num_tokens, dtype=torch.long, device=device)

        embeds = self.imap_token_embedding(token_types)
        if self.encode_position:
            pos_embeds = self.imap_pos_layer(self.imap_pos_fts.to(device))
            embeds = embeds + pos_embeds
        else:
            pos_embeds = None
        embeds = self.ft_fusion_layer(embeds)

        embeds = einops.repeat(embeds, 'n d -> b n d', b=batch_size)
        if pos_embeds is not None:
            pos_embeds = einops.repeat(pos_embeds, 'n d -> b n d', b=batch_size)

        return embeds, pos_embeds

class GpsCompassEmbedding(nn.Module):
    def __init__(self, hidden_size) -> None:
        super().__init__()
        self.gps_layer = nn.Sequential(
            nn.Linear(2, hidden_size), # (location: x, y)
            nn.LayerNorm(hidden_size)
        )
        self.compass_layer = nn.Sequential(
            nn.Linear(2, hidden_size), # (heading: sin, cos)
            nn.LayerNorm(hidden_size)
        )

    def forward(self, gps, compass):
        gps_embeds = self.gps_layer(gps)
        compass_fts = torch.cat(
            [torch.cos(compass), torch.sin(compass)], -1,
        )
        compass_embeds = self.compass_layer(compass_fts)
        embeds = gps_embeds + compass_embeds
        return embeds


class NavImapSingleTransformer(NavILBaseModel):
    def build_model(self, model_config):
        super().build_model(model_config)
        
        if model_config.STATE_ENCODER.learnable_step_embedding:
            self.step_embedding = nn.Embedding(
                model_config.STATE_ENCODER.max_steps,
                model_config.hidden_size,
            )
        else:
            self.step_embedding = PositionalEncoding(
                model_config.hidden_size,
                max_len=model_config.STATE_ENCODER.max_steps,
            )

        encoder_layer = TransformerEncoderLayer(
            model_config.hidden_size, 
            model_config.STATE_ENCODER.num_attention_heads, 
            dim_feedforward=model_config.STATE_ENCODER.intermediate_size, 
            dropout=model_config.STATE_ENCODER.dropout_prob,
            activation=model_config.STATE_ENCODER.hidden_act, 
            normalize_before=False
        )
        self.state_encoder = TransformerEncoder(
            encoder_layer, 
            model_config.STATE_ENCODER.num_hidden_layers, 
            norm=None, batch_first=True
        )
        if model_config.STATE_ENCODER.add_pos_attn:
            self.obs_pos_layer = GpsCompassEmbedding(model_config.hidden_size)

        self.imap_embedding = ImapEmbedding(model_config)

        if model_config.infer_visual_feature_loss > 0:
            dim_vis = 2048  # TODO: CLIP resnet50
            if model_config.infer_depth_feature:
                dim_vis += 2048 # depth resnet50
            self.vis_pred_layer = nn.Sequential(
                nn.Dropout(model_config.dropout_rate),
                nn.Linear(model_config.hidden_size, dim_vis),
            )
        
        if model_config.infer_local_map_loss > 0:
            output_size = model_config.pred_map_nchannels * (model_config.pred_map_image_size**2)
            input_size = model_config.pred_map_input_token_dim * (model_config.MAP_ENCODER.imap_size ** 2)
            self.map_pred_layer = nn.Sequential(
                nn.Linear(model_config.hidden_size, model_config.pred_map_input_token_dim),
                nn.ReLU(),
                nn.Flatten(start_dim=1),
                nn.Dropout(model_config.dropout_rate),
                nn.Linear(input_size, output_size),
                nn.Sigmoid(),
            )

        if 'infer_sem_label_loss' in self.model_config and self.model_config.infer_sem_label_loss > 0:
            # TODO: only valid for concat encoder
            input_size = model_config.RGB_ENCODER.output_size + model_config.DEPTH_ENCODER.output_size
            self.sem_pred_layer = nn.Sequential(
                nn.Dropout(model_config.dropout_rate),
                nn.Linear(input_size, 21*2+1)
            )     

    def forward(self, batch, compute_loss=False, return_imap_embeds=False):
        batch = self.prepare_batch(batch)
        batch_size, num_steps, _ = batch['gps'].size()

        # (batch, num_steps, dim)
        stepid_embeds = self.step_embedding(batch['step_ids'])
        inputs = self.encode_step_obs(batch, step_embeddings=stepid_embeds)
        if isinstance(inputs, dict):
            input_embeds = inputs['fused_embeds']
        else:
            input_embeds = inputs

        imap_embeds, imap_pos_embeds = self.imap_embedding(batch_size)
        if 'init_map_embeds' in batch:
            imap_embeds = batch['init_map_embeds']
            
        if self.model_config.STATE_ENCODER.add_pos_attn:
            obs_pos_embeds = self.obs_pos_layer(batch['gps'], batch['compass'])     

        if compute_loss and self.model_config.infer_visual_feature_loss > 0:
            if self.model_config.encoder_type == 'concat':
                masked_input_embeds = torch.zeros_like(input_embeds)
                dim_vis = self.step_input_sizes['depth'] + self.step_input_sizes['rgb']
                masked_input_embeds[:, :, dim_vis: dim_vis+32] = self.gps_embedding(batch['infer_gps'])
                compass_fts = torch.concat(
                    [torch.cos(batch['infer_compass']), torch.sin(batch['infer_compass'])],
                    -1,
                )
                masked_input_embeds[:, :, dim_vis+32: dim_vis+64] = self.compass_embedding(compass_fts)
            elif self.model_config.encoder_type == 'add':
                masked_input_embeds = self.gps_embedding(batch['infer_gps']) + \
                                      self.compass_embedding(torch.concat([torch.cos(batch['infer_compass']),
                                          torch.sin(batch['infer_compass'])], -1))
                masked_input_embeds = self.ft_fusion_layer(masked_input_embeds)

            if self.model_config.STATE_ENCODER.add_pos_attn:
                masked_obs_pos_embeds = self.obs_pos_layer(
                    batch['infer_gps'], batch['infer_compass']
                )

            target_visual_features = batch['infer_visual_features']
            seq_visual_features = batch['rgb_features']
            if self.model_config.infer_depth_feature:
                seq_visual_features = torch.cat([seq_visual_features, batch['depth_features']], dim=-1)
            seq_visual_features = F.normalize(seq_visual_features, p=2, dim=-1)

            ntokens = imap_embeds.size(1) + 2
            attn_masks = torch.zeros(ntokens, ntokens).bool().to(imap_embeds.device)
            attn_masks[:-1, -1] = True  # imap and obs tokens should not see the masked token
            attn_masks[-1, -2] = True   # the masked token should not see the current observation

        logits = []
        vis_pred_losses, map_pred_losses, sem_pred_losses = [], [], []
        if return_imap_embeds:
            last_imap_embeds = [None for _ in range(batch_size)]
        for t in range(num_steps):
            t_input_embeds = torch.cat(
                [imap_embeds, input_embeds[:, t:t+1]], dim=1
            )
            if self.model_config.STATE_ENCODER.add_pos_attn:
                t_pos_embeds = torch.cat(
                    [imap_pos_embeds, obs_pos_embeds[:, t:t+1]], dim=1
                )
                if compute_loss and self.model_config.infer_visual_feature_loss > 0:
                    t_pos_embeds = torch.cat(
                        [t_pos_embeds, masked_obs_pos_embeds[:, t:t+1]], dim=1
                    )
            else:
                t_pos_embeds = None

            if compute_loss and self.model_config.infer_visual_feature_loss > 0:
                hiddens = self.state_encoder(
                    torch.cat([t_input_embeds, masked_input_embeds[:, t:t+1]], dim=1),
                    pos=t_pos_embeds, mask=attn_masks,
                )
                imap_embeds = hiddens[:, :-2]
                obs_hiddens = hiddens[:, -2]
                masked_obs_hiddens = hiddens[:, -1]
                # (batch_size, dim)
                masked_visual_preds = self.vis_pred_layer(masked_obs_hiddens)
                if self.model_config.infer_visual_feature_loss_type == 'mse':
                    vis_pred_losses.append(F.mse_loss(
                        masked_visual_preds,
                        target_visual_features[:, t], reduction='none'
                    ).mean(dim=1))
                elif self.model_config.infer_visual_feature_loss_type == 'nce':
                    masked_visual_preds = F.normalize(masked_visual_preds, p=2, dim=-1)
                    pos_sim_scores = torch.einsum(
                        'bd,bd->b', F.normalize(target_visual_features[:, t], p=2, dim=-1), 
                        masked_visual_preds
                    )
                    neg_sim_scores = torch.einsum(
                        'btd,bd->bt', seq_visual_features, masked_visual_preds
                    )
                    neg_sim_scores.masked_fill_(batch['demonstration'] == -100, -float('inf'))
                    sim_scores = torch.cat([pos_sim_scores.unsqueeze(1), neg_sim_scores], 1)
                    sim_scores = sim_scores / 0.1
                    # print(t, sim_scores)
                    vis_pred_losses.append(F.cross_entropy(
                        sim_scores, 
                        torch.zeros(sim_scores.size(0), dtype=torch.long, device=sim_scores.device),
                        reduction='none'
                    ))
                    # print(vis_pred_losses[-1])
            else:
                hiddens, layer_attn_weights = self.state_encoder(
                    t_input_embeds, pos=t_pos_embeds, return_attn_weights=True
                )
                imap_embeds = hiddens[:, :-1]
                obs_hiddens = hiddens[:, -1]

            if self.model_config.infer_local_map_loss > 0:
                pred_maps = self.map_pred_layer(imap_embeds).view(
                    batch_size, self.model_config.pred_map_nchannels, -1
                )
                if self.model_config.infer_local_map_loss_type == 'mse':
                    t_map_pred_loss = F.mse_loss(
                        pred_maps, batch['infer_local_maps'][:, t], reduction='none'
                    ).view(batch_size, -1).mean(1)
                elif self.model_config.infer_local_map_loss_type == 'clf':
                    t_map_pred_loss = F.binary_cross_entropy(
                        pred_maps, batch['infer_local_maps'][:, t], reduction='none'
                    ).view(batch_size, -1).mean(1)
                map_pred_losses.append(t_map_pred_loss)

            if 'infer_sem_label_loss' in self.model_config and self.model_config.infer_sem_label_loss > 0:
                # TODO: only valid for concat encoder
                sem_input_hiddens = inputs[:, t, :self.model_config.DEPTH_ENCODER.output_size+self.model_config.RGB_ENCODER.output_size]
                pred_semantics = torch.sigmoid(self.sem_pred_layer(sem_input_hiddens))
                sem_labels = batch['sem_features'][:, t, :-1]
                sem_clf_loss = F.binary_cross_entropy(
                    pred_semantics[:, :21], (sem_labels > 0).float(), reduction='none'
                ).mean(1)
                sem_mse_loss = F.mse_loss(
                    pred_semantics[:, 21:], batch['sem_features'][:, t], reduction='none'
                ).mean(1)
                sem_pred_losses.append(sem_clf_loss + sem_mse_loss)

            add_objgoal = self.model_config.get('encoder_add_objgoal', True)
            if not add_objgoal:
                obs_hiddens = obs_hiddens + inputs['objectgoal']
            step_logits = self.action_distribution(obs_hiddens)
            logits.append(step_logits)

            if return_imap_embeds:
                for ib in range(batch_size):
                    if last_imap_embeds[ib] is None and batch['demonstration'][ib, t] == 0:
                        last_imap_embeds[ib] = imap_embeds[ib].data.cpu().numpy()

        logits = torch.stack(logits, 1)
        
        if compute_loss:
            act_loss = self.compute_loss(logits, batch)
            loss_dict = {'overall': act_loss, 'action_loss': act_loss}
            
            if self.model_config.infer_visual_feature_loss > 0:
                vis_pred_losses = torch.stack(vis_pred_losses, 1) # (batch, nsteps)
                vis_pred_masks = (batch['inflection_weight'] > 0).float()
                vis_pred_loss = torch.mean(
                    torch.sum(vis_pred_losses * vis_pred_masks, 1) / \
                        torch.sum(vis_pred_masks, 1)
                )
                # print(torch.sum(vis_pred_losses * vis_pred_masks, 0))
                loss_dict['overall'] = loss_dict['overall'] + vis_pred_loss * self.model_config.infer_visual_feature_loss
                loss_dict['vis_pred_loss'] = vis_pred_loss

            if self.model_config.infer_local_map_loss > 0:
                map_pred_masks = (batch['inflection_weight'] > 0).float()
                map_pred_losses = torch.stack(map_pred_losses, 1) # (batch, nsteps)
                map_pred_loss = torch.mean(
                    torch.sum(map_pred_losses * map_pred_masks, 1) / \
                        torch.sum(map_pred_masks, 1)
                )
                loss_dict['overall'] = loss_dict['overall'] + map_pred_loss * self.model_config.infer_local_map_loss
                loss_dict['map_pred_loss'] = map_pred_loss

            if 'infer_sem_label_loss' in self.model_config and self.model_config.infer_sem_label_loss > 0:
                sem_pred_losses = torch.stack(sem_pred_losses, 1)
                sem_pred_masks = (batch['inflection_weight'] > 0).float()
                sem_pred_loss = torch.mean(
                    torch.sum(sem_pred_losses * sem_pred_masks, 1) / torch.sum(sem_pred_masks, 1)
                )
                loss_dict['overall'] = loss_dict['overall'] + sem_pred_loss * self.model_config.infer_sem_label_loss
                loss_dict['sem_pred_loss'] = sem_pred_loss

            # for k, v in loss_dict.items():
            #     print(k, v.item())
            return loss_dict, logits

        if return_imap_embeds:
            for ib in range(batch_size):
                if last_imap_embeds[ib] is None:
                    last_imap_embeds[ib] = imap_embeds[ib].data.cpu().numpy()

            return last_imap_embeds

        return logits

offline_bc/models/onav_imap_models.py
cshizhe/onav_rim
import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F

from habitat_baselines.rl.models.rnn_state_encoder import RNNStateEncoder

from offline_bc.models.transformer import (
    triangular_mask, PositionalEncoding,
    TransformerEncoderLayer,
    TransformerEncoder,
)
from offline_bc.models.onav_base import NavILBaseModel
from offline_bc.utils.logger import LOGGER



class NavILRecurrentNet(NavILBaseModel):
    r"""Reimplementation of habitat_baselines.il.env_based.policy.resnet_policy.ObjectNavILPolicy
    """

    def build_model(self, model_config):
        super().build_model(model_config)
        self.state_encoder = RNNStateEncoder(
            input_size=model_config.hidden_size,
            hidden_size=model_config.hidden_size,
            num_layers=model_config.STATE_ENCODER.num_recurrent_layers,
            rnn_type=model_config.STATE_ENCODER.rnn_type,
        )

    @property
    def num_recurrent_layers(self):
        return self.state_encoder.num_recurrent_layers

    def forward(self, batch, compute_loss=False):
        batch = self.prepare_batch(batch)
        batch_size, num_steps, _ = batch['gps'].size()

        inputs = self.encode_step_obs(batch)
        if isinstance(inputs, dict):
            input_embeds = inputs['fused_embeds']
        else:
            input_embeds = inputs

        init_hidden_states = torch.zeros(
            self.num_recurrent_layers,
            batch_size,
            self.model_config.hidden_size
        ).to(self.device)

        x, _ = self.state_encoder(input_embeds, init_hidden_states, seq_single_episode=True)

        add_objgoal = self.model_config.get('encoder_add_objgoal', True)
        if not add_objgoal:
            x = x + inputs['objectgoal'].unsqueeze(1)
        logits = self.action_distribution(x)    # (N, T, D)

        if compute_loss:
            loss = self.compute_loss(logits, batch)
            return loss, logits

        return logits


class NavILTransformer(NavILBaseModel):
    r"""Reimplementation of habitat_baselines.il.env_based.policy.resnet_policy.ObjectNavILPolicy
    """
    def build_model(self, model_config):
        super().build_model(model_config)

        if model_config.STATE_ENCODER.learnable_step_embedding:
            self.step_embedding = nn.Embedding(
                model_config.STATE_ENCODER.max_steps,
                model_config.hidden_size,
            )
        else:
            self.step_embedding = PositionalEncoding(
                model_config.hidden_size,
                max_len=model_config.STATE_ENCODER.max_steps,
            )

        encoder_layer = TransformerEncoderLayer(
            model_config.hidden_size, 
            model_config.STATE_ENCODER.num_attention_heads, 
            dim_feedforward=model_config.STATE_ENCODER.intermediate_size, 
            dropout=model_config.STATE_ENCODER.dropout_prob,
            activation=model_config.STATE_ENCODER.hidden_act, 
            normalize_before=False
        )
        self.state_encoder = TransformerEncoder(
            encoder_layer, 
            model_config.STATE_ENCODER.num_hidden_layers, 
            norm=None, batch_first=True
        )

    def forward(self, batch, compute_loss=False):
        batch = self.prepare_batch(batch)
        batch_size, num_steps, _ = batch['gps'].size()

        # (batch, num_steps, dim)
        stepid_embeds = self.step_embedding(batch['step_ids'])
        inputs = self.encode_step_obs(batch, step_embeddings=stepid_embeds)
        if isinstance(inputs, dict):
            input_embeds = inputs['fused_embeds']
        else:
            input_embeds = inputs

        # (len_target, len_source): True denotes disallowing attention
        causal_attn_masks = triangular_mask(num_steps, self.device, diagonal_shift=1)
        hiddens = self.state_encoder(
            input_embeds, mask=causal_attn_masks, 
        )

        add_objgoal = self.model_config.get('encoder_add_objgoal', True)
        if not add_objgoal:
            hiddens = hiddens + inputs['objectgoal'].unsqueeze(1)
        logits = self.action_distribution(hiddens)    # (N, T, D)

        if compute_loss:
            loss = self.compute_loss(logits, batch)
            return loss, logits

        return logits

offline_bc/models/onav_vis_models.py
cshizhe/onav_rim
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
"""
DETR Transformer class.

Copy-paste from torch.nn.Transformer with modifications:
    * positional encodings are passed in MHattention
    * extra LN at the end of encoder is removed
    * decoder returns a stack of activations from all decoding layers
"""
import copy
import math
from typing import Optional, List

import torch
import torch.nn.functional as F
from torch import nn, Tensor


class Transformer(nn.Module):

    def __init__(self, d_model=512, nhead=8, num_encoder_layers=6,
                 num_decoder_layers=6, dim_feedforward=2048, dropout=0.1,
                 activation="relu", normalize_before=False,
                 return_intermediate_dec=False):
        super().__init__()

        encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward,
                                                dropout, activation, normalize_before)
        encoder_norm = nn.LayerNorm(d_model) if normalize_before else None
        self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)

        decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward,
                                                dropout, activation, normalize_before)
        decoder_norm = nn.LayerNorm(d_model)
        self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm,
                                          return_intermediate=return_intermediate_dec)

        self._reset_parameters()

        self.d_model = d_model
        self.nhead = nhead

    def _reset_parameters(self):
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

    def forward(self, src, mask, query_embed, pos_embed):
        # flatten NxCxHxW to HWxNxC
        bs, c, h, w = src.shape
        src = src.flatten(2).permute(2, 0, 1)
        pos_embed = pos_embed.flatten(2).permute(2, 0, 1)
        query_embed = query_embed.unsqueeze(1).repeat(1, bs, 1)
        mask = mask.flatten(1)

        tgt = torch.zeros_like(query_embed)
        memory = self.encoder(src, src_key_padding_mask=mask, pos=pos_embed)
        hs = self.decoder(tgt, memory, memory_key_padding_mask=mask,
                          pos=pos_embed, query_pos=query_embed)
        return hs.transpose(1, 2), memory.permute(1, 2, 0).view(bs, c, h, w)


class TransformerEncoder(nn.Module):

    def __init__(self, encoder_layer, num_layers, norm=None, batch_first=True):
        super().__init__()
        self.layers = _get_clones(encoder_layer, num_layers)
        self.num_layers = num_layers
        self.norm = norm
        self.batch_first = batch_first

    def forward(self, src,
                mask: Optional[Tensor] = None,
                src_key_padding_mask: Optional[Tensor] = None,
                pos: Optional[Tensor] = None,
                return_attn_weights=False):
        output = src
        if self.batch_first:
            output = output.transpose(0, 1)
            if (pos is not None) and (len(pos.size()) == 3):
                pos = pos.transpose(0, 1)

        layer_attn_weights = []
        for layer in self.layers:
            output = layer(output, src_mask=mask,
                           src_key_padding_mask=src_key_padding_mask, pos=pos)
            if return_attn_weights:
                layer_attn_weights.append(output[1])
            output = output[0]

        if self.norm is not None:
            output = self.norm(output)

        if self.batch_first:
            output = output.transpose(0, 1)

        if return_attn_weights:
            return output, layer_attn_weights
        return output


class TransformerDecoder(nn.Module):

    def __init__(self, decoder_layer, num_layers, norm=None, return_intermediate=False, batch_first=True):
        super().__init__()
        self.layers = _get_clones(decoder_layer, num_layers)
        self.num_layers = num_layers
        self.norm = norm
        self.return_intermediate = return_intermediate
        self.batch_first = batch_first

    def forward(self, tgt, memory,
                tgt_mask: Optional[Tensor] = None,
                memory_mask: Optional[Tensor] = None,
                tgt_key_padding_mask: Optional[Tensor] = None,
                memory_key_padding_mask: Optional[Tensor] = None,
                pos: Optional[Tensor] = None,
                query_pos: Optional[Tensor] = None):
        output = tgt
        if self.batch_first:
            output = output.transpose(0, 1) # (length, batch, hidden_dim)
            memory = memory.transpose(0, 1)
            # TODO: pos embedding (b, d), (b, m, d)

        intermediate = []

        for layer in self.layers:
            output = layer(output, memory, tgt_mask=tgt_mask,
                           memory_mask=memory_mask,
                           tgt_key_padding_mask=tgt_key_padding_mask,
                           memory_key_padding_mask=memory_key_padding_mask,
                           pos=pos, query_pos=query_pos)
            if self.return_intermediate:
                intermediate.append(self.norm(output))

        if self.norm is not None:
            output = self.norm(output)
            if self.return_intermediate:
                intermediate.pop()
                intermediate.append(output)

        if self.batch_first:
            output = output.transpose(0, 1)

        if self.return_intermediate:
            return torch.stack(intermediate)

        return output


class TransformerEncoderLayer(nn.Module):

    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
                 activation="relu", normalize_before=False):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        # Implementation of Feedforward model
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

        self.activation = _get_activation_fn(activation)
        self.normalize_before = normalize_before

    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
        return tensor if pos is None else tensor + pos

    def forward_post(self,
                     src,
                     src_mask: Optional[Tensor] = None,
                     src_key_padding_mask: Optional[Tensor] = None,
                     pos: Optional[Tensor] = None):
        q = k = self.with_pos_embed(src, pos)
        src2, attn_weights = self.self_attn(
            q, k, value=src, attn_mask=src_mask,
            key_padding_mask=src_key_padding_mask,
            need_weights=True, # average_attn_weights=True
        )
        src = src + self.dropout1(src2)
        src = self.norm1(src)
        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
        src = src + self.dropout2(src2)
        src = self.norm2(src)
        return src, attn_weights

    def forward_pre(self, src,
                    src_mask: Optional[Tensor] = None,
                    src_key_padding_mask: Optional[Tensor] = None,
                    pos: Optional[Tensor] = None):
        src2 = self.norm1(src)
        q = k = self.with_pos_embed(src2, pos)
        src2, attn_weights = self.self_attn(
            q, k, value=src2, attn_mask=src_mask,
            key_padding_mask=src_key_padding_mask,
            need_weights=True, average_attn_weights=True
        )
        src = src + self.dropout1(src2)
        src2 = self.norm2(src)
        src2 = self.linear2(self.dropout(self.activation(self.linear1(src2))))
        src = src + self.dropout2(src2)
        return src, attn_weights

    def forward(self, src,
                src_mask: Optional[Tensor] = None,
                src_key_padding_mask: Optional[Tensor] = None,
                pos: Optional[Tensor] = None):
        if self.normalize_before:
            return self.forward_pre(src, src_mask, src_key_padding_mask, pos)
        return self.forward_post(src, src_mask, src_key_padding_mask, pos)


class TransformerDecoderLayer(nn.Module):

    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
                 activation="relu", normalize_before=False):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        # Implementation of Feedforward model
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)

        self.activation = _get_activation_fn(activation)
        self.normalize_before = normalize_before

    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
        return tensor if pos is None else tensor + pos

    def forward_post(self, tgt, memory,
                     tgt_mask: Optional[Tensor] = None,
                     memory_mask: Optional[Tensor] = None,
                     tgt_key_padding_mask: Optional[Tensor] = None,
                     memory_key_padding_mask: Optional[Tensor] = None,
                     pos: Optional[Tensor] = None,
                     query_pos: Optional[Tensor] = None):
        q = k = self.with_pos_embed(tgt, query_pos)
        tgt2 = self.self_attn(q, k, value=tgt, attn_mask=tgt_mask,
                              key_padding_mask=tgt_key_padding_mask)[0]
        tgt = tgt + self.dropout1(tgt2)
        tgt = self.norm1(tgt)
        tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt, query_pos),
                                   key=self.with_pos_embed(memory, pos),
                                   value=memory, attn_mask=memory_mask,
                                   key_padding_mask=memory_key_padding_mask)[0]
        tgt = tgt + self.dropout2(tgt2)
        tgt = self.norm2(tgt)
        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
        tgt = tgt + self.dropout3(tgt2)
        tgt = self.norm3(tgt)
        return tgt

    def forward_pre(self, tgt, memory,
                    tgt_mask: Optional[Tensor] = None,
                    memory_mask: Optional[Tensor] = None,
                    tgt_key_padding_mask: Optional[Tensor] = None,
                    memory_key_padding_mask: Optional[Tensor] = None,
                    pos: Optional[Tensor] = None,
                    query_pos: Optional[Tensor] = None):
        tgt2 = self.norm1(tgt)
        q = k = self.with_pos_embed(tgt2, query_pos)
        tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask,
                              key_padding_mask=tgt_key_padding_mask)[0]
        tgt = tgt + self.dropout1(tgt2)
        tgt2 = self.norm2(tgt)
        tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt2, query_pos),
                                   key=self.with_pos_embed(memory, pos),
                                   value=memory, attn_mask=memory_mask,
                                   key_padding_mask=memory_key_padding_mask)[0]
        tgt = tgt + self.dropout2(tgt2)
        tgt2 = self.norm3(tgt)
        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
        tgt = tgt + self.dropout3(tgt2)
        return tgt

    def forward(self, tgt, memory,
                tgt_mask: Optional[Tensor] = None,
                memory_mask: Optional[Tensor] = None,
                tgt_key_padding_mask: Optional[Tensor] = None,
                memory_key_padding_mask: Optional[Tensor] = None,
                pos: Optional[Tensor] = None,
                query_pos: Optional[Tensor] = None):
        if self.normalize_before:
            return self.forward_pre(tgt, memory, tgt_mask, memory_mask,
                                    tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos)
        return self.forward_post(tgt, memory, tgt_mask, memory_mask,
                                 tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos)



def _get_clones(module, N):
    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])


def build_transformer(args):
    return Transformer(
        d_model=args.hidden_dim,
        dropout=args.dropout,
        nhead=args.nheads,
        dim_feedforward=args.dim_feedforward,
        num_encoder_layers=args.enc_layers,
        num_decoder_layers=args.dec_layers,
        normalize_before=args.pre_norm,
        return_intermediate_dec=True,
    )


def _get_activation_fn(activation):
    """Return an activation function given a string"""
    if activation == "relu":
        return F.relu
    if activation == "gelu":
        return F.gelu
    if activation == "glu":
        return F.glu
    raise RuntimeError(F"activation should be relu/gelu, not {activation}.")


class PositionalEncoding(nn.Module):
    '''
        Transformer-style positional encoding with wavelets
    '''

    def __init__(self, dim_embed, max_len=500):
        super().__init__()

        pe = torch.zeros(max_len, dim_embed)
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp((torch.arange(0, dim_embed, 2, dtype=torch.float) *
                -(math.log(10000.0) / dim_embed)))
        pe[:, 0::2] = torch.sin(position.float() * div_term)
        pe[:, 1::2] = torch.cos(position.float() * div_term)
        
        self.pe = pe # size=(max_len, dim_embed)
        self.dim_embed = dim_embed

    def forward(self, step_ids):
        if step_ids.device != self.pe.device:
            self.pe = self.pe.to(step_ids.device)
        return self.pe[step_ids]


def triangular_mask(size, device, diagonal_shift=1):
    '''
    generate upper triangular matrix filled with ones
    '''
    square = torch.triu(torch.ones(size, size, device=device).bool(), diagonal=diagonal_shift)
    # square = square.masked_fill(square == 1., float('-inf'))
    return square

offline_bc/models/transformer.py
cshizhe/onav_rim
"""
Copyright (c) Microsoft Corporation.
Licensed under the MIT license.

optimizer learning rate scheduling helpers
"""
from math import ceil


def noam_schedule(step, warmup_step=4000):
    """ original Transformer schedule"""
    if step <= warmup_step:
        return step / warmup_step
    return (warmup_step ** 0.5) * (step ** -0.5)


def warmup_linear(step, warmup_step, tot_step):
    """ BERT schedule """
    if step < warmup_step:
        return step / warmup_step
    return max(0, (tot_step-step)/(tot_step-warmup_step))

def warmup_inverse_sqrt(step, warmup_step, tot_step):
    """Decay the LR based on the inverse square root of the update number.
    We also support a warmup phase where we linearly increase the learning rate
    from some initial learning rate (``--warmup-init-lr``) until the configured
    learning rate (``--lr``). Thereafter we decay proportional to the number of
    updates, with a decay factor set to align with the configured learning rate.

    During warmup::

      lrs = torch.linspace(cfg.warmup_init_lr, cfg.lr, cfg.warmup_updates)
      lr = lrs[update_num]

    After warmup::

      decay_factor = cfg.lr * sqrt(cfg.warmup_updates)
      lr = decay_factor / sqrt(update_num)
    """
    if step < warmup_step:
        return step / warmup_step
    else:
        return warmup_step**0.5 * step**-0.5


def get_lr_sched(global_step, opts):
    # learning rate scheduling
    if opts.lr_sched == 'linear':
        func = warmup_linear
    elif opts.lr_sched == 'inverse_sqrt':
        func = warmup_inverse_sqrt
    else:
        raise NotImplementedError(f'invalid lr scheduler {opts.lr_sched}')

    lr_this_step = opts.learning_rate * func(
        global_step, opts.warmup_steps, opts.num_train_steps
    )
    if lr_this_step <= 0:
        lr_this_step = 1e-8
    return lr_this_step

offline_bc/optim/sched.py
cshizhe/onav_rim
import numpy as np

import torch
import torch.nn as nn

import torchvision
import torchvision.transforms as T

import clip
from habitat_baselines.il.common.encoders.resnet_encoders import VlnResnetDepthEncoder
from habitat_baselines.il.env_based.policy.rednet import load_rednet


class CLIPEncoder(object):
    def __init__(self, device, batch_size=64) -> None:
        # load clip model
        clip_model, clip_preprocess = clip.load('RN50', device=device)
        clip_model = clip_model.visual
        clip_model.attnpool = nn.Sequential(
            nn.AdaptiveAvgPool2d(output_size=(1, 1)),
            nn.Flatten()
        )
        clip_model.eval()
        clip_preprocess = T.Compose([
            T.ToTensor(),
            clip_preprocess.transforms[4] # normalize
        ])
        self.model = clip_model
        self.preprocess = clip_preprocess
        self.device = device
        self.batch_size = batch_size

    def extract_fts(self, rgb_images):
        '''
        images: list of PIL.Image (preprocessed 224 x 224 x 3)
        '''
        fts = []
        for i in range(0, len(rgb_images), self.batch_size):
            # (batch, 3, 224, 224)
            inputs = torch.stack(
                [self.preprocess(x) for x in rgb_images[i: i+self.batch_size]], 0
            ).to(self.device)
            # (batch, 2048)
            fts.append(self.model(inputs).data.cpu().numpy())
        fts = np.concatenate(fts, 0)
        return fts

class ResnetEncoders(object):
    def __init__(self, device, backbones=['resnet50'], batch_size=64) -> None:
        self.device = device
        self.preprocess = T.Compose([
            T.ToTensor(),
            T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
        self.models = {
            backbone: getattr(torchvision.models, backbone)(pretrained=True).to(self.device) \
                for backbone in backbones
        }
        for name, model in self.models.items():
            model.fc = nn.Identity()
            model.eval()
        self.batch_size = batch_size

    def extrat_fts(self, rgb_images):
        '''
        images: list of PIL.Image (preprocessed 224 x 224 x 3)
        '''
        fts = {name: [] for name in self.models.keys()}
        for name, model in self.models.items():
            for i in range(0, len(rgb_images), self.batch_size):
                # (batch, 3, 224, 224)
                inputs = torch.stack(
                    [self.preprocess(x) for x in rgb_images[i: i+self.batch_size]], 0
                ).to(self.device)
                # (batch, 2048)
                fts[name].append(model(inputs).data.cpu().numpy())
        fts = {name: np.concatenate(value, 0) for name, value in fts.items()}
        return fts

class DepthEncoder(object):
    def __init__(self, device, observation_space, ckpt_file, backbone, batch_size=64) -> None:
        depth_encoder = VlnResnetDepthEncoder(
            observation_space,
            output_size=256, # useless
            checkpoint=ckpt_file,
            backbone=backbone,
            trainable=False,
        )
        depth_encoder.eval()
        self.model = depth_encoder.to(device)
        self.device = device
        self.batch_size = batch_size
    
    def extract_fts(self, depth_images):
        '''
        depth_images: Tensor (batch, height, width, 1)
        '''
        fts = []
        for i in range(0, len(depth_images), self.batch_size):
            ft = self.model.visual_encoder(
                {'depth': depth_images[i: i+self.batch_size]}
            ) # (batch, 128, 4, 4)
            fts.append(ft.flatten(start_dim=1).data.cpu().numpy())
        fts = np.concatenate(fts, 0)
        return fts

class SemanticPredictor(object):
    def __init__(self, device, ckpt_file, num_classes, batch_size=64) -> None:
        semantic_predictor = load_rednet(
            device,
            ckpt=ckpt_file,
            resize=True, # since we train on half-vision
            num_classes=num_classes
        )
        semantic_predictor.eval()
        self.model = semantic_predictor
        self.device = device
        self.batch_size = batch_size

    def predict(self, rgb_images, depth_images):
        preds = []
        for i in range(0, len(rgb_images), self.batch_size):
            preds.append(self.model(
                rgb_images[i: i+self.batch_size], depth_images[i: i+self.batch_size]
            ).data.cpu().numpy())
        preds = np.concatenate(preds, 0)
        return preds


offline_bc/preprocess/encoders.py
cshizhe/onav_rim
'''
For each demo episode, extract:
- RGB images: T x (224, 224, 3)
- Depth features extracted by a pretrained Resnet50 (DDPPO): (T, 2048)
- Predicted semantics by a pretrained Rednet segmentor: T x (224, 224, 28)
'''

import os
import argparse
import json
import jsonlines
import numpy as np
import cv2
from PIL import Image
from tqdm import tqdm, trange
import collections

import lmdb
import msgpack
import msgpack_numpy
msgpack_numpy.patch()

from habitat_baselines.config.default import get_config
from habitat_baselines.common.environments import NavRLEnv

import torch
import torch.nn as nn

import torchvision.transforms as T

from offline_bc.preprocess.encoders import (
    CLIPEncoder,
    ResnetEncoders,
    DepthEncoder,
    SemanticPredictor,
)

ACTION_MAPS = {
    'STOP': 0, 'MOVE_FORWARD': 1, 'TURN_LEFT': 2, 'TURN_RIGHT': 3, 'LOOK_UP': 4, 'LOOK_DOWN': 5
}

category_to_task_category_id = {
    'chair': 0,
    'table': 1,
    'picture': 2,
    'cabinet': 3,
    'cushion': 4,
    'sofa': 5,
    'bed': 6,
    'chest_of_drawers': 7,
    'plant': 8,
    'sink': 9,
    'toilet': 10,
    'stool': 11,
    'towel': 12,
    'tv_monitor': 13,
    'shower': 14,
    'bathtub': 15,
    'counter': 16,
    'fireplace': 17,
    'gym_equipment': 18,
    'seating': 19,
    'clothes': 20
}
category_to_mp3d_category_id = {
    'chair': 3,
    'table': 5,
    'picture': 6,
    'cabinet': 7,
    'cushion': 8,
    'sofa': 10,
    'bed': 11,
    'chest_of_drawers': 13,
    'plant': 14,
    'sink': 15,
    'toilet': 18,
    'stool': 19,
    'towel': 20,
    'tv_monitor': 22,
    'shower': 23,
    'bathtub': 25,
    'counter': 26,
    'fireplace': 27,
    'gym_equipment': 33,
    'seating': 34,
    'clothes': 38
}


LABEL_MAP = {}
for k, v in category_to_mp3d_category_id.items():
    LABEL_MAP[v] = category_to_task_category_id[k]


def extract_demo_obs_and_fts(args):
    # build env
    config = get_config(args.cfg_file)
    config.defrost()

    cfg = config.TASK_CONFIG
    cfg.DATASET.CONTENT_SCENES = [args.scene_id]

    cfg.DATASET.MAX_REPLAY_STEPS = 2000
    cfg.ENVIRONMENT.MAX_EPISODE_STEPS = 2000

    cfg.ENVIRONMENT.ITERATOR_OPTIONS.CYCLE = False
    cfg.ENVIRONMENT.ITERATOR_OPTIONS.SHUFFLE = False
    cfg.ENVIRONMENT.ITERATOR_OPTIONS.MAX_SCENE_REPEAT_EPISODES = -1
    cfg.ENVIRONMENT.ITERATOR_OPTIONS.MAX_SCENE_REPEAT_STEPS = -1

    orientation = np.deg2rad(args.camera_init_pitch)
    cfg.SIMULATOR.DEPTH_SENSOR.HFOV = args.camera_hfov
    cfg.SIMULATOR.DEPTH_SENSOR.POSITION = [0, args.camera_height, 0]
    cfg.SIMULATOR.DEPTH_SENSOR.ORIENTATION = [orientation, 0, 0]
    cfg.SIMULATOR.RGB_SENSOR.HFOV = args.camera_hfov
    cfg.SIMULATOR.RGB_SENSOR.POSITION = [0, args.camera_height, 0]
    cfg.SIMULATOR.RGB_SENSOR.ORIENTATION = [orientation, 0, 0]
    cfg.SIMULATOR.SEMANTIC_SENSOR.HFOV = args.camera_hfov
    cfg.SIMULATOR.SEMANTIC_SENSOR.POSITION = [0, args.camera_height, 0]
    cfg.SIMULATOR.SEMANTIC_SENSOR.ORIENTATION = [orientation, 0, 0]
    cfg.SIMULATOR.AGENT_0.SENSORS.append('SEMANTIC_SENSOR')

    config.freeze()

    env = NavRLEnv(config)
    num_episodes = len(env.episodes)
    print('num episodes:', num_episodes)

    # build encoders
    device = torch.device('cuda:0')
    torch.set_grad_enabled(False)

    rgb_preprocess = T.Compose([
        T.Resize(args.image_size),
        T.CenterCrop(size=(args.image_size, args.image_size))
    ])

    if args.encode_rgb_clip:
        rgb_clip_encoder = CLIPEncoder(device, batch_size=args.batch_size)
    if args.encode_rgb_resnet:
        rgb_resnet_encoders = ResnetEncoders(device, ['resnet18', 'resnet50'], batch_size=args.batch_size)
    if args.encode_depth:
        depth_encoder = DepthEncoder(
            device, env.observation_space,
            config.MODEL.DEPTH_ENCODER.ddppo_checkpoint,
            config.MODEL.DEPTH_ENCODER.backbone,
            batch_size=args.batch_size
        )
    if args.predict_semantic:
        semantic_predictor = SemanticPredictor(
            device, 
            config.MODEL.SEMANTIC_ENCODER.rednet_ckpt,
            config.MODEL.SEMANTIC_ENCODER.num_classes,
            batch_size=args.batch_size
        )

    # extract observations and features
    if args.save_rgb:
        os.makedirs(os.path.join(args.outdir, 'rgb_images'), exist_ok=True)
        rgb_lmdb_env = lmdb.open(
            os.path.join(args.outdir, 'rgb_images', args.scene_id),
            map_size=int(1024**4)
        )
    if args.save_semantic:
        os.makedirs(os.path.join(args.outdir, 'sem_images'), exist_ok=True)
        sem_lmdb_env = lmdb.open(
            os.path.join(args.outdir, 'sem_images', args.scene_id),
            map_size=int(1024**4)
        )
    if args.save_semantic_fts:
        os.makedirs(os.path.join(args.outdir, 'sem_fts'), exist_ok=True)
        sem_ft_lmdb_env = lmdb.open(
            os.path.join(args.outdir, 'sem_fts', args.scene_id),
            map_size=int(1024**4)
        )

    rgb_encoder_names = []
    rgb_ft_lmdb_envs = {}
    if args.encode_rgb_clip:
        rgb_encoder_names.append('clip')
    if args.encode_rgb_resnet:
        rgb_encoder_names.extend(['resnet18', 'resnet50'])
    for name in rgb_encoder_names:
        os.makedirs(os.path.join(args.outdir, 'rgb_fts', name), exist_ok=True)
        rgb_ft_lmdb_envs[name] = lmdb.open(
            os.path.join(args.outdir, 'rgb_fts', name, args.scene_id),
            map_size=int(1024**4)
        )

    if args.encode_depth:
        os.makedirs(os.path.join(args.outdir, 'depth_fts'), exist_ok=True)
        depth_ft_lmdb_env = lmdb.open(
            os.path.join(args.outdir, 'depth_fts', args.scene_id),
            map_size=int(1024**4)
        )

    if args.predict_semantic:
        os.makedirs(os.path.join(args.outdir, 'semantic_preds'), exist_ok=True)
        sem_preds_lmdb_env = lmdb.open(
            os.path.join(args.outdir, 'semantic_preds', args.scene_id),
            map_size=int(1024**4)
        )

    os.makedirs(os.path.join(args.outdir, 'meta_infos_lmdb'), exist_ok=True)
    meta_info_lmdb_env = lmdb.open(
        os.path.join(args.outdir, 'meta_infos_lmdb', args.scene_id),
        map_size=int(1024**4)
    )

    for _ in trange(num_episodes):
        observations = env.reset()
        episode_id = env.current_episode.episode_id
        # rgb <class 'numpy.ndarray'> (480, 640, 3) uint8 0 255
        # depth <class 'numpy.ndarray'> (480, 640, 1) float32 0.0 1.0
        # objectgoal <class 'numpy.ndarray'> (1,) int64 5 5
        # compass <class 'numpy.ndarray'> (1,) float32 1.0365224e-09 1.0365224e-09
        # gps <class 'numpy.ndarray'> (2,) float32 -0.0 -0.0
        # demonstration <class 'int'> 3
        # inflection_weight <class 'float'> 1.0

        demo_actions = [
            x.action if x.action is not None else 'STOP' \
                for x in env.current_episode.reference_replay[1:]
        ]
        demo_actions = [ACTION_MAPS[x] for x in demo_actions]

        episode_obs = {'reward': [], 'info': []}

        for action in demo_actions:
            for k, v in observations.items():
                episode_obs.setdefault(k, [])
                episode_obs[k].append(v)
            observations, reward, done, info = env.step(action=action)
            episode_obs['reward'].append(reward)
            episode_obs['info'].append(info)
            if action == 'STOP':
                break

        # remove collision steps
        if not args.keep_collision_steps:
            collide_step_ids = set()
            gps = episode_obs['gps']
            compass = episode_obs['compass']
            for t in range(len(gps)-1):
                # look up and down also result in the same gps and compass
                if all(np.isclose(gps[t], gps[t+1])) and all(np.isclose(compass[t], compass[t+1])) and demo_actions[t] in [1, 2, 3]:
                    collide_step_ids.add(t)
                    
            for k, v in episode_obs.items():
                episode_obs[k] = [v[idx] for idx in range(len(v)) if idx not in collide_step_ids]
            demo_actions = [x for idx, x in enumerate(demo_actions) if idx not in collide_step_ids]

        resized_rgb_images = [rgb_preprocess(Image.fromarray(x)) for x in episode_obs['rgb']]

        rgb_fts = {}
        if args.encode_rgb_clip:
            rgb_fts['clip'] = rgb_clip_encoder.extract_fts(resized_rgb_images)
        if args.encode_rgb_resnet:
            rgb_fts.update(rgb_resnet_encoders.extrat_fts(resized_rgb_images))

        if args.encode_depth:
            depth_images = torch.stack([torch.from_numpy(x) for x in episode_obs['depth']]).to(device)
            depth_fts = depth_encoder.extract_fts(depth_images) # (batch, dim_ft)

        if args.predict_semantic:
            rgb_images = torch.from_numpy(np.stack(episode_obs['rgb'], 0)).to(device)
            sem_preds = semantic_predictor.predict(rgb_images, depth_images) # (batch, height, width)
        
        # save to files
        episode_key = episode_id.encode('ascii')

        if args.save_rgb:  
            for t, x in enumerate(resized_rgb_images):
                x = np.array(x)
                _, x = cv2.imencode('.png', x, [cv2.IMWRITE_PNG_COMPRESSION, 3])
                txn = rgb_lmdb_env.begin(write=True)
                txn.put(('%s:%04d'%(episode_id, t)).encode('ascii'), x)
                txn.commit()

        if args.save_semantic:
            for t, x in enumerate(episode_obs['semantic']):
                _, x = cv2.imencode('.png', x, [cv2.IMWRITE_PNG_COMPRESSION, 3])
                txn = sem_lmdb_env.begin(write=True)
                txn.put(('%s:%04d'%(episode_id, t)).encode('ascii'), x)
                txn.commit()

        if args.save_semantic_fts:
            sem_fts = np.zeros((len(episode_obs['semantic']), len(LABEL_MAP)), dtype=np.float32)
            for t, x in enumerate(episode_obs['semantic']):
                x = x.flatten()
                label_counter = collections.Counter(x)
                npixels = len(x)
                for label, count in label_counter.most_common():
                    if label in LABEL_MAP:
                        sem_fts[t, LABEL_MAP[label]] = count / npixels
            txn = sem_ft_lmdb_env.begin(write=True)
            txn.put(episode_key, msgpack.packb(sem_fts))
            txn.commit()
        
        if args.encode_rgb_clip or args.encode_rgb_resnet:
            for k, v in rgb_fts.items():
                txn = rgb_ft_lmdb_envs[k].begin(write=True)
                txn.put(episode_key, msgpack.packb(v))
                txn.commit()

        if args.encode_depth:
            txn = depth_ft_lmdb_env.begin(write=True)
            txn.put(episode_key, msgpack.packb(depth_fts))
            txn.commit()

        if args.predict_semantic:            
            sem_preds = np.concatenate([
                np.array(x) for x in sem_preds
            ], 0)
            _, sem_preds = cv2.imencode(
                '.png', sem_preds, 
                [cv2.IMWRITE_PNG_COMPRESSION, 3]
            )
            txn = sem_preds_lmdb_env.begin(write=True)
            txn.put(episode_key, sem_preds)
            txn.commit()

        meta_info = {
            'episode_id': episode_key.decode('ascii'),
            'objectgoal': episode_obs['objectgoal'][0].item(),
            'object_category': env.current_episode.object_category,
            'compass': np.stack(episode_obs['compass'], 0).tolist(),
            'gps': np.stack(episode_obs['gps'], 0).tolist(),
            'demonstration': demo_actions,
            # 'demonstration': np.array(episode_obs['demonstration']).tolist(),
            # 'inflection_weight': np.array(episode_obs['inflection_weight']).tolist(),
            'reward': np.array(episode_obs['reward']).tolist(),
            'info': episode_obs['info'],
        }
        txn = meta_info_lmdb_env.begin(write=True)
        txn.put(episode_key, msgpack.packb(meta_info))
        txn.commit()

    if args.save_rgb:
        rgb_lmdb_env.close()
    if args.save_semantic:
        sem_lmdb_env.close()
    if args.save_semantic_fts:
        sem_ft_lmdb_env.close()
    if args.encode_rgb_clip or args.encode_rgb_resnet:
        for k, v in rgb_ft_lmdb_envs.items():
            v.close()
    if args.encode_depth:
        depth_ft_lmdb_env.close()
    if args.predict_semantic:
        sem_preds_lmdb_env.close()
    meta_info_lmdb_env.close()

        
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--cfg_file', help='habitat_baselines config file', 
        default='habitat_baselines/config/objectnav/il_ddp_objectnav.yaml'
    )
    parser.add_argument('--scene_id', help='scene id in mp3d', default='gZ6f7yhEvPG')

    parser.add_argument('--camera_hfov', type=int, default=79, help='degree')
    parser.add_argument('--camera_height', type=float, default=0.88, help='meter')
    parser.add_argument('--camera_init_pitch', type=int, default=0,
                        help='+lookup, -lookdown (degrees)')
    # parser.add_argument('--agent_height', type=float, default=0.88)
    # parser.add_argument('--agent_radius', type=float, default=0.18)

    parser.add_argument('--save_rgb', action='store_true', default=False)
    parser.add_argument('--save_semantic', action='store_true', default=False)
    parser.add_argument('--save_semantic_fts', action='store_true', default=False)

    parser.add_argument('--encode_depth', action='store_true', default=False)
    parser.add_argument('--encode_rgb_clip', action='store_true', default=False)
    parser.add_argument('--encode_rgb_resnet', action='store_true', default=False)
    parser.add_argument('--predict_semantic', action='store_true', default=False)

    parser.add_argument('--image_size', type=int, default=224)
    parser.add_argument('--batch_size', type=int, default=128)

    parser.add_argument('--keep_collision_steps', action='store_true', default=False)

    parser.add_argument('--outdir', help='output directory')
    
    args = parser.parse_args()

    extract_demo_obs_and_fts(args)


if __name__ == '__main__':
    main()

offline_bc/preprocess/extract_demo_observations.py
cshizhe/onav_rim
import os
import sys
import json
import argparse
import time
from collections import defaultdict
from tqdm import tqdm

import torch
import torch.nn.functional as F
import torch.distributed as dist

from utils.logger import LOGGER, TB_LOGGER, RunningMeter, add_log_to_file
from utils.save import ModelSaver, save_training_meta
from utils.misc import NoOp, set_dropout, set_random_seed, set_cuda, wrap_model
from utils.distributed import all_gather

from optim import get_lr_sched
from optim.misc import build_optimizer

from offline_bc.config.default import get_config
from offline_bc.data.dataset import NavDemoDataset, collate_fn
from offline_bc.data.loader import build_dataloader

from offline_bc.models.onav_vis_models import (
    NavILRecurrentNet, NavILTransformer
)
from offline_bc.models.onav_imap_models import (
    NavImapSingleTransformer
)

dataset_factory = {
    'NavDemoDataset': (NavDemoDataset, collate_fn),
}
model_factory = {
    'NavILRecurrentNet': NavILRecurrentNet,
    'NavILTransformer': NavILTransformer,
    'NavImapSingleTransformer': NavImapSingleTransformer,
}

def main(config):
    config.defrost()
    if isinstance(config.DATASET.num_ft_views, str):
        config.MODEL.num_ft_views = int(config.DATASET.num_ft_views.split('_')[0])
    else:
        config.MODEL.num_ft_views = config.DATASET.num_ft_views
    default_gpu, n_gpu, device = set_cuda(config)
    # config.freeze()

    if default_gpu:
        LOGGER.info(
            'device: {} n_gpu: {}, distributed training: {}'.format(
                device, n_gpu, bool(config.local_rank != -1)
            )
        )
 
    seed = config.SEED
    if config.local_rank != -1:
        seed += config.rank
    set_random_seed(seed)

    # load data training set
    dataset_class, dataset_collate_fn = dataset_factory[config.DATASET.dataset_class]

    trn_dataset = dataset_class(**config.DATASET)
    trn_data_loader, trn_pre_epoch = build_dataloader(
        trn_dataset, dataset_collate_fn, True, config
    )
    LOGGER.info(f'#num_steps_per_epoch: {len(trn_data_loader)}')
    config.num_train_steps = len(trn_data_loader) * config.num_epochs
    config.freeze()

    if len(config.DATASET.val_scene_ids) > 0:
        val_dataset = dataset_class(**config.DATASET, validation=True)
        val_data_loader, _ = build_dataloader(
            val_dataset, dataset_collate_fn, False, config
        )
    else:
        val_data_loader = None

    # setup loggers
    if default_gpu:
        save_training_meta(config)
        TB_LOGGER.create(os.path.join(config.output_dir, 'logs'))
        pbar = tqdm(total=config.num_train_steps)
        model_saver = ModelSaver(os.path.join(config.output_dir, 'ckpts'))
        add_log_to_file(os.path.join(config.output_dir, 'logs', 'log.txt'))
    else:
        LOGGER.disabled = True
        pbar = NoOp()
        model_saver = NoOp()

    # Prepare model
    model_class = model_factory[config.MODEL.model_class]
    model = model_class(config.MODEL, device)

    LOGGER.info("Model: nweights %d nparams %d" % (model.num_parameters))
    LOGGER.info("Model: trainable nweights %d nparams %d" % (model.num_trainable_parameters))
    # for k, v in model.named_parameters():
    #     if v.requires_grad:
    #         print(k, v.size())

    if config.resume_file:
        checkpoint = torch.load(config.resume_file, map_location=lambda storage, loc: storage)
        LOGGER.info('resume: #params %d' % (len(checkpoint)))
        model.load_state_dict(checkpoint, strict=True)
        # model.load_state_dict(checkpoint, strict=False)
    
    model.train()
    set_dropout(model, config.MODEL.dropout_rate)
    model = wrap_model(model, device, config.local_rank)
    global_step = 0

    # Prepare optimizer
    optimizer = build_optimizer(model, config)
    if config.resume_optimizer:
        optimizer_state = torch.load(config.resume_optimizer)
        print('load optimizer step: %d, weights: %d' % (
            optimizer_state['step'], len(optimizer_state['optimizer']))
        )
        optimizer.load_state_dict(optimizer_state['optimizer'])
        global_step = optimizer_state['step']

    LOGGER.info(f"***** Running training with {config.world_size} GPUs *****")
    LOGGER.info("  Batch size = %d", config.train_batch_size if config.local_rank == -1 else config.train_batch_size * config.world_size)
    LOGGER.info("  Accumulate steps = %d", config.gradient_accumulation_steps)
    LOGGER.info("  Num steps = %d", config.num_train_steps)

    # to compute training statistics
    meteors = {
        key: RunningMeter(key) for key in ['loss', 'acc']
    }
    
    start_time = time.time()
    # quick hack for amp delay_unscale bug
    optimizer.zero_grad()
    optimizer.step()

    for epoch_id in range(config.num_epochs):
        # In distributed mode, calling the set_epoch() method at the beginning of each epoch
        trn_pre_epoch(epoch_id)

        for step, batch in enumerate(trn_data_loader):
            # forward pass
            loss, logits = model(batch, compute_loss=True)
            if isinstance(loss, dict):
                loss_dict = loss
                for loss_key, loss_value in loss.items():
                    if loss_key == 'overall': continue
                    if loss_key not in meteors:
                        meteors[loss_key] = RunningMeter(loss_key)
                    meteors[loss_key](loss_value.item())
                loss = loss['overall']
            
            # backward pass
            if config.gradient_accumulation_steps > 1: # average loss 
                loss = loss / config.gradient_accumulation_steps
            loss.backward()

            meteors['loss'](loss.item())
            if isinstance(logits, tuple):
                logits, subgoal_logits = logits
                if 'subgoal_action_acc' not in meteors:
                    meteors['subgoal_action_acc'] = RunningMeter('subgoal_action_acc')
                acc = torch.mean(
                    (subgoal_logits[:, :3].max(dim=1)[1].data.cpu() == batch['subgoal_actions'])[batch['subgoal_actions'] != -100].float()
                )
                meteors['subgoal_action_acc'](acc.item())
                if 'subgoal_turn_acc' not in meteors:
                    meteors['subgoal_turn_acc'] = RunningMeter('subgoal_turn_acc')
                acc = torch.mean(
                    (subgoal_logits[:, 3:25].max(dim=1)[1].data.cpu() == batch['subgoal_turn_degrees'])[batch['subgoal_actions'] == 2].float()
                )
                meteors['subgoal_turn_acc'](acc.item())
                if 'subgoal_goto_mae' not in meteors:
                    meteors['subgoal_goto_mae'] = RunningMeter('subgoal_goto_mae')
                meteors['subgoal_goto_mae'](loss_dict['subgoal_goto_loss'].sqrt().item() * config.DATASET.max_goto_dist)
            
            acc = torch.mean(
                (logits.max(dim=-1)[1].data.cpu() == batch['demonstration'])[batch['inflection_weight'] > 0].float()
            )
            meteors['acc'](acc.item())
            TB_LOGGER.add_scalar('step/loss', loss.item(), global_step)
            TB_LOGGER.add_scalar('step/acc', acc.item(), global_step)

            # optimizer update and logging
            if (global_step + 1) % config.gradient_accumulation_steps == 0:
                global_step += 1

                # learning rate scheduling
                lr_this_step = get_lr_sched(global_step, config)
                for param_group in optimizer.param_groups:
                    param_group['lr'] = lr_this_step
                TB_LOGGER.add_scalar('lr', lr_this_step, global_step)

                # log loss
                # NOTE: not gathered across GPUs for efficiency
                TB_LOGGER.log_scalar_dict({ll.name: ll.val for ll in meteors.values()})
                TB_LOGGER.step()

                # update model params
                if config.grad_norm != -1:
                    grad_norm = torch.nn.utils.clip_grad_norm_(
                        model.parameters(), config.grad_norm
                    )
                    # print(global_step, grad_norm)
                    # for k, v in model.named_parameters():
                    #     if v.grad is not None:
                    #         v = torch.norm(v).data.item()
                    #         print(k, v)
                    TB_LOGGER.add_scalar('grad_norm', grad_norm, global_step)
                optimizer.step()
                optimizer.zero_grad()
                pbar.update(1)

            if global_step % config.log_steps == 0:
                # monitor training throughput
                LOGGER.info(f'==============Epoch {epoch_id} Step {global_step}===============')
                LOGGER.info(', '.join(['%s:%.4f'%(ll.name, ll.val) for ll in meteors.values()]))
                LOGGER.info('===============================================')

            if global_step % config.valid_steps == 0:
                model_saver.save(model, global_step, optimizer=optimizer)  
                if val_data_loader is not None:
                    LOGGER.info(f'-----Epoch {epoch_id} Step {global_step}: start validation-----')
                    validate(model, val_data_loader, TB_LOGGER)

            if global_step > config.num_train_steps:
                break

        if global_step > config.num_train_steps:
            break

    if global_step % config.valid_steps != 0:
        LOGGER.info(f'==============Epoch {epoch_id} Step {global_step}===============')
        LOGGER.info(', '.join(['%s:%.4f'%(ll.name, ll.val) for ll in meteors.values()]))
        LOGGER.info('===============================================')
        model_saver.save(model, global_step, optimizer=optimizer)
        if val_data_loader is not None:
            LOGGER.info(f'-----Epoch {epoch_id} Step {global_step}: start validation-----')
            validate(model, val_data_loader, TB_LOGGER)
    

@torch.no_grad()
def validate(model, data_loader, TB_LOGGER=None):
    model.eval()

    st = time.time()
    val_loss, n_correct, n_total = 0, 0, 0
    n_stop_correct, n_stop_total = 0, 0
    for i, batch in enumerate(data_loader):
        logits = model(batch, compute_loss=False).data.cpu()
        preds = logits.max(dim=-1)[1]
        labels = batch['demonstration']
        # print(preds)
        # print(labels)
        loss = F.cross_entropy(logits.permute(0, 2, 1), labels, reduction='sum', ignore_index=-100)
        val_loss += loss.item()
        n_correct += (preds == labels)[labels != -100].float().sum().item()
        n_total += (labels != -100).float().sum().item()
        n_stop_correct += (preds == labels)[labels == 0].float().sum().item()
        n_stop_total += (labels == 0).float().sum().item()
    val_loss = sum(all_gather(val_loss))
    n_correct = sum(all_gather(n_correct))
    n_total = sum(all_gather(n_total))
    n_stop_correct = sum(all_gather(n_stop_correct))
    n_stop_total = sum(all_gather(n_stop_total))

    val_loss /= n_total
    acc = n_correct / n_total
    stop_acc = n_stop_correct / n_stop_total
    val_log = {'loss': val_loss, 'acc': acc, 'stop_acc': stop_acc}

    tot_time = time.time() - st
    LOGGER.info(f"validation finished in {int(tot_time)} seconds, "
                f"loss: {val_loss:.4f}, acc: {acc*100:.2f}, stop_acc: {stop_acc*100:.2f}")

    if TB_LOGGER is not None:
        TB_LOGGER.log_scalar_dict(
            {f'valid/{k}': v for k, v in val_log.items()}
        )

    model.train()


def build_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--exp-config",
        type=str,
        required=True,
        help="path to config yaml containing info about experiment",
    )
    parser.add_argument(
        "opts",
        default=None,
        nargs=argparse.REMAINDER,
        help="Modify config options from command line",
    )
    args = parser.parse_args()

    config = get_config(args.exp_config, args.opts)
    
    if os.path.exists(config.output_dir) and os.listdir(config.output_dir):
        LOGGER.warning(
            "Output directory ({}) already exists and is not empty.".format(
                config.output_dir
            )
        )

    return config

if __name__ == '__main__':
    config = build_args()
    main(config)

offline_bc/train_models.py
cshizhe/onav_rim
"""
Distributed tools
"""
import os
from pathlib import Path
from pprint import pformat
import pickle

import torch
import torch.distributed as dist


def set_local_rank(opts) -> int:
    if os.environ.get("LOCAL_RANK", "") != "":
        opts.local_rank = int(os.environ["LOCAL_RANK"])
    elif os.environ.get("SLURM_LOCALID", "") != "":
        opts.local_rank = int(os.environ["SLURM_LOCALID"])
    else:
        opts.local_rank = -1
    return opts.local_rank


def load_init_param(opts):
    """
    Load parameters for the rendezvous distributed procedure
    """
    # num of gpus per node
    # WARNING: this assumes that each node has the same number of GPUs
    if os.environ.get("SLURM_NTASKS_PER_NODE", "") != "":
        num_gpus = int(os.environ['SLURM_NTASKS_PER_NODE'])
    else:
        num_gpus = torch.cuda.device_count()

    # world size
    if os.environ.get("WORLD_SIZE", "") != "":
        world_size = int(os.environ["WORLD_SIZE"])
    elif os.environ.get("SLURM_JOB_NUM_NODES", ""):
        num_nodes = int(os.environ["SLURM_JOB_NUM_NODES"])
        world_size = num_nodes * num_gpus
    else:
        raise RuntimeError("Can't find any world size")
    opts.world_size = world_size

    # rank
    if os.environ.get("RANK", "") != "":
        # pytorch.distributed.launch provide this variable no matter what
        opts.rank = int(os.environ["RANK"])
    elif os.environ.get("SLURM_PROCID", "") != "":
        opts.rank = int(os.environ["SLURM_PROCID"])
    else:
        if os.environ.get("NODE_RANK", "") != "":
            opts.node_rank = int(os.environ["NODE_RANK"])
        elif os.environ.get("SLURM_NODEID", "") != "":
            opts.node_rank = int(os.environ["SLURM_NODEID"])
        else:
            raise RuntimeError("Can't find any rank or node rank")

        opts.rank = opts.local_rank + node_rank * num_gpus

    init_method = "env://" # need to specify MASTER_ADDR and MASTER_PORT
    
    return {
        "backend": "nccl",
        "init_method": init_method,
        "rank": opts.rank,
        "world_size": world_size,
    }


def init_distributed(opts):
    init_param = load_init_param(opts)
    rank = init_param["rank"]
    print(f"Init distributed {init_param['rank']} - {init_param['world_size']}")

    dist.init_process_group(**init_param)


def is_default_gpu(opts) -> bool:
    return opts.local_rank == -1 or dist.get_rank() == 0


def is_dist_avail_and_initialized():
    if not dist.is_available():
        return False
    if not dist.is_initialized():
        return False
    return True

def get_world_size():
    if not is_dist_avail_and_initialized():
        return 1
    return dist.get_world_size()

def all_gather(data):
    """
    Run all_gather on arbitrary picklable data (not necessarily tensors)
    Args:
        data: any picklable object
    Returns:
        list[data]: list of data gathered from each rank
    """
    world_size = get_world_size()
    if world_size == 1:
        return [data]

    # serialized to a Tensor
    buffer = pickle.dumps(data)
    storage = torch.ByteStorage.from_buffer(buffer)
    tensor = torch.ByteTensor(storage).to("cuda")

    # obtain Tensor size of each rank
    local_size = torch.tensor([tensor.numel()], device="cuda")
    size_list = [torch.tensor([0], device="cuda") for _ in range(world_size)]
    dist.all_gather(size_list, local_size)
    size_list = [int(size.item()) for size in size_list]
    max_size = max(size_list)

    # receiving Tensor from all ranks
    # we pad the tensor because torch all_gather does not support
    # gathering tensors of different shapes
    tensor_list = []
    for _ in size_list:
        tensor_list.append(torch.empty((max_size,), dtype=torch.uint8, device="cuda"))
    if local_size != max_size:
        padding = torch.empty(size=(max_size - local_size,), dtype=torch.uint8, device="cuda")
        tensor = torch.cat((tensor, padding), dim=0)
    dist.all_gather(tensor_list, tensor)

    data_list = []
    for size, tensor in zip(size_list, tensor_list):
        buffer = tensor.cpu().numpy().tobytes()[:size]
        data_list.append(pickle.loads(buffer))

    return data_list


def reduce_dict(input_dict, average=True):
    """
    Args:
        input_dict (dict): all the values will be reduced
        average (bool): whether to do average or sum
    Reduce the values in the dictionary from all processes so that all processes
    have the averaged results. Returns a dict with the same fields as
    input_dict, after reduction.
    """
    world_size = get_world_size()
    if world_size < 2:
        return input_dict
    with torch.no_grad():
        names = []
        values = []
        # sort the keys so that they are consistent across processes
        for k in sorted(input_dict.keys()):
            names.append(k)
            values.append(input_dict[k])
        values = torch.stack(values, dim=0)
        dist.all_reduce(values)
        if average:
            values /= world_size
        reduced_dict = {k: v for k, v in zip(names, values)}
    return reduced_dict



offline_bc/utils/distributed.py
cshizhe/onav_rim
import os
import random
import numpy as np
from typing import Tuple, Union, Dict, Any

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

from .distributed import init_distributed, set_local_rank
from .logger import LOGGER


def set_random_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

def set_dropout(model, drop_p):
    for name, module in model.named_modules():
        # we might want to tune dropout for smaller dataset
        if name.startswith('map_encoder'): continue
        if isinstance(module, torch.nn.Dropout):
            if module.p != drop_p:
                module.p = drop_p
                LOGGER.info(f'{name} set to {drop_p}')
            

def set_cuda(opts) -> Tuple[bool, int, torch.device]:
    """
    Initialize CUDA for distributed computing
    """
    set_local_rank(opts)

    if not torch.cuda.is_available():
        assert opts.local_rank == -1, opts.local_rank
        return True, 0, torch.device("cpu")

    # get device settings
    if opts.local_rank != -1:
        init_distributed(opts)
        torch.cuda.set_device(opts.local_rank)
        device = torch.device("cuda", opts.local_rank)
        n_gpu = 1
        default_gpu = dist.get_rank() == 0
        if default_gpu:
            LOGGER.info(f"Found {dist.get_world_size()} GPUs")
    else:
        default_gpu = True
        device = torch.device("cuda")
        n_gpu = torch.cuda.device_count()

    return default_gpu, n_gpu, device


def wrap_model(
    model: torch.nn.Module, device: torch.device, local_rank: int
) -> torch.nn.Module:
    model.to(device)

    if local_rank != -1:
        model = DDP(model, device_ids=[local_rank], find_unused_parameters=True)
        # At the time of DDP wrapping, parameters and buffers (i.e., model.state_dict()) 
        # on rank0 are broadcasted to all other ranks.
    elif torch.cuda.device_count() > 1:
        LOGGER.info("Using data parallel")
        model = torch.nn.DataParallel(model)

    return model


class NoOp(object):
    """ useful for distributed training No-Ops """
    def __getattr__(self, name):
        return self.noop

    def noop(self, *args, **kwargs):
        return


offline_bc/utils/misc.py
cshizhe/onav_rim
import numpy as np
import torch

def pad_tensors(tensors, lens=None, pad=0):
    """B x [T, ...] torch tensors"""
    if lens is None:
        lens = [t.size(0) for t in tensors]
    max_len = max(lens)
    bs = len(tensors)
    hid = list(tensors[0].size()[1:])
    size = [bs, max_len] + hid

    dtype = tensors[0].dtype
    output = torch.zeros(*size, dtype=dtype)
    if pad:
        output.data.fill_(pad)
    for i, (t, l) in enumerate(zip(tensors, lens)):
        output.data[i, :l, ...] = t.data
    return output

def pad_tensors_wgrad(tensors, lens=None, value=0):
    """B x [T, ...] torch tensors"""
    if lens is None:
        lens = [t.size(0) for t in tensors]
    max_len = max(lens)
    batch_size = len(tensors)
    hid = list(tensors[0].size()[1:])

    device = tensors[0].device
    dtype = tensors[0].dtype

    output = []
    for i in range(batch_size):
        if lens[i] < max_len:
            tmp = torch.cat(
                [tensors[i], torch.zeros([max_len-lens[i]]+hid, dtype=dtype).to(device) + value],
                dim=0
            )
        else:
            tmp = tensors[i]
        output.append(tmp)
    output = torch.stack(output, 0)
    return output


def gen_seq_masks(seq_lens, max_len=None):
    """
    Args:
        seq_lens: list or nparray int, shape=(N, )
    Returns:
        masks: nparray, shape=(N, L), padded=0
    """
    seq_lens = np.array(seq_lens)
    if max_len is None:
        max_len = max(seq_lens)
    if max_len == 0:
        return np.zeros((len(seq_lens), 0), dtype=np.bool)
    batch_size = len(seq_lens)
    masks = np.arange(max_len).reshape(-1, max_len).repeat(batch_size, 0)
    masks = masks < seq_lens.reshape(-1, 1)
    return masks


def extend_neg_masks(masks, dtype=None):
    """
    mask from (N, L) into (N, 1(H), 1(L), L) and make it negative
    """
    if dtype is None:
        dtype = torch.float
    extended_masks = masks.unsqueeze(1).unsqueeze(2)
    extended_masks = extended_masks.to(dtype=dtype)
    extended_masks = (1.0 - extended_masks) * -10000.0
    return extended_masks
offline_bc/utils/ops.py
cshizhe/onav_rim
"""
Copyright (c) Microsoft Corporation.
Licensed under the MIT license.

saving utilities
"""
import json
import os
import torch


def save_training_meta(args):
    os.makedirs(os.path.join(args.output_dir, 'logs'), exist_ok=True)
    os.makedirs(os.path.join(args.output_dir, 'ckpts'), exist_ok=True)

    with open(os.path.join(args.output_dir, 'logs', 'training_config.yaml'), 'w') as writer:
        args_str = args.dump()
        print(args_str, file=writer)


class ModelSaver(object):
    def __init__(self, output_dir, prefix='model_step', suffix='pt'):
        self.output_dir = output_dir
        self.prefix = prefix
        self.suffix = suffix

    def save(self, model, step, optimizer=None):
        output_model_file = os.path.join(self.output_dir,
                                 f"{self.prefix}_{step}.{self.suffix}")
        state_dict = {}
        for k, v in model.state_dict().items():
            if k.startswith('module.'):
                k = k[7:]
            if isinstance(v, torch.Tensor):
                state_dict[k] = v.cpu()
            else:
                state_dict[k] = v
        torch.save(state_dict, output_model_file)
        if optimizer is not None:
            dump = {'step': step, 'optimizer': optimizer.state_dict()}
            if hasattr(optimizer, '_amp_stash'):
                pass  # TODO fp16 optimizer
            torch.save(dump, f'{self.output_dir}/train_latest_state.pt')


offline_bc/utils/save.py
cshizhe/onav_rim