<filename>ET/alfred/utils/eval_util.py
import os
import json
import queue
import torch
import shutil
import filelock
import numpy as np

from PIL import Image
from termcolor import colored

from alfred.gen import constants
from alfred.env.thor_env import ThorEnv
from alfred.nn.enc_visual import FeatureExtractor
from alfred.utils import data_util, model_util


def setup_scene(env, traj_data, reward_type='dense', test_split=False):
    '''
    intialize the scene and agent from the task info
    '''
    # scene setup
    scene_num = traj_data['scene']['scene_num']
    object_poses = traj_data['scene']['object_poses']
    dirty_and_empty = traj_data['scene']['dirty_and_empty']
    object_toggles = traj_data['scene']['object_toggles']
    scene_name = 'FloorPlan%d' % scene_num
    env.reset(scene_name, silent=True)
    env.restore_scene(object_poses, object_toggles, dirty_and_empty)
    # initialize to start position
    env.step(dict(traj_data['scene']['init_action']))
    # setup task for reward
    if not test_split:
        env.set_task(traj_data, reward_type=reward_type)


def load_agent(model_path, dataset_info, device):
    '''
    load a pretrained agent and its feature extractor
    '''
    learned_model, _ = model_util.load_model(model_path, device)
    model = learned_model.model
    model.eval()
    model.args.device = device
    extractor = FeatureExtractor(
        archi=dataset_info['visual_archi'],
        device=device,
        checkpoint=dataset_info['visual_checkpoint'],
        compress_type=dataset_info['compress_type'])
    return model, extractor


def load_object_predictor(args):
    if args.object_predictor is None:
        return None
    return FeatureExtractor(
        archi='maskrcnn', device=args.device,
        checkpoint=args.object_predictor, load_heads=True)


def worker_loop(
        evaluate_function,
        dataset,
        trial_queue,
        log_queue,
        args,
        cuda_device=None,
        hierarchical=False):
    '''
    evaluation loop
    '''
    if cuda_device:
        torch.cuda.set_device(cuda_device)
        args.device = 'cuda:{}'.format(cuda_device)
    # start THOR
    env = ThorEnv(x_display=args.x_display)
    # master may ask to evaluate different models
    model_path_loaded = None
    object_predictor = load_object_predictor(args)

    if args.num_workers == 0:
        num_success, num_trials_done, num_trials = 0, 0, trial_queue.qsize()
    try:
        while True:
            print(colored(f'Remaining {trial_queue.qsize()}...', 'yellow'))
            trial_uid, dataset_idx, model_path = trial_queue.get(timeout=3)
            if model_path != model_path_loaded:
                if model_path_loaded is not None:
                    del model, extractor
                    torch.cuda.empty_cache()
                if hierarchical:
                    model_path_sg = '/home/ubuntu/bak/ET/logs/sg_et_non_navi/model_13.pth'
                    #model_path_sg = '/home/ubuntu/bak/ET/logs/try/model_19.pth'
                    model_path_navi = '/home/ubuntu/bak/ET/logs/sg_model_navi_v2/model_13.pth'
                    model_sg, extractor = load_agent(
                        model_path_sg,
                        dataset.dataset_info,
                        args.device)
                    model_navi, _ = load_agent(
                        model_path_navi,
                        dataset.dataset_info,
                        args.device)
                    dataset.vocab_translate = model_sg.vocab_out
                    model_path_loaded = model_path
                else:
                    model, extractor = load_agent(
                        model_path,
                        dataset.dataset_info,
                        args.device)
                    dataset.vocab_translate = model.vocab_out
                    model_path_loaded = model_path
            if hierarchical:
                log_entry = evaluate_function(
                    env, model_sg, dataset, extractor, trial_uid, dataset_idx, args,
                    object_predictor)
            else:
                log_entry = evaluate_function(
                    env, model, dataset, extractor, trial_uid, dataset_idx, args, 
                    object_predictor)
            if (args.debug or args.num_workers == 0) and 'success' in log_entry:
                if 'subgoal_action' in log_entry:
                    trial_type = log_entry['subgoal_action']
                else:
                    trial_type = 'full task'
                print(colored('Trial {}: {} ({})'.format(
                    trial_uid, 'success' if log_entry['success'] else 'fail',
                    trial_type), 'green' if log_entry['success'] else 'red'))
            if args.num_workers == 0 and 'success' in log_entry:
                num_trials_done += 1
                num_success += int(log_entry['success'])
                print('{:4d}/{} trials are done (current SR = {:.1f})'.format(
                    num_trials_done, num_trials, 100 * num_success / num_trials_done))
            log_queue.put((log_entry, trial_uid, model_path))
    except queue.Empty:
        pass
    # stop THOR
    env.stop()


