Spaces:
Paused
Paused
| from typing import Any, Dict, Union | |
| import numpy as np | |
| import sapien | |
| import torch | |
| from mani_skill.agents.robots import SO100, Fetch, Panda | |
| from mani_skill.envs.sapien_env import BaseEnv | |
| from mani_skill.envs.tasks.tabletop.pick_cube_cfgs import PICK_CUBE_CONFIGS | |
| from mani_skill.sensors.camera import CameraConfig | |
| from mani_skill.utils import sapien_utils | |
| from mani_skill.utils.building import actors | |
| from mani_skill.utils.registration import register_env | |
| from mani_skill.utils.scene_builder.table import TableSceneBuilder | |
| from mani_skill.utils.structs.pose import Pose | |
| from mani_skill.utils import common, sapien_utils | |
| from mani_skill.utils.structs import Actor | |
| #Robomme | |
| import matplotlib.pyplot as plt | |
| from mani_skill.utils.geometry.rotation_conversions import ( | |
| euler_angles_to_matrix, | |
| matrix_to_quaternion, | |
| ) | |
| from .utils import * | |
| from .utils.difficulty import normalize_robomme_difficulty | |
| from .utils.subgoal_evaluate_func import static_check | |
| from .utils import subgoal_language | |
| from .utils.object_generation import spawn_fixed_cube, build_board_with_hole | |
| from .utils import reset_panda | |
| from .utils import subgoal_evaluate_func | |
| from ..logging_utils import logger | |
| PICK_CUBE_DOC_STRING = """**Task Description:** | |
| A simple task where the objective is to grasp a red cube with the {robot_id} robot and move it to a target goal position. This is also the *baseline* task to test whether a robot with manipulation | |
| capabilities can be simulated and trained properly. Hence there is extra code for some robots to set them up properly in this environment as well as the table scene builder. | |
| **Randomizations:** | |
| - the cube's xy position is randomized on top of a table in the region [0.1, 0.1] x [-0.1, -0.1]. It is placed flat on the table | |
| - the cube's z-axis rotation is randomized to a random angle | |
| - the target goal position (marked by a green sphere) of the cube has its xy position randomized in the region [0.1, 0.1] x [-0.1, -0.1] and z randomized in [0, 0.3] | |
| **Success Conditions:** | |
| - the cube position is within `goal_thresh` (default 0.025m) euclidean distance of the goal position | |
| - the robot is static (q velocity < 0.2) | |
| """ | |
| class InsertPeg(BaseEnv): | |
| _sample_video_link = "https://github.com/haosulab/ManiSkill/raw/main/figures/environment_demos/PickCube-v1_rt.mp4" | |
| SUPPORTED_ROBOTS = [ | |
| "panda", | |
| "fetch", | |
| "xarm6_robotiq", | |
| "so100", | |
| "widowxai", | |
| ] | |
| agent: Union[Panda] | |
| goal_thresh = 0.025 | |
| cube_spawn_half_size = 0.05 | |
| cube_spawn_center = (0, 0) | |
| _clearance = 0.01 | |
| def __init__(self, *args, robot_uids="panda_wristcam", robot_init_qpos_noise=0,seed=0,Robomme_video_episode=None,Robomme_video_path=None, | |
| **kwargs): | |
| self.use_demonstrationwrapper=False | |
| self.demonstration_record_traj=False | |
| self.robot_init_qpos_noise = robot_init_qpos_noise | |
| if robot_uids in PICK_CUBE_CONFIGS: | |
| cfg = PICK_CUBE_CONFIGS[robot_uids] | |
| else: | |
| cfg = PICK_CUBE_CONFIGS["panda"] | |
| self.cube_half_size = cfg["cube_half_size"] | |
| self.goal_thresh = cfg["goal_thresh"] | |
| self.cube_spawn_half_size = cfg["cube_spawn_half_size"] | |
| self.cube_spawn_center = cfg["cube_spawn_center"] | |
| self.max_goal_height = cfg["max_goal_height"] | |
| self.sensor_cam_eye_pos = cfg["sensor_cam_eye_pos"] | |
| self.sensor_cam_target_pos = cfg["sensor_cam_target_pos"] | |
| self.human_cam_eye_pos = cfg["human_cam_eye_pos"] | |
| self.human_cam_target_pos = cfg["human_cam_target_pos"] | |
| self.seed = seed | |
| self._hb_generator = torch.Generator() | |
| self._hb_generator.manual_seed(int(self.seed)) | |
| self.robomme_failure_recovery = bool( | |
| kwargs.pop("robomme_failure_recovery", False) | |
| ) | |
| self.robomme_failure_recovery_mode = kwargs.pop( | |
| "robomme_failure_recovery_mode", None | |
| ) | |
| if isinstance(self.robomme_failure_recovery_mode, str): | |
| self.robomme_failure_recovery_mode = ( | |
| self.robomme_failure_recovery_mode.lower() | |
| ) | |
| normalized_robomme_difficulty = normalize_robomme_difficulty( | |
| kwargs.pop("difficulty", None) | |
| ) | |
| if normalized_robomme_difficulty is not None: | |
| self.difficulty = normalized_robomme_difficulty | |
| else: | |
| seed_mod = seed % 3 | |
| if seed_mod == 0: | |
| self.difficulty = "easy" | |
| elif seed_mod == 1: | |
| self.difficulty = "medium" | |
| else: | |
| self.difficulty = "hard" | |
| self.restore_flag=False | |
| super().__init__(*args, robot_uids=robot_uids, **kwargs) | |
| def _default_sensor_configs(self): | |
| pose = sapien_utils.look_at( | |
| eye=self.sensor_cam_eye_pos, target=self.sensor_cam_target_pos | |
| ) | |
| camera_eye=[0.3,0,0.4] | |
| camera_target =[0,0,-0.2] | |
| pose = sapien_utils.look_at( | |
| eye=camera_eye, target=camera_target | |
| ) | |
| return [CameraConfig("base_camera", pose, 256, 256, np.pi / 2, 0.01, 100)] | |
| def _default_human_render_camera_configs(self): | |
| pose = sapien_utils.look_at( | |
| eye=self.human_cam_eye_pos, target=self.human_cam_target_pos | |
| ) | |
| return CameraConfig("render_camera", pose, 512, 512, 1, 0.01, 100) | |
| def _load_agent(self, options: dict): | |
| super()._load_agent(options, sapien.Pose(p=[-0.615, 0, 0])) | |
| def _load_scene(self, options: dict): | |
| self.table_scene = TableSceneBuilder( | |
| self, robot_init_qpos_noise=0 | |
| ) | |
| self.table_scene.build() | |
| length_tensor = torch.rand(1, generator=self._hb_generator) | |
| radius_tensor = torch.rand(1, generator=self._hb_generator) | |
| # self.length = (0.05 + (0.0 - 0.0) * length_tensor).item() | |
| # self.radius = (0.01 + (0.0 - 0.0) * radius_tensor).item() | |
| self.length = (0.05 + (0.01 - 0.01) * length_tensor).item() | |
| self.radius = (0.01 + (0.005 - 0.005) * radius_tensor).item() | |
| # Create 3 identical pegs with different x-axis coordinates | |
| self.pegs = [] | |
| self.peg_heads = [] | |
| self.peg_tails = [] | |
| self._peg_initial_poses = [] | |
| offsets = [0.1,0,-0.1] # X-axis differences for the 3 pegs | |
| # Sample a single pair of colors so all pegs share the same appearance per seed. | |
| peg_head_color = torch.rand(3, generator=self._hb_generator).tolist() | |
| # Use complementary tail color so head/tail are contrasting. | |
| peg_tail_color = [1.0 - c for c in peg_head_color] | |
| for offset in offsets: | |
| peg_spawn_translation = np.array([self.length / 2 , -0.15-offset, self.radius], dtype=np.float32) | |
| #initial_yaw = (torch.rand(1, generator=self._hb_generator).item() * 2 * np.pi) - np.pi | |
| initial_yaw = 0 | |
| yaw_angles = torch.tensor([[0.0, 0.0, initial_yaw]], dtype=torch.float32) | |
| yaw_matrix = euler_angles_to_matrix(yaw_angles, convention="XYZ") | |
| yaw_quat = matrix_to_quaternion(yaw_matrix)[0].detach().cpu().numpy().tolist() | |
| peg_initial_pose = sapien.Pose( | |
| p=peg_spawn_translation.tolist(), | |
| q=yaw_quat, | |
| ) | |
| peg, peg_head, peg_tail = build_peg( | |
| self, | |
| length=self.length, | |
| radius=self.radius, | |
| initial_pose=peg_initial_pose, | |
| name=f"peg_{len(self.pegs)}", | |
| head_color=peg_head_color, | |
| tail_color=peg_tail_color, | |
| ) | |
| self.pegs.append(peg) | |
| self.peg_heads.append(peg_head) | |
| self.peg_tails.append(peg_tail) | |
| self._peg_initial_poses.append(peg_initial_pose) | |
| # Randomly select one peg from the 3 pegs | |
| random_peg_idx = int(torch.randint(0, 3, (1,), generator=self._hb_generator).item()) | |
| random_peg_idx=0 | |
| self.peg = self.pegs[random_peg_idx] | |
| self.peg_head = self.peg_heads[random_peg_idx] | |
| self.peg_tail = self.peg_tails[random_peg_idx] | |
| self._peg_initial_pose = self._peg_initial_poses[random_peg_idx] | |
| self.box=build_box_with_hole(self,inner_radius=self.radius*1.7,outer_radius=self.radius*4,depth=self.length,center=[0,0]) | |
| self.reset_in_proecess=False | |
| def _initialize_episode(self, env_idx: torch.Tensor, options: dict): | |
| with torch.device(self.device): | |
| self.end_steps=None | |
| self.table_scene.initialize(env_idx) | |
| # Reset highlight state at the start of each episode | |
| self._insert_highlight_start = None | |
| self._insert_highlight_active = False | |
| if not hasattr(self, "pegs"): | |
| return | |
| base_translation = (0, 0) | |
| x_jitter_2 = (torch.rand(1, generator=self._hb_generator).item() - 0.5) * 0.2 | |
| y_jitter_2 = (torch.rand(1, generator=self._hb_generator).item() - 0.5) * 0.2 | |
| # x_jitter_2=0 | |
| # y_jitter_2=0 | |
| box_translation = [base_translation[0] + x_jitter_2, base_translation[1] + y_jitter_2, self.radius * 4] | |
| box_yaw = np.pi / 2 + (torch.rand(1, generator=self._hb_generator).item() * 2 - 1) * np.radians(20) | |
| box_angles = torch.tensor([[0.0, 0.0, box_yaw]], dtype=torch.float32) | |
| box_matrix = euler_angles_to_matrix(box_angles, convention="XYZ") | |
| box_quat = matrix_to_quaternion(box_matrix)[0].detach().cpu().numpy().tolist() | |
| self.box.set_pose(sapien.Pose(p=box_translation, q=box_quat)) | |
| box_xy = np.array(box_translation[:2], dtype=np.float32) | |
| sampled_xy_positions = [] | |
| max_sampling_attempts = 512 | |
| #Initialize all 3 pegs with constrained random placements | |
| for i, peg in enumerate(self.pegs): | |
| candidate_xy = None | |
| for _ in range(max_sampling_attempts): | |
| x_sample = (torch.rand(1, generator=self._hb_generator).item() * 0.4) - 0.2 | |
| y_sample = (torch.rand(1, generator=self._hb_generator).item() * 0.6) - 0.3 | |
| sampled_xy = np.array([x_sample, y_sample], dtype=np.float32) | |
| if np.linalg.norm(sampled_xy - box_xy) <= self.radius * 6: | |
| continue | |
| if any(np.linalg.norm(sampled_xy - prev_xy) <= self.length * 1.5 for prev_xy in sampled_xy_positions): | |
| continue | |
| candidate_xy = sampled_xy | |
| break | |
| if candidate_xy is None: | |
| raise RuntimeError("Failed to sample peg positions satisfying placement constraints.") | |
| yaw_value = (torch.rand(1, generator=self._hb_generator).item() * 2 - 1) * np.radians(45) | |
| yaw_angles = torch.tensor([[0.0, 0.0, yaw_value]], dtype=torch.float32) | |
| yaw_matrix = euler_angles_to_matrix(yaw_angles, convention="XYZ") | |
| yaw_quat = matrix_to_quaternion(yaw_matrix)[0].detach().cpu().numpy().tolist() | |
| pose = sapien.Pose(p=[float(candidate_xy[0]), float(candidate_xy[1]), 0.0], q=yaw_quat) | |
| peg.set_pose(pose) | |
| sampled_xy_positions.append(candidate_xy) | |
| # Store initial poses for all pegs | |
| self.peg_init_poses = [] | |
| for peg in self.pegs: | |
| pose = peg.pose | |
| pose_p = np.asarray(pose.p, dtype=np.float32).reshape(-1).copy() | |
| pose_q = np.asarray(pose.q, dtype=np.float32).reshape(-1).copy() | |
| self.peg_init_poses.append(sapien.Pose(p=pose_p, q=pose_q)) | |
| # robomme-v2.7/robomme/robomme_env/PickPeg.py:243 | |
| pose = self.peg.pose | |
| pose_p = np.asarray(pose.p, dtype=np.float32).reshape(-1).copy() | |
| pose_q = np.asarray(pose.q, dtype=np.float32).reshape(-1).copy() | |
| self.peg_init_pose = sapien.Pose(p=pose_p, q=pose_q) | |
| self.peg_init_pose = sapien.Pose(p=pose_p, q=pose_q) | |
| # Define task list, each task contains a dictionary with function, name, demonstration flag, and optional failure_func | |
| obj_sample = torch.randint(0, 2, (1,), generator=self._hb_generator) | |
| dir_sample = torch.randint(0, 2, (1,), generator=self._hb_generator) | |
| self.obj_flag = -1 if obj_sample.item() == 0 else 1 | |
| self.direction = -1 if dir_sample.item() == 0 else 1 | |
| # if self.seed<30: | |
| # self.obj_flag=-1 | |
| # self.direction=1 | |
| # elif self.seed<60: | |
| # self.obj_flag=-1 | |
| # self.direction=-1 | |
| # elif self.seed<90: | |
| # self.obj_flag=1 | |
| # self.direction=1 | |
| # #elif self.seed<60: | |
| # else: | |
| # self.obj_flag=1 | |
| # self.direction=-1 | |
| # self.obj_flag=1 | |
| # self.direction=-1 | |
| qpos = np.array( | |
| [ | |
| 0.0, | |
| 0, | |
| 0, | |
| -np.pi * 4 / 8, | |
| 0, | |
| np.pi * 2 / 4, | |
| np.pi / 4, | |
| 0.04, | |
| 0.04, | |
| ], | |
| dtype=np.float32, | |
| ) | |
| self.agent.reset(qpos) | |
| if self.obj_flag==-1: | |
| self.grasp_target=self.peg_head | |
| self.insert_target=self.peg_tail | |
| else: | |
| self.grasp_target=self.peg_tail | |
| self.insert_target=self.peg_head | |
| agent_x = self.agent.robot.pose.p.tolist()[0][0] | |
| head_x = float(self.peg_head.pose.p.tolist()[0][0]) | |
| tail_x = float(self.peg_tail.pose.p.tolist()[0][0]) | |
| logger.debug(f"agent_x: {agent_x}, head_x: {head_x}, tail_x: {tail_x}") | |
| near_link = self.peg_head if abs(head_x - agent_x) <= abs(tail_x - agent_x) else self.peg_tail | |
| self.grasp_target_distance = "near" if self.grasp_target is near_link else "far" | |
| logger.debug(f"grasp_target_distance: {self.grasp_target_distance}") | |
| self.insert_way="left" if self.direction == -1 else "right" | |
| tasks = [ | |
| { | |
| "func": lambda: is_A_pickup_notB(self, self.grasp_target, self.insert_target), | |
| "name": f"Pick up the peg by grasping the {self.grasp_target_distance} end", | |
| "subgoal_segment":f"Pick up the peg by grasping the {self.grasp_target_distance} end at <>", | |
| "choice_label": "pick up the peg by grasping one end", | |
| "demonstration": True, | |
| "failure_func": lambda: is_A_pickup_notB(self, self.insert_target, self.grasp_target), | |
| "solve": lambda env, planner: grasp_and_lift_peg_side(env, planner, env.grasp_target), | |
| "segment":self.grasp_target | |
| }, | |
| { | |
| "func": lambda: is_A_insert_notB(self, self.insert_target, self.grasp_target, self.box,direction=self.direction), | |
| "name": f"Insert the peg from the {self.insert_way} side of the box", | |
| "subgoal_segment":f"Insert the peg from the {self.insert_way} side of the box at <>", | |
| "choice_label": f"insert the peg from the {self.insert_way} side", | |
| "demonstration": True, | |
| "failure_func": None, | |
| "solve": lambda env, planner: insert_peg(env, planner, direction=self.direction,obj=self.obj_flag,insert_obj=self.insert_target), | |
| "segment":self.box | |
| }, | |
| { | |
| "func": lambda:reset_check(self), | |
| "name": "NO RECORD", | |
| "subgoal_segment":f"NO RECORD", | |
| "demonstration": True, | |
| "failure_func": None, | |
| "specialflag":"reset pegs", | |
| "solve": lambda env, planner: [solve_strong_reset(env,planner)], | |
| }, | |
| { | |
| "func": lambda: static_check(self, timestep=int(self.elapsed_steps), static_steps=100), | |
| "name": "NO RECORD", | |
| "subgoal_segment":"NO RECORD", | |
| "demonstration": True, | |
| "failure_func": None, | |
| "solve": lambda env, planner: [solve_hold_obj(env, planner, static_steps=100,close=False)], | |
| }, | |
| { | |
| "func": lambda: is_A_pickup_notB(self, self.grasp_target, self.insert_target), | |
| "name": f"Pick up the peg by grasping the {self.grasp_target_distance} end", | |
| "subgoal_segment":f"Pick up the peg by grasping the {self.grasp_target_distance} end at <>", | |
| "choice_label": "pick up the peg by grasping one end", | |
| "demonstration": False, | |
| "failure_func": lambda: [ | |
| is_A_pickup_notB(self, self.insert_target, self.grasp_target), | |
| is_any_obj_pickup(self, [head for i, head in enumerate(self.peg_heads) if self.pegs[i] is not self.peg] + | |
| [tail for i, tail in enumerate(self.peg_tails) if self.pegs[i] is not self.peg]) | |
| ], | |
| "solve": lambda env, planner: grasp_and_lift_peg_side(env, planner, env.grasp_target), | |
| "segment":self.grasp_target | |
| }, | |
| { | |
| "func": lambda: is_A_insert_notB(self, self.insert_target, self.grasp_target,self.box,direction=self.direction,mark_end_flag=True), | |
| "name": f"Insert the peg from the {self.insert_way} side", | |
| "subgoal_segment":f"Insert the peg from the {self.insert_way} side at <>", | |
| "choice_label": f"insert the peg from the {self.insert_way} side", | |
| "demonstration": False, | |
| "failure_func": lambda: [ | |
| is_A_insert_notB(self, self.grasp_target, self.insert_target, self.box), | |
| is_A_insert_notB(self, self.insert_target, self.grasp_target, self.box, direction=-self.direction), | |
| is_any_obj_pickup(self, [head for i, head in enumerate(self.peg_heads) if self.pegs[i] is not self.peg] + | |
| [tail for i, tail in enumerate(self.peg_tails) if self.pegs[i] is not self.peg]) | |
| ], | |
| "solve": lambda env, planner: insert_peg(env, planner, direction=self.direction,obj=self.obj_flag,insert_obj=self.insert_target,cut_retreat=True), | |
| "segment":self.box | |
| }, | |
| ] | |
| # Store task list for RecordWrapper use | |
| self.task_list = tasks | |
| def evaluate(self,solve_complete_eval=False): | |
| timestep = self.elapsed_steps | |
| self.successflag=torch.tensor([False]) | |
| self.failureflag = torch.tensor([False]) | |
| # Use encapsulated sequence task check function | |
| if(self.use_demonstrationwrapper==False):# change subgoal after planner ends during recording | |
| if solve_complete_eval==True: | |
| allow_subgoal_change_this_timestep=True | |
| else: | |
| allow_subgoal_change_this_timestep=False | |
| else:# during demonstration, video needs to call evaluate(solve_complete_eval), video ends and flag changes in demonstrationwrapper | |
| if solve_complete_eval==True or self.demonstration_record_traj==False: | |
| allow_subgoal_change_this_timestep=True | |
| else: | |
| allow_subgoal_change_this_timestep=False | |
| all_tasks_completed, current_task_name, task_failed,_ = sequential_task_check(self, self.task_list,allow_subgoal_change_this_timestep=allow_subgoal_change_this_timestep) | |
| if self.end_steps!=None:# truncate tail, also truncate tail in planner | |
| logger.debug( | |
| "elapsed_steps=%s, end_steps=%s", | |
| self.elapsed_steps, | |
| self.end_steps, | |
| ) | |
| if int(getattr(self, "elapsed_steps", 0))>=self.end_steps+3: | |
| self.successflag = torch.tensor([True]) | |
| if task_failed: | |
| self.failureflag = torch.tensor([True]) | |
| logger.debug(f"Task failed: {current_task_name}") | |
| # If static_check succeeds or all tasks completed, set success flag | |
| if all_tasks_completed and not task_failed: | |
| self.successflag = torch.tensor([True]) | |
| return { | |
| "success": self.successflag, | |
| "fail": self.failureflag, | |
| } | |
| def compute_dense_reward(self, obs: Any, action: torch.Tensor, info: Dict): | |
| tcp_to_obj_dist = torch.linalg.norm( | |
| self.agent.tcp_pose.p - self.agent.tcp_pose.p, axis=1 | |
| ) | |
| reaching_reward = 1 - torch.tanh(5 * tcp_to_obj_dist) | |
| reward = reaching_reward*0 | |
| return reward | |
| def compute_normalized_dense_reward( | |
| self, obs: Any, action: torch.Tensor, info: Dict | |
| ): | |
| return self.compute_dense_reward(obs=obs, action=action, info=info) / 5 | |
| #Robomme | |
| def step(self, action: Union[None, np.ndarray, torch.Tensor, Dict]): | |
| obs, reward, terminated, truncated, info = super().step(action) | |
| cur_step = int(self.elapsed_steps[0].item()) | |
| if self.reset_in_proecess==True: | |
| for i, peg in enumerate(self.pegs): | |
| peg.set_pose(self.peg_init_poses[i]) | |
| if peg.dof > 0: | |
| zero = np.zeros(peg.dof) | |
| peg.set_qpos(zero) | |
| peg.set_qvel(zero) | |
| logger.debug("reset peg!") | |
| if is_A_insert_notB(self, self.insert_target, self.grasp_target, self.box,direction=self.direction): | |
| self.start_step=cur_step | |
| color=sapien.render.RenderMaterial( | |
| base_color=sapien_utils.hex2rgba("#FFD289"), roughness=0.5, specular=0.5) | |
| if getattr(self, "start_step", None) is not None : | |
| if cur_step <= self.start_step + 20 and cur_step>= self.start_step: | |
| color=[1.0, 0.0, 0.0, 1.0] | |
| highlight_obj( | |
| self, | |
| self.box, | |
| start_step= 0, | |
| end_step= 99999, | |
| cur_step=cur_step, | |
| disk_radius=0.015, | |
| disk_half_length=0.055, | |
| highlight_color=color,) | |
| return obs, reward, terminated, truncated, info | |