Semantic segmentation parity by tkelestemur · Pull Request #1283 · google-deepmind/mujoco_warp

tkelestemur · 2026-04-12T23:42:27Z

Summary

This PR adds first-class semantic segmentation to MJWarp's renderer.

Adds mjw.get_semantic_segmentation(...), returning per-pixel (object_id, object_type) pairs aligned with MuJoCo segmentation semantics.
Preserves the existing mjw.get_segmentation(...) API and buffer for backward compatibility.
Reports regular geom hits as (geom_id, mjOBJ_GEOM) and flex hits as (flex_id, mjOBJ_FLEX), instead of reducing flex hits to the legacy -2 sentinel only.
Exposes the new API from the public package surface and adds test coverage for the new renderer/accessor behavior.
Adds a dedicated semantic segmentation notebook demo, including flex examples, RGB/depth rendering, legacy segmentation, semantic object ids, semantic object types, and semantic flex ids.

PS: The new notebook imports mujoco_warp directly from the local checkout and disables IPython autoreload for the session, since Warp kernels require file-backed Python source in notebook environments.

Aside from the regular tests, I've also run pytest contrib/kernel_analyzer/kernel_analyzer/ast_analyzer_test.py.

thowell · 2026-04-13T15:00:37Z

@tkelestemur thank you for contributing this feature to mujoco warp!

@StafaH @btaba would it make sense to combine this feature (in a potentially breaking way) with the existing segmentation api?

@tkelestemur instead of adding a separate notebook for this feature, can we add an example to the existing tutorial notebook?

StafaH · 2026-04-13T19:00:28Z

Hi @tkelestemur, thanks for the PR!

In this case I believe this new feature should not be implemented in MJWarp, and can be implemented at the framework level that you use downstream. The existing segmentation output provides enough information to perform post processing downstream e.g. by implementing semantic segmentation as an extension in Mjlab with a kernel that converts the existing integer id to another value inplace.

tkelestemur · 2026-04-14T00:00:19Z

@StafaH The existing segmentation buffer is sufficient to reconstruct semantic output for geom hits and background, but not for flex hits, because all flex hits are collapsed to -2 and lose the underlying flex_id.

StafaH · 2026-04-14T20:05:19Z

I agree, we should fix existing segmentation output. I think what @thowell mentioned is correct, we should focus on having the existing segmentation output match MuJoCo exactly instead of making a seperate output, and we can update the tests to check that this segmentation output matches exactly the C MuJoCo output (similar to what was done for depth recently).

WDYT @thowell?

@tkelestemur this might also require opening an issue and a PR in mjlab to fix any changes that happen downstream there.

tkelestemur · 2026-04-14T20:52:56Z

@StafaH Sounds good -- I'll update this PR to use the existing segmentation API and also open an issue and a PR in mjlab.

tkelestemur · 2026-04-15T14:03:43Z

@thowell I've removed the extra notebook and added a section to the current notebook.

@StafaH I've reverted the new segmentation api and updated the current one. Soon, I'll open a PR on the mjlab side.

thowell · 2026-04-15T22:20:12Z

-    seg = rc.seg_data.numpy()
-    self.assertTrue(np.any(seg >= 0), "Expected geom hits from auto-detected seg")
-    self.assertGreater(np.unique(seg).shape[0], 1)
+    seg = wp.zeros((1, 32, 32, 2), dtype=int)


can we create seg from the render context?

Updated to size the segmentation output from rc.cam_res instead of hardcoding the resolution. Is that what you meant?

lets change this back to what you had before wp.zeros((1, 32, 32, 2), dtype=int). thanks!

thowell · 2026-04-15T22:22:34Z

        "media.show_image(depth_grid, cmap='gray', vmin=0, vmax=1)"
      ]
    },
