update custom robot docs, qol property about checking if gpu sim is e…

…nabled instead of checking device, demo robot can test gpu backend (#302)

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.

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()

mani_skill/envs/tasks/tabletop/pick_single_ycb.py

@@ -166,13 +166,15 @@ 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),
         )
 
     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(

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()