[BUG FIX] Fix render (Genesis-Embodied-AI#1537)

* Fix support of RayTracer on MacOS.
* Fix Rasterizer video recording.
* Fix RayTracer camera initialization.
* Fix support of OSMesa backend for Rasterizer.
* Improve exception handling for 'Rasterizer.render_camera'.
* Enable antialiasing by default for BatchRender.

Author: duburcqa

genesis/engine/scene.py

@@ -1,5 +1,6 @@
 import os
 import pickle
+import sys
 import time
 
 import numpy as np
@@ -526,7 +527,7 @@ def add_camera(
         focus_dist=None,
         GUI=False,
         spp=256,
-        denoise=True,
+        denoise=None,
         env_idx=None,
     ):
         """
@@ -564,14 +565,16 @@ def add_camera(
             Samples per pixel. Only available when using RayTracer renderer. Defaults to 256.
         denoise : bool
             Whether to denoise the camera's rendered image. Only available when using the RayTracer renderer. Defaults
-            to True. If OptiX denoiser is not available in your platform, consider enabling the OIDN denoiser option
-            when building the RayTracer.
+            to True on Linux, otherwise False. If OptiX denoiser is not available in your platform, consider enabling
+            the OIDN denoiser option when building the RayTracer.
 
         Returns
         -------
         camera : genesis.Camera
             The created camera object.
         """
+        if denoise is None:
+            denoise = sys.platform != "darwin"
         return self._visualizer.add_camera(
             res, pos, lookat, up, model, fov, aperture, focus_dist, GUI, spp, denoise, env_idx
         )

genesis/ext/pyrender/offscreen.py

@@ -173,16 +173,22 @@ def render(
             else:
                 if flags & RenderFlags.ENV_SEPARATE:
                     gs.raise_exception("'env_separate_rigid=True' not supported on this platform.")
-                if normal:
-                    gs.raise_exception("'normal=True' not supported on this platform.")
                 renderer.render(scene, flags, seg_node_map)
+                glBindFramebuffer(GL_READ_FRAMEBUFFER, 0)
+                glReadBuffer(GL_FRONT)
                 if depth:
-                    depth = renderer.read_depth_buf()
+                    z_near = scene.main_camera_node.camera.znear
+                    z_far = scene.main_camera_node.camera.zfar
+                    if z_far is None:
+                        z_far = -1.0
+                    depth_arr = renderer.jit.read_depth_buf(self.viewport_height, self.viewport_width, z_near, z_far)
+                    depth_arr = renderer._resize_image(depth_arr, antialias=not seg)
                 if flags & RenderFlags.DEPTH_ONLY:
-                    retval = (depth,)
+                    retval = (depth_arr,)
                 else:
-                    color = renderer.read_color_buf()
-                    retval = (color, depth)
+                    color_arr = renderer.jit.read_color_buf(self.viewport_height, self.viewport_width, rgba=False)
+                    color_arr = renderer._resize_image(color_arr, antialias=not seg)
+                    retval = (color_arr, depth_arr) if depth else (color_arr,)
         else:
             retval = ()
 
@@ -207,7 +213,16 @@ def get_program(self, vertex_shader, fragment_shader, geometry_shader=None, defi
             if env_separate_rigid:
                 flags |= RenderFlags.ENV_SEPARATE
             glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
-            normal_arr, *_ = renderer.render(scene, flags, is_first_pass=False, force_skip_shadows=True)
+
+            if self._platform.supports_framebuffers():
+                normal_arr, *_ = renderer.render(scene, flags, is_first_pass=False, force_skip_shadows=True)
+            else:
+                glBindFramebuffer(GL_READ_FRAMEBUFFER, 0)
+                glReadBuffer(GL_FRONT)
+                renderer.render(scene, flags, is_first_pass=False, force_skip_shadows=True)
+                normal_arr = renderer.jit.read_color_buf(self.viewport_height, self.viewport_width, rgba=False)
+                normal_arr = renderer._resize_image(normal_arr, antialias=not seg)
+
             retval = (*retval, normal_arr)
 
             renderer._program_cache = old_cache

genesis/ext/pyrender/viewer.py

@@ -1063,11 +1063,12 @@ def _save_image(self):
         filename = self._get_save_filename(["png", "jpg", "gif", "all"])
         if filename is not None:
             self.viewer_flags["save_directory"] = os.path.dirname(filename)
-            imageio.imwrite(filename, self._renderer.read_color_buf())
+            data = self._renderer.jit.read_color_buf(*self._viewport_size, rgba=False)
+            imageio.imwrite(filename, img_arr)
 
     def _record(self):
         """Save another frame for the GIF."""
-        data = self._renderer.read_color_buf()
+        data = self._renderer.jit.read_color_buf(*self._viewport_size, rgba=False)
         if not np.all(data == 0.0):
             self.video_recorder.write_frame(data)
 

genesis/vis/camera.py

@@ -130,6 +130,13 @@ def __init__(
             self._focus_dist = np.linalg.norm(np.asarray(lookat) - np.asarray(pos))
 
     def build(self):
+        n_envs = max(self._visualizer.scene.n_envs, 1)
+        self._multi_env_pos_tensor = torch.empty((n_envs, 3), dtype=gs.tc_float, device=gs.device)
+        self._multi_env_lookat_tensor = torch.empty((n_envs, 3), dtype=gs.tc_float, device=gs.device)
+        self._multi_env_up_tensor = torch.empty((n_envs, 3), dtype=gs.tc_float, device=gs.device)
+        self._multi_env_transform_tensor = torch.empty((n_envs, 4, 4), dtype=gs.tc_float, device=gs.device)
+        self._multi_env_quat_tensor = torch.empty((n_envs, 4), dtype=gs.tc_float, device=gs.device)
+
         self._envs_offset = torch.as_tensor(self._visualizer._scene.envs_offset, dtype=gs.tc_float, device=gs.device)
 
         self._batch_renderer = self._visualizer.batch_renderer
@@ -154,7 +161,12 @@ def build(self):
                 self._other_stacked = self._visualizer._context.env_separate_rigid
 
         self._is_built = True
-        self.setup_initial_env_poses()
+        self.set_pose(
+            transform=self._initial_transform, pos=self._initial_pos, lookat=self._initial_lookat, up=self._initial_up
+        )
+        # FIXME: For some reason, it is necessary to update the camera twice...
+        if self._raytracer is not None:
+            self._raytracer.update_camera(self)
 
     def attach(self, rigid_link, offset_T):
         """
@@ -296,8 +308,8 @@ def _batch_render(
         depth=False,
         segmentation=False,
         normal=False,
+        antialiasing=True,
         force_render=False,
-        antialiasing=False,
     ):
         """
         Render the camera view with batch renderer.
@@ -371,7 +383,7 @@ def render(
         normal_arr : np.ndarray
             The rendered surface normal(s).
         """
-        rgb_arr, depth_arr, seg_arr, seg_color_arr, normal_arr = None, None, None, None, None
+        rgb_arr, depth_arr, seg_arr, seg_color_arr, seg_idxc_arr, normal_arr = None, None, None, None, None, None
 
         if self._batch_renderer is not None:
             rgb_arr, depth_arr, seg_idxc_arr, normal_arr = self._batch_render(
@@ -521,28 +533,6 @@ def render_pointcloud(self, world_frame=True):
         point_cloud = point_cloud[:, :3].reshape((*depth_arr.shape, 3))
         return point_cloud, mask
 
-    @gs.assert_built
-    def setup_initial_env_poses(self):
-        """
-        Setup the camera poses for multiple environments.
-        """
-        if self._initial_transform is not None:
-            assert self._initial_transform.shape == (4, 4)
-            self._initial_pos, self._initial_lookat, self._initial_up = gu.T_to_pos_lookat_up(self._initial_transform)
-        else:
-            self._initial_transform = gu.pos_lookat_up_to_T(self._initial_pos, self._initial_lookat, self._initial_up)
-
-        n_envs = max(self._visualizer.scene.n_envs, 1)
-        self._multi_env_pos_tensor = self._initial_pos.repeat((n_envs, 1))
-        self._multi_env_lookat_tensor = self._initial_lookat.repeat((n_envs, 1))
-        self._multi_env_up_tensor = self._initial_up.repeat((n_envs, 1))
-        self._multi_env_transform_tensor = self._initial_transform.repeat((n_envs, 1, 1))
-        self._multi_env_quat_tensor = _T_to_quat_for_madrona(self._multi_env_transform_tensor)
-
-        self._rasterizer.update_camera(self)
-        if self._raytracer is not None:
-            self._raytracer.update_camera(self)
-
     @gs.assert_built
     def set_pose(self, transform=None, pos=None, lookat=None, up=None, env_idx=None):
         """

genesis/vis/rasterizer.py

@@ -69,36 +69,37 @@ def render_camera(self, camera, rgb=True, depth=False, segmentation=False, norma
             self._context.buffer.clear()
 
             # Render
-            if rgb or depth or normal:
-                retval = self._renderer.render(
-                    self._context._scene,
-                    self._camera_targets[camera.uid],
-                    camera_node=self._camera_nodes[camera.uid],
-                    env_separate_rigid=self._context.env_separate_rigid,
-                    rgb=rgb,
-                    normal=normal,
-                    seg=False,
-                    depth=depth,
-                    plane_reflection=rgb and self._context.plane_reflection,
-                    shadow=rgb and self._context.shadow,
-                )
-
-            if segmentation:
-                seg_idxc_rgb_arr, *_ = self._renderer.render(
-                    self._context._scene,
-                    self._camera_targets[camera.uid],
-                    camera_node=self._camera_nodes[camera.uid],
-                    env_separate_rigid=self._context.env_separate_rigid,
-                    rgb=False,
-                    normal=False,
-                    seg=True,
-                    depth=False,
-                    plane_reflection=False,
-                    shadow=False,
-                )
-
-            # Unset the context
-            self._renderer.make_uncurrent()
+            try:
+                if rgb or depth or normal:
+                    retval = self._renderer.render(
+                        self._context._scene,
+                        self._camera_targets[camera.uid],
+                        camera_node=self._camera_nodes[camera.uid],
+                        env_separate_rigid=self._context.env_separate_rigid,
+                        rgb=rgb,
+                        normal=normal,
+                        seg=False,
+                        depth=depth,
+                        plane_reflection=rgb and self._context.plane_reflection,
+                        shadow=rgb and self._context.shadow,
+                    )
+
+                if segmentation:
+                    seg_idxc_rgb_arr, *_ = self._renderer.render(
+                        self._context._scene,
+                        self._camera_targets[camera.uid],
+                        camera_node=self._camera_nodes[camera.uid],
+                        env_separate_rigid=self._context.env_separate_rigid,
+                        rgb=False,
+                        normal=False,
+                        seg=True,
+                        depth=False,
+                        plane_reflection=False,
+                        shadow=False,
+                    )
+            finally:
+                # Unset the context
+                self._renderer.make_uncurrent()
         else:
             # Render
             if rgb or depth or normal:

genesis/vis/raytracer.py

@@ -604,7 +604,7 @@ def add_camera(self, camera):
         if camera_model == "pinhole":
             self._cameras[camera_name] = LuisaRenderPy.PinholeCamera(
                 name=camera_name,
-                pose=self.get_transform(camera.transform),
+                pose=self.get_transform(np.eye(4)),
                 film=LuisaRenderPy.Film(resolution=camera.res),
                 filter=LuisaRenderPy.Filter(),
                 spp=camera.spp,
@@ -613,7 +613,7 @@ def add_camera(self, camera):
         elif camera_model == "thinlens":
             self._cameras[camera_name] = LuisaRenderPy.ThinLensCamera(
                 name=camera_name,
-                pose=self.get_transform(camera.transform),
+                pose=self.get_transform(np.eye(4)),
                 film=LuisaRenderPy.Film(resolution=camera.res),
                 filter=LuisaRenderPy.Filter(),
                 spp=camera.spp,

tests/test_render.py

@@ -119,7 +119,7 @@ def test_segmentation(segmentation_level, particle_mode):
 
 @pytest.mark.required
 @pytest.mark.flaky(reruns=3, condition=(sys.platform == "darwin"))
-def test_batched_offscreen_rendering(show_viewer, tol):
+def test_batched_offscreen_rendering(tmp_path, show_viewer, tol):
     scene = gs.Scene(
         vis_options=gs.options.VisOptions(
             # rendered_envs_idx=(0, 1, 2),
@@ -247,6 +247,7 @@ def test_batched_offscreen_rendering(show_viewer, tol):
     )
     scene.build(n_envs=3, env_spacing=(2.0, 2.0))
 
+    cam.start_recording()
     for _ in range(7):
         dofs_lower_bound, dofs_upper_bound = robot.get_dofs_limit()
         qpos = dofs_lower_bound + (dofs_upper_bound - dofs_lower_bound) * torch.rand(robot.n_qs)
@@ -270,6 +271,7 @@ def test_batched_offscreen_rendering(show_viewer, tol):
 
         for i in range(3):
             assert_allclose(steps_rgb_arrays[0][i], steps_rgb_arrays[1][i], tol=tol)
+    cam.stop_recording(save_to_filename=(tmp_path / "video.mp4"))
 
 
 @pytest.mark.required