Fix: actor show/hide visual not working on gpu sim

mani_skill2/envs/sapien_env.py

@@ -446,12 +446,15 @@ def _get_obs_with_sensor_data(self, info: Dict) -> OrderedDict:
             obj.hide_visual()
         self._scene.update_render()
         self.capture_sensor_data()
-        return OrderedDict(
+        obs = OrderedDict(
             agent=self._get_obs_agent(),
             extra=self._get_obs_extra(info),
             sensor_param=self.get_sensor_params(),
             sensor_data=self.get_sensor_obs(),
         )
+        for obj in self._hidden_objects:
+            obj.show_visual()
+        return obs
 
     @property
     def robot_link_ids(self):
@@ -969,9 +972,6 @@ def render_human(self):
                 self._viewer.set_camera_pose(
                     self._human_render_cameras["render_camera"].camera.global_pose
                 )
-
-        for obj in self._hidden_objects:
-            obj.show_visual()
         if physx.is_gpu_enabled() and self._scene._gpu_sim_initialized:
             self.physx_system.sync_poses_gpu_to_cpu()
         self._viewer.render()
@@ -981,8 +981,6 @@ def render_rgb_array(self, camera_name: str = None):
         """Returns an RGB array / image of size (num_envs, H, W, 3) of the current state of the environment.
         This is captured by any of the registered human render cameras. If a camera_name is given, only data from that camera is returned.
         Otherwise all camera data is captured and returned as a single batched image"""
-        for obj in self._hidden_objects:
-            obj.show_visual()
         self._scene.update_render()
         images = []
         # TODO (stao): refactor this code either into ManiSkillScene class and/or merge the code, it's pretty similar?
@@ -1015,14 +1013,16 @@ def render_sensors(self):
         """
         Renders all sensors that the agent can use and see and displays them
         """
-        images = []
         for obj in self._hidden_objects:
             obj.hide_visual()
+        images = []
         self._scene.update_render()
         self.capture_sensor_data()
         sensor_images = self.get_sensor_obs()
         for sensor_images in sensor_images.values():
             images.extend(observations_to_images(sensor_images))
+        for obj in self._hidden_objects:
+            obj.show_visual()
         return tile_images(images)
 
     def render(self):

mani_skill2/utils/structs/actor.py

@@ -150,6 +150,9 @@ def has_collision_shapes(self):
             > 0
         )
 
+    # NOTE (arth): when using hide/show_visual on gpu sim, should call
+    #       hide_visual, then soon after show_visual, such that
+    #       poses are not incorrectly updated in between
     def hide_visual(self):
         """
         Hides this actor from view. In CPU simulation the visual body is simply set to visibility 0
@@ -162,17 +165,14 @@ def hide_visual(self):
         if self.hidden:
             return
         if physx.is_gpu_enabled():
-            # TODO (stao): fix hiding visuals
-            pass
-            # self.last_pose = self.px.cuda_rigid_body_data.torch()[
-            #     self._body_data_index, :7
-            # ].clone()
-            # temp_pose = self.pose.raw_pose
-            # temp_pose[..., :3] += 99999
-            # self.pose = temp_pose
-            # self.px.gpu_apply_rigid_dynamic_data()
-            # self.px.gpu_fetch_rigid_dynamic_data()
-            # print("HIDE", self.pose.raw_pose[0, :3])
+            self.before_hide_pose = self.px.cuda_rigid_body_data.torch()[
+                self._body_data_index, :7
+            ].clone()
+            temp_pose = self.pose.raw_pose
+            temp_pose[..., :3] += 99999
+            self.pose = temp_pose
+            self.px.gpu_apply_rigid_dynamic_data()
+            self.px.gpu_fetch_rigid_dynamic_data()
         else:
             self._objs[0].find_component_by_type(
                 sapien.render.RenderBodyComponent
@@ -184,8 +184,8 @@ def show_visual(self):
         if not self.hidden:
             return
         if physx.is_gpu_enabled():
-            if hasattr(self, "last_pose"):
-                self.pose = self.last_pose
+            if hasattr(self, "before_hide_pose"):
+                self.pose = self.before_hide_pose
                 self.px.gpu_apply_rigid_dynamic_data()
                 self.px.gpu_fetch_rigid_dynamic_data()
         else: