93 lines
4.0 KiB
Python
93 lines
4.0 KiB
Python
|
|
from .base_task import Base_task
|
|
from .blocks_stack_easy import blocks_stack_easy
|
|
from .utils import *
|
|
import sapien
|
|
|
|
class gpt_blocks_stack_easy(blocks_stack_easy):
|
|
def play_once(self):
|
|
# Retrieve actor objects and data
|
|
block1 = self.actor_name_dic['block1']
|
|
block2 = self.actor_name_dic['block2']
|
|
block1_target_pose = self.actor_name_dic['block1_target_pose']
|
|
|
|
block1_data = self.actor_data_dic['block1_data']
|
|
block2_data = self.actor_data_dic['block2_data']
|
|
block1_target_pose_data = self.actor_data_dic['block1_target_pose']
|
|
|
|
# Determine which arm to use for block1
|
|
block1_pose = self.get_actor_goal_pose(block1, block1_data, 0)
|
|
if block1_pose[0] > 0:
|
|
arm_tag_block1 = "right"
|
|
arm_tag_block2 = "left"
|
|
else:
|
|
arm_tag_block1 = "left"
|
|
arm_tag_block2 = "right"
|
|
|
|
# Pick up block1 and move it to the target position
|
|
self.pick_and_place_block(block1, block1_data, block1_target_pose, block1_target_pose_data, arm_tag_block1)
|
|
|
|
# Determine which arm to use for block2
|
|
block2_pose = self.get_actor_goal_pose(block2, block2_data, 0)
|
|
if block2_pose[0] > 0:
|
|
arm_tag_block2 = "right"
|
|
else:
|
|
arm_tag_block2 = "left"
|
|
|
|
# Pick up block2 and place it on block1
|
|
self.pick_and_place_block(block2, block2_data, block1, block1_data, arm_tag_block2)
|
|
|
|
def pick_and_place_block(self, block, block_data, target_block, target_block_data, arm_tag):
|
|
# Get the pre-grasp and target grasp poses
|
|
pre_grasp_pose = self.get_grasp_pose_to_grasp_object(endpose_tag=arm_tag, actor=block, actor_data=block_data, pre_dis=0.09)
|
|
target_grasp_pose = self.get_grasp_pose_to_grasp_object(endpose_tag=arm_tag, actor=block, actor_data=block_data, pre_dis=0)
|
|
|
|
# Move to the pre-grasp pose
|
|
self.move_arm_to_pose(arm_tag, pre_grasp_pose)
|
|
|
|
# Move to the target grasp pose
|
|
self.move_arm_to_pose(arm_tag, target_grasp_pose)
|
|
|
|
# Close the gripper to grasp the block
|
|
self.close_gripper(arm_tag)
|
|
|
|
# Lift the block up
|
|
self.move_arm_to_pose(arm_tag, pre_grasp_pose)
|
|
|
|
# Get the target pose for placing the block
|
|
target_pose = self.get_actor_goal_pose(target_block, target_block_data, 0)
|
|
target_approach_direction = self.world_direction_dic['top_down']
|
|
pre_place_pose = self.get_grasp_pose_from_goal_point_and_direction(block, block_data, endpose_tag=arm_tag, actor_functional_point_id=0, target_point=target_pose, target_approach_direction=target_approach_direction, pre_dis=0.09)
|
|
target_place_pose = self.get_grasp_pose_from_goal_point_and_direction(block, block_data, endpose_tag=arm_tag, actor_functional_point_id=0, target_point=target_pose, target_approach_direction=target_approach_direction, pre_dis=0)
|
|
|
|
# Move to the pre-place pose
|
|
self.move_arm_to_pose(arm_tag, pre_place_pose)
|
|
|
|
# Move to the target place pose
|
|
self.move_arm_to_pose(arm_tag, target_place_pose)
|
|
|
|
# Open the gripper to place the block
|
|
self.open_gripper(arm_tag)
|
|
|
|
# Move the arm away to avoid collision
|
|
avoid_collision_pose = self.get_avoid_collision_pose(arm_tag)
|
|
self.move_arm_to_pose(arm_tag, avoid_collision_pose)
|
|
|
|
def move_arm_to_pose(self, arm_tag, pose):
|
|
if arm_tag == "left":
|
|
self.left_move_to_pose_with_screw(pose)
|
|
elif arm_tag == "right":
|
|
self.right_move_to_pose_with_screw(pose)
|
|
|
|
def close_gripper(self, arm_tag):
|
|
if arm_tag == "left":
|
|
self.close_left_gripper()
|
|
elif arm_tag == "right":
|
|
self.close_right_gripper()
|
|
|
|
def open_gripper(self, arm_tag):
|
|
if arm_tag == "left":
|
|
self.open_left_gripper()
|
|
elif arm_tag == "right":
|
|
self.open_right_gripper()
|