Render tests and cleanup by StafaH · Pull Request #1021 · google-deepmind/mujoco_warp

StafaH · 2026-01-14T08:08:19Z

This PR includes the following changes:

Tests for bvh, render_context, and render
Ruff format cleanup
Moving create_render_context into io.py to avoid circular import

Testspeed now also prints out render context information (only when render is called):

Loading model from: benchmark/franka_emika_panda/scene.xml...
  nbody: 12 nv: 9 ngeom: 23 nu: 8 is_sparse: False
  broadphase: NXN broadphase_filter: PLANE|SPHERE|OBB
  solver: NEWTON cone: PYRAMIDAL iterations: 100 iterative linesearch iterations: 50
  integrator: IMPLICITFAST graph_conditional: True
Data
  nworld: 8192 naconmax: 393216 njmax: 64
Render Context
  shadows: False textures: True nlight: 2 bvh_ngeom: 15 ncam: 1 cam_res: [[64 64]]

Rolling out 100 steps at dt = 0.005...

Summary for 8192 parallel rollouts

Total JIT time: 0.01 s
Total simulation time: 6.31 s
Total steps per second: 129,879
Total realtime factor: 649.39 x
Total time per step: 7699.49 ns
Total converged worlds: 8192 / 8192

thowell · 2026-01-14T11:55:09Z

+      render_depth=True,
+    )
+    render.render(m, d, rc)
+    rgb = rc.rgb_data.numpy().copy()


is copy necessary?

thowell · 2026-01-14T12:09:05Z

thanks for this great contribution @StafaH!

thowell · 2026-01-14T20:38:29Z

lgtm. thanks!

* 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 4 commits January 14, 2026 01:36

ruff

429678a

Move render script to contrib

aec5538

Add render tests

2e38e38

Move create render context into io.py

ac7a26a

StafaH requested review from btaba, erikfrey and thowell January 14, 2026 08:08