gen/scripts/evaluate_trajectories.py

import os
import sys
sys.path.append(os.path.join(os.environ['ALFRED_ROOT']))
sys.path.append(os.path.join(os.environ['ALFRED_ROOT'], 'gen'))

import time
import multiprocessing as mp
import json
import random
import shutil
import argparse
import numpy as np
import pandas as pd
from collections import OrderedDict
from datetime import datetime

import constants
from agents.deterministic_planner_agent import DeterministicPlannerAgent
from env.thor_env import ThorEnv
from game_states.task_game_state_full_knowledge import TaskGameStateFullKnowledge
from utils.video_util import VideoSaver
from utils.dataset_management_util import load_successes_from_disk, load_fails_from_disk


# params
RAW_IMAGES_FOLDER = 'raw_images/'
DATA_JSON_FILENAME = 'traj_data.json'

# video saver
video_saver = VideoSaver()

# structures to help with constraint enforcement.
goal_to_required_variables = {"pick_and_place_simple": {"pickup", "receptacle", "scene"},
                              "pick_two_obj_and_place": {"pickup", "receptacle", "scene"},
                              "look_at_obj_in_light": {"pickup", "receptacle", "scene"},
                              "pick_clean_then_place_in_recep": {"pickup", "receptacle", "scene"},
                              "pick_heat_then_place_in_recep": {"pickup", "receptacle", "scene"},
                              "pick_cool_then_place_in_recep": {"pickup", "receptacle", "scene"},
                              "pick_and_place_with_movable_recep": {"pickup", "movable", "receptacle", "scene"}}
goal_to_pickup_type = {'pick_heat_then_place_in_recep': 'Heatable',
                       'pick_cool_then_place_in_recep': 'Coolable',
                       'pick_clean_then_place_in_recep': 'Cleanable'}
goal_to_receptacle_type = {'look_at_obj_in_light': "Toggleable"}
goal_to_invalid_receptacle = {'pick_heat_then_place_in_recep': {'Microwave'},
                              'pick_cool_then_place_in_recep': {'Fridge'},
                              'pick_clean_then_place_in_recep': {'SinkBasin'},
                              'pick_two_obj_and_place': {'CoffeeMachine', 'ToiletPaperHanger', 'HandTowelHolder'}}

scene_id_to_objs = {}
obj_to_scene_ids = {}
scenes_for_goal = {g: [] for g in constants.GOALS}
scene_to_type = {}


