[BUG FIX] Fix heterogeneous entities in the Madrona batch renderer

[Kashu7100]

Summary

Heterogeneous entities — objects whose geometry variant only exists in a subset of environments — were drawn in every environment by the Madrona batch renderer. This is the Madrona-path counterpart to #2958, which fixed the same class of bug in the pyrender rasterizer.

Root cause (and why it differs from #2958)

The pyrender fix in #2958 had two parts: (a) stop compacting per-env poses, and (b) add a per-(primitive, env) active_envs visibility mask.

The Madrona retriever (GenesisGeomRetriever.retrieve_rigid_state_torch) already feeds dense per-env poses — vgeoms_state.pos transposed to [num_worlds, num_vgeoms] — and never compacts them, so part (a) is already correct here. The only missing piece is part (b): per-environment visibility. Without it, a variant vgeom that exists only in envs X, Y is still drawn in all worlds (the Madrona renderer draws every geom in every world).

Change

GenesisGeomRetriever.retrieve_rigid_meshes_static now emits an optional geom_env_mask of shape [num_worlds, num_vgeoms] (1 = visible, 0 = hidden), built from each vgeom's existing active_envs_mask (None ⇒ active in all envs). This is the direct Madrona analog of the active_envs mask added to the pyrender rasterizer in #2958.

• Omitted entirely for homogeneous scenes → preserves legacy behavior and stays back-compatible with Madrona builds that don't yet accept the kwarg.
• No pose changes needed (the dense per-env poses are already aligned); inactive (world, geom) entries are simply masked out and never drawn / never enter the BVH.

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

Dependency

Requires the matching renderer-side change in gs-madrona that consumes geom_env_mask: Genesis-Embodied-AI/gs-madrona#53 (adds the optional geom_env_mask kwarg → per-(world, geom) renderable disable). With an older gs-madrona, the mask is only emitted for heterogeneous scenes, so homogeneous workloads are unaffected; heterogeneous scenes need both PRs.

Testing / follow-up

• py_compile clean; pre-commit (ruff) passes.
• Functional validation needs a CUDA + Madrona environment. Suggested follow-up: a test_render-style test mirroring test_rasterizer_renders_heterogeneous_entities from [BUG FIX] Render each heterogeneous variant in its own environment. #2958 but driving the batch renderer (assert a variant only covers pixels in its active envs). Happy to add it.

Note

Draft — opened for review of the approach alongside gs-madrona#53.

🤖 Generated with Claude Code

[Kashu7100]

[Kashu7100]

video_old.mp4

video_new.mp4

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 12d4450d45

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Gate geom_env_mask on the installed Madrona API

In a default Linux install this repo still pins gs-madrona==0.0.7.post2 in pyproject.toml, while the renderer-side API that accepts geom_env_mask lives in a separate gs-madrona change. For any heterogeneous scene this line adds a new kwarg that the pinned MadronaBatchRenderer binding does not accept, so BatchRenderer.build() fails before rendering instead of preserving the previous behavior; please bump the dependency in this change or only emit the key when the installed adapter supports it.

Useful? React with 👍 / 👎.

[github-actions]

🔴 Benchmark Regression Detected ➡️ Report