Skip to content

Commit

Permalink
Fix: show and hide visual broken on gpu sim
Browse files Browse the repository at this point in the history
  • Loading branch information
arth-shukla committed Mar 2, 2024
1 parent ea84760 commit 20a2d04
Show file tree
Hide file tree
Showing 3 changed files with 21 additions and 22 deletions.
14 changes: 7 additions & 7 deletions mani_skill2/envs/sapien_env.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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()
Expand All @@ -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?
Expand Down Expand Up @@ -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):
Expand Down
6 changes: 4 additions & 2 deletions mani_skill2/envs/scenes/tasks/sequential_task.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
),
Expand Down Expand Up @@ -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):
Expand Down
23 changes: 10 additions & 13 deletions mani_skill2/utils/structs/actor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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:
Expand Down

0 comments on commit 20a2d04

Please sign in to comment.