def sample_task_params(succ_traj, full_traj, fail_traj,
                       goal_candidates, pickup_candidates, movable_candidates, receptacle_candidates, scene_candidates,
                       inject_noise=10):
    # Get the current conditional distributions of all variables (goal/pickup/receptacle/scene).
    goal_weight = [(1 / (1 + np.random.randint(0, inject_noise + 1) + succ_traj.loc[
        (succ_traj['pickup'].isin(pickup_candidates) if 'pickup' in goal_to_required_variables[c] else True) &
        (succ_traj['movable'].isin(movable_candidates) if 'movable' in goal_to_required_variables[c] else True) &
        (succ_traj['receptacle'].isin(receptacle_candidates) if 'receptacle' in goal_to_required_variables[c] else True)
        & (succ_traj['scene'].isin(scene_candidates) if 'scene' in goal_to_required_variables[c] else True)]
            ['goal'].tolist().count(c)))  # Conditional.
                   * (1 / (1 + succ_traj['goal'].tolist().count(c)))  # Prior.
                   for c in goal_candidates]
    goal_probs = [w / sum(goal_weight) for w in goal_weight]

    pickup_weight = [(1 / (1 + np.random.randint(0, inject_noise + 1) +
                           sum([succ_traj.loc[
                                    succ_traj['goal'].isin([g]) &
                                    (succ_traj['movable'].isin(movable_candidates)
                                     if 'movable' in goal_to_required_variables[g] else True) &
                                    (succ_traj['receptacle'].isin(receptacle_candidates)
                                     if 'receptacle' in goal_to_required_variables[g] else True) &
                                    (succ_traj['scene'].isin(scene_candidates)
                                     if 'scene' in goal_to_required_variables[g] else True)]
                                ['pickup'].tolist().count(c) for g in goal_candidates])))
                     * (1 / (1 + succ_traj['pickup'].tolist().count(c)))
                     for c in pickup_candidates]
    pickup_probs = [w / sum(pickup_weight) for w in pickup_weight]

    movable_weight = [(1 / (1 + np.random.randint(0, inject_noise + 1) +
                            sum([succ_traj.loc[
                                     succ_traj['goal'].isin([g]) &
                                     (succ_traj['pickup'].isin(pickup_candidates)
                                      if 'pickup' in goal_to_required_variables[g] else True) &
                                     (succ_traj['receptacle'].isin(receptacle_candidates)
                                      if 'receptacle' in goal_to_required_variables[g] else True) &
                                     (succ_traj['scene'].isin(scene_candidates)
                                      if 'scene' in goal_to_required_variables[g] else True)]
                                 ['movable'].tolist().count(c) for g in goal_candidates])))
                      * (1 / (1 + succ_traj['movable'].tolist().count(c)))
                      for c in movable_candidates]
    movable_probs = [w / sum(movable_weight) for w in movable_weight]

    receptacle_weight = [(1 / (1 + np.random.randint(0, inject_noise + 1) +
                               sum([succ_traj.loc[
                                        succ_traj['goal'].isin([g]) &
                                        (succ_traj['pickup'].isin(pickup_candidates)
                                         if 'pickup' in goal_to_required_variables[g] else True) &
                                        (succ_traj['movable'].isin(movable_candidates)
                                         if 'movable' in goal_to_required_variables[g] else True) &
                                        (succ_traj['scene'].isin(scene_candidates)
                                         if 'scene' in goal_to_required_variables[g] else True)]
                                    ['receptacle'].tolist().count(c) for g in goal_candidates])))
                         * (1 / (1 + succ_traj['receptacle'].tolist().count(c)))
                         for c in receptacle_candidates]
    receptacle_probs = [w / sum(receptacle_weight) for w in receptacle_weight]
    scene_weight = [(1 / (1 + np.random.randint(0, inject_noise + 1) +
                          sum([succ_traj.loc[
                                   succ_traj['goal'].isin([g]) &
                                   (succ_traj['pickup'].isin(pickup_candidates)
                                    if 'pickup' in goal_to_required_variables[g] else True) &
                                   (succ_traj['movable'].isin(movable_candidates)
                                    if 'movable' in goal_to_required_variables[g] else True) &
                                   (succ_traj['receptacle'].isin(receptacle_candidates)
                                    if 'receptacle' in goal_to_required_variables[g] else True)]
                               ['scene'].tolist().count(c) for g in goal_candidates])))
                    * (1 / (1 + succ_traj['scene'].tolist().count(c)))
                    for c in scene_candidates]
    scene_probs = [w / sum(scene_weight) for w in scene_weight]

    # Calculate the probability difference between each value and the maximum so we can iterate over them to find a
    # next-best candidate to sample subject to the constraints of knowing which will fail.
    diffs = [("goal", goal_candidates[idx], goal_probs[idx] - min(goal_probs))
             for idx in range(len(goal_candidates)) if len(goal_candidates) > 1]
    diffs.extend([("pickup", pickup_candidates[idx], pickup_probs[idx] - min(pickup_probs))
                  for idx in range(len(pickup_candidates)) if len(pickup_candidates) > 1])
    diffs.extend([("movable", movable_candidates[idx], movable_probs[idx] - min(movable_probs))
                  for idx in range(len(movable_candidates)) if len(movable_candidates) > 1])
    diffs.extend([("receptacle", receptacle_candidates[idx], receptacle_probs[idx] - min(receptacle_probs))
                  for idx in range(len(receptacle_candidates)) if len(receptacle_candidates) > 1])
    diffs.extend([("scene", scene_candidates[idx], scene_probs[idx] - min(scene_probs))
                  for idx in range(len(scene_candidates)) if len(scene_candidates) > 1])

    # Iteratively pop the next biggest difference until we find a combination that is valid (e.g., not already
    # flagged as impossible by the simulator).
    variable_value_by_diff = {}
    diffs_as_keys = []  # list of diffs; index into list will be used as key values.
    for _, _, diff in diffs:
        already_keyed = False
        for existing_diff in diffs_as_keys:
            if np.isclose(existing_diff, diff):
                already_keyed = True
                break
        if not already_keyed:
            diffs_as_keys.append(diff)
    for variable, value, diff in diffs:
        key = None
        for kidx in range(len(diffs_as_keys)):
            if np.isclose(diffs_as_keys[kidx], diff):
                key = kidx
        if key not in variable_value_by_diff:
            variable_value_by_diff[key] = []
        variable_value_by_diff[key].append((variable, value))

    for key, diff in sorted(enumerate(diffs_as_keys), key=lambda x: x[1], reverse=True):
        variable_value = variable_value_by_diff[key]
        random.shuffle(variable_value)
        for variable, value in variable_value:

            # Select a goal.
            if variable == "goal":
                gtype = value
                # print("sampled goal '%s' with prob %.4f" % (gtype, goal_probs[goal_candidates.index(gtype)]))
                _goal_candidates = [gtype]

                _pickup_candidates = pickup_candidates[:]
                _movable_candidates = movable_candidates[:]
                _receptacle_candidates = receptacle_candidates[:]
                _scene_candidates = scene_candidates[:]

            # Select a pickup object.
            elif variable == "pickup":
                pickup_obj = value
                # print("sampled pickup object '%s' with prob %.4f" %
                #       (pickup_obj,  pickup_probs[pickup_candidates.index(pickup_obj)]))
                _pickup_candidates = [pickup_obj]

                _goal_candidates = goal_candidates[:]
                _movable_candidates = movable_candidates[:]
                _receptacle_candidates = receptacle_candidates[:]
                _scene_candidates = scene_candidates[:]

            # Select a movable object.
            elif variable == "movable":
                movable_obj = value
                # print("sampled movable object '%s' with prob %.4f" %
                #       (movable_obj,  movable_probs[movable_candidates.index(movable_obj)]))
                _movable_candidates = [movable_obj]
                _goal_candidates = [g for g in goal_candidates if g == 'pick_and_place_with_movable_recep']

                _pickup_candidates = pickup_candidates[:]
                _receptacle_candidates = receptacle_candidates[:]
                _scene_candidates = scene_candidates[:]

            # Select a receptacle.
            elif variable == "receptacle":
                receptacle_obj = value
                # print("sampled receptacle object '%s' with prob %.4f" %
                #       (receptacle_obj, receptacle_probs[receptacle_candidates.index(receptacle_obj)]))
                _receptacle_candidates = [receptacle_obj]

                _goal_candidates = goal_candidates[:]
                _pickup_candidates = pickup_candidates[:]
                _movable_candidates = movable_candidates[:]
                _scene_candidates = scene_candidates[:]

            # Select a scene.
            else:
                sampled_scene = value
                # print("sampled scene %s with prob %.4f" %
                #       (sampled_scene, scene_probs[scene_candidates.index(sampled_scene)]))
                _scene_candidates = [sampled_scene]

                _goal_candidates = goal_candidates[:]
                _pickup_candidates = pickup_candidates[:]
                _movable_candidates = movable_candidates[:]
                _receptacle_candidates = receptacle_candidates[:]
            # Perform constraint propagation to determine whether this is a valid assignment.
            propagation_finished = False
            while not propagation_finished:
                assignment_lens = (len(_goal_candidates), len(_pickup_candidates), len(_movable_candidates),
                                   len(_receptacle_candidates), len(_scene_candidates))
                # Constraints on goal.
                _goal_candidates = [g for g in _goal_candidates if
                                    (g not in goal_to_pickup_type or
                                     len(set(_pickup_candidates).intersection(  # Pickup constraint.
                                        constants.VAL_ACTION_OBJECTS[goal_to_pickup_type[g]])) > 0)
                                    and (g not in goal_to_receptacle_type or
                                         np.any([r in constants.VAL_ACTION_OBJECTS[goal_to_receptacle_type[g]]
                                                for r in _receptacle_candidates]))  # Valid by goal receptacle const.
                                    and (g not in goal_to_invalid_receptacle or
                                         len(set(_receptacle_candidates).difference(
                                            goal_to_invalid_receptacle[g])) > 0)  # Invalid by goal receptacle const.
                                    and len(set(_scene_candidates).intersection(
                                        scenes_for_goal[g])) > 0  # Scene constraint
                                    ]

                # Define whether to consider constraints for each role based on current set of candidate goals.
                pickup_constrained = np.any(["pickup" in goal_to_required_variables[g] for g in _goal_candidates])
                movable_constrained = np.any(["movable" in goal_to_required_variables[g] for g in _goal_candidates])
                receptacle_constrained = np.any(["receptacle" in goal_to_required_variables[g]
                                                 for g in _goal_candidates])
                scene_constrained = np.any(["scene" in goal_to_required_variables[g] for g in _goal_candidates])

                # Constraints on pickup obj.
                _pickup_candidates = [p for p in _pickup_candidates if
                                      np.any([g not in goal_to_pickup_type or
                                              p in constants.VAL_ACTION_OBJECTS[goal_to_pickup_type[g]]
                                              for g in _goal_candidates])  # Goal constraint.
                                      and (not movable_constrained or
                                           np.any([p in constants.VAL_RECEPTACLE_OBJECTS[m]
                                                  for m in _movable_candidates]))  # Movable constraint.
                                      and (not receptacle_constrained or
                                           np.any([r in constants.VAL_ACTION_OBJECTS["Toggleable"] or
                                                  p in constants.VAL_RECEPTACLE_OBJECTS[r]
                                                  for r in _receptacle_candidates]))  # Receptacle constraint.
                                      and (not scene_constrained or
                                           np.any([s in obj_to_scene_ids[constants.OBJ_PARENTS[p]]
                                                   for s in _scene_candidates])) # Scene constraint
                                      ]
                # Constraints on movable obj.
                _movable_candidates = [m for m in _movable_candidates if
                                       'pick_and_place_with_movable_recep' in _goal_candidates  # Goal constraint
                                       and (not pickup_constrained or
                                            np.any([p in constants.VAL_RECEPTACLE_OBJECTS[m]
                                                   for p in _pickup_candidates]))  # Pickup constraint.
                                       and (not receptacle_constrained or
                                            np.any([r in constants.VAL_RECEPTACLE_OBJECTS and
                                                   m in constants.VAL_RECEPTACLE_OBJECTS[r]
                                                   for r in _receptacle_candidates]))  # Receptacle constraint.
                                       and (not scene_constrained or
                                            np.any([s in obj_to_scene_ids[constants.OBJ_PARENTS[m]]
                                                    for s in _scene_candidates]))  # Scene constraint
                                       ]
                # Constraints on receptacle obj.
                _receptacle_candidates = [r for r in _receptacle_candidates if
                                          np.any([(g not in goal_to_receptacle_type or
                                                   r in constants.VAL_ACTION_OBJECTS[goal_to_receptacle_type[g]]) and
                                                  (g not in goal_to_invalid_receptacle or
                                                  r not in goal_to_invalid_receptacle[g])
                                                  for g in _goal_candidates])  # Goal constraint.
                                          and (not receptacle_constrained or
                                               r in constants.VAL_ACTION_OBJECTS["Toggleable"] or
                                               np.any([p in constants.VAL_RECEPTACLE_OBJECTS[r]
                                                       for p in _pickup_candidates]))  # Pickup constraint.
                                          and (not movable_constrained or
                                               r in constants.VAL_ACTION_OBJECTS["Toggleable"] or
                                               np.any([m in constants.VAL_RECEPTACLE_OBJECTS[r]
                                                       for m in _movable_candidates]))  # Movable constraint.
                                          and (not scene_constrained or
                                               np.any([s in obj_to_scene_ids[constants.OBJ_PARENTS[r]]
                                                       for s in _scene_candidates]))  # Scene constraint
                                          ]
                # Constraints on scene.
                _scene_candidates = [s for s in _scene_candidates if
                                     np.any([s in scenes_for_goal[g]
                                             for g in _goal_candidates])  # Goal constraint.
                                     and (not pickup_constrained or
                                          np.any([obj_to_scene_ids[constants.OBJ_PARENTS[p]]
                                                  for p in _pickup_candidates]))  # Pickup constraint.
                                     and (not movable_constrained or
                                          np.any([obj_to_scene_ids[constants.OBJ_PARENTS[m]]
                                                  for m in _movable_candidates]))  # Movable constraint.
                                     and (not receptacle_constrained or
                                          np.any([obj_to_scene_ids[constants.OBJ_PARENTS[r]]
                                                  for r in _receptacle_candidates]))  # Receptacle constraint.
                                     ]
                if assignment_lens == (len(_goal_candidates), len(_pickup_candidates), len(_movable_candidates),
                                       len(_receptacle_candidates), len(_scene_candidates)):
                    propagation_finished = True

            candidate_lens = {"goal": len(_goal_candidates), "pickup": len(_pickup_candidates),
                              "movable": len(_movable_candidates), "receptacle": len(_receptacle_candidates),
                              "scene": len(_scene_candidates)}
            if candidate_lens["goal"] == 0:
                # print("Goal over-constrained; skipping")
                continue
            if np.all([0 in [candidate_lens[v] for v in goal_to_required_variables[g]] for g in _goal_candidates]):
                continue

            # Ensure some combination of the remaining constraints is not in failures and is not already populated
            # by the target number of repeats.
            failure_ensured = True
            full_ensured = True
            for g in _goal_candidates:
                pickup_iter = _pickup_candidates if "pickup" in goal_to_required_variables[g] else ["None"]
                for p in pickup_iter:
                    movable_iter = _movable_candidates if "movable" in goal_to_required_variables[g] else ["None"]
                    for m in movable_iter:
                        receptacle_iter = _receptacle_candidates if "receptacle" in goal_to_required_variables[g] \
                            else ["None"]
                        for r in receptacle_iter:
                            scene_iter = _scene_candidates if "scene" in goal_to_required_variables[g] else ["None"]
                            for s in scene_iter:
                                if (g, p, m, r, s) not in fail_traj:
                                    failure_ensured = False
                                if (g, p, m, r, s) not in full_traj:
                                    full_ensured = False
                                if not failure_ensured and not full_ensured:
                                    break
                            if not failure_ensured and not full_ensured:
                                break
                        if not failure_ensured and not full_ensured:
                            break
                    if not failure_ensured and not full_ensured:
                        break
                if not failure_ensured and not full_ensured:
                    break
            if failure_ensured:
                continue
            if full_ensured:
                continue

            if candidate_lens["goal"] > 1 or np.any([np.any([candidate_lens[v] > 1
                                                             for v in goal_to_required_variables[g]])
                                                     for g in _goal_candidates]):
                task_sampler = sample_task_params(succ_traj, full_traj, fail_traj,
                                                  _goal_candidates, _pickup_candidates, _movable_candidates,
                                                  _receptacle_candidates, _scene_candidates)
                sampled_task = next(task_sampler)
                if sampled_task is None:
                    continue
            else:
                g = _goal_candidates[0]
                p = _pickup_candidates[0] if "pickup" in goal_to_required_variables[g] else "None"
                m = _movable_candidates[0] if "movable" in goal_to_required_variables[g] else "None"
                r = _receptacle_candidates[0] if "receptacle" in goal_to_required_variables[g] else "None"
                s = _scene_candidates[0] if "scene" in goal_to_required_variables[g] else "None"
                sampled_task = (g, p, m, r, int(s))

            yield sampled_task

    yield None  # Discovered that there are no valid assignments remaining.