+    {


i think we can simplify this and make it similar to the depth render example (essentially a comment plus 2 lines of code and then an additional segmentation image to display.

Simplified the tutorial example to extend the existing rendering cell with a small segmentation extraction/display block. Lmk how it looks.

thowell · 2026-04-15T22:23:10Z

the updated pr looks great! added a few comments. please let us known if you have any questions. thanks!

tkelestemur · 2026-04-16T00:33:52Z

Thanks @thowell I've addressed your reviews. I also opened an downstream PR for mjlab here: mujocolab/mjlab#911

thowell · 2026-04-16T15:29:10Z

+        "segmentation_grid = segmentation_data.numpy()[..., 0].reshape(4, 4, CAM_RES[1], CAM_RES[0])\n",
+        "segmentation_grid = segmentation_grid.transpose(0, 2, 1, 3)\n",
+        "segmentation_grid = segmentation_grid.reshape(4 * CAM_RES[0], 4 * CAM_RES[1])\n",
+        "media.show_image(segmentation_grid, cmap='viridis', vmin=-1, vmax=max(1, int(segmentation_grid.max())))\n"


please confirm that this line runs as expected

ValueError: Type int32 is not a valid media data type (uint or float).

converting the array to floats produces

tkelestemur · 2026-04-16T18:14:42Z

sorry forgot to enable render_seg. This is what I'm getting right now:

so should be good to go.

StafaH

This is shaping up nicely. Left some comments.

StafaH · 2026-04-17T01:26:35Z

Also might be worth having the seg test xml include 1 flex object for completeness

tkelestemur · 2026-04-17T16:02:55Z

@StafaH thanks for the reviews, I think I addressed all of them.

In f4bf916, I also dded a small flex object to the synthetic segmentation XML for completeness, and the test now asserts the model includes one flex before checking the render-context segmentation setup.

Let me know if this is enough.

StafaH

Thanks @tkelestemur, left one last comment, but otherwise LGTM

thowell · 2026-04-22T13:31:04Z

@tkelestemur some of the checks are not passing, please take a look. might just need to sync with main. thanks!

…parity

thowell · 2026-04-22T17:01:02Z

there is a segmentation fault with one of the checks, please take a look. thanks! https://github.com/google-deepmind/mujoco_warp/actions/runs/24784925329/job/72532705830?pr=1283

tkelestemur · 2026-04-22T17:06:53Z

@thowell on it!

tkelestemur · 2026-04-22T19:04:49Z

@thowell I pushed a fix for the Linux CPU segfault in 98d6411.

Root cause was that build_flex_bvh() creates a grouped wp.Mesh, but we were looking up roots with the BVH path. For flexes the correct lookup is mesh_get_group_root, and in this Warp version that is only exposed as a kernel builtin, not a normal Python host API. So the final change keeps the surface area small:

add a private _compute_mesh_group_roots helper just to call the mesh-specific builtin
switch flex BVH root lookup over to that helper
fix flex_group_root layout in create_render_context() to match the renderer indexing ([worldid, flexid])

I considered a smaller nworld == 1 -> -1 workaround, but that would only paper over the single-world case instead of fixing the underlying mesh-vs-BVH mismatch. I also considered splitting flexes into one mesh per world, but that would be a much broader change.

I rechecked the targeted regressions locally on macOS and on Linux CPU (l4):

render_util_test.py::test_get_segmentation_preserves_flex_ids
render_test.py -k segmentation_matches_mujoco (skips on headless Linux as expected)
io_test.py -k segmentation_from_camera_output

All of those pass with this change, so this should address the CI segfault without broadening the PR beyond the failing path.

StafaH · 2026-04-22T19:35:05Z


+# Warp exposes mesh group-root lookup as a kernel builtin in this version.
+@wp.kernel
+def _compute_mesh_group_roots(


nit: the _ is not necessary

@tkelestemur can we update this as @StafaH suggests? thanks!

done in 0fdb3e6

…parity

thowell · 2026-04-27T22:06:35Z

@tkelestemur thank you for this contribution!

tkelestemur added 6 commits April 11, 2026 18:24

Add semantic segmentation render parity

5890a38

Add semantic segmentation tutorial example

e20288a

Format io test skip message

c26825e

Fix API import ordering

3e97ff6

Add dedicated semantic segmentation notebook

2e3a61f

Fix notebook local import setup

907a7ee

thowell requested a review from StafaH April 13, 2026 14:51

thowell requested a review from btaba April 13, 2026 15:00

tkelestemur added 2 commits April 14, 2026 21:57

Update segmentation API and tutorial

73a9282

Merge upstream/main into tarik/semantic-segmentation-parity

0ebe870

tkelestemur mentioned this pull request Apr 15, 2026

Keep segmentation compatible with typed mujoco-warp mujocolab/mjlab#911

Merged

thowell reviewed Apr 15, 2026

View reviewed changes

Comment thread mujoco_warp/_src/io_test.py Outdated

thowell reviewed Apr 15, 2026

View reviewed changes

tkelestemur added 3 commits April 15, 2026 20:16

Address latest PR review comments

9caa705

Revert synthetic zero-nefc test change

8c530d5

Restore io test to match main

9b9feaf

thowell reviewed Apr 16, 2026

View reviewed changes

tkelestemur added 2 commits April 16, 2026 13:46

Fix tutorial segmentation display dtype

3898820

Enable segmentation rendering in tutorial

0346fc0

StafaH reviewed Apr 16, 2026

View reviewed changes

Comment thread mujoco_warp/_src/render_util.py Outdated

Comment thread mujoco_warp/_src/render_util.py Outdated

Comment thread mujoco_warp/_src/io_test.py Outdated

Comment thread mujoco_warp/_src/render_test.py Outdated

Comment thread mujoco_warp/_src/render_test.py Outdated

tkelestemur added 2 commits April 17, 2026 11:48

address StafaH's reviews

423cfe4

add flex segmentation test

f4bf916

StafaH approved these changes Apr 17, 2026

View reviewed changes

Comment thread mujoco_warp/_src/render_test.py Outdated

Use render buffer directly in segmentation test

f77e747

thowell approved these changes Apr 21, 2026

View reviewed changes

Merge branch 'google-deepmind:main' into tarik/semantic-segmentation-…

c1a64ae

…parity

Fix flex mesh group roots for segmentation

98d6411

StafaH reviewed Apr 22, 2026

View reviewed changes

tkelestemur and others added 2 commits April 25, 2026 11:08

remove leading _

0fdb3e6

Merge branch 'google-deepmind:main' into tarik/semantic-segmentation-…

b41b305

…parity

thowell merged commit 5f63341 into google-deepmind:main Apr 27, 2026
10 checks passed

tkelestemur deleted the tarik/semantic-segmentation-parity branch April 27, 2026 23:18

kevinzakka mentioned this pull request Apr 28, 2026

Upgrade mujoco to 3.8 and mujoco-warp to 3.8.0 mujocolab/mjlab#952

Merged

Conversation

tkelestemur commented Apr 12, 2026

Summary

Uh oh!

thowell commented Apr 13, 2026

Uh oh!

StafaH commented Apr 13, 2026

Uh oh!

tkelestemur commented Apr 14, 2026

Uh oh!

StafaH commented Apr 14, 2026

Uh oh!

tkelestemur commented Apr 14, 2026

Uh oh!

tkelestemur commented Apr 15, 2026

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

thowell commented Apr 15, 2026

Uh oh!

tkelestemur commented Apr 16, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tkelestemur commented Apr 16, 2026

Uh oh!

StafaH left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

StafaH commented Apr 17, 2026

Uh oh!

tkelestemur commented Apr 17, 2026

Uh oh!

StafaH left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

thowell commented Apr 22, 2026

Uh oh!

thowell commented Apr 22, 2026

Uh oh!

tkelestemur commented Apr 22, 2026

Uh oh!

tkelestemur commented Apr 22, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

thowell commented Apr 27, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants