| """ |
| MultiStepDemonstrationWrapper: Wraps DemonstrationWrapper to provide waypoint step interface. |
| |
| Each step(action) receives action = waypoint_p(3) + rpy(3) + gripper_action(1), total 7 dimensions. |
| Internally converts RPY to quat then calls move_to_pose_with_screw and close_gripper/open_gripper via planner_denseStep, |
| where PatternLock/RouteStick will force skip close_gripper/open_gripper. |
| Returns obs as dictionary-of-lists, and reward/terminated/truncated as the last step value. |
| Caller must ensure scripts/ is in sys.path to import planner_fail_safe. |
| """ |
| import numpy as np |
| import sapien |
| import torch |
| import gymnasium as gym |
|
|
| from ..robomme_env.utils import planner_denseStep |
| from ..robomme_env.utils.rpy_util import rpy_xyz_to_quat_wxyz_torch |
| from ..robomme_env.utils.planner_fail_safe import ScrewPlanFailure |
|
|
| DATASET_SCREW_MAX_ATTEMPTS = 3 |
| DATASET_RRT_MAX_ATTEMPTS = 3 |
|
|
|
|
| class RRTPlanFailure(RuntimeError): |
| """Raised when move_to_pose_with_RRTStar returns -1 (planning failed).""" |
|
|
|
|
| class MultiStepDemonstrationWrapper(gym.Wrapper): |
| """ |
| Wraps DemonstrationWrapper. step(action) interprets action as |
| (waypoint_p, rpy, gripper_action) total 7 dims, internally converts RPY to quat, |
| executes planning via planner_denseStep, and returns last-step signals. |
| """ |
|
|
| def __init__(self, env, gui_render=True, vis=True, **kwargs): |
| super().__init__(env) |
| self._planner = None |
| self._gui_render = gui_render |
| self._vis = vis |
| self.action_space = gym.spaces.Box( |
| low=-np.inf, high=np.inf, shape=(7,), dtype=np.float64 |
| ) |
|
|
| @staticmethod |
| def _batch_to_steps(batch): |
| obs_batch, reward_batch, terminated_batch, truncated_batch, info_batch = batch |
| n = int(reward_batch.numel()) |
| steps = [] |
| obs_keys = list(obs_batch.keys()) |
| info_keys = list(info_batch.keys()) |
| for idx in range(n): |
| obs = {k: obs_batch[k][idx] for k in obs_keys} |
| info = {k: info_batch[k][idx] for k in info_keys} |
| reward = reward_batch[idx] |
| terminated = terminated_batch[idx] |
| truncated = truncated_batch[idx] |
| steps.append((obs, reward, terminated, truncated, info)) |
| return steps |
|
|
| @staticmethod |
| def _flatten_info_batch(info_batch): |
| return {k: v[-1] if isinstance(v, list) and v else v for k, v in info_batch.items()} |
|
|
| @staticmethod |
| def _take_last_step_value(value): |
| if isinstance(value, torch.Tensor): |
| if value.numel() == 0 or value.ndim == 0: |
| return value |
| return value.reshape(-1)[-1] |
| if isinstance(value, np.ndarray): |
| if value.size == 0 or value.ndim == 0: |
| return value |
| return value.reshape(-1)[-1] |
| if isinstance(value, (list, tuple)): |
| return value[-1] if value else value |
| return value |
|
|
| def _get_planner(self): |
| if self._planner is not None: |
| return self._planner |
| from ..robomme_env.utils.planner_fail_safe import ( |
| FailAwarePandaArmMotionPlanningSolver, |
| FailAwarePandaStickMotionPlanningSolver, |
| ) |
|
|
| env_id = self.env.unwrapped.spec.id |
| base_pose = self.env.unwrapped.agent.robot.pose |
| if env_id in ("PatternLock", "RouteStick"): |
| self._planner = FailAwarePandaStickMotionPlanningSolver( |
| self.env, |
| debug=False, |
| vis=self._vis, |
| base_pose=base_pose, |
| visualize_target_grasp_pose=False, |
| print_env_info=False, |
| joint_vel_limits=0.3, |
| ) |
| else: |
| self._planner = FailAwarePandaArmMotionPlanningSolver( |
| self.env, |
| debug=False, |
| vis=self._vis, |
| base_pose=base_pose, |
| visualize_target_grasp_pose=True, |
| print_env_info=False, |
| ) |
| return self._planner |
|
|
| def _current_tcp_p(self): |
| current_pose = self.env.unwrapped.agent.tcp.pose |
| p = current_pose.p |
| if hasattr(p, "cpu"): |
| p = p.cpu().numpy() |
| p = np.asarray(p).flatten() |
| return p |
|
|
| def _no_op_step(self): |
| """Execute one step using current qpos + gripper, without moving arm, only to get observation.""" |
| robot = self.env.unwrapped.agent.robot |
| qpos = robot.get_qpos().cpu().numpy().flatten() |
| arm = qpos[:7] |
| gripper = float(qpos[7]) if len(qpos) > 7 else 0.0 |
| action = np.hstack([arm, gripper]) |
| return self.env.step(action) |
|
|
| def step(self, action): |
| """Execute waypoint step and return last-step signals for reward/terminated/truncated.""" |
| action = np.asarray(action, dtype=np.float64).flatten() |
| if action.size < 7: |
| raise ValueError(f"action must have at least 7 elements, got {action.size}") |
| waypoint_p = action[:3] |
| rpy = action[3:6] |
| gripper_action = float(action[6]) |
|
|
| |
| rpy_t = torch.as_tensor(rpy, dtype=torch.float64) |
| waypoint_q = rpy_xyz_to_quat_wxyz_torch(rpy_t).numpy() |
|
|
| pose = sapien.Pose(p=waypoint_p, q=waypoint_q) |
| planner = self._get_planner() |
| is_stick_env = self.env.unwrapped.spec.id in ("PatternLock", "RouteStick") |
|
|
| current_p = self._current_tcp_p() |
| dist = np.linalg.norm(current_p - waypoint_p) |
|
|
| collected_steps = [] |
| |
| |
| move_steps = -1 |
| for attempt in range(1, DATASET_SCREW_MAX_ATTEMPTS + 1): |
| try: |
| result = planner_denseStep._collect_dense_steps( |
| planner, lambda: planner.move_to_pose_with_screw(pose) |
| ) |
| except ScrewPlanFailure as exc: |
| print(f"[MultiStep] screw planning failed (attempt {attempt}/{DATASET_SCREW_MAX_ATTEMPTS}): {exc}") |
| continue |
| |
| if isinstance(result, int) and result == -1: |
| print(f"[MultiStep] screw planning returned -1 (attempt {attempt}/{DATASET_SCREW_MAX_ATTEMPTS})") |
| continue |
|
|
| move_steps = result |
| break |
|
|
| if move_steps == -1: |
| print(f"[MultiStep] screw planning exhausted; fallback to RRT* (max {DATASET_RRT_MAX_ATTEMPTS} attempts)") |
| for attempt in range(1, DATASET_RRT_MAX_ATTEMPTS + 1): |
| try: |
| result = planner_denseStep._collect_dense_steps( |
| planner, lambda: planner.move_to_pose_with_RRTStar(pose) |
| ) |
| except Exception as exc: |
| print(f"[MultiStep] RRT* planning failed (attempt {attempt}/{DATASET_RRT_MAX_ATTEMPTS}): {exc}") |
| continue |
|
|
| if isinstance(result, int) and result == -1: |
| print(f"[MultiStep] RRT* planning returned -1 (attempt {attempt}/{DATASET_RRT_MAX_ATTEMPTS})") |
| continue |
|
|
| move_steps = result |
| break |
|
|
| if move_steps == -1: |
| raise RRTPlanFailure("Both screw and RRTStar planning exhausted.") |
| collected_steps.extend(move_steps) |
|
|
| |
| if not is_stick_env: |
| if gripper_action == -1: |
| if hasattr(planner, "close_gripper"): |
| result = planner_denseStep.close_gripper(planner) |
| if result != -1: |
| collected_steps.extend(self._batch_to_steps(result)) |
| elif gripper_action == 1: |
| if hasattr(planner, "open_gripper"): |
| result = planner_denseStep.open_gripper(planner) |
| if result != -1: |
| collected_steps.extend(self._batch_to_steps(result)) |
|
|
| obs_batch, reward_batch, terminated_batch, truncated_batch, info_batch = planner_denseStep.to_step_batch( |
| collected_steps |
| ) |
| info_flat = self._flatten_info_batch(info_batch) |
| return ( |
| obs_batch, |
| self._take_last_step_value(reward_batch), |
| self._take_last_step_value(terminated_batch), |
| self._take_last_step_value(truncated_batch), |
| info_flat, |
| ) |
|
|
| def reset(self, **kwargs): |
| self._planner = None |
| return self.env.reset(**kwargs) |
|
|
| def close(self): |
| self._planner = None |
| return self.env.close() |
|
|
