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)
