[FEATURE] Add support of heterogeneous entities in batch renderer.

genesis/vis/batch_renderer.py

@@ -103,6 +103,18 @@ def retrieve_rigid_meshes_static(self):
         args["geom_sizes"] = np.ones((self.n_vgeoms, 3), dtype=np.float32)
         args["enabled_geom_groups"] = self.default_enabled_geom_groups
 
+        # Heterogeneous entities: a vgeom variant may exist in only a subset of envs
+        # (vgeom.active_envs_mask). Emit an optional [num_worlds, num_geoms] visibility
+        # mask (1 = visible, 0 = hidden) so each variant renders only where it is active;
+        # omitted for homogeneous scenes (legacy + back-compat with older Madrona).
+        if any(vgeom.active_envs_mask is not None for vgeom in vgeoms):
+            num_worlds = max(self.rigid_solver.n_envs, 1)
+            geom_env_mask = np.ones((num_worlds, self.n_vgeoms), dtype=np.int32)
+            for vgeom in vgeoms:
+                if vgeom.active_envs_mask is not None:
+                    geom_env_mask[:, vgeom.idx] = vgeom.active_envs_mask.detach().cpu().numpy()
+            args["geom_env_mask"] = geom_env_mask
+
         # Retrieve material data
         geom_mat_ids = []
         num_materials = 0

pyproject.toml

