def create_images_from_dataset(variant): from op3.core import logger copy_to_save_file(logger.get_snapshot_dir()) train_path = get_module_path() + '/ec2_data/{}.h5'.format( variant['dataset']) num_samples = 100 train_dataset, _ = load_dataset(train_path, train=False, batchsize=1, size=num_samples, static=False) models_and_type = [] for a_model in variant['models']: m = load_model(a_model["saved_model_args"], train_dataset.action_dim, a_model["K"]) m_type = a_model['model_type'] models_and_type.append((m, m_type)) image_indices = list(range(num_samples)) for idx in image_indices: frames, actions = train_dataset[idx] frames = frames.unsqueeze(0) actions = actions.unsqueeze(0) create_image(models_and_type, frames, actions, logger.get_snapshot_dir() + "/image_{}".format(idx), variant['T'], variant['plot_type'])
def get_mse_from_dataset(variant): from op3.core import logger copy_to_save_file(logger.get_snapshot_dir()) train_path = get_module_path() + '/ec2_data/{}.h5'.format( variant['dataset']) num_samples = 100 train_dataset, _ = load_dataset(train_path, train=False, batchsize=1, size=num_samples, static=False) models_and_type = [] for a_model in variant['models']: m = load_model(a_model["saved_model_args"], train_dataset.action_dim, a_model["K"]) m_type = a_model['model_type'] models_and_type.append((m, m_type)) # image_indices = list(range(50)) batch_indices = np.arange(0, num_samples, 4) #bs=4 all_mse = [] for i in range(len(batch_indices) - 1): start_idx, end_idx = batch_indices[i], batch_indices[i + 1] frames, actions = train_dataset[start_idx:end_idx] #(bs, T, 3, D, D) # pdb.set_trace() # frames = frames.unsqueeze(0) # actions = actions.unsqueeze(0) mse = get_mse(models_and_type, frames, actions, variant['T']) #(M, bs, T), torch tensors all_mse.append(mse.permute(1, 0, 2)) #(bs, M, T) all_mse = torch.stack(all_mse, dim=0) #(I/bs, bs, M, T) all_mse = ptu.get_numpy( all_mse.view(-1, len(models_and_type), variant['T'])) #(I, M, T), numpy array now np.save(logger.get_snapshot_dir() + '/computed_mse.npy', all_mse) mean_vals = np.mean(all_mse, axis=0) #(M, T) std_vals = np.std(all_mse, axis=0) #(M, T) for i in range(len(models_and_type)): if models_and_type[i][1] == 'next_step' or 'rprp_pred': plt.errorbar(range(1, variant['T']), mean_vals[i][1:], std_vals[i][1:], label='{}'.format(models_and_type[i][1]), capsize=5) else: plt.errorbar(range(0, variant['T']), mean_vals[i], std_vals[i], label='{}'.format(models_and_type[i][1]), capsize=5) # plt.legend(bbox_to_anchor=(0.4, 0.8), loc="upper right") plt.legend(loc='center left', bbox_to_anchor=(1, 0.5)) # plt.yscale('log') plt.savefig(logger.get_snapshot_dir() + '/relative_mse.png', bbox_inches="tight")
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 test_epoch( self, epoch, save_reconstruction=True, train=True, batches=1, ): schedule = self.scheduler_class.get_schedule(epoch, is_train=True) loss_schedule = self.scheduler_class.get_loss_schedule(schedule) self.model.eval() losses, log_probs, kles, mses = [], [], [], [] dataloader = self.train_dataset.dataloader if train else self.test_dataset.dataloader for batch_idx, tensors in enumerate(dataloader): true_images, actions = self.prepare_tensors( tensors) #(B,T,3,D,D), (B,T,A) or None colors, masks, final_recon, total_loss, total_kle_loss, total_clog_prob, mse, cur_hidden_state = \ self.model.forward(true_images, actions, initial_hidden_state=None, schedule=schedule, loss_schedule=loss_schedule) # For DataParallel total_loss = total_loss.mean() total_clog_prob = total_clog_prob.mean() total_kle_loss = total_kle_loss.mean() mse = mse.mean() losses.append(total_loss.item()) log_probs.append(total_clog_prob.item()) kles.append(total_kle_loss.item()) mses.append(mse.mean().item()) # if batch_idx == 0 and save_reconstruction: # quicksave(true_images[0], colors[:,0], masks[:,0], schedule=schedule, # file_name=logger.get_snapshot_dir()+"/{}_{}.png".format('train' if train else 'val', epoch), quicksave_type="full") if batch_idx == 0 and save_reconstruction: quicksave( true_images[0], colors[0], masks[0], schedule=schedule, file_name=logger.get_snapshot_dir() + "/{}_{}.png".format('train' if train else 'val', epoch), quicksave_type="full") if batch_idx >= batches - 1: break stats = OrderedDict([("test/Log Prob", np.mean(log_probs)), ("test/KL", np.mean(kles)), ("test/loss", np.mean(losses)), ("test/mse", np.mean(mses))]) return stats
def create_mse_graphs(variant): from op3.core import logger copy_to_save_file(logger.get_snapshot_dir()) if variant['dataset'] == 'cloth': file_path = "/nfs/kun1/users/rishiv/Research/op3_exps/07-06-images-cloth/07-06-images-cloth_2019_07_06_06_15_12_0000--s-93491/computed_mse.npy" elif variant['dataset'] == 'poke': file_path = "/nfs/kun1/users/rishiv/Research/op3_exps/07-06-images-poke/07-06-images-poke_2019_07_06_05_29_35_0000--s-93758/computed_mse.npy" elif variant['dataset'] == 'solid': file_path = '/nfs/kun1/users/rishiv/Research/op3_exps/07-06-images-solid/07-06-images-solid_2019_07_06_05_41_24_0000--s-13353/computed_mse.npy' else: raise ValueError("Invalid dataset given") all_mse = np.load(file_path) # (I, M, T) # Models are rprp, rprp_pred, static, next_step in that order T = all_mse.shape[-1] # static_mse = all_mse[:, 2] #(I, T) next_step_mse = all_mse[:, 3] #(I, T) rprp = all_mse[:, 0] #(I, T) rprpr_next_step = all_mse[:, 1] #(I, T) plt.figure() plt.errorbar(range(T), np.mean(static_mse, axis=0), np.std(static_mse, axis=0), label='Static Iodine', capsize=5, color='blue', linestyle='--') plt.errorbar(range(T), np.mean(rprp, axis=0), np.std(rprp, axis=0), label='OP3', capsize=5, color='green', linestyle='--') # plt.errorbar(range(1, T), np.mean(next_step_mse, axis=0)[1:], np.std(static_mse, axis=0)[1:], label='Sequence', capsize=5, color='blue', linestyle='-') # plt.errorbar(range(1, T), np.mean(rprpr_next_step, axis=0)[1:], np.std(rprpr_next_step, axis=0)[1:], label='OP3_pred', capsize=5, color='green', linestyle='-') plt.legend(loc='center left', bbox_to_anchor=(1, 0.5)) plt.savefig(logger.get_snapshot_dir() + '/comparison.png', bbox_inches="tight")
def analyze_mse(variant): from op3.core import logger copy_to_save_file(logger.get_snapshot_dir()) if variant['dataset'] == 'cloth': file_path = "/nfs/kun1/users/rishiv/Research/op3_exps/07-03-images-cloth/07-03-images-cloth_2019_07_03_01_40_45_0000--s-9047/computed_mse.npy" elif variant['dataset'] == 'poke': file_path = "/nfs/kun1/users/rishiv/Research/op3_exps/07-04-images-poke/07-04-images-poke_2019_07_04_00_09_54_0000--s-46658/computed_mse.npy" elif variant['dataset'] == 'solid': file_path = '/nfs/kun1/users/rishiv/Research/op3_exps/07-06-images-solid/07-06-images-solid_2019_07_06_05_41_24_0000--s-13353/computed_mse.npy' else: raise ValueError("Invalid dataset given") all_mse = np.load(file_path) #(I, M, T) #Models are rprp, static, next_step in that order T = all_mse.shape[-1] rprp_vs_static = all_mse[:, 1] - all_mse[:, 0] #(I, T) rprp_vs_next_step = all_mse[:, 2] - all_mse[:, 0] #(I, T) print("RPRP: mean={}, std={}. \nStatic: mean={}. std={}".format( np.mean(all_mse[:, 0], axis=0), np.std(all_mse[:, 0], axis=0), np.mean(all_mse[:, 1], axis=0), np.mean(all_mse[:, 1], axis=0))) plt.figure() plt.errorbar(range(T), np.mean(rprp_vs_static, axis=0), np.std(rprp_vs_static, axis=0), label='rprp vs static', capsize=5) plt.savefig(logger.get_snapshot_dir() + '/rprp_vs_static.png') plt.figure() plt.errorbar(range(1, T), np.mean(rprp_vs_next_step, axis=0)[1:], np.std(rprp_vs_next_step, axis=0)[1:], label='rprp vs next_step', capsize=5) plt.savefig(logger.get_snapshot_dir() + '/rprp_vs_next_step.png')
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 save_model(self, prefix=""): torch.save( self.model.state_dict(), open( osp.join(logger.get_snapshot_dir(), '{}_params.pkl'.format(prefix)), "wb"))
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)