[FEATURE] Add New Features for Madrona Batch Rendering

examples/rigid/single_franka_batch_render.py

@@ -13,16 +13,20 @@ def main():
     parser.add_argument("-b", "--n_envs", type=int, default=3)
     parser.add_argument("-s", "--n_steps", type=int, default=2)
     parser.add_argument("-r", "--render_all_cameras", action="store_true", default=False)
-    parser.add_argument("-o", "--output_dir", type=str, default="img_output/test")
+    parser.add_argument("-o", "--output_dir", type=str, default="data/test")
     parser.add_argument("-u", "--use_rasterizer", action="store_true", default=False)
     parser.add_argument("-d", "--debug", action="store_true", default=False)
+    parser.add_argument("-l", "--seg_level", type=str, default="link")
     args = parser.parse_args()
 
     ########################## init ##########################
     gs.init(backend=gs.cpu if args.cpu else gs.gpu)
 
     ########################## create a scene ##########################
     scene = gs.Scene(
+        vis_options=gs.options.VisOptions(
+            segmentation_level=args.seg_level,
+        ),
         renderer=gs.options.renderers.BatchRenderer(
             use_rasterizer=args.use_rasterizer,
         ),
@@ -61,21 +65,31 @@ def main():
         fov=45,
         GUI=args.vis,
     )
+    cam_2 = scene.add_camera(
+        res=(512, 512),
+        pos=(0.0, 0.1, 5.0),
+        lookat=(0.0, 0.0, 0.0),
+        fov=45,
+        GUI=args.vis,
+    )
+
     scene.add_light(
-        pos=[0.0, 0.0, 1.5],
-        dir=[1.0, 1.0, -2.0],
-        directional=1,
-        castshadow=1,
+        pos=(0.0, 0.0, 1.5),
+        dir=(1.0, 1.0, -2.0),
+        color=(1.0, 0.0, 0.0),
+        directional=True,
+        castshadow=True,
         cutoff=45.0,
         intensity=0.5,
     )
     scene.add_light(
-        pos=[4, -4, 4],
-        dir=[-1, 1, -1],
-        directional=0,
-        castshadow=1,
-        cutoff=45.0,
-        intensity=0.5,
+        pos=(4, -4, 4),
+        dir=(0, 0, -1),
+        directional=False,
+        castshadow=True,
+        cutoff=80.0,
+        intensity=1.0,
+        attenuation=0.1,
     )
 
     ########################## build ##########################
@@ -91,11 +105,19 @@ def main():
         if args.debug:
             debug_cam.render()
         if args.render_all_cameras:
-            rgba, depth, _, _ = scene.render_all_cameras(rgb=True, depth=True)
-            exporter.export_frame_all_cameras(i, rgb=rgba, depth=depth)
+            color, depth, seg, normal = scene.render_all_cameras(
+                rgb=True, depth=i % 2 == 1, segmentation=i % 2 == 1, normal=True
+            )
+            exporter.export_frame_all_cameras(i, rgb=color, depth=depth, segmentation=seg, normal=normal)
         else:
-            rgba, depth, _, _ = cam_1.render(rgb=True, depth=True)
-            exporter.export_frame_single_camera(i, cam_1.idx, rgb=rgba, depth=depth)
+            color, depth, seg, normal = cam_1.render(
+                rgb=False,
+                depth=True,
+                segmentation=True,
+                colorize_seg=True,
+                normal=False,
+            )
+            exporter.export_frame_single_camera(i, cam_1.idx, rgb=seg, depth=depth, segmentation=None, normal=normal)
     if args.debug:
         debug_cam.stop_recording("debug_cam.mp4")
 

genesis/engine/scene.py

