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)