def get_model_paths(args):
    '''
    check which models need to be evaluated
    '''
    model_paths = []
    if args.eval_range is None:
        # evaluate only the latest checkpoint
        model_paths.append(
            os.path.join(constants.ET_LOGS, args.exp, args.checkpoint))
    else:
        # evaluate a range of epochs
        for model_epoch in range(*args.eval_range):
            model_path = os.path.join(
                constants.ET_LOGS, args.exp,
                'model_{:02d}.pth'.format(model_epoch))
            if os.path.exists(model_path):
                model_paths.append(model_path)
    for idx, model_path in enumerate(model_paths):
        if os.path.islink(model_path):
            model_paths[idx] = os.readlink(model_path)
    if len(model_paths) == 0:
        raise ValueError('No models are found for evaluation')
    return model_paths


def get_sr(eval_json, eval_type):
    if eval_type == 'task':
        return eval_json['results']['all']['success']['success_rate']
    sr_sum, sr_count = 0, 0
    for _, sg_dict in sorted(eval_json['results'].items()):
        sr_sum += sg_dict['successes']
        sr_count += sg_dict['evals']
    return sr_sum / sr_count


def overwrite_eval_json(results, model_path, eval_type, args, dataset_id):
    '''
    append results to an existing eval.json or create a new one
    '''
    # eval_json: eval_epoch / eval_split / {'subgoal','task'} / {'normal','fast_epoch'}
    # see if eval.json file alredy existed
    eval_json = {}
    eval_json_path = os.path.join(os.path.dirname(model_path), 'eval.json')
    lock = filelock.FileLock(eval_json_path + '.lock')
    with lock:
        if os.path.exists(eval_json_path):
            with open(eval_json_path, 'r') as eval_json_file:
                eval_json = json.load(eval_json_file)
    eval_epoch = os.path.basename(model_path)
    if eval_epoch not in eval_json:
        eval_json[eval_epoch] = {}
    if dataset_id not in eval_json[eval_epoch]:
        eval_json[eval_epoch][dataset_id] = {}
    if eval_type not in eval_json[eval_epoch][dataset_id]:
        eval_json[eval_epoch][dataset_id][eval_type] = {}
    eval_mode = 'normal' if not args.fast_epoch else 'fast_epoch'
    if eval_mode in eval_json[eval_epoch][dataset_id][eval_type]:
        print('WARNING: the evaluation was already done')
        sr_new = len(results['successes']) / (len(results['successes']) + len(results['failures']))
        prev_res = eval_json[eval_epoch][dataset_id][eval_type][eval_mode]
        sr_old = len(prev_res['successes']) / (len(prev_res['successes']) + len(prev_res['failures']))
        print('Previous success rate = {:.1f}, new success rate = {:.1f}'.format(
            100 * sr_old, 100 * sr_new))
    eval_json[eval_epoch][dataset_id][eval_type][eval_mode] = results
    # make a backup copy of eval.json file before writing
    save_with_backup(eval_json, eval_json_path, lock)
    print('Evaluation is saved to {}'.format(eval_json_path))


