def test_trajectory_plan():
d = load_model_and_data()
controller = ZMQController(
config=d['config'],
model_dy=d['model_dy'],
action_function=d['action_function'],
observation_function=d['observation_function'],
visual_observation_function=d['visual_observation_function'],
planner=d['planner'],
debug=True,
zmq_enabled=True,
camera_info=d['camera_info'],
)
plan_msg_file = "/home/manuelli/data/key_dynam/hardware_experiments/demonstrations/stable/2020-07-09-20-27-09_push_right_blue_tapes/plan_msg.p"
plan_msg = load_pickle(plan_msg_file, encoding='latin1')
plan_msg = zmq_utils.convert(plan_msg)
controller._on_plan_msg(plan_msg)
traj_goal = controller._state_dict['plan'].get_trajectory_goal(15, 5)
print("traj_goal.shape", traj_goal.shape)
plan_msg_file = "/home/manuelli/data/key_dynam/hardware_experiments/demonstrations/stable/2020-07-09-20-27-09_push_right_blue_tapes/plan_msg.p"
plan_msg = load_pickle(plan_msg_file, encoding='latin1')
plan_msg = zmq_utils.convert(plan_msg)
for i in range(4):
compute_control_action_msg = {
'type': "COMPUTE_CONTROL_ACTION",
'data': plan_msg['data']['plan_data'][i]
}
controller._on_compute_control_action(compute_control_action_msg,
visualize=True)
input("press Enter to continue")
def load_model_state_dict(model_folder=None):
models_root = os.path.join(get_data_root(),
"dev/experiments/drake_pusher_slider_v2")
# model_name = "dataset_2020-04-23-20-45-12-697915_T_aug_random_velocity_1000_angled_cam/trained_models/dynamics/DD_3D/2020-05-12-12-03-05-252242_DD_3D_spatial_z_n_his_2"
model_name = "dataset_2020-04-23-20-45-12-697915_T_aug_random_velocity_1000_angled_cam/trained_models/dynamics/DD_3D/2020-05-11-19-44-35-085478_DD_3D_n_his_2"
model_folder = os.path.join(models_root, model_name)
model_dy = model_builder.load_dynamics_model_from_folder(
model_folder)['model_dy']
# load dense descriptor model
metadata = load_pickle(os.path.join(model_folder, 'metadata.p'))
model_dd_file = metadata['model_file']
model_dd = torch.load(model_dd_file)
model_dd = model_dd.eval()
model_dd = model_dd.cuda()
spatial_descriptor_data = load_pickle(
os.path.join(model_folder, 'spatial_descriptors.p'))
return {
'model_dy': model_dy,
'model_dd': model_dd,
'spatial_descriptor_data': spatial_descriptor_data,
'metadata': metadata
}
def load_model_state_dict(model_folder=None):
# load dynamics model
#
# model_folder = "/home/manuelli/data/key_dynam/dev/experiments/drake_pusher_slider_box_on_side/dataset_dps_box_on_side_600/trained_models/dynamics/DD_3D/2020-05-15-04-40-10-770703_DD_3D_all_z_n_his_2"
models_root = os.path.join(
get_data_root(),
"dev/experiments/drake_pusher_slider_box_on_side/dataset_dps_box_on_side_600/trained_models/dynamics"
)
# model_name = "DD_3D/2020-05-15-00-54-26-961701_DD_3D_n_his_2"
# model_name = "DD_3D/2020-05-15-04-40-10-770703_DD_3D_all_z_n_his_2"
model_name = "DD_3D/2020-05-15-02-07-44-204479_DD_3D_spatial_z_n_his_2"
model_folder = os.path.join(models_root, model_name)
model_dy = model_builder.load_dynamics_model_from_folder(
model_folder)['model_dy']
# load dense descriptor model
metadata = load_pickle(os.path.join(model_folder, 'metadata.p'))
model_dd_file = metadata['model_file']
model_dd = torch.load(model_dd_file)
model_dd = model_dd.eval()
model_dd = model_dd.cuda()
spatial_descriptor_data = load_pickle(
os.path.join(model_folder, 'spatial_descriptors.p'))
return {
'model_dy': model_dy,
'model_dd': model_dd,
'spatial_descriptor_data': spatial_descriptor_data,
'metadata': metadata
}
def load_model(model_train_dir, strict=True):
"""
Helper function to load dynamics model and vision model
"""
model_dy_dict = model_builder.load_dynamics_model_from_folder(
model_train_dir, strict=strict)
model_dy = model_dy_dict['model_dy']
model_dy = model_dy.eval()
model_dy = model_dy.cuda()
model_name = model_dy_dict['model_name']
# load dense descriptor model
metadata = load_pickle(os.path.join(model_train_dir, 'metadata.p'))
model_dd_file = metadata['model_file']
model_dd = torch.load(model_dd_file)
model_dd = model_dd.eval()
model_dd = model_dd.cuda()
spatial_descriptor_data = load_pickle(
os.path.join(model_train_dir, 'spatial_descriptors.p'))
return {
'model_dy': model_dy,
'model_dd': model_dd,
'spatial_descriptor_data': spatial_descriptor_data,
'metadata': metadata,
'model_name': model_name,
}
def load_model_state_dict(model_folder=None):
models_root = "/home/manuelli/data/key_dynam/dev/experiments/20/dataset_correlle_mug-small_many_colors_600/trained_models/dynamics"
# model_name = "DD_3D/2020-06-04-19-18-48-274487_DD_3D_z_state_n_his_2_no_T_aug"
model_name = "DD_3D/2020-06-05-00-25-39-676089_DD_3D_all_n_his_2"
# model_name = "DD_3D/2020-06-05-15-25-01-580144_DD_3D_all_z_n_his_2"
model_folder = os.path.join(models_root, model_name)
model_dy = model_builder.load_dynamics_model_from_folder(model_folder)['model_dy']
# load dense descriptor model
metadata = load_pickle(os.path.join(model_folder, 'metadata.p'))
model_dd_file = metadata['model_file']
model_dd = torch.load(model_dd_file)
model_dd = model_dd.eval()
model_dd = model_dd.cuda()
spatial_descriptor_data = load_pickle(os.path.join(model_folder, 'spatial_descriptors.p'))
return {'model_dy': model_dy,
'model_dd': model_dd,
'spatial_descriptor_data': spatial_descriptor_data,
'metadata': metadata}
def load_dataset(
config, # e.g. experiments/07/config.yaml
episodes_config, # dict, e.g. real_push_box.yaml
episodes_root, # str: root of where all the logs are stored
precomputed_data_root=None,
max_num_episodes=None,
load_image_episode=True):
multi_episode_dict = dict()
for counter, episode_name in enumerate(episodes_config['episodes']):
# this is for debugging purposes
if (max_num_episodes is not None) and counter >= max_num_episodes:
break
episode_processed_dir = os.path.join(episodes_root, episode_name,
"processed")
# load the DenseCorrespondence Episode that handles image observations
# if requested
dc_episode = None
if load_image_episode:
####
precomputed_data = None
precomputed_data_file = None
if precomputed_data_root is not None:
# replace .h5 filename with .p for pickle file
precomputed_data_file = os.path.join(
precomputed_data_root, "%s.p" % (episode_name))
precomputed_data_file_hdf5 = os.path.join(
precomputed_data_root, "%s.h5" % (episode_name))
if os.path.isfile(precomputed_data_file):
precomputed_data = load_pickle(precomputed_data_file)
else:
raise ValueError("file doesn't exist: %s" %
(precomputed_data_file))
dc_episode = dc_episode_reader.DynamicSpartanEpisodeReader(
config=None,
root_dir=episode_processed_dir,
name=episode_name,
precomputed_data=precomputed_data,
precomputed_data_file=precomputed_data_file)
episode = DynamicSpartanEpisodeReader(
config,
episode_processed_dir,
name=episode_name,
downsample_rate=config['dataset']['downsample_rate'],
downsample_idx=0, # hardcoded for now
dc_episode_reader=dc_episode,
)
multi_episode_dict[episode_name] = episode
return multi_episode_dict
def debug():
save_dir = os.path.join(get_project_root(),
'sandbox/mpc/push_right_box_horizontal')
save_dir = "/home/manuelli/data/key_dynam/sandbox/2020-07-07-20-09-54_push_right_box_horizontal"
plan_msg = load_pickle(os.path.join(save_dir, 'plan_msg.p'),
encoding='latin1')
plan_msg = zmq_utils.convert(plan_msg)
compute_control_action_msg = load_pickle(os.path.join(
save_dir, 'compute_control_action_msg.p'),
encoding='latin1')
compute_control_action_msg = zmq_utils.convert(compute_control_action_msg)
K_matrix = None
T_world_camera = None
if 'K_matrix' in plan_msg:
K_matrix = plan_msg['K_matrix']
if 'T_world_camera' in plan_msg:
T_world_camera = plan_msg['T_world_camera']
d = load_model_and_data(
K_matrix=K_matrix,
T_world_camera=T_world_camera,
)
controller = ZMQController(
config=d['config'],
model_dy=d['model_dy'],
action_function=d['action_function'],
observation_function=d['observation_function'],
visual_observation_function=d['visual_observation_function'],
planner=d['planner'],
debug=True,
zmq_enabled=False,
camera_info=d['camera_info'],
)
controller._on_plan_msg(plan_msg)
controller._on_compute_control_action(compute_control_action_msg)
def drake_sim_dataset_loader(
precomputed_vision_data_dir, # precomputed_vision_data_dir
dataset_root, # location of original dataset
max_num_episodes=None,
):
spatial_descriptor_data = load_pickle(
os.path.join(precomputed_vision_data_dir, 'spatial_descriptors.p'))
metadata = load_pickle(
os.path.join(precomputed_vision_data_dir, 'metadata.p'))
descriptor_keypoints_root = os.path.join(precomputed_vision_data_dir,
'descriptor_keypoints')
multi_episode_dict = DrakeSimEpisodeReader.load_dataset(
dataset_root,
max_num_episodes=max_num_episodes,
descriptor_keypoints_root=descriptor_keypoints_root)
return {
'spatial_descriptors_data': spatial_descriptor_data,
'metadata': metadata,
'multi_episode_dict': multi_episode_dict,
}
def select_spatially_separated_descriptors(K=5, # number of reference descriptors
output_dir=None,
visualize=False):
raise ValueError("deprecated")
multi_episode_dict = exp_utils.load_episodes()['multi_episode_dict']
model_file = exp_utils.get_DD_model_file()
confidence_scores_folder = os.path.join(get_data_root(),
"dev/experiments/09/descriptor_confidence_scores/2020-03-25-19-57-26-556093_constant_velocity_500/2020-03-30-14-21-13-371713")
# folder = "dev/experiments/07/descriptor_confidence_scores/real_push_box/2020-03-10-15-57-43-867147"
folder = confidence_scores_folder
folder = os.path.join(get_data_root(), folder)
data_file = os.path.join(folder, 'data.p')
data = load_pickle(data_file)
heatmap_values = data['heatmap_values']
scoring_func = keypoint_selection.create_scoring_function(gamma=3)
score_data = keypoint_selection.score_heatmap_values(heatmap_values,
scoring_func=scoring_func)
sorted_idx = score_data['sorted_idx']
metadata_file = os.path.join(folder, 'metadata.p')
metadata = load_pickle(metadata_file)
camera_name = metadata['camera_name']
keypoint_idx = keypoint_selection.select_spatially_separated_keypoints(sorted_idx,
metadata['indices'],
position_diff_threshold=30,
K=K,
verbose=False)
ref_descriptors = metadata['ref_descriptors'][keypoint_idx] # [K, D]
spatial_descriptors_data = score_data
spatial_descriptors_data['spatial_descriptors'] = ref_descriptors
spatial_descriptors_data['spatial_descriptors_idx'] = keypoint_idx
save_pickle(spatial_descriptors_data, os.path.join(folder, 'spatial_descriptors.p'))
def load_episodes_from_config(config):
"""
Loads episodes using the path specified in the config
:param config:
:type config:
:return:
:rtype:
"""
data_path = config["dataset"]["data_path"]
if not os.path.isabs(data_path):
data_path = os.path.join(get_project_root(), data_path)
# load the data
print("loading data from disk . . . ")
raw_data = load_pickle(data_path)
print("finished loading data")
episodes = PyMunkEpisodeReader.load_pymunk_episodes_from_raw_data(raw_data)
return episodes
def create_pusher_slider_keypoint_dataset(config=None):
# load some previously generated data
project_root = get_project_root()
if config is None:
config_file = os.path.join(project_root, "experiments/02/config.yaml")
config = load_yaml(config_file)
action_function = ActionFunctionFactory.pusher_velocity
obs_function = ObservationFunctionFactory.pusher_pose_slider_keypoints(
config)
DATA_PATH = os.path.join(
project_root,
"test_data/pusher_slider_10_episodes/2019-10-22-21-30-02-536750.p")
raw_data = load_pickle(DATA_PATH)
episodes = PyMunkEpisodeReader.load_pymunk_episodes_from_raw_data(raw_data)
# create MultiEpisodeDataset
dataset = MultiEpisodeDataset(config,
action_function=action_function,
observation_function=obs_function,
episodes=episodes)
episode = dataset.get_random_episode()
data_0 = episode.get_observation(0)
data_1 = episode.get_observation(1)
print("time 0", data_0["sim_time"])
print("time 1", data_1["sim_time"])
# episode_name = episodes.keys()[0]
# episode = episodes[episode_name]
# data = episode.data
# print("episode.data.keys()", episode.data.keys())
# print("test ", type(data["trajectory"][0].keys()))
# print("test ", data["trajectory"][0].keys())
return dataset, config
def main():
dataset_name, multi_episode_dict = load_episodes()
## Load Model
model_name, model_file = get_DD_model_file()
model = torch.load(model_file)
model = model.cuda()
model = model.eval()
# make this unique
output_dir = os.path.join(get_data_root(),
"dev/experiments/09/precomputed_vision_data",
"dataset_%s" % (dataset_name),
"model_name_%s" % (model_name),
get_current_YYYY_MM_DD_hh_mm_ss_ms())
camera_name = "camera_1_top_down"
episode_name = "2020-05-13-21-55-01-487901_idx_33"
episode_idx = 22
# compute descriptor confidence scores
if True:
print("\n\n---------Computing Descriptor Confidence Scores-----------")
metadata_file = os.path.join(output_dir, 'metadata.p')
if os.path.isfile(metadata_file):
answer = input(
"metadata.p file already exists, do you want to overwrite it? y/n"
)
if answer == "y":
os.rmdir(output_dir)
print("removing existing file and continuing")
else:
print("aborting")
quit()
compute_descriptor_confidences(
multi_episode_dict,
model,
output_dir,
batch_size=10,
num_workers=20,
model_file=model_file,
camera_name=camera_name,
num_ref_descriptors=50,
num_batches=10,
episode_name_arg=episode_name,
episode_idx=episode_idx,
)
if True:
metadata_file = os.path.join(output_dir, 'metadata.p')
metadata = load_pickle(metadata_file)
# metadata_file = "/media/hdd/data/key_dynam/dev/experiments/09/precomputed_vision_data/dataset_2020-03-25-19-57-26-556093_constant_velocity_500/model_name_2020-04-07-14-31-35-804270_T_aug_dataset/2020-04-09-20-51-50-624799/metadata.p"
# metadata = load_pickle(metadata_file)
print("\n\n---------Precomputing Descriptor Keypoints-----------")
descriptor_keypoints_output_dir = os.path.join(output_dir,
"descriptor_keypoints")
precompute_descriptor_keypoints(
multi_episode_dict,
model,
descriptor_keypoints_output_dir,
ref_descriptors_metadata=metadata,
batch_size=10,
num_workers=20,
)
if True:
confidence_score_data_file = os.path.join(output_dir, 'data.p')
confidence_score_data = load_pickle(confidence_score_data_file)
print(
"\n\n---------Selecting Spatially Separated Keypoints-----------")
score_and_select_spatially_separated_keypoints(
metadata,
confidence_score_data=confidence_score_data,
K=5,
position_diff_threshold=25,
output_dir=output_dir,
)
print("Data saved at: ", output_dir)
print("Finished Normally")
def load_dataset(
dataset_root, # str: folder containing dataset
load_image_data=True,
descriptor_images_root=None, # str: (optional) folder containing hdf5 files with descriptors
descriptor_keypoints_root=None,
max_num_episodes=None, # int, max num episodes to load
precomputed_data_root=None,
):
"""
:param dataset_root: folder should contain
- config.yaml
- metadata.yaml
- <episode_name.p>
- <episode_name.h5>
:type dataset_root:
:return:
:rtype:
"""
if load_image_data:
from key_dynam.dense_correspondence.dc_drake_sim_episode_reader import DCDrakeSimEpisodeReader
metadata = load_yaml(DrakeSimEpisodeReader.metadata_file(dataset_root))
multi_episode_dict = dict()
episode_names = list(metadata['episodes'].keys())
episode_names.sort() # sort the keys
num_episodes = len(episode_names)
# optionally don't read all episodes
if (max_num_episodes is not None) and (max_num_episodes > 0):
# compute the number of episodes to read, in sorted order
num_episodes = int(min(len(episode_names), max_num_episodes))
for idx in range(num_episodes):
episode_name = episode_names[idx]
val = metadata['episodes'][episode_name]
# load non image data
non_image_data_file = os.path.join(dataset_root,
val['non_image_data_file'])
assert os.path.isfile(
non_image_data_file), "File doesn't exist: %s" % (
non_image_data_file)
non_image_data = load_pickle(non_image_data_file)
dc_episode_reader = None
if load_image_data:
# load image data
image_data_file = os.path.join(dataset_root,
val['image_data_file'])
assert os.path.isfile(
image_data_file), "File doesn't exist: %s" % (
image_data_file)
descriptor_image_data_file = None
if descriptor_images_root is not None:
descriptor_image_data_file = os.path.join(
descriptor_images_root, val['image_data_file'])
assert os.path.isfile(
descriptor_image_data_file
), "File doesn't exist: %s" % (descriptor_image_data_file)
descriptor_keypoints_data = None
descriptor_keypoints_hdf5_file = None
if descriptor_keypoints_root is not None:
# replace .h5 filename with .p for pickle file
descriptor_keypoints_data_file = val[
'image_data_file'].split(".")[0] + ".p"
descriptor_keypoints_data_file = os.path.join(
descriptor_keypoints_root,
descriptor_keypoints_data_file)
descriptor_keypoints_hdf5_file = os.path.join(
descriptor_keypoints_root, val['image_data_file'])
if os.path.isfile(descriptor_keypoints_data_file):
descriptor_keypoints_data = load_pickle(
descriptor_keypoints_data_file)
else:
assert os.path.isfile(
descriptor_keypoints_hdf5_file
), "File doesn't exist: %s" % (
descriptor_keypoints_hdf5_file)
####
precomputed_data = None
precomputed_data_file = None
if precomputed_data_root is not None:
# replace .h5 filename with .p for pickle file
precomputed_data_file = val['image_data_file'].split(
".")[0] + ".p"
precomputed_data_file = os.path.join(
precomputed_data_root, precomputed_data_file)
if os.path.isfile(precomputed_data_file):
precomputed_data = load_pickle(precomputed_data_file)
else:
raise ValueError("file doesn't exist: %s" %
(precomputed_data_file))
dc_episode_reader = DCDrakeSimEpisodeReader(
non_image_data,
image_data_file,
descriptor_image_data_file=descriptor_image_data_file,
descriptor_keypoints_data=descriptor_keypoints_data,
descriptor_keypoints_data_file=
descriptor_keypoints_hdf5_file,
precomputed_data=precomputed_data,
precomputed_data_file=precomputed_data_file)
episode_reader = DrakeSimEpisodeReader(
non_image_data=non_image_data,
episode_name=episode_name,
dc_episode_reader=dc_episode_reader)
multi_episode_dict[episode_name] = episode_reader
return multi_episode_dict
def main(dataset_name):
# sample from specific image
dataset_paths = None
episode_name = None
camera_name = None
episode_idx = None
if dataset_name == "dps_box_on_side_600":
camera_name = "camera_angled"
episode_name = "2020-05-13-21-55-01-487901_idx_33"
episode_idx = 22
dataset_paths = exp_18_utils.get_dataset_paths(dataset_name)
elif dataset_name == "correlle_mug-small_single_color_600":
camera_name = "camera_1_top_down"
episode_name = "2020-06-02-14-15-27-898104_idx_56"
episode_idx = 18
dataset_paths = exp_20_utils.get_dataset_paths(dataset_name)
elif dataset_name == "correlle_mug-small_many_colors_600":
camera_name = "camera_1_top_down"
episode_name = "2020-06-03-15-48-50-165064_idx_56"
episode_idx = 20
dataset_paths = exp_20_utils.get_dataset_paths(dataset_name)
else:
raise ValueError("unknown dataset")
dataset_root = dataset_paths['dataset_root']
dataset_name = dataset_paths['dataset_name']
## Load Model
model_name, model_file = get_DD_model_file(dataset_name)
model_train_dir = os.path.dirname(model_file)
print("model_train_dir", model_train_dir)
print("model_file", model_file)
model = torch.load(model_file)
model = model.cuda()
model = model.eval()
multi_episode_dict = DCDrakeSimEpisodeReader.load_dataset(
dataset_root, max_num_episodes=None)
output_dir = os.path.join(
model_train_dir,
'precomputed_vision_data/descriptor_keypoints/dataset_%s/' %
(dataset_name))
# compute descriptor confidence scores
if False:
print("\n\n---------Computing Descriptor Confidence Scores-----------")
metadata_file = os.path.join(output_dir, 'metadata.p')
if os.path.isfile(metadata_file):
answer = input(
"metadata.p file already exists, do you want to overwrite it? y/n\n"
)
if answer == "y":
shutil.rmtree(output_dir)
print("removing existing file and continuing")
else:
print("aborting")
quit()
set_seed(0)
compute_descriptor_confidences(
multi_episode_dict,
model,
output_dir,
batch_size=10,
num_workers=20,
model_file=model_file,
camera_name=camera_name,
num_ref_descriptors=50,
num_batches=10,
episode_name_arg=episode_name,
episode_idx=episode_idx,
)
if False:
confidence_score_data_file = os.path.join(output_dir, 'data.p')
confidence_score_data = load_pickle(confidence_score_data_file)
metadata_file = os.path.join(output_dir, 'metadata.p')
metadata = load_pickle(metadata_file)
print(
"\n\n---------Selecting Spatially Separated Keypoints-----------")
score_and_select_spatially_separated_keypoints(
metadata,
confidence_score_data=confidence_score_data,
K=4,
position_diff_threshold=15,
output_dir=output_dir,
)
# visualize descriptors
if False:
metadata_file = os.path.join(output_dir, 'metadata.p')
metadata = load_pickle(metadata_file)
episode_name = metadata['episode_name']
episode_idx = metadata['episode_idx']
camera_name = metadata['camera_name']
episode = multi_episode_dict[episode_name]
data = episode.get_image_data(camera_name, episode_idx)
rgb = data['rgb']
uv = metadata['indices']
print("uv.shape", uv.shape)
color = [0, 255, 0]
draw_reticles(rgb, uv[:, 0], uv[:, 1], label_color=color)
save_file = os.path.join(output_dir, 'sampled_descriptors.png')
plt.figure()
plt.imshow(rgb)
plt.savefig(save_file)
plt.show()
# visualize spatially separated descriptors
if False:
metadata_file = os.path.join(output_dir, 'metadata.p')
metadata = load_pickle(metadata_file)
spatial_descriptor_file = os.path.join(output_dir,
'spatial_descriptors.p')
spatial_descriptors_data = load_pickle(spatial_descriptor_file)
des_idx = spatial_descriptors_data['spatial_descriptors_idx']
episode_name = metadata['episode_name']
episode_idx = metadata['episode_idx']
camera_name = metadata['camera_name']
episode = multi_episode_dict[episode_name]
data = episode.get_image_data(camera_name, episode_idx)
rgb = data['rgb']
uv = metadata['indices']
print("uv.shape", uv.shape)
color = [0, 255, 0]
draw_reticles(rgb, uv[des_idx, 0], uv[des_idx, 1], label_color=color)
save_file = os.path.join(output_dir,
'spatially_separated_descriptors.png')
plt.figure()
plt.imshow(rgb)
plt.savefig(save_file)
plt.show()
if True:
metadata_file = os.path.join(output_dir, 'metadata.p')
metadata = load_pickle(metadata_file)
# metadata_file = "/media/hdd/data/key_dynam/dev/experiments/09/precomputed_vision_data/dataset_2020-03-25-19-57-26-556093_constant_velocity_500/model_name_2020-04-07-14-31-35-804270_T_aug_dataset/2020-04-09-20-51-50-624799/metadata.p"
# metadata = load_pickle(metadata_file)
print("\n\n---------Precomputing Descriptor Keypoints-----------")
descriptor_keypoints_output_dir = os.path.join(output_dir,
"descriptor_keypoints")
precompute_descriptor_keypoints(multi_episode_dict,
model,
descriptor_keypoints_output_dir,
ref_descriptors_metadata=metadata,
batch_size=8,
num_workers=20,
camera_names=[camera_name])
print("Data saved at: ", output_dir)
print("Finished Normally")
import numpy as np
from key_dynam.utils.utils import load_pickle, save_pickle, save_yaml
from key_dynam.utils import meshcat_utils
from key_dynam.utils import transform_utils
data_file = "/home/manuelli/data/key_dynam/hardware_experiments/closed_loop_rollouts/stable/2020-07-10-22-16-08_long_push_on_long_side/mpc_rollouts/2020-07-10-22-19-03-591910/data.p"
data = load_pickle(data_file)
pts = data['plan']['plan_data'][-1]['dynamics_model_input_data'][
'visual_observation']['pts_W']
print("pts\n", pts)
centroid = np.mean(pts, axis=0)
pts_centered = pts - centroid
save_data = {'object_points': pts_centered.tolist()}
save_file = "object_points_master.yaml"
save_yaml(save_data, save_file)
# do some meshcat debug
vis = meshcat_utils.make_default_visualizer_object()
meshcat_utils.visualize_points(vis,
"object_points_centered",
pts_centered,
color=[0, 0, 255],
size=0.01)
meshcat_utils.visualize_points(vis,
"object_points_world",
pts,
color=[0, 255, 0],
size=0.01)
def load_model_and_data():
dataset_name = "push_box_hardware"
# model_name = "DD_2D/2020-06-24-22-22-58-234812_DD_3D_n_his_2" # this model is actually 3D
# model_name = "DD_3D/2020-06-25-00-49-29-679042_DD_3D_n_his_2_T_aug"
# model_name = "DD_3D/2020-06-25-00-39-29-020621_DD_3D_n_his_2"
model_name = "DD_3D/2020-07-02-17-59-21-362337_DD_3D_n_his_2_T_aug"
train_dir = "/home/manuelli/data/key_dynam/dev/experiments/22/dataset_push_box_hardware/trained_models/dynamics"
train_dir = os.path.join(train_dir, model_name)
ckpt_file = os.path.join(train_dir, "net_best_dy_state_dict.pth")
config = load_yaml(os.path.join(train_dir, 'config.yaml'))
state_dict = torch.load(ckpt_file)
# build dynamics model
model_dy = build_dynamics_model(config)
# print("state_dict.keys()", state_dict.keys())
model_dy.load_state_dict(state_dict)
model_dy = model_dy.eval()
model_dy = model_dy.cuda()
spatial_descriptor_data = load_pickle(
os.path.join(train_dir, 'spatial_descriptors.p'))
metadata = load_pickle(os.path.join(train_dir, 'metadata.p'))
# load the dataset
dataset_paths = get_dataset_paths(dataset_name)
dataset_root = dataset_paths['dataset_root']
episodes_config = dataset_paths['episodes_config']
precomputed_vision_data_root = DD_utils.get_precomputed_data_root(
dataset_name)['precomputed_data_root']
# descriptor_keypoints_root = os.path.join(precomputed_vision_data_root, 'descriptor_keypoints')
descriptor_keypoints_root = os.path.join(precomputed_vision_data_root,
'descriptor_keypoints')
multi_episode_dict = DynamicSpartanEpisodeReader.load_dataset(
config=config,
episodes_config=episodes_config,
episodes_root=dataset_paths['dataset_root'],
load_image_episode=True,
precomputed_data_root=descriptor_keypoints_root,
max_num_episodes=None)
visual_observation_function = PrecomputedVisualObservationFunctionFactory.function_from_config(
config,
keypoint_idx=spatial_descriptor_data['spatial_descriptors_idx'])
action_function = ActionFunctionFactory.function_from_config(config)
observation_function = ObservationFunctionFactory.function_from_config(
config)
dataset = MultiEpisodeDataset(
config,
action_function=action_function,
observation_function=observation_function,
episodes=multi_episode_dict,
visual_observation_function=visual_observation_function,
phase="valid", # this means no data augmentation
)
return {
"model_dy": model_dy,
'dataset': dataset,
'config': config,
"multi_episode_dict": multi_episode_dict,
'spatial_descriptor_data': spatial_descriptor_data,
}
def load_model_and_data():
dataset_name = "push_box_hardware"
# model_name = "DD_2D/2020-06-24-22-22-58-234812_DD_3D_n_his_2" # this model is actually 3D
# model_name = "DD_3D/2020-06-25-00-49-29-679042_DD_3D_n_his_2_T_aug"
# model_name = "DD_3D/2020-06-25-00-39-29-020621_DD_3D_n_his_2"
model_name = "DD_3D/2020-07-02-17-59-21-362337_DD_3D_n_his_2_T_aug"
train_dir = "/home/manuelli/data/key_dynam/dev/experiments/22/dataset_push_box_hardware/trained_models/dynamics"
train_dir = os.path.join(train_dir, model_name)
ckpt_file = os.path.join(train_dir, "net_best_dy_state_dict.pth")
config = load_yaml(os.path.join(train_dir, 'config.yaml'))
state_dict = torch.load(ckpt_file)
# build dynamics model
model_dy = build_dynamics_model(config)
# print("state_dict.keys()", state_dict.keys())
model_dy.load_state_dict(state_dict)
model_dy = model_dy.eval()
model_dy = model_dy.cuda()
spatial_descriptor_data = load_pickle(
os.path.join(train_dir, 'spatial_descriptors.p'))
metadata = load_pickle(os.path.join(train_dir, 'metadata.p'))
# build dense-descriptor model
model_dd_file = metadata['model_file']
model_dd = torch.load(model_dd_file)
model_dd = model_dd.eval()
model_dd = model_dd.cuda()
# load the dataset
dataset_paths = get_dataset_paths(dataset_name)
dataset_root = dataset_paths['dataset_root']
episodes_config = dataset_paths['episodes_config']
precomputed_vision_data_root = DD_utils.get_precomputed_data_root(
dataset_name)['precomputed_data_root']
# descriptor_keypoints_root = os.path.join(precomputed_vision_data_root, 'descriptor_keypoints')
descriptor_keypoints_root = os.path.join(precomputed_vision_data_root,
'descriptor_keypoints')
multi_episode_dict = DynamicSpartanEpisodeReader.load_dataset(
config=config,
episodes_config=episodes_config,
episodes_root=dataset_paths['dataset_root'],
load_image_episode=True,
precomputed_data_root=descriptor_keypoints_root,
max_num_episodes=None)
visual_observation_function = PrecomputedVisualObservationFunctionFactory.function_from_config(
config,
keypoint_idx=spatial_descriptor_data['spatial_descriptors_idx'])
action_function = ActionFunctionFactory.function_from_config(config)
observation_function = ObservationFunctionFactory.function_from_config(
config)
dataset = MultiEpisodeDataset(
config,
action_function=action_function,
observation_function=observation_function,
episodes=multi_episode_dict,
visual_observation_function=visual_observation_function,
phase="valid", # this means no data augmentation
)
#### PLANNER #######
planner = None
# make a planner config
planner_config = copy.copy(model_dy.config)
config_tmp = load_yaml(
os.path.join(get_project_root(),
'experiments/exp_22_push_box_hardware/config_DD_3D.yaml'))
planner_config['mpc'] = config_tmp['mpc']
if PLANNER_TYPE == "random_shooting":
planner = RandomShootingPlanner(planner_config)
elif PLANNER_TYPE == "mppi":
planner = PlannerMPPI(planner_config)
else:
raise ValueError("unknown planner type: %s" % (PLANNER_TYPE))
return {
"model_dy": model_dy,
'model_dd': model_dd,
'dataset': dataset,
'config': config,
"multi_episode_dict": multi_episode_dict,
'spatial_descriptor_data': spatial_descriptor_data,
'planner': planner,
'observation_function': observation_function,
'action_function': action_function,
}
def train_dynamics(config,
train_dir, # str: directory to save output
multi_episode_dict, # multi_episode_dict
):
use_precomputed_keypoints = config['dataset']['visual_observation']['enabled'] and config['dataset']['visual_observation']['descriptor_keypoints']
# set random seed for reproduction
set_seed(config['train']['random_seed'])
st_epoch = config['train']['resume_epoch'] if config['train']['resume_epoch'] > 0 else 0
tee = Tee(os.path.join(train_dir, 'train_st_epoch_%d.log' % st_epoch), 'w')
tensorboard_dir = os.path.join(train_dir, "tensorboard")
if not os.path.exists(tensorboard_dir):
os.makedirs(tensorboard_dir)
writer = SummaryWriter(log_dir=tensorboard_dir)
# save the config
save_yaml(config, os.path.join(train_dir, "config.yaml"))
action_function = ActionFunctionFactory.function_from_config(config)
observation_function = ObservationFunctionFactory.function_from_config(config)
datasets = {}
dataloaders = {}
data_n_batches = {}
for phase in ['train', 'valid']:
print("Loading data for %s" % phase)
datasets[phase] = MultiEpisodeDataset(config,
action_function=action_function,
observation_function=observation_function,
episodes=multi_episode_dict,
phase=phase)
dataloaders[phase] = DataLoader(
datasets[phase], batch_size=config['train']['batch_size'],
shuffle=True if phase == 'train' else False,
num_workers=config['train']['num_workers'], drop_last=True)
data_n_batches[phase] = len(dataloaders[phase])
use_gpu = torch.cuda.is_available()
# compute normalization parameters if not starting from pre-trained network . . .
'''
define model for dynamics prediction
'''
model_dy = build_visual_dynamics_model(config)
K = config['vision_net']['num_ref_descriptors']
print("model_dy.vision_net._reference_descriptors.shape", model_dy.vision_net._ref_descriptors.shape)
print("model_dy.vision_net.descriptor_dim", model_dy.vision_net.descriptor_dim)
print("model_dy #params: %d" % count_trainable_parameters(model_dy))
camera_name = config['vision_net']['camera_name']
W = config['env']['rgbd_sensors']['sensor_list'][camera_name]['width']
H = config['env']['rgbd_sensors']['sensor_list'][camera_name]['height']
diag = np.sqrt(W**2 + H**2) # use this to scale the loss
# sample reference descriptors unless using precomputed keypoints
if not use_precomputed_keypoints:
# sample reference descriptors
episode_names = list(datasets["train"].episode_dict.keys())
episode_names.sort()
episode_name = episode_names[0]
episode = datasets["train"].episode_dict[episode_name]
episode_idx = 0
camera_name = config["vision_net"]["camera_name"]
image_data = episode.get_image_data(camera_name, episode_idx)
des_img = torch.Tensor(image_data['descriptor'])
mask_img = torch.Tensor(image_data['mask'])
ref_descriptor_dict = sample_descriptors(des_img,
mask_img,
config['vision_net']['num_ref_descriptors'])
model_dy.vision_net._ref_descriptors.data = ref_descriptor_dict['descriptors']
model_dy.vision_net.reference_image = image_data['rgb']
model_dy.vision_net.reference_indices = ref_descriptor_dict['indices']
else:
metadata_file = os.path.join(get_data_root(), config['dataset']['descriptor_keypoints_dir'], 'metadata.p')
descriptor_metadata = load_pickle(metadata_file)
# [32, 2]
ref_descriptors = torch.Tensor(descriptor_metadata['ref_descriptors'])
# [K, 2]
ref_descriptors = ref_descriptors[:K]
model_dy.vision_net._ref_descriptors.data = ref_descriptors
model_dy.vision_net._ref_descriptors_metadata = descriptor_metadata
# this is just a sanity check
assert model_dy.vision_net.num_ref_descriptors == K
print("reference_descriptors", model_dy.vision_net._ref_descriptors)
# criterion
criterionMSE = nn.MSELoss()
l1Loss = nn.L1Loss()
# optimizer
params = model_dy.parameters()
lr = float(config['train']['lr'])
optimizer = optim.Adam(params, lr=lr, betas=(config['train']['adam_beta1'], 0.999))
# setup scheduler
sc = config['train']['lr_scheduler']
scheduler = ReduceLROnPlateau(optimizer,
mode='min',
factor=sc['factor'],
patience=sc['patience'],
threshold_mode=sc['threshold_mode'],
cooldown= sc['cooldown'],
verbose=True)
if use_gpu:
print("using gpu")
model_dy = model_dy.cuda()
print("model_dy.vision_net._ref_descriptors.device", model_dy.vision_net._ref_descriptors.device)
print("model_dy.vision_net #params: %d" %(count_trainable_parameters(model_dy.vision_net)))
best_valid_loss = np.inf
global_iteration = 0
epoch_counter_external = 0
try:
for epoch in range(st_epoch, config['train']['n_epoch']):
phases = ['train', 'valid']
epoch_counter_external = epoch
writer.add_scalar("Training Params/epoch", epoch, global_iteration)
for phase in phases:
model_dy.train(phase == 'train')
meter_loss_rmse = AverageMeter()
step_duration_meter = AverageMeter()
# bar = ProgressBar(max_value=data_n_batches[phase])
loader = dataloaders[phase]
for i, data in enumerate(loader):
step_start_time = time.time()
global_iteration += 1
with torch.set_grad_enabled(phase == 'train'):
n_his, n_roll = config['train']['n_history'], config['train']['n_rollout']
n_samples = n_his + n_roll
if DEBUG:
print("global iteration: %d" %(global_iteration))
# visual_observations = data['visual_observations']
visual_observations_list = data['visual_observations_list']
observations = data['observations']
actions = data['actions']
if use_gpu:
observations = observations.cuda()
actions = actions.cuda()
# states, actions = data
assert actions.size(1) == n_samples
B = actions.size(0)
loss_mse = 0.
# compute the output of the visual model for all timesteps
visual_model_output_list = []
for visual_obs in visual_observations_list:
# visual_obs is a dict containing observation for a single
# time step (of course across a batch however)
# visual_obs[<camera_name>]['rgb_tensor'] has shape [B, 3, H, W]
# probably need to cast input to cuda
dynamics_net_input = None
if use_precomputed_keypoints:
# note precomputed descriptors stored on disk are of size
# K = 32. We need to trim it down to the appropriate size
# [B, K_disk, 2] where K_disk is num keypoints on disk
keypoints = visual_obs[camera_name]['descriptor_keypoints']
# [B, 32, 2] where K is num keypoints
keypoints = keypoints[:,:K]
if DEBUG:
print("keypoints.shape", keypoints.shape)
dynamics_net_input = keypoints.flatten(start_dim=1)
else:
out_dict = model_dy.vision_net.forward(visual_obs)
# [B, vision_model_out_dim]
dynamics_net_input = out_dict['dynamics_net_input']
visual_model_output_list.append(dynamics_net_input)
# concatenate this into a tensor
# [B, n_samples, vision_model_out_dim]
visual_model_output = torch.stack(visual_model_output_list, dim=1)
# cast this to float so it can be concatenated below
visual_model_output = visual_model_output.type_as(observations)
if DEBUG:
print('visual_model_output.shape', visual_model_output.shape)
print("observations.shape", observations.shape)
print("actions.shape", actions.shape)
# states is gotten by concatenating visual_observations and observations
# [B, n_samples, vision_model_out_dim + obs_dim]
states = torch.cat((visual_model_output, observations), dim=-1)
# state_cur: B x n_his x state_dim
state_cur = states[:, :n_his]
if DEBUG:
print("states.shape", states.shape)
for j in range(n_roll):
if DEBUG:
print("n_roll j: %d" %(j))
state_des = states[:, n_his + j]
# action_cur: B x n_his x action_dim
action_cur = actions[:, j : j + n_his] if actions is not None else None
# state_pred: B x state_dim
# state_pred: B x state_dim
input = {'observation': state_cur,
'action': action_cur,
}
if DEBUG:
print("state_cur.shape", state_cur.shape)
print("action_cur.shape", action_cur.shape)
state_pred = model_dy.dynamics_net(input)
# normalize by diag to ensure the loss is in [0,1] range
loss_mse_cur = criterionMSE(state_pred/diag, state_des/diag)
loss_mse += loss_mse_cur / n_roll
# l1Loss
loss_l1 = l1Loss(state_pred, state_des)
# update state_cur
# state_pred.unsqueeze(1): B x 1 x state_dim
# state_cur: B x n_his x state_dim
state_cur = torch.cat([state_cur[:, 1:], state_pred.unsqueeze(1)], 1)
meter_loss_rmse.update(np.sqrt(loss_mse.item()), B)
step_duration_meter.update(time.time() - step_start_time)
if phase == 'train':
optimizer.zero_grad()
loss_mse.backward()
optimizer.step()
if (i % config['train']['log_per_iter'] == 0) or (global_iteration % config['train']['log_per_iter'] == 0):
log = '%s [%d/%d][%d/%d] LR: %.6f' % (
phase, epoch, config['train']['n_epoch'], i, data_n_batches[phase],
get_lr(optimizer))
log += ', rmse: %.6f (%.6f)' % (
np.sqrt(loss_mse.item()), meter_loss_rmse.avg)
log += ', step time %.6f' %(step_duration_meter.avg)
step_duration_meter.reset()
print(log)
# log data to tensorboard
# only do it once we have reached 100 iterations
if global_iteration > 100:
writer.add_scalar("Params/learning rate", get_lr(optimizer), global_iteration)
writer.add_scalar("Loss_MSE/%s" %(phase), loss_mse.item(), global_iteration)
writer.add_scalar("L1/%s" %(phase), loss_l1.item(), global_iteration)
writer.add_scalar("L1_fraction/%s" %(phase), loss_l1.item()/diag, global_iteration)
writer.add_scalar("RMSE average loss/%s" %(phase), meter_loss_rmse.avg, global_iteration)
if phase == 'train' and i % config['train']['ckp_per_iter'] == 0:
save_model(model_dy, '%s/net_dy_epoch_%d_iter_%d' % (train_dir, epoch, i))
log = '%s [%d/%d] Loss: %.6f, Best valid: %.6f' % (
phase, epoch, config['train']['n_epoch'], meter_loss_rmse.avg, best_valid_loss)
print(log)
if phase == 'valid':
if config['train']['lr_scheduler']['enabled']:
scheduler.step(meter_loss_rmse.avg)
# print("\nPhase == valid")
# print("meter_loss_rmse.avg", meter_loss_rmse.avg)
# print("best_valid_loss", best_valid_loss)
if meter_loss_rmse.avg < best_valid_loss:
best_valid_loss = meter_loss_rmse.avg
save_model(model_dy, '%s/net_best_dy' % (train_dir))
writer.flush() # flush SummaryWriter events to disk
except KeyboardInterrupt:
# save network if we have a keyboard interrupt
save_model(model_dy, '%s/net_dy_epoch_%d_keyboard_interrupt' % (train_dir, epoch_counter_external))
writer.flush() # flush SummaryWriter events to disk
def load_model_and_data(
K_matrix=None,
T_world_camera=None,
):
dataset_name = "push_box_hardware"
model_name = "DD_3D/2020-07-02-17-59-21-362337_DD_3D_n_his_2_T_aug"
train_dir = os.path.join(
get_data_root(),
"dev/experiments/22/dataset_push_box_hardware/trained_models/dynamics")
# train_dir = "/home/manuelli/data/key_dynam/dev/experiments/22/dataset_push_box_hardware/trained_models/dynamics"
train_dir = os.path.join(train_dir, model_name)
ckpt_file = os.path.join(train_dir, "net_best_dy_state_dict.pth")
train_config = load_yaml(os.path.join(train_dir, 'config.yaml'))
state_dict = torch.load(ckpt_file)
# build dynamics model
model_dy = build_dynamics_model(train_config)
# print("state_dict.keys()", state_dict.keys())
model_dy.load_state_dict(state_dict)
model_dy = model_dy.eval()
model_dy = model_dy.cuda()
# load the dataset
dataset_paths = get_dataset_paths(dataset_name)
dataset_root = dataset_paths['dataset_root']
episodes_config = dataset_paths['episodes_config']
spatial_descriptor_data = load_pickle(
os.path.join(train_dir, 'spatial_descriptors.p'))
metadata = load_pickle(os.path.join(train_dir, 'metadata.p'))
ref_descriptors = spatial_descriptor_data['spatial_descriptors']
ref_descriptors = torch_utils.cast_to_torch(ref_descriptors).cuda()
# dense descriptor model
model_dd_file = metadata['model_file']
model_dd = torch.load(model_dd_file)
model_dd = model_dd.eval()
model_dd = model_dd.cuda()
camera_name = train_config['dataset']['visual_observation_function'][
'camera_name']
camera_info = None
if (T_world_camera is not None) and (K_matrix is not None):
camera_info = {
"K": K_matrix,
'T_world_camera': T_world_camera,
}
else:
camera_info = get_spartan_camera_info(camera_name)
camera_info['camera_name'] = camera_name
visual_observation_function = \
VisualObservationFunctionFactory.descriptor_keypoints_3D(config=train_config,
camera_name=camera_name,
model_dd=model_dd,
ref_descriptors=ref_descriptors,
K_matrix=camera_info['K'],
T_world_camera=camera_info['T_world_camera'],
)
action_function = ActionFunctionFactory.function_from_config(train_config)
observation_function = ObservationFunctionFactory.function_from_config(
train_config)
#### PLANNER #######
planner = None
# make a planner config
planner_config = copy.copy(train_config)
config_tmp = load_yaml(
os.path.join(get_project_root(),
'experiments/exp_22_push_box_hardware/config_DD_3D.yaml'))
planner_config['mpc'] = config_tmp['mpc']
if PLANNER_TYPE == "random_shooting":
planner = RandomShootingPlanner(planner_config)
elif PLANNER_TYPE == "mppi":
planner = PlannerMPPI(planner_config)
else:
raise ValueError("unknown planner type: %s" % (PLANNER_TYPE))
return {
"model_dy": model_dy,
'config': train_config,
'spatial_descriptor_data': spatial_descriptor_data,
'action_function': action_function,
'observation_function': observation_function,
'visual_observation_function': visual_observation_function,
'planner': planner,
'camera_info': camera_info,
}
def run_precompute_descriptors_pipeline(multi_episode_dict, # dict
model, # dense descriptor model file
model_file=None,
output_dir=None, # str where to save data
episode_name=None, # optional for descriptor sampling
camera_name=None, # which camera to compute descriptors for
episode_idx=None, # optional for descriptor sampling
visualize=True,
K=5,
position_diff_threshold=20,
seed=0,
):
assert model_file is not None
assert camera_name is not None
#
# ## Load Model
# model_train_dir = os.path.dirname(model_file)
#
# print("model_train_dir", model_train_dir)
# print("model_file", model_file)
# model = torch.load(model_file)
# model = model.cuda()
# model = model.eval()
#
# output_dir = os.path.join(model_train_dir,
# 'precomputed_vision_data/descriptor_keypoints/dataset_%s/' % (dataset_name))
# compute descriptor confidence scores
set_seed(seed)
if True:
print("\n\n---------Computing Descriptor Confidence Scores-----------")
metadata_file = os.path.join(output_dir, 'metadata.p')
if os.path.isfile(metadata_file):
answer = input("metadata.p file already exists, do you want to overwrite it? y/n\n")
if answer == "y":
shutil.rmtree(output_dir)
print("removing existing file and continuing")
else:
print("aborting")
quit()
compute_descriptor_confidences(multi_episode_dict,
model,
output_dir,
batch_size=10,
num_workers=20,
model_file=model_file,
camera_name=camera_name,
num_ref_descriptors=50,
num_batches=10,
episode_name_arg=episode_name,
episode_idx=episode_idx,
)
if True:
confidence_score_data_file = os.path.join(output_dir, 'data.p')
confidence_score_data = load_pickle(confidence_score_data_file)
metadata_file = os.path.join(output_dir, 'metadata.p')
metadata = load_pickle(metadata_file)
print("\n\n---------Selecting Spatially Separated Keypoints-----------")
score_and_select_spatially_separated_keypoints(metadata,
confidence_score_data=confidence_score_data,
K=K,
position_diff_threshold=position_diff_threshold,
output_dir=output_dir,
)
# visualize descriptors
if True:
metadata_file = os.path.join(output_dir, 'metadata.p')
metadata = load_pickle(metadata_file)
episode_name = metadata['episode_name']
episode_idx = metadata['episode_idx']
camera_name = metadata['camera_name']
episode = multi_episode_dict[episode_name]
data = episode.get_image_data(camera_name, episode_idx)
rgb = data['rgb']
uv = metadata['indices']
print("uv.shape", uv.shape)
color = [0, 255, 0]
draw_reticles(rgb, uv[:, 0], uv[:, 1], label_color=color)
save_file = os.path.join(output_dir, 'sampled_descriptors.png')
plt.figure()
plt.imshow(rgb)
plt.savefig(save_file)
if visualize:
plt.show()
# visualize spatially separated descriptors
if True:
metadata_file = os.path.join(output_dir, 'metadata.p')
metadata = load_pickle(metadata_file)
spatial_descriptor_file = os.path.join(output_dir, 'spatial_descriptors.p')
spatial_descriptors_data = load_pickle(spatial_descriptor_file)
des_idx = spatial_descriptors_data['spatial_descriptors_idx']
episode_name = metadata['episode_name']
episode_idx = metadata['episode_idx']
camera_name = metadata['camera_name']
episode = multi_episode_dict[episode_name]
data = episode.get_image_data(camera_name, episode_idx)
rgb = data['rgb']
uv = metadata['indices']
print("uv.shape", uv.shape)
color = [0, 255, 0]
draw_reticles(rgb, uv[des_idx, 0], uv[des_idx, 1], label_color=color)
save_file = os.path.join(output_dir, 'spatially_separated_descriptors.png')
plt.figure()
plt.imshow(rgb)
plt.savefig(save_file)
if visualize:
plt.show()
if True:
metadata_file = os.path.join(output_dir, 'metadata.p')
metadata = load_pickle(metadata_file)
print("\n\n---------Precomputing Descriptor Keypoints-----------")
descriptor_keypoints_output_dir = os.path.join(output_dir, "descriptor_keypoints")
precompute_descriptor_keypoints(multi_episode_dict,
model,
descriptor_keypoints_output_dir,
ref_descriptors_metadata=metadata,
batch_size=8,
num_workers=20,
camera_names=[camera_name]
)
print("Data saved at: ", output_dir)
print("Finished Normally")
import os
from key_dynam.utils.utils import load_pickle
from key_dynam.dataset.drake_sim_episode_reader import DrakeSimEpisodeReader
data_file = "/home/manuelli/data/key_dynam/dev/experiments/18/data/dps_box_on_side_600/2020-05-13-21-53-45-823302_idx_0.p"
non_image_data = load_pickle(data_file)
episode = DrakeSimEpisodeReader(non_image_data)
data = episode.get_data(0)
print("data.keys()", data.keys())
print("data['observation']", data['observation'].keys())
def DD_3D_dynamics(dataset_name):
from key_dynam.training.train_dynamics_pusher_slider_precomputed_keypoints import train_dynamics
from key_dynam.experiments.drake_pusher_slider import DD_utils
def load_config():
return load_yaml(
os.path.join(
get_project_root(),
'experiments/exp_22_push_box_hardware/config_DD_3D.yaml'))
dataset_paths = get_dataset_paths(dataset_name)
dataset_root = dataset_paths['dataset_root']
episodes_config = dataset_paths['episodes_config']
precomputed_vision_data_root = DD_utils.get_precomputed_data_root(
dataset_name)['precomputed_data_root']
# descriptor_keypoints_root = os.path.join(precomputed_vision_data_root, 'descriptor_keypoints')
descriptor_keypoints_root = os.path.join(precomputed_vision_data_root,
'descriptor_keypoints')
config = load_config()
multi_episode_dict = DynamicSpartanEpisodeReader.load_dataset(
config=config,
episodes_config=episodes_config,
episodes_root=dataset_paths['dataset_root'],
load_image_episode=True,
precomputed_data_root=descriptor_keypoints_root,
max_num_episodes=None)
experiment_save_root = get_experiment_save_root(dataset_name)
# experiment_save_root = os.path.join(get_data_root(), 'sandbox')
# standard
if True:
TRAIN = True
spatial_descriptor_data = load_pickle(
os.path.join(precomputed_vision_data_root,
'spatial_descriptors.p'))
metadata = load_pickle(
os.path.join(precomputed_vision_data_root, 'metadata.p'))
train_dir = None
config = load_config()
config['train']['n_history'] = 2
config['train']['random_seed'] = 1
config['dataset']['visual_observation_function'][
'camera_name'] = dataset_paths['main_camera_name']
config['dataset'][
'precomputed_data_root'] = precomputed_vision_data_root
suffix = "_DD_3D_n_his_2"
model_name = get_current_YYYY_MM_DD_hh_mm_ss_ms() + suffix
if TRAIN:
train_dir = os.path.join(experiment_save_root,
'trained_models/dynamics/DD_3D/',
model_name)
os.makedirs(train_dir)
visual_observation_function = PrecomputedVisualObservationFunctionFactory.function_from_config(
config,
keypoint_idx=spatial_descriptor_data['spatial_descriptors_idx']
)
train_dynamics(
config=config,
train_dir=train_dir,
multi_episode_dict=multi_episode_dict,
visual_observation_function=visual_observation_function,
metadata=metadata,
spatial_descriptors_data=spatial_descriptor_data)
# data_aug
if False:
TRAIN = True
spatial_descriptor_data = load_pickle(
os.path.join(precomputed_vision_data_root,
'spatial_descriptors.p'))
metadata = load_pickle(
os.path.join(precomputed_vision_data_root, 'metadata.p'))
train_dir = None
config = load_config()
config['train']['n_history'] = 2
config['train']['random_seed'] = 1
config['dataset']['visual_observation_function'][
'camera_name'] = dataset_paths['main_camera_name']
config['dataset']['data_augmentation']['enabled'] = True
config['dataset'][
'precomputed_data_root'] = precomputed_vision_data_root
suffix = "_DD_3D_n_his_2_T_aug"
model_name = get_current_YYYY_MM_DD_hh_mm_ss_ms() + suffix
if TRAIN:
train_dir = os.path.join(experiment_save_root,
'trained_models/dynamics/DD_3D/',
model_name)
os.makedirs(train_dir)
visual_observation_function = PrecomputedVisualObservationFunctionFactory.function_from_config(
config,
keypoint_idx=spatial_descriptor_data['spatial_descriptors_idx']
)
train_dynamics(
config=config,
train_dir=train_dir,
multi_episode_dict=multi_episode_dict,
visual_observation_function=visual_observation_function,
metadata=metadata,
spatial_descriptors_data=spatial_descriptor_data)
