[AO migration] async rendering

examples/tutorials/async_rendering.py

@@ -0,0 +1,111 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+r"""
+Simple example showing how to overlap physics and rendering in Habitat-Sim.  The general
+pattern is to call sim.start_async_render_and_step_physics instead of sim.step_physics for
+the *first* step_physics call and then retrieve observations using sim.get_sensor_observations_async_finish
+instead of sim.get_sensor_observations
+
+
+Known limitations/issues:
+    * Creating and destroying simulation instances leaks GPU memory.  On some drivers
+    OpenGL leaks GPU memory when the background thread used for rendering is destroyed.
+    The workaround is to use sim.close(destroy=False) then sim.reconfigure(new_cfg) instead of
+    sim = habitat_sim.Simulator(new_cfg) as this will keep the background thread alive.
+
+    * Segfaults with headed Linux builds.  We've seen this segfault for headed builds.  There
+    isn't any reason why the headed Linux build should segfault so we've left it enabled
+    under that this is a driver bug and it may work with other driver/OS combinations.
+
+    * Semantic Sensor rendering does not work when there is an RGB mesh and a semantic mesh, like
+    in MP3D and Gibson.  When there is a single mesh for RGB and semantics, a semantic sensor works.
+
+    * Sensors that aren't CameraSensors, i.e. FisheyeSensor or EquirectangularSensor, are not
+    supported
+
+    * Adding or deleting objects during the async render is not supported.  This will lead to
+    an assert checking that the main process owns the OpenGL failing.  You can wait on the
+    async render to finish and transfer the context back to the main thread by calling
+    sim.renderer.acquire_gl_context().  Note that by default the OpenGL context is transferred
+    back to the main thread by calling sim.get_sensor_observations_async_finish() by default.
+"""
+
+import habitat_sim
+
+
+def main():
+    backend_cfg = habitat_sim.SimulatorConfiguration()
+    backend_cfg.scene_id = (
+        "data/scene_datasets/habitat-test-scenes/skokloster-castle.glb"
+    )
+    backend_cfg.enable_physics = True
+
+    # Leaving the context with the background thread can improve
+    # performance, but makes it so that sim.get_sensor_observations_async_finish()
+    # no longer returns the context back to the main thread. If you set this to
+    # true and get context assert failures, call sim.renderer.acquire_gl_context()
+    # to move the OpenGL context back to the main thread.
+    # The default is False
+    backend_cfg.leave_context_with_background_renderer = False
+
+    agent_cfg = habitat_sim.AgentConfiguration(
+        sensor_specifications=[habitat_sim.CameraSensorSpec()]
+    )
+    cfg = habitat_sim.Configuration(backend_cfg, [agent_cfg])
+
+    # Number of physics steps per render
+    n_physics_steps_per_render = 4
+    # Render at 30 hz
+    physics_step_time = 1.0 / (30 * n_physics_steps_per_render)
+
+    sim = habitat_sim.Simulator(cfg)
+
+    for i in range(n_physics_steps_per_render):
+        # Always call sim.step_physics when not using async rendering
+        if i == 0:
+            sim.start_async_render_and_step_physics(physics_step_time)
+        else:
+            sim.step_physics(physics_step_time)
+
+        # Call sim.renderer.acquire_gl_context() if you need to create or destroy an object.
+        # Note that this will block until rendering is done.
+
+    obs = sim.get_sensor_observations_async_finish()
+    # To change back to normal rendering, use
+    #     obs = sim.get_sensor_observations()
+
+    # You can add/remove objects after the call to get_sensor_observations_async_finish()
+    # if backend_cfg.leave_context_with_background_renderer was left as False, it that as
+    # true, you'd need to call
+    #     sim.renderer.acquire_gl_context()
+    # Calling acquire_gl_context() is a noop if the main thread already has the OpenGL context
+
+    backend_cfg = habitat_sim.SimulatorConfiguration()
+    backend_cfg.scene_id = "data/scene_datasets/habitat-test-scenes/apartment_1.glb"
+    backend_cfg.enable_physics = True
+    new_cfg = habitat_sim.Configuration(backend_cfg, [agent_cfg])
+
+    # To create a new instance of the simulator to swap scene or similar,
+    # do not delete the simulator instance as this may leak GPU memory
+    # Instead, call sim.close(destroy=False) as this will close/delete
+    # everything except the OpenGL context and the background render
+    # thread.
+    sim.close(destroy=False)
+    # Then use reconfigure.
+    sim.reconfigure(new_cfg)
+    # If
+    #     sim.close()
+    #     sim = habitat_sim.Simulator(new_cfg)
+    # was used, that could end up leaking memory
+
+    sim.start_async_render_and_step_physics(physics_step_time)
+    obs = sim.get_sensor_observations_async_finish()  # noqa: F841
+
+    # Call close with destroy=True here because this example is over :)
+    sim.close(destroy=True)
+
+
+if __name__ == "__main__":
+    main()