def save_with_backup(obj, file_path, lock):
    with lock:
        # make a backup copy of results file before writing and swipe it after
        file_path_back = file_path + '.back'
        file_path_back_back = file_path + '.back.back'
        # always remove the second backup
        if os.path.exists(file_path_back_back):
            os.remove(file_path_back_back)
        # rename the first backup to the second one
        if os.path.exists(file_path_back):
            os.rename(file_path_back, file_path_back_back)
        # put the original file as the first backup
        if os.path.exists(file_path):
            os.rename(file_path, file_path_back)
        # write the updated content to the file path
        with open(file_path, 'w') as file_opened:
            json.dump(obj, file_opened, indent=2, sort_keys=True)


def disable_lock_logs():
    lock_logger = filelock.logger()
    lock_logger.setLevel(30)


def extract_rcnn_pred(class_idx, obj_predictor, env, verbose=False):
    '''
    extract a pixel mask using a pre-trained MaskRCNN
    '''
    rcnn_pred = obj_predictor.predict_objects(Image.fromarray(env.last_event.frame))
    class_name = obj_predictor.vocab_obj.index2word(class_idx)
    candidates = list(filter(lambda p: p.label == class_name, rcnn_pred))
    if verbose:
        visible_objs = [
            obj for obj in env.last_event.metadata['objects']
            if obj['visible'] and obj['objectId'].startswith(class_name + '|')]
        print('Agent prediction = {}, detected {} objects (visible {})'.format(
            class_name, len(candidates), len(visible_objs)))
    if len(candidates) > 0:
        if env.last_interaction[0] == class_idx and env.last_interaction[1] is not None:
            # last_obj['id'] and class_name + '|' in env.last_obj['id']:
            # do the association based selection
            last_center = np.array(env.last_interaction[1].nonzero()).mean(axis=1)
            cur_centers = np.array(
                [np.array(c.mask[0].nonzero()).mean(axis=1) for c in candidates])
            distances = ((cur_centers - last_center)**2).sum(axis=1)
            index = np.argmin(distances)
            mask = candidates[index].mask[0]
        else:
            # do the confidence based selection
            index = np.argmax([p.score for p in candidates])
            mask = candidates[index].mask[0]
    else:
        mask = None
    return mask


def agent_step(
        model, input_dict, vocab, prev_action, env, args, num_fails, obj_predictor):
    '''
    environment step based on model prediction
    '''
    # forward model
    with torch.no_grad():
        m_out = model.step(input_dict, vocab, prev_action=prev_action)
    m_pred = model_util.extract_action_preds(
        m_out, model.pad, vocab['action_low'], clean_special_tokens=False)[0]
    action = m_pred['action']
    if args.debug:
        print("Predicted action: {}".format(action))

    mask = None
    obj = m_pred['object'][0][0] if model_util.has_interaction(action) else None
    if obj is not None:
        # get mask from a pre-trained RCNN
        assert obj_predictor is not None
        class_name = obj_predictor.vocab_obj.index2word(obj)
        mask = extract_rcnn_pred(
            obj, obj_predictor, env, args.debug)
        m_pred['mask_rcnn'] = mask
    # remove blocking actions
    action = obstruction_detection(
        action, env, m_out, model.vocab_out, args.debug)
    m_pred['action'] = action

    # use the predicted action
    episode_end = (action == constants.STOP_TOKEN)
    api_action = None
    # constants.TERMINAL_TOKENS was originally used for subgoal evaluation
    target_instance_id = ''
    if not episode_end:
        step_success, _, target_instance_id, err, api_action = env.va_interact(
            action, interact_mask=mask, smooth_nav=args.smooth_nav, debug=args.debug)
        env.last_interaction = (obj, mask)
        if not step_success:
            print(err) #####
            num_fails += 1
            if num_fails >= args.max_fails:
                if args.debug:
                    print("Interact API failed {} times; latest error '{}'".format(
                        num_fails, err))
                episode_end = True
    return episode_end, str(action), num_fails, target_instance_id, api_action


def expert_step(action, masks, model, input_dict, vocab, prev_action, env, args, test_mode=False):
    '''
    environment step based on expert action
    '''
    mask = masks.pop(0).float().numpy()[0] if model_util.has_interaction(
        action) else None
    # forward model
    if not args.no_model_unroll and model is not None:
        with torch.no_grad():
            model.step(input_dict, vocab, prev_action=prev_action)
        prev_action = (action if not args.no_teacher_force else None)
    # execute expert action
    step_success, _, _, err, api_action = env.va_interact(
        action, interact_mask=mask, smooth_nav=args.smooth_nav, debug=args.debug)
    if not step_success:
        print("expert initialization failed")
        if test_mode:
            return True, prev_action, api_action
        else:
            return True, prev_action

    if test_mode:
        return False, prev_action, api_action
    else:
        # update transition reward
        _, _ = env.get_transition_reward()
        return False, prev_action


def get_observation(event, extractor):
    '''
    get environment observation
    '''
    frames = extractor.featurize([Image.fromarray(event.frame)], batch=1)
    return frames


def load_language(
        dataset, task, dataset_key, model_args, extractor, subgoal_idx=None, navi=False,
        test_split=False):
    '''
    load language features from the dataset and unit-test the feature extractor
    '''
    # load language features
    feat_numpy = dataset.load_features(task, subgoal_idx=subgoal_idx, navi=navi)

    # test extractor with the frames
    if not test_split:
        frames_expert = dataset.load_frames(dataset_key)
        model_util.test_extractor(task['root'], extractor, frames_expert)
    if not test_split and 'frames' in dataset.ann_type:
        # frames will be used as annotations
        feat_numpy['frames'] = frames_expert
    _, input_dict, _ = data_util.tensorize_and_pad(
        [(task, feat_numpy)], model_args.device, dataset.pad)
    return input_dict


def load_expert_actions(dataset, task, dataset_key, subgoal_idx):
    '''
    load actions and masks for expert initialization
    '''
    expert_dict = dict()
    expert_dict['actions'] = [
        a['discrete_action'] for a in task['plan']['low_actions']
        if a['high_idx'] < subgoal_idx]
    expert_dict['masks'] = dataset.load_masks(dataset_key)
    return expert_dict


def read_task_data(task, subgoal_idx=None):
    '''
    read data from the traj_json
    '''
    # read general task info
    repeat_idx = task['repeat_idx']
    task_dict = {'repeat_idx': repeat_idx,
                 'type': task['task_type'],
                 'task': '/'.join(task['root'].split('/')[-3:-1])}
    # read subgoal info
    if subgoal_idx is not None:
        task_dict['subgoal_idx'] = subgoal_idx
        task_dict['subgoal_action'] = task['plan']['high_pddl'][
            subgoal_idx]['discrete_action']['action']
    return task_dict


def obstruction_detection(action, env, m_out, vocab_out, verbose):
    '''
    change 'MoveAhead' action to a turn in case if it has failed previously
    '''
    if action != 'MoveAhead_25':
        return action
    if env.last_event.metadata['lastActionSuccess']:
        return action
    dist_action = m_out['action'][0][0].detach().cpu()
    idx_rotateR = vocab_out.word2index('RotateRight_90')
    idx_rotateL = vocab_out.word2index('RotateLeft_90')
    action = 'RotateLeft_90' if dist_action[idx_rotateL] > dist_action[idx_rotateR] else 'RotateRight_90'
    if verbose:
        print("Blocking action is changed to: {}".format(action))
    return action

<filename>ET/alfred/utils/eval_util.py
import os
import json
import queue
import torch
import shutil
import filelock
import numpy as np

from PIL import Image
from termcolor import colored

from alfred.gen import constants
from alfred.env.thor_env import ThorEnv
from alfred.nn.enc_visual import FeatureExtractor
from alfred.utils import data_util, model_util


