File size: 15,380 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 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 |
Before writing the code for the task "TASK_NAME_TEMPLATE". Here are some APIs that are defined. Please confirm that you understand these APIs.
"""
class Task():
"""Base Task class."""
def __init__(self):
self.ee = Suction
self.mode = 'train'
self.sixdof = False
self.primitive = primitives.PickPlace()
self.oracle_cams = cameras.Oracle.CONFIG
# Evaluation epsilons (for pose evaluation metric).
self.pos_eps = 0.01
self.rot_eps = np.deg2rad(15)
# Workspace bounds.
self.pix_size = 0.003125
self.bounds = np.array([[0.25, 0.75], [-0.5, 0.5], [0, 0.3]])
self.zone_bounds = np.copy(self.bounds)
self.goals = []
self.lang_goals = []
self.task_completed_desc = "task completed."
self.progress = 0
self._rewards = 0
self.assets_root = None
def reset(self, env):
if not self.assets_root:
raise ValueError('assets_root must be set for task, '
'call set_assets_root().')
self.goals = []
self.lang_goals = []
self.progress = 0 # Task progression metric in range [0, 1].
self._rewards = 0 # Cumulative returned rewards.
# -------------------------------------------------------------------------
# Oracle Agent
# -------------------------------------------------------------------------
def oracle(self, env):
"""Oracle agent."""
OracleAgent = collections.namedtuple('OracleAgent', ['act'])
def act(obs, info):
"""Calculate action."""
# Oracle uses perfect RGB-D orthographic images and segmentation masks.
_, hmap, obj_mask = self.get_true_image(env)
# Unpack next goal step.
objs, matches, targs, replace, rotations, _, _, _ = self.goals[0]
# Match objects to targets without replacement.
if not replace:
# Modify a copy of the match matrix.
matches = matches.copy()
# Ignore already matched objects.
for i in range(len(objs)):
object_id, (symmetry, _) = objs[i]
pose = p.getBasePositionAndOrientation(object_id)
targets_i = np.argwhere(matches[i, :]).reshape(-1)
for j in targets_i:
if self.is_match(pose, targs[j], symmetry):
matches[i, :] = 0
matches[:, j] = 0
# Get objects to be picked (prioritize farthest from nearest neighbor).
nn_dists = []
nn_targets = []
for i in range(len(objs)):
object_id, (symmetry, _) = objs[i]
xyz, _ = p.getBasePositionAndOrientation(object_id)
targets_i = np.argwhere(matches[i, :]).reshape(-1)
if len(targets_i) > 0:
targets_xyz = np.float32([targs[j][0] for j in targets_i])
dists = np.linalg.norm(
targets_xyz - np.float32(xyz).reshape(1, 3), axis=1)
nn = np.argmin(dists)
nn_dists.append(dists[nn])
nn_targets.append(targets_i[nn])
# Handle ignored objects.
else:
nn_dists.append(0)
nn_targets.append(-1)
order = np.argsort(nn_dists)[::-1]
# Filter out matched objects.
order = [i for i in order if nn_dists[i] > 0]
pick_mask = None
for pick_i in order:
pick_mask = np.uint8(obj_mask == objs[pick_i][0])
# Erode to avoid picking on edges.
# pick_mask = cv2.erode(pick_mask, np.ones((3, 3), np.uint8))
if np.sum(pick_mask) > 0:
break
# Trigger task reset if no object is visible.
if pick_mask is None or np.sum(pick_mask) == 0:
self.goals = []
self.lang_goals = []
print('Object for pick is not visible. Skipping demonstration.')
return
# Get picking pose.
pick_prob = np.float32(pick_mask)
pick_pix = utils.sample_distribution(pick_prob)
# For "deterministic" demonstrations on insertion-easy, use this:
# pick_pix = (160,80)
pick_pos = utils.pix_to_xyz(pick_pix, hmap,
self.bounds, self.pix_size)
pick_pose = (np.asarray(pick_pos), np.asarray((0, 0, 0, 1)))
# Get placing pose.
targ_pose = targs[nn_targets[pick_i]]
obj_pose = p.getBasePositionAndOrientation(objs[pick_i][0])
if not self.sixdof:
obj_euler = utils.quatXYZW_to_eulerXYZ(obj_pose[1])
obj_quat = utils.eulerXYZ_to_quatXYZW((0, 0, obj_euler[2]))
obj_pose = (obj_pose[0], obj_quat)
world_to_pick = utils.invert(pick_pose)
obj_to_pick = utils.multiply(world_to_pick, obj_pose)
pick_to_obj = utils.invert(obj_to_pick)
place_pose = utils.multiply(targ_pose, pick_to_obj)
# Rotate end effector?
if not rotations:
place_pose = (place_pose[0], (0, 0, 0, 1))
place_pose = (np.asarray(place_pose[0]), np.asarray(place_pose[1]))
return {'pose0': pick_pose, 'pose1': place_pose}
return OracleAgent(act)
# -------------------------------------------------------------------------
# Reward Function and Task Completion Metrics
# -------------------------------------------------------------------------
def reward(self):
"""Get delta rewards for current timestep.
Returns:
A tuple consisting of the scalar (delta) reward, plus `extras`
dict which has extra task-dependent info from the process of
computing rewards that gives us finer-grained details. Use
`extras` for further data analysis.
"""
reward, info = 0, {}
# Unpack next goal step.
objs, matches, targs, _, _, metric, params, max_reward = self.goals[0]
# Evaluate by matching object poses.
if metric == 'pose':
step_reward = 0
for i in range(len(objs)):
object_id, (symmetry, _) = objs[i]
pose = p.getBasePositionAndOrientation(object_id)
targets_i = np.argwhere(matches[i, :]).reshape(-1)
for j in targets_i:
target_pose = targs[j]
if self.is_match(pose, target_pose, symmetry):
step_reward += max_reward / len(objs)
print(f"object {i} match with target {j} rew: {step_reward}")
break
# Evaluate by measuring object intersection with zone.
elif metric == 'zone':
zone_pts, total_pts = 0, 0
obj_pts, zones = params
for zone_idx, (zone_pose, zone_size) in enumerate(zones):
# Count valid points in zone.
for obj_idx, obj_id in enumerate(obj_pts):
pts = obj_pts[obj_id]
obj_pose = p.getBasePositionAndOrientation(obj_id)
world_to_zone = utils.invert(zone_pose)
obj_to_zone = utils.multiply(world_to_zone, obj_pose)
pts = np.float32(utils.apply(obj_to_zone, pts))
if len(zone_size) > 1:
valid_pts = np.logical_and.reduce([
pts[0, :] > -zone_size[0] / 2, pts[0, :] < zone_size[0] / 2,
pts[1, :] > -zone_size[1] / 2, pts[1, :] < zone_size[1] / 2,
pts[2, :] < self.zone_bounds[2, 1]])
# if zone_idx == matches[obj_idx].argmax():
zone_pts += np.sum(np.float32(valid_pts))
total_pts += pts.shape[1]
step_reward = max_reward * (zone_pts / total_pts)
# Get cumulative rewards and return delta.
reward = self.progress + step_reward - self._rewards
self._rewards = self.progress + step_reward
# Move to next goal step if current goal step is complete.
if np.abs(max_reward - step_reward) < 0.01:
self.progress += max_reward # Update task progress.
self.goals.pop(0)
if len(self.lang_goals) > 0:
self.lang_goals.pop(0)
return reward, info
def done(self):
"""Check if the task is done or has failed.
Returns:
True if the episode should be considered a success, which we
use for measuring successes, which is particularly helpful for tasks
where one may get successes on the very last time step, e.g., getting
the cloth coverage threshold on the last alllowed action.
However, for bag-items-easy and bag-items-hard (which use the
'bag-items' metric), it may be necessary to filter out demos that did
not attain sufficiently high reward in external code. Currently, this
is done in `main.py` and its ignore_this_demo() method.
"""
return (len(self.goals) == 0) or (self._rewards > 0.99)
# return zone_done or defs_done or goal_done
# -------------------------------------------------------------------------
# Environment Helper Functions
# -------------------------------------------------------------------------
def is_match(self, pose0, pose1, symmetry):
"""Check if pose0 and pose1 match within a threshold."""
# Get translational error.
diff_pos = np.float32(pose0[0][:2]) - np.float32(pose1[0][:2])
dist_pos = np.linalg.norm(diff_pos)
# Get rotational error around z-axis (account for symmetries).
diff_rot = 0
if symmetry > 0:
rot0 = np.array(utils.quatXYZW_to_eulerXYZ(pose0[1]))[2]
rot1 = np.array(utils.quatXYZW_to_eulerXYZ(pose1[1]))[2]
diff_rot = np.abs(rot0 - rot1) % symmetry
if diff_rot > (symmetry / 2):
diff_rot = symmetry - diff_rot
return (dist_pos < self.pos_eps) and (diff_rot < self.rot_eps)
def get_random_pose(self, env, obj_size):
"""Get random collision-free object pose within workspace bounds."""
# Get erosion size of object in pixels.
max_size = np.sqrt(obj_size[0] ** 2 + obj_size[1] ** 2)
erode_size = int(np.round(max_size / self.pix_size))
_, hmap, obj_mask = self.get_true_image(env)
# Randomly sample an object pose within free-space pixels.
free = np.ones(obj_mask.shape, dtype=np.uint8)
for obj_ids in env.obj_ids.values():
for obj_id in obj_ids:
free[obj_mask == obj_id] = 0
free[0, :], free[:, 0], free[-1, :], free[:, -1] = 0, 0, 0, 0
free = cv2.erode(free, np.ones((erode_size, erode_size), np.uint8))
# if np.sum(free) == 0:
# return None, None
if np.sum(free) == 0:
# avoid returning None, None
# return None, None
pix = (obj_mask.shape[0] // 2, obj_mask.shape[1] // 2)
else:
pix = utils.sample_distribution(np.float32(free))
pos = utils.pix_to_xyz(pix, hmap, self.bounds, self.pix_size)
pos = (pos[0], pos[1], obj_size[2] / 2)
theta = np.random.rand() * 2 * np.pi
rot = utils.eulerXYZ_to_quatXYZW((0, 0, theta))
return pos, rot
def get_lang_goal(self):
if len(self.lang_goals) == 0:
return self.task_completed_desc
else:
return self.lang_goals[0]
def get_reward(self):
return float(self._rewards)
# -------------------------------------------------------------------------
# Helper Functions
# -------------------------------------------------------------------------
def fill_template(self, template, replace):
"""Read a file and replace key strings."""
full_template_path = os.path.join(self.assets_root, template)
with open(full_template_path, 'r') as file:
fdata = file.read()
for field in replace:
for i in range(len(replace[field])):
fdata = fdata.replace(f'{field}{i}', str(replace[field][i]))
alphabet = string.ascii_lowercase + string.digits
rname = ''.join(random.choices(alphabet, k=16))
tmpdir = tempfile.gettempdir()
template_filename = os.path.split(template)[-1]
fname = os.path.join(tmpdir, f'{template_filename}.{rname}')
with open(fname, 'w') as file:
file.write(fdata)
return fname
def get_random_size(self, min_x, max_x, min_y, max_y, min_z, max_z):
"""Get random box size."""
size = np.random.rand(3)
size[0] = size[0] * (max_x - min_x) + min_x
size[1] = size[1] * (max_y - min_y) + min_y
size[2] = size[2] * (max_z - min_z) + min_z
return tuple(size)
"""""
# Environment Class
def add_object(self, urdf, pose, category='rigid'):
"""List of (fixed, rigid, or deformable) objects in env."""
fixed_base = 1 if category == 'fixed' else 0
obj_id = pybullet_utils.load_urdf(
p,
os.path.join(self.assets_root, urdf),
pose[0],
pose[1],
useFixedBase=fixed_base)
self.obj_ids[category].append(obj_id)
return obj_id
"""
=========
Note that the objects need to obey physics and not collide with each other, and the object goal poses need to be above the table with lower bound x=0.25, y=-0.5 and upper bound x=0.75, y=0.5. When there are multiple objects for a multi-step pick-and-place task, there are often multiple subgoals. Once the task and environment are generated, an agent with a pick and place primitive will follow the defined goal to accomplish the tasks.
Additionally, make sure you understand and summarize the ``self.goals`` variables, which has a list of 8-tuple with (objs, matches, targ_poses, replace, rotations, metric, params, step_max_reward, symmetries).
- objs (List of obj_id): object ID.
- matches (Binary Matrix): a binary matrix that denotes which object is matched with which target. This matrix has dimension len(objs) x len(targs).
- targ_poses (List of Poses [(translation, rotation)] ): a list of target poses of tuple (translation, rotation). Don't pass in object IDs such as `bowls[i-1][0]` or `[stands[i][0]]`.
- replace (Boolean): whether each object can match with one unique target. This is important if we have one target and multiple objects. If it's set to be false, then any object matching with the target will satisfy.
- rotations (Boolean): whether the placement action has a rotation degree of freedom.
- metric (`pose` or `zone`): `pose` or `zone` that the object needs to be transported to. Example: `pose`.
- params (List of (zone_target, zone_size)): a list of (zone_target, zone_size) for each zone if the metric is `zone`.
- step_max_reward (float): subgoal reward threshold.
- symmetries: the radians that the object is symmetric around z axis.
|