[Feature] Parallel render viewer (#383)

* work * Update types.py * initial work to support parallel render in GUI * Update sapien_env.py * bug fixes * w * bug fixes * modify which scene is at true 0, 0 * ww * bug fixes * Delete ppo_demo.py * bug fixes for static objects and lighting and single objects * Update quickstart.md * docs and fixes * Update teaser.png * bug fix for gpu sim recording trajectories * Update roll_ball.py
haosulab · Jun 21, 2024 · 417ded2 · 417ded2
1 parent 799d0c3
commit 417ded2
Show file tree

Hide file tree

Showing 16 changed files with 228 additions and 91 deletions.
diff --git a/docs/source/teaser.png b/docs/source/teaser.png
diff --git a/docs/source/user_guide/getting_started/images/parallel_gui_render.png b/docs/source/user_guide/getting_started/images/parallel_gui_render.png
diff --git a/docs/source/user_guide/getting_started/quickstart.md b/docs/source/user_guide/getting_started/quickstart.md
@@ -4,7 +4,7 @@ ManiSkill is a robotics simulator built on top of SAPIEN. It provides a standard
 
 ## Gym Interface
 
-Here is a basic example of how to make a ManiSkill task following the interface of [Gymnasium](https://gymnasium.farama.org/) and run a random policy.
+Here is a basic example of how to run a ManiSkill task following the interface of [Gymnasium](https://gymnasium.farama.org/) and execute a random policy.
 
 ```python
 import gymnasium as gym
@@ -35,8 +35,10 @@ Changing `num_envs` to a value > 1 will automatically turn on the GPU simulation
 You can also run the same code from the command line to demo random actions
 
 ```bash
-python -m mani_skill.examples.demo_random_action -e PickCube-v1 # run headless
-python -m mani_skill.examples.demo_random_action -e PickCube-v1 --render-mode="human" # run with A GUI
+# run headless / without a display
+python -m mani_skill.examples.demo_random_action -e PickCube-v1
+# run with A GUI
+python -m mani_skill.examples.demo_random_action -e PickCube-v1 --render-mode="human"
 ```
 
 Running with `render_mode="human"` will open up a GUI shown below that you can use to interactively explore the scene, pause/play the script, teleport objects around, and more.
@@ -110,7 +112,7 @@ for i in range(200):
 env.close()
 ```
 
-Note that as long as the GPU simulation is being used, all values returned by `env.step` and `env.reset` are batched and are torch tensors. With CPU simulation (`num_envs=1`), we keep to the standard gymnasium format and return unbatched numpy values. In general however, to make programming easier by default everything inside ManiSkill is kept as torch CPU/GPU tensors whenever possible, with the only exceptions being the return values of `env.step` and `env.reset`.
+Note that all values returned by `env.step` and `env.reset` are batched and are torch tensors. Whether GPU or CPU simulation is used then determines what device the tensor is on (CUDA or CPU).
 
 To benchmark the parallelized simulation, you can run 
 
@@ -132,6 +134,28 @@ which will look something like this
 <source src="https://github.com/haosulab/ManiSkill/raw/main/docs/source/_static/videos/mani_skill_gpu_sim-PickCube-v1-num_envs=16-obs_mode=state-render_mode=sensors.mp4" type="video/mp4">
 </video>
 
+We further support opening a GUI to view all parallel environments at once, and you can also turn on ray-tracing for more photo-realism. Note that this feature is not useful for any practical purposes (for e.g. machine learning) apart from generating cool demonstration videos and so it is not well optimized.
+
+Turning the parallel GUI render on simply requires adding the argument `parallel_gui_render_enabled` to `gym.make` as so
+
+```python
+import gymnasium as gym
+import mani_skill.envs
+
+env = gym.make(
+    "PickCube-v1",
+    obs_mode="state",
+    control_mode="pd_joint_delta_pos",
+    num_envs=16,
+    parallel_gui_render_enabled=True,
+    shader_dir="rt-fast" # optionally set this argument for more photo-realistic rendering
+)
+```
+
+This will then open up a GUI that looks like so:
+```{figure} images/parallel_gui_render.png
+```
+
 
 ## Task Instantiation Options
 

diff --git a/mani_skill/envs/sapien_env.py b/mani_skill/envs/sapien_env.py
@@ -85,6 +85,10 @@ class BaseEnv(gym.Env):
             we use the cpu sim backend. Can also be "cpu" or "gpu" to force usage of a particular sim backend. Note that if this is "cpu", num_envs
             can only be equal to 1.
 
+        parallel_gui_render_enabled (bool): By default this is False. If True, when using env.render_human() which opens a viewer/GUI,
+            the viewer will show all parallel environments. This is only really useful for generating cool videos showing
+            all environments at once but it is not recommended otherwise as it slows down simulation and rendering.
+
     Note:
         `sensor_cfgs` is used to update environement-specific sensor configurations.
         If the key is one of sensor names (e.g. a camera), the value will be applied to the corresponding sensor.
@@ -147,12 +151,14 @@ def __init__(
         sim_cfg: Union[SimConfig, dict] = dict(),
         reconfiguration_freq: int = None,
         sim_backend: str = "auto",
+        parallel_gui_render_enabled: bool = False,
     ):
         self.num_envs = num_envs
         self.reconfiguration_freq = reconfiguration_freq if reconfiguration_freq is not None else 0
         self._reconfig_counter = 0
         self._custom_sensor_configs = sensor_configs
         self._custom_human_render_camera_configs = human_render_camera_configs
+        self._parallel_gui_render_enabled = parallel_gui_render_enabled
         self.robot_uids = robot_uids
         if self.SUPPORTED_ROBOTS is not None:
             if robot_uids not in self.SUPPORTED_ROBOTS:
@@ -180,9 +186,12 @@ def __init__(
             if obs_mode in ["sensor_data", "rgb", "rgbd", "pointcloud"]:
                 raise RuntimeError("""Currently you cannot use ray-tracing while running simulation with visual observation modes. You may still use
                 env.render_rgb_array() or the RecordEpisode wrapper to save videos of ray-traced results""")
-            if num_envs > 1:
+            if num_envs > 1 and parallel_gui_render_enabled == False:
                 raise RuntimeError("""Currently you cannot run ray-tracing on more than one environment in a single process""")
 
+        assert not parallel_gui_render_enabled or (obs_mode not in ["sensor_data", "pointcloud", "rgb", "depth", "rgbd"]), \
+            "Parallel rendering from parallel cameras is only supported when the gui/viewer is not used. parallel_gui_render_enabled must be False if using parallel rendering. If True only state based observations are supported."
+
         # TODO (stao): move the merge code / handling union typed arguments outside here so classes inheriting BaseEnv only get
         # the already parsed sim config argument
         if isinstance(sim_cfg, SimConfig):
@@ -897,8 +906,8 @@ def _setup_scene(self):
             scene_grid_length = int(np.ceil(np.sqrt(self.num_envs)))
             for scene_idx in range(self.num_envs):
                 scene_x, scene_y = (
-                    scene_idx % scene_grid_length,
-                    scene_idx // scene_grid_length,
+                    scene_idx % scene_grid_length - scene_grid_length // 2,
+                    scene_idx // scene_grid_length - scene_grid_length // 2,
                 )
                 scene = sapien.Scene(
                     systems=[self.physx_system, sapien.render.RenderSystem()]
@@ -918,7 +927,7 @@ def _setup_scene(self):
                 sapien.Scene([self.physx_system, sapien.render.RenderSystem()])
             ]
         # create a "global" scene object that users can work with that is linked with all other scenes created
-        self.scene = ManiSkillScene(sub_scenes, sim_cfg=self.sim_cfg, device=self.device)
+        self.scene = ManiSkillScene(sub_scenes, sim_cfg=self.sim_cfg, device=self.device, parallel_gui_render_enabled=self._parallel_gui_render_enabled)
         self.physx_system.timestep = 1.0 / self._sim_freq
 
     def _clear(self):
@@ -1019,10 +1028,21 @@ def _setup_viewer(self):
 
         Called by `self._reconfigure`
         """
-        # TODO (stao): handle GPU parallel sim rendering code:
-        if physx.is_gpu_enabled():
-            self._viewer_scene_idx = 0
-        # CAUTION: `set_scene` should be called after assets are loaded.
+        # commented code below is for a different parallel render system in the GUI but it does not support ray tracing
+        # instead to show parallel envs in the GUI they are all spawned into the same sub scene and offsets are auto
+        # added / subtracted from object poses.
+        # if self.num_envs > 1:
+        #     side = int(np.ceil(self.num_envs ** 0.5))
+        #     idx = np.arange(self.num_envs)
+        #     offsets = np.stack([idx // side, idx % side, np.zeros_like(idx)], axis=1) * self.sim_cfg.spacing
+        #     self.viewer.set_scenes(self.scene.sub_scenes, offsets=offsets)
+        #     vs = self.viewer.window._internal_scene  # type: ignore
+        #     cubemap = self.scene.sub_scenes[0].render_system.get_cubemap()
+        #     if cubemap is not None:  # type: ignore [sapien may return None]
+        #         vs.set_cubemap(cubemap._internal_cubemap)
+        #     else:
+        #         vs.set_ambient_light([0.5, 0.5, 0.5])
+        # else:
         self._viewer.set_scene(self.scene.sub_scenes[0])
         control_window: sapien.utils.viewer.control_window.ControlWindow = (
             sapien_utils.get_obj_by_type(

diff --git a/mani_skill/envs/scene.py b/mani_skill/envs/scene.py
@@ -1,5 +1,7 @@
-from typing import Dict, List, Tuple, Union
+from functools import cached_property
+from typing import Dict, List, Optional, Tuple, Union
 
+import numpy as np
 import sapien
 import sapien.physx as physx
 import sapien.render
@@ -33,6 +35,7 @@ def __init__(
         sim_cfg: SimConfig = SimConfig(),
         debug_mode: bool = True,
         device: Device = None,
+        parallel_gui_render_enabled: bool = False,
     ):
         if sub_scenes is None:
             sub_scenes = [sapien.Scene()]
@@ -77,6 +80,9 @@ def __init__(
         non-unique between episode resets in order to be easily rebuilt and deallocate old queries. This essentially acts as a way
         to invalidate the cached queries."""
 
+        self.parallel_gui_render_enabled: bool = parallel_gui_render_enabled
+        """Whether rendering all parallel scenes in the viewer/gui is enabled"""
+
     # -------------------------------------------------------------------------- #
     # Functions from sapien.Scene
     # -------------------------------------------------------------------------- #
@@ -375,14 +381,15 @@ def add_point_light(
         shadow_near=0.1,
         shadow_far=10.0,
         shadow_map_size=2048,
-        scene_idxs=None,
+        scene_idxs: Optional[List[int]] = None,
     ):
-        lighting_scenes = (
-            self.sub_scenes
-            if scene_idxs is None
-            else [self.sub_scenes[i] for i in scene_idxs]
-        )
-        for scene in lighting_scenes:
+        if scene_idxs is None:
+            scene_idxs = list(range(len(self.sub_scenes)))
+        for scene_idx in scene_idxs:
+            if self.parallel_gui_render_enabled:
+                scene = self.sub_scenes[0]
+            else:
+                scene = self.sub_scenes[scene_idx]
             entity = sapien.Entity()
             light = sapien.render.RenderPointLightComponent()
             entity.add_component(light)
@@ -391,7 +398,11 @@ def add_point_light(
             light.shadow_near = shadow_near
             light.shadow_far = shadow_far
             light.shadow_map_size = shadow_map_size
-            light.pose = sapien.Pose(position)
+            if self.parallel_gui_render_enabled:
+                light.pose = sapien.Pose(position + self.scene_offsets_np[scene_idx])
+            else:
+                light.pose = sapien.Pose(position)
+
             scene.add_entity(entity)
         return light
 
@@ -405,14 +416,15 @@ def add_directional_light(
         shadow_near=-10.0,
         shadow_far=10.0,
         shadow_map_size=2048,
-        scene_idxs=None,
+        scene_idxs: Optional[List[int]] = None,
     ):
-        lighting_scenes = (
-            self.sub_scenes
-            if scene_idxs is None
-            else [self.sub_scenes[i] for i in scene_idxs]
-        )
-        for scene in lighting_scenes:
+        if scene_idxs is None:
+            scene_idxs = list(range(len(self.sub_scenes)))
+        for scene_idx in scene_idxs:
+            if self.parallel_gui_render_enabled:
+                scene = self.sub_scenes[0]
+            else:
+                scene = self.sub_scenes[scene_idx]
             entity = sapien.Entity()
             light = sapien.render.RenderDirectionalLightComponent()
             entity.add_component(light)
@@ -422,10 +434,19 @@ def add_directional_light(
             light.shadow_far = shadow_far
             light.shadow_half_size = shadow_scale
             light.shadow_map_size = shadow_map_size
+            if self.parallel_gui_render_enabled:
+                light_position = position + self.scene_offsets_np[scene_idx]
+            else:
+                light_position = position
             light.pose = sapien.Pose(
-                position, sapien.math.shortest_rotation([1, 0, 0], direction)
+                light_position, sapien.math.shortest_rotation([1, 0, 0], direction)
             )
             scene.add_entity(entity)
+            if self.parallel_gui_render_enabled:
+                # for directional lights adding multiple does not make much sense
+                # and for parallel gui rendering setup accurate lighting does not matter as it is only
+                # for demo purposes
+                break
         return
 
     def add_spot_light(
@@ -439,14 +460,15 @@ def add_spot_light(
         shadow_near=0.1,
         shadow_far=10.0,
         shadow_map_size=2048,
-        scene_idxs=None,
+        scene_idxs: Optional[List[int]] = None,
     ):
-        lighting_scenes = (
-            self.sub_scenes
-            if scene_idxs is None
-            else [self.sub_scenes[i] for i in scene_idxs]
-        )
-        for scene in lighting_scenes:
+        if scene_idxs is None:
+            scene_idxs = list(range(len(self.sub_scenes)))
+        for scene_idx in scene_idxs:
+            if self.parallel_gui_render_enabled:
+                scene = self.sub_scenes[0]
+            else:
+                scene = self.sub_scenes[scene_idx]
             entity = sapien.Entity()
             light = sapien.render.RenderSpotLightComponent()
             entity.add_component(light)
@@ -457,8 +479,12 @@ def add_spot_light(
             light.shadow_map_size = shadow_map_size
             light.inner_fov = inner_fov
             light.outer_fov = outer_fov
+            if self.parallel_gui_render_enabled:
+                light_position = position + self.scene_offsets_np[scene_idx]
+            else:
+                light_position = position
             light.pose = sapien.Pose(
-                position, sapien.math.shortest_rotation([1, 0, 0], direction)
+                light_position, sapien.math.shortest_rotation([1, 0, 0], direction)
             )
             scene.add_entity(entity)
         return
@@ -569,6 +595,23 @@ def get_pairwise_contact_forces(
         """
         return self.get_pairwise_contact_impulses(obj1, obj2) / self.px.timestep
 
+    @cached_property
+    def scene_offsets(self):
+        """torch tensor of shape (num_envs, 3) representing the offset of each scene in the world frame"""
+        return torch.tensor(
+            np.array(
+                [self.px.get_scene_offset(sub_scene) for sub_scene in self.sub_scenes]
+            ),
+            device=self.device,
+        )
+
+    @cached_property
+    def scene_offsets_np(self):
+        """numpy array of shape (num_envs, 3) representing the offset of each scene in the world frame"""
+        return np.array(
+            [self.px.get_scene_offset(sub_scene) for sub_scene in self.sub_scenes]
+        )
+
     # -------------------------------------------------------------------------- #
     # Simulation state (required for MPC)
     # -------------------------------------------------------------------------- #

diff --git a/mani_skill/envs/tasks/tabletop/pick_cube.py b/mani_skill/envs/tasks/tabletop/pick_cube.py
@@ -12,6 +12,7 @@
 from mani_skill.utils.registration import register_env
 from mani_skill.utils.scene_builder.table import TableSceneBuilder
 from mani_skill.utils.structs.pose import Pose
+from mani_skill.utils.structs.types import SimConfig
 
 
 @register_env("PickCube-v1", max_episode_steps=50)