def train_vae(variant): from op3.core import logger seed = int(variant['seed']) torch.manual_seed(seed) np.random.seed(seed) random.seed(seed) ######Dataset loading###### train_path = get_module_path() + '/data/datasets/{}.h5'.format( variant['dataset']) test_path = train_path bs = variant['training_args']['batch_size'] train_size = 100 if variant['debug'] == 1 else None static = (variant['schedule_args']['schedule_type'] == 'static_iodine' ) # Boolean train_dataset, max_T = load_dataset(train_path, train=True, batchsize=bs, size=train_size, static=static) test_dataset, _ = load_dataset(test_path, train=False, batchsize=bs, size=100, static=static) print(logger.get_snapshot_dir()) ######Model loading###### op3_args = variant["op3_args"] m = op3_model.create_model_v2(op3_args, op3_args['det_repsize'], op3_args['sto_repsize'], action_dim=train_dataset.action_dim) if variant['dataparallel']: m = torch.nn.DataParallel(m) m.cuda() ######Training###### scheduler = TrainingScheduler(**variant["schedule_args"], max_T=max_T) t = OP3Trainer(train_dataset, test_dataset, m, scheduler, **variant["training_args"]) save_period = variant['save_period'] for epoch in range(variant['num_epochs']): should_save_imgs = (epoch % save_period == 0) train_stats = t.train_epoch(epoch) test_stats = t.test_epoch(epoch, train=False, batches=1, save_reconstruction=should_save_imgs) t.test_epoch(epoch, train=True, batches=1, save_reconstruction=should_save_imgs) for k, v in {**train_stats, **test_stats}.items(): logger.record_tabular(k, v) logger.dump_tabular() t.save_model()
def load_model(variant, action_size, k): op3_args = variant["op3_args"] op3_args['K'] = k if "action_dim" in op3_args.keys(): action_size = op3_args["action_dim"] m = op3_model.create_model_v2(op3_args, op3_args['det_repsize'], op3_args['sto_repsize'], action_dim=action_size) model_file = get_module_path( ) + '/exps/realworld_exps/saved_models/{}'.format(variant['model_file']) state_dict = torch.load(model_file) new_state_dict = OrderedDict() for k, v in state_dict.items(): name = k if 'module.' in k: name = k[7:] # remove 'module.' of dataparallel new_state_dict[name] = v m.load_state_dict(new_state_dict) m.cuda() m.eval_mode = True return m
def load_model(variant, action_size, k): ######Start Model loading###### op3_args = variant["op3_args"] op3_args['K'] = k if "action_dim" in op3_args.keys(): action_size = op3_args["action_dim"] m = iodine_v2.create_model_v2(op3_args, op3_args['det_repsize'], op3_args['sto_repsize'], action_dim=action_size) model_file = variant["model_file"] state_dict = torch.load(model_file) new_state_dict = OrderedDict() for k, v in state_dict.items(): name = k if 'module.' in k: name = k[7:] # remove 'module.' of dataparallel new_state_dict[name] = v m.load_state_dict(new_state_dict) m.cuda() m.eval_mode = True return m
def main(variant): from op3.core import logger #copy_to_save_file(logger.get_snapshot_dir()) seed = int(variant['seed']) torch.manual_seed(seed) np.random.seed(seed) module_path = get_module_path() ######Start goal info loading###### n_goals = 1 if variant['debug'] == 1 else 10 goal_idxs = range(n_goals) actions_lst = [] stats = { 'mse': [], 'correct': [], 'max_pos': [], 'max_rgb': [], 'actions': [] } goal_counter = 0 structure, n_goal_obs = variant['structure'] ######End goal info loading###### ######Start Model loading###### op3_args = variant["op3_args"] # op3_args['K'] = n_goal_obs + 2 op3_args['K'] = 1 m = iodine_v2.create_model_v2(op3_args, op3_args['det_repsize'], op3_args['sto_repsize'], action_dim=0) # model_file = variant['model_file'] model_file = module_path + '/data/saved_models/{}.pkl'.format( variant['model_file']) state_dict = torch.load(model_file) new_state_dict = OrderedDict() for k, v in state_dict.items(): name = k if 'module.' in k: name = k[7:] # remove 'module.' of dataparallel new_state_dict[name] = v m.load_state_dict(new_state_dict) m.cuda() m.eval_mode = True ######End Model loading###### ######Start planning execution###### for i, goal_idx in enumerate(goal_idxs): goal_file = module_path + '/data/goals/stack_goals/manual_constructions/{}/{}_1.png'.format( structure, goal_idx) goal_image = imageio.imread(goal_file)[:, :, :3] # (D,D,3), numpy # goal_image = ptu.from_numpy(np.moveaxis(goal_image, 2, 0)).unsqueeze(0).float()[:, :3] / 255. # (1,3,D,D) true_data = np.load( module_path + '/data/goals/stack_goals/manual_constructions/{}/{}.p'.format( structure, goal_idx), allow_pickle=True) # pdb.set_trace() env = BlockEnv(n_goal_obs) logging_directory = "{}/{}_{}".format(logger.get_snapshot_dir(), structure, goal_idx) mpc = Stage1_MPC(m, logging_directory) cost_class = Cost(logging_directory, **variant['cost_args']) cem_process = Stage1_CEM(score_actions_class=cost_class, **variant['cem_args']) filter_goal_params = None if variant["filter_goal_image"]: filter_goal_params = {"n_objects": n_goal_obs} single_stats = mpc.run_plan(goal_image, env, cem_process, n_goal_obs, true_data, filter_goal_image=filter_goal_params) for k, v in single_stats.items(): stats[k].append(v) with open(logger.get_snapshot_dir() + '/results_rolling.pkl', 'wb') as f: pickle.dump(stats, f) with open(logger.get_snapshot_dir() + '/results_final.pkl', 'wb') as f: pickle.dump(stats, f) aggregate_stats = {} for k, v in stats.items(): if k != 'actions': aggregate_stats[k] = float(np.mean(v)) aggregate_stats["individual_correct"] = stats["correct"] json.dump(aggregate_stats, open(logger.get_snapshot_dir() + '/results_stats.json', 'w'))
def main(variant): from op3.core import logger #copy_to_save_file(logger.get_snapshot_dir()) seed = int(variant['seed']) torch.manual_seed(seed) np.random.seed(seed) module_path = get_module_path() ######Start Model loading###### op3_args = variant["op3_args"] op3_args['K'] = 4 m = op3_model.create_model_v2(op3_args, op3_args['det_repsize'], op3_args['sto_repsize'], action_dim=4) model_file = module_path + '/exps/pickplace_exps/saved_models/{}.pkl'.format( variant['model_file']) state_dict = torch.load(model_file) new_state_dict = OrderedDict() for k, v in state_dict.items(): name = k if 'module.' in k: name = k[7:] # remove 'module.' of dataparallel new_state_dict[name] = v m.load_state_dict(new_state_dict) m.cuda() m.eval_mode = True ######End Model loading###### ######Start goal info loading###### goal_idxs = range(variant["goal_start_end_range"][0], variant["goal_start_end_range"][1]) stats = { 'mse': [], 'correct': [], 'actions': [], 'first_finished_plan_steps': [] } goal_folder = module_path + '/data/goals/pickplace_goals/objects_seed_{}/'.format( variant['number_goal_objects']) num_seed_frames = 3 aggregate_stats = {} goals_tried = 0 ######End goal info loading###### ######Start planning execution###### for goal_idx in goal_idxs: env = BlockPickAndPlaceEnv( num_objects=1, num_colors=None, img_dim=64, include_z=False) # Note num_objects & num_colors do not matter ####Load goal and starting info with open(goal_folder + 'goal_data.pkl', 'rb') as f: goal_dict = pickle.load(f) goal_image = goal_dict["goal_image"][goal_idx] # (D,D,3) np # frames = goal_dict["frames"][i] # (T,D,D,3) np actions = goal_dict["actions"][goal_idx] # (T-1,6) np seed_actions = env._post_process_actions(actions) # (T-1,4) np goal_env_info = goal_dict["goal_env_info"][goal_idx] starting_env_info = goal_dict["starting_env_info"][goal_idx] #####Get seed steps env.set_env_info(starting_env_info) seed_frames = [env.get_observation()] if num_seed_frames > 1: seed_actions = seed_actions[:num_seed_frames - 1] for an_action in seed_actions: seed_frames.append(env.step(an_action)) else: seed_actions = None seed_frames = np.array(seed_frames) # (T,D,D,3) np if env.get_block_locs(check_all_in_bounds=True) is False: continue goals_tried += 1 #####Set up mpc logging_directory = "{}/goal_{}".format(logger.get_snapshot_dir(), goal_idx) cost_class = Cost(logging_directory, **variant['cost_args']) cem_process = Stage3_CEM(logging_dir=logging_directory, score_actions_class=cost_class, **variant['cem_args']) mpc = Stage3_MPC(m, logging_directory, variant['accuracy_threshold']) single_stats = mpc.run_plan( goal_image, env, seed_frames, seed_actions, cem_process, num_actions_to_take=variant["num_actions_to_take"], planning_horizon=variant["num_action_to_take_per_plan"], true_data=goal_env_info, filter_goal_image=False) # filter_goal_image={"n_objects": variant['number_goal_objects']}) for k, v in single_stats.items(): stats[k].append(v) with open(logger.get_snapshot_dir() + '/results_rolling.pkl', 'wb') as f: pickle.dump(stats, f) for k, v in stats.items(): if k != 'actions': aggregate_stats[k] = float(np.nanmean(v)) aggregate_stats["individual_correct"] = stats["correct"] aggregate_stats["num_goals_tried"] = goals_tried aggregate_stats["individual_first_finished_plan_steps"] = stats[ "first_finished_plan_steps"] json.dump(aggregate_stats, open(logger.get_snapshot_dir() + '/results_stats.json', 'w')) with open(logger.get_snapshot_dir() + '/results_final.pkl', 'wb') as f: pickle.dump(stats, f)