class Thor(threading.Thread): def __init__(self, queue, train_eval="train"): Thread.__init__(self) self.action_queue = queue self.mask_rcnn = None self.env = None self.train_eval = train_eval self.controller_type = "oracle" def run(self): while True: action, reset, task_file = self.action_queue.get() try: if reset: self.reset(task_file) else: self.step(action) finally: self.action_queue.task_done() def init_env(self, config): self.config = config screen_height = config['env']['thor']['screen_height'] screen_width = config['env']['thor']['screen_width'] smooth_nav = config['env']['thor']['smooth_nav'] save_frames_to_disk = config['env']['thor']['save_frames_to_disk'] if not self.env: self.env = ThorEnv(player_screen_height=screen_height, player_screen_width=screen_width, smooth_nav=smooth_nav, save_frames_to_disk=save_frames_to_disk) self.controller_type = self.config['controller']['type'] self._done = False self._res = () self._feedback = "" self.expert = HandCodedThorAgent(self.env, max_steps=200) self.prev_command = "" self.load_mask_rcnn() def load_mask_rcnn(self): # load pretrained MaskRCNN model if required if 'mrcnn' in self.config['controller'][ 'type'] and not self.mask_rcnn: model_path = os.path.join( os.environ['ALFRED_ROOT'], self.config['mask_rcnn']['pretrained_model_path']) self.mask_rcnn = load_pretrained_model(model_path) def set_task(self, task_file): self.task_file = task_file self.traj_root = os.path.dirname(task_file) with open(task_file, 'r') as f: self.traj_data = json.load(f) def reset(self, task_file): assert self.env assert self.controller_type self.set_task(task_file) # scene setup scene_num = self.traj_data['scene']['scene_num'] object_poses = self.traj_data['scene']['object_poses'] dirty_and_empty = self.traj_data['scene']['dirty_and_empty'] object_toggles = self.traj_data['scene']['object_toggles'] scene_name = 'FloorPlan%d' % scene_num self.env.reset(scene_name) self.env.restore_scene(object_poses, object_toggles, dirty_and_empty) # recording save_frames_path = self.config['env']['thor']['save_frames_path'] self.env.save_frames_path = os.path.join( save_frames_path, self.traj_root.replace('../', '')) # initialize to start position self.env.step(dict( self.traj_data['scene']['init_action'])) # print goal instr task_desc = get_templated_task_desc(self.traj_data) print("Task: %s" % task_desc) # print("Task: %s" % (self.traj_data['turk_annotations']['anns'][0]['task_desc'])) # setup task for reward class args: pass args.reward_config = os.path.join(os.environ['ALFRED_ROOT'], 'agents/config/rewards.json') self.env.set_task(self.traj_data, args, reward_type='dense') # set controller self.controller_type = self.config['controller']['type'] self.goal_desc_human_anns_prob = self.config['env'][ 'goal_desc_human_anns_prob'] load_receps = self.config['controller']['load_receps'] debug = self.config['controller']['debug'] if self.controller_type == 'oracle': self.controller = OracleAgent( self.env, self.traj_data, self.traj_root, load_receps=load_receps, debug=debug, goal_desc_human_anns_prob=self.goal_desc_human_anns_prob) elif self.controller_type == 'oracle_astar': self.controller = OracleAStarAgent( self.env, self.traj_data, self.traj_root, load_receps=load_receps, debug=debug, goal_desc_human_anns_prob=self.goal_desc_human_anns_prob) elif self.controller_type == 'mrcnn': self.controller = MaskRCNNAgent( self.env, self.traj_data, self.traj_root, pretrained_model=self.mask_rcnn, load_receps=load_receps, debug=debug, goal_desc_human_anns_prob=self.goal_desc_human_anns_prob, save_detections_to_disk=self.env.save_frames_to_disk, save_detections_path=self.env.save_frames_path) elif self.controller_type == 'mrcnn_astar': self.controller = MaskRCNNAStarAgent( self.env, self.traj_data, self.traj_root, pretrained_model=self.mask_rcnn, load_receps=load_receps, debug=debug, goal_desc_human_anns_prob=self.goal_desc_human_anns_prob, save_detections_to_disk=self.env.save_frames_to_disk, save_detections_path=self.env.save_frames_path) else: raise NotImplementedError() # zero steps self.steps = 0 # reset expert state self.expert.reset(task_file) self.prev_command = "" # return intro text self._feedback = self.controller.feedback self._res = self.get_info() return self._feedback def step(self, action): if not self._done: # take action self.prev_command = str(action) self._feedback = self.controller.step(action) self._res = self.get_info() if self.env.save_frames_to_disk: self.record_action(action) self.steps += 1 def get_results(self): return self._res def record_action(self, action): txt_file = os.path.join(self.env.save_frames_path, 'action.txt') with open(txt_file, 'a+') as f: f.write("%s\r\n" % str(action)) def get_info(self): won = self.env.get_goal_satisfied() pcs = self.env.get_goal_conditions_met() goal_condition_success_rate = pcs[0] / float(pcs[1]) acs = self.controller.get_admissible_commands() # expert action if self.train_eval == "train": game_state = { 'admissible_commands': acs, 'feedback': self._feedback, 'won': won } expert_actions = ["look"] try: if not self.prev_command: self.expert.observe(game_state['feedback']) else: next_action = self.expert.act(game_state, 0, won, self.prev_command) if next_action in acs: expert_actions = [next_action] except HandCodedAgentTimeout: print("Expert Timeout") except Exception as e: print(e) traceback.print_exc() else: expert_actions = [] training_method = self.config["general"]["training_method"] if training_method == "dqn": max_nb_steps_per_episode = self.config["rl"]["training"][ "max_nb_steps_per_episode"] elif training_method == "dagger": max_nb_steps_per_episode = self.config["dagger"]["training"][ "max_nb_steps_per_episode"] else: raise NotImplementedError self._done = won or self.steps > max_nb_steps_per_episode return (self._feedback, self._done, acs, won, goal_condition_success_rate, expert_actions) def get_last_frame(self): return self.env.last_event.frame[:, :, ::-1] def get_exploration_frames(self): return self.controller.get_exploration_frames()
def run_rollouts(cls, model, task_queue, results, args, validation=False): env = ThorEnv() while True: if validation: task, seen = task_queue.get() else: task = task_queue.get() if task is None: break # reset model model.reset() # setup scene traj_data = model.load_task_json(task) r_idx = task['repeat_idx'] cls.setup_scene(env, traj_data, r_idx, args) feat = model.featurize([traj_data], load_frames=False, load_mask=False) curr_rollout = [] done = False fails = 0 total_reward = 0 num_steps = 0 while not done and num_steps < args.max_steps: # extract visual features curr_image = Image.fromarray(np.uint8(env.last_event.frame)) feat['frames'] = model.resnet.featurize([curr_image], batch=1).unsqueeze(0) # forward model out = model.step(feat) pred = model.sample_pred(out, greedy=validation) # monitor resource usage monitor = start_monitor( path=args.dout, note="validation" if validation else "rollout" + f" step={num_steps}") # # check if <<stop>> was predicted # if pred['action_low'] == "<<stop>>": # print("\tpredicted STOP") # break # get action and mask action = pred['action_low'] mask = pred['action_low_mask'] if cls.has_interaction( action) else None # use predicted action and mask (if available) to interact with the env t_success, _, _, err, _ = env.va_interact( action, interact_mask=mask, smooth_nav=args.smooth_nav, debug=args.debug) if not t_success: fails += 1 if fails >= args.max_fails: break # next time-step reward, done = env.get_transition_reward() total_reward += reward num_steps += 1 if not validation: curr_rollout.append({ 'frames': feat['frames'].cpu().detach().numpy(), 'lang_goal_instr_data': feat['lang_goal_instr'].data.cpu().detach().numpy(), 'lang_goal_instr_batch': feat['lang_goal_instr'].batch_sizes.cpu().detach( ).numpy(), 'lang_goal_instr_sorted': feat['lang_goal_instr'].sorted_indices.cpu().detach( ).numpy() if feat['lang_goal_instr'].sorted_indices is not None else None, 'lang_goal_instr_unsorted': feat['lang_goal_instr'].unsorted_indices.cpu().detach( ).numpy() if feat['lang_goal_instr'].unsorted_indices is not None else None, 'action_dist': pred['action_low_dist'].cpu().detach().numpy(), 'action_mask_dist': pred['action_low_mask_dist'].cpu().detach().numpy(), 'action_idx': pred['action_low_idx'].cpu().detach().numpy(), 'action_mask_idx': pred['action_low_mask_idx'].cpu().detach().numpy(), 'reward': np.array([reward]) }) stop_monitor(monitor) if validation: # check if goal was satisfied goal_satisfied = env.get_goal_satisfied() # goal_conditions pcs = env.get_goal_conditions_met() goal_condition_success_rate = pcs[0] / float(pcs[1]) # SPL path_len_weight = len(traj_data['plan']['low_actions']) s_spl = (1 if goal_satisfied else 0) * min( 1., path_len_weight / float(num_steps)) pc_spl = goal_condition_success_rate * min( 1., path_len_weight / float(num_steps)) # path length weighted SPL plw_s_spl = s_spl * path_len_weight plw_pc_spl = pc_spl * path_len_weight # log success/fails log_entry = { 'trial': traj_data['task_id'], 'type': traj_data['task_type'], 'repeat_idx': int(r_idx), 'seen': seen, 'goal_instr': traj_data['turk_annotations']['anns'][r_idx]['task_desc'], 'goal_satisfied': goal_satisfied, 'completed_goal_conditions': int(pcs[0]), 'total_goal_conditions': int(pcs[1]), 'goal_condition_success': float(goal_condition_success_rate), 'success_spl': float(s_spl), 'path_len_weighted_success_spl': float(plw_s_spl), 'goal_condition_spl': float(pc_spl), 'path_len_weighted_goal_condition_spl': float(plw_pc_spl), 'path_len_weight': int(path_len_weight), 'reward': float(total_reward) } results.put(log_entry) else: results.put(curr_rollout) env.stop()