def create_images_from_dataset(variant):
from op3.core import logger
train_path = get_module_path() + '/data/datasets/{}.h5'.format(
variant['dataset'])
num_samples = 5 # Number of videos to visualize
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'],
image_type='normal')
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 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 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 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 copy_to_save_file(dir_str):
base = get_module_path()
shutil.copytree(base + '/rlkit/torch/iodine',
dir_str + '/saved_torch_iodine_files')
shutil.copy2(base + '/examples/mpc_v2/stage1/stage1_mpc.py',
dir_str + '/saved_stage1_mpc.py')
def copy_to_save_file(dir_str):
base = get_module_path()
shutil.copytree(base + '/rlkit/torch/iodine',
dir_str + '/saved_torch_iodine_files')
shutil.copy2(base + '/examples/mpc_v2/real_world/visualize_datasets_v2.py',
dir_str + '/visualize_datasets_v2.py')
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)
