| import numpy as np | |
| from cliport.tasks.task import Task | |
| from cliport.utils import utils | |
| import pybullet as p | |
| class BlockInsertion(Task): | |
| """pick up the L-shaped red block and place it into the L-shaped fixture.""" | |
| def __init__(self): | |
| super().__init__() | |
| self.max_steps = 3 | |
| self.lang_template = "put the L shape block in the L shape hole" | |
| self.task_completed_desc = "done with insertion." | |
| self.additional_reset() | |
| def get_random_pose(self, env, obj_size): | |
| pose = super().get_random_pose(env, obj_size) | |
| pos, rot = pose | |
| rot = utils.eulerXYZ_to_quatXYZW((0, 0, np.pi / 2)) | |
| return pos, rot | |
| def reset(self, env): | |
| super().reset(env) | |
| """Add L-shaped block.""" | |
| size = (0.1, 0.1, 0.04) | |
| urdf = 'insertion/ell.urdf' | |
| pose = self.get_random_pose(env, size) | |
| block_id = env.add_object(urdf, pose) | |
| """Add L-shaped fixture to place block.""" | |
| size = (0.1, 0.1, 0.04) | |
| urdf = 'insertion/fixture.urdf' | |
| targ_pose = self.get_random_pose(env, size) | |
| env.add_object(urdf, targ_pose, 'fixed') | |
| self.add_goal(objs=[block_id], matches=np.int32([[1]]), targ_poses=[targ_pose], replace=False, | |
| rotations=False, metric='pose', params=None, step_max_reward=1, symmetries=[2 * np.pi], | |
| language_goal=self.lang_template) | |