[StafaH]

Adds an OptiX hardware ray tracing backend to the batch renderer, built on the warp-optix 0.3.0 release.

Usage

• create_render_context(..., backend=...) accepts:
  • "warp" (default): the existing Warp BVH renderer.
  • "optix": OptiX. If it is unavailable, warns with the reason and uses "warp".
• OptiX is available when warp-optix is installed (pip install mujoco-warp[optix]), the GPU can create an OptiX context, and the model has no flexes or splats.
• The chosen backend is stored on rc.backend.
• warp_optix is imported only when backend="optix" is requested, so importing mujoco_warp does not load it.

How it works

1. Meshes and heightfields are built once as triangle acceleration structures shared by all worlds.
2. Every enabled geom of every world is an instance in one instance acceleration structure, with worlds laid out in a cubic grid. Analytic geoms instance a padded unit cube whose intersection program reuses the intersections in ray.py. Planes are unbounded, so they are intersected directly in the ray generation program.
3. refit_bvh rewrites the instance transforms and refits the instance acceleration structure each step.
4. The render megakernel is compiled as the OptiX ray generation program, so shading, textures, shadows and outputs are shared with the Warp backend and only the traversal in cast_ray differs. Capping it at 72 registers was the fastest setting.
5. A few megakernel settings differ by backend, tuned on the benchmarks below:
   • Rays are traced in row-major order. OptiX orders its launch indices for coherence itself, so the 8-row tiling from render improvements #1959 only scattered the pixel writes (9% slower).
   • The skybox texture is loaded on a miss rather than before the trace (3 to 4% faster, 7.5% at 2x2 supersampling). The Warp backend keeps the early load, which is up to 43% faster there.

Benchmarks

Refit + render throughput (refit_bvh followed by render, in world-steps per second) from mjwarp-testspeed --function=render, median of 3 runs.

• RTX PRO 6000 Blackwell, Warp 1.18.0, warp-optix 0.3.0, 64x64 cameras, run-to-run spread under 0.5%.
• The physics step between frames is not timed.
• "main" is the Warp BVH renderer on main at eb81a03, which includes the render improvements from render improvements #1959.
• The last three scenes are menagerie models with an added camera (not committed).

Scene	Worlds	main	OptiX	Speedup
primitives (RGB + depth)	8192	2.62M	3.25M	1.2x
unitree_g1_hfield_render (RGB)	8192	197k	4.10M	21x
aloha_clutter_vision (4 cameras, RGB + depth + seg, shadows)	2048	13.8k	284k	20x
Franka Panda (RGB + depth)	8192	434k	3.66M	8.4x
Boston Dynamics Spot (RGB + depth)	8192	312k	3.95M	13x
Shadow Hand (RGB + depth)	8192	551k	3.30M	6.0x

• The Warp backend shares the restructured cast_ray. Against main it is -1 to +2% on primitives and Shadow Hand and 41 to 51% faster on the other four scenes. On Spot the two match with render_skybox=False, so the gap comes from how main's cast_ray compiles together with the early skybox load from render improvements #1959.
• cloth_render uses flexes, so it runs on the Warp backend either way.

Limitations

• OptiX launches cannot be recorded into CUDA graphs. cli.unroll launches the OptiX backend directly, and code that captures render in a graph should keep the default backend="warp".
• rays() needs a Warp render context and raises for an OptiX one.
• Packing worlds into a grid moves geometry up to about 1.4 km from the origin (8192 G1 worlds), which costs float32 precision. Against the Warp backend, 1 to 3% of RGB pixels and about 0.01% of segmentation pixels differ, mostly at silhouettes. Single-world renders match to within one RGB level.
• Coincident geoms are not tie-broken toward the higher geom index as they are with the Warp BVH.
• Flexes and splats are not supported yet and use the Warp backend.
• warp-optix only publishes Python 3.12 wheels, so the optix extra is limited to Python 3.12.

[thowell]

a few suggestions:

1. Keep OptiX strictly lazy-imported so it has zero import-time impact when off by default (render_optix.py, bvh.py, render.py, render_util.py)

Issue: bvh.py, render.py, and render_util.py all import render_optix at module top level, and render_optix.py executes import warp_optix as wo at import time. Consequently, every import mujoco_warp unconditionally attempts to import warp_optix (and initialize a Warp module for render_optix) even when the default backend="bvh" is used. If warp_optix is present in an environment with an incompatible warp-lang version or missing libnvoptix.so.1 driver, raising AttributeError or OSError/RuntimeError on import will break import mujoco_warp for users not using OptiX.

Recommendation: Avoid importing warp_optix at module import time. Load warp_optix lazily inside a helper (called only when backend="optix" is requested in create_render_context), and only import render_optix when backend == "optix" or rc.optix is not None.

# mujoco_warp/_src/render_optix.py
import functools


@functools.cache
def _get_wo():
  try:
    import warp_optix as wo  # noqa: PLC0415

    return wo
  except ImportError as e:
    raise ImportError(
      "backend='optix' requires warp_optix (https://github.com/nvtw/otk-pyoptix) and warp-lang>=1.17"
    ) from e

---

2. Reuse make_ray_primitive for ray_geom in ray.py

Issue: ray.py adds make_ray_primitive, which duplicates the 6-way analytic primitive dispatch (PLANE, SPHERE, CAPSULE, ELLIPSOID, CYLINDER, BOX) already in ray_geom (ray.py:806-827, which has # TODO(team): static loop unrolling to remove unnecessary branching).

Recommendation: Define ray_geom using make_ray_primitive over all analytic primitive GeomTypes instead of maintaining two separate primitive dispatch functions in ray.py.

# mujoco_warp/_src/ray.py
_PRIMITIVE_TYPES = (
  int(GeomType.PLANE),
  int(GeomType.SPHERE),
  int(GeomType.CAPSULE),
  int(GeomType.ELLIPSOID),
  int(GeomType.CYLINDER),
  int(GeomType.BOX),
)


def make_ray_primitive(geom_ray_types: Tuple[int, ...] = _PRIMITIVE_TYPES) -> wp.Function:
  ...


ray_geom = make_ray_primitive()

---

3. Use mjm.geom_aabb (or m.geom_aabb) instead of recomputing per-type primitive half-extents (render_optix.py)

Issue: render_optix.py:create_scene manually computes local-frame AABB half-extents per GeomType (SPHERE, CAPSULE, CYLINDER, ELLIPSOID, BOX), duplicating the local AABB half-sizes already stored in mjm.geom_aabb[gi, 3:] and m.geom_aabb[..., gi, 1] (computed in set_const._set_geom_size). In addition, building static per-geom GASes from host mjm does not reflect per-world m.geom_size / m.geom_aabb domain randomization.

Recommendation: Use mjm.geom_aabb[gi, 3:] directly on host, or build a single unit-cube custom-primitive GAS ([-1.001, 1.001]^3) and scale the instance transform in write_instances by m.geom_aabb[worldid % geom_aabb.shape[0], gi, 1].

# Using mjm.geom_aabb[gi, 3:] on host:
  for gi in rc.enabled_geom_ids.numpy():
    gtype = mjm.geom_type[gi]
    if gtype in _MESH_TYPES or gtype == GeomType.PLANE:
      continue
    half = mjm.geom_aabb[gi, 3:]
    bounds = wp.array(np.concatenate([-half, half])[None] * 1.001, dtype=float)
    gas.append(wo.build_accel(ctx, [wo.aabbs(bounds)], compact=True))
    primitive_handles[gi] = gas[-1].handle

---

4. Share the write_instances launch and IAS update between create_scene and refit (render_optix.py)

Issue: create_scene and refit duplicate the wp.launch(write_instances, ...) call (with the same 11 inputs and instances output), wo.update_accel(ctx, ias), and rc.bvh_id = ias.handle.

Recommendation: Extract a shared helper (or call refit once rc.optix is populated) so write_instances is only launched from one place.

def _update_instances(
  scene: OptixScene,
  rc,
  geom_type: wp.array,
  geom_dataid: wp.array2d,
  geom_matid: wp.array2d,
  geom_rgba: wp.array2d,
  mat_rgba: wp.array2d,
  geom_xpos: wp.array2d,
  geom_xmat: wp.array2d,
):
  wp.launch(
    scene.write_instances,
    dim=scene.instances.shape,
    inputs=[
      geom_type,
      geom_dataid,
      geom_matid,
      geom_rgba,
      mat_rgba,
      geom_xpos,
      geom_xmat,
      rc.enabled_geom_ids,
      scene.mesh_traversable,
      scene.hfield_traversable,
      scene.primitive_traversable,
    ],
    outputs=[scene.instances],
  )
  if scene.ias is None:
    scene.ias = wo.build_accel(
      scene.ctx, [wo.instances(scene.instances)], flags=wo.constants.OPTIX_BUILD_FLAG_ALLOW_UPDATE
    )
  scene.ias = wo.update_accel(scene.ctx, scene.ias)
  rc.bvh_id = scene.ias.handle

---

5. Avoid unnecessary global memory loads on the default BVH path for meshes, heightfields, and spheres (ray.py, render.py, render_optix.py)

Issue: In render.py:cast_ray (the default backend="bvh" path), ray_primitive(gtype, geom_xpos_in[worldid, gi], geom_xmat_in[worldid, gi], geom_size[worldid % geom_size.shape[0], gi], ...) is called unconditionally for all local_id < bvh_ngeom. This eagerly loads geom_xpos_in, geom_xmat_in, and geom_size from global memory even when gtype is MESH, SDF, or HFIELD (where geom_size is unused and geom_xpos_in/geom_xmat_in are loaded again in ray_mesh_with_bvh), and loads geom_xmat_in (9 floats) for SPHERE (which only needs geom_xpos_in and geom_size[...][0] in both the BVH and OptiX backends).

Recommendation: Pass geom_size, geom_xpos_in, geom_xmat_in, worldid, and gi into ray_primitive (with # Model:, # Data in:, and # In: headers) so each branch only reads the global memory arrays it needs.

# mujoco_warp/_src/ray.py
def make_ray_primitive(geom_ray_types: Tuple[int, ...] = _PRIMITIVE_TYPES) -> wp.Function:
  """Build a ray-primitive intersection specialized to the geom types present in the scene."""

  @wp.func
  def ray_primitive(
    # Model:
    geom_size: wp.array2d[wp.vec3],
    # Data in:
    geom_xpos_in: wp.array2d[wp.vec3],
    geom_xmat_in: wp.array2d[wp.mat33],
    # In:
    worldid: int,
    gi: int,
    gtype: int,
    pnt: wp.vec3,
    vec: wp.vec3,
  ) -> Tuple[float, wp.vec3]:
    d = float(-1.0)
    n = wp.vec3(0.0, 0.0, 0.0)
    if wp.static(int(GeomType.PLANE) in geom_ray_types):
      if gtype == GeomType.PLANE:
        d, n = ray_plane(
          geom_xpos_in[worldid, gi], geom_xmat_in[worldid, gi], geom_size[worldid % geom_size.shape[0], gi], pnt, vec
        )
    if wp.static(int(GeomType.SPHERE) in geom_ray_types):
      if gtype == GeomType.SPHERE:
        r = geom_size[worldid % geom_size.shape[0], gi][0]
        d, n = ray_sphere(geom_xpos_in[worldid, gi], r * r, pnt, vec)
    if wp.static(int(GeomType.ELLIPSOID) in geom_ray_types):
      if gtype == GeomType.ELLIPSOID:
        d, n = ray_ellipsoid(
          geom_xpos_in[worldid, gi], geom_xmat_in[worldid, gi], geom_size[worldid % geom_size.shape[0], gi], pnt, vec
        )
    if wp.static(int(GeomType.CAPSULE) in geom_ray_types):
      if gtype == GeomType.CAPSULE:
        d, n = ray_capsule(
          geom_xpos_in[worldid, gi], geom_xmat_in[worldid, gi], geom_size[worldid % geom_size.shape[0], gi], pnt, vec
        )
    if wp.static(int(GeomType.CYLINDER) in geom_ray_types):
      if gtype == GeomType.CYLINDER:
        d, n = ray_cylinder(
          geom_xpos_in[worldid, gi], geom_xmat_in[worldid, gi], geom_size[worldid % geom_size.shape[0], gi], pnt, vec
        )
    if wp.static(int(GeomType.BOX) in geom_ray_types):
      if gtype == GeomType.BOX:
        d, _, n = ray_box(
          geom_xpos_in[worldid, gi], geom_xmat_in[worldid, gi], geom_size[worldid % geom_size.shape[0], gi], pnt, vec
        )
    return d, n

  return ray_primitive

---

6. Skip surface evaluation and plane tests on shadow-ray hits (first_hit=True) and remove duplicate plane backface check (render_optix.py)

Issue: In render_optix.py:make_cast_ray:

• When first_hit=True (shadow rays), compute_lighting only checks shadow_geom_id != -1 and discards dist, normal, u, v, face, and mesh_id. However, when hit.t != wo.MISS, cast_ray still loads geom_type, geom_dataid, hfield_bvh_id/mesh_bvh_id, and geom_xmat_in from global memory, evaluates wp.mesh_eval_face_normal, and loops over all planes.
• geom_ray_types is passed to make_cast_ray but never used to eliminate the triangle vs. primitive hit branches via wp.static.
• In the plane loop, geom_xpos_in, geom_xmat_in, and geom_size are loaded and ray_plane is called before checking if rgba[3] == 0.0: continue. Also, culled = cull_backfaces and wp.dot(ray_dir_world, n) > 0.0 duplicates the front-face check already performed inside ray_plane (if lvec[2] > -MJ_MINVAL: return -1.0, wp.vec3()).

Recommendation: Return immediately on hit.t != wo.MISS when wp.static(first_hit) is True, guard the triangle hit branch with wp.static(any(t in geom_ray_types for t in _MESH_TYPES)), check rgba[3] == 0.0 before ray_plane, remove the duplicate culled check, and return early on the first plane hit when wp.static(first_hit) is True.

    if hit.t != wo.MISS:
      if wp.static(first_hit):
        return 0, hit.t, normal, u, v, face, mesh_id
      geom_id = int(hit.instance_id) - worldid * geom_type.shape[0]
      dist = hit.t
      gtype = geom_type[geom_id]
      if wp.static(any(t in geom_ray_types for t in _MESH_TYPES)):
        if gtype == GeomType.MESH or gtype == GeomType.SDF or gtype == GeomType.HFIELD:
          face = int(hit.primitive_index)
          dataid = geom_dataid[worldid % geom_dataid.shape[0], geom_id]
          mesh_id = dataid
          if wp.static(int(GeomType.HFIELD) in geom_ray_types):
            if gtype == GeomType.HFIELD:
              mesh = hfield_bvh_id[dataid]
              mesh_id = -1
            else:
              mesh = mesh_bvh_id[dataid]
          else:
            mesh = mesh_bvh_id[dataid]
          normal = wp.normalize(geom_xmat_in[worldid, geom_id] @ wp.mesh_eval_face_normal(mesh, face))
          bary = wp.optix_hit_object_get_triangle_barycentrics()
          u = 1.0 - bary[0] - bary[1]
          v = bary[0]
        else:
          normal = wp.vec3(
            wp.uint32_to_float(wp.optix_hit_object_get_attribute_0()),
            wp.uint32_to_float(wp.optix_hit_object_get_attribute_1()),
            wp.uint32_to_float(wp.optix_hit_object_get_attribute_2()),
          )
      else:
        normal = wp.vec3(
          wp.uint32_to_float(wp.optix_hit_object_get_attribute_0()),
          wp.uint32_to_float(wp.optix_hit_object_get_attribute_1()),
          wp.uint32_to_float(wp.optix_hit_object_get_attribute_2()),
        )

    # Planes are unbounded, so they cannot live in a world cell; intersect them directly.
    for i in range(wp.static(len(planes))):
      gi = wp.static(planes[i])
      rgba, _ = geom_color(geom_matid, geom_rgba, mat_rgba, worldid, gi)
      if rgba[3] == 0.0:
        continue
      d, n = ray_plane(
        geom_xpos_in[worldid, gi],
        geom_xmat_in[worldid, gi],
        geom_size[worldid % geom_size.shape[0], gi],
        ray_origin_world,
        ray_dir_world,
      )
      if d >= 0.0 and d < dist:
        if wp.static(first_hit):
          return gi, d, n, 0.0, 0.0, -1, -1
        geom_id, dist, normal, u, v, face, mesh_id = gi, d, n, 0.0, 0.0, -1, -1

---

7. Clamp tmax >= 0.0 in make_cast_ray and check d < wp.optix_get_ray_tmin() before backface dot product in intersect_primitive (render_optix.py)

Issue:

• In make_cast_ray, if a ray origin lies outside [-extent, extent] along an axis and points away from that bound, (bound - ray_origin_world[axis]) / ray_dir_world[axis] is negative, passing tmax < tmin (0.0) to wo.closest_hit (optixTraverse).
• In intersect_primitive, wp.optix_get_instance_id() is called twice, and wp.dot(vec, n) is computed before checking d >= wp.optix_get_ray_tmin() even on missed rays (d == -1.0).

Recommendation: Clamp tmax = wp.max(0.0, tmax) in make_cast_ray, and in intersect_primitive cache inst_id and check d < wp.optix_get_ray_tmin() before computing wp.dot(vec, n).

# In make_cast_ray:
    tmax = max_dist
    for axis in range(3):
      if ray_dir_world[axis] != 0.0:
        bound = wp.where(ray_dir_world[axis] > 0.0, wp.static(extent), -wp.static(extent))
        tmax = wp.min(tmax, (bound - ray_origin_world[axis]) / ray_dir_world[axis])
    tmax = wp.max(0.0, tmax)

# In intersect_primitive:
    ngeom = geom_type.shape[0]
    inst_id = int(wp.optix_get_instance_id())
    worldid = inst_id // ngeom
    gi = inst_id - worldid * ngeom
    pnt = wp.optix_get_world_ray_origin() - _world_offset(worldid, wp.static(grid), wp.static(spacing))
    vec = wp.optix_get_world_ray_direction()
    d, n = ray_primitive(geom_size, geom_xpos_in, geom_xmat_in, worldid, gi, geom_type[gi], pnt, vec)
    if d < wp.optix_get_ray_tmin():
      return
    if wp.static(cull_backfaces):
      if wp.dot(vec, n) > 0.0:
        return
    wp.optix_report_intersection(
      d, wp.uint32(0), wp.float_to_uint32(n[0]), wp.float_to_uint32(n[1]), wp.float_to_uint32(n[2])
    )

---

8. Avoid Any type annotations (types.py, render_optix.py) and fix [optix] extra reference in docstring (render_util.py)

Issue: types.py and render_optix.py add Any type annotations (optix: Optional[Any] = None), and render_util.py documents backend="optix" (requires mujoco-warp[optix]) although pyproject.toml does not define an optix optional-dependency extra.

Recommendation: Define OptixScene in types.py (or type optix: Optional[OptixScene] = None without Any) and update the backend docstring in create_render_context to reference warp_optix.

# mujoco_warp/_src/types.py
@dataclasses.dataclass
class OptixScene:
  ctx: object
  gas: list
  mesh_traversable: wp.array
  hfield_traversable: wp.array
  primitive_traversable: wp.array
  instances: wp.array2d
  ias: object
  extent: float
  planes: tuple
  grid: int
  write_instances: wp.Kernel
  pipeline: Optional[object] = None

---

9. Use an inline XML fixture, pre-fill output buffers, and cover hfield in render_optix_test.py

Issue: render_optix_test.py loads external XML files ("primitives.xml" with 126 geoms and "mug/mug.xml"), steps the simulation twice per test case, does not pre-fill rc.rgb_data / rc.depth_data before mjw.render, and does not test GeomType.HFIELD (hfield_traversable).

Recommendation: Use an inline test_data.fixture(xml="""...""") scene that includes primitives, a mesh, and an hfield, step (m, d) once and render with both backends, parameterize over nworld: 1, 2 with distinct worlds when nworld == 2, and pre-fill output buffers before calling mjw.render.

@absltest.skipIf(render_optix.wo is None, "requires warp_optix")
class RenderOptixTest(parameterized.TestCase):
  @parameterized.parameters(1, 2)
  def test_matches_bvh(self, nworld):
    mjm, _, m, d = test_data.fixture(
      xml="""
      <mujoco>
        <asset>
          <hfield name="hf" nrow="4" ncol="4" size="1 1 0.2 0.1"
                  elevation="0 0.2 0.4 0 0.1 0.5 0.3 0.1 0.2 0.4 0.6 0.2 0 0.1 0.2 0"/>
          <mesh name="tet" vertex="0 0 0 0.3 0 0 0 0.3 0 0 0 0.3" face="0 2 1 0 1 3 1 2 3 0 3 2"/>
        </asset>
        <worldbody>
          <light pos="0 -1 3" dir="0 0.3 -1" castshadow="true"/>
          <camera pos="0 -2.2 1.2" xyaxes="1 0 0 0 0.5 1"/>
          <geom type="plane" size="3 3 0.1" rgba="0.8 0.8 0.8 1"/>
          <geom type="hfield" hfield="hf" pos="0 0.5 0" rgba="0.4 0.7 0.4 1"/>
          <body pos="-0.4 0 0.3">
            <freejoint/>
            <geom size="0.15" rgba="1 0 0 1"/>
            <geom type="capsule" pos="0.3 0 0" size="0.08 0.1" rgba="0 1 0 1"/>
            <geom type="box" pos="0.6 0 0" size="0.1 0.1 0.1" rgba="0 0 1 1"/>
            <geom type="cylinder" pos="0 -0.3 0" size="0.08 0.1" rgba="1 1 0 1"/>
            <geom type="ellipsoid" pos="0.3 -0.3 0" size="0.12 0.08 0.1" rgba="0 1 1 1"/>
            <geom type="mesh" mesh="tet" pos="0.6 -0.3 0" rgba="1 0 1 1"/>
          </body>
        </worldbody>
      </mujoco>
      """,
      nworld=nworld,
    )
    if nworld == 2:
      qpos = d.qpos.numpy()
      qpos[1, 0] += 0.15
      d.qpos.assign(qpos)
      mjw.kinematics(m, d)

    outs = {}
    for backend in ("bvh", "optix"):
      rc = mjw.create_render_context(
        mjm,
        nworld=nworld,
        cam_res=(48, 48),
        render_rgb=True,
        render_depth=True,
        render_seg=True,
        use_shadows=True,
        backend=backend,
      )
      rc.rgb_data.fill_(wp.uint32(0xFFFFFFFF))
      rc.depth_data.fill_(wp.inf)
      mjw.refit_bvh(m, d, rc)
      mjw.render(m, d, rc)
      outs[backend] = (rc.rgb_data.numpy(), rc.depth_data.numpy(), rc.seg_data.numpy())

    rgb, depth, seg = outs["bvh"]
    rgb_optix, depth_optix, seg_optix = outs["optix"]
    self.assertGreater(np.unique(seg).size, 4)
    np.testing.assert_array_less(np.abs(rgb.view(np.uint8).astype(int) - rgb_optix.view(np.uint8)), 2)
    np.testing.assert_allclose(depth, depth_optix, atol=1e-4)
    np.testing.assert_array_equal(seg, seg_optix)
    if nworld == 2:
      self.assertFalse(np.allclose(depth[0], depth[1]))

[StafaH]

Thanks for the review. This moves the backend to the released NVIDIA/warp-optix and its new API, as an optix extra.

1. Lazy import: done.
2. ray_geom = make_ray_primitive(): done.
3. Bounds: one unit-cube GAS, scaled per instance by geom_aabb.
4. create_scene and refit share _update_instances.
5. Kept the hoisted loads: per-branch loads dropped BVH primitives from 2.57M to 2.46M steps/s (8192 worlds).
6. Done: any-hit shadow rays with early return, wp.static guards, plane alpha check first, duplicate cull removed.
7. Done.
8. Done: OptixScene is in types.py, Any removed, docstring fixed.
9. Done: inline scene with an hfield, nworld 1 and 2, pre-filled outputs. It caught BVH testing hfield shadows one-sided (meshes are two-sided), now fixed.

[thowell]

thanks for updating the pr @StafaH!

3 recommendations and then i think we are good to merge

1. Keep warp-optix out of pyproject.toml and uv.lock until Linux wheels are on PyPI

Adding optix with [tool.uv.sources] warp-optix = {git = "https://github.com/NVIDIA/warp-optix"} causes uv.lock to upgrade mujoco-warp's base warp-lang dependency on Python 3.12 from 1.15.0 to 1.19.0 (requiring CUDA 13 drivers) even when --extra optix is not requested (uv sync). It also causes the documented developer setup (uv sync --all-extras) on Python 3.12 to clone warp-optix from Git and compile pyoptix and warp-optix C++/CMake extensions from source, while pip install mujoco-warp[optix] on Linux ignores [tool.uv.sources] and fails on PyPI because warp-optix publishes no Linux wheels yet.

Recommendation: Revert the pyproject.toml and uv.lock changes while the backend is experimental and warp-optix lacks Linux wheels on PyPI, and point users to the explicit warp-optix install command in _get_wo().

@functools.cache
def _get_wo():
  try:
    import warp_optix as wo  # noqa: PLC0415
  except ImportError as e:
    raise ImportError(
      "backend='optix' requires warp_optix (pip install git+https://github.com/NVIDIA/warp-optix)"
    ) from e
  return wo

2. Mark GeomType.PLANE instances inactive in the IAS and scale analytic primitives from geom_size

In write_instances, GeomType.PLANE falls into the else branch, keeping traversable = primitive_traversable whenever the scene has any analytic primitive and scaling the unit-cube GAS by geom_aabb[..., gi, 1] = (1e10, 1e10, 5e9). Setting visibilityMask = 0 only masks ray traversal—optixAccelBuild still inserts every instance with a non-zero traversable into the IAS BVH, creating nworld 2e10-wide bounds overlapping every world cell. In addition, scaling analytic primitives from geom_aabb instead of geom_size requires callers randomizing m.geom_size (as in render_optix_test.py) to manually update m.geom_aabb as well, whereas intersect_primitive and the BVH backend read m.geom_size directly.

Recommendation: Set traversable = wp.uint64(0) for GeomType.PLANE so optixAccelBuild ignores plane instances, and compute primitive half-extents directly from geom_size (removing geom_aabb from write_instances / refit and the manual m.geom_aabb update in render_optix_test.py).

    if gtype == GeomType.PLANE:
      traversable = wp.uint64(0)
    elif gtype == GeomType.HFIELD:
      traversable = hfield_traversable[geom_dataid[worldid % geom_dataid.shape[0], gi]]
      sbt_offset = _TRIANGLE_SBT_OFFSET
    elif gtype == GeomType.MESH or gtype == GeomType.SDF:
      traversable = mesh_traversable[geom_dataid[worldid % geom_dataid.shape[0], gi]]
      sbt_offset = _TRIANGLE_SBT_OFFSET
    else:
      s = geom_size[worldid % geom_size.shape[0], gi]
      if gtype == GeomType.SPHERE:
        half = wp.vec3(s[0], s[0], s[0])
      elif gtype == GeomType.CAPSULE:
        half = wp.vec3(s[0], s[0], s[0] + s[1])
      elif gtype == GeomType.CYLINDER:
        half = wp.vec3(s[0], s[0], s[1])
      else:
        half = s
      linear = linear @ wp.diag(half)

3. Guard zero-geom / zero-ray scenes and --cpu test runs

When rc.bvh_ngeom == 0, _update_instances passes an empty (nworld, 0) buffer to wo.Accel, and when rc.total_rays == 0, wo.launch raises ValueError("dim must contain one to three positive integers") (unlike wp.launch, which treats a zero dimension as a no-op). Also, warp_optix imports cleanly without a GPU, so RenderOptixTest fails instead of skipping under pytest --cpu or on CPU-only runners when warp_optix is installed.

Recommendation: Early-return on rc.bvh_ngeom == 0 in _update_instances and dim[0] == 0 or dim[1] == 0 in render_optix.launch, and include not wp.get_device().is_cuda in RenderOptixTest's skip condition.

# In render_optix.py:
def _update_instances(...):
  if rc.bvh_ngeom == 0:
    return
  ...

def launch(rc, dim: Tuple[int, int], inputs: list):
  if dim[0] == 0 or dim[1] == 0:
    return
  _get_wo().launch(rc.optix.pipeline, dim, inputs)

# In render_optix_test.py:
@absltest.skipIf(
  importlib.util.find_spec("warp_optix") is None or not wp.get_device().is_cuda,
  "requires warp_optix and a CUDA device",
)
class RenderOptixTest(parameterized.TestCase):
  ...

[thowell]

Thanks for the updates @StafaH! Two remaining recommendations regarding making OptiX opt-in (backend="warp" by default) and lazy-importing warp_optix:

1. Default backend to "warp" so OptiX is opt-in and off by default (render_util.py, cli.py, README.md, ray_test.py, render_test.py, render_util_test.py)

• Issue: create_render_context defaults to backend: str | None = None, which automatically enables "optix" whenever warp-optix is installed. Because the OptiX backend cannot be passed to mjw.ray/mjw.rays, cannot be recorded in wp.ScopedCapture(), clips geometry and ray segments outside [-extent, extent]^3, and does not apply RAY_TOL_REL tie-breaking on near-coincident geoms, enabling it by default breaks existing create_render_context(mjm) callers when warp-optix is installed (requiring 13 backend="warp" workarounds in ray_test.py, render_util_test.py, and render_test.py, while test_coincident_geoms_with_roundoff_resolve_to_higher_index in render_test.py:925 omits backend="warp" and fails under OptiX).
• Recommendation: Set the default to backend: str = "warp" in create_render_context and RENDER_BACKEND in cli.py so OptiX is opt-in (backend="optix"), update README.md, and drop the backend="warp" workarounds in ray_test.py, render_util_test.py, and render_test.py.

# mujoco_warp/_src/render_util.py
def create_render_context(
  ...
  deterministic: bool = False,
  backend: str = "warp",
) -> RenderContext:
  ...
  if backend not in ("warp", "optix"):
    raise ValueError(f"backend must be 'warp' or 'optix', got {backend!r}")
  if backend == "optix":
    reason = ""
    wo = render_optix.get_wo()
    if wo is None:
      reason = "warp-optix is not installed (pip install mujoco-warp[optix])"
    elif rc.bvh_nflexgeom or rc.splat_count:
      reason = "the OptiX backend does not support flexes or splats yet"
    elif nworld * mjm.ngeom >= 1 << 24:
      reason = "OptiX instance ids hold 24 bits, so nworld * ngeom must be below 2**24"
    else:
      device = wp.get_device()
      try:
        wo.get_context(device)
      except Exception as e:
        reason = f"OptiX is not supported on device {device.alias} ({e})"
    if reason:
      warnings.warn(f"backend='optix' is unavailable, using backend='warp': {reason}")
      backend = "warp"
  rc.backend = backend

2. Lazy-import warp_optix only when backend="optix" is used (render_optix.py, render.py, render_util.py)

• Issue: render.py, render_optix.py, and render_util.py execute import warp_optix as wo at module top level on every import mujoco_warp. When warp-optix is installed, importing it loads pyoptix (optix._optix), which unconditionally calls cudaFree(0) (initializing a CUDA runtime context on device 0 at import time before wp.init() or device selection) and optixInit(), and registers all OptiX builtins with Warp even when OptiX is not used.
• Recommendation: Lazy-import warp_optix inside a cached helper in render_optix.py and call it only when backend == "optix" (and in render.py only when optix is not None).

# mujoco_warp/_src/render_optix.py
import functools


@functools.cache
def get_wo():
  try:
    import warp_optix as wo  # noqa: PLC0415

    return wo
  except ImportError:
    return None

# mujoco_warp/_src/render.py
def _make_cast_ray(geom_ray_types: Tuple[int], first_hit: bool = False, optix: OptixScene | None = None) -> wp.Function:
  """Build closest-hit or shadow-ray traversal specialized to the scene and backend."""
  ray_primitive = make_ray_primitive(geom_ray_types)
  wo = render_optix.get_wo() if optix is not None else None
  ...

[StafaH]

Thanks @thowell, both done:

1. backend now defaults to "warp" in create_render_context and --render_backend, so OptiX is opt-in with backend="optix". I removed the backend="warp" workarounds from the tests and updated the README (0d75988).
2. warp_optix is now lazy-imported through a cached render_optix.get_wo(), called only on the OptiX path, so import mujoco_warp no longer loads it (a247618).

[thowell]

awesome work @StafaH! 🚀