[Feature] Massively improve success rate of pd joint pos / delta pos …

…conversion to pd ee delta pos/pose controller (#549)

* work

* Update record.py

* Update replay_trajectory.py

mani_skill/examples/motionplanning/panda/run.py

@@ -57,13 +57,18 @@ def main(args):
     pbar = tqdm(range(args.num_traj))
     seed = 0
     successes = []
+    solution_episode_lengths = []
+    failed_motion_plans = 0
     passed = 0
     while True:
         res = solve(env, seed=seed, debug=False, vis=True if args.vis else False)
         if res == -1:
             success = False
+            failed_motion_plans += 1
         else:
             success = res[-1]["success"].item()
+            elapsed_steps = res[-1]["elapsed_steps"].item()
+            solution_episode_lengths.append(elapsed_steps)
         successes.append(success)
         if args.only_count_success and not success:
             seed += 1
@@ -76,7 +81,15 @@ def main(args):
             if args.save_video:
                 env.flush_video()
             pbar.update(1)
-            pbar.set_postfix(dict(success_rate=np.mean(successes)))
+            pbar.set_postfix(
+                dict(
+                    success_rate=np.mean(successes),
+                    failed_motion_plan_rate=failed_motion_plans / (seed + 1),
+                    avg_episode_length=np.mean(solution_episode_lengths),
+                    max_episode_length=np.max(solution_episode_lengths),
+                    # min_episode_length=np.min(solution_episode_lengths)
+                )
+            )
             seed += 1
             passed += 1
             if passed == args.num_traj:

mani_skill/examples/motionplanning/panda/solutions/peg_insertion_side.py

@@ -1,6 +1,6 @@
 import gymnasium as gym
 import numpy as np
-import sapien.core as sapien
+import sapien
 
 from mani_skill.envs.tasks import PegInsertionSideEnv
 from mani_skill.examples.motionplanning.panda.motionplanner import \
@@ -36,8 +36,8 @@ def solve(env: PegInsertionSideEnv, seed=None, debug=False, vis=False):
         base_pose=env.unwrapped.agent.robot.pose,
         visualize_target_grasp_pose=vis,
         print_env_info=False,
-        joint_vel_limits=0.5,
-        joint_acc_limits=0.5,
+        joint_vel_limits=0.75,
+        joint_acc_limits=0.75,
     )
     env = env.unwrapped
     FINGER_LENGTH = 0.025
@@ -60,12 +60,13 @@ def solve(env: PegInsertionSideEnv, seed=None, debug=False, vis=False):
     # Reach
     # -------------------------------------------------------------------------- #
     reach_pose = grasp_pose * (sapien.Pose([0, 0, -0.05]))
-    planner.move_to_pose_with_screw(reach_pose)
-
+    res = planner.move_to_pose_with_screw(reach_pose)
+    if res == -1: return res
     # -------------------------------------------------------------------------- #
     # Grasp
     # -------------------------------------------------------------------------- #
-    planner.move_to_pose_with_screw(grasp_pose)
+    res = planner.move_to_pose_with_screw(grasp_pose)
+    if res == -1: return res
     planner.close_gripper()
 
     # -------------------------------------------------------------------------- #
@@ -76,17 +77,20 @@ def solve(env: PegInsertionSideEnv, seed=None, debug=False, vis=False):
     insert_pose = env.goal_pose * peg_init_pose.inv() * grasp_pose
     offset = sapien.Pose([-0.01 - env.peg_half_sizes[0, 0], 0, 0])
     pre_insert_pose = insert_pose * (offset)
-    planner.move_to_pose_with_screw(pre_insert_pose)
+    res = planner.move_to_pose_with_screw(pre_insert_pose)
+    if res == -1: return res
     # refine the insertion pose
     for i in range(3):
         delta_pose = env.goal_pose * (offset) * env.peg.pose.inv()
         pre_insert_pose = delta_pose * pre_insert_pose
