def score_and_select_spatially_separated_keypoints(
metadata, # metadata.p,
confidence_score_data, # data.p file
K, # number of reference descriptors
position_diff_threshold, # threshold in pixels
output_dir,
visualize=False,
multi_episode_dict=None, # needed if you want to visualize
):
"""
Scores keypoints according to their confidence.
Selects the top keypoints that are "spatially separated"
Saves out a file 'spatial_descriptors.p' that records this data
"""
data = confidence_score_data
heatmap_values = data['heatmap_values']
scoring_func = create_scoring_function(gamma=3)
score_data = score_heatmap_values(heatmap_values,
scoring_func=scoring_func)
sorted_idx = score_data['sorted_idx']
keypoint_idx = select_spatially_separated_keypoints(
sorted_idx,
metadata['indices'],
position_diff_threshold=position_diff_threshold,
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(output_dir, 'spatial_descriptors.p'))
def save_data(
self,
save_dir=None,
):
"""
Saves data from the
- PlanContainer
- OnlineEpisode
"""
if save_dir is None:
save_dir = os.path.join(
utils.get_data_root(),
'hardware_experiments/closed_loop_rollouts/sandbox',
utils.get_current_YYYY_MM_DD_hh_mm_ss_ms())
if not os.path.exists(save_dir):
os.makedirs(save_dir)
print("saving MPC rollout data at: %s" % (save_dir))
save_data = {
'episode': self._state_dict['episode'].get_save_data(),
'plan': self._state_dict['plan'].get_save_data(),
}
save_file = os.path.join(save_dir, "data.p")
utils.save_pickle(save_data, save_file)
print("done saving data")
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 train_dynamics(
config,
train_dir, # str: directory to save output
multi_episode_dict=None,
spatial_descriptors_idx=None,
metadata=None,
spatial_descriptors_data=None,
):
assert multi_episode_dict is not None
# assert spatial_descriptors_idx is not None
# 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"))
if metadata is not None:
save_pickle(metadata, os.path.join(train_dir, 'metadata.p'))
if spatial_descriptors_data is not None:
save_pickle(spatial_descriptors_data,
os.path.join(train_dir, 'spatial_descriptors.p'))
training_stats = dict()
training_stats_file = os.path.join(train_dir, 'training_stats.yaml')
# load the data
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 . . .
'''
Build model for dynamics prediction
'''
model_dy = build_dynamics_model(config)
camera_name = config['vision_net']['camera_name']
# criterion
criterionMSE = nn.MSELoss()
l1Loss = nn.L1Loss()
smoothL1 = nn.SmoothL1Loss()
# 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 = None
if config['train']['lr_scheduler']['enabled']:
if config['train']['lr_scheduler']['type'] == "ReduceLROnPlateau":
scheduler = ReduceLROnPlateau(optimizer,
mode='min',
factor=sc['factor'],
patience=sc['patience'],
threshold_mode=sc['threshold_mode'],
cooldown=sc['cooldown'],
verbose=True)
elif config['train']['lr_scheduler']['type'] == "StepLR":
step_size = config['train']['lr_scheduler']['step_size']
gamma = config['train']['lr_scheduler']['gamma']
scheduler = StepLR(optimizer, step_size=step_size, gamma=gamma)
else:
raise ValueError("unknown scheduler type: %s" %
(config['train']['lr_scheduler']['type']))
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
valid_loss_type = config['train']['valid_loss_type']
global_iteration = 0
counters = {'train': 0, 'valid': 0}
epoch_counter_external = 0
loss = 0
index_map = get_object_and_robot_state_indices(config)
object_state_indices = torch.LongTensor(index_map['object_indices'])
robot_state_indices = torch.LongTensor(index_map['robot_indices'])
object_state_shape = config['dataset']['object_state_shape']
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:
# only validate at a certain frequency
if (phase == "valid") and (
(epoch % config['train']['valid_frequency']) != 0):
continue
model_dy.train(phase == 'train')
average_meter_container = dict()
step_duration_meter = AverageMeter()
# bar = ProgressBar(max_value=data_n_batches[phase])
loader = dataloaders[phase]
for i, data in enumerate(loader):
loss_container = dict() # store the losses for this step
step_start_time = time.time()
global_iteration += 1
counters[phase] += 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))
print("n_samples", n_samples)
# [B, n_samples, obs_dim]
observations = data['observations']
visual_observations_list = data[
'visual_observations_list']
# [B, n_samples, action_dim]
actions = data['actions']
B = actions.shape[0]
if use_gpu:
observations = observations.cuda()
actions = actions.cuda()
# compile the visual observations
# 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
# [B, -1, 3]
keypoints = visual_obs[camera_name][
'descriptor_keypoints_3d_world_frame']
# [B, K, 3] where K = len(spatial_descriptors_idx)
keypoints = keypoints[:, spatial_descriptors_idx]
B, K, _ = keypoints.shape
# [B, K*3]
keypoints_reshape = keypoints.reshape([B, K * 3])
if DEBUG:
print("keypoints.shape", keypoints.shape)
print("keypoints_reshape.shape",
keypoints_reshape.shape)
visual_model_output_list.append(keypoints_reshape)
visual_model_output = None
if len(visual_model_output_list) > 0:
# concatenate this into a tensor
# [B, n_samples, vision_model_out_dim]
visual_model_output = torch.stack(
visual_model_output_list, dim=1)
else:
visual_model_output = torch.Tensor(
) # empty tensor
# states, actions = data
assert actions.shape[1] == n_samples
# cast this to float so it can be concatenated below
visual_model_output = visual_model_output.type_as(
observations)
# we don't have any visual observations, so states are observations
# 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]
# [B, n_his, state_dim]
state_init = states[:, :n_his]
# We want to rollout n_roll steps
# actions = [B, n_his + n_roll, -1]
# so we want action_seq.shape = [B, n_roll, -1]
action_start_idx = 0
action_end_idx = n_his + n_roll - 1
action_seq = actions[:, action_start_idx:
action_end_idx, :]
if DEBUG:
print("states.shape", states.shape)
print("state_init.shape", state_init.shape)
print("actions.shape", actions.shape)
print("action_seq.shape", action_seq.shape)
# try using models_dy.rollout_model instead of doing this manually
rollout_data = rollout_model(state_init=state_init,
action_seq=action_seq,
dynamics_net=model_dy,
compute_debug_data=False)
# [B, n_roll, state_dim]
state_rollout_pred = rollout_data['state_pred']
# [B, n_roll, state_dim]
state_rollout_gt = states[:, n_his:]
if DEBUG:
print("state_rollout_gt.shape",
state_rollout_gt.shape)
print("state_rollout_pred.shape",
state_rollout_pred.shape)
# the loss function is between
# [B, n_roll, state_dim]
state_pred_err = state_rollout_pred - state_rollout_gt
# [B, n_roll, object_state_dim]
object_state_err = state_pred_err[:, :,
object_state_indices]
B, n_roll, object_state_dim = object_state_err.shape
# [B, n_roll, *object_state_shape]
object_state_err_reshape = object_state_err.reshape(
[B, n_roll, *object_state_shape])
# num weights
J = object_state_err_reshape.shape[2]
weights = model_dy.weight_matrix
assert len(
weights) == J, "len(weights) = %d, but J = %d" % (
len(weights), J)
# loss mse object, note the use of broadcasting semantics
# [B, n_roll]
object_state_loss_mse = weights * torch.pow(
object_state_err_reshape, 2).sum(dim=-1)
object_state_loss_mse = object_state_loss_mse.mean()
l2_object = (weights * torch.norm(
object_state_err_reshape, dim=-1)).mean()
l2_object_final_step = (weights * torch.norm(
object_state_err_reshape[:, -1], dim=-1)).mean()
# [B, n_roll, robot_state_dim]
robot_state_err = state_pred_err[:, :,
robot_state_indices]
robot_state_loss_mse = torch.pow(robot_state_err,
2).sum(dim=-1).mean()
loss_container[
'object_state_loss_mse'] = object_state_loss_mse
loss_container[
'robot_state_loss_mse'] = robot_state_loss_mse
loss_container['l2_object'] = l2_object
loss_container[
'l2_object_final_step'] = l2_object_final_step
# total loss
loss = object_state_loss_mse + robot_state_loss_mse
loss_container['loss'] = loss
for key, val in loss_container.items():
if not key in average_meter_container:
average_meter_container[key] = AverageMeter()
average_meter_container[key].update(val.item(), B)
step_duration_meter.update(time.time() - step_start_time)
if phase == 'train':
optimizer.zero_grad()
loss.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 += ', l2: %.6f' % (loss_container['l2'].item())
# log += ', l2_final_step: %.6f' %(loss_container['l2_final_step'].item())
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_train/%s" % (phase),
loss.item(), global_iteration)
for loss_type, loss_obj in loss_container.items():
plot_name = "Loss/%s/%s" % (loss_type, phase)
writer.add_scalar(plot_name, loss_obj.item(),
counters[phase])
# only plot the weights if we are in the train phase . . . .
if phase == "train":
for i in range(len(weights)):
plot_name = "Weights/%d" % (i)
writer.add_scalar(plot_name,
weights[i].item(),
counters[phase])
if phase == 'train' and global_iteration % 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'],
average_meter_container[valid_loss_type].avg,
best_valid_loss)
print(log)
# record all average_meter losses
for key, meter in average_meter_container.items():
writer.add_scalar("AvgMeter/%s/%s" % (key, phase),
meter.avg, epoch)
if phase == "train":
if (scheduler is not None) and (
config['train']['lr_scheduler']['type']
== "StepLR"):
scheduler.step()
if phase == 'valid':
if (scheduler is not None) and (
config['train']['lr_scheduler']['type']
== "ReduceLROnPlateau"):
scheduler.step(
average_meter_container[valid_loss_type].avg)
if average_meter_container[
valid_loss_type].avg < best_valid_loss:
best_valid_loss = average_meter_container[
valid_loss_type].avg
training_stats['epoch'] = epoch
training_stats['global_iteration'] = counters['valid']
save_yaml(training_stats, training_stats_file)
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 train_dynamics(
config,
train_dir, # str: directory to save output
multi_episode_dict=None,
visual_observation_function=None,
metadata=None,
spatial_descriptors_data=None,
):
assert multi_episode_dict is not None
# assert spatial_descriptors_idx is not None
# 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"))
if metadata is not None:
save_pickle(metadata, os.path.join(train_dir, 'metadata.p'))
if spatial_descriptors_data is not None:
save_pickle(spatial_descriptors_data,
os.path.join(train_dir, 'spatial_descriptors.p'))
training_stats = dict()
training_stats_file = os.path.join(train_dir, 'training_stats.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,
visual_observation_function=visual_observation_function)
print("len(datasets[phase])", len(datasets[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 . . .
if False:
dataset = datasets["train"]
data = dataset[0]
print("data['observations_combined'].shape",
data['observations_combined'].shape)
print("data.keys()", data.keys())
print("data['observations_combined']",
data['observations_combined'][0])
print("data['observations_combined'].shape",
data['observations_combined'].shape)
print("data['actions'].shape", data['actions'].shape)
print("data['actions']\n", data['actions'])
quit()
'''
Build model for dynamics prediction
'''
model_dy = build_dynamics_model(config)
if config['dynamics_net'] == "mlp_weight_matrix":
raise ValueError("can't use weight matrix with standard setup")
# criterion
criterionMSE = nn.MSELoss()
l1Loss = nn.L1Loss()
smoothL1 = nn.SmoothL1Loss()
# 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 = None
if config['train']['lr_scheduler']['enabled']:
if config['train']['lr_scheduler']['type'] == "ReduceLROnPlateau":
scheduler = ReduceLROnPlateau(optimizer,
mode='min',
factor=sc['factor'],
patience=sc['patience'],
threshold_mode=sc['threshold_mode'],
cooldown=sc['cooldown'],
verbose=True)
elif config['train']['lr_scheduler']['type'] == "StepLR":
step_size = config['train']['lr_scheduler']['step_size']
gamma = config['train']['lr_scheduler']['gamma']
scheduler = StepLR(optimizer, step_size=step_size, gamma=gamma)
else:
raise ValueError("unknown scheduler type: %s" %
(config['train']['lr_scheduler']['type']))
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
valid_loss_type = config['train']['valid_loss_type']
global_iteration = 0
counters = {'train': 0, 'valid': 0}
epoch_counter_external = 0
loss = 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:
# only validate at a certain frequency
if (phase == "valid") and (
(epoch % config['train']['valid_frequency']) != 0):
continue
model_dy.train(phase == 'train')
average_meter_container = dict()
step_duration_meter = AverageMeter()
# bar = ProgressBar(max_value=data_n_batches[phase])
loader = dataloaders[phase]
for i, data in enumerate(loader):
loss_container = dict() # store the losses for this step
step_start_time = time.time()
global_iteration += 1
counters[phase] += 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))
print("n_samples", n_samples)
# [B, n_samples, obs_dim]
states = data['observations_combined']
# [B, n_samples, action_dim]
actions = data['actions']
B = actions.shape[0]
if use_gpu:
states = states.cuda()
actions = actions.cuda()
# state_cur: B x n_his x state_dim
# state_cur = states[:, :n_his]
# [B, n_his, state_dim]
state_init = states[:, :n_his]
# We want to rollout n_roll steps
# actions = [B, n_his + n_roll, -1]
# so we want action_seq.shape = [B, n_roll, -1]
action_start_idx = 0
action_end_idx = n_his + n_roll - 1
action_seq = actions[:, action_start_idx:
action_end_idx, :]
if DEBUG:
print("states.shape", states.shape)
print("state_init.shape", state_init.shape)
print("actions.shape", actions.shape)
print("action_seq.shape", action_seq.shape)
# try using models_dy.rollout_model instead of doing this manually
rollout_data = rollout_model(state_init=state_init,
action_seq=action_seq,
dynamics_net=model_dy,
compute_debug_data=False)
# [B, n_roll, state_dim]
state_rollout_pred = rollout_data['state_pred']
# [B, n_roll, state_dim]
state_rollout_gt = states[:, n_his:]
if DEBUG:
print("state_rollout_gt.shape",
state_rollout_gt.shape)
print("state_rollout_pred.shape",
state_rollout_pred.shape)
# the loss function is between
# [B, n_roll, state_dim]
state_pred_err = state_rollout_pred - state_rollout_gt
# everything is in 3D space now so no need to do any scaling
# all the losses would be in meters . . . .
loss_mse = criterionMSE(state_rollout_pred,
state_rollout_gt)
loss_l1 = l1Loss(state_rollout_pred, state_rollout_gt)
loss_l2 = torch.norm(state_pred_err, dim=-1).mean()
loss_smoothl1 = smoothL1(state_rollout_pred,
state_rollout_gt)
loss_smoothl1_final_step = smoothL1(
state_rollout_pred[:, -1], state_rollout_gt[:, -1])
# compute losses at final step of the rollout
mse_final_step = criterionMSE(
state_rollout_pred[:, -1], state_rollout_gt[:, -1])
l2_final_step = torch.norm(state_pred_err[:, -1],
dim=-1).mean()
l1_final_step = l1Loss(state_rollout_pred[:, -1],
state_rollout_gt[:, -1])
loss_container['mse'] = loss_mse
loss_container['l1'] = loss_l1
loss_container['mse_final_step'] = mse_final_step
loss_container['l1_final_step'] = l1_final_step
loss_container['l2_final_step'] = l2_final_step
loss_container['l2'] = loss_l2
loss_container['smooth_l1'] = loss_smoothl1
loss_container[
'smooth_l1_final_step'] = loss_smoothl1_final_step
# compute the loss
loss = 0
for key, val in config['loss_function'].items():
if val['enabled']:
loss += loss_container[key] * val['weight']
loss_container['loss'] = loss
for key, val in loss_container.items():
if not key in average_meter_container:
average_meter_container[key] = AverageMeter()
average_meter_container[key].update(val.item(), B)
step_duration_meter.update(time.time() - step_start_time)
if phase == 'train':
optimizer.zero_grad()
loss.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 += ', l2: %.6f' % (loss_container['l2'].item())
log += ', l2_final_step: %.6f' % (
loss_container['l2_final_step'].item())
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_train/%s" % (phase),
loss.item(), global_iteration)
for loss_type, loss_obj in loss_container.items():
plot_name = "Loss/%s/%s" % (loss_type, phase)
writer.add_scalar(plot_name, loss_obj.item(),
counters[phase])
if phase == 'train' and global_iteration % 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'],
average_meter_container[valid_loss_type].avg,
best_valid_loss)
print(log)
# record all average_meter losses
for key, meter in average_meter_container.items():
writer.add_scalar("AvgMeter/%s/%s" % (key, phase),
meter.avg, epoch)
if phase == "train":
if (scheduler is not None) and (
config['train']['lr_scheduler']['type']
== "StepLR"):
scheduler.step()
if phase == 'valid':
if (scheduler is not None) and (
config['train']['lr_scheduler']['type']
== "ReduceLROnPlateau"):
scheduler.step(
average_meter_container[valid_loss_type].avg)
if average_meter_container[
valid_loss_type].avg < best_valid_loss:
best_valid_loss = average_meter_container[
valid_loss_type].avg
training_stats['epoch'] = epoch
training_stats['global_iteration'] = counters['valid']
save_yaml(training_stats, training_stats_file)
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 precompute_transporter_keypoints(
multi_episode_dict,
model_kp,
output_dir, # str
batch_size=10,
num_workers=10,
camera_names=None,
model_file=None,
):
assert model_file is not None
metadata = dict()
metadata['model_file'] = model_file
save_yaml(metadata, os.path.join(output_dir, 'metadata.yaml'))
start_time = time.time()
log_freq = 10
device = next(model_kp.parameters()).device
model_kp = model_kp.eval() # make sure model is in eval mode
image_data_config = {
'rgb': True,
'mask': True,
'depth_int16': True,
}
# build all the dataset
datasets = {}
dataloaders = {}
for episode_name, episode in multi_episode_dict.items():
single_episode_dict = {episode_name: episode}
config = model_kp.config
# need to do this since transporter type data sampling only works
# with tuple_size = 1
dataset_config = copy.deepcopy(config)
dataset_config['dataset']['use_transporter_type_data_sampling'] = False
datasets[episode_name] = ImageTupleDataset(
dataset_config,
single_episode_dict,
phase="all",
image_data_config=image_data_config,
tuple_size=1,
compute_K_inv=True,
camera_names=camera_names)
dataloaders[episode_name] = DataLoader(datasets[episode_name],
batch_size=batch_size,
num_workers=num_workers,
shuffle=False)
episode_counter = 0
num_episodes = len(multi_episode_dict)
for episode_name, dataset in datasets.items():
episode_counter += 1
print("\n\n")
episode = multi_episode_dict[episode_name]
hdf5_file = None
try:
hdf5_file = os.path.basename(episode.image_data_file)
except AttributeError:
hdf5_file = "%s.h5" % (episode.name)
hdf5_file_fullpath = os.path.join(output_dir, hdf5_file)
str_split = hdf5_file_fullpath.split(".")
assert len(str_split) == 2
pickle_file_fullpath = str_split[0] + ".p"
# print("episode_name", episode_name)
# print("hdf5_file_fullpath", hdf5_file_fullpath)
# print("pickle_file_fullpath", pickle_file_fullpath)
if os.path.isfile(hdf5_file_fullpath):
os.remove(hdf5_file_fullpath)
if os.path.isfile(pickle_file_fullpath):
os.remove(pickle_file_fullpath)
episode_keypoint_data = dict()
episode_start_time = time.time()
with h5py.File(hdf5_file_fullpath, 'w') as hf:
for i, data in enumerate(dataloaders[episode_name]):
data = data[0]
rgb_crop_tensor = data['rgb_crop_tensor'].to(device)
crop_params = data['crop_param']
depth_int16 = data['depth_int16']
key_tree_joined = data['key_tree_joined']
# print("\n\n i = %d, idx = %d, camera_name = %s" %(i, data['idx'], data['camera_name']))
depth = depth_int16.float() * 1.0 / DEPTH_IM_SCALE
if (i % log_freq) == 0:
log_msg = "computing [%d/%d][%d/%d]" % (
episode_counter, num_episodes, i + 1,
len(dataloaders[episode_name]))
print(log_msg)
B = rgb_crop_tensor.shape[0]
_, H, W, _ = data['rgb'].shape
kp_pred = None
kp_pred_full_pixels = None
with torch.no_grad():
kp_pred = model_kp.predict_keypoint(rgb_crop_tensor)
# [B, n_kp, 2]
kp_pred_full_pixels = transporter_utils.map_cropped_pixels_to_full_pixels_torch(
kp_pred, crop_params)
xy = kp_pred_full_pixels.clone()
xy[:, :, 0] = (xy[:, :, 0]) * 2.0 / W - 1.0
xy[:, :, 1] = (xy[:, :, 1]) * 2.0 / H - 1.0
# debug
# print("xy[0,0]", xy[0,0])
# get depth values
kp_pred_full_pixels_int = kp_pred_full_pixels.type(
torch.LongTensor)
z = pdc_utils.index_into_batch_image_tensor(
depth.unsqueeze(1),
kp_pred_full_pixels_int.transpose(1, 2))
z = z.squeeze(1)
K_inv = data['K_inv']
pts_camera_frame = pdc_torch_utils.pinhole_unprojection(
kp_pred_full_pixels, z, K_inv)
# print("pts_camera_frame.shape", pts_camera_frame.shape)
pts_world_frame = pdc_torch_utils.transform_points_3D(
data['T_W_C'], pts_camera_frame)
# print("pts_world_frame.shape", pts_world_frame.shape)
for j in range(B):
keypoint_data = {}
# this goes from [-1,1]
keypoint_data['xy'] = torch_utils.cast_to_numpy(xy[j])
keypoint_data['uv'] = torch_utils.cast_to_numpy(
kp_pred_full_pixels[j])
keypoint_data['uv_int'] = torch_utils.cast_to_numpy(
kp_pred_full_pixels_int[j])
keypoint_data['z'] = torch_utils.cast_to_numpy(z[j])
keypoint_data[
'pos_world_frame'] = torch_utils.cast_to_numpy(
pts_world_frame[j])
keypoint_data[
'pos_camera_frame'] = torch_utils.cast_to_numpy(
pts_camera_frame[j])
# save out some data in both hdf5 and pickle format
for key, val in keypoint_data.items():
save_key = key_tree_joined[
j] + "/transporter_keypoints/%s" % (key)
hf.create_dataset(save_key, data=val)
episode_keypoint_data[save_key] = val
save_pickle(episode_keypoint_data, pickle_file_fullpath)
print("duration: %.3f seconds" %
(time.time() - episode_start_time))
def main():
d = load_model_and_data()
model_dy = d['model_dy']
dataset = d['dataset']
config = d['config']
multi_episode_dict = d['multi_episode_dict']
planner = d['planner']
planner_config = planner.config
idx_dict = get_object_and_robot_state_indices(config)
object_indices = idx_dict['object_indices']
robot_indices = idx_dict['robot_indices']
n_his = config['train']['n_history']
# save_dir = os.path.join(get_project_root(), 'sandbox/mpc/', get_current_YYYY_MM_DD_hh_mm_ss_ms())
save_dir = os.path.join(get_project_root(),
'sandbox/mpc/push_right_box_horizontal')
print("save_dir", save_dir)
if not os.path.exists(save_dir):
os.makedirs(save_dir)
# rotate
# episode_names = dataset.get_episode_names()
# print("len(episode_names)", len(episode_names))
# episode_name = episode_names[0]
# start_idx = 1
# n_roll = 15
# # straight + rotate
# episode_name = "2020-06-29-21-04-16"
# print('episode_name', episode_name)
# start_idx = 1
# n_roll = 15
# this is a nice straight push . . .
# push with box in horizontal position
episode_name = "2020-06-29-22-03-45"
start_idx = 2
n_roll = 10
# # validation set episodes
# episode_names = dataset.get_episode_names()
# print("len(episode_names)", len(episode_names))
# episode_name = episode_names[1]
# start_idx = 2
# n_roll = 15
camera_name = "d415_01"
episode = multi_episode_dict[episode_name]
print("episode_name", episode_name)
vis = meshcat_utils.make_default_visualizer_object()
vis.delete()
idx_list = list(range(start_idx, start_idx + n_roll + 1))
idx_list_GT = idx_list
goal_idx = idx_list[-1]
print("idx_list", idx_list)
# visualize ground truth rollout
if True:
for display_idx, episode_idx in enumerate(idx_list):
visualize_episode_data_single_timestep(
vis=vis,
dataset=dataset,
episode=episode,
camera_name=camera_name,
episode_idx=episode_idx,
display_idx=episode_idx,
)
data_goal = dataset._getitem(episode,
goal_idx,
rollout_length=1,
n_history=1)
states_goal = data_goal['observations_combined'][0]
z_states_goal = model_dy.compute_z_state(states_goal)['z']
print("states_goal.shape", states_goal.shape)
print("z_states_goal.shape", z_states_goal.shape)
##### VISUALIZE PREDICTED ROLLOUT ##########
data = dataset._getitem(episode, start_idx, rollout_length=n_roll)
states = data['observations_combined'].unsqueeze(0)
z = model_dy.compute_z_state(states)['z']
actions = data['actions'].unsqueeze(0)
idx_range_model_dy_input = data['idx_range']
print("data.keys()", data.keys())
print("data['idx_range']", data['idx_range'])
# z_init
z_init = z[:, :n_his]
# actions_init
action_start_idx = 0
action_end_idx = n_his + n_roll - 1
action_seq = actions[:, action_start_idx:action_end_idx]
print("action_seq GT\n", action_seq)
with torch.no_grad():
rollout_data = rollout_model(state_init=z_init.cuda(),
action_seq=action_seq.cuda(),
dynamics_net=model_dy,
compute_debug_data=False)
# [B, n_roll, state_dim]
# state_rollout_pred = rollout_data['state_pred']
z_rollout_pred = rollout_data['state_pred'].squeeze()
print("z_rollout_pred.shape", z_rollout_pred.shape)
if True:
for idx in range(len(z_rollout_pred)):
display_idx = data['idx_range'][idx + n_his]
visualize_model_prediction_single_timestep(
vis,
config,
z_pred=z_rollout_pred[idx],
display_idx=display_idx)
print("z_rollout_pred.shape", z_rollout_pred.shape)
# compute loss when rolled out using GT action sequence
eval_indices = object_indices
obs_goal = z_states_goal[object_indices].cuda()
reward_data = planner_utils.evaluate_model_rollout(
state_pred=rollout_data['state_pred'],
obs_goal=obs_goal,
eval_indices=eval_indices,
terminal_cost_only=planner_config['mpc']['mppi']['terminal_cost_only'],
p=planner_config['mpc']['mppi']['cost_norm'])
print("reward_data using action_seq_GT\n", reward_data['reward'])
##### MPC ##########
data = dataset._getitem(episode,
start_idx,
rollout_length=0,
n_history=config['train']['n_history'])
state_cur = data['observations_combined'].cuda()
z_state_cur = model_dy.compute_z_state(state_cur)['z']
action_his = data['actions'][:(n_his - 1)].cuda()
print("z_state_cur.shape", state_cur.shape)
print("action_his.shape", action_his.shape)
# don't seed with nominal actions just yet
action_seq_rollout_init = None
set_seed(SEED)
mpc_out = planner.trajectory_optimization(
state_cur=z_state_cur,
action_his=action_his,
obs_goal=obs_goal,
model_dy=model_dy,
action_seq_rollout_init=action_seq_rollout_init,
n_look_ahead=n_roll,
eval_indices=object_indices,
rollout_best_action_sequence=True,
verbose=True,
add_grid_action_samples=True,
)
print("\n\n------MPC output-------\n\n")
print("action_seq:\n", mpc_out['action_seq'])
mpc_state_pred = mpc_out['state_pred']
# current shape is [n_roll + 1, state_dim] but really should be
# [n_roll, state_dim] . . . something is up
print("mpc_state_pred.shape", mpc_state_pred.shape)
print("mpc_out['action_seq'].shape", mpc_out['action_seq'].shape)
print("n_roll", n_roll)
# visualize
for idx in range(n_roll):
episode_idx = start_idx + idx + 1
visualize_model_prediction_single_timestep(vis,
config,
z_pred=mpc_state_pred[idx],
display_idx=episode_idx,
name_prefix="mpc",
color=[255, 0, 0])
######## MPC w/ dynamics model input builder #############
print("\n\n-----DynamicsModelInputBuilder-----")
# dynamics model input builder
online_episode = OnlineEpisodeReader(no_copy=True)
ref_descriptors = d['spatial_descriptor_data']['spatial_descriptors']
ref_descriptors = torch_utils.cast_to_torch(ref_descriptors).cuda()
K_matrix = episode.image_episode.camera_K_matrix(camera_name)
T_world_camera = episode.image_episode.camera_pose(camera_name, 0)
visual_observation_function = \
VisualObservationFunctionFactory.descriptor_keypoints_3D(config=config,
camera_name=camera_name,
model_dd=d['model_dd'],
ref_descriptors=ref_descriptors,
K_matrix=K_matrix,
T_world_camera=T_world_camera,
)
input_builder = DynamicsModelInputBuilder(
observation_function=d['observation_function'],
visual_observation_function=visual_observation_function,
action_function=d['action_function'],
episode=online_episode)
compute_control_action_msg = dict()
compute_control_action_msg['type'] = "COMPUTE_CONTROL_ACTION"
for i in range(n_his):
episode_idx = idx_range_model_dy_input[i]
print("episode_idx", episode_idx)
# add image information to
data = add_images_to_episode_data(episode, episode_idx, camera_name)
online_episode.add_data(copy.deepcopy(data))
compute_control_action_msg['data'] = data
# hack for seeing how much the history matters .. .
# online_episode.add_data(copy.deepcopy(data))
# save informatin for running zmq controller
save_pickle(compute_control_action_msg,
os.path.join(save_dir, 'compute_control_action_msg.p'))
goal_idx = idx_list_GT[-1]
goal_data = add_images_to_episode_data(episode, goal_idx, camera_name)
goal_data['observations']['timestamp_system'] = time.time()
plan_msg = {
'type': "PLAN",
'data': [goal_data],
'n_roll': n_roll,
'K_matrix': K_matrix,
'T_world_camera': T_world_camera,
}
save_pickle(plan_msg, os.path.join(save_dir, "plan_msg.p"))
print("len(online_episode)", len(online_episode))
# use this to construct input
# verify it's the same as what we got from using the dataset directly
idx = online_episode.get_latest_idx()
mpc_input_data = input_builder.get_dynamics_model_input(idx,
n_history=n_his)
# print("mpc_input_data\n", mpc_input_data)
state_cur_ib = mpc_input_data['states'].cuda()
action_his_ib = mpc_input_data['actions'].cuda()
z_state_cur_ib = model_dy.compute_z_state(state_cur_ib)['z']
set_seed(SEED)
mpc_out = planner.trajectory_optimization(
state_cur=z_state_cur_ib,
action_his=action_his_ib,
obs_goal=obs_goal,
model_dy=model_dy,
action_seq_rollout_init=None,
n_look_ahead=n_roll,
eval_indices=object_indices,
rollout_best_action_sequence=True,
verbose=True,
add_grid_action_samples=True,
)
# visualize
for idx in range(n_roll):
episode_idx = start_idx + idx + 1
visualize_model_prediction_single_timestep(
vis,
config,
z_pred=mpc_out['state_pred'][idx],
display_idx=episode_idx,
name_prefix="mpc_input_builder",
color=[255, 255, 0])
def compute_descriptor_confidences(multi_episode_dict,
model,
output_dir, # str
batch_size=10,
num_workers=10,
model_file=None,
ref_descriptors=None,
episode_name_arg=None,
episode_idx=None,
camera_name=None,
num_ref_descriptors=None,
localization_type="spatial_expectation", # ['spatial_expectation', 'argmax']
num_batches=None,
):
"""
Computes confidence scores for different reference descriptors.
Saves two files
metadata.p: has information about reference descriptors, etc.
data.p: descriptor confidence scores
"""
if not os.path.isdir(output_dir):
os.makedirs(output_dir)
start_time = time.time()
log_freq = 10
device = next(model.parameters()).device
model.eval() # make sure model is in eval mode
# build all the dataset
config = None
dataset = ImageDataset(config, multi_episode_dict, phase="all")
dataloader = DataLoader(dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=True)
# sample ref descriptors
if ref_descriptors is None:
if episode_name_arg is None:
episode_list = list(multi_episode_dict.keys())
episode_list.sort()
episode_name_arg = episode_list[0]
if episode_idx is None:
episode_idx = 0
# can't be left blank
assert camera_name is not None
episode = multi_episode_dict[episode_name_arg]
data = dataset._getitem(episode,
episode_idx,
camera_name)
rgb_tensor = data['rgb_tensor'].unsqueeze(0).to(device)
# compute descriptor image from model
with torch.no_grad():
out = model.forward(rgb_tensor)
des_img = out['descriptor_image']
# get image mask and cast it to a tensor on the appropriate
# device
img_mask = episode.get_image(camera_name,
episode_idx,
type="mask")
img_mask = torch.tensor(img_mask).to(des_img.device)
ref_descriptors_dict = sample_descriptors(des_img.squeeze(), img_mask, num_ref_descriptors)
ref_descriptors = ref_descriptors_dict['descriptors']
ref_indices = ref_descriptors_dict['indices']
print("ref_descriptors_dict\n", ref_descriptors_dict)
# save metadata in dict
metadata = {'model_file': model_file,
'ref_descriptors': ref_descriptors.cpu().numpy(), # [N, D]
'indices': ref_indices.cpu().numpy(),
'episode_name': episode_name_arg,
'episode_idx': episode_idx,
'camera_name': camera_name}
metadata_file = os.path.join(output_dir, 'metadata.p')
save_pickle(metadata, metadata_file)
scores = dict()
heatmap_value_list = []
for i, data in enumerate(dataloader):
if num_batches is not None and i > num_batches:
break
rgb_tensor = data['rgb_tensor'].to(device)
key_tree_joined = data['key_tree_joined']
if (i % log_freq) == 0:
log_msg = "computing %d" % (i)
print(log_msg)
# don't use gradients
with torch.no_grad():
out = model.forward(rgb_tensor)
# [B, D, H, W]
des_img = out['descriptor_image']
B, _, H, W = rgb_tensor.shape
heatmap_values = None
if localization_type == "spatial_expectation":
sigma_descriptor_heatmap = 5 # default
try:
sigma_descriptor_heatmap = model.config['network']['sigma_descriptor_heatmap']
except:
pass
# print("ref_descriptors.shape", ref_descriptors.shape)
# print("des_img.shape", des_img.shape)
d = get_spatial_expectation(ref_descriptors,
des_img,
sigma=sigma_descriptor_heatmap,
type='exp',
compute_heatmap_values=True,
return_heatmap=True,
)
# [B, K]
heatmap_values = d['heatmap_values']
else:
raise ValueError("unknown localization type: %s" % (localization_type))
heatmap_value_list.append(heatmap_values)
heatmap_values_tensor = torch.cat(heatmap_value_list)
heatmap_values_np = heatmap_values_tensor.cpu().numpy()
save_data = {'heatmap_values': heatmap_values_np}
data_file = os.path.join(output_dir, 'data.p')
save_pickle(save_data, data_file)
print("total time to compute descriptors: %.3f seconds" % (time.time() - start_time))
def precompute_descriptor_keypoints(multi_episode_dict,
model,
output_dir, # str
ref_descriptors_metadata,
batch_size=10,
num_workers=10,
localization_type="spatial_expectation", # ['spatial_expectation', 'argmax']
compute_3D=True, # in world frame
camera_names=None,
):
if not os.path.isdir(output_dir):
os.makedirs(output_dir)
start_time = time.time()
log_freq = 10
device = next(model.parameters()).device
model = model.eval() # make sure model is in eval mode
# build all the dataset
datasets = {}
dataloaders = {}
for episode_name, episode in iteritems(multi_episode_dict):
single_episode_dict = {episode_name: episode}
config = None
datasets[episode_name] = ImageDataset(config,
single_episode_dict,
phase="all",
camera_names=camera_names)
dataloaders[episode_name] = DataLoader(datasets[episode_name],
batch_size=batch_size,
num_workers=num_workers,
shuffle=False)
# K = num_ref_descriptors
metadata = ref_descriptors_metadata
ref_descriptors = torch.Tensor(metadata['ref_descriptors'])
ref_descriptors = ref_descriptors.cuda()
K, _ = ref_descriptors.shape
metadata_file = os.path.join(output_dir, 'metadata.p')
save_pickle(metadata, metadata_file)
episode_counter = 0
num_episodes = len(multi_episode_dict)
for episode_name, dataset in iteritems(datasets):
episode_counter += 1
print("\n\n")
episode = multi_episode_dict[episode_name]
hdf5_file = None
try:
hdf5_file = os.path.basename(episode.image_data_file)
except AttributeError:
hdf5_file = "%s.h5" % (episode.name)
hdf5_file_fullpath = os.path.join(output_dir, hdf5_file)
str_split = hdf5_file_fullpath.split(".")
assert len(str_split) == 2
pickle_file_fullpath = str_split[0] + ".p"
# print("hdf5_file_fullpath", hdf5_file_fullpath)
# print("pickle_file_fullpath", pickle_file_fullpath)
if os.path.isfile(hdf5_file_fullpath):
os.remove(hdf5_file_fullpath)
if os.path.isfile(pickle_file_fullpath):
os.remove(pickle_file_fullpath)
dataloader = dataloaders[episode_name]
episode_keypoint_data = dict()
episode_start_time = time.time()
with h5py.File(hdf5_file_fullpath, 'w') as hf:
for i, data in enumerate(dataloaders[episode_name]):
rgb_tensor = data['rgb_tensor'].to(device)
key_tree_joined = data['key_tree_joined']
if (i % log_freq) == 0:
log_msg = "computing [%d/%d][%d/%d]" % (episode_counter, num_episodes, i + 1, len(dataloader))
print(log_msg)
# don't use gradients
tmp_time = time.time()
with torch.no_grad():
out = model.forward(rgb_tensor)
# [B, D, H, W]
des_img = out['descriptor_image']
B, _, H, W = rgb_tensor.shape
# [B, N, 2]
batch_indices = None
preds_3d = None
if localization_type == "spatial_expectation":
sigma_descriptor_heatmap = 5 # default
try:
sigma_descriptor_heatmap = model.config['network']['sigma_descriptor_heatmap']
except:
pass
# print("ref_descriptors.shape", ref_descriptors.shape)
# print("des_img.shape", des_img.shape)
d = get_spatial_expectation(ref_descriptors,
des_img,
sigma=sigma_descriptor_heatmap,
type='exp',
return_heatmap=True,
)
batch_indices = d['uv']
# [B, K, H, W]
if compute_3D:
# [B*K, H, W]
heatmaps_no_batch = d['heatmap_no_batch']
# [B, H, W]
depth = data['depth_int16'].to(device)
# expand depth images and convert to meters, instead of mm
# [B, K, H, W]
depth_expand = depth.unsqueeze(1).expand([B, K, H, W]).reshape([B * K, H, W])
depth_expand = depth_expand.type(torch.FloatTensor) / constants.DEPTH_IM_SCALE
depth_expand = depth_expand.to(heatmaps_no_batch.device)
pred_3d = get_integral_preds_3d(heatmaps_no_batch,
depth_images=depth_expand,
compute_uv=True)
pred_3d['uv'] = pred_3d['uv'].reshape([B, K, 2])
pred_3d['xy'] = pred_3d['xy'].reshape([B, K, 2])
pred_3d['z'] = pred_3d['z'].reshape([B, K])
elif localization_type == "argmax":
# localize descriptors
best_match_dict = get_argmax_l2(ref_descriptors,
des_img)
# [B, N, 2]
# where N is num_ref_descriptors
batch_indices = best_match_dict['indices']
else:
raise ValueError("unknown localization type: %s" % (localization_type))
print("computing keypoints took", time.time() - tmp_time)
tmp_time = time.time()
# iterate over elements in the batch
for j in range(B):
keypoint_data = {} # dict that stores information to save out
# [N,2]
# indices = batch_indices[j].cpu().numpy()
# key = key_tree_joined[j] + "/descriptor_keypoints"
# hf.create_dataset(key, data=indices)
# keypoint_indices_dict[key] = indices
# stored 3D keypoint locations (in both camera and world frame)
if pred_3d is not None:
# key_3d_W = key_tree_joined[j] + "/descriptor_keypoints_3d_world_frame"
# key_3d_C = key_tree_joined[j] + "/descriptor_keypoints_3d_camera_frame"
# T_W_C = data['T_world_camera'][j].cpu().numpy()
# K_matrix = data['K'][j].cpu().numpy()
T_W_C = torch_utils.cast_to_numpy(data['T_world_camera'][j])
K_matrix = torch_utils.cast_to_numpy(data['K'][j])
uv = torch_utils.cast_to_numpy(pred_3d['uv'][j])
xy = torch_utils.cast_to_numpy(pred_3d['xy'][j])
z = torch_utils.cast_to_numpy(pred_3d['z'][j])
# [K, 3]
# this is in camera frame
pts_3d_C = pdc_utils.pinhole_unprojection(uv, z, K_matrix)
# hf.create_dataset(key_3d_C, data=pts_3d_C)
# keypoint_indices_dict[key_3d_C] = pts_3d_C
# project into world frame
pts_3d_W = transform_utils.transform_points_3D(transform=T_W_C,
points=pts_3d_C)
# hf.create_dataset(key_3d_W, data=pts_3d_W)
# keypoint_indices_dict[key_3d_W] = pts_3d_W
keypoint_data['xy'] = torch_utils.cast_to_numpy(xy)
keypoint_data['uv'] = torch_utils.cast_to_numpy(uv)
keypoint_data['z'] = torch_utils.cast_to_numpy(z)
keypoint_data['pos_world_frame'] = torch_utils.cast_to_numpy(pts_3d_W)
keypoint_data['pos_camera_frame'] = torch_utils.cast_to_numpy(pts_3d_C)
# save out some data in both hdf5 and pickle format
for key, val in keypoint_data.items():
save_key = key_tree_joined[j] + "/descriptor_keypoints/%s" % (key)
hf.create_dataset(save_key, data=val)
episode_keypoint_data[save_key] = val
print("saving to disk took", time.time() - tmp_time)
# save_pickle(keypoint_indices_dict, pickle_file_fullpath)
save_pickle(episode_keypoint_data, pickle_file_fullpath)
print("duration: %.3f seconds" % (time.time() - episode_start_time))
print("total time to compute descriptors: %.3f seconds" % (time.time() - start_time))