def print_successes(succ_traj):
    print("###################################\n")
    print("Successes: ")
    print(succ_traj)
    print("\n##################################")


def main(args):
    # create env and agent
    env = ThorEnv()

    game_state = TaskGameStateFullKnowledge(env)
    agent = DeterministicPlannerAgent(thread_id=0, game_state=game_state)

    fail_traj = [
        # valid_seen
        'data/json_feat_2.1.0/pick_clean_then_place_in_recep-ButterKnife-None-Drawer-30/trial_T20190908_052007_212776',
        'data/json_feat_2.1.0/pick_clean_then_place_in_recep-Ladle-None-DiningTable-27/trial_T20190911_131350_027076',
        'data/json_feat_2.1.0/pick_two_obj_and_place-SprayBottle-None-Toilet-429/trial_T20190907_033843_707544',
        'data/json_feat_2.1.0/pick_heat_then_place_in_recep-AppleSliced-None-GarbageCan-11/trial_T20190906_211222_551053',
        'data/json_feat_2.1.0/pick_heat_then_place_in_recep-AppleSliced-None-Fridge-26/trial_T20190907_234650_780308',
        'data/json_feat_2.1.0/look_at_obj_in_light-CD-None-DeskLamp-320/trial_T20190907_224451_655673',
        'data/json_feat_2.1.0/pick_clean_then_place_in_recep-Mug-None-CoffeeMachine-26/trial_T20190907_233056_022802',
        'data/json_feat_2.1.0/pick_cool_then_place_in_recep-Bowl-None-Cabinet-20/trial_T20190909_112912_154874',
        'data/json_feat_2.1.0/pick_cool_then_place_in_recep-Pan-None-StoveBurner-18/trial_T20190906_181023_216915',
        'data/json_feat_2.1.0/look_at_obj_in_light-AlarmClock-None-DeskLamp-323/trial_T20190909_044715_250790',
        'data/json_feat_2.1.0/pick_two_obj_and_place-TissueBox-None-Toilet-426/trial_T20190909_044602_338962',
        'data/json_feat_2.1.0/pick_heat_then_place_in_recep-Plate-None-Cabinet-1/trial_T20190908_070946_578973',
        'data/json_feat_2.1.0/pick_heat_then_place_in_recep-TomatoSliced-None-SinkBasin-15/trial_T20190909_044151_889150',
        'data/json_feat_2.1.0/pick_heat_then_place_in_recep-AppleSliced-None-SinkBasin-30/trial_T20190909_001852_804129',
        'data/json_feat_2.1.0/pick_heat_then_place_in_recep-BreadSliced-None-DiningTable-26/trial_T20190908_023536_977733',
        'data/json_feat_2.1.0/pick_clean_then_place_in_recep-Plate-None-Fridge-5/trial_T20190909_112854_740612',
        'data/json_feat_2.1.0/pick_cool_then_place_in_recep-BreadSliced-None-Microwave-11/trial_T20190909_102940_350963',
        'data/json_feat_2.1.0/pick_cool_then_place_in_recep-PotatoSliced-None-SinkBasin-13/trial_T20190907_073447_193556',
        'data/json_feat_2.1.0/pick_and_place_with_movable_recep-AppleSliced-Pan-CounterTop-8/trial_T20190909_002950_338362',
        'data/json_feat_2.1.0/pick_two_obj_and_place-CD-None-Drawer-304/trial_T20190907_230232_832893',
        'data/json_feat_2.1.0/pick_two_obj_and_place-LettuceSliced-None-Fridge-1/trial_T20190906_181501_970690',
        'data/json_feat_2.1.0/pick_and_place_with_movable_recep-Fork-Cup-DiningTable-23/trial_T20190906_205421_980666',
        'data/json_feat_2.1.0/pick_and_place_simple-Mug-None-SideTable-329/trial_T20190909_032318_169393',
        'data/json_feat_2.1.0/pick_heat_then_place_in_recep-BreadSliced-None-Fridge-19/trial_T20190908_070424_700178',
        'data/json_feat_2.1.0/pick_cool_then_place_in_recep-Pan-None-DiningTable-7/trial_T20190908_232648_241836',
        'data/json_feat_2.1.0/pick_clean_then_place_in_recep-SoapBar-None-Toilet-428/trial_T20190907_041436_588106',
        'data/json_feat_2.1.0/pick_heat_then_place_in_recep-PotatoSliced-None-GarbageCan-15/trial_T20190908_042418_987263',
        'data/json_feat_2.1.0/look_at_obj_in_light-CD-None-DeskLamp-314/trial_T20190907_114323_767231',
        'data/json_feat_2.1.0/pick_cool_then_place_in_recep-Plate-None-Cabinet-7/trial_T20190907_231206_512424',
        'data/json_feat_2.1.0/pick_and_place_simple-SprayBottle-None-Toilet-418/trial_T20190906_190942_706186',
        'data/json_feat_2.1.0/pick_and_place_simple-Candle-None-Toilet-429/trial_T20190908_052232_887934',
        'data/json_feat_2.1.0/pick_cool_then_place_in_recep-Tomato-None-CounterTop-12/trial_T20190909_063754_216554',
        'data/json_feat_2.1.0/pick_heat_then_place_in_recep-TomatoSliced-None-GarbageCan-16/trial_T20190909_004537_873323',
        'data/json_feat_2.1.0/pick_and_place_simple-Book-None-SideTable-329/trial_T20190908_050633_745514',
        'data/json_feat_2.1.0/pick_and_place_with_movable_recep-TomatoSliced-Bowl-Fridge-20/trial_T20190907_060455_935544',
        'data/json_feat_2.1.0/pick_and_place_simple-Mug-None-SideTable-329/trial_T20190909_032340_274147',
        'data/json_feat_2.1.0/pick_and_place_simple-HandTowel-None-BathtubBasin-408/trial_T20190908_053502_505422',
        'data/json_feat_2.1.0/pick_clean_then_place_in_recep-DishSponge-None-BathtubBasin-427/trial_T20190906_234735_610018',
        'data/json_feat_2.1.0/pick_and_place_simple-Pencil-None-Dresser-330/trial_T20190909_071128_012892',
        'data/json_feat_2.1.0/pick_clean_then_place_in_recep-SoapBar-None-Toilet-410/trial_T20190906_201106_979461',
        'data/json_feat_2.1.0/pick_clean_then_place_in_recep-Cloth-None-Drawer-423/trial_T20190908_140701_251653',
        'data/json_feat_2.1.0/pick_and_place_with_movable_recep-PotatoSliced-Pot-DiningTable-21/trial_T20190907_113142_479158',
        'data/json_feat_2.1.0/pick_and_place_with_movable_recep-Ladle-Plate-Cabinet-16/trial_T20190907_070303_578496',
        'data/json_feat_2.1.0/pick_and_place_simple-ToiletPaper-None-ToiletPaperHanger-406/trial_T20190908_122858_883968',
        'data/json_feat_2.1.0/pick_heat_then_place_in_recep-Mug-None-CoffeeMachine-1/trial_T20190907_222640_487432',
        'data/json_feat_2.1.0/pick_and_place_with_movable_recep-Pen-Mug-Dresser-311/trial_T20190909_154212_097094',
        'data/json_feat_2.1.0/pick_cool_then_place_in_recep-Plate-None-Cabinet-27/trial_T20190906_173120_350651',
        'data/json_feat_2.1.0/look_at_obj_in_light-Statue-None-DeskLamp-216/trial_T20190909_193045_208933',
        'data/json_feat_2.1.0/pick_heat_then_place_in_recep-TomatoSliced-None-CounterTop-2/trial_T20190910_180358_375701',
        'data/json_feat_2.1.0/pick_and_place_with_movable_recep-TomatoSliced-Pan-Fridge-16/trial_T20190908_134410_550530',
        'data/json_feat_2.1.0/pick_and_place_simple-CellPhone-None-Bed-324/trial_T20190907_233817_198528',
        'data/json_feat_2.1.0/pick_and_place_with_movable_recep-Ladle-Plate-Cabinet-16/trial_T20190907_070152_814652',
        'data/json_feat_2.1.0/pick_clean_then_place_in_recep-ButterKnife-None-Drawer-2/trial_T20190908_121728_511866',
        'data/json_feat_2.1.0/pick_cool_then_place_in_recep-Pan-None-StoveBurner-23/trial_T20190906_215856_374127',
        'data/json_feat_2.1.0/pick_cool_then_place_in_recep-AppleSliced-None-GarbageCan-25/trial_T20190908_193525_796595',

        # valid_unseen
        'data/json_feat_2.1.0/pick_and_place_simple-SaltShaker-None-Cabinet-10/trial_T20190906_191501_563086',
        'data/json_feat_2.1.0/pick_cool_then_place_in_recep-TomatoSliced-None-CounterTop-10/trial_T20190909_061838_159982',
        'data/json_feat_2.1.0/pick_cool_then_place_in_recep-LettuceSliced-None-GarbageCan-10/trial_T20190909_213642_049125',
        'data/json_feat_2.1.0/pick_heat_then_place_in_recep-Potato-None-GarbageCan-10/trial_T20190907_161817_163233',
        'data/json_feat_2.1.0/pick_clean_then_place_in_recep-Pan-None-CounterTop-10/trial_T20190908_032434_013084',
        'data/json_feat_2.1.0/look_at_obj_in_light-Laptop-None-FloorLamp-219/trial_T20190906_181009_893506',
        'data/json_feat_2.1.0/pick_clean_then_place_in_recep-LettuceSliced-None-GarbageCan-10/trial_T20190918_143430_997524',
        'data/json_feat_2.1.0/pick_cool_then_place_in_recep-BreadSliced-None-Microwave-10/trial_T20190906_224718_745017',
        'data/json_feat_2.1.0/pick_and_place_simple-SaltShaker-None-Drawer-10/trial_T20190909_021613_077537',
        'data/json_feat_2.1.0/pick_clean_then_place_in_recep-SoapBar-None-CounterTop-424/trial_T20190907_074106_050405',
        'data/json_feat_2.1.0/pick_and_place_simple-SaltShaker-None-Drawer-10/trial_T20190909_021650_880235',
        'data/json_feat_2.1.0/pick_cool_then_place_in_recep-Bread-None-CounterTop-10/trial_T20190908_091747_866951',
        'data/json_feat_2.1.0/pick_heat_then_place_in_recep-Mug-None-Cabinet-10/trial_T20190909_021100_341887',
        'data/json_feat_2.1.0/pick_heat_then_place_in_recep-TomatoSliced-None-SinkBasin-10/trial_T20190909_025529_706272',
        'data/json_feat_2.1.0/look_at_obj_in_light-Box-None-FloorLamp-219/trial_T20190907_102952_001375',
        'data/json_feat_2.1.0/pick_heat_then_place_in_recep-Egg-None-GarbageCan-10/trial_T20190908_113523_123938',
        'data/json_feat_2.1.0/pick_and_place_with_movable_recep-Glassbottle-Plate-SinkBasin-10/trial_T20190908_123236_460615',
        'data/json_feat_2.1.0/pick_cool_then_place_in_recep-Pan-None-CounterTop-10/trial_T20190908_114656_768805',
        'data/json_feat_2.1.0/pick_clean_then_place_in_recep-LettuceSliced-None-GarbageCan-10/trial_T20190908_115403_911841',
        'data/json_feat_2.1.0/pick_cool_then_place_in_recep-Bread-None-CounterTop-10/trial_T20190908_091835_825830',
        'data/json_feat_2.1.0/pick_and_place_with_movable_recep-Spoon-Pan-SinkBasin-10/trial_T20190908_092900_630408',
        'data/json_feat_2.1.0/pick_and_place_with_movable_recep-Egg-Bowl-CounterTop-10/trial_T20190912_190125_832581',
        'data/json_feat_2.1.0/pick_clean_then_place_in_recep-Plate-None-CounterTop-10/trial_T20190908_213420_728917',
        'data/json_feat_2.1.0/pick_two_obj_and_place-CD-None-Safe-308/trial_T20190907_051013_060265',
        'data/json_feat_2.1.0/pick_cool_then_place_in_recep-LettuceSliced-None-GarbageCan-10/trial_T20190909_213855_654145',
        'data/json_feat_2.1.0/pick_and_place_with_movable_recep-Pencil-Bowl-Shelf-308/trial_T20190909_232711_494127',
        'data/json_feat_2.1.0/look_at_obj_in_light-Pencil-None-DeskLamp-308/trial_T20190908_220545_153480',
        'data/json_feat_2.1.0/pick_two_obj_and_place-CD-None-Safe-308/trial_T20190907_050942_897916',
        'data/json_feat_2.1.0/pick_clean_then_place_in_recep-Knife-None-CounterTop-10/trial_T20190909_110445_675754',
        'data/json_feat_2.1.0/pick_and_place_with_movable_recep-Fork-Cup-SinkBasin-10/trial_T20190907_144323_122735',
        'data/json_feat_2.1.0/pick_and_place_with_movable_recep-Fork-Cup-CounterTop-10/trial_T20190918_231203_583927',
        'data/json_feat_2.1.0/pick_cool_then_place_in_recep-Bread-None-CounterTop-10/trial_T20190908_091811_414150,'
        'data/json_feat_2.1.0/pick_cool_then_place_in_recep-Pan-None-CounterTop-10/trial_T20190908_114545_244903',
        'data/json_feat_2.1.0/pick_and_place_with_movable_recep-KeyChain-Box-ArmChair-219/trial_T20190909_042101_443154',
        'data/json_feat_2.1.0/pick_heat_then_place_in_recep-Apple-None-GarbageCan-10/trial_T20190908_145050_918567',
        'data/json_feat_2.1.0/pick_heat_then_place_in_recep-Potato-None-GarbageCan-10/trial_T20190907_161745_664033',
        'data/json_feat_2.1.0/pick_heat_then_place_in_recep-BreadSliced-None-CounterTop-10/trial_T20190908_025303_695053',
        'data/json_feat_2.1.0/look_at_obj_in_light-CD-None-DeskLamp-308/trial_T20190908_141958_463362',
        'data/json_feat_2.1.0/pick_and_place_simple-Pencil-None-Shelf-308/trial_T20190908_122154_042763',
        'data/json_feat_2.1.0/pick_cool_then_place_in_recep-Mug-None-Cabinet-10/trial_T20190909_121559_082363',
        'data/json_feat_2.1.0/pick_heat_then_place_in_recep-BreadSliced-None-CounterTop-10/trial_T20190908_025420_586121',
        'data/json_feat_2.1.0/pick_heat_then_place_in_recep-Potato-None-GarbageCan-10/trial_T20190907_161853_945788',
        'data/json_feat_2.1.0/pick_clean_then_place_in_recep-Plate-None-CounterTop-10/trial_T20190908_213356_017769',
        'data/json_feat_2.1.0/pick_and_place_with_movable_recep-KeyChain-Box-ArmChair-219/trial_T20190918_185132_685516',
        'data/json_feat_2.1.0/pick_heat_then_place_in_recep-Mug-None-CoffeeMachine-10/trial_T20190907_171850_960211',
        'data/json_feat_2.1.0/pick_heat_then_place_in_recep-Mug-None-Cabinet-10/trial_T20190909_021200_669381',
        'data/json_feat_2.1.0/pick_heat_then_place_in_recep-Mug-None-Cabinet-10/trial_T20190909_021247_306737',
        'data/json_feat_2.1.0/pick_and_place_with_movable_recep-ButterKnife-Cup-SinkBasin-10/trial_T20190909_025308_903206',
        'data/json_feat_2.1.0/pick_heat_then_place_in_recep-Tomato-None-GarbageCan-10/trial_T20190908_225046_020282',
        'data/json_feat_2.1.0/pick_heat_then_place_in_recep-AppleSliced-None-CounterTop-10/trial_T20190910_122059_929600',
    ] # trajectories failed with GT information

    from tqdm import tqdm
    import pickle
    import string

    root = '../data/json_2.1.0/'

    # GT: GT information
    # predTaskArgs: GT information w/ pred. task arguments (FILM)
    mode = 'GT'
    #mode = 'predTaskArgs'

    # save_root = '../data/json_2.1.0_generated'
    save_root = '../data/json_2.1.0_' + mode
    if mode in ['predTaskArgs']:
        film_args_dicts = {
            'valid_seen': pickle.load(open('../../../hdd/FILM/models/instructions_processed_LP/instruction2_params_val_seen_916_noappended.p', 'rb')),
            'valid_unseen': pickle.load(open('../../../hdd/FILM/models/instructions_processed_LP/instruction2_params_val_unseen_916_noappended.p', 'rb')),
        }

    for split in ['valid_seen', 'valid_unseen']:
        if mode in ['predTaskArgs']:
            film_args_dict = film_args_dicts[split]
            task_type_dict = {2: 'pick_and_place_simple',
                              5: 'look_at_obj_in_light',
                              1: 'pick_and_place_with_movable_recep',
                              3: 'pick_two_obj_and_place',
                              6: 'pick_clean_then_place_in_recep',
                              4: 'pick_heat_then_place_in_recep',
                              0: 'pick_cool_then_place_in_recep'}

        for ep in tqdm(os.listdir(os.path.join(root, split)), desc=split):
            for trial in os.listdir(os.path.join(root, split, ep)):
                # if not any(trial in ft for ft in fail_traj):
                #     continue

                traj_root = os.path.join(root, split, ep, trial)
                traj_data = json.load(open(traj_root+'/traj_data.json', 'r'))
                
                init_action_new = {
                    'action': 'TeleportFull',
                    'position': {
                        'x': traj_data['scene']['init_action']['x'],
                        'y': traj_data['scene']['init_action']['y'],
                        'z': traj_data['scene']['init_action']['z'],
                    },
                    'rotation': dict(x=0, y=traj_data['scene']['init_action']['rotation'], z=0),
                    'horizon': traj_data['scene']['init_action']['horizon'],
                    'standing': True,
                }
                traj_data['scene']['init_action'] = init_action_new

                print(traj_data['task_type'])
                print(traj_data['pddl_params'])

                if mode in ['GT']: # GT information
                    pass
                elif mode in ['predTaskArgs']: # GT information w/ pred. task arguments (FILM)
                    goal_statement = traj_data['turk_annotations']['anns'][0]['task_desc'].lower()
                    goal_statement = ''.join(c for c in goal_statement if c not in set(string.punctuation))

                    # move a bottle of hand soap from the sink to the toilet
                    # {'task_type': 2,
                    #  'mrecep_target': None,
                    #  'object_target': 'SoapBottle',
                    #  'parent_target': 'Toilet',
                    #  'sliced': 0}
                    task_type = task_type_dict[film_args_dict[goal_statement]['task_type']]
                    traj_data['pddl_params']['task_type'] = task_type

                    mrecep_target = film_args_dict[goal_statement]['mrecep_target']
                    mrecep_target = '' if mrecep_target is None else mrecep_target
                    traj_data['pddl_params']['mrecep_target'] = mrecep_target

                    object_sliced = film_args_dict[goal_statement]['sliced']
                    traj_data['pddl_params']['object_sliced'] = False if object_sliced == 0 else True

                    parent_target = film_args_dict[goal_statement]['parent_target']
                    if parent_target is None:
                        assert task_type == 'look_at_obj_in_light', print('Invalid film task arguments:', film_args_dict[goal_statement])
                        traj_data['pddl_params']['parent_target'] = ''
                        traj_data['pddl_params']['toggle_target'] = 'DeskLamp' # TODO: should be FloorLamp as well
                    else:
                        traj_data['pddl_params']['parent_target'] = parent_target
                        traj_data['pddl_params']['toggle_target'] = ''

                try:
                    setup_data_dict()
                    #print([o['objectId'] for o in env.last_event.metadata['objects']])

                    #######################################################################################################
                    scene = traj_data['scene']
                    objs = traj_data['scene']['object_poses']
                    seed = traj_data['scene']['random_seed']
                    agent.reset(seed=seed, scene=scene, objs=objs, traj_data=traj_data)
                    #######################################################################################################
                    #print([o['objectId'] for o in env.last_event.metadata['objects']])

                    #######################################################################################################
                    seed = traj_data['scene']['random_seed']
                    scene = traj_data['scene']
                    info = None
                    agent.setup_problem(dict(seed=seed, info=info, traj_data=traj_data), scene=scene)
                    #######################################################################################################
                    #print([o['objectId'] for o in env.last_event.metadata['objects']])

                    #######################################################################################################
                    print("Performing reset via thor_env API")
                    sampled_scene = traj_data['scene']['floor_plan']
                    env.reset(sampled_scene)

                    print("Performing restore via thor_env API")
                    object_poses = traj_data['scene']['object_poses']
                    object_toggles = traj_data['scene']['object_toggles']
                    dirty_and_empty = traj_data['scene']['dirty_and_empty']
                    env.restore_scene(object_poses, object_toggles, dirty_and_empty)
                    env.step(traj_data['scene']['init_action'])
                    #######################################################################################################

                    terminal = False
                    while not terminal and agent.current_frame_count <= constants.MAX_EPISODE_LENGTH:
                        action_dict = agent.get_action(None)
                        agent.step(action_dict)
                        reward, terminal = agent.get_reward()

                    # dump_data_dict(os.path.join(save_root, split, ep, trial))
                    print('Success:', os.path.join(split, ep, trial))
                except Exception as e:
                    print('Fail:', os.path.join(ep, trial), e)
                # save_video()