@@ -451,77 +451,86 @@ def link_entities(
         parent_link._child_idxs.append(child_link.idx)
 
     @gs.assert_unbuilt
-    def add_light(
+    def add_mesh_light(
         self,
-        *,
         morph: Morph | None = None,
         color: ArrayLike | None = (1.0, 1.0, 1.0, 1.0),
         intensity: float = 20.0,
         revert_dir: bool | None = False,
         double_sided: bool | None = False,
-        beam_angle: float | None = 180.0,
-        pos: ArrayLike | None = None,
-        dir: ArrayLike | None = None,
-        directional: bool | None = None,
-        castshadow: bool | None = None,
-        cutoff: float | None = None,
+        cutoff: float | None = 180.0,
     ):
         """
-        Add a light to the scene.
-
-        Warning
-        -------
-        The signature of this method is different depending on the renderer being used, i.e.:
-        - RayTracer: 'add_light(self, morph, color, intensity, revert_dir, double_sided, beam_angle)'
-        - BatchRender: 'add_ligth(self, pos, dir, intensity, directional, castshadow, cutoff)'
-        - Rasterizer: **Unsupported**
+        Add a mesh light to the scene. Only supported by RayTracer.
 
         Parameters
         ----------
         morph : gs.morphs.Morph
-            The morph of the light. Must be an instance of `gs.morphs.Primitive` or `gs.morphs.Mesh`. Only supported by
-            RayTracer.
+            The morph of the light. Must be an instance of `gs.morphs.Primitive` or `gs.morphs.Mesh`.
         color : tuple of float, shape (3,)
-            The color of the light, specified as (r, g, b). Only supported by RayTracer.
+            The color of the light, specified as (r, g, b).
         intensity : float
             The intensity of the light.
         revert_dir : bool
             Whether to revert the direction of the light. If True, the light will be emitted towards the mesh's inside.
-            Only supported by RayTracer.
         double_sided : bool
-            Whether to emit light from both sides of surface. Only supported by RayTracer.
-        beam_angle : float
-            The beam angle of the light. Only supported by RayTracer.
+            Whether to emit light from both sides of surface.
+        cutoff : float
+            The cutoff angle of the light in degrees. Range: [0.0, 180.0].
+        """
+        if not isinstance(self.renderer_options, gs.renderers.RayTracer):
+            gs.raise_exception(
+                "add_mesh_light() is only supported when using RayTracer. "
+                "If you are using BatchRenderer, please use add_light()."
+            )
+
+        if not isinstance(morph, (gs.morphs.Primitive, gs.morphs.Mesh)):
+            gs.raise_exception("Light morph only supports `gs.morphs.Primitive` or `gs.morphs.Mesh`.")
+        mesh = gs.Mesh.from_morph_surface(morph, gs.surfaces.Plastic(smooth=False))
+        self._visualizer.add_mesh_light(mesh, color, intensity, morph.pos, morph.quat, revert_dir, double_sided, cutoff)
+
+    @gs.assert_unbuilt
+    def add_light(
+        self,
+        pos: ArrayLike,
+        dir: ArrayLike,
+        color: ArrayLike = (1.0, 1.0, 1.0),
+        intensity: float = 1.0,
+        directional: bool = False,
+        castshadow: bool = True,
+        cutoff: float = 45.0,
+        attenuation: float = 0.0,
+    ):
+        """
+        Add a light to the scene for batch renderer.
+
+        Parameters
+        ----------
         pos : tuple of float, shape (3,)
-            The position of the light, specified as (x, y, z). Only supported by BatchRenderer.
+            The position of the light, specified as (x, y, z).
         dir : tuple of float, shape (3,)
-            The direction of the light, specified as (x, y, z). Only supported by BatchRenderer.
+            The direction of the light, specified as (x, y, z).
+        color : tuple of float, shape (3,)
+            The color of the light, specified as (r, g, b).
         intensity : float
-            The intensity of the light. Only supported by BatchRenderer.
+            The intensity of the light.
         directional : bool
-            Whether the light is directional. Only supported by BatchRenderer.
+            Whether the light is directional.
         castshadow : bool
-            Whether the light casts shadows. Only supported by BatchRenderer.
+            Whether the light casts shadows.
         cutoff : float
-            The cutoff angle of the light in degrees. Only supported by BatchRenderer.
+            The cutoff angle of the light in degrees. Range: (0.0, 90.0).
+        attenuation : float
+            The attenuation factor of the light.
+            Light intensity will attenuate by distance with (1 / (1 + attenuation * distance ^ 2))
         """
-        if self._visualizer.batch_renderer is not None:
-            if any(map(lambda e: e is None, (pos, dir, intensity, directional, castshadow, cutoff))):
-                gs.raise_exception("Input arguments do not complain with expected signature when using 'BatchRenderer'")
-
-            self.visualizer.add_light(pos, dir, intensity, directional, castshadow, cutoff)
-        elif self.visualizer.raytracer is not None:
-            if any(map(lambda e: e is None, (morph, color, intensity, revert_dir, double_sided, beam_angle))):
-                gs.raise_exception("Input arguments do not complain with expected signature when using 'RayTracer'")
-            if not isinstance(morph, (gs.morphs.Primitive, gs.morphs.Mesh)):
-                gs.raise_exception("Light morph only supports `gs.morphs.Primitive` or `gs.morphs.Mesh`.")
-
-            mesh = gs.Mesh.from_morph_surface(morph, gs.surfaces.Plastic(smooth=False))
-            self.visualizer.raytracer.add_mesh_light(
-                mesh, color, intensity, morph.pos, morph.quat, revert_dir, double_sided, beam_angle
+        if not isinstance(self.renderer_options, gs.renderers.BatchRenderer):
+            gs.raise_exception(
+                "add_light() is only supported when using BatchRenderer. "
+                "If you are using Raytracer, please use add_mesh_light()."
             )
-        else:
-            gs.raise_exception("Adding lights is only supported by 'RayTracer' and 'BatchRenderer'.")
+
+        self.visualizer.add_light(pos, dir, color, intensity, directional, castshadow, cutoff, attenuation)
 
     @gs.assert_unbuilt
     def add_sensor(self, sensor_options: "SensorOptions"):
@@ -541,6 +550,8 @@ def add_camera(
         GUI=False,
         spp=256,
         denoise=None,
+        near=0.05,
+        far=100.0,
         env_idx=None,
         debug=False,
     ):
@@ -581,6 +592,12 @@ def add_camera(
             Whether to denoise the camera's rendered image. Only available when using the RayTracer renderer. Defaults
             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.
+        near: float
+            Distance from camera center to near plane in meters.
+            Only available when using rasterizer in Rasterizer and BatchRender renderer. Defaults to 0.05.
+        far: float
+            Distance from camera center to far plane in meters.
+            Only available when using rasterizer in Rasterizer and BatchRender renderer. Defaults to 100.0.
         debug : bool
             Whether to use the debug camera. It enables to create cameras that can used to monitor / debug the
             simulation without being part of the "sensors". Their output is rendered by the usual simple Rasterizer
@@ -596,7 +613,7 @@ def add_camera(
         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, debug
+            res, pos, lookat, up, model, fov, aperture, focus_dist, GUI, spp, denoise, near, far, env_idx, debug
         )
 
     @gs.assert_unbuilt
@@ -1106,7 +1123,16 @@ def draw_debug_path(self, qposs, entity, link_idx=-1, density=0.3, frame_scaling
             )
 
     @gs.assert_built
-    def render_all_cameras(self, rgb=True, depth=False, normal=False, segmentation=False, force_render=False):
+    def render_all_cameras(
+        self,
+        rgb=True,
+        depth=False,
+        segmentation=False,
+        colorize_seg=False,
+        normal=False,
+        antialiasing=False,
+        force_render=False,
+    ):
         """
         Render the scene for all cameras using the batch renderer.
 
@@ -1116,10 +1142,12 @@ def render_all_cameras(self, rgb=True, depth=False, normal=False, segmentation=F
             Whether to render the rgb image.
         depth : bool, optional
             Whether to render the depth image.
-        normal : bool, optional
-            Whether to render the normal image.
         segmentation : bool, optional
             Whether to render the segmentation image.
+        normal : bool, optional
+            Whether to render the normal image.
+        antialiasing : bool, optional
+            Whether to apply anti-aliasing.
         force_render : bool, optional
             Whether to force render the scene.
 
@@ -1131,7 +1159,12 @@ def render_all_cameras(self, rgb=True, depth=False, normal=False, segmentation=F
         if self._visualizer.batch_renderer is None:
             gs.raise_exception("Method only supported by 'BatchRenderer'")
 
-        return self._visualizer.batch_renderer.render(rgb, depth, normal, segmentation, force_render)
+        rgb_out, depth_out, seg_out, normal_out = self._visualizer.batch_renderer.render(
+            rgb, depth, segmentation, normal, antialiasing, force_render
+        )
+        if segmentation and colorize_seg:
+            seg_out = tuple(self._visualizer.batch_renderer.colorize_seg_idxc_arr(seg) for seg in seg_out)
+        return rgb_out, depth_out, seg_out, normal_out
 
     @gs.assert_built
     def clear_debug_object(self, object):
@@ -1382,3 +1415,18 @@ def fem_solver(self):
     def pbd_solver(self):
         """The scene's `pbd_solver`, managing all the `PBDEntity` in the scene."""
         return self._sim.pbd_solver
+
+    @property
+    def segmentation_idx_dict(self):
+        """
+        Returns a dictionary mapping segmentation indices to scene entities.
+
+        In the segmentation map:
+        - Index 0 corresponds to the background (-1).
+        - Indices > 0 correspond to scene elements, which may be represented as:
+            - `entity_id`
+            - `(entity_id, link_id)`
+            - `(entity_id, link_id, geom_id)`
+          depending on the material type and the configured segmentation level.
+        """
+        return self._visualizer.segmentation_idx_dict

genesis/options/surfaces.py

@@ -67,8 +67,8 @@ class Surface(Options):
     double_sided : bool | None, optional
         Whether to render both sides of the surface. Useful for non-watertight 2D objects. Defaults to True for Cloth
         material and False for others.
-    beam_angle : float
-        The beam angle of emission. Defaults to 180.0.
+    cutoff : float
+        The cutoff angle of emission. Defaults to 180.0.
     normal_diff_clamp : float, optional
         Controls the threshold for computing surface normals by interpolating vertex normals.
     recon_backend : str, optional
@@ -99,8 +99,8 @@ class Surface(Options):
     vis_mode: Optional[str] = None
     smooth: bool = True
     double_sided: Optional[bool] = None
-    beam_angle: float = 180
-    normal_diff_clamp: float = 180
+    cutoff: float = 180.0
+    normal_diff_clamp: float = 180.0
     recon_backend: str = "splashsurf"
     generate_foam: bool = False
     foam_options: Optional[FoamOptions] = None