def setup_scene(env, traj_data, reward_type='dense', test_split=False):
    '''
    intialize the scene and agent from the task info
    '''
    # scene setup
    scene_num = traj_data['scene']['scene_num']
    object_poses = traj_data['scene']['object_poses']
    dirty_and_empty = traj_data['scene']['dirty_and_empty']
    object_toggles = traj_data['scene']['object_toggles']
    scene_name = 'FloorPlan%d' % scene_num
    env.reset(scene_name, silent=True)
    env.restore_scene(object_poses, object_toggles, dirty_and_empty)
    # initialize to start position
    env.step(dict(traj_data['scene']['init_action']))
    # setup task for reward
    if not test_split:
        env.set_task(traj_data, reward_type=reward_type)


def load_agent(model_path, dataset_info, device):
    '''
    load a pretrained agent and its feature extractor
    '''
    learned_model, _ = model_util.load_model(model_path, device)
    model = learned_model.model
    model.eval()
    model.args.device = device
    extractor = FeatureExtractor(
        archi=dataset_info['visual_archi'],
        device=device,
        checkpoint=dataset_info['visual_checkpoint'],
        compress_type=dataset_info['compress_type'])
    return model, extractor


def load_object_predictor(args):
    if args.object_predictor is None:
        return None
    return FeatureExtractor(
        archi='maskrcnn', device=args.device,
        checkpoint=args.object_predictor, load_heads=True)


def worker_loop(
        evaluate_function,
        dataset,
        trial_queue,
        log_queue,
        args,
        cuda_device=None,
        hierarchical=False):
    '''
    evaluation loop
    '''
    if cuda_device:
        torch.cuda.set_device(cuda_device)
        args.device = 'cuda:{}'.format(cuda_device)
    # start THOR
    env = ThorEnv(x_display=args.x_display)
    # master may ask to evaluate different models
    model_path_loaded = None
    object_predictor = load_object_predictor(args)

    if args.num_workers == 0:
        num_success, num_trials_done, num_trials = 0, 0, trial_queue.qsize()
    try:
        while True:
            print(colored(f'Remaining {trial_queue.qsize()}...', 'yellow'))
            trial_uid, dataset_idx, model_path = trial_queue.get(timeout=3)
            if model_path != model_path_loaded:
                if model_path_loaded is not None:
                    del model, extractor
                    torch.cuda.empty_cache()
                if hierarchical:
                    model_path_sg = '/home/ubuntu/bak/ET/logs/sg_et_non_navi/model_13.pth'
                    #model_path_sg = '/home/ubuntu/bak/ET/logs/try/model_19.pth'
                    model_path_navi = '/home/ubuntu/bak/ET/logs/sg_model_navi_v2/model_13.pth'
                    model_sg, extractor = load_agent(
                        model_path_sg,
                        dataset.dataset_info,
                        args.device)
                    model_navi, _ = load_agent(
                        model_path_navi,
                        dataset.dataset_info,
                        args.device)
                    dataset.vocab_translate = model_sg.vocab_out
                    model_path_loaded = model_path
                else:
                    model, extractor = load_agent(
                        model_path,
                        dataset.dataset_info,
                        args.device)
                    dataset.vocab_translate = model.vocab_out
                    model_path_loaded = model_path
            if hierarchical:
                log_entry = evaluate_function(
                    env, model_sg, dataset, extractor, trial_uid, dataset_idx, args,
                    object_predictor)
            else:
                log_entry = evaluate_function(
                    env, model, dataset, extractor, trial_uid, dataset_idx, args, 
                    object_predictor)
            if (args.debug or args.num_workers == 0) and 'success' in log_entry:
                if 'subgoal_action' in log_entry:
                    trial_type = log_entry['subgoal_action']
                else:
                    trial_type = 'full task'
                print(colored('Trial {}: {} ({})'.format(
                    trial_uid, 'success' if log_entry['success'] else 'fail',
                    trial_type), 'green' if log_entry['success'] else 'red'))
            if args.num_workers == 0 and 'success' in log_entry:
                num_trials_done += 1
                num_success += int(log_entry['success'])
                print('{:4d}/{} trials are done (current SR = {:.1f})'.format(
                    num_trials_done, num_trials, 100 * num_success / num_trials_done))
            log_queue.put((log_entry, trial_uid, model_path))
    except queue.Empty:
        pass
    # stop THOR
    env.stop()


def get_model_paths(args):
    '''
    check which models need to be evaluated
    '''
    model_paths = []
    if args.eval_range is None:
        # evaluate only the latest checkpoint
        model_paths.append(
            os.path.join(constants.ET_LOGS, args.exp, args.checkpoint))
    else:
        # evaluate a range of epochs
        for model_epoch in range(*args.eval_range):
            model_path = os.path.join(
                constants.ET_LOGS, args.exp,
                'model_{:02d}.pth'.format(model_epoch))
            if os.path.exists(model_path):
                model_paths.append(model_path)
    for idx, model_path in enumerate(model_paths):
        if os.path.islink(model_path):
            model_paths[idx] = os.readlink(model_path)
    if len(model_paths) == 0:
        raise ValueError('No models are found for evaluation')
    return model_paths


def get_sr(eval_json, eval_type):
    if eval_type == 'task':
        return eval_json['results']['all']['success']['success_rate']
    sr_sum, sr_count = 0, 0
    for _, sg_dict in sorted(eval_json['results'].items()):
        sr_sum += sg_dict['successes']
        sr_count += sg_dict['evals']
    return sr_sum / sr_count


def overwrite_eval_json(results, model_path, eval_type, args, dataset_id):
    '''
    append results to an existing eval.json or create a new one
    '''
    # eval_json: eval_epoch / eval_split / {'subgoal','task'} / {'normal','fast_epoch'}
    # see if eval.json file alredy existed
    eval_json = {}
    eval_json_path = os.path.join(os.path.dirname(model_path), 'eval.json')
    lock = filelock.FileLock(eval_json_path + '.lock')
    with lock:
        if os.path.exists(eval_json_path):
            with open(eval_json_path, 'r') as eval_json_file:
                eval_json = json.load(eval_json_file)
    eval_epoch = os.path.basename(model_path)
    if eval_epoch not in eval_json:
        eval_json[eval_epoch] = {}
    if dataset_id not in eval_json[eval_epoch]:
        eval_json[eval_epoch][dataset_id] = {}
    if eval_type not in eval_json[eval_epoch][dataset_id]:
        eval_json[eval_epoch][dataset_id][eval_type] = {}
    eval_mode = 'normal' if not args.fast_epoch else 'fast_epoch'
    if eval_mode in eval_json[eval_epoch][dataset_id][eval_type]:
        print('WARNING: the evaluation was already done')
        sr_new = len(results['successes']) / (len(results['successes']) + len(results['failures']))
        prev_res = eval_json[eval_epoch][dataset_id][eval_type][eval_mode]
        sr_old = len(prev_res['successes']) / (len(prev_res['successes']) + len(prev_res['failures']))
        print('Previous success rate = {:.1f}, new success rate = {:.1f}'.format(
            100 * sr_old, 100 * sr_new))
    eval_json[eval_epoch][dataset_id][eval_type][eval_mode] = results
    # make a backup copy of eval.json file before writing
    save_with_backup(eval_json, eval_json_path, lock)
    print('Evaluation is saved to {}'.format(eval_json_path))


def save_with_backup(obj, file_path, lock):
    with lock:
        # make a backup copy of results file before writing and swipe it after
        file_path_back = file_path + '.back'
        file_path_back_back = file_path + '.back.back'
        # always remove the second backup
        if os.path.exists(file_path_back_back):
            os.remove(file_path_back_back)
        # rename the first backup to the second one
        if os.path.exists(file_path_back):
            os.rename(file_path_back, file_path_back_back)
        # put the original file as the first backup
        if os.path.exists(file_path):
            os.rename(file_path, file_path_back)
        # write the updated content to the file path
        with open(file_path, 'w') as file_opened:
            json.dump(obj, file_opened, indent=2, sort_keys=True)


def disable_lock_logs():
    lock_logger = filelock.logger()
    lock_logger.setLevel(30)


def extract_rcnn_pred(class_idx, obj_predictor, env, verbose=False):
    '''
    extract a pixel mask using a pre-trained MaskRCNN
    '''
    rcnn_pred = obj_predictor.predict_objects(Image.fromarray(env.last_event.frame))
    class_name = obj_predictor.vocab_obj.index2word(class_idx)
    candidates = list(filter(lambda p: p.label == class_name, rcnn_pred))
    if verbose:
        visible_objs = [
            obj for obj in env.last_event.metadata['objects']
            if obj['visible'] and obj['objectId'].startswith(class_name + '|')]
        print('Agent prediction = {}, detected {} objects (visible {})'.format(
            class_name, len(candidates), len(visible_objs)))
    if len(candidates) > 0:
        if env.last_interaction[0] == class_idx and env.last_interaction[1] is not None:
            # last_obj['id'] and class_name + '|' in env.last_obj['id']:
            # do the association based selection
            last_center = np.array(env.last_interaction[1].nonzero()).mean(axis=1)
            cur_centers = np.array(
                [np.array(c.mask[0].nonzero()).mean(axis=1) for c in candidates])
            distances = ((cur_centers - last_center)**2).sum(axis=1)
            index = np.argmin(distances)
            mask = candidates[index].mask[0]
        else:
            # do the confidence based selection
            index = np.argmax([p.score for p in candidates])
            mask = candidates[index].mask[0]
    else:
        mask = None
    return mask


def agent_step(
        model, input_dict, vocab, prev_action, env, args, num_fails, obj_predictor):
    '''
    environment step based on model prediction
    '''
    # forward model
    with torch.no_grad():
        m_out = model.step(input_dict, vocab, prev_action=prev_action)
    m_pred = model_util.extract_action_preds(
        m_out, model.pad, vocab['action_low'], clean_special_tokens=False)[0]
    action = m_pred['action']
    if args.debug:
        print("Predicted action: {}".format(action))

    mask = None
    obj = m_pred['object'][0][0] if model_util.has_interaction(action) else None
    if obj is not None:
        # get mask from a pre-trained RCNN
        assert obj_predictor is not None
        class_name = obj_predictor.vocab_obj.index2word(obj)
        mask = extract_rcnn_pred(
            obj, obj_predictor, env, args.debug)
        m_pred['mask_rcnn'] = mask
    # remove blocking actions
    action = obstruction_detection(
        action, env, m_out, model.vocab_out, args.debug)
    m_pred['action'] = action

    # use the predicted action
    episode_end = (action == constants.STOP_TOKEN)
    api_action = None
    # constants.TERMINAL_TOKENS was originally used for subgoal evaluation
    target_instance_id = ''
    if not episode_end:
        step_success, _, target_instance_id, err, api_action = env.va_interact(
            action, interact_mask=mask, smooth_nav=args.smooth_nav, debug=args.debug)
        env.last_interaction = (obj, mask)
        if not step_success:
            print(err) #####
            num_fails += 1
            if num_fails >= args.max_fails:
                if args.debug:
                    print("Interact API failed {} times; latest error '{}'".format(
                        num_fails, err))
                episode_end = True
    return episode_end, str(action), num_fails, target_instance_id, api_action


def expert_step(action, masks, model, input_dict, vocab, prev_action, env, args, test_mode=False):
    '''
    environment step based on expert action
    '''
    mask = masks.pop(0).float().numpy()[0] if model_util.has_interaction(
        action) else None
    # forward model
    if not args.no_model_unroll and model is not None:
        with torch.no_grad():
            model.step(input_dict, vocab, prev_action=prev_action)
        prev_action = (action if not args.no_teacher_force else None)
    # execute expert action
    step_success, _, _, err, api_action = env.va_interact(
        action, interact_mask=mask, smooth_nav=args.smooth_nav, debug=args.debug)
    if not step_success:
        print("expert initialization failed")
        if test_mode:
            return True, prev_action, api_action
        else:
            return True, prev_action

    if test_mode:
        return False, prev_action, api_action
    else:
        # update transition reward
        _, _ = env.get_transition_reward()
        return False, prev_action


