From 20a2d04e4c7ca28eddb736b6e3f0e931b3f5a6b0 Mon Sep 17 00:00:00 2001
From: arth-shukla <44527268+arth-shukla@users.noreply.github.com>
Date: Sat, 2 Mar 2024 00:50:06 -0800
Subject: [PATCH] Fix: show and hide visual broken on gpu sim

---
 mani_skill2/envs/sapien_env.py                | 14 +++++------
 .../envs/scenes/tasks/sequential_task.py      |  6 +++--
 mani_skill2/utils/structs/actor.py            | 23 ++++++++-----------
 3 files changed, 21 insertions(+), 22 deletions(-)

diff --git a/mani_skill2/envs/sapien_env.py b/mani_skill2/envs/sapien_env.py
index f1912793d..d84e3afcb 100644
--- a/mani_skill2/envs/sapien_env.py
+++ b/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):
diff --git a/mani_skill2/envs/scenes/tasks/sequential_task.py b/mani_skill2/envs/scenes/tasks/sequential_task.py
index 21fd44888..87bce984f 100644
--- a/mani_skill2/envs/scenes/tasks/sequential_task.py
+++ b/mani_skill2/envs/scenes/tasks/sequential_task.py
@@ -56,7 +56,7 @@ class SequentialTaskEnv(SceneManipulationEnv):
     SUPPORTED_ROBOTS = ["fetch"]
     agent: Fetch
     sim_cfg = SimConfig(
-        spacing=20,
+        spacing=30,
         gpu_memory_cfg=GPUMemoryConfig(
             found_lost_pairs_capacity=2**25, max_rigid_patch_count=2**18
         ),
@@ -509,14 +509,16 @@ def _register_human_render_cameras(self):
         return robot_camera_config
 
     def render_cameras(self):
-        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([self.render_rgb_array()] + images)
 
     def render_rgb_array(self):
diff --git a/mani_skill2/utils/structs/actor.py b/mani_skill2/utils/structs/actor.py
index bc52ba650..1f64ea3f7 100644
--- a/mani_skill2/utils/structs/actor.py
+++ b/mani_skill2/utils/structs/actor.py
@@ -162,17 +162,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 +181,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: