Render by StafaH · Pull Request #836 · google-deepmind/mujoco_warp

StafaH · 2025-11-04T21:09:05Z

No description provided.

hartikainen · 2025-11-04T21:25:32Z

Hey Mustafa! This is super neat. What do I need to test this? What warp version should I use?

thowell · 2025-11-05T15:27:01Z

@StafaH awesome work!!

btaba · 2025-11-05T23:34:36Z

There is a circular import between render_context and bvh, should be resolved. Ownership between the two modules can also be improved

StafaH · 2025-11-06T05:07:30Z

Thanks @thowell, I went through all the comments and cleaned everything up.

@btaba I can hack it in place with an if TYPE_CHECKING, but the correct way here would probably be to move the create_render_context to either render or io, since that's where the circularity comes from. Open to either, @thowell should we keep rendering contained to one render.py or have the create function be called make_render_context and live alongside make_data etc?

@hartikainen there are warp changes that will make the renderer possible, currently the only way to test is to compile from source using the PR/fork found in the warp repo

thowell · 2025-11-06T10:14:56Z

@erikfrey what is your preference for where make_render_context is located?

thowell · 2025-11-06T10:25:34Z

@StafaH

we should add render_test.py with a few tests. we can also add todos for tests we think would be valuable to eventually implement.
we should add something to the readme, maybe under features: Renderer: ray tracing? or maybe a separate table for renderer under 'Compatibility' that lists current features? @erikfrey what are your thoughts/recommendations?
it would be great to add a benchmark with rendering to benchmark/benchmark.py and TODOs for additional benchmarks that we think would be valuable.

erikfrey · 2025-11-12T18:44:50Z

@thowell my first instinct would be that create_render_context live in io.py, but that's not a blocker to get this PR merged. LGTM for now.

hartikainen · 2025-11-12T20:40:36Z

For what it's worth, I played around with the renderer a bit. I'm noticing some weird behavior from it from a simple multi-object scene.

Here's the model I'm using:

Details

<mujoco model="cubes">

  <visual>
    <map znear="0.01" zfar="200" />
    <quality shadowsize="8192" />
    <headlight diffuse=".8 .8 .8" ambient=".2 .2 .2" specular="1 1 1" />
    <global azimuth="140" elevation="-20" />
  </visual>

  <asset>
    <texture type="skybox" builtin="gradient" rgb1="1 1 1" rgb2="1 1 1" width="800" height="800" />
    <texture name="grid" type="2d" builtin="checker" width="512" height="512" rgb1=".1 .2 .3" rgb2=".2 .3 .4"/>
    <material name="grid" texture="grid" texrepeat="1 1" texuniform="true" reflectance=".2"/>
  </asset>

  <worldbody>
    <geom name="floor" size="0 0 .05" type="plane" material="grid" contype="1" conaffinity="1"/>
    <camera pos="3 0 2" xyaxes="0 1 0 -0.6 0 1"/>
    <light pos="0 0 3" dir="0 0 -1" directional="true" />

    <geom name="table" pos="0.8 0.0 0.5" type="box" size="0.5 1.0 0.01" quat="1.0 0.0 0.0 0.0" contype="1" conaffinity="1"/>

    <replicate count="5" offset="0.1 0 0">
        <replicate count="5" offset="0 0.2 0">
            <body name="sphere" pos="0.6 -0.4 0.8">
                <freejoint name="sphere"/>
                <geom name="geom_sphere" type="sphere" size="0.05" contype="1" conaffinity="1" rgba="1 0 0 1"/>
            </body>
        </replicate>
    </replicate>

    <replicate count="5" offset="0.1 0 0">
        <replicate count="5" offset="0 0.2 0">
            <body name="capsule" pos="0.6 -0.4 0.95">
                <freejoint name="capsule"/>
                <geom name="geom_capsule" type="capsule" size="0.04 0.05" contype="1" conaffinity="1" rgba="0 1 0 1"/>
            </body>
        </replicate>
    </replicate>

    <replicate count="5" offset="0.1 0 0">
        <replicate count="5" offset="0 0.2 0">
            <body name="ellipsoid" pos="0.6 -0.4 1.1">
                <freejoint name="ellipsoid"/>
                <geom name="geom_ellipsoid" type="ellipsoid" size="0.05 0.04 0.03" contype="1" conaffinity="1" rgba="0 0 1 1"/>
            </body>
        </replicate>
    </replicate>

    <replicate count="5" offset="0.1 0 0">
        <replicate count="5" offset="0 0.2 0">
            <body name="cylinder" pos="0.6 -0.4 1.2">
                <freejoint name="cylinder"/>
                <geom name="geom_cylinder" type="cylinder" size="0.04 0.04" contype="1" conaffinity="1" rgba="1 1 0 1"/>
            </body>
        </replicate>
    </replicate>

    <replicate count="5" offset="0.1 0 0">
        <replicate count="5" offset="0 0.2 0">
            <body name="box1" pos="0.6 -0.4 1.4">
                <freejoint name="joint1"/>
                <geom name="geom1" type="box" size="0.03 0.03 0.03" contype="1" conaffinity="1" rgba="0 1 1 1"/>
            </body>
        </replicate>
    </replicate>
  </worldbody>
</mujoco>

Here's what the scene looks when rendered on the standard MJC viewer:

And here's with when renderer with the warp renderer. First, with width=160 and height=120:

Details

```console python -m mujoco_warp.render \ --nworld=4 \ --width=160 \ --height=120 \ --rgb=true \ --depth=true \ --textures=true \ --shadows=false \ --cam=0 \ --depth_scale=8.0 \ ./benchmark/100_cubes/scene.xml Loading model from: benchmark/100_cubes/scene.xml... Model: ncam=1 nlight=1 ngeom=127 Render: 640x480 rgb=True depth=True Rendering... Saved RGB to: debug.png Saved depth to: debug_depth.png ```

With width=640 and height=480:

Details

```console python -m mujoco_warp.render \ --nworld=4 \ --width=640 \ --height=480 \ --rgb=true \ --depth=true \ --textures=true \ --shadows=false \ --cam=0 \ --depth_scale=8.0 \ ./benchmark/100_cubes/scene.xml Loading model from: benchmark/100_cubes/scene.xml... Model: ncam=1 nlight=1 ngeom=127 Render: 640x480 rgb=True depth=True Rendering... Saved RGB to: debug.png Saved depth to: debug_depth.png ```

There are two, I believe, separate issues in the above images. First, it seems as if some black, noisy plane is being rendered right in front of the camera. Second, the bottom-right corner has a completely black strip to it. The reason I believe these are two separate issues is that I see the black strips on other models too, even when the depth rendering otherwise looks good. For example, if I remove all the objects except the table from the above scene, I get:

Here's some system information in case that matters:

$ uv pip list | grep "mujoco\|warp"
mujoco                  3.3.8.dev831015711
mujoco-mjx              3.3.8
mujoco-warp             0.0.1              /home/user/google-deepmind/mujoco_warp
warp-lang               1.11.0.dev20251111
$ /usr/bin/lspci | grep "NVIDIA"
00:03.0 3D controller: NVIDIA Corporation AD104GL [L4] (rev a1)
00:04.0 3D controller: NVIDIA Corporation AD104GL [L4] (rev a1)

hartikainen · 2025-11-12T20:52:54Z

Digging a bit deeper into my issue above, I think the weird noise is caused by ellipsoid geoms. If I modify the above scene by commenting out the ellipsoids, things seem fine:

Details

modified   benchmark/100_cubes/scene.xml
@@ -38,14 +38,14 @@
         </replicate>
     </replicate>
 
