import numpy as np import os import pybullet as p import random from cliport.tasks import primitives from cliport.tasks.grippers import Spatula from cliport.tasks.task import Task from cliport.utils import utils import numpy as np from cliport.tasks.task import Task from cliport.utils import utils class InsertAndStack(Task): """Insert the ell into the fixture and then stack blocks on top of it.""" def __init__(self): super().__init__() self.max_steps = 10 self.metric = 'pose' self.lang_template = "insert the ell into the fixture and then stack blocks on top of it" self.task_completed_desc = "done insert-and-stack." def reset(self, env): super().reset(env) # Add ell ell_size = (0.1, 0.1, 0.1) ell_pose = self.get_random_pose(env, ell_size) ell_urdf = 'insertion/ell.urdf' ell_id = env.add_object(ell_urdf, ell_pose) self.color_random_bright(ell_id) # Add fixture fixture_size = (0.12, 0.12, 0.1) fixture_pose = self.get_random_pose(env, fixture_size) fixture_urdf = 'insertion/fixture.urdf' fixture_id = env.add_object(fixture_urdf, fixture_pose) self.color_random_bright(fixture_id) # Add blocks block_size = (0.04, 0.04, 0.04) block_urdf = 'stacking/block.urdf' blocks = [] for _ in range(3): # We want 3 blocks block_pose = self.get_random_pose(env, block_size) block_id = env.add_object(block_urdf, block_pose) self.color_random_bright(block_id) # Randomly color the blocks blocks.append(block_id) # Define the zone zone_size = (0.1, 0.1, 0.1) zone_pose = self.get_random_pose(env, zone_size) zone_urdf = 'zone/zone.urdf' env.add_object(zone_urdf, zone_pose, 'fixed') # Zone is static # Set task goals objs = [ell_id] + blocks goal_poses = [fixture_pose] * (len(objs)) self.add_goal(objs=objs, matches=np.ones((len(objs), 1)), targ_poses=goal_poses, replace=False, rotations=True, metric='pose', params=None, step_max_reward=1, language_goal=self.lang_template)