Spaces:
Runtime error
Runtime error
File size: 2,533 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 |
import os
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils
class AlignBoxCorner(Task):
"""Pick up the randomly sized box and align one of its corners to the L-shaped marker on the tabletop."""
def __init__(self):
super().__init__()
self.max_steps = 3
self.lang_template = "align the brown box with the green corner"
self.task_completed_desc = "done with alignment"
self.additional_reset()
def reset(self, env):
super().reset(env)
# Generate randomly shaped box.
box_size = self.get_random_size(0.05, 0.15, 0.05, 0.15, 0.01, 0.06)
# Add corner.
dimx = (box_size[0] / 2 - 0.025 + 0.0025, box_size[0] / 2 + 0.0025)
dimy = (box_size[1] / 2 + 0.0025, box_size[1] / 2 - 0.025 + 0.0025)
corner_template = 'corner/corner-template.urdf'
replace = {'DIMX': dimx, 'DIMY': dimy}
# IMPORTANT: REPLACE THE TEMPLATE URDF
corner_urdf = self.fill_template(corner_template, replace)
corner_size = (box_size[0], box_size[1], 0)
corner_pose = self.get_random_pose(env, corner_size)
env.add_object(corner_urdf, corner_pose, 'fixed')
# Add possible placing poses.
theta = utils.quatXYZW_to_eulerXYZ(corner_pose[1])[2]
fip_rot = utils.eulerXYZ_to_quatXYZW((0, 0, theta + np.pi))
pose1 = (corner_pose[0], fip_rot)
alt_x = (box_size[0] / 2) - (box_size[1] / 2)
alt_y = (box_size[1] / 2) - (box_size[0] / 2)
alt_pos = (alt_x, alt_y, 0)
alt_rot0 = utils.eulerXYZ_to_quatXYZW((0, 0, np.pi / 2))
alt_rot1 = utils.eulerXYZ_to_quatXYZW((0, 0, 3 * np.pi / 2))
pose2 = utils.multiply(corner_pose, (alt_pos, alt_rot0))
pose3 = utils.multiply(corner_pose, (alt_pos, alt_rot1))
# Add box.
box_template = 'box/box-template.urdf'
# IMPORTANT: REPLACE THE TEMPLATE URDF
box_urdf = self.fill_template(box_template, {'DIM': np.float32(box_size)})
box_pose = self.get_random_pose(env, box_size)
box_id = env.add_object(box_urdf, box_pose)
self.color_random_brown(box_id)
# Goal: box is aligned with corner (1 of 4 possible poses).
self.add_goal(objs=[box_id], matches=np.int32([[1, 1, 1, 1]]), targ_poses=[corner_pose, pose1, pose2, pose3], replace=False,
rotations=True, metric='pose', params=None, step_max_reward=1, symmetries=[2 * np.pi])
self.lang_goals.append(self.lang_template)
|