def get_observation(event, extractor):
    '''
    get environment observation
    '''
    frames = extractor.featurize([Image.fromarray(event.frame)], batch=1)
    return frames


def load_language(
        dataset, task, dataset_key, model_args, extractor, subgoal_idx=None, navi=False,
        test_split=False):
    '''
    load language features from the dataset and unit-test the feature extractor
    '''
    # load language features
    feat_numpy = dataset.load_features(task, subgoal_idx=subgoal_idx, navi=navi)

    # test extractor with the frames
    if not test_split:
        frames_expert = dataset.load_frames(dataset_key)
        model_util.test_extractor(task['root'], extractor, frames_expert)
    if not test_split and 'frames' in dataset.ann_type:
        # frames will be used as annotations
        feat_numpy['frames'] = frames_expert
    _, input_dict, _ = data_util.tensorize_and_pad(
        [(task, feat_numpy)], model_args.device, dataset.pad)
    return input_dict


def load_expert_actions(dataset, task, dataset_key, subgoal_idx):
    '''
    load actions and masks for expert initialization
    '''
    expert_dict = dict()
    expert_dict['actions'] = [
        a['discrete_action'] for a in task['plan']['low_actions']
        if a['high_idx'] < subgoal_idx]
    expert_dict['masks'] = dataset.load_masks(dataset_key)
    return expert_dict


def read_task_data(task, subgoal_idx=None):
    '''
    read data from the traj_json
    '''
    # read general task info
    repeat_idx = task['repeat_idx']
    task_dict = {'repeat_idx': repeat_idx,
                 'type': task['task_type'],
                 'task': '/'.join(task['root'].split('/')[-3:-1])}
    # read subgoal info
    if subgoal_idx is not None:
        task_dict['subgoal_idx'] = subgoal_idx
        task_dict['subgoal_action'] = task['plan']['high_pddl'][
            subgoal_idx]['discrete_action']['action']
    return task_dict


def obstruction_detection(action, env, m_out, vocab_out, verbose):
    '''
    change 'MoveAhead' action to a turn in case if it has failed previously
    '''
    if action != 'MoveAhead_25':
        return action
    if env.last_event.metadata['lastActionSuccess']:
        return action
    dist_action = m_out['action'][0][0].detach().cpu()
    idx_rotateR = vocab_out.word2index('RotateRight_90')
    idx_rotateL = vocab_out.word2index('RotateLeft_90')
    action = 'RotateLeft_90' if dist_action[idx_rotateL] > dist_action[idx_rotateR] else 'RotateRight_90'
    if verbose:
        print("Blocking action is changed to: {}".format(action))
    return action

intialize the scene and agent from the task info # scene setup # initialize to start position # setup task for reward load a pretrained agent and its feature extractor evaluation loop # start THOR # master may ask to evaluate different models #model_path_sg = '/home/ubuntu/bak/ET/logs/try/model_19.pth' # stop THOR check which models need to be evaluated # evaluate only the latest checkpoint # evaluate a range of epochs append results to an existing eval.json or create a new one # eval_json: eval_epoch / eval_split / {'subgoal','task'} / {'normal','fast_epoch'} # see if eval.json file alredy existed # make a backup copy of eval.json file before writing # make a backup copy of results file before writing and swipe it after # always remove the second backup # rename the first backup to the second one # put the original file as the first backup # write the updated content to the file path extract a pixel mask using a pre-trained MaskRCNN # last_obj['id'] and class_name + '|' in env.last_obj['id']: # do the association based selection # do the confidence based selection environment step based on model prediction # forward model # get mask from a pre-trained RCNN # remove blocking actions # use the predicted action # constants.TERMINAL_TOKENS was originally used for subgoal evaluation ##### environment step based on expert action # forward model # execute expert action # update transition reward get environment observation load language features from the dataset and unit-test the feature extractor # load language features # test extractor with the frames # frames will be used as annotations load actions and masks for expert initialization read data from the traj_json # read general task info # read subgoal info change 'MoveAhead' action to a turn in case if it has failed previously