def generate_pretrain_imprint_data(envs, grasp_object_dataset, grasp_dataset):
env_pool = Pool(envs)
output_dir = f"{grasp_dataset.directory_path}grippers_imprint/"
if not exists(output_dir):
mkdir(output_dir)
pbar = tqdm(total=len(grasp_object_dataset),
desc='Generating pretrain imprint dataset')
loader = DataLoader(grasp_object_dataset,
batch_size=grasp_object_dataset.batch_size,
shuffle=False,
collate_fn=lambda batch: batch)
for grasp_objects in loader:
grasp_object_urdfs = dicts_get(grasp_objects, 'urdf_path')
gripper_urdfs = [
splitext(grasp_object_urdf)[0] + "_imprint.urdf"
for grasp_object_urdf in grasp_object_urdfs
]
left_finger_tsdf_paths = [
splitext(grasp_object_urdf)[0] + "_imprint_left_tsdf.npy"
for grasp_object_urdf in grasp_object_urdfs
]
right_finger_tsdf_paths = [
splitext(grasp_object_urdf)[0] + "_imprint_right_tsdf.npy"
for grasp_object_urdf in grasp_object_urdfs
]
grasp_object_tsdf_paths = dicts_get(grasp_objects, "tsdf_path")
update_grasp_dataset(grasp_results=env_pool.map(
exec_fn=lambda env, args: env.simulate_grasp.remote(*args),
iterable=zip(grasp_object_urdfs, gripper_urdfs)),
left_finger_tsdf_paths=left_finger_tsdf_paths,
right_finger_tsdf_paths=right_finger_tsdf_paths,
grasp_object_tsdf_paths=grasp_object_tsdf_paths,
grasp_dataset=grasp_dataset)
pbar.update(len(grasp_object_urdfs))
def create_fingers(self, grasp_objects):
left_fingers, right_fingers = torch.empty(
(0, 1, 40, 40, 40)), torch.empty((0, 1, 40, 40, 40))
grasp_object_tsdfs = dicts_get(grasp_objects, 'tsdf')
self.net.eval()
go_batch = grasp_object_tsdf.to(self.device)
lf_emb_batch = 2 * \
torch.rand((len(grasp_object_tsdfs),
self.net.embedding_dim), device=self.device) - 1
rf_emb_batch = 2 * \
torch.rand((len(grasp_object_tsdfs),
self.net.embedding_dim), device=self.device) - 1
lf_emb_batch.requires_grad_()
rf_emb_batch.requires_grad_()
finger_emb_opt = Adam([lf_emb_batch, rf_emb_batch], lr=self.lr)
for i in range(self.finetune_steps + 1):
finger_emb_opt.zero_grad()
if i != self.finetune_steps:
lf_batch = self.net.decoder(lf_emb_batch)
rf_batch = self.net.decoder(rf_emb_batch)
loss = -self.net.evaluate_fingers(go_batch, lf_batch, rf_batch)
if self.optim_grasp_success_only:
loss[0].backward()
else:
loss.mean().backward()
finger_emb_opt.step()
else:
left_fingers = cat((left_fingers, lf_batch.detach().cpu()),
dim=0)
right_fingers = cat((right_fingers, rf_batch.detach().cpu()),
dim=0)
return left_fingers, right_fingers
def create_fingers(self, grasp_objects):
left_fingers, right_fingers = torch.empty(
(0, 1, 40, 40, 40)), torch.empty((0, 1, 40, 40, 40))
grasp_object_tsdfs = dicts_get(grasp_objects, 'tsdf')
for grasp_object_index, grasp_object_tsdf in enumerate(
grasp_object_tsdfs):
self.net.eval()
go_batch = grasp_object_tsdf.unsqueeze(0).to(self.device)
with no_grad():
lf_batch, rf_batch = self.net.create_fingers(go_batch)
lf_batch.requires_grad_()
rf_batch.requires_grad_()
finger_optim = Adam([lf_batch, rf_batch], lr=1e-5)
for i in range(self.finetune_steps + 1):
finger_optim.zero_grad()
if i != self.finetune_steps:
fitness_pred = self.net.evaluate_fingers(
go_batch, lf_batch, rf_batch)
generator_loss = -fitness_pred
if self.optim_grasp_success_only:
generator_loss[0].backward()
else:
generator_loss.mean().backward()
finger_optim.step()
else:
left_fingers = cat((left_fingers, lf_batch.detach().cpu()),
dim=0)
right_fingers = cat(
(right_fingers, rf_batch.detach().cpu()), dim=0)
# dump_grasp_tuple(lf_batch.detach().clone(), go_batch.detach().clone(), rf_batch.detach().clone(), os.path.join(self.base_path, f"{grasp_object_index}_step_{i}.obj"))
self.net.generator_net = deepcopy(self.gn_net_orig)
return left_fingers, right_fingers
def create_fingers(self, grasp_objects):
left_fingers, right_fingers = torch.empty(
(0, 1, 40, 40, 40)), torch.empty((0, 1, 40, 40, 40))
grasp_object_tsdfs = dicts_get(grasp_objects, 'tsdf')
for grasp_object_index, grasp_object_tsdf in enumerate(
grasp_object_tsdfs):
self.net.eval()
go_batch = grasp_object_tsdf.unsqueeze(0).to(self.device)
gn_optim = Adam(self.net.generator_net.parameters(), lr=1e-5)
for i in range(self.finetune_steps + 1):
gn_optim.zero_grad()
lf_batch, rf_batch = self.net.create_fingers(go_batch)
if i != self.finetune_steps:
fitness_pred = self.net.evaluate_fingers(
go_batch, lf_batch, rf_batch)
generator_loss = -fitness_pred
if self.optim_grasp_success_only:
generator_loss[0, 0].backward()
else:
generator_loss.mean().backward()
gn_optim.step()
else:
left_fingers = cat((left_fingers, lf_batch.detach().cpu()),
dim=0)
right_fingers = cat(
(right_fingers, rf_batch.detach().cpu()), dim=0)
self.net.generator_net = deepcopy(self.gn_net_orig)
return left_fingers, right_fingers
def create_fingers(self, grasp_objects):
grasp_object_tsdfs = stack(dicts_get(grasp_objects,
'tsdf')).to(self.device)
with no_grad():
left_fingers, right_fingers = self.net.create_fingers(
grasp_object_tsdfs)
left_fingers = left_fingers.cpu()
right_fingers = right_fingers.cpu()
return left_fingers, right_fingers
def create_fingers(self, grasp_objects):
left_fingers, right_fingers = torch.empty(
(0, 1, 40, 40, 40)), torch.empty((0, 1, 40, 40, 40))
grasp_object_tsdfs = dicts_get(grasp_objects, 'tsdf')
for grasp_object_index, grasp_object_tsdf in enumerate(
grasp_object_tsdfs):
# Create starting fingers using generator network
self.net.eval()
go_batch = grasp_object_tsdf.unsqueeze(0).to(self.device)
with torch.no_grad():
lf_batch, rf_batch = self.net.create_fingers(go_batch)
lf_batch.requires_grad_()
rf_batch.requires_grad_()
finger_optimizer = Adam([lf_batch, rf_batch], lr=self.lr)
# optimize fingers for all objects
for i in range(self.finetune_steps):
finger_optimizer.zero_grad()
# collect loss one by one
sum_robustness_loss = None
for robustness_angle in self.robustness_angles:
# TODO: verify that the tsdf out has value 1
rotated_object = torch.tensor(
rotate(go_batch.cpu().numpy(),
angle=robustness_angle,
axes=(2, 3),
reshape=False,
mode="constant",
cval=1), # rotate in x-y plane
device=self.device)
fitness_pred = self.net.evaluate_fingers(
rotated_object, lf_batch, rf_batch)
if sum_robustness_loss is None:
sum_robustness_loss = -fitness_pred
else:
sum_robustness_loss += -fitness_pred
robustness_loss = sum_robustness_loss / \
len(self.robustness_angles)
if self.optim_grasp_success_only:
robustness_loss[0].backward()
else:
robustness_loss.mean().backward()
finger_optimizer.step()
left_fingers = cat((left_fingers, lf_batch.detach().cpu()), dim=0)
right_fingers = cat((right_fingers, rf_batch.detach().cpu()),
dim=0)
return left_fingers, right_fingers
def extend_grasp_dataset(epoch: int, object_dataset: ObjectDataset,
grasp_dataset: GraspDataset, net: GripperDesigner,
env_pool: Pool, logger: Logger,
grippers_directory: str, new_fingers_count: int):
start_extend_grasp_dataset = time()
new_grasp_stats = {
'base_connected': [],
'failed_gripper_creation': [],
'single_connected_component': [],
'grasp_scores': [],
}
newly_added_indices = list(
) # list of indices for newly added datapoints in cotrain dataset
while len(new_grasp_stats['grasp_scores']) < new_fingers_count:
grasp_object = object_dataset.sample()
print(f'\t1. Sampled {len(grasp_object)} grasp objects')
grasp_object_tsdfs = stack(dicts_get(grasp_object,
'tsdf')).to(net.device,
non_blocking=True)
with no_grad():
left_finger_gpu, right_finger_gpu = \
net.create_fingers(grasp_object_tsdfs)
print('\t2. Generated fingers')
left_finger = left_finger_gpu.cpu().clone()
right_finger = right_finger_gpu.cpu().clone()
# 2. Evaluate new fingers
start_grasp_eval_time = time()
gripper_output_paths = [
'{}/{:06}'.format(grippers_directory,
len(grasp_dataset) + i)
for i in range(len(grasp_object))
]
grasp_results = env_pool.map(exec_fn=lambda env, args: env.
compute_finger_grasp_score.remote(*args),
iterable=zip(
left_finger, right_finger,
dicts_get(grasp_object, 'urdf_path'),
gripper_output_paths))
print('\t3. Evaluated fingers ({:.01f}s)'.format(
float(time() - start_grasp_eval_time)))
# Collect stats
grasp_scores = dicts_get(grasp_results, 'score')
new_grasp_stats['grasp_scores'].extend(list(filter(None,
grasp_scores)))
new_grasp_stats['base_connected'].extend(
dicts_get(grasp_results, 'base_connected'))
new_grasp_stats['failed_gripper_creation'].extend(
list(
filter(None, dicts_get(grasp_results,
'created_grippers_failed'))))
new_grasp_stats['single_connected_component'].extend(
list(
filter(None,
dicts_get(grasp_results,
'single_connected_component'))))
# 4. Add new grasp results to grasp result training dataset
tmp_new_indices = grasp_dataset.extend([{
'grasp_object_tsdf_path':
grasp_object_path,
'left_finger_tsdf':
left_finger,
'right_finger_tsdf':
right_finger,
'grasp_result':
np.array(list(grasp_score.values()))
} for grasp_object_path, left_finger, right_finger, grasp_score in zip(
dicts_get(grasp_object, 'tsdf_path'), left_finger, right_finger,
grasp_scores) if grasp_score is not None])
newly_added_indices.extend(tmp_new_indices)
log_data = {
'Grasp_Scores/Success':
np.mean(dicts_get(new_grasp_stats['grasp_scores'], 'success')),
'Grasp_Scores/Stability_1':
np.mean(dicts_get(new_grasp_stats['grasp_scores'], 'stability_1')),
'Grasp_Scores/Stability_2':
np.mean(dicts_get(new_grasp_stats['grasp_scores'], 'stability_2')),
'Grasp_Scores/Stability_3':
np.mean(dicts_get(new_grasp_stats['grasp_scores'], 'stability_3')),
'Grasp_Scores/Stability_4':
np.mean(dicts_get(new_grasp_stats['grasp_scores'], 'stability_4')),
'Finger_Stats/Single_Connected_Component':
1.0 if len(new_grasp_stats['single_connected_component']) == 0 else
np.mean(new_grasp_stats['single_connected_component']),
'Finger_Stats/Base_Connected':
1.0 if len(new_grasp_stats['base_connected']) == 0 else np.mean(
new_grasp_stats['base_connected']),
'Finger_Stats/Mesh_Dump_Failures':
0.0 if len(new_grasp_stats['failed_gripper_creation']) == 0 else
np.mean(new_grasp_stats['failed_gripper_creation']),
}
pprint(log_data)
logger.log(data=log_data, step=epoch)
print(
f'Time taken for extend grasp dataset: {float(time()-start_extend_grasp_dataset):.1f}s'
)
return newly_added_indices
def generate_pretrain_data(envs,
grasp_object_dataset,
grasp_dataset,
voxel_size,
total=1000000):
manual_seed(time())
env_pool = Pool(envs)
output_dir = f'{grasp_dataset.directory_path}grippers/'
if not exists(output_dir):
mkdir(output_dir)
pbar = tqdm(total=total,
desc='Generating pretrain dataset',
dynamic_ncols=True)
pbar.update(len(grasp_dataset))
async_create_shapenet_gripper = ray.remote(create_shapenet_gripper)
while True:
pbar.set_description('Sampling grasp objects')
grasp_objects = grasp_object_dataset.sample()
pbar.set_description('Sampling left fingers')
left_fingers = grasp_object_dataset.sample()
pbar.set_description('Sampling right fingers')
right_fingers = grasp_object_dataset.sample()
pbar.set_description('Creating grippers')
grippers = ray.get([
async_create_shapenet_gripper.remote(
left_finger_tsdf=left_finger_tsdf,
right_finger_tsdf=right_finger_tsdf,
output_prefix=output_prefix,
voxel_size=voxel_size)
for left_finger_tsdf, right_finger_tsdf, output_prefix in zip(
dicts_get(left_fingers, 'tsdf'),
dicts_get(right_fingers, 'tsdf'),
[
'{}{:09d}'.format(output_dir,
len(grasp_dataset) + i)
for i in range(len(grasp_objects))
],
)
])
if len(list(filter(None, grippers))) == 0:
continue
# filter out invalid grippers
grasp_objects, grippers = zip(
*((grasp_object, gripper)
for grasp_object, gripper in zip(grasp_objects, grippers)
if gripper))
left_finger_tsdfs = dicts_get(grippers, 'left_finger_tsdf')
right_finger_tsdfs = dicts_get(grippers, 'right_finger_tsdf')
gripper_urdf_paths = dicts_get(grippers, 'gripper_urdf_path')
grasp_object_tsdf_paths = dicts_get(grasp_objects, "tsdf_path")
grasp_object_urdfs = dicts_get(grasp_objects, 'urdf_path')
pbar.set_description('Simulating grasps')
grasp_results = env_pool.map(
exec_fn=lambda env, args: env.simulate_grasp.remote(*args),
iterable=zip(grasp_object_urdfs, gripper_urdf_paths,
left_finger_tsdfs, right_finger_tsdfs))
grasp_data = [
{
'grasp_object_tsdf_path': grasp_object_tsdf_path,
'left_finger_tsdf': left_finger_tsdf,
'right_finger_tsdf': right_finger_tsdf,
'grasp_result': array(list(grasp_score.values()))
} for grasp_object_tsdf_path, left_finger_tsdf, right_finger_tsdf,
grasp_score in zip(grasp_object_tsdf_paths, left_finger_tsdfs,
right_finger_tsdfs, grasp_results)
if grasp_score is not None
]
grasp_dataset.extend(grasp_data, log=False)
pbar.update(len(grasp_data))
if len(grasp_dataset) > total:
exit()
retval = finger_generator.create_fingers(grasp_objects)
gripper_output_paths = [
'{}/{:06}'.format(grippers_directory, i + len(results))
for i in range(len(grasp_objects))
]
if type(finger_generator) == ImprintFingerGenerator:
left_fingers, right_fingers, gripper_urdf_paths = retval
for gripper_urdf_path, target_gripper_urdf_path in \
zip(gripper_urdf_paths, gripper_output_paths):
# copy collision meshes over
prefix = splitext(gripper_urdf_path)[0]
copyfile(prefix + '_right_collision.obj',
target_gripper_urdf_path + '_right_collision.obj')
copyfile(prefix + '_left_collision.obj',
target_gripper_urdf_path + '_left_collision.obj')
else:
left_fingers, right_fingers = retval
grasp_results = env_pool.map(
exec_fn=lambda env, args: env.compute_finger_grasp_score.
remote(*args),
iterable=zip(left_fingers, right_fingers,
dicts_get(grasp_objects,
'urdf_path'), gripper_output_paths,
[False] * len(gripper_output_paths)))
results.extend(grasp_results)
metrics = acc_results(results)
print_results(metrics, name=finger_generator_config["desc"])
save_path = os.path.join(grippers_directory, "val_results.npz")
np.savez(save_path, **metrics)
print(f"saved results at {save_path}")