Spaces:
Runtime error
Runtime error
File size: 3,952 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 |
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils
import pybullet as p
import random
class StackBlockPyramidSeq(Task):
"""Stacking Block Pyramid Sequence base class."""
def __init__(self):
super().__init__()
self.max_steps = 12
self.lang_template = "put the {pick} block on {place}"
self.task_completed_desc = "done stacking block pyramid."
self.additional_reset()
def reset(self, env):
super().reset(env)
# Add base.
# x, y, z dimensions for the asset size
base_size = (0.05, 0.15, 0.005)
base_urdf = 'stacking/stand.urdf'
base_pose = self.get_random_pose(env, base_size)
env.add_object(base_urdf, base_pose, 'fixed')
# Block colors.
colors, color_names = utils.get_colors(self.mode)
# Add blocks.
# x, y, z dimensions for the asset size
block_size = (0.04, 0.04, 0.04)
block_urdf = 'stacking/block.urdf'
objs = []
for i in range(6):
block_pose = self.get_random_pose(env, block_size)
block_id = env.add_object(block_urdf, block_pose, color=colors[i])
objs.append(block_id)
# Associate placement locations for goals.
place_pos = [(0, -0.05, 0.03), (0, 0, 0.03),
(0, 0.05, 0.03), (0, -0.025, 0.08),
(0, 0.025, 0.08), (0, 0, 0.13)]
targs = [(utils.apply(base_pose, i), base_pose[1]) for i in place_pos]
# Goal: make bottom row.
self.add_goal(objs=[objs[0]], matches=np.ones((1, 1)), targ_poses=[targs[0]], replace=False,
rotations=True, metric='pose', params=None, step_max_reward=1 / 6, symmetries=[np.pi/2])
self.lang_goals.append(self.lang_template.format(pick=color_names[0],
place="the lightest brown block"))
self.add_goal(objs=[objs[1]], matches=np.ones((1, 1)), targ_poses=[targs[1]], replace=False,
rotations=True, metric='pose', params=None, step_max_reward=1 / 6, symmetries=[np.pi/2])
self.lang_goals.append(self.lang_template.format(pick=color_names[1],
place="the middle brown block"))
self.add_goal(objs=[objs[2]], matches=np.ones((1, 1)), targ_poses=[targs[2]], replace=False,
rotations=True, metric='pose', params=None, step_max_reward=1 / 6, symmetries=[np.pi/2])
self.lang_goals.append(self.lang_template.format(pick=color_names[2],
place="the darkest brown block"))
# Goal: make middle row.
self.add_goal(objs=[objs[3]], matches=np.ones((1, 1)), targ_poses=[targs[3]], replace=False,
rotations=True, metric='pose', params=None, step_max_reward=1 / 6, symmetries=[np.pi/2])
self.lang_goals.append(self.lang_template.format(pick=color_names[3],
place=f"the {color_names[0]} and {color_names[1]} blocks"))
self.add_goal(objs=[objs[4]], matches=np.ones((1, 1)), targ_poses=[targs[4]], replace=False,
rotations=True, metric='pose', params=None, step_max_reward=1 / 6, symmetries=[np.pi/2])
self.lang_goals.append(self.lang_template.format(pick=color_names[4],
place=f"the {color_names[1]} and {color_names[2]} blocks"))
# Goal: make top row.
self.add_goal(objs=[objs[5]], matches=np.ones((1, 1)), targ_poses=[targs[5]], replace=False,
rotations=True, metric='pose', params=None, step_max_reward=1 / 6, symmetries=[np.pi/2])
self.lang_goals.append(self.lang_template.format(pick=color_names[5],
place=f"the {color_names[3]} and {color_names[4]} blocks")) |