def create_dirs(gtype, pickup_obj, movable_obj, receptacle_obj, scene_num):
    task_id = 'trial_T' + datetime.now().strftime("%Y%m%d_%H%M%S_%f")
    save_name = '%s-%s-%s-%s-%d' % (gtype, pickup_obj, movable_obj, receptacle_obj, scene_num) + '/' + task_id

    constants.save_path = os.path.join(constants.DATA_SAVE_PATH, save_name, RAW_IMAGES_FOLDER)
    if not os.path.exists(constants.save_path):
        os.makedirs(constants.save_path)

    print("Saving images to: " + constants.save_path)
    return task_id


def save_video():
    images_path = constants.save_path + '*.png'
    video_path = os.path.join(constants.save_path.replace(RAW_IMAGES_FOLDER, ''), 'video.mp4')
    video_saver.save(images_path, video_path)


def setup_data_dict():
    constants.data_dict = OrderedDict()
    constants.data_dict['task_id'] = ""
    constants.data_dict['task_type'] = ""
    constants.data_dict['scene'] = {'floor_plan': "", 'random_seed': -1, 'scene_num': -1, 'init_action': [],
                                    'object_poses': [], 'dirty_and_empty': None, 'object_toggles': []}
    constants.data_dict['plan'] = {'high_pddl': [], 'low_actions': []}
    constants.data_dict['images'] = []
    constants.data_dict['template'] = {'task_desc': "", 'high_descs': []}
    constants.data_dict['pddl_params'] = {'object_target': -1, 'object_sliced': -1,
                                          'parent_target': -1, 'toggle_target': -1,
                                          'mrecep_target': -1}
    constants.data_dict['dataset_params'] = {'video_frame_rate': -1}
    constants.data_dict['pddl_state'] = []


