47 lines
1.7 KiB
Python
47 lines
1.7 KiB
Python
|
|
from .base_task import Base_task
|
|
from .utils import *
|
|
import sapien
|
|
|
|
class diverse_bottles_pick(Base_task):
|
|
def setup_demo(self,is_test = False,**kwags):
|
|
super()._init(**kwags)
|
|
self.create_table_and_wall()
|
|
self.load_robot()
|
|
self.setup_planner()
|
|
self.load_camera(kwags.get('camera_w', 640),kwags.get('camera_h', 480))
|
|
self.pre_move()
|
|
if is_test:
|
|
self.id_list = [2*i + 1 for i in range(11)]
|
|
else:
|
|
self.id_list = [2*i for i in range(11)]
|
|
self.left_bottle_target_position = [-0.06,-0.105, 0.92]
|
|
self.right_bottle_target_position = [0.06,-0.105, 0.92]
|
|
self.load_actors(f"./task_config/scene_info/{self.task_name[4:]}.json")
|
|
self.step_lim = 400
|
|
|
|
def pre_move(self):
|
|
render_freq = self.render_freq
|
|
self.render_freq=0
|
|
|
|
self.together_close_gripper(save_freq=None)
|
|
self.together_open_gripper(save_freq=None)
|
|
|
|
self.render_freq = render_freq
|
|
|
|
def play_once(self):
|
|
pass
|
|
|
|
def check_success(self):
|
|
red_target = [-0.06,-0.105]
|
|
green_target = [0.06,-0.105]
|
|
eps = 0.03
|
|
bottle1 = self.actor_name_dic['bottle1']
|
|
bottle2 = self.actor_name_dic['bottle2']
|
|
bottle1_pose = bottle1.get_pose().p
|
|
bottle2_pose = bottle2.get_pose().p
|
|
if bottle1_pose[2] < 0.78 or bottle2_pose[2] < 0.78:
|
|
self.actor_pose = False
|
|
return abs(bottle1_pose[0]-red_target[0])<eps and abs(bottle1_pose[1]-red_target[1])<eps and bottle1_pose[2]>0.89 and\
|
|
abs(bottle2_pose[0]-green_target[0])<eps and abs(bottle2_pose[1]-green_target[1])<eps and bottle2_pose[2]>0.89
|