@@ -67,7 +67,7 @@ dependencies = [
     # Used for loading raytracing special texture images used by LuisaRender
     "OpenEXR",
     # Native batch renderer specifically designed for Genesis
-    "gs-madrona==0.0.7.post2; platform_system=='Linux' and (platform_machine=='x86_64' or platform_machine=='AMD64')",
+    "gs-madrona==0.0.8; platform_system=='Linux' and (platform_machine=='x86_64' or platform_machine=='AMD64')",
     # Used for mesh simplification
     "fast_simplification>=0.1.12",
     # Surface reconstruction library for particle data from SPH simulations

tests/test_render.py

@@ -80,20 +80,25 @@ def skip_if_not_installed(renderer_type):
 
 
 @pytest.mark.required
-@pytest.mark.parametrize("renderer_type", [RENDERER_TYPE.RASTERIZER])
-@pytest.mark.parametrize("env_separate_rigid", [True, False])
-@pytest.mark.parametrize("n_envs", [1, 3])
-def test_rasterizer_renders_heterogeneous_entities(n_envs, env_separate_rigid, show_viewer, png_snapshot, renderer):
-    # A heterogeneous entity instantiates a different morph variant per environment (a sphere in some, the duck mesh in
-    # others). Each variant must render only in its own environments, both in the per-env batched path
-    # (env_separate_rigid + n_envs > 1, env_idx >= 0) and in the combined single-image path (env_idx == -1). A
-    # homogeneous box renders alongside to confirm ordinary entities are unaffected by the per-env masking.
+@pytest.mark.parametrize(
+    "renderer_type",
+    [RENDERER_TYPE.RASTERIZER, RENDERER_TYPE.BATCHRENDER_RASTERIZER, RENDERER_TYPE.BATCHRENDER_RAYTRACER],
+)
+@pytest.mark.parametrize("n_envs", [0, 4])
+def test_renders_heterogeneous_entities(n_envs, show_viewer, png_snapshot, renderer_type, renderer):
+    # A heterogeneous entity instantiates a different morph variant per environment (a sphere in some envs, the duck
+    # mesh in others). Each variant must render only in the environments it is active in - and this must hold across
+    # every supported renderer. The rasterizer composites all environments into a single image (the env_idx == -1
+    # draw-all path), while the Madrona batch renderer returns one image per environment; both must apply the same
+    # per-environment visibility. A homogeneous green box renders alongside to confirm ordinary entities are
+    # unaffected by the masking. The captured snapshot is the ground truth.
     CAM_RES = (160, 120)
+    IS_BATCHRENDER = renderer_type in (RENDERER_TYPE.BATCHRENDER_RASTERIZER, RENDERER_TYPE.BATCHRENDER_RAYTRACER)
 
     scene = gs.Scene(
         renderer=renderer,
         vis_options=gs.options.VisOptions(
-            env_separate_rigid=env_separate_rigid,
+            env_separate_rigid=False,
             shadow=False,
         ),
         show_viewer=show_viewer,
@@ -130,28 +135,56 @@ def test_rasterizer_renders_heterogeneous_entities(n_envs, env_separate_rigid, s
         fov=55.0,
         GUI=show_viewer,
     )
+    # The batch renderer has no built-in lighting, so add explicit lights; the rasterizer lights itself.
+    if IS_BATCHRENDER:
+        scene.add_light(
+            pos=(0.0, 0.0, 1.5),
+            dir=(1.0, 1.0, -2.0),
+            directional=True,
+            castshadow=True,
+            cutoff=45.0,
+            intensity=0.5,
+        )
+        scene.add_light(
+            pos=(4.0, -4.0, 4.0),
+            dir=(-1.0, 1.0, -1.0),
+            directional=False,
+            castshadow=True,
+            cutoff=45.0,
+            intensity=0.5,
+        )
 
     scene.build(n_envs=n_envs)
 
     if sys.platform == "darwin" and scene.visualizer.is_software:
         # Small discrepancies between different hardware due the different physics integration
         png_snapshot.extension._std_err_threshold = 3.0
 
-    # The sphere variant fills the first environments and the duck the last, placed on opposite sides so the duck
-    # (yellow) occupies a known half of the combined image. The box (homogeneous) sits at the back of each environment.
-    box.set_pos(np.array([[0.0, 1.4, 0.15], [0.4, 1.4, 0.15], [-0.4, 1.4, 0.15]])[:n_envs])
-    heterogeneous.set_pos(np.array([[-1.2, 0.3, 0.3], [-1.2, -0.3, 0.3], [1.2, 0.0, 0.0]])[:n_envs])
+    # The sphere variant fills the first environments and the duck the rest, placed on opposite sides so the duck
+    # (yellow) occupies a known half of the image. The box (homogeneous) sits at the back of every environment.
+    box_pos = np.array([[0.0, 1.4, 0.15], [0.4, 1.4, 0.15], [-0.4, 1.4, 0.15], [0.2, 1.4, 0.15]])
+    het_pos = np.array([[-1.2, 0.3, 0.3], [-1.2, -0.3, 0.3], [1.2, 0.3, 0.3], [1.2, -0.3, 0.3]])
+    if n_envs == 0:
+        box.set_pos(box_pos[0])
+        heterogeneous.set_pos(het_pos[0])
+    else:
+        box.set_pos(box_pos[:n_envs])
+        heterogeneous.set_pos(het_pos[:n_envs])
 
     rgb = tensor_to_array(camera.render(rgb=True)[0])
-    # The combined draw-all path is taken unless rendering is both per-env-separated and multi-env.
-    per_env = env_separate_rigid and n_envs > 1
+    # The batch renderer returns one image per environment; the rasterizer composites all environments into one image.
+    per_env = IS_BATCHRENDER and n_envs > 0
     if per_env:
         assert rgb.shape == (n_envs, *CAM_RES[::-1], 3)
         frames = list(rgb)
     else:
         assert rgb.shape == (*CAM_RES[::-1], 3)
         frames = [rgb]
 
+    def duck_yellow(img):
+        # Yellow-pixel mask isolating the (textured) duck variant from the plane, box and white sphere.
+        return (img[..., 0].astype(int) > 120) & (img[..., 1].astype(int) > 120) & (img[..., 2].astype(int) < 110)
+
     # The homogeneous green box renders in every frame.
     for frame in frames:
         assert ((frame[..., 1].astype(int) - frame[..., 0].astype(int)) > 40).sum() > 0
@@ -161,15 +194,22 @@ def test_rasterizer_renders_heterogeneous_entities(n_envs, env_separate_rigid, s
         for i in range(len(frames)):
             for j in range(i + 1, len(frames)):
                 assert (frames[i] != frames[j]).any()
+        # The duck variant is active in only a subset of environments, so it must not leak into the others: at least
+        # one environment must show it and at least one must show none of it.
+        counts = [int(duck_yellow(frame).sum()) for frame in frames]
+        assert max(counts) > 3, f"duck variant never rendered; yellow per env={counts}"
+        assert min(counts) <= max(counts) // 8, f"duck variant leaked into inactive envs; yellow per env={counts}"
     elif n_envs > 1:
-        # Combined image: the duck variant (yellow) belongs to a single environment on one side, so it must not be
+        # Combined image: the duck variant (yellow) belongs to the environments on one side, so it must not be
         # duplicated into the sphere environments on the other side.
         half = rgb.shape[1] // 2
-        yellow = (rgb[..., 0].astype(int) > 120) & (rgb[..., 1].astype(int) > 120) & (rgb[..., 2].astype(int) < 110)
+        yellow = duck_yellow(rgb)
         left, right = int(yellow[:, :half].sum()), int(yellow[:, half:].sum())
-        assert max(left, right) > 3 * min(left, right), (
-            f"duck must render only in its environment; yellow={left, right}"
-        )
+        assert max(left, right) > 3 * min(left, right), f"duck must render only in its env; yellow={left, right}"
+    else:
+        # Single environment: the duck variant is inactive there and must be masked out entirely; only the sphere shows.
+        n_yellow = int(duck_yellow(rgb).sum())
+        assert n_yellow <= 3, f"masked duck variant must not render; yellow={n_yellow}"
 
     for frame in frames:
         assert rgb_array_to_png_bytes(frame) == png_snapshot

tests/utils.py

@@ -44,7 +44,7 @@
 DEFAULT_BRANCH_NAME = "main"
 
 HUGGINGFACE_ASSETS_REVISION = "990a727788f11e34ad006c69bf769303b20cb11c"
-HUGGINGFACE_SNAPSHOT_REVISION = "3e0fb746f2af048825c4a5c1b86ff61832906bf7"
+HUGGINGFACE_SNAPSHOT_REVISION = "e1340f06fabeb7e0453ffe90cced02dcf51cfd78"
 
 MESH_EXTENSIONS = (".mtl", *MESH_FORMATS, *GLTF_FORMATS, *USD_FORMATS)
 IMAGE_EXTENSIONS = (".png", ".jpg")