From 25005a201552f86931767ac9eecd0716995c4754 Mon Sep 17 00:00:00 2001
From: Stone Tao <stonezt2019@gmail.com>
Date: Wed, 1 May 2024 12:16:49 -0700
Subject: [PATCH] update custom robot docs, qol property about checking if gpu
 sim is enabled instead of checking device, demo robot can test gpu backend
 (#302)

---
 docs/source/user_guide/tutorials/custom_robots.md |  5 ++++-
 mani_skill/envs/sapien_env.py                     | 10 ++++++++++
 mani_skill/envs/tasks/tabletop/pick_single_ycb.py |  4 +++-
 mani_skill/examples/demo_robot.py                 |  7 ++++++-
 4 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/docs/source/user_guide/tutorials/custom_robots.md b/docs/source/user_guide/tutorials/custom_robots.md
index e3b76c241..10e7d71bd 100644
--- a/docs/source/user_guide/tutorials/custom_robots.md
+++ b/docs/source/user_guide/tutorials/custom_robots.md
@@ -131,10 +131,11 @@ class ANYmalC(BaseAgent):
 
 The keyframe can also specify `qvel` values as well. Using that keyframe you can set the robot to the given pose, qpos, qvel and you can get the desired predefined keyframe
 
-Running the script again should have correctly loaded poses and joint positions. The script by default picks the first keyframe defined. You can add more and select them shown below.
+Running the script again should have correctly loaded poses and joint positions. The script by default picks the first keyframe defined. You can add more and select them shown below. You may also want to check if it is working on the GPU simulation as well
 
 ```bash
 python test.py -r "my_panda"
+python test.py -r "my_panda" -b "gpu" # use gpu simulation
 python test.py -r "my_panda" -k "name_of_keyframe_to_show"
 ```
 
@@ -241,6 +242,7 @@ To try this out, simply run the following and unpause the simulation when you ar
 
 ```bash
 python -m test.py -r "my_panda" -c "pd_joint_delta_pos" --random-actions
+python -m test.py -r "my_panda" -c "pd_joint_delta_pos" -b "gpu" --random-actions
 ```
 
 <video preload="auto" controls="True" width="100%">
@@ -253,6 +255,7 @@ You can also test the stability of the robot you modelled by trying and set all
 
 ```bash
 python -m test.py -r "my_panda" -c "pd_joint_pos" --keyframe-actions
+python -m test.py -r "my_panda" -c "pd_joint_pos" -b "gpu" --keyframe-actions
 ```
 
 If the robot is staying still in the given keyframe, it is working well.
diff --git a/mani_skill/envs/sapien_env.py b/mani_skill/envs/sapien_env.py
index 5e34fd536..9411e636b 100644
--- a/mani_skill/envs/sapien_env.py
+++ b/mani_skill/envs/sapien_env.py
@@ -156,6 +156,11 @@ def __init__(
         self.robot_uids = robot_uids
         if self.SUPPORTED_ROBOTS is not None:
             assert robot_uids in self.SUPPORTED_ROBOTS
+
+        if physx.is_gpu_enabled() and num_envs == 1 and (sim_backend == "auto" or sim_backend == "cpu"):
+            logger.warn("GPU simulation has already been enabled on this process, switching to GPU backend")
+            sim_backend == "gpu"
+
         if num_envs > 1 or sim_backend == "gpu":
             if not physx.is_gpu_enabled():
                 physx.enable_gpu()
@@ -285,6 +290,11 @@ def single_observation_space(self):
     def observation_space(self):
         return batch_space(self.single_observation_space, n=self.num_envs)
 
+    @property
+    def gpu_sim_enabled(self):
+        """Whether the gpu simulation is enabled. A wrapper over physx.is_gpu_enabled()"""
+        return physx.is_gpu_enabled()
+
     @property
     def _default_sim_config(self):
         return SimConfig()
diff --git a/mani_skill/envs/tasks/tabletop/pick_single_ycb.py b/mani_skill/envs/tasks/tabletop/pick_single_ycb.py
index 26fc68562..ec2f7ee67 100644
--- a/mani_skill/envs/tasks/tabletop/pick_single_ycb.py
+++ b/mani_skill/envs/tasks/tabletop/pick_single_ycb.py
@@ -166,6 +166,7 @@ def evaluate(self):
             obj_to_goal_pos=obj_to_goal_pos,
             is_obj_placed=is_obj_placed,
             is_robot_static=is_robot_static,
+            is_grasping=self.agent.is_grasping(self.obj),
             success=torch.logical_and(is_obj_placed, is_robot_static),
         )
 
@@ -173,6 +174,7 @@ def _get_obs_extra(self, info: Dict):
         obs = dict(
             tcp_pose=self.agent.tcp.pose.raw_pose,
             goal_pos=self.goal_site.pose.p,
+            is_grasping=info["is_grasping"],
         )
         if "state" in self.obs_mode:
             obs.update(
@@ -190,7 +192,7 @@ def compute_dense_reward(self, obs: Any, action: torch.Tensor, info: Dict):
         reaching_reward = 1 - torch.tanh(5 * tcp_to_obj_dist)
         reward = reaching_reward
 
-        is_grasped = self.agent.is_grasping(self.obj)
+        is_grasped = info["is_grasping"]
         reward += is_grasped
 
         obj_to_goal_dist = torch.linalg.norm(
diff --git a/mani_skill/examples/demo_robot.py b/mani_skill/examples/demo_robot.py
index 0fd456c76..47b129aa9 100644
--- a/mani_skill/examples/demo_robot.py
+++ b/mani_skill/examples/demo_robot.py
@@ -13,6 +13,7 @@
 def parse_args(args=None):
     parser = argparse.ArgumentParser()
     parser.add_argument("-r", "--robot-uid", type=str, default="panda", help="The id of the robot to place in the environment")
+    parser.add_argument("-b", "--sim-backend", type=str, default="auto", help="Which simulation backend to use. Can be 'auto', 'cpu', 'gpu'")
     parser.add_argument("-c", "--control-mode", type=str, default="pd_joint_pos", help="The control mode to use. Note that for new robots being implemented if the _controller_configs is not implemented in the selected robot, we by default provide two default controllers, 'pd_joint_pos' and 'pd_joint_delta_pos' ")
     parser.add_argument("-k", "--keyframe", type=str, help="The name of the keyframe of the robot to display")
     parser.add_argument("--shader", default="default", type=str, help="Change shader used for rendering. Default is 'default' which is very fast. Can also be 'rt' for ray tracing and generating photo-realistic renders. Can also be 'rt-fast' for a faster but lower quality ray-traced renderer")
@@ -37,8 +38,8 @@ def main():
         control_mode=args.control_mode,
         robot_uids=args.robot_uid,
         shader_dir=args.shader,
-        sim_cfg=SimConfig(sim_freq=100, scene_cfg=SceneConfig(solver_position_iterations=50)),
         render_mode="human",
+        sim_backend=args.sim_backend,
     )
     env.reset(seed=0)
     env: BaseEnv = env.unwrapped
@@ -58,6 +59,10 @@ def main():
         if kf.qvel is not None:
             env.agent.robot.set_qvel(kf.qvel)
         env.agent.robot.set_pose(kf.pose)
+    if env.gpu_sim_enabled:
+        env.scene._gpu_apply_all()
+        env.scene.px.gpu_update_articulation_kinematics()
+        env.scene._gpu_fetch_all()
     viewer = env.render()
     viewer.paused = True
     viewer = env.render()