RoboTwin_image/envs/gpt_pick_apple_messy.py
2025-07-02 03:13:07 +00:00

41 lines
1.5 KiB
Python

from .base_task import Base_task
from .pick_apple_messy import pick_apple_messy
from .utils import *
import sapien
class gpt_pick_apple_messy(pick_apple_messy):
def play_once(self):
# Retrieve the actor and actor_data
apple = self.actor_name_dic['apple']
apple_data = self.actor_data_dic['apple_data']
# Get the apple's pose
apple_pose = self.get_actor_functional_pose(apple, apple_data)
# Determine which arm to use based on the apple's x coordinate
if apple_pose[0] > 0:
arm_tag = "right"
move_function = self.right_move_to_pose_with_screw
close_gripper_function = self.close_right_gripper
else:
arm_tag = "left"
move_function = self.left_move_to_pose_with_screw
close_gripper_function = self.close_left_gripper
# Get the grasp pose for the apple
pre_grasp_pose = self.get_grasp_pose_to_grasp_object(endpose_tag=arm_tag, actor=apple, actor_data=apple_data, pre_dis=0.09)
target_grasp_pose = self.get_grasp_pose_to_grasp_object(endpose_tag=arm_tag, actor=apple, actor_data=apple_data, pre_dis=0)
# Move the arm to the pre-grasp pose
move_function(pre_grasp_pose)
# Move the arm to the grasp pose
move_function(target_grasp_pose)
# Close the gripper to grasp the apple
close_gripper_function()
# Lift the apple by moving back to the pre-grasp pose
move_function(pre_grasp_pose)