From da6e056a4b49bd4f5e31f57c1b140ca72f97485d Mon Sep 17 00:00:00 2001
From: Stone Tao <stonezt2019@gmail.com>
Date: Tue, 16 Jul 2024 11:39:05 -0700
Subject: [PATCH] Update motionplanner.py

---
 .../motionplanning/panda/motionplanner.py     | 37 ++++++++++---------
 1 file changed, 19 insertions(+), 18 deletions(-)

diff --git a/mani_skill/examples/motionplanning/panda/motionplanner.py b/mani_skill/examples/motionplanning/panda/motionplanner.py
index 2608fb58b..52fe24f6c 100644
--- a/mani_skill/examples/motionplanning/panda/motionplanner.py
+++ b/mani_skill/examples/motionplanning/panda/motionplanner.py
@@ -25,7 +25,8 @@ def __init__(
         joint_acc_limits=0.9,
     ):
         self.env = env
-        self.env_agent: BaseAgent = self.env.unwrapped.agent
+        self.base_env: BaseEnv = env.unwrapped
+        self.env_agent: BaseAgent = self.base_env.agent
         self.robot = self.env_agent.robot
         self.joint_vel_limits = joint_vel_limits
         self.joint_acc_limits = joint_acc_limits
@@ -33,7 +34,7 @@ def __init__(
         self.base_pose = to_sapien_pose(base_pose)
 
         self.planner = self.setup_planner()
-        self.control_mode = self.env.unwrapped.control_mode
+        self.control_mode = self.base_env.control_mode
 
         self.debug = debug
         self.vis = vis
@@ -42,10 +43,13 @@ def __init__(
         self.gripper_state = OPEN
         self.grasp_pose_visual = None
         if self.vis and self.visualize_target_grasp_pose:
-            self.grasp_pose_visual = build_panda_gripper_grasp_pose_visual(
-                env.unwrapped.scene
-            )
-            self.grasp_pose_visual.set_pose(env.unwrapped.agent.tcp.pose)
+            if "grasp_pose_visual" not in self.base_env.scene.actors:
+                self.grasp_pose_visual = build_panda_gripper_grasp_pose_visual(
+                    self.base_env.scene
+                )
+            else:
+                self.grasp_pose_visual = self.base_env.scene.actors["grasp_pose_visual"]
+            self.grasp_pose_visual.set_pose(self.base_env.agent.tcp.pose)
         self.elapsed_steps = 0
 
         self.use_point_cloud = False
@@ -56,11 +60,11 @@ def render_wait(self):
         if not self.vis or not self.debug:
             return
         print("Press [c] to continue")
-        viewer = self.env.unwrapped.render_human()
+        viewer = self.base_env.render_human()
         while True:
             if viewer.window.key_down("c"):
                 break
-            self.env.unwrapped.render_human()
+            self.base_env.render_human()
 
     def setup_planner(self):
         link_names = [link.get_name() for link in self.robot.get_links()]
@@ -93,7 +97,7 @@ def follow_path(self, result, refine_steps: int = 0):
                     f"[{self.elapsed_steps:3}] Env Output: reward={reward} info={info}"
                 )
             if self.vis:
-                self.env.unwrapped.render_human()
+                self.base_env.render_human()
         return obs, reward, terminated, truncated, info
 
     def move_to_pose_with_RRTConnect(
@@ -105,7 +109,7 @@ def move_to_pose_with_RRTConnect(
         result = self.planner.plan_qpos_to_pose(
             np.concatenate([pose.p, pose.q]),
             self.robot.get_qpos().cpu().numpy()[0],
-            time_step=self.env.unwrapped.control_timestep,
+            time_step=self.base_env.control_timestep,
             use_point_cloud=self.use_point_cloud,
             wrt_world=True,
         )
@@ -129,14 +133,14 @@ def move_to_pose_with_screw(
         result = self.planner.plan_screw(
             np.concatenate([pose.p, pose.q]),
             self.robot.get_qpos().cpu().numpy()[0],
-            time_step=self.env.unwrapped.control_timestep,
+            time_step=self.base_env.control_timestep,
             use_point_cloud=self.use_point_cloud,
         )
         if result["status"] != "Success":
             result = self.planner.plan_screw(
                 np.concatenate([pose.p, pose.q]),
                 self.robot.get_qpos().cpu().numpy()[0],
-                time_step=self.env.unwrapped.control_timestep,
+                time_step=self.base_env.control_timestep,
                 use_point_cloud=self.use_point_cloud,
             )
             if result["status"] != "Success":
@@ -163,7 +167,7 @@ def open_gripper(self):
                     f"[{self.elapsed_steps:3}] Env Output: reward={reward} info={info}"
                 )
             if self.vis:
-                self.env.unwrapped.render_human()
+                self.base_env.render_human()
         return obs, reward, terminated, truncated, info
 
     def close_gripper(self, t=6):
@@ -181,7 +185,7 @@ def close_gripper(self, t=6):
                     f"[{self.elapsed_steps:3}] Env Output: reward={reward} info={info}"
                 )
             if self.vis:
-                self.env.unwrapped.render_human()
+                self.base_env.render_human()
         return obs, reward, terminated, truncated, info
 
     def add_box_collision(self, extents: np.ndarray, pose: sapien.Pose):
@@ -206,10 +210,7 @@ def clear_collisions(self):
         self.use_point_cloud = False
 
     def close(self):
-        if self.grasp_pose_visual is not None:
-            if not physx.is_gpu_enabled():
-                self.grasp_pose_visual.remove_from_scene()
-
+        pass
 
 from transforms3d import quaternions