Merge branch 'main' into yiling/250713_aux_kernels

Author: duburcqa

genesis/ext/pyrender/jit_render.py

@@ -826,6 +826,10 @@ def update_normal(self, node, vertices):
             return self._update_normal_flat(vertices.reshape((-1, 3, 3)))
 
     def update_buffer(self, buffer_updates):
+        # Early return if nothing to do
+        if not buffer_updates:
+            return
+
         updates = np.zeros((len(buffer_updates), 3), dtype=np.int64)
         buffers = []
         for idx, (id, data) in enumerate(buffer_updates.items()):

genesis/ext/pyrender/viewer.py

@@ -400,8 +400,6 @@ def __init__(
         self.pending_offscreen_camera = None
         self.offscreen_result = None
 
-        self.pending_buffer_updates = {}
-
         # Starting the viewer would raise an exception if the OpenGL context is invalid for some reason. This exception
         # must be caught in order to implement some fallback mechanism. One may want to start the viewer from the main
         # thread while the running loop would be running on a background thread. However, this approach is not possible
@@ -655,10 +653,6 @@ def render_offscreen(self, camera_node, render_target, depth=False, seg=False, n
             self.render_flags["depth"] = False
         return self.offscreen_result
 
-    def update_buffers(self):
-        self._renderer.jit.update_buffer(self.pending_buffer_updates)
-        self.pending_buffer_updates.clear()
-
     def wait_until_initialized(self):
         self._initialized_event.wait()
 
@@ -671,7 +665,10 @@ def draw_offscreen(self):
 
         # Make OpenGL context current
         self.switch_to()
-        self.update_buffers()
+
+        # Update the context if not already done before
+        self._renderer.jit.update_buffer(self.gs_context.buffer)
+        self.gs_context.buffer.clear()
 
         self.offscreen_results = []
         self.render_flags["offscreen"] = True
@@ -696,7 +693,10 @@ def on_draw(self):
 
         # Make OpenGL context current
         self.switch_to()
-        self.update_buffers()
+
+        # Update the context if not already done before
+        self._renderer.jit.update_buffer(self.gs_context.buffer)
+        self.gs_context.buffer.clear()
 
         # Render the scene
         self.clear()

genesis/vis/rasterizer.py

@@ -64,11 +64,9 @@ def render_camera(self, camera, rgb=True, depth=False, segmentation=False, norma
             # Set the context
             self._renderer.make_current()
 
-            # Update the buffers
-            if self._buffer_updates is None:
-                gs.raise_exception("Buffers were not set before rendering. Please call 'update_scene' method.")
-
-            self._context.jit.update_buffer(self._buffer_updates)
+            # Update the context if not already done before
+            self._context.jit.update_buffer(self._context.buffer)
+            self._context.buffer.clear()
 
             # Render
             if rgb or depth or normal:
@@ -100,39 +98,37 @@ def render_camera(self, camera, rgb=True, depth=False, segmentation=False, norma
             # Unset the context
             self._renderer.make_uncurrent()
         else:
-            # Update the buffers
-            self._context.jit.update_buffer(self._buffer_updates)
-
             # Render
             if rgb or depth or normal:
-                retval = self._viewer._pyrender_viewer.render_offscreen(
+                retval = self._viewer.render_offscreen(
                     self._camera_nodes[camera.uid],
                     self._camera_targets[camera.uid],
                     depth=depth,
                     normal=normal,
                 )
 
             if segmentation:
-                seg_idxc_rgb_arr, _ = self._viewer._pyrender_viewer.render_offscreen(
+                seg_idxc_rgb_arr, _ = self._viewer.render_offscreen(
                     self._camera_nodes[camera.uid],
                     self._camera_targets[camera.uid],
                     depth=False,
                     normal=False,
                     seg=True,
                 )
 
-        if rgb:
+        if segmentation:
+            seg_idxc_arr = self._context.seg_idxc_rgb_arr_to_idxc_arr(seg_idxc_rgb_arr)
+
+        if rgb or depth or normal:
             rgb_arr = retval[0]
         if depth:
             depth_arr = retval[1]
         if normal:
             normal_arr = retval[2]
-        if segmentation:
-            seg_idxc_arr = self._context.seg_idxc_rgb_arr_to_idxc_arr(seg_idxc_rgb_arr)
         return rgb_arr, depth_arr, seg_idxc_arr, normal_arr
 
     def update_scene(self):
-        self._buffer_updates = self._context.update()
+        self._context.update()
 
     def destroy(self):
         for node in self._camera_nodes.values():
@@ -144,13 +140,13 @@ def destroy(self):
 
         if self._offscreen and self._renderer is not None:
             try:
-                self._renderer._platform.make_current()
+                self._renderer.make_current()
                 self._renderer.delete()
             except OpenGL.error.GLError:
                 pass
             del self._renderer
-            gc.collect()
             self._renderer = None
+            gc.collect()
 
     @property
     def viewer(self):

genesis/vis/rasterizer_context.py

@@ -35,6 +35,8 @@ def __init__(self, options):
         self.rendered_envs_idx = options.rendered_envs_idx
         self.env_separate_rigid = options.env_separate_rigid
 
+        self.buffer = dict()
+
         # nodes
         self.world_frame_node = None
         self.link_frame_nodes = dict()
@@ -831,31 +833,31 @@ def clear_debug_objects(self):
         self.clear_external_nodes()
 
     def update(self):
-        buffer_updates = dict()
+        # Early return if already updated previously
+        if self._t >= self.scene._t:
+            return
 
-        if self._t >= self.scene._t:  # already updated
-            return buffer_updates
-        else:
-            self._t = self.scene._t
+        self._t = self.scene._t
 
         # clear up all dynamic nodes
         self.clear_dynamic_nodes()
 
         # update variables not used in simulation
         self.visualizer.update_visual_states()
 
+        # Reset scene bounds to trigger recomputation. They are involved in shadow map
         self._scene._bounds = None
-        self.update_link_frame(buffer_updates)
-        self.update_tool(buffer_updates)
-        self.update_rigid(buffer_updates)
-        self.update_contact(buffer_updates)
-        self.update_avatar(buffer_updates)
-        self.update_mpm(buffer_updates)
-        self.update_sph(buffer_updates)
-        self.update_pbd(buffer_updates)
-        self.update_fem(buffer_updates)
-
-        return buffer_updates
+
+        self.buffer.clear()
+        self.update_link_frame(self.buffer)
+        self.update_tool(self.buffer)
+        self.update_rigid(self.buffer)
+        self.update_contact(self.buffer)
+        self.update_avatar(self.buffer)
+        self.update_mpm(self.buffer)
+        self.update_sph(self.buffer)
+        self.update_pbd(self.buffer)
+        self.update_fem(self.buffer)
 
     def add_light(self, light):
         # light direction is light pose's -z frame
@@ -875,7 +877,7 @@ def create_node_seg(self, seg_key, seg_node):
             self.seg_node_map[seg_node] = self.seg_idxc_to_idxc_rgb(seg_idxc)
 
     def remove_node_seg(self, seg_node):
-        self.seg_node_map.pop(seg_node)
+        self.seg_node_map.pop(seg_node, None)
 
     def generate_seg_vars(self):
         # seg_key: same as entity/link/geom's idx

genesis/vis/viewer.py

@@ -144,7 +144,8 @@ def update(self, auto_refresh=None):
             self.update_following()
 
         with self.lock:
-            self._pyrender_viewer.pending_buffer_updates |= self.context.update()
+            # Update context
+            self.context.update()
 
             # Refresh viewer by default if and if this is possible
             if auto_refresh is None:
@@ -158,6 +159,9 @@ def update(self, auto_refresh=None):
         if self._max_FPS is not None:
             self.rate.sleep()
 
+    def render_offscreen(self, camera_node, render_target, depth=False, seg=False, normal=False):
+        return self._pyrender_viewer.render_offscreen(camera_node, render_target, depth, seg, normal)
+
     def set_camera_pose(self, pose=None, pos=None, lookat=None):
         """
         Set viewer camera pose.

genesis/vis/visualizer.py

@@ -130,15 +130,16 @@ def reset(self):
 
         self._context.reset()
 
-        # Need to update viewer once here, because otherwise camera will update scene if render is called right after
-        # build, which will lead to segfault.
-        if self._viewer is not None:
-            if self._viewer.is_alive():
-                self._viewer.update(auto_refresh=True)
-
         if self._raytracer is not None:
             self._raytracer.reset()
 
+        if self.viewer_lock is not None:
+            for camera in self._cameras:
+                self._rasterizer.render_camera(camera)
+
+            if self._viewer is not None:
+                self._viewer.update(auto_refresh=True)
+
     def build(self):
         self._context.build(self._scene)
 
@@ -155,15 +156,8 @@ def build(self):
         for camera in self._cameras:
             camera._build()
 
-        if self._cameras:
-            # need to update viewer once here, because otherwise camera will update scene if render is called right
-            # after build, which will lead to segfault.
-            if self._viewer is not None:
-                self._viewer.update(auto_refresh=True)
-            else:
-                # viewer creation will compile rendering kernels if viewer is not created, render here once to compile
-                self._rasterizer.update_scene()
-                self._rasterizer.render_camera(self._cameras[0])
+        # Make sure that the viewer is fully compiled and in a clean state
+        self.reset()
 
     def update(self, force=True, auto=None):
         if force:  # force update