Format fix.

habitat-lab/habitat/core/batch_renderer_helper.py

@@ -0,0 +1,62 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Meta Platforms, 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.
+
+
+
+from typing import TYPE_CHECKING, List, Optional, OrderedDict
+from habitat.core.batch_renderer import BatchRenderer
+import numpy as np
+from omegaconf import DictConfig
+
+
+class BatchRendererHelper:
+    r"""Provides :ref:`VectorEnv` with batch rendering capability.
+
+    Instead of individually rendering their environment, the worker simulators include their keyframe into observations.
+    The BatchRenderedVectorEnv then provides these observations to a batch renderer to produce all visual sensor observations simultaneously.
+
+    Refer to the BatchRenderer class.
+    """
+    _config: "DictConfig"
+    _num_envs: int
+    _batch_renderer: BatchRenderer
+
+    def __init__(self, config: "DictConfig", num_envs: int) -> None:
+        r"""Initialize batch renderer.
+
+        :param config: Base configuration.
+        :param num_envs: Number of environments to render."""
+        self._config = config
+        self._num_envs = num_envs
+        self._batch_renderer = BatchRenderer(config, self.num_envs)
+
+    def debug_render(
+        self, mode: str = "human", *args, **kwargs
+    ) -> Optional[np.ndarray]:
+        r"""Creates a tiled image from observations rendered during the last post_step call.
+        Use for debugging and testing only."""
+        images = self._batch_renderer.copy_output_to_image()
+        tile = tile_images(images)
+        if mode == "human":
+            from habitat.core.utils import try_cv2_import
+
+            cv2 = try_cv2_import()
+            cv2.imshow("BatchRenderedVectorEnv", tile[:, :, ::-1])
+            cv2.waitKey(1)
+            return None
+        elif mode == "rgb_array":
+            return tile
+        else:
+            raise NotImplementedError
+
+    def post_step(self, observations) -> List[OrderedDict]:
+        r"""
+        Renders observations for all environments by consuming keyframe observations.
+
+        :param observations: List of observations for each environment.
+        :return: List of rendered observations for each environment.
+        """
+        return self._batch_renderer.post_step(observations)