-    <replicate count="5" offset="0.1 0 0">
-        <replicate count="5" offset="0 0.2 0">
-            <body name="ellipsoid" pos="0.6 -0.4 1.1">
-                <freejoint name="ellipsoid"/>
-                <geom name="geom_ellipsoid" type="ellipsoid" size="0.05 0.04 0.03" contype="1" conaffinity="1" rgba="0 0 1 1"/>
-            </body>
-        </replicate>
-    </replicate>
+    <!-- <replicate count="5" offset="0.1 0 0"> -->
+    <!--     <replicate count="5" offset="0 0.2 0"> -->
+    <!--         <body name="ellipsoid" pos="0.6 -0.4 1.1"> -->
+    <!--             <freejoint name="ellipsoid"/> -->
+    <!--             <geom name="geom_ellipsoid" type="ellipsoid" size="0.05 0.04 0.03" contype="1" conaffinity="1" rgba="0 0 1 1"/> -->
+    <!--         </body> -->
+    <!--     </replicate> -->
+    <!-- </replicate> -->
 
     <replicate count="5" offset="0.1 0 0">
         <replicate count="5" offset="0 0.2 0">

The weird black strip still remains.

StafaH · 2025-11-13T06:00:00Z

Hi @hartikainen, Thanks for doing some early testing! Yes you are correct, the issues you are seeing are because ellipsoids haven't been implemented yet in the primitives, I'll be adding that soon. The black strip is because of the patch tiles we use, but this won't be an issue in the next batch of changes since we stop using those patch tiles altogether. For now the workaround would be to use a square image size until that fix is in.

* Integrate madrona_warp into mjwarp * Fix bug for <8 views * Move mesh construction to SAH, tile pixels * Change to registry to pass jax io test * Remove dist from bvh_query_ray * Add render context * Update testspeed * Add render context registry wrapper * Fix circular import * fix ray box and ambient lighting * Cleanup rc options in kernel * Small fix to static nlight * ray capsule fix * Fix testspeed merge artifact * Cleanup before PR into mjwarp * Remove global fovy, use cam 0 values * Cleanup bvh/render/ray * Fix capsule calc comment * Cleanup docstrings * Cleanup testspeed for render * Fix bug when no texid is set

* Render (#836) * Integrate madrona_warp into mjwarp * Fix bug for <8 views * Move mesh construction to SAH, tile pixels * Change to registry to pass jax io test * Remove dist from bvh_query_ray * Add render context * Update testspeed * Add render context registry wrapper * Fix circular import * fix ray box and ambient lighting * Cleanup rc options in kernel * Small fix to static nlight * ray capsule fix * Fix testspeed merge artifact * Cleanup before PR into mjwarp * Remove global fovy, use cam 0 values * Cleanup bvh/render/ray * Fix capsule calc comment * Cleanup docstrings * Cleanup testspeed for render * Fix bug when no texid is set * Render hetero (#873) * Integrate madrona_warp into mjwarp * Fix bug for <8 views * Move mesh construction to SAH, tile pixels * Change to registry to pass jax io test * Remove dist from bvh_query_ray * Add render context * Update testspeed * Add render context registry wrapper * Fix circular import * fix ray box and ambient lighting * Cleanup rc options in kernel * Small fix to static nlight * ray capsule fix * Fix testspeed merge artifact * Cleanup before PR into mjwarp * Remove global fovy, use cam 0 values * Cleanup bvh/render/ray * Fix capsule calc comment * Cleanup docstrings * Cleanup testspeed for render * Fix bug when no texid is set * hetero camera support * Add heightfield support * Add flex mesh building * Add correct flex 2d mesh building * Fixes for downstream use * Rename output buffers to match mujoco style * Add flex support, improve render script * Add selective camera rendering support via cam_active parameter * Fix render output reading, fix bg color * Add more optimzed hfield mesh building * Add visual franka benchmark, fix merge artifacts * Add 3d flex rendering support * Add ellipsoid and cylinder primitives * Cleanup flex 2d and 3d rendering * Cleanup of bvh code * Fix 2d flex side triangle update * More cleanup for PR * More cleanup of naming, paranthesis, casts * Add correct frustrum calculation for rays * Ruff check * Fix naming, add todo * Update kernel arguments using analyzer * Remove comment * Remove visual assets * Add back camera * Remove stray 1, dim --------- Co-authored-by: Kevin Zakka <kevinarmandzakka@gmail.com> * Render merge ray normal (#979) * Add ray with normal (#940) * Add ray with normal * Ruff * Propagate ray normal changes * ruff * Fix tuple type check * Update ray to return -1 for non-hits * Ray normal fixes (#960) * Add ray with normal * Ruff * Propagate ray normal changes * ruff * Fix tuple type check * Update ray to return -1 for non-hits * Small fixes * Ruff format * Simplify assert for vec3 * Add comment and adjust docstring * Fix bugs in ray hfield impl (#967) * Merge main into render branch (#980) * Add ray with normal (#940) * Add ray with normal * Ruff * Propagate ray normal changes * ruff * Fix tuple type check * Update ray to return -1 for non-hits * Ray normal fixes (#960) * Add ray with normal * Ruff * Propagate ray normal changes * ruff * Fix tuple type check * Update ray to return -1 for non-hits * Small fixes * Ruff format * Simplify assert for vec3 * Add comment and adjust docstring * Fix bugs in ray hfield impl (#967) * clamp pair_friction with minmu (#969) * add --kernel_cache_dir pytest flag for warp kernel cache (#968) * fix _NCONMAX info (#965) * reuse xnorm (#959) * Expose rays through the public API. (#975) --------- Co-authored-by: Taylor Howell <taylorhowell@google.com> Co-authored-by: Kevin Zakka <kevinzakka@users.noreply.github.com> * Fix render aspect ratio bug (#1002) * Update to new cam proj field * Fix frustrum aspect ratio * Render tests and cleanup (#1021) * ruff * Move render script to contrib * Add render tests * Move create render context into io.py * Cleanup bvh test, ruff * cleanup render_context_test * More cleanup * Render Cleanup 2 (#1037) * Update context to read from new mujoco fields * Add anyhit ray mesh query for fast shadows * Cleanup tests, msgs, comments * Render Cleanup 3 (#1063) * Fix naming, move mesh bvh build to bvh module * Add batched cam intrinsic for randomization * Fix indexing in render context * ruff * Render Cleanup 4 (#1064) * Fix failing tests with new batched cam fields * Move hfield bvh build to bvh module * Fix docstring * Render Cleanup 5 (#1075) * Fix flex rendering for multiple flex * Add flex bvh test * Render Cleanup - Textures (#1080) * Add cuda texture sampling * Fix textures and mesh sample, cleanup * ruff * Clean up texture processing kernel * Fix flipped texture uv * Remove nested texture kernel * Render Cleanup - Ray (#1081) * Add bvh accelerated ray * Move ray_bvh into ray * Remove ray_bvh_test and merge with ray_test * Use absolute import * Cleanup public API for render/bvh * Render Cleanup - types and io (#1103) * Move RenderContext to types * Cleanup flex fields and refit kernel * Use array in render context type * Fix all render related tests after refactor * Move packing to render util * Cleanup render file * Cleanup based on pr comments * Fix error with rc init * Fix flex render error (#1108) * Update _render_megakernel decorator (#1109) * Fix error for downstream mjlab * Fix ray none check * Change nested kernel to wp.kernel * Remove TODO * Revert uv.lock * Post merge cleanup and fixes * Add guards for bleeding edge features * Fix guards for mjtProjection * Fix failing test, post merge artifact * Fix failing tests for ray/render * Add more guards for mujoco < 3.4.1 * Ruff * Disable flex cpu test, fix leading dim * Add newline --------- Co-authored-by: Kevin Zakka <kevinarmandzakka@gmail.com> Co-authored-by: Taylor Howell <taylorhowell@google.com> Co-authored-by: Kevin Zakka <kevinzakka@users.noreply.github.com>

StafaH added 15 commits October 31, 2025 15:05

Integrate madrona_warp into mjwarp

3d43f0f

Fix bug for <8 views

7ad9758

Move mesh construction to SAH, tile pixels

38264da

Change to registry to pass jax io test

f726110

Remove dist from bvh_query_ray

17edff2

Add render context

4f1f787

Update testspeed

e3389ba

Add render context registry wrapper

326773e

Fix circular import

583eb42

fix ray box and ambient lighting

8e8e7a9

Cleanup rc options in kernel

ff27710

Small fix to static nlight

8f17e89

ray capsule fix

954e965

Fix testspeed merge artifact

3fcd577

Cleanup before PR into mjwarp

d09c3a4

Remove global fovy, use cam 0 values

1109851