def dump_data_dict(path):
    os.makedirs(path, exist_ok=True)
    with open(os.path.join(path, DATA_JSON_FILENAME), 'w') as fp:
        json.dump(constants.data_dict, fp, sort_keys=True, indent=4)


def delete_save(in_parallel):
    save_folder = constants.save_path.replace(RAW_IMAGES_FOLDER, '')
    if os.path.exists(save_folder):
        try:
            shutil.rmtree(save_folder)
        except OSError as e:
            if in_parallel:  # another thread succeeded at this task while this one failed.
                return False
            else:
                raise e  # if we're not running in parallel, this is an actual.
    return True


def parallel_main(args):
    procs = [mp.Process(target=main, args=(args,)) for _ in range(args.num_threads)]
    try:
        for proc in procs:
            proc.start()
            time.sleep(0.1)
    finally:
        for proc in procs:
            proc.join()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    # settings
    parser.add_argument('--force_unsave', action='store_true', help="don't save any data (for debugging purposes)")
    parser.add_argument('--debug', action='store_true')
    parser.add_argument('--save_path', type=str, default="dataset/new_trajectories", help="where to save the generated data")
    parser.add_argument('--x_display', type=str, required=False, default=constants.X_DISPLAY, help="x_display id")
    parser.add_argument("--just_examine", action='store_true', help="just examine what data is gathered; don't gather more")
    parser.add_argument("--in_parallel", action='store_true', help="this collection will run in parallel with others, so load from disk on every new sample")
    parser.add_argument("-n", "--num_threads", type=int, default=1, help="number of processes for parallel mode")
    parser.add_argument('--json_file', type=str, default="", help="path to json file with trajectory dump")

    # params
    parser.add_argument("--repeats_per_cond", type=int, default=3)
    parser.add_argument("--trials_before_fail", type=int, default=5)
    parser.add_argument("--async_load_every_n_samples", type=int, default=10)

    parse_args = parser.parse_args()

    if parse_args.in_parallel and parse_args.num_threads > 1:
        parallel_main(parse_args)
    else:
        main(parse_args)