Spaces:
Runtime error
Runtime error
File size: 5,269 Bytes
8fc2b4e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 |
import os
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils
import pybullet as p
class PackingSeenGoogleObjectsSeq(Task):
""": Place the specified objects in the brown box following the order prescribed in the language
instruction at each timestep."""
def __init__(self):
super().__init__()
self.max_steps = 6
self.lang_template = "pack the {obj} in the brown box"
self.task_completed_desc = "done packing objects."
self.object_names = self.get_object_names()
self.additional_reset()
def get_object_names(self):
return utils.google_all_shapes
def reset(self, env):
super().reset(env)
# object names
object_names = self.object_names[self.mode]
# Add container box.
zone_size = self.get_random_size(0.2, 0.35, 0.2, 0.35, 0.05, 0.05)
zone_pose = self.get_random_pose(env, zone_size)
container_template = 'container/container-template_DIM_HALF.urdf'
replace = {'DIM': zone_size, 'HALF': (zone_size[0] / 2, zone_size[1] / 2, zone_size[2] / 2)}
container_urdf = self.fill_template(container_template, replace)
env.add_object(container_urdf, zone_pose, 'fixed')
margin = 0.01
min_object_dim = 0.08
bboxes = []
# Split container space with KD trees.
stack_size = np.array(zone_size)
stack_size[0] -= 0.01
stack_size[1] -= 0.01
root_size = (0.01, 0.01, 0) + tuple(stack_size)
root = utils.TreeNode(None, [], bbox=np.array(root_size))
utils.KDTree(root, min_object_dim, margin, bboxes)
# Add Google Scanned Objects to scene.
object_ids = []
bboxes = np.array(bboxes)
scale_factor = 5
object_template = 'google/object-template_FNAME_COLOR_SCALE.urdf'
chosen_objs, repeat_category = self.choose_objects(object_names, len(bboxes))
object_descs = []
for i, bbox in enumerate(bboxes):
size = bbox[3:] - bbox[:3]
max_size = size.max()
position = size / 2. + bbox[:3]
position[0] += -zone_size[0] / 2
position[1] += -zone_size[1] / 2
shape_size = max_size * scale_factor
pose = self.get_random_pose(env, size)
# Add object only if valid pose found.
if pose[0] is not None:
# Initialize with a slightly tilted pose so that the objects aren't always erect.
slight_tilt = utils.q_mult(pose[1], (-0.1736482, 0, 0, 0.9848078))
ps = ((pose[0][0], pose[0][1], pose[0][2]+0.05), slight_tilt)
object_name = chosen_objs[i]
object_name_with_underscore = object_name.replace(" ", "_")
mesh_file = os.path.join(self.assets_root,
'google',
'meshes_fixed',
f'{object_name_with_underscore}.obj')
texture_file = os.path.join(self.assets_root,
'google',
'textures',
f'{object_name_with_underscore}.png')
try:
replace = {'FNAME': (mesh_file,),
'SCALE': [shape_size, shape_size, shape_size],
'COLOR': (0.2, 0.2, 0.2)}
urdf = self.fill_template(object_template, replace)
box_id = env.add_object(urdf, ps)
object_ids.append((box_id, (0, None)))
texture_id = p.loadTexture(texture_file)
p.changeVisualShape(box_id, -1, textureUniqueId=texture_id)
p.changeVisualShape(box_id, -1, rgbaColor=[1, 1, 1, 1])
object_descs.append(object_name)
except Exception as e:
print("Failed to load Google Scanned Object in PyBullet")
print(object_name_with_underscore, mesh_file, texture_file)
print(f"Exception: {e}")
self.set_goals(object_descs, object_ids, repeat_category, zone_pose, zone_size)
for i in range(480):
p.stepSimulation()
def choose_objects(self, object_names, k):
repeat_category = None
return np.random.choice(object_names, k, replace=False), repeat_category
def set_goals(self, object_descs, object_ids, repeat_category, zone_pose, zone_size):
# Random picking sequence.
num_pack_objs = np.random.randint(1, len(object_ids))
object_ids = object_ids[:num_pack_objs]
true_poses = []
for obj_idx, (object_id, _) in enumerate(object_ids):
true_poses.append(zone_pose)
self.add_goal(objs=[object_id], matches=np.int32([[1]]), targ_poses=[zone_pose], replace=False,
rotations=True, metric='zone', params=[(zone_pose, zone_size)], step_max_reward=1 / len(object_ids))
self.lang_goals.append(self.lang_template.format(obj=object_descs[obj_idx]))
# Only mistake allowed.
self.max_steps = len(object_ids)+1
|