habitat_sim/simulator.py

@@ -79,6 +79,7 @@ class Simulator(SimulatorBackend):
     _previous_step_time: float = attr.ib(
         default=0.0, init=False
     )  # track the compute time of each step
+    _async_draw_agent_ids: Optional[Union[int, List[int]]] = None
     __last_state: Dict[int, AgentState] = attr.ib(factory=dict, init=False)
 
     @staticmethod
@@ -119,7 +120,16 @@ def __attrs_post_init__(self) -> None:
         self._sanitize_config(self.config)
         self.__set_from_config(self.config)
 
-    def close(self) -> None:
+    def close(self, destroy: bool = True) -> None:
+        r"""Close the simulator instance.
+
+        :param destroy: Whether or not to force the OpenGL context to be
+        destroyed if async rendering was used.  If async rendering wasn't used,
+        this has no effect.
+        """
+        if self.renderer is not None:
+            self.renderer.acquire_gl_context()
+
         for agent_sensorsuite in self.__sensors:
             for sensor in agent_sensorsuite.values():
                 sensor.close()
@@ -135,13 +145,13 @@ def close(self) -> None:
 
         self.__last_state.clear()
 
-        super().close()
+        super().close(destroy)
 
     def __enter__(self) -> "Simulator":
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        self.close()
+        self.close(destroy=True)
 
     def seed(self, new_seed: int) -> None:
         super().seed(new_seed)
@@ -311,6 +321,60 @@ def initialize_agent(
         self.__last_state[agent_id] = agent.state
         return agent
 
+    def start_async_render_and_step_physics(
+        self, dt: float, agent_ids: Union[int, List[int]] = 0
+    ):
+        if self._async_draw_agent_ids is not None:
+            raise RuntimeError(
+                "start_async_render_and_step_physics was already called.  "
+                "Call get_sensor_observations_async_finish before calling this again.  "
+                "Use step_physics to step physics additional times."
+            )
+
+        self._async_draw_agent_ids = agent_ids
+        if isinstance(agent_ids, int):
+            agent_ids = [agent_ids]
+
+        for agent_id in agent_ids:
+            agent_sensorsuite = self.__sensors[agent_id]
+            for _sensor_uuid, sensor in agent_sensorsuite.items():
+                sensor._draw_observation_async()
+
+        self.renderer.start_draw_jobs()
+        self.step_physics(dt)
+
+    def get_sensor_observations_async_finish(
+        self,
+    ) -> Union[
+        Dict[str, Union[ndarray, "Tensor"]],
+        Dict[int, Dict[str, Union[ndarray, "Tensor"]]],
+    ]:
+        if self._async_draw_agent_ids is None:
+            raise RuntimeError(
+                "get_sensor_observations_async_finish was called before calling start_async_render_and_step_physics."
+            )
+
+        agent_ids = self._async_draw_agent_ids
+        self._async_draw_agent_ids = None
+        if isinstance(agent_ids, int):
+            agent_ids = [agent_ids]
+            return_single = True
+        else:
+            return_single = False
+
+        self.renderer.wait_draw_jobs()
+        # As backport. All Dicts are ordered in Python >= 3.7
+        observations: Dict[int, Dict[str, Union[ndarray, "Tensor"]]] = OrderedDict()
+        for agent_id in agent_ids:
+            agent_observations: Dict[str, Union[ndarray, "Tensor"]] = {}
+            for sensor_uuid, sensor in self.__sensors[agent_id].items():
+                agent_observations[sensor_uuid] = sensor._get_observation_async()
+
+            observations[agent_id] = agent_observations
+        if return_single:
+            return next(iter(observations.values()))
+        return observations
+
     @overload
     def get_sensor_observations(self, agent_ids: int = 0) -> ObservationDict:
         ...
@@ -452,7 +516,7 @@ def step_filter(self, start_pos: Vector3, end_pos: Vector3) -> Vector3:
         return end_pos
 
     def __del__(self) -> None:
-        self.close()
+        self.close(destroy=True)
 
     def step_physics(self, dt: float, scene_id: int = 0) -> None:
         self.step_world(dt)
@@ -497,16 +561,23 @@ def __init__(self, sim: Simulator, agent: Agent, sensor_id: str) -> None:
                     resolution[0], resolution[1], 4, dtype=torch.uint8, device=device
                 )
         else:
+            size = self._sensor_object.framebuffer_size
             if self._spec.sensor_type == SensorType.SEMANTIC:
                 self._buffer = np.empty(
                     (self._spec.resolution[0], self._spec.resolution[1]),
                     dtype=np.uint32,
                 )
+                self.view = mn.MutableImageView2D(
+                    mn.PixelFormat.R32UI, size, self._buffer
+                )
             elif self._spec.sensor_type == SensorType.DEPTH:
                 self._buffer = np.empty(
                     (self._spec.resolution[0], self._spec.resolution[1]),
                     dtype=np.float32,
                 )
+                self.view = mn.MutableImageView2D(
+                    mn.PixelFormat.R32F, size, self._buffer
+                )
             else:
                 self._buffer = np.empty(
                     (
@@ -516,6 +587,11 @@ def __init__(self, sim: Simulator, agent: Agent, sensor_id: str) -> None:
                     ),
                     dtype=np.uint8,
                 )
+                self.view = mn.MutableImageView2D(
+                    mn.PixelFormat.RGBA8_UNORM,
+                    size,
+                    self._buffer.reshape(self._spec.resolution[0], -1),
+                )
 
         noise_model_kwargs = self._spec.noise_model_kwargs
         self._noise_model = make_sensor_noise_model(
@@ -529,7 +605,25 @@ def __init__(self, sim: Simulator, agent: Agent, sensor_id: str) -> None:
         )
 
     def draw_observation(self) -> None:
-        # sanity check:
+        # see if the sensor is attached to a scene graph, otherwise it is invalid,
+        # and cannot make any observation
+        if not self._sensor_object.object:
+            raise habitat_sim.errors.InvalidAttachedObject(
+                "Sensor observation requested but sensor is invalid.\
+                 (has it been detached from a scene node?)"
+            )
+        self._sim.renderer.draw(self._sensor_object, self._sim)
+
+    def _draw_observation_async(self) -> None:
+        if (
+            self._spec.sensor_type == SensorType.SEMANTIC
+            and self._sim.get_active_scene_graph()
+            is not self._sim.get_active_semantic_scene_graph()
+        ):
+            raise RuntimeError(
+                "Async drawing doesn't support semantic rendering when there are multiple scene graphs"
+            )
+        # TODO: sync this path with renderer changes as above (render from sensor object)
 
         # see if the sensor is attached to a scene graph, otherwise it is invalid,
         # and cannot make any observation
@@ -539,10 +633,42 @@ def draw_observation(self) -> None:
                  (has it been detached from a scene node?)"
             )
 
-        self._sim.renderer.draw(self._sensor_object, self._sim)
+        # get the correct scene graph based on application
+        if self._spec.sensor_type == SensorType.SEMANTIC:
+            if self._sim.semantic_scene is None:
+                raise RuntimeError(
+                    "SemanticSensor observation requested but no SemanticScene is loaded"
+                )
+            scene = self._sim.get_active_semantic_scene_graph()
+        else:  # SensorType is DEPTH or any other type
+            scene = self._sim.get_active_scene_graph()
 
-    def get_observation(self) -> Union[ndarray, "Tensor"]:
+        # now, connect the agent to the root node of the current scene graph
+
+        # sanity check is not needed on agent:
+        # because if a sensor is attached to a scene graph,
+        # it implies the agent is attached to the same scene graph
+        # (it assumes backend simulator will guarantee it.)
 
+        agent_node = self._agent.scene_node
+        agent_node.parent = scene.get_root_node()
+
+        # get the correct scene graph based on application
+        if self._spec.sensor_type == SensorType.SEMANTIC:
+            scene = self._sim.get_active_semantic_scene_graph()
+        else:  # SensorType is DEPTH or any other type
+            scene = self._sim.get_active_scene_graph()
+
+        render_flags = habitat_sim.gfx.Camera.Flags.NONE
+
+        if self._sim.frustum_culling:
+            render_flags |= habitat_sim.gfx.Camera.Flags.FRUSTUM_CULLING
+
+        self._sim.renderer.enqueue_async_draw_job(
+            self._sensor_object, scene, self.view, render_flags
+        )
+
+    def get_observation(self) -> Union[ndarray, "Tensor"]:
         tgt = self._sensor_object.render_target
 
         if self._spec.gpu2gpu_transfer:
@@ -556,25 +682,22 @@ def get_observation(self) -> Union[ndarray, "Tensor"]:
 
                 obs = self._buffer.flip(0)  # type: ignore[union-attr]
         else:
-            size = self._sensor_object.framebuffer_size
 
             if self._spec.sensor_type == SensorType.SEMANTIC:
-                tgt.read_frame_object_id(
-                    mn.MutableImageView2D(mn.PixelFormat.R32UI, size, self._buffer)
-                )
+                tgt.read_frame_object_id(self.view)
             elif self._spec.sensor_type == SensorType.DEPTH:
-                tgt.read_frame_depth(
-                    mn.MutableImageView2D(mn.PixelFormat.R32F, size, self._buffer)
-                )
+                tgt.read_frame_depth(self.view)
             else:
-                tgt.read_frame_rgba(
-                    mn.MutableImageView2D(
-                        mn.PixelFormat.RGBA8_UNORM,
-                        size,
-                        self._buffer.reshape(self._spec.resolution[0], -1),
-                    )
-                )
+                tgt.read_frame_rgba(self.view)
+
+            obs = np.flip(self._buffer, axis=0)
 
+        return self._noise_model(obs)
+
+    def _get_observation_async(self) -> Union[ndarray, "Tensor"]:
+        if self._spec.gpu2gpu_transfer:
+            obs = self._buffer.flip(0)
+        else:
             obs = np.flip(self._buffer, axis=0)
 
         return self._noise_model(obs)

src/CMakeLists.txt

@@ -47,6 +47,11 @@ option(
   "Build Habitat-Sim with VHACD convex hull decomposition and voxelization engine enabled -- Requires VHACD"
   OFF
 )
+option(
+  BUILD_WITH_BACKGROUND_RENDERER
+  "Build Habitat-Sim with async rendering support.  This will be forced to OFF when building emscripten"
+  ON
+)
 option(BUILD_TEST "Build test binaries" OFF)
 option(USE_SYSTEM_ASSIMP "Use system Assimp instead of a bundled submodule" OFF)
 option(USE_SYSTEM_EIGEN "Use system Eigen instead of a bundled submodule" OFF)
@@ -125,6 +130,15 @@ if(BUILD_PTEX_SUPPORT)
   message("Building with ptex support")
 endif()
 
+if(CORRADE_TARGET_EMSCRIPTEN)
+  if(BUILD_WITH_BACKGROUND_RENDERER)
+    message("Emscriptem build, disabling background rendering")
+  endif()
+  set(BUILD_WITH_BACKGROUND_RENDERER OFF CACHE BOOL "BUILD_WITH_BACKGROUND_RENDERER"
+                                               FORCE
+  )
+endif()
+
 # add subdirectories
 add_subdirectory(esp)