[optimize] optimize GPU pipeline efficiency by indexing GPU tensor first, consider asset_option use_mesh_materials

yzqin · yzqin · commit a3e0f7cb849b · 2023-08-11T13:29:23.000-07:00
diff --git a/example/isaacgym/run_isaacgym_viz.py b/example/isaacgym/run_isaacgym_viz.py
@@ -32,11 +32,11 @@ def wrapped_create_sim(self: vec_task.VecTask, compute_device: int, graphics_dev
 vec_task.VecTask.create_sim = wrapped_create_sim
 
 # Create web visualizer
-create_isaac_visualizer(port=6000, host="localhost", keep_default_viewer=True, max_env=1)
+create_isaac_visualizer(port=6000, host="localhost", keep_default_viewer=True, max_env=4)
 
 # Create the environment and step the simulation as normal
 device = "cuda"  # or "cpu"
-num_env = 2
+num_env = 8
 rl_device = device
 envs = isaacgymenvs.make(
     seed=0,
diff --git a/sim_web_visualizer/base_visualizer_client.py b/sim_web_visualizer/base_visualizer_client.py
@@ -39,9 +39,6 @@ def add_default_scene_elements(self):
         self.viz["/Lights/AmbientLight/<object>"].set_property("intensity", 0.8)
         self.viz["/Lights/PointLightNegativeX"].delete()
         self.viz["/Lights/PointLightPositiveX"].delete()
-        # self.viz["/Lights/PointLightNegativeX/<object>"].set_property("intensity", 1)
-        # TODO: cast shadow does not work
-        # self.viz["/Lights/PointLightNegativeX/<object>"].set_property("castShadow", False)
 
     def delete_all(self):
         print("Deleting all previous scene asset.")
@@ -50,10 +47,11 @@ def delete_all(self):
 
     @staticmethod
     def dry_load_asset(filename, collapse_fixed_joints: bool,
-                       replace_cylinder_with_capsule=False) -> AssetResource:
+                       replace_cylinder_with_capsule=False, use_mesh_materials=False) -> AssetResource:
         if filename.endswith(".urdf"):
             resource = load_urdf_with_yourdfpy(filename, collapse_fixed_joints,
-                                               replace_cylinder_with_capsule=replace_cylinder_with_capsule)
+                                               replace_cylinder_with_capsule=replace_cylinder_with_capsule,
+                                               use_mesh_materials=use_mesh_materials)
         elif filename.endswith(".xml"):
             resource = load_mjcf_with_dmc(filename, collapse_fixed_joints)
         else:
diff --git a/sim_web_visualizer/isaac_visualizer_client.py b/sim_web_visualizer/isaac_visualizer_client.py
@@ -128,7 +128,8 @@ def load_asset(sim: gymapi.Sim, rootpath: str, filename: str,
             # Load asset into the memory as asset resource but not forward it into the web viewer
             # The asset will only be loaded into web viewer after a `create_actor` call
             resource = self.dry_load_asset(filename, collapse_fixed_joints,
-                                           replace_cylinder_with_capsule=options.replace_cylinder_with_capsule)
+                                           replace_cylinder_with_capsule=options.replace_cylinder_with_capsule,
+                                           use_mesh_materials=options.use_mesh_materials)
             self.asset_resource_map[id(asset)] = resource
             return asset
 
@@ -353,7 +354,7 @@ def draw_viewer(viewer: Union[MimicViewer, gymapi.Viewer], sim: gymapi.Sim, rend
 
             if USE_GPU_PIPELINE:
                 self.original_gym.refresh_rigid_body_state_tensor(self.sim)
-                rigid_body_state = self.rigid_body_state_tensor.cpu().numpy()[:len(self.env_map), :, :7]
+                rigid_body_state = self.rigid_body_state_tensor[:len(self.env_map), :, :7].cpu().numpy()
                 pos = rigid_body_state[:, :, :3]
                 quat = np.concatenate([rigid_body_state[..., 6:7], rigid_body_state[..., 3:6]], axis=-1)  # xyzw -> wxyz
                 qs = quaternion.as_quat_array(quat)
diff --git a/sim_web_visualizer/parser/urdf.py b/sim_web_visualizer/parser/urdf.py
@@ -16,9 +16,9 @@
 
 
 def load_urdf_with_yourdfpy(urdf_path: str, collapse_fixed_joints: bool,
-                            replace_cylinder_with_capsule=False) -> AssetResource:
+                            replace_cylinder_with_capsule=False, use_mesh_materials=False) -> AssetResource:
     robot = urdf.URDF.load(urdf_path)
-    material_map = robot._material_map
+    material_map = robot._material_map if not use_mesh_materials else {}
 
     # Dry run data for faster loading
     offline_data_dict = {}