-        planner.move_to_pose_with_screw(pre_insert_pose)
+        res = planner.move_to_pose_with_screw(pre_insert_pose)
+        if res == -1: return res
 
     # -------------------------------------------------------------------------- #
     # Insert
     # -------------------------------------------------------------------------- #
     res = planner.move_to_pose_with_screw(insert_pose * (sapien.Pose([0.05, 0, 0])))
+    if res == -1: return res
     planner.close()
     return res
 

mani_skill/trajectory/replay_trajectory.py

@@ -26,6 +26,7 @@
 from mani_skill.trajectory.merge_trajectory import merge_h5
 from mani_skill.utils import common, gym_utils, io_utils, wrappers
 from mani_skill.utils.structs.link import Link
+from mani_skill.utils.structs.pose import Pose
 
 
 def qpos_to_pd_joint_delta_pos(controller: PDJointPosController, qpos):
@@ -112,10 +113,10 @@ def from_pd_joint_pos_to_ee(
         "root_translation:root_aligned_body_rotation",
         "root_translation",
     ], "Currently only support the 'root_translation:root_aligned_body_rotation' ee control frame for delta pose control and 'root_translation' ee control frame for delta pos control"
-    ori_ee_link = ori_env.agent.robot.links_map[arm_controller.ee_link.name]
+    ori_env.agent.robot.links_map[arm_controller.ee_link.name]
     ee_link: Link = arm_controller.ee_link
     pos_only = arm_controller.config.frame == "root_translation"
-
+    pin_model = ori_controller.articulation.create_pinocchio_model()
     info = {}
 
     for t in range(n):
@@ -130,15 +131,28 @@ def from_pd_joint_pos_to_ee(
             ori_action_dict.copy(), device=env.device
         )  # do not in-place modify
         ori_env.step(ori_action)
-        flag = True
 
+        # NOTE (stao): for high success rate of pd joint pos to pd ee delta pos/pose control, we need to use the current target qpos of the original
+        # environment controller to compute the target ee pose to try and reach. this is because if we attempt to reach the original envs ee link pose
+        # we may fall short and fail.
+        full_qpos = ori_controller.articulation.get_qpos()
+        full_qpos[
+            :, ori_arm_controller.active_joint_indices
+        ] = ori_arm_controller._target_qpos
+        pin_model.compute_forward_kinematics(full_qpos.cpu().numpy()[0])
+        target_ee_pose_pin = Pose.create(
+            ori_controller.articulation.pose.sp
+            * pin_model.get_link_pose(arm_controller.ee_link.index)
+        )
+        flag = True
         for _ in range(4):
             delta_q = [1, 0, 0, 0]
             if "root_translation" in arm_controller.config.frame:
-                delta_position = ori_ee_link.pose.p - ee_link.pose.p
+                delta_position = target_ee_pose_pin.p - ee_link.pose.p
             if "root_aligned_body_rotation" in arm_controller.config.frame:
-                delta_q = (ee_link.pose * ori_ee_link.pose.inv()).q.cpu().numpy()[0]
+                delta_q = (ee_link.pose.sp * target_ee_pose_pin.sp.inv()).q
             delta_pose = sapien.Pose(delta_position.cpu().numpy()[0], delta_q)
+
             arm_action = delta_pose_to_pd_ee_delta(
                 arm_controller, delta_pose, pos_only=pos_only
             )
@@ -162,7 +176,6 @@ def from_pd_joint_pos_to_ee(
 
             if flag:
                 break
-
     return info
 
 

mani_skill/utils/wrappers/record.py

@@ -354,7 +354,9 @@ def reset(
             self._trajectory_buffer = None
         if self.save_trajectory:
             state_dict = self.base_env.get_state_dict()
-            action = common.batch(self.single_action_space.sample())
+            action = common.batch(
+                self.env.get_wrapper_attr("single_action_space").sample()
+            )
             first_step = Step(
                 state=common.to_numpy(common.batch(state_dict)),
                 observation=common.to_numpy(common.batch(obs)),