diff --git a/mani_skill2/envs/sapien_env.py b/mani_skill2/envs/sapien_env.py index d84e3afcb..e9e2dfadc 100644 --- a/mani_skill2/envs/sapien_env.py +++ b/mani_skill2/envs/sapien_env.py @@ -446,15 +446,12 @@ def _get_obs_with_sensor_data(self, info: Dict) -> OrderedDict: obj.hide_visual() self._scene.update_render() self.capture_sensor_data() - obs = OrderedDict( + return 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): @@ -965,6 +962,8 @@ def _setup_viewer(self): ) def render_human(self): + for obj in self._hidden_objects: + obj.show_visual() if self._viewer is None: self._viewer = Viewer() self._setup_viewer() @@ -981,6 +980,8 @@ 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? @@ -1007,6 +1008,8 @@ def render_rgb_array(self, camera_name: str = None): return None if len(images) == 1: return images[0] + for obj in self._hidden_objects: + obj.hide_visual() return tile_images(images) def render_sensors(self): @@ -1021,8 +1024,6 @@ def render_sensors(self): 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/utils/structs/actor.py b/mani_skill2/utils/structs/actor.py index ee4e87d59..ecb453c24 100644 --- a/mani_skill2/utils/structs/actor.py +++ b/mani_skill2/utils/structs/actor.py @@ -150,9 +150,6 @@ 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 @@ -246,4 +243,7 @@ def pose(self, arg1: Union[Pose, sapien.Pose, Array]) -> None: self._objs[0].pose = to_sapien_pose(arg1) def set_pose(self, arg1: Union[Pose, sapien.Pose]) -> None: - self.pose = arg1 + if physx.is_gpu_enabled() and self.hidden: + self.before_hide_pose = vectorize_pose(arg1) + else